# 📊 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:

| dtype | $e^{z}$ tràn khi $z$ vượt | vì giá trị lớn nhất xấp xỉ |
| --- | --- | --- |
| float16 | ~11.09 | 65504 |
| float32 | ~88.72 | 3.4e38 |
| float64 | ~709.78 | 1.8e308 |

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`:

| $z_i$ | `exp(z)` float32 | $z_i - 1002$ | `exp(z - 1002)` |
| --- | --- | --- | --- |
| 1000 | inf | -2 | 0.1353 |
| 1001 | inf | -1 | 0.3679 |
| 1002 | inf | 0 | 1.0000 |

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ệ.

```mermaid
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:

```python
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.

```python
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**

```python
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`.

---

## Contents

Part of [🧠 Khóa 5: Toán cho Machine Learning — tích chập, softmax và lan truyền ngược cho deep learning](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho.md)

- **🖼️ Tích chập, shape và lan truyền ngược**
  - [🔲 Tích chập 1D và 2D: kernel, stride, padding](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/tich-chap-1d-2d-kernel-stride-padding.md)
  - [📐 Kiểm tra shape và vector hóa tích chập bằng im2col](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/kiem-tra-shape-va-vector-hoa-tich-chap-im2col.md) (sign-in required)
  - [🔁 Lan truyền ngược qua lớp tuyến tính và lớp tích chập](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/lan-truyen-nguoc-qua-lop-tuyen-tinh-va-lop-tich-chap.md)
  - [✅ Checkpoint: tích chập và backprop](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/checkpoint-tich-chap-va-backprop.md)
- **📉 Softmax, cross-entropy và numerical stability**
  - [📊 Softmax và log-sum-exp: ổn định số học](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/softmax-va-log-sum-exp-on-dinh-so-hoc.md) — this lesson
  - [🎯 Cross-entropy và đạo hàm của loss theo từng tham số](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/cross-entropy-va-dao-ham-loss-theo-tham-so.md)
  - [🧪 Vanishing/exploding gradient, numerical stability và gradient checking](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/vanishing-exploding-gradient-numerical-stability-gradient-checking.md)
  - [✅ Checkpoint: loss, gradient và numerical stability](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/checkpoint-loss-gradient-va-numerical-stability.md)
- **🚀 Capstone: forward/backward cho mạng 2 lớp**
  - [🏗️ Project: forward pass và kiểm tra shape cho mạng 2 lớp](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/project-forward-pass-va-kiem-tra-shape-mang-2-lop.md)
  - [⚙️ Project: backward pass, gradient checking và training loop](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/project-backward-pass-gradient-checking-va-training-loop.md) (premium content)
  - [✅ Checkpoint: capstone forward/backward](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/checkpoint-capstone-forward-backward.md) (premium content)
- [Course cheatsheet](https://everwise.study/courses/khoa-5-toan-cho-machine-learning-tich-chap-softmax-va-lan-truyen-nguoc-cho/cheatsheet.md) (premium content)