Research map / Architectures and efficiency / Attention and transformers
The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
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.
Shows transformer blocks (embedding, QK, VO, FFN, LayerNorm) have a consistent sharpness ordering and proposes Blockwise LR, which rescales per-block learning rates to speed up AdamW and Adam-mini LLM pre-training by about 2x.
Problem
Transformer building blocks (embedding, normalization, attention QK/VO, FFN) differ in curvature, but existing layerwise learning-rate methods do not transfer to transformers; the paper asks how to exploit block-type sharpness disparity to accelerate LLM pre-training.
Approach
Empirically and theoretically establish a block-type sharpness ordering S(Embed) << S(QK) < S(FFN) < S(VO) << S(Norm) that emerges early and persists, then keep the LayerNorm LR at the base value and scale LRs of flatter block types up by tuned ratios, switching after warmup.
Key results
- Blockwise LR applied to AdamW reaches lower terminal loss and nearly 2x speedup over tuned AdamW for GPT-2 and LLaMA (0.12B to 2B) on OpenWebText, MiniPile, and C4.
- On LLaMA (1.1B) OpenWebText, Blockwise LR gives a 1.92x speedup; zero-shot downstream accuracy improves on all seven tasks (e.g., ARC_C 22.87 to 25.34, SCIQ 77.60 to 81.60).
- Adam-mini with Blockwise LR achieves about 2x speedup and 2x memory saving versus AdamW; Lion with Blockwise LR also yields about 2x speedup.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Blockwise LR (block-type LR scaling, LN kept at base LR) added to AdamW | AdamW pre-training of LLaMA/GPT-2 · OpenWebText, MiniPile, C4 pre-training (0.12B-2B) | terminal loss and training speed: nearly 2x speedup; lower terminal loss | Helped |
| Blockwise LR added to AdamW | LLaMA (1.1B) downstream accuracy · 0-shot lm-evaluation-harness, 7 tasks | accuracy: higher on all 7 tasks (e.g., ARC_E +1.6, WinoGrande +4.7) | Helped |
| Blockwise LR added to Adam-mini | Adam-mini · LLaMA pre-training | training speed: 2x speedup; combined 2x speedup and 2x memory saving vs AdamW | Helped |
| Blockwise LR added to Lion | Lion · LLaMA (0.25B) on OpenWebText | training speed: 2x speedup | Helped |
| Blockwise LR applied only to Embed blocks (ablation) | AdamW, LLaMA (0.25B) · OpenWebText 50k steps | terminal loss: -0.016 | Helped |
| Blockwise LR applied to Embed and FFN (ablation) | AdamW, LLaMA (0.25B) · OpenWebText 50k steps | terminal loss: -0.043 | Helped |
| Blockwise LR applied to Embed, FFN, QK, and VO (all non-norm blocks) (ablation) | AdamW, LLaMA (0.25B) · OpenWebText 50k steps | terminal loss: -0.050 | Helped |
| Doubling LR of LN blocks only (ablation) | AdamW, LLaMA (0.25B) · OpenWebText 50k steps | terminal loss: +0.003 | Hurt |
| Blockwise LR applied to FFN blocks alone vs Embed alone (ablation) | AdamW, LLaMA (0.25B) · OpenWebText 50k steps | terminal loss: FFN gives largest gain (~0.027 over Embed-only as stated) | Helped |
| Switching blockwise LR on at end of warmup vs at start | AdamW Blockwise LR · training | terminal loss | No clear effect |
| Sharpness disparity principle among block types (Embed < QK < FFN < VO << Norm) | GPT-2 and LLaMA pre-training · OpenWebText, sharpness estimates | average block sharpness ordering | Helped |
| Layer-wise learning rate adjustment based on layer-level sharpness | Transformers · LLaMA (0.25B) OpenWebText | layer-level sharpness trend | No clear effect |
Surprising findings
- Doubling the LayerNorm learning rate worsens terminal loss (+0.003), whereas increasing other block LRs helps.
- The sharpness disparity does not show a consistent pattern across layers, only across block types.
Limitations
- Adjusting ratios are manually tuned (four hyperparameters) rather than derived from the sharpness ratios; Adam-mini ratios do not follow the sharpness rule.
- Theoretical results are bounds with problem-dependent constants and are restricted to single-head attention and within-layer comparisons; the QK vs VO ordering is deferred to prior work.
- Experiments are limited to decoder-only LLMs up to 2B parameters and do not cover non-LLM tasks.
Experimental setup
Datasets: OpenWebText MiniPile C4
Benchmarks: ARC_E ARC_C PIQA HellaSwag OBQA WinoGrande SCIQ
Metrics: terminal validation loss speedup zero-shot accuracy
Baselines: AdamW Adam-mini Lion
Builds on: AdamW Adam-mini Sophia-style diagonal Hessian estimation edge of stability WSD scheduler
Closest papers by meaning
- Eureka-Moments in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problems (ICML 2024 · similarity 0.88)
- The Curse of Depth in Large Language Models (NeurIPS 2025 · similarity 0.88)
- Transformer Block Coupling and its Correlation with Generalization in LLMs (ICLR 2025 · similarity 0.88)
- Preparing Lessons for Progressive Training on Language Models (AAAI 2024 · similarity 0.88)
- HybridNorm: Towards Stable and Efficient Transformer Training via Hybrid Normalization (NeurIPS 2025 · similarity 0.87)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.87)
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.