PaperlayerView exact revision on arXiv

Source-linked AI summary

Second-Moment Stochastic Approximation Methods

Tao Jiang, Lin Xiao

arXiv:2609.36600v1math.OCcs.AIstat.ML

TL;DR

Classical stochastic approximation methods use first-moment estimates, while many matrix-valued stochastic problems arise in deep-learning optimization. This paper derives second-moment methods through optimal preconditioning and analyzes them in two stages; the practical methods converge almost surely to a neighborhood whose size depends on estimator bias and variance.

  • Problem

    Stochastic matrix problems arise in deep-learning optimization, while conventional equation-solving methods require accurate evaluations that are unavailable when only random mappings can be accessed.

    Source: §1 · pp. 1–2

  • Method

    The paper derives second-moment stochastic approximation methods from optimal preconditioning and analyzes conceptual exact-moment methods before bounding practical methods’ deviations using moment-estimator bias and variance.

    Source: §7 · pp. 31–32

  • Results

    The practical methods converge almost surely to a neighborhood of a target satisfying the aiming condition, with neighborhood size depending on estimator bias and variance.

    Source: §7 · pp. 32–33

  • Takeaways & Limitations

    The framework unifies analysis across methods without specifying exact moment estimators, while allowing convergence arbitrarily close to the target when estimator bias and variance diminish.

    Source: §7 · pp. 32–33

    The general theorem guarantees convergence to a neighborhood rather than exactly to the target, and persistent estimator bias can prevent convergence to the target itself.

    Source: §7 · pp. 32–33

Abstract

from arXiv · show

Classical stochastic approximation methods rely on estimators of the first moment (mean) of a random regression function. We study methods that employ estimators of both the first and the second moments, which include modern deep-learning optimizers such as Adam and Muon as special cases. We derive second-moment stochastic approximation methods through the lens of optimal preconditioning for solving matrix equations, and develop a two-stage framework for their convergence analysis. The first stage focuses on the analysis of conceptual (impractical) methods that rely on the exact first and second moments. In the second stage, we replace the exact moments with their respective estimators, and invoke Dvoretzky's theorem to show that the resulting practical methods converge almost surely to a neighborhood of the target solution. The size of the neighborhood depends on the biases and variances of the first- and second-moment estimators. We derive concrete bounds for Muon and a spectral variant of Adam that determine the radius of their neighborhood of convergence.

1 Introduction

The paper extends stochastic approximation for matrix equations from first-moment estimates to methods that also use second moments, motivated by noisy optimization of deep-learning loss functions. It derives a unified preconditioning framework and a two-stage convergence analysis for these methods.

1 Introduction

  • Stochastic matrix equations arise in deep-learning training, where solving M(X) = 0 corresponds to finding a stationary point of the expected loss.The framework focuses on one matrix variable and can be applied block-wise to models with multiple matrices.Source: §1 · p. 1

1.1 Second-moment stochastic approximation

  • Second-moment methods combine estimates of the mean and second moment of the random mapping to construct preconditioned updates.The paper interprets Adam and Muon within this framework, even though the simplified Muon update does not explicitly use a second-moment estimator.Source: §1 · p. 2, §1.1 · pp. 3–4

1.2 A two-stage framework for convergence analysis

  • A two-stage analysis studies exact-moment conceptual methods first, then practical methods using estimators whose bias and variance determine their convergence-neighborhood size.Under the stated bounded-deviation condition, the practical methods converge almost surely to a neighborhood of the target; the paper also contrasts the first-moment estimator trade-offs of RMSProp and Adam.Source: §1.2 · p. 4

1.3 Outline

  • The paper outlines derivations of preconditioners and blending functions for Muon, ASGO, and Adam, followed by convergence analysis and numerical experiments with Muon and ASGO.The analysis also includes decoupled weight decay and derives estimator-dependent deviation bounds for Muon and ASGO.Source: §1.3 · p. 5

2 Second-moment preconditioning

The paper derives second-moment preconditioners for stochastic matrix equations by minimizing expected distance to a solution, then simplifies them into several optimizer families. The derivations connect the preconditioner structure and noise assumptions to the resulting update rules.

2 Second-moment preconditioning

  • Second-moment preconditioning avoids the costly Jacobian-based tensor required by a full Newton preconditioner, using smaller preconditioners instead.The paper considers elementwise, left- or right-sided fused, and two-sided matrix preconditioning.Source: §2 · p. 5

2.1 The optimal fused preconditioner and its approximation

  • The optimal fused preconditioner minimizes expected distance to a solution, yielding a convex quadratic in the preconditioner with coefficients from the first and second moments.Because the optimum depends on the unknown solution, the paper approximates it by absorbing the error magnitude into a step size; the resulting practical algorithm is simplified because its matrix operations are costly.Source: §2 · pp. 5–6, §2.1 · p. 7

2.2 Derivation of Muon and its variants

  • Muon follows from a simplification that assumes signal fractions are relatively uniform across directions, producing an update based on the orthogonal factor of the first moment.This simplification retains scaling invariance and avoids explicitly estimating the second moment; the method can use an online estimate of the first moment.Source: §2.2 · p. 8

2.3 Derivation of ASGO and its variants

  • When signal fractions are nonuniform, ASGO retains second-moment structure, with the paper stating its update can be more effective than Muon’s under that condition.A related two-sided construction uses left and right second moments with matrix powers of −1/4 to retain scaling invariance; Shampoo instantiates this approach.Source: §2.3 · pp. 9–10

2.4 Element-wise preconditioning

  • Elementwise preconditioning yields Adam and related methods, whose examples mainly differ in their first-moment estimators while sharing Adam’s second-moment estimator.The framework also allows customized, model-specific, and data-driven estimators.Source: §2.4 · p. 11

3 A general framework for convergence analysis

The framework analyzes second-moment stochastic approximation methods, first establishing convergence conditions and rates for exact-moment conceptual methods and then examining practical methods with estimated updates. It uses an extension of Dvoretzky’s theorem to show almost-sure convergence of broad classes of practical methods to a neighborhood of the target, with the neighborhood size related to estimator errors.

3.1 Aiming condition

  • The generalized aiming condition makes the blended update a strong pseudogradient, while bounded second-moment updates make the growth condition automatic for Muon, ASGO, and Adam variants.Unlike the classical first-moment condition, it is not simply connected to convexity; when the update is bounded, global aiming requires λ > 0.Source: §3.1 · p. 13

3.2 Analysis of conceptual algorithms

  • Under Assumption A and the stated step-size conditions, the conceptual method’s iterates converge to X*.The analysis uses a one-step improvement bound and a deterministic convergence lemma.Source: §3.2 · p. 15
  • O(1/t) squared-distance convergence holds for α_t = α/(t+1) under the theorem’s step-size bounds, including a lower bound on α.For methods with L = 0, the stated condition reduces to λ < 2µ, and the upper bound becomes α ≤ 1/λ.Source: §3.2 · p. 16
  • O(1/t^p) squared-distance convergence holds with α_t = α/(t+1)^p for p in (1/2,1), removing the lower bound on α required for the O(1/t) result.The step-size constant still has the theorem’s stated upper bound.Source: §3.2 · p. 17

3.3 Convergence of practical algorithms

  • With bounded practical-update deviation and the theorem’s other conditions, practical iterates converge almost surely to a neighborhood of X* whose radius is governed by µ and the deviation bound.The deviation depends on bias and variance of both moment estimators and their numerical conditioning through the blending map; persistent estimator bias can prevent convergence to X* itself, and the proof invokes an extension of Dvoretzky’s theorem [26].Source: §3.3 · pp. 18–19
  • Theorem 3.5 combines prior extensions of Dvoretzky’s theorem and permits time-varying coefficients that need not converge to zero.The paper also provides a self-contained proof in Appendix A.Source: §3.3 · p. 19
  • Practical second-moment stochastic approximation methods converge almost surely to a neighborhood of the target, with its radius governed by the error bound ε̄ and μ.Theorem 3.4 establishes this guarantee for a broad class of methods.Source: §3.3 · p. 22
  • The error bound ε_t depends on the blending function’s sensitivity to perturbations and on biases and variances of both moment estimators.The paper next derives concrete bounds for Muon and ASGO.Source: §3.3 · p. 22

4 Deviation bound for Muon

The Muon analysis bounds the deviation between practical and conceptual updates in terms of estimator error and matrix conditioning. This deviation determines the size of the neighborhood around the target where the algorithm converges.

4 Deviation bound for Muon

  • The analysis assumes the estimator’s expected matrix has the same full row rank as the true first moment, with rank m almost surely.Under this assumption and finite estimator variance, Theorem 4.1 gives the deviation bound.Source: §4 · p. 23
  • Muon’s deviation bound accounts for estimator bias and variance and the conditioning of both the estimated and true first moments; its asymptotic bound controls the convergence-neighborhood radius.The conditioning terms involve the smallest positive singular values of the relevant matrices.Source: §4 · p. 23
  • Orthogonal-factor stability depends on input perturbation and smallest singular values: well-conditioned matrices yield stable factors, while near-rank deficiency amplifies errors.The stated lemma covers full-row-rank wide matrices and unitarily invariant norms.Source: §4 · p. 23
  • The deviation is decomposed into a variance-induced error from orthogonalizing the estimator and a perturbation from its bias relative to the true first moment.The two terms are bounded separately.Source: §4 · p. 24
  • For Muon, the growth condition holds with L = 0 and C = √m, and the bounded-variance condition holds with b = m.The argument uses the bound on the Frobenius norm of Orth(M(X)).Source: §4 · p. 25

5 Deviation bound for ASGO

The ASGO analysis bounds the practical update’s deviation from its conceptual counterpart under positive-definiteness conditions on the second moment and its estimator. The bound depends on estimator bias, variance, and numerical conditioning.

5 Deviation bound for ASGO

  • The second-moment matrix and its estimator must be symmetric positive definite, with uniform positive lower bounds on their eigenvalues.The theorem’s bound also assumes finite estimator variance.Source: §5 · pp. 25, 27
  • ASGO’s deviation bound is governed by the bias and variance of both moment estimators and by their numerical conditioning.The bound concerns the blending function W^-1/2M and the conditioning of its matrix square-root inverse.Source: §5 · pp. 25–26
  • The eigenvalue lower-bound assumption is compatible with practical algorithms that add a small offset ϵI to their second-moment estimators.Under the stated assumptions, the growth condition holds with L = 0 and C = √m.Source: §5 · p. 28

6 Numerical experiments

Experiments on a 1.81-billion-parameter LLaMA-style model evaluate the aiming condition for ASGO and Muon, then compare their losses across smoothing factors. The aiming condition holds empirically for most of training, while the relative loss performance of the methods changes with the smoothing factor.

6 Numerical experiments

  • The experiments train a 1.81-billion-parameter LLaMA-style model on FineWeb-Edu for 17.5K steps across 32 NVIDIA H100 GPUs.Training uses 2048K-token batches and 20 tokens per parameter; Muon and ASGO optimize matrix parameters, while AdamW optimizes non-matrix parameters.Source: §6 · pp. 28–29

6.1 Empirical evidence of the aiming condition

  • The aiming condition for ASGO’s update direction holds empirically throughout training for the whole parameter group and most sub-blocks.For Muon, the update direction is positively aligned with Xt − Xfinal except for a short period at the beginning; this alignment persists when restricted to parameter sub-blocks.Source: §6.1 · pp. 29–30

6.2 Effect of the smoothing factor

  • At β1 ≤ 0.3, ASGO achieves lower final loss, whereas at β1 > 0.5, Muon outperforms ASGO.The comparison follows sweeps of learning rate, effective weight decay, and β2; the best β2 is 0.95 for every tested β1.Source: §6.2 · pp. 30–31
  • For both Muon and ASGO, training loss decreases monotonically as β1 increases, with marginal gains diminishing as β1 approaches 1.The reported trend is consistent with increasing β1 mainly reducing variance in the first-moment estimator.Source: §6.2 · p. 31

7 Conclusion and discussion

The paper unifies second-moment stochastic approximation methods through optimal preconditioning and develops a two-stage convergence analysis. It highlights scaling invariance and almost-sure convergence guarantees, while noting that practical methods may converge only to a neighborhood of the target.

7 Conclusion and discussion

  • A common optimal-preconditioning framework derives algorithms including Adam, Muon, and their variants from different matrix products and simplifications.The framework interprets the methods as simplifications of a shared preconditioner under different assumptions about noise structure.Source: §7 · pp. 31–32
  • The two-stage analysis studies exact-moment conceptual methods first, then practical methods through bounds on their deviation from the conceptual updates.The deviation bounds depend on the bias and variance of the first- and second-moment estimators without fixing their forms.Source: §7 · p. 32
  • Second-moment methods inherit scaling invariance and can work when the expected regression function is non-smooth or discontinuous.The paper presents this as an advantage over first-moment methods.Source: §7 · p. 32
  • The statistical preconditioning perspective supports almost-sure convergence analysis using classical stochastic-approximation results.The paper identifies this guarantee as important for large-scale, costly deep-learning training.Source: §7 · p. 32
  • The practical-method guarantee is convergence almost surely to a neighborhood of a point satisfying the aiming condition, not necessarily to that point itself.Diminishing estimator bias and variance can make the neighborhood arbitrarily small; persistent bias can prevent exact convergence.Source: §7 · pp. 32–33
  • Diminishing regularization is proposed as a strategy that can be strong far from the target and unnecessary near it under the stated aiming condition.The discussion says this strategy has been effective for Adam and Muon.Source: §7 · p. 33
  • Combining second-moment preconditioners with Jacobian-based preconditioners is identified as a promising direction for addressing randomness and curvature together.The discussion also points to variational-inequality algorithms as possible routes to faster convergence rates with second-moment preconditioners.Source: §7 · p. 33

A Proof of the extended Dvoretzky’s theorem

The proof extends a deterministic recursion lemma and applies it to a rescaled squared distance. It converts a martingale cross term into a scalar martingale difference before invoking the lemma.

A Proof of the extended Dvoretzky’s theorem

  • The proof builds on Derman and Sacks’ approach by extending their key lemma for deterministic recursions of multiple sequences.The appendix states the lemma for nonnegative sequences and applies it in the convergence proof.Source: A Proof of the extended Dvoretzky’s… · pp. 33–34
  • The proof converts a vector martingale cross term into a scalar martingale difference, then applies the extended lemma to a rescaled squared distance.The coefficient of the martingale term is measurable with respect to the filtration and has Frobenius norm at most one.Source: A Proof of the extended Dvoretzky’s… · p. 34
  • After fixing a probability-one sample point, the argument uses a positive multiplicative sequence and the lemma to establish the theorem.The proof constructs the sequence after the recursion coefficients satisfy eventual bounds, then derives the conclusion from its limiting behavior.Source: A Proof of the extended Dvoretzky’s… · pp. 34–35

B An example with persistent estimator bias

A one-dimensional deterministic example shows that persistent bias in the practical update can satisfy the theorem’s assumptions yet prevent convergence to the target. The iterates instead converge to a nonzero offset.

B An example with persistent estimator bias

C Perturbation bound for matrix square-root inverse

The appendix bounds perturbations of the matrix square-root inverse using its integral representation. It applies the resolvent identity, norm inequalities, and lower bounds on the matrices’ eigenvalues.

C Perturbation bound for matrix square-root inverse

  • The proof begins from an integral representation of the matrix square-root inverse function.This representation is cited as the basis for the perturbation argument.Source: C Perturbation bound for matrix… · p. 36
  • For symmetric positive matrices, the proof uses the resolvent identity and sub-multiplicativity of matrix norms to bound the difference.The passage introduces the argument for two such matrices A and B.Source: C Perturbation bound for matrix… · p. 36
  • Lower bounds A ⪰ λ_A I and B ⪰ λ_B I bound the inverse norms, which are substituted into the integral estimate.The proof then evaluates the integral to obtain the perturbation bound.Source: C Perturbation bound for matrix… · p. 36

Related papers

Loading 2609.36600v1…