Technical Summary: Muonp – Muon with Fractional Spectral Powers

Problem Statement

The Muon optimizer has emerged as a powerful tool for training deep neural networks, particularly large language models (LLMs). It operates by replacing the gradient matrix $G = USV^\top$ with its polar factor $UV^\top$, effectively flattening the singular value spectrum. This "spectrum flattening" promotes balanced learning across singular directions and prevents dominant modes from monopolizing training updates.

However, the authors identify a critical limitation in Muon's approach: full flattening discards all singular-value information. While this uniformity is beneficial for broad feature discovery during pretraining, it may be suboptimal for adaptation tasks (such as finetuning) where the magnitude of singular values encodes which residual directions are most worth changing. The central question addressed by this work is: Can we preserve Muon's matrix-level spectral benefits while retaining the magnitude structure necessary for effective adaptation?

Methodology

Theoretical Formulation: Muonp

The authors introduce Muonp, a generalization of the Muon optimizer that updates weights using fractional spectral powers:
$$\text{Update} \propto US^pV^\top$$
where $p \in (0, 1)$ is a rational number.

• Geometric Interpretation: The authors prove (Theorem 2.1) that this update rule corresponds to the steepest-descent direction under the Schatten $q$-norm, where $q = 1 + 1/p$.
• Interpolation:
  • $p = 0$: Corresponds to standard Muon ($UV^\top$), maximizing uniformity.
  • $p = 1$: Corresponds to standard gradient descent ($USV^\top = G$).
  • $0 < p < 1$: Interpolates between the two, retaining partial magnitude information.

Algorithmic Challenge and Solution

Computing $US^pV^\top$ directly requires Singular Value Decomposition (SVD), which is computationally prohibitive for large models. The authors address the challenge of approximating this update without explicit SVD:

1. Impossibility of Univariate Iteration: The authors prove (Lemma 2.2) that no fixed univariate polynomial iteration $S_{n+1} = f(S_n)$ can converge to $S^p$ for arbitrary rational $p$. Unlike the polar factor ($p=0$), fractional powers require the algorithm to "remember" the original matrix.
2. Bivariate Polynomial Recurrence: To solve this, the authors derive low-degree odd bivariate recurrences. The algorithm iteratively computes matrices $G_n = US_nV^\top$ where the singular values evolve via $S_{n+1} = f(S_n, S)$.
   • Theorem 2.3 demonstrates that for any odd bivariate polynomial $f(x, y)$, the term $Uf(S_n, S)V^\top$ can be computed using only matrix multiplications and transposes, preserving the computational structure of Muon.
   • Theorem 2.4 establishes the existence of such polynomials that converge to $S^p$ for any rational $p \in (0, 1)$.
3. Explicit Construction: The paper provides explicit formulas for these polynomials (e.g., $f(x, y) = x + c(y - x^3)$ for $p=1/3$). The coefficients are chosen to ensure $S^p$ is an attracting fixed point.

Implementation Details

• Complexity: The computational complexity remains $O(\omega k N)$, where $\omega$ is matrix multiplication complexity, $k$ is the number of multiplications per step (dependent on polynomial degree), and $N$ is the number of iterations. Muonp differs from Muon only by a constant factor in the number of matrix multiplications.
• Stability: The method is numerically stable when implemented in bfloat16, making it suitable for modern hardware.
• Curriculum Compatibility: Since Muon ($p=0$) is a special case of Muonp, switching between them during training (e.g., pretraining with Muon and finetuning with Muonp) requires no state reset, only a change in the polynomial approximation.

Key Contributions

1. Muonp Optimizer: Introduction of a generalized optimizer computing $US^pV^\top$ for any rational $p \in (0, 1)$ without explicit SVD.
2. Theoretical Limits: Proof that fractional spectral powers cannot be computed via fixed univariate polynomial iterations, necessitating a novel bivariate approach.
3. Constructive Algorithm: Derivation of low-degree odd bivariate recurrences that approximate fractional powers while maintaining the matrix-multiplication-only structure of Muon.
4. Empirical Validation: Demonstration that Muonp significantly outperforms Muon and AdamW in finetuning scenarios across math reasoning, coding, and downstream tasks.
5. Spectral Geometry Analysis: Insight that the optimal spectral geometry is regime-dependent: uniform updates ($p=0$) benefit pretraining, while partial preservation ($0 < p < 1$) benefits finetuning.

Experimental Results

Finetuning Performance

Muonp demonstrates consistent improvements over Muon and AdamW when finetuning Llama3-1B and Llama3-3B models:

• Math & Code Reasoning: On GSM8K, Numina, and OpenCodeInstruct, Muonp improves pass@8 accuracy and reduces loss compared to Muon.
• Perplexity: Muonp achieves lower validation perplexity on Fineweb and Pile of Law datasets.
• Downstream Tasks: Muonp yields higher accuracy across a suite of reasoning benchmarks (PIQA, Winogrande, LAMBADA, etc.) compared to both Muon and AdamW.
• Optimizer Mismatch: Even when pretraining with Muon (matching the optimizer), Muonp outperforms Muon during finetuning, suggesting the gain is fundamental to the spectral geometry rather than just optimizer consistency.

Pretraining vs. Finetuning Dichotomy

The paper reveals a distinct dichotomy in performance based on the training phase:

• Pretraining: Muon ($p=0$) outperforms Muonp. The authors argue that pretraining benefits from "isotropization" to discover broad features and prevent dominant directions from monopolizing updates.
• Finetuning: Muonp ($0 < p < 1$) outperforms Muon. In this regime, the model has a low intrinsic dimension, and retaining singular value information helps focus adaptation on the most relevant residual directions.
• Curriculum Experiments: Switching from Muon to Muonp in the final steps of pretraining (a curriculum) yields immediate loss reductions, validating the hypothesis that increasing $p$ adapts the effective learning rate for smaller singular directions.

Runtime

Muonp maintains a runtime comparable to Muon. For instance, with $p=1/3$ (cubic recurrence), Muonp is slightly faster than Muon (which uses a quintic recurrence for $p=0$ in some implementations), with step times around 100–120ms on H100 GPUs.

Significance and Claims

The paper positions Muonp as a principled generalization of the Muon optimizer. Its primary significance lies in:

1. Bridging the Gap: It offers a systematic way to interpolate between the uniformity of Muon and the magnitude sensitivity of gradient descent, resolving the trade-off between "spectrum flattening" and "magnitude retention."
2. Regime-Dependent Optimization: It provides evidence that the "best" spectral geometry is not universal; pretraining requires full flattening, while finetuning requires partial flattening.
3. Practical Efficiency: It achieves these theoretical benefits without incurring the cost of SVD, maintaining the hardware-friendly, matrix-multiplication-only structure essential for scaling to billion-parameter models.
4. Future Directions: The authors suggest that automatically selecting the optimal $p$ for specific tasks and integrating Muonp into pretraining curricula are promising avenues for future research.

The work concludes that preserving the singular spectrum can bring significant gains in adaptation scenarios, and Muonp provides a computationally efficient, principled mechanism to achieve this.