Xavier & He Init — Sai từ đầu thì gradient chết ngay

Bạn đã bao giờ train một deep network và nhận được NaN loss ngay 10 step đầu tiên chưa? Hoặc nhìn training accuracy nằm chết dí ở mức random guessing hàng giờ liền mà loss không nhúc nhích? Chín trên mười lần, thủ phạm không phải learning rate hay dataset — mà là cách bạn khởi tạo weight ban đầu. Xavier và He initialization là hai "công thức chuẩn" giữ cho mạng sâu sống sót đủ lâu để thực sự học được, ngăn gradient biến mất thành số 0 hoặc bùng nổ thành infinity trước khi model kịp nhìn thấy mini-batch thứ hai.

Vấn đề

Hãy tưởng tượng bạn đang cố truyền tin nhắn qua một hành lang dài 100 phòng nối tiếp nhau. Nếu bạn thì thầm (weight khởi tạo quá nhỏ), giọng nói chết trước khi tới phòng thứ 5 — gradient trở nên không phân biệt được với 0, weight ngừng cập nhật, model chết ngay từ vạch xuất phát. Nếu bạn hét qua loa phóng thanh (weight quá lớn), tiếng vọng khuếch đại thành tiếng ồn inh tai — gradient overflow lên infinity, loss thành NaN, training run sụp đổ tức thì.

Trước năm 2010, các nhà nghiên cứu đúng nghĩa không thể train neural network sâu hơn 3-4 lớp một cách ổn định. Họ khởi tạo weight từ phân phối đều U[-1,1] hoặc Gaussian chuẩn N(0,1), rồi bất lực nhìn training curve hoặc nằm chết ở baseline accuracy (vanishing gradients) hoặc phóng thẳng lên infinity ngay backward pass đầu tiên (exploding gradients). "Thảm hoạ gradient flow" này là nút thắt cổ chai chính ngăn deep learning trở thành thực tế.

Và bước đột phá không phải optimizer mới hay activation function lạ — mà là nhận ra rằng variance của những con số ngẫu nhiên ban đầu cần được hiệu chuẩn theo kiến trúc mạng, cụ thể là số neuron đầu vào (fan-in) tại mỗi kết nối.

Ý tưởng cốt lõi

Trực giác cốt lõi đơn giản đến bất ngờ: giữ variance không đổi khi dữ liệu đi qua các lớp.

Khi tín hiệu đi qua một lớp tính y = Wx + b, variance đầu ra nên xấp xỉ bằng variance đầu vào. Nếu Var(y) co lại mỗi lớp, thông tin mờ dần thành nhiễu trắng. Nếu nó phình ra, tín hiệu bùng nổ thành hỗn loạn. Xavier (Glorot) và He initialization đơn giản là scale weight ngẫu nhiên ban đầu sao cho Var(output) ≈ Var(input), tính toán dựa trên số đầu vào mỗi neuron nhận và activation function theo sau.

Hãy nghĩ như áp suất nước trong hệ thống ống dẫn. Nếu bạn chia một ống chính thành 100 nhánh nhỏ (fan-in cao), bạn phải giảm áp suất mỗi nhánh tương ứng, nếu không lực tổng hợp sẽ vỡ mối nối. Xavier cân bằng bằng cách scale theo trung bình điều hoà của input và output size: weight được rút từ phân phối với variance $2/(n_\text{in} + n_\text{out})$. He initialization nhận ra rằng ReLU giết mất nửa tín hiệu (set giá trị âm thành 0), giảm variance đi một nửa, nên bù lại bằng cách nhân đôi scale: variance = $2/n_\text{in}$.

Vậy thôi. Không có phép thuật — chỉ là khớp mức độ ngẫu nhiên với hình học của mạng. Bằng cách rút weight từ phân phối có variance tỉ lệ với 1/n (n là số kết nối đầu vào), bạn đảm bảo tín hiệu gradient không bốc hơi cũng không nổ tung khi backpropagate qua hàng chục, thậm chí hàng trăm lớp.

Phá hiểu lầm: Đây không chỉ đơn giản là làm weight "nhỏ". Weight 0.01 có thể quá lớn nếu fan-in là 10,000, và weight 0.5 hoàn toàn ổn nếu fan-in là 4. Việc hiệu chuẩn là tương đối so với kết nối của lớp, không phải co weight lại một cách tuyệt đối.

Tại sao nó hoạt động

Toán đằng sau chỉ là theo dõi cẩn thận cách variance lan truyền qua phép nhân ma trận và hàm phi tuyến.

Xét một lớp tuyến tính với ma trận weight W. Nếu mỗi phần tử được khởi tạo độc lập từ phân phối trung bình 0, variance σ², và input x có variance 1, thì output y = Wx có variance xấp xỉ $n_\text{in} \times \sigma^2$ (với $n_\text{in}$ là fan-in). Lý do: variance cộng dồn khi bạn tổng các biến ngẫu nhiên độc lập.

Để bảo toàn cường độ tín hiệu (Var(y) = 1), ta giải:

n_in × σ² = 1  →  σ² = 1/n_in

Tuy nhiên, cần tính đến cả backward pass (gradient chảy ngược) và activation function:

• Với activation đối xứng (tanh, sigmoid) — bảo toàn variance gần như bằng nhau cả hai chiều — Xavier/Glorot đề xuất dùng trung bình fan_in và fan_out để cân bằng:

σ = sqrt(2 / (n_in + n_out))

• Với ReLU — zero ra khoảng nửa input (giá trị âm thành 0) — số kết nối active thực tế giảm một nửa, nên variance phải nhân đôi để bù:

σ = sqrt(2 / n_in)

Trong PyTorch, đoạn pseudo-code trông như sau:

# Xavier/Glorot (cho tanh, sigmoid, softmax)
W = torch.randn(fan_in, fan_out) * sqrt(2.0 / (fan_in + fan_out))

# He/Kaiming (cho ReLU, GELU, SiLU)
W = torch.randn(fan_in, fan_out) * sqrt(2.0 / fan_in)

Việc scale này đảm bảo trong quá trình backpropagation, variance của gradient cũng ổn định qua từng lớp, ngăn chặn sự suy giảm hoặc tăng trưởng theo cấp số nhân — thứ đã khiến deep network trước 2010 không thể train được.

Ý nghĩa thực tế

Không có các scheme khởi tạo này, các kiến trúc hiện đại như ResNet (152+ lớp) và Transformer (100+ lớp) đơn giản là không train được. Chúng sẽ hoặc output nhiễu ngẫu nhiên liên tục (vanishing gradients) hoặc overflow thành NaN (exploding gradients) trước khi xử lý xong batch thứ hai.

So sánh các chiến lược khởi tạo:

Ai đang dùng: Đây là mặc định trong mọi framework hiện đại. torch.nn.Linear của PyTorch dùng Kaiming/He uniform initialization mặc định. TensorFlow/Keras cung cấp glorot_uniform và he_normal. Khi bạn gọi model.init_weights() trong HuggingFace Transformers, bạn đang gọi chính xác các scheme này. Mọi model production — từ GPT-4 đến Llama đến BERT — đều dựa trên nền tảng toán học này để bắt đầu training ở vùng parameter space khả thi.

Các giới hạn: Initialization chỉ giải quyết vấn đề "bước đầu tiên" — đảm bảo mạng khởi đầu trên một con đường có thể đi được, nhưng không thể sửa kiến trúc tệ, dữ liệu kém, hay learning rate sai. Nó cũng giả định input features được normalize tương đối; nếu input có scale khác nhau nhiều hoặc mạng dùng skip connections (ResNet) hay normalization layers (BatchNorm/LayerNorm), các cơ chế này tương tác với initialization để ổn định training dynamics thêm.

Đào sâu hơn

• Paper gốc:

  • "Understanding the difficulty of training deep feedforward neural networks" (Glorot & Bengio, AISTATS 2010) — giới thiệu Xavier/Glorot initialization
  • "Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification" (He et al., ICCV 2015) — giới thiệu He initialization cho ReLU

• Cùng cụm:

  • Backpropagation — Chain rule chạy ngược, đó là tất cả
  • Activation Functions — ReLU, GELU, SiLU — Tại sao neuron cần "bật/tắt"
  • Residual Connections — Đường cao tốc cho gradient
  • Dropout — Giết neuron ngẫu nhiên để mạng mạnh hơn

• Đọc tiếp:

  • BatchNorm vs LayerNorm — Tại sao Transformer chọn LayerNorm (kỹ thuật bổ trợ để ổn định training)
  • Gradient Descent & Adam — Tại sao Adam thắng SGD (điều xảy ra sau khi khởi tạo xong)
  • Sequence Modeling — Các primitives này hoạt động như thế nào trong Transformer
  • Training at Scale — Khởi tạo tương tác với distributed training như thế nào

• External:

  • PyTorch documentation: torch.nn.init — chi tiết implementation thực tế