Gradient đi qua một lớp transformer

1Định nghĩa

Gradient qua một lớp Transformer là cách đạo hàm từ phía trên đi xuyên ngược một lớp để ra hai thứ: gradient cho lớp phía dưới và gradient cho từng ma trận trọng số của lớp này. Trong bước backward, gradient đến đầu ra của lớp Transformer hiện tại rồi tách làm hai ở mỗi phép cộng tắt (residual): một nhánh đi thẳng về phía đầu vào của lớp, nhánh còn lại đi ngược vào khối MLP rồi tới khối attention.

Tại mỗi phép nhân ma trận bên trong, hệ thống làm hai phép tính cùng lúc: nhân gradient với ma trận trọng số chuyển vị để đẩy tiếp về lớp trước (grad_input), và nhân với vector kích hoạt đã lưu từ forward để ra gradient của chính ma trận trọng số (grad_weight) giao cho AdamW.

Ví dụ đời thường: Gradient qua một lớp giống cách phân bổ trách nhiệm trong một phòng ban khi sản phẩm cần chỉnh sửa: một phần phản hồi được chuyển thẳng lên phòng ban trước đó qua đường ưu tiên (nhánh residual), phần còn lại chia cho tổ nội dung (MLP) và tổ tra cứu (attention) để mỗi tổ tự biết cần sửa quy trình của mình bao nhiêu. Khác phòng ban thật, mỗi tổ vừa chuyển phản hồi lên trên vừa giữ lại phần dành cho chính mình.

2Luồng

Đây là ô 3C, chi tiết của ô 3B cho một lớp. Thứ tự đi ngược là đảo của forward: residual cuối, MLP, norm, rồi residual đầu, attention, norm. Mỗi mũi tên mang hai thứ: gradient đi tiếp xuống lớp dưới và gradient trọng số dừng lại chờ bộ tối ưu.

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

3C3 · Learn from the Error

grad norm của layer 1 … layer N

3C · Per-Layer Gradients

Gom gradient theo từng lớp và đo độ lớn để thấy mỗi lớp nhận tín hiệu sửa mạnh hay yếu.

Nhận gì

Gradient của mọi tham số do ô 3B tính (mũi tên đi xuống).

Làm gì, vì sao

Gom gradient theo từng lớp và đo độ lớn của chúng (grad norm: độ dài tổng của gradient). Gradient quá lớn hoặc gần bằng 0 ở một lớp là dấu hiệu huấn luyện có vấn đề. Ô này là cách quan sát và không thêm phép tính mới cho mô hình.

Ra gì

Gradient được giữ nguyên và đi tiếp xuống ô 3D. Các con số grad norm dùng để theo dõi và để cắt bớt gradient quá lớn (clipping).

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

Thanh ngang ở mỗi hàng layer 1 … layer N là grad norm của lớp đó; thanh dài hơn nghĩa là gradient lớn hơn. Ở mô hình 2 lớp của lab, lớp 0 là 0.9951 và lớp 1 là 0.6791.

Lưu ý

Grad norm là chỉ báo chẩn đoán, không phải mục tiêu. So sánh giữa các lớp chỉ có nghĩa khi cùng cách đo.

Đọc bài 3C · Per-Layer Gradients↗Giai đoạn 5 trong Bài Lab: Backward & gradient từng lớp↗

flowchart TD
  step1["Gradient ở đầu ra lớp"]
  step2["Rẽ nhánh ở residual cuối"]
  step3["Ngược qua MLP và norm"]
  step4["Ngược qua attention và nhập nhánh"]
  step5["Giao grad_weight cho bộ tối ưu"]
  step1 --> step2
  step2 --> step3
  step3 --> step4
  step4 --> step5

3Giải pháp

Vị trí trên sơ đồ: ô 3C "Gradient từng lớp" trong cột 3 · Học từ sai số, phóng to một đoạn của ô 3B (backward). Nó chỉ nói về một lớp; mô hình có L lớp thì ô này lặp L lần từ lớp trên cùng xuống lớp dưới cùng.

Vấn đề cần giải. Ta đã biết backward đi ngược qua đồ thị. Câu hỏi cụ thể: bên trong một lớp có residual, norm, attention, MLP; gradient đi qua từng phần đó theo đường nào, và mỗi ma trận trọng số nhận gradient ở đâu?

Khái niệm tối thiểu. Với y = x·W (x là hàng d số, W là ma trận d × m), nếu biết gradient theo y là dy thì: gradient theo đầu vào dx = dy nhân W chuyển vị (cùng shape với x), và gradient theo trọng số dW = x chuyển vị nhân dy (cùng shape với W). Dùng x ở công thức dW, nên x phải còn trong bộ nhớ hoặc phải tính lại. Với y = x + F(x), gradient theo y được sao cho cả nhánh thẳng và nhánh qua F, rồi hai phần gradient theo x được cộng lại.

Ví dụ có số (minh họa, tính tay). Khối y = x + x·W với x = [2, 4], W = [[0,5, 0], [0, −0,5]] (nhánh F(x) = x·W). Forward: x·W = [1, −2], y = [3, 2]. Để đơn giản, coi loss L = y0 + y1 = 5, nên dy = [1, 1].

Backward qua khối đó. Nhánh thẳng chuyển dy = [1, 1] nguyên vẹn về x. Nhánh qua F: dx_nhánh = dy nhân W chuyển vị = [1 × 0,5 + 1 × 0; 1 × 0 + 1 × (−0,5)] = [0,5; −0,5]. Cộng hai nhánh: dx = [1,5; 0,5]. Gradient trọng số: dW = x chuyển vị nhân dy = [[2 × 1, 2 × 1], [4 × 1, 4 × 1]] = [[2, 2], [4, 4]].

Kiểm tra bằng sai phân hữu hạn: nhích x0 từ 2 lên 2,001 thì y0 = 2,001 + 2,001 × 0,5 = 3,0015 nên L tăng thêm 0,0015 và tỉ số 0,0015 / 0,001 = 1,5, khớp dx0. Tương tự x1 cho 0,5 và nhích W00 từ 0,5 lên 0,501 làm L tăng 0,002, tỉ số 2,0 khớp dW00. Phép kiểm tra này cho thấy công thức đúng ở điểm đang xét.

Cách đọc. Nhánh residual đóng góp phần 1 (dx = dy + phần qua F), nên ngay cả khi F rất nhỏ hay rất lớn, gradient vẫn có một phần chắc chắn đến được x; đó là lý do gradient không biến mất sau nhiều lớp. Thứ tự đi ngược trong một lớp đảo ngược forward: cộng residual cuối, MLP, norm, cộng residual đầu, attention, norm. Mỗi phép tuyến tính tạo một cặp (grad_input đi tiếp, grad_weight dừng lại). Điều không thể kết luận: gradient lớn không có nghĩa là lớp đó 'quan trọng' nhất, chỉ là loss nhạy với nó tại điểm này.

Thực thi. Autograd lo việc nối các đạo hàm thành phần; gradient accumulation cộng đóng góp từ nhiều micro-batch trước một lần cập nhật; activation checkpointing tính lại một số giá trị trung gian để tiết kiệm bộ nhớ. Hãy chuẩn hóa loss nhất quán giữa các micro-batch như bài backward.

Để kiểm tra gradient của cả mô hình bằng sai phân hữu hạn, mở giai đoạn 5 · Backward & gradient từng lớp của Bài Lab tổng hợp. Trước khi chạy, hãy đoán ma trận trọng số nào có gradient cùng shape với nó và vì sao.

4Profiling

• Tính đúng của gradient

  Dùng mô hình rất nhỏ và sai phân hữu hạn như ví dụ để đối chiếu gradient từng tensor. Kiểm tra shape của mỗi grad_weight trùng shape trọng số. Khớp tại một số điểm cho thấy cài đặt đúng ở những điểm đó, không thay thế một bộ kiểm thử đầy đủ.

• Thời gian và bộ nhớ backward

  Đo bộ nhớ của các tensor được giữ lại cho từng khối, thời gian tính lại và lượng giao tiếp giữa thiết bị. Tách thời gian của attention, MLP và norm để xem khối nào chiếm nhiều nhất; theo dõi chuẩn gradient từng lớp và NaN.

• Gradient theo độ sâu

  In chuẩn gradient của một tensor tiêu biểu ở lớp đầu, giữa và cuối trên cùng batch. Nếu chuẩn ở lớp đầu nhỏ hơn hàng triệu lần, có thể có vấn đề khởi tạo, norm hoặc kiểu số. Đây là chẩn đoán, không phải kết luận về 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ữ

grad_input và grad_weight

Tại mỗi phép nhân ma trận, backward tạo hai kết quả: gradient theo đầu vào (để tiếp tục đi ngược) và gradient theo trọng số (để bộ tối ưu cập nhật).

Ma trận chuyển vị

Ma trận đã đổi hàng thành cột; gradient đầu vào của y = x·W là gradient đầu ra nhân với W chuyển vị.

Rẽ nhánh residual

Tại phép cộng y = x + F(x), gradient của y được sao nguyên cho cả hai nhánh: nhánh thẳng về x và nhánh qua F.

Saved activations

Các giá trị của forward (ví dụ x) cần để tính grad_weight; thiếu chúng phải tính lại.

Chain rule

Quy tắc nối các đạo hàm thành phần; trong một lớp có nhiều đoạn, mỗi đoạn nhân vào gradient đang đi ngược.

7Bài liên quan

8Nguồn