MuonGS: Relaxing Weight-Norm Control to Accelerate LLM Pretraining

In language model pretraining, weight decay and explicit weight-norm constraints are widely used to improve optimizer performance, but there has been a lack of a unified theoretical explanation for why this control works and how much control is needed. A recent paper provides an in-depth analysis of norm control in the Muon optimizer and proposes a new method, MuonGS, which accelerates training by relaxing weight-norm control.

A New Theoretical Perspective on Weight-Norm Control

The research team introduces the "cumulative effective learning rate" $C_T = \sum_{t<T} \eta_t / |W_t|_F$ as an analytical tool, 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, the cumulative quantity over any training interval can constrain the relative displacement of the weights. This framework connects weight-norm control with the optimizer's exploration capability: the larger the cumulative effective learning rate, the greater the relative magnitude by which the weights can move during training, potentially leading to better exploration.

The analysis reveals an interesting phenomenon: the cumulative effective learning rate of original Muon can only achieve logarithmic growth during training, which may limit the model's exploration capability and help explain its slower progress in practical training. In contrast, MuonW and MuonH can guarantee at least linear accumulation under given conditions.

Furthermore, the authors point out that if the weight norm grows slowly enough, the cumulative effective learning rate can exceed the logarithmic worst case, opening up space for designing optimizers with fewer restrictions.

Cumulative effective learning rate: The "effective learning rate" refers to the nominal learning rate divided by the weight norm, reflecting the relative update magnitude of parameters at a unit norm scale. Since the weight norm changes during training, under a fixed nominal learning rate, the actual exploration step size may be compressed. The cumulative quantity $C_T$ sums the relative step sizes over the entire training process, characterizing the total relative distance the weights can move. This metric is directly related to the optimizer's exploration capability: the slower $C_T$ grows, the more constrained the relative displacement of the weights, and the harder it is for the model to escape its current region. Original Muon achieves only logarithmic growth, meaning the effective learning rate decays severely in later stages, providing a theoretical basis for explaining its slower training progress.

MuonGS: Removing the Radial Component with Gram-Schmidt Projection

Based on the above insights, the paper proposes MuonGS. This method uses Gram-Schmidt (GS) projection to remove the radial component from each Muon update without fixing the weight norm. This design preserves the benefits of weight-norm control while avoiding the limitations imposed by hard constraints.

Theoretical analysis shows that, under bounded update directions and a constant nominal learning rate, after $t$ steps the weight norm grows at most proportionally to $\sqrt{t}$. Therefore, MuonGS guarantees $C_T = \Omega(\sqrt{T})$, meaning the cumulative effective learning rate can continue to grow in later stages of training. For smooth scale-invariant objectives, the authors also establish a stationarity guarantee of $O(T^{-1/2})$.

Gram-Schmidt (GS) projection is used here to remove the component of the Muon update vector along the current weight direction. Specifically, given the weight $W_t$ and the update direction $U_t$, the GS projection computes the projection of $U_t$ onto the direction of $W_t$ and subtracts it, yielding an update component orthogonal to $W_t$. In this way, the update only changes the direction of the weights without directly changing their norm, thereby avoiding hard norm constraints. Unlike fixed-norm methods, GS projection does not force the weight norm to be constant, allowing it to grow slowly. Theoretical analysis shows that this design makes the weight norm grow at most proportionally to $\sqrt{t}$, thereby guaranteeing the cumulative effective learning rate $C_T = \Omega(\sqrt{T})$, which can continue to grow in later stages.

Pretraining Experimental Validation

The research conducted pretraining experiments on dense models with 160M and 400M parameters and on an MoE model with 135M activated parameters and 1B total parameters. The results show that, under long training cycles, MuonGS outperforms MuonH on validation perplexity. In particular, on dense models, compared with MuonH, an estimated token-equivalent speedup of about 15% can be obtained when matching validation loss.

This acceleration effect means that, under the same computational budget, MuonGS can utilize data more efficiently or achieve lower loss. For large-scale pretraining, this efficiency improvement has practical value.

Token-equivalent speedup is a commonly used metric for measuring optimizer efficiency: the percentage of tokens saved by a new method compared with a baseline when reaching the same validation loss. For example, a 15% token-equivalent speedup means that MuonGS can reach the loss level that MuonH reaches with 100% of the tokens using about 85% of the tokens. This metric directly reflects the improvement in data efficiency and has significant economic value in large-scale pretraining. The experiments were conducted on 160M and 400M dense models and on an MoE model with 135M activated parameters and 1B total parameters, covering different scales and architectures and enhancing the generalizability of the conclusions.

Significance and Outlook

The proposal of MuonGS not only provides a new design approach for optimizers but also deepens the understanding of weight-norm control. Through theoretical analysis and experimental validation, the authors show that relaxing norm control can accelerate training without sacrificing performance. In the future, this method is expected to be integrated into existing large-model training pipelines, further advancing improvements in pretraining efficiency.

Paper title: MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control

(This article is based on the paper abstract; please refer to the original text for specific details.)

One-line Summary

MuonGS accelerates large-model pretraining by relaxing weight-norm control, with theoretical proof that the cumulative effective learning rate continues to grow.

Paragraph Summary

This article analyzes the role of weight-norm control in the Muon optimizer and proposes the cumulative effective learning rate $C_T$ as a theoretical tool, finding that original Muon can achieve only logarithmic accumulation, which may limit exploration capability. To address this, MuonGS is proposed, using Gram-Schmidt projection to remove the radial component from updates without fixing the weight norm. It is theoretically proven that the weight norm grows as √t and $C_T=\Omega(\sqrt{T})$. In pretraining on 160M and 400M dense models and on an MoE model with 135M activated parameters and 1B total parameters, MuonGS outperforms MuonH on validation perplexity under long training cycles, and on dense models an estimated token-equivalent speedup of about 15% is obtained when matching loss.