Bài 8: Activation function: Softmax cho multi-class

Sau bài học, bạn sẽ:

• Viết được công thức softmax \( \text{softmax}(\mathbf{z})_i = e^{z_i} / \sum_{j} e^{z_j} \) và giải thích vì sao đầu ra là probability distribution.
• Phân biệt được softmax (multi-class, 1 trong K) với sigmoid (binary / multi-label).
• Hiểu và implement được stable softmax — trừ \( \max(\mathbf{z}) \) trước khi exp để tránh overflow.
• Dùng được temperature để điều khiển độ sharpness của distribution (sampling LLM, knowledge distillation, RL).
• Biết quy tắc quan trọng trong PyTorch: không apply softmax thủ công khi dùng nn.CrossEntropyLoss; truyền raw logits.
• Hiểu vì sao softmax thuộc về output layer chứ không phải hidden layer.

Bài này khép lại chuỗi 4 activation B5-B8 (Sigmoid, ReLU, Tanh, Softmax) và mở đường sang B9-B10 về loss function.

Trong một network phân loại \( K \) class, output layer cho ra vector \( \mathbf{z} \in \mathbb{R}^K \) — gọi là logits. Logits chưa phải xác suất: phần tử có thể âm, có thể lớn hơn 1, không cộng lại bằng 1.

Softmax là phép biến đổi chuẩn để chuyển \( \mathbf{z} \) thành probability distribution \( \hat{\mathbf{p}} \in [0, 1]^K \) với \( \sum_i \hat{p}_i = 1 \). Mỗi \( \hat{p}_i \) đọc như xác suất sample thuộc class \( i \).

Ví dụ điển hình: phân loại iris (3 class), MNIST (10 class), ImageNet (1000 class), language model dự đoán token tiếp theo (vocab size 32k-256k).

Với \( \mathbf{z} = (z_1, z_2, \dots, z_K) \in \mathbb{R}^K \):

\[ \text{softmax}(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}, \quad i = 1, 2, \dots, K \]

Diễn giải từng bước:

• Exponentiate: \( e^{z_i} \) — đưa mọi giá trị về dương và làm nổi bật phần tử lớn (vì \( e^x \) tăng nhanh).
• Normalize: chia cho tổng \( \sum_j e^{z_j} \) — đưa tổng các phần tử về 1.

Ví dụ với \( \mathbf{z} = (1, 2, 3) \):

• \( e^1 \approx 2.718, \; e^2 \approx 7.389, \; e^3 \approx 20.086 \).
• Tổng \( \approx 30.193 \).
• \( \text{softmax}(\mathbf{z}) \approx (0.0900, 0.2447, 0.6652) \).
• Tổng \( 0.0900 + 0.2447 + 0.6652 = 0.9999 \approx 1 \).

Phần tử lớn nhất nhận xác suất cao nhất, nhưng các phần tử khác vẫn nhận được phần dương — softmax là phiên bản "soft" của argmax.

• Range: \( \text{softmax}(\mathbf{z})_i \in (0, 1) \) với mọi \( i \). Cận không bao giờ đạt được vì \( e^{z_i} > 0 \).
• Probability distribution: \( \sum_{i=1}^{K} \text{softmax}(\mathbf{z})_i = 1 \).
• Monotonic theo \( z_i \): tăng \( z_i \) thì \( \hat{p}_i \) tăng, các \( \hat{p}_j \) khác (\( j \ne i \)) giảm. Cross-neuron dependent.
• Translation invariance: cộng cùng một hằng số \( c \) vào mọi \( z_i \) không đổi kết quả: \[ \text{softmax}(\mathbf{z} + c)_i = \frac{e^{z_i + c}}{\sum_j e^{z_j + c}} = \frac{e^c \cdot e^{z_i}}{e^c \sum_j e^{z_j}} = \text{softmax}(\mathbf{z})_i \] Đây là nền tảng cho stable softmax (mục 7).
• Không scale invariant: nhân \( \mathbf{z} \) với hằng số làm thay đổi distribution — nền tảng cho temperature (mục 8).

Hai activation đều cho ra giá trị \( \in (0, 1) \) đọc như xác suất, nhưng khác về cách tính giữa các neuron:

• Sigmoid (B5): per-neuron, độc lập. Mỗi output tính riêng từ pre-activation của neuron đó: \( \sigma(z_i) = 1 / (1 + e^{-z_i}) \). Tổng \( K \) sigmoid không bằng 1.
• Softmax: cross-neuron, phụ thuộc lẫn nhau. Output thứ \( i \) phụ thuộc tất cả \( z_j \) vì mẫu số là tổng \( \sum_j e^{z_j} \). Tổng \( K \) output bằng 1.

Trường hợp đặc biệt \( K = 2 \): softmax suy biến về sigmoid. Cho \( \mathbf{z} = (z_1, z_2) \):

\[ \text{softmax}(\mathbf{z})_1 = \frac{e^{z_1}}{e^{z_1} + e^{z_2}} = \frac{1}{1 + e^{-(z_1 - z_2)}} = \sigma(z_1 - z_2) \]

Nói cách khác, softmax 2 class chỉ có 1 bậc tự do — chính bằng sigmoid của hiệu hai logit. Đó là lý do binary classification dùng một neuron sigmoid thay vì 2 neuron softmax.

Hai bài toán nghe giống nhau nhưng activation và loss khác nhau:

• Multi-class (1 trong K): mỗi sample thuộc đúng một class. Label là one-hot vector. Ví dụ: iris (setosa / versicolor / virginica), MNIST (chữ số 0-9), ImageNet (1 class trên 1000).
  • Activation output layer: softmax (\( K \) neuron).
  • Loss: cross-entropy (B10).
• Multi-label: mỗi sample có thể có nhiều label cùng lúc. Label là binary vector nhiều bit bật. Ví dụ: ảnh chứa cả "dog" và "outdoor", bài viết được tag "ai", "machine-learning", "python" cùng lúc.
  • Activation output layer: sigmoid trên từng neuron độc lập (\( K \) neuron).
  • Loss: binary cross-entropy trên từng output (B10).

Nhầm hai bài toán này dẫn đến hai lỗi thường gặp: dùng softmax cho multi-label (mạng bị ép tổng xác suất = 1, không thể bật nhiều label cùng lúc), hoặc dùng sigmoid + cross-entropy cho multi-class (loss tính sai vì các class không bị ép cạnh tranh).

Implementation ngây thơ np.exp(z) / np.sum(np.exp(z)) bị tràn số khi \( z_i \) lớn:

• \( z_i = 1000 \): \( e^{1000} \) overflow trong float32 (max \( \approx 3.4 \times 10^{38} \), tương đương \( e^{88.7} \)) và float64 (max \( \approx 1.8 \times 10^{308} \), tương đương \( e^{709.8} \)) → kết quả inf, chia inf / inf ra NaN.
• \( z_i \) rất âm: \( e^{z_i} \) underflow về 0, không phải vấn đề trừ khi mọi \( z_i \) đều underflow.

Trick: dùng tính chất translation invariance ở mục 4. Trừ \( \max(\mathbf{z}) \) ra trước khi exp:

\[ \text{softmax}(\mathbf{z})_i = \frac{e^{z_i - \max(\mathbf{z})}}{\sum_{j=1}^{K} e^{z_j - \max(\mathbf{z})}} \]

Sau khi trừ max:

• Phần tử lớn nhất bằng 0, \( e^0 = 1 \) — luôn an toàn.
• Các phần tử khác \( \le 0 \), \( e^{\le 0} \in (0, 1] \) — không overflow.
• Mẫu số \( \ge 1 \), không chia cho 0.

Toán học tương đương chính xác với công thức gốc, chỉ khác cách viết. Mọi implementation production (PyTorch, TensorFlow, JAX, NumPy scipy.special.softmax) đều làm như vậy. Khi tự viết, đừng dùng phiên bản ngây thơ.

Chia logits cho một hằng số dương \( T \) (gọi là temperature) trước khi softmax:

\[ \text{softmax}(\mathbf{z} / T)_i = \frac{e^{z_i / T}}{\sum_{j=1}^{K} e^{z_j / T}} \]

Tham số \( T \) điều khiển độ sharpness của distribution:

• \( T = 1 \): softmax tiêu chuẩn.
• \( T > 1 \): distribution mềm hơn (smoother, less confident). Hiệu giữa các \( z_i / T \) nhỏ đi → xác suất các class gần nhau hơn.
• \( T < 1 \): distribution sắc hơn (sharper, more confident). Hiệu giữa các \( z_i / T \) lớn hơn → class top trội rõ.
• \( T \to 0^+ \): tiệm cận argmax (one-hot). Class lớn nhất nhận xác suất 1, các class khác nhận 0.
• \( T \to \infty \): tiệm cận uniform. Mọi class nhận xác suất \( 1/K \).

Minh hoạ với \( \mathbf{z} = (1, 2, 3) \):

• \( T = 0.5 \): \( \text{softmax}(\mathbf{z}/0.5) = \text{softmax}((2, 4, 6)) \approx (0.0159, 0.1173, 0.8668) \) — sắc, class 3 chiếm 87%.
• \( T = 1 \): \( \approx (0.0900, 0.2447, 0.6652) \) — chuẩn, class 3 chiếm 67%.
• \( T = 2 \): \( \text{softmax}((0.5, 1, 1.5)) \approx (0.1863, 0.3072, 0.5065) \) — mềm, class 3 chỉ chiếm 51%.
• \( T = 100 \): \( \approx (0.3300, 0.3333, 0.3367) \) — gần uniform.

• LLM sampling (S4 — LLM & GenAI): mỗi step language model cho ra distribution trên vocab. temperature=0 tương đương greedy / argmax (luôn chọn token có xác suất cao nhất). temperature=1 sample theo distribution gốc. temperature=1.5 tăng diversity nhưng dễ ra văn bản kỳ quặc. Các API OpenAI, Anthropic, Hugging Face đều expose tham số này.
• Knowledge distillation (Hinton, Vinyals, Dean 2015): teacher model tính softmax với \( T \) cao (vd. \( T = 4 \)) → distribution mềm chứa nhiều thông tin về quan hệ giữa các class ("dog gần với cat hơn là với truck"). Student model học theo distribution mềm này, hiệu quả hơn học theo one-hot label.
• Reinforcement Learning — exploration vs exploitation: policy network dùng softmax trên action values; \( T \) cao khuyến khích explore (thử action mới), \( T \) thấp tập trung exploit (chọn action tốt nhất).
• Calibration: model neural thường overconfident; điều chỉnh \( T > 1 \) sau training (Guo và cộng sự 2017 — "On calibration of modern neural networks") làm probability đọc đúng hơn theo nghĩa thống kê.

Khi inference chỉ cần nhãn dự đoán (không cần probability), chú ý:

\[ \arg\max_i \text{softmax}(\mathbf{z})_i = \arg\max_i z_i \]

Vì softmax monotonic theo \( z_i \) — class có logit lớn nhất cũng là class có probability lớn nhất. Hệ quả: không cần compute softmax nếu chỉ cần dự đoán label:

logits = model(x)              # shape (batch, K)
pred = logits.argmax(dim=-1)   # không cần softmax

Tiết kiệm chút compute. Chỉ cần softmax khi: in probability cho user đọc, tính loss (mà như mục 13 sẽ thấy, PyTorch đã gộp luôn trong CrossEntropyLoss), hoặc dùng để sampling (LLM).

Softmax \( K \to K \) nên đạo hàm là một ma trận Jacobian \( K \times K \). Với \( \hat{p}_i = \text{softmax}(\mathbf{z})_i \):

\[ \frac{\partial \hat{p}_i}{\partial z_j} = \begin{cases} \hat{p}_i (1 - \hat{p}_i) & \text{nếu } i = j \\ -\hat{p}_i \, \hat{p}_j & \text{nếu } i \ne j \end{cases} \]

Viết gọn: \( \partial \hat{p}_i / \partial z_j = \hat{p}_i (\delta_{ij} - \hat{p}_j) \), với \( \delta_{ij} \) là Kronecker delta.

Jacobian này phức tạp, nhưng cặp softmax + cross-entropy loss rút gọn thành một biểu thức rất gọn. Với label \( \mathbf{y} \) one-hot và cross-entropy \( L = -\sum_i y_i \log \hat{p}_i \):

\[ \frac{\partial L}{\partial z_i} = \hat{p}_i - y_i \]

Gradient trên logits chính là sai khác giữa xác suất dự đoán và label. Đơn giản, ổn định, dễ implement — chi tiết đầy đủ ở B11 (Backpropagation). Đó cũng là lý do PyTorch gộp softmax vào trong loss function (mục 13).

Softmax có hai lý do thực dụng khiến nó không phù hợp cho hidden layer:

• Ép tổng bằng 1: ép phụ thuộc giữa các neuron trong cùng layer (tăng cái này phải giảm cái khác). Hidden representation lý tưởng cần linh hoạt — nhiều feature có thể cùng "bật" mạnh, không cạnh tranh zero-sum.
• Saturation: nếu một \( z_i \) lớn, \( \hat{p}_i \to 1 \) và các \( \hat{p}_j \to 0 \). Đạo hàm về \( z_j \) tỉ lệ với \( \hat{p}_j (1 - \hat{p}_j) \to 0 \). Gradient tắt giống saturation của sigmoid — không tốt cho mạng sâu.

Mặc định cho hidden layer của feed-forward và CNN là ReLU hoặc biến thể (LeakyReLU, GELU). Softmax chỉ xuất hiện ở:

• Output layer của multi-class classifier.
• Attention trong Transformer (B-Series 3 — Transformer): softmax tính trên một chiều cụ thể (sequence length) để chuyển attention scores thành weight tổng bằng 1, không phải softmax trên hidden dimension.
• Mixture-of-Experts gating: gate network output qua softmax để chọn expert.

Quy tắc quan trọng nhất trong bài: trong PyTorch, nn.CrossEntropyLoss đã built-in softmax. Input của nó là logits thô, không phải probability.

import torch
import torch.nn as nn

logits = model(x)              # shape (batch, K) — KHÔNG apply softmax
target = torch.tensor([2, 0, 1])  # class index, không phải one-hot
loss = nn.CrossEntropyLoss()(logits, target)

Lỗi cực kỳ phổ biến của người mới:

# SAI — apply softmax 2 lần
probs = torch.softmax(model(x), dim=-1)
loss = nn.CrossEntropyLoss()(probs, target)  # CrossEntropyLoss tự softmax lần nữa

Hậu quả: loss vẫn ra số nhưng gradient sai, training hội tụ chậm hoặc không học được. Không có warning. Cách phát hiện: loss đầu epoch quá thấp (vd. \( \approx 1.0 \) với 10 class trong khi expected \( \approx \ln 10 \approx 2.3 \) lúc init random).

Lý do PyTorch gộp: kết hợp \( \log \) và softmax thành log-sum-exp trick — tính \( \log \sum_j e^{z_j} \) ổn định bằng cách trừ max ngay trong log, tránh tính probability rồi mới log (khi probability gần 0, \( \log \) ra \( -\infty \)). Công thức tích hợp:

\[ -\log \text{softmax}(\mathbf{z})_{y} = -z_y + \log \sum_{j=1}^{K} e^{z_j} \]

với \( y \) là class index đúng. Vế phải chỉ cần một logsumexp stable, không bao giờ chạm inf hay NaN.

Quy tắc nhớ: để raw logits cho loss function. Chỉ apply softmax khi inference cần probability đọc.

Phép \( \log(\text{softmax}(\mathbf{z})) \) tính trực tiếp sẽ không stable: softmax có thể trả về \( 0 \) (underflow) và \( \log(0) = -\infty \). Cách stable là gộp \( \log \) và softmax thành một bước:

\[ \log \text{softmax}(\mathbf{z})_i = z_i - \log \sum_{j=1}^{K} e^{z_j} = (z_i - m) - \log \sum_{j=1}^{K} e^{z_j - m} \]

với \( m = \max(\mathbf{z}) \) — chính là log-sum-exp trick.

PyTorch cung cấp:

• torch.nn.functional.log_softmax(z, dim=-1) — phiên bản stable.
• nn.LogSoftmax(dim=-1) — module, dùng trong nn.Sequential.
• nn.NLLLoss — Negative Log Likelihood, expect input đã qua log_softmax.

Hai cách viết tương đương:

# Cách 1: gộp luôn
loss = nn.CrossEntropyLoss()(logits, target)

# Cách 2: tách
log_probs = nn.functional.log_softmax(logits, dim=-1)
loss = nn.NLLLoss()(log_probs, target)

Cách 1 phổ biến hơn vì gọn. Cách 2 hữu ích khi cần log-probability cho mục đích khác (vd. distillation cần \( \log p_{\text{student}} \) để tính KL divergence với \( p_{\text{teacher}} \)).

PyTorch 2.x cung cấp các API sau (chú ý tham số dim):

import torch
import torch.nn as nn
import torch.nn.functional as F

# Batch 3 sample, 4 class
logits = torch.tensor([
    [1.0, 2.0, 3.0, 4.0],
    [0.5, 1.5, 0.0, 2.0],
    [3.0, 3.0, 3.0, 3.0],
])

# Functional — dim=-1 nghĩa là softmax theo chiều cuối (class dimension)
probs = torch.softmax(logits, dim=-1)
print(probs.sum(dim=-1))   # tensor([1.0000, 1.0000, 1.0000])

# Module — dùng trong nn.Sequential nếu thật sự cần
act = nn.Softmax(dim=-1)
probs2 = act(logits)

# Log-softmax stable
log_probs = F.log_softmax(logits, dim=-1)

Chú ý tham số dim: với tensor shape (batch, num_classes), gần như luôn là dim=-1 (chiều cuối). Nhầm dim=0 sẽ softmax theo batch — sai nghĩa và sai gradient.

Workflow chuẩn:

• Training: layer cuối là nn.Linear(hidden, K) — output raw logits, không Softmax. Loss là nn.CrossEntropyLoss() nhận thẳng logits.
• Inference cần probability: probs = torch.softmax(logits, dim=-1).
• Inference chỉ cần nhãn: pred = logits.argmax(dim=-1) — bỏ qua softmax.

Implement softmax stable với NumPy và test với logits nhỏ / lớn:

import numpy as np


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


def softmax_stable(z, temperature=1.0):
    z = np.asarray(z, dtype=np.float64) / temperature
    z = z - z.max()                # trừ max để stable
    e = np.exp(z)
    return e / e.sum()


# Test 1: logits nhỏ
print(softmax_stable([1, 2, 3]))
# [0.09003057 0.24472847 0.66524096]

# Test 2: logits lớn — naive sẽ overflow
print(softmax_stable([1000, 1001, 1002]))
# [0.09003057 0.24472847 0.66524096]  — kết quả giống hệt!

# Naive trên cùng input
print(softmax_naive([1000, 1001, 1002]))
# RuntimeWarning: overflow in exp
# [nan nan nan]

# Verify sum = 1
print(softmax_stable([1, 2, 3]).sum())   # 1.0

Quan sát: softmax_stable([1, 2, 3]) và softmax_stable([1001, 1002, 1003]) ra kết quả giống hệt — đó là translation invariance ở mục 4.

Demo temperature scaling:

z = [1, 2, 3]
for T in [0.5, 1.0, 2.0, 100.0]:
    p = softmax_stable(z, temperature=T)
    print(f"T = {T:5.1f}  p = {np.round(p, 4)}")
# T =   0.5  p = [0.0159 0.1173 0.8668]
# T =   1.0  p = [0.09   0.2447 0.6652]
# T =   2.0  p = [0.1863 0.3072 0.5065]
# T = 100.0  p = [0.33   0.3333 0.3367]

Cùng workflow trong PyTorch:

import torch

z = torch.tensor([1.0, 2.0, 3.0])
print(torch.softmax(z, dim=-1))
# tensor([0.0900, 0.2447, 0.6652])
print(torch.softmax(z, dim=-1).sum())    # tensor(1.)

# Stable với logits lớn
big = torch.tensor([1000.0, 1001.0, 1002.0])
print(torch.softmax(big, dim=-1))
# tensor([0.0900, 0.2447, 0.6652])  — không NaN

1. Tính tay \( \text{softmax}((1, 2, 3)) \), làm tròn 4 chữ số. Đối chiếu với output ở mục 3.
2. Chứng minh tính đối xứng / translation invariance: với mọi \( c \), \( \text{softmax}(\mathbf{z} + c \cdot \mathbf{1}) = \text{softmax}(\mathbf{z}) \). Áp dụng: tính nhẩm \( \text{softmax}((1001, 1002, 1003)) \) khi đã biết \( \text{softmax}((1, 2, 3)) \).
3. Chứng minh \( K = 2 \): \( \text{softmax}((z_1, z_2))_1 = \sigma(z_1 - z_2) \).
4. Implement softmax_naive và softmax_stable. Test với z = [1000, 1001, 1002], in ra warning của naive và verify stable không warning.
5. Plot 3 đường softmax với \( T \in \{0.5, 1, 2\} \) trên cùng input \( \mathbf{z} = (1, 2, 3) \) bằng matplotlib (bar chart). Đối chiếu với bảng ở mục 8.
6. Chứng minh \( \text{softmax}(\mathbf{z}/T) \to \) one-hot tại argmax khi \( T \to 0^+ \) (với điều kiện argmax duy nhất). Gợi ý: viết \( e^{z_i / T} = e^{(z_i - z_{\max}) / T} \cdot e^{z_{\max} / T} \) rồi chia cho mẫu.
7. Trong PyTorch, viết một classifier MNIST 10 class với layer cuối là nn.Linear(128, 10). Hỏi: nên apply nn.Softmax sau layer cuối hay không? Vì sao? (Đáp: không, để raw logits cho nn.CrossEntropyLoss.)
8. Cho logits z = torch.tensor([2.0, 1.0, 0.1]). So sánh kết quả 3 cách: (a) torch.log(torch.softmax(z, -1)), (b) F.log_softmax(z, -1), (c) tự implement bằng log-sum-exp. Cả 3 ra số giống nhau ở case này; thử lại với z = torch.tensor([100.0, 200.0, 300.0]) để thấy (a) bị -inf còn (b)(c) vẫn ổn.

• Softmax \( \text{softmax}(\mathbf{z})_i = e^{z_i} / \sum_j e^{z_j} \) biến vector logits thành probability distribution trên \( K \) class: mỗi phần tử \( \in (0, 1) \), tổng bằng 1.
• Dùng ở output layer cho multi-class classification (1 trong K). Multi-label thì dùng sigmoid trên từng output.
• Trường hợp \( K = 2 \): softmax suy biến về sigmoid của hiệu \( z_1 - z_2 \).
• Stable softmax: trừ \( \max(\mathbf{z}) \) trước khi exp — toán học tương đương, tránh overflow. Mọi implementation production đều làm vậy.
• Temperature: chia logits cho \( T \) trước khi softmax. \( T > 1 \) mềm hơn, \( T < 1 \) sắc hơn, \( T \to 0 \) argmax, \( T \to \infty \) uniform. Dùng cho LLM sampling, knowledge distillation, RL exploration, calibration.
• Khi chỉ cần dự đoán nhãn (không probability), bỏ softmax: \( \arg\max_i \text{softmax}(\mathbf{z})_i = \arg\max_i z_i \).
• Đạo hàm softmax: Jacobian \( K \times K \), \( \partial \hat{p}_i / \partial z_j = \hat{p}_i (\delta_{ij} - \hat{p}_j) \). Cặp softmax + cross-entropy rút gọn thành \( \partial L / \partial z_i = \hat{p}_i - y_i \).
• Softmax không dùng cho hidden layer: ép tổng = 1 không tự nhiên, và gradient tắt khi một class lấn át.
• Quy tắc PyTorch quan trọng: nn.CrossEntropyLoss đã built-in softmax → input là raw logits, KHÔNG apply softmax thủ công rồi truyền cho loss. Tương tự, nn.NLLLoss đi cặp với log_softmax.
• Inference cần probability: torch.softmax(logits, dim=-1). Inference chỉ cần nhãn: logits.argmax(dim=-1).