Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Muon+: Towards Better Muon via One Additional Normalization Step

Ruijie Zhang, Yequan Zhao, Ziyue Liu, Zhengyang Wang, Zheng Zhang

TL;DR

This work proposes a simple yet effective enhancement to Muon, namely Muon+, which introduces an additional normalization step after orthogonalization and extends the token-to-parameter (T2P) ratio to an industrial level of $\approx 200$.

Abstract

The Muon optimizer has demonstrated promising performance in pre-training large language models through gradient (or momentum) orthogonalization. In this work, we propose a simple yet effective enhancement to Muon, namely Muon+, which introduces an additional normalization step after orthogonalization. We demonstrate the effectiveness of Muon+ through extensive pre-training experiments across a wide range of model scales and architectures. Our evaluation includes GPT-style models ranging from 130M to 774M parameters and LLaMA-style models ranging from 60M to 1B parameters. We comprehensively evaluate the effectiveness of Muon+ in the compute-optimal training regime and further extend the token-to-parameter (T2P) ratio to an industrial level of $\approx 200$. Experimental results show that Muon+ provides a consistent boost on training and validation perplexity over Muon. We provide our code here: https://github.com/K1seki221/MuonPlus.

Muon+: Towards Better Muon via One Additional Normalization Step

TL;DR

This work proposes a simple yet effective enhancement to Muon, namely Muon+, which introduces an additional normalization step after orthogonalization and extends the token-to-parameter (T2P) ratio to an industrial level of .

Abstract

The Muon optimizer has demonstrated promising performance in pre-training large language models through gradient (or momentum) orthogonalization. In this work, we propose a simple yet effective enhancement to Muon, namely Muon+, which introduces an additional normalization step after orthogonalization. We demonstrate the effectiveness of Muon+ through extensive pre-training experiments across a wide range of model scales and architectures. Our evaluation includes GPT-style models ranging from 130M to 774M parameters and LLaMA-style models ranging from 60M to 1B parameters. We comprehensively evaluate the effectiveness of Muon+ in the compute-optimal training regime and further extend the token-to-parameter (T2P) ratio to an industrial level of . Experimental results show that Muon+ provides a consistent boost on training and validation perplexity over Muon. We provide our code here: https://github.com/K1seki221/MuonPlus.
Paper Structure (26 sections, 5 equations, 4 figures, 15 tables, 1 algorithm)

This paper contains 26 sections, 5 equations, 4 figures, 15 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Muon Optimizer
  4. Remark on the role of normalization.
  5. The Muon+ Method
  6. Experiments
  7. Pre-training GPT
  8. Pre-training LLaMA
  9. Overtraining GPT and LLaMA
  10. Overtraining GPT
  11. Overtraining LLaMA
  12. Ablation Study
  13. Performance under Different Learning Rates
  14. Impact of Different Normalization Directions
  15. Ablation for Polar Methods $\mathrm{Ortho}(\cdot)$
  16. ...and 11 more sections

Figures (4)

  • Figure 1: Pre-training GPT and LLaMA models at scales ranging from 130M to 1B parameters under compute-optimal settings. Quantitative results are provided in Section \ref{['sec:experiments']}. Muon+ consistently outperforms Muon across all runs. We also conduct overtraining experiments for both GPT and LLaMA; the results are presented in Section \ref{['sec:overtrain']}.
  • Figure 2: Training loss curves under overtraining for GPT-Base and LLaMA-350M.
  • Figure 3: Validation perplexity sweep for LLaMA models under different settings. Here "none (baseline)" is the standard Muon optimizer; "row", "col", "row_col" and "col_row" indicate different normalization directions in Muon+.
  • Figure 4: Validation perplexity sweep for GPT models under different settings.