Research map / Architectures and efficiency / Attention and transformers
Taming Transformer Without Using Learning Rate Warmup
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes why Transformer training crashes without LR warmup (spectral energy concentration of W_q^T W_k causing malignant entropy collapse) and proposes AdamW^2, which caps the learning rate via a Weyl-inequality-motivated smooth update rule, matching warmup baselines on ViT, Swin, and GPT.
Problem
Transformers often diverge or crash during training without learning rate warmup or architectural modifications; the underlying cause of this instability is not well understood.
Approach
Identifies spectral energy concentration of W_q^T W_k as the cause of malignant attention entropy collapse, and proposes a smooth weight-updating rule motivated by Weyl's inequality that bounds the learning rate by the ratio sigma_1(W_{t-1})/sigma_1(grad W_t), implemented as AdamW^2 with power-iteration norm estimates.
Key results
- AdamW^2 without warmup reaches 80.58 on ViT-B vs 80.22 for AdamW with warmup on ImageNet.
- AdamW^2 without warmup reaches 83.14 on Swin-S and 83.44 on Swin-B vs 83.02 and 83.48 with warmup.
- AdamW^2 achieves GPT-S validation loss 2.840 vs 2.848 for AdamW with warmup.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use AdamW² without learning-rate warmup instead of AdamW with warmup | ViT-B · ImageNet (ViT-B) | accuracy: 80.22 -> 80.58 | Helped |
| Use AdamW² without learning-rate warmup instead of AdamW with warmup | ViT-L · ImageNet (ViT-L) | accuracy: 81.65 -> 81.82 | Helped |
| Use AdamW² without learning-rate warmup instead of AdamW with warmup | GPT-S · OpenWebText (GPT-S) | loss: 2.848 -> 2.840 | Helped |
| Use AdamW² without learning-rate warmup instead of AdamW with warmup | Swin-S · ImageNet (Swin-S) | accuracy: 83.02 -> 83.14 | Helped |
| Use AdamW² without learning-rate warmup instead of AdamW with warmup | Swin-B · ImageNet (Swin-B) | accuracy: 83.48 -> 83.44 | Hurt |
| Vary the smoothing threshold τ (ablation) | AdamW² · GPT and Swin-Transformer | performance | No clear effect |
Surprising findings
- Benign entropy collapse (attention close to identity) is argued not to crash training, while only the malignant mode does.
Limitations
- Theoretical analysis relies on simplifying assumptions (symmetric PSD W, iid Gaussian inputs).
- Gains over warmup baseline are small and within noise in the reported table.
- Tau hyperparameter requires choice; only ablated on Swin-S and GPT-S.
Experimental setup
Datasets: ImageNet OpenWebText
Benchmarks: ImageNet OpenWebText
Metrics: top-1 accuracy validation loss
Baselines: AdamW (with warmup)
Builds on: AdamW Weyl's inequality Adam
Closest papers by meaning
- Linear attention is (maybe) all you need (to understand Transformer optimization) (ICLR 2024 · similarity 0.88)
- Small-scale proxies for large-scale Transformer training instabilities (ICLR 2024 · similarity 0.88)
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.87)
- Infinite Limits of Multi-head Transformer Dynamics (NeurIPS 2024 · similarity 0.86)
- Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping (NeurIPS 2026 · similarity 0.86)
- Learning Rate Transfer for Hybrid Transformer-SSM Architectures (NeurIPS 2026 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.