PaperlayerView exact revision on arXiv

Source-linked AI summary

MeqMuon: Matrix-Equilibrating Muon for LLM Pretraining

Chang-Wei Shi, Xu Wang, Wu-Jun Li

arXiv:2609.35701v1cs.LG

TL;DR

LLM scale increases pretraining costs, while row-wise normalization cannot accommodate all imbalance patterns in update matrices. MeqMuon tailors normalization to balance rows and columns, and reports better convergence than AdamW, Muon, and other baselines.

  • Problem

    Row-wise normalization alone cannot accommodate update matrices with different row- and column-magnitude imbalance patterns, a concern as LLM pretraining costs grow.

    Source: §1 · pp. 1–2

  • Method

    MeqMuon automatically tailors normalization to balance row and column magnitudes and replaces AdamW updates for remaining parameters with normalized momentum.

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

  • Results

    MeqMuon achieves better convergence performance than AdamW, Muon, and other baselines, with the best final validation PPL at every reported model scale.

    Source: §4 · p. 8, §5 · p. 10

  • Takeaways & Limitations

    Replacing AdamW updates with normalized momentum eliminates second-moment storage and reduces optimizer-state memory usage.

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

    The paper focuses on 1D and 2D parameters; higher-dimensional tensors are reshaped or partitioned into 2D matrices in Muon implementations.

    Source: §2 · p. 2

Abstract

from arXiv · show

The success of large language models (LLMs) has been accompanied by continued growth in model size and pretraining costs. Muon offers high accuracy and training efficiency in LLM pretraining. Recent work introduces row-wise normalization into Muon to balance update magnitudes and improve pretraining performance. However, row-wise normalization alone cannot accommodate different imbalance patterns in update matrices. In this paper, we propose an improved Muon optimizer, called \underline{m}atrix-\underline{eq}uilibrating Muon~(MeqMuon), for LLM pretraining. MeqMuon balances both row and column magnitudes through normalization that can be automatically tailored to different imbalance patterns without manual intervention. Moreover, MeqMuon eliminates the need to store AdamW's second-moment estimates, reducing optimizer-state memory usage. Empirical results demonstrate that MeqMuon achieves better convergence performance than AdamW, Muon, and other baselines in LLM pretraining.

1 INTRODUCTION

Growing LLM scale makes pretraining costly and increases the importance of efficient optimizer design. The section identifies limits of row-wise normalization and introduces MeqMuon’s matrix-equilibrating approach and reported convergence gains.

1 INTRODUCTION

  • LLM growth raises pretraining costs, making optimizer design important for improving performance and reducing training cost.The cited examples include DeepSeek-V4-Pro at 1.6 trillion parameters and 33 trillion pretraining tokens, and Kimi K3 at 2.8 trillion parameters.Source: §1 · p. 1
  • Muon orthogonalizes momentum matrices for hidden-layer 2D weights, while updating remaining parameters with AdamW.AdamW uses first- and second-moment estimates for coordinate-wise adaptive updates.Source: §1 · p. 1
  • Muon’s orthogonalized updates can be imbalanced predominantly by row or column, while row-wise normalization alone cannot accommodate both patterns.The section also identifies imbalance in both row and column magnitudes for unorthogonalized momentum matrices of remaining 2D parameters.Source: §1 · p. 2
  • MeqMuon balances row and column magnitudes with normalization tailored automatically to different imbalance patterns, and reports better convergence than AdamW, Muon, and other baselines.It replaces AdamW updates for remaining parameters with normalized momentum, eliminating second-moment storage and reducing optimizer-state memory.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries frame LLM pretraining as optimization over parameter tensors, focusing on 1D and 2D parameters. They review AdamW and Muon update rules and describe common Muon scaling and weight-decay modifications.

2 PRELIMINARIES

  • The paper focuses on 1D and 2D parameters; higher-dimensional tensors are reshaped or partitioned into 2D matrices in Muon implementations.The optimization notation uses L for training loss and θ for the collection of model parameter tensors.Source: §2 · p. 2
  • Adam maintains exponential moving averages of minibatch-gradient first and second moments, which support coordinate-wise adaptive updates.AdamW modifies Adam with decoupled weight decay.Source: §2 · p. 2
  • Muon accumulates a momentum matrix, orthogonalizes it, and uses the result to update hidden-layer 2D weights; other parameters use AdamW.Orthogonalization is defined over nonzero singular values and approximated in practice with Newton–Schulz iterations.Source: §2 · p. 3
  • Common large-scale Muon modifications add decoupled weight decay and RMS alignment, with a shape factor compensating for matrix-shape effects on update RMS.The RMS alignment coefficient sets the target update RMS approximately to that of AdamW; cited implementations use 0.2 or 0.18.Source: §2 · p. 3

3 METHOD

Section 3 characterizes distinct row- and column-magnitude imbalance patterns in Muon updates, then introduces MeqMuon’s normalization tailored to those patterns. MeqMuon applies adaptive one-sided normalization to hidden-layer weights and different normalization schemes to other parameter types.

3.1 IMBALANCE PATTERNS

  • Row and column CVs measure magnitude imbalance as standard deviation divided by mean, with zero rows and columns excluded from calculation.RMS values at or below δ = 10^-7 are treated as zero; larger CVs indicate greater imbalance.Source: §3.1 · pp. 3–4
  • After five NS iterations, q_proj, k_proj, v_proj, gate_proj, and up_proj show row imbalance, while o_proj and down_proj show column imbalance.These contrasting patterns occur within the same Llama-350M Transformer layer.Source: §3.1 · p. 4, Table 1 · p. 4
  • During NS iterations, both CVs approach zero for square matrices, but for rectangular matrices one approaches zero while the other remains large.For v_proj, (γr, γc) changes from (0.673, 0.773) at k = 0 to (0.195, 0.041) at k = 5; exact orthogonalization constrains only one side of rectangular matrices.Source: §3.1 · p. 4, Figure 1 · p. 5
  • Unorthogonalized momentum matrices for embed_tokens and lm_head have substantial row imbalance and non-negligible column imbalance.No orthogonality condition enforces equal row or column magnitudes for these matrices.Source: §3.1 · p. 4, Table 2 · p. 5
  • Because updates can be row-, column-, or two-sided imbalanced, row-wise normalization alone cannot accommodate all observed patterns.NorMuon applies row-wise normalization to address row imbalance in orthogonalized updates.Source: §3.1 · p. 4

3.2 MEQMUON

  • MeqMuon uses adaptive row- or column-wise normalization for hidden-layer 2D weights, two-sided normalization for other 2D parameters, and global-RMS normalization for 1D parameters.For hidden-layer weights, it selects the normalization direction independently for each matrix at every optimization step.Source: §3.2 · pp. 5–6
  • Table 3 shows that tailored normalization reduces the larger CV of orthogonalized updates to nearly zero and substantially reduces both CVs for remaining 2D parameters.For orthogonalized updates, the other CV remains small.Source: §3.2 · p. 6, Table 3 · p. 6

4 EXPERIMENTS

Across the evaluated models, MeqMuon achieves the best final validation PPL and generally stays below Muon and NorMuon during training. It also reduces optimizer-state memory and selects normalization directions in line with matrix geometry.

Experimental setup

  • Experiments use Llama, SmolLM2, and Qwen2 models trained from random initialization on English C4, with a token budget set to 20 times model parameters.Source: §4 · pp. 6–7

Convergence performance

Memory overhead

  • On Qwen2-0.5B, MeqMuon reduces optimizer-state memory by 21.6% compared with Muon and NorMuon.Source: §4 · p. 8, Table 5 · p. 8

Normalization direction

  • Tall gate_proj and up_proj matrices always select row-wise normalization, while wide down_proj matrices always select column-wise normalization.Source: §4 · p. 9, Table 6 · p. 9

Ablation study

  • At both model scales, variants modifying either hidden-layer 2D weights or remaining 2D parameters alone beat Muon on validation PPL, while MeqMuon achieves the lowest PPL.Source: §4 · p. 9, Table 7 · p. 9

5 CONCLUSION

MeqMuon adapts normalization across row and column magnitudes and removes second-moment storage. In LLM pretraining, it achieves better convergence performance than AdamW, Muon, and other baselines.

  • MeqMuon balances row and column magnitudes through normalization tailored automatically to different imbalance patterns.Source: §5 · p. 9
  • MeqMuon eliminates second-moment storage by replacing AdamW updates with normalized momentum for remaining parameters.Source: §5 · p. 9
  • MeqMuon achieves better convergence performance than AdamW, Muon, and other baselines in LLM pretraining.Source: §5 · p. 10

A ADDITIONAL IMBALANCE PATTERNS

The additional imbalance analysis reports row and column coefficients of variation for selected update matrices across three model configurations.

A ADDITIONAL IMBALANCE PATTERNS

B MODEL ARCHITECTURES

The model-architecture description distinguishes attention types and weight sharing across model families; accompanying tables cover architectures and selected imbalance statistics.

B MODEL ARCHITECTURES

Related papers

Loading 2609.35701v1…