Free Heavy-Tailed Lunch for Muon: A Theoretical Justification of Empirical Success
arXiv:2606.14560v1 [math.OC] 12 Jun 2026
Florian Hübler1,2,
∗
Thomas Pethick
Suvrit Sra2,3
1 Department of Computer Science, ETH Zurich, Switzerland 2 Department of Mathematics, Technical University of Munich, Germany 3 Munich Center for Machine Learning (MCML)
Abstract Non-Euclidean optimisation methods with matrix-valued updates, such as M UON and S CION, have recently shown strong empirical performance for training Transformer models, yet their theoretical advantages over Euclidean methods remain poorly understood. We address this gap in the heavy-tailed non-convex regime, where stochastic gradients have bounded p-th central moments, p ∈ (1, 2]. We show that certain non-Euclidean methods achieve optimal sample complexity under stronger stationarity measures, while Euclidean methods incur additional dimension-dependent costs. As a consequence, for UON finds m × n matrices, M p an ε-stationary point in nuclear norm within O min{m, n} ∆ε12L σε p−1 samples, absorbing heavy-tailed noise without extra dimension dependence, unlike Euclidean methods. We further prove this sample complexity, including its dimension dependence, is optimal for all first-order methods under nuclear-norm stationarity. Experiments on large language models support our theory. Surprisingly, our results suggest that other Schatten geometries beyond the spectral geometry of M UON can perform competitively in certain settings.
1
Introduction
We consider the unconstrained stochastic optimisation problem min {F (x) := E [f (x, ξ)]}, x∈B
(1)
where ξ is distributed according to some unknown distribution and B is a Banach space such as Rd or Rm×n . This setting has been the focus of a large body of work in optimisation due to its prevalence in machine learning and data-driven optimisation (Bottou et al., 2018). While A DAM (Kingma and Ba, 2014) has been the de-facto standard optimisation algorithm for training Transformer-based models for the last decade, a recent line of literature has proposed a new class of methods based on steepest descent in general normed spaces that are empirically competitive with A DAM and its variants. The general form of these algorithms is given by the normalised steepest descent (Boyd and Vandenberghe, 2004) update xt+1 ← xt + η argmin ⟨gt , d⟩ ,
(2)
∥d∥≤1
where η > 0 is the stepsize, gt some gradient estimator and ∥·∥ a norm on B (Pethick et al., 2025b). The most prominent example is M UON (Tuddenham et al., 2022; Jordan et al., 2024b), which uses (1) the spectral norm ∥·∥ = ∥·∥S∞ , (2) the momentum gradient estimator, and (3) Newton-Schulz ∗ Corresponding author: [email protected]
Preprint.
iterations (Kovarik, 1970) to estimate the argmin. This variant achieves impressive results on the modded nanoGPT speedrun (Jordan et al., 2024a) and was used to train recent trillion parameter language models such as Kimi K2 (Team et al., 2025) and DeepSeek V4 (DeepSeek-AI, 2026). From a theoretical perspective, recent works have established convergence of such algorithms in the non-convex setting. These guarantees are typically of the form 2 2 ∆1 Lσ samples, (3) E [∥∇F (xT )∥∗ ] ≤ ε after T = O ρ · ε4 where F is L-smooth and the stochastic gradients have bounded variance σ 2 , ∥·∥∗ is the dual norm and ρ := maxv∈B∗ ∥v∥∗/∥v∥∗,2 is the norm-equivalence constant between the dual norm and a Hilbert norm on B∗ (e.g., Pethick et al., 2025b; Shen et al., 2025). Importantly, (3) guarantees E [∥∇F (xT )∥∗ ] ≤ ε, which we denote ε-∥·∥∗ -stationary point. For example, for M UON the dual norm is the Schatten-1 norm ∥·∥S1 and thus a strictly stronger measure than the Frobenius norm typically used in Euclidean analyses. While informative, these guarantees face two difficulties in capturing the potential benefits of non-Euclidean methods in stochastic settings. First, the appearance of the norm-equivalence constant ρ can make the bound effectively dimensiondependent, thereby obscuring potential advantages over Euclidean methods. This term arises fromPthe treatment of the gradient noise. Indeed, the analysis must control terms of the form E[∥ (∇F (xt ) − ∇f (xt , ξt ))∥∗ ]. Such bounds are straightforward in Hilbert spaces, where the inner-product structure is available, but much harder for general norms. Existing analyses therefore use the norm equivalence ∥·∥∗ ≤ ρ ∥·∥∗,2 , before controlling the noise in the Euclidean geometry. p In particular, for M UON’s nuclear norm on Rm×n , we get ρ = min {m, n}. Second, these results rely on the bounded variance assumption, which can be overly restrictive. Indeed, empirical evidence from image classification (Simsekli et al., 2019; Battash et al., 2024), language modelling (J. Zhang et al., 2020a; Ahn et al., 2024), and reinforcement learning (RL) (Garg et al., 2021) suggests that stochastic gradients can exhibit heavy-tailed behaviour. These observations suggest that heavy-tailed noise provides a more faithful model of stochastic gradients in modern machine learning. Consequently, if one seeks to understand the empirical advantage of non-Euclidean optimisation methods, it is natural to analyse them in this regime. These two limitations naturally raise the following question: Can non-Euclidean methods such as M UON admit provable improvements over Euclidean methods in stochastic regimes, particularly under heavy-tailed noise? To study this question, we depart from the usual Euclidean noise assumption and instead make the heavy-tailed assumption directly in the native geometry: for some tail index p ∈ (1, 2] we assume p E [∥∇f (x, ξ) − ∇F (x)∥∗ ] ≤ σ p . This choice is the root of the main technical challenges, but also what enables sharper, geometryaware guarantees. Contributions. In this work we show that non-Euclidean methods can achieve strict improvements over their Euclidean counterparts in the heavy-tailed setting. (i) We show that normalised steepest descent methods with momentum can achieve optimal rates for strictly stronger convergence measures. More precisely, we provide the optimal p ∆1 L σ p−1 sample complexity bound O ε2 ε for reaching an ε-∥·∥∗ -stationary point, whenever the dual norm ∥·∥∗ is p-uniformly smooth (see Definition 2.1). Moreover, we provide this guarantee under the weak (L0 , L1 )-smoothness assumption. (ii) Based on this result, we obtain improved convergence guarantees for M UON, whose dual norm ∥·∥S1 is not uniformly smooth. We show that M UON achieves the sample complexity p p 3p−2 ∆1 L σ p−1 ∆1 L σ p−1 O min {m, n} · 2 vs. O min {m, n} 2(p−1) · 2 ε ε ε ε obtained by Euclidean algorithms when corrected for the stationarity measure. This shows that M UON can absorb heavy-tailed noise essentially without extra dimensional cost, whereas classical Euclidean methods can match the same stationarity measure only at the price of a strictly worse, p-dependent dimension factor. 2
(iii) We provide a first-order sample complexity lower-bound of p ∆1 L σ p−1 Ω min {m, n} · 2 ε ε for ∥·∥S1 -stationarity. In particular, M UON is first-order optimal and there cannot exist a dimension-free guarantee for ∥·∥S1 -stationarity. As a side result, we provide tools from functional analysis to control the gradient noise for general norms, sidestepping the usual reduction to Euclidean martingale bounds through norm equivalence. These tools may be of independent interest for the analysis of general-norm algorithms beyond the heavy-tailed setting.
2
Preliminaries
Notation. Throughout, (B, ∥·∥) denotes a Banach space with dual (B ∗ , ∥·∥∗ ). We denote the duality pairing by ⟨x∗ , x⟩ := x∗ (x) for x ∈ B, x∗ ∈ B ∗ . To keep consistency with the typical notation in Rd we abuse notation2 by writing ∇F (x) for the Frechet derivative DF (x) ∈ B ∗ and ⟨∇F (x), d⟩ for the duality pairing between the Frechet derivative and d ∈ B. Algorithm. We study the stochastic normalised steepest descent method with momentum, mt ← βt mt−1 + (1 − βt )∇f (xt , ξt ), xt+1 ← xt + ηt argmin ⟨mt , d⟩ ,
(N SSD -M)
∥d∥≤1 i.i.d.
where ξ1 , ξ2 , . . . ∼ ξ are independent copies of ξ, ∇f (xt , ξt ) is a stochastic gradient oracle of ∇F (xt ), β ∈ [0, 1) is the momentum parameter and m0 := ∇f (x1 , ξ1 ). When choosing (B, ∥·∥) = (Rm×n , ∥·∥S∞ ) this corresponds to an idealised version of M UON with an exact lmo, and (B, ∥·∥) = (Rd , ∥·∥ℓ∞ ) to S IGNUM. We consider the following standard assumptions. Assumption 1. There exists ∆1 < ∞ such that F (x1 ) − inf x∈B F (x) ≤ ∆1 . For ease of exposition, we present the results for the case of L-smoothness as defined below. Our results readily also hold for the weaker notion of (L0 , L1 )-smoothness and can be found in Section C. Assumption 2. There exists L ≥ 0 such that for all x, y ∈ B, ∥∇F (x) − ∇F (y)∥∗ ≤ L ∥x − y∥. Finally, as mentioned in the introduction, we consider the heavy-tailed setting where the stochastic gradient oracle has bounded p-th central moment for some p ∈ (1, 2]. Assumption 3. There exists p ∈ (1, 2], σ ≥ 0 such that the stochastic gradient oracle is unbiased and has bounded p-th central moment, i.e., for all x ∈ B, E [∇f (x, ξ)] = ∇F (x), 2.1
p
E [∥∇f (x, ξ) − ∇F (x)∥∗ ] ≤ σ p .
and
Moment Bounds for General Norms
We first summarise the necessary technical tools to control the gradient approximation error for general norms, sidestepping the usual reduction to Euclidean martingale bounds through norm equivalence. Therefore we first recall the following key geometric definitions of general norms. Definition 2.1 (Ball et al., 1994, Equation (2.7)). Let (X, ∥·∥) be a Banach space. i) The p-smoothness constant of (X, ∥·∥) is defined as p
p
p
p
q
q
Sp (X) := inf {S ≥ 1 | ∀ x, y ∈ X : ∥x + y∥ + ∥x − y∥ ≤ 2 ∥x∥ + 2S p ∥y∥ }, and we call (X, ∥·∥) p-uniformly smooth if Sp (X) < ∞. ii) The q-convexity constant of (X, ∥·∥) is defined as q
q
Kq (X) := inf {K ≥ 1 | ∀ x, y ∈ X : 2 ∥x∥ + 2/K q ∥y∥ ≤ ∥x + y∥ + ∥x − y∥ }, and we call (X, ∥·∥) q-uniformly convex if Kq (X) < ∞. 2 Note that the gradient is in general not well defined in Banach spaces, as it is defined by the Riesz representation of
DF (x), which only exists in Hilbert spaces.
3
The following strong duality holds between the convexity and smoothness constants. Lemma 2.2 (Ball et al., 1994, Lemma 5). Let (X, ∥·∥) be a Banach space and 1 < p ≤ 2 ≤ q < ∞ conjugate exponents, i.e., 1/p + 1/q = 1. Then Sp (X ∗ ) = Kq (X). The above duality relationship allows us to verify the smoothness assumption on the dual space by checking a convexity property of the primal space, which is often more natural. Next let us get some intuition of the above definition for the special cases of Rd and Rm×n Example 2.3 (Hirzallah and Kittaneh, 2002). Consider the ℓr space (X, ∥·∥) = (Rd , ∥·∥r ) or the Schatten-r space (X, ∥·∥) = (Rm×n , ∥·∥Sr ). In both cases the following holds. i) If 1 < r ≤ 2, then X is r-uniformly smooth with Sr (X) = 1. ii) If 2 ≤ r < ∞, then X is r-uniformly convex with Kr (X) = 1. Finally we provide the main technical tool which will allow us to control the gradient approximation error for general norms, sidestepping the usual reduction to Euclidean martingale bounds through norm equivalence. Theorem 2.4 (c.f., Wenzel, 2005, Theorem 3.2). Every martingale difference sequence Y1 , . . . , Yn in B ∗ satisfies " n p# n X X p E Yt ≤ Sp (B ∗ )p · E [∥Yt ∥∗ ] . t=1
t=1
∗
We note that the above inequality is vacuous if B ∗ is not p-uniformly smooth, i.e., if Sp (B ∗ ) = ∞.
3
Main Result
In this section we present our theoretical results. In Section 3.1 we provide convergence guarantees for N SSD -M for uniformly convex norms under heavy-tailed noise. In Section 3.2 we extend these results to the case of M UON. Finally, in Section 3.3 we provide a lower bound showing that the obtained dimension dependence is first-order optimal for nuclear-stationarity guarantees. 3.1
Guarantees for Uniformly Convex Norms
We are now ready to prove the convergence guarantee for uniformly convex norms under heavytailed noise. For ease of exposition we will present it under the standard L-smoothness assumption, extended results for (L0 , L1 )-smoothness and the corresponding proofs can be found in Section C.1. Theorem 3.1 (Uniformly-Convex Norm Guarantee). Suppose that Assumptions 1 to 3 hold. Furthermore let q = p/(p−1) and assume that (B, ∥·∥) is q-uniformly convex with q-convexity constant K. Then the iterates generated by N SSD -M with parameters ( ( )) r p 3p−2 p ∆1 (1 − β) ∆1 L − 2p−1 , βt ≡ β = 1 − min 1, max T , ηt ≡ LT σ2 K 2 T satisfy r p−1 3p−2 T p ∆1 L ∆1 L σK 1X E [∥∇F (xt )∥∗ ] ≤ 9 + 14 (σK) p−1 + 10 p−1 . T t=1 T T T 2p−1 This shows that N SSD -M achieves the optimal heavy-tailed sample complexity in arbitrary, possibly infinite-dimensional Banach spaces. Importantly, the stationary measure ∥·∥∗ can in general be stronger than in the Euclidean case. This becomes particularly transparent for Schatten and ℓ norms. Corollary 3.2. Let (B, ∥·∥) = (Rm×n , ∥·∥Sr ) or (Rd , ∥·∥ℓr ), where r ≤ p/(p−1). Then, in the setting of Theorem 3.1, N SSD -M satisfies p ∆1 L σ p−1 E [∥∇F (xτ )∥∗ ] ≤ ε after T = O samples, (4) ε2 ε where τ ∼ Unif([T ]). Here ∥·∥∗ = ∥·∥Sr′ and ∥·∥ℓr′ , where r′ = r/(r−1), for the Schatten and ℓ case respectively. 4
In particular, when choosing r = 2, Corollary 3.2 exactly reconstructs the existing heavy-tailed results for the Euclidean case (Hübler et al., 2025; Z. Liu and Zhou, 2025). For 2 < r ≤ p/(p−1), (r−2)/2r the result is strictly stronger, as ∥·∥∗ ≥ ∥·∥2 ≥ 1/ρ ∥·∥∗ , where ρ = min {m, n} in Schatten-, (r−2)/2r and ρ = d in ℓ-geometry. A stationary-measure comparable result from the Euclidean setting hence has the form p 3p−2 ∆1 L σ p−1 p−1 E [∥∇F (xτ )∥∗ ] ≤ ε after T = O ρ · 2 samples. ε ε This multiplicative factor can be considerable, as already for moderate values such as p = 1.5 and r = p/(p−1) we have ρ(3p−2)/(p−1) ≈ 104 in the Schatten-, and ≈ 107 in the ℓ-geometry for modern transformers (Chowdhery et al., 2023). The main technical novelty in the proof of Theorem 3.1 is based on the careful application of Theorem 2.4 to derive the following deviation bound on our gradient estimator. Lemma 3.3. Suppose that Assumptions 2 and 3 hold, and let q = p/(p−1). Assume that (B, ∥·∥) is q-uniformly convex with q-convexity constant K. Then the iterates generated by N SSD -M with ηt ≡ η and βt ≡ β satisfy T X
E [∥∇F (xt ) − mt ∥∗ ] ≤
t=1
p−1 σK ηLT + σKT (1 − β) p + . 1−β 1−β
The above Lemma allows us to control the convergence of N SSD -M directly in its native geometry, without having to reduce to the Euclidean case through norm equivalence, thereby avoiding the dimension dependence. In the special case of the Euclidean geometry, the above Lemma exactly reconstructs previous results (Cutkosky and Mehta, 2020). The extended version of this Lemma and its proof can be found in Lemma C.2, we provide a proof sketch below. Proof Sketch of Lemma 3.3. First note that the momentum deviation can be unrolled to mt −∇F (xt ) = a(m1 −∇F (x1 ))+
t X
bτ (∇f (xτ , ξτ ) − ∇F (xτ ))+
t X
cτ (∇F (xτ −1 ) − ∇F (xτ )),
τ =2
τ =2
for some a, bτ , cτ . Here the last term can be controlled deterministically using smoothness, and the first term can be controlled after taking expectation. For the middle term let us denote γt := ∇f (xt , ξt ) − ∇F (xt ) and note that γ1 , . . . , γt is a martingale difference sequence in B ∗ . Furthermore, by Lemma 2.2, we know that B ∗ is p-uniformly smooth with Sp (B ∗ ) = K. Hence we can apply Theorem 2.4 to get # " t " t p #1/p X X bτ (∇f (xτ , ξτ ) − ∇F (xτ )) ≤E bτ (∇f (xτ , ξτ ) − ∇F (xτ )) E τ =2
τ =1
∗
≤K
t X
∗
!1/p p bpτ E [∥γτ ∥∗ ]
τ =1
!1/p ≤ σK
X
bpτ
.
τ =1
Plugging in the corresponding values of a, bτ , cτ wraps up the proof. Remark 3.4. It is well-known that weight-decay is often required in practice to efficiently train large language models (Team et al., 2025). For normalised steepest descent based methods such as N SSD -M, this regularisation has a natural constrained interpretation as a conditional gradient step for optimisation over the corresponding norm ball (Pethick et al., 2025b). We extend our convergence analysis to this constrained stochastic conditional gradient setting in Section D. 3.2
Guarantees for M UON
While the previous results established convergence of N SSD -M for a broad range of norms, they do not cover M UON as the Schatten-∞ norm is not uniformly convex. Instead, we will show in this section that M UON can automatically adapt to the optimal Schatten-r geometry, thereby achieving a 5
reduced dimension dependence without any modification to the algorithm. A more general version under (L0 , L1 )-smoothness and the corresponding proofs can be found in Section C.2. The main technical idea is to use equivalence of norms only to get from the nuclear norm to the Schatten-p norm, where p is the tail index, instead of going all the way to the Frobenius norm. This allows us to heavily reduce the dimension dependence as the noise becomes more heavy-tailed, while still controlling a strictly stronger stationarity measure than the Frobenius norm. Let us now state the result. Theorem 3.5 (Convergence of M UON under Heavy-Tailed Noise). Consider B = Rm×n and sup(p−1)/p pose Assumptions 1 to 3 hold. Furthermore let κ := min {m, n} . Then the iterates generated by M UON with parameters ( ( )) r p p ∆1 L 3p−2 − 2p−1 ∆1 (1 − β) , βt ≡ 1 − min 1, max ,T ηt ≡ LT σ 2 κ2 T satisfy p−1 ! 3p−2 r (p−1) p T 1X min {m, n} p σ ∆1 L ∆1 Lσ p−1 + 10 . E ∥∇F (xt )∥S1 ≤ 9 + 14 min {m, n} p−1 T t=1 T T T 2p−1
When dropping non-leading terms, the sample complexity guarantee for M UON is hence given by p ∆1 L σ p−1 O min {m, n} · 2 . ε ε In comparison, the existing heavy-tailed Euclidean guarantees (Hübler et al., 2025; Z. Liu and Zhou, 2025) imply a nuclear norm sample complexity guarantee of p 3p−2 ∆1 L σ p−1 , (5) O min {m, n} 2(p−1) · 2 ε ε which is strictly worse in its dimension dependence. To illustrate, for p = 1.5, modern transformer 3p−2 based models (Chowdhery et al., 2023) satisfy min {m, n} ≈ 104 but min {m, n} 2(p−1) ≈ 1010 , reducing the dimension dependence by 6 orders of magnitude. Hence, Theorem 3.5 shows that M UON can automatically adapt to the optimal Schatten-p geometry, thereby avoiding the worst-case rank-dependence of existing heavy-tailed guarantees using the Euclidean geometry. Remark 3.6. For simplicity of presentation, Theorem 3.5 is stated for a single matrix. However, the extended result Theorem C.6 directly applies to neural networks consisting of a collection of weight matrices {Wl }l∈[D] , by equipping Rm1 ,n1 × . . . × RmD ,nD with the norm maxP l∈[D] ∥Wl ∥S∞ (Bernstein and Newhouse, 2024). In this case, the dimension dependence becomes l∈[D] min {ml , nl }. A further extension to full networks including vector-valued weights is discussed in Remark C.9. 3.3
First-Order Lower-Bounds for Nuclear-Stationarity
Finally we theoretically examine whether the remaining dimension-dependence in the guarantee for M UON can be avoided. We show that, in the above setting, no first-order algorithm can achieve a better dimension-dependence than M UON, when convergence is measured in ∥·∥S1 . We make use of the sample complexity framework of Arjevani et al. (2023). More precisely, we consider possibly randomised first-order algorithms A ∈ Arand which, given a starting point x0 , random seed r and access to a stochastic gradient oracle ∇f (·, ξ), produce a sequence of random iterates x1 , x2 , . . . . For ∆1 , L ≥ 0, we consider the function class F∆1 ,L (m, n) := F : Rm×n → R ∥∇F (x) − ∇F (y)∥∗ ≤ L ∥x − y∥ and F (0)−inf F (x) ≤ ∆1 x
and, for σ ≥ 0, p ∈ (1, 2] the oracle class Op,σ of gradient oracles satisfying Assumption 3. We will derive lower-bounds on the distributional complexity mrand ε (∥·∥∗ , ∆1 , L, p, σ) := sup
inf inf {T ∈ N | E [∥∇F (xT )∥∗ ] ≤ ε},
sup
O∈Op,σ PF ∈P[F∆1 ,L (m,n)] A∈Arand
where P[F∆1 ,L (m, n)] denotes the set of all probability distributions over F∆1 ,L (m, n). Lowerbounds on this distributional complexity imply lower-bounds on the worst-case complexity (Nemirovskij and Yudin, 1983). 6
Theorem 3.7 (First-Order Nuclear-Stationarity Lower-Bound). Let ∆1 , L, σ > 0, p ∈ (1, 2] and √ assume ε < min σ4 , ∆1 L/2 . Then, for all m, n ∈ N≥2 , we have p min {m, n} σ p−1 mrand ∥·∥ , ∆ , L, p, σ ≥ · . 1 ε S1 2 4ε We give an informal proof sketch, the complete proof can be found in Section E. Therefore let n (ei )m bases of Rm and Rn respectively, and denote Ei := ei fi⊤ ∈ Rm×n . i=1 and (fi )i=1 be standard p Pκ κ 4ε p−1 ∈ (0, 1], s ∈ {±1} be a sign vector and Ss := i=1 si Ei , where Furthermore let q := σ κ := min {m, n}. Now we construct the family of hard instances as Fs (x) :=
2ε L 2 ∥x∥F − ⟨Ss , x⟩F , 2κ κ
∇fs (x, ξ) :=
L 2ε x − BsI EI , κ q
where ξ = (B, I) with B ∼ Bern(q) and I ∼ Unif([κ]) independently. Starting at zero, one can then roughly show that the algorithm must discover the sign of each si to reach an ε-∥·∥S1 -stationary point. However, in order to discover the sign of one si , (B, I) = (1, i) must occur, which happens with probability κq . Hence, by a coupon collector type argument, we need at least κq iterations in expectation to discover the sign of all si , which gives the desired lower-bound after showing that Fs and ∇fs satisfy the required smoothness and noise properties. We note that the above lower-bound matches Theorem 3.5 in its dimension and σ dependence, but not in its ε dependence. We will now present a lower-bound that addresses this gap, matching the exact dependence. It comes at the cost of a high-dimensional construction, which is known to be necessary to achieve the optimal ε dependence (Arjevani et al., 2023). For simplicity, we only St prove a version for zero-respecting algorithms, i.e., algorithms which satsify supp(xt+1 ) ⊆ τ =1 supp(gτ ), the extension to randomised follows with the randomised rotation argument (Arjevani et al., 2023). √ Theorem 3.8. Let c := 1/300, ∆1 , L, σ > 0, p ∈ (1, 2] and assume 0 < ε ≤ c min σ, ∆1 L . Then, for all m, n ∈ N≥2 with max {m, n} ≥ min {m, n} ∆ε12L we have p ∆1 L cσ p−1 2 mzr . ε ∥·∥S1 , ∆1 , L, p, σ ≥ c min {m, n} · ε2 ε In particular, M UON is first-order optimal for ∥·∥S1 -stationarity. The formal proof can be found in Section E.1 and is based on embedding κ instances of the Rd construction (Arjevani et al., 2023) into matrices of size κd × κ. The dimension requirement follows from the fact that the Rd construction requires d ≳ ∆ε12L to achieve the optimal ε dependence.
4
Experiments
This section aims to empirically validate the theoretical findings of this paper. Since it was empirically observed that the heavy-tailed gradient distribution occur in language modelling (J. Zhang et al., 2020b; Ahn et al., 2024), we focus on this task.3 Experimental Details. We pre-train the 70m parameter architecture of the Pythia series of models (Biderman et al., 2023) from scratch on the FineWeb-Edu dataset4 (Penedo et al., 2024) for Chinchilla-optimal 1.4B tokens (Hoffmann et al., 2022). Following previous work, we choose the sequence length to be 2048 and the batch size to be 512 (Biderman et al., 2023). Additionally we use the standard hyper-parameters β = 0.95, λ = 0.1 for MUON and use the WSD learning rate schedule (Zhai et al., 2022; Hägele et al., 2024; Hu et al., 2024) with a warmup of 5% and a cooldown period of 20%. We tune learning rates across the grid [0.005, 0.0075, 0.01, 0.0125, 0.015, 0.03, 0.06] and repeat the experiments across the seeds 0, 1, 2. Further details can be found in Section F. Norm Change in Noise Assumption. Section 3 analyses N SSD -M directly in its native geometry rather than reducing to the Euclidean case. In particular, we consider the quantity p
p := E [∥∇f (x, ξ) − ∇F (x)∥∗ ] σ∥·∥ ,p ∗
rather than
3 Our code is available on Github: https://github.com/fhueb/ht-schatten-experiments 4 https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu
7
p
p := E [∥∇f (x, ξ) − ∇F (x)∥F ] σF,p
22.5 20.0
15.0
22.5
Mean Ratio Worst-Case Ratio
20.0 17.5 S1 / S2
S1 / S2
17.5
Weight Family att.qkv att.out mlp.up mlp.down
12.5
Mean Ratio Worst-Case Ratio
15.0 12.5
10.0
10.0
7.5
7.5
5.0
5.0
(a) Ratios at Initialisation
Weight Family att.qkv att.out mlp.up mlp.down
(b) Ratios at Final Checkpoint
Figure 1: Ratios σS1 ,p /σF,p for p = 1.5 across all weight matrices trained by MUON. Colours indicate weight families, ordered by layer p depth within. While non-trivial, the noise ratios stay considerably below the worst-case ratio min {m, n} ≈ 22.5. which would be consistent with previous work. As noted by Shen et al. (2025), such a norm change in assumptions may implicitly introduce dimension dependencies. In the case of M UON, i.e., ∥·∥∗ = p min {m, n}σF,p . We therefore empirically ∥·∥S1 , we have the worst-case guarantee σS1 ,p ≤ calculate the ratio σS1 ,p/σF,p for p = 1.5 for all weight matrices trained by M UON, as shown in Figure 1. We observe that, while non-trivial, the noise ratios are far below the worst-case factor p min {m, n} ≈ 22.5. Furthermore the ratio is significantly lower at initialisation than at the final checkpoint, which may further explain the observation that M UON performs particularly well at the start of training (J. Liu et al., 2025). We repeat the same experiment while adding the input and output embedding layer, but measure the noise in ℓ-geometry, reporting a [0, 1] normalised ratio in Figure 2. Matrices that were empirically observed to be more efficiently trained by M UON have a lower Schatten ratio, while we observe that embeddings have a lower ℓ ratio. This possibly explains why embedding layers are better trained by A DAM. For additional details we refer to Section F and Shen et al. (2025), who conduct similar experiments for the smoothness constant. Impact of Norm Choice. To confirm the theoretical insights from Table 1: Final perplexity Theorem 3.1, we compare N SSD -M in different Schatten-r geome- for Schatten-r geometries. tries, r ∈ {1, 8/7, 4/3, 3/2, 2, 3, 4, 8, ∞}. Here r = 2 corresponds to the r Perplexity ↓ standard Euclidean geometry, while r = ∞ corresponds to the spectral geometry of M UON. Our theory suggests that increasing r should 1 88.89 ± 0.48 8/7 initially improve performance: larger r corresponds to a stronger dual 86.50 ± 0.91 4/3 stationarity measure until r ≈ p/p−1. Beyond this point, the stronger 67.61 ± 5.98 3/2 60.57 ± 4.48 stationarity competes with additional dimension dependence it intro2 40.18 ± 0.09 duces (Theorem C.6), so the performance may plateau or decrease. 3
34.51 ± 0.13
This is precisely the behaviour observed in Table 1. Performance im4 33.54 ± 0.04 proves sharply from r = 1 to r = 3, after which it essentially plateaus. 8 33.02 ± 0.01 These results are consistent with a tail index of approximately p = 1.5 ∞ 33.34 ± 0.02 for the gradient noise. We also observe that the spectral geometry of M UON, corresponding to r = ∞, substantially outperforms the Euclidean geometry, in line with Theorem 3.5. Finally, although intermediate Schatten geometries perform competitively, they require a full SVD, whereas M UON can use Newton–Schulz iterations. Hence M UON, which automatically adapts to the optimal geometry while being computationally efficient, is still preferable.
5
Related Work
Steepest Descent in General Norms. While classical S GD can be seen as steepest descent in the Euclidean ℓ2 norm, early works have proposed to consider general-norm (normalised) steepest descent algorithms (Boyd and Vandenberghe, 2004, Section 9.4). In particular, Kelner et al. (2014) proposed considering ℓ∞ -steepest descent for specific graph tasks, and Carlson et al. (2015a) suggest considering S∞ -steepest descent for Restricted Boltzmann Machines and Deep Learning (Carlson et al., 2015b). Independently, Tuddenham et al. (2022) proposed normalised S∞ -steepest descent with momentum for machine learning tasks under the name O RTHOGONAL -S GDM. This algorithm first 8
0.8 0.6
Weight Family att.qkv att.out mlp.up mlp.down embd
1.0
Mean Ratio
0.8
Normalized 1/ 2
Normalized S1/ S2
1.0
0.4 0.2 0.0
0.6
Weight Family att.qkv att.out mlp.up mlp.down embd
Mean Ratio
0.4 0.2 0.0
(a) Ratios in S1 .
(b) Ratios in ℓ1 .
Figure 2: Normalised noise ratios including the embedding and unembedding layer at initialisation. A normalised value of 0 corresponds to σ1 = σ2 , a value of 1 corresponds to σ1 = ρσ2 , where ρ is the worst-case ratio. Interestingly, we find that the ratio is significantly smaller for the embedding and unembedding layers for σℓ1 /σℓ2 possibly explaining why the ℓ∞ -norm is favourable over Schatten-∞ norm in practice for those layers.
orthogonalises each stochastic gradient and then applies momentum to the resulting orthogonalised directions. Neither approach was widely adopted until recently, when Bernstein and Newhouse (2024) proposed using Newton-Schulz iterations (Kovarik, 1970) to efficiently compute the orthogonalisation step. The blog (Jordan et al., 2024b) popularised the normalised S∞ -steepest descent approach under the name M UON, by changing the order of orthogonalisation and momentum in O RTHOGONAL -S GDM. Recently, Pethick et al. (2025b) and J. Liu et al. (2025) suggest M UON with decoupled weight decay as a natural extension from the theoretical and empirical perspective respectively. Theory of M UON. Motivated by the strong empirical performance of M UON, a rapidly growing line of literature studies its convergence properties. J. Li and Hong (2025) and Shen et al. (2025) analyse M UON under L-smoothness and Euclidean bounded-variance, and later relate its behaviour to lowrank and approximately blockwise-diagonal Hessian structure. Going beyond the S∞ geometry of M UON, Kovalev (2025) and Pethick et al. (2025b) linked this family of algorithms to linear minimisation oracle (lmo) based algorithms, and establish convergence for arbitrary norms in finitedimensional spaces. Guarantees for M UON with decoupled weight decay (Pethick et al., 2025b) and extensions to Nesterov momentum (L. Chen et al., 2025; Sfyraki and Wang, 2025) draw a close connection to stochastic conditional gradient (Mokhtari et al., 2020) and Lion-K. While the above works make their noise assumptions in the Euclidean space, Kovalev and Borodich (2025) consider the weighted-Euclidean assumption E[∥∇f (x, ξ) − ∇F (x)∥2Σ−1 ] ≤ ∥Σ∥S1 , where Σ is a self-adjoint positive definite operator and ∥v∥2Σ−1 = v ⊤ Σ−1 v. Heavy-Tailed Noise. A separate line of work suggests that bounded variance may be too restrictive for modern deep learning, and instead suggests studying stochastic gradients with bounded p-th central moments for p ∈ (1, 2] (Nemirovskij and Yudin, 1983; J. Zhang et al., 2020a; Garg et al., 2021). In Euclidean geometry, a large body of work establishes convergence of gradient-clipping based methods under heavy-tailed noise (e.g., J. Zhang et al., 2020a; Cutkosky and Mehta, 2021; Sadiev et al., 2023). Recent works establish optimal sample complexities for N SSD -M without the need for gradient-clipping under heavy-tailed noise (Hübler et al., 2025; Z. Liu and Zhou, 2025). Beyond Euclidean spaces, Cutkosky and Mehta (2021) provide guarantees for 2-uniformly convex Banach spaces, but require an additional gradient clipping step. This extension covers Schattenr norms for r ∈ (1, 2], but not the Schatten-∞ geometry required for M UON. Complementary literature examines heavy-tailed data distributions (Kim et al., 2026; Kunstner and Bach, 2026). Concurrent Work. While preparing this paper, we became aware of the concurrent preprint (Y. Zhang and Lin, 2026). The authors analyse M UON under heavy-tailed noise and (L0 , L1 )smoothness, where both assumptions are stated in the Frobenius norm. This results in a 3p−2 p O ρ p−1 ∆ε12L σεF p−1 convergence guarantee, which does not improve over the stationarity adjusted guarantee of Euclidean algorithms (see Equation (5)). In contrast, our work considers general norms beyond S∞ , assumptions stated in the native optimisation geometry and obtains a strictly im9
proved dimension dependence. Additionally we provide lower-bounds showing that our dimension dependence is first-order optimal.
6
Conclusion and Limitations
We provide convergence guarantees for normalised steepest descent with momentum in general Banach spaces under heavy-tailed noise. In particular, for M UON, we obtain a nuclearstationarity guarantee with substantially improved dimension dependence compared to the corresponding stationarity-adjusted Euclidean guarantees. We further prove that the remaining dimension dependence is first-order optimal and therefore cannot be removed. Our experiments connect these results to empirical observations in language-model training: varying the Schatten-r descent geometry produces the predicted improvement from Euclidean geometry to larger r, and weight matrices trained by MUON exhibit favourable Schatten noise ratios, while embeddings appear better aligned with an ℓ∞ -type geometry. Our work has several limitations. The empirical evaluation is restricted to a 70M-parameter model trained on a single dataset, so the observed phenomena should be validated across larger variants. On the theoretical side, our guarantees are in expectation, and extending them to high-probability bounds under heavy-tailed noise remains open. Finally, the analysis studies idealised normalised steepest-descent updates, whereas practical M UON implementations rely on approximate Newton–Schulz iterations.
Acknowledgments and Disclosure of Funding FH gratefully acknowledges financial support from the ETH research grant and Swiss National Science Foundation (SNSF) Project Funding No. 200021-207343, the Alexander von Humboldt Foundation, and the European Union’s Horizon Europe research and innovation programme under grant agreements No. 101120237 (ELIAS) and No. 101070617 (ELSA). Views and opinions expressed are however those of the author only and do not necessarily reflect those of the European Union or European Commission. Neither the European Union nor the granting authority can be held responsible for them. SS gratefully acknowledges generous support from the Alexander von Humboldt Foundation.
References Kovarik, Zdislav (1970). “Some iterative methods for improving orthonormality”. In: SIAM Journal on Numerical Analysis 7.3, pp. 386–389. Nemirovskij, Arkadij Semenovič and David Borisovich Yudin (1983). “Problem Complexity and Method Efficiency in Optimization”. In. Ball, Keith, Eric A. Carlen, and Elliott H. Lieb (Dec. 1994). “Sharp Uniform Convexity and Smoothness Inequalities for Trace Norms”. In: Inventiones mathematicae 115.1, pp. 463–482. ISSN: 1432-1297. DOI: 10.1007/BF01231769. Hirzallah, Omar and Fuad Kittaneh (2002). “Non-Commutative Clarkson Inequalities for Unitarily Invariant Norms”. In: Pacific J. Math 202.2, pp. 363–369. Boyd, Stephen and Lieven Vandenberghe (2004). Convex Optimization. Cambridge: Cambridge University Press. Wenzel, Jörg (2005). “Strong Martingale Type and Uniform Smoothness”. In: Journal of Convex Analysis 12.1, pp. 159–171. Bhatia, Rajendra (2013). Matrix Analysis. Springer Science & Business Media. Kelner, Jonathan A. et al. (2014). “An Almost-Linear-Time Algorithm for Approximate Max Flow in Undirected Graphs, and Its Multicommodity Generalizations”. In: Proceedings of the 2014 Annual ACM-SIAM Symposium on Discrete Algorithms (SODA), pp. 217–226. DOI: 10.1137/1. 9781611973402.16. Kingma, Diederik P. and Jimmy Ba (2014). “Adam: A Method for Stochastic Optimization”. In: arXiv preprint arXiv:1412.6980. Mendel, Manor and Assaf Naor (2014). “Nonlinear Spectral Calculus and Super-Expanders”. In: Publications mathématiques de l’IHÉS 119, pp. 1–95.
10
Carlson, David, Volkan Cevher, and Lawrence Carin (May 2015a). “Stochastic Spectral Descent for Restricted Boltzmann Machines”. In: Proceedings of the Eighteenth International Conference on Artificial Intelligence and Statistics. Ed. by Guy Lebanon and S. V. N. Vishwanathan. Vol. 38. Proceedings of Machine Learning Research. San Diego, California, USA: PMLR, pp. 111–119. Carlson, David et al. (2015b). “Preconditioned Spectral Descent for Deep Learning”. In: Advances in Neural Information Processing Systems. Ed. by C. Cortes et al. Vol. 28. Curran Associates, Inc. URL: https : / / proceedings . neurips . cc / paper _ files / paper / 2015 / file / f50a6c02a3fc5a3a5d4d9391f05f3efc-Paper.pdf. Bottou, Léon, Frank E. Curtis, and Jorge Nocedal (2018). “Optimization Methods for Large-Scale Machine Learning”. In: SIAM Review 60.2, pp. 223–311. DOI: 10.1137/16M1080173. eprint: https://doi.org/10.1137/16M1080173. URL: https://doi.org/10.1137/16M1080173. Simsekli, Umut, Levent Sagun, and Mert Gurbuzbalaban (June 2019). “A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networks”. In: Proceedings of the 36th International Conference on Machine Learning. Ed. by Kamalika Chaudhuri and Ruslan Salakhutdinov. Vol. 97. Proceedings of Machine Learning Research. PMLR, pp. 5827–5837. Balles, Lukas, Fabian Pedregosa, and Nicolas Le Roux (2020). “The Geometry of Sign Gradient Descent”. In: arXiv preprint arXiv:2002.08056. arXiv: 2002.08056. Cutkosky, Ashok and Harsh Mehta (July 2020). “Momentum Improves Normalized SGD”. In: Proceedings of the 37th International Conference on Machine Learning. Ed. by Hal Daumé III and Aarti Singh. Vol. 119. Proceedings of Machine Learning Research. PMLR, pp. 2260–2268. Mokhtari, Aryan, Hamed Hassani, and Amin Karbasi (2020). “Stochastic Conditional Gradient Methods: From Convex Minimization to Submodular Maximization”. In: Journal of Machine Learning Research 21.105, pp. 1–49. Zhang, Bohang et al. (2020). “Improved Analysis of Clipping Algorithms for Non-Convex Optimization”. In: Advances in Neural Information Processing Systems. Ed. by H. Larochelle et al. Vol. 33. Curran Associates, Inc., pp. 15511–15521. Zhang, Jingzhao et al. (2020a). “Why are Adaptive Methods Good for Attention Models?” In: Advances in Neural Information Processing Systems. Ed. by H. Larochelle et al. Vol. 33. Curran Associates, Inc., pp. 15383–15393. Zhang, Jingzhao et al. (2020b). “Why Gradient Clipping Accelerates Training: A Theoretical Justification for Adaptivity”. In: International Conference on Learning Representations. Cutkosky, Ashok and Harsh Mehta (2021). “High-Probability Bounds for Non-Convex Stochastic Optimization with Heavy Tails”. In: Advances in Neural Information Processing Systems. Ed. by M. Ranzato et al. Vol. 34. Curran Associates, Inc., pp. 4883–4895. Garg, Saurabh et al. (July 2021). “On Proximal Policy Optimization’s Heavy-Tailed Gradients”. In: Proceedings of the 38th International Conference on Machine Learning. Ed. by Marina Meila and Tong Zhang. Vol. 139. Proceedings of Machine Learning Research. PMLR, pp. 3610–3619. Zhao, Shen-Yi, Yin-Peng Xie, and Wu-Jun Li (Feb. 2021). “On the Convergence and Improvement of Stochastic Normalized Gradient Descent”. In: Science China Information Sciences 64.3, p. 132103. ISSN: 1869-1919. DOI: 10.1007/s11432-020-3023-7. Crawshaw, Michael et al. (2022). “Robustness to Unbounded Smoothness of Generalized SignSGD”. In: Advances in Neural Information Processing Systems. Ed. by S. Koyejo et al. Vol. 35. Curran Associates, Inc., pp. 9955–9968. Hoffmann, Jordan et al. (2022). “An Empirical Analysis of Compute-Optimal Large Language Model Training”. In: Advances in Neural Information Processing Systems. Ed. by S. Koyejo et al. Vol. 35. Curran Associates, Inc., pp. 30016–30030. Tuddenham, Mark, Adam Prügel-Bennett, and Jonathan Hare (2022). “Orthogonalising Gradients to Speed up Neural Network Optimisation”. In: arXiv preprint arXiv:2202.07052. arXiv: 2202. 07052. Zhai, Xiaohua et al. (2022). “Scaling vision transformers”. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12104–12113. Arjevani, Yossi et al. (May 2023). “Lower Bounds for Non-Convex Stochastic Optimization”. In: Mathematical Programming 199.1, pp. 165–214. ISSN: 1436-4646. DOI: 10 . 1007 / s10107 022-01822-7. Biderman, Stella et al. (2023). “Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling”. In: International Conference on Machine Learning. PMLR, pp. 2397–2430. Chowdhery, Aakanksha et al. (2023). “PaLM: Scaling Language Modeling with Pathways”. In: Journal of machine learning research 24.240, pp. 1–113.
11
Faw, Matthew et al. (July 2023). “Beyond Uniform Smoothness: A Stopped Analysis of Adaptive SGD”. In: Proceedings of Thirty Sixth Conference on Learning Theory. Ed. by Gergely Neu and Lorenzo Rosasco. Vol. 195. Proceedings of Machine Learning Research. PMLR, pp. 89–160. Kunstner, Frederik et al. (2023). “Noise Is Not the Main Factor Behind the Gap Between Sgd and Adam on Transformers, But Sign Descent Might Be”. In: The Eleventh International Conference on Learning Representations. URL: https://openreview.net/forum?id=a65YK0cqH8g. Li, Haochuan, Alexander Rakhlin, and Ali Jadbabaie (2023). “Convergence of Adam Under Relaxed Assumptions”. In: Thirty-Seventh Conference on Neural Information Processing Systems. Sadiev, Abdurakhmon et al. (July 2023). “High-Probability Bounds for Stochastic Optimization and Variational Inequalities: The Case of Unbounded Variance”. In: Proceedings of the 40th International Conference on Machine Learning. Ed. by Andreas Krause et al. Vol. 202. Proceedings of Machine Learning Research. PMLR, pp. 29563–29648. Ahn, Kwangjun et al. (2024). “Linear attention is (maybe) all you need (to understand Transformer optimization)”. In: The Twelfth International Conference on Learning Representations. URL : https://openreview.net/forum?id=0uI5415ry7. Battash, Barak, Lior Wolf, and Ofir Lindenbaum (Feb. 2024). “Revisiting the Noise Model of Stochastic Gradient Descent”. In: Proceedings of The 27th International Conference on Artificial Intelligence and Statistics. Ed. by Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li. Vol. 238. Proceedings of Machine Learning Research. PMLR, pp. 4780–4788. URL: https : //proceedings.mlr.press/v238/battash24a.html. Bernstein, Jeremy and Laker Newhouse (2024). “Old Optimizer, New Norm: An Anthology”. In: arXiv preprint arXiv:2409.20325. Hägele, Alexander et al. (2024). “Scaling laws and compute-optimal training beyond fixed training durations”. In: Advances in Neural Information Processing Systems 37, pp. 76232–76264. Hu, Shengding et al. (2024). “MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies”. In: First Conference on Language Modeling. URL: https : / / openreview.net/forum?id=3X2L2TFr0f. Hübler, Florian et al. (May 2024). “Parameter-Agnostic Optimization under Relaxed Smoothness”. In: Proceedings of The 27th International Conference on Artificial Intelligence and Statistics. Ed. by Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li. Vol. 238. Proceedings of Machine Learning Research. PMLR, pp. 4861–4869. URL: https://proceedings.mlr.press/v238/ hubler24a.html. Jordan, Keller et al. (2024a). modded-nanogpt: Speedrunning the NanoGPT baseline. URL: https: //github.com/KellerJordan/modded-nanogpt. Jordan, Keller et al. (2024b). “Muon: An Optimizer for Hidden Layers in Neural Networks”. In: URL https://kellerjordan. github. io/posts/muon 6.3, p. 4. Penedo, Guilherme et al. (2024). “The fineweb datasets: Decanting the web for the finest text data at scale”. In: Advances in Neural Information Processing Systems 37, pp. 30811–30849. Vankov, Daniil et al. (2024). “Optimizing (L0 ,L1 )-Smooth Functions by Gradient Methods”. In: arXiv preprint arXiv:2410.10800. arXiv: 2410.10800. Chen, Lizhang, Jonathan Li, and Qiang Liu (2025). “Muon Optimizes Under Spectral Norm Constraints”. In: arXiv preprint arXiv:2506.15054. arXiv: 2506.15054. Chezhegov, Savelii et al. (2025). “Convergence of Clipped-SGD for Convex (L_0, L_1)-Smooth Optimization with Heavy-Tailed Noise”. In: arXiv preprint arXiv:2505.20817. Hübler, Florian, Ilyas Fatkhullin, and Niao He (May 2025). “From Gradient Clipping to Normalization for Heavy Tailed SGD”. In: Proceedings of The 28th International Conference on Artificial Intelligence and Statistics. Ed. by Yingzhen Li et al. Vol. 258. Proceedings of Machine Learning Research. PMLR, pp. 2413–2421. URL: https : / / proceedings . mlr . press / v258 / hubler25a.html. Kovalev, Dmitry (2025). “Understanding Gradient Orthogonalization for Deep Learning Via NonEuclidean Trust-Region Optimization”. In: arXiv preprint arXiv:2503.12645. arXiv: 2503 . 12645. Kovalev, Dmitry and Ekaterina Borodich (2025). “Non-Euclidean SGD for Structured Optimization: Unified Analysis and Improved Rates”. In: arXiv preprint arXiv:2511.11466. arXiv: 2511. 11466. Li, Jiaxiang and Mingyi Hong (2025). “A Note on the Convergence of Muon”. In: arXiv preprint arXiv:2502.02900. arXiv: 2502.02900. Liu, Jingyuan et al. (2025). “Muon Is Scalable for LLM Training”. In: arXiv preprint arXiv:2502.16982. arXiv: 2502.16982.
12
Liu, Zijian and Zhengyuan Zhou (2025). “Nonconvex Stochastic Optimization under Heavy-Tailed Noises: Optimal Convergence without Gradient Clipping”. In: The Thirteenth International Conference on Learning Representations. URL: https : / / openreview . net / forum ? id = NKotdPUc3L. Pethick, Thomas et al. (2025a). “Generalized Gradient Norm Clipping & Non-Euclidean (L0 , L1 )Smoothness”. In: The Thirty-Ninth Annual Conference on Neural Information Processing Systems. Pethick, Thomas et al. (July 2025b). “Training Deep Learning Models with Norm-Constrained LMOs”. In: Proceedings of the 42nd International Conference on Machine Learning. Ed. by Aarti Singh et al. Vol. 267. Proceedings of Machine Learning Research. PMLR, pp. 49069–49104. Riabinin, Artem et al. (2025). “Gluon: Making Muon & Scion Great Again! (Bridging Theory and Practice of Lmo-Based Optimizers for Llms)”. In: arXiv preprint arXiv:2505.13416. arXiv: 2505.13416. Sfyraki, Maria-Eleni and Jun-Kun Wang (2025). “Lions and Muons: Optimization Via Stochastic Frank-Wolfe”. In: arXiv preprint arXiv:2506.04192. arXiv: 2506.04192. Shen, Wei et al. (2025). “On the Convergence Analysis of Muon”. In: arXiv preprint arXiv:2505.23737. arXiv: 2505.23737. Team, Kimi et al. (2025). “Kimi k2: Open agentic intelligence”. In: arXiv preprint arXiv:2507.20534. DeepSeek-AI (2026). DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence. Technical report. URL: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/ main/DeepSeek_V4.pdf. Kim, Juno et al. (2026). “Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory”. In: arXiv preprint arXiv:2603.26554. Kunstner, Frederik and Francis Bach (2026). “Scaling Laws for Gradient Descent and Sign Descent for Linear Bigram Models under Zipf’s Law”. In: The Thirty-ninth Annual Conference on Neural Information Processing Systems. URL: https://openreview.net/forum?id=VUbwLjLkws. Zhang, Yubo and Junhong Lin (2026). “On Convergence of Muon for Nonconvex Stochastic Optimization Under Generalized Smoothness”. In: TechRxiv 2026.0227. DOI: 10.36227/techrxiv. 177220100.05629800/v1.
13
Contents 1
Introduction
1
2
Preliminaries
3
2.1
3
3
Moment Bounds for General Norms . . . . . . . . . . . . . . . . . . . . . . . . .
Main Result
4
3.1
Guarantees for Uniformly Convex Norms . . . . . . . . . . . . . . . . . . . . . .
4
3.2
Guarantees for M UON . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
5
3.3
First-Order Lower-Bounds for Nuclear-Stationarity . . . . . . . . . . . . . . . . .
6
4
Experiments
7
5
Related Work
8
6
Conclusion and Limitations
10
A Technical Lemmas
15
B (L0 , L1 )-Smoothness
15
B.1 Related Work on (L0 , L1 )-Smoothness . . . . . . . . . . . . . . . . . . . . . . . .
15
B.2 Properties of (L0 , L1 )-Smoothness in Banach Spaces . . . . . . . . . . . . . . . .
15
C Proofs of Main Results
15
C.1 Proofs of Section 3.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
15
C.2 Proofs of Section 3.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
19
C.2.1
Banach Space Setting . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
19
C.2.2
Application to M UON . . . . . . . . . . . . . . . . . . . . . . . . . . . .
21
D Extension to Stochastic Conditional Gradient
22
E Lower Bound Proofs
25
E.1 High-Dimensional Lower-Bound . . . . . . . . . . . . . . . . . . . . . . . . . . . F Experimental Details
27 30
14
A
Technical Lemmas
Lemma A.1 (see Mendel and Naor, 2014, Equation (167)). Let (X, ∥·∥) be a Banach space. i) For all 1 < p′ ≤ p ≤ 2 we have Sp′ (X) ≤ Sp (X). ii) For all 2 ≤ q ≤ q ′ < ∞ we have Kq′ (X) ≤ Kq (X). Lemma A.2 (see Bhatia, 2013). Let 1 ≤ p ≤ q ≤ ∞. Then the Schatten-p and Schatten-q norms satisfy ∥A∥Sq ≤ ∥A∥Sp ≤ ρ ∥A∥Sq , 1
where ρ = min {m, n} p
B
− q1
, for all A ∈ Rm×n .
(L0 , L1 )-Smoothness
We will derive our results for the weaker smoothness notion of (L0 , L1 )-smoothness. This section contains the related work, and required definitions and properties of this smoothness notion in general Banach spaces. B.1
Related Work on (L0 , L1 )-Smoothness
Motivated by observations in language modelling, J. Zhang et al. (2020b) propose considering (L0 , L1 )-smoothness, allowing the smoothness to grow affinely with the gradient norm. A long line of works analyse diverse algorithms in this setting (e.g., Crawshaw et al., 2022; Faw et al., 2023; H. Li et al., 2023; Chezhegov et al., 2025). In Euclidean geometry, B. Zhang et al. (2020) and Zhao et al. (2021) prove the convergence of N SSD -M, and Z. Liu and Zhou (2025) extend this result to the heavy-tailed setting. Beyond Euclidean geometry, Balles et al. (2020) provide guarantees for the ℓ∞ -geometry in the deterministic setting, and Pethick et al. (2025a) provide convergence guarantees in arbitrary finite-dimensional spaces in the stochastic setting by adding gradient-clipping, while Riabinin et al. (2025)5 concurrently prove convergence without the clipping step. B.2
Properties of (L0 , L1 )-Smoothness in Banach Spaces
Definition B.1 ((L0 , L1 )-smoothness). Let L0 , L1 ≥ 0. A function F : B → R is said to be (L0 , L1 )-smooth if it is differentiable and, for all x, y ∈ B, ∥∇F (x) − ∇F (y)∥∗ ≤ (L0 + L1 ∥∇F (x)∥∗ )A(L1 ∥x − y∥) ∥x − y∥ , c
where A(c) := e c−1 . Lemma B.2 (Smoothness Upper-Bound). Let F : B → R be an (L0 , L1 )-smooth function. Then, for all x, y ∈ B, 2
F (y) ≤ F (x) + ⟨∇F (x), y − x⟩ + (L0 + L1 ∥∇F (x)∥∗ )B(L1 ∥x − y∥) ∥x − y∥ , c
. where B(c) = e −1−c c2 Proof. This directly follows with the same proof as (Vankov et al., 2024, Lemma 2.5) by noting that all used inequalities also hold in general Banach spaces.
C
Proofs of Main Results
C.1
Proofs of Section 3.1
This section contains generalisations of the main results from Section 3 and their proofs. We will derive all results for the weaker (L0 , L1 )-smoothness notion popularised by J. Zhang et al. (2020b). 5 We note that this preprint has a typo. While their variance assumption (Assumption 2) is stated in the dual norm, the authors use the inner product structure of the Frobenius norm in the proof of Theorem 2. In particular, a fix requires switching to σF and introduces a ρ2 dependence.
15
Assumption 4 ((L0 , L1 )-smoothness). We assume the objective function F is (L0 , L1 )-smooth. We first derive a slightly improved descent lemma over (Pethick et al., 2025a). Lemma C.1 (Descent Lemma). Suppose Assumptions 1 and 4 hold. Furthermore let T ∈ N≥1 and consider the iterates (xt ) generated by N SSD -M. Then T X
T T X X ηt − L1 ηt2 Bt ∥∇F (xt )∥∗ ≤ ∆1 + L0 ηt2 Bt + 2 ηt ∥∇F (xt ) − mt ∥∗ ,
t=1
t=1
t=1
c
. where Bt := B(L1 ηt ), B(c) = e −1−c c2 Proof. This proof generalises (Hübler et al., 2024, Lemma 12) to Banach spaces. Using Lemma B.2 we get F (xt+1 ) − F (xt ) ≤ ⟨∇F (xt ), xt+1 − xt ⟩ + (L0 + L1 ∥∇F (xt )∥∗ )B(L1 ηt ) ∥xt+1 − xt ∥
2
= ηt ⟨∇F (xt ), lmo(mt )⟩ + (L0 + L1 ∥∇F (xt )∥∗ )B(L1 ηt )ηt2 . Furthermore ⟨∇F (xt ), lmo(mt )⟩ = ⟨∇F (xt ) − mt , lmo(mt )⟩ + ⟨mt , lmo(mt )⟩ ≤ ∥∇F (xt ) − mt ∥∗ ∥lmo(mt )∥ − ∥mt ∥∗ ≤ 2 ∥∇F (xt ) − mt ∥∗ − ∥∇F (xt )∥∗ , where we used the definition of lmo and the dual norm in the first, and the triangle inequality in the second step. Combining yields F (xt+1 ) − F (xt ) ≤ ηt (2 ∥∇F (xt ) − mt ∥∗ − ∥∇F (xt )∥∗ ) + (L0 + L1 ∥∇F (xt )∥∗ )B(L1 ηt )ηt2 = L1 ηt2 B(L1 ηt ) − ηt ∥∇F (xt )∥∗ + 2ηt ∥∇F (xt ) − mt ∥∗ + L0 B(L1 ηt )ηt2 . Summing up and re-arranging finishes the proof. Next we derive the main technical lemma of our analysis, which bounds the deviation of the momentum from the true gradient for general norms, without requiring a conversion to the Euclidean norm. Lemma C.2. Suppose Assumptions 3 and 4 hold and that (B ∗ , ∥·∥∗ ) is p′ -uniformly smooth with smoothness constant Sp′ (B ∗ ) ≤ S. Then the iterates generated by N SSD -M satisfy !1/r t t t X X X r E [∥∇F (xt ) − mt ∥∗ ] ≤ σS bτ,t + L0 cτ,t + L1 cτ,t E [∥∇F (xτ )∥∗ ] , τ =1
τ =2
τ =2
′
where r := min {p, p } and, for easier readability, we define b1,t :=
t Y
and for τ ≥ 2,
βκ ,
bτ,t := (1 − βτ )
κ=2
t Y
βκ ,
cτ,t := ητ −1 Aτ −1
κ=τ +1
t Y
βκ ,
κ=τ
L1 ητ
where Aτ = A(L1 ητ ) = e L1 ητ−1 . In particular, for constant parameters ηt ≡ η and βt ≡ β, we get T X
E [∥∇F (xt ) − mt ∥∗ ] ≤
t=1
T r−1 σS ηAT ηA X + σST (1 − β) r + L0 + L1 E [∥∇F (xt )∥∗ ] , 1−β 1−β 1 − β t=1
where A ≡ A(L1 η). Proof. First note that the standard decomposition of µt := mt −∇F (xt ) into a noise and smoothness term (Cutkosky and Mehta, 2020), µt =
t X
bτ,t (∇f (xτ , ξτ ) − ∇F (xτ )) +
t Y t X τ =2 κ=τ
τ =1
16
βκ (∇F (xτ −1 ) − ∇F (xτ )),
is norm independent and hence also holds in general Banach spaces. Next, we have ∥∇F (xτ ) − ∇F (xτ −1 )∥∗ ≤ (L0 + L1 ∥∇F (xτ )∥∗ )A(L1 ητ −1 ) ∥xτ − xτ −1 ∥ ≤ Lτ Aτ −1 ητ −1 , where we use the slight abuse of notation Lτ := L0 + L1 ∥∇F (xτ )∥∗ . Hence, by our definition of bτ,t and cτ,t in the statement, the triangle inequality yields # " t ! t t X X Y bτ,t (∇f (xτ , ξτ ) − ∇F (xτ )) + βκ E [∥∇F (xτ −1 ) − ∇F (xτ )∥∗ ] E [∥µt ∥∗ ] ≤ E τ =1
" ≤E
t X
# bτ,t (∇f (xτ , ξτ ) − ∇F (xτ ))
τ =1
κ=τ
τ =2
∗
+
t X
cτ,t E [Lτ ] .
τ =2
∗
Now note that (∇f (xτ , ξτ ) − ∇F (xτ ))tτ =1 is a martingale difference sequence in B ∗ , and, by Lemma A.1, B ∗ is also r-uniformly smooth with Sr (B ∗ ) = Sp′ (B ∗ ) ≤ S. Hence, we may apply Theorem 2.4 to get # " t " t r #1/r X X bτ,t (∇f (xτ , ξτ ) − ∇F (xτ )) bτ,t (∇f (xτ , ξτ ) − ∇F (xτ )) ≤E E τ =1
τ =1
∗
≤S
t X
∗
!1/r r brτ,t E [∥∇f (xτ , ξτ ) − ∇F (xτ )∥∗ ]
τ =1
≤ σS
t X
!1/r brτ,t
,
τ =1
where we used r ≤ p and Assumption 3 in the last step. For the constant parameter version note that t X
brτ,t = β r(t−1) + (1 − β)r
τ =1 t X
t X
β r(t−τ ) ≤ β r(t−1) +
τ =2
cτ,t = ηA
t−1 X
βτ ≤
τ =1
τ =2
(1 − β)r ≤ β r(t−1) + (1 − β)r−1 1 − βr
ηA , 1−β
where A ≡ A(L1 η), and hence σS
T t X X t=1
!1/r brτ,t
τ =1
L0
T X t X
r−1 1 ≤ σS + T (1 − β) r 1−β
cτ,t ≤ L0
t=1 τ =2
L1
T X t X
cτ,t E [∥∇F (xτ )∥∗ ] = L1
t=1 τ =2
ηAT 1−β T T X X τ =2
≤ L1 ≤
! cτ,t E [∥∇F (xτ )∥∗ ]
t=τ
T X ηA E [∥∇F (xτ )∥∗ ] 1−β τ =2
T L1 ηA X E [∥∇F (xt )∥∗ ] , 1 − β t=1
Combining the three bounds finishes the proof. Finally we are able to derive the main convergence result for N SSD -M in general Banach spaces. Theorem C.3. Suppose Assumptions 1, 3 and 4 hold and that (B, ∥·∥) is q-uniformly convex with constant Kq (B) ≤ K. Furthermore let p′ be the conjugate exponent of q, i.e., 1/q + 1/p′ = 1, and 17
r = min {p, p′ }. Then the iterates generated by N SSD -M with s 1 − β ∆1 (1 − β) , ηt ≡ η = min , 5L1 L0 T ( ( )) r r r ∆1 L0 3r−2 5L1 ∆1 2r−1 − 2r−1 βt ≡ β = 1 − min 1, max , ,T σ2 K 2 T 2σKT satisfy T
1X ∆1 L1 E [∥∇F (xt )∥∗ ] ≤ 12 +9 T t=1 T
r
r
∆1 L0 ∆1 L1 (σK) r−1 + 14 T T r−1 ! r 3r−2 ∆1 L0 (σK) (r−1) 10σK + 14 + r−1 . T T 2r−1
r−1 ! 2r−1
In particular, the sample complexity to reach an ε-∥·∥∗ -stationary point is at most ∆1 L 1 ∆1 L0 ∆1 L1 O + + ε ε2 ε
σK ε
r r−1
∆1 L 0 + ε2
σK ε
r r−1
+
σK ε
! 2r−1 r−1 .
Proof. Applying Lemma C.1 yields η
T X
E [∥∇F (xt )∥∗ ] ≤ ∆1 + L0 η 2 BT + L1 η 2 B
T X t=1
t=1
E [∥∇F (xt )∥∗ ] + 2
T X
ηt E [∥∇F (xt ) − mt ∥∗ ]
t=1
r−1 2A 1 ≤ ∆1 + L 0 η 2 T B + + 2σηK + T (1 − β) r 1−β 1−β X T 2A + L1 η 2 B + E [∥∇F (xt )∥∗ ] . 1 − β t=1 Next we calculate 2A L1 η A(5 − β) 5 L1 η 2 B + ≤η ≤ η 1−β 1−β 2 9
(6)
L1 η where we used B ≤ A/2 in the first, and our choice of η, which implies 1−β ≤ 15 and A ≤ 10/9, in the last inequality. Hence, re-arranging further yields T
r−1 1X 6L0 ηA 5σK 9∆1 E [∥∇F (xt )∥∗ ] ≤ + + + 5σK(1 − β) r T t=1 4ηT (1 − β) T (1 − β) s r−1 ∆1 L0 12∆1 L1 5σK ≤9 + + + 5σK(1 − β) r T (1 − β) T (1 − β) T (1 − β)
(7)
where we used that 1/ min {a, b} ≤ 1/a + 1/b and the definition of η in the last step. Next we upper bound the remaining 4 terms by our choice of momentum. To simplify exposition, let us denote r r r ∆1 L0 3r−2 5L1 ∆1 2r−1 := u := , v , w := T − 2r−1 σ2 K 2 T 2σKT 1 such that 1 − β = min {1, max {u, v, w}}. First note that, due to √1−β ≤ min 11,√u ≤ 1 + √1u , { }
s 9
r−1 ! 3r−2 r r ∆1 L0 ∆1 L0 ∆1 L0 (σK) (r−1) ≤9 +9 . T (1 − β) T T
18
For the second term we use the same argument as before, this time using v, to get r−1 r−1 ! 2r−1 ! 2r−1 r r r 2r−1 ∆1 L1 2 ∆1 L1 ∆1 L1 (σK) r−1 12∆1 L1 ∆1 L1 (σK) r−1 ≤ 12 +12 ≤ 12 +7 T (1 − β) T 5 T T T 1 For the third term note that 1−β ≤ w1 and hence
5σK 5σK ≤ r−1 . T (1 − β) T 2r−1 Finally, for the last term we use (1 − β)α ≤ uα + v α + wα to get r−1 r−1 r−1 r−1 ∆1 L0 3r−2 5L1 ∆1 2r−1 r 5σK(1 − β) ≤ 5σK + 5σK + 5σKT − 2r−1 σ2 K 2 T 2σKT r−1 r−1 ! 3r−2 ! 2r−1 r r σK ∆1 L0 (σK) r−1 5∆1 L1 (σK) r−1 =5 +5 + 5 r−1 T 2T T 2r−1 r−1 Plugging these bounds into (7) and using 5 52 2r−1 ≤ 7 finishes the proof. Remark C.4. The above result can readily be extended to the parameter-agnostic setting by using the same arguments as in (Hübler et al., 2024) to handle the non-constant parameters and error accumulation before ηt ≲ 1/L1 . C.2
Proofs of Section 3.2
In Theorem C.3 we provide guarantees for arbitrary p′ -uniformly smooth dual norms ∥·∥∗ , but do not achieve optimal rates whenever p′ < p. In this section we will use an orthogonal approach, where we instead assume that ∥·∥∗ ≤ ρ ∥·∥p for some p-uniformly smooth norm ∥·∥p . Such a constant always exists in the finite dimensional case, and will allow us to achieve optimal rates. We first derive a general result for N SSD -M in Section C.2.1, before specialising it to M UON in Section C.2.2. C.2.1
Banach Space Setting
Therefore we first formulate Lemma C.2 for this approach. Lemma C.5. Suppose Assumptions 3 and 4 hold, and that ∥·∥p ≤ ∥·∥∗ ≤ ρ ∥·∥p , where (B ∗ , ∥·∥p ) is p-uniformly smooth with smoothness constant Sp (B ∗ , ∥·∥p ) ≤ S. Then the iterates generated by N SSD -M with constant parameters ηt ≡ η and βt ≡ β satisfy T X
E [∥∇F (xt ) − mt ∥∗ ] ≤
t=1
T p−1 ρσS ηAT ηA X E [∥∇F (xt )∥∗ ] , +ρσST (1−β) p +L0 +L1 1−β 1−β 1 − β t=1
L1 η
where A ≡ A(L1 η) = e L1 η−1 . We note that such ρ always exists in finite dimensional, but may not exist in general Banach spaces. Proof. We proceed exactly as in the proof of Lemma C.2, but replace " t # " t r #1/r X X E bτ,t (∇f (xτ , ξτ ) − ∇F (xτ )) ≤E bτ,t (∇f (xτ , ξτ ) − ∇F (xτ )) τ =1
τ =1
∗
∗
with " E
t X τ =1
# bτ,t (∇f (xτ , ξτ ) − ∇F (xτ ))
≤ ρE ∗
≤ ρE
t X
bτ,t (∇f (xτ , ξτ ) − ∇F (xτ ))
τ =1
p
t X
p
bτ,t (∇f (xτ , ξτ ) − ∇F (xτ ))
τ =1
The proof closes the same way as before by using the p-uniform smoothness of ∥·∥p . 19
1/p p
.
Now we are ready to prove the norm-constant version of the convergence result for N SSD -M under heavy-tailed noise and (L0 , L1 )-smoothness. Theorem C.6. Suppose Assumptions 1, 3 and 4 hold. Further assume that ∥·∥p ≤ ∥·∥∗ ≤ ρ ∥·∥p , where (B ∗ , ∥·∥p ) is p-uniformly smooth with smoothness constant Sp (B ∗ , ∥·∥p ) ≤ S. Then the iterates generated by N SSD -M with s 1 − β ∆1 (1 − β) ηt ≡ η = min , , 5L1 L0 T ( ( )) p 3p−2 p p ∆1 L0 5L1 ∆1 2p−1 − 2p−1 βt ≡ β = 1 − min 1, max , ,T σ 2 ρ2 S 2 T 2σρST satisfy p−1 r p ! 3p−2 T ∆1 L 1 ∆1 L0 (ρσS) p−1 1X ∆1 L0 E [∥∇F (xt )∥∗ ] ≤ 9 + 12 + 14 T t=1 T T T p−1 p ! 2p−1 ∆1 L1 (ρσS) p−1 ρσS + 14 + 10 p−1 . T T 2p−1
In particular, the sample complexity to reach an ε-∥·∥∗ -stationary point is at most
∆1 L 1 ∆1 L0 ∆1 L1 O + + ε ε2 ε
ρσS ε
p p−1
∆1 L 0 + ε2
ρσS ε
p p−1
+
ρσS ε
! 2p−1 p−1 .
h i p We note that Assumption 3 could be replaced with the Sp version E ∥∇f (x, ξ) − ∇F (x)∥p ≤ σ p for all x ∈ B, which would allow for smaller σ values. However, we choose to state the result with the stronger Assumption 3 for better comparability with the previous results for N SSD -M. Proof. For notational conciseness denote κ := ρσS. We proceed as in the proof of Theorem C.3, but replacing Lemma C.2 by Lemma C.5, to get T X p−1 2A 2κη E [∥∇F (xt )∥∗ ] ≤ ∆1 + L0 η 2 T B + η + + 2κηT (1 − β) p 1−β 1−β t=1 X T 2A + L1 η 2 B + E [∥∇F (xt )∥∗ ] . 1 − β t=1 2A By Equation (6) L1 η 2 B + 1−β ≤ 59 η and re-arranging thus yields T
p−1 1X 9∆1 6L0 ηA 5κ E [∥∇F (xt )∥∗ ] ≤ + + + 5κ(1 − β) p . T t=1 4ηT (1 − β) T (1 − β) q ∆1 (1−β) Using η = min 1−β , we further obtain 5L1 L0 T
T
1X E [∥∇F (xt )∥∗ ] ≤ 9 T t=1
s
p−1 ∆1 L0 12∆1 L1 5κ + + + 5κ(1 − β) p . T (1 − β) T (1 − β) T (1 − β)
Now denote p
U :=
∆1 L0 κ p−1 T
p−1 ! 3p−2
p
,
∆1 L1 κ p−1 T
V :=
20
p−1 ! 2p−1
,
κ
W := T
p−1 2p−1
With the same arguments as in the proof of Theorem C.3 we derive the upper bounds s r ∆1 L0 ∆1 L0 9 ≤9 + 9U T (1 − β) T ∆1 L1 12∆1 L1 ≤ 12 + 7V T (1 − β) T 5κ ≤ 5W T (1 − β) p−1
5κ(1 − β) p ≤ 5U + 7V + 5W on the remaining terms. Combining these bounds yields r T ∆1 L0 ∆1 L1 1X E [∥∇F (xt )∥∗ ] ≤ 9 + 12 + 14U + 14V + 10W, T t=1 T T which finishes the proof. C.2.2
Application to M UON
Next we extend the previous result to M UON, which is an instance of N SSD -M with ∥·∥ = ∥·∥S∞ , i.e., mt+1 ← βt mt−1 + (1 − βt )∇f (xt , ξt ) xt+1 ← xt + ηt argmin⟨mt , d⟩,
(M UON)
∥d∥S∞
where argmin∥d∥S ⟨mt , d⟩ = U V is the polar decomposition. In order to analyse M UON, we need ∞ to upper-bound the nuclear norm ∥·∥S1 ≤ ρ ∥·∥p with a p-uniformly smooth norm ∥·∥p , to use the machinery from Section C.2.1. In particular, we want to choose a p-uniformly smooth norm that has the smallest norm equivalence constant ρ to S1 . We motivate our choice of p-uniformly smooth norm with the following lemma, which shows that the Schatten-p norm is the best possible choice in the above sense. Lemma C.7. Let ∥·∥ be a p-uniformly smooth norm on Rm×n with smoothness constant Sp (∥·∥) ≤ S. Furthermore assume the normalisation uv ⊤ = 1 for all u ∈ Rm , v ∈ Rn with ∥u∥2 = ∥v∥2 = 1. Then we have ( p−1 ) ∥A∥S1 min {m, n} p sup ≥ max 1, . S A∈Rm×n \{0} ∥A∥ In particular, the Schatten-p norm achieves this lower-bound exactly. Proof. Denote d := min {m, n} and define C :=
∥A∥S1 . A∈Rm×n \{0} ∥A∥ sup
n m n Furthermore let Ei := ei fi⊤ ∈ Rm×n , where (ei )m i=1 , (fi )i=1 are the standard bases of R and R respectively, for i = 1, . . . , d, and ε1 , . . . , εd be independent Rademacher random variables. Define k X
Mk :=
εi Ei .
i=1
Then, by definition of the smoothness constant Sp (Rm×n ) (see Definition 2.1), we have p
p
∥Mk−1 + Ek ∥ + ∥Mk−1 − Ek ∥ p p ≤ ∥Mk−1 ∥ + S p ∥Ek ∥ . 2 p Hence, by induction and ∥Ek ∥ = 1, we get E [∥Md ∥ ] ≤ dS p . On the other hand note that p
E [∥Mk ∥ | ε1 , . . . , εk−1 ] =
C≥
∥Md ∥S1 d = ∥Md ∥ ∥Md ∥ 21
and hence
dp p ≤ E [∥Md ∥ ] ≤ dS p . Cp
Rearranging yields the claim. Finally we can state the (L0 , L1 )-smooth version of Theorem 3.5. (p−1)
Corollary C.8. Suppose Assumptions 1, 3 and 4 hold and let κ := min {m, n} p . Then the iterates generated by M UON with s 1 − β ∆1 (1 − β) ηt ≡ η = min , , 5L1 L0 T ( ( )) p p p ∆1 L0 3p−2 5L1 ∆1 2p−1 − 2p−1 βt ≡ β = 1 − min 1, max , ,T κ2 σ 2 T 2κσT satisfy p−1 ! 3p−2 p ∆1 L0 ∆1 L1 min {m, n}∆1 L0 σ p−1 + 12 + 14 T T T p−1 ! (p−1) p 2p−1 min {m, n}∆1 L1 σ p−1 min {m, n} p σ + 14 . + 10 p−1 T T 2p−1
T 1X E ∥∇F (xt )∥S1 ≤ 9 T t=1
r
In particular, the sample complexity to reach an ε-∥·∥S1 -stationary point is at most p p σ 2p−1 (2p−1) ∆1 L1 ∆1 L0 σ p−1 ∆1 L0 ∆1 L1 σ p−1 p−1 p O +ι 2 +ι , + +ι ε ε2 ε ε ε ε ε where ι := min {m, n}. Proof. By Lemma A.2 we have ∥A∥S1 ≤ κ ∥A∥Sp for all A ∈ Rm×n , and the Schatten-p norm is p-uniformly smooth with smoothness constant Sp (Rm×n ) ≤ 1 (Example 2.3). Hence we may apply Theorem C.6 with (B, ∥·∥) = (Rm×n , ∥·∥S∞ ), ∥·∥p = ∥·∥Sp , ρ ← κ and S = 1. Remark C.9. Extending Remark 3.6, we note that Theorem C.6 also extends to full networks including vector valued weights by considering (Pethick et al., 2025b) the Banach space B = Rm1 ,n1 × . . . × RmD ,nD × Rd1 × . . . × RdK , ∥·∥ = max max ∥Wl ∥S∞ , max ∥vl ∥ℓ∞ . l∈[D]
l∈[K]
In this case, vector-valued weights are trained with S IGNUM, which is widely considered to be an idealised version of A DAM (Kunstner et al., 2023). The corresponding dimension dependence is given by X X min {ml , nl } + dl . l∈[D]
D
l∈[K]
Extension to Stochastic Conditional Gradient
Let D ⊂ B be convex and bounded, with diameter DD := sup ∥x − y∥ . x,y∈D
Consider min F (x), x∈D
22
and assume that the linear minimisation oracle over D is exact, i.e., each subproblem inf u∈D ⟨ϕ, u⟩ admits a minimiser. We consider the stochastic conditional gradient (SCG) iteration mt ← βt mt−1 + (1 − βt )∇f (xt , ξt ), vt ∈ argmin ⟨mt , u⟩ , (SCG) u∈D
xt+1 ← (1 − ηt )xt + ηt vt , where m0 := ∇f (x1 , ξ1 ). We measure stationarity by the Frank-Wolfe gap gD (x) := sup ⟨∇F (x), x − u⟩ . u∈D
Lemma D.1 (Constrained descent bound). Suppose Assumption 2 holds. Let µt := mt − ∇F (xt ). Then, for every t ≥ 1, 2 . ηt E [gD (xt )] ≤ E [F (xt ) − F (xt+1 )] + ηt DD E [∥µt ∥∗ ] + L2 ηt2 DD Proof. For any u ∈ D, smoothness and xt+1 = xt + ηt (vt − xt ) give F (xt+1 ) ≤ F (xt ) + ηt ⟨∇F (xt ), vt − xt ⟩ + L2 ηt2 ∥vt − xt ∥
2 2
≤ F (xt ) + ηt ⟨∇F (xt ), u − xt ⟩ + ηt ⟨µt , u − vt ⟩ + L2 ηt2 ∥vt − xt ∥ , where we used ⟨mt , vt ⟩ ≤ ⟨mt , u⟩. Since u, vt , xt ∈ D, 2 ηt ⟨∇F (xt ), xt − u⟩ ≤ F (xt ) − F (xt+1 ) + ηt DD ∥µt ∥∗ + L2 ηt2 DD . Taking expectations and the supremum over u proves the claim.
The proof relies on an error control Lemma analogue to Lemma C.2. Fortunately, the argument in Lemma C.2 only depends on the momentum recursion in the update rule, which remains unchanged for SCG. The only new ingredient in the following proof is that the iterate difference is controlled by the domain diameter rather than the unconstrained update magnitude. Lemma D.2 (Momentum error for SCG). Suppose Assumptions 2 and 3 hold, and that (B, ∥·∥) is q-uniformly convex with q-convexity constant K. Let p′ be the conjugate exponent of q and set r := min {p, p′ }. For constant parameters ηt ≡ η, and βt ≡ β ∈ [0, 1), r−1 PT 1 r + ηLDD T . t=1 E [∥mt − ∇F (xt )∥∗ ] ≤ σK 1−β + σKT (1 − β) 1−β Proof. Writing δt := ∇f (xt , ξt ) − ∇F (xt ), we have µt := mt − ∇F (xt ) = βt µt−1 + (1 − βt )δt + βt (∇F (xt−1 ) − ∇F (xt )). For constant β, unrolling gives µt =
t X
bτ,t δτ +
τ =1
t X
cτ,t (∇F (xτ −1 ) − ∇F (xτ )),
τ =2
where b1,t = β t−1 , bτ,t = (1 − β)β t−τ for τ ≥ 2, and cτ,t = β t−τ +1 . By Assumption 2, ∥∇F (xτ −1 ) − ∇F (xτ )∥∗ ≤ L ∥xτ −1 − xτ ∥ . The SCG update implies xt − xt−1 = η(vt−1 − xt−1 ) with xt−1 , vt−1 ∈ D, hence ∥xt − xt−1 ∥ ≤ ηDD . Therefore " t # t X X E [∥µt ∥∗ ] ≤ E bτ,t δτ + ηLDD cτ,t . τ =1
τ =2
∗
Since (δτ )tτ =1 is a martingale difference sequence in B ∗ , the same application of Theorem 2.4 and lemma A.1 as in Lemma C.2 yields " t # !1/r t X X E bτ,t δτ ≤ σK brτ,t . τ =1
τ =1
∗
Summing over t and using the same geometric-series estimates as in Lemma C.2 proves the claim.
23
We are now ready to state the convergence theorem for SCG. Theorem D.3 (Frank-Wolfe gap bound for SCG). Suppose Assumptions 2 and 3 hold, and that (B, ∥·∥) is q-uniformly convex with q-convexity constant K. Let p′ be the conjugate exponent of q, r := min {p, p′ }, and ∆D := F (x1 ) − inf F (x). x∈D
Consider SCG with r D (1−β) , ηt ≡ η := min 1, ∆LD 2 T D r r DL 3r−2 , T − 2r−1 βt ≡ β := 1 − min 1, max σ∆ . 2K2T Then PT 1
t=1 E [gD (xt )] ≤
T
∆D T + 3DD
q
∆D L + 3DD T
r−1 r ! 3r−2 ∆D L(σK) r−1 + 2DD T
σK r−1 . T 2r−1
Proof. Summing Lemma D.1 and telescoping gives PT PT 2 T. η t=1 E [gD (xt )] ≤ ∆D + ηDD t=1 E [∥µt ∥∗ ] + L2 η 2 DD Plugging in Lemma D.2 yields 1 T
Let η̄ :=
r−1 2 2 ηLDD LηDD ∆D σK r + t=1 E [gD (xt )] ≤ ηT + DD T (1−β) + DD σK(1 − β) 1−β + 2 .
PT
q
∆D (1−β) 1 1 2 T , so that η = min {1, η̄}. We have that η ≤ 1 + η̄ and η ≤ η̄. Consequently, LDD
∆D ∆D ηT ≤ T + DD
q
∆D L T (1−β) ,
2 ηLDD 1−β ≤ DD
q
∆D L T (1−β) ,
2 LηDD ≤ D2D 2
q
∆D L(1−β) ≤ D2D T
q
∆D L T .
Next, let U :=
r−1 r ! 3r−2 r−1 ∆D L(σK) , T
W :=
σK r−1 , 2r−1 T
u :=
∆D L σ2 K 2 T
r 3r−2 ,
so that 1 − β = min {1, max {u, w}}. Hence √1 ≤ min 11,√u ≤ 1 + √1u , 1−β { }
which implies 2DD
q
∆D L DD T (1−β) + 2
q
∆D L ≤ 3DD T
q
∆D L + 2DD U. T
Moreover, since 1 − β ≥ w and β ≤ 1, σK DD T (1−β) ≤ DD W,
and, using (1 − β)α ≤ uα + wα for α = r−1 r , r−1
DD σK(1 − β) r ≤ DD U + DD W. Combining the bounds gives 1 T
PT
∆D t=1 E [gD (xt )] ≤ T + 3DD
which is exactly the stated claim. 24
q
∆D L + 3DD U + 2DD W, T
w := T
r − 2r−1
,
E
Lower Bound Proofs
In this section we carry the missing proof of Theorem 3.7. Therefore let us recite the definitions of Section 3.3 for the reader’s convenience. Let m, n ∈ N≥2 and set κ := min {m, n} ≥ 2 n m n Further let (ei )m i=1 and (fj )j=1 are the standard bases of R and R respectively, and
Ei := ei fi⊤ ∈ Rm×n κ
For a sign vector s ∈ {±1} , define Ss :=
κ X
si Ei .
i=1
Fix δ > 0 to be specified later and q :=
2δ σ
p p−1
. Then the hard instance is given by
L δ L δ 2 ∥x∥F − ⟨Ss , x⟩F , ∇fs (x, ξ) x − BsI EI , 2κ κ κ q where B ∼ Bern(q) and I ∼ Unif([κ]) independently. Fs (x) :=
Now we first verify that Fs and ∇fs satisfy the required properties. 2
δ Lemma E.1. Assume δ ≤ σ/2. Then the function Fsis L-smooth with F (0) − inf x Fs (x) ≤ 2L and p the gradient oracle ∇fs is unbiased and satisfies E ∥∇fs (x, ξ) − ∇Fs (x)∥S1 ≤ σ p .
Proof. The smoothness is immediate by noting that ∥∇Fs (x) − ∇Fs (y)∥S1 = L κ ∥x − y∥S1 ≤ L ∥x − y∥S∞ for all x, y ∈ Rm×n . To prove the bounded initialisation gap, we complete the square to get 2 L δ2 δ Fs (x) = x − Ss − ∥Ss ∥F , 2κ L 2Lκ F 2
δ . Unbiasedness again directly follows from the independence of B and I and hence Fs (x) ≥ − 2L and the definition of q. Finally, we have " # " #! p p δ δ δ p E ∥∇f (x, ξ) − ∇F (x)∥S1 = E Ss − BsI EI ≤ 2p−1 δ p + E BsI EI κ q q S1 S1
and
" E
# p p p δ δ σp δ δp p BsI EI =q E ∥sI EI ∥S1 = q = p−1 = p . q q q q 2 S1
Hence the claim follows from our assumption δ ≤ σ2 . Additionally we need the following technical lemmata. Pκ Lemma E.2. For every A ∈ Rm×n , ∥A∥S1 ≥ i=1 |Aii |. Pκ Proof. Let ωi := sgn(Aii ) for i = 1, . . . , κ and U := i=1 ωi Ei . Then ∥U ∥S∞ ≤ 1 and, by duality Pκ of the nuclear and operator norm, i=1 |Aii | = ⟨U, A⟩F ≤ ∥U ∥S∞ ∥A∥S1 = ∥A∥S1 . κ
For a sign vector s ∈ {±1} we denote the version which flips the sign at index i as s(i) = (s1 , . . . , si−1 , −si , si+1 , . . . , sκ ). The following lemma says that, as long as (Bt , It ) = (1, i) did not yet happen, the iterates of an algorithm run on Fs and Fs(i) are identical. κ Lemma E.3. Let A be an optimization algorithm, i ∈ [κ] and s ∈ {±1} . Let B1 , . . . , BT and I1 , . . . , IT be independent and run A on Fs and Fs(i) to obtain iterates x1 , g1 , . . . , xT , gT and x′1 , g1′ , . . . , x′T , gT′ respectively. Furthermore let Ai :=
T \
{(It , Bt ) ̸= (i, 1)}.
t=1
Then we have 1Ai (xt , gt ) = 1Ai (x′t , gt′ ) for all t ∈ [T ]. 25
Proof. This follows by induction. For t = 1 the claim is trivial as x1 = x′1 . Now assume the claim holds for some t ∈ [T − 1]. Then, on the event Ai , we have ∇fs (xt , ξt ) = ∇fs(i) (x′t , ξt ) and hence xt+1 = x′t+1 . κ
Finally we need a technical lemma which roughly states that, for s ∼ Unif ({±1} ), as long as the algorithm has not seen an index i, si is still uniformly distributed conditioned on the history. Lemma E.4. Let A be an optimization algorithm with internal randomness r. Let s ∼ κ Unif ({±1} ) and denote the sequence generated by A on Fs as (x1 , g1 ), . . . , (xT , gT ). Furthermore define the observed randomness HT := σ({r, g1 , . . . , gT }) and Ai :=
T \
{(It , Bt ) ̸= (i, 1)}.
t=1
Finally assume that r, s, B1 , I1 , . . . , BT , IT are mutually independent. Then 1Ai E [si | HT ] = 0
a.s.
In particular, for any HT -measurable ν and Borel-measurable function ψ : R × {±1} → R, we have ψ(ν, 1) + ψ(ν, −1) 1Ai E [ψ(ν, si ) | HT ] = 1Ai . 2 Proof. Fix i ∈ [κ], and define the sigma-algebra containing all primitive randomness except si as Ki := σ(r, I1 , B1 , . . . , IT , BT , s1 , . . . , si−1 , si+1 , . . . , sκ ). Furthermore let M := (r, g1 , . . . , gT ) be the trajectory on Fs , and M (±1) the counterfactual trajectories on Fs(±1) , where s(±1) = (s1 , . . . , si−1 , ±1, si+1 , . . . , sκ ). Now let Z be any bounded HT -measurable random variable. Since HT = σ(M ), the Doob-Dynkin Lemma supplies a bounded Borel-measurable function h : R → R such that Z = h(M ). Hence h(M (1) ) − h(M (−1) ) E [1Ai Zsi ] = E [E [1Ai Zsi | Ki ]] = E 1Ai , 2 where we used Ai ∈ Ki , M = M (si ) and that si is independent of Ki . By Lemma E.3 we have 1Ai M (1) = 1Ai M (−1) and hence E [1Ai Zsi ] = 0. The definition of conditional expectation hence implies E [1Ai si | HT ] = 0, which finishes the first part. For the second statement note that ψ(ν, x) = α(ν) + β(ν)x, where α(ν) =
ψ(ν, 1) + ψ(ν, −1) , 2
β(ν) =
ψ(ν, 1) − ψ(ν, −1) 2
and thus 1Ai E [ψ(ν, si ) | HT ] = 1Ai α(ν) + β(ν)E [1Ai si | HT ] = 1Ai α(ν), where we used the previous result in the last step. This finishes the proof. This finally allows us to provide the lower-bound on E ∥∇Fs (xT )∥S1 . Importantly, we will choose κ s to be uniformly distributed over {±1} , which allows us to leverage the fact that some Fs are hard to differentiate as highlighted by Lemma E.3. κ Proposition E.5. Let s ∼ Unif ({±1} ). Then, for every possibly randomised algorithm A and every T ∈ N≥1 , we have q T E ∥∇Fs (xT )∥S1 ≥ δ 1 − . κ Proof. Denote the sequence generated by A on Fs as (x1 , g1 ), . . . , (xT , gT ) and define HT as in the previous lemma. For i ∈ [κ] denote νi := ⟨Ei , xT ⟩ and note that νi is HT -measurable. Moreover, we have Lνi − δsi ⟨Ei , ∇Fs (xT )⟩ = κ 26
and hence, by Lemma E.2, ∥∇Fs (xT )∥S1 ≥
κ X
|⟨Ei , ∇Fs (xT )⟩| =
i=1
κ X |Lνi − δsi |
κ
i=1
.
Taking expectation we obtain
E ∥∇Fs (xT )∥S1
X κ κ X |Lνi − δsi | |Lνi − δsi | ≥ ≥ E 1Ai E κ κ i=1 i=1
and, by Lemma E.4, |Lνi − δsi | |Lνi − δ| + |Lνi + δ| |Lνi − δsi | E 1Ai HT ≥ E 1Ai = E 1Ai E κ κ 2κ i +δ| ≥ κδ , and E [1Ai ] = Pr(Ai ) = (1 − κq )T to close the proof. Finally, note that |Lνi −δ|+|Lν 2κ
Finally we provide the proof of Theorem 3.7. Proof of Theorem 3.7. Let δ := 2ε which satisfies δ ≤ σ2 by our assumption on ε. By Lemma E.1, the function Fs satisfies 2ε2 δ2 = ≤ ∆1 F (0) − inf F s (x) ≤ 2L L x∈Rm×n κ by our assumption on ε. Hence Fs ∈ F∆1 ,L and ∇fs ∈ Op,σ . Furthermore, for s ∼ Unif ({±1} ), Proposition E.5 implies q T E ∥∇Fs (x̂T )∥S1 ≥ δ 1 − κ T and hence, for the output to satisfy E ∥∇Fs (x̂T )∥S1 ≤ ε, we need 2ε(1 − q/κ) ≤ ε. This is equivalent to q T q 1 1 log(2) 1− ≤ log ≤ ⇔ T log 1 − ⇔T ≥ κ 2 κ 2 − log 1 − κq Since κ ≥ 2 we have q/κ ≤ 1/2 and thus we can use − log(1 − x) ≤ 2x log(2), for x ∈ [0, 1/2], to get q q − log 1 − ≤ 2 log(2) . κ κ Combining yields p log (2)κ κ σ p−1 , T ≥ = · 2 log (2)q 2 4ε finishing the proof. E.1
High-Dimensional Lower-Bound
In this section we provide a first-order lower-bound for ∥·∥S1 -stationarity of the form p ∆1 L σ p−1 Ω min {m, n} 2 . ε ε As vector valued functions are a special case, it is well-known such lower-bounds require a highdimensional construction (Arjevani et al., 2023). Let us first recite the construction for the Euclidean case on Rd from J. Zhang et al. (2020a) and Arjevani et al. (2023). For simplicity we will only consider the case of zero-respecting algorithms, the extension to randomised algorithms follows from the same techniques as in (Arjevani et al., 2023). Therefore we first introduce the Euclidean Rd -version of the function class as n o (2) F∆1 ,L (d) := F : Rd → R ∥∇F (x) − ∇F (y)∥2 ≤ L ∥x − y∥2 , F (0) − inf F (x) ≤ ∆1 x
d
and the Euclidean R -version of the oracle class as ( (2) Op,σ (d) :=
∇f : Rd × Ξ → Rd
E [∇f (x, ξ)] = ∇F (x), p
E [∥∇f (x, ξ) − ∇F (x)∥2 ] ≤ σ p , ∀x ∈ Rd 27
) .
Lemma E.6 (J. Zhang et al., 2020a;Arjevani et al., 2023). Let c := 1/300 and p ∈ (1, 2]. For all √ ∆1 , L, σ > 0 and 0 < ε ≤ c min σ ∆1 L there exists a dimension d ≤ ∆1 Lε−2 , a function (2) (2) F ∈ F∆1 ,L (d) and a gradient oracle ∇f ∈ Op,σ (d) such that every possibly randomised zerorespecting algorithm A satisfies p ∆1 L cσ p−1 . E [∥∇F (xt )∥2 ] > ε for all t ≤ c2 2 ε ε In particular, ∇f (x, ξ) has the structure
Z − 1 h(x), θ p ε p−1 . where ∥h(x)∥2 ≤ 46ε, and ξ ∼ Bernoulli(θ) with θ = cσ ∇f (x, ξ) = ∇F (x) +
For (Rm×n , ∥·∥S∞ ), the corresponding classes are defined as ) ( F (0) − inf F (X) ≤ ∆1 , m×n X , F∆1 ,L (m, n) := F : R →R ∥∇F (X) − ∇F (Y )∥S1 ≤ L ∥X − Y ∥S∞ ) ( E [∇f (X, ξ)] = ∇F (X), . Op,σ (m, n) := ∇f : Rm×n × Ξ → Rm×n p E ∥∇f (X, ξ) − ∇F (X)∥S1 ≤ σ p Next we define the quantities which will be used to embed the hard instance from Rd into Rm×n . Therefore let d ∈ N≥1 , κ ∈ N≥1 , and define m := κd and n := κ. Now, for i ∈ [κ] and v ∈ Rd , let X Ri := {(i − 1)d + 1, . . . , id}, Ei (v) := vj−(i−1)d · ej fi⊤ ∈ Rm×n . j∈Ri
That is, Ei (v) places v in the rows Ri of its i-th column and is zero everywhere else. Furthermore let Pi be the inverse operation, i.e., X Pi (A) := ⟨ej fi⊤ , A⟩ej−(i−1)d ∈ Rd j∈Ri
extracts the rows Ri of the i-th column of A. Then the following properties hold. Lemma E.7. Let v1 , . . . , vκ ∈ Rd , and A ∈ Rm×n . Then κ X
= max ∥vi ∥2 ,
Ei (vi )
i=1
i∈[κ]
Ei (vi )
=
i=1
S∞
Proof. The columns of
κ X
S1
X
∥vi ∥2 ,
and
∥Pi (A)∥2 ≤ ∥A∥S∞ .
i∈[κ]
Pκ
κ X
i=1 Ei (vi ) have disjoint row support, so
!⊤ Ei (vi )
i=1
κ X
! Ei (vi )
2 2 = Diag ∥v1 ∥2 , . . . , ∥vκ ∥2
i=1
and the singular values are hence ∥v1 ∥2 , . . . , ∥vκ ∥2 , which proves the first two identities. For the last claim note that ∥Pi (A)∥2 ≤ ∥Afi ∥2 ≤ ∥A∥S∞ . (2)
This allows us to embed the hard instance from Rd into Rm×n as follows. Let F (2) ∈ F (2)
∆1 ,L(2)
(d)
(2) and ∇f (2) ∈ Op,σ(2) (d) be the hard instance from the Euclidean case. Then we define
F (X) :=
κ X
F (2) (Pi (X)),
∇f (X, ξ) :=
i=1
κ X
Ei ∇f (2) (Pi (X), ξi ) ,
(8)
i=1
i.i.d.
where ξi ∼ Bernoulli(θ) are independent. Next we prove that this embedding — after appropriate rescaling — preserves the required properties. 28
(2) σ (2) := ε (2) := Proposition E.8. Let ∆1 = ∆κ1 , L(2) := L and F, ∇f be defined κ,ε κ, σ κ1/p √ as above. Furthermore let c := 1/300 and assume 0 < ε ≤ c min σ, ∆1 L . Then F ∈ F∆1 ,L (m, n) and ∇f ∈ Op,σ (m, n).
(2)
≤ c min σ
Proof. It is straight-forward to check ε
(2)
,
q
(2) ∆1 L(2)
, which we require to apply
Lemma E.6. Hence the Euclidean hard instance exists and we calculate κ X (2) F (0) − inf F (X) = F (2) (0) − inf F (2) (x) ≤ κ∆1 = ∆1 , X
and, for X, Y ∈ R
m×n
x
i=1
, κ X
∥∇F (X) − ∇F (Y )∥S1 =
Ei ∇F (2) (Pi (X)) − ∇F (2) (Pi (Y ))
i=1
S1
κ Lemma E.7 X ≤ ∇F (2) (Pi (X)) − ∇F (2) (Pi (Y ))
≤ L(2)
i=1 κ X
2
∥Pi (X) − Pi (Y )∥2
i=1 Lemma E.7
≤
L(2) κ ∥X − Y ∥S∞
= L ∥X − Y ∥S∞ . This shows F ∈ F∆1 ,L (m, n). For the oracle properties, note that κ κ h i X X Ei E ∇f (2) (Pi (X), ξi ) = Ei ∇F (2) (Pi (X)) = ∇F (X). E [∇f (X, ξ)] = i=1
i=1
To prove the bounded p-th central moment, we make use of the explicit structure of ∇f (X, ξ) to get κ p X p Ei ∇F (2) (Pi (X)) − ∇f (2) (Pi (X), ξi ) E ∥∇F (X) − ∇f (X, ξ)∥S1 = E i=1
= E
κ X
S1
Ei
i=1
p ξi . h(Pi (X)) 1− θ S1
By Lemma E.7 we have κ κ κ X X X ξi ξi ξi = 1− Ei h(Pi (X)) ∥h(Pi (X))∥2 ≤ 46ε(2) 1− . 1− θ θ θ i=1 i=1 i=1 S1 Pκ Now let N := i=1 ξi ∼ Binomial(κ, θ), and p (2) p−1 p ε p−1 ε q := κθ = κ = ≤1 cσ cσ (2) by our assumption ε ≤ cσ. Next, Pκ κ ξi 1X ξi N 1− ≤ 1 + i=1 = 1 + κ i=1 θ κθ q and hence p E ∥∇F (X) − ∇f (X, ξ)∥S1
(9)
!p # p " X κ ξi ≤ E 1− θ i=1 p (12) N p ≤ (46ε) E 1 + q p N p ≤ 2(46ε) 1 + E . qp
(9),(10)
29
46ε κ
(10)
(11)
(12)
Finally we calculate E [N p ] ≤ E N 2 = κθ(1 − θ) + (κθ)2 = q(1 − θ) + q 2 ≤ 2q, which yields p N p p ≤ 2(46ε) 1 + 2q 1−p ≤ 6(46c) σ p ≤ σ p , 2(46ε) 1 + E p q p
and hence the claim. Now we are ready to prove the first-order lower-bound for ∥·∥S1 -stationarity. Theorem E.9 (First-Order Lower-Bound for ∥·∥S1 -Stationarity). Let c := 1/300, κ ∈ N≥1 and √ p ∈ (1, 2]. For all ∆1 , L, σ > 0 and 0 < ε ≤ c min σ, ∆1 L there exists m ≤ κ ∆ε12L , a function F ∈ F∆1 ,L (m, κ) and a gradient oracle ∇f ∈ Op,σ (m, κ) such that every possibly randomised zero-respecting algorithm A satisfies p ∆1 L cσ p−1 . E ∥∇F (xt )∥S1 > ε for all t ≤ c2 · min {m, n} 2 ε ε In particular, there does not exist a dimension-free first-order guarantee in ∥·∥S1 -stationarity, and M UON is first-order optimal. Proof. Let d ≤ ∆1 Lε−2 be the dimension from Lemma E.6 and set m := κd. Furthermore let F, ∇f be the hard instance defined in Equation (8). By Proposition E.8 we have F ∈ F∆1 ,L (m, κ) and ∇f ∈ Op,σ (m, κ). Now let A be a zero-respecting algorithm and denote the sequence generated by A on F, ∇f as (i) (Xt )t∈N≥1 . Let i ∈ [κ] and note that, by the independence of ξ1 , . . . , ξκ , the sequence xt := Pi (Xt ) ∈ Rd corresponds to a randomised zero-respecting algorithm on F (2) , ∇f (2) . In particular, we have p (2) (2) h i cσ (2) p−1 (i)) (2) (2) 2 ∆1 L E ∇F xt >ε for all t ≤ c · . 2 (ε(2) )2 ε(2) (2)
(2) := ε (2) := σ this is equivalent to with our choices ∆1 = ∆κ1 , L(2) := L κ,ε κ, σ κ1/p p h i ε ∆1 L cσ p−1 (i)) > E ∇F (2) xt for all t ≤ c2 · κ 2 . κ ε ε 2 Finally, by Lemma E.7 we have κ κ h X X (i)) (i)) = Ei ∇F (2) xt E ∇F (2) xt E ∥∇F (Xt )∥S1 = E i=1
S1
i=1
i
> ε,
2
which finishes the proof. Theorem E.9 is a direct corollary of the previous result, by noting that the case m < n can be handled by considering the transposed version of the above proof, and unnecessary dimensions can be padded with zeros.
F
Experimental Details
In this section we provide further details and extensions of our experiments in Section 4. Further Details. We follow the usual convention of training all 2-d parameters but the input- and output-embeddings with M UON, while all other parameters are trained with A DAM W. Computational Resources. All experiments were conducted on an internal cluster, using 2 H100 GPUs. Overall the experiments required 606 GPU hours, out of which 311 were spent on preliminary experiments, 283 on the Schatten-r experiments, and an additional 12 GPU hours to calculate the gradient noise. 30
0.5
0.5 0.4
0.3
S1
S1
0.4
Weight Family att.qkv att.out mlp.up mlp.down embd
Weight Family att.qkv att.out mlp.up mlp.down embd
0.3
0.2
0.2
0.1
0.1
0.0
0.0
(a) σS1 at Init
50
1
40
60
Weight Family att.qkv att.out mlp.up mlp.down embd
50 40 1
60
(b) σS1 at Final
30
Weight Family att.qkv att.out mlp.up mlp.down embd
30
20
20
10
10
0
0
(c) σℓ1 at Init
(d) σℓ1 at Final
Figure 3: Plain noise values. Licenses. The FineWeb-Edu dataset (Penedo et al., 2024) is distributed under the Open Data Commons Attribution License (ODC-By) v1.0. Our use complies with this license and adheres to CommonCrawl’s Terms of Use. Details on Figure 1 and 2. We consider the checkpoint at initialisation and at the end of training for the optimal learning rate η = 0.015, and seed 0. For both checkpoints, we first calculate the true gradient ∇F (xi ), i ∈ {1, T } by going over the whole 1.4B token training dataset {ξ1 , . . . , ξT }. In a second pass we then calculate, for each weight matrix x(j) T
σSp 1 ,p =
1 X p ∥∂ (j) F (xi ) − ∂x(j) f (xi , ξτ )∥S1 , T τ =1 x
p σF,p =
1 X p ∥∂ (j) F (xi ) − ∂x(j) f (xi , ξτ )∥F , T τ =1 x
T
where ξ1 , . . . , ξT are the mini-batches of size 512. We then calculate the ratio σS1 ,p/σF,p for all weight matrices trained by M UON and report them in Figure 1. We additionally visualise the plain σS1 ,p and σℓ1 ,p in Figure 3. Details on Table 1. To prevent confounding from the choice of learning rates, we additionally provide the full grid of results for the learning rate and Schatten-r sweep in Figure 4. Importantly, none of the optimal learning rates lay on the end points of the grid.
31
90.52 ±0.36
88.89 ±0.48
91.25 ±3.28
92.04 ±0.64
99.68 ±7.93
109.79 ±0.45
126.17 ±2.51
100+
8/7
88.06 ±0.22
86.50 ±0.91
88.33 ±3.05
88.42 ±0.90
93.66 ±8.18
113.76 ±8.27
116.63 ±3.03
90
4/3
79.56 ±0.31
75.41 ±0.30
73.56 ±3.89
71.45 ±2.32
67.61 ±5.98
83.61 ±6.94
92.45 ±1.94
1.5
73.19 ±0.38
67.87 ±1.09
65.20 ±0.91
60.57 ±4.48
62.59 ±6.89
72.64 ±12.17
171.33 ±128.16
2
47.79 ±0.58
43.52 ±0.21
41.45 ±0.12
40.63 ±0.21
40.18 ±0.09
40.35 ±0.37
42.47 ±0.28
70
3
37.93 ±0.22
36.09 ±0.24
35.31 ±0.14
34.79 ±0.09
34.58 ±0.06
34.51 ±0.13
35.55 ±0.10
60
4
35.54 ±0.16
34.36 ±0.14
33.97 ±0.05
33.74 ±0.04
33.60 ±0.04
33.54 ±0.04
34.32 ±0.04
33.83 ±0.21
33.36 ±0.04
33.18 ±0.04
33.12 ±0.03
33.02 ±0.01
33.09 ±0.07
34.94 ±0.04
50
8
33.60 ±0.07
33.45 ±0.09
33.38 ±0.10
33.35 ±0.04
33.34 ±0.02
82.36 ±4.11
168.03 ±6.81
5
0.01
0.03
0.06
0.005
0.007
5 0.012
Learning rate
0.015
80
Mean final perplexity
Schatten r
1
40
Figure 4: Full result grid for the Schatten-r sweep. The value represents the mean final perplexity ± its standard deviation over the seeds 0, 1, 2. Optimal learning rates per geometry are marked with black borders, the overall best result with an orange border.
32