Clarifying Shampoo: Adapting Spectral Descent to Stochasticity and the Parameter Trajectory
This work clarifies optimizer behavior for machine learning practitioners, but it is incremental as it builds on existing methods like Shampoo and Muon.
The paper tackled the unclear relationship and data efficiency between Shampoo and Muon optimizers in neural networks, showing through experiments on language models that Shampoo achieves higher token efficiency than Muon, mirroring Adam's advantage over Signum.
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 language models, 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.