AdamW: cập nhật tham số từ gradient

1Định nghĩa

Bước cập nhật tham số (optimizer step) là bước thực sự sửa các ma trận trọng số sau khi đã có gradient từ lượt backward. Trong huấn luyện LLM, bộ tối ưu tiêu chuẩn là AdamW: thay vì trừ thẳng gradient thô vốn dao động mạnh, AdamW giữ cho mỗi tham số hai số nhớ, m (trung bình trượt của gradient, thể hiện hướng đi trung bình) và v (trung bình trượt của gradient bình phương, thể hiện độ biến thiên), để tự co giãn bước nhảy cho từng phần tử.

Song song, AdamW áp dụng suy giảm trọng số tách rời (decoupled weight decay) kéo nhẹ trọng số về gần 0, rồi nhân với tốc độ học (learning rate) theo lịch của vòng lặp huấn luyện. Vì phải giữ thêm m và v cho mọi tham số được huấn luyện, trạng thái AdamW chiếm một phần lớn ngân sách bộ nhớ lúc huấn luyện, và được bỏ đi khi xuất mô hình để phục vụ suy luận.

Ví dụ đời thường: AdamW giống một người lái tàu kinh nghiệm chỉnh bánh lái theo quán tính thay vì giật cục theo từng cơn gió lẻ: người đó giữ đà hướng đi của các nhịp trước (m), vừa tự giảm tay lái ở những khúc đường đang rung lắc mạnh (v) để tàu tiến về đích êm. Khác người lái, bộ tối ưu làm điều đó riêng cho từng tham số, hàng tỉ ‘bánh lái’ cùng lúc.

2Luồng

Đây là ô 3D. Mũi tên vào mang gradient (cùng shape trọng số); mũi tên ra mang trọng số mới và hai tensor m, v được giữ lại cho bước sau. Vòng lặp đóng lại ở ô 3E: trọng số mới đi ngược lên cột 2 cho lô kế tiếp.

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

3D3 · Learn from the Error

m, v moments · weight decay · θ ← θ − lr · update

3D · AdamW Update

AdamW đổi gradient thành bước cập nhật riêng cho từng tham số rồi sửa trọng số.

Nhận gì

Gradient từ ô 3B (đã đi qua ô 3C) và trọng số hiện tại.

Làm gì, vì sao

AdamW (bộ tối ưu, optimizer) giữ trung bình trượt của gradient (m) và của gradient bình phương (v) để chia độ dài bước theo từng tham số, rồi trừ đi một bước nhỏ nhân learning rate (lr). Suy giảm trọng số (weight decay) được trừ riêng. Gradient chỉ hướng; optimizer quyết định bước dài bao nhiêu.

Ra gì

Trọng số mới cùng m, v mới, đi xuống ô 3E.

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

Đọc nhãn: θ ← θ − lr · update, với θ là trọng số và lr là learning rate. Ở bước đầu tiên mỗi tham số đổi gần đúng lr theo dấu của gradient (lab: khoảng 0.0098 đến 0.0101 khi lr = 0.01).

Lưu ý

m và v chiếm thêm hai lần bộ nhớ của trọng số. Loss giảm sau một bước chưa chứng minh các bước sau cũng giảm; lr quá lớn vẫn làm loss tăng.

Đọc bài 3D · AdamW Update↗Giai đoạn 6 trong Bài Lab: AdamW↗

flowchart TD
  step1["Gradient hợp lệ"]
  step2["Cập nhật m và v"]
  step3["Hiệu chỉnh và tính bước"]
  step4["Áp weight decay và ghi trọng số"]
  step1 --> step2
  step2 --> step3
  step3 --> step4

3Giải pháp

Vị trí trên sơ đồ: ô 3D "AdamW" trong cột 3 · Học từ sai số, ngay sau backward (ô 3B và 3C) và trước ô 3E (vòng lặp). Đây là nơi duy nhất trong sơ đồ mà trọng số thực sự thay đổi.

Vấn đề cần giải. Backward cho biết hướng giảm loss tại điểm hiện tại. Cách đơn giản nhất là trọng số mới = trọng số cũ − lr × gradient (SGD). Cách này có hai nhược điểm: gradient của từng batch dao động mạnh vì mỗi batch khác nhau; và các tham số có độ lớn gradient rất khác nhau nhưng bị bước nhảy cùng hệ số lr. AdamW giải quyết cả hai bằng cách làm mượt và chuẩn hóa bước theo từng tham số.

Khái niệm tối thiểu. Với mỗi tham số θ ở bước t, gradient g. Bộ tối ưu giữ m và v, khởi đầu 0. m = β₁ × m + (1 − β₁) × g (trung bình trượt của g), v = β₂ × v + (1 − β₂) × g² (trung bình trượt của g bình phương). Hiệu chỉnh: m̂ = m / (1 − β₁^t), v̂ = v / (1 − β₂^t). Cập nhật: θ ← θ × (1 − lr × λ) − lr × m̂ / (căn v̂ + ε), trong đó λ là hệ số suy giảm trọng số.

Ví dụ có số, bước đầu tiên (minh họa, tính tay). Dùng gradient của bài backward cho ma trận LM head U. Chọn lr = 0,01, β₁ = 0,9, β₂ = 0,999, ε = 1e-8, λ = 0,01. Phần tử U[0, cá] = 0,4 có g = −0,8265. Bước 1: m = 0,1 × g = −0,08265; v = 0,001 × g² = 0,000683. Hiệu chỉnh: m̂ = m / 0,1 = −0,8265; v̂ = v / 0,001 = 0,6831 nên căn v̂ = 0,8265. Tỉ số m̂ / căn v̂ = −1. Cập nhật: θ = 0,4 × (1 − 0,0001) − 0,01 × (−1) = 0,39996 + 0,01 = 0,40996.

Phần tử khác: U[1, pad] = 0 có g = +0,2326, nhỏ hơn khoảng 3,6 lần. Hiệu chỉnh cho tỉ số m̂ / căn v̂ = +1 nên θ = 0 − 0,01 = −0,01. Hai tham số có gradient khác nhau khoảng 3,6 lần nhưng bước đầu đều dịch khoảng 0,01 theo hướng ngược dấu gradient (phần suy giảm trọng số rất nhỏ nên chỉ làm lệch ở chữ số thập phân thứ tư). Trong khi đó SGD với cùng lr = 0,01 sẽ dịch 0,0083 và 0,0023, tỉ lệ theo gradient. Ở những bước sau, độ lớn bước phụ thuộc cả lịch sử gradient, nên sẽ không còn luôn bằng lr.

Tác dụng của bias correction. Nếu bỏ hiệu chỉnh, tỉ số ở bước đầu là m / căn v = −0,08265 / 0,02613 ≈ −3,16, tức là bước lớn gấp 3,16 lần lr. Hiệu chỉnh bù việc m, v bắt đầu từ 0 nên thấp hơn giá trị thực, và dần biến mất khi t lớn.

Cách đọc kết quả. Nếu cập nhật cả 14 phần tử của U theo cách trên với h giữ nguyên, logit của cá tăng và xác suất của cá từ 0,1735 lên khoảng 0,1823; loss từ 1,752 xuống khoảng 1,702. Đây là số minh họa của một bước trên một vị trí, không phải đường cong học của một mô hình thật. Điều có thể kết luận: một bước đúng hướng làm loss tại điểm này giảm. Điều không thể kết luận: nhiều bước sẽ luôn giảm, hay lr = 0,01 là phù hợp cho mô hình thật.

Thực thi và chi phí bộ nhớ. Dùng cài đặt AdamW của framework với nhóm tham số rõ ràng (ví dụ không áp weight decay lên norm và bias nếu recipe yêu cầu); bản fused giảm overhead. Mỗi tham số cần thêm hai tensor m và v cùng shape, thường lưu bằng kiểu số 32 bit: với một tỉ tham số là 8 GB cho riêng m và v, chưa tính trọng số và gradient (xem ZeRO để chia nhỏ phần này). Lịch learning rate đổi theo bước (warmup rồi suy giảm) theo recipe. Không đổi siêu tham số chỉ để rút ngắn thời gian bước mà bỏ qua hội tụ; với mixed precision cần kiểm tra kiểu số của m, v và bản sao trọng số chính.

Để chạy một bước AdamW trên gradient thật và quan sát m, v, mở giai đoạn 6 · AdamW của Bài Lab tổng hợp. Hãy dự đoán trước độ lớn bước đầu của một tham số (gần lr?) rồi đối chiếu.

4Profiling

• Độ lớn cập nhật và chuẩn gradient

  Kiểm tra trước một bước trên vài tham số bằng tay như ví dụ, rồi theo dõi trọng số thay đổi bao nhiêu mỗi bước (tỉ lệ độ lớn cập nhật so với độ lớn trọng số) và chuẩn gradient theo thời gian. Quan sát tràn số, NaN hoặc phân kỳ (loss tăng liên tục). Các chỉ báo này cho biết huấn luyện có ổn định không, không cho biết mô hình tốt đến đâu.

• Bộ nhớ và thời gian của bộ tối ưu

  Đo kích thước m, v, bản sao trọng số chính (nếu dùng mixed precision) và thời gian chạy bước cập nhật trong tổng thời gian một bước. So fused với không fused trên cùng recipe, cùng phần cứng, và đừng kết luận về tốc độ trên mô hình lớn từ một mô hình nhỏ.

• Tiếp tục huấn luyện có đúng không

  Lưu và nạp lại cả trọng số lẫn m, v và bước t, rồi so các bước kế tiếp với một lần chạy liền mạch. Hai chuỗi giá trị phải trùng nhau trong sai số số học, nếu các nguồn ngẫu nhiên khác (thứ tự dữ liệu, dropout) cũng được khôi phục.

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ữ

First và second moments (m, v)

m là trung bình trượt của gradient (hướng đi), v là trung bình trượt của gradient bình phương (độ biến thiên); mỗi tham số có một cặp m, v riêng.

Bias correction

Chia m cho (1 − β₁^t) và v cho (1 − β₂^t) để bù việc m, v khởi đầu bằng 0 làm chúng thấp ở những bước đầu.

Learning rate (lr)

Hệ số điều khiển độ dài bước cập nhật; đổi theo lịch (warmup rồi giảm dần) trong vòng huấn luyện.

Decoupled weight decay

Kéo trọng số về 0 bằng một bước riêng θ ← θ(1 − lr × λ), không trộn vào gradient hay moment như Adam gốc.

β₁, β₂, ε

β₁ (thường 0,9) và β₂ (thường 0,999) là hệ số nhớ của m và v; ε (thường 1e-8) tránh chia cho 0.

7Bài liên quan

8Nguồn