Learning Rate Scheduling — Warmup rồi decay, tại sao cần cả hai

Training một model như GPT-4 tốn hàng trăm triệu đô, nhưng chỉ cần sai một hyperparameter—learning rate—là toàn bộ investment biến thành vô dụng. Learning rate scheduling (warmup rồi decay) chính là chiếc thermostat giữ cho cỗ máy đắt đỏ này không bị nổ máy lúc khởi động, cũng không chạy tốc độ tối đa rồi bay qua điểm đích mà không thể dừng lại.

Vấn đề

Learning rate cố định (constant LR) thất bại ở hai giai đoạn quan trọng nhất của training LLM quy mô lớn.

Ở giai đoạn đầu (cold start), gradient từ backpropagation là hỗn loạn và cực kỳ lớn do weight khởi tạo ngẫu nhiên. Với learning rate cao (ví dụ 1e-3), model "explode" ngay trong 100-1000 bước đầu—loss trở thành NaN hoặc spike đến mức không recover được. Đây là lý do các run training LLM sớm thường chết yểu trong đêm đầu tiên, đốt cháy hàng nghìn đô tiền GPU mà không có checkpoint hữu dụng nào được lưu lại.

Ở giai đoạn cuối, khi model đã gần converge, learning rate cao khiến nó "nảy quanh" điểm optimal như một quả bóng không thể lăn vào đáy chén—nó cứ đập vào hai vách bên và phản hồi. Điều này lãng phí 20% compute cuối cùng (tương đương hàng triệu đô với model 175B parameters training trên 300B tokens) để tạo ra noise thay vì refinement.

Ý tưởng cốt lõi

Hãy tưởng tượng training như lái xe xuyên lục địa. Bạn không thể nổ máy và đạp ga 100km/h ngay từ chỗ đỗ—bạn cần làm nóng máy (warmup). Và bạn không thể đâm sầm vào đích ở tốc độ tối đa—bạn cần nhả ga dần (decay) để coast vào đúng điểm dừng.

Warmup bắt đầu từ learning rate gần như bằng 0 và tăng linearly trong 2,000–4,000 steps (hoặc ~0.1% tổng số tokens cho LLM lớn). Điều này ngăn "cold start shock"—những gradient nhiễu loạn đầu tiên không thể phá hủy representation mới chớm hình thành.

Decay (thường là cosine annealing hoặc exponential) giảm dần LR sau khi đạt peak. Khi model đi vào các "thung lũng" hẹp hơn trong loss landscape, bước chân cần nhỏ lại để lăn vào điểm sâu nhất thay vì bật ra.

Đó là tất cả. Không cần toán học phức tạp—chỉ cần hiểu gradient descent cần tốc độ khác nhau cho các giai đoạn khác nhau, giống như lái xe trong thành phố vs đường cao tốc vs đỗ xe.

Tại sao nó hoạt động

Warmup bảo vệ các layer đầu. Ở bước 0, attention heads là ngẫu nhiên. Một gradient update độ lớn 0.1 có thể xoay ma trận projection 90 độ, phá hủy mọi pattern đang hình thành. Warmup giới hạn các bước đầu ở scale 1e-6, cho phép model "cảm nhận" phân phối dữ liệu mà không bị co giật.

Cosine decay tuân theo công thức: lr(t) = lr_max × 0.5 × (1 + cos(π × t / T)). Tạo ra đường cong chữ S mượt mà xuống gần 0. Khác với step decay (giảm đột ngột mỗi epoch), cosine annealing co hẹp learning rate liên tục, khớp với việc thung lũng loss ngày càng hẹp khi tiến gần minimum.

Các nghiên cứu quy mô lớn (Chinchilla) chỉ ra peak LR cần scale với kích thước model: model lớn hơn cần peak LR cao hơn một chút nhưng warmup dài hơn để ổn định lượng parameter tăng thêm. Ví dụ, GPT-3 175B dùng peak LR ~6e-4 với 375M tokens warmup (~0.12% tổng data), trong khi model nhỏ hơn có thể dùng 1e-4.

Mã giả PyTorch cho scheduler điển hình:

# Linear warmup rồi cosine decay
if step < warmup_steps:
    lr = peak_lr * (step / warmup_steps)
else:
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    lr = peak_lr * 0.5 * (1 + math.cos(math.pi * progress))

Ý nghĩa thực tế

Với GPT-3 175B, OpenAI dùng 375M tokens warmup. Bỏ qua đoạn này, training diverge ngay. Bỏ qua decay, final perplexity kém hơn 5-10%, tương đương với việc training model nhỏ hơn 20%—mất hàng chục triệu đô tiền compute cho kết quả tương đương model 140B.

Các giới hạn: Warmup quá dài lãng phí compute (bạn đang train ở LR suboptimal). Warmup quá ngắn rồi risk divergence. Curve cosine giả định training length cố định; nếu bất ngờ extend training, bạn cần adjust schedule (restart hoặc kéo dài period) hoặc LR sẽ xuống quá thấp quá sớm.

Ai dùng: Mọi LLM major—GPT-4, Llama 3, Gemini, Claude. Được tích hợp sẵn trong DeepSpeed, Megatron-LM, và PyTorch schedulers.

Đào sâu hơn

• Paper gốc: "Attention Is All You Need" (Vaswani et al., 2017) — giới thiệu LR schedule với warmup 4,000 steps cho Transformer đầu tiên; "Training Compute-Optimal Large Language Models" (Hoffmann et al., 2022 — Chinchilla) — thảo luận scaling của optimal LR với model size.

• Cùng cụm (training-at-scale):

  • Zero Optimizer — Cách tiết kiệm bộ nhớ optimizer state để train model nghìn tỷ parameter
  • Data Parallelism — Chia batch cho nhiều GPU, kết hợp với LR schedule để scaling
  • Mixed Precision Training — Tương tác giữa FP16/BF16 và learning rate (loss scaling)
  • Gradient Checkpointing — Đánh đổi compute lấy memory, ảnh hưởng đến effective batch size và LR
  • Scaling Laws — Quy luật Chinchilla về cách chọn optimal LR theo model size

• Đọc tiếp (Level 1):

  • Training Efficiency — Các kỹ thuật nâng cao để tối ưu throughput và convergence ngoài LR scheduling