PaperlayerView exact revision on arXiv

Source-linked AI summary

Sign-Based Optimizers Are Effective Under Heavy-Tailed Noise

Dingzhi Yu, Hongyi Tao, Yuanyu Wan, Luo Luo, Lijun Zhang

arXiv:2602.07425v2cs.LGcs.CLmath.OC

TL;DR

The paper asks why sign-based optimizers can outperform AdamW when LLM gradients are heavy-tailed. It introduces generalized heavy-tailed noise conditions and convergence analyses for vector and matrix sign methods, finding sharp guarantees and empirical agreement with LLM pretraining behavior.

  • Problem

    The paper addresses the limited theoretical understanding of why Lion and Muon outperform AdamW for LLM training under heavy-tailed gradient noise.

  • Method

    The paper develops generalized heavy-tailed noise models and convergence analyses for SignSGD, Lion, Muon, and Muonlight under generalized smoothness.

  • Results

    The analyses establish convergence rates matching or surpassing previous best-known bounds and extend rigorous heavy-tailed-stochasticity analysis to matrix optimization.

  • Takeaways & Limitations

    The results provide theoretical justification for sign-based optimizers as a natural fit for noisy, heavy-tailed gradients in LLM training.

  • Takeaways & Limitations

    The matrix noise model extends prior formulations and is empirically validated, but its stated scope is the matrix optimization setting under the proposed assumptions.

Abstract

from arXiv · show

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.

1. Introduction

The paper studies why sign-based optimizers outperform AdamW in LLM training by analyzing heavy-tailed stochastic gradients. It develops generalized convergence theory for vector and matrix sign methods and tests whether the proposed noise model reflects practice.

  • Heavy-tailed stochastic gradients are widely observed in language modeling, challenging the finite-variance assumptions common in optimization theory.The paper models noise with a finite pth moment for p ∈(1, 2].
  • New convergence theories cover SignSGD, Lion, Muon, and Muonlight under generalized smoothness.The reported rates match or surpass previous state-of-the-art bounds.
  • The paper introduces a generalized heavy-tailed noise framework for vector and matrix optimization.The framework allows noise magnitude to grow with the gradient norm and is described as weaker than existing assumptions.
  • The analysis characterizes sign operations as robustifiers against heavy-tailed stochasticity, supporting their theoretical advantage over SGD and AdamW.The paper connects this mechanism to element-wise sign updates and matrix orthogonalization.
  • LLM pretraining experiments on GPT2 provide empirical evidence that the proposed noise model reflects heavy-tailed stochasticity in practice.The experiments are presented as validation of the theoretical findings.

2. Related Work

Related work establishes heavy-tailed noise as a setting where clipping, normalization, and adaptive methods have been studied, while sign-based methods span vector and matrix optimization. The paper positions its analysis relative to these strands.

  • Heavy-tailed noise: Prior theory studies heavy-tailed noise through finite pth-moment assumptions and analyzes clipped SGD and related stochastic behavior.The cited literature also examines how heavy-tailedness may explain optimizer and architecture performance.
  • Sign-based methods: vector optimizers: SignSGD established sign-based vector optimization, with later work targeting communication efficiency, noise robustness, and improved convergence.Lion extended this line with two momentum buffers and achieved empirical success.
  • Sign-based methods: matrix optimizers: Muon treats matrix optimization as matrix sign descent implemented through Newton–Schulz iteration.Muonlight adds Nesterov momentum and learning-rate alignment to Muon.
  • How to handle heavy-tailed stochasticity?: Existing work shows that clipping or normalization can address heavy-tailed noise, while Adam and AdaGrad may converge poorly under such noise.Clipped variants are reported to mitigate this issue.

3. Vector Sign Optimizer: SignSGD and Lion

This section develops convergence guarantees for SignSGD and Lion under generalized heavy-tailed noise, using generalized smoothness assumptions. The results show that sign-based vector optimizers remain effective across broader noise and objective classes while improving several prior guarantees.

  • 3.1. Notations and Assumptions: The analysis studies vector optimization under lower-bounded, generalized smooth, and unbiased stochastic-gradient assumptions.The generalized smoothness model includes coordinate-wise smoothness as a special case when l1 = 0.
  • 3.1. Notations and Assumptions: The new (σ0, σ1)-heavy-tailed noise condition allows noise levels to grow with the gradient norm and is supported by LLM pretraining evidence.It generalizes the classic coordinate-wise model, which sets σ1 = 0, and is intended to better reflect observed stochasticity.
  • 3.2. Convergence Theory for SignSGD: SignSGD updates the model with the sign of a momentum buffer and remains effective under heavy-tailed gradient noise.The algorithm maintains a momentum buffer from batched stochastic gradients before applying the signed update.
  • 3.2. Convergence Theory for SignSGD: SignSGD’s convergence guarantee is dimension-independent and applies for any l1 ≥ 0, σ1 ≥ 0, and p ∈ (1, 2].The result also matches the ℓ2-norm lower bound and recovers the O(ϵ^-4) lower bound when p = 2.
  • 3.3. Convergence Theory for Lion: Lion extends SignSGD with separate momentum roles and decoupled weight decay while retaining the same convergence rate as SignSGD.Lion uses different β values for stored and update momentum, and the authors provide a rigorous heavy-tailed convergence guarantee.
  • 3.3. Convergence Theory for Lion: Compared with prior Lion analysis, the result removes explicit dimensional dependence and relaxes requirements on momentum, weight decay, initialization, and iterate bounds.The stated iterate bound is uniform in time, unlike the earlier bound that grows over time.

4. Matrix Sign Optimizer: Muon and Muonlight

The paper develops Muon and Muonlight for matrix optimization under generalized heavy-tailed noise, establishing convergence guarantees that cover p∈(1,2] and improve dimension-dependent bounds. Muonlight matches Muon’s convergence bounds while also providing trajectory stability under an initialization condition.

  • 4.1. Notations and Assumptions: The proposed matrix heavy-tailed noise model allows p∈(1,2] and permits the pth noise moment to grow with the gradient.When V1=0 and p=2, it recovers an adaptive variance assumption; its matrix formulation is intended to capture structural noise properties.
  • 4. Matrix Sign Optimizer: Muon and Muonlight: Muon uses orthogonalized momentum and the matrix sign operator, with convergence analyzed under matrix generalized smoothness and generalized heavy-tailed noise.The matrix sign operator is defined from the SVD, and the framework assumes a lower-bounded objective, matrix smoothness, unbiasedness, and matrix heavy-tailed noise.
  • 4.2. Convergence Theory for Muon: Theorem 3 provides the first Muon convergence guarantee for the general heavy-tailed regime and yields a rate without explicit dimensional dependence.The result also gives a sharper p=2 bound than previous dimension-dependent bounds, using the nuclear norm rather than dimension-dependent factors.
  • 4.3. Convergence Theory for Muonlight: Muonlight adds Nesterov momentum and decoupled weight decay to Muon, while supporting general momentum parameters (β1, β2).These modifications are motivated by variance reduction and training-trajectory stabilization, respectively.
  • 4.3. Convergence Theory for Muonlight: Muonlight achieves the same convergence bounds as Muon under generalized heavy-tailed noise.Under ∥X1∥op≤1/(3λ), its iterates satisfy ∥Xt∥op≤2/(3λ) for all t∈[T].

5. How Sign Operator Works

The paper explains sign-based optimization through non-Euclidean gradient normalization and develops concentration tools tailored to the resulting ℓ1 and nuclear-norm geometry. These tools control heavy-tailed noise terms and yield sublinear convergence analysis for vector and matrix methods.

  • 5. How Sign Operator Works: Applying Hölder’s inequality and von Bahr–Esseen-type concentration to the cumulative noise term yields the optimal sublinear rate under the heavy-tailed assumptions.The cumulative curvature and initial-noise terms admit straightforward bounds, leaving cumulative noise as the main analytical challenge.
  • 5. How Sign Operator Works: The sign operator performs coordinate-wise normalization, making signed-gradient methods normalized steepest descent with respect to the ℓ∞ norm.This differs geometrically from normalized gradient methods based on the ℓ2 norm.
  • 5. How Sign Operator Works: Signed-gradient analysis shifts both the convergence criterion and error term from the ℓ2 norm to the ℓ1 norm.This non-Euclidean geometry changes the corresponding analysis.
  • 5. How Sign Operator Works: Standard Euclidean concentration inequalities do not directly control the critical non-Euclidean noise terms, motivating new concentration techniques.The analysis identifies von Bahr–Esseen and Hölder-based approaches as inadequate or cumbersome for ℓ1-norm control.
  • 5.2. New Martingale Concentration Inequalities: Lemma 3 bounds the ℓ1 norm of a vector martingale through coordinate-wise variance, producing a term directly controllable under Assumption 4a.The proof uses regret analysis for diagonal AdaGrad and is applied with gk=β^(t−k)nk.
  • 5.2. New Martingale Concentration Inequalities: For matrix methods, the paper introduces nuclear-norm concentration and a one-sided Shampoo regret analysis to avoid dimension-dependent factors.This enables the term cancellations needed to handle heavy-tailed matrix noise.

6. Empirical Study

The empirical study tests whether the proposed heavy-tailed noise models reflect LLM training and evaluates sign-based optimizers in nanoGPT pretraining on C4. The experiments report faster training for Lion and Muon than AdamW.

  • Noise-model validation: The proposed generalized heavy-tailed noise model closely reflects stochasticity observed in LLM training.Figures 1 and 2 provide the supporting empirical evidence.
  • Optimizer efficiency: 1.07× speedup over AdamW is achieved by Lion in nanoGPT pretraining on C4.Figure 3 compares training and validation behavior alongside zoom-in speedup results.
  • Optimizer efficiency: 1.32× speedup over AdamW is achieved by Muon in nanoGPT pretraining on C4.The figure reports training loss, validation loss, and validation accuracy in its first three panels.

7. Conclusion

The paper concludes that generalized heavy-tailed noise assumptions support sharp convergence guarantees for sign-based optimizers and align with LLM pretraining dynamics. Its theory and experiments together support the reported practical superiority of these methods.

  • Conclusion: Generalized heavy-tailed noise assumptions allow noise magnitude to scale linearly with the gradient norm and empirically reflect LLM training dynamics.The paper presents this model as a more accurate description of real-world LLM stochasticity.
  • Conclusion: Sharp convergence guarantees for SignSGD, Lion, Muon, and Muonlight match or surpass the best-known bounds.The guarantees are established under the generalized heavy-tailed regime.
  • Conclusion: The theory indicates provable advantages of sign-based methods over NSGD and AdamW under heavy-tailed stochasticity.LLM pretraining experiments corroborate the practical superiority of sign-based optimizers.

Appendix A. Provable Complexity Improvement of Sign Gradient Descent over Normalized Gradient Descent

This appendix compares sign-based methods with normalized gradient descent through geometry-aware complexity bounds. It argues that vector and matrix sign methods can obtain provable convergence advantages, especially when gradients are dense or curvature is low-rank.

  • Comparison framework: Sign-based methods and NSGD are compared as steepest-descent or linear-minimization methods under distinct geometric constraints.Both methods incorporate momentum, and NSGD is treated as a robust heavy-tailed-noise baseline.
  • Vector optimization: Vector density functions quantify gradient geometry for comparing SignSGD and Lion with NSGD.The trajectory-wide density is defined using the minimum density along the optimization trajectory.
  • Vector optimization: R2 is bounded by 5, while dense gradients keep the dominant ratio R1 controlled in the cited empirical settings.The resulting comparison uses gradient-density properties observed in language modeling and computer vision experiments.
  • Vector optimization: When curvature is axis-aligned, SignSGD and Lion achieve a dimension-wise speedup of factor d over NSGD.This conclusion follows when the curvature density satisfies ϕ∞(l0) ≈1.
  • Matrix optimization: Matrix density functions compare Muon and Muonlight with MNSGD using the distribution of singular values.The matrix density captures spectral flatness through Schatten-norm relationships.
  • Complexity comparison: Table 1 compares complexity bounds for finding an ϵ-stationary point and presents sign-based algorithms against normalized-gradient lower bounds.The appendix frames the comparison as a provable advantage derived from problem geometry.
  • Smoothness conditions: The appendix’s smoothness discussion states that the ℓ∞-norm formulation matches sign-descent geometry, whereas alternative ℓ2-based formulations do not.The cited generalized smoothness model is also described as weaker than the compared gradient-curvature bounds.

B.3. Heavy-Tailed Noise: Assumption 4a

This appendix develops supporting lemmas for the generalized heavy-tailed noise framework and its convergence analysis. It also records variants, examples, and technical stability results used for vector sign-based optimization.

  • Noise assumptions: Assumption 4a is introduced as a relaxation of classic coordinate-wise heavy-tailed noise, with noise allowed to depend on gradient magnitude.The paper states that the model is meaningful theoretically and aligns with LLM pretraining empirically.
  • Noise assumptions: The separable-function example shows that Assumption 4a cannot be satisfied when σ1 = 0, while σ1 > 0 permits specified coordinate parameters.The construction uses centered independent coordinate noise and p-th moment bounds.
  • Matrix noise: Assumption 4c relaxes the matrix noise condition by replacing the full matrix V1 with the scalar coefficient ∥V1∥op.The paper states that this variant is more amenable to empirical verification.
  • Matrix noise: The appendix introduces matrix noise terms that preserve p-th-moment homogeneity and capture anisotropic structure through weighted norms.The framework is related to assumptions involving noise covariance and principal directions.
  • Technical tools: Weight-decay stability results require bounded initialization and step-size conditions involving the smoothness parameter and decay coefficient.The stated updates include both the sign direction and weight decay.

C.3. Proof of Theorem 1

The proof establishes convergence by combining descent, momentum-error, noise, and trajectory-curvature bounds under the stated step-size conditions. It concludes with a convergence rate expressed as O(·).

  • The proof controls each update’s movement using the sign-based step and the step-size restriction η ≤ 1/∥l1∥∞.
  • Momentum error is decomposed into initial, cumulative, and trajectory-curvature noise components before being bounded separately.
  • The resulting bounds for At, Bt, and Ct are combined and summed over iterations to obtain the main inequality.
  • The proof imposes an additional condition for one inequality and then combines it with the preceding bound.
  • The theorem concludes with a convergence rate of O(·).

C.4. Proof of Theorem 2

The proof extends the descent and error-decomposition strategy to the second theorem, using matrix norm inequalities, martingale arguments, and stability bounds. It combines these ingredients to derive the theorem’s result under the stated assumptions.

  • The proof bounds update movement through ∥xt+1 − xt∥∞ and the step-size constraint, enabling application of the descent analysis.
  • The error analysis separates initial noise, cumulative noise, and trajectory curvature, with weight decay requiring a separate treatment.
  • Assumption 4b is maintained as the operative assumption, while its implication of Assumption 4c is used in the analysis.
  • For matrix variables, the proof uses trace and matrix-norm properties, msign properties, matrix Cauchy–Schwarz, and a matrix descent lemma.
  • The proof combines the resulting bounds with stability, projection, norm-duality, and regularization arguments to complete the convergence analysis.
  • Conditional independence and unbiasedness yield martingale-difference sequences used to control mini-batch noise.

D.3. Proof of Theorem 3

The proof analyzes the matrix sign method by rewriting momentum and Newton–Schulz updates in matrix-sign form, then bounding optimization error, stochastic noise, and trajectory curvature. It concludes by selecting parameters to obtain the stated convergence rate.

  • The matrix proof defines optimization error, stochastic noise, and gradient-change terms as Et, Nt, and St.
  • Under the exact Newton–Schulz oracle assumption, the update becomes Ot = msign(Mt) because msign is invariant to input scale.
  • The error recursion separates momentum error, gradient variation, and stochastic noise before bounding their contributions.
  • The trajectory-curvature analysis uses matrix smoothness and bounds on consecutive iterate differences.
  • The bounds for At, Bt, and Ct are combined, and parameters B, β, and η are chosen according to the theorem’s conditions.
  • The resulting argument yields a convergence rate of O(·).

E.3. Experimental Setup and Methodology

The experiments verify heavy-tailed stochastic-gradient noise and assess the proposed assumptions during nanoGPT pretraining on C4. They estimate the tail index and examine noise behavior across checkpoints and gradient coordinates.

  • The nanoGPT experiments train on C4 and sample stochastic-gradient noise across mini-batches using several optimizers.
  • The experiments use PyTorch DDP across four NVIDIA Pro6000 GPUs, with 10k training steps and global batch size 512.
  • Heavy-tailed noise is assessed with coordinate-wise histograms and Q-Q plots, where Gaussian behavior corresponds to alignment with the red diagonal.
  • The heavy-tail index p is estimated with the estimator of Simsekli et al. (2019).
  • Assumptions 4a and 4b are approximated empirically, with Figure 5 plotting |∇jf|p against E[|gj − ∇jf|p] and Figure 6 using ∥∇f∥p.
  • Checkpoint analysis tracks the heavy-tail coefficient and empirical gradient-noise norms throughout pretraining.
Loading 2602.07425v2…