Muon과 Mamba의 만남: 상태 공간 모델 (State Space Models)을 위한 스펙트럼 최적화 (Spectral
요약
Muon 최적화 도구를 Mamba-2 상태 공간 모델(SSM)에 적용하여 성능을 분석한 연구입니다. 특정 가중치 그룹, 특히 출력 투영(output projection)에 Muon을 적용했을 때 토큰 효율성이 크게 향상됨을 확인했습니다.
핵심 포인트
- Muon은 Newton-Schulz 반복법을 통해 가중치 행렬을 직교화하는 최적화 도구임
- Mamba-2 모델의 출력 투영에 Muon을 적용할 때 가장 높은 성능 향상을 보임
- Muon 적용 시 토큰 효율성이 개선되며, 이는 계산 최적점을 지난 후에도 유지됨
- Muon은 투영의 조건수(condition number)를 낮추는 효과가 있음
Muon은 Newton-Schulz 반복법 (iteration)을 통해 각 가중치 행렬 (weight matrix)에 대한 업데이트를 직교화 (orthogonalize)하는 최신 최적화 도구 (optimizer)로, 스펙트럼 노름 (spectral norm) 하에서 가장 가파른 경사 하강법 (steepest descent)을 수행합니다. Muon에 대한 거의 모든 증거는 Transformer 모델에서 도출되었으며, 상태 공간 모델 (State Space Models)에서의 동작은 거의 보고된 바가 없습니다. 우리는 Muon으로 훈련할 가중치 그룹 (weight groups)만을 변경하는 통제된 프로토콜 하에서 Mamba-2 130M을 대상으로 Muon과 AdamW를 비교합니다. 그 이점은 국소적 (localized)입니다. 출력 투영 (output projection)에만 Muon을 적용하는 것이 입력 투영 (input projection)이나 양쪽 모두에 적용하는 것보다 성능이 뛰어납니다. 이 장점은 주로 토큰 효율성 (token efficiency) 측면에서 나타납니다. 이는 두 개의 코퍼스 (corpora)와 두 개의 토큰 예산 (token budgets)에서 유지되며, 계산 최적점 (compute-optimal point)을 훨씬 지나 훈련이 계속될 때도 지속됩니다. 컨디셔닝 (Conditioning)은 이 이득을 설명하지 못합니다. Muon은 훈련하는 투영의 조건수 (condition number)를 낮추지만, 더 나은 조건수를 가진 입력 투영이 도움이 되는 것은 아닙니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.LG의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기