[ICLR 2025] HTMuon: Beyond Orthogonality with Heavy-Tailed Spectral Correction

TL;DR

HTMuon is a breakthrough in matrix-based optimization that fixes a critical flaw in the popular Muon optimizer: the "all-ones" singular value update. By applying a power-law correction to the momentum spectrum, HTMuon creates heavy-tailed updates that significantly improve LLM pretraining efficiency and generalization, achieving a ~1.0 PPL reduction on LLaMA models over the original Muon.

The "All-Ones" Fallacy: Why Orthogonality Isn't Enough

Recent LLM training breakthroughs (like Moonshot's Kimi K2 and GLM-4.5) have pivoted toward Muon, a matrix-based optimizer. Muon's secret sauce is its orthogonalization step: it takes the momentum matrix and forces all its singular values to be exactly 1.0.

While this stabilizes early training, the authors of HTMuon expose a hidden cost:

1. Noise Amplification: Setting all singular values to 1 treats small, noise-dominated directions with the same "importance" as the principal signal directions.
2. Spectral Suppression: It prevents the model from developing heavy-tailed empirical spectral density (ESD). According to Heavy-Tailed Self-Regularization (HT-SR) theory, the best-performing models naturally evolve heavy-tailed weights. By forcing a flat spectrum, Muon inadvertently limits the model's final capacity.

Methodology: The Power of $p$

The design of HTMuon is elegant and grounded in mathematical intuition. Instead of Muon’s $O_t=U_t{V}_t^{o}p$ (which is a power of 0) or SGDM's $O_t=M_t$ (a power of 1), HTMuon operates in the optimal middle ground:

$O_t=U_t{\Sigma }_t^p{V}_t^{o}p,p\in (0,1)$

By choosing $p=0.125$, the optimizer retains the matrix-based ability to capture parameter correlations while allowing the "signal" components (large singular values) to lead the update slightly more than the "noise" (small singular values).

Figure: The Spectrum Comparison. Note how Muon NS implicitly outperforms Muon SVD because it fails to reach a perfectly flat spectrum, hinting that a "perfect" all-ones spectrum is actually suboptimal.

Theoretical Depth: Schatten-q Norms

The authors provide a rigorous foundation by showing that HTMuon is the explicit solution to steepest descent under a Schatten-q norm constraint. This generalizes Muon (which corresponds to Schatten-$\infty$) and places HTMuon in a broader class of mathematically sound "Proximity Operators."

Experimental Results: Setting New SOTAs

The performance gains are remarkably consistent across tasks:

1. LLaMA Pretraining: On the C4 dataset, HTMuon reduces PPL by 0.98 for 135M models.
2. Larger Scales: In LLaMA-1B experiments, HTMuon outperforms Muon (14.17 vs 14.33 PPL).
3. Image Classification: For ResNet-50 on CIFAR-100, HTMuon reaches 80.16% accuracy, surpassing Muon's 79.85%.

Table: Comparison across LLaMA scales. HTMuon shows widening leads as the model size increases.

Speed vs. Performance: The HTMuon-NS Variant

To address the computational cost of SVD, the authors introduced HTMuon-NS. By using an iterative Newton-Schulz root routine to approximate the fractional power, they achieved $p$-power updates with only minimal overhead (approx. 6-11% additional wall-clock time), which is quickly recovered by faster convergence.

Critical Analysis & Takeaways

HTMuon represents a shift from "stability at all costs" to "stability through correct inductive bias."

Core Insights:

• Generalization is in the Tail: The layer-wise analysis confirms that HTMuon produces models with smaller $\alpha$ (PL exponent), meaning more heavy-tailed weights and better generalization.
• Plug-and-Play: HTMuon can be added to other Muon variants (like NorMuon) for additive gains.
• Limitations: While the $p=0.125$ default is robust, the computational SVD overhead is still non-zero. Future hardware-accelerated SVD or deeper Newton-Schulz optimizations will be key to making this the default for 100B+ parameter models.

In conclusion, HTMuon proves that the "shape" of our updates matters just as much as their magnitude. By respecting the natural heavy-tailed tendencies of neural networks, we can train smarter, not just longer.