Chuẩn hóa vector và đường cộng residual trong transformer

1Định nghĩa

Chuẩn hóa (normalization) và đường cộng tắt (residual connection) là khung xương bao quanh mỗi khối con trong lớp Transformer, giúp tín hiệu và đạo hàm đi qua hàng chục lớp mà không bùng nổ hay triệt tiêu. Ở cấu trúc Pre-Norm phổ biến, vector ẩn x đi vào nhánh con được chuẩn hóa biên độ trước (bằng RMSNorm hoặc LayerNorm), đưa qua khối attention hoặc MLP, rồi cộng kết quả vào chính vector x ban đầu.

Nhờ phép cộng y = x + F(Norm(x)), ở lượt backward gradient luôn có một đường thẳng về các lớp đầu mà không phải chui qua mọi phép nhân ma trận. Khi chạy trên GPU, phép cộng residual và phép chuẩn hóa của nhánh kế tiếp thường được gộp trong một kernel để tiết kiệm một lượt đọc ghi bộ nhớ.

Ví dụ đời thường: Chuẩn hóa và đường tắt giống một đường ống dẫn nước qua tòa nhà cao tầng: van điều áp ở đầu mỗi tầng (norm) giữ áp lực nước vừa phải để ống không vỡ, còn ống chính chạy thẳng từ tầng một lên tầng thượng (đường residual) bảo đảm dòng nước gốc luôn thông suốt dù các nhánh lọc phụ ở từng tầng có phức tạp tới đâu. Khác đường ống thật, ở đây mỗi nhánh lọc còn cộng thêm phần nước đã xử lý vào ống chính.

2Luồng

Đây là ô 2C. Mũi tên chính (đường residual) mang vector x đi thẳng từ lớp này sang lớp kế; nhánh phụ đi qua norm và khối con rồi cộng ngược vào. Khi backward, gradient cũng đi theo đường thẳng này.

Sơ đồ kiến trúc huấn luyện · 22 ô (1A → 5C)

Kiến trúc huấn luyện LLM · Ánh xạ với Bài Lab tổng hợp

Thứ tự đọc: 1A → 1B ở cột 1, 2A → 2D ở cột 2, 3A → 3E ở cột 3; mũi tên next batch từ 3E quay về cột 1 và mũi tên load weights đi xuống băng 5 (5A → 5C). Cột 4 không phải bước kế tiếp: 4A, 4B tiết kiệm bộ nhớ, còn 4C–4G là các phương pháp huấn luyện thay thế nhau. Di chuột hoặc bấm vào từng ô để đọc pop-over.

Tất cả (22)1. Dữ liệu & nhãn2. Forward3. Học từ sai số4. Quy mô & phương pháp5. Suy luận

Ẩn khung overlay

1A1B1C2A2B2C2D3A3B3C3D3E4A4B4C4D4E4F4G5A5B5C

2C2 · Forward through the Transformer

x + Attn(Norm(x)) · x + MLP(Norm(x)) · × N layers

2C · Norm & Residual

Chuẩn hóa biên độ vector và cộng đường tắt residual quanh mỗi khối con, lặp N lớp.

Nhận gì

Đầu vào và đầu ra của từng khối con: attention (ô 2A) và MLP (ô 2B), theo mũi tên đi xuống.

Làm gì, vì sao

Norm (chuẩn hóa) đưa độ lớn vector về mức ổn định trước mỗi khối con. Residual (đường cộng tắt) cộng đầu vào trở lại đầu ra: x + Attn(Norm(x)). Nhờ đường tắt, thông tin và gradient đi thẳng qua nhiều lớp mà không bị bóp nghẹt.

Ra gì

Vector [T, d] mới của lớp. Lặp N lần (N lớp) rồi đi xuống ô 2D.

Ví dụ, cách đọc ô

Đọc nhãn: x + Attn(Norm(x)) rồi x + MLP(Norm(x)) là một lớp; × N layers là xếp N lớp giống nhau, mỗi lớp có trọng số riêng.

Lưu ý

Residual không có tham số riêng. Chỗ đặt norm (trước hay sau khối con) là lựa chọn thiết kế; ô này vẽ kiểu norm trước.

Đọc bài 2C · Norm & Residual↗Giai đoạn 3 trong Bài Lab: Forward: embedding, attention, MLP, norm → logits↗

flowchart TD
  step1["Vector ẩn x"]
  step2["Chuẩn hóa nhánh"]
  step3["Attention hoặc MLP"]
  step4["Cộng residual"]
  step1 --> step2
  step2 --> step3
  step3 --> step4

3Giải pháp

Vị trí trên sơ đồ: ô 2C "Norm & residual" trong cột 2 · Forward qua Transformer. Dù xếp thứ ba trong cột, về thứ tự chạy nó xuất hiện hai lần trong mỗi lớp: một lần trước attention (ô 2A) và một lần trước MLP (ô 2B), cùng hai phép cộng residual sau mỗi khối đó.

Vấn đề cần giải. Qua nhiều lớp, độ lớn của vector có thể phình ra hoặc co lại, khiến các phép tính phía sau (đặc biệt softmax và hàm kích hoạt) làm việc ở vùng kém ổn định, và gradient đi ngược về cũng bị khuếch đại hoặc triệt tiêu. Cần hai công cụ: giữ biên độ ổn định (norm) và giữ một đường thông suốt cho thông tin (residual).

Khái niệm tối thiểu. Một vector x gồm d số. 'Trung bình bình phương' (mean square) là (x1² + x2² + ... + xd²) / d; căn bậc hai của nó là RMS, đo độ lớn điển hình của các phần tử. RMSNorm(x) = γ ⊙ x / sqrt(mean(x²) + ε), trong đó γ là vector d số học được và ε rất nhỏ. Mỗi hàng của bảng (T × d) được chuẩn hóa riêng, không phụ thuộc hàng khác.

Ví dụ có số (minh họa, d = 2, γ = [1, 1], bỏ ε cho gọn). Vector x = [3, 4]. Trung bình bình phương = (9 + 16) / 2 = 12,5; RMS = 3,5355. Vector chuẩn hóa = [3 / 3,5355; 4 / 3,5355] = [0,8485; 1,1314]. Kiểm tra: trung bình bình phương của kết quả = (0,72 + 1,28) / 2 = 1, tức độ lớn đã được đưa về cỡ 1. Nếu nhân x lên 10 thành [30, 40], kết quả chuẩn hóa vẫn là [0,8485; 1,1314]: RMSNorm không phụ thuộc vào thang đo tổng thể của vector.

Residual với số. Giả sử khối con (attention hoặc MLP) nhận vector đã chuẩn hóa và cho ra F = [0,5; −0,5]. Cộng vào vector gốc chưa chuẩn hóa: y = x + F = [3 + 0,5; 4 − 0,5] = [3,5; 3,5]. Điểm quan trọng: x đi thẳng sang y mà không bị chuẩn hóa hay nhân ma trận, nên khối con chỉ học 'cần sửa thêm bao nhiêu' thay vì phải tái tạo toàn bộ vector.

Vì sao đường thẳng giúp học được. Ở backward, đạo hàm của y theo x gồm phần đi thẳng (hệ số 1) cộng phần đi qua khối con. Phần đi thẳng luôn tồn tại, nên gradient từ lớp trên về lớp dưới không bị nhân liên tiếp với hàng chục ma trận (xem gradient qua một lớp). Bên ngoài ví dụ này còn có các biến thể mới như Post-norm, Sandwich, QK-norm và residual cải tiến, được so sánh ở Normalization và Residual.

Điều có thể và không thể kết luận. Đã chuẩn hóa nghĩa là độ lớn của vector ổn định, chứ không có nghĩa là mô hình học tốt. Không đổi RMSNorm thành LayerNorm chỉ như một tối ưu của engine: đó là đổi phép tính của mô hình, và checkpoint huấn luyện với phép nào phải chạy bằng đúng phép đó.

Thực thi. Phép cộng residual và phép norm của nhánh sau đều đọc và ghi cùng một vector nên có thể gộp vào một kernel để giảm một vòng đọc ghi bộ nhớ, miễn giữ đúng thứ tự toán học. Sai khác nhỏ về sai số số học do gộp là bình thường, xem độ chính xác và sai số.

Để chạy chuẩn hóa trên các vector thật của một forward, mở giai đoạn 3 · Forward (khối norm và residual) của Bài Lab tổng hợp. Hãy tự dự đoán: sau RMSNorm với γ = 1, trung bình bình phương của mỗi hàng xấp xỉ bằng bao nhiêu?

4Profiling

• Khớp số học

  So đầu ra và gradient với cài đặt tham chiếu, bao gồm cả vector giá trị rất nhỏ (gần 0, dễ chia cho 0 nếu thiếu ε) và rất lớn (dễ tràn nếu dùng kiểu số ít bit). Với γ = 1, trung bình bình phương của đầu ra mỗi hàng phải xấp xỉ 1. Khớp tham chiếu chứng minh cài đặt đúng, không chứng minh mô hình học tốt hơn.

• Lưu lượng bộ nhớ

  Đo số vòng đọc ghi và độ trễ của phiên bản gộp và chưa gộp ở các shape đại diện, cùng kiểu số và cùng phần cứng. Chỉ báo giảm thời gian cả bước (step), chứ không chỉ thời gian của riêng kernel norm, vì phần norm thường nhỏ so với toàn bước.

• Độ lớn theo lớp

  Theo dõi RMS của vector residual sau từng lớp trên một batch cố định. Nếu độ lớn tăng rất nhanh theo độ sâu hoặc xuất hiện NaN, hãy kiểm tra ε, kiểu số và thứ tự norm. Con số này là chỉ báo chẩn đoán, không phải thước đo chất lượng.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

6Thuật ngữ

Root Mean Square Normalization (RMSNorm)

Chia vector cho căn bậc hai của trung bình bình phương các phần tử (cộng ε nhỏ để tránh chia cho 0) rồi nhân với hệ số học được γ.

LayerNorm

Trước khi chia, trừ đi trung bình của vector và chia cho độ lệch chuẩn; thường có thêm bias học được.

Pre-norm

Chuẩn hóa trước khi vào attention hoặc MLP, rồi cộng kết quả vào đường residual chưa chuẩn hóa; thứ tự cụ thể phụ thuộc mô hình.

Residual connection

Phép cộng y = x + F(Norm(x)) tạo đường tắt cho tín hiệu và gradient đi thẳng xuyên qua khối.

Epsilon (ε)

Số rất nhỏ cộng vào mẫu số để không chia cho 0 khi vector gần như toàn số 0.

7Bài liên quan

8Nguồn