Paper deep dive
Second-Moment Stochastic Approximation Methods
Tao Jiang, Lin Xiao
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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.
Tags
Links
- Source: https://arxiv.org/abs/2609.36600v1
- Canonical: https://arxiv.org/abs/2609.36600v1
Trouble viewing inline? Open PDF directly →
Full Text
135,936 characters extracted from source content.
Expand or collapse full text
Second-Moment Stochastic Approximation Methods Tao Jiang* Lin Xiao Note: Fundamental AI Research (FAIR), Meta Superintelligence Labs (taojiang@meta.com, linx@meta.com). September 28, 2026 Abstract 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 Consider the problem of solving a matrix equation M(X)=0M(X)=0 where M is a mapping from m×nR^m× n (space of m by n matrices) to itself. We focus on the stochastic setting M(X)=ξ[G(X,ξ)],M(X)=E_ξ[G(X,ξ)], (1) where G is a random regression function and ξ is a random variable with some unknown distribution. This is equivalent to solving a vector equation ℳ(x)=0M(x)=0 where x∈mnx ^mn is a vector of length mnmn and ℳ:mn→mnM:R^mn ^mn is the vectorization of M. However, the matrix form arises naturally from many applications and allows us to use matrix factorization methods such as the singular value decomposition (SVD) to construct better preconditioners, which is the focus of this paper. Our motivation comes from studying algorithms for training large-scale deep learning models (1, 2, e.g.,). In such applications, one need to minimize a loss function of the form f(X)=ξ[F(X,ξ)]f(X)=E_ξ[F(X,ξ)] where X denotes the model parameters, ξ represents random samples from a dataset, and F(⋅,ξ):m×n→F(·,ξ):R^m× n is the per-sample loss function. In general, the model parameters consist of a large set of matrices of various sizes. To simplify presentation, we focus on the setting of a single matrix variable, but the methods we study can be applied block-wise to every matrix in a more complex model. In this context, we have G(X,ξ)=∇F(X,ξ)G(X,ξ)=∇ F(X,ξ) and M(X)=∇f(X)M(X)=∇ f(X), thus solving the matrix equation M(X)=0M(X)=0 corresponds to finding a stationary point of f. There is a large literature on solving systems of nonlinear equations (3, 4, e.g.,), and most methods require accurate evaluation of M(⋅)M(·) and for Newton’s method also its Jacobian [5]. None of them can be applied to the stochastic setting where we only have access to the random map G(⋅,ξ)G(·,ξ). In this setting, we resort to the stochastic approximation (SA) method of Robbins and Monro [6]: Xt+1=Xt−αtG(Xt,ξt),X_t+1=X_t- _tG(X_t, _t), (2) where ξt _t is a realization of the random variable ξ at iteration t and αt>0 _t>0 is an appropriately chosen step size. Since G(Xt,ξt)G(X_t, _t) is a random estimator of M(Xt)M(X_t), we can interpret (2) as a practical variant of the conceptual method Xt+1=Xt−αtM(Xt)X_t+1=X_t- _tM(X_t). In general, we can employ a more sophisticated estimator of M(Xt)M(X_t), e.g., the exponential moving average (EMA) of G(Xt,ξt)G(X_t, _t), M^t=βtM^t−1+(1−βt)G(Xt,ξt), M_t= _t M_t-1+(1- _t)G(X_t, _t), (3) where βt∈[0,1) _t∈[0,1) is a smoothing factor. This leads to the SA method Xt+1=Xt−αtM^t.X_t+1=X_t- _t M_t. The use of EMA, also known as the momentum, was analyzed for stochastic convex optimization as early as in [7], and became very popular in deep learning following the work of Sutskever et al. [8]. In the rest of this paper, we use M^t M_t to denote any estimator of M(Xt)M(X_t), including G(Xt,ξt)G(X_t, _t) itself and its EMA as special cases. Since M^t M_t is an estimate of the first moment M(Xt)M(X_t), we call such algorithms first-moment SA methods. In this paper, we study stochastic approximation methods that also make use of the second moment of the random mapping G(X,ξ)G(X,ξ). The rationale behind using the second moment in SA parallels that of using second-order derivatives in optimization. In deterministic optimization, we use second-order derivatives (Hessian of the loss function, corresponding to Jacobian of M in our context) to construct preconditioners to combat anisotropic curvature. In stochastic approximation, the presence of randomness/noise renders the second-order derivatives hard to evaluate and less effective. On the other hand, the second moment of G(X,ξ)G(X,ξ) carries statistical information that can be exploited to better combat randomness/noise. Specifically, we can use online estimators of the second moment to construct preconditioned SA methods. 1.1 Second-moment stochastic approximation There are several different ways to define the second moment for the matrix-valued map G(X,ξ)G(X,ξ), depending on the type of matrix product used. The three apparent choices are ξ[G(X,ξ)⊙G(X,ξ)] _ξ [G(X,ξ) G(X,ξ) ] ∈m×n, ^m× n, (4a) ξ[G(X,ξ)G(X,ξ)T] _ξ [G(X,ξ)G(X,ξ)^T ] ∈m×m, ^m× m, (4b) orξ[G(X,ξ)TG(X,ξ)] or _ξ [G(X,ξ)^TG(X,ξ) ] ∈n×n, ^n× n, (4c) where the first one uses the Hadamard (element-wise) product, denoted as ⊙ , and the second and third definitions use matrix multiplication in different orders. For convenience, we use W(X)W(X) to denote the second moment of G(X,ξ)G(X,ξ) regardless of the specific definition, which should be self-evident from the context. To cope with the ambiguity in the size of W under different definitions, we use dim(W) (W) to denote its dimension. Specifically, we have dim(W)=m×n (W)=m× n, m×m× m and n×n× n, respectively, for the three cases listed above. Conceptually, a general second-moment method takes the form Xt+1=Xt−αtϕ(M(Xt),W(Xt)),X_t+1=X_t- _tφ (M(X_t),W(X_t) ), (5) where ϕ:m×n×dim(W)→m×nφ:R^m× n×R (W) ^m× n is a blending function that combines the first and second moments to form the update direction. Since M(Xt)M(X_t) and W(Xt)W(X_t) cannot be evaluated exactly in practice, we need to replace them with estimators M^t M_t and W^t W_t respectively, leading to the practical second-moment SA method Xt+1=Xt−αtϕ(M^t,W^t).X_t+1=X_t- _tφ ( M_t, W_t ). (6) Below we give several concrete examples of such methods. For the first moment, they all use the EMA estimator M^t M_t in (3) with a constant smoothing factor β1 _1, which we do not repeat in the descriptions below. To simplify the presentation, we use the notation GtG_t for G(Xt,ξt)G(X_t, _t). • Adam [9] can be interpreted through the first definition of the second moment (4a). It uses the EMA of Gt⊙GtG_t G_t, with a smoothing factor β2 _2, as the estimator for W(Xt)W(X_t), i.e., W^t=β2W^t−1+(1−β2)Gt⊙Gt. W_t= _2 W_t-1+(1- _2)G_t G_t. Given the estimates M^t M_t and W^t W_t, and a prespecified ϵ>0ε>0, Adam follows the update rule Xt+1 X_t+1 =Xt−αtM^t⊘W^t+ϵ, =X_t- _t M_t W_t+ε, (7) where ⊘ denotes element-wise division of two matrices and the square-root is taken element-wise. In this case, the blending function ϕ(M,W)=M⊘W+ϵφ(M,W)=M W+ε for some ϵ>0ε>0. • The RMSProp method [10] is a precursor to Adam, which uses the simple estimator M^t=Gt M_t=G_t but the same EMA estimator for W^t W_t and the same blending function ϕφ as for Adam. • ASGO (Adaptive Structured Gradient Optimization) [11] can be interpreted through the definition of W(X)W(X) under matrix multiplication in (4b) if m≤nm≤ n and (4c) if m>nm>n. Specifically, for m≤nm≤ n, it approximates the second moment W(Xt)W(X_t) with the EMA of GtGtTG_tG_t^T and adopts a matrix variant of Adam’s update rule: W^t W_t =β2W^t−1+(1−β2)GtGtT, = _2 W_t-1+(1- _2)G_tG_t^T, (8) Xt+1 X_t+1 =Xt−αt(W^t+ϵI)−1/2M^t. =X_t- _t( W_t+ε I)^-1/2 M_t. Corresponding to the general framework (6), it has ϕ(M,W)=(W+ϵI)−1/2Mφ(M,W)=(W+ε I)^-1/2M. • Muon [12] employs the update rule Xt+1=Xt−αtNewtonSchulz5(M^t),X_t+1=X_t- _t\, NewtonSchulz5( M_t), (9) where NewtonSchulz5(⋅) NewtonSchulz5(·) is a carefully tuned Newton-Schulz iterative method [13, 14, 15, 16] to approximate the nearest semi-orthogonal matrix to M^t M_t, denoted as Orth(M^t) Orth( M_t). Suppose the SVD of M^t M_t is M^t=UtΣtVtT M_t=U_t _tV_t^T where Ut∈m×rU_t ^m× r and Vt∈n×rV_t ^n× r (with r=minm,nr= \m,n\) have orthonormal columns, it holds that Orth(M^t)=UtVtT Orth( M_t)=U_tV_t^T. We focus on the ideal Muon method Xt+1=Xt−αtOrth(M^t),X_t+1=X_t- _t\, Orth( M_t), where the blending function ϕ(M,W)=Orth(M)φ(M,W)= Orth(M). Notice that Muon does not use a second-moment estimator. Nevertheless, we will show that it can be derived from our framework of second-moment preconditioning, even though the simplified form, under favorable noise structure, does not involve the second moment explicitly. Many other modern deep-learning optimizers can also be cast into the same framework, including SignSGD [17], Signum [16], Lion [18], AdaGrad [19], Nadam [20], Adam-mini [21], AdEMAMix [22], Shampoo [23] and SOAP [24]. We omit the details here due to space limitations. Despite the apparent differences between the forms of these methods, we show that most of them can be derived or interpreted from a unified framework of second-moment preconditioning. The details are given in Section 2. 1.2 A two-stage framework for convergence analysis Accompanying our general formulation of second-moment SA methods, we develop a two-stage framework for their convergence analysis. • Stage 1. We analyze the conceptual algorithm (5) under the following pair of assumptions: ⟨ϕ(M(X),W(X)),X−X∗⟩ φ (M(X),W(X) ),\,X-X_* ≥μ‖X−X∗‖F2, ≥μ\|X-X_*\|_F^2, (10) ‖ϕ(M(X),W(X))‖F \|φ (M(X),W(X) ) \|_F ≤L‖X−X∗‖F+C, ≤ L\|X-X_*\|_F+C, where X∗X_* is a solution to M(X)=0M(X)=0, ∥⋅∥F\|·\|_F denotes the matrix Frobenius norm, and μ, L and C are nonnegative constants. The first inequality is an extension of the original assumption by Robbins and Monro [6], which played a fundamental role in the classical SA literature (25, see, e.g.). The second inequality, which bounds the magnitude of ϕφ, automatically holds for many second-moment methods including Adam and Muon. Therefore, our main focus is on the first inequality, which we call the “aiming” condition. • Stage 2. To analyze the practical method (6), we need an additional assumption on how the update ϕ(M^t,W^t)φ( M_t, W_t) deviates from its nominal value ϕ(M(Xt),W(Xt))φ(M(X_t),W(X_t)). Specifically, we assume ϵt:=‖[ϕ(M^t,W^t)]−ϕ(M(Xt),W(Xt))‖F _t:= \|E [φ( M_t, W_t) ]-φ (M(X_t),W(X_t) ) \|_F (11) is bounded (having finite limit superior) almost surely. This assumption allows us to invoke Dvoretzky’s theorem [26, 27, 28] to show that the practical method converges almost surely to a neighborhood of X∗X_*, and the size of the neighborhood depends on the asymptotic upper bound on ϵt _t. Then we derive concrete bounds on ϵt _t in the specific cases of Muon and ASGO, which fittingly depend on the bias and variance of both estimators M^t M_t and W^t W_t. The two-stage procedure outlined above provides a general framework that can unify the analysis of many existing algorithms as well as guide the development of new ones. For example, RMSProp and Adam share the same blending function ϕφ and second-moment estimator but differ in their first-moment estimators. In particular, RMSProp uses GtG_t as its first-moment estimator which has no bias but high variance, while Adam uses an EMA of GtG_t which has bias but low variance. Therefore, they share the same analysis except that the bias and variance of their respective first-moment estimators lead to different bounds on ϵt _t in (11), which dictate the size of the neighborhood around X∗X_* they converge to. As an example of developing new variants, AdEMAMix [22] combines two EMAs of GtG_t with different smoothing factors as the first-moment estimator. We can also employ more sophisticated estimators, such as double exponential moving average (29, e.g.,) of GtG_t as the first-moment estimator and that of Gt⊙GtG_t G_t or GtGtTG_tG_t^T as the second-moment estimator. This gives rise to a plethora of algorithms that would be challenging to analyze separately one by one. With our framework, we first identify the underlying blending function ϕφ and check if it satisfies the assumptions in (10), and then characterize how the bias and variance of the first- and second-moment estimators translate into a bound on ϵt _t defined in (11). 1.3 Outline In Section 2, we present a general framework of second-moment preconditioning for solving stochastic matrix equations and derive the blending functions for Muon, ASGO and Adam. In Section 3, we present details of the two-stage framework for convergence analysis outlined in Section 1.2. Our analysis is based on a slightly more general form of (5) that also include decoupled weight decay [30], which is an essential ingredient for practical methods. Our theory indeed supports its necessity. In Sections 4 and Section 5, we derive concrete bounds for the deviation ϵt _t defined in (11) for Muon and ASGO respectively, and show that they depend on the bias and variance of the first- and second-moment estimators. In Section 6, we present numerical experiments on training large language models with Muon and ASGO to illustrates some implications of our theory. Finally, we conclude in Section 7 with comments on related work and the limitations of our approach. 2 Second-moment preconditioning For solving the matrix equation M(X)=0M(X)=0 where M maps m×nR^m× n to itself, a full preconditioner based on its Jacobian (Newton’s method) requires an m×n×m×nm× n× m× n tensor, which is very expensive to compute and store for even moderate sizes of m and n. In the stochastic case with M(X)=ξ[G(X,ξ)]M(X)=E_ξ[G(X,ξ)], we present a framework of second-moment preconditioning that does not involve the Jacobian and employs much smaller preconditioners. Different matrix products lead to different definitions of the second moment listed in (4). First, with the matrix Hadamard product, we have the preconditioned update rule Xt+1=Xt−Pt⊙Gt,X_t+1=X_t-P_t G_t, (12) where Pt∈m×nP_t ^m× n is an elementwise or “diagonal” preconditioner. This leads to the second-moment definition in (4a). With left matrix multiplication, we have Xt+1=Xt−PtGt,X_t+1=X_t-P_tG_t, (13) where Pt∈m×mP_t ^m× m can be much smaller in size than XtX_t if m<nm<n. In the case of m>nm>n, we can switch the order of matrix multiplication, i.e., work with the update rule Xt+1=Xt−GtPt,X_t+1=X_t-G_tP_t, (14) where Pt∈n×nP_t ^n× n. These two cases lead to the definitions of the second moment in (4b) and (4c) respectively. Since PtP_t may have a smaller size than XtX_t itself, we call it a fused preconditioner. With any of the update rules discussed above, we want to find the optimal preconditioner PtP_t that can move Xt+1X_t+1 as fast as possible towards some X∗X_* satisfying M(X∗)=0M(X_*)=0. In the stochastic setting, we minimize the expected distance t[‖Xt+1−X∗‖F2],E_t\! [\|X_t+1-X_*\|_F^2 ], (15) where t[⋅]E_t[·] denotes the conditional expectation [⋅|X0,ξ0,X1,…,ξt−1,Xt].E\! [\,·\, |\,X_0, _0,X_1,…, _t-1,X_t ]. (16) By substituting (12), (13) or (14) into (15), we notice that the expected distance is a convex quadratic function of PtP_t, whose coefficients consist of the first and second moments of GtG_t, namely M(Xt)M(X_t) and W(Xt)W(X_t), as well as the unknown quantity Xt−X∗X_t-X_*. Therefore we can express the optimal preconditioner in terms of these quantities, albeit not computable. Next we can approximate the optimal preconditioner by removing the explicit dependence on Xt−X∗X_t-X_* and absorbing its magnitude with a tunable step size αt _t (which decreases to zero as XtX_t approaches X∗X_*). This leads to a conceptual algorithm of the form (5). Then we replace M(Xt)M(X_t) and W(Xt)W(X_t) with their estimators M^t M_t and W^t W_t respectively to obtain practical second-moment SA methods. In Section 2.1, we derive a conceptual second-moment method based on the fused preconditioner in (13). Then we make further simplifications that lead to the update rules of Muon and ASGO in Sections 2.2 and 2.3 respectively. Derivations for (14) mirror those for (13) and thus are omitted. In Section 2.4, we present results for the elementwise preconditioner (12), which lead to the update rule for Adam and several variants. 2.1 The optimal fused preconditioner and its approximation In this section, we focus on the case m≤nm≤ n and the preconditioned SA method (13). Assuming that there exists an X∗X_* satisfying M(X∗)=0M(X_*)=0, we would like to choose PtP_t to minimize the expected distance from Xt+1X_t+1 to X∗X_*. First, we use the update rule (13) to obtain ‖Xt+1−X∗‖F2 \|X_t+1-X_*\|_F^2 =‖Xt−PtGt−X∗‖F2 =\|X_t-P_tG_t-X_*\|_F^2 =‖Xt−X∗‖F2−2Tr((Xt−X∗)TPtGt)+Tr(PtGtGtTPtT). =\|X_t-X_*\|_F^2-2 Tr ((X_t-X_*)^TP_tG_t )+ Tr (P_tG_tG_t^TP_t^T ). Taking the conditional expectation defined in (16) yields t[‖Xt+1−X∗‖F2] _t [ \|X_t+1-X_ \|_F^2 ] =‖Xt−X∗‖F2−2Tr((Xt−X∗)TPtt[Gt])+Tr(Ptt[GtGtT]PtT) = \|X_t-X_ \|_F^2-2 Tr ((X_t-X_ )^TP_tE_t[G_t] )+ Tr (P_tE_t[G_tG_t^T]P_t^T ) =‖Xt−X∗‖F2−2Tr((Xt−X∗)tPtMt)+Tr(PtWtPtT), = \|X_t-X_ \|_F^2-2 Tr ((X_t-X_ )^tP_tM_t )+ Tr (P_tW_tP_t^T ), where we used the linearity of expectation and replaced t[Gt]E_t[G_t] and t[GtGtT]E_t[G_tG_t^T] with the simpler notations MtM_t and WtW_t respectively. Clearly this is a convex quadratic function of PtP_t. Assume that WtW_t is positive definite. Then the expected distance is strictly convex in PtP_t and has a unique minimizer Pt⋆P_t . Setting the gradient with respect to PtP_t equal to zero, i.e., −2(Xt−X∗)MtT+2PtWt=0,-2(X_t-X_ )M_t^T+2P_tW_t=0, we obtain the unique optimal preconditioner Pt⋆=(Xt−X∗)MtTWt−1.P_t =(X_t-X_ )M_t^TW_t^-1. (17) An important property of Pt⋆P_t is that the resulting update Pt⋆GtP_t G_t is scaling-invariant, meaning that it does not change if GtG_t is multiplied by any nonzero scalar. We must approximate Pt⋆P_t because X∗X_ is unknown. A simple idea is to replace the matrix Xt−X∗X_t-X_ with a scalar step size αt≈‖Xt−X∗‖ _t≈ \|X_t-X_ \| to account for its magnitude. However, the remaining part MtTWt−1M_t^TW_t^-1 is a n×mn× m matrix, which does not match PtP_t’s size of m×m× m. Therefore, we need a more nuanced approximation. For this purpose, we assume that MtM_t has full row rank and its SVD is Mt=UtΣtVtT,M_t=U_t _tV_t^T, (18) where Ut∈m×mU_t ^m× m and Vt∈n×mV_t ^n× m have orthonormal columns and Σt∈+m×m _t ^m× m_+ is the diagonal matrix of singular values. Then MtMtTM_tM_t^T is positive definite and the SVD in (18) gives MtT(MtMtT)−1/2=VtUtT.M_t^T(M_tM_t^T)^-1/2=V_tU_t^T. Consequently, Pt⋆ P_t =(Xt−X∗)MtT(MtMtT)−1/2(MtMtT)1/2Wt−1 =(X_t-X_ )M_t^T(M_tM_t^T)^-1/2(M_tM_t^T)^1/2W_t^-1 =(Xt−X∗)VtUtT(MtMtT)1/2Wt−1. =(X_t-X_ )V_tU_t^T(M_tM_t^T)^1/2W_t^-1. Notice that multiplication by VtUtTV_tU_t^T is a partial isometry, satisfying ‖(Xt−X∗)VtUtT‖F=‖(Xt−X∗)Vt‖F≤‖Xt−X∗‖F,\|(X_t-X_*)V_tU_t^T\|_F=\|(X_t-X_*)V_t\|_F≤\|X_t-X_*\|_F, with equality when m=nm=n. Therefore we approximate Pt⋆P_t with Pt=αt(MtMtT)1/2Wt−1,P_t= _t(M_tM_t^T)^1/2W_t^-1, where αt _t is a tunable step size, roughly proportional to ‖(Xt−X∗)VtUtT‖F\|(X_t-X_*)V_tU_t^T\|_F. Another benefit of the above approximation is that the overall update, PtGt=αt(MtMtT)1/2Wt−1Gt,P_tG_t= _t(M_tM_t^T)^1/2W_t^-1G_t, retains the scaling-invariant property of Pt⋆GtP_t G_t (against positive scaling of GtG_t). This is not the case, for example, by simply choosing Pt=αtWt−1P_t= _tW_t^-1. For construction of a conceptual algorithm, we are most interested in its average behavior. Therefore, we further replace GtG_t with its expectation MtM_t to arrive at Xt+1=Xt−αt(MtMtT)1/2Wt−1Mt,X_t+1=X_t- _t(M_tM_t^T)^1/2W_t^-1M_t, (19) which belongs to the general class of (5). We can immediately instantiate this conceptual algorithm by replacing MtM_t and WtW_t with their estimators M^t M_t and W^t W_t respectively: Xt+1=Xt−αt(M^tM^tT)1/2W^t−1M^t,X_t+1=X_t- _t( M_t M_t^T)^1/2 W_t^-1 M_t, for example, using a pair of EMA estimators M^t M_t =β1M^t−1+(1−β1)Gt, = _1 M_t-1+(1- _1)G_t, (20) W^t W_t =β2W^t−1+(1−β2)GtGtT. = _2 W_t-1+(1- _2)G_tG_t^T. However, each step of this algorithm requires four matrix multiplications, a matrix square-root and a matrix inverse, which is computationally very expensive. Therefore we propose to further simplify it. In the following, we present two approaches that differ in our assumptions on the noise structure in GtG_t, which lead to Muon and ASGO respectively. 2.2 Derivation of Muon and its variants As discussed in Section 2.1, an important property of the optimal update Pt⋆GtP_t G_t or Pt⋆MtP_t M_t is its scaling invariance, i.e., it does not change if we multiply GtG_t by an arbitrary positive scalar. We would like to retain this property with any further simplification. Given the stochastic mapping Gt:=G(Xt,ξt)G_t:=G(X_t, _t), we interpret Mt:=t[Gt]M_t:=E_t[G_t] as the signal and Gt−MtG_t-M_t as the noise. Then we have the following decomposition of the second moment into signal and noise powers: Wt:=t[GtGtT]=MtMtT+t[(Gt−Mt)(Gt−Mt)T].W_t:=E_t[G_tG_t^T]=M_tM_t^T+E_t[(G_t-M_t)(G_t-M_t)^T]. (21) Again, assume that WtW_t is positive definite. To measure the relative strength of the signal and noise, we define a signal fraction matrix St:=(MtMtT)1/2Wt−1(MtMtT)1/2.S_t:=(M_tM_t^T)^1/2W_t^-1(M_tM_t^T)^1/2. (22) Since Wt−MtMtTW_t-M_tM_t^T is positive semidefinite, the signal fraction matrix StS_t is positive semidefinite and all its eigenvalues lie in the interval [0,1][0,1]. If in addition MtM_t has full row rank, then StS_t is positive definite and its eigenvalues lie in (0,1](0,1]. Eigenvalues near zero or one correspond to low or high signal-to-noise ratios along the direction of their associated eigenvectors. With the definition of StS_t, we can write the update in the conceptual algorithm (19) as αt(MtMtT)1/2Wt−1Mt=αtSt(MtMtT)−1/2Mt. _t(M_tM_t^T)^1/2W_t^-1M_t= _tS_t(M_tM_t^T)^-1/2M_t. (23) If the signal fractions are relatively uniform along different directions, i.e., St≈γtImS_t≈ _tI_m for some γt∈(0,1] _t∈(0,1], then we can approximate the update as αtγt(MtMtT)−1/2Mt _t _t(M_tM_t^T)^-1/2M_t and assimilate γt _t into the tunable step size αt _t. This leads to the following simplified algorithm Xt+1=Xt−αt(MtMtT)−1/2Mt.X_t+1=X_t- _t(M_tM_t^T)^-1/2M_t. (24) This method retains the scaling-invariance property. Moreover, in light of the SVD in (18), we have (MtMtT)−1/2Mt=UtVtT=Orth(Mt)(M_tM_t^T)^-1/2M_t=U_tV_t^T= Orth(M_t). Therefore, it can be rewritten as Xt+1=Xt−αtOrth(Mt).X_t+1=X_t- _t Orth(M_t). (25) In practice, we replace MtM_t with an online estimator. Below are a few concrete examples. • The stochastic spectral descent (SSD) method [31] uses GtG_t as an estimator for MtM_t in (25). • As discussed in Section 1.1, Muon [12] uses an EMA of GtG_t as M^t M_t. The authors of [12] also proposed a variant that uses a combination of the EMA and GtG_t as an estimator, i.e., M^t M_t =βM^t−1+(1−β)Gt, =β M_t-1+(1-β)G_t, (26) Xt+1 X_t+1 =Xt−αtOrth(γtM^t+(1−γt)Gt), =X_t- _t Orth ( _t M_t+(1- _t)G_t ), where β,γt∈[0,1]β, _t∈[0,1]. This estimator is also called a quasi-hyperbolic momentum [32]. • MARS [33] first computes an estimator G^t G_t using a variance-reduction technique (34, e.g.,) and then uses the EMA of G^t G_t to form an estimator forMtM_t in (25). 2.3 Derivation of ASGO and its variants The class of algorithms presented in Section 2.2 has the advantage of not requiring an estimator for the second moment explicitly. But their derivation relies on the assumption that St≈γtImS_t≈ _tI_m for some γt∈(0,1] _t∈(0,1], where StS_t is the signal fraction matrix defined in (22). When this assumption does not hold, we need to derive a simplified version of (23) that captures the structure of StS_t. For this purpose, we define the following matrix Rt=(MtMtT)1/2Wt−1/2,R_t=(M_tM_t^T)^1/2W_t^-1/2, which satisfies St=RtRtTS_t=R_tR_t^T. Notice that RtR_t is not the matrix square-root of StS_t, but the decomposition St=RtRtTS_t=R_tR_t^T is unique up to unitary transformations. Then we can write (23) as αt(MtMtT)1/2Wt−1Mt _t(M_tM_t^T)^1/2W_t^-1M_t =αtRtWt−1/2Mt. = _tR_tW_t^-1/2M_t. Instead of dropping the full signal-fraction matrix StS_t (approximating it with a scaled identity matrix), here we drop RtR_t, which is “half” of StS_t, and retain the other half RtTR_t^T within Wt−1/2MtW_t^-1/2M_t. This leads to the following simplified conceptual algorithm: Xt+1=Xt−αtWt−1/2Mt.X_t+1=X_t- _tW_t^-1/2M_t. (27) Compared with (24), we replaced (MtMtT)−1/2 (M_tM_t^T )^-1/2 with Wt−1/2W_t^-1/2. According to the decomposition in (21), WtW_t takes into account both the signal power MtMtTM_tM_t^T and also the noise power. Therefore, the update (27) can be more effective than (24) under nonuniform signal fractions. Again, the update Wt−1/2MtW_t^-1/2M_t is invariant under positive scaling of GtG_t. Now we can derive practical methods by replacing MtM_t and WtW_t in (27) with online estimators. The immediate choice is to use the two EMA estimators in (20). Below are two more variations. • ASGO [11] adds a small perturbation ϵIε I to the second-moment estimator, i.e., Xt+1=Xt−αt(W^t+ϵI)−1/2M^t,X_t+1=X_t- _t( W_t+ε I)^-1/2 M_t, where ϵ>0ε>0. This is consistent with our assumption that WtW_t is positive definite. • Leon [35] employs a more sophisticated second-moment estimator: Xt+1=Xt−αt(γtW^t+(1−γt)M^tM^tT)−1/2M^t,X_t+1=X_t- _t( _t W_t+(1- _t) M_t M_t^T)^-1/2 M_t, where γt∈[0,1] _t∈[0,1]. It is an interpolation between Muon (as γt→0 _t→ 0) and ASGO (as γt→1 _t→ 1). Two-sided preconditioning and the Shampoo optimizer As discussed in the beginning of Section 2, when m>nm>n, it is natural to derive second-moment SA methods with right matrix multiplication as in (14). In addition, we can also consider two-sided preconditioning: Xt+1=Xt−LtGtRt,X_t+1=X_t-L_tG_tR_t, where Lt∈m×mL_t ^m× m and Rt∈n×nR_t ^n× n are the left and right preconditioners. Following the derivations above, we can construct the following algorithm, Xt+1=Xt−αtWt−1/4MtQt−1/4,X_t+1=X_t- _tW_t^-1/4M_tQ_t^-1/4, where Wt=t[GtGtT]W_t=E_t[G_tG_t^T] and Qt=t[GtTGt]Q_t=E_t[G_t^TG_t] are the left and right second moments, appearing in (4b) and (4c) respectively. Setting the two matrix powers to be −1/4-1/4 instead of −1/2-1/2 is necessary to retain the scaling-invariance property of the update. The Shampoo optimizer [23] is an instantiation of this conceptual method where MtM_t is replaced by the simple estimator GtG_t and WtW_t and QtQ_t are replaced by EMAs of GtGtTG_tG_t^T and GtTGtG_t^TG_t respectively. Obviously, a simple variant can be obtained by replacing MtM_t by an EMA of GtG_t. 2.4 Element-wise preconditioning In this section, we focus on element-wise preconditioning based on the update rule Xt+1=Xt−Pt⊙Gt,X_t+1=X_t-P_t G_t, where Pt∈m×nP_t ^m× n has the same size as XtX_t and GtG_t. The derivation of the optimal preconditioner and the ensuing simplifications follow the same steps presented in Sections 2.1, 2.2 and 2.3, only with simpler algebra. Some results in this case were obtained in our previous work [36], without relying on the concept of second-moment preconditioning. Here we omit the details and only summarize the main results. First, following similar steps as in Section 2.1, we obtain a conceptual second-moment method Xt+1=Xt−αt|Mt|⊘Wt⊙Mt,X_t+1=X_t- _t|M_t| W_t M_t, (28) where Mt=t[Gt]M_t=E_t[G_t] and Wt=t[Gt⊙Gt]W_t=E_t[G_t G_t]. Since |Mt|=Mt⊙Mt|M_t|= M_t M_t, this method has clear correspondence with (19). The signal fraction matrix in this case is also defined element-wise: St:=(Mt⊙Mt)⊘Wt.S_t:= (M_t M_t ) W_t. All elements of the SiF matrix StS_t take values between 00 and 11, with 00 indicating GtG_t being pure noise and 11 indicating pure signal. Since |Mt|⊙Mt=(Mt⊙Mt)⊙sign(Mt)|M_t| M_t=(M_t M_t) sign(M_t), we can write (28) as Xt+1=Xt−αtSt⊙sign(Mt).X_t+1=X_t- _tS_t sign(M_t). (29) The following two simplifications parallel those in Sections 2.2 and 2.3 respectively. Sign gradient method and variants When the signal fractions are relatively uniform, i.e., St≈γtm×nS_t≈ _t1_m× n for some γt>0 _t>0, we can assimilate γt _t into the tunable step size αt _t and arrive at Xt+1=Xt−αtsign(Mt).X_t+1=X_t- _t\, sign(M_t). (30) With different estimators for the first moment MtM_t, we recover several popular practical methods: • SignSGD [17] simply uses GtG_t as an estimator of MtM_t. • Signum [16] employs an EMA of GtG_t as an estimator of MtM_t. • Lion [18] can be interpreted as using a combination of GtG_t and its EMA as an estimator for MtM_t, essentially replacing Orth(⋅) Orth(·) in (26) with the sign(⋅) sign(·) function. Adam and variants When the signal fractions are not uniform, we should take account of their effect but also need to compensate for the inaccuracy and volatility brought by any practical estimators. A reasonable compromise is to only factor in their square root (mapping them closer to uniform). More concretely, we would like to keep the square root of StS_t in (29) instead of replacing it with a uniform constant as done in (30). This leads to Xt+1 X_t+1 =Xt−αtSt⊙sign(Mt) =X_t- _t S_t sign(M_t) =Xt−αtMt⊘Wt. =X_t- _tM_t W_t. Replacing MtM_t and WtW_t with respective estimators, we recover several practical methods: • Adam [9] uses EMA estimators for both MtM_t and WtW_t but with an ϵ>0ε>0 added to the denominator, as shown in (7). We can also view W^t+ϵ W_t+ε as an estimator for WtW_t. • RMSProp [10] uses GtG_t to estimate MtM_t but the same second-moment estimator as Adam. AdaGrad [19] can be interpreted as using a different second-moment estimator. • Nadam [20] uses a combination of GtG_t and its EMA as the estimator for MtM_t, similar to the one in (26). AdEMAMix [22] combines two EMAs of GtG_t with different smoothing factors. The above examples mainly differ in their choice of first-moment estimators and share the same second-moment estimator as Adam. Exploring more varieties of the second-moment estimator could also be fruitful. Beyond the generic EMA estimators, our framework also facilitates the adoption of more customized estimators that can be model-specific and data-driven. 3 A general framework for convergence analysis In this section, we present a general framework for the convergence analysis of second-moment SA methods. For full generality, our framework also covers decoupled weight decay [30], which has proven to be a very effective ingredient in deep learning optimizers, especially for improving the performance of both Adam [30] and Muon [37]. Specifically, decoupled weight decay means including an additional damping term λXtλ X_t, where λ≥0λ≥ 0, to the update in (6): Xt+1 X_t+1 =Xt−αt(ϕ(M^t,W^t)+λXt) =X_t- _t (φ( M_t, W_t)+λ X_t ) =(1−αtλ)Xt−αtϕ(M^t,W^t). =(1- _tλ)X_t- _tφ( M_t, W_t). (31) The damping term can be interpreted through the classical Halpern iteration [38] or, in the context of machine learning, understood as the consequence of quadratic regularization. As explained in Section 1.2, our framework consists of two stages. First, we analyze the conceptual counterpart of (31): Xt+1=(1−αtλ)Xt−αtϕ(Mt,Wt).X_t+1=(1- _tλ)X_t- _tφ (M_t,W_t ). (32) (Recall that MtM_t and WtW_t are simple notations for M(Xt)M(X_t) and W(Xt)W(X_t) respectively.) The analysis relies on an aiming condition for the conceptual update ϕ(Mt,Wt)φ(M_t,W_t), which we introduce in Section 3.1. Then we prove the convergence of (32) and characterize its rate of convergence in Section 3.2. In the second stage, presented in Section 3.3, we analyze the practical algorithms (31) by assuming that the deviation between the expected practical update t[ϕ(M^t,W^t)]E_t[φ( M_t, W_t)] and the conceptual update ϕ(Mt,Wt)φ(M_t,W_t) is bounded almost surely. 3.1 Aiming condition For the classical Robbins-Monro method (2), the following is a typical set of assumptions that guarantee almost sure iterate convergence: • The random variables ξtt≥0\ _t\_t≥ 0 are independent and G(X,ξ)G(X,ξ) has bounded variance; • The step sizes αtt≥0\ _t\_t≥ 0 are nonnegative and satisfy ∑t=0∞αt=∞ _t=0^∞ _t=∞ and ∑t=0∞αt2<∞ _t=0^∞ _t^2<∞; • Suppose M(X∗):=ξ∗[G(X∗,ξ∗)]=0M(X_ ):=E_ _ [G(X_ , _ )]=0 and there exists L,μ>0L,μ>0 such that for all X∈m×nX ^m× n: ⟨M(X),X−X∗⟩ M(X),\,X-X_ ≥μ‖X−X∗‖F2, ≥μ \|X-X_ \|_F^2, (33) ‖M(X)‖F \|M(X) \|_F ≤L‖X−X∗‖F. ≤ L \|X-X_ \|_F. (34) Under these conditions, the sequence Xt\X_t\ generated by (2) converges to X∗X_ almost surely [39, 40]. Specifically, condition (33) ensures that the update direction M(X)=ξ[G(X,ξ)]M(X)=E_ξ[G(X,ξ)] always aims towards the target X∗X_ . For this reason, we call it the aiming condition. Condition (34) controls the growth of ‖M(X)‖F\|M(X)\|_F as X deviates from X∗X_*. Together, they guarantee that moving along the update direction −M(X)-M(X) reduces the distance to X∗X_* for sufficiently small αt _t. X∇F∇ F∇F(X)∇ F(X)X∗X_ L(X−X∗)L(X-X_ )μ(X−X∗)μ(X-X_ ) Figure 1: One-dimensional depiction of the gradient of a function F that satisfies both the aiming condition (35) and the growth condition (36) but is neither monotone nor Lipschitz continuous. In the context of optimization, i.e., if M is the gradient of a cost function F:m×n→F:R^m× n , the two conditions (33) and (34) become ⟨∇F(X),X−X∗⟩ ∇ F(X),\,X-X_ ≥μ‖X−X∗‖F2, ≥μ \|X-X_ \|_F^2, (35) ‖∇F(X)‖F \|∇ F(X) \|_F ≤L‖X−X∗‖F, ≤ L \|X-X_ \|_F, (36) which are consequences of F being μ-strongly convex and L-smooth respectively (41, see, e.g.,). Notice that these two conditions are much weaker than strong convexity and smoothness because they are anchored at a specific point X∗X_*. Figure 1 depicts a one-dimensional gradient mapping ∇F(X)∇ F(X) that is sandwiched between two linear functions L(X−X∗)L(X-X_*) and μ(X−X∗)μ(X-X_*), satisfying both the aiming and growth conditions. However, ∇F(X)∇ F(X) is neither monotone (implying that F is not convex) nor Lipschitz-continuous (thus F is nonsmooth). To analyze second-moment SA methods, we extend the aiming condition (33) and the growth condition (34) to work with the blending function ϕφ as well as decoupled weight decay. Assumption A (Aiming and growth conditions). There exist X∗∈m×nX_* ^m× n, μ>0μ>0, and L,C≥0L,C≥ 0 such that, for all X∈m×nX ^m× n, ⟨ϕ(M(X),W(X))+λX,X−X∗⟩ φ (M(X),W(X) )+λ X,\,X-X_* ≥μ‖X−X∗‖F2, ≥μ\|X-X_*\|_F^2, (37) ‖ϕ(M(X),W(X))‖F \|φ (M(X),W(X) ) \|_F ≤L‖X−X∗‖F+C. ≤ L \|X-X_ \|_F+C. (38) Unlike the first-moment aiming condition (33), the generalization in (37) does not have a simple connection with convexity (when interpreting M as the gradient mapping of a cost function) due to the presence of ϕφ and W. It is easy to construct examples that satisfy one but fail the other (36, Appendix A). Nonetheless, it is closely related to the definition of strong pseudogradient (42, Section 2.2.3). A matrix mapping D(X)D(X) is a strong pseudogradient relative to a Lyapunov function V(X)V(X) if there exists γ>0γ>0 such that ⟨D(X),∇V(X)⟩≥γV(X). D(X),∇ V(X) ≥γ V(X). Condition (37) implies that ϕ(M(X),W(X))+λXφ(M(X),W(X))+λ X is a strong pseudogradient with respect to V(X)=(1/2)‖X−X∗‖F2V(X)=(1/2)\|X-X_*\|_F^2 with parameter γ=2μγ=2μ. The generalized growth condition (38) offers more flexibility than its classical counterpart (34). When C>0C>0, we can set L=0L=0 if ϕφ is a bounded mapping, which is the case for sign(M(X)) sign(M(X)) and Orth(M(X)) Orth(M(X)). More generally, the composite mapping ϕ(M(X),W(X))φ(M(X),W(X)) may remain bounded even if ϕφ itself is not. As an example, consider the update mapping in (27), ϕ(M(X),W(X))=W(X)−1/2M(X).φ(M(X),W(X))=W(X)^-1/2M(X). While ϕ(Y,Z)=Y−1/2Zφ(Y,Z)=Y^-1/2Z is not bounded in general, the fact that M(X)M(X) and W(X)W(X) are the first and second moments of the same random variable guarantees that ϕ(M(X),W(X))φ(M(X),W(X)) is always bounded. The same conclusion is true for Adam’s update map ϕ(M(X),W(X))=M(X)⊘W(X).φ (M(X),W(X) )=M(X) W(X). Indeed, having a bounded update map is a common property for the general class of second-moment SA method we consider in this paper. It is rooted from the derivation of the optimal preconditioner in Section 2.1, where the optimal update Pt⋆GtP_t G_t is scaling invariant. This property persisted in all our approximation and simplifications in deriving Muon, ASGO, Adam and their variants. As a consequence, the growth condition (38) holds automatically for all of them, even if the first-moment M(X)M(X) does not satisfy (34). This is a major advantage of second-moment SA methods over the classical first-moment methods. Remark 3.1. When ϕ(M(X),W(X))φ(M(X),W(X)) is bounded by a constant, we need λ>0λ>0 in order for the aiming condition (37) to hold globally. This provides an explanation for the use of decoupled weight decay [30] in second-moment methods such as Adam and Muon. While it is difficult to characterize the exact range of λ that guarantees the global aiming condition, in practice we only need this condition to hold locally within a sufficiently large neighborhood of X∗X_*. Example 3.1 (Aiming condition of a matrix polynomial for Muon). Assume m≤nm≤ n and fix some X∗∈m×nX_* ^m× n. Consider the matrix polynomial M(X)=((X−X∗)(X−X∗)T)p(X−X∗),M(X)= ((X-X_*)(X-X_*)^T )^p(X-X_*), where p≥−1/2p≥-1/2. Given the SVD of X−X∗=UΣVTX-X_*=U V^T where U∈m×mU ^m× m and V∈n×mV ^n× m, we have M(X)=U(Σ2p+1)VTM(X)=U ( ^2p+1 )V^T and thus Orth(M(X))=UVT Orth(M(X))=UV^T. This leads to ⟨Orth(M(X)),X−X∗⟩=Tr(Σ)=‖X−X∗‖tr. Orth(M(X)),X-X_* = Tr( )=\|X-X_*\|_tr. where ‖X−X∗‖tr\|X-X_*\|_tr denotes the trace norm (nuclear norm) of X−X∗X-X_*. We deduce that the aiming condition (37) with λ=0λ=0 and any μ>0μ>0 holds in the following neighborhood of X∗X_*: X∈m×n:‖X−X∗‖F2≤1μ‖X−X∗‖tr. \X ^m× n:\|X-X_*\|_F^2≤ 1μ\|X-X_*\|_tr \. The global aiming condition holds with sufficiently large λ>0λ>0. Meanwhile, the growth condition (38) holds with L=0L=0 and C=mC= m, because ‖Orth(M(X))‖F2=‖UVT‖F2≤m\| Orth(M(X))\|_F^2=\|UV^T\|_F^2≤ m. 3.2 Analysis of conceptual algorithms In this section, we prove convergence and establish the rates of convergence for the conceptual method (32) under Assumption A. Our first result is the following one-step contraction property. Lemma 3.1. Suppose Assumption A holds, αt≥0 _t≥ 0, and αtλ≤1 _tλ≤ 1 for all t≥0t≥ 0. Then the sequence Xt\X_t\ generated by (32) satisfies the following one-step improvement bound for all t≥0t≥ 0, ‖Xt+1−X∗‖F2≤(1+αt2(3L2+2λμ)−2αtμ)‖Xt−X∗‖F2+3αt2(λ2‖X∗‖F2+C2). \|X_t+1-X_ \|_F^2≤ (1+ _t^2(3L^2+2λμ)-2 _tμ ) \|X_t-X_ \|_F^2+3 _t^2(λ^2 \|X_* \|_F^2+C^2). (39) Proof. First, we notice that the aiming condition (37) in Assumption A is equivalent to ⟨Xt−X∗,ϕ(Mt,Wt)+λX∗⟩≥(μ−λ)‖Xt−X∗‖F2. X_t-X_*,\,φ (M_t,W_t )+λ X_* ≥(μ-λ)\|X_t-X_*\|_F^2. (40) To see this, we simply add and subtract λXtλ X_t inside the inner product: ⟨Xt−X∗,ϕ(Mt,Wt)+λX∗⟩ X_t-X_*,\,φ (M_t,W_t )+λ X_* =⟨Xt−X∗,ϕ(Mt,Wt)+λXt−λXt+λX∗⟩ = X_t-X_*,\,φ (M_t,W_t )+λ X_t-λ X_t+λ X_* =⟨Xt−X∗,ϕ(Mt,Wt)+λXt⟩−λ‖Xt−X∗‖F2, = X_t-X_*,\,φ (M_t,W_t )+λ X_t -λ\|X_t-X_*\|^2_F, which together with (37) implies (40). Given the conceptual update rule (32), we have ‖Xt+1−X∗‖F2 \|X_t+1-X_ \|^2_F =‖(1−αtλ)Xt−αtϕ(Mt,Wt)−X∗‖F2 = \|(1- _tλ)X_t- _tφ (M_t,W_t )-X_* \|^2_F =‖(1−αtλ)(Xt−X∗)−αt(ϕ(Mt,Wt)+λX∗)‖F2 = \|(1- _tλ)(X_t-X_*)- _t (φ (M_t,W_t )+λ X_* ) \|^2_F =(1−αtλ)2‖Xt−X∗‖F2−2αt(1−αtλ)⟨Xt−X∗,ϕ(Mt,Wt)+λX∗⟩ =(1- _tλ)^2 \|X_t-X_* \|_F^2-2 _t(1- _tλ) X_t-X_*,φ (M_t,W_t )+λ X_* +αt2‖ϕ(Mt,Wt)+λX∗‖F2. + _t^2 \|φ (M_t,W_t )+λ X_* \|^2_F. (41) For the second term on the right-hand side, we use the assumption αtλ≤1 _tλ≤ 1 and (40) to obtain −2αt(1−αtλ)⟨Xt−X∗,ϕ(Mt,Wt)+λX∗⟩≤−2αt(1−αtλ)(μ−λ)‖Xt−X∗‖F2. -2 _t(1- _tλ) X_t-X_*,φ (M_t,W_t )+λ X_* ≤-2 _t(1- _tλ)(μ-λ)\|X_t-X_*\|_F^2. For the third term, applying the triangle inequality and the growth condition (38) yields ‖ϕ(Mt,Wt)+λX∗‖F2 \|φ (M_t,W_t )+λ X_* \|^2_F ≤(‖ϕ(Mt,Wt)‖F+λ‖X∗‖F)2 ≤ ( \|φ (M_t,W_t ) \|_F+λ \|X_* \|_F )^2 ≤(L‖Xt−X∗‖+C+λ‖X∗‖F)2 ≤(L \|X_t-X_ \|+C+λ \|X_* \|_F)^2 ≤3L2‖Xt−X∗‖F2+3C2+3λ2‖X∗‖F2, ≤ 3L^2 \|X_t-X_ \|_F^2+3C^2+3λ^2 \|X_* \|_F^2, where the last inequality is due to the fact (a+b+c)2≤3a2+3b2+3c2(a+b+c)^2≤ 3a^2+3b^2+3c^2. Substituting these bounds into (41), we get: ‖Xt+1−X∗‖F2 \|X_t+1-X_ \|^2_F ≤(1−αtλ)2‖Xt−X∗‖F2−2αt(1−αtλ)(μ−λ)‖Xt−X∗‖F2 ≤(1- _tλ)^2 \|X_t-X_* \|_F^2-2 _t(1- _tλ)(μ-λ)\|X_t-X_*\|_F^2 +3αt2L2‖Xt−X∗‖F2+3αt2C2+3αt2λ2‖X∗‖F2 +3 _t^2L^2 \|X_t-X_ \|_F^2+3 _t^2C^2+3 _t^2λ^2 \|X_* \|_F^2 =(1−2αtμ+2αt2λμ−αt2λ2+3αt2L2)‖Xt−X∗‖F2+3αt2(λ2‖X∗‖F2+C2). = (1-2 _tμ+2 _t^2λμ- _t^2λ^2+3 _t^2L^2 ) \|X_t-X_ \|_F^2+3 _t^2(λ^2 \|X_* \|_F^2+C^2). Dropping the negative coefficient −αt2λ2- _t^2λ^2 for ‖Xt−X∗‖F2 \|X_t-X_ \|_F^2 yields the desired result. ∎ We now establish convergence for the conceptual algorithm under Asspumption A and appropriate assumptions on the step sizes. Theorem 3.1 (Convergence of conceptual method). Suppose the step size sequence αtt≥0\ _t\_t≥ 0 satisfies αt>0,αt→0,αtλ≤1,∑t=0∞αt=∞. _t>0, _t→ 0, _tλ≤ 1, _t=0^∞ _t=∞. (42) Then under Assumption A, the sequence Xt\X_t\ generated by (32) satisfies ‖Xt−X∗‖F→0 \|X_t-X_ \|_F→ 0. The proof for Theorem 3.1 is based on the following lemma in [42]. Lemma 3.2 ((42, Section 2.2, Lemma 6’)). Let dtt≥0\d_t\_t≥ 0 be a nonnegative sequence satisfying dt+1≤(1+at)dt−bth(dt)+ct,d_t+1≤(1+a_t)d_t-b_th(d_t)+c_t, (43) where at,bt,ct≥0a_t,b_t,c_t≥ 0 satisfy ∑t=0∞bt=∞,at→0,ct→0,atbt→0,ctbt→0, _t=0^∞b_t=∞, a_t→ 0, c_t→ 0, a_tb_t→ 0, c_tb_t→ 0, and h is a convex function satisfying h(d)>0∀d>0,h(0)=0,h(d′)≥h(d)∀d′≥d≥0.h(d)>0 ∀ d>0, h(0)=0, h(d )≥ h(d) ∀ d ≥ d≥ 0. Then dt→0d_t→ 0. Proof of Theorem 3.1. Define dt=‖Xt−X∗‖F2d_t= \|X_t-X_ \|_F^2. Lemma 3.1 implies the following recursion: dt+1 d_t+1 =(1+αt2(3L2+2λμ))dt−2αtμdt+3αt2(λ2‖X∗‖F2+C2). = (1+ _t^2(3L^2+2λμ) )d_t-2 _tμ\,d_t+3 _t^2(λ^2 \|X_* \|_F^2+C^2). which is in the form of (43) with h being the identity map and at=αt2(3L2+2λμ),bt=2αtμ,ct=3αt2(λ2‖X∗‖F2+C2).a_t= _t^2(3L^2+2λμ), b_t=2 _tμ, c_t=3 _t^2(λ^2 \|X_* \|_F^2+C^2). Notice that the sequences ata_t, btb_t, and ctc_t are all non-negative and the assumption on αt _t in (42) guarantees that they satisfy all the requirements in Lemma 3.2. Moreover, the identity map is convex, monotonically increasing and attains 00 at 00. Therefore, we may apply the lemma to conclude that dt=‖Xt−X∗‖F2→0d_t= \|X_t-X_ \|_F^2→ 0. ∎ Remark 3.2. Lemma 3.2 is a deterministic counterpart of the Robbins-Siegmund lemma [40], which is a powerful tool in establishing almost sure convergence of stochastic approximation methods. We used the Robbins-Siegmund lemma in our previous work [36] where the conceptual method is also a stochastic algorithm. In this paper, we work with deterministic conceptual algorithms (which have more flexibility in modeling different algorithms) and thus resort to Lemma 3.2. Our next two results concern the rate of convergence of the conceptual method. First, for stepsizes αt _t following a 1/t1/t schedule, we establish an asymptotic (1/t)O(1/t) rate of convergence in squared distance ‖Xt−X∗‖F2 \|X_t-X_* \|^2_F, as given in the following theorem. Theorem 3.2. Consider the algorithm (32) with the step size rule αt=αt+1 _t= αt+1. Suppose Assumption A holds and α satisfies 12μ<α≤2μ3L2+2λμ 12μ<α≤ 2μ3L^2+2λμ where we assume 3L2+2λμ<4μ23L^2+2λμ<4μ^2. Then as t→∞t→∞, ‖Xt−X∗‖F2≤α2A2μα−11t+1+(1(t+1)2+1(t+1)2μα), \|X_t-X_* \|^2_F≤ α^2A2μα-1 1t+1+O ( 1(t+1)^2+ 1(t+1)^2μα ), where A=(3L2+2λμ)B+3λ2‖X∗‖F2+3C2,B=‖X0−X∗‖F2+12(λ2‖X∗‖F2+C2)α2π2.A=(3L^2+2λμ)B+3λ^2 \|X_* \|_F^2+3C^2, B= \|X_0-X_ \|_F^2+ 12 (λ^2 \|X_* \|_F^2+C^2 )α^2π^2. As discussed in Section 3.1, we have L=0L=0 for a broad class of second-moment SA methods. In this case, the assumption 3L2+2λμ<4μ23L^2+2λμ<4μ^2 boils down to λ<2μλ<2μ, and the upper bound on α becomes α≤1/λα≤ 1/λ which translates into αtλ≤1 _tλ≤ 1 for all t≥0t≥ 0. The proof of Theorem 3.2 is based on Lemma 3.1 and a classical result of Chung [43]. Lemma 3.3 ((43, Lemma 1)). Suppose dt\d_t\ is a sequence of real numbers such that for all t≥1t≥ 1, dt+1≤(1−at)dt+btp+1,d_t+1≤ (1- at )d_t+ bt^p+1, (44) where a>p>0a>p>0 and b>0b>0. Then dt≤ba−p1tp+(1tp+1+1ta).d_t≤ ba-p 1t^p+O ( 1t^p+1+ 1t^a ). In order to apply Chung’s lemma, we need to first derive a uniform upper bound on ‖Xt−X∗‖F2 \|X_t-X_ \|_F^2 for any t≥0t≥ 0, which will be incorporated into the constant b in the lemma. To this end, we start with the one-step improvement inequality (39): ‖Xt+1−X∗‖F2≤(1+αt(αt(3L2+2λμ)−2μ))‖Xt−X∗‖F2+3αt2(λ2‖X∗‖F2+C2). \|X_t+1-X_ \|_F^2≤ (1+ _t ( _t(3L^2+2λμ)-2μ ) ) \|X_t-X_ \|_F^2+3 _t^2(λ^2 \|X_* \|_F^2+C^2). If αt≤2μ3L2+2λμ _t≤ 2μ3L^2+2λμ, we can bound the coefficient of ‖Xt−X∗‖F2 \|X_t-X_ \|_F^2 by 11. Therefore, ‖Xt−X∗‖F2 \|X_t-X_ \|_F^2 ≤‖Xt−1−X∗‖F2+3αt−12(λ2‖X∗‖F2+C2) ≤ \|X_t-1-X_ \|_F^2+3 _t-1^2(λ^2 \|X_* \|_F^2+C^2) ≤‖X0−X∗‖F2+3(λ2‖X∗‖F2+C2)∑s=0t−1αs2, ≤ \|X_0-X_ \|_F^2+3(λ^2 \|X_* \|_F^2+C^2) _s=0^t-1 _s^2, (45) where the last term is bounded if ∑s=0∞αs2<∞ _s=0^∞ _s^2<∞, which is true for αs=αs+1 _s= αs+1. Proof of Theorem 3.2. Under the step size rule αt=αt+1 _t= αt+1 with α≤2μ3L2+lλμα≤ 2μ3L^2+lλμ, we have αt≤2μ3L2+2λμ _t≤ 2μ3L^2+2λμ for all t≥0t≥ 0 and ∑s=0∞αs2≤α2π26 _s=0^∞ _s^2≤ α^2π^26. Therefore, the inequality (45) implies the following bound: ‖Xt−X∗‖F2≤‖X0−X∗‖F2+12(λ2‖X∗‖F2+C2)α2π2=B. \|X_t-X_ \|_F^2≤ \|X_0-X_ \|_F^2+ 12 (λ^2 \|X_* \|_F^2+C^2 )α^2π^2=B. (46) Substituting (46) and αt=αt+1 _t= αt+1 into (39) and noticing the definition of A, we obtain ‖Xt+1−X∗‖F2 \|X_t+1-X_ \|_F^2 ≤(1−2μαt)‖Xt−X∗‖F2+αt2((3L2+2λμ)B+3λ2‖X∗‖F2+3C2) ≤ (1-2μ _t ) \|X_t-X_ \|_F^2+ _t^2 ((3L^2+2λμ)B+3λ^2 \|X_* \|_F^2+3C^2 ) =(1−2μαt+1)‖Xt−X∗‖F2+α2(t+1)2A. = (1- 2μαt+1 ) \|X_t-X_ \|_F^2+ α^2(t+1)^2A. Let dt+1=‖Xt−X∗‖F2d_t+1= \|X_t-X_ \|_F^2, a=2μαa=2μα, p=1p=1 and b=α2Ab=α^2A. The above inequality can be written in the form of (44) with shifted index t→t+1t→ t+1. Applying Lemma 3.3 gives the desired result. ∎ Theorem 3.2 requires the lower bound α>1/(2μ)α>1/(2μ) to obtain the (1/t)O(1/t) rate of convergence. However, it can be hard to obtain a precise estimate of μ in order to satisfy this condition. We can remove this requirement by working with the step size rule αt=α(t+1)p _t= α(t+1)^p with some p∈(12,1)p∈( 12,1). In this case we have an asymptotic (1/tp)O(1/t^p) rate, as stated in the following theorem. Theorem 3.3. Consider the algorithm (32) with step size rule αt=α(t+1)p _t= α(t+1)^p where p∈(12,1)p∈ ( 12,1 ) and the constant α satisfies 0<α≤2μ3L2+2λμ0<α≤ 2μ3L^2+2λμ. Then under Assumption A, we have lim supt→∞(t+1)p‖Xt−X∗‖F2≤αAp2μ, _t→∞\,(t+1)^p \|X_t-X_ \|_F^2≤ α A_p2μ, (47) where Ap=(3L2+2λμ)Bp+3λ2‖X∗‖F2+3C2,Bp=‖X0−X∗‖F2+α2(λ2‖X∗‖F2+C2)∑s=1∞1s2p.A_p=(3L^2+2λμ)B_p+3λ^2 \|X_* \|_F^2+3C^2, B_p= \|X_0-X_ \|_F^2+α^2 (λ^2 \|X_* \|_F^2+C^2 ) _s=1^∞ 1s^2p. The proof is based on another lemma of Chung in the same paper[43], restated below. Lemma 3.4 ((43, Lemma 4)). Suppose dt\d_t\ is a sequence of real numbers such that for all t≥1t≥ 1, dt+1≤(1−attp)dt+btq,d_t+1≤ (1- a_tt^p )d_t+ bt^q, (48) where 0<p<1,p<q,at≥a>0,b>00<p<1,p<q,a_t≥ a>0,b>0. Then lim supt→∞tq−pdt≤ba. _t→∞\,t^q-pd_t≤ ba. Proof of Theorem 3.3. Under the step size rule αt=α(t+1)p _t= α(t+1)^p with α≤2μ3L2+2λμα≤ 2μ3L^2+2λμ and p∈(12,1)p∈( 12,1), we have αt≤2μ3L2+2λμ _t≤ 2μ3L^2+2λμ for all t≥0t≥ 0 and ∑s=0∞αs2=α2∑s=1∞1s2p<∞ _s=0^∞ _s^2=α^2 _s=1^∞ 1s^2p<∞. Therefore, (45) holds and implies the following upper bound for ‖Xt−X∗‖F2 \|X_t-X_* \|_F^2 for all t≥0t≥ 0: ‖Xt−X∗‖F2≤‖X0−X∗‖F2+3α2(λ2‖X∗‖F2+C2)∑s=0∞1(s+1)2p=Bp. \|X_t-X_ \|_F^2≤ \|X_0-X_ \|_F^2+3α^2(λ^2 \|X_* \|_F^2+C^2) _s=0^∞ 1(s+1)^2p=B_p. Substituting the inequality above into (39) and using αt=α(t+1)p _t= α(t+1)^p gives ‖Xt+1−X∗‖F2 \|X_t+1-X_ \|_F^2 ≤(1−2μαt)‖Xt−X∗‖F2+αt2((3L2+2λμ)Bp+3λ2‖X∗‖F2+3C2) ≤ (1-2μ _t ) \|X_t-X_ \|_F^2+ _t^2 ((3L^2+2λμ)B_p+3λ^2 \|X_* \|_F^2+3C^2 ) =(1−2μα(t+1)p)‖Xt−X∗‖F2+α2(t+1)2pAp. = (1- 2μα(t+1)^p ) \|X_t-X_ \|_F^2+ α^2(t+1)^2pA_p. Let dt+1=‖Xt−X∗‖F2d_t+1= \|X_t-X_ \|_F^2, a=2μαa=2μα, b=α2Apb=α^2A_p and q=2pq=2p. Then the last display becomes (48) with shifted index t→t+1t→ t+1 and applying Lemma 3.4 yields the desired result. ∎ 3.3 Convergence of practical algorithms Moving on to the second stage of our analysis, we focus on algorithms of the form Xt+1=(1−αtλ)Xt−αtϕ(M^t,W^t).X_t+1=(1- _tλ)X_t- _tφ( M_t, W_t). (49) We assume that the estimators M^t M_t and W^t W_t are constructed from X0,…,XtX_0,…,X_t and G0,…,GtG_0,…,G_t. Consequently, M^t M_t, W^t W_t and ϕ(M^t,W^t)φ( M_t, W_t) are measurable with respect to the filtration ℱt+1F_t+1 where ℱt:=σ(X0,ξ0,X1,…,ξt−1,Xt).F_t:=σ(X_0, _0,X_1,…, _t-1,X_t). (50) We use t[⋅]E_t[·] to denote the conditional expectation [⋅∣ℱt]E[\,· _t] and define the conditional variance Vart[Z]:=t[‖Z−t[Z]‖F2]. Var_t[Z]:=E_t [\|Z-E_t[Z]\|_F^2 ]. (51) Our main assumption is that the expected deviation of the practical update from its conceptual counterpart, defined as ϵt:=‖t[ϕ(M^t,W^t)]−ϕ(Mt,Wt)‖F, _t:= \|E_t [φ ( M_t, W_t ) ]-φ (M_t,W_t ) \|_F, (52) remains bounded almost surely for all t≥0t≥ 0. In general, ϵt _t depends on the bias and variance of both estimators M^t M_t and W^t W_t as well as their numerical conditioning amplified by the blending map ϕφ. We will derive concrete bounds on ϵt _t for Muon and ASGO in Sections 4 and 5 respectively. Assumption B (Practical update deviation bound). There exists a constant ϵ¯≥0 ε≥ 0 such that the practical update deviation ϵt _t defined in (52) satisfies lim supt→∞ϵt≤ϵ¯almost surely. _t→∞ _t≤ ε surely. The following theorem states that the iterates generated by (49) converge almost surely to a neighborhood of X∗X_ (which is specified in Assumption A) and the radius of the neighborhood is governed by μ and ϵ¯ ε. In fact, a persistent bias in the estimators can prevent almost-sure convergence to X∗X_* itself. We give a simple example in Appendix B. Theorem 3.4 (Convergence of practical algorithms). Suppose that both Assumptions A and B hold, the initialization X0X_0 satisfies [‖X0‖F2]<∞E [\|X_0\|_F^2 ]<∞, and the following conditions are satisfied: (a) Bounded variance: there is a constant b such that for all t≥0t≥ 0, Vart[ϕ(M^t,W^t)]≤balmost surely; Var_t [φ( M_t, W_t) ]≤ b surely; (53) (b) The step sizes αt\ _t\ are deterministic and satisfy αt≥0,∑t=0∞αt=∞,∑t=0∞αt2<∞. _t≥ 0, _t=0^∞ _t=∞, _t=0^∞ _t^2<∞. Then the iterates generated by (49) satisfy lim supt→∞‖Xt−X∗‖F≤ϵ¯μ,almost surely. _t→∞\|X_t-X_*\|_F≤ εμ, surely. (54) The proof of Theorem 3.4 relies on the following extension of Dvoretzky’s theorem [26]. Theorem 3.5 (An extension of Dvoretzky’s theorem). Let (Ω,ℱ,P)( ,F,P) be a probability space with a filtration (ℱt)t≥0(F_t)_t≥ 0, and fix X∗∈m×nX_* ^m× n. Let XtX_t and Tt(X0,…,Xt)T_t(X_0,…,X_t) be finite ℱtF_t-measurable random variables in m×nR^m× n that satisfy, for every t≥0t≥ 0, Xt+1=Tt(X0,…,Xt)+Yt,X_t+1=T_t(X_0,…,X_t)+Y_t, (55) where Yt∈m×nY_t ^m× n is ℱt+1F_t+1-measurable. Suppose that, outside one null set of events, i.e., almost surely, the following inequality holds for all t≥t0t≥ t_0 where t0t_0 is a deterministic integer, ‖Tt(X0,…,Xt)−X∗‖F2≤maxat,(1+bt)‖Xt−X∗‖F2+ct−ht,\|T_t(X_0,…,X_t)-X_*\|_F^2≤ \a_t,(1+b_t)\|X_t-X_*\|_F^2+c_t-h_t \, (56) where ata_t, btb_t, ctc_t and hth_t be finite ℱtF_t-measurable real-valued random variables. Assume that ata_t, btb_t and ctc_t are nonnegative and ht≥0h_t≥ 0 for all sufficiently large t. In addition, there exists a constant a∞∈[0,∞)a_∞∈[0,∞) such that outside the same null set, lim supt→∞at≤a∞,∑t=0∞bt<∞,∑t=0∞ct<∞,∑t=0∞ht=∞. _t→∞\,a_t\,≤\,a_∞, _t=0^∞b_t<∞, _t=0^∞c_t<∞, _t=0^∞h_t=∞. (57) Finally, suppose ∑t=0∞[‖Yt‖F2]<∞,and[Yt∣ℱt]=0almost surely for every t. _t=0^∞E [\|Y_t\|_F^2 ]<∞, [Y_t _t]=0 surely for every t. (58) Then lim supt→∞‖Xt−X∗‖F2≤a∞almost surely. _t→∞\|X_t-X_*\|_F^2≤ a_∞ surely. Remark 3.3. Dvoretzky’s original theorem [26] uses three deterministic coefficient sequences at,bt,ht\a_t\,\b_t\,\h_t\ without ct\c_t\ and assumes at→0a_t→ 0. In the same paper, Dvoretzky extends the same conclusion to filtration dependent coefficients at(ℱt)a_t(F_t), bt(ℱt)b_t(F_t) and ht(ℱt)h_t(F_t). Derman and Sacks [27] introduce an additional summable nonnegative term ct\c_t\ to the recursion and extend the theorem to multidimensional random vectors. Venter [28] further lifts the at→0a_t→ 0 restriction, allowing at=a>0a_t=a>0 to be a positive constant sequence. The present form in Theorem 3.5 combines and generalizes these extensions. In particular, we permit at\a_t\ to be a time-varying sequence not necessarily converging to zero. This was documented as a generalization in Dvoretzky’s paper [26] without proof. For complete reference, we provide a self-contained proof of Theorem 3.5 in Appendix A. Proof of Theorem 3.4. First we can decompose the update in (49) as Xt+1 X_t+1 =(1−αtλ)Xt−αtϕ(M^t,W^t) =(1- _tλ)X_t- _tφ ( M_t, W_t ) =(1−αtλ)Xt−αtt[ϕ(M^t,W^t)]+αt(t[ϕ(M^t,W^t)]−ϕ(M^t,W^t)), =(1- _tλ)X_t- _tE_t [φ ( M_t, W_t ) ]+ _t (E_t [φ ( M_t, W_t ) ]-φ ( M_t, W_t ) ), which matches the formula (55) with Tt(X0,…,Xt) T_t(X_0,…,X_t) =(1−αtλ)Xt−αtt[ϕ(M^t,W^t)], =(1- _tλ)X_t- _tE_t [φ( M_t, W_t) ], Yt Y_t =αt(t[ϕ(M^t,W^t)]−ϕ(M^t,W^t)). = _t (E_t [φ( M_t, W_t) ]-φ( M_t, W_t) ). The above construction automatically implies [Yt∣ℱt]=0E[Y_t _t]=0 for every t≥0t≥ 0. By the bounded variance assumption (53) and the assumption that αt2\ _t^2\ is summable, we have ∑t=0∞[‖Yt‖F2]=∑t=0∞αt2Vart[ϕ(M^t,W^t)] _t=0^∞E [\|Y_t\|_F^2 ]= _t=0^∞ _t^2\, Var_t [φ( M_t, W_t) ] ≤b∑t=0∞αt2<∞. ≤ b _t=0^∞ _t^2<∞. Therefore, the conditions in (58) are satisfied. The squared distance between Tt(X0,…,Xt)T_t(X_0,…,X_t) and X∗X_* can be expanded as ‖Tt(X0,…,Xt)−X∗‖F2 \|T_t(X_0,…,X_t)-X_* \|_F^2 =‖(1−αtλ)(Xt−X∗)−αt(t[ϕ(M^t,W^t)]+λX∗)‖F2 = \|(1- _tλ)(X_t-X_*)- _t (E_t[φ( M_t, W_t)]+λ X_* ) \|_F^2 =(1−αtλ)2‖Xt−X∗‖F2+αt2‖t[ϕ(M^t,W^t)]+λX∗‖F2 =(1- _tλ)^2\|X_t-X_*\|_F^2~+~ _t^2 \|E_t[φ( M_t, W_t)]+λ X_* \|_F^2 −2αt(1−αtλ)⟨Xt−X∗,t[ϕ(M^t,W^t)]+λX∗⟩. -2 _t(1- _tλ) X_t-X_*,\,E_t[φ( M_t, W_t)]+λ X_* . (59) For the cross term, we proceed by adding and subtracting ϕ(Mt,Wt)φ (M_t,W_t ) inside the inner product: ⟨t[ϕ(M^t,W^t)]+λX∗,Xt−X∗⟩ _t\! [φ ( M_t, W_t ) ]+λ X_*,\,X_t-X_ = = ⟨t[ϕ(M^t,W^t)]+λX∗+ϕ(Mt,Wt)−ϕ(Mt,Wt),Xt−X∗⟩ _t\! [φ ( M_t, W_t ) ]+λ X_*+φ (M_t,W_t )-φ (M_t,W_t ),\,X_t-X_ = = ⟨ϕ(Mt,Wt)+λX∗,Xt−X∗⟩+⟨t[ϕ(M^t,W^t)]−ϕ(Mt,Wt),Xt−X∗⟩ φ (M_t,W_t )+λ X_*,\,X_t-X_* + _t [φ ( M_t, W_t ) ]-φ (M_t,W_t ),\,X_t-X_* ≥ ≥ (μ−λ)‖Xt−X∗‖F2−‖t[ϕ(M^t,W^t)]−ϕ(Mt,Wt)‖F‖Xt−X∗‖F (μ-λ)\|X_t-X_*\|_F^2- \|E_t [φ ( M_t, W_t ) ]-φ (M_t,W_t ) \|_F \|X_t-X_ \|_F = = (μ−λ)‖Xt−X∗‖F2−ϵt‖Xt−X∗‖F, (μ-λ)\|X_t-X_*\|_F^2- _t \|X_t-X_ \|_F, (60) where in the inequality we used the equivalent aiming condition (40) and the Cauchy-Shwartz inequality, and in the last equality we used the definition of ϵt _t in (52). For the second squared term in (59), we apply the triangle inequality followed by Young’s inequality: ‖t[ϕ(M^t,W^t)]+λX∗‖F2≤2‖t[ϕ(M^t,W^t)]‖F2+2λ2‖X∗‖F2. \|E_t[φ( M_t, W_t)]+λ X_* \|_F^2≤ 2 \|E_t[φ( M_t, W_t)] \|_F^2+2λ^2\|X_*\|_F^2. Next, we decompose the conditional mean into the conceptual update and its deviation to obtain ‖t[ϕ(M^t,W^t)]‖F2 \|E_t[φ( M_t, W_t)] \|_F^2 ≤2‖ϕ(Mt,Wt)‖F2+2‖t[ϕ(M^t,W^t)]−ϕ(Mt,Wt)‖F2 ≤ 2\|φ(M_t,W_t)\|_F^2+2 \|E_t[φ( M_t, W_t)]-φ(M_t,W_t) \|_F^2 ≤2(L‖Xt−X∗‖F+C)2+2ϵt2 ≤ 2 (L\|X_t-X_*\|_F+C )^2+2 _t^2 ≤4L2‖Xt−X∗‖F2+4C2+2ϵt2, ≤ 4L^2\|X_t-X_*\|_F^2+4C^2+2 _t^2, where we used the growth bound (38) and definition of ϵt _t. Define the random variable δt:=8C2+4ϵt2+2λ2‖X∗‖F2. _t:=8C^2+4 _t^2+2λ^2\|X_*\|_F^2. (61) Then we have ‖t[ϕ(M^t,W^t)]+λX∗‖F2≤8L2‖Xt−X∗‖F2+δt. \|E_t[φ( M_t, W_t)]+λ X_* \|_F^2≤ 8L^2\|X_t-X_*\|_F^2+ _t. (62) Since ∑t=0∞αt2<∞ _t=0^∞ _t^2<∞, we have αt→0 _t→ 0. Because the step sizes αt\ _t\ are deterministic, there exists a deterministic t0t_0 such that 0≤1−αtλ≤10≤ 1- _tλ≤ 1 holds for all t≥t0t≥ t_0. For every t≥t0t≥ t_0, substituting the bounds in (60) and (62) into (59) yields ‖Tt(X0,…,Xt)−X∗‖F2 \|T_t(X_0,…,X_t)-X_*\|_F^2 ≤(1−αtλ)2‖Xt−X∗‖F2−2αt(1−αtλ)(μ−λ)‖Xt−X∗‖F2 ≤(1- _tλ)^2\|X_t-X_*\|_F^2-2 _t(1- _tλ)(μ-λ)\|X_t-X_*\|_F^2 +2αt(1−αtλ)ϵt‖Xt−X∗‖F+8αt2L2‖Xt−X∗‖F2+αt2δt +2 _t(1- _tλ) _t\|X_t-X_*\|_F+8 _t^2L^2\|X_t-X_*\|_F^2+ _t^2 _t ≤(1+αt2(2λμ+8L2))‖Xt−X∗‖F2−2αtμ‖Xt−XF‖F2 ≤ (1+ _t^2(2λμ+8L^2) )\|X_t-X_*\|_F^2-2 _tμ\|X_t-X_F\|_F^2 +2αtϵt‖Xt−X∗‖F+αt2δt. +2 _t _t\|X_t-X_*\|_F+ _t^2 _t. Young’s inequality gives 2ϵt‖Xt−X∗‖F≤ϵt2μ+μ‖Xt−X∗‖F2.2 _t\|X_t-X_*\|_F≤ _t^2μ+μ\|X_t-X_*\|_F^2. Consequently, ‖Tt(X0,…,Xt)−X∗‖F2≤(1+αt2(2λμ+8L2))‖Xt−X∗‖F2−αt(μ‖Xt−X∗‖F2−ϵt2μ)+αt2δt.\|T_t(X_0,…,X_t)-X_*\|_F^2≤ (1+ _t^2(2λμ+8L^2) )\|X_t-X_*\|_F^2- _t (μ\|X_t-X_*\|_F^2- _t^2μ )+ _t^2 _t. Let ϵ¯ ε be the constant specified in Assumption B and fix an arbitrary number η>0η>0. When ‖Xt−X∗‖F2>ϵ¯2/μ2+η \|X_t-X_ \|_F^2> ε^2/μ^2+η, we deduce from the above inequality that ‖Tt(X0,…,Xt)−X∗‖F2≤(1+αt2(2λμ+8L2))‖Xt−X∗‖F2−μαt(ϵ¯2μ2+η−ϵt2μ2)+αt2δt. \|T_t(X_0,…,X_t)-X_ \|_F^2≤ (1+ _t^2(2λμ+8L^2) ) \|X_t-X_ \|_F^2-μ _t ( ε^2μ^2+η- _t^2μ^2 )+ _t^2 _t. Otherwise, i.e., when ‖Xt−X∗‖F2≤ϵ¯2/μ2+η \|X_t-X_ \|_F^2≤ ε^2/μ^2+η, we have ‖Tt(X0,…,Xt)−X∗‖F2≤(1+αt2(2λμ+8L2))(ϵ¯2μ2+η)+αtϵt2μ+αt2δt. \|T_t(X_0,…,X_t)-X_ \|_F^2≤ (1+ _t^2(2λμ+8L^2) ) ( ε^2μ^2+η )+ _t _t^2μ+ _t^2 _t. We can combine these two cases as ‖Tt(X0,…,Xt)−X∗‖F2≤maxat,(1+bt)‖Xt−X∗‖F2+ct−ht, \|T_t(X_0,…,X_t)-X_ \|_F^2≤ \a_t,(1+b_t) \|X_t-X_ \|_F^2+c_t-h_t \, with the definitions at:=(1+αt2(2λμ+8L2))(ϵ¯2μ2+η)+αtϵt2μ+αt2δt,a_t:= (1+ _t^2(2λμ+8L^2) ) ( ε^2μ^2+η )+ _t _t^2μ+ _t^2 _t, and bt:=αt2(2λμ+8L2),ct:=αt2δt,ht:=μαt(ϵ¯2μ2+η−ϵt2μ2).b_t:= _t^2(2λμ+8L^2), c_t:= _t^2 _t, h_t:=μ _t ( ε^2μ^2+η- _t^2μ^2 ). Thus (56) holds for every t≥t0t≥ t_0 where t0t_0 is the smallest t such that αt≤1/λ _t≤ 1/λ. To apply Theorem 3.5, it remains to verify the conditions in (57). Assumption B implies that ϵt _t, and hence δt _t through the definition in (61), are bounded for sufficiently large t. Since αt→0 _t→ 0, we have a∞=limt→∞at=ϵ¯2μ2+η<∞.a_∞= _t→∞a_t= ε^2μ^2+η<∞. Moreover, we have ∑t=0∞bt=(2λμ+8L2)∑t=0∞αt2<∞,∑t=0∞ct=∑t=0∞αt2δt<∞almost surely. _t=0^∞b_t=(2λμ+8L^2) _t=0^∞ _t^2<∞, _t=0^∞c_t= _t=0^∞ _t^2 _t<∞ surely. The assumption lim supt→∞ϵt≤ϵ¯ _t→∞ _t≤ ε also implies that there exists a finite integer tηt_η such that ϵt2μ2≤ϵ¯2μ2+η2,∀t≥tη. _t^2μ^2≤ ε^2μ^2+ η2, ∀\,t≥ t_η. Therefore, for every t≥tηt≥ t_η, ht=μαt(ϵ¯2μ2+η−ϵt2μ2)≥μη2αt≥0,h_t=μ _t ( ε^2μ^2+η- _t^2μ^2 )≥ μη2 _t≥ 0, and hence ∑t=tη∞ht≥μη2∑t=tη∞αt=∞. _t=t_η^∞h_t≥ μη2 _t=t_η^∞ _t=∞. The omitted prefix contains only finitely many bounded terms, therefore ∑t=0∞ht=∞ _t=0^∞h_t=∞ almost surely. With all conditions satisfied, we can apply Theorem 3.5 to obtain lim supt→∞‖Xt−X∗‖F2≤ϵ¯2μ2+ηalmost surely. _t→∞\|X_t-X_*\|_F^2≤ ε^2μ^2+η surely. We apply this conclusion with η=1/kη=1/k for k=1,2,3,…k=1,2,3,… and intersect the resulting countable family of probability-one events. As k→∞k→∞, this proves lim supt→∞‖Xt−X∗‖F2≤ϵ¯2μ2almost surely. _t→∞\|X_t-X_*\|_F^2≤ ε^2μ^2 surely. Taking square roots on both sides of the above inequality gives (54). ∎ Theorem 3.4 provides convergence guarantee for a broad class of practical second-moment SA methods. To apply it to any particular method, we need to characterize the sequence ϵt _t defined in (52) under specific algorithmic settings. In general, it is determined by the sensitivity of the blending function ϕφ under perturbation as well as the biases and variances of both the first- and second-moment estimators. In the next two sections, we present concrete bounds on ϵt _t for Muon and ASGO respectively. 4 Deviation bound for Muon We consider (the ideal version of) Muon with decoupled weight decay: Xt+1=(1−αtλ)Xt−αtOrth(M^t),X_t+1=(1- _tλ)X_t- _t Orth( M_t), (63) where M^t∈m×n M_t ^m× n is an estimator of the first moment M(Xt)=t[G(Xt,ξt)]M(X_t)=E_t[G(X_t, _t)]. Recall that for a matrix M of rank r>0r>0 with compact SVD M=UΣVTM=U V^T where U∈m×rU ^m× r and V∈n×rV ^n× r, we have Orth(M)=UVT. Orth(M)=UV^T. Throughout this section, we assume m≤nm≤ n and the following assumption. Assumption C (Full-rank first-moment estimator). For every t≥0t≥ 0, we have t[‖M^t‖F]<∞E_t [\| M_t\|_F ]<∞ almost surely and Pr(rank(M^t)=m∣ℱt)=1,rank(t[M^t])=rank(Mt)=malmost surely. (rank( M_t)=m _t )=1, (E_t[ M_t])=rank(M_t)=m surely. Under the above assumption, we have the following bound on the deviation defined in (52). Theorem 4.1 (Deviation bound for Muon). Suppose Assumption C holds and t[‖M^t‖F2]<∞E_t [\| M_t\|_F^2 ]<∞ almost surely for every t≥0t≥ 0. Then the following holds almost surely for every t≥0t≥ 0: ϵt:=‖t[Orth(M^t)]−Orth(Mt)‖F≤3σ^tVart[M^t]+4σ^t+σt‖t[M^t]−Mt‖F, _t:= \|E_t[ Orth( M_t)]- Orth(M_t) \|_F~≤~ 3 σ_t Var_t[ M_t]+ 4 σ_t+ _t \|E_t[ M_t]-M_t \|_F, (64) where σ^t σ_t and σt _t are the smallest positive singular values of t[M^t]E_t[ M_t] and MtM_t respectively. The bound in (64) accounts for the estimator’s bias and variance, and the numerical conditioning of both MtM_t and t[M^t]E_t[ M_t] through their smallest singular values. According to Theorem 3.4, its asymptotic value (more precisely limit superior) of ϵt _t bound the size of the neighborhood around X∗X_* where the algorithm (63) converges to. The key to the proof of Theorem 4.1 hinges on the sensitivity of the mapping Orth(⋅) Orth(·) to perturbations in the input matrix. The following lemma characterizes this stability in terms of unitarily invariant norms |||⋅||| | -1.07639pt | -1.07639pt |· | -1.07639pt | -1.07639pt |. Lemma 4.1 (Wide matrix variant of (44, Theorem 2)). Suppse m≤nm≤ n and A,A~∈m×nA, A ^m× n have full row rank. Let Q=Orth(A)Q= Orth(A), Q~=Orth(A~) Q= Orth( A), and σm _m and σ~m σ_m be the smallest singular values of A and A~ A respectively. Then for any unitarily invariant norm |||⋅||| | -1.07639pt | -1.07639pt |· | -1.07639pt | -1.07639pt |, ‖|Q−Q~|‖≤(2σm+σ~m+1maxσm,σ~m)‖|A−A~|‖. | -1.07639pt | -1.07639pt |Q- Q | -1.07639pt | -1.07639pt |≤ ( 2 _m+ σ_m+ 1 \ _m, σ_m\ ) | -1.07639pt | -1.07639pt |A- A | -1.07639pt | -1.07639pt |. (65) Lemma 4.1 establishes that the stability of Orth(⋅) Orth(·) is governed by the input error ‖|A−A~|‖ | -1.07639pt | -1.07639pt |A- A | -1.07639pt | -1.07639pt |, scaled inversely by the inputs’ smallest singular values. Consequently, well-conditioned matrices yield stable orthogonal factors, whereas near-rank deficiency severely amplifies perturbation errors. This lemma is the wide matrix variant of Theorem 2 in [44], which assumes that A and A~ A are tall matrices with m≥nm≥ n. The proofs for the two cases are similar so we omit the one here. In order to derive the bound (64), we first decompose the deviation ϵt _t as follows: ϵt:= _t:= ‖t[Orth(M^t)]−Orth(Mt)‖F \|E_t[ Orth( M_t)]- Orth(M_t) \|_F = = ‖t[Orth(M^t)]−Orth(t[M^t])+Orth(t[M^t])−Orth(Mt)‖F \|E_t[ Orth( M_t)]- Orth(E_t[ M_t])+ Orth(E_t[ M_t])- Orth(M_t) \|_F ≤ ≤ ‖t[Orth(M^t)]−Orth(t[M^t])‖F+‖Orth(t[M^t])−Orth(Mt)‖F. \|E_t[ Orth( M_t)]- Orth(E_t[ M_t]) \|_F+ \| Orth(E_t[ M_t])- Orth(M_t) \|_F. (66) We bound the two terms separately. The following lemma shows that the first term captures the error induced by the variance of the estimator M^t M_t, amplified by the conditioning of Orth(⋅) Orth(·). Lemma 4.2. Suppose Assumption C holds and t[‖M^t‖F2]<∞E_t [\| M_t\|_F^2 ]<∞ almost surely for all t≥0t≥ 0. Then, ‖t[Orth(M^t)]−Orth(t[M^t])‖F≤3σ^tVart[M^t]almost surely. \|E_t[ Orth( M_t)]- Orth(E_t[ M_t]) \|_F≤ 3 σ_t Var_t[ M_t] surely. where σ^t σ_t is the smallest positive singular values of t[M^t]E_t[ M_t]. Proof. First, we use the linearity of expectation and Jensen’s inequality to obtain ‖t[Orth(M^t)]−Orth(t[M^t])‖F \|E_t[ Orth( M_t)]- Orth(E_t[ M_t]) \|_F =‖t[Orth(M^t)−Orth(t[M^t])]‖F = \|E_t[ Orth( M_t)- Orth(E_t[ M_t])] \|_F ≤t[‖Orth(M^t)−Orth(t[M^t])‖F]. _t [ \| Orth( M_t)- Orth(E_t[ M_t]) \|_F ]. Assumption C allows us to apply Lemma 4.1 so that almost surely for every t≥0t≥ 0: t[‖Orth(M^t)−Orth(t[M^t])‖F] _t [ \| Orth( M_t)- Orth(E_t[ M_t]) \|_F ] ≤t[(2σ~t+σ^t+1maxσ~t,σ^t)‖M^t−t[M^t]‖F] _t [ ( 2 σ_t+ σ_t+ 1 \ σ_t, σ_t\ ) \| M_t-E_t[ M_t] \|_F ] ≤3σ^tt[‖M^t−t[M^t]‖F], ≤ 3 σ_tE_t [ \| M_t-E_t[ M_t] \|_F ], where σ~t σ_t and σ^t σ_t are the smallest positive singular values of M^t M_t and t[M^t]E_t[ M_t] respectively. To further simplify the bound, recall that Frobenius norm is the square-root of the sum of the squares of all elements, which allows us to rewrite the upper bound as follows: 3σ^tt[‖M^t−t[M^t]‖F]=3σ^tt[∑i,j(M^t(i,j)−t[M^t(i,j)])2]. 3 σ_tE_t [ \| M_t-E_t[ M_t] \|_F ]= 3 σ_tE_t [ _i,j ( M_t(i,j)-E_t [ M_t(i,j) ] )^2 ]. Applying Jensen’s inequality again to interchange expectation and the square root, we obtain: 3σ^tt[∑i,j(M^t(i,j)−t[M^t(i,j)])2] 3 σ_tE_t [ _i,j ( M_t(i,j)-E_t[ M_t(i,j)] )^2 ] ≤3σ^t∑i,jt[(M^t(i,j)−t[M^t(i,j)])2] ≤ 3 σ_t _i,jE_t [ ( M_t(i,j)-E_t[ M_t(i,j)] )^2 ] =3σ^tVart[M^t]. = 3 σ_t Var_t[ M_t]. Combining the chain of inequalities, we obtain the desired result. ∎ The second term on the right-hand side of (66) accounts for the perturbation of Orth(⋅) Orth(·) due to the bias of M^t M_t relative to MtM_t. Lemma 4.3. Under Assumption C, the following error bound holds almost surely for every t: ‖Orth(t[M^t])−Orth(Mt)‖F≤4σ^t+σt‖t[M^t]−Mt‖F. \| Orth(E_t[ M_t])- Orth(M_t) \|_F≤ 4 σ_t+ _t \|E_t[ M_t]-M_t \|_F. where σ^t σ_t and σt _t are the smallest positive singular values of t[M^t]E_t[ M_t] and MtM_t respectively. Proof. In this case, we can apply Lemma 4.1 directly to obtain ‖Orth(t[M^t])−Orth(Mt)‖F \| Orth(E_t[ M_t])- Orth(M_t) \|_F ≤(2σ^t+σt+1maxσ^t,σt)‖t[M^t]−Mt‖F. ≤ ( 2 σ_t+ _t+ 1 \ σ_t, _t\ ) \|E_t[ M_t]-M_t \|_F. Moreover, since maxσ^t,σt≥12(σ^t+σt) \ σ_t, _t\≥ 12( σ_t+ _t), the coefficient can be further upper bounded as (2σ^t+σt+1maxσ^t,σt)≤4σ^t+σt, ( 2 σ_t+ _t+ 1 \ σ_t, _t\ )≤ 4 σ_t+ _t, which gives the desired result. ∎ Theorem 4.1 is a direct consequence of (66) and the combination of Lemma 4.2 and 4.3. In order to apply Theorem 3.4 to Muon, we need to check Assumptions A and B as well as the bounded variance condition (53). For Assumption A, we gave an examples of M(X)M(X) that satisfies the aiming condition at the end of Section 3.1 and provide empirical evidences for deep learning in Section 6. In addition, the fact ‖Orth(M(X))‖F≤m\| Orth(M(X))\|_F≤ m implies that the growth condition (38) holds with L=0L=0 and C=mC= m. For Assumption B, it is sufficient to assume that the first-moment estimator M^t M_t has bounded bias and variance, then Theorem 4.1 implies that the expected deviation ϵt _t is bounded almost surely. Finally, we have Vart[Orth(M^t)]≤t[‖Orth(M^t)‖F2]≤m, Var_t [ Orth( M_t) ] _t [\| Orth( M_t)\|_F^2 ]≤ m, which means that the bounded variance condition (53) holds with b=mb=m. 5 Deviation bound for ASGO We consider ASGO [11] with decoupled weight decay: Xt+1=(1−λαt)Xt−αtW^t−1/2M^t,X_t+1=(1-λ _t)X_t- _t W_t^-1/2 M_t, (67) where M^t M_t is an estimator for the first-moment Mt=t[Gt]M_t=E_t[G_t] and W^t W_t is an estimator for the second moment Wt=t[GtGtT]W_t=E_t[G_tG_t^T]. We assume that both the second-moment WtW_t and its estimator W^t W_t are positive definite with uniform lower bounds on their eigenvalues. Assumption D (Positive-definiteness of the second moment). The second-moment estimator W^t W_t is symmetric and positive definite for all t≥0t≥ 0. There are deterministic constants λm,λ^m>0 _m, λ_m>0 such that Wt⪰λmIW_t _mI and W^t⪰λ^mI W_t λ_mI hold with probability one for every t≥0t≥ 0. For ASGO, the deviation defined in (52) specializes to ϵt=∥t[W^t−1/2M^t]−Wt−1/2Mt∥F _t= \|E_t[ W_t^-1/2 M_t]-W_t^-1/2M_t \|_F, for which we have the following bound. Theorem 5.1 (Deviation bound for ASGO). Suppose Assumption D holds and t[‖M^t‖F2]<∞E_t [\| M_t\|_F^2 ]<∞ almost surely. Then the following bound holds almost surely for each t≥0t≥ 0: ϵt _t ≤λ^m−3/22Vart[W^t](‖t[M^t]‖F+Vart[M^t]) ~≤~ λ_m^-3/22 Var_t [ W_t ] ( \|E_t[ M_t] \|_F+ Var_t[ M_t] ) (68) +(λmλ^m1/2+λ^mλm1/2)−1∥t[W^t]−Wt∥F∥t[M^t]∥2+λm−1/2∥t[M^t]−Mt∥F, ~+ ( _m λ_m^1/2+ λ_m _m^1/2 )^-1 \|E_t[ W_t]-W_t \|_F \|E_t[ M_t] \|_2~+~ _m^-1/2 \|E_t[ M_t]-M_t \|_F, where λ^m λ_m and λm _m are the uniform lower bound on the eigenvalues of W^t W_t and WtW_t respectively. Theorem 5.1 establishes a deviation bound that is governed by the bias, variance, and numerical conditioning of the estimators M^t M_t and W^t W_t. The following lemma captures the numerical conditioning of the matrix square-root inverse, as part of the blending function ϕ(M,W)=W−1/2Mφ(M,W)=W^-1/2M. Lemma 5.1. Let A,B∈m×mA,B ^m× m be symmetric positive definite matrices with smallest eigenvalues λA>0 _A>0 and λB>0 _B>0 respectively. Then ∥A−1/2−B−1/2∥F≤(λAλB1/2+λBλA1/2)−1∥B−A∥F. \|A^-1/2-B^-1/2 \|_F≤ ( _A _B^1/2+ _B _A^1/2 )^-1\|B-A\|_F. (69) The proof of Lemma 5.1 is given in Appendix C. To prove Theorem 5.1, we follow a similar strategy for Muon by first decomposing the deviation ϵt _t as follows: ϵt:= _t:= ∥t[W^t−1/2M^t]−Wt−1/2Mt∥F \|E_t [ W_t^-1/2 M_t ]-W_t^-1/2M_t \|_F = = ∥t[W^t−1/2M^t]−t[W^t]−1/2t[M^t]+t[W^t]−1/2t[M^t]−Wt−1/2Mt∥F \|E_t [ W_t^-1/2 M_t ]-E_t [ W_t ]^-1/2E_t[ M_t]+E_t[ W_t]^-1/2E_t[ M_t]-W_t^-1/2M_t \|_F ≤ ≤ ∥t[W^t−1/2M^t]−t[W^t]−1/2t[M^t]∥F+∥t[W^t]−1/2t[M^t]−Wt−1/2Mt∥F. \|E_t[ W_t^-1/2 M_t]-E_t[ W_t]^-1/2E_t[ M_t] \|_F+ \|E_t[ W_t]^-1/2E_t[ M_t]-W_t^-1/2M_t \|_F. (70) We bound these two terms separately through the next two lemmas. Lemma 5.2. Suppose Assumption D holds and t[‖M^t‖F2]<∞E_t [\| M_t\|_F^2 ]<∞ and t[‖W^t‖F2]<∞E_t [\| W_t\|_F^2 ]<∞ almost surely for every t≥0t≥ 0. Then we have for every t≥0t≥ 0, almost surely, ∥t[W^t−1/2M^t]−t[W^t]−1/2t[M^t]∥F≤λ^m−3/22Vart[W^t](∥t[M^t]∥F+Vart[M^t]), \|E_t[ W_t^-1/2 M_t]-E_t[ W_t]^-1/2E_t[ M_t] \|_F≤ λ_m^-3/22 Var_t [ W_t ] ( \|E_t[ M_t] \|_F+ Var_t[ M_t] ), where λ^m λ_m is the uniform lower bound on the eigenvalue of W^t W_t in Assumption D. Proof. First, by definition of the conditional expectation, we have t[W^t]−1/2t[M^t]=t[t[W^t]−1/2M^t],E_t[ W_t]^-1/2E_t[ M_t]=E_t[E_t[ W_t]^-1/2 M_t], which allows us to rewrite the error term into a form that can apply Jensen’s inequality: ∥t[W^t−1/2M^t]−t[W^t]−1/2t[M^t]∥F \|E_t[ W_t^-1/2 M_t]-E_t[ W_t]^-1/2E_t[ M_t] \|_F =∥t[(W^t−1/2−t[W^t]−1/2)M^t]∥F = \|E_t [ ( W_t^-1/2-E_t[ W_t]^-1/2 ) M_t ] \|_F ≤t[∥(W^t−1/2−t[W^t]−1/2)M^t∥F]. _t [ \| ( W_t^-1/2-E_t[ W_t]^-1/2 ) M_t \|_F ]. Since the Frobenius norm is submultiplicative, we can further upper bound the norm of a matrix product by the product of matrix norms and then apply the Cauchy-Schwarz inequality: t[∥(W^t−1/2−t[W^t]−1/2)M^t∥F] _t [ \| ( W_t^-1/2-E_t[ W_t]^-1/2 ) M_t \|_F ] ≤t[∥W^t−1/2−t[W^t]−1/2∥F∥M^t∥F] _t [ \| W_t^-1/2-E_t[ W_t]^-1/2 \|_F \| M_t \|_F ] ≤t[∥W^t−1/2−t[W^t]−1/2∥F2]t[‖M^t‖F2]. ≤ E_t [ \| W_t^-1/2-E_t[ W_t]^-1/2 \|_F^2 ] E_t [ \| M_t \|_F^2 ]. (71) Next we bound the two expected squared Frobenius norms separately. For t[∥W^t−1/2−t[W^t]−1/2∥F2]E_t [ \| W_t^-1/2\!-E_t[ W_t]^-1/2 \|_F^2 ], we apply Lemma 5.1 together with the fact λm≥λ^m _m≥ λ_m to obtain t[∥W^t−1/2−t[W^t]−1/2∥F2] _t [ \| W_t^-1/2-E_t[ W_t]^-1/2 \|_F^2 ] ≤t[(12λ^m−3/2∥W^t−t[W^t]∥F)2] _t\! [ ( 12 λ_m^-3/2 \| W_t-E_t[ W_t] \|_F )^2 ] (72) =λ^m−34t[‖W^t−t[W^t]‖F2] = λ_m^-34E_t\! [ \| W_t-E_t[ W_t] \|_F^2 ] =λ^m−34Vart[W^t]. = λ_m^-34 Var_t[ W_t]. For the second-moment term t[‖M^t‖F2]E_t [ \| M_t \|_F^2 ], we have the mean-variance decomposition t[‖M^t‖F2]=‖t[M^t]‖F2+Vart[M^t].E_t [\| M_t\|_F^2 ]=\|E_t[ M_t]\|_F^2+ Var_t[ M_t]. (73) Substituting the identity (73) and the bound (72) into (71), we get ∥t[W^t−1/2M^t]−t[W^t]−1/2t[M^t]∥F \|E_t[ W_t^-1/2 M_t]-E_t[ W_t]^-1/2E_t[ M_t] \|_F ≤λ^m−34Vart[W^t]‖t[M^t]‖F2+Vart[M^t] ≤ λ_m^-34 Var_t[ W_t] \|E_t[ M_t] \|_F^2+ Var_t[ M_t] ≤λ^m−3/22Vart[W^t](‖t[M^t]‖F+Vart[M^t]), ≤ λ_m^-3/22 Var_t [ W_t ] ( \|E_t[ M_t] \|_F+ Var_t[ M_t] ), where the last inequality follows from sub-additivity of the square-root function. ∎ The next lemma bounds the second term on the right side of (70). Lemma 5.3. Suppose Assumption D holds and t[‖W^t‖F]<∞E_t [\| W_t\|_F ]<∞ and t[‖M^t‖F]<∞E_t [\| M_t\|_F ]<∞ for every t≥0t≥ 0. Then, almost surely for every t≥0t≥ 0, ∥t[W^t]−1/2t[M^t]−Wt−1/2Mt∥F \|E_t[ W_t]^-1/2E_t[ M_t]-W_t^-1/2M_t \|_F ≤ ≤~ (λmλ^m1/2+λ^mλm1/2)−1∥t[W^t]−Wt∥F∥t[M^t]∥2+λm−1/2∥t[M^t]−Mt∥F, ( _m λ_m^1/2+ λ_m _m^1/2 )^-1 \|E_t[ W_t]-W_t \|_F \|E_t[ M_t] \|_2~+~ _m^-1/2 \|E_t[ M_t]-M_t \|_F, where λm _m and λ^m λ_m are the uniform lower bounds for the eigenvalues of WtW_t and W^t W_t respectively. Proof. Adding and subtracting Wt−1/2t[M^t]W_t^-1/2E_t[ M_t] inside the norm and rearranging the terms, we get ∥t[W^t]−1/2t[M^t]−Wt−1/2Mt∥F \|E_t[ W_t]^-1/2E_t[ M_t]-W_t^-1/2M_t \|_F = =~ ∥t[W^t]−1/2t[M^t]−Wt−1/2t[M^t]+Wt−1/2t[M^t]−Wt−1/2Mt∥F \|E_t[ W_t]^-1/2E_t[ M_t]-W_t^-1/2E_t[ M_t]+W_t^-1/2E_t[ M_t]-W_t^-1/2M_t \|_F = =~ ∥(t[W^t]−1/2−Wt−1/2)t[M^t]+Wt−1/2(t[M^t]−Mt)∥F. \| (E_t[ W_t]^-1/2-W_t^-1/2 )E_t[ M_t]+W_t^-1/2 (E_t[ M_t]-M_t ) \|_F. Applying the triangle inequality followed by the submultiplicativity of matrix norms, we get ∥(t[W^t]−1/2−Wt−1/2)t[M^t]+Wt−1/2(t[M^t]−Mt)∥F \| (E_t[ W_t]^-1/2-W_t^-1/2 )E_t[ M_t]+W_t^-1/2 (E_t[ M_t]-M_t ) \|_F ≤ ≤~ ∥(t[W^t]−1/2−Wt−1/2)t[M^t]∥F+∥Wt−1/2(t[M^t]−Mt)∥F \| (E_t[ W_t]^-1/2-W_t^-1/2 )E_t[ M_t] \|_F+ \|W_t^-1/2 (E_t[ M_t]-M_t ) \|_F ≤ ≤~ ∥t[W^t]−1/2−Wt−1/2∥F∥t[M^t]∥2+∥Wt−1/2∥2∥t[M^t]−Mt∥F. \|E_t[ W_t]^-1/2-W_t^-1/2 \|_F \|E_t[ M_t] \|_2+ \|W_t^-1/2 \|_2 \|E_t[ M_t]-M_t \|_F. (74) Again we bound the two resulting terms separately. For the first one, we use Lemma 5.1 to obtain ∥t[W^t]−1/2−Wt−1/2∥F≤(λmλ^m1/2+λ^mλm1/2)−1∥t[W^t]−Wt∥F \|E_t[ W_t]^-1/2-W_t^-1/2 \|_F≤ ( _m λ_m^1/2+ λ_m _m^1/2 )^-1 \|E_t[ W_t]-W_t \|_F For the second term on the right-hand side of (74), we have ∥Wt−1/2∥2=λmin(Wt)−1/2≤λm−1/2. \|W_t^-1/2 \|_2= _ (W_t )^-1/2≤ _m^-1/2. Substituting the above two bounds into (74), we get ∥t[W^t]−1/2−Wt−1/2∥F∥t[M^t]∥2+∥Wt−1/2∥2∥t[M^t]−Mt∥F \|E_t[ W_t]^-1/2-W_t^-1/2 \|_F \|E_t[ M_t] \|_2+ \|W_t^-1/2 \|_2 \|E_t[ M_t]-M_t \|_F ≤ ≤~ (λmλ^m1/2+λ^mλm1/2)−1∥t[W^t]−Wt∥F∥t[M^t]∥2+λm−1/2∥t[M^t]−Mt∥F. ( _m λ_m^1/2+ λ_m _m^1/2 )^-1 \|E_t[ W_t]-W_t \|_F \|E_t[ M_t] \|_2+ _m^-1/2 \|E_t[ M_t]-M_t \|_F. This finishes the proof. ∎ Theorem 5.1 is a direct consequence of (70) and the combination of Lemmas 5.2 and 5.3. In order to apply Theorem 3.4 to ASGO, we need Assumption B to hold, which requires Assumption D as well as t[‖M^t‖F2]E_t [\| M_t\|_F^2 ] being bounded. The uniform lower bound on the eigenvalues in Assumption D is not a limiting restriction because most practical algorithms add a small offset ϵIε I to their second-moment estimators. For Assumption A, the growth condition (38) holds with L=0L=0 and C=mC= m because ∥Wt−1/2Mt∥F2=Tr(Wt−1MtMtT)=Tr(t[GtGtT]−1t[Gt]t[Gt]T)≤m. \|W_t^-1/2M_t \|_F^2= Tr (W_t^-1M_tM_t^T )= Tr (E_t[G_tG_t^T]^-1E_t[G_t]E_t[G_t]^T )≤ m. We provide empirical evidence for the aiming condition (37) in Section 6. 6 Numerical experiments In this section, we present numerical experiments to illustrate some characteristics of second-moment SA methods captured by our theory. We focus on training a large language model using ASGO [11] and Muon [12], listed in the equations (8) and (9) respectively. Both methods can be interpreted through the optimal fused preconditioner in Section 2.1, but with simplifications assuming different noise structures, presented in Sections 2.2 and 2.3 respectively. Figure 2: Empirical evaluation of the aiming condition for training a large language model with Muon (β1=0.95)( _1=0.95). The y axis shows the absolute values of μt(Dt) _t(D_t) at each iteration. Left: μt _t computed for the whole group with various directions Dt=GtD_t=G_t, M^t M_t, Orth(M^t) Orth( M_t) and Orth(M^t)+λXt Orth( M_t)+λ X_t. Right: μt _t computed for Dt=Orth(M^t)+λXtD_t= Orth( M_t)+λ X_t but over various parameter groups. All experiments are conducted to train a LLaMA-style model [45] with 1.81 Billion parameters (with bf16 precision). The model has 2525 layers, 20482048 hidden dimensions, 1616 attention heads, 8/38/3 MLP intermediate expansion and 40964096 sequence length on the FineWeb-Edu dataset [46]. Training is done with a batch size of 2 million (2048K) tokens, distributed across 32 NVIDIA H100 GPUs, and run for 17.5K steps to match a Chinchilla-optimal 20 tokens-per-parameter (TPP) [47]. We apply Muon and ASGO to optimize the matrix parameters within the model. Non-matrix parameters are optimized using AdamW (α=0.003α=0.003, weight decay λ=0.1λ=0.1, β1=0.9 _1=0.9, β2=0.95 _2=0.95). We employ a standard cosine decay learning rate schedule with a 1% linear warmup, decaying to 1% of the peak learning rate α at the final step. All optimal hyperparameter configurations are evaluated across three independent random seeds governing model parameter initialization and data loading using the PyTorch package[48]. 6.1 Empirical evidence of the aiming condition Figure 3: Empirical evaluation of the aiming condition for training a large language model with ASGO (β1=0.9,β2=0.95)( _1=0.9,\, _2=0.95). The y axis shows the absolute values of μt(Dt) _t(D_t). Left: μt _t computed for the whole group with various directions Dt=GtD_t=G_t, M^t M_t, W^t−1/2M^t W_t^-1/2 M_t and W^t−1/2M^t+λXt W_t^-1/2 M_t+λ X_t. Right: μt _t computed for Dt=W^t−1/2M^t+λXtD_t= W_t^-1/2 M_t+λ X_t but over various parameter groups. Our first set of experiments demonstrate that the aiming condition (37) holds empirically during most of the training process. For each training trajectory, we use the final iterate XfinalX_final to approximate X∗X_*, and record the the inner product between ϕ(M^t,W^t)+λXtφ( M_t, W_t)+λ X_t and Xt−XfinalX_t-X_final. We then normalize the inner product by squared norm of Xt−XfinalX_t-X_final to obtain the empirical estimate of the parameter μ in (37). More generally, we define μt(Dt):=⟨Dt,Xt−Xfinal⟩‖Xt−Xfinal‖2. _t(D_t):= D_t,\,X_t-X_final \|X_t-X_final \|^2. The aiming condition corresponding to the direction DtD_t holds empirically if μt(Dt)>0 _t(D_t)>0. Figure 2 shows the absolute values of μ(Dt)μ(D_t) for Muon (β1=0.95 _1=0.95), where the solid line depicts μt _t for μt>0 _t>0 and dotted line plots |μt|| _t| when μt<0 _t<0. In addition to Dt=Orth(M^t)+λXtD_t= Orth( M_t)+λ X_t, we also plot the corresponding values for the gradient GtG_t, the momentum M^t M_t, and the update Orth(M^t) Orth( M_t). The left panel shows that the update direction Orth(M^t)+λXt Orth( M_t)+λ X_t is positively aligned with Xt−XfinalX_t-X_final throughout training except for a short period in the beginning. In contrast, GtG_t, M^t M_t and Orth(M^t) Orth( M_t) come into positive alignment with Xt−XfinalX_t-X_final in a much later stage, implying the corresponding aiming conditions hold in a much smaller neighborhood of XfinalX_final. The right panel shows |μt|| _t| for Orth(M^t)+λXt Orth( M_t)+λ X_t computed over various parameter groups: matrix parameters updated by Muon, scalar parameters updated by AdamW, attention and MLP blocks at layers 66, 1212, and 1818, the embedding layer, the language-model head, and the normalization parameters. It demonstrates that the positive alignment persists when the update direction is restricted to parameter sub-blocks, providing empirical evidence for a stronger, blockwise form of the aiming condition. Figure 3 plots empirical |μt|| _t| for ASGO (β1=0.9,β2=0.95)( _1=0.9,\, _2=0.95). We observe that the aiming condition for Dt=(W^t+ϵI)−1/2M^t+λXtD_t=( W_t+ε I)^-1/2 M_t+λ X_t holds empirically throughout the training process for the whole parameter group and most of the sub-blocks. Figure 4: Loss curves of Muon and ASGO with different β1 _1’s, averaged over 3 random seeds. 6.2 Effect of the smoothing factor Our main result in Theorem 3.4 states that the algorithms converge almost surely to a neighborhood of some solution X∗X_*, and the radius of the neighborhood is determined by the deviation of the expected practical update from its conceptual counterpart. Theorems 4.1 and 5.1 give concrete bounds on the deviation for Muon and ASGO respectively, both of which depend on the bias and variance of M^t M_t and W^t W_t. We note that the result on Muon does not have explicit dependence on W^t W_t, as a result of our simplification under the assumption that the signal fraction matrix StS_t is approximately diagonal (see Section 2.2). This is a good approximation, for example, if the noise in GtG_t is relatively small (high signal-to-noise ratio). The instrument for adjusting the bias-variance tradeoff of an EMA estimators is its smoothing factor. Increasing the smoothing factor in general reduces the variance but potentially increases the bias (depending on the smoothness of the quantity under estimation). In the second set of exepriments, we examine the performance of Muon and ASGO under different values of the smoothing factor β1∈0.0,0.1,0.3,0.9,0.95 _1∈\0.0,0.1,0.3,0.9,0.95\. For each β1 _1, we extensively sweep the learning rate, effective weight decay, and the second-moment parameter β2 _2. The best value for β2 _2 is 0.950.95 for every β1 _1. After identifying the optimal hyperparameters, we repeat the training experiments across three independent random seeds across model initialization and data loading. Figure 4 shows that for both Muon and ASGO, the training loss decreases monotonically as β1 _1 increases, with marginal gains diminishing as β1 _1 approaches 1. Since increasing β1 _1 mainly reduces the variance of M^t M_t, this indicates that the dominant factor in the bias-variance tradeoff is the variance of M^t M_t; see the bounds in Theorems 4.1 and 5.1. Figure 5 (left) reveals a performance crossover: at β1≤0.3 _1≤ 0.3, ASGO achieves a lower final loss, whereas at β1>0.5 _1>0.5, Muon outperforms ASGO. This is consistent with our theory that Muon is a good approximation of the optimal second-moment method under high signal-to-noise ratio, which mostly corresponds to large β1 _1 for EMA estimators. Arguably ASGO employs a better approximation of the optimal preconditioner across all scenarios as a conceptual method. But with noisy, practical estimators, it is more prone to bad numerical conditioning caused by forming the Gram matrix GtGtTG_tG_t^T, whose condition number squares that of GtG_t. The worse numerical conditioning of ASGO can be observed by comparing the bounds in (64) and (68). Figure 5 (right) depicts a microscopic crossover for β1=0 _1=0: Muon decreases the loss faster than ASGO in the beginning of the training, but ASGO catches up and has a lower loss during the rest of the training process. We can also interpret this phenomenon through the lens of signal fractions. When XtX_t is far from X∗X_*, the random regression function GtG_t has relatively high signal-to-noise ratio, thus the preconditioner of Muon can be very effective. As XtX_t getting closer to X∗X_*, the signal-to-noise ratio drops and ASGO becomes more effective in the later stage of training. This phenomenon persists for small values of β1 _1 (Figure 6, left). However, with β1≥0.5 _1≥ 0.5, the EMA estimator has significantly smaller variance and we observe that the microscopic crossover disappears (Figure 6, right) and Muon obtains a lower loss than ASGO throughout the training process. Figure 5: Left: Final loss of Muon and ASGO versus different smooth factors β1 _1. Right: Loss curves of Muon and ASGO with β1=0.0 _1=0.0. Figure 6: Loss curves of Muon and ASGO, averaged over 3 seeds. Left: β1=0.1 _1=0.1. Right: β1=0.9 _1=0.9. 7 Conclusion and discussion We study a general class of second-moment stochastic approximation (SA) methods. They are derived from a common framework of optimal preconditioning for solving stochastic matrix equations, and different matrix products and simplifications lead to popular algorithms such as Adam, Muon, and their many variants. A unified approach of convergence analysis was developed based on a two-stage procedure: the first stage focuses on the analysis of conceptual methods that work with exact first and second moments; the second stage establishes convergence of practical methods by bounding the their deviation from the conceptual counterparts. The deviation bounds depend on the bias and variance of the first- and second-moment estimators without specifying their exact forms, which provides great flexibility and generality of our framework. The traditional approach of preconditioning for solving nonlinear system of equations is bases on computing or approximating the Jacobian (corresponding to second-order methods in optimization). Many recent and concurrent works focus on extending this approach to the stochastic setting in order to analyze modern deep-learning optimizers and develop new ones (49, 50, 51, 52, 53, e.g.,). In contrast, we introduce a statistical prospect based on second-moment preconditioning, which tackles the fundamental challenge of working with randomness in the stochastic setting (just as second-order derivatives combat anisotropic curvature). This approach has several advantages: • It provides a unified view of many algorithms of different apparent forms, interpreting them as results of simplifying the same optimal preconditioner based on different assumptions on the structure of the noise (through the signal-fraction matrix). • A major advantage of second-moment SA methods over first-moment methods is the scaling-invariant property inherited from the optimal preconditioner. As a consequence, they work well even if the expected regression function (first-moment) is non-smooth or discontinuous. • The statistical prospect allows us to leverage many classical results in the stochastic approximation literature to establish almost sure convergence, which is essential to support the ever-growing scale and cost of training modern deep learning models. In the following, we discuss some limitations of our approach as well as opportunities for future development. • Our result in Theorem 3.4 guarantees almost sure convergence to a neighborhood of some X∗X_* satisfying the aiming condition, not exactly to X∗X_*. This is a tradeoff with the generality of our two-stage analysis framework, which does not specify particular first- and second-moment estimators. We give an example in Appendix B illustrating that a persistent bias can prevent convergence to X∗X_* itself. On the other hand, the deviation bounds in (64) and (68) indicates that the iterates may converge arbitrarily close to X∗X_* if the estimators have diminishing bias and variance. Even with persistent bias, it is possible to specify nuanced aiming conditions for particular estimators that lead to almost convergence to a singleton. • The aiming condition (37) is in the weak form of a variational inequality (54, 55, e.g.,). There is a large literature on efficient algorithms for solving variational inequalities, both in the deterministic setting (56, 57, 58, 59, e.g.,) and the stochastic setting (60, 61, 62, 63, 64, 65, e.g.,). Many of them may be applied in the context of this paper, potentially obtaining fast convergence rates together with second-moment preconditioners. • With λ>0λ>0, the target point X∗X_* in Assumption A is a solution to the regularized variational inequality (37), which may no longer satisfy the original matrix equation M(X)=0M(X)=0. In the context of machine learning, the regularized solution may have better generalization performance [66]. However, with unlimited amount of data in the modern era of deep learning, regularization becomes insignificant or even harming for generalization, but remains necessary for improving training speed and stability. Therefore, a sensible strategy is to apply diminishing regularization strength, i.e., replace the constant λ with a sequence λt→0 _t→ 0. This is consistent with our theory that the aiming condition (37) requires a sufficiently large λt _t when XtX_t is far from X∗X_* (see Remark 3.1) but can holds without regularization near X∗X_* (see Example 3.1). This strategy has been shown to be very effective for both Adam [67] and Muon [68], and an analysis that supports our reasoning here is given in [68]. • It is highly desirable to combine the second-moment methods proposed in this paper with Jacobian-based preconditioners to combat randomness and curvature simultaneously. In fact, these two aspects are not independent of each other and often intertwined. For example, second-moment estimators also appear in stochastic Gauss-Newton methods (69, 70, e.g.,), where they are interpreted as good approximation of the Hessian when the objective gap is small (e.g., 71, Chapter 10). However, in Gauss-Newton methods, the second moment estimators appear in the denominator without the square root. Therefore, they do not possess the scaling invariant property, which is a natural consequence from the perspective of second-moment preconditioning. The recent work of Newton-Muon method [72] combines elements of both second-order derivates and second moment. This is a very promising direction for future research. Appendix A Proof of the extended Dvoretzky’s theorem Our proof is based on the approach of Derman and Sacks [27]. We first present a slightly general version of their key lemma on a deterministic recursion of multiple sequences. Lemma A.1 (Extension of (27, Lemma 1)). Let vt\v_t\, at\a_t\, bt\b_t\, ct\c_t\, and ht\h_t\ be sequences of nonnegative real numbers such that a∞:=lim supt→∞at<∞,∑t=0∞bt<∞,∑t=0∞ct<∞,∑t=0∞ht=∞.a_∞:= _t→∞a_t<∞, _t=0^∞b_t<∞, _t=0^∞c_t<∞, _t=0^∞h_t=∞. (75) Suppose that there is a nonnegative integer T0T_0 such that for all t≥T0t≥ T_0, vt+1≤maxat,(1+bt)vt+ct−ht.v_t+1≤ \a_t,(1+b_t)v_t+c_t-h_t\. (76) Then lim supt→∞vt≤a∞. _t→∞v_t≤ a_∞. Proof of Theorem 3.5. In this proof, we write TtT_t for Tt(X0,…,Xt)T_t(X_0,…,X_t). The proof converts the vector martingale cross term into a scalar martingale difference and then applies Lemma A.1 to a rescaled squared distance. Define wt:=2⟨Tt−X∗,Yt⟩1+‖Tt−X∗‖F2.w_t:= 2 T_t-X_*,Y_t 1+\|T_t-X_*\|_F^2. The coefficient of YtY_t in wtw_t is ℱtF_t-measurable and has Frobenius norm at most one. Hence [wt∣ℱt]=0,[wt2]≤‖Yt‖F2.E[w_t _t]=0, [w_t^2] \|Y_t\|_F^2. The partial sums of ∑twt _tw_t therefore form an L2L^2-bounded martingale, so ∑twt _tw_t converges almost surely. Moreover, Tonelli’s theorem and (58) give ∑t=0∞‖Yt‖F2<∞almost surely. _t=0^∞\|Y_t\|_F^2<∞ surely. Since |wt|≤‖Yt‖F|w_t|≤\|Y_t\|_F, this also gives ∑twt2<∞ _tw_t^2<∞ almost surely. Intersect these two probability-one events with the common probability-one event in the statement of the theorem, and fix an arbitrary sample point in the resulting event. From this point onward, all random matrices and coefficients denote their values at this fixed sample point. Thus they are ordinary numerical sequences satisfying every displayed recursion, inequality, limit, and summability condition used below. Since wt→0w_t→ 0 and hth_t is eventually nonnegative, choose an integer N≥t0N≥ t_0 such that |wt|≤1/2|w_t|≤ 1/2 and ht≥0h_t≥ 0 for every t≥Nt≥ N. Define pt:=∏j=Nt−1(1+wj),t≥N,p_t:= _j=N^t-1(1+w_j), t≥ N, where the empty product gives pN=1p_N=1. Every ptp_t is positive. Because ∑twt _tw_t converges and ∑twt2<∞ _tw_t^2<∞, the expansion log(1+wt)=wt+O(wt2) (1+w_t)=w_t+O(w_t^2) shows that pt⟶p∞for some p∞∈(0,∞).p_t p_∞ some p_∞∈(0,∞). The definition of wtw_t gives the exact identity 1+‖Xt+1−X∗‖F2=(1+‖Tt−X∗‖F2)(1+wt)+‖Yt‖F2.1+\|X_t+1-X_*\|_F^2= (1+\|T_t-X_*\|_F^2 )(1+w_t)+\|Y_t\|_F^2. Using pt+1=pt(1+wt)p_t+1=p_t(1+w_t) and the recursive upper bound (56), we obtain 1+‖Xt+1−X∗‖F2pt+1 1+\|X_t+1-X_*\|_F^2p_t+1 =1+‖Tt−X∗‖F2pt+‖Yt‖F2pt+1 = 1+\|T_t-X_*\|_F^2p_t+ \|Y_t\|_F^2p_t+1 ≤max1+atpt+‖Yt‖F2pt+1,(1+bt)1+‖Xt−X∗‖F2pt+ctpt+‖Yt‖F2pt+1−htpt−btpt ≤ \ 1+a_tp_t+ \|Y_t\|_F^2p_t+1,\,(1+b_t) 1+\|X_t-X_*\|_F^2p_t+ c_tp_t+ \|Y_t\|_F^2p_t+1- h_tp_t- b_tp_t \ ≤max1+atpt+‖Yt‖F2pt+1,(1+bt)1+‖Xt−X∗‖F2pt+ctpt+‖Yt‖F2pt+1−htpt. ≤ \ 1+a_tp_t+ \|Y_t\|_F^2p_t+1,\,(1+b_t) 1+\|X_t-X_*\|_F^2p_t+ c_tp_t+ \|Y_t\|_F^2p_t+1- h_tp_t \. We now verify the hypotheses of Lemma A.1. The state sequence (1+‖Xt−X∗‖F2)/ptt≥N\(1+\|X_t-X_*\|_F^2)/p_t\_t≥ N and all the coefficients in the last display are nonnegative. Since pt→p∞∈(0,∞)p_t→ p_∞∈(0,∞) and ∑t‖Yt‖F2<∞ _t\|Y_t\|_F^2<∞, lim supt→∞(1+atpt+‖Yt‖F2pt+1)≤1+a∞p∞. _t→∞ ( 1+a_tp_t+ \|Y_t\|_F^2p_t+1 )≤ 1+a_∞p_∞. The multiplicative sequence remains btb_t, while ∑t=N∞(ctpt+‖Yt‖F2pt+1)<∞,∑t=N∞htpt=∞. _t=N^∞ ( c_tp_t+ \|Y_t\|_F^2p_t+1 )<∞, _t=N^∞ h_tp_t=∞. The lemma therefore yields lim supt→∞1+‖Xt−X∗‖F2pt≤1+a∞p∞. _t→∞ 1+\|X_t-X_*\|_F^2p_t≤ 1+a_∞p_∞. Finally, multiplying by pt→p∞p_t→ p_∞ and subtracting one gives lim supt→∞‖Xt−X∗‖F2≤a∞. _t→∞\|X_t-X_*\|_F^2≤ a_∞. The fixed sample point was arbitrary in an event of probability one, which proves the theorem. ∎ Appendix B An example with persistent estimator bias We give a simple example to demonstrate that the assumptions of Theorem 3.4 do not guarantee Xt→X∗X_t→ X_* when the practical update has a persistent bias. Consider the following one-dimensional case with X∗:=0,M(X):=X,M^t:=Mt+ϵ,ϕ(M,W):=M,αt:=1t+2.X_*:=0, M(X):=X, M_t:=M_t+ε, φ(M,W):=M, _t:= 1t+2. Here ϵ>0ε>0 is constant. Take X0=0X_0=0 and λ=0λ=0, and let all variables be deterministic. We now verify the assumptions in Theorem 3.4. First, the aiming and growth Assumption A hold with μ=1,L=1,C=0μ=1,L=1,C=0. The error bound Assumption B is true with ϵt=ϵ _t=ε. Indeed, t‖X0‖F2<∞E_t \|X_0 \|_F^2<∞ is true because we initialize X0X_0 at 00. The required conditions are: (a) Bounded variance under ϕφ: The conditional variance is expressed as: t[‖ϕ(M^t,W^t)−t[ϕ(M^t,W^t)]‖F2]=0,E_t [ \|φ ( M_t, W_t )-E_t [φ ( M_t, W_t ) ] \|_F^2 ]=0, satisfying (53) in condition (a) with b=0b=0. (b) The step sizes are deterministic and satisfy 1t+2≥0,∀t≥0,∑t=0∞1t+2=∞,∑t=0∞1(t+2)2<∞. 1t+2≥ 0, ∀ t≥ 0, _t=0^∞ 1t+2=∞, _t=0^∞ 1(t+2)^2<∞. Now we show that XtX_t does not converge to X∗X_*. For t≥1t≥ 1, the chosen step sizes and practical update rule give Xt+ϵ X_t+ε =Xt−1−αt−1(Xt−1+ϵ)+ϵ =X_t-1- _t-1(X_t-1+ε)+ε =(1−αt−1)(Xt−1+ϵ) =(1- _t-1)(X_t-1+ε) =∏s=0t−1(1−αs)(X0+ϵ). = _s=0^t-1(1- _s)(X_0+ε). By construction, ∏s=0t−1(1−αs)=∏s=0t−1s+1s+2=1t+1→0. _s=0^t-1(1- _s)= _s=0^t-1 s+1s+2= 1t+1→ 0. Hence Xt→−ϵ≠0=X∗X_t→-ε≠ 0=X_*. Exact convergence fails even though all stated assumptions hold. Appendix C Perturbation bound for matrix square-root inverse Proof of Lemma 5.1. We use the following integral representation of the matrix square-root inverse function (13, Chapter 5) A−1/2=1π∫s=0∞s−1/2(A+sI)−1ds.A^-1/2= 1π _s=0^∞s^-1/2(A+sI)^-1ds. For two symmetric positive matrices A and B, we have A−1/2−B−1/2 A^-1/2-B^-1/2 =1π∫s=0∞s−1/2((A+sI)−1−(B+sI)−1)ds = 1π _s=0^∞s^-1/2 ((A+sI)^-1-(B+sI)^-1 )ds =1π∫s=0∞s−1/2(A+sI)−1(B−A)(B+sI)−1ds, = 1π _s=0^∞s^-1/2(A+sI)^-1(B-A)(B+sI)^-1ds, where the second line follows from the resolvent identity. Using the sub-multiplicative property of matrix norms, we get: ∥A−1/2−B−1/2∥F≤1π∫s=0∞s−1/2∥(A+sI)−1∥2∥B−A∥F∥(B+sI)−1∥2ds. \|A^-1/2-B^-1/2 \|_F≤ 1π _s=0^∞s^-1/2 \|(A+sI)^-1 \|_2 \|B-A \|_F \|(B+sI)^-1 \|_2ds. Since A⪰λAIA _AI and B⪰λBIB _BI, we can derive bounds for the two operator norms: ‖(A+sI)−1‖2≤1λA+s,‖(B+sI)−1‖2≤1λB+s.\|(A+sI)^-1\|_2≤ 1 _A+s, \|(B+sI)^-1\|_2≤ 1 _B+s. Substituting these results into the previous upper bound and evaluating the integral, we get ∥A−1/2−B−1/2∥F \|A^-1/2-B^-1/2 \|_F ≤‖B−A‖Fπ∫s=0∞s−1/2(λA+s)−1(λB+s)−1ds ≤ \|B-A\|_Fπ _s=0^∞s^-1/2( _A+s)^-1( _B+s)^-1\,ds =(λAλB1/2+λBλA1/2)−1‖B−A‖F, = ( _A _B^1/2+ _B _A^1/2 )^-1\|B-A\|_F, as desired. ∎ References [1] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. http://w.deeplearningbook.org. [2] Christopher M. Bishop and Hugh Bishop. Deep Learning: Foundations and Concepts. Springer, 2024. https://w.bishopbook.com. [3] J. M. Ortega and W. C. Rheinboldt. Iterative Solution of Nonlinear Equations in Several Variables. Classics in Applied Mathematics. SIAM, Philadelphia, PA, 2000. [4] C. T. Kelley. Solving Nonlinear Equations with Iterative Methods: Solvers and Examples in Julia. Society for Industrial and Applied Mathematics, 2022. [5] C. T. Kelley. Solving Nonlinear Equations with Newton’s Method. Society for Industrial and Applied Mathematics, Philadelphia, PA, 2003. [6] Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathematical Statistics, 22(3):400–407, 1951. [7] A. M. Gupal and L. T. Bazhenov. A stochastic analog of the conjugate gradient method. Cybernetics, 8(1):138–140, 1972. [8] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In Proceedings of the 30th International Conference on Machine Learning, pages 1139–1147, Atlanta, Georgia, USA, 17–19 Jun 2013. [9] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proceedings of International Conference on Learning Representations (ICLR), 2015. arXiv:1412.6980. [10] Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-RMSProp: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning, 2012. [11] Kang An, Yuxing Liu, Rui Pan, Yi Ren, Shiqian Ma, Donald Goldfarb, and Tong Zhang. ASGO: Adaptive structured gradient optimization. Preprint, arXiv:2503.20762, 2025. [12] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. https://kellerjordan.github.io/posts/muon/. [13] Nicholas J. Higham. Functions of Matrices. Society for Industrial and Applied Mathematics, 2008. [14] Zdislav Kovarik. Some iterative methods for improving orthonormality. SIAM Journal on Numerical Analysis, 7(3):386–389, 1970. [15] Å. Björck and C. Bowie. An iterative algorithm for computing the best estimate of an orthogonal matrix. SIAM Journal on Numerical Analysis, 8(2):358–364, 1971. [16] Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. signSGD: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pages 560–569. PMLR, 2018. [17] Václav Fabian. Stochastic approximation methods. Czechoslovak Mathematical Journal, 10(1):123–159, 1960. [18] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36:49205–49233, 2023. [19] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(Jul):2121–2159, 2011. [20] Timothy Dozat. Incorporating Nesterov Momentum into Adam. In Proceedings of the 4th International Conference on Learning Representations (ICLR), 2016. Workshop track. [21] Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P Kingma, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. Preprint, arXiv:2406.16793, 2024. [22] Matteo Pagliardini, Pierre Ablin, and David Grangier. The AdEMAMix optimizer: Better, faster, older. Preprint, arXiv:2409.03137, 2024. [23] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1842–1850. PMLR, 10–15 Jul 2018. [24] Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham M. Kakade. SOAP: Improving and stabilizing shampoo using adam for language modeling. In The Thirteenth International Conference on Learning Representations, 2025. [25] M. T. Wasan. Stochastic Approximation. Cambridge University Press, 1969. [26] Aryeh Dvoretzky. On stochastic approximation. In Proceedings of the Third Berkeley Symposium on Mathematical Statistics and Probability, volume 1, pages 39–55. University of California Press, 1956. [27] C. Derman and J. Sacks. On Dvoretzky’s Stochastic Approximation Theorem. The Annals of Mathematical Statistics, 30(2):601 – 606, 1959. [28] J. H. Venter. On Dvoretzky Stochastic Approximation Theorems. The Annals of Mathematical Statistics, 37(6):1534 – 1544, 1966. [29] Patrick G. Mulloy. Smoothing data with faster moving averages. Technical Analysis of Stock & Commodities, 12(1), January 1994. [30] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), 2019. [31] David Carlson, Volkan Cevher, and Lawrence Carin. Stochastic spectral descent for restricted boltzmann machines. In Artificial intelligence and statistics, pages 111–119. PMLR, 2015. [32] Jerry Ma and Denis Yarats. Quasi-hyperbolic momentum and adam for deep learning. ArXiv, abs/1810.06801, 2018. [33] Huizhuo Yuan, Yifeng Liu, Shuang Wu, Xun Zhou, and Quanquan Gu. Mars: Unleashing the power of variance reduction for training large models. Preprint, arXiv:2411.10438, 2024. [34] Ashok Cutkosky and Francesco Orabona. Momentum-based variance reduction in non-convex sgd. Advances in neural information processing systems, 32, 2019. [35] Ruichen Jiang, Zakaria Mhammedi, Mehryar Mohri, and Aryan Mokhtari. Adaptive matrix online learning through smoothing with guarantees for nonsmooth nonconvex optimization. Preprint, arXiv:2602.08232, 2026. [36] Tao Jiang and Lin Xiao. Stochastic approximation with block coordinate optimal stepsizes. Preprint, arXiv:2507.08963, 2025. [37] Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for LLM training. Preprint, arXiv:2502.16982, 2025. [38] B Halpern. Fixed points of nonexpanding maps. Bulletin of the American Mathematical Society, 73:959, 1976. [39] David J. Sakrison. Stochastic approximation: A recursive method for solving regression problems. In A.V. Balakrishnan, editor, Advances in Communication Systems, volume 2, pages 51–106. Elsevier, 1966. [40] H. Robbins and D. Siegmund. A convergence theorem for non negative almost supermartingales and some applications. In Jagdish S. Rustagi, editor, Optimizing Methods in Statistics, pages 233–257. Academic Press, 1971. [41] Yurii Nesterov. Introductory Lecture on Convex Optimization: A Basic Course. Kluwer Academic Publishers, 2004. [42] Boris T. Polyak. Introduction to Optimization. Optimization Software, Inc., Publications Division, New York, 1987. Translated from the Russian by Igor V. Bakelman and Gennady M. Kuritsky. [43] Kai Lai Chung. On a stochastic approximation method. The Annals of Mathematical Statistics, pages 463–483, 1954. [44] Ren-Cang Li. New perturbation bounds for the unitary polar factor. SIAM Journal on Matrix Analysis and Applications, 16(1):327–332, 1995. [45] Llama Team at Meta. The Llama 3 herd of models. Preprint, arXiv:2407.21783, 2024. [46] G. Penedo, H. Kydlícek, L. Ben Allal, A. Lozhkov, M. Mitchell, C. A. Raffel, L. von Werra, and T. Wolf. The FineWeb datasets: Decanting the web for the finest text data at scale. In Advances in Neural Information Processing Systems (NeurIPS), volume 37, 2024. [47] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, J. W. Rae, O. Vinyals, and L. Sifre. Training compute-optimal large language models. In Advances in Neural Information Processing Systems, volume 35, 2022. [48] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Z. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. [49] Qi Deng, Yi Cheng, and Guanghui Lan. Optimal adaptive and accelerated stochastic gradient descent, 2018. [50] Antonio Orvieto and Lin Xiao. An adaptive stochastic gradient method with non-negative Gauss-Newton stepsizes. Preprint, arXiv:2407.04358, 2024. [51] Weijie Su. Isotropic curvature model for understanding deep learning optimization: Is gradient orthogonalization optimal?, 2025. [52] Natalie Abreu, Nikhil Vyas, Sham M. Kakade, and Depen Morwani. The potential of second-order optimization for LLMs: A study with full Gauss-Newton. In The Fourteenth International Conference on Learning Representations (ICLR), 2026. [53] Dmitry Kovalev. SGD with adaptive preconditioning: Unified analysis and momentum acceleration. In The Fourteenth International Conference on Learning Representations, 2026. [54] David Kinderlehrer and Guido Stampacchia. An introduction to variational inequalities and their applications. Academic Press, New York, 1980. [55] Francisco Facchinei and Jong-Shi Pang. Finite-Dimensional Variational Inequalities and Complementarity Problems, Volume I & I. Springer, New York, 2003. [56] G. M. Korpelevich. The extragradient method for finding saddle points and other problems. Ekonomika i Matematicheskie Metody, 12:747–756, 1976. [57] Arkadi Nemirovski. Prox-method with rate of convergence o(1/t) for variational inequalities with lipschitz continuous monotone operators and smooth convex-concave saddle point problems. SIAM Journal on Optimization, 15(1):229–251, 2004. [58] Yurii Nesterov. Dual extrapolation and its applications to solving variational inequalities and related problems. Mathematical Programming, 109(2–3):319–344, 2007. [59] Yu. Malitsky. Projected reflected gradient methods for monotone variational inequalities. SIAM Journal on Optimization, 25(1):502–520, 2015. [60] Anatoli Juditsky, Arkadi Nemirovski, and Claire Tauvel. Solving variational inequalities with Stochastic Mirror-Prox algorithm. Stochastic Systems, 1(1):17–58, 2011. [61] Yunmei Chen, Guanghui Lan, and Yuyuan Ouyang. Accelerated schemes for a class of variational inequalities. Mathematical Programming, 165(1):113–149, 2017. [62] A. N. Iusem, A. Jofré, R. I. Oliveira, and P. Thompson. Extragradient method with variance reduction for stochastic variational inequalities. SIAM Journal on Optimization, 27(2):686–724, 2017. [63] Farzad Yousefian, A. Nedić, and Uday V. Shanbhag. On smoothing, regularization, and averaging in stochastic approximation methods for stochastic variational inequality problems. Mathematical Programming, 165(1):391–431, 2017. [64] Georgios Kotsalis, Guanghui Lan, and Tianjiao Li. Simple and optimal methods for stochastic variational inequalities, I: Operator extrapolation. SIAM Journal on Optimization, 32(3):2041–2073, 2022. [65] Yao Ji, Guanghui Lan, and Jason Zhu. Computation of strong solutions to stochastic variational inequalities. arXiv:2609.04188, September 2026. [66] Vladimir N. Vapnik. Statistical Learning Theory. Wiley, New York, 1998. [67] Aaron Defazio. Why gradients rapidly increase near the end of training. arXiv:2506.02285, June 2025. [68] Anuj Apte. Scale weight decay and train better. arXiv:2607.23777, July 2026. [69] Yann Le Cun, Léon Bottou, Genevieve B. Orr, and Klaus-Robert Müller. Efficient backprop. In Neural Networks, Tricks of the Trade, Lecture Notes in Computer Science LNCS 1524. Springer Verlag, 1998. [70] Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223–311, 2018. [71] Jorge Nocedal and Stephen J. Wright. Numerical Optimization. Springer Series in Operations Research and Financial Engineering. Springer, New York, NY, 2nd edition, 2006. [72] Zhehang Du and Weijie Su. The Newton-Muon optimizer. Preprint, arXiv:2604.01472, April 2026.