샴푸와 뮤온: 스펙트럴 하강의 적응과 토큰 효율성 비교

초록

본 논문은 행렬 구조를 활용하는 최적화기인 Shampoo와 Muon을 이론적·실험적으로 비교한다. Shampoo가 Muon보다 토큰 효율성이 높으며, 이는 Adam이 Signum보다 우수한 이유와 유사한 메커니즘—즉, 행렬 형태의 적응(왼·오른 변환)과 스펙트럴 사인 연산의 결합—에 기인한다는 것을 보인다. 또한 Shampoo 업데이트가 시간 평균된 반직교성을 기대값으로 갖는다는 새로운 해석을 제시한다.

상세 분석

이 논문은 최근 주목받는 행렬 기반 최적화기인 Shampoo와 Muon을 기존의 요소별 적응형 옵티마이저인 Adam과 Signum과 구조적으로 대응시키는 새로운 프레임워크를 제시한다. 먼저 Adam‑Signum 관계를 재조명하면서, Adam의 업데이트를 “적응(스케일링) ⊙ sign” 형태로 분해하고, 적응 부분이 토큰 효율성을 크게 향상시킨다는 점을 강조한다. 이와 유사하게 Shampoo와 Muon 사이에서도 Shampoo 업데이트를 Muon(행렬 사인) ⊙ 양측 적응 행렬(L, R) 형태로 분해한다는 식(10)을 도출한다. 여기서 L과 R은 각각 왼쪽·오른쪽 Kronecker‑팩터의 역제곱근(또는 p‑제곱근)이며, 이는 Adam에서의 element‑wise variance‑adaptation에 해당한다.

이론적 분석에서는 스펙트럴 GD( SpectralGD )를 정의하고, β₁=β₂=ε=0, p=¼(또는 ½)일 때 Shampoo가 정확히 SpectralGD, 즉 G의 폴라 분해에서 얻어지는 반직교 행렬 U Vᵀ와 동일함을 보인다. 따라서 Shampoo는 Muon에 비해 두 개의 추가 적응 행렬을 통해 스케일링을 보정함으로써, stochastic gradient의 변동성을 더 효과적으로 억제한다. 논문은 또한 KL‑Shampoo와 기존 Shampoo(1/4, 1/2) 사이의 차이를 KL‑다이버전스 최소화 관점에서 설명하고, KL‑Shampoo가 기대값에 대한 최적화된 적응을 제공하지만 배치 크기에 민감함을 실험적으로 확인한다.

실험 부분에서는 Llama‑3 기반 언어 모델을 다양한 토큰 예산(1×, 8× Chinchilla), 모델 규모(320 M, 1.5 B), 배치 크기(64, 256)에서 훈련시킨다. 모든 옵티마이저는 동일한 코드베이스와 grafting 전략(Adam 기반 스케일링)으로 구현되어, 차이가 순수히 알고리즘적 차이에 기인함을 보장한다. 결과는 Shampoo(특히 1/2 버전)가 Muon보다 평균 5‑30% 적은 검증 perplexity를 달성함을 보여준다. 배치 크기가 커질수록 행렬 기반 옵티마이저의 최적 배치 규모가 더 크게 나타나, 기존 연구와 일치하게 “critical batch size”가 높아짐을 확인한다. 또한, Muon에 다양한 레이어‑와이드 스케일링(클래식, Moonlight)을 적용했지만, grafting이 가장 좋은 성능을 제공한다는 점을 강조한다.

마지막으로, Shampoo 업데이트가 시간 평균된 반직교성을 기대값으로 갖는다는 새로운 해석을 제시한다. 이는 기존의 “분산 적응”이나 “화이트닝” 해석과 달리, 행렬 구조 자체가 스펙트럴 사인 연산을 통해 자연스럽게 반직교성을 유지하면서도 EMA를 통해 stochasticity를 보정한다는 의미이다. 따라서 Shampoo의 장점은 행렬 형태에만 국한된 것이 아니라, 적응 행렬 L, R이 제공하는 스케일링 메커니즘에 크게 의존한다는 결론에 도달한다.