mlsanityData to 14 Oct 2026

Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions

arXiv 2605.28961 ↗First posted 27 May 2026stat.ML cs.LG math.OCcited by 0

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

  1. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,509
  2. 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,073
  3. 2023Efficient Streaming Language Models with Attention Sinkscited by 2,186
  4. 2020GLU Variants Improve Transformercited by 1,804
  5. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,631
  6. 2019Generating Long Sequences with Sparse Transformerscited by 1,546
  7. 2023Sparse Autoencoders Find Highly Interpretable Features in Language Modelscited by 1,420
  8. 2022Toy Models of Superpositioncited by 971
  9. 2021GLaM: Efficient Scaling of Language Models with Mixture-of-Expertscited by 897
  10. 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 867
  11. 2019Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedcited by 856
  12. 2016Language Modeling with Gated Convolutional Networkscited by 702
33 more
  1. 2022Mixture-of-Experts with Expert Choice Routingcited by 623
  2. 2025Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attentioncited by 407
  3. 2022ST-MoE: Designing Stable and Transferable Sparse Expert Modelscited by 399
  4. 2024Massive Activations in Large Language Modelscited by 294
  5. 2025Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Freecited by 283
  6. 2023Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Timecited by 252
  7. 2023PowerInfer: Fast Large Language Model Serving with a Consumer-grade GPUcited by 217
  8. 2021MoEfication: Transformer Feed-forward Layers are Mixtures of Expertscited by 152
  9. 2022The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformerscited by 122
  10. 2019Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Modelcited by 117
  11. 2016A Variational Perspective on Accelerated Methods in Optimizationcited by 110
  12. 2024Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Modelscited by 80
  13. 2024How Does Critical Batch Size Scale in Pre-training?cited by 73
  14. 2018Understanding the Acceleration Phenomenon via High-Resolution Differential Equationscited by 69
  15. 2018On the insufficiency of existing momentum schemes for Stochastic Optimizationcited by 49
  16. 2025Muon Outperforms Adam in Tail-End Associative Memory Learningcited by 46
  17. 2019Understanding the Role of Momentum in Stochastic Gradient Methodscited by 44
  18. 2016Stochastic Heavy Ballcited by 42
  19. 2025Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wastefulcited by 41
  20. 2021SGD in the Large: Average-case Analysis, Asymptotics, and Stepsize Criticalitycited by 35
  21. 2025Scaling Laws for Gradient Descent and Sign Descent for Linear Bigram Models under Zipf's Lawcited by 24
  22. 2025Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Trainingcited by 23
  23. 2022Trajectory of Mini-Batch Momentum: Batch Size Saturation and Convergence in High Dimensionscited by 19
  24. 2023When and Why Momentum Accelerates SGD:An Empirical Studycited by 19
  25. 2023The Marginal Value of Momentum for Small Learning Rate SGDcited by 18
  26. 2021Dynamics of Stochastic Momentum Methods on Large-scale, Quadratic Modelscited by 16
  27. 2025Mixture of Sparse Attention: Content-Based Learnable Sparse Attention via Expert-Choice Routingcited by 14
  28. 2022On the fast convergence of minibatch heavy ball momentumcited by 13
  29. 2025Dimension-adapted Momentum Outscales SGDcited by 12
  30. 2026Logarithmic-time Schedules for Scaling Language Models with Momentumcited by 9
  31. 2020Almost sure convergence rates for Stochastic Gradient Descent and Stochastic Heavy Ballcited by 7
  32. 2026Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Lawscited by 6
  33. 2025Class Imbalance in Anomaly Detection: Learning from an Exactly Solvable Modelcited by 1

Built on by

Not cited yet.