PaperlayerView exact revision on arXiv

Source-linked AI summary

MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads

Linkai Ma, Xinyu Luo, Mengbo Wang, Ananth Grama, Petros Drineas, Brian Bullins

arXiv:2610.02705v1cs.LGcs.AI

TL;DR

Standard Muon implementations use AdamW for embedding and language-model head layers, without a unified principled norm-aware treatment for both. MuonIO matches their geometries through transposition to produce one normalized-vector update, and the paper reports lower optimizer-state cost and improved validation perplexity. The analysis also identifies a limitation in how the head’s softmax invariance is handled.

  • Problem

    Standard Muon implementations route embedding and language-model head layers to AdamW, and a principled norm-aware approach covering both with one update had not been established.

    Source: §1 · pp. 1–2

  • Method

    MuonIO uses the 1→2 norm for embeddings and the 2→∞ norm for the language-model head, whose transposition identity yields a shared normalized-vector update.

    Source: §1.1 · p. 2

  • Results

    MuonIO halves persistent optimizer state and improves validation perplexity over a strong Muon baseline across 60M, 130M, and 1B LLaMA-family pretraining.

    Source: §1.1 · p. 2

  • Takeaways & Limitations

    A single geometry-matched Muon-style update for both vocabulary-facing matrices is reported to reduce their persistent optimizer state while improving validation perplexity.

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

    The 2→∞ head analysis does not explicitly account for softmax invariance to constant shifts of all logits.

    Source: §6 · p. 9

Abstract

from arXiv · show

The Muon optimizer derives its update rule for hidden linear layers by solving a local linearization of the loss penalized by the spectral norm, motivated by an RMS-stability argument for dense linear layers. Standard Muon implementations, however, exclude the input (embedding table) and output (language model head) layers from this principled treatment, for which they use AdamW instead. We present MuonIO, a single Muon-style update for both of these layers. For the language model head $\mathbf{L} \in \mathbb{R}^{V \times d}$, we motivate the use of the $2\to\infty$ operator norm, due to the Lipschitz continuity of the softmax output geometry, while for the embedding table $\mathbf{E} \in \mathbb{R}^{d \times V}$, we draw on the $1 \to 2$ operator norm, based on the one-hot input geometry identified by Bernstein & Newhouse (2025). The identity $\lVert\mathbf{L}\rVert_{2\to\infty}=\lVert\mathbf{L}^\top\rVert_{1\to2}$ then puts both matrices in the same vocabulary-oriented geometry: MuonIO applies a single normalized-vector rule, which appears as column normalization for $\mathbf{E}$ and row normalization for $\mathbf{L}$. Empirical evaluations demonstrate the effectiveness of our approach, with MuonIO reducing I/O optimizer state memory by 50% and I/O update FLOPs by $\sim$46% compared to Muon for 1B LLaMA pretraining on C4, while also improving validation perplexity.

1 INTRODUCTION

MuonIO addresses the unexplained use of AdamW for embedding and language-model head layers by giving both a principled, shared Muon-style update. It aligns their geometries through transposition and reports improved validation perplexity alongside lower optimizer costs.

1 INTRODUCTION

  • Standard Muon routes embedding and language-model head layers to AdamW, while a principled norm-aware update for both vocabulary-facing layers had not been established.Prior work identified the one-hot embedding geometry and proposed normalization updates on other grounds, but did not provide a unified treatment.Source: §1 · pp. 1–2

1.1 OUR CONTRIBUTION

  • MuonIO uses the 1→2 geometry for embeddings and the 2→∞ geometry for the language-model head; transposition makes their vocabulary-oriented updates the same normalized-vector rule.For weight-tied models, where L = E⊤, the rule acts directly on the shared parameter.Source: §1.1 · p. 2
  • MuonIO halves persistent optimizer state and improves validation perplexity over a strong Muon baseline across 60M, 130M, and 1B LLaMA-family pretraining.AdamW uses two optimizer moments per parameter, whereas MuonIO keeps one; the two I/O matrices each contain Vd entries.Source: §1.1 · p. 2

1.2 RELATED WORK

2 BACKGROUND AND NOTATION

This section introduces norm-penalized steepest descent and the operator-norm framework used to motivate updates for different transformer layers. It distinguishes dense hidden layers, one-hot embeddings, and the probability-facing language-model head.

2.1 NOTATION

  • The notation section defines matrix norms, the Frobenius inner product, row and column indexing, identity matrices, and the vocabulary index set.It also introduces the probability simplex over V outcomes.Source: §2.1 · p. 2

2.2 STEEPEST DESCENT UNDER A NORM

  • For a weight-matrix norm, steepest descent minimizes a penalized local linearization of the loss; Frobenius geometry recovers the ordinary gradient step.With spectral-norm geometry, the solution is the polar-factor update used by Muon.Source: §2.2 · p. 3

2.3 OPERATOR NORMS

  • An operator norm measures the largest output norm over unit-norm inputs, and the section introduces dual-norm identities and row- and column-normalization directions.The 1→2-penalized problem has a column-normalized solution.Source: §2.3 · p. 3

2.4 LAYER TYPES AND THEIR GEOMETRIES

  • Hidden layers map dense activations, so the spectral norm bounds output perturbations; embeddings map one-hot token inputs to dense vectors.For hidden layers, the bound is ∥∆y∥2 ≤ ∥∆W∥2∥x∥2.Source: §2.4 · pp. 3–4
  • The language-model head maps hidden states to logits and then probabilities, so its perturbation is assessed through total variation of the output distribution rather than logits alone.The corresponding analysis is developed in Subsection 3.2.Source: §2.4 · p. 4

3 MUONIO FOR VOCABULARY-FACING LAYERS

MuonIO identifies a geometry for each vocabulary-facing layer, relates the two geometries by transposition, and applies one update rule to both. The embedding geometry comes from one-hot lookup, while the language-model-head geometry is tied to stability of its probability output.

3.1 EMBEDDING GEOMETRY FROM PRIOR WORK

  • The embedding table's 1→2 operator norm measures the largest change to any token embedding, and its established geometry supplies one half of MuonIO.Bernstein & Newhouse (2024) identify the 1→RMS norm, equivalent to 1→2 up to a scaling factor.Source: §3.1 · p. 4

3.2 LANGUAGE-MODEL-HEAD GEOMETRY

  • The 2→∞ norm gives sharp uniform control of worst-case total-variation perturbation at fixed hidden-state norm, making it the natural penalty for lm_head.The constant in the global softmax Lipschitz bound cannot be improved, although its supremum is not attained by any finite pair of distinct inputs.Source: §3.2 · p. 5

3.3 TRANSPOSITION SYMMETRY

  • Transposition makes the two layer geometries equivalent: the lm_head 2→∞ norm equals the transposed embedding 1→2 norm.With tied weights, column normalization of E is exactly row normalization of L = E⊤, so MuonIO acts once on the shared parameters.Source: §3.3 · p. 5

3.4 ONE MUONIO UPDATE

  • In vocabulary-oriented form, both layers use the same 2→∞ penalty, the maximum row 2-norm, and the unsmoothed solution row-normalizes the gradient.In the original storage orientation, this normalizes columns of E and rows of L; the algorithm uses a momentum estimate and numerical safeguards.Source: §3.4 · p. 6

4 EXPERIMENTS

Experiments compare MuonIO and alternative I/O-layer treatments in C4 language-model pretraining, measuring validation perplexity, optimizer cost, and I/O-gradient dynamics. MuonIO improves perplexity over standard-Muon configurations while reducing I/O-layer optimizer costs versus AdamW.

4.1 EXPERIMENTAL SETUP

  • Experiments compare MuonIO, Ember, and Polar Express IO with standard Muon and Polar Express hidden updates on C4 LLaMA models at 60M, 130M, and 1B scales.The models are trained from scratch with untied embeddings and heads; the study uses final validation perplexity, optimizer cost, and I/O-gradient norms. Configurations jointly tune their distinct learning rates by parameter group.Source: §4 · p. 6, §4.1 · pp. 6–7, Table 1 · p. 7

4.2 OPTIMIZER PERFORMANCES

  • With standard Muon hidden updates, MuonIO achieves the lowest mean validation perplexity among four configurations at every scale; with Polar Express, it lowers mean perplexity at 130M and 1B but slightly raises it at 60M.Table 2 reports final C4 validation perplexity, with lower values better. For 60M and 130M, mean and sample standard deviation use the three of five seeds with the lowest final validation perplexity; for 1B, they use three seeds.Source: §4.2 · pp. 7–8, Table 2 · p. 7

4.3 IO VERSUS ADAMW: EXACT OPTIMIZER COST

  • For one V × d I/O matrix, MuonIO requires roughly half the FLOPs and half the optimizer state of AdamW.Table 3 reports per-update costs; its counting convention and operation-by-operation derivation are in Appendix C.Source: Table 3 · p. 8, §4.3 · p. 8

4.4 GRADIENT NORM DYNAMICS

  • Across three gradient measures, MuonIO reduces I/O-layer gradients faster than Ember and Muon Tuned, with the clearest separation for the language-model head’s maximum absolute entry.The 130M comparison covers the embedding table and language-model head; each plotted curve represents one run.Source: §4.4 · p. 9, Figure 2 · p. 9

5 CONCLUSION

MuonIO applies one Muon-style update to the input embedding table and output language-model head, using geometry matched to each layer. Experiments show improved loss and perplexity alongside lower optimizer costs for these layers.

5 CONCLUSION

  • MuonIO applies a single Muon-style update to the embedding table and language-model head, using 1→2 geometry for embeddings and 2→∞ geometry for the head.The embedding geometry is equivalent to 1→RMS up to a fixed d^-1/2 factor; the head geometry follows from total-variation stability of softmax.Source: §5 · p. 9
  • For vocabulary-facing layers, MuonIO is cheaper than AdamW in optimizer-state memory and FLOPs.The experiments also report improved final loss and perplexity.Source: §5 · p. 9

6 LIMITATIONS AND FUTURE WORK

The analysis uses total variation to measure changes in the language-model head’s output distribution and controls them with the 2→∞ norm.

6 LIMITATIONS AND FUTURE WORK

  • The head’s total-variation bound does not account explicitly for softmax invariance to common shifts of all logits.The authors suggest studying shift-invariant formulations and other distribution-change measures, including KL divergence.Source: §6 · p. 9

A RELATED WORK

Related work includes embedding geometry, vocabulary-wise normalization, and other special treatments for token-facing matrices, but uses differing motivations and formulations.

A RELATED WORK

  • Bernstein & Newhouse (2025) derive column normalization for embeddings without empirical evaluation, while Scion’s operator-norm framework includes column and row normalization options for different motivations.Scion motivates its output option by controlling the largest output coordinate rather than softmax stability.Source: A RELATED WORK · p. 11
  • SCALE uses column normalization across layers, whereas Lau & Su (2026) derive a family including token normalization from vocabulary-reordering and hidden-space equivariance requirements.After transposition, SCALE’s final-layer normalization resembles one component of MuonIO, but SCALE is otherwise a distinct algorithm.Source: A RELATED WORK · p. 11
  • Ember treats embeddings and language-model heads specially using factorized second moments rather than a Muon-style update.For untied heads, its rule also centers momentum, removes common logit shifts, and selects scale separately.Source: A RELATED WORK · p. 12

B PROOFS

The proofs establish that the softmax map has supremum sensitivity 1 in the stated norm ratio, although no finite pair attains equality. They transfer this bound to the language-model head by realizing logit pairs with a matrix update.

Softmax bound

  • For even V, the proof constructs a sign vector with zero mean and unit deviations, giving ∥J(z)u∥1 = 1.The odd-V construction balances the sign groups by adjusting one coordinate's probability mass.Source: B PROOFS · p. 12
  • For odd V, doubling one coordinate's mass balances the positive and negative sign groups and again yields ∥J(z)u∥1 = 1.This is described as the even case with V + 1 uniform coordinates, two positive coordinates merged.Source: B PROOFS · p. 12
  • The supremum of the softmax ratio is 1: the constructed ratios approach 1, and the upper bound is supplied by Lemma 3.4.For odd V, the construction approaches the same limit by sending the last logit to −M and taking M → ∞.Source: B PROOFS · p. 12

Non-attainment

  • No finite distinct logit pair attains equality, because equality would force a zero derivative despite the proof deriving a strictly positive one.The proof concludes that the inequality is strict for every finite pair.Source: B PROOFS · p. 12

Language-model head bound

  • Every distinct logit pair can be realized by an lm_head update, so its ratio has supremum 1, but no finite update triple attains it.The non-attainment follows from the strict inequality for finite logit pairs.Source: B PROOFS · pp. 12–13

C EXACT FLOP AND OPTIMIZER-STATE ACCOUNTING

The accounting contrasts AdamW’s two moment buffers and 13N per-matrix FLOPs with MuonIO’s single momentum buffer and fused row-normalized update. MuonIO avoids materializing the update matrix.

AdamW accounting

MuonIO accounting

D EXPERIMENT SETTINGS AND HYPERPARAMETERS

The experiments tune optimizer settings on C4 validation data and report 18 selected LLaMA pretraining runs. The settings specify shared training conditions, scale-dependent batching, and model configurations.

Hyperparameter tuning

  • Hyperparameters are selected through Bayesian sweeps for Polar Express variants and grid searches for Muon, Muon Tuned, MuonIO, and Ember.The Polar Express sweeps allow up to 100 trials; I/O learning-rate grids are independently searched for Muon Tuned, MuonIO, and Ember.Source: D EXPERIMENT SETTINGS AND… · pp. 14–15

Selected hyperparameters

  • The 18 selected C4 runs report hidden, embedding, output-head, and gains/biases learning rates, alongside method-specific momentum or Adam coefficients.For MuonIO and PE IO, βI/O is the I/O momentum coefficient and their I/O updates use no second-moment coefficient; Ember instead uses an I/O second-moment coefficient without a first-moment coefficient.Source: Table 6 · p. 15, D EXPERIMENT SETTINGS AND… · p. 15

Shared training settings

  • All runs use C4, batch size 512, sequence length 256, bfloat16, cosine decay to 0.1 of peak learning rate, and 0.1 warmup.The embedding and language-model head weights are untied.Source: D EXPERIMENT SETTINGS AND… · p. 15

Batching

  • Batching differs by model scale: the 60M, 130M, and 1B runs use distinct microbatch, accumulation, and GPU configurations.MuonIO at 1B uses (64, 2, 4), while most 1B runs use (64, 8, 1).Source: D EXPERIMENT SETTINGS AND… · p. 15

Model configuration

  • The 60M, 130M, and 1B models use hidden widths 512, 768, and 2048, with 8, 12, and 24 transformer layers, respectively.The models use 8, 12, and 32 attention heads, and feed-forward widths 1376, 2048, and 5461, respectively.Source: D EXPERIMENT SETTINGS AND… · p. 15

Related papers

Loading 2610.02705v1…