Momentum Methods

Follow this topic Ask about this topic

Momentum

2 papers in the last four weeks, down 33% on the four weeks before. 0.0% of all new papers.

Jul 13Week of Sep 28

Latest papers 43

Oct 1, 2026cs.LG

Muon meets Tamed Langevin: Momentum Preconditioning beyond Convex and gradient-Lipschitz Potentials

We consider the problem of sampling from Gibbs distributions on matrix spaces whose potential energies are neither convex nor globally gradient-Lipschitz. We introduce a family of non-quadratic kinetic energies that lead to a new underdamped Langevin system with momentum preconditioning, in which the gradient of the kinetic energy acts as a smooth spectral taming of the momentum. We prove that, under these relaxed assumptions on the potential, the resulting dynamics leaves the target Gibbs measure invariant, and we establish exponential convergence to equilibrium in a weighted total variation distance. Finally, we show that the corresponding Euler-Maruyama discretization admits moment bounds that are uniform in time, without any modification of the potential gradient, which ensures the stability of the resulting sampling algorithm.

Nikolaos Makras, Sotirios Sabanis

School of Mathematics, University of Edinburgh, UK · National Technical University of Athens, Greece · Athena/Archimedes Research Centre, Greece

Markov Chain Monte Carlo Momentum Methods Langevin Dynamics

Oct 1, 2026math.OC

Optimal Momentum Methods for Stochastic Multilevel Compositional Optimization

This paper investigates stochastic multi-level optimization where the objective is a nested composition of several smooth non-convex functions. We assume that only stochastic estimates of the gradient and function values for each level are accessible. Consequently, obtaining an accurate estimate of the overall gradient is challenging due to the nested structure. To address this, we employ a momentum-based estimator with mini-batches to track the function values of each level, which are subsequently used to construct momentum gradient estimators. We establish an optimal sample complexity of

\mathcal{O}(ε^{-4})

for finding an

ε

-stationary point, avoiding the stronger average smoothness assumption commonly relied upon in prior literature. Furthermore, by employing a normalization technique, we attain the same rate without requiring problem-dependent constants to set hyperparameters. To achieve the optimal rate without mini-batches, we further develop a batch-free method that incorporates a first-order approximation and a clipping technique for function value estimation. Finally, we validate the effectiveness of our proposed methods through experiments on risk-averse portfolio optimization and hierarchical tilted empirical risk minimization.

Wei Jiang, Rui Yan, Sifan Yang +3

School of Computer Science and Engineering, Nanjing University of Science and Technology, Nanjing, China · School of Artificial Intelligence, Nanjing University, Nanjing, China · School of Software Technology, Zhejiang University, Ningbo, China

Stochastic Optimization Momentum Methods Nonconvex Stochastic Optimization Stochastic Optimization Convergence

Sep 2, 2026stat.ML

Momentum in large-batch training: Polyak enlarges the critical batch size, Nesterov improves data efficiency

We study when and how momentum improves large-batch training in the one-pass regime, using power-law kernel regression as a tractable setting. We first characterize risk stability through the critical learning rate, defined as the largest learning rate for stable training, and obtain

η_{\mathrm{SGD}}^{\mathrm{crit}}\eqsim 1

,

η_{\mathrm{Polyak}}^{\mathrm{crit}}\eqsim \min\{1,B(1-ρ)\}

, and

η_{\mathrm{Nesterov}}^{\mathrm{crit}}\eqsim \min\{1,B^β(1-ρ)\}

, where

B

is the batch size,

ρ

is the momentum factor, and

β>1

is the capacity exponent. Within this admissible region, we derive scaling laws for the full risk dynamics, capturing the progression from an early transient, through power-law decay, to a noise floor. We then minimize the final-step risk over the admissible learning rates and momentum factors under a fixed data budget, yielding a three-regime batch-size phase diagram that reveals how the role of momentum changes with batch size. Notably, Polyak enlarges the critical batch size, the largest batch size preserving the best small-batch data-scaling exponent, thereby enabling greater parallelism without sacrificing data efficiency. In contrast, Nesterov achieves better data efficiency in the large-batch regime because its look-ahead mechanism suppresses noise accumulation. Numerical experiments validate the predicted stability boundaries, risk dynamics, and batch-size phase diagram.

Jia-Nan Wang, Zixun Huang, Kairui Li +1

Peking University · University of Pennsylvania · AI for Science Institute, Beijing

Momentum Methods Gradient Descent

Aug 19, 2026cs.LG

Activation-Keyed Momentum: An Anisotropic Momentum Update via the Delta Rule

Most modern optimizers form their momentum as an exponential moving average (EMA) of past gradients, forgetting every direction at one fixed rate. However, the inputs a deep network sees during training can be highly anisotropic, with a few directions queried frequently while most are seen rarely. Preconditioning methods address this anisotropy by wrapping extra processing around this buffer and leave the momentum update itself unchanged. We propose Activation-Keyed Momentum (AK-Momentum), which builds direction-awareness into the momentum update rule. The gradient of a linear layer splits into an input activation that acts as a key and an output-side error that acts as a value. Keying on that activation, AK-Momentum updates the momentum buffer by the canonical delta rule, so each direction is forgotten at a rate set by how often it appears. We prove that it is a valid momentum, that it applies the input-side curvature correction without matrix inversion, and that it clears stale directions faster than EMA under both a fixed and a drifting optimum. It is a drop-in replacement for the momentum buffer of any optimizer, its coefficient transfers across widths under

μ

P, and its extra compute stays between

22.2\%

and

25.0\%

of a gated-MLP block's linear cost with no persistent memory. In FineWeb-Edu pretraining, AdamW with AK-Momentum (AK-AdamW) reaches AdamW's validation loss in up to

46.39 \pm 4.32\%

fewer steps at 67M and

22.12 \pm 0.80\%

at 370M over three seeds, and the gain persists at 1B on a Chinchilla-optimal budget. A Muon baseline tuned under the same protocol sits above AK-AdamW at both language-model scales, and the gain holds for SGD, ResNet-18, and ViT-Tiny on CIFAR-10. Training-time diagnostics confirm the predicted mechanism, better gradient tracking and healthier input directions.

Euijin Hong, Guannan Qu

Electrical and Computer Engineering, Carnegie Mellon University

Deep Learning Optimization Momentum Methods Neural Network Optimization Accelerated Gradient Methods

Aug 13, 2026cs.LG

Momentum as Residual-Driven Multiplier Correction for Deep Learning Optimization

Momentum-based optimizers are widely used in modern deep learning, yet the relations among momentum recursion, update geometry, and acceleration remain only partially understood. We develop an

\textbf{A}

DMM-

\textbf{I}

nspired

\textbf{M}

omentum (AIM) framework based on residual-penalty variable splitting, which interprets momentum as a multiplier-like correction driven by the splitting residual. AIM recovers the exponential moving average of gradients from an ADMM-style multiplier update and separates two mechanisms that are usually intertwined in practical optimizers: the residual penalty determines the update geometry, whereas the approximation of the objective-related subproblem determines the acceleration form. Building on AIM, we propose

\textbf{R}

elativistic

\textbf{A}

daptive gradient

\textbf{D}

escent with

\textbf{A}

ccelerated

\textbf{R}

esidual (RADAR), which combines relativistic adaptive geometry, decoupled residual correction, and second-order momentum filtering to improve the update direction and momentum estimation. We establish stochastic convergence through a variance-perturbed Lyapunov drift analysis. Experiments on supervised vision learning, language modeling, and reinforcement learning show that RADAR achieves consistent improvements over strong adaptive optimizer baselines.

Zhixin Ren, Yau Lyu, Congrong Li +2

Department of Mathmatical Science, Tsinghua University · School of Vehicle and Mobility, Tsinghua University · College of AI, Tsinghua University

Deep Learning Optimization Momentum Methods Adaptive Gradient Methods Accelerated Gradient Methods

Aug 6, 2026math.OC

An Inertial Block Proximal Linearized Method with Adaptive Momentum for Nonconvex and Nonsmooth Optimization

In this paper, we consider a class of multiblock nonconvex nonsmooth optimization problems, which covers many applications such as the analysis of pre-earthquake anomalies and machine learning. To solve this class of problems, we propose the inertial block proximal linearized method with two-phase adaptive momentum (IBPL

^+

-TP). Compared to the current methods, our method possesses three main advantages: (1) it introduces a two-phase adaptive momentum strategy to effectively update the extrapolation parameters, (2) it allows using two different extrapolation points to accelerate the convergence, (3) it allows the extrapolation parameters of these two extrapolation points to be independent of and unconstrained by all other parameters. While maintaining the above advantages, we prove that our method ensures the monotonic convergence of the objective function of this class of problems, and we also prove that the sequence generated by our method globally converges to a critical point, as well as establish the convergence rate of our method. To demonstrate the effectiveness of our method, we apply it to solve two nonconvex and nonsmooth machine learning problems, namely sparse nonnegative matrix factorization with

\ell_0

-constraints and sparse nonnegative CP decomposition with

\ell_0

-constraints. The numerical experimental results on solving these problems show that our method outperforms several state-of-the-art methods.

Weifeng Yang

Yunnan Earthquake Administration, Kunming 650224, China

Momentum Methods Nonconvex Optimization

Jul 25, 2026math.OC

Nesterov acceleration in optimizing over probability measures

Optimization over probability measures has become an increasingly important paradigm in modern machine learning, scientific computing, and uncertainty quantification. Motivated by Nesterov's accelerated gradient method in Euclidean space, we develop Heavy-ball and Nesterov acceleration methods over the probability measure space

\mathcal{P}_2

and establish non-asymptotic convergence guarantees that match their Euclidean counterparts. In particular, we derive convergence rates with respect to both the number of iterations and the number of particles used to represent the underlying probability distributions. Extending accelerated optimization from Euclidean space to probability measures is challenging. The natural notion of momentum requires concepts such as tangent bundles of the set of probability space and they are hard to operate numerically. To overcome these difficulties, we introduce two complementary lifting procedures. The first lifts probability measures to phase space through a Hamiltonian formulation, introducing momentum variables into the dynamics. The second lifts probability measures to a common Hilbert space, restoring the linear structure required for convergence analysis while simultaneously yielding executable particle dynamics. Together, these two complementary lifting procedures provide a systematic methodology for designing, analyzing, and implementing momentum-based accelerated optimization methods over probability measure spaces.

Jiaqi Tang, Qin Li, Wilfrid Gangbo

Department of Statistics, University of Wisconsin–Madison, Madison, WI 53706, USA · Department of Mathematics, University of Wisconsin–Madison, Madison, WI 53706, USA · Department of Mathematics, University of California, Los Angeles, Los Angeles, CA 90095, USA

Momentum Methods Gradient Descent

Jul 16, 2026cs.LG

Muse: Representation Geometry of Muon Beyond Normalized Momentum

Muon-style optimizers apply a polar map to matrix momentum, but their updates also depend on the representation of each parameter block before orthogonalization. We study this representation choice as a form of optimizer geometry and introduce {\method}, a family of Muon-style optimizers that shares the same momentum rule and Newton--Schulz backend across native, nearest-square, skinny, and vector representations. Each Frobenius-isometric representation induces a distinct polar steepest-descent geometry, in which the shorter matrix dimension determines the number of supported singular channels, the pullback scaling, and the constants in stochastic nonconvex convergence bounds. In a teacher--student model, curvature collapse and an isotropic Marchenko--Pastur spectral profile connect early-stage dissipation to the represented nuclear-to-squared-Frobenius norm ratio. Pretraining experiments on LLaMA2-130M and LLaMA2-600M, together with fixed-momentum diagnostics, show that balanced non-native representations can match the performance of the native representation, whereas reducing the shorter dimension weakens the scaling and singular-channel support, leading to behavior that increasingly resembles normalized momentum.

Da Chang, Qiankun Shi, Lvgang Zhang +4

1Pengcheng Laboratory · 2Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences +3

Representation Geometry Momentum Methods Muon Optimizer Matrix Optimization

Jun 25, 2026cs.LG

Heavy-Ball Q-Learning with Residual Weighting Correction

This paper proposes a corrected heavy-ball Q-learning method for reinforcement learning (RL) and establishes convergence of its deterministic mean dynamics. It also identifies conditions under which the method is theoretically guaranteed to converge faster than standard Q-learning. The same construction is then extended to Q-learning with linear function approximation, where analogous convergence and acceleration statements are derived for the corresponding corrected fixed point. The sampled stochastic versions are treated through conditional-mean recursions and, in the stated linear-function-approximation setting, finite-time bounds. The analysis is based on a switched linear system (SLS) representation of Q-learning algorithms and on the joint spectral radius (JSR) of the associated switching families. This SLS viewpoint is not commonly used in standard analyses of Q-learning, and it provides a complementary framework and new insight into how heavy-ball momentum can accelerate Q-learning.

Donghwan Lee

Department of Electrical Engineering Korea Advanced Institute of Science and Technology (KAIST) Daejeon 34141, South Korea

Momentum Methods Q-Learning

Jun 17, 2026cs.LG

Compute Efficiency and Serial Runtime Tradeoffs for Stochastic Momentum Methods

Stochastic momentum methods such as heavy ball (HB), Nesterov momentum, and variants of Accelerated SGD (ASGD) [Kidambi et al., 2018] are widely used in modern training, but their stochastic benefits depend on two distinct quantities: serial runtime, the number of iterations needed to reach a target accuracy, and compute efficiency (CE), the inverse total gradient-query or FLOP cost. Larger batches reduce serial runtime without hurting CE only when the contraction gap grows linearly with batch size. We study stochastic HB and ASGD for consistent linear regression with Gaussian covariates and prove finite-dimensional, discrete-time lower bounds on their batch-size tradeoffs. Our first result shows that HB does not improve the CE frontier over SGD for arbitrary spectra; rather, it preserves SGD-level CE over a larger batch-size window, allowing larger batches to reduce serial runtime until HB reaches its deterministic accelerated scale. This window can be a factor

\sqrtκ

larger than the SGD critical batch size. For ASGD, the picture is more spectrum-dependent: for rapidly decaying power-law spectra, ASGD improves small-batch CE over HB/SGD, but as batch size grows it trades this CE advantage for improved serial runtime. Synthetic linear-regression experiments verify these qualitative regimes, including near-overlap of ASGD and HB for slowly decaying spectra and the predicted CE--serial tradeoff for rapidly decaying spectra.

Depen Morwani, Alexandru Meterez, Pranav Nair +1

1Harvard University · 2Kempner Institute at Harvard University · †: Work done while at Harvard.

Stochastic Optimization Momentum Methods Stochastic Optimization Convergence Accelerated Gradient Methods Stochastic Gradient Descent

Jun 15, 2026cs.LG

CacheMuon: Using Temporal Preconditioning To Approximate Polar Factor

Muon is an optimizer that computes updates using the polar factor of the momentum matrix and has shown strong empirical performance across a range of training settings. A key component of Muon is the Newton-Schulz iteration used to compute this polar factor. Although this avoids the cost of an exact singular value decomposition, it remains expensive in practice because it is applied at every optimization step. At the same time, the momentum matrix changes smoothly over training, suggesting strong temporal correlation in the corresponding polar factors. In this paper, we exploit this structure and propose CacheMuon, a temporal preconditioning method that reuses information from previous optimization steps to approximate the polar factor at the current step. This reduces redundant orthogonalization computation across iterations. We analyze CacheMuon as an inexact Muon update, with error controlled by fresh-solver error and cache staleness. Empirically, CacheMuon provides a controllable quality-efficiency frontier: conservative thresholds closely match fresh Muon on language-model and vision training while reducing orthogonalization FLOPs, whereas more aggressive thresholds yield larger arithmetic savings at the cost of modest validation-quality degradation.

Bishnu Dev, Sushil Bohara, Martin Takáč +1

Mohamed bin Zayed University of Artificial Intelligence, Abu Dhabi, UAE · Mohamed bin Zayed University of Artificial Intelligence

Momentum Methods Neural Network Optimization Muon Optimizer

Jun 7, 2026math.OC

OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality

Orthogonalized momentum updates, as used in Muon-style optimizers, have recently shown strong empirical stability in large-scale deep learning. However, most current orthogonalized methods are still paired with fixed, externally scheduled, or otherwise open-loop magnitude rules, so their scale is not directly calibrated from the realized optimization trajectory. Motivated by the closed-loop perspective behind Lipschitz-free and noise-adaptive methods, we propose OptMuon, a family of adaptive momentum orthogonalization methods for stochastic nonconvex optimization. OptMuon combines Muon-style polar-factor directions with a trajectory-dependent AdaGrad-Norm-type coefficient schedule, so that the update magnitude is determined by the observed gradient and momentum history rather than by a prescribed Lipschitz-dependent rule. The schedule does not use the smoothness constant, the variance level, or the bounded-gradient constant in parameter selection, and its running-maximum correction prevents isolated gradient spikes from causing excessive coefficient collapse. Under lower-boundedness, unbiased stochastic gradients with bounded variance, smoothness, and an almost-sure bounded stochastic-gradient condition, we prove two complementary expected-stationarity guarantees. OptMuon-A achieves the noise-adaptive rate

\tilde{\mathcal O}(T^{-1/2}+σ^{1/2}T^{-1/4})

under average smoothness, while OptMuon-I achieves

\tilde{\mathcal O}(T^{-1/2}+σ^{1/3}T^{-1/3})

under individual smoothness. In the zero-noise regime, both bounds automatically reduce to a nearly optimal deterministic first-order rate

\tilde{\mathcal O}(T^{-1/2})

without manual hyperparameter retuning. These results show that closed-loop scalar adaptation can be combined with Muon-style momentum orthogonalization while retaining noise adaptivity and zero-noise optimality up to logarithmic factors.

Ganzhao Yuan

Faculty of Computer Science and Artificial Intelligence Shenzhen University of Advanced Technology (SUAT), China

Momentum Methods Nonconvex Stochastic Optimization Muon Optimizer Stochastic Optimization Convergence Adaptive Gradient Methods

Jun 2, 2026cs.LG

Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering

Muon has recently demonstrated strong empirical performance in large language model training, but the theoretical role of momentum in Muon remains unclear. Existing analyses of Muon either remove momentum to study spectral updates in isolation, or retain momentum without explaining why it improves empirical performance. Our work bridges this gap by showing momentum in Muon acts as a spectral filter. Under a structured signal-plus-perturbation gradient model, we prove that momentum suppresses perturbations while preserving the dominant signal, thereby enlarging the spectral gap between them. This enlarged gap stabilizes the singular subspaces of the matrix passed to Muon's orthogonalization step, making the resulting update more reliable. We further show that applying momentum before orthogonalization achieves provably stronger alignment with the signal component of the gradient than either reversing this order or simply removing momentum. Experiments across diverse tasks, including LLM pretraining, support our theoretical analysis. More broadly, our theory offers a starting point for understanding the benefits of momentum in other matrix-based optimizers.

Xianliang Li, Zihan Zhang, Weiyang Liu +1

1The Institute of Statistical Mathematics · 2The Graduate Institute for Advanced Studies, SOKENDAI · 3National Institute of Informatics +3

Deep Learning Optimization Momentum Methods Muon Optimizer Spectral Methods Gradient Descent Dynamics

Jun 2, 2026cs.LG

Spectral Scaling Laws of Muon

Orthonormalized update rules have rapidly become a leading choice of optimizer for training large language models, with recent open-source state-of-the-art models adopting Muon. To keep these updates tractable, Muon performs the orthonormalization with the Newton--Schulz (NS) iteration. Since NS is only approximate, directions with small singular values fail to be orthonormalized. In Muon, NS is applied to the momentum matrix at every step, yet little is known about how the singular value spectrum of these momentum matrices behaves during training, or how that behavior changes with model size. We present the first systematic study of this question. Tracking singular value quantiles of the momentum buffer across layers in models ranging from 77M to 2.8B parameters, we observe a consistent picture: after a short burn-in, the quantiles stabilize at a value determined by the layer type and model size. These stabilization values follow remarkably clean power laws in model size, with layer-dependent exponents. Layers up to mid-late depth scale very mildly with model size

M

(around

M^{-0.25}

), so the standard 5-step NS configuration used at academic scale will continue to orthonormalize them at much larger scales. Some of the late layers, however, scale much more aggressively (up to

M^{-0.96}

) and will fall into the NS failure regime at frontier scale unless one uses more NS iterations or better-tuned coefficients. NS iterations are computationally expensive at scale; our laws give practitioners a principled, layer-aware recipe for choosing the minimum NS configuration that still orthonormalizes the directions that matter -- avoiding unnecessary computation without sacrificing update quality.

Gagik Magakyan, Pablo Parrilo, Asuman Ozdaglar

MIT

Newton-Schulz Iteration Momentum Methods Muon Optimizer Second-Order Optimization Neural Network Training Dynamics LLM Training

Jun 1, 2026cs.LG

Fast Generalization after Interpolation via Critically Damped Momentum Optimization

A central problem in machine learning is that models can achieve near-perfect training performance while generalizing substantially less well to unseen examples. This gap is especially acute in high-dimensional, low-sample regimes, where many interpolating solutions exist and optimization must implicitly select among minima with different generalization properties. Following recent theoretical advances on optimization dynamics near the interpolation threshold, we note that the two-regime structure of risk minimization, with loss minimization followed by complexity minimization, motivates a biphasic optimization schedule. We thus theoretically demonstrate that GROKtimizer, a biphasic strategy that combines rapid convergence to interpolation with Critically Damped Momentum (CDM)-based post-interpolation norm minimization, offers a natural solution for selecting low-norm interpolating solutions. Under a local quadratic model of the post-interpolation basin, GROKtimizer provides a quadratic speedup over classical gradient descent, with provable optimality among first-order optimizers. To showcase the applicability of our method, we evaluate GROKtimizer on several synthetic benchmarks common in the classical grokking literature and on various real-world datasets. Finally, we reconcile our findings with the flat-minima hypothesis, highlighting the importance of post-interpolation dynamics in the construction of high-quality, generalizing models.

Luca Muscarnera, Silas Ruhrberg Estévez, Yuanzhang Xiao +1

Neural Network Generalization Momentum Methods Neural Network Optimization Implicit Regularization

May 27, 2026stat.ML

Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions

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.

Katie Everett, Elliot Paquette

Google DeepMind & MIT · McGill University & Mila

Momentum Methods Gradient Descent Dynamics Stochastic Gradient Descent

May 27, 2026cs.LG

Outer-Momentum Restarting in High-Dimensional Two-Phase Optimization

Communication-efficient distributed optimizers such as DiLoCo reduce synchronization costs by letting workers perform many local updates before aggregating their progress with an outer momentum optimizer. Recent theory suggests that the outer optimizer acts on an effective spectrum induced by the inner optimization loop, and that the choice of outer momentum controls how progress from local updates is accumulated across communication rounds. We study periodic restarting of the outer momentum as a simple complementary mechanism for controlling this outer memory. In a linearized squared-loss model where prediction-space residuals evolve under the empirical NTK, we derive a mode-wise restart contraction showing that resets exploit phase cancellation by discarding stale momentum while preserving inner-loop progress. Toy experiments verify the predicted contraction behavior, and language-model pretraining experiments show that periodic restarts widen the stable range of outer learning rates and momentum values across communication periods.

Kristi Topollai, Allan Ma, Tolga Dimlioglu +2

New York University, New York, USA

Momentum Methods Communication-Efficient Distributed Training Distributed Optimization

May 27, 2026cs.LG

Stochastic Gradient Descent with Momentum is Algorithmically Stable

Stochastic gradient descent with momentum (SGDM) is one of the most widely used optimization algorithms in machine learning. While optimization properties of SGDM have been extensively studied in the literature, it remains insufficiently understood whether and when SGDM can generalize well to unseen data. In particular, it has been conjectured that while momentum accelerates training, it may degrade generalization. In this paper, we close this gap by developing a comprehensive generalization analysis of SGDM through the lens of algorithmic stability. More specifically, we introduce a generalized SGDM framework that encompasses both Polyak's and Nesterov's momentum schemes, and establish tight on-average model stability bounds for smooth and convex problems. Notably, the obtained bounds exploit small optimization error bounds along the trajectory, apply to any momentum parameter in the interval

[0, 1)

, and do not require the commonly assumed Lipschitzness of loss functions. We further derive optimization error bounds for the generalized SGDM, and combine them with our generalization analyses to obtain optimal excess population risk bounds for SGDM with both Polyak's and Nesterov's momentum.

Yunwen Lei, Zimeng Wang, Xiaoming Yuan

Department of Mathematics, The University of Hong Kong, Hong Kong, China · Department of Mathematics and Mathematical Statistics, Umeå University, Sweden

Momentum Methods Algorithmic Stability Stochastic Optimization Convergence Accelerated Gradient Methods Stochastic Gradient Descent Generalization Bounds

May 23, 2026cs.LG

Momentum Streams for Optimizer-Inspired Transformers

The residual update of a pre-norm Transformer layer admits an interpretation as one step of a first-order optimizer acting on a surrogate token energy, wherein the attention and MLP sublayers function as gradient oracles. Based on this observation, we build a family of optimizer-inspired Transformers (triple-momentum, Adam/AdamW, Muon, SOAP) and compare them under matched compute. In our main pretraining experiment, the triple-momentum TMMFormer achieves the lowest validation loss, outperforming the vanilla Transformer and prior architectural variants. A controlled ablation and supporting theory show that momentum, not preconditioning, is the main source of the gain. We further show that TMMFormer and other momentum-based designs reach flatter minima than the vanilla Transformer, which leads to less forgetting and better generalization.

Jingchu Gai, Nai-Chieh Huang, Jiayun Wu

Carnegie Mellon University

Deep Learning Optimization Transformer Momentum Methods Neural Network Optimization Flat Minima

May 22, 2026stat.ML

Move on Muon : A Hamiltonian probability gradient flow perspective of Muon optimizer

We develop a gradient flow on the space of probability measures defined on matrix-valued parameters induced by regularized Muon, an analytically smoothed version of the idealized Muon optimizer. The key observation is that the regularized orthogonalization map is the gradient of a smooth Fenchel-dual smoothing of the nuclear norm. This identifies the (regularized) Muon update as a mirror/prox step in the update variable, with momentum acting as the dual coordinate. We use this structure to lift Muon from a single matrix parameter to finite-particle probability objectives of the form

J(ρ)=R\left(\int F d ρ\right)

, a setting motivated by mean-field descriptions of neural-network training, and derive the inertial continuous-time limit. Using this structure, we derive the finite-particle continuous-time limit under the inertial scaling of step size and momentum, and then pass to a phase-space mean-field equation over probability laws on parameter-momentum pairs. The resulting flow can be shown to be a damped Hamiltonian probability dynamics whose kinetic energy is induced by the regularized Muon mirror potential. We prove an exact Hamiltonian dissipation identity, showing that the Hamiltonian energy decreases monotonically. While the target objective itself need not be monotone along the inertial Muon dynamics, under additional gradient-dominance, bounded-momentum, and curvature/alignment assumptions, we obtain continuous and discrete-time exponential convergence rates for the objective gap. We also study the well-posedness of the mean-field limit equation and establish propagation of chaos guarantees for the interacting particle system. Finally, we extend the formulation to Hilbert-valued feature maps on product matrix spaces, yielding a blockwise Muon probability flow applicable to smooth transformer mixture-of-experts models.

Aratrika Mustafi, Soumya Mukherjee, Bharath K. Sriperumbudur

Department of Statistics, Pennsylvania State University

Momentum Methods Muon Optimizer Mean-Field Theory Neural Network Training Dynamics

May 19, 2026cs.LG

Ada2MS: A Hybrid Optimization Algorithm Based on Exponential Mixing of Elementwise and Global Second-Moment Estimates

Optimization algorithms are core methods by which machine learning models iteratively minimize loss functions, update parameters, learn from data, and improve performance. Momentum SGD and AdamW represent two important optimization paradigms. AdamW produces stable updates and usually has strong robustness across training scenarios, but its generalization performance is sometimes weaker than that of momentum methods. Momentum SGD can often obtain better generalization after careful tuning, but it is more sensitive to gradient-scale variation and hyperparameter settings. To balance the strengths and weaknesses of the two paradigms, this paper proposes Ada2MS, an optimization algorithm that achieves a smooth transition between AdamW-like behavior and momentum-SGD-like behavior through continuous exponential interpolation between elementwise second-moment estimates and global second-moment estimates. On the visual tasks evaluated in this study, Ada2MS obtains competitive results under a unified optimizer-comparison protocol. The code will be released at https://github.com/mengzhu0308/Ada2MS

Meng Zhu, Quan Xiao, Weidong Min

School of Information Management and Mathematics, Jiangxi University of Finance and Economics, Nanchang 330032, China · School of Mathematics and Computer Science, Nanchang University, Nanchang 330031, China · Institute of Metaverse, Nanchang University, Nanchang 330031, China +1

Deep Learning Optimization Momentum Methods Adaptive Gradient Methods

May 19, 2026cs.LG

LionMuon: Alternating Spectral and Sign Descent for Efficient Training

In large-scale optimization, the cheapness and effectiveness of update steps are the most crucial factors for a successful optimizer. Sign-based optimizers like Lion or Signum produce cheap per-step updates, whereas Muon's spectral matrix-sign update gives a much stronger direction at a substantially higher per-step cost. In this work, we propose LionMuon, which retains the effectiveness of Muon steps while considerably cutting the averaged iteration cost, similar to sign-based methods. It alternates between Lion's and Muon's updates on a fixed period P, sharing a single dual-EMA momentum buffer between them. The optimizer state memory therefore matches Lion and is exactly half of AdamW's. A simpler single-EMA variant, SignMuon, by itself already outperforms pure Muon. At P = 2, LionMuon Pareto-dominates Muon, Lion, Signum, and AdamW on every dataset and architecture we tested at 124M model size, reaching lower validation loss at lower compute, and the same advantage persists at 355M and 720M scale. On the theory side, we prove sharp complexity bounds under heavy-tailed noise which are governed by period-averaged smoothness and noise that interpolate between Muon's and Lion's constants. These bounds predict the compute-optimal period and the conditions under which LionMuon outruns Muon and Lion. Code: https://github.com/brain-lab-research/lion-muon

Arman Bolatov, Artem Riabinin, Nikita Kornilov +4

Mohamed bin Zayed University of Artificial Intelligence (MBZUAI) · Basic Research of Artificial Intelligence Laboratory (BRAIn Lab) · Applied Artificial Intelligence Institute +1

Stochastic Optimization Sign-Based Optimization Momentum Methods Neural Network Optimization Muon Optimizer Stochastic Optimization Convergence

May 18, 2026cs.LG

Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration

Accelerating stochastic gradient methods with classical momentum schemes, such as Polyak's heavy ball, has proven highly successful in training large-scale machine learning models, particularly when combined with the hardware acceleration of large mini-batch computations. Yet, the effect of classical momentum on stochastic mini-batch optimization has been poorly understood theoretically, with prior works requiring strong noise assumptions and extremely large mini-batches. In this work, we develop a general theory of stochastic momentum acceleration for optimizing over quadratics in the interpolation regime, a popular abstraction for studying deep learning dynamics which also includes classical methods such as randomized Kaczmarz and coordinate descent. Our framework encompasses both heavy ball and Nesterov-style momentum, allows for arbitrary mini-batch sizes, and makes minimal assumptions on the stochastic noise. In particular, we show that acceleration from classical momentum is directly proportional to the gradient mini-batch size (up to a natural saturation point), thereby enabling perfect parallelization of mini-batch computations. Our theory also provides a simple choice for the momentum parameter, which is shown to be effective empirically.

Sachin Garg, Michał Dereziński

University of Michigan

Momentum Methods Stochastic Optimization Convergence Accelerated Gradient Methods Stochastic Gradient Descent

May 18, 2026cs.LG

Ringmaster LMO: Asynchronous Linear Minimization Oracle Momentum Method

Muon has recently emerged as a strong alternative to AdamW for training neural networks, with encouraging large-scale pretraining results and growing evidence that matrix-structured updates can be faster in practice. Yet Muon, and more generally Linear Minimization Oracle (LMO) based methods, are typically used synchronously. This is problematic in heterogeneous distributed systems, where workers complete gradient computations at different speeds and synchronous training must repeatedly wait for slower workers. In this work, we introduce Ringmaster LMO, an asynchronous LMO-based momentum method for unconstrained stochastic nonconvex optimization. Our method builds on the delay-thresholding idea of Ringmaster ASGD. For SGD-type methods, Ringmaster ASGD achieves optimal time complexity by discarding overly stale gradients. Ringmaster LMO extends this mechanism to general LMO-based updates. We establish convergence guarantees under generalized

(L_0, L_1)

-smoothness and further develop a parameter-agnostic variant with decreasing stepsizes and adaptive delay thresholds. Finally, we translate our iteration guarantees into time complexity bounds under heterogeneous worker computation times. In the classical Euclidean smooth setting, these bounds recover the optimal time complexity of Ringmaster ASGD. Experiments on stochastic quadratic problems and NanoChat language-model pretraining show that the advantages of Ringmaster LMO grow with system heterogeneity and that the method outperforms strong synchronous and asynchronous baselines.

Abdurakhmon Sadiev, Artavazd Maranjyan, Ivan Ilin +1

KAUST

Momentum Methods Distributed Optimization Nonconvex Stochastic Optimization Stochastic Optimization Convergence

May 14, 2026math.OC

Stochastic Compositional Optimization via Hybrid Momentum Frank--Wolfe

Stochastic compositional optimization minimizes objectives of the form

\min_{\bm{x} \in \mathcal{X}} F(\bm{f}(\bm{x}), \bm{x})

, where

\bm{f}

is accessible only through noisy stochastic queries. Existing methods for this problem assume that the outer function

F

is continuously differentiable, which excludes many practically important applications such as robust max-of-losses, Conditional Value-at-Risk, and norm regularizers. We propose the Hybrid Momentum Stochastic Frank--Wolfe algorithm, which drops the smoothness assumption on

F

. By combining a momentum-based Jacobian tracker with a Taylor-corrected function tracker, the algorithm feeds an entire stochastic linearization -- rather than a single gradient -- into a generalized linear minimization oracle. We establish an

\mathcal{O}(K^{-1/4})

convergence rate in the generalized Frank--Wolfe gap for non-convex objectives with

L_F

-Lipschitz outer functions, matching the optimal complexity for projection-free single-sample stochastic methods under expected smoothness. The analysis extends to heavy-tailed noise oracles with bounded

r

-th moments for

r \in (1, 2]

and recovers the deterministic rates of Vladarean et al (2023) as the noise vanishes.

El Mahdi Chayti

Machine Learning and Optimization Laboratory (MLO) EPFL, Switzerland

Stochastic Optimization Momentum Methods Nonconvex Stochastic Optimization Constrained Optimization

May 13, 2026cs.LG

Spectral Flattening Is All Muon Needs: How Orthogonalization Controls Learning Rate and Convergence

Muon orthogonalizes the momentum buffer before each update, replacing its singular values with ones via Newton-Schulz iterations. This simple change lets Muon tolerate far larger learning rates and converge faster than other optimizers, but why? We show that the mechanism is spectral flattening, and develop two results around it. First, we prove that Muon's maximal stable step size scales with the average singular value of the gradient rather than the largest, which bottlenecks standard gradient descent. Second, we recast Muon as a preconditioned gradient method and show, under a Kronecker-factored curvature model, that it improves the effective convergence factor, with the improvement controlled by the spectrum of the gradient covariance. Extensive experiments validate both results: Muon remains stable at learning rates that cause SGD to diverge within the first few iterations, and reaches accuracy milestones several epochs earlier even at identical step sizes. Taken together, our results offer a principled, geometric explanation for Muon's empirical success.

Tien-Phat Nguyen, Truong Nguyen, Minh-Phuc Truong +3

Hanoi University of Science and Technology Hanoi, Vietnam · Indiana University 107 S. Indiana Ave, Bloomington, IN 47405, USA · Monash University Clayton, VIC 3800, Australia

Momentum Methods Gradient Descent Muon Optimizer Spectral Methods Stochastic Gradient Descent

May 13, 2026math.OC

Adam-SHANG: A Convergent Adam-Type Method for Stochastic Smooth Convex Optimization

We propose Adam-SHANG, a Lyapunov-guided Adam-type method that couples momentum, adaptive preconditioning, and a curvature-aware correction through a more stable lagged-preconditioner update. For stochastic smooth convex optimization, we prove convergence in expectation under an admissible stepsize condition that can always be satisfied by a conservative spectral bound, without imposing global monotonicity on the second-moment sequence. To obtain a less conservative practical rule, we introduce a computable trace-ratio stepsize, motivated by a local coordinatewise alignment condition. The same structural update is also tested beyond the convex setting with simplified parameters. Experiments validate the predicted stochastic decay and show competitive training performance against Adam and AdamW on deep learning tasks.

Yaxin Yu, Long Chen, Minfu Feng

School of Mathematics Sichuan University Chengdu, China · Department of Mathematics University of California, Irvine Irvine, CA, USA

Stochastic Optimization Momentum Methods Stochastic Optimization Convergence Adaptive Gradient Methods

May 11, 2026cs.LG

Refresh-Scaling the Memory of Balanced Adam

Recent evidence suggests that Adam performs robustly when its momentum parameters are tied,

β_1=β_2

, reducing the optimizer to a single remaining parameter. However, how this parameter should be set remains poorly understood. We argue that, in balanced Adam,

β

should not be treated as a dimensionless constant: it defines a statistical memory horizon

H_β=(1-β)^{-1}

. In terms of the effective learning horizon

T_{\mathrm{ES}}

, estimated from the validation trajectory, we study the refresh count

R_β=(1-β)T_{\mathrm{ES}}

, which measures how many times Adam renews its internal statistics during the useful phase of training. Across 11 vision and language experiments, we find that choosing

β

so that

R_β\approx1000

selects different

β

values depending on the training scale, yet improves robustness over the best fixed-beta baseline. Compared with the strongest fixed choice

β=0.944

, the refresh rule improves worst-case robustness, reducing the maximum relative gap in validation loss by 33.4%, while bringing all 11 runs within 1% of their validation oracle. These results suggest that the remaining hyperparameter of balanced Adam is more naturally viewed as a memory-scale variable than as a fixed constant. This provides a simple budget-aware perspective on optimizer scaling and opens a path toward treating Adam's momentum as part of the learning dynamics rather than as a static default.

Alberto Fernández-Hernández, Cristian Pérez-Corral, Jose I. Mestre +2

Universitat Polit`ecnica de Val`encia, Val`encia, Spain · Universitat Jaume I, Castell´o de la Plana, Spain

Deep Learning Optimization Momentum Methods Adaptive Gradient Methods Hyperparameter Optimization

May 7, 2026math.OC

Muon with Nesterov Momentum: Heavy-Tailed Noise and (Randomized) Inexact Polar Decomposition

Most first-order optimizers treat matrix-valued parameters as vectors, ignoring the intrinsic geometry of hidden-layer weights in neural networks. Muon addresses this mismatch by updating along the polar factor of a momentum matrix, but its theoretical understanding has lagged behind practice. In particular, practical implementations incorporate Nesterov momentum, compute the polar factor only approximately, and operate with stochastic gradients that may be heavy-tailed. We close this gap by developing a convergence theory for Muon with Nesterov momentum and inexact polar decomposition in non-convex matrix optimization under heavy-tailed noise. Our analysis builds on a unified framework for inexact polar decomposition that captures practical iterative approximations such as Newton-Schulz and quantifies how their errors propagate through the optimization dynamics. Under this framework, we establish an optimal iteration and sample complexity of

O \left(\varepsilon^{\frac{-(3α-2)}{(α-1)}} \right)

for finding an

\varepsilon

-stationary point, where

α\in(1,2]

denotes the heavy-tail index. For the inexact-polar setting with

σ_1=0

, we also provide guarantees that do not require prior knowledge of

α

. We analyze a randomized low-rank polar decomposition that is substantially more efficient than full-space methods while remaining compatible with our theory. Numerical experiments further demonstrate the effectiveness of the proposed inexact and randomized variants.

Sayantan Choudhury, Xiaoran Cheng, Martin Takáč +2

MBZUAI · Penn State University · Georgia Institute of Technology +1

Stochastic Optimization Momentum Methods Gradient Descent Nonconvex Stochastic Optimization Muon Optimizer Matrix Optimization+1

May 7, 2026cs.LG

A Rod Flow Model for Adam at the Edge of Stability

Cohen et al. (arXiv:2207.14484) observed that adaptive gradient methods such as Adam operate at the edge of stability. While there has been significant work on continuous-time modeling of gradient descent at the edge of stability, extending these models to momentum methods remains underdeveloped. In the gradient descent setting, Regis et al. (arXiv:2602.01480) introduced rod flow, which models consecutive iterates as an extended one-dimensional object -- a "rod." Here we extend rod flow to Adam by working in the joint phase space of parameters and first moment

(w, m)

and treating the second moment

ν

as a smooth auxiliary variable. We also develop rod flows for heavy ball momentum, Nesterov momentum, and scalar and per-component versions of RMSProp, Adam, and NAdam. For all eight optimizers, we empirically evaluate rod flow on representative machine learning architectures, where it tracks the discrete iterates through the edge-of-stability regime significantly more accurately than the corresponding stable flow.

Eric Regis, Sinho Chewi

Yale University

Edge of Stability Momentum Methods Adaptive Gradient Methods Gradient Descent Dynamics