Muon Optimizer 정리
원 출처: Muon: An optimizer for hidden layers in neural networks (Keller Jordan)
대규모 확장: Muon is Scalable for LLM Training (Moonshot AI)
모멘텀을 직교화해서 업데이트한다는 Muon의 핵심 아이디어를 수식 중심으로 정리했습니다.
Muon이란
- 신경망의 은닉층 2D 가중치 행렬을 위한 옵티마이저다. AdamW를 대체하거나 보완하는 용도로 설계됐다
- 경사하강법이나 Adam은 모멘텀을 원소 단위로 스케일링하는 반면, Muon은 모멘텀 전체를 직교화해서 업데이트한다
이름 그대로 “MomentUm Orthogonalized by Newton-schulz”의 약자다.
업데이트 규칙
Muon의 업데이트는 세 줄로 끝난다.
\[\begin{aligned} M_t &= \mu M_{t-1} + \nabla L(W_{t-1}) &\quad (1) \\ O_t &= \mathrm{NS5}(M_t) &\quad (2) \\ W_t &= W_{t-1} - \eta_t O_t &\quad (3) \end{aligned}\]모멘텀 누적
$M_t$는 $t$ 시점에서의 모멘텀 값이다. 이전 모멘텀과 loss 그래디언트를 더한 것, 즉 $t$ 시점까지 누적된 그래디언트다.
식 (1)을 펼쳐 보면 지수이동평균(EMA) 형태라는 것이 드러난다.
\[M_t = \nabla L(W_{t-1}) + \mu \nabla L(W_{t-2}) + \mu^2 \nabla L(W_{t-3}) + \cdots\]최근 그래디언트일수록 가중치가 크고, 과거로 갈수록 $\mu$의 거듭제곱만큼 감쇠된다. 여기까지는 일반적인 모멘텀과 같다.
Newton-Schulz 직교화
$O_t$는 Newton-Schulz 5차 반복을 통해 직교 행렬로 근사한 결과다. 이 연산을 통해 모멘텀의 특이값들을 전부 1로 맞춰서, 방향은 유지하고 스케일은 모든 방향에서 균등하게 만든다.
SGD나 Adam 같은 옵티마이저는 특이값 분해(SVD) 연산을 하지 않는다. 즉 $M_t$를 계산하고 그대로 파라미터 업데이트에 사용한다. Muon은 여기에 한 단계를 더 넣는 것이다.
정확히 무엇을 근사하는지 보면, $M_t$의 특이값 분해가 $M_t = U\Sigma V^\top$일 때 다음이 성립한다.
\[(M_t M_t^\top)^{-1/2} M_t = U V^\top\]즉 특이값 행렬 $\Sigma$가 사라지고 $UV^\top$만 남는데, 이것이 바로 특이값을 전부 1로 만든 직교 행렬이다. Newton-Schulz 반복은 이 $UV^\top$를 SVD 없이 근사한다.
가중치 업데이트
최종적으로 $t-1$ 시점의 가중치에서 직교화된 모멘텀에 learning rate를 곱해 빼줌으로써 새로운 가중치를 얻는다. 식 (3)이 이 부분이다.
다른 옵티마이저와의 비교
세 방식의 차이는 “무엇을 빼는가” 하나로 정리된다.
| 옵티마이저 | 업데이트 식 | 차이점 |
|---|---|---|
| 일반 SGD | $W_t = W_{t-1} - \eta_t \nabla L(W_{t-1})$ | 원본 그래디언트를 그대로 사용 |
| SGD + Momentum | $W_t = W_{t-1} - \eta_t M_t$ | 그래디언트 대신 누적 모멘텀 $M_t$ 사용 |
| Muon | $W_t = W_{t-1} - \eta_t O_t$ | 모멘텀 $M_t$를 한 번 더 직교화한 $O_t$를 사용 |
왜 직교화를 하는가
Muon의 목적은 다음 흐름으로 이해하면 된다.
- 실제 신경망의 그래디언트/모멘텀 행렬은 경험적으로 특이값이 불균형하다
- 이걸 교정하려면 특이값을 전부 균일하게(1로) 맞추고 싶다
- 정석대로 하려면 SVD로 정확히 분해해서 특이값만 1로 바꾸면 된다
- 그런데 SVD는 계산 비용이 크고 GPU에서 비효율적이다
- 그래서 Newton-Schulz(NS5) 반복법으로 근사해서 비슷한 효과를 낸다
NS5는 $f(x) = ax + bx^3 + cx^5$ 형태의 홀수차 다항식을 반복 적용하는 방식이다. 계수를 $a = 3.4445$, $b = -4.7750$, $c = 2.0315$로 두면 작은 초기 특이값에 대해서도 빠르게 수렴하도록 고정점이 1 근처에 놓인다. 행렬 곱셈만으로 구성되므로 GPU에서 효율적으로 돌아간다.
직관적으로 보면, 직교화는 업데이트 행렬을 등방적으로 만들어서 가중치가 몇 개의 지배적인 방향으로만 학습되는 것을 막는다.
장점
- 직교화를 통해 모든 방향의 업데이트 크기를 균등화하기 때문에 수렴 속도가 빠르다
- 특정 방향의 그래디언트가 폭주하거나 사라지는 현상이 줄어든다. Muon을 적용하면 attention logit의 스케일이 별도의 clipping 없이도 안정적으로 유지된다
- 대규모 배치 학습에 유리하다
- AdamW 대비 약 2배의 학습 효율 개선
참고: 대규모 학습으로의 확장
원래 Muon은 작은 규모의 언어모델에서 좋은 결과를 보였지만 대규모 확장성은 검증되지 않았다. Moonshot AI의 기술 보고서는 이를 위한 두 가지 기법을 제시한다.
- Weight decay 추가: 규모를 키우면 가중치와 레이어 출력의 RMS가 계속 커져 bf16의 고정밀 범위를 벗어나면서 성능 이득이 줄어든다. AdamW의 weight decay를 도입해 해결한다
- 파라미터별 업데이트 스케일 조정: 하이퍼파라미터 튜닝 없이 바로 쓸 수 있게 하고 학습 안정성도 크게 개선한다
이 두 가지를 적용한 결과, scaling law 실험에서 Muon은 AdamW와 같은 성능에 도달하는 데 약 52%의 학습 FLOPs만 필요했다. 앞에서 말한 “AdamW 대비 약 2배 효율”이 이 수치다. 이를 바탕으로 5.7T 토큰으로 학습한 3B/16B MoE 모델 Moonlight를 공개했다.
댓글남기기