Per-example gradients: a new frontier for understanding and improving optimizers
Not on the map. Why?
Abstract (arXiv)
Training algorithms in deep learning usually treat a mini-batch of samples as a single object; they average gradients over the mini-batch, and then process the average in various ways. Computing other statistics beyond the average may have been seen as prohibitively resource intensive in automatic differentiation (AD) frameworks. We show that this is not the case. Generally, gradient statistics can be implemented through a surgery of the AD graph, which, in some cases, incur almost no computational and memory overheads compared to the mini-batch gradient computation. Additionally, we show that in certain classes of models, including transformers, JAX's vectorization transformation offers a viable implementation for prototyping and experimentation. We then revise our understanding of two nonlinear operations in optimization through the lens of per-example gradient transformations. We first study signSGD and show that the optimal placement of the sign operation in the gradient processing chain is crucial to success and can be predicted with a simple signal-to-noise ratio argument. Next we study per-example variations of the Adam preconditioner, and show that optimization is best served when the preconditioner is dominated by the mean rather than the variance of the gradient distribution - in contrast to conventional wisdom. Overall we demonstrate that per-example gradient information enables new analyses and possibilities for algorithm design.
Key facts. Cited by 2 papers in the map (data through 14 Oct 2026). Builds on 32 of its 45 references that are in the map. First posted on arXiv 30 Sep 2025.
Builds on
32 of 45 references.
References · 32
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
- 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,599
- 2022Training Compute-Optimal Large Language Modelscited by 3,019
- 2018signSGD: Compressed Optimisation for Non-Convex Problemscited by 571
- 2021Gradient Descent on Neural Networks Typically Occurs at the Edge of Stabilitycited by 427
- 2018An Empirical Model of Large-Batch Trainingcited by 281
- 2018Measuring the Effects of Data Parallelism on Neural Network Trainingcited by 193
- 2023DiLoCo: Distributed Low-Communication Training of Language Modelscited by 121
- 2023Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Becited by 114
- 2019BackPACK: Packing more into backpropcited by 78
- 2022Differentially Private Optimization on Large Model at Small Costcited by 65
- 2024Understanding Optimization in Deep Learning with Central Flowscited by 64
20 more
- 2022Fast Finite Width Neural Tangent Kernelcited by 62
- 2016Efficient Orthogonal Parametrisation of Recurrent Neural Networks Using Householder Reflectionscited by 57
- 2023Benchmarking Neural Network Training Algorithmscited by 52
- 2022Scalable and Efficient Training of Large Convolutional Neural Networks with Differential Privacycited by 48
- 2024The Elements of Differentiable Programmingcited by 43
- 2015Efficient Per-Example Gradient Computationscited by 42
- 2020The Geometry of Sign Gradient Descentcited by 41
- 2020Understanding Why Neural Networks Generalize Well Through GSNR of Parameterscited by 40
- 2025Accelerating Neural Network Training: An Analysis of the AlgoPerf Competitioncited by 32
- 2020A Study of Gradient Variance in Deep Learningcited by 28
- 2024Adaptive Methods through the Lens of SDEs: Theoretical Insights on the Role of Noisecited by 23
- 2025Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networkscited by 22
- 2024High dimensional analysis reveals conservative sharpening and a stochastic edge of stabilitycited by 18
- 2024DrJAX: Scalable and Differentiable MapReduce Primitives in JAXcited by 9
- 2021Cockpit: A Practical Debugging Tool for the Training of Deep Neural Networkscited by 7
- 2024Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformerscited by 7
- 2024Exact Risk Curves of signSGD in High-Dimensions: Quantifying Preconditioning and Noise-Compression Effectscited by 7
- 2023Adam Accumulation to Reduce Memory Footprints of both Activations and Gradients for Large-scale DNN Trainingcited by 5
- 2024Batch size invariant Adamcited by 5
- 2023Target-independent XLA optimization using Reinforcement Learningcited by 1
Built on by
Citing papers per year: 2025: 1; 2026: 1.
2 papers.