Florian Hübler* Department of Computer Science ETH Zurich, Switzerland [email protected]
Kai Lion* Department of Computer Science ETH Zurich, Switzerland [email protected]
Antonio Orvieto ELLIS Institute Tübingen, MPI-IS Tübingen AI Center, Germany [email protected]
Niao He Department of Computer Science ETH Zurich, Switzerland [email protected]
Abstract Matrix-aware optimizers such as Muon and Muown have recently shown strong empirical performance for pre-training Transformers. In particular, Muown separates each weight matrix into row magnitudes and an un-normalized direction variable, updating the former with Adam and the latter with Muon. We show that the directional update of Muown is equivalent to a Riemannian step on the normalized directions, while the magnitude of the un-normalized parameterization only modulates the angular step size. This explains the step-size stability of Muown and suggests making the angular step size explicit. The resulting method, AngularMuown, optimizes directly over the normalized directions and uses a schedulable angular multiplier decoupled from the radial magnitude update. AngularMuown improves over Muown and, at the time of writing, a preliminary version is leading the per-optimizer category of the modded nanoGPT speedrunning competition. Further experiments on Qwen2-0.5B, and 1.1B parameter mixture-of-experts models confirm the algorithm scales beyond small models. An implementation of the algorithm is available at https://github.com/fhueb/angular-muown.
1
Introduction
For over a decade, Adam (Kingma and Ba, AdamW 3.3 2015) has been the default optimizer for deep NorMuon Muown learning. Recently, a family of matrix-aware op3.2 AngularMuown timizers that exploit the two-dimensional struc3.1 ture of weight matrices has gained traction, particularly for Transformers whose parameters are 3.0 dominated by dense matrices. Muon, one of the simplest and most effective variants, performs 2.9 2.5B 5B 7.5B 10B normalized steepest-descent updates in the specTokens tral norm (Jordan et al., 2024). This realizes Figure 1: Validation loss of a 124M Transformer the per-step size prescribed by the spectral concheckpoint-branching pre-trained on FineWebdition for stable feature learning (Yang et al., Edu. AngularMuown substantially improves 2023), but it only bounds the change ∥∆W∥S∞ over tuned Muown, NorMuon, and AdamW ∗ across the token budget. Equal contribution.
Validation Loss
arXiv:2606.23637v1 [cs.LG] 22 Jun 2026
Muown Implicitly Performs Angular Step-size Decay
1
and leaves the growth of ∥W∥S∞ along the trajectory uncontrolled. Muown (Lion et al., 2026) addresses this shortcoming by re-parameterizing each weight into its row-magnitudes and directions, i.e., W = Diag(g/ ∥R∥row )R, where ∥A∥row ∈ Rm denotes the vector of Euclidean row norms of A ∈ Rm×n . This corresponds to an optimizer-internal weight normalization (Salimans and Kingma, 2016), which separates a per-row magnitude g from a direction R. The magnitudes are trained with Adam and the direction with Muon. Beyond improving performance, this simple change makes training remarkably robust to large stepsizes that would otherwise cause Muon to diverge. We argue that this robustness is not incidental but geometric. The directional update on R, conducted in the ambient space Rm×n , can equivalently be written as a Riemannian update on the row-oblique manifold M = {A ∈ Rm×n | ∥A∥row = 1}. In this view, the directional update step becomes a per-row angular displacement whose size scales inversely with the row norm. As the row norms grow during training, this angular step is silently annealed, which explains Muown’s stepsize stability. It is also exactly the role weight decay plays implicitly elsewhere, modulating the angular step size by controlling the weight norm (Van Laarhoven, 2017; Kosson et al., 2024). Making this implicit structure explicit yields AngularMuown. Instead of updat- Algorithm 1 Simplified AngularMuown ing the direction R in the ambient space, Require: Gradient G = ∇W L(W), momentum buffer we explicitly keep it normalized and perM, learning rate η, momentum β1 , β2 , shape scale form Riemannian Muon updates on U = sm,n , angular step multiplier κt Diag(∥R∥row )−1 R. In particular, the directional stepsize now governs the angle of the W = Diag(g)U ▷ Reparameterize internally ⊤ directional update and can be scheduled exgradg ← diag GU plicitly. The method thus preserves the specgradU ← Diag(g) G − Diag gradg U tral update that makes Muon effective while M ← β1 M + gradU ▷ Update U with Muon exposing the angular step size as an explicit design choice. A preliminary version of AnO ← Orth(gradU +β1M) gularMuown performed strongly in the perU ← RowNormalize U − η κt sm,n O optimizer category of the modded nanoGPT g ← Adam(g, η, β1 , β2 ) ▷ Update g with Adam speedrun, suggesting that the Riemannian interpretation is not only explanatory but also W ← Diag(g) U ▷ Re-compose practically useful.1 Prior public disclosure. The preliminary version of AngularMuown was publicly disclosed in our May 8, 2026 modded-nanoGPT speedrun submission. That submission used Muown’s internal decomposition into a per-row gain g and stored direction v, and introduced a scheduled storeddirection norm, explicitly described there as angular step-size attenuation: increasing ∥v∥ while recomposing through v/∥v∥ reduces the effective angular displacement on the unit sphere. The present work formalizes this mechanism as implicit Riemannian angular step-size decay and replaces the stored-norm gauge schedule by an explicit angular multiplier κt . Contributions. In this work we identify a key implicit property of Muown which potentially explains its success, make it explicit and provide strong empirical support for the resulting algorithm AngularMuown. 1
As of June 16th 2026, this preliminary version is the leading per-optimizer result, see Muown-RowNormControl in Figure 2: https://github.com/KellerJordan/modded-nanogpt/tree/ 23f758f123df4b8dd5b5fe64c4c9070f6ef33b52/records/track_3_optimization#notable-results-history.
2
• We show that Muown implicitly performs Riemannian updates on the direction, and that its stepsize stability arises from an implicit annealing of the angular step size by row-norm growth of the direction. Building on this view, we introduce AngularMuown, which turns the angular step size into an explicit, schedulable quantity that is independent of the radial update (Sections 2 and 3). • We conduct extensive experiments on transformer-based architectures that show a consistent improvement of AngularMuown compared to AdamW, (Nor-)Muon, and Muown. In particular, we observe a significant speed-up over all baselines across dense 124M and 500M, as well as 1.1B mixture-of-experts models (Section 4). • We provide theoretical support for the directional update. In particular, in the non-convex setting, we show that an ε-stationary point is reached in at most O min {m, n}∆1 Lσ 2 ε−4 gradient evaluations under reasonable assumptions (Section 3.5).
2
Motivation
For normalized or scale-invariant weights, the direction alone de∆R = ηO termines the feature the row (or neuron) encodes. Given this obRlarge servation, the meaningful effective step size is the per-step angular θlarge change of that row. While the row-norm does not impact the feature representation, it does impact optimization dynamics by modulating how strongly a fixed Euclidean update rotates that direction. A Rsmall θsmall large norm makes the same update induce a smaller angular change, and a small norm induces a larger one. Thus, the norm acts as an implicit controller of the angular step size. This perspective explains part of the effect of weight decay in normalized networks. Rather Figure 2: Impact of weight than acting only as classical regularization, weight decay counteracts norm on angular update size. norm growth and thereby prevents angular updates from collapsing, yielding an equilibrium of angular updates (Van Laarhoven, 2017; Hoffer et al., 2018; Kosson et al., 2024). Muown inherits this mechanism in a more direct form. Its internal parameterization W = Diag(g/ ∥R∥row )R = Diag(g)U, with U = Diag(∥R∥row )−1 R, represents the row magnitude as g whereas ∥R∥row is a gauge variable. Changing it does not change W, but it does change the optimization dynamics. Since Muown updates R in the ambient space, the same additive update to R induces a smaller rotation when ∥R∥row is large, as illustrated in Figure 2. Mathematically, the internal directional update Rt+1 ← e t+1 ← Ut − η Diag(∥Rt ∥ )−1 O, making this dependence Rt − ηO can equivalently be written as U row explicit. This provides a simple explanation for Muown’s stepsize stability: growth of the unnormalized direction norm silently anneals the angular update. More importantly, it reveals that Muown’s parameterization already exposes the angular step size as the natural quantity to control, opening up a new opportunity: rather than letting the angular schedule emerge implicitly from the growth of the row norm, we can schedule it explicitly. AngularMuown realizes this by optimizing directly over the row-normalized directions and replacing Muown’s implicit norm-induced annealing with an explicit angular multiplier, making weight decay redundant in the process. 3
3
AngularMuown
In this section, we discuss the motivation and different parts of AngularMuown in more detail. We first make the geometry that is implicit in Muown explicit: its directional update is a Riemannian steepest-descent step on the oblique manifold (Sections 3.1 and 3.2), which turns the Euclidean Muon step into a per-row angular displacement modulated by the row scale (Section 3.3). Promoting this modulation to an explicit schedule yields AngularMuown (Section 3.4).
3.1
The Oblique-Manifold Parameterization
Muown reparameterizes each weight as W = Diag(g/ ∥R∥row ) R (Salimans and Kingma, 2016), separating a per-row scale g ∈ Rm from a direction matrix R ∈ Rm×n . Since only the row-normalized Diag(1/ ∥R∥row )R enters the weight, the directional degrees of freedom are exactly the matrices with unit rows, i.e. the oblique manifold Mob := {U ∈ Rm×n : ∥U∥row = 1m } = (Sn−1 )m , the product of m unit spheres embedded in Rm×n (Absil et al., 2008; Boumal, 2023). We therefore write the parameterization directly on the manifold, U ∈ Mob ,
W(g, U) = Diag(g) U,
g ∈ Rm ,
(1)
which matches Muown under U = Diag(1/ ∥R∥row )R and identifies g with the row magnitudes ∥W∥row . The tangent space at U collects matrices with rows orthogonal to those of U, TU Mob = {V : diag VU⊤ = 0}, where diag(·) extracts the diagonal. Under the Frobenius metric, the orthogonal projection onto TU Mob removes the per-row radial component, and we return to Mob by the row-normalization retraction: ProjU (A) := A − Diag diag AU⊤ U, RetrU (V) := Diag(1/ ∥U + V∥row ) (U + V). (2) The projection subtracts ⟨Ai , Ui ⟩ Ui from each row Ai , and the retraction is exactly the row renormalization Muown performs when recomputing the effective weight (Absil et al., 2008).
3.2
Muown Implicitly Performs Riemannian Optimization
Differentiating through (1) splits the gradient G = ∇W L(W) into a radial and a directional part. The chain rule gives the radial gradient ∇g L = diag GU⊤ ∈ Rm , with (∇g L)i = ⟨Gi , Ui ⟩, and the Euclidean directional gradient Diag(g)G. Projecting the latter onto TU Mob yields the Riemannian gradient of U 7→ L(Diag(g)U), gradU L = ProjU (Diag(g)G) = Diag(g) ProjU (G) = Diag(g)(G − Diag(∇g L) U),
(3)
the second equality holding because Diag(g) rescales rows and the projection (2) acts row-wise. These are exactly the decoupled gradients Muown forms internally: the radial part drives the row scales, the tangent part the directions. Muown updates the directions by feeding (3) through the Muon pipeline (momentum, Nesterov, spectral orthogonalization) and renormalizing. Spectral orthogonalization O = argmin∥O∥S ≤1 ⟨V, O⟩, ∞ computed by Newton–Schulz, returns the normalized spectral steepest-descent direction (Boyd and Vandenberghe, 2004; Bernstein and Newhouse, 2024). Composing it with the retraction (2) shows the Muown direction update is a retracted, spectrally-conditioned Riemannian (quasi-)gradient step on Mob , Ut+1 = RetrUt (ηt sm,n Ot ),
Ot = argmin ⟨gradUt L + βMt , O⟩, ∥O∥S∞ ≤1
4
Table 1: Three regimes for the per-row angular step size θt,i , all sharing the Euclidean tangential with O⊥ step at,i := ηt sm,n O⊥ t,i t,i = ProjUt,i (Ot,i ). Weight decay and Muown both modulate 2
the angle through a norm νt,i in the denominator, and differ only in that norm’s dynamics: decay pins it at a rotational equilibrium (Van Laarhoven, 2017; Kosson et al., 2024), whereas Muown’s unregularized stored norm grows, annealing the angle as a side effect. AngularMuown fixes νt,i ≡ 1 and moves the schedule into the numerator, decoupling it from any norm. Regime
tan θt,i
Governing norm νt,i
Late-training angle
Weight decay (Muon, λ > 0)
at,i /νt,i
plateau, θt,i → θ∞ > 0
Implicit (Muown)
at,i /νt,i
∥Wt,i ∥2 → ν∞ (equilibrium) √ rt,i ∼ t (unregularized growth)
Explicit (AngularMuown)
κt at,i
νt,i ≡ 1 (decoupled)
scheduled, θt,i ∼ κt
anneals, θt,i ∼ t−1/2
with the radius g optimized in parallel by Adam, matching the ℓ∞ geometry singled out for the row scales in (Lion et al., 2026), and sm,n the shape scale of Section 3.4. Two details depart from a textbook Riemannian method: the momentum Mt is accumulated in the ambient space rather than parallel-transported, and Ot is retracted directly without re-projecting onto TUt Mob , with the retraction absorbing any residual radial component. Muown, though derived from a purely row-magnitude argument, thus already is a Riemannian optimizer on Mob .
3.3
From Euclidean Steps to Angular Step Sizes
Section 2 singled out the per-row angle as the effective step size of a direction update. The obliquemanifold view makes this angle intrinsic: each update in (3) rotates a unit row Ui along its sphere, so its proper step size is the swept angle, not the Euclidean displacement ∥ξi ∥2 . The following elementary fact converts the Euclidean step the optimizer computes into that angle. Proposition 3.1 (Angular step size). Let Ui ∈ Sn−1 be a unit row and ξi ∈ Rn a proposed step, split into radial and tangential parts ξi = ⟨ξi , Ui ⟩ Ui + ξi⊥ with ξi⊥ = ProjUi (ξi ). The angle θi between Ui and the retracted row RetrUi (ξi ) satisfies tan θi =
ξi⊥ 2 . 1 + ⟨ξi , Ui ⟩
(4)
Only the tangential part turns the row—the radial part merely stretches it and is erased by the renormalization—so for a purely tangential step (⟨ξi , Ui ⟩ = 0) the rotation reduces to tan θi = ∥ξi ∥2 . Proof. The retraction rescales Ui + ξi to unit length without changing its direction, so θi is the angle between Ui and Ui + ξi . Since ∥Ui ∥2 = 1, this vector splits orthogonally as Ui + ξi = (1 + ⟨ξi , Ui ⟩) Ui + ξi⊥ : its component along Ui (length 1 + ⟨ξi , Ui ⟩) and its perpendicular component (length ξi⊥ 2 ) are the adjacent and opposite sides of the right triangle subtending θi , and their ratio is tan θi . The proposition lets us read off the per-step rotation in each parameterization. For the small steps taken in practice the radial denominator 1 + ⟨ξi , Ui ⟩ ≈ 1, so tan θi ≈ ξi⊥ 2 and only the tangential part contributes, mirroring the projection in (3). AngularMuown and Muown use the same Euclidean update −ηt sm,n Ot,i , but feed a different step ξi into the retraction. AngularMuown retracts the unit row directly, so ξi = −ηt κt sm,n Ot,i . Muown keeps the unnormalized direction and 5
adds the update to it, Rt+1,i = Rt,i − ηt sm,n Ot,i , with unit row Ut,i = Rt,i /rt,i and rt,i = ∥Rt,i ∥2 . Factoring out the current norm, η sm,n Rt+1,i = rt,i Ut,i − trt,i Ot,i , and as normalization discards the positive prefactor rt,i , the next unit row Ut+1,i = Rt+1,i / ∥Rt+1,i ∥2 is exactly RetrUt,i (ξi ) with the rescaled step ξi = −(ηt sm,n /rt,i ) Ot,i : the larger the stored norm, the smaller the step the unit row actually sees. Substituting each ξi into Proposition 3.1 gives tan θt,i = ηt κt sm,n O⊥ t,i | {z }2
versus
AngularMuown
tan θt,i = ηt sm,n O⊥ rt,i . t,i 2 | {z }
(5)
Muown
where O⊥ t,i = ProjUt,i (Ot,i ). AngularMuown’s rotation is fixed by the multiplier κt and is independent of the row scale, whereas Muown’s carries the extra factor 1/rt,i inherited from the growing stored norm. As the row norms drift upward under Muon, the mechanism Muown exposes, the diagonal multiplier 1/rt,i acts as a per-row angular schedule that decays as a side effect of the row-scale dynamics. AngularMuown replaces this implicit, scale-coupled decay with an explicit, independently chosen schedule.
3.4
The AngularMuown Update
3.4.1
Angular Learning-Rate Multiplier
Section 3.3 suggests an obvious degree of freedom: rather than letting the angular schedule emerge as the side effect 1/rt,i of row-norm growth, we prescribe it. We keep U ∈ Mob at unit row norm and scale the directional step by an angular learning-rate multiplier κt ∈ (0, 1], recovering Muown’s implicit decay when κt ∝ 1/rt,i while opening the design space to schedules chosen on their own merit. Motivated by theory, we use an inverse-polynomial, κpoly = (1 + c (t − tw )+ )−p , t held at 1 for t ≤ tw warm-up steps. Here c > 0 and p > 0 are hyperparameters setting the decay rate and shape, but we observe that c = 0.001, p = 1 work uniformly. In particular, κpoly decays without t reference to a fixed endpoint and is therefore compatible with warmup-stable-decay, horizonless training (Hu et al., 2024). Crucially κt acts only on the directional step, while the row scales g keep the base rate ηt through Adam, so the radial and angular dynamics are scheduled independently. The shape scale sm,n calibrates each directional step on the manifold. 3.4.2
Shape Scaling for Dimension Independence
The remaining factor is the shape scale sm,n ,p which calibrates each directional step on the manifold. We use the spectral-condition choice sm,n = max(1, m/n) by default such that the angular update is dimension-independent. By (5), the per-row angle obeys tan θt,i = ηt κt sm,n O⊥ t,i , so the 2 choice of sm,n sets how a fixed multiplier κt translates into an actual rotation across layers of different shapes. The dimension dependence enters through O⊥ t,i 2 Ot has unit singular values, = min(m, n). p such that ∥Ot ∥F p
2
and the orthogonalized update
Consequently, the energy of a typical max(1, m/n). The spectral-condition choice sm,n =
row is roughly ∥Ot,i ∥2 ≈ min(m, n)/m = 1/ p max(1, m/n) (Yang et al., 2023) is exactly the reciprocal of this factor, so sm,n O⊥ t,i 6
2
= Θ(1)
3.5
[deg]
Validation Loss
3.4 3.3 Muon ( = 0) Muon ( = 0.1) NorMuon Muown AngularMuown ( t 1t ) AngularMuown ( t 1)
3.2 3.1
0.5B
1.0B
1.5B
Tokens
2.0B
2.5B
(a) Loss curves.
2.00 1.75 1.50 1.25 1.00 0.75 0.50 0.25 0.00
Muon ( = 0) Muon ( = 0.1) NorMuon Muown AngularMuown ( t 1t ) AngularMuown ( t 1)
0.5B
1.0B
1.5B
Tokens
2.0B
2.5B
(b) Angular displacement.
Figure 3: Analysis on the 124M model on 2.5B tokens. Weight decay is denoted by λ, AngularMuown’s angular multiplier as κt . The reported runs correspond to the optimally tuned learning rate. (a) Validation loss curves across optimizers. (b) Mean angular displacement θ. regardless of m and n, rendering the angular update dimension-independent: a single schedule κt then induces the same angular step size on every row regardless of its shape. By contrast, the Adamp √ RMS-matching choice sm,n ∝ max(m, n) (Liu et al., 2025) leaves a residual sm,n O⊥ = Θ( n) t,i 2 that grows with the fan-in, so the same κt would produce systematically larger rotations on wider rows. In particular this choice of shape scaling gives our angular stepsize ηt κt a physical meaning: by Proposition 3.1 with ⟨ξi , Ui ⟩ ≈ 0 the direction is updated by tan θi = ηt κt sm,n O⊥ ≈ ηt κt . t,i 2
Since, for small stepsize, we have arctan(x) ≈ x + O(x3 ) this corresponds to a directional update of ◦ ◦ ≈ 180 π ηt κt ≈ 57ηt κt degrees. Algorithm 1 collects these ingredients. What AngularMuown adds over Muown is control. The angular and radial geometries that Muown entangles through the single row norm rt,i become explicit and independently schedulable, which we show in Section 4 improves perplexity. Moreover, our shape scaling improves hyperparameter transferability over Muown’s by removing implicit shape dependence.
3.5
Convergence Guarantee
Finally we provide a convergence guarantee for the directional update of an idealized version of AngularMuown. The idealized version keeps the orthogonalized update in the tangent space, see Algorithm 2. In the following we denote the Riemannian gradient of L : M → R at U ∈ M as grad L(U). The formal assumptions can be found in Appendix B. Theorem 3.2 (Convergence Guarantee). Let Assumptions 1, 2, and 3 hold, and denote r := min {m, n}. Then Algorithm 2 with parameters ( ( )) r r ∆1 (1 − β) ∆ L 2 1 η= , β = 1 − min 1, max T − /3 , LT rσ 2 T satisfies T
1X E [∥grad L(Ut )∥F ] ≤ 4 T t=1
r
1/4 √ ∆1 L r∆1 Lσ 2 rσ + 4 1/3 . +7 T T T
The proof closely follows Muown’s proof, but requires some care to handle the retraction and can be found in Appendix B. The result implies a leading sample complexity of O min {m, n}∆1 Lσ 2 ε−4 . 7
κt ∝ 1/t
κt ≡ 1
4e-3 2.862 6e-3 2.792 8e-3 2.749 1e-2 2.721 2e-2 2.662 4e-2 2.648
2.700 2.694 2.698 2.708 2.767 2.873
67M 124M 300M best LR
3.3
NorMuon 2.713 2.701 2.695 2.692 2.718 2.749
3.2
Validation Loss
η
AngularMuown
3.1 3.0 2.9 28
(a) Qwen2-0.5B learning-rate sweep.
27
26
25
Learning Rate
24
23
(b) Width transfer.
Figure 4: Validation loss across learning rates. (a) Qwen2-0.5B trained by AngularMuown with polynomial angular decay (κt ∝ 1/t), without a decay (κt ≡ 1), and NorMuon. (b) Width transfer for AngularMuown at model widths 512, 768, and 1280. Surprisingly and contrary to previous observations, we observe that AngularMuown actually requires the polynomially decaying stepsize suggested by theory. The proof also works for general retractions and exponential maps, however we stick to our specific choices for ease of exposition.
4
Experiments
In this section, we evaluate AngularMuown on language-model pre-training with modern architectures on FineWeb-Edu. We conduct experiments to validate the empirical performance of AngularMuown across different transformer architectures and sizes. Furthermore we conduct an angular displacement comparison between algorithms to highlight its efficiency, and examine the aglorithms learning rate transfer. Setup. All experiments are conducted on nodes with 4 GH200 GPUs with a total of approximately 10’000 GPU hours, including preliminary experiments. Unless specified otherwise, we use PolarExpress (Amsel et al., 2025) for orthogonalization, the WSD schedule (Hu et al., 2024) with 100 warmup steps, a cooldown ratio of 0.1, sequence length 1024, and batch size 512. PlainLM. We adopt the architecture from Ajroldi (2024), which is based on a nanoGPT (Karpathy, 2022) implementation modified to include recent architectural improvements such as RoPE (Su et al., 2024), RMSNorm normalization (Zhang and Sennrich, 2019), and SwiGLU activations (Shazeer, 2020). To compare performance, we sweep the learning rate of AdamW, NorMuon, and Muown across [5e-4, 1e-3, 2e-3, 4e-3, 8e-3]. AngularMuown was swept across [5e-3, 1e-2, 2e-2, 4e-2, 8e-2]. Other parameters are set to their usual values and can be found in Table 3. We start decay branches after 2.5B, 5B, 7.5B, and 10B tokens and the results can be found in Figure 1. AngularMuown outperforms all baselines for each horizon and achieves a speedup of approximately 2× over AdamW, and 1.5× over NorMuon and Muown. Next we empirically examine the angular displacement discussed in Section 3.3. Therefore we train the 124M model for Chinchilla-optimal 2.5B tokens (Hoffmann et al., 2022) with Muon with and without weight decay, NorMuon, Muown, as well as AngularMuown with (κt ∝ 1/t) and without (κt ≡ 1) angular decay. We perform a learning-rate sweep over [1e-3, 2e-3, 4e-3,
8
6e-3, 8e-3, 1e-2] and report the optimal run for each optimizer. Figure 3b makes AngularMuown’s angular mechanism visible: with a constant step without angular decay the per-step rotation stays roughly flat, whereas every decaying-angle method—including AngularMuown— anneals θ smoothly towards zero. Figure 3b also corroborates the three regimes mentioned in Table 1. First, adding weight decay to Muon (λ = 0.1) lifts its angular displacement to a sustained, nonvanishing plateau instead of letting it decay. Second, Muown results in an annealed angular update, while AngularMuown allows for arbitrary annealing. At a well-tuned learning rate this annealing lets AngularMuown take large, aggressive directional steps and reach the lowest perplexity of all methods (Figure 3a). Finally we study the learning rate transfer across widths. Following Pethick et al. (2025), we train our 124M base architecture for 5100 iterations with model widths 512, 768, and 1280 (i.e., total parameter count of 67M, 124M, and 300M). For each width, we run AngularMuown with learning rates η = 2−8 , . . . , 2−3 and report the final validation loss in Figure 4b. We note that the optimal learning rate perfectly transfers across widths, validating the effectiveness of the scale factor discussed in Section 3.4.2. Qwen2-0.5B. The same pattern transfers to a Qwen2-0.5B architecture (Figure 4a). At its welltuned learning rate, the polynomial schedule attains the best loss overall and improves over NorMuon for every η ≥ 2×10−2 , whereas the unscheduled variant degrades sharply as η grows. Making the angular step size explicit thus shifts the optimum towards larger, more aggressive steps, consistent with the angular-annealing view of Section 2. Mixture-of-Experts. Moving beyond dense transformer models, we extend our experimental study to Mixture-of-Experts (MoE) models (Shazeer et al., 2017) using the widely adopted Megatron-LM codebase (Shoeybi et al., 2019). We consider a 1.1B DeepSeek-V3-style MoE (Liu et al., 2024) with 370M active parameters (1.1B-A370M), following the architectural choices of the Moonlight recipe for Muon (Liu et al., 2025). The improvement of AngularMuown over Muon persists in this sparse model setup. For AngularMuown, some of the best points lie at the edge of the tested grid, so the reported MoE losses may be conservative. We kept the learning-rate grids comparable across algorithms to avoid giving either method a larger tuning budget.
5
η 3e-4 6e-4 1e-3 2e-3 4e-3 8e-3
AngularMuown Muon p p max(m, n) max(1, m/n) 2.6124 2.5278 2.4887 2.4807 — —
— — 2.6952 2.5901 2.5198 2.4914
2.5297 2.5255 2.5263 2.5353 — —
Table 2: Learning rate sweep on a DeepSeekV3-style MoE-model (Liu et al., 2025) with 1.1B-A370M parameters. For AngularMuown, we ablate the spectral condition p scaling sm,n = max(1, m/n) p and the RMSmatching scaling sm,n = max(m, n).
Related Work
Matrix-aware optimizers. Early methods exploiting the matrix structure of feedforward layers include Shampoo (Gupta et al., 2018) and Soap (Vyas et al., 2025). Another notable work is Muon (Jordan et al., 2024), which performs spectral-norm steepest descent via Newton-Schulz orthogonalization. Descendants of Muon include NorMuon (Li et al., 2025), which applies neuronwise adaptive scaling of the update after orthogonalization, and Muown (Lion et al., 2026) which splits each weight into row magnitudes and a direction component internally. While Muown updates
9
the direction in the ambient space, leaving the row-scales to drift freely, our parameterization fixes them to unit norm. Weight normalization and reparameterization. Salimans and Kingma (2016) introduce the weight-norm parameterization to decouple magnitude learning from the direction, making the direction component of the weights scale-invariant in the sense that any positive rescaling of the direction component does not alter the underlying function. Hoffer et al. (2018) present a modification of the parameterization to fix the learnable magnitude to a constant value. Scale invariance, angular step size, and weight decay. For scale-invariant neural networks, the effective step size is the angular rate of change, which is governed implicitly through the weight norm (Van Laarhoven, 2017; Hoffer et al., 2018). Kosson et al. (2024) formalize this as a rotational equilibrium in which weight decay balances norm growth so the angular update reaches a stable value. They propose a rotational wrapper whose main goal is to provide an update whose average angular update matches the optimizer’s predicted equilibrium rotation. In their analysis, this equilibrium rotation is induced by the choice of the learning rate in conjunction with the weight decay strength. These works focus on the implicitly induced angular step size behavior. While Van Laarhoven (2017) focuses on the implicit weight norm dynamics, Kosson et al. (2024) attempt to mimic the induced rotational equilibrium throughout training by enforcing the predicted average angular rotation directly. In contrast, AngularMuown makes the angular step size an explicit and independently scheduled quantity, decoupled from the radial magnitude update and its schedule, which renders weight decay redundant on the directions. Concurrent work. Very recently we became aware of the concurrent blog post (Hägele et al., 2026). Similar to this work, Hägele et al. (2026) keep the directional factor of a magnitude-direction parameterization on a fixed norm. The main overlap with AngularMuown is therefore not the magnitude-direction decomposition itself, which is already used by Muown (Lion et al., 2026), but the fixed-norm treatment of the direction and the resulting direct control of angular updates. The methods differ in their geometry: AngularMuown uses the row-wise Riemannian gradient gradU L = Diag(g) G − Diag(diag(GU ⊤ ))U , and introduces an explicit angular multiplier κt . By contrast, Hägele et al. (2026) use the ambient chain-rule gradient for the on-sphere direction, GW c = c c Diag(γrow ) G Diag(γcol ) for W = Diag(γrow )W Diag(γcol ), and project the direction W back to the chosen sphere after its update. Their reported experiments mainly use endpoint-dependent linear decay of the learning rates, whereas AngularMuown is designed as a drop-in replacement for the standard WSD schedule, with the angular schedule controlled separately through κt .
6
Conclusion and Limitations
We observe that Muown’s stored row norm implicitly anneals the angular step size. Based on this observation we propose AngularMuown, which makes this mechanism explicit by optimizing directly on the row-oblique manifold and scheduling the angular multiplier explicitly. Across the studied language-model pre-training settings, this explicit control yields considerable improvements over AdamW, NorMuon, and Muown. The main limitation is that our evaluation is centered on language-model pre-training and limited to medium model sizes. A broader study of architectures, training regimes, and scaling to 100B and larger parameter models remains open.
10
Acknowledgments This work was supported under project ID a0184 as part of the Swiss AI Initiative, through a small grant from the ETH Domain and computational resources provided by the Swiss National Supercomputing Centre (CSCS) under the Alps infrastructure. Kai Lion is supported by Swiss National Science Foundation (SNSF) Sinergia Funding No. 216600. Florian Hübler acknowledges financial support from the ETH research grant and Swiss National Science Foundation (SNSF) Project Funding No. 200021-207343. Antonio Orvieto acknowledges the financial support of the Hector Foundation. Niao He is supported by an ETH research grant funded through the ETH Zurich Foundation and by an SNSF Starting Grant.
References Boyd, Stephen and Lieven Vandenberghe (2004). Convex Optimization. Cambridge: Cambridge University Press. Absil, P.-A., R. Mahony, and R. Sepulchre (2008). Optimization Algorithms on Matrix Manifolds. Princeton, NJ: Princeton University Press. Kingma, Diederik P. and Jimmy Ba (2015). “Adam: A Method for Stochastic Optimization”. In: Proc. Int. Conf. on Learning Representations (ICLR). Salimans, Tim and Durk P Kingma (2016). “Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks”. In: Proc. Neural Information Processing Systems (NeurIPS). Shazeer, Noam, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean (2017). “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-ofExperts Layer”. In: Proc. Int. Conf. on Learning Representations (ICLR). Van Laarhoven, Twan (2017). “L2 regularization versus batch and weight normalization”. In: arXiv preprint arXiv:1706.05350. Gupta, Vineet, Tomer Koren, and Yoram Singer (Oct. 2018). “Shampoo: Preconditioned Stochastic Tensor Optimization”. In: Proceedings of the 35th International Conference on Machine Learning. Ed. by Jennifer Dy and Andreas Krause. Vol. 80. Proceedings of Machine Learning Research. PMLR, pp. 1842–1850. url: https://proceedings.mlr.press/v80/gupta18a.html. Hoffer, Elad, Ron Banner, Itay Golan, and Daniel Soudry (2018). “Norm matters: efficient and accurate normalization schemes in deep networks”. In: Advances in Neural Information Processing Systems. Ed. by S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett. Vol. 31. Curran Associates, Inc. url: https : / / proceedings . neurips . cc / paper _ files/paper/2018/file/a0160709701140704575d499c997b6ca-Paper.pdf. Shoeybi, Mohammad, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro (2019). “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism”. In: arXiv preprint arXiv:1909.08053. Zhang, Biao and Rico Sennrich (2019). “Root mean square layer normalization”. In: Proc. Neural Information Processing Systems (NeurIPS). 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. Shazeer, Noam (2020). “GLU Variants Improve Transformer”. In: arXiv preprint arXiv:2002.05202. Hoffmann, Jordan, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, 11
Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack W. Rae, and Laurent Sifre (2022). “Training compute-optimal large language models”. In: Proc. Neural Information Processing Systems (NeurIPS). Karpathy, Andrej (Dec. 2022). nanoGPT: The simplest, fastest repository for training/finetuning medium-sized GPTs. GitHub repository. README notes an update/deprecation (Nov 2025). Accessed 2026-02-06. url: https://github.com/karpathy/nanoGPT. Boumal, Nicolas (2023). An Introduction to Optimization on Smooth Manifolds. Cambridge University Press. Yang, Greg, James B Simon, and Jeremy Bernstein (2023). “A Spectral Condition for Feature Learning”. In: arXiv preprint arXiv:2310.17813. Ajroldi, Niccolò (2024). plainLM: Language Model Pretraining in PyTorch. https://github.com/ Niccolo-Ajroldi/plainLM. Bernstein, Jeremy and Laker Newhouse (2024). “Old Optimizer, New Norm: An Anthology”. In: arXiv preprint arXiv:2409.20325. Hu, Shengding, Yuge Tu, Xu Han, Ganqu Cui, Chaoqun He, Weilin Zhao, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Xinrong Zhang, Zhen Leng Thai, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, dahai li, Zhiyuan Liu, and Maosong Sun (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. Jordan, Keller, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein (2024). Muon: An optimizer for hidden layers in neural networks. url: https : / / kellerjordan.github.io/posts/muon/. Kosson, Atli, Bettina Messmer, and Martin Jaggi (21–27 Jul 2024). “Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks”. In: Proceedings of the 41st International Conference on Machine Learning. Ed. by Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp. Vol. 235. Proceedings of Machine Learning Research. PMLR, pp. 25333–25369. url: https://proceedings. mlr.press/v235/kosson24a.html. Liu, Aixin, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. (2024). “DeepSeek-V3 Technical Report”. In: arXiv preprint arXiv:2412.19437. Penedo, Guilherme, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, 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. Su, Jianlin, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu (2024). “RoFormer: Enhanced transformer with Rotary Position Embedding”. In: Neurocomputing 568, p. 127063. issn: 0925-2312. doi: 10.1016/j.neucom.2023.127063. url: https://www.sciencedirect. com/science/article/pii/S0925231223011864. Yang, An, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jianxin Yang, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Xuejing Liu, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yuqiong Liu, Zeyu Cui, Zhenru Zhang,
12
Zhifang Guo, and Zhihao Fan (2024). Qwen2 Technical Report. arXiv: 2407 . 10671 [cs.CL]. url: https://arxiv.org/abs/2407.10671. Amsel, Noah, David Persson, Christopher Musco, and Robert M Gower (2025). “The polar express: Optimal matrix sign methods and their application to the muon algorithm”. In: arXiv preprint arXiv:2505.16932. Li, Zichong, Liming Liu, Chen Liang, Weizhu Chen, and Tuo Zhao (2025). “NorMuon: Making Muon more efficient and scalable”. In: arXiv preprint arXiv:2510.05491. Liu, Jingyuan, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. (2025). “Muon is scalable for llm training”. In: arXiv preprint arXiv:2502.16982. Pethick, Thomas, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher (13–19 Jul 2025). “Training Deep Learning Models with Norm-Constrained LMOs”. In: Proceedings of the 42nd International Conference on Machine Learning. Ed. by Aarti Singh, Maryam Fazel, Daniel Hsu, Simon Lacoste-Julien, Felix Berkenkamp, Tegan Maharaj, Kiri Wagstaff, and Jerry Zhu. Vol. 267. Proceedings of Machine Learning Research. PMLR, pp. 49069–49104. url: https://proceedings.mlr.press/v267/pethick25a.html. Vyas, Nikhil, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade (2025). “SOAP: Improving and Stabilizing Shampoo using Adam for Language Modeling”. In: International Conference on Learning Representations. Ed. by Y. Yue, A. Garg, N. Peng, F. Sha, and R. Yu. Vol. 2025, pp. 93423–93444. url: https://proceedings.iclr.cc/ paper_files/paper/2025/file/e988664070e9591f93fdcf605f7dc623- Paper- Conference. pdf. Hägele, Alexander, Atli Kosson, Alejandro Hernández-Cano, and Martin Jaggi (2026). Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors. url: https://haeggee.github.io/posts/magnitude-direction-decoupling. Lion, Kai, Florian Hübler, Bingcong Li, Antonio Orvieto, and Niao He (2026). “Muown: Row-Norm Control for Muon Optimization”. In: arXiv preprint arXiv:2605.10797.
13
Table 3: Optimizer hyper-parameters, underlines values denote the tuned learning rates. Optimizer
Learning rate η
β1
β2
λ
AdamW 5e-4, 1e-3, 2e-3, 4e-3, 8e-3 0.9 0.95 0.1 NorMuon 5e-4, 1e-3, 2e-3, 4e-3, 8e-3 0.95 0.95 — — Muown 5e-4, 1e-3, 2e-3, 4e-3, 8e-3 0.95 — AngularMuown 5e-3, 1e-2, 2e-2, 4e-2, 8e-2 0.95 — —
A
Orth. decay_scale steps — 5 5 5
— — — 0.001
Experimental Details
Common setup. We adopt the experimental setup of Ajroldi (2024), which is based on a nanoGPT (Karpathy, 2022) implementation augmented with recent architectural improvements such as RoPE (Su et al., 2024), RMSNorm normalization (Zhang and Sennrich, 2019), and SwiGLU activations (Shazeer, 2020). All models are pre-trained on FineWeb-Edu (Penedo et al., 2024) with a warmupstable-decay learning-rate schedule (Hu et al., 2024), and we report perplexity (equivalently validation loss) at a fixed token budget. Following Muown, the matrix-aware optimizer acts on all 2D hidden weight matrices (every 2D weight except the token embeddings and the LM head); the remaining parameters are handled by AdamW. For AngularMuown the row scales g are updated with Adam under the base learning rate ηt , while the directional update orthogonalizes the projected momentum with Newton–Schulz and rescales it by the angular learning-rate multiplier κt (the radial and angular dynamics are therefore scheduled independently, cf. Section 3). 124M checkpoint-branching (Figure 1). The hyperparameters used by the different optimizers can be found in Table 3. 124M learning-rate sweep (Figure 3). We train a 124M-parameter model for 2.5B tokens and sweep a log-spaced learning-rate grid (Figure 3a), comparing Muon, NorMuon (Li et al., 2025), Muown, and AngularMuown both with and without its angular schedule. Muon uses weight decay λ = 0.1 and all other optimizers use λ = 0, the settings found optimal in the Muown study. AngularMuown’s angular schedule is the inverse-polynomial κpoly = (1 + c (t − tw )+ )−p , using the t −3 default values c = 10 and p = 1. The angular-displacement curves in Figure 3b each correspond to the best-tuned learning rate per optimizer (lowest final validation loss). The per-step row rotation θ = arccos⟨u, u′ ⟩ is read directly from the stored weight rows before and after each step. Qwen2-0.5B sweep (Table 4a). To probe transfer beyond the architecture above, we additionally train a Qwen2-0.5B (Yang et al., 2024) model without weight decay (λ = 0). We sweep η ∈ {4, 6, 8} × 10−3 ∪ {1, 2, 4} × 10−2 and compare AngularMuown with its polynomial angular schedule against the constant-angular variant (κt ≡ 1) and NorMuon. MoE Experiments (Table 2). The model under consideration is a downscaled variant of Moonshot’s Moonlight-16B-A3B (Liu et al., 2025) reduced to 14 transformer layers (1 dense + 13 MoE), hidden size 1024, and 32 experts (top-6, sigmoid router with learned expert bias) plus one shared expert, for 1.1B total and 370M active parameters. We keep Moonlight’s DeepSeek-V3-style MoE design but replace MLA with standard grouped-query attention (16 heads, 4 KV groups) to avoid latent-projection special cases in the optimizer comparison. Training runs in patched Megatron-LM (Shoeybi et al., 2019) on FineWeb-Edu (sample-100BT, GPT-NeoX tokenizer, sequence length 4096) 14
Algorithm 2 Idealized Directional AngularMuown Update Require: U1 ∈ M, M0 ← grad L(U1 , ξ1 ), stepsize η, momentum β Mt ← βMt−1 + (1 − β)gradL(Ut , ξt ) Ot ← argminD∈TUt M,∥D∥S
∞
≤1 ⟨Mt , D⟩
▷ Tangent-space aware orthogonalization
Ut+1 ← RetrUt (ηOt )
▷ Retraction
at a global batch of 256 (≈1.05M tokens/step) for 30B tokens in total. For the optimizer settings, we mostly follow Liu et al., p 2025: Muon and AngularMuown are set with momentum 0.95 with Nesterov, shape scaling of max(m, n), 5 Newton–Schulz steps, and weight decay 0.1 for Muon. Moreover, routers, embeddings, norms, and biases are handled by Adam p in all cases. For AngularMuown, we additionally consider the shape scaling factor of sm,n = max(1, m/n), which results in a dimension-independent rotation.
B
Proof of Theorem 3.2
In this section we provide the missing proof for Section 3.5. The proof follows the classical arguments from (Cutkosky and Mehta, 2020). Notation. Let m, n ∈ N≥1 and for matrices A ∈ Rm×n let Ai ∈ Rn denote the i-th row of A. Let M = {A ∈ Rm×n | ∥A∥row = 1} denote the row-oblique manifold, TU M := A ∈ Rm×n ⟨Ai , Ui ⟩ = 0 for all i ∈ [m] the tangent space at U ∈ M, and ProjU (A) = A − Diag diag AU⊤ U the orthogonal projector onto TU M. Furthermore we denote the tangent-space restricted dual norm as ∥A∥U,∗ :=
max
⟨A, V⟩,
V∈TU M,∥V∥S∞ ≤1
and, for U ∈ M, A ∈ TU M, use the retraction RetrU (A) = ΠM (U + A). Now we first introduce the required assumptions. Assumption 1. The objective function is lower bounded with L(U1 ) − inf U∈M L(U) ≤ ∆1 . Assumption 2. For all U ∈ M and A ∈ TU M with V := RetrU (A) we have LL ∥A∥2S∞ , 2 ∥grad L(U) − grad L(V)∥S1 ≤ Lg ∥A∥S∞ ,
L(V) ≤ L(U) + ⟨grad L(U), A⟩ +
and denote L := max {LL , Lg }. Assumption 3. The gradient oracle is unbiased and has finite variance, i.e., for all U ∈ M, h i E [grad L(U, ξ)] = grad L(U) and E ∥grad L(U, ξ) − grad L(U)∥2F ≤ σ 2 . Additionally we require the following standard inequalities. Lemma B.1. For all U ∈ M and A ∈ TU M we have p ∥A∥F ≤ ∥A∥U,∗ ≤ ∥A∥S1 ≤ min {m, n} ∥A∥F . 15
Now we are ready to provide the descent lemma. Lemma B.2 (Descent Lemma). Let Assumptions 1 and 2 hold. Then the iterates generated by Algorithm 2 satisfy T X
T
∥grad L(Ut )∥Ut ,∗ ≤
t=1
X ∆1 ηLT + +2 ∥Mt − grad L(Ut )∥Ut ,∗ η 2 t=1
Proof. The proof arguments follow (Cutkosky and Mehta, 2020), we include it for completeness. By Assumption 2 we have L(Ut+1 ) − L(Ut ) ≤ η⟨grad L(Ut ), Ot ⟩ +
η 2 LL η2L ∥O∥2S∞ ≤ η⟨Mt , Ot ⟩ − η⟨Et , Ot ⟩ + , 2 2
where Et := Mt − grad L(Ut ). The Hölder-inequality and ∥Mt ∥Ut ,∗ ≥ ∥grad L(Ut )∥Ut ,∗ − ∥Et ∥Ut ,∗ further yield ⟨Mt , Ot ⟩ − ⟨Et , Ot ⟩ ≤ − ∥Mt ∥Ut ,∗ + ∥Et ∥Ut ,∗ ≤ − ∥grad L(Ut )∥Ut ,∗ + 2 ∥Et ∥Ut ,∗ . Summing up and using Assumption 1 yields the claim. Next we derive the noise bound, which follows the same arguments as Muown’s (Lion et al., 2026). We include it for completeness. Lemma B.3 (Deviation Bound). Let Assumptions 2 and 3 hold, and denote r := min {m, n}. Then the iterates generated by Algorithm 2 satisfy √ T X p √ rσ ηLT + rσT 1 − β + . E ∥Et ∥S1 ≤ 1−β 1−β t=1
Proof. For notational conciseness, define Gt := grad L(Ut , ξt ),
Zt := Gt − grad L(Ut ),
St := grad L(Ut−1 ) − grad L(Ut ).
Then we can unroll Mt := βMt−1 + (1 − β)Gt into Et = β t−1 E1 + (1 − β)
t X
β t−τ (Gτ − grad L(Uτ )) +
τ =2
t X
β t−τ +1 Sτ .
τ =2
Next, Assumption 2 yields ∥St ∥S1 ≤ ηLg ∥Ot−1 ∥S∞ ≤ ηL and thus t X τ =2
β t−τ +1 Sτ
≤
t X
β t−τ +1 ∥Sτ ∥S1 ≤
τ =2
S1
ηL . 1−β
Next note that Gτ − grad L(Uτ ) is a martingale difference sequence and thus (Cutkosky and Mehta, 2020) implies # " t t X X √ (1 − β)β t−τ (Gτ − grad L(Uτ )) ≤ rE (1 − β)β t−τ (Gτ − grad L(Uτ )) E τ =2
τ =2
S1
F
v u t uX √ ≤ r(1 − β)t β 2(t−τ ) σ 2 ≤
√
16
τ =1
p rσ 1 − β.
Combining the above with our initialization M1 = G1 yields √ T X p √ rσ ηLT E ∥Et ∥S1 ≤ + rσT 1 − β + 1−β 1−β t=1
and thus the claim. Finally we provide the proof for Theorem 3.2. Proof of Theorem 3.2. The proof follows the same steps as Muown’s proof (Lion et al., 2026, Appendix B.2). By Lemma B.2 we have T T h i ∆ h i X X ηLT 1 E ∥grad L(Ut )∥Ut ,∗ ≤ + +2 E ∥Mt − grad L(Ut )∥Ut ,∗ η 2 t=1 t=1 p √ ηLT ∆1 1 +3 + 2σ r +T 1−β , ≤ η 1−β 1−β
where we used Lemma B.3 in the second step. Dividing by T and our choice of stepsize implies s √ T i √ p 2σ r 1X h ∆1 L E ∥grad L(Ut )∥Ut ,∗ ≤ 4 + + 2σ r 1 − β. T T (1 − β) T (1 − β) t=1
Our choice of momentum guarantees s r 1/4 ∆1 L ∆1 L r∆1 Lσ 2 ≤ + T (1 − β) T T 1/4 √ √ p rσ r∆1 Lσ 2 σ r 1 − β ≤ 1/3 + T T √ √ σ r rσ ≤ 1/3 , T (1 − β) T and thus
r 1/4 √ T i 1X h ∆1 L r∆1 Lσ 2 rσ E ∥grad L(Ut )∥Ut ,∗ ≤ 4 +7 + 4 1/3 . T T T T t=1
This proves a stronger statement, but for ease of presentation we use Lemma B.1 to convert ∥·∥Ut ,∗ to the Frobenius norm, which finishes the proof.
17