Muon 옵티마이저 완전 정리: Adam을 넘어서는 행렬 기반 학습법과 대규모 모델 확장·안정화 기법

핵심 메시지

• Adam은 가중치를 좌표 하나하나로만 보고 행렬 구조를 무시하는데, 바로 그 지점에 비효율이 숨어 있다.
• Muon은 모멘텀 업데이트의 특잇값(스펙트럼)을 모두 1에 가깝게 맞춰(직교화) 모든 방향으로 고르게 학습을 진행한다.
• 실제 계산은 값비싼 SVD 대신 Newton-Schulz 반복(행렬 곱만 사용)으로 근사한다.
• 가중치 감쇠와 RMS 매칭을 더하면 소규모 실험용 옵티마이저가 대규모 학습에서도 그대로 작동한다.
• Moonlight 실험에서 Muon은 같은 손실을 AdamW의 약 52% 연산량으로 달성했다.

쉽게 이해하기

학습이란 무작위 가중치에서 쓸모 있는 모델로 옮겨가는 과정이며, 그 한 걸음의 크기와 방향을 정하는 것이 옵티마이저다. 지난 10년간 그 역할은 거의 언제나 Adam 또는 AdamW가 맡아 왔다. Adam은 가중치마다 학습률을 독립적으로 조정하고, 각 값의 기울기가 얼마나 흔들리는지를 추적해 일관된 기울기에는 과감한 걸음을, 노이즈가 큰 기울기에는 조심스러운 걸음을 준다. 그래서 별다른 튜닝 없이도 잘 작동한다.

문제는 가중치가 원래 행과 열, 구조를 가진 2차원 행렬이라는 점이다. Adam은 이 행렬을 그저 독립적인 숫자들의 나열로 보고 좌표별로만 크기를 다시 맞춘다. 행렬을 벡터에 적용하면 어떤 방향은 크게 늘리고 어떤 방향은 눌러 타원이 되는데, 이 늘림과 눌림의 정도가 특잇값(SVD의 Σ)이다. 실제 신경망의 모멘텀 행렬은 이 특잇값 분포가 크게 치우쳐 있어, 매 업데이트가 소수의 지배적인 방향으로만 세게 밀고 나머지 방향은 거의 건드리지 못한다.

Muon의 아이디어는 이 치우친 스펙트럼을 먼저 '고르게' 만드는 것이다. 모든 특잇값을 1로 맞춰 타원을 완전한 원으로 바꾸면 모든 방향이 같은 크기의 힘을 받는다. 수학적으로는 U·Vᵀ에 해당하는 극분해(polar factor)로, 방향(U, V)은 그대로 두고 치우친 크기(Σ)만 버린다. 이름은 'Newton-Schulz로 직교화한 모멘텀(Momentum orthogonalized by Newton-Schulz)'이며, 가중치 자체가 아니라 업데이트 걸음만 이렇게 처리한다.

정확한 극분해는 매 스텝 SVD를 요구해 너무 비싸므로, Muon은 홀수 5차 다항식을 약 다섯 번 반복 적용하는 Newton-Schulz 방법으로 근사한다. 이 다항식은 특잇값을 1 근처의 좁은 띠로 모으면서 방향 U, V는 건드리지 않고, 필요한 연산은 현대 가속기가 특히 빠르게 처리하는 행렬 곱뿐이라 bfloat16으로도 돌릴 수 있다. 또한 Muon은 은닉층의 큰 2차원 행렬에만 적용하고 임베딩·출력층·정규화·편향 같은 1차원 파라미터는 여전히 AdamW가 맡는 하이브리드다.

대규모에서는 두 가지 보정이 필요했다. 긴 학습에서 가중치의 스펙트럼 노름이 폭주하지 않도록 가중치 감쇠를 넣고, Muon 업데이트의 크기를 AdamW의 전형적 크기에 맞추는 RMS 매칭을 더하는 것이다. 이렇게 하면 하이퍼파라미터를 다시 튜닝하지 않아도 대규모에서 그대로 작동한다. Moonshot AI의 Moonlight 실험(총 160억·활성 30억 파라미터 MoE, 5.7조 토큰)에서 Muon 곡선은 AdamW 아래에 놓였고, 같은 손실 목표에 도달하는 데 약 52%의 연산으로 충분했다.

주요 인사이트

• Adam이 '고장 난' 것이 아니라 애초에 행렬 단위로 생각하지 않기 때문에 구조적 비효율을 보지 못한다는 점이 핵심이다.
• 옵티마이저의 개선이 곧 '방향을 고르게 탐색하는 것'이라는 관점은, 학습된 가중치의 특잇값 엔트로피가 더 높아진다는 관찰과도 일치한다.
• 프런티어 규모에서 어텐션의 Q·K 행렬을 그대로 직교화하면 로짓이 폭주해 소프트맥스가 붕괴하는데, 이때는 오히려 Muon이 AdamW보다 불안정해진다.
• 해결책 QK-clip은 스텝마다 각 헤드의 최대 로짓을 확인해 임계값(100)을 넘는 헤드의 Q·K 가중치만 골라 축소하는, 헤드 단위의 정밀한 사후 보정이다.
• 2024년 speed-run의 호기심거리였던 Muon이 2026년 NVIDIA Megatron Core의 정식 지원과 프런티어 모델 학습에까지 이르렀다는 점은 연구 아이디어의 확장 속도를 보여준다.

자주 묻는 질문

관련 AI 소식