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

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

## 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 (most cited first, 30 of 44)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [Dueling Network Architectures for Deep Reinforcement Learning](https://mlsanity.com/paper/1511.06581/) (2015)
- [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://mlsanity.com/paper/2204.06745/) (2022)
- [Carbon Emissions and Large Neural Network Training](https://mlsanity.com/paper/2104.10350/) (2021)
- [Large Language Models Can Be Strong Differentially Private Learners](https://mlsanity.com/paper/2110.05679/) (2021)
- [Analyzing and Improving the Training Dynamics of Diffusion Models](https://mlsanity.com/paper/2312.02696/) (2023)
- [Opacus: User-Friendly Differential Privacy Library in PyTorch](https://mlsanity.com/paper/2109.12298/) (2021)
- [Fine-Tuning Language Models with Just Forward Passes](https://mlsanity.com/paper/2305.17333/) (2023)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networks](https://mlsanity.com/paper/1901.06053/) (2019)
- [Measuring the Effects of Data Parallelism on Neural Network Training](https://mlsanity.com/paper/1811.03600/) (2018)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [A Bayesian Perspective on Generalization and Stochastic Gradient Descent](https://mlsanity.com/paper/1710.06451/) (2017)
- [Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformers](https://mlsanity.com/paper/2401.11605/) (2024)
- [Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model](https://mlsanity.com/paper/1907.04164/) (2019)
- [Adam Can Converge Without Any Modification On Update Rules](https://mlsanity.com/paper/2208.09632/) (2022)
- [Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Be](https://mlsanity.com/paper/2304.13960/) (2023)

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

- [Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs](https://mlsanity.com/paper/2502.15938/) (2025)
- [Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Training](https://mlsanity.com/paper/2505.23971/) (2025)
- [Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws](https://mlsanity.com/paper/2602.14208/) (2026)
- [Per-example gradients: a new frontier for understanding and improving optimizers](https://mlsanity.com/paper/2510.00236/) (2025)
- [Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent](https://mlsanity.com/paper/2602.03001/) (2026)
- [Adaptive Forgetting for Nonstationary Optimization: Towards Robust EEG Decoding](https://mlsanity.com/paper/2609.24233/) (2026)
- [How Does Local Landscape Geometry Evolve in Language Model Pre-Training?](https://mlsanity.com/paper/2609.39767/) (2026)

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