MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control

MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control

Athanasios Glentis∗

Dawei Li∗

Mingyi Hong

arXiv:2610.37647v1 [cs.LG] 27 Oct 2026

University of Minnesota {glent007,li004678,mhong}@umn.edu

Abstract Weight decay and explicit weight-norm constraints improve optimizer performance in language model pretraining, yet a unified understanding of why such control helps and how much is needed remains incomplete. We analyze norm control in Muon over the P entire training trajectory using the accumulated effective learning rate, CT = t<T ηt /∥Wt ∥F , where ηt is the nominal learning rate and ηt /∥Wt ∥F is the effective learning rate relative to the weight norm. For bounded update directions, accumulation over any training interval bounds relative weight displacement. We show that vanilla Muon can accumulate only logarithmically, potentially limiting exploration and helping explain its slower empirical progress, whereas MuonW and MuonH guarantee at least linear accumulation under the stated conditions. Going beyond bounded-norm control, we show that sufficiently slow norm growth guarantees faster accumulation than the logarithmic worst case, opening up less restrictive optimizer designs, with more relaxed weight-norm control. Building on this connection between norm growth and exploration, we propose MuonGS, which uses Gram–Schmidt (GS) projection to remove the radial component of each Muon update without fixing the weight norm. With bounded update directions and a constant nominal √ learning rate, we prove that the weight norm grows at most√ proportionally to t after t steps. Consequently, MuonGS guarantees CT = Ω( T ), with accumulation sustained even in later training. Further, for smooth scale-invariant objectives we establish its O(T −1/2 ) stationarity guarantee. In pretraining experiments with 160M- and 400M-parameter dense models as well as a 135M active- and 1B total- parameter MoE model, MuonGS achieves lower validation perplexity than MuonH at long training horizons. In particular, for the dense model it achieves an estimated token-equivalent speedup up to about 15% at matched validation loss to MuonH.

1

Introduction

Optimizer design plays an important role in large language model (LLM) pretraining, shaping learning dynamics and model quality [Liu et al., 2025, Wen et al., 2026]. Understanding the principles behind successful optimizers is therefore essential for improving training. Weight-norm control is a common component of effective optimizers for LLM pretraining. AdamW and the weight-decay variant of Muon, MuonW, counter norm growth through decoupled shrinkage [Loshchilov and Hutter, 2017, Liu et al., 2025]. Its strength must be tuned, and the resulting weight scale depends on the optimization dynamics. Explicit constraints provide more direct control, as in fixed-radius optimization [Kodryan et al., 2022], unit-norm weight vectors in ∗ Equal contribution, alphabetical order.

Preprint.

nGPT [Loshchilov et al., 2024], and fixed weight and update norms in MuonH, the Hyperball variant of Muon [Wen et al., 2026]. These constraints prescribe weight scales throughout training and leave less freedom for their evolution. These different norm controls raise a common question: When is norm control needed, why does it help, and how much is sufficient? Existing analyses explain important effects of weight norms on optimization [Zhang et al., 2018, D’Angelo et al., 2024, Liu et al., 2026], but a unified basis for choosing among these controls remains incomplete. In particular, are the restrictions imposed by existing norm-control methods essential, or can weaker control deliver better optimization performance? A common rationale for controlling weight norms is to prevent updates from becoming too small relative to the current weights. Following prior work [D’Angelo et al., 2024, Liu et al., 2026], the effective learning rate is the ratio ηt /∥Wt ∥F of the nominal learning rate ηt to the Frobenius norm of the current weight matrix Wt . It determines the step size relative to the weight norm, up to the norm of the update direction. For a Muon update Wt+1 = Wt − ηt Ot , ∥Wt+1 − Wt ∥F ηt = ∥Ot ∥F . ∥Wt ∥F ∥Wt ∥F Keeping the weight norm below R ensures ηt /∥Wt ∥F ≥ ηt /R. When update norms are fixed, this protects relative movement in each step from shrinking due to norm growth. But must norm control preserve relative movement at every step, or is it sufficient to preserve exploration capacity over the training trajectory as a whole? To examine this question, we begin with the measure of net displacement ∥WT − W0 ∥F . Clearly, the same displacement produces a smaller relative change in the parameters when their norm is larger. P We therefore normalize by the average weight norm r̄T = T −1 t<T ∥Wt ∥F , which summarizes the typical scale along the trajectory. This ratio measures net movement relative to that scale and is unchanged by a common rescaling of all iterates. For nonzero iterates, a constant nominal learning rate η, and ∥Ot ∥F ≤ M , X ∥WT − W0 ∥F M ηT η ≤ ≤M . (1) r̄T r̄T ∥Wt ∥F t<T

The first bound sums update lengths; the second uses the fact that the average inverse norm is at least the inverse average norm. Clearly, sufficient accumulation of effective learning rates is therefore a necessary condition for substantial movement. Our contribution. We develop a global perspective on norm control for Muon-type algorithms by examining exploration over entire training trajectories. This perspective identifies less restrictive normcontrol conditions and guides the design of MuonGS, a new optimizer with theoretical guarantees and strong empirical performance. Central to our analysis is the accumulated effective learning rate, defined for general learning-rate schedules by CT =

T −1 X

ηt . ∥Wt ∥F t=0

Our main claim is that continued exploration requires sufficient growth of CT over the training intervals of interest. Further, we extend the displacement argument to arbitrary intervals and learningrate schedules, that for bounded update directions, net weight displacement from iteration s to T , relative to the larger endpoint norm, is at most M (CT − Cs ). A large CT may come mostly from early iterations. Continued exploration after iteration s requires sufficient additional accumulation, CT − Cs , and if this increment is small, relative weight displacement over [s, T ] must also be small. We next compare how different norm-control mechanisms affect this accumulation. When the learning rate ηt = η is held constant throughout training, vanilla Muon can exhibit linear norm growth and only logarithmic accumulation, CT = Θ(log T ). MuonW guarantees CT = Ω(T ) under the stated conditions, while fixed-norm control gives CT = Θ(T ). These results show that weight decay and fixed-norm constraints guarantee at least linear accumulation, compared with vanilla Muon’s logarithmic worst case. A direct consequence of our analysis is that bounded weight norms are sufficient for these linear accumulation guarantees, but stronger than needed to improve on vanilla Muon’s logarithmic worst case. We show that this requirement can be relaxed to sufficiently slow 2

radial

radial

Muon

Muon

radial part

MuonW MuonGS ∥W ∥F = rt

radial part

−ηOt

−ηλWt

−ηλWt

MuonGS

tangent space at Wt

et −η O

MuonH

MuonW Wt

−ηOt et −η O

tangent space at Wt

MuonH

∥W ∥F = rt

Wt 0

0

(a) small ∥Wt ∥F

(b) large ∥Wt ∥F

Figure 1: One step of vanilla Muon, MuonW, MuonH and MuonGS from the same Wt and the same direction Ot , for a small (left) and a large (right) weight norm. Vanilla Muon’s radial part (dashed) grows the norm at first order; MuonH retracts to ∥W ∥F = R (dotted); MuonGS keeps the tangential part only, so the norm changes at second order. MuonW’s decay −ηλWt scales with ∥Wt ∥F but the actual step depends on the size of Wt ; the decay therefore hardly slows down the norm growth for a small norm but can shrink the norm for a large one.

norm growth, which still guarantees faster accumulation (as compared to vanilla Muon) even as weight norms increase and effective learning rates decrease. This observation opens up a broader family of Muon variants that control norm growth without requiring bounded weight norms. Can this additional flexibility be used to design optimizers that sustain accumulation throughout training and achieve better practical performance? To investigate this question, we propose MuonGS, which uses Gram–Schmidt (GS) projection to remove the radial component of Muon’s update. The resulting update is orthogonal to the current weights in the Frobenius inner product. This orthogonality eliminates the first-order increase in the weight norm, so the norm grows only at second order in directions √ the step size. With bounded √ and a constant nominal learning rate, this guarantees O( t) norm growth and CT = Ω( T ), with the same accumulation order over later training windows. MuonGS thus improves vanilla Muon’s worst-case accumulation guarantee without fixing a radius. For smooth scale-invariant objectives, we further established the convergence rate for MuonGS. We evaluate MuonGS on FineWeb pretraining with modernized GPT-style models of approximately 160M and 400M parameters. At longer training horizons, MuonGS achieves lower validation perplexity than MuonH. Fitting MuonH’s loss as a function of the training-token budget gives an estimated token-equivalent speedup of up to about 15% at matched validation loss. In addition, MuonGS obtains lower validation perplexity than MuonH when pretraining a modern MoE model, having about 135M active and 1B total parameters, on Ultra-FineWeb. These results support MuonGS as a practical optimizer and demonstrate the value of relaxing norm-control requirements. Together with our theoretical guarantees, these findings provide a positive answer to the question above by showing that weaker norm control can indeed sustain accumulation throughout training while improving empirical performance. Our contributions are threefold. • We establish a global perspective on norm control in Muon, in the sense that we assess exploration over entire training trajectories through the growth of CT and its increments over training intervals. Our analysis bounds relative weight displacement and, under scale-invariance and Lipschitz assumptions, attainable loss change. At a constant nominal learning rate, it distinguishes vanilla Muon’s logarithmic worst case from the linear guarantees of MuonW and fixed-norm control. It also explains when norm growth can offset a larger nominal learning rate. 3

√ √ • We√propose MuonGS and prove O( t) weight-norm growth, CT = Ω( T ), and C2T − CT = Ω( T ) for bounded directions and a constant nominal learning rate. These guarantees improvement on vanilla Muon’s logarithmic worst-case accumulation without requiring bounded weight norms. The bound on C2T − CT further shows that accumulation is sustained over later training windows, preserving a necessary condition for continued exploration. • We demonstrate MuonGS’s effectiveness in dense transformer pretraining at 160M and 400M parameter scale, with lower validation perplexity than MuonH at longer horizons and an estimated token-equivalent speedup of up to about 15% at matched validation loss. In addition, we show MuonGS obtains lower validation perplexity than MuonH when pretraining a modern MoE model having about 135M active and 1B total parameters.

2

Related Works

Weight Norm Control. Weight-norm control approaches have been commonly used in neural network training through reparameterization, shrinkage or constraints. Early works like weight normalization controlled weight geometry [Salimans and Kingma, 2016]. AdamW then decoupled weight decay from adaptive gradient preconditioning [Loshchilov and Hutter, 2017], while AdamP and SGDP [Heo et al., 2020] subsequently slowed norm growth by removing the radial component of the update for weights they detect as scale-invariant. Later methods introduced norm targets and per-matrix upper bounds [Loshchilov, 2023, Franke et al., 2024], followed by explicit constraints for LLM pretraining [Loshchilov et al., 2024, Wen et al., 2026, An et al., 2026, Hägele et al., 2026, Xie et al., 2026]. However, these methods do not describe how fast the norm grows, nor do they analyze what impact the increased weight norm has on the overall algorithm performance. Effective Learning Rate. The effective learning rate describes how parameter scale modifies the effect of a nominal learning rate. In scale-invariant settings, Van Laarhoven [2017] derived effective learning rates proportional to η/∥W ∥2 for SGD and η/∥W ∥ for RMSProp and Adam under their respective scaling assumptions. Later works related weight norms and decay to effective learning-rate schedules [Hoffer et al., 2018, Zhang et al., 2018], set up equivalences between weight decay and learning-rate schedules [Li and Arora, 2019], and studied related concepts such as spherical equilibrium, periodic behavior, and optimizer-dependent angular dynamics [Wan et al., 2021, Lobacheva et al., 2021, Roburin et al., 2022, Kosson et al., 2023]. In LLM pretraining, the effective learning rate also helps explain why weight decay improves training [D’Angelo et al., 2024] and why SGD underperforms Adam [Glentis et al., 2026]. More recently, Liu et al. [2026] observed that effective learning rate determines the loss dynamics, while Amin et al. [2026] established stability results for effective learning-rate schedules. However, they do not explore how much the effective learning rate can accumulate over training, or relate this accumulation to exploration capacity.

3

Norm Control Via Accumulated Effective Learning Rate

We begin with formalizing the global perspective on norm control by relating Muon’s exploration over training intervals to the accumulated effective learning rate CT . Section 3.1 starts from what exploration means for a weight matrix of an LLM, derives the effective learning rate as the quantity that sets the size of a Muon step relative to the weights, and shows that the sum of effective learning rates over a training interval bounds how far training can explore in that interval. Section 3.2 compares how fast this sum grows under the three forms in which Muon is run in practice, without norm control, with decoupled weight decay, and with a fixed norm. Then it further answers whether a larger nominal learning rate can substitute for norm control. Proofs are collected in Appendix B. 3.1

Exploration and accumulated effective learning rate

Consider a weight matrix Wt ∈ Rm×n of a transformer block updated by Muon [Jordan et al., 2024], and denote p = min{m, n}. Vanilla Muon updates it as Wt+1 = Wt − ηt Ot , Ot = Orthogonalize(Ut+1 ), (2) where ηt ≥ 0 is the nominal learning rate and Ut+1 is the gradient momentum. Exact orthogonalization replaces each nonzero singular value of Ut+1 by one, giving p √ ∥Ot ∥F = rank(Ut+1 ) ≤ p. (3) 4

Thus a full-rank Muon direction has norm

√

p, independent of the gradient magnitude2 .

Write rt = ∥Wt ∥F for the weight norm and assume rt > 0 throughout training. We study exploration through the net movement of a weight matrix relative to its scale. Over an interval [s, T ], the quantity ∥WT − Ws ∥F rT measures how far the weights have moved relative to their final magnitude. To connect this net movement to Muon’s updates, first compare each step with the weights on which it acts: p ∥Wt+1 − Wt ∥F ηt √ ηt = rank(Ut+1 ) ≤ p . (4) rt rt rt Muon supplies the dimension-dependent update scale, while ηt /rt sets its size relative to the weights. We therefore use ηt /rt as the effective learning rate relative to the weight norm [D’Angelo et al., 2024, Liu et al., 2026]. Norm growth can reduce this rate even when the nominal learning rate is unchanged. Over a training interval [s, T ], these relative updates accumulate: T −1 T −1 X ∥Wt+1 − Wt ∥F √ X ηt ≤ p . rt r t=s t=s t

(5)

This leads to the accumulated effective learning rate C0 = 0,

CT =

T −1 X

ηt . r t=0 t

(6)

The sum CT captures how much relative updating can accumulate along the trajectory. The following proposition connects this quantity to the net displacement that motivated the analysis. Proposition 3.1 (Necessary accumulation for relative displacement). For the updates above and integers 0 ≤ s < T , ∥WT − Ws ∥F √ (7) ≤ p(CT − Cs ). max{rs , rT } In addition, define F (W ) = f (W/∥W ∥F ) and suppose f is L-Lipschitz on the unit Frobenius sphere with respect to Frobenius distance. Then √ |F (WT ) − F (Ws )| ≤ 2L p(CT − Cs ). (8) The first bound (7) implies that√a weight change of at least a fraction δ > 0 of the larger endpoint norm requires CT − Cs ≥ δ/ p for any weight matrix. The second bound (8) further assumes a scale-invariant loss, which means the loss depends only on its direction W/∥W ∥F because its output is normalized. These scale-invariant matrices are common in query and key projections of general transformers. Bound (8) states that under the scale-invariance and √ Lipschitz assumptions with L > 0, a loss decrease of at least ∆ > 0 requires CT − Cs ≥ ∆/(2L p). Note that the Lipschitz assumption is mild as L is the Lipschitz constant of the loss restricted to the unit sphere, which is finite for every continuously differentiable loss. With the two bounds in Proposition 3.1, small accumulation therefore limits both relative movement and, under these additional assumptions, loss reduction. We further note that Proposition 3.1 serves as a necessary condition that a lower-bound threshold is needed for sufficient exploration of the algorithm. It does not claim that larger exploration capacity can achieve lower loss. Proposition 3.1 holds for every interval [s, T ], not only for the whole interval [0, T ]. This is particularly relevant to later training because a large CT may come mainly from early steps. Substantial relative weight movement after step s requires sufficient additional accumulation CT − Cs , so continued exploration requires CT to keep growing. The next section therefore compares norm controls by two metrics: the relative growth CT − C1 with respect to T , as well as the later stage relative growth C2T − CT over the window [T, 2T ]. Note that we use the relative growth of CT because it is not sensitive to initialization scale or model size, which is not the case if we only consider CT alone. See more discussion in Appendix D. 2 Practical implementations use Newton-Schulz iterations [Jordan et al., 2024] and multiply the update by a constant that depends on the shape of the matrix [Liu et al., 2025]. These implementations yield a constant-scale change in the bound of √ ∥Ot ∥F that can absorbed into ηt or p. We further note that all theoretical results below bound through ∥Ot ∥F , so we can use a constant M to bound ∥Ot ∥F for general cases.

5

3.2

Accumulation under three Muon norm controls

Given the metric of CT , we now compare how CT grows under three forms of norm control for Muon, that is, leaving the norm uncontrolled, applying decoupled weight decay, or constraining the norm explicitly. They correspond to vanilla Muon [Jordan et al., 2024], MuonW [Liu et al., 2025], and the norm constraint used by MuonH [Wen et al., 2026], respectively. To isolate their effect, take a constant nominal learning rate ηt = η > 0. • Vanilla Muon follows (2). • MuonW uses Wt+1 = (1 − ηλ)Wt − ηOt ,

λ > 0,

0 < ηλ < 1.

(9)

• For fixed-norm control, let r0 = R > 0 and Wt+1 = R

Wt − ηDt , ∥Wt − ηDt ∥F

(10)

where Dt is the update direction before the weights are renormalized to norm R. Fixed norm uses Dt = Ot , while MuonH additionally fixes the update norm by setting Dt = ROt /∥Ot ∥F . The accumulation result below uses only the fixed weight norm, i.e., Dt = Ot . The next proposition characterizes the accumulated effective learning rate under these three cases. Proposition 3.2 (Accumulation under √ norm control). Assume a constant nominal learning rate ηt = η > 0, and let Rλ = max{r0 , p/λ}. For every integer T ≥ 1: (i) Vanilla Muon satisfies √  η pT  1 , CT ≥ √ log 1 + p r0

√ r0 + 2η p T 1 C2T − CT ≥ √ log . √ p r0 + η p T

(ii) MuonW satisfies CT ≥ ηT /Rλ and C2T − CT ≥ ηT /Rλ . (iii) A fixed norm R gives CT = ηT /R and C2T − CT = ηT /R. Further, for vanilla Muon, the lower bound expressed in (i) is tight as there exists a function and an initialization such that the worst-case rate is attainable, i.e., CT = Θ(log T ) and C2T − CT → √ log 2/ p. Proposition 3.2 shows why norm control matters for accumulation. Without it, the norm can grow linearly in t, so the effective learning rate can decay like 1/t and √ CT can grow as slowly as log T . Over a later window [T, 2T ], the bound in (i) tends to (log 2)/ p as T grows, so the accumulation of vanilla Muon over later windows can stay bounded. In comparison, weight decay and a fixed norm keep the norm at most Rλ or R, so the effective learning rate stays at least η/Rλ or η/R, and CT and C2T − CT grow at least linearly. By Proposition 3.1, bounded accumulation over later windows limits how far the weights can move late in training and, under the scale-invariance and Lipschitz assumptions, how much the loss can still decrease. This offers an explanation for vanilla Muon’s slower progress in training [Liu et al., 2025] and for the improvements obtained with weight decay or fixed-norm constraints. Figure 2 shows the three forms on the example (see Appendix B.3) that shows the bounds of Proposition 3.2(i) are tight: F (W ) = exp(tr W ) from W0 = −4I2 , where every Muon direction is Ot = I2 , so each update√−ηOt points away from the origin and the norm grows as fast as possible. Vanilla Muon has rt = 2(4 + ηt) and CT = Θ(log T ), whereas MuonW and fixed-norm control √ yield CT = Θ(T ). The distinction is clearer over later windows: C2T − CT approaches log 2/ 2 for vanilla Muon but grows linearly with T under the other two controls. The example in Proposition 3.2 shows that the difference in worst-case growth orders can occur in an actual Muon run. In practical LLM pretraining, we also observe such numerical differences in accumulation within the same training budget, though the gap is smaller. (See Figure 3). The results so far fix the nominal learning rate. A natural question is whether a larger nominal learning rate can make up for the small accumulation of vanilla Muon. It does not, because a larger 6

Vanilla Muon

MuonW

MuonH (fixed norm)

(a) Accumulation over the full run

(b) Accumulation over later windows

103 102

102

C2T − CT

CT

101 100

101 100

10−1

10−1

10−2

10−2 100

101

102

103

104

105

log 2/√ 2

101

100

102

103

104

Window start T (interval [T, 2T))

Training horizon T

Figure 2: An attainable gap in accumulation growth. Exact Muon is run on F (W ) = exp(tr W ) from W0 = −4I2 for 105 steps, with √ η = 0.04 and momentum coefficient β = 0.95. MuonW uses λ = 0.2, and MuonH uses R = 4 2. All methods have the same orthogonalized direction I2 , which MuonH rescales √ to Dt = 4I2 . Left: total accumulation. Right: accumulation over [T, 2T ); the dotted line is log 2/ 2. Vanilla Muon

MuonW (WD=0.1)

(a) Accumulation over the full run 100

100

10−1

C2T − CT

CT

10−1 10−2 10−3

10−2 10−3 10−4

10−4 100

MuonH

(b) Accumulation over later windows

101

102

103

104

105

100

101

102

103

104

Window start T (interval [T, 2T])

Training horizon T

Figure 3: Accumulate effective learning rate CT vs training horizon T . Results from training a 160M GPT-style model on FineWeb for 100,000 steps (about 50B tokens) with a constant learning rate of 0.005. Left: total accumulation CT . Right: accumulation over [T, 2T ).

η also makes the norm grow faster, roughly in proportion to η, so the effective learning rates after the first step barely increase (Appendix C makes this precise). Figure 4 shows this in GPT-160M pretraining: raising the learning rate 16-fold moves C20,000 of vanilla Muon only from 0.165 to 0.190, whereas with weight decay C20,000 grows with the learning rate. The best MuonW run for each weight decay (stars) reaches C20,000 between 0.3 and 0.4, a level that vanilla Muon does not reach at any learning rate, and a lower perplexity (21.98–22.02, versus 22.49 for the best vanilla Muon run). To summarize, our analysis links norm control to exploration through accumulated effective learning rates. Weight decay and fixed-norm constraints sustain accumulation, while rapid norm growth can limit it and offset a larger nominal learning rate.

4

MuonGS for LLM Pretraining

The preceding analysis links the benefit of bounded weight norms to the accumulation they preserve. This suggests relaxing the bounded-norm requirement, since sufficiently slow norm growth can still guarantee faster accumulation than vanilla Muon’s logarithmic worst case. For example, at a constant 7

Muon / MuonW: C20,000 versus learning rate WD=0 WD=0.00625

WD=0.0125 WD=0.025

Lowest PPL per WD

22.0595

0.50 22.0078

C20,000

22.0281

0.40

22.0238

22.0542

22.1057

0.30

22.0572

21.9789

22.0241

22.0596

22.1068 22.1738 22.2619

0.20 22.7940

0.01

22.6901

22.6226

22.6014

22.4934

1

0.02

3 8284 0.02

0.04

4142

0.01

22.5313

22.5464

5 6568

0.08

0.05

Matrix learning rate (log scale)

22.6223

22.7649

3137

0.16

0.11

Figure 4: Accumulate effective learning rate CT vs learning rate scale η. Results from training a 160M GPT-style model on FineWeb for 20,000 steps (about 10.5B tokens). Each point represents a learning-rate and weight-decay setting; solid curves connect settings with the same weight decay (WD). WD= 0 is vanilla Muon, and positive WD corresponds to MuonW. Numbers beside the points report perplexity. Stars mark the lowest perplexity among the tested learning rates for each WD; the dashed line connects these stars for MuonW. nominal learning rate η > 0, the bound rt ≤ r0 + cta with c > 0 and 0 ≤ a < 1 gives C2T − CT ≥

ηT = Ω(T 1−a ). r0 + c(2T )a

Accumulation can therefore grow over later training windows even when weight norms increase and effective learning rates decrease. This opens up a broader family of Muon variants and brings us back to the question posed in the introduction. Can this additional flexibility be used to design optimizers that sustain accumulation throughout training and achieve better practical performance? We address this question with MuonGS, which controls norm growth by removing the radial component of Muon’s update. We analyze its accumulation and optimization properties, then evaluate its performance in LLM pretraining. 4.1

Removing the radial component of the Muon update

To construct an update that limits weight-norm growth, we first examine how a vanilla Muon step changes the squared weight norm. ∥Wt − ηt Ot ∥2F = rt2 − 2ηt ⟨Wt , Ot ⟩F + ηt2 ∥Ot ∥2F .

(11)

The term −2ηt ⟨Wt , Ot ⟩F is linear in the step size and increases the squared weight norm when ⟨Wt , Ot ⟩F < 0. Muon controls the size of Ot through matrix orthogonalization, but the update may still point partly along Wt . This radial component is responsible for the first-order norm change. We therefore seek a modified direction that eliminates the first-order norm change while remaining as close as possible to Muon’s proposed update. Giving priority to the norm-control requirement leads to the following simple bilevel formulation for fixed Wt ̸= 0 and Ot min

D∈Rm×n

subject to

1 ∥D − Ot ∥2F 2 1 D ∈ arg min ⟨Wt , Z⟩F2 . Z∈Rm×n 2 8

Algorithm 1 MuonGS for one weight matrix (the projection is line 4). Require: learning rate ηt , momentum β, initial weights W0 , U0 = 0 1: for t = 0, 1, 2, . . . do 2: Ut+1 ← βUt + ∇f (Wt ) ▷ gradient momentum 3: Ot ← Orthogonalize(Ut+1 ) ▷ Newton–Schulz, as in Muon et ← Ot − ⟨Wt , Ot ⟩ Wt 4: O ▷ remove the component along Wt ∥Wt ∥2F et 5: Wt+1 ← Wt − ηt O 6: end for

The lower level minimizes the squared first-order norm change, so its minimizers are precisely the directions orthogonal to Wt . The upper level selects the closest such direction to Ot . The unique solution is the Gram–Schmidt (GS) projection onto the tangent space of the sphere through Wt . We call the resulting method MuonGS, with the update et = Ot − ⟨Wt , Ot ⟩F Wt , O rt2

et . Wt+1 = Wt − ηt O

(12)

The projection gives the solution in closed form, so no additional iterative optimization is required. Since Ot is treated as given, it also applies when the Muon direction incorporates momentum and approximate matrix orthogonalization. et ⟩F = 0 and ∥O et ∥F ≤ √p. Substituting into (11) leaves The projection gives ⟨Wt , O 2 et ∥2F . rt+1 = rt2 + ηt2 ∥O

(13)

Thus the norm may increase, but its squared increment is only quadratic in the step size. AdamP applies the same tangent-space projection to the Adam update, removing its radial component to suppress momentum-induced norm growth on weights it detects as scale-invariant, i.e., when the cosine similarity between the gradient and the weight falls below a small threshold [Heo et al., 2020]. MuonGS applies this projection directly to Muon’s matrix updates, without a scale-invariance test. More importantly, beyond this algorithmic distinction, our analysis links the resulting normgrowth control to exploration over entire training trajectories. As shown next, √ for bounded update √ directions and a constant nominal learning rate, MuonGS guarantees r = O( t), C = Ω( T ), t T √ and C2T − CT = Ω( T ). These guarantees hold without assuming a scale-invariant objective. 4.2

Norm growth and accumulation

We now quantify how the projection’s control of norm √ growth sustains accumulation over training. Starting from r0 > 0 and using the bound ∥Ot ∥F ≤ p, the following theorem controls weight-norm growth under arbitrary learning-rate schedules. For a constant learning rate, it then gives lower bounds on both CT and C2T − CT , showing that accumulation is sustained during later training as well as over the full trajectory. Theorem 4.1 (MuonGS growth and accumulation). Under (12), all iterates remain nonzero and rT2 = r02 +

T −1 X

et ∥ 2 ≤ r 2 + p ηt2 ∥O F 0

t=0

T −1 X

ηt2 .

(14)

t=0

For ηt = η > 0, 2 CT ≥ pη

q  √ 2 2 r0 + pη T − r0 = Ω( T ),

√ ηT C2T − CT ≥ p 2 = Ω( T ). r0 + 2pη 2 T

(15) (16)

√ MuonGS therefore limits the norm to O( t) and guarantees at least square-root accumulation, including over later windows. It improves the worst-case guarantee of vanilla Muon without enforcing 9

Table 1: Norm and accumulation guarantees at a constant positive learning rate. Control Vanilla Muon Weight decay Fixed norm GS

Norm bound

CT

O(t) O(1) R √ O( t)

Ω(log T ) Ω(T ) ηT√/R Ω( T )

C2T − CT Ω(1) Ω(T ) ηT√/R Ω( T )

a fixed radius. Thus the accumulation preserved by MuonGS is available during later training, rather than coming only from early updates. Table 1 summarizes the comparison with vanilla Muon and bounded-norm controls. The preceding bounds describe the relative updating available along a trajectory. To connect this accumulation to optimization, we now consider a fixed scale-invariant matrix objective F (W ) = f (W/∥W ∥F ),

Ot = Orthogonalize(∇F (Wt )),

(17)

where f is twice continuously differentiable on the unit Frobenius sphere. Proposition 4.2 (Stationarity for exact current-gradient MuonGS). Under (17), MuonGS with r0 > 0 and a constant learning rate η > 0 satisfies, for every even T ≥ 2, K ∥ grad f (Wt /rt )∥F ≤ √ , T /2≤t<T T where grad f is the sphere gradient and K < ∞ is independent of T . min

(18)

Here the current-gradient direction provides descent alignment, while GS controls the weight norm and the resulting effective learning rates. Together they guarantee an iterate with a small spheregradient norm in the second half of training, even though the weight norm can keep growing.

5

Experiments

In this section we evaluate the MuonGS optimizer on the task of LLM pre-training. Since the results from Wen et al. [2026] indicate that the MuonH outperforms MuonW in typical pre-training settings, we focus our main comparison to MuonH. Note that both MuonGS and MuonH share the favorable property of not needing to tune a weight decay coefficient as opposed to MuonW. We include some additional results comparing with Muon and MuonW in Appendix A.1. Below we present our main results using dense GPT-style models, but include preliminary results on MoE model training in Section 5.3. 5.1

Experimental settings

Architecture. We use the modernized GPT-style architecture from the Modded-NanoGPT Optimization Benchmark [Jordan, 2026] at two scales, approximately 160M and 400M parameters. The smaller model, which corresponds to the NanoGPT model used in Jordan [2026], has 12 layers, hidden width 768, and 6 attention heads; the other is a scaled-up version and has 24 layers, hidden width 1024, and 8 heads, with head dimension 128 for both sizes. Both models use RMSNorm, rotary positional embeddings, QK-Norm, squared-ReLU MLPs, and untied input and output embeddings. Pre-training setup. We pre-train on the 100B-token subset of the FineWeb dataset [Penedo et al., 2024] using a sequence length of 1024 and a global batch of 512 sequences. Our experiments span approximately 3–25B training tokens for the 160M model and 8–63B for the 400M model, corresponding to roughly 1–8 Chinchilla budgets [Hoffmann et al., 2022] at both scales. We use BF16 mixed-precision training. Optimizers. We compare MuonGS as described in Algorithm 1 with MuonH and Muon/MuonW. All methods use Nesterov momentum 0.95 for hidden matrix weights. AdamW is used for token embeddings, output-head weights, and all 1D parameters, with (β1 , β2 ) = (0.8, 0.95), ϵ = 10−10 , weight decay 10−3 , and respective max learning rates 0.7, 0.004, and 0.015. Matrix and auxiliary learning rates decay linearly to zero without warmup after constant phases covering 15% and 30% of training, respectively. 10

1× Tokens / Chinchilla

2× Tokens / Chinchilla

4× Tokens / Chinchilla

8× Tokens / Chinchilla

23.20

Learning rate

8

04

82

0.

02

4

02 0.

01 0.

41

0.

0.

0.

01

5

71

00

70

0.

00

8

04

82

0.

0.

02

4

02 0.

01

41

0.

0.

01

71

00

0.

5

04

28

8 02

0.

70

02

0.

0.

4 01 41 0. 01 0.

00

71

70 00 0.

0.

7 65

05

8

04

0.

82

0.

0.

02

4

02 0.

01

41

0.

01

71

0.

70 0.

20.80

21.60

22.80

24.50

21.00

21.80

23.00

25.00

Learning rate

Learning rate

Learning rate

4× Tokens / Chinchilla

8× Tokens / Chinchilla

GPT-400M 1× Tokens / Chinchilla

2× Tokens / Chinchilla 16.50

18.40

15.80

17.15 16.40

18.30

17.10

Learning rate

Learning rate

6

8

53

82

02

02

0.

Learning rate

0.

4 41

0.

01

01

71

0.

70

28

8 02

0.

00

02

0.

5

4 01 41 0. 01 0.

00

5 71 00 70 0. 00 0.

0.

8

82 02 0.

36

02 0.

35

4

41 01 0.

0.

01 0.

15.70

16.20 00

71 70

00

0.

16.30

17.05

18.20

0.

Validation perplexity (PPL)

MuonH

GPT-160M

00

Validation perplexity (PPL)

MuonGS

3 00

0.

5 1 00 707 00 0.

0.

14

4 01

0.

02 828 02 0.

0.

Learning rate

Figure 5: Validation perplexity versus learning rate for MuonGS and MuonH across the two model sizes and Chinchilla ration 1×–8×. Stars and dashed horizontal lines indicate the best observed perplexity in each panel. √ Tuning and metrics. We tune learning rates on a multiplicative grid with ratio 2 similar to Wen et al. [2026], separately for all optimizers, model sizes and training durations. For MuonW we additionally tune the weight decay as shown in Appendix A.1. We use as metric the cross-entropy loss Lval over about 10M tokens from the FineWeb validation set, equivalently expressed as perplexity PPL = exp(Lval ). 5.2

Benchmarking Results

Learning Rate Sweeps. From Figure 5 we observe that, for extended training horizons (overtraining), MuonGS achieves lower perplexity than MuonH. We note that in modern LLM pre-training extending to beyond one Chinchilla training tokens, i.e.,. more than 20 tokens per model parameter, is common practice [Grattafiori et al., 2024, Yang et al., 2025]. Therefore, MuonGS’s favorable scaling with training horizon is particularly relevant to modern LLM pre-training. Speedup Estimation. Using the best results for each algorithm from the sweeps above, we estimate the training-budget speedup of MuonGS over MuonH following the method used in Wen et al. [2025]. Specifically, for each model size N , we fit MuonH’s validation loss as a function of training-token budget D by nonlinear least squares: b N (D) = αN D−BN + βN . L b N (DH ) = LGS and report For a MuonGS run achieving loss LGS at budget DGS , we solve L DH /DGS as the estimated speedup, where DH is the estimated training-token budget MuonH would require to match the validation loss achieved by MuonGS. From Figure 6 we observe that MuonGS achieves a token-equivalent speedup over MuonH up to about 7-15% for longer training horizons. 5.3 MoE experiments In this section we include preliminary results comparing MuonGS to MuonH on a modern Mixtureof-Experts (MoE) architecture. The model incorporates grouped-query attention [Ainslie et al., 2023], gated attention [Qiu et al., 2026] and attention residuals (AttnRes) [Team et al., 2026]. It has 32 layers, hidden width 512, and 64 routed experts with top-7 routing plus one shared expert per layer, yielding approximately 135M active non-embedding parameters and 1B total parameters. We train on Ultra-FineWeb [Wang et al., 2025] for 10.8B and 21.6B tokens (i.e., 80 and 160 times the active parameters), using a batch size of 1,024 sequences of length 4,096. Learning rates are tuned on a √ 2-spaced grid, with 2% warmup followed by linear decay to zero. Both MuonGS and MuonH use momentum 0.95 with Nesterov acceleration. 11

GPT-400M Token-equivalent speedup over MuonH

Token-equivalent speedup over MuonH

GPT-160M 1.20×

1.147×

1.15×

1.10×

1.068× 1.05×

1.013× 1.00×

0.95×

0.976×

1×

2×

4×

8×

1.20×

1.139×

1.15×

1.10×

1.074×

1.05×

1.00×

1.026×

0.999×

0.95× 1×

2×

Tokens / Chinchilla

4×

8×

Tokens / Chinchilla

Figure 6: Estimated token-equivalent speedup of MuonGS over MuonH for GPT-160M (left) and GPT-400M (right) across Chinclilla ratios 1×–8×. 135M Active / 1B Total MoE MuonGS

MuonH

Validation perplexity (PPL)

10.8B tokens

21.6B tokens 11.70

13.10

11.65 11.60

13.00

11.55 11.50

12.90

11.45 12.80

11.40

0 0.

2

8 82

2

0 0.

04 0.

01 0.

Learning rate

14

4 01 0.

02

0.

Learning rate

8

82 02 0.

Figure 7: Validation perplexity versus learning rate for MuonGS and MuonH for an 135M active / 1B total MoE model for two different training durations. Stars and dashed horizontal lines indicate the best observed perplexity in each panel. From Figure 7, we observe that MuonGS achieves lower best validation perplexity than MuonH at both training budgets, making it a promising method for modern MoE optimization.

6

Conclusion

We reexamined the role of weight norm control for Muon over the entire training trajectory instead of a single step. In particular, we use the accumulated effective learning rate to measure how far the weights can move relative to their norm over any training interval, and shows that continued exploration requires it to keep growing. We further show that vanilla Muon can only guarantee the accumulation to grow logarithmically, while weight norm control approaches such as weight decay and fixed-norm control guarantee a linear accumulation growth, which explains the role of weight norm control. Based on this finding, we relax the weight norm control to a larger range of approaches. In particular, we propose MuonGS by simply removing the radial component of√the Muon update. Theoretically, we show that MuonGS guarantees an accumulation of at least T , sufficient for continued exploration. Numerically, MuonGS outperforms MuonW and MuonH at long training horizons, providing a speedup of up to 15% at matched validation loss. Our accumulation guarantees are worst-case bounds, and the accumulation guarantee provides a necessary but not sufficient condition for continued exploration. Which rate of norm growth works best for a given training budget remains an open question.

12

References Joshua Ainslie, James Lee-Thorp, Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 conference on empirical methods in natural language processing, pages 4895–4901, 2023. Hasan Amin, Wei-Kai Chang, and Rajiv Khanna. When does scale-invariant optimization become unstable? an exact schedule law with weight decay. arXiv preprint arXiv:2609.09116, 2026. Kang An, Jiaxiang Li, Donald Goldfarb, and Shiqian Ma. Demystifying manifold constraints in llm pre-training. arXiv preprint arXiv:2605.04418, 2026. Francesco D’Angelo, Maksym Andriushchenko, Aditya Varre, and Nicolas Flammarion. Why do we need weight decay in modern deep learning? In Advances in Neural Information Processing Systems, volume 37, 2024. URL https://proceedings.neurips.cc/paper_files/paper/ 2024/hash/29496c942ed6e08ecc469f4521ebfff0-Abstract-Conference.html. Jörg K Franke, Michael Hefenbrock, Gregor Koehler, and Frank Hutter. Improving deep learning optimization through constrained parameter regularization. Advances in Neural Information Processing Systems, 37:8984–9025, 2024. Athanasios Glentis, Dawei Li, Chung-Yiu Yau, and Mingyi Hong. Revisiting the adam-sgd gap in llm pre-training: The role of large effective learning rates. arXiv preprint arXiv:2605.17787, 2026. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. Alexander Hägele, Alejandro Hernández-Cano, Atli Kosson, and Martin Jaggi. Improving neural network training by decoupling the magnitude and direction of weight vectors. arXiv preprint arXiv:2606.25971, 2026. Byeongho Heo, Sanghyuk Chun, Seong Joon Oh, Dongyoon Han, Sangdoo Yun, Gyuwan Kim, Youngjung Uh, and Jung-Woo Ha. Adamp: Slowing down the slowdown for momentum optimizers on scale-invariant weights. arXiv preprint arXiv:2006.08217, 2020. Elad Hoffer, Ron Banner, Itay Golan, and Daniel Soudry. Norm matters: efficient and accurate normalization schemes in deep networks. Advances in Neural Information Processing Systems, 31, 2018. Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models (2022). arXiv preprint arXiv:2203.15556, 2022. Keller Jordan. Modded-nanogpt optimization benchmark, 2026. URL https://github.com/ KellerJordan/modded-nanogpt/tree/master/records/track_3_optimization. Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https://kellerjordan. github. io/posts/muon, 6(3):4, 2024. Maxim Kodryan, Ekaterina Lobacheva, Maksim Nakhodnov, and Dmitry Vetrov. Training scaleinvariant neural networks on the sphere can happen in three regimes. Advances in Neural Information Processing Systems, 35:14058–14070, 2022. Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay balances learning across neural networks. arXiv preprint arXiv:2305.17212, 2023. Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. arXiv preprint arXiv:1910.07454, 2019. Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for llm training. arXiv preprint arXiv:2502.16982, 2025. 13

Zihan Liu, Ruiheng Zheng, Shaobo Zhang, Changxin Tian, Kunlong Chen, Zhiqiang Zhang, and Lei Wu. Effective learning rate governs loss dynamics in language model pretraining. arXiv preprint arXiv:2608.24814, 2026. URL https://arxiv.org/abs/2608.24814. Ekaterina Lobacheva, Maxim Kodryan, Nadezhda Chirkova, Andrey Malinin, and Dmitry P Vetrov. On the periodic behavior of neural network training with batch normalization and weight decay. Advances in Neural Information Processing Systems, 34:21545–21556, 2021. Ilya Loshchilov. Weight norm control. arXiv preprint arXiv:2311.11446, 2023. Ilya Loshchilov and Frank Hutter. arXiv:1711.05101, 2017.

Decoupled weight decay regularization.

arXiv preprint

Ilya Loshchilov, Cheng-Ping Hsieh, Simeng Sun, and Boris Ginsburg. nGPT: Normalized transformer with representation learning on the hypersphere. arXiv preprint arXiv:2410.01131, 2024. Guilherme Penedo, Hynek Kydlíček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale, 2024. Zihan Qiu, Zekun Wang, Bo Zheng, Zeyu Huang, Kaiyue Wen, Songlin Yang, Rui Men, Le Yu, Fei Huang, Suozhi Huang, et al. Gated attention for large language models: Non-linearity, sparsity, and attention-sink-free. Advances in Neural Information Processing Systems, 38:100092–100118, 2026. Simon Roburin, Yann de Mont-Marin, Andrei Bursuc, Renaud Marlet, Patrick Perez, and Mathieu Aubry. Spherical perspective on learning with normalization layers. Neurocomputing, 487:66–74, 2022. Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems, 29, 2016. Kimi Team, Guangyu Chen, Yu Zhang, Jianlin Su, Weixin Xu, Siyuan Pan, Yaoyu Wang, Yucheng Wang, Guanduo Chen, Bohong Yin, et al. Attention residuals. arXiv preprint arXiv:2603.15031, 2026. Twan Van Laarhoven. L2 regularization versus batch and weight normalization. arXiv preprint arXiv:1706.05350, 2017. Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamics: Learning dynamics of normalized neural network using sgd and weight decay. Advances in Neural Information Processing Systems, 34:6380–6391, 2021. Yudong Wang, Zixuan Fu, Jie Cai, Peijun Tang, Hongya Lyu, Yewei Fang, Zhi Zheng, Jie Zhou, Guoyang Zeng, Chaojun Xiao, et al. Ultra-fineweb: Efficient data filtering and verification for high-quality llm training data. arXiv preprint arXiv:2505.05427, 2025. Kaiyue Wen, David Hall, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them. arXiv preprint arXiv:2509.02046, 2025. Kaiyue Wen, Xingyu Dang, Kaifeng Lyu, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them ii: Hyperball optimization. arXiv preprint arXiv:2606.16899, 2026. Tian Xie, Haoming Luo, Haoyu Tang, Yiwen Hu, Jason Klein Liu, Qingnan Ren, Yang Wang, Wayne Xin Zhao, Rui Yan, Bing Su, et al. Controlled llm training on spectral sphere. arXiv preprint arXiv:2601.08393, 2026. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. arXiv preprint arXiv:1810.12281, 2018. 14

A

Additional Experiments

A.1

Experiments with Muon/MuonW

In Figure 8 we include a comparison of MuonH and the proposed MuonGS over Muon and MuonW (with varying weight decay). We observe that Muon, not using any type of weight norm control, performs noticeably worse than the other baselines. MuonGS obtains the best perplexity, followed by MuonH and MuonW.

MuonGS MuonH

Muon (WD=0) MuonW (WD=0.00625)

MuonW (WD=0.0125) MuonW (WD=0.025)

GPT-160M

GPT-400M 18.30

22.80

Validation perplexity (PPL)

Validation perplexity (PPL)

18.20 22.60

22.40

22.20

22.00

18.10 18.00 17.90 17.80 17.70

21.80

7 0. 0

0.

56 5

04

8 28 2

0.

02

4

Learning rate

0. 0

0.

0. 0

14 1

01 0.

1 70 7 00

05 0. 0

16

31 0.

0.

11

7

08 0.

56 5 0. 0

8

04 0.

02

28 2 0. 0

4 0. 0

0.

14 1

1

01 0.

70 7 00 0.

0.

00

0. 0

35 3

05

6

17.60

Learning rate

Figure 8: Validation perplexity vs learning rate for MuonGS, MuonH, Muon and MuonW with varying weight decay. Both models are trained for 20k steps (about 10.5B tokens) on FineWeb.

We also compare the evolution of the norm of the combined matrix weights, ∥wt ∥2 , during training. As shown in Figure 9, MuonH maintains a constant norm, while Muon without weight decay exhibits considerably larger norm growth. MuonW’s norm initially increases and subsequently declines, whereas MuonGS exhibits more moderate, gradual norm growth without weight decay or an explicit fixed-norm constraint.

GPT-160M: matrix-weight norms MuonGS (LR=0.0282843) MuonH (LR=0.01)

GPT-400M: matrix-weight norms

Muon (LR=0.04) MuonW (WD=0.0125, LR=0.04)

MuonGS (LR=0.02) MuonH (LR=0.00707107)

Muon (LR=0.0282843) MuonW (WD=0.0125, LR=0.0282843)

5000 4000 4000

‖wt‖2

‖wt‖2

3000 2000

3000 2000

1000

1000

0 0

5k

10k

15k

20k

Training steps

0

0

5k

10k

15k

20k

Training steps

Figure 9: Weight-norm evolution for the 160M (left) and 400M (right) models over 20k training steps, using each method’s best learning rate. Here, wt concatenates the matrix weights updated by Muon variants, excluding auxiliary AdamW parameters. MuonW uses weight decay 0.0125 which was found to be the best among the ones tested as shown 8 . 15

B

Proofs and Examples

B.1

Proof of Proposition 3.1

Proof. For nonzero matrices X, Y let d(X, Y ) =

∥X − Y ∥F . max{∥X∥F , ∥Y ∥F }

We first show that d(X, Y ) ≤ d(X, Z) + d(Z, Y ). Write x = ∥X∥F , y = ∥Y ∥F , z = ∥Z∥F , and assume x ≥ y. If z ≤ x, both denominators on the right are at most x, and the claim follows from the triangle inequality for ∥ · ∥F . If z > x, use the inversion ι(X) = X/∥X∥2F , which satisfies ∥ι(X) − ι(Y )∥F = ∥X − Y ∥F /(xy). Then d(X, Y ) = y∥ι(X) − ι(Y )∥F ,

d(X, Z) = x∥ι(X) − ι(Z)∥F ,

d(Z, Y ) = y∥ι(Z) − ι(Y )∥F ,

and the claim follows from the triangle inequality for ι and y ≤ x. Applying this inequality along the iterates and then (4), T −1 T −1 X X ∥Wt+1 − Wt ∥F ηt ∥Ot ∥F ∥WT − Ws ∥F √ ≤ ≤ ≤ p (CT − Cs ), max{rs , rT } max{rt , rt+1 } rt t=s t=s

which is (7). For (8), let x ≥ y again; since x − y ≤ ∥X − Y ∥F , 1 X Y ∥X − Y ∥F 1  ∥X − Y ∥F + x − y − ≤ + ∥Y ∥F − ≤ 2 d(X, Y ). = x y F x y x x Hence |F (WT ) − F (Ws )| ≤ L∥WT /rT − Ws /rs ∥F ≤ 2L d(WT , Ws ), and (7) gives (8). B.2

Proof of Proposition 3.2

√ Proof. Vanilla Muon. Fix integers 0 ≤ √ s < u. The triangle inequality gives √ rt+1 ≤ rt + ηt p, so at a constant learning rate rt ≤ r0 + η p t. The function x 7→ η/(r0 + η p x) is decreasing, so its value at an integer t is at least its integral over [t, t + 1]. Summing over s ≤ t < u gives Cu − C s =

u−1 X

η ≥ r t=s t

Z u s

√ r0 + η p u 1 η dx = √ log √ √ . r0 + η p x p r0 + η p s

Taking (s, u) = (0, T ) and (s, u) = (T, 2T ) give (i). √ √ MuonW. For 0 ≤ ηt λ ≤ 1, rt+1√≤ (1 − ηt λ)rt + ηt p. Since p ≤ λRλ , the bound rt ≤ Rλ implies rt+1 ≤ Rλ − ηt (λRλ − p) ≤ Rλ . Therefore by induction rt ≤ Rλ for all t, so Cu − C s =

u−1 X

u−1 ηt 1 X ≥ ηt . r Rλ t=s t=s t

At a constant learning rate, the steps 0 ≤ t < T and T ≤ t < 2T each contribute ηT to the last sum, which gives (ii). Fixed norm. Renormalization keeps rt = R as long as the weight before renormalization is nonzero, Pu−1 so Cu − Cs = R−1 t=s ηt , which gives (iii) in the same way. The tightness of (i) is shown by the example of Appendix B.3. B.3

An exact Muon example

Let F (W ) = exp(tr W ) on Rp×p and W0 = −bIp with b > 0. The gradient ∇F (W ) = exp(tr W )Ip is a positive multiple of Ip , at every W , and so is the momentum√(also with Nesterov’s variant). Exact orthogonalization therefore gives Ot = Ip with ∥Ot ∥F = p at every step. For every method below, we choose a constant nominal learning rate ηt = η > 0. 16

Vanilla Muon.

By induction, Wt = −(b + ηt)Ip ,

rt =

√

p (b + ηt),

η 1 η =√ . rt p b + ηt

(19)

The applied update −ηOt is a positive multiple of Wt , so every step points directly √ √ away from the origin. In particular, −⟨Wt , Ot ⟩F /rt = p, and the norm bound rt+1 ≤ rt + η p in Appendix B.2 holds with equality. We now calculate Cu − Cs : Cu − Cs =

u−1 1 X η η =√ . r p t=s b + ηt t=s t

u−1 X

Since the function x 7→ 1/(b + x) is decreasing, we have   Z η(t+1) η dx 1 1 0≤ − ≤η − . b + ηt b+x b + ηt b + η(t + 1) ηt √ Summing over s ≤ t < u and dividing by p gives   1 b + ηu η 1 1 η 0 ≤ Cu − Cs − √ log ≤√ − ≤ √ . (20) p b + ηs p b + ηs b + ηu b p √ Since r0 = b p, the logarithms for (s, u) = (0, T ) and (T, 2T √) are exactly the two bounds of Proposition 3.2(i). So C exceeds the first bound by at most η/(b p) = C1 , C2T − CT exceeds the T √ second by at most η/( p (b + ηT )) → 0, and, as T → ∞, log T CT = √ + O(1), p

log 2 C2T − CT −→ √ . p

(21)

With weight decay and q = 1 − ηλ ∈ (0, 1), the update Wt+1 =

MuonW and fixed norm. qWt − ηIp gives

1 t 1  √ + b− q, rt = p bt . (22) λ λ √ √ Since bt → 1/λ geometrically, η/rt = ηλ/ p + O(q t ). Summing gives CT = ηλT / p + O(1) √ and C2T − CT = ηλT / p + O(q T ). Wt = −bt Ip ,

bt =

With a fixed norm R = r0 , the weight before renormalization is −(b + η)Ip for Dt = Ot , and −b(1 + η)Ip for MuonH, whose rescaled direction is Dt = ROt /∥Ot ∥F = bIp . Both renormalize to W0 , so Wt = W0 and CT = C2T − CT = ηT /r0 . Numerical illustration in Figure 2. In Figure 2, we choose p = 2, b = 4, η = 0.04, λ = 0.2, momentum coefficient β = 0.95 and 105 steps. The corresponding weight norms are thus √ √ √ rtvanilla = 2 (4 + 0.04t), rtW = 2 (5 − 0.992t ), rtH = 4 2. (23) B.4

Proof of Theorem 4.1

et ⟩F = 0 and ∥O et ∥2 = ∥Ot ∥2 − ⟨Wt , Ot ⟩2 /rt2 ≤ p. Hence Proof. The projection gives ⟨Wt , O F F F 2 et ∥2 = r2 + η 2 ∥O et ∥2F , rt+1 = ∥Wt − ηt O F t t

so the iterates stay nonzero, and every subsequent projection is defined. Summing over t gives rT2 = r02 +

T −1 X

et ∥F2 ≤ r02 + p ηt2 ∥O

t=0

T −1 X t=0

This proves the radius bound (14). 17

ηt2 .

p

For ηt = η, (14) gives rt ≤ T −1 X

r02 + pη 2 t, which increases with t; comparing the sum with an integral, Z T

q  η dx 2 2 2T − r . p = r + pη 0 0 pη r02 + pη 2 t r02 + pη 2 x 0 t=0 p For fixed r0 , p and η, the right-hand side equals 2 T /p + O(1), proving (15). p For T ≤ t < 2T , the radius bound gives rt ≤ r02 + 2pη 2 T . Hence CT ≥

η

≥

p

C2T − CT =

2T −1 X t=T

and the right-hand side equals B.5

p

η ηT ≥p 2 , rt r0 + 2pη 2 T

T /(2p) + O(1), which proves (16).

Proof of Proposition 4.2

Proof. Write W̄t = Wt /rt and Gt = grad f (W̄t ). The chain rule gives ∇F (Wt ) = Gt /rt , so Ot = Orthogonalize(Gt ). Since Gt is tangent, ⟨Gt , W̄t ⟩F = 0, and et ⟩F = ⟨Gt , Ot ⟩F = ∥Gt ∥∗ . ⟨Gt , O Here ∥ · ∥∗ denotes the nuclear norm. When Gt ̸= 0, Cauchy–Schwarz gives 1≤

∥Gt ∥∗ et ∥F ≤ ∥Ot ∥F ≤ √p. ≤ ∥O ∥Gt ∥F

If Gt = 0 for some t < T , all later iterates are stationary and (18) holds. Otherwise (14) gives r02 + η 2 t ≤ rt2 ≤ r02 + pη 2 t,

0 ≤ t ≤ T.

(24)

Let L1 and L2 bound the sphere-gradient norm and the operator norm of the Riemannian Hessian of f . These bounds are finite because f is C 2 on the compact unit sphere. At a nonstationary step, set u=

et O , e t ∥F ∥O

x=

et ∥F η∥O , rt

θ = arctan x.

Orthogonality gives W̄t+1 = W̄t cos θ − u sin θ. Taylor’s theorem along this great-circle arc yields f (W̄t+1 ) ≤ f (W̄t ) − θ⟨Gt , u⟩F + Since x − arctan x =

L2 2 θ . 2

Rx 2 s /(1 + s2 ) ds ≤ x2 /2, using θ ≤ x and 0 ≤ ⟨Gt , u⟩F ≤ L1 , we obtain 0

f (W̄t+1 ) ≤ f (W̄t ) − x⟨Gt , u⟩F +

L1 + L 2 2 η η2 x ≤ f (W̄t ) − ∥Gt ∥∗ + B 2 2 rt rt

(25)

where B = p(L1 + L2 )/2. If a stationary iterate occurs before T , the claimed minimum is zero. Otherwise, the radius bounds in (24) imply T −1 X t=T /2

η2 ≤ 1, rt2

T −1 X t=T /2

η ηT ≥ p 2 . rt 2 r0 + pη 2 T

Summing the one-step estimate (25) over interval T /2 ≤ t < T and using ∥Gt ∥F ≤ ∥Gt ∥∗ gives p
2 r02 + pη 2 T min ∥Gt ∥F ≤ max f − min f + B = O(T −1/2 ), ηT T /2≤t<T where the extrema are over the unit sphere and all constants are independent of T . 18

C

Dependence on the Nominal Learning Rate

Section 3.2 asks whether a larger nominal learning rate can make up for the small accumulation of vanilla Muon. This section shows that it cannot beyond the first step. In particular, we vary the constant learning rate η while keeping the initialization and the norm-control parameters fixed. Note that the directions Ot and the norms rt may depend on η. The first term C1 = η/r0 is linear in η for every method, so we compare the accumulation after the first step, CT (η) − Cs (η) =

T −1 X

η , r (η) t=s t

1 ≤ s < T.

Vanilla Muon. For vanilla Muon, a larger learning rate enlarges every update but may also accelerate norm growth. The next proposition shows that the two effects offset each other when the updates keep an outward component and, with a weaker bound, whenever they have no inward component. Proposition C.1 (Learning-rate compensation). Consider vanilla Muon with a constant nominal learning rate η, and let 1 ≤ s < T . (i) If ⟨Wt , Ot ⟩F ≥ v > 0, 0 ≤ t < T, rt √ then r0 + ηvt ≤ rt ≤ r0 + η p t for 0 ≤ t ≤ T . Moreover, −

T −1 T −1 X 1 1X1 1 ≤ . √ ≤ CT − Cs ≤ r /η + p t r /η + vt v t=s t t=s 0 t=s 0

T −1 X

(26)

(27)

(ii) If ⟨Wt , Ot ⟩F ≤ 0 and ∥Ot ∥F ≥ µ > 0 for 0 ≤ t < T , then rt2 ≥ r02 + µ2 η 2 t for 0 ≤ t ≤ T . Moreover, T −1 T −1 X η 1 X 1 p √ . CT − C s ≤ (28) ≤ µ t=s t r02 + µ2 η 2 t t=s Proof. (i) Consider the direction along Wt /rt . The norm rt+1 is at least the component of Wt+1 = Wt − ηOt along this direction, which gives   Wt rt+1 ≥ , Wt − ηOt = rt − η⟨Wt , Ot ⟩F /rt ≥ rt + ηv. rt F √ Meanwhile, the triangle inequality gives rt+1 ≤ rt + η p. Summing these increments gives (26). Dividing η by these bounds and summing over s ≤ t < T gives (27). (ii) Since ⟨Wt , Ot ⟩F ≤ 0, 2 rt+1 = rt2 − 2η⟨Wt , Ot ⟩F + η 2 ∥Ot ∥2F ≥ rt2 + µ2 η 2 .

Summing gives the norm bound, and dividing η by it gives (28) in the same way. Proposition C.1 shows why a larger learning rate cannot raise the accumulation of vanilla Muon after the first step: the norm grows in proportion to η. In (i), every update has an outward component, 2 2 and rt ≥ ηvt; √ in (ii), tangent updates enlarge the norm through the second-order term η ∥Ot ∥F , and rt ≥ µη t. A larger η thus enlarges the steps and the norm by the same factor, so the effective learning rates after the first step are bounded independently of η: η 1 ≤ rt vt

in (i),

η 1 ≤ √ rt µ t

in (ii),

t ≥ 1.

Summing them gives the last bounds in (27) and (28): whatever the value of η, CT − Cs stays PT −1 PT −1 below v −1 t=s 1/t in (i) and below µ−1 t=s t−1/2 in (ii). The middle sums in (27) and (28) do increase with η, but they converge to these limits, so a larger η can bring CT − Cs at most up to a level that does not depend on η. Only the first step, taken before the norm has grown, gains from a larger η, as C1 = η/r0 grows linearly. 19

The same holds when ηPscales a fixed schedule P ηt 2= ηgt with g0 > 0 as in practical training. In particular, with St = Qt = j<t gj and j<t gj , the effective learning rates satisfy ηt /rt ≤ √ gt /(vSt ) in (i) and ηt /rt ≤ gt /(µ Qt ) in (ii) for t ≥ 1, again independent of η. In contrast, norm control approaches bound the norm rt by a quantity that does not depend on η, so the accumulation grows linearly with η, exactly for a fixed norm and as a guarantee for MuonW. Specifically, at a fixed radius R, CT − Cs = η(T√− s)/R; and for MuonW with √fixed λ > 0 and 0 < ηλ < 1, the recursion rt+1 ≤ (1 − ηλ)rt + η p gives rt ≤ Rλ = max{r0 , p/λ}, and hence PT −1 CT − Cs ≥ η(T − s)/Rλ . With a schedule, η(T − s) becomes η t=s gt in both, for MuonW when ηλgt ≤ 1.

D

Dependence on Scale and the Matrix Size

In the previous sections, the value of CT depends on the initialization scale, the matrix size and on how Muon’s updates are scaled, while the orders of growth that our bounds guarantee do not. Values can therefore be compared between runs that share the initialization and the update scale, as in Figures 3 and 4, whose runs differ only in the norm control, η and λ. Comparisons across initializations or widths need the normalization below. Below we assume ηt = η > 0. Scale and growth.

Since C1 = η/r0 , T −1 X r0 CT = . C1 r t=0 t

(29)

Thus the factor C1 carries the learning rate, the initial norm, and any constant that an implementation maintains throughout training. The ratio CT /C1 measures accumulation in units of the first step and depends only on the normalized norms rt /r0 . Because C1 does not depend on T , CT and CT /C1 grow at the same rate in T . Normalized accumulation bounds.

Let

√ η p . r0 This quantity bounds the initial update size relative to the weight norm. By Proposition 3.2 and Theorem 4.1, we have α=

vanilla Muon: MuonGS: MuonW: fixed norm:

CT log(1 + αT ) , ≥ C1 α   CT 2 p 1 + α2 T − 1 , ≥ 2 C1 α CT T ≥ , C1 max{1, α/(ηλ)} CT = T. C1

(30) (31) (32) (33)

The first two follow by comparing the sum of r0 /rt with an integral, and √ the third from Rλ /r0 = max{1, α/(ηλ)}. For every α > 0, the right-hand sides grow like log T , T , T and T . The value of α changes their constants and the time at which they stop growing linearly, about αT ≈ 1 for vanilla Muon and α2 T ≈ 1 for MuonGS, but not their orders in T . These expressions show that matching α gives the same normalized bounds for vanilla Muon and MuonGS across matrix sizes. MuonW may additionally require matching ηλ. Independence of matrix size for Muon. With the usual initialization, √ a matrix of shape md × nd 2 with independent N (0, σ 2 /nd ) entries has r0,d ≈ md σ 2 , so r0,d ≈ σ md , while the Muon update p √ size is pd = min{md , nd }. Therefore, r η pd √ αd = η pd /r0,d ≈ σ md 20

is asymptotically independent of d. Consequently, at a fixed ratio of p/m, the normalized Muon -type bounds (30)-(32) are asymptotically independent of absolute matrix size at a fixed horizon.

21

Related documents

• MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control#1469317
• The Effectiveness and Durability of Pretraining Data Filtering in Protein Language Models#1401906
• The Effectiveness and Durability of Pretraining Data Filtering in Protein Language Models#1381256
• Rank Measures of Fine-Tuning Updates Need a Noise Baseline#1386416
• CoLU: Confidence-Weighted Layerwise Unlearning for Large Language Models#1386693