Sign-Based Optimizers Are Effective Under Heavy-Tailed Noise

AI summary

Overview

Research area: Machine learning optimization theory, specifically the convergence analysis of sign-based stochastic optimizers (SignSGD, Lion, Muon, Muonlight) under heavy-tailed gradient noise, with motivation drawn from large language model (LLM) pretraining.

Technical level: Advanced. The paper is a theoretical optimization paper built on non-Euclidean martingale concentration inequalities, generalized smoothness assumptions, and coordinate-wise/matrix-wise moment conditions, supplemented by LLM pretraining experiments.

Scope: The paper provides new convergence guarantees for vector and matrix sign-descent optimizers under a newly proposed generalized heavy-tailed noise model, and validates that model empirically on GPT2 pretraining.

What This Paper Is About

Sign-based optimizers such as Lion and Muon have empirically outperformed AdamW in training large language models, but there has been no theory explaining why. The authors argue that the answer lies in heavy-tailed gradient noise—a regime, commonly observed in language modeling, where only a finite p-th moment (with p ∈ (1, 2]) of the noise exists rather than a finite variance. The goal is to prove convergence guarantees for sign-based methods under this more realistic noise model, thereby justifying their advantage over variance-adapted methods like Adam and AdaGrad.

Key Contributions

1. A generalized heavy-tailed noise framework for both vector and matrix variables, in which the noise level is allowed to grow with the gradient norm (rather than being a fixed constant), a condition the authors show is strictly weaker than standard assumptions and is supported by LLM pretraining evidence.

2. Sharp convergence rates of O(T^{-(p-1)/(3p-2)}) for SignSGD, Lion, Muon, and Muonlight, derived under generalized smooth function classes. These rates match or surpass previous best-known bounds, including dimension-free improvements over prior work.

3. A characterization of the sign operator's robustness mechanism, offering a theoretical justification for why sign-based optimizers outperform SGD, Adam, and AdaGrad in LLM training, and why their advantage diminishes in computer-vision tasks with more concentrated gradient noise.

4. New vector and matrix martingale concentration inequalities in non-Euclidean norms, presented as the key technical tools enabling the analysis and noted as potentially of independent interest.

Main Findings

• SignSGD converges under heavy tails: Under the paper's assumptions, SignSGD achieves (1/T) Σ E[‖∇f(x_t)‖₁] ≤ O((Δ_f ‖l₀‖₁)^{(p-1)/(3p-2)} ‖σ₀‖₁^{p/(3p-2)} (BT)^{-(p-1)/(3p-2)}), implying a complexity of O(ε^{-(3p-2)/(p-1)}) to reach an ε-stationary point.

• Optimality of the rate: Because ‖∇f(x)‖₁ ≥ ‖∇f(x)‖₂, this complexity matches the ℓ₂-norm lower bound in Liu and Zhou (2025, Theorem 3.3), and recovers the known O(ε^{-4}) lower bound (Arjevani et al., 2023) when p = 2.

• Dimension-free results: Unlike prior work (Jiang et al., 2025b; Kornilov et al., 2025), the result does not explicitly depend on dimension d and holds for any l₁ ≥ 0, σ₁ ≥ 0 and p ∈ (1, 2].

• Improvement over Kornilov et al. (2025): When the curvature vector l₀ is sparse and axis-aligned and σ₀ is sparse, Theorem 1 achieves a complexity improvement of d^{(3p-2)/(2(p-1))}; even in Kornilov et al.'s best case, their bound carries an extra d^{(2-p)/(2(p-1))} factor. The geometry quantities are φ_l := ‖l₀‖₁/l₀ ∈ [1, d] and φ_σ := ‖σ₀‖₁/‖σ₀‖_p ∈ [1, d^{1-1/p}].

• First guarantee for Lion under heavy tails: Lion matches SignSGD's rate, since Lion reduces to SignSGD when β₁ = β₂. The paper states this is the first rigorous theoretical guarantee for Lion under heavy-tailed noise, and that its O(ε^{-(3p-2)/(p-1)}) complexity is optimal.

• Strict improvement over the finite-variance Lion analysis of Jiang and Zhang (2025): That work gives O(√d T^{-1/4}); the new result removes the explicit dimensional dependence, relaxes the requirement β₁ ∈ [1-√(1-β₂), 1-(1-β₂)²] to β₁ ∈ [1-(1-β₂)^{(p-1)/p}, 1], relaxes the weight-decay condition λ ≤ 1/(2ηT) to λ ≤ (1-2^{-1/T})/η, and replaces the initialization condition ‖x₁‖∞ ≤ η with the more general ‖x₁‖∞ ≤ 1/(3λ). It also proves a uniform bound ‖x_t‖∞ ≤ 2/(3λ) for all t, tighter than the time-growing ‖x_t‖∞ ≤ ηt bound in Jiang and Zhang (2025).

• First rigorous analysis of matrix optimization under heavy-tailed stochasticity: Theorem 3 gives Muon the rate (1/T) Σ E[‖∇f(X_t)‖*] ≤ O((Δ_f ‖L₀‖)^{(p-1)/(3p-2)} ‖V₀‖_^{p/(3p-2)} (BT)^{-(p-1)/(3p-2)}), free of explicit dimensional dependence—described as an advance over existing Muon analyses (Li and Hong, 2025; Shen et al., 2025; Chang et al., 2025; Huang et al., 2025), which the paper says are restricted to a simpler regime.

• New assumption classes: Assumption 2a is a coordinate-wise generalized smoothness condition (degenerating to standard coordinate-wise smoothness when l₁ = 0), and Assumption 2b is described as the first formulation of a matrix generalized smoothness model. When L₁ = 0, Assumption 2b recovers classical matrix L-smoothness as in An et al. (2025, Assumption 2) and Kovalev and Borodich (2025, Assumption 1).

• The standard heavy-tailed model can fail: The authors demonstrate that the classic noise model with σ₁ = 0 may fail even on linear regression tasks.

• Empirical validation on GPT2 pretraining: The paper reports that LLM pretraining experiments on GPT2 confirm the theoretical insights and that the proposed noise models align with practice. The specific configurations, model sizes, datasets, and measured numbers are not reported in the portion of the paper available here.

Methodology in Plain English

The authors start from an empirical observation: in language modeling, stochastic gradients are not just noisy—their noise follows a heavy-tailed distribution (linked to Zipf's Law), so only a finite p-th moment exists for some p ∈ (1, 2], rather than a finite variance. Under this regime, SGD can diverge and standard convergence theory breaks down, while Adam and AdaGrad have no rigorous convergence theory and have been shown to potentially converge worse.

To analyze sign-based methods, the authors replace the usual constant-variance noise assumption with a condition that the expected p-th power of the noise is bounded by a constant term plus a term proportional to the p-th power of the coordinate-wise gradient. This lets the allowed noise scale with the gradient, which is strictly weaker than assuming the gradient itself has a p-th moment. They pair this with "generalized smoothness" assumptions—ones where the local curvature can itself depend on the gradient magnitude—stated separately for vectors and matrices.

The analysis then tracks the optimization progress using weighted norms tailored to sign updates: an l-weighted vector norm for SignSGD and Lion, and an L-weighted matrix norm for Muon and Muonlight. The central technical hurdle is controlling the martingale noise terms, which the authors handle with newly derived vector and matrix martingale concentration inequalities in non-Euclidean norms. The final step combines these bounds with a careful treatment of the momentum buffers, the matrix sign (computed via Newton–Schulz iteration, assumed exact), and—for Lion—the weight-decay dynamics, to produce step-size and momentum schedules that yield the stated rates.

Why This Matters

Impact on research. The paper supplies a theory-level explanation for an empirical phenomenon that has driven a shift in LLM training practice: sign-based optimizers beating AdamW. It also extends heavy-tailed analysis from vectors to matrices for the first time, and introduces noise and smoothness assumptions that the authors show can be validated empirically. The claim that adaptive methods like Adam and AdaGrad may converge poorly under heavy-tailed noise (Chezhegov et al., 2025) contrasts sharply with these new sign-based guarantees, framing a clear theoretical distinction between the two families.

Real-world applications:

• Large language model pretraining, where heavy-tailed gradients are prevalent and Muonlight has already been used at massive scale.
• LLM post-training/fine-tuning, cited as another area where sign-based optimizers have shown speedups.
• Memory- and communication-constrained training, since sign updates (discussed alongside zeroth-order and error-feedback variants) reduce the information exchanged per step.
• Robust stochastic optimization generally, for any setting where gradient noise is heavy-tailed—financial modeling, bandit-style problems, and other domains where outlier-heavy noise degrades variance-based methods.

Industry relevance. Choosing an optimizer directly affects the compute cost and wall-clock time of training runs that can cost enormous resources. A theoretical argument that sign-based methods are naturally suited to the noise profile of language modeling provides a principled basis for adopting them, and the paper's dimension-free bounds and explicit hyperparameter schedules (batch size, momentum, learning rate, weight decay) are the kind of guidance practitioners need to set up these optimizers.

Future Directions

• Extending the analysis to Muonlight in full detail. The abstract claims analysis of both Muon and Muonlight; the available text establishes Theorem 3 for Muon and describes Muonlight as incorporating Nesterov momentum and learning rate alignment, but the corresponding convergence statement is not present in the content available here.
• Closing the gap between the matrix bound and the matrix lower bound. The paper notes the absence of explicit dimensional dependence in Theorem 3 as an advance; whether the matrix rate is optimal in the same tight sense as the vector rate (matching Liu and Zhou's ℓ₂ lower bound) is not resolved in the material provided.
• Relaxing the exact-matrix-sign assumption. The analysis assumes zero numerical error in the Newton–Schulz iteration; the authors say this is for simplicity rather than necessity and defer details, leaving the effect of finite-iteration error open.
• Connecting the noise model more tightly to measured LLM gradient statistics. The generalized condition allows noise to grow with gradient norm; quantifying those growth constants across model scales, architectures, and training stages would test how broadly the framework applies.

Target Audience

This paper is best suited to optimization theorists and machine learning researchers working on stochastic optimization, convergence analysis, and heavy-tailed statistical models, as well as to LLM training practitioners and infrastructure engineers evaluating optimizer choices. Readers need comfort with convergence rates, moment conditions, norm-weighted analysis, and martingale arguments to follow the technical content; those primarily interested in practical recipes can focus on the algorithm boxes and the hyperparameter schedules in Theorems 1–3.

Authors’ abstract

While adaptive gradient methods are the workhorse of modern machine learning, sign-based optimization algorithms such as Lion and Muon have recently demonstrated superior empirical performance over AdamW in training large language models (LLM). However, a theoretical understanding of why sign-based updates outperform variance-adapted methods remains elusive. In this paper, we aim to bridge the gap between theory and practice through the lens of heavy-tailed gradient noise, a phenomenon frequently observed in language modeling tasks. Theoretically, we introduce a novel generalized heavy-tailed noise condition that captures the behavior of LLMs more accurately than standard finite variance assumptions. Under this noise model, we establish sharp convergence rates of SignSGD and Lion for generalized smooth function classes, matching or surpassing previous best-known bounds. Furthermore, we extend our analysis to Muon and Muonlight, providing what is, to our knowledge, the first rigorous analysis of matrix optimization under heavy-tailed stochasticity. These results offer a strong theoretical justification for the empirical superiority of sign-based optimizers, showcasing that they are naturally suited to handle the noisy gradients associated with heavy tails. Empirically, LLM pretraining experiments validate our theoretical insights and confirm that our proposed noise models are well-aligned with practice.