Muon Optimizer

• Muon은 그래디언트의 이동 평균(momentum)을 직교화하여 파라미터를 업데이트하는 옵티마이저
• 모든 특잇값을 1에 가깝게 만드는 Newton-Schulz 반복 기반 직교화 알고리즘
• AdamW 대비 약 2배의 계산 효율을 보이며, NanoGPT 속도 경쟁에서 압도적 성과를 기록한 차세대 옵티마이저

해당 개념이 필요한 이유

• SGD-momentum과 Adam의 업데이트는 Transformer 2D 파라미터에서 condition number가 매우 높음 (거의 low-rank)
• 이는 소수의 노이즈 방향이 최적화를 지배하게 만드는 문제
• 직교화는 “크기는 작지만 학습에 중요한 희소 방향”의 스케일을 효과적으로 키워줌

AS-IS: Adam/SGD-momentum의 업데이트

gradient → momentum → update

업데이트 행렬의 특잇값 분포가 극심하게 편중됨. 몇 개의 큰 특잇값이 업데이트를 지배하고, 나머지 방향은 무시됨.

TO-BE: Muon의 업데이트

gradient → momentum → Newton-Schulz 직교화 → update
                      (모든 singular value → 1)

방향 정보는 보존하면서 크기를 정규화. 모든 방향이 동등한 스케일로 업데이트에 기여.

핵심 작동 원리

1. 그래디언트 계산
2. Momentum 적용 (이동 평균)
3. Newton-Schulz 반복으로 momentum 행렬 직교화 → 모든 singular value를 1에 가깝게
4. 직교화된 업데이트로 파라미터 갱신

기존 Orthogonal-SGDM(Tuddenham 2022)은 직교화 후 momentum을 적용했지만, Muon은 momentum 후 직교화 순서로 바꾸어 경험적으로 더 좋은 성능을 달성했다. 또한 SVD 대신 Newton-Schulz 반복을 사용해 계산 비용을 절감했다.

autoresearch에서의 MuonAdamW 하이브리드

autoresearch의 train.py는 파라미터 유형별로 옵티마이저를 분리한다:

파라미터 유형	옵티마이저	Learning Rate
Embedding 레이어	AdamW	0.6
Unembedding (lm_head)	AdamW	0.004
레이어별 스칼라	AdamW	0.5
2D 행렬 (어텐션/MLP)	Muon	orthogonalization 기반

여기에 Cautious weight decay를 추가: 그래디언트와 파라미터의 곱이 ≥ 0일 때만 weight decay를 적용하여 불필요한 정규화를 방지한다.

스케일링과 최신 동향

• GLM-4.5 (355B), KIMI Moonshot (1T+) 등 초대형 모델 학습에 이미 실전 배치
• Moonlight: 3B/16B MoE 모델을 5.7T 토큰으로 학습, 기존 대비 훨씬 적은 FLOPs로 동등 성능
• Turbo-Muon: spectral preconditioning으로 Newton-Schulz 단계 계산 비용 절감
• AdaMuon: element-wise adaptivity + 직교 업데이트 결합, 대규모에서 Adam 대비 40%+ 효율 향상
• Block-wise Orthogonalization (ICML 2025): 행렬을 독립 타일로 분할 후 개별 직교화, 16× tensor parallel 가능

참고 문서

• Muon: An optimizer for hidden layers — Keller Jordan blog
• Muon is Scalable for LLM Training (arXiv)
• The Muon Optimizer Explained: Why Orthogonal Gradients Work
• Towards Understanding Orthogonalization in Muon (ICML 2025)