# Muon Optimizer

With architecture and data held fixed, a checkpoint-only reconstruction attack recovers training images best from Muon-trained models and worst from Adam-trained models, making the optimizer choice a control axis for training-data privacy.

With architecture and data held fixed, a checkpoint-only reconstruction attack recovers training images best from Muon-trained models and worst from Adam-trained models, making the optimizer choice a control axis for training-data privacy.

MuonIO targets the embedding table E and language-model head L that standard Muon excludes, motivating a 1→2 operator norm for E and a 2→∞ operator norm for L, and using the identity ‖L‖_2→∞=‖Lᵀ‖_1→2 to place both matrices in one vocabulary-oriented geometry under a single normalized-vector rule (column normalization for E, row normalization for L); in 1B LLaMA pretraining on C4 it reduces I/O optimizer-state memory by 50% and I/O update FLOPs by about 46% versus Muon while improving validation perplexity.

MuonIO targets the embedding table E and language-model head L that standard Muon excludes, motivating a 1→2 operator norm for E and a 2→∞ operator norm for L, and using the identity ‖L‖_2→∞=‖Lᵀ‖_1→2 to place both matrices in one vocabulary-oriented geometry under a single normalized-vector rule (column normalization for E, row normalization for L); in 1B LLaMA pretraining on C4 it reduces I/O optimizer-state memory by 50% and I/O update FLOPs by about 46% versus Muon while improving validation perplexity.

This work studies the mechanism of Muon's spectral orthogonalization with a tractable factual-recall model: with S subjects and R relations and S>R, gradient flow learns relations before subjects with a learning-time ratio of Θ̃(√(S/R)), while spectral gradient flow reduces that ratio to Θ̃(1) and makes subject and relation errors decay as 1/(T log T) and exp(-poly(T)) respectively; gradient flow and spectral gradient flow are equivariant under orthogonal transformations of token embeddings whereas Sign gradient flow is not, so different orthonormal embeddings can yield no feature separation, a large separation phase, or even a reversed learning order.

This work studies the mechanism of Muon's spectral orthogonalization with a tractable factual-recall model: with S subjects and R relations and S>R, gradient flow learns relations before subjects with a learning-time ratio of Θ̃(√(S/R)), while spectral gradient flow reduces that ratio to Θ̃(1) and makes subject and relation errors decay as 1/(T log T) and exp(-poly(T)) respectively; gradient flow and spectral gradient flow are equivariant under orthogonal transformations of token embeddings whereas Sign gradient flow is not, so different orthonormal embeddings can yield no feature separation, a large separation phase, or even a reversed learning order.

For Muon under momentum accumulation and a tuned finite-step Newton–Schulz update, the work proves that full-batch training of a sufficiently wide two-layer ReLU network reaches any target empirical squared loss ε with high probability when the constant learning rate scales as (1-μ)√ε, with a hitting-time bound of O((1-μ)^-1ε^-1/2) and a sufficient width independent of both target accuracy and momentum.

For Muon under momentum accumulation and a tuned finite-step Newton–Schulz update, the work proves that full-batch training of a sufficiently wide two-layer ReLU network reaches any target empirical squared loss ε with high probability when the constant learning rate scales as (1-μ)√ε, with a hitting-time bound of O((1-μ)^-1ε^-1/2) and a sufficient width independent of both target accuracy and momentum.