mlsanityData to 14 Oct 2026

Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformers

arXiv 2411.00999 ↗First posted 1 Nov 2024cs.LG stat.MLcited by 7

Not on the map. Why?

Abstract (arXiv)

Per-example gradient norms are a vital ingredient for estimating gradient noise scale (GNS) with minimal variance. Observing the tensor contractions required to compute them, we propose a method with minimal FLOPs in 3D or greater tensor regimes by simultaneously computing the norms while computing the parameter gradients. Using this method we are able to observe the GNS of different layers at higher accuracy than previously possible. We find that the total GNS of contemporary transformer models is predicted well by the GNS of only the normalization layers. As a result, focusing only on the normalization layer, we develop a custom kernel to compute the per-example gradient norms while performing the LayerNorm backward pass with zero throughput overhead. Tracking GNS on only those layers, we are able to guide a practical batch size schedule that reduces training time by 18% on a Chinchilla-optimal language model.

Key facts. Cited by 7 papers in the map (data through 14 Oct 2026). Builds on 44 of its 66 references that are in the map. First posted on arXiv 1 Nov 2024.

Builds on

44 of 66 references.

References · 44

  1. 2014Adam: A Method for Stochastic Optimizationcited by 37,240
  2. 2020Language Models are Few-Shot Learnerscited by 32,419
  3. 2015Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shiftcited by 7,005
  4. 2016Layer Normalizationcited by 5,669
  5. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
  6. 2023QLoRA: Efficient Finetuning of Quantized LLMscited by 3,418
  7. 2020The Pile: An 800GB Dataset of Diverse Text for Language Modelingcited by 2,411
  8. 2018Spectral Normalization for Generative Adversarial Networkscited by 1,636
  9. 2019Root Mean Square Layer Normalizationcited by 1,480
  10. 2019Energy and Policy Considerations for Deep Learning in NLPcited by 1,254
  11. 2016Optimization Methods for Large-Scale Machine Learningcited by 1,150
  12. 2019On the Convergence of Adam and Beyondcited by 841
32 more
  1. 2015Dueling Network Architectures for Deep Reinforcement Learningcited by 785
  2. 2022GPT-NeoX-20B: An Open-Source Autoregressive Language Modelcited by 713
  3. 2021Carbon Emissions and Large Neural Network Trainingcited by 567
  4. 2021Large Language Models Can Be Strong Differentially Private Learnerscited by 412
  5. 2023Analyzing and Improving the Training Dynamics of Diffusion Modelscited by 402
  6. 2021Opacus: User-Friendly Differential Privacy Library in PyTorchcited by 383
  7. 2023Fine-Tuning Language Models with Just Forward Passescited by 345
  8. 2018An Empirical Model of Large-Batch Trainingcited by 281
  9. 2019A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networkscited by 214
  10. 2018Measuring the Effects of Data Parallelism on Neural Network Trainingcited by 193
  11. 2023Small-scale proxies for large-scale Transformer training instabilitiescited by 183
  12. 2023Stabilizing Transformer Training by Preventing Attention Entropy Collapsecited by 161
  13. 2022High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representationcited by 150
  14. 2017A Bayesian Perspective on Generalization and Stochastic Gradient Descentcited by 135
  15. 2024Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformerscited by 118
  16. 2019Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Modelcited by 117
  17. 2022Adam Can Converge Without Any Modification On Update Rulescited by 114
  18. 2023Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Becited by 114
  19. 2023Cerebras-GPT: Open Compute-Optimal Language Models Trained on the Cerebras Wafer-Scale Clustercited by 109
  20. 2020Estimating Example Difficulty Using Variance of Gradientscited by 97
  21. 2021How to Train BERT with an Academic Budgetcited by 80
  22. 2024Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Modelscited by 80
  23. 2019BackPACK: Packing more into backpropcited by 78
  24. 2019First Exit Time Analysis of Stochastic Gradient Descent Under Heavy-Tailed Gradient Noisecited by 44
  25. 2015Efficient Per-Example Gradient Computationscited by 42
  26. 2021Batch size-invariance for policy optimizationcited by 35
  27. 2016Coupling Adaptive Batch Sizes with Learning Ratescited by 30
  28. 2020A Study of Gradient Variance in Deep Learningcited by 28
  29. 2023CLIPA-v2: Scaling CLIP Training with 81.1% Zero-shot ImageNet Accuracy within a 10,000 Budget; An Extra 4,000 Unlocks 81.8% Accuracycited by 24
  30. 2019Efficient Per-Example Gradient Computations in Convolutional Neural Networkscited by 17
  31. 2016Big Batch SGD: Automated Inference using Adaptive Batch Sizescited by 13
  32. 2018The Effect of Network Width on the Performance of Large-batch Trainingcited by 5

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 3; 2026: 4.

7 papers.

  1. 2025Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMscited by 35
  2. 2025Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Trainingcited by 23
  3. 2026Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Lawscited by 6
  4. 2025Per-example gradients: a new frontier for understanding and improving optimizerscited by 2
  5. 2026Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descentcited by 2
  6. 2026Adaptive Forgetting for Nonstationary Optimization: Towards Robust EEG Decodingcited by 0
  7. 2026How Does Local Landscape Geometry Evolve in Language Model Pre-Training?cited by 0