Muon과 Mamba의 만남: 상태 공간 모델 (State Space Models)을 위한 스펙트럼 최적화 (Spectral

요약

Muon 최적화 도구를 Mamba-2 상태 공간 모델(SSM)에 적용하여 성능을 분석한 연구입니다. 특정 가중치 그룹, 특히 출력 투영(output projection)에 Muon을 적용했을 때 토큰 효율성이 크게 향상됨을 확인했습니다.

핵심 포인트

• Muon은 Newton-Schulz 반복법을 통해 가중치 행렬을 직교화하는 최적화 도구임
• Mamba-2 모델의 출력 투영에 Muon을 적용할 때 가장 높은 성능 향상을 보임
• Muon 적용 시 토큰 효율성이 개선되며, 이는 계산 최적점을 지난 후에도 유지됨
• Muon은 투영의 조건수(condition number)를 낮추는 효과가 있음