Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions
Not on the map. Why?
Abstract (arXiv)
Existing theory of momentum assumes that gradients arrive at every parameter at a roughly constant rate, an assumption violated in practice by heavy-tailed data distributions and modern architectures. We theoretically analyze the dynamics of two tractable models of momentum under sparse updates: a least squares model with sparse inputs and a logistic regression model with a rare class. Both admit exact closed-form second-moment dynamics whose high-dimensional limits we characterize across three scaling exponents for sparsity, batch size, and momentum decay. The phase structure on both problems is governed by the ratio of two intrinsic timescales: a momentum retention timescale (how many active updates the buffer survives) and a learning timescale (how many active updates it takes to reduce the squared error). When learning is much slower than retention, the limit matches SGD; when learning is faster, the system is unstable; where the timescales coincide, we recover classical heavy-ball dynamics. The oscillatory dynamics occur at different momentum values for different token sparsity, creating a spectral conflict for global momentum across token frequencies.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 45 of its 54 references that are in the map. First posted on arXiv 27 May 2026.
Builds on
45 of 54 references.
References · 45
- 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,509
- 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,073
- 2023Efficient Streaming Language Models with Attention Sinkscited by 2,186
- 2020GLU Variants Improve Transformercited by 1,804
- 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,631
- 2019Generating Long Sequences with Sparse Transformerscited by 1,546
- 2023Sparse Autoencoders Find Highly Interpretable Features in Language Modelscited by 1,420
- 2022Toy Models of Superpositioncited by 971
- 2021GLaM: Efficient Scaling of Language Models with Mixture-of-Expertscited by 897
- 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 867
- 2019Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedcited by 856
- 2016Language Modeling with Gated Convolutional Networkscited by 702
33 more
- 2022Mixture-of-Experts with Expert Choice Routingcited by 623
- 2025Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attentioncited by 407
- 2022ST-MoE: Designing Stable and Transferable Sparse Expert Modelscited by 399
- 2024Massive Activations in Large Language Modelscited by 294
- 2025Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Freecited by 283
- 2023Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Timecited by 252
- 2023PowerInfer: Fast Large Language Model Serving with a Consumer-grade GPUcited by 217
- 2021MoEfication: Transformer Feed-forward Layers are Mixtures of Expertscited by 152
- 2022The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformerscited by 122
- 2019Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Modelcited by 117
- 2016A Variational Perspective on Accelerated Methods in Optimizationcited by 110
- 2024Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Modelscited by 80
- 2024How Does Critical Batch Size Scale in Pre-training?cited by 73
- 2018Understanding the Acceleration Phenomenon via High-Resolution Differential Equationscited by 69
- 2018On the insufficiency of existing momentum schemes for Stochastic Optimizationcited by 49
- 2025Muon Outperforms Adam in Tail-End Associative Memory Learningcited by 46
- 2019Understanding the Role of Momentum in Stochastic Gradient Methodscited by 44
- 2016Stochastic Heavy Ballcited by 42
- 2025Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wastefulcited by 41
- 2021SGD in the Large: Average-case Analysis, Asymptotics, and Stepsize Criticalitycited by 35
- 2025Scaling Laws for Gradient Descent and Sign Descent for Linear Bigram Models under Zipf's Lawcited by 24
- 2025Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Trainingcited by 23
- 2022Trajectory of Mini-Batch Momentum: Batch Size Saturation and Convergence in High Dimensionscited by 19
- 2023When and Why Momentum Accelerates SGD:An Empirical Studycited by 19
- 2023The Marginal Value of Momentum for Small Learning Rate SGDcited by 18
- 2021Dynamics of Stochastic Momentum Methods on Large-scale, Quadratic Modelscited by 16
- 2025Mixture of Sparse Attention: Content-Based Learnable Sparse Attention via Expert-Choice Routingcited by 14
- 2022On the fast convergence of minibatch heavy ball momentumcited by 13
- 2025Dimension-adapted Momentum Outscales SGDcited by 12
- 2026Logarithmic-time Schedules for Scaling Language Models with Momentumcited by 9
- 2020Almost sure convergence rates for Stochastic Gradient Descent and Stochastic Heavy Ballcited by 7
- 2026Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Lawscited by 6
- 2025Class Imbalance in Anomaly Detection: Learning from an Exactly Solvable Modelcited by 1
Built on by
Not cited yet.