Bài 17: Optimizer cơ bản: SGD và Adam

Sau bài học, bạn sẽ:

• Viết được update rule của SGD, SGD-momentum, Nesterov, AdaGrad, RMSProp, Adam, AdamW.
• Hiểu vì sao momentum giúp vượt qua "rãnh hẹp" và adaptive LR của Adam giúp train ổn định với mọi scale của gradient.
• Viết train step chuẩn zero_grad → loss.backward → optimizer.step và biết hệ quả khi quên zero_grad.
• Chọn được optimizer hợp lý: SGD+momentum cho CNN/ImageNet, Adam cho NLP, AdamW cho Transformer modern.
• Biết cách set learning rate khác cho từng layer (param groups), save/load state để resume training.
• Có cảm nhận sơ bộ về Lion, Sophia, 8-bit Adam — các hướng tối ưu mới cho LLM.

Bài này nối tiếp Bài 16 — nn.Module và là phần cuối ghép trước khi vào Bài 18 — Training Loop: model (B16) + loss (B12, B13) + autograd (B15) + optimizer (bài này) = một step train hoàn chỉnh.

Series 1 Bài 27 đã giới thiệu gradient descent. Update rule cơ bản:

\[ \theta \leftarrow \theta - \eta \, \nabla L(\theta) \]

Trong đó \( \theta \) là tập tham số (weight, bias), \( \eta \) là learning rate, \( \nabla L(\theta) \) là gradient của loss theo parameter. Ý nghĩa hình học: di chuyển ngược hướng gradient với bước dài \( \eta \) — vì gradient chỉ hướng tăng nhanh nhất của \( L \), đi ngược lại sẽ giảm \( L \) nhanh nhất (trong vùng tuyến tính địa phương).

Trong deep learning, \( \nabla L(\theta) \) tính trên toàn bộ training set là quá đắt. Thực tế dùng mini-batch: lấy một batch nhỏ \( B \) sample, ước lượng gradient bằng trung bình trên batch:

\[ g_t = \frac{1}{|B|} \sum_{i \in B} \nabla L_i(\theta) \]

Đây là gốc của tên gọi Stochastic Gradient Descent (SGD) — gradient ước lượng từ một subset ngẫu nhiên thay vì toàn bộ data. Mọi optimizer trong bài này đều dựa trên \( g_t \) này; khác biệt nằm ở cách dùng \( g_t \) để cập nhật \( \theta \).

Optimizer trong PyTorch (torch.optim) là class wrap hai thứ:

1. Update rule: hàm thuần đại số đi từ (parameter, gradient) sang parameter mới.
2. State nội tại: buffer cần lưu giữa các step — momentum buffer, moving average của \( g^2 \) cho Adam, step counter cho bias correction...

Khởi tạo cần list parameter và hyperparameter:

import torch
import torch.optim as optim

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# hoặc
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# hoặc
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)

model.parameters() trả về iterator của các nn.Parameter (xem B16). Optimizer giữ reference tới chính những tensor đó — không copy. Khi gọi optimizer.step(), nó đọc p.grad đã được autograd điền sẵn, áp update rule, ghi đè p.data.

Bốn method quan trọng:

• optimizer.step(): thực hiện một update step trên mọi parameter.
• optimizer.zero_grad(): zero .grad của mọi parameter — bắt buộc trước backward() mới vì autograd cộng dồn (xem B15).
• optimizer.state_dict() / load_state_dict(): serialize/restore state để checkpoint.
• optimizer.param_groups: list các nhóm parameter với hyperparameter riêng (xem mục 16).

SGD vanilla là phiên bản đơn giản nhất, không state nội tại:

\[ \theta_{t+1} \leftarrow \theta_t - \eta \cdot g_t \]

Trong đó \( g_t \) là gradient của loss trên mini-batch ở step \( t \). Khởi tạo trong PyTorch:

optimizer = optim.SGD(model.parameters(), lr=0.01)

Đặc điểm:

• Không nhớ history. Mỗi step chỉ phụ thuộc gradient hiện tại.
• Dao động mạnh trong vùng cong (curvature) không đều — gradient hai trục có scale khác nhau làm bước đi zig-zag.
• Dễ kẹt ở local minimum hoặc plateau (vùng gradient rất nhỏ).
• Cần learning rate tinh chỉnh kỹ và thường kèm LR schedule (giảm dần) để hội tụ tốt.

Vanilla SGD hiếm khi được dùng trực tiếp trong deep learning hiện đại — gần như luôn kèm momentum (mục 5).

Momentum tích luỹ "đà" theo hướng gradient nhất quán. State là vector \( v_t \) cùng shape với parameter:

\[ v_t = \mu \, v_{t-1} + g_t \]

\[ \theta_{t+1} \leftarrow \theta_t - \eta \, v_t \]

Hệ số \( \mu \) (momentum coefficient) thường set \( 0.9 \). Khi gradient nhiều step liên tiếp cùng hướng, \( v_t \) tích lại lớn dần → bước đi dài hơn theo hướng đó. Khi gradient đổi hướng (oscillation), thành phần cũ trong \( v_t \) triệt tiêu thành phần mới → bước đi nhỏ lại theo hướng đang dao động.

Lợi ích:

• Tăng tốc ở vùng có gradient đều — trung bình bước đi lớn hơn vanilla SGD.
• Giảm dao động ở vùng "rãnh hẹp" (ravine) — Adam paper minh hoạ rất rõ bằng quadratic ill-conditioned.
• Vượt qua local minimum cạn — đà giúp đi qua plateau và saddle point.

PyTorch:

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

Lưu ý quy ước PyTorch: v_t = momentum * v_{t-1} + g_t rồi theta -= lr * v_t. Một số tài liệu khác viết v_t = momentum * v_{t-1} + lr * g_t; theta -= v_t — hai dạng cho cùng kết quả nếu giữ \( \mu, \eta \) cố định, nhưng khi thay đổi LR sẽ ảnh hưởng khác nhau lên state cũ. Đọc doc nếu đối chiếu với paper.

Nesterov accelerated gradient (Nesterov 1983, áp dụng cho NN bởi Sutskever 2013) là variant của momentum với một phép "nhìn trước":

\[ v_t = \mu \, v_{t-1} + \nabla L(\theta_t - \eta \, \mu \, v_{t-1}) \]

Khác biệt: gradient được tính tại vị trí nếu chỉ áp momentum cũ chứ không phải tại \( \theta_t \). Trực giác: trước khi quyết bước đi tiếp theo, hãy xem mình sẽ trôi tới đâu nếu chỉ có momentum, rồi mới tính gradient ở chỗ đó. Nhờ "look-ahead", Nesterov điều chỉnh sớm hơn khi sắp vượt qua minimum.

PyTorch thực hiện reformulation của Sutskever để khỏi cần đánh giá gradient ở vị trí khác:

optimizer = optim.SGD(model.parameters(), lr=0.01,
                      momentum=0.9, nesterov=True)

Yêu cầu: phải có momentum > 0 và dampening = 0. Trong thực hành Nesterov thường ổn định hơn momentum thường khoảng vài phần trăm trên CNN — không khác biệt lớn, nhưng "free" nếu bạn đã dùng momentum.

AdaGrad (Duchi et al. 2011) là optimizer đầu tiên dùng learning rate riêng cho từng parameter, dựa trên lịch sử gradient. State là tổng tích luỹ gradient bình phương:

\[ G_t = G_{t-1} + g_t^2 \]

\[ \theta_{t+1} \leftarrow \theta_t - \frac{\eta}{\sqrt{G_t} + \epsilon} \cdot g_t \]

Phép tính trên là element-wise: mỗi parameter có \( G_t \) riêng. Parameter mà gradient lớn liên tục → \( G_t \) lớn → LR hiệu dụng nhỏ; parameter có gradient nhỏ → \( G_t \) nhỏ → LR hiệu dụng lớn. Ý tưởng rất tốt cho dữ liệu sparse (NLP với embedding rất nhiều feature hiếm).

Vấn đề: \( G_t \) chỉ tăng, không giảm. Theo thời gian \( \sqrt{G_t} \to \infty \), LR hiệu dụng \( \to 0 \), parameter ngừng cập nhật. Train deep network nhiều epoch sẽ "chết" giữa chừng. PyTorch vẫn có optim.Adagrad nhưng ít dùng cho NN sâu hiện đại; RMSProp và Adam là sửa lỗi này.

RMSProp (Tieleman & Hinton 2012, công bố qua slide course Coursera, không paper chính thức) thay tổng tích luỹ bằng exponentially weighted moving average của \( g^2 \):

\[ v_t = \rho \, v_{t-1} + (1 - \rho) \, g_t^2 \]

\[ \theta_{t+1} \leftarrow \theta_t - \frac{\eta}{\sqrt{v_t} + \epsilon} \cdot g_t \]

Với \( \rho \) thường \( 0.9 \) hoặc \( 0.99 \). EMA cho \( v_t \) "quên" gradient cũ — chỉ giữ statistic gần đây. Nhờ đó LR hiệu dụng không tiến tới 0 mà bám theo scale gradient hiện tại của từng parameter.

RMSProp ổn định và hiệu quả với RNN — Karpathy và nhiều tác giả 2014–2016 dùng cho language model. Adam (mục 9) lấy ý tưởng cốt lõi của RMSProp và bổ sung thêm momentum cho gradient.

optimizer = optim.RMSprop(model.parameters(), lr=1e-3, alpha=0.99)

Adam (Adaptive Moment Estimation, Kingma & Ba 2014, arXiv:1412.6980) kết hợp momentum (như SGD-momentum) và adaptive LR per parameter (như RMSProp). State gồm hai moving average:

\[ m_t = \beta_1 \, m_{t-1} + (1 - \beta_1) \, g_t \]

\[ v_t = \beta_2 \, v_{t-1} + (1 - \beta_2) \, g_t^2 \]

\( m_t \) là ước lượng first moment (mean) của gradient; \( v_t \) là second raw moment (uncentered variance) — chính là EMA của \( g^2 \) như RMSProp.

Bias correction — đầu training \( m_0 = v_0 = 0 \) làm \( m_t, v_t \) bị "lệch về 0". Sửa:

\[ \hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1 - \beta_2^t} \]

Update:

\[ \theta_{t+1} \leftarrow \theta_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \]

Hyperparameter default (theo paper, cũng là default PyTorch): \( \beta_1 = 0.9 \), \( \beta_2 = 0.999 \), \( \epsilon = 10^{-8} \), \( \eta = 10^{-3} \).

Tại sao Adam phổ biến:

• Adaptive LR per parameter → ít nhạy với LR khởi đầu so với SGD.
• Momentum \( m_t \) làm bước đi mượt như SGD-momentum.
• Bias correction quan trọng ở vài chục step đầu — không có nó, \( m_t \) và \( v_t \) quá nhỏ làm update sai scale.
• Default hyperparameter chạy tốt cho rất nhiều bài toán mà không cần tune.

optimizer = optim.Adam(model.parameters(), lr=1e-3,
                       betas=(0.9, 0.999), eps=1e-8)

L2 regularization cộng \( \frac{\lambda}{2} \|\theta\|^2 \) vào loss, tương đương cộng \( \lambda \theta \) vào gradient. Trong SGD vanilla, hai cách này thực sự tương đương. Trong Adam, không: khi cộng \( \lambda \theta \) vào \( g_t \), thành phần đó cũng đi qua \( m_t, v_t, \hat{v}_t \), và chia cho \( \sqrt{\hat{v}_t} \) — parameter có gradient bình thường lớn sẽ bị regularization yếu, parameter nhỏ bị regularization mạnh quá mức. Hệ quả: Adam + L2 thường không ổn định và underperform khi so với SGD + L2.

AdamW (Loshchilov & Hutter 2017, arXiv:1711.05101) sửa bằng cách tách hẳn weight decay khỏi gradient:

\[ \theta_{t+1} \leftarrow \theta_t - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \, \theta_t \right) \]

Term \( \lambda \theta_t \) không vào \( m_t, v_t \) — nó là một bước decay trực tiếp lên parameter. AdamW giờ là default cho:

• Transformer modern (BERT, GPT, ViT, LLaMA...).
• Hầu hết LLM pretraining và fine-tuning.
• HuggingFace Transformers Trainer mặc định AdamW.

optimizer = optim.AdamW(model.parameters(),
                        lr=1e-3, weight_decay=0.01)

Cảnh báo: optim.Adam(..., weight_decay=x) trong PyTorch vẫn là L2 cũ. Muốn decoupled, dùng AdamW rõ ràng.

Một train step chuẩn trong PyTorch luôn theo bộ ba:

optimizer.zero_grad()          # 1. reset grad
out = model(x)
loss = criterion(out, y)
loss.backward()                # 2. tính grad qua autograd
optimizer.step()               # 3. update parameter

Giải thích từng dòng:

• zero_grad() set p.grad = 0 (hoặc p.grad = None với set_to_none=True) cho mọi p trong optimizer.param_groups. Vì autograd cộng dồn (B15), nếu không zero thì gradient của batch cũ vẫn nằm đó.
• loss.backward() traverse computation graph ngược, điền p.grad cho từng leaf có requires_grad=True.
• optimizer.step() đọc p.grad, áp update rule (SGD / Adam / AdamW...), ghi p.data.

Thứ tự ba dòng có thể đảo nhẹ: nhiều codebase đặt zero_grad() ngay sau step() của step trước thay vì đầu step mới — tương đương về kết quả, nhưng có ích khi cần truy cập p.grad để log trước khi reset.

Đây là pitfall phổ biến nhất khi mới viết training loop. Bỏ quên zero_grad():

# SAI — quên zero_grad
for x, y in loader:
    out = model(x)
    loss = criterion(out, y)
    loss.backward()            # grad cộng dồn từ batch trước
    optimizer.step()

Hệ quả: ở step \( t \), p.grad thực ra bằng tổng gradient của tất cả các batch từ đầu epoch. Update sai về scale và sai về hướng. Triệu chứng quan sát được:

• Loss tăng nhanh hoặc oscillate dữ dội.
• NaN / Inf sau vài chục step.
• Train tưởng "không hội tụ" trong khi thật ra LR effective cứ lớn dần.

Sửa: thêm optimizer.zero_grad() đầu mỗi iteration. Có một trường hợp cố ý không zero — gradient accumulation để simulate batch lớn (xem B15 mục 12) — nhưng vẫn phải zero sau \( k \) batch.

Mẹo nhỏ: optimizer.zero_grad(set_to_none=True) (default từ PyTorch 2.0) đặt p.grad = None thay vì tensor zero. Nhanh hơn (không phải fill 0) và tiết kiệm bộ nhớ một chút; backward sẽ tự cấp phát lại.

Thực tế các benchmark và codebase công khai cho thấy phân chia tương đối rõ:

• SGD + momentum + LR schedule: thường cho final accuracy nhỉnh hơn Adam khoảng \( 0.1 \)–\( 0.5\% \) trên image classification (ResNet, EfficientNet trên ImageNet). Lý do được tranh luận trong nhiều paper — một số quy cho việc Adam tìm điểm minimum "sharp" hơn, generalize kém hơn SGD vốn tìm điểm "flat".
• Adam: default cho NLP, RL, GAN, mạng có loss landscape khó (gradient scale rất khác giữa các layer). Hầu hết paper Transformer pre-AdamW dùng Adam.
• AdamW: default cho Transformer modern, LLM, ViT. Khi có weight decay là yếu tố thực sự cần (regularization mạnh trên model lớn), AdamW vượt rõ Adam-with-L2.
• Switch SGD ở cuối training: một số recipe (Keskar & Socher 2017, "Improving Generalization Performance by Switching from Adam to SGD") đề xuất train Adam giai đoạn đầu cho nhanh, switch SGD-momentum giai đoạn cuối để hội tụ tốt hơn.

Quy tắc heuristic cho người mới: thử Adam (hoặc AdamW) trước với LR 1e-3. Nếu kết quả không đạt và bạn nghi ngờ generalize kém, đổi sang SGD + momentum 0.9 + LR schedule. Đừng tốn thời gian tinh chỉnh SGD nếu chưa có baseline Adam.

Learning rate \( \eta \) là hyperparameter quan trọng nhất — nhạy hơn nhiều so với batch size, optimizer hay weight decay.

Giá trị start thường dùng:

• SGD: \( 0.01 \)–\( 0.1 \) cho image classification, \( 0.001 \)–\( 0.01 \) cho NLP.
• Adam / AdamW: \( 10^{-4} \)–\( 10^{-3} \). Pretraining LLM thường \( 10^{-4} \)–\( 6 \cdot 10^{-4} \); fine-tuning LLM thường \( 10^{-5} \)–\( 5 \cdot 10^{-5} \).

Triệu chứng khi LR sai:

• Quá lớn: loss tăng vọt, NaN sau vài step, hoặc "divergence" (loss tăng dần thay vì giảm). Giảm LR 10 lần và thử lại.
• Quá nhỏ: loss giảm nhưng cực chậm, mỗi epoch chỉ giảm vài phần trăm so với expected. Tăng LR 3–10 lần.
• Vừa đủ: loss giảm đều, không NaN; sau vài epoch validation cũng giảm.

Kỹ thuật thực dụng:

• LR finder (Smith 2015): train một epoch ngắn với LR tăng dần theo log scale, plot loss vs LR, chọn LR ngay trước điểm loss bắt đầu tăng. Có sẵn trong fastai.
• LR schedule: giảm LR theo thời gian — StepLR, CosineAnnealing, ReduceLROnPlateau. Cosine + warmup là chuẩn cho Transformer.
• Warmup: vài trăm step đầu tăng LR từ 0 lên LR target. Ổn định pretraining Transformer.

LR schedule là chủ đề lớn — sẽ đi sâu trong Bài 18 (Training Loop) và bài dành cho LR schedule.

Weight decay là regularization phạt \( \|\theta\|^2 \) của parameter, kéo parameter về gần 0. Đây là dạng L2 regularization (xem Series 2 — Regularization).

optimizer = optim.AdamW(model.parameters(),
                        lr=1e-3, weight_decay=0.01)

Giá trị thường dùng:

• \( 10^{-4} \): mặc định cho CNN classification.
• \( 10^{-2} \): default AdamW trong HuggingFace cho Transformer fine-tuning.
• \( 0.1 \) trở lên: pretraining LLM lớn (LLaMA dùng 0.1 cho một số layer).

Lưu ý quan trọng: không nên apply weight decay cho bias và normalization layer (LayerNorm, BatchNorm). Lý do: bias không có gradient overfit (mỗi neuron có 1 bias), normalization param là scale/shift của activation — decay về 0 sẽ làm "tắt" layer đó. Recipe chuẩn:

decay_params, no_decay_params = [], []
for name, p in model.named_parameters():
    if p.dim() < 2 or "bias" in name or "norm" in name.lower():
        no_decay_params.append(p)
    else:
        decay_params.append(p)

optimizer = optim.AdamW(
    [{"params": decay_params, "weight_decay": 0.01},
     {"params": no_decay_params, "weight_decay": 0.0}],
    lr=1e-3,
)

Đây là pattern param groups — mục tiếp theo.

Optimizer PyTorch chấp nhận một list dict thay vì iterator parameter. Mỗi dict có params riêng và override hyperparameter chung:

optimizer = optim.Adam([
    {"params": model.encoder.parameters(), "lr": 1e-5},
    {"params": model.head.parameters(),    "lr": 1e-3},
], betas=(0.9, 0.999), weight_decay=0.0)

Use case:

• Transfer learning: backbone pretrained giữ LR nhỏ (\( 10^{-5} \)) để không phá feature đã học; head mới khởi tạo random dùng LR lớn (\( 10^{-3} \)) để fit nhanh.
• Discriminative fine-tuning (Howard & Ruder, ULMFiT 2018): LR tăng dần từ layer thấp lên layer cao.
• Tách weight decay như mục 15: bias và norm dùng weight_decay=0, ma trận weight dùng weight_decay=0.01.
• Freeze layer bằng cách set lr=0 hoặc đơn giản hơn là p.requires_grad = False.

Khi truy cập optimizer.param_groups sau khi khởi tạo, bạn thấy list dict với toàn bộ hyperparameter (kể cả default). Có thể thay đổi LR runtime: optimizer.param_groups[0]["lr"] = new_lr — đây là cách LRScheduler hoạt động bên trong.

Adam và SGD-momentum có state nội tại (\( m_t, v_t, v \) momentum buffer). Khi checkpoint training để resume, lưu cả model state lẫn optimizer state:

# Save
torch.save({
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "loss": loss.item(),
}, "checkpoint.pt")

# Load
ckpt = torch.load("checkpoint.pt")
model.load_state_dict(ckpt["model_state_dict"])
optimizer.load_state_dict(ckpt["optimizer_state_dict"])
start_epoch = ckpt["epoch"] + 1

Vì sao phải load cả optimizer state:

• Nếu chỉ load model, momentum buffer / \( m_t, v_t \) reset về 0. Adam mất "đà" và bias correction lại từ \( t = 0 \) — training resume không khớp với "tiếp tục" liền mạch.
• Có thể quan sát bằng cách so sánh loss curve: resume không load optimizer state thường có một "vết nẩy" loss ở điểm resume.

Khi load model lên device khác (CPU → GPU hoặc ngược lại), phải đảm bảo state tensor cùng device. torch.load(map_location=...) giải quyết phần model; optimizer state cũng cần move qua device tương ứng — nhiều khi cách an toàn nhất là khởi tạo optimizer sau khi model.to(device), rồi load optimizer state.

SGD và Adam phủ phần lớn use case, nhưng có vài hướng mới đáng biết khi train model lớn:

• Lion (Chen et al., Google 2023, arXiv:2302.06675): chỉ dùng dấu của gradient, momentum buffer duy nhất. Tiết kiệm 1/2 RAM optimizer so với Adam (chỉ 1 state thay vì 2). Báo cáo accuracy comparable hoặc nhỉnh hơn AdamW trên ViT và một số LLM. Hyperparameter mặc định: lr ≈ 1e-4 (nhỏ hơn Adam), weight_decay ≈ 1.0.
• Sophia (Liu et al., Stanford 2023, arXiv:2305.14342): xấp xỉ Hessian rẻ qua sampling, dùng diagonal Hessian để precondition gradient. Báo cáo tăng tốc 2x trên GPT-2 scale so với AdamW. Cộng đồng vẫn đang verify ở scale lớn.
• 8-bit Adam (Dettmers et al. 2021, bitsandbytes): lượng tử state \( m_t, v_t \) xuống 8-bit, giữ master copy 32-bit cho weight. Tiết kiệm khoảng 75% RAM optimizer so với Adam fp32 — quan trọng khi train model 7B+ trên GPU consumer. Dùng qua bitsandbytes.optim.AdamW8bit.

Quan điểm thực dụng: với 95% project, AdamW (hoặc SGD-momentum cho CNN) là đủ. Lion / Sophia / 8-bit chỉ cân nhắc khi đã có baseline AdamW tốt và muốn squeeze thêm 10–20% RAM hoặc tốc độ — không phải optimization đầu tiên cần thử.

Train MLP 1 hidden layer trên iris, so sánh SGD và Adam:

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

torch.manual_seed(0)
X, y = load_iris(return_X_y=True)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=0)

scaler = StandardScaler().fit(X_tr)
X_tr = torch.tensor(scaler.transform(X_tr), dtype=torch.float32)
X_te = torch.tensor(scaler.transform(X_te), dtype=torch.float32)
y_tr = torch.tensor(y_tr, dtype=torch.long)
y_te = torch.tensor(y_te, dtype=torch.long)


def make_model():
    torch.manual_seed(0)                       # fix init cho fair compare
    return nn.Sequential(
        nn.Linear(4, 16),
        nn.ReLU(),
        nn.Linear(16, 3),
    )


def train(optimizer_cls, lr, epochs=200, **kwargs):
    model = make_model()
    criterion = nn.CrossEntropyLoss()
    optimizer = optimizer_cls(model.parameters(), lr=lr, **kwargs)
    losses = []
    for ep in range(epochs):
        optimizer.zero_grad()
        out = model(X_tr)
        loss = criterion(out, y_tr)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    with torch.no_grad():
        acc = (model(X_te).argmax(1) == y_te).float().mean().item()
    return losses, acc


sgd_losses, sgd_acc = train(optim.SGD, lr=0.01, momentum=0.9)
adam_losses, adam_acc = train(optim.Adam, lr=1e-3)

print(f"SGD+mom  acc = {sgd_acc:.4f}, final loss = {sgd_losses[-1]:.4f}")
print(f"Adam     acc = {adam_acc:.4f}, final loss = {adam_losses[-1]:.4f}")

Kết quả điển hình trên iris: cả hai đều đạt 95%+ accuracy. Adam giảm loss nhanh hơn vài chục epoch đầu; SGD-momentum bắt kịp ở cuối. Trên bài toán nhỏ và "dễ" như iris, khác biệt giữa optimizer là rất ít.

Demo hệ quả khi quên zero_grad():

model = make_model()
opt = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

for ep in range(20):
    # KHÔNG zero_grad
    out = model(X_tr)
    loss = criterion(out, y_tr)
    loss.backward()
    opt.step()
    if ep < 5 or ep % 5 == 0:
        print(f"ep {ep:2d}  loss = {loss.item():.4f}")

Loss thường tăng / oscillate / nổ NaN sau vài epoch — gradient cộng dồn làm bước update phình lên.

Param groups — backbone LR thấp, head LR cao:

model = nn.Sequential(
    nn.Linear(4, 16),   # "backbone"
    nn.ReLU(),
    nn.Linear(16, 3),   # "head"
)
backbone = list(model[:2].parameters())
head = list(model[2:].parameters())

opt = optim.Adam([
    {"params": backbone, "lr": 1e-5},
    {"params": head,     "lr": 1e-3},
])
print(opt.param_groups[0]["lr"], opt.param_groups[1]["lr"])
# 1e-05 0.001

1. Train MLP iris (kiến trúc như mục 19) với optim.SGD(lr=0.01) không momentum và optim.Adam(lr=1e-3). Chạy 200 epoch, ghi lại loss curve và test accuracy. Optimizer nào hội tụ nhanh hơn?
2. Thêm momentum=0.9 vào SGD. So sánh loss curve với SGD vanilla ở câu 1. Quan sát: momentum giảm dao động hay tăng tốc nhiều hơn?
3. Thay Adam bằng AdamW(lr=1e-3, weight_decay=1e-4). Chạy 200 epoch. So sánh weight norm sum(p.norm()**2 for p in model.parameters()) cuối training giữa Adam và AdamW.
4. Implement SGD bằng tay (không dùng optim.SGD): với mỗi p trong model.parameters(), viết with torch.no_grad(): p -= lr * p.grad sau loss.backward(). So sánh kết quả với optim.SGD(lr=...).
5. Implement SGD-momentum bằng tay: lưu buffer v cho từng parameter, update theo công thức mục 5. Verify khớp optim.SGD(momentum=0.9).
6. Demo "quên zero_grad" như mục 19. In ra p.grad.norm() mỗi epoch để thấy nó tăng vô hạn.
7. Tách parameter thành 2 group: bias (p.dim() < 2) weight_decay=0, weight matrix weight_decay=0.01. Khởi tạo AdamW với 2 group. Verify optimizer.param_groups có đúng giá trị.
8. Save model + optimizer state vào file "ckpt.pt" sau epoch 100. Khởi tạo model + optimizer mới, load state, train tiếp 100 epoch nữa. So sánh loss curve với train liền mạch 200 epoch — phải khớp gần như tuyệt đối.
9. Thử LR sai: optim.Adam(lr=1.0) và optim.Adam(lr=1e-8). Quan sát triệu chứng (NaN vs gần như không học) sau vài chục epoch.

• Optimizer wrap update rule + state. Update rule biến \( (\theta, g) \) thành \( \theta \) mới.
• SGD vanilla: \( \theta \leftarrow \theta - \eta g \). Đơn giản, ít dùng độc lập.
• SGD + momentum: \( v_t = \mu v_{t-1} + g_t \), \( \theta \leftarrow \theta - \eta v_t \). Default \( \mu = 0.9 \). Giảm dao động, tăng tốc.
• Nesterov: variant với "look-ahead", bật bằng nesterov=True.
• AdaGrad: LR adaptive per parameter qua tổng \( g^2 \). Vấn đề: LR \( \to 0 \).
• RMSProp: thay tổng bằng EMA của \( g^2 \), fix AdaGrad.
• Adam: momentum (\( m_t \)) + RMSProp (\( v_t \)) + bias correction. Default \( \beta_1 = 0.9, \beta_2 = 0.999, \epsilon = 10^{-8}, \eta = 10^{-3} \).
• AdamW: weight decay decoupled khỏi gradient. Default cho Transformer modern.
• Train step: zero_grad → backward → step. Quên zero_grad → gradient cộng dồn → loss nổ.
• Chọn optimizer: SGD-momentum cho CNN/ImageNet; Adam/AdamW cho NLP, Transformer, đa số trường hợp.
• LR là hyperparameter nhạy nhất. SGD: 0.01–0.1. Adam: \( 10^{-4} \)–\( 10^{-3} \). LR schedule + warmup là chuẩn cho Transformer.
• Weight decay 0.01 là default AdamW; không apply cho bias và norm layer.
• Param groups cho LR / weight_decay khác nhau theo layer — phục vụ transfer learning và recipe decay-no-decay.
• Save cả optimizer.state_dict() khi checkpoint để resume training đúng.
• Lion, Sophia, 8-bit Adam: hướng tối ưu mới cho LLM, cân nhắc khi đã có baseline AdamW.