mlsanityData to 7 Oct 2026

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training

arXiv 2502.19002 ↗First posted 26 Feb 2025cs.LG cs.AI math.OCcited by 27

Not on the map. Why?

Abstract (arXiv)

Transformers consist of diverse building blocks, such as embedding layers, normalization layers, self-attention mechanisms, and point-wise feedforward networks. Thus, understanding the differences and interactions among these blocks is important. In this paper, we uncover a clear Sharpness Disparity across these blocks, which emerges early in training and intriguingly persists throughout the training process. Motivated by this finding, we propose Blockwise Learning Rate (LR), a strategy that tailors the LR to each block's sharpness, accelerating large language model (LLM) pre-training. By integrating Blockwise LR into AdamW, we consistently achieve lower terminal loss and nearly 2× speedup compared to vanilla AdamW. We demonstrate this acceleration across GPT-2 and LLaMA, with model sizes ranging from 0.12B to 2B and datasets of OpenWebText, MiniPile, and C4. Finally, we incorporate Blockwise LR into Adam-mini (Zhang et al., 2024), a recently proposed memory-efficient variant of Adam, achieving a combined 2× speedup and 2× memory saving. These results underscore the potential of exploiting the sharpness disparity to improve LLM training.

Key facts. Cited by 27 papers in the map (data through 7 Oct 2026). Builds on 47 of its 63 references that are in the map. First posted on arXiv 26 Feb 2025.

Builds on

47 of 63 references.

References · 47

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  4. 2023GPT-4 Technical Reportcited by 18,132
  5. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,491
  6. 2023LLaMA: Open and Efficient Foundation Language Modelscited by 13,958
  7. 2019RoBERTa: A Robustly Optimized BERT Pretraining Approachcited by 13,266
  8. 2016Layer Normalizationcited by 5,579
  9. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
  10. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
  11. 2022Training Compute-Optimal Large Language Modelscited by 2,927
  12. 2019XLNet: Generalized Autoregressive Pretraining for Language Understandingcited by 2,801
35 more
  1. 2020The Pile: An 800GB Dataset of Diverse Text for Language Modelingcited by 2,346
  2. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
  3. 2019Root Mean Square Layer Normalizationcited by 1,425
  4. 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
  5. 2020On Layer Normalization in the Transformer Architecturecited by 906
  6. 2015Optimizing Neural Networks with Kronecker-factored Approximate Curvaturecited by 776
  7. 2024MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategiescited by 694
  8. 2024TinyLlama: An Open-Source Small Language Modelcited by 604
  9. 2024GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projectioncited by 447
  10. 2023Symbolic Discovery of Optimization Algorithmscited by 423
  11. 2021Gradient Descent on Neural Networks Typically Occurs at the Edge of Stabilitycited by 413
  12. 2022Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfercited by 272
  13. 2023Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-trainingcited by 260
  14. 2024SOAP: Improving and Stabilizing Shampoo using Adamcited by 179
  15. 2022Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Modelscited by 172
  16. 2020The Break-Even Point on Optimization Trajectories of Deep Neural Networkscited by 149
  17. 2018Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasiscited by 143
  18. 2016A Kronecker-factored approximate Fisher matrix for convolution layerscited by 138
  19. 2024Adam-mini: Use Fewer Learning Rates To Gain Morecited by 134
  20. 2024Why Transformers Need Adam: A Hessian Perspectivecited by 123
  21. 2022Adaptive Gradient Methods at the Edge of Stabilitycited by 88
  22. 2018Training Tips for the Transformer Modelcited by 85
  23. 2024Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Modelscited by 78
  24. 2024Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspectivecited by 75
  25. 2023The MiniPile Challenge for Data-Efficient Language Modelscited by 74
  26. 2024Scaling Exponents Across Parameterizations and Optimizerscited by 71
  27. 2022Make Sharpness-Aware Minimization Stronger: A Sparsified Perturbation Approachcited by 70
  28. 2024MARS: Unleashing the Power of Variance Reduction for Training Large Modelscited by 59
  29. 2024Understanding Optimization in Deep Learning with Central Flowscited by 57
  30. 2023Saddle-to-Saddle Dynamics in Diagonal Linear Networkscited by 54
  31. 2024Deconstructing What Makes a Good Optimizer for Language Modelscited by 40
  32. 2024Does SGD really happen in tiny subspaces?cited by 35
  33. 2024Understanding the Expressive Power and Mechanisms of Transformer for Sequence Modelingcited by 25
  34. 2024Improving Generalization and Convergence by Enhancing Implicit Regularizationcited by 14
  35. 2024What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysiscited by 9

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 7; 2026: 20.

27 papers.

Sort:

By area:

  1. 2025Completed Hyperparameter Transfer across Modules, Width, Depth, Batch and Durationcited by 22
  2. 2025Towards Quantifying the Hessian Structure of Neural Networkscited by 13
  3. 2026Mano: Restriking Manifold Optimization for LLM Trainingcited by 12
  4. 2026HTMuon: Improving Muon via Heavy-Tailed Spectral Correctioncited by 11
  5. 2026Learning Rate Matters: Vanilla LoRA May Suffice for LLM Fine-tuningcited by 9
  6. 2026Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancementcited by 8
  7. 2025GradPower: Powering Gradients for Faster Language Model Pre-Trainingcited by 6
  8. 2026Weight Decay Improves Language Model Plasticitycited by 4
  9. 2026Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Lawscited by 4
  10. 2026Why Muon Outperforms Adam: A Curvature Perspectivecited by 4
  11. 2025Noise-Adaptive Layerwise Learning Rates: Accelerating Geometry-Aware Optimization for Deep Neural Network Trainingcited by 3
  12. 2025BSFA: Leveraging the Subspace Dichotomy to Accelerate Neural Network Trainingcited by 3
15 more
  1. 2025Controlling changes to attention logitscited by 3
  2. 2026AMUSE: Anytime Muon with Stable Gradient Evaluationcited by 3
  3. 2026Negligible in Size, Significant in Effect: On Scale Vectors in Large Language Modelscited by 3
  4. 2026SwiftVLM: Efficient Vision-Language Model Inference via Cross-Layer Token Bypasscited by 2
  5. 2026One LR Doesn't Fit All: Heavy-Tail Guided Layerwise Learning Rates for LLMscited by 2
  6. 2026Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in Language Model Pre-Trainingcited by 2
  7. 2026On the Nonlinearity of Learning Rate Scaling for LLM Trainingcited by 2
  8. 2026Effective Learning Rate Governs Loss Dynamics in Language Model Pretrainingcited by 2
  9. 2025FOAM: Blocked State Folding for Memory-Efficient LLM Trainingcited by 1
  10. 2026Revealing Modular Gradient Noise Imbalance in LLMs: Calibrating Adam via Signal-to-Noise Ratiocited by 0
  11. 2026Rethinking Bregman Divergences in Kronecker-Factored Optimizerscited by 0
  12. 2026Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretrainingcited by 0
  13. 2026SOLAR: A State-Driven Online Learning Rate Scheduler for LLM Pretrainingcited by 0
  14. 2026Normalize-Then-Precondition: A Hierarchical Approach to Marginal Scale and Interaction Geometry for LLM Trainingcited by 0
  15. 2026How Does Local Landscape Geometry Evolve in Language Model Pre-Training?cited by 0