Vanishing/Exploding gradient và khởi tạo trọng số

Bạn xây một mạng 20 lớp, bấm huấn luyện, và… loss gần như không nhúc nhích. Hoặc ngược lại, sau vài trăm bước loss bỗng thành NaN. Hai triệu chứng này đến từ cùng một gốc rễ: gradient biến mất (vanishing) hoặc bùng nổ (exploding) khi lan truyền ngược qua nhiều lớp. Bài này giải thích nguyên nhân bằng toán và thí nghiệm, rồi đi qua các "liều thuốc": hàm kích hoạt phù hợp, khởi tạo trọng số đúng công thức, BatchNorm, skip connection và gradient clipping.

Số trang (ví dụ "trang 81") trong bài chỉ tới slide của bài giảng Deep Learning – Computer Vision – NLP (Thang Nguyen, Viet Nguyen); bài này giải thích chi tiết, kiểm chứng và đính chính các slide đó.

Vanishing gradient: vì sao các lớp đầu "chết"?

Nhắc lại quy tắc chuỗi (bài Backpropagation): gradient tại lớp đầu là tích của rất nhiều đạo hàm cục bộ, mỗi lớp góp một thừa số:

\[ \frac{\partial L}{\partial x} = \frac{\partial L}{\partial a}\cdot\frac{\partial a}{\partial b}\cdot\frac{\partial b}{\partial c}\cdots\frac{\partial g}{\partial x} \]

• Nếu các thừa số thường nhỏ hơn 1, tích co lại theo cấp số nhân: gradient ở các lớp đầu gần bằng 0, các lớp này gần như không học (hình "R.I.P." ở trang 118). Mạng sâu nhưng chỉ vài lớp cuối thực sự làm việc.
• Nếu các thừa số thường lớn hơn 1, tích phình to theo cấp số nhân: gradient bùng nổ.

Mỗi thừa số gồm hai phần: đạo hàm của hàm kích hoạt và độ lớn của trọng số. Vì vậy hai "thủ phạm" chính là activation và cách khởi tạo trọng số.

Thủ phạm 1: hàm kích hoạt sigmoid

\[ S(z) = \frac{1}{1+e^{-z}}, \qquad S'(z) = S(z)\,\big(1 - S(z)\big) \]

Đồ thị đạo hàm ở trang 119 cho thấy \(S'(z)\) đạt cực đại chỉ 0,25 tại \(z = 0\) và gần như bằng 0 khi \(|z| > 5\) (vùng bão hoà). Qua 10 lớp sigmoid, kể cả trong trường hợp tốt nhất, gradient bị nhân với \(0{,}25^{10} \approx 9{,}5\times10^{-7}\). Ngoài ra sigmoid còn hai nhược điểm: đầu ra luôn dương (không zero-centered) khiến cập nhật trọng số bị zigzag, và phép exp tốn kém.

Tanh tốt hơn (đạo hàm cực đại bằng 1, zero-centered) nhưng vẫn bão hoà ở hai đầu.

Liều thuốc 1: ReLU và các biến thể

ReLU

\[ R(z) = \max(0, z), \qquad R'(z) = \begin{cases} 1 & z > 0 \\ 0 & z < 0 \end{cases} \]

Với đầu vào dương, đạo hàm bằng đúng 1: gradient đi qua không bị co lại (trang 120). ReLU còn rẻ (chỉ một phép so sánh) và tạo kích hoạt thưa (nhiều neuron bằng 0). Nhược điểm: neuron có đầu vào luôn âm nhận gradient 0 mãi mãi, gọi là dying ReLU, hay xảy ra khi learning rate quá lớn. Tại \(z = 0\) đạo hàm không xác định; các framework quy ước bằng 0.

Các biến thể

Biểu đồ Swish ở trang 123 minh hoạ vai trò của \(\beta\): \(\beta\) nhỏ (0,1) gần như tuyến tính; \(\beta\) lớn (10) gần giống ReLU; \(\beta = 1\) nằm giữa.

Exploding gradient: mặt còn lại của đồng xu

Dấu hiệu (trang 124):

• Tham số của mô hình tăng rất nhanh (tăng theo cấp số nhân).
• Tham số hoặc loss trở thành NaN trong lúc huấn luyện.
• Loss vẫn lớn dù đã huấn luyện lâu, hoặc nhảy lên xuống thất thường.

Nguyên nhân:

• Khởi tạo trọng số không phù hợp (quá lớn), dẫn đến cập nhật trọng số quá lớn.
• Learning rate ban đầu quá cao.
• Hàm loss không phù hợp, cho giá trị loss rất lớn, kéo theo cập nhật trọng số lớn (ví dụ MSE với mục tiêu có giá trị hàng nghìn chưa chuẩn hoá).

Exploding gradient đặc biệt hay gặp ở mạng hồi quy (RNN) xử lý chuỗi dài, nơi cùng một ma trận trọng số bị nhân lặp lại hàng trăm lần.

Tổng hợp các giải pháp

Theo trang 125, kèm theo những gì ta đã học:

1. Activation phù hợp: ReLU và các biến thể.
2. Batch Normalization: giữ activation mỗi lớp ở thang đo ổn định (xem Batch Normalization, Dropout và Receptive Field).
3. Regularization: weight decay (L2) giữ trọng số nhỏ.
4. Khởi tạo trọng số đúng cách: Xavier, He.
5. Gradient clipping (chống exploding): theo giá trị hoặc theo chuẩn (norm).
6. Skip connection: đường tắt của ResNet cho gradient đi thẳng về các lớp đầu (xem Các kiến trúc CNN kinh điển).

Khởi tạo trọng số (weight initialization)

Hiệu năng của CNN/NN phụ thuộc đáng kể vào cách khởi tạo (trang 126). Khởi tạo tốt giúp: kết quả tái lập được (cố định seed), mô hình hội tụ nhanh hơn, và tránh vanishing/exploding gradient ngay từ bước đầu.

Khởi tạo bằng 0 hoặc hằng số: đừng bao giờ làm

Nếu mọi trọng số trong một lớp bằng nhau, mọi neuron của lớp đó nhận cùng đầu vào, tính ra cùng đầu ra, nhận cùng gradient, và được cập nhật giống hệt nhau. Chúng mãi mãi là bản sao của nhau, lớp 512 neuron hoạt động như 1 neuron. Đây là lỗi không phá vỡ được tính đối xứng (fail to break symmetry), luôn dẫn đến kết quả tồi (trang 127: "We don't do that here"). Lưu ý: bias khởi tạo bằng 0 thì không sao, vì trọng số ngẫu nhiên đã phá đối xứng.

Khởi tạo ngẫu nhiên: đúng hướng nhưng phải đúng "liều"

Lấy trọng số ngẫu nhiên từ phân phối chuẩn (random normal) hoặc phân phối đều (random uniform) giải quyết được vấn đề đối xứng, nhưng độ lớn sai có thể gây vanishing/exploding (trang 128). Thí nghiệm sau đưa 1.000 mẫu qua mạng 10 lớp, mỗi lớp 500 neuron, và in độ lệch chuẩn của activation sau từng lớp:

import torch

torch.manual_seed(0)
x0 = torch.randn(1000, 500)                  # 1000 mẫu, 500 đặc trưng


def run(act, init_std, layers=10, width=500):
    x = x0
    stds = []
    for _ in range(layers):
        W = torch.randn(width, width) * init_std(width)
        x = act(x @ W)
        stds.append(x.std().item())
    return " ".join(f"{s:.3g}" for s in stds)


tanh, relu = torch.tanh, torch.relu
print("tanh, N(0, 0.01) :", run(tanh, lambda n: 0.01))
print("tanh, N(0, 1)    :", run(tanh, lambda n: 1.0))
print("tanh, Xavier     :", run(tanh, lambda n: (2 / (n + n)) ** 0.5))
print("ReLU, Xavier     :", run(relu, lambda n: (2 / (n + n)) ** 0.5))
print("ReLU, He         :", run(relu, lambda n: (2 / n) ** 0.5))

Xavier / Glorot initialization (cho sigmoid, tanh)

Ý tưởng: giữ phương sai của tín hiệu không đổi khi đi qua mỗi lớp, cả chiều xuôi lẫn chiều ngược. Với neuron \(y = \sum_{i=1}^{n_{in}} w_i x_i\) (trọng số và đầu vào độc lập, trung bình 0):

\[ \mathrm{Var}(y) = n_{in}\,\mathrm{Var}(w)\,\mathrm{Var}(x) \]

Muốn \(\mathrm{Var}(y) = \mathrm{Var}(x)\) cần \(\mathrm{Var}(w) = 1/n_{in}\); xét thêm chiều backward cần \(1/n_{out}\); Glorot & Bengio (2010) lấy trung bình hai điều kiện:

\[ \mathrm{Var}(W) = \frac{2}{n_{in} + n_{out}} \]

• Xavier normal: \(W \sim \mathcal N\!\left(0,\ \sigma^2\right)\) với \(\sigma = \sqrt{\dfrac{2}{n_{in} + n_{out}}}\).
• Xavier uniform: \(W \sim U[-x, x]\) với \(x = \sqrt{\dfrac{6}{n_{in} + n_{out}}}\) (phân phối đều trên \([-x, x]\) có phương sai \(x^2/3\)).

He / Kaiming initialization (cho ReLU và biến thể)

ReLU đặt một nửa số đầu ra về 0, làm phương sai giảm một nửa sau mỗi lớp (đúng như dòng "ReLU, Xavier" ở trên). He et al. (2015) bù lại bằng cách nhân đôi:

\[ \mathrm{Var}(W) = \frac{2}{n_{in}} \]

• He normal: \(W \sim \mathcal N(0, \sigma^2)\) với \(\sigma = \sqrt{2/n_{in}}\).
• He uniform: \(W \sim U\!\left(-\sqrt{6/n_{in}},\ +\sqrt{6/n_{in}}\right)\).

Với lớp conv, \(n_{in}\) (fan-in) \(= C_{in}\times k_h\times k_w\): mỗi neuron đầu ra nhận ngần ấy đầu vào.

Khởi tạo trong PyTorch

PyTorch đã khởi tạo sẵn nn.Linear và nn.Conv2d bằng kaiming_uniform_ với tham số \(a = \sqrt5\), tương đương \(U(-1/\sqrt{n_{in}},\ 1/\sqrt{n_{in}})\) (lớp Linear(500, 100) có trọng số trong khoảng ±0,0447). Mặc định này dùng được cho đa số trường hợp, nhưng các kiến trúc như ResNet trong torchvision tự khởi tạo lại bằng He normal. Cách tự đặt:

import torch.nn as nn


def init_weights(m):
    if isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.BatchNorm2d):
        nn.init.ones_(m.weight)      # gamma = 1
        nn.init.zeros_(m.bias)       # beta = 0
    elif isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        nn.init.zeros_(m.bias)


model.apply(init_weights)            # áp dụng đệ quy cho mọi lớp con

(mode="fan_out" giữ ổn định phương sai ở chiều backward; cả hai mode đều phổ biến.)

Gradient clipping

Ngay cả khi mọi thứ đúng, đôi khi một batch "lạ" sinh ra gradient cực lớn và đẩy trọng số văng xa khỏi vùng tốt (trang 131: một bước nhảy vọt ra khỏi vách đá của mặt loss). Gradient clipping chặn độ lớn gradient trước khi cập nhật:

• Theo giá trị (by value): kẹp từng phần tử vào \([-c, c]\). Đơn giản nhưng làm đổi hướng vector gradient.
• Theo chuẩn (by norm): nếu \(\|g\| > \tau\) thì co cả vector \(g \leftarrow g\cdot\dfrac{\tau}{\|g\|}\). Giữ nguyên hướng, chỉ giảm độ dài. Đây là cách được khuyên dùng.

for images, labels in train_loader:
    loss = criterion(model(images), labels)
    optimizer.zero_grad()
    loss.backward()
    # đặt GIỮA backward() và step(); trả về norm trước khi clip, hữu ích để ghi log
    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    # hoặc: torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)
    optimizer.step()

Ví dụ: gradient \([30, 40]\) có norm 50; clip với max_norm=5 cho \([3, 4]\), cùng hướng, độ dài 5. Clip theo value 1.0 với \([30, -0.5]\) cho \([1, -0.5]\), hướng đã thay đổi. Khi dùng mixed precision (torch.amp), phải gọi scaler.unscale_(optimizer) trước khi clip.

Checklist khi huấn luyện "có vấn đề"

• Loss không giảm: in norm gradient từng lớp; nếu lớp đầu ~0 thì nghi vanishing (dùng ReLU, BN, skip connection, kiểm tra init).
• Loss thành NaN: giảm learning rate 3–10 lần; thêm gradient clipping; kiểm tra dữ liệu có NaN/Inf; kiểm tra phép log(0) hay chia 0 trong loss tự viết.
• Loss dao động mạnh: learning rate quá lớn hoặc batch quá nhỏ; thêm warmup.

# In norm gradient của từng lớp sau loss.backward()
for name, p in model.named_parameters():
    if p.grad is not None:
        print(f"{name:40s} {p.grad.norm().item():.2e}")

Tóm tắt

• Gradient lớp đầu là tích nhiều thừa số: nhỏ hơn 1 thì biến mất, lớn hơn 1 thì bùng nổ.
• Sigmoid có đạo hàm tối đa 0,25: tránh dùng ở lớp ẩn; dùng ReLU, LeakyReLU, SiLU, GELU.
• Không khởi tạo trọng số bằng hằng số; Xavier cho tanh/sigmoid (\(2/(n_{in}+n_{out})\)), He cho ReLU (\(2/n_{in}\)).
• Gradient clipping theo norm, đặt giữa backward() và step().
• BatchNorm và skip connection là hai "vũ khí" quan trọng nhất giúp mạng sâu huấn luyện được.

Bài tiếp theo trong chặng: Dataset, DataLoader và augmentation. Đọc thêm: Transfer Learning và Data Augmentation, hai kỹ thuật giúp bạn đạt kết quả tốt với ít dữ liệu.

Tài liệu tham khảo

• Glorot & Bengio (2010). Understanding the difficulty of training deep feedforward neural networks.
• He et al. (2015). Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification (He init, PReLU).
• Ramachandran et al. (2017). Searching for Activation Functions (Swish); Hendrycks & Gimpel (2016). GELU.
• Pascanu et al. (2013). On the difficulty of training recurrent neural networks (gradient clipping).
• CS231n: Neural Networks Part 2: Setting up the Data and the Loss (phần weight initialization).
• Tài liệu PyTorch: torch.nn.init, torch.nn.utils.clip_grad_norm_.