Clarifying Shampoo: Adapting Spectral Descent to Stochasticity and the Parameter Trajectory

Abstract: Optimizers leveraging the matrix structure in neural networks, such as Shampoo and Muon, are more data-efficient than element-wise algorithms like Adam and Signum. While in specific settings, Shampoo and Muon reduce to spectral descent analogous to how Adam and Signum reduce to sign descent, their general relationship and relative data efficiency under controlled settings remain unclear. Through extensive experiments on LLMs, we demonstrate that Shampoo achieves higher token efficiency than Muon, mirroring Adam's advantage over Signum. We show that Shampoo's update applied to weight matrices can be decomposed into an adapted Muon update. Consistent with this, Shampoo's benefits can be exclusively attributed to its application to weight matrices, challenging interpretations agnostic to parameter shapes. This admits a new perspective that also avoids shortcomings of related interpretations based on variance adaptation and whitening: rather than enforcing semi-orthogonality as in spectral descent, Shampoo's updates are time-averaged semi-orthogonal in expectation.

• The paper clarifies that Shampoo adapts spectral descent to stochastic deep learning by leveraging two-sided matrix preconditioning.
• It shows that applying Shampoo exclusively to 2D weight matrices enhances token efficiency compared to elementwise optimizers like Adam and Signum.
• Empirical and theoretical analyses highlight that tuning hyperparameters and utilizing time-averaged semi-orthogonality are key to improved model performance.

Structural Analysis of Shampoo: Adapting Spectral Descent for Stochastic Deep Learning

Introduction and Motivation

This work rigorously examines the mechanics of Shampoo, a class of matrix preconditioned optimizers, in the context of both stochastic optimization and parameter trajectory adaptation. The core motivation is to clarify how Shampoo relates to its non-adaptive counterparts (Spectral Descent, Muon) and elementwise approaches (Adam, Signum), and to delineate where its empirical advantages originate, focusing specifically on data efficiency and update structure.

Analogy: Adam/Signum vs. Shampoo/Muon

The authors construct a formal analogy: Adam is to Signum as Shampoo is to Muon. This analogy operates at both the algorithmic and empirical level. Adam can be written as a sign-based update (Signum) augmented by elementwise adaptation, whereas Shampoo can be viewed as a spectral sign-based update (Muon) with matrix adaptations left and right-multiplied. This duality is visualized and motivated with the geometry of steepest descent under different norms and parameter geometries.

Figure 1: Structural and empirical analogy between Adam/Signum and Shampoo/Muon, exhibiting improved token efficiency by augmenting sign-based methods with adaptation matrices.

Empirical results robustly show that Shampoo outperforms Muon across a broad array of LLM scaling regimes and batch sizes, just as Adam outperforms Signum. All matrix-structured variants, including KL-Shampoo and Shampoo$^{1/2}$, demonstrate superior token efficiency compared to AdamW, especially at larger batch sizes.

Decomposition of the Shampoo Update

Shampoo's update for a weight matrix $W_t$ is decomposed structurally as

$P_t^{-p} G_t Q_t^{-p}$

with $P_t, Q_t$ being EMAs of left and right covariance matrices of the gradient, exponent $p$ typically $1/4$ or $1/2$. The key insight is that this decomposes into a Muon update (matrix-polar normalization of the EMA) sandwiched by two adaptation matrices, paralleling Adam's sign-and-adapt structure. The adaptation matrices relax the semi-orthogonality constraint (enforced by Muon) to account for both stochasticity and the optimizer's evolution along the parameter trajectory.

Empirical Analysis: Localizing Shampoo's Benefits

The work presents strong evidence that the gains of Shampoo are almost entirely localized to its application on weight matrices, not on 1D parameters or embeddings. Adapting the preconditioner shape to the structural role of the parameters—rather than defaulting to shape-agnostic full-matrix variants—yields the strongest empirical results. Notably, preconditioning embeddings or 1D slices yields degraded performance compared to restricting Shampoo purely to 2D weight matrices.

Ablations: One-sided vs Two-sided Preconditioning

The empirical study on one-sided versus two-sided preconditioning further demonstrates that while left-sided (output-activated) preconditioning yields better results than right-sided, the full two-sided update used by Shampoo remains superior in terms of token efficiency. Thus, the classic two-sided structure of Shampoo is not just a legacy artifact but confers clear benefits in high-dimensional neural optimization landscapes.

Adaptation to Both Stochasticity and Parameter Trajectory

A major theoretical contribution is formalizing how Shampoo's adaptation cannot be explained by standard variance adaptation or whitening alone. The EMA statistics in Shampoo induce a time-averaged semi-orthogonality constraint in expectation, generalizing spectral descent for stochastic and nonstationary settings. This result is substantiated by analyzing the limiting behavior of KL-Shampoo and relating its optimality condition to a moving-window matrix whitening problem coupled across training iterates.

Figure 3: In the full-batch regime, adaptive methods like RMSProp and Shampoo$^{1/2}$ still outperform non-adaptive counterparts, indicating the adaptation to parameter trajectory is effective even without stochasticity.

Hyperparameter Trends and Optimization Landscape

The extensive hyperparameter study finds that the optimal betas for Adam and Shampoo are correlated and shrink with batch size, and that the small but highly-tuned $\epsilon$ values materially affect the optimizer ranking. For matrix variants, batch size increases favor Shampoo/Muon, whereas AdamW and Signum's relative performance peaks at smaller batch sizes. These findings clarify many contradictions in prior optimizer benchmark studies.

Figure 2: AdamW validation perplexity over the hyperparameter search-space, highlighting regions of instability and optimality in $(\alpha, \beta_1, \beta_2)$.

Figure 4: Validation perplexity for Signum and Muon as a function of learning rate and momentum, showing Muon's enhanced robustness over the search space.

Figure 5: Validation perplexity across Shampoo variant hyperparameters, indicating marked improvements with proper choice of ($\alpha$, $\beta_1$, $\beta_2$).

Theoretical Implications

The authors' formalism exposes foundational limitations in common interpretations of adaptive methods. Variance adaptation does not capture the update structure; whitening is only meaningful when defined in a matrix sense with time-averaging over the parameter path. Spectral Descent emerges as the limiting case of Shampoo in absence of time-averaging and stochasticity, and KL-Shampoo provides a tractable relaxation that interpolates between these regimes.

Practical Implications and Future Directions

Practically, these findings have direct consequences for the implementation and tuning of Shampoo and related matrix preconditioners:

• Apply Shampoo exclusively to 2D weight matrices; avoid embeddings and 1D parameters.
• Use two-sided preconditioning for maximal data efficiency.
• Tune $\epsilon$ and adaptation hyperparameters carefully; defaults can invert relative optimizer performance.
• Grafting (modulating update magnitude with AdamW) remains beneficial for robust training, especially in large-scale distributed scenarios.

Theoretically, understanding why $S_\infty$ (spectral) geometry outperforms $\ell_\infty$ and $\ell_2$ geometries in large-scale neural network optimization remains open, with possible connections to robustness under heavy-tailed loss surfaces and class/contact imbalance. The unification of time-averaged and stochastic adaptation foreshadows the development of a more general modular duality theory for deep learning optimizers.

Conclusion

This study clarifies the empirical and theoretical basis of Shampoo, identifying its core advantage as an adaptation of spectral descent that accounts for both stochastic gradients and the evolution of the parameter trajectory. The matrix preconditioner learns a time-averaged orthogonalization structure, tailored to neural weight matrices, which cannot be reduced to simple variance adaptation or classic whitening. These insights not only rationalize prior empirical discrepancies but provide a foundation for principled design of future matrix-adaptive optimizers and more robust automated hyperparameter transfer protocols.

---