Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions

2026-05-27

AI Summary

This work addresses the breakdown of conventional momentum theory, which assumes near-constant gradient update frequencies across all parameters—an assumption invalidated under sparse data and modern architectures. By constructing least-squares and rare-class logistic regression models with sparse inputs, the authors characterize momentum dynamics under sparse updates through high-dimensional asymptotic analysis, closed-form second-moment dynamics, and multi-scale timescale modeling. The study reveals that the ratio between the momentum retention timescale and the learning timescale governs the system’s phase structure and identifies a spectral conflict in global momentum arising from varying token sparsity levels. Furthermore, it delineates a dynamic phase diagram shaped jointly by sparsity, batch size, and momentum decay exponent, recovering classical heavy-ball dynamics precisely within the timescale-matching regime.

Abstract

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.

Problem

• stochastic momentum
• sparse updates
• high dimensions
• timescale conflict
• heavy-tailed data

Innovation

• stochastic momentum
• sparse updates
• high-dimensional dynamics
• timescale analysis
• heavy-ball method