Softmax và log-sum-exp: ổn định số học

By Better Wise

Module 1 khép lại với forward/backward của lớp tích chập: bạn đã có một hàm nhận ảnh và trả về logits $z \in \mathbb{R}^{C}$. Bước cuối của mạng phân loại là biến logits thành phân bố xác suất, và đó chính là chỗ những con số trông vô hại bỗng biến thành inf rồi nan. Bài này chỉ giải quyết một việc: viết softmax sao cho nó không bao giờ nổ, kể cả khi logits lệch nhau vài bậc độ lớn.

Sau bài học này bạn có thể:

• Viết softmax "naive" và chỉ ra chính xác dòng code gây tràn số.
• Chứng minh bất biến dịch chuyển (shift invariance) và cài đặt trick trừ max.
• Cài logsumexp và log_softmax, giải thích vì sao chúng tránh cả overflow lẫn underflow.
• Áp dụng tất cả vào logits của mạng 2 lớp trong capstone.

Định nghĩa softmax và chỗ nó chết vì exp

Với vector logits $z$, softmax cho phân bố xác suất:

$$ \text{softmax}(z)i = \frac{e^{z_i}}{\sum{j} e^{z_j}} $$

Về mặt toán học thì không có gì sai. Về mặt số học thì có hai cái chết nằm ngay trong công thức: exp tràn khi $z_i$ lớn, và mẫu số bằng 0 khi mọi $z_i$ nhỏ.

Giới hạn tràn của exp rất thấp:

Logits trong một mạng thật hoàn toàn có thể vượt các ngưỡng này, đặc biệt sau vài epoch khi trọng số phình lên. Xem ví dụ với float32:

Dòng đầu tiên cho ra inf / inf = nan. Không có exception nào được ném, mô hình chỉ đơn giản là học được toàn số nan sau vài bước.

Chiều ngược lại cũng hỏng: với $z = [-1000, -1001, -1002]$, cả ba exp đều underflow về 0, mẫu số bằng 0, kết quả là 0 / 0 = nan. Softmax cần chính xác cả hai đầu.

Warning: nan từ softmax không lan ra từ từ, nó lan ngay lập tức: gradient trở thành nan, toàn bộ trọng số phía trước thành nan. Khi debugging, hãy in logits trước khi nghi ngờ bất cứ thứ gì khác.

Trick trừ max: bất biến dịch chuyển

Chìa khóa là một tính chất đại số đơn giản. Với mọi hằng số $c$:

$$ \text{softmax}(z - c)_i = \frac{e^{z_i - c}}{\sum_j e^{z_j - c}} = \frac{e^{-c} e^{z_i}}{e^{-c} \sum_j e^{z_j}} = \text{softmax}(z)_i $$

Tỉ số không đổi vì tử và mẫu cùng nhân $e^{-c}$. Vậy ta được tự do dịch logits mà không đổi phân bố, và nên chọn $c = m = \max_j z_j$. Khi đó số mũ lớn nhất bằng đúng $0$, mọi exp nằm trong $(0, 1]$, không thể tràn. Underflow thì vẫn có thể xảy ra ở các logits rất âm, nhưng lúc đó mẫu số luôn có ít nhất một hạng tử bằng 1, nên phép chia vẫn hợp lệ.

flowchart TD
    Z["🔢<br/>logits z"] --> M["📉<br/>m = max(z)"]
    M --> S["➖<br/>z - m"]
    S --> E["🚀<br/>exp(z - m)"]
    E --> Sum["Σ<br/>tổng mẫu số"]
    Sum --> P["🎯<br/>p = exp / tổng"]

Cài đặt, kèm bản naive để so sánh:

import numpy as np

def softmax_naive(z):
    e = np.exp(z)
    return e / e.sum()

def softmax_stable(z):
    m = np.max(z)
    e = np.exp(z - m)
    return e / e.sum()

z = np.array([1000.0, 1001.0, 1002.0], dtype=np.float32)
softmax_naive(z)    # [nan, nan, nan]  vì exp(1000) = inf
softmax_stable(z)   # [0.0900, 0.2447, 0.6652]

Bài tập 1. Chứng minh $\text{softmax}(z - c) = \text{softmax}(z)$ mà không cần giả thiết $c = \max z$. Từ đó trả lời: nếu chỉ cần "một hằng số đủ lớn", tại sao người ta luôn chọn max?

Xem đáp án

Chứng minh bằng cách đặt nhân tử chung $e^{-c}$ ở tử và mẫu như trên, rút gọn là xong. Bất kỳ $c$ nào cũng đúng về mặt toán học, kể cả $c = 0$.

Chọn $c = \max z$ vì nó là giá trị nhỏ nhất đảm bảo không có số mũ nào dương, tức là không tốn thêm gì mà chặn được cả hai đầu: max cho exp lớn nhất bằng 1 (không tràn), và mẫu số có ít nhất một hạng tử bằng 1 (không chia cho 0). Chọn $c$ lớn tùy ý thì các exp còn lại underflow về 0 nhanh hơn, mất precision vô ích.

Bài tập 2. Trong capstone, logits của một mẫu là [12.0, 8.0, 0.5] và bạn đang chạy float16. Bản naive trả về gì, bản ổn định trả về gì?

Xem đáp án

float16 tràn từ khoảng 11.09, mà exp(12) ≈ 162754, nên bản naive cho inf / inf = nan.

Bản ổn định trừ 12: [0, -4, -11.5], exp thành [1, 0.0183, 1.01e-5], tổng 1.0183, xác suất [0.9820, 0.0180, 1.0e-5]. Lưu ý rằng đáp án đúng về mặt lý thuyết cũng là như vậy, chỉ khác là ta tính được nó.

log-sum-exp và log_softmax: làm việc trong miền log

Softmax ổn định rồi, nhưng đầu ra của nó sắp được đưa vào cross-entropy, và ở đó ta cần log của xác suất. Tính log(softmax(z)) là một cái bẫy: nếu một xác suất underflow về 0 thì log(0) = -inf, và toàn bộ loss thành inf. Câu trả lời là gộp phép log vào trước khi chia:

$$ \text{logsumexp}(z) = m + \log \sum_j e^{z_j - m}, \qquad \log \text{softmax}(z)_i = z_i - \text{logsumexp}(z) $$

Hàm logsumexp chính là log của mẫu số, nhưng được tính trong miền log. Nó chỉ dùng một lần log trên một tổng luôn dương (vì có hạng tử $e^{0} = 1$), nên không bao giờ chạm log(0). Và công thức $\log p_i = z_i - \text{logsumexp}(z)$ không hề chia, không hề lấy log của xác suất, nên nó ổn định ở cả hai đầu.

def logsumexp(z):
    m = np.max(z)
    return m + np.log(np.sum(np.exp(z - m)))

def log_softmax(z):
    m = np.max(z)
    return (z - m) - np.log(np.sum(np.exp(z - m)))

z = np.array([1000.0, 1001.0, 1002.0], dtype=np.float32)
log_softmax(z)   # [-2.4076, -1.4076, -0.4076]

z_neg = np.array([-1000.0, -1001.0, -1002.0], dtype=np.float32)
softmax_naive(z_neg)   # [nan, nan, nan]  vì tổng mẫu số = 0
softmax_stable(z_neg)  # [0.6652, 0.2447, 0.0900]

Một tính chất đáng nhớ vì bài sau dùng lại: $\nabla_z \text{logsumexp}(z) = \text{softmax}(z)$. Đạo hàm của một hàm ổn định về số học lại là một hàm ổn định, đó là lý do người ta thích viết loss qua logsumexp.

Tip: Với mini-batch trong capstone, max phải lấy theo từng hàng, không phải max toàn ma trận: m = z.max(axis=1, keepdims=True). keepdims=True là bắt buộc để broadcasting trừ đúng cột.

Bài tập 3. Với z = [-1000, -1001, -1002] ở float32, giải thích từng bước vì sao bản naive cho nan, và bản ổn định cho kết quả nào. Kết quả đó có khớp với softmax của [0, -1, -2] không?

Xem đáp án

Naive: exp(-1000), exp(-1001), exp(-1002) đều underflow về đúng 0.0 trong float32. Tổng mẫu số bằng 0, phép chia 0/0 cho nan.

Ổn định: $m = -1000$, dịch thành [0, -1, -2], exp thành [1, 0.3679, 0.1353], tổng 1.5032, xác suất [0.6652, 0.2447, 0.0900]. Khớp chính xác với softmax của [0, -1, -2], đúng như bất biến dịch chuyển dự đoán.

Bài tập 4. Viết log_softmax cho cả batch (N, C) và dùng nó để tính negative log-likelihood của nhãn đúng mà không bao giờ gọi np.log lên xác suất.

Xem đáp án

def log_softmax_batch(z):            # z: (N, C)
    m = z.max(axis=1, keepdims=True)
    shifted = z - m
    return shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True))

def nll(logits, y):                  # y: (N,) nhãn đúng
    logp = log_softmax_batch(logits)
    return -logp[np.arange(len(y)), y].mean()

Điểm mấu chốt: logp được tính trực tiếp, không có bước nào lấy log của một xác suất đã underflow. Bài sau ta sẽ lấy đạo hàm của chính biểu thức này theo logits và thuđược một công thức gọn bất ngờ: gradient của NLL theo logits chỉ là softmax(z) - one_hot(y), không có exp nào xuất hiện thêm.

Ba cái bẫy số học khi viết softmax

Bẫy 1: quên keepdims=True. z.max(axis=1) trả về shape (N,). Trừ vào z shape (N, C) vẫn chạy được nhờ broadcasting, nhưng NumPy align từ phải sang trái, nên nó trừ theo cột chứ không theo hàng. Kết quả không báo lỗi, chỉ sai âm thầm. Luôn dùng keepdims=True rồi kiểm tra m.shape == (N, 1).

Bẫy 2: trộn float32 và float64. np.exp trên mảng float32 trả về float32, ngưỡng tràn chỉ khoảng 88.7. Một logit bằng 90 đã đủ cho inf nếu bạn quên trừ max. Kiểm tra z.dtype trước khi tin vào kết quả.

Bẫy 3: lấy log sau softmax. Xác suất nhỏ hơn khoảng 1e-38 (float32) underflow về 0, log(0) cho -inf, và loss thành inf dù mô hình chỉ "khá tự tin sai". Đó là lý do trong capstone ta luôn có một hàm log_softmax riêng và không bao giờ gọi np.log lên đầu ra của softmax.

Quy tắc ngón tay: nếu trong code có exp mà không có phép trừ max đứng trước, hãy coi đó là bug cho tới khi chứng minh được ngược lại.

Tự kiểm tra trước khi sang cross-entropy

Với mọi hằng số c, softmax(z + c) bằng _____, nên dịch logits không đổi phân phối. Sau khi trừ max, phần tử lớn nhất của logits dịch là _____, nên exp của nó đúng bằng 1 và mẫu số luôn nằm trong khoảng _____. Bản log-sum-exp ổn định tính bằng m cộng với log của tổng exp(z trừ _____). log_softmax(z) bằng z trừ đi _____, nhờ đó không bao giờ phải lấy log của một xác suất đã underflow. Trong float32, exp của _____ tràn thành inf, còn exp của -1000 underflow về 0.

Front: Softmax nhận logits và trả về gì? Back: Một phân phối xác suất: các phần tử không âm và tổng bằng 1.

Front: Công thức softmax? Back: softmax(z)_i = exp(z_i) / tổng_j exp(z_j).

Front: Vì sao softmax naive cho nan với logits lớn? Back: exp(logit) tràn thành inf, rồi inf / inf = nan.

Front: Trick trừ max làm gì với logits? Back: Tính z' = z - max(z) rồi softmax trên z'; kết quả không đổi.

Front: Tính chất nào khiến trick trừ max hợp lệ? Back: Bất biến dịch chuyển: softmax(z + c) = softmax(z) với mọi hằng số c.

Front: Sau khi trừ max, exp lớn nhất bằng bao nhiêu? Back: Bằng 1, vì logit lớn nhất đã thành 0 và exp(0) = 1.

Front: Sau khi trừ max, mẫu số của softmax nằm trong khoảng nào? Back: Từ 1 đến C, với C là số lớp, nên không thể tràn.

Front: Log-sum-exp được viết ổn định thế nào? Back: logsumexp(z) = m + log(tổng exp(z_i - m)) với m = max(z).

Front: log_softmax(z) rút gọn bằng gì? Back: z - logsumexp(z), tính trực tiếp không qua bước softmax.

Front: Vì sao không nên lấy log lên đầu ra của softmax? Back: Xác suất rất nhỏ underflow về 0, log(0) = -inf làm loss nổ.

Front: Vì sao lấy max theo từng hàng thay vì max toàn ma trận? Back: Mỗi hàng có phần tử lớn nhất bằng 1 và mẫu số trong [1, C], giữ sai số làm tròn nhỏ; max toàn cục vẫn an toàn nhưng kém chính xác hơn.

Front: Với z = [-1000, -1001, -1002] ở float32, bản naive cho gì? Back: nan, vì cả ba exp underflow về 0 và mẫu số bằng 0.

Front: Softmax ổn định của [-1000, -1001, -1002] bằng bao nhiêu? Back: Xấp xỉ [0.6652, 0.2447, 0.0900], trùng với softmax của [0, -1, -2].

Front: Ngưỡng tràn của exp trên float32 khoảng bao nhiêu? Back: Khoảng 88.7, nên logit từ 90 trở lên đã đủ gây inf nếu không trừ max.

Front: Phép kiểm tra nhanh nhất cho một hàm softmax tự viết? Back: Cho logits cộng thêm 1000 và kiểm tra đầu ra không đổi và không có nan.

Chốt lại: một phép dịch, không phải một phép tính mới

Ba điều cần mang sang bài cross-entropy:

• Softmax chỉ là chuẩn hóa exp, và nó bất biến với phép dịch. Trừ max không thay đổi toán học, chỉ thay đổi nơi sai số xuất hiện.
• Sau khi trừ max, mọi thứ nằm trong vùng an toàn: exp lớn nhất bằng 1, mẫu số trong [1, C], không có inf và không có 0/0.
• Khi cần xác suất log, hãy tính log_softmax trực tiếp bằng log-sum-exp thay vì log(softmax(...)). Đây chính là biểu thức mà bài sau sẽ đạo hàm để ra softmax(z) - one_hot(y), từ đó đi tiếp tới đạo hàm của loss theo từng tham số W và b.