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

arXiv: https://arxiv.org/abs/2502.19002 · HTML: https://mlsanity.com/paper/2502.19002/

## Key facts

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

## Builds on (most cited first, 30 of 47)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [TinyLlama: An Open-Source Small Language Model](https://mlsanity.com/paper/2401.02385/) (2024)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Symbolic Discovery of Optimization Algorithms](https://mlsanity.com/paper/2302.06675/) (2023)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models](https://mlsanity.com/paper/2208.06677/) (2022)
- [The Break-Even Point on Optimization Trajectories of Deep Neural Networks](https://mlsanity.com/paper/2002.09572/) (2020)
- [Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis](https://mlsanity.com/paper/1806.03884/) (2018)
- [A Kronecker-factored approximate Fisher matrix for convolution layers](https://mlsanity.com/paper/1602.01407/) (2016)

## Built on by (most cited first, 30 of 30)

- [Completed Hyperparameter Transfer across Modules, Width, Depth, Batch and Duration](https://mlsanity.com/paper/2512.22382/) (2025)
- [Towards Quantifying the Hessian Structure of Neural Networks](https://mlsanity.com/paper/2505.02809/) (2025)
- [Mano: Restriking Manifold Optimization for LLM Training](https://mlsanity.com/paper/2601.23000/) (2026)
- [HTMuon: Improving Muon via Heavy-Tailed Spectral Correction](https://mlsanity.com/paper/2603.10067/) (2026)
- [Learning Rate Matters: Vanilla LoRA May Suffice for LLM Fine-tuning](https://mlsanity.com/paper/2602.04998/) (2026)
- [GradPower: Powering Gradients for Faster Language Model Pre-Training](https://mlsanity.com/paper/2505.24275/) (2025)
- [Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement](https://mlsanity.com/paper/2602.22681/) (2026)
- [Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws](https://mlsanity.com/paper/2602.14208/) (2026)
- [Why Muon Outperforms Adam: A Curvature Perspective](https://mlsanity.com/paper/2606.04662/) (2026)
- [Weight Decay Improves Language Model Plasticity](https://mlsanity.com/paper/2602.11137/) (2026)
- [AMUSE: Anytime Muon with Stable Gradient Evaluation](https://mlsanity.com/paper/2605.22432/) (2026)
- [Negligible in Size, Significant in Effect: On Scale Vectors in Large Language Models](https://mlsanity.com/paper/2605.26895/) (2026)
- [Noise-Adaptive Layerwise Learning Rates: Accelerating Geometry-Aware Optimization for Deep Neural Network Training](https://mlsanity.com/paper/2510.14009/) (2025)
- [BSFA: Leveraging the Subspace Dichotomy to Accelerate Neural Network Training](https://mlsanity.com/paper/2510.25244/) (2025)
- [Controlling changes to attention logits](https://mlsanity.com/paper/2511.21377/) (2025)
- [One LR Doesn't Fit All: Heavy-Tail Guided Layerwise Learning Rates for LLMs](https://mlsanity.com/paper/2605.22297/) (2026)
- [Effective Learning Rate Governs Loss Dynamics in Language Model Pretraining](https://mlsanity.com/paper/2608.24814/) (2026)
- [SwiftVLM: Efficient Vision-Language Model Inference via Cross-Layer Token Bypass](https://mlsanity.com/paper/2602.03134/) (2026)
- [Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in Language Model Pre-Training](https://mlsanity.com/paper/2605.26484/) (2026)
- [On the Nonlinearity of Learning Rate Scaling for LLM Training](https://mlsanity.com/paper/2606.29158/) (2026)
- [FOAM: Blocked State Folding for Memory-Efficient LLM Training](https://mlsanity.com/paper/2512.07112/) (2025)
- [Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks](https://mlsanity.com/paper/2606.00340/) (2026)
- [Revealing Modular Gradient Noise Imbalance in LLMs: Calibrating Adam via Signal-to-Noise Ratio](https://mlsanity.com/paper/2605.05794/) (2026)
- [Rethinking Bregman Divergences in Kronecker-Factored Optimizers](https://mlsanity.com/paper/2606.00542/) (2026)
- [Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining](https://mlsanity.com/paper/2608.28442/) (2026)
- [Momentum in large-batch training: Polyak enlarges the critical batch size, Nesterov improves data efficiency](https://mlsanity.com/paper/2609.02728/) (2026)
- [SOLAR: A State-Driven Online Learning Rate Scheduler for LLM Pretraining](https://mlsanity.com/paper/2609.34681/) (2026)
- [Normalize-Then-Precondition: A Hierarchical Approach to Marginal Scale and Interaction Geometry for LLM Training](https://mlsanity.com/paper/2609.36692/) (2026)
- [How Does Local Landscape Geometry Evolve in Language Model Pre-Training?](https://mlsanity.com/paper/2609.39767/) (2026)
- [Which Geometry on Which Layer? Measured Sign-Spectral Assignment in Mixed-Optimizer Training](https://mlsanity.com/paper/2610.17285/) (2026)

## 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.