Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions

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.

• The paper introduces a unified high-dimensional framework that models the interplay of momentum and sparsity via closed-form ODEs in both least squares and logistic regression.
• It reveals distinct phase regimes—SGD-like, heavy-ball oscillatory, and resonance—governed by the ratio of momentum retention to learning timescales with explicit learning rate constraints.
• The study uncovers a spectral conflict from global momentum schedules impacting rare and frequent features, highlighting the need for adaptive, sparsity-aware optimization.

Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions

Overview and Motivation

The paper “Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions” (2605.28961) analyzes the interplay between momentum-based stochastic optimization and heavy-tailed, sparsely updated data distributions, motivated by practical regimes encountered in modern deep networks and LLMs. Unlike classical momentum theory, which assumes homogeneous gradient arrival rates for each parameter, this work incorporates power-law data-induced sparsity (e.g., Zipfian token distributions in language, gated routing in mixtures-of-experts) and systematically studies their consequences for the dynamics and stability of momentum optimizers in high-dimensional settings.

Figure 1: Power-law token frequencies result in gradient updates that are sparse and highly non-uniform across parameters, directly impacting the effective timescales for learning and momentum retention.

Model Framework and Phase Structure

The theoretical framework is constructed around two analytically tractable cases: a least squares (LS) regression model with Bernoulli-gated (i.e., sparsely activated) inputs, and a logistic regression (LR) model distinguishing a rare class from a dominant background. Both are studied in a high-dimensional limit ($d\to\infty$), under generic power-law co-scaling of the effective sparsity $p$, batch size $B$, and momentum decay rate $\varepsilon=1-\beta$:

• $p\sim d^{-\kappa}$: token/class sparsity exponent
• $B\sim d^{\sigma}$: batch exponent
• $\varepsilon\sim d^{-\gamma}$: momentum exponent
• $\eta\sim d^{-\alpha}$: learning-rate exponent (phase-dependent)

The phase structure of the dynamics is governed by the competition and ratio between two intrinsic timescales:

• Momentum retention timescale: inverse rate at which the momentum memory buffer decays due to inactivity on rare features.
• Learning timescale: inverse rate at which squared error (or other loss) decays, determined by the effective arrival of signal-carrying gradients.

This structure is formalized by characterizing the ratio $\Delta=$ (retention timescale) / (learning timescale), with $\Delta\sim 1$ (the “resonance line”) identifying the critical regime where classical heavy-ball oscillatory behavior can be retrieved. Elsewhere, $p$0 or $p$1, the system either reduces to standard SGD-like exponential decay or is unstable.

Figure 3: The phase diagram in the $p$2 plane demarcates sample complexity and dominant ODE phase structure for regimes with differing batch size, sparsity, and momentum co-scaling (left); sample-efficiency heatmap under a fixed budget of nonzero gradients (right).

Main Analytical Results

Unified High-Dimensional Dynamics

Closed-form ODE systems are derived for the evolution of the second moments of parameter error and momentum, tracking risk $p$3, momentum energy, and error-momentum correlation. In the LS case these ODEs are exact due to isotropic Gaussian closure.

Stability and Learning Rate Constraints

A sharp Routh–Hurwitz criterion yields explicit learning-rate scalings required for mean-square stability in each regime. Two distinct constraints emerge:

• Noise-limited regime: $p$4, learning constrained by stochastic gradient noise.
• Correlation-limited regime (above resonance): $p$5, learning rate must decay with increasing momentum retention to maintain stability and avoid oscillatory instability.

Canonical Limiting ODEs

On the basis of these analyses, the large-$p$6 dynamics collapse to two canonical forms (see Theorem statements in the paper):

• Below resonance ($p$7): Scalar exponential contraction as in standard high-dimensional SGD, with a phase-dependent learning clock.
• Above resonance ($p$8): Deterministic 2D (heavy-ball) ODE that captures Nesterov/Polyak oscillatory mechanisms.

In all cases, the precise phase is dictated by the position relative to the resonance line, and the vertical/horizontal structure (batch vs momentum) is not interchangeable. The heavy-ball regime and SGD regime are mutually exclusive except on the critical line.

Figure 2: Empirical convergence of Monte Carlo dynamics and Main ODE risk trajectories to the universal limit ODE as $p$9 increases, across representative dense/sparse regimes. The resonance transition mediates between fast converging SGD-like and heavy-ball ODE-like dynamics.

Heterogeneity and the Spectral Conflict

A core insight produced by the analysis is that oscillatory “acceleration” via classical momentum can only exist simultaneously at one effective sparsity level, since different tokens/parameters, distributed across a power-law, require mutually incompatible momentum schedules to be at resonance. In practical terms: for a large-vocabulary LLM, using a single global momentum $B$0 induces a “spectral conflict”—rare and common tokens cannot both satisfy the resonance condition. Increasing batch size helps concentrate the distribution, but does not fully resolve this conflict, particularly in scenarios where batch or data scales shift during fine-tuning or transfer.

Extension to Logistic Regression: Noise-Floor and Non-Square Losses

Extending the analysis to the logistic regression case reveals two critical departures from least squares:

• The noise in the gradient never fully disappears due to persistent random label-flipping in the rare class—there exists an irreducible loss floor even at the population optimum for $B$1.
• The phase diagram and resonance structure essentially remain, but the noise floor introduces a bias-variance tradeoff: more aggressive (smaller $B$2) momentum attenuates the noise floor, but at the cost of slower convergence.

Figure 5: In logistic regression, a trade-off appears between limiting loss (noise-driven bias) and convergence time along different sparsity-momentum regimes; the loss floor and bias structure are absent in pure least squares.

Numerical Verification

Numerical solutions of the high-dimensional ODEs and direct Monte Carlo experiments with finite-$B$3 models validate the theoretical phase diagrams, convergence times, limiting risk floors, and the tightness of phase boundaries.

Figure 4: At the triple point (intersection of dense-sparse and resonance boundaries), high-dimensional ODEs and limit ODEs converge closely for all risk and momentum observables, sweeping through the control parameter $B$4.

Implications and Future Directions

Theoretical Implications

• Momentum acceleration for rare features is fundamentally incompatible with the same for frequent features when using global hyperparameters. This implies that apparent “optimal” batch plus momentum schedules measured on average might leave a huge fraction of parameters in non-accelerated, slow-converging, or unstable regimes.
• Heavy-ball acceleration is not guaranteed in regimes with high stochasticity or extreme sparsity: momentum needs to be carefully scaled—blindly increasing $B$5 for rare features can destabilize rare-parameter sub-dynamics.

Practical and Architectural Consequences

• Adaptive and sparsity-aware optimization: Future momentum methods might need to assign individual momentum/learning rate clocks to different parameter subgroups, as is possible in embedding layers or MoE setups, or dynamically vary clocks in time.
• LLM scaling: When scaling vocabulary, model size, or batch independently, practitioners should be aware that changing effective sparsity shifts where parameters “sit” in the $B$6 plane, potentially degrading the overall optimization regime.
• Fine-tuning and forgetting: A decrease in batch size after pretraining can drop previously-dense features below the resonance line, exposing them to higher noise and less stable momentum dynamics, plausibly contributing to catastrophic forgetting.

Limitations

The analysis is restricted to isotropic covariance models and cannot fully address the subtleties of ill-conditioning present in real neural network training. The effect of non-trivial spectrum, as well as more complex non-linear optimizers, remains an open domain.

Conclusion

This work introduces a detailed, principled understanding of how data-induced sparsity and momentum hyperparameters interact to structure the high-dimensional dynamics of optimization in modern deep learning settings. The dichotomy between SGD and heavy-ball-like ODEs is shown to be governed by the resonance condition tied to per-parameter sparsity. The identification of a “spectral conflict” for global momentum schedules underscores a critical challenge in large-scale model optimization, motivating further work in sparsity- and adaptation-aware training protocols.

Figure 8: Dense above-resonance regime—main ODE trajectories rapidly converge to the canonical 2D heavy-ball limit, verifying the reduction to deterministic oscillatory behavior.

---

References

• "Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions" (2605.28961)