# Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions

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

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

- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Sparse Autoencoders Find Highly Interpretable Features in Language Models](https://mlsanity.com/paper/2309.08600/) (2023)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models](https://mlsanity.com/paper/2401.06066/) (2024)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Language Modeling with Gated Convolutional Networks](https://mlsanity.com/paper/1612.08083/) (2016)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [Massive Activations in Large Language Models](https://mlsanity.com/paper/2402.17762/) (2024)
- [Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free](https://mlsanity.com/paper/2505.06708/) (2025)
- [Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time](https://mlsanity.com/paper/2310.17157/) (2023)
- [PowerInfer: Fast Large Language Model Serving with a Consumer-grade GPU](https://mlsanity.com/paper/2312.12456/) (2023)
- [MoEfication: Transformer Feed-forward Layers are Mixtures of Experts](https://mlsanity.com/paper/2110.01786/) (2021)
- [The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers](https://mlsanity.com/paper/2210.06313/) (2022)
- [Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model](https://mlsanity.com/paper/1907.04164/) (2019)
- [A Variational Perspective on Accelerated Methods in Optimization](https://mlsanity.com/paper/1603.04245/) (2016)
- [Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models](https://mlsanity.com/paper/2402.19449/) (2024)
- [How Does Critical Batch Size Scale in Pre-training?](https://mlsanity.com/paper/2410.21676/) (2024)
- [Understanding the Acceleration Phenomenon via High-Resolution Differential Equations](https://mlsanity.com/paper/1810.08907/) (2018)
- [On the insufficiency of existing momentum schemes for Stochastic Optimization](https://mlsanity.com/paper/1803.05591/) (2018)
- [Muon Outperforms Adam in Tail-End Associative Memory Learning](https://mlsanity.com/paper/2509.26030/) (2025)
- [Understanding the Role of Momentum in Stochastic Gradient Methods](https://mlsanity.com/paper/1910.13962/) (2019)
- [Stochastic Heavy Ball](https://mlsanity.com/paper/1609.04228/) (2016)

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