PaperlayerView exact revision on arXiv

Source-linked AI summary

MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control

Athanasios Glentis, Dawei Li, Mingyi Hong

arXiv:2610.37647v1cs.LGcs.AImath.OC

TL;DR

The paper asks how and why norm control helps Muon optimization, and how much control is needed. It analyzes accumulated effective learning rate and proposes MuonGS; at long horizons, MuonGS achieves lower validation perplexity than MuonH, with an estimated token-equivalent speedup of up to about 15% at matched validation loss in dense models.

  • Problem

    A unified basis for choosing among norm-control methods remains incomplete, including whether existing restrictions are essential or weaker control can improve optimization.

    Source: §1 · p. 2

  • Method

    The paper analyzes norm control across training trajectories using accumulated effective learning rate and proposes MuonGS, which removes the radial component of Muon updates through Gram–Schmidt projection.

    Source: §1 · p. 2, §4.1 · p. 9

  • Results

    Up to about 15% estimated token-equivalent speedup at matched validation loss to MuonH; MuonGS also achieves lower validation perplexity than MuonH at longer horizons in dense and MoE pretraining.

    Source: §1 · p. 4

  • Takeaways & Limitations

    MuonGS’s accumulation guarantee improves on vanilla Muon’s logarithmic worst case without requiring bounded weight norms, while sustaining accumulation in later training.

    Source: §1 · p. 4, §6 · p. 12

    The best norm-growth rate for a given training budget remains open, and accumulation guarantees are worst-case and necessary but not sufficient for continued exploration.

    Source: §6 · p. 12

Abstract

from arXiv · show

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 entire training trajectory using the \emph{accumulated effective learning rate}, $C_T=\sum_{t<T}\eta_t/\|W_t\|_F$, where $\eta_t$ is the nominal learning rate and $\eta_t/\|W_t\|_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. Based on this insight, 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 $\sqrt{t}$ after $t$ steps. Consequently, MuonGS guarantees $C_T=\Omega(\sqrt{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

The paper studies whether Muon’s norm-control requirements can be relaxed while preserving exploration over training trajectories. It uses accumulated effective learning rate to analyze norm-control methods and motivate MuonGS, then evaluates the optimizer theoretically and in pretraining.

1 Introduction

  • The paper asks whether existing norm-control restrictions are essential, since a unified basis for choosing among controls remains incomplete.Weight-norm control is used in LLM pretraining through weight decay and explicit constraints, which prescribe scales and leave less freedom for their evolution.Source: §1 · pp. 1–2
  • For bounded update directions, accumulated effective learning rate bounds relative weight displacement over any training interval; continued exploration requires sufficient additional accumulation on that interval.A large total accumulation may come mostly from early iterations, so the increment C_T − C_s matters for later exploration.Source: §1 · p. 2
  • With a constant learning rate, vanilla Muon can have Θ(log T) accumulation, while MuonW guarantees Ω(T) and fixed-norm control gives Θ(T).The passage states these results under their stated conditions.Source: §1 · p. 2
  • Sufficiently slow norm growth can yield faster accumulation than vanilla Muon’s logarithmic worst case without requiring bounded weight norms.This broader condition motivates optimizer designs with less restrictive norm control.Source: §1 · pp. 2–3
  • MuonGS removes the radial update component, making the update orthogonal to the weights and limiting norm growth to second order in the step size.The paper says its accumulation guarantee improves on vanilla Muon’s worst case without fixing a radius.Source: §1 · p. 3
  • At long horizons, MuonGS achieves lower validation perplexity than MuonH and up to about 15% token-equivalent speedup at matched validation loss.MuonGS also achieves lower validation perplexity than MuonH in pretraining a modern MoE model with about 135M active and 1B total parameters.Source: §1 · p. 4

2 Related Works

Prior work addresses weight-norm control through reparameterization, shrinkage, and constraints, and studies how parameter scale affects effective learning rates. The paper identifies trajectory-wide accumulation and its relation to exploration as gaps in that work.

2 Related Works

  • Weight-norm methods range from weight normalization and decoupled decay to radial-update removal, norm targets, upper bounds, and explicit constraints.The passage notes that prior methods do not describe norm-growth rates or analyze their impact on overall algorithm performance.Source: §2 · p. 4
  • The related-work gap is trajectory-wide analysis of how norm growth and effective-learning-rate accumulation affect exploration and algorithm performance.The paper distinguishes this gap from prior studies of norm growth and effective learning-rate schedules.Source: §2 · p. 4
  • Prior effective-learning-rate work studies parameter-scale effects, decay equivalences, optimizer dynamics, and LLM pretraining, but not accumulation over training or its relation to exploration capacity.The passage cites work on effective learning rates for SGD, RMSProp, and Adam under their respective scaling assumptions.Source: §2 · p. 4

3 Norm Control Via Accumulated Effective Learning Rate

The section uses accumulated effective learning rate to relate norm control to exploration over training intervals. It shows that weight decay and fixed-norm control sustain accumulation, unlike vanilla Muon’s logarithmic worst case.

3.1 Exploration and accumulated effective learning rate

  • The accumulated effective learning rate sums η_t/∥W_t∥_F, the step size relative to the current weight norm, across a training interval.Muon’s update direction sets a dimension-dependent scale, while norm growth can reduce the effective rate even when the nominal rate stays unchanged.Source: §3.1 · pp. 4–5
  • For bounded update directions, relative weight displacement over any interval is bounded by its additional accumulation, C_T − C_s.Under scale-invariant, Lipschitz loss assumptions, the same accumulation also bounds attainable loss decrease; insufficient accumulation limits both movement and loss reduction.Source: §3.1 · p. 5

3.2 Accumulation under three Muon norm controls

  • Vanilla Muon can have Θ(log T) total accumulation and bounded accumulation over later windows, while MuonW and fixed-norm control guarantee linear growth.The example demonstrates this gap in an exact Muon run; practical pretraining also shows numerical differences, though smaller.Source: §3.2 · p. 6
  • In a tight example, vanilla Muon’s C_2T − C_T approaches log 2/√p, while MuonW and fixed-norm control grow linearly over that later window.Figure 2 contrasts total accumulation with accumulation over [T, 2T) for the three controls.Source: §3.2 · p. 6, Figure 2 · p. 7
  • Raising vanilla Muon’s learning rate 16-fold moves C_20,000 only from 0.165 to 0.190, while the best MuonW runs reach 0.3–0.4 and lower perplexity.The best MuonW runs report perplexity 21.98–22.02 versus 22.49 for the best vanilla Muon run.Source: §3.2 · p. 7

4 MuonGS for LLM Pretraining

The section develops MuonGS as a less restrictive way to control norm growth while sustaining effective-learning-rate accumulation over training. It uses a radial-removing projection and establishes accumulation and stationarity guarantees under stated conditions.

4 MuonGS for LLM Pretraining

  • With norm growth bounded by r_t ≤ r_0 + c t^a for 0 ≤ a < 1, accumulation over later windows is Ω(T^(1−a)), motivating norm control without a fixed radius.MuonGS is analyzed from r_0 > 0 with bounded update directions and a constant positive nominal learning rate.Source: §4 · p. 8, §4.2 · pp. 9–10

4.1 Removing the radial component of the Muon update

  • MuonGS projects Muon’s update onto the tangent space at the current weights, removing its radial component so the squared norm increment is only quadratic in the step size.The projection is the closest update to Muon’s proposed direction among directions orthogonal to the weights, and requires no additional iterative optimization.Source: §4.1 · p. 9

4.2 Norm growth and accumulation

  • For a twice continuously differentiable scale-invariant objective and constant learning rate, the stationarity result guarantees an iterate with small sphere-gradient norm in the second half of training, even as the norm grows.The proof uses current-gradient descent alignment alongside GS control of weight norm and effective learning rates.Source: §4.2 · p. 10

5 Experiments

The experiments compare MuonGS with MuonH on dense GPT-style models and a modern MoE model. MuonGS achieves lower validation perplexity at extended dense-model training horizons and at both tested MoE budgets.

5.1 Experimental settings

  • The dense experiments use modernized GPT-style models of approximately 160M and 400M parameters, pretrained on FineWeb across roughly 1–8 Chinchilla budgets.Learning rates are tuned separately by optimizer, model size, and training duration; validation perplexity is the reported metric.Source: §5.1 · pp. 10–11

5.2 Benchmarking Results

  • At extended training horizons, MuonGS achieves lower validation perplexity than MuonH on the dense GPT-style models.Figure 5 compares validation perplexity across learning rates, two model sizes, and Chinchilla ratios of 1×–8×.Source: §5.2 · p. 11, Figure 5 · p. 11
  • MuonGS achieves a token-equivalent speedup of about 7–15% over MuonH for longer training horizons.The estimate uses MuonH’s fitted validation loss as a function of training-token budget to calculate the budget needed to match MuonGS’s loss.Source: §5.2 · p. 11, Figure 6 · p. 12

5.3 MoE experiments

  • MuonGS achieves lower best validation perplexity than MuonH at both tested training budgets for the 135M-active, 1B-total-parameter MoE model.The model is pretrained on Ultra-FineWeb; Figure 7 compares two training durations.Source: §5.3 · pp. 11–12, Figure 7 · p. 12

6 Conclusion

The conclusion frames MuonGS as a less restrictive approach to norm control, using removal of Muon’s radial update component. The paper reports improved long-horizon results but notes that its accumulation guarantees do not ensure exploration.

6 Conclusion

  • MuonGS removes the radial component of Muon updates and is reported to outperform MuonW and MuonH at long training horizons.The conclusion reports up to 15% speedup at matched validation loss.Source: §6 · p. 12
  • Accumulation guarantees are worst-case bounds and provide a necessary, but not sufficient, condition for continued exploration.The rate of norm growth that works best for a given training budget remains an open question.Source: §6 · p. 12

A.1 Experiments with Muon/MuonW

The appendix compares MuonGS with Muon and MuonW, and tracks how each method’s matrix-weight norm evolves. It also gives proof steps for a normalized distance bound.

A.1 Experiments with Muon/MuonW

  • MuonGS obtains the best perplexity, followed by MuonH and MuonW, while uncontrolled Muon performs noticeably worse.Figure 8 compares these optimizers, including MuonW with varying weight decay, after 20k steps on FineWeb.Source: §A.1 · p. 15, Figure 8 · p. 15
  • MuonGS exhibits moderate, gradual matrix-weight norm growth without weight decay or an explicit fixed-norm constraint.MuonH maintains a constant norm, Muon’s norm grows considerably, and MuonW’s norm first increases and then declines.Source: §A.1 · p. 15, Figure 9 · p. 15

B.1 Proof of Proposition 3.1

  • The proof establishes the triangle inequality for d(X,Y), defined by the Frobenius distance divided by the larger endpoint norm.It uses the Frobenius triangle inequality when the intermediate norm is no larger than the larger endpoint, and inversion otherwise.Source: §B.1 · p. 16
  • Applying the distance inequality along the iterates supports a bound on objective change for the scale-normalized weights.The stated bound is |F(WT) − F(Ws)| ≤ 2L d(WT, Ws).Source: §B.1 · p. 16

B.2 Proof of Proposition 3.2

The proof compares accumulation under vanilla Muon, MuonW, and fixed-norm control, then gives an exact example illustrating the different growth orders. Under constant learning rate, vanilla Muon’s accumulation can be logarithmic, while fixed-norm control yields linear accumulation.

B.2 Proof of Proposition 3.2

  • For constant learning rate, vanilla Muon’s norm is bounded above by its initial norm plus a term linear in the step count, yielding a logarithmic lower bound on accumulated effective learning rate.The proof sums the decreasing lower bound η/(r0 + η√p t) over the interval.Source: §B.2 · p. 16
  • MuonW preserves the norm bound r_t ≤ Rλ by induction when 0 ≤ η_tλ ≤ 1 and √p ≤ λRλ.At constant learning rate, the proof uses this bound to establish linear accumulation over successive intervals.Source: §B.2 · p. 16
  • Fixed-norm renormalization keeps r_t = R, giving linear accumulated effective learning rate at constant learning rate.This holds as long as the weight before renormalization is nonzero.Source: §B.2 · p. 16

B.3 An exact Muon example

  • In the exact example, vanilla Muon’s update points directly away from the origin, attaining the norm-growth bound used in the proof.The example sets F(W)=exp(tr W), W0=−bI_p, and uses a constant nominal learning rate.Source: §B.3 · pp. 16–17
  • For that example, vanilla Muon has C_T = log T/√p + O(1), while fixed-norm control gives C_T = ηT/r0.The fixed-norm example also has C_2T − C_T = ηT/r0.Source: §B.3 · p. 17

B.4 Proof of Theorem 4.1

The proof establishes a radius bound and derives a stationarity estimate using the sphere gradient, tangent updates, and a one-step Taylor bound.

B.4 Proof of Theorem 4.1

  • The proof establishes the radius bound identified as (14).Source: §B.4 · p. 17

B.5 Proof of Proposition 4.2

  • The chain rule relates the gradient at the normalized iterate to the gradient at the original weight, and orthogonalization yields a tangent update.Source: §B.5 · p. 18
  • Compactness of the unit sphere and C2 smoothness make the sphere-gradient and Riemannian-Hessian bounds finite.Source: §B.5 · p. 18
  • A one-step estimate follows by applying Taylor’s theorem along the great-circle arc between successive normalized iterates.The proof sums this estimate over the latter half of the training interval.Source: §B.5 · p. 18

C Dependence on the Nominal Learning Rate

The section shows that increasing the nominal learning rate does not raise post-first-step accumulation for vanilla Muon: the larger steps are offset by norm growth. By contrast, norm-control approaches can make accumulation grow linearly with the learning rate.

C Dependence on the Nominal Learning Rate

D Dependence on Scale and the Matrix Size

The section normalizes accumulation by its first-step value to compare runs across initialization scales and matrix sizes. It reports that normalized Muon-type bounds are asymptotically independent of absolute matrix size at fixed p/m and horizon.

D Dependence on Scale and the Matrix Size

Related papers

Loading 2610.37647v1…