클래리파잉 샴푸: 확률성과 파라미터 궤적에 적응하는 스펙트럴 디센트

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

신경망의 행렬 구조를 활용하는 최적화 알고리즘인 Shampoo와 Muon은 Adam 및 Signum과 같은 요소별 알고리즘보다 데이터 효율성이 더 높습니다. 특정 조건에서 Shampoo와 Muon은 Adam과 Signum이 부호 강하(sign descent)로 수렴하는 것처럼 스펙트럴 디센트와 유사하게 동작하지만, 일반적인 관계와 통제된 환경에서의 상대적인 데이터 효율성은 여전히 명확하지 않습니다. 언어 모델에 대한 광범위한 실험을 통해 Shampoo가 Muon보다 더 높은 토큰 효율성을 달성하며, 이는 Adam이 Signum보다 우수한 성능을 보이는 것과 유사함을 보여줍니다. Shampoo의 가중치 행렬에 적용되는 업데이트가 변형된 Muon 업데이트로 분해될 수 있음을 보여주었습니다. 이러한 결과와 일관되게, Shampoo의 장점은 가중치 행렬에 적용될 때만 나타나며, 이는 파라미터의 모양에 무관한 해석에 도전합니다. 이를 통해 스펙트럴 디센트에서 요구되는 반직교성(semi-orthogonality)을 강제하는 대신, Shampoo의 업데이트는 기대값에서 시간 평균화된 반직교성을 갖는다는 새로운 관점을 제시하며, 이는 관련 해석의 단점을 피합니다.

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.