Tính gradient để biết tham số cần thay đổi ra sao

1Định nghĩa

Backward pass (lan truyền ngược) là bước tính xem mỗi tham số của mô hình nên tăng hay giảm bao nhiêu để làm giảm hàm mất mát. Kết quả là gradient: một tensor cùng shape với tham số, trong đó mỗi phần tử là độ nhạy của loss đối với đúng tham số đó. Quá trình đi ngược từ giá trị loss ở đầu ra về các lớp Transformer và bảng embedding bằng quy tắc chuỗi (chain rule) của đạo hàm.

Để tính đạo hàm tại mỗi lớp (gradient từng lớp), hệ thống phải đọc lại các activation trung gian đã tạo ra ở lượt forward, hoặc tính lại chúng tại chỗ nếu bật activation checkpointing. Đầu ra là các tensor gradient để giao cho bộ tối ưu AdamW cập nhật trọng số.

Ví dụ đời thường: Backward giống buổi họp rút kinh nghiệm truy ngược từ khâu thành phẩm về đầu dây chuyền khi món ăn bị mặn: bếp trưởng lần lượt hỏi khâu nêm cuối, khâu nấu sốt rồi khâu ướp thịt đã góp bao nhiêu vào vị mặn đó, và ghi phiếu điều chỉnh riêng cho từng khâu (gradient). Giống như buổi họp thật, muốn truy đúng thì phải còn nhớ lại công thức và lượng muối đã dùng ở từng khâu, tương ứng với các activation đã lưu.

2Luồng

Đây là ô 3B; chiều mũi tên ngược với forward. Mũi tên vào mang gradient của loss theo đầu ra của một bước; mũi tên ra mang gradient theo đầu vào bước đó và theo trọng số của nó. Chi tiết trong một lớp là ô 3C.

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

3B3 · Learn from the Error

loss ← logits ← layers ← embedding · chain rule

3B · Backward Pass

Lan truyền ngược từ loss qua các lớp để biết mỗi tham số nên tăng hay giảm bao nhiêu.

Nhận gì

Gradient theo logits từ ô 3A (mũi tên đi xuống).

Làm gì, vì sao

Đi ngược từ loss qua đầu LM head, các lớp, tới embedding bằng quy tắc chuỗi (chain rule), để biết mỗi tham số nên tăng hay giảm bao nhiêu thì loss giảm. Nó đi ngược đúng thứ tự của forward.

Ra gì

Gradient của mọi tham số. Chúng đi xuống ô 3C để xem theo từng lớp và vào ô 3D để cập nhật.

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

Đọc nhãn từ trái sang phải: loss ← logits ← layers ← embedding. Mũi tên chạy ngược chiều forward.

Lưu ý

Gradient chỉ chỉ hướng dốc xuống ngay tại điểm hiện tại, chưa nói đi bao xa. Backward cần dùng lại các kết quả trung gian của forward (xem ô 4B).

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

flowchart TD
  step1["Loss và gradient theo logits"]
  step2["Qua LM head và bảng embedding"]
  step3["Qua từng lớp Transformer"]
  step4["Gradient của từng tham số"]
  step1 --> step2
  step2 --> step3
  step3 --> step4

3Giải pháp

Vị trí trên sơ đồ: ô 3B "Backward" trong cột 3 · Học từ sai số, ngay sau ô 3A (loss). Nó dùng đúng đồ thị tính toán của ô 2A đến 2D theo chiều ngược lại, và đưa kết quả cho ô 3D (AdamW). Cách gradient đi qua một lớp cụ thể là nội dung ô 3C (gradient từng lớp).

Vấn đề cần giải. Mô hình có hàng tỉ tham số. Ta cần biết với mỗi tham số, nếu nhích nó một chút thì loss thay đổi thế nào, nhưng thử từng tham số bằng cách chạy lại forward là không khả thi. Backward cho cả triệu độ nhạy đó chỉ với chi phí cỡ vài lượt forward.

Khái niệm tối thiểu. Đạo hàm của loss theo một tham số w nói: nếu w tăng một lượng nhỏ δ thì loss đổi khoảng (đạo hàm) × δ. Đạo hàm dương nghĩa là tăng w làm loss tăng, nên cần giảm w. Quy tắc chuỗi: nếu loss phụ thuộc vào u và u phụ thuộc vào w, thì đạo hàm theo w bằng đạo hàm theo u nhân đạo hàm của u theo w. Backward áp dụng quy tắc này lần lượt qua từng phép toán của forward theo thứ tự ngược.

Ví dụ nhỏ nhất (tính tay). Loss = (w × x − y)² với w = 2, x = 3, y = 4: w × x − y = 2 nên loss = 4. Gradient theo w = 2 × (w × x − y) × x = 2 × 2 × 3 = 12, dương nên cần giảm w. Với bước nhỏ 0,01: w mới = 2 − 0,01 × 12 = 1,88 và loss mới = (1,88 × 3 − 4)² = 2,6896, thấp hơn 4. Đó là toàn bộ ý tưởng; mạng lớn chỉ lặp lại hàng tỉ lần.

Ví dụ trong mô hình (số minh họa tiếp nối bài loss). Tại vị trí “ăn”, xác suất của 7 token là [0,116; 0,192; 0,105; 0,116; 0,174; 0,105; 0,192] và đáp án là cá (chỉ số 4). Gradient của loss theo logits là p trừ vector one-hot của đáp án: dz = [0,116; 0,192; 0,105; 0,116; −0,827; 0,105; 0,192]. Lưu ý tổng của dz bằng 0, và chỉ có thành phần của đáp án âm: tăng logit của cá làm loss giảm, tăng logit của mọi token khác làm loss tăng.

Đi tiếp ngược qua LM head. Với h = [1, 2] là vector đầu vào của LM head, gradient của ma trận U (shape 2 × 7) là tích ngoài của h và dz: hàng 0 bằng 1 × dz và hàng 1 bằng 2 × dz = [0,233; 0,384; 0,211; 0,233; −1,653; 0,211; 0,384]. Cột của cá có gradient âm lớn nhất, nghĩa là trọng số nối tới cá nên tăng lên. Gradient theo h là U nhân dz = [−0,286; 0,063]; nó được chuyển tiếp xuống các lớp Transformer bên dưới.

Cách đọc đầu ra và giới hạn. Gradient của một tensor có cùng shape với tensor đó; phần tử nào có độ lớn lớn là phần tử mà loss nhạy nhất ở điểm hiện tại. Gradient là thông tin cục bộ: nó cho hướng giảm loss ngay tại điểm này, không đảm bảo một bước dài theo hướng đó vẫn giảm. Vì vậy sau backward còn cần learning rate và bộ tối ưu. Cũng đừng đọc độ lớn gradient như tầm quan trọng của tham số cho mọi dữ liệu.

Thực thi. PyTorch autograd dựng đồ thị các phép cần đạo hàm trong forward rồi chạy ngược. Backward cần các activation của forward, nên phải giữ chúng trong bộ nhớ hoặc tính lại (activation checkpointing); với phép nhân ma trận, backward gồm hai phép nhân (một cho gradient đầu vào, một cho gradient trọng số), nên chi phí tính toán thường xấp xỉ gấp đôi forward theo ước lượng lý thuyết, còn số đo thật phụ thuộc phần cứng và kernel. Gradient accumulation cộng gradient của nhiều micro-batch trước một lần cập nhật; khi huấn luyện trên nhiều thiết bị, gradient của các phần dữ liệu cần được kết hợp theo chiến lược chọn (xem ZeRO).

Chuẩn hóa loss nhất quán giữa các micro-batch, nếu không các batch có số token hợp lệ khác nhau sẽ đóng góp sai trọng số. Ví dụ: micro-batch A có 3 token hợp lệ với loss trung bình 2,0 (tổng 6,0); micro-batch B có 1 token hợp lệ với loss 4,0. Trung bình đúng trên mỗi token = (6,0 + 4,0) / 4 = 2,5. Trung bình hai số trung bình = (2,0 + 4,0) / 2 = 3,0, đã thổi phồng token duy nhất của B.

Để tính gradient cho mọi tensor tham số và kiểm tra 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. Hãy đoán dấu gradient của cột đáp án trước khi chạy.

4Profiling

• Tính đúng của gradient

  Dùng bài toán nhỏ hoặc sai phân hữu hạn để đối chiếu: nhích một tham số lên 0,001, đo loss đổi bao nhiêu rồi chia cho 0,001, kết quả phải gần gradient do backward tính. Ví dụ nhích phần tử U[0, cá] lên 0,001 thì loss giảm khoảng 0,00083, tỉ số khoảng −0,826, khớp với −0,8265 ở ví dụ trên. Đây là phép kiểm tra tại một điểm, không chứng minh mọi gradient đúng.

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

  Đo riêng bộ nhớ của các tensor được giữ lại, thời gian tính lại (nếu có checkpointing) và lượng giao tiếp giữa các thiết bị. Theo dõi chuẩn gradient (độ lớn tổng) theo bước và số lần xuất hiện NaN; chuẩn tăng vọt hoặc NaN là dấu hiệu cần dừng và kiểm tra trước khi cho bộ tối ưu cập nhật.

• Nhất quán khi cộng dồn

  Chạy cùng một batch với một lần backward lớn và với nhiều micro-batch cộng dồn, rồi so gradient. Hai kết quả phải trùng nhau trong sai số số học nếu loss được chuẩn hóa theo cùng số token hợp lệ. Nếu lệch, kiểm tra mẫu số của loss.

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ữ

Gradient

Tập đạo hàm của loss theo từng tham số; cùng shape với tham số, dấu cho biết hướng loss tăng.

Chain rule (quy tắc chuỗi)

Đạo hàm của hàm hợp bằng tích các đạo hàm thành phần; cho phép đi ngược từng bước từ loss về tham số.

Saved activations

Các giá trị trung gian của forward cần dùng để tính gradient; được giữ lại hoặc tính lại.

Gradient accumulation

Cộng gradient của nhiều micro-batch trước khi cập nhật trọng số một lần, để mô phỏng batch lớn khi bộ nhớ có hạn.

Finite differences

Cách kiểm tra gradient: nhích tham số một lượng rất nhỏ, đo loss thay đổi bao nhiêu rồi chia cho lượng nhích.

7Bài liên quan

8Nguồn