블로그 원문에는 일시 정지하거나 단계를 따라가고 회전할 수 있는 인터랙티브 도표도 있다.

AdamW는 편리한 기본 선택이다. 하지만 행렬이 단순한 숫자 배열이 아니라 변환이라는 점에서 출발하면, Muon의 업데이트 규칙을 직관적으로 유도할 수 있다. 이 글은 경사하강법에서 시작해 행렬의 출력 변화를 제한하는 방법, 최근 기울기를 기억하는 방법, 계산을 줄이는 다항식 근사를 차례로 연결한다.

손실과 기울기에서 업데이트까지

모델은 조정 가능한 매개변수 θ를 가진 함수다. 학습 데이터에서 예측과 목표값의 차이를 손실로 계산하고, 손실을 줄이는 방향으로 가중치를 바꾼다. 언어 모델에서는 앞선 문맥을 바탕으로 다음 토큰의 확률을 예측하며, 실제로 관측된 토큰에 낮은 확률을 부여하면 교차 엔트로피 손실이 커진다. 낮은 손실은 선택한 목적 함수에서의 개선을 뜻하지만, 유용한 목적 함수를 정하는 일은 별개의 문제다.

한 층의 가중치 행렬 W만 바꾸고 나머지는 고정하면, 기울기 행렬 G의 각 원소는 해당 가중치를 조금 바꿨을 때 손실이 어떻게 변하는지를 나타낸다. 경사하강법은 학습률 α를 곱한 음의 기울기 방향으로 이동한다. 다만 이 방향은 현재 위치의 국소 정보일 뿐이며, 이동 폭이 너무 크면 실제 손실이 증가할 수 있다.

행렬 업데이트 D가 손실에 미치는 1차 변화는 기울기와 업데이트의 프로베니우스 내적 ⟨G,D⟩F로 근사할 수 있다. 이 값이 더 음수일수록 예측상 손실이 더 크게 감소한다. 근사는 국소적이므로 업데이트 크기에도 제한을 둬야 한다.

가중치 크기 대신 층의 출력 변화를 제한하기

프로베니우스 노름으로 업데이트 크기를 제한하면 행렬 원소의 변화를 한데 모아 측정한다. 그러나 층의 입력 x에 대해 업데이트가 출력에 주는 변화는 Dx다. 따라서 모든 길이 1인 입력 가운데 가장 큰 출력 변화를 제한하는 스펙트럴 노름을 사용할 수도 있다. 이 제한은 해당 층의 단위 입력당 최악의 출력 변화를 통제하지만, 신경망 전체의 손실 곡률까지 통제하지는 않는다.

이 기준에서 업데이트를 찾는 문제는 다음과 같다.

min_{||D||op ≤ η} ⟨G,D⟩F

특이값 분해(SVD)로 기울기 행렬을 입력 방향, 출력 방향, 각 방향의 크기로 나누면 최선의 업데이트는 각 활성 입력 방향을 대응하는 출력 방향의 반대편으로 길이 η만큼 보낸다. 그 결과 업데이트 방향은 기울기의 활성 특이값을 모두 1로 바꾼 행렬, 즉 극분해 인자 polar(G)가 된다. 여기서 1은 가운데 대각 성분의 값이며, 결과 행렬 자체가 항등행렬이라는 뜻은 아니다.

학습 중에는 미니배치마다 기울기가 달라질 수 있으므로, Muon은 먼저 최근 기울기를 지수 이동 평균 Mₜ으로 합친 뒤 극분해 방향을 구한다. 평균과 극분해는 순서를 바꾸면 결과가 달라진다. 업데이트 규칙의 이상화된 핵심은 Wₜ₊₁ = Wₜ − ηₜ polar(Mₜ)다. 이 규칙은 Mₜ를 기준으로 한 국소 점수를 최적화하므로, 현재 미니배치의 손실을 반드시 낮추는 것은 아니다.

극분해 방향을 다항식으로 근사하기

매 학습 단계마다 SVD를 계산하는 비용을 피하려면 특이 방향을 직접 구하지 않고 특이값만 바꾸는 행렬 연산이 필요하다. X가 SVD 형태를 가질 때 XXᵀX는 특이 방향을 유지하면서 특이값을 세제곱한다. 이를 바탕으로 뉴턴–슐츠 반복을 구성하면 정규화된 행렬의 양의 특이값을 1에 가깝게 만들 수 있다.

고전적인 반복식은 Xₖ₊₁ = 3/2 Xₖ − 1/2 XₖXₖᵀXₖ다. 다만 작은 초기 특이값은 목표에 가까워지는 데 많은 반복이 필요할 수 있다. 원래 Muon 실험에서는 시험한 모델에서 5회의 내부 반복으로 충분했다고 보고했으며, 계산에 맞춰 조정한 다음 다항식을 사용한다.

q(s)=3.4445s−4.7750s³+2.0315s⁵

이 다항식은 작은 값에서 빠르게 키우지만 q(1)=0.701이므로 1을 고정점으로 두지 않는다. 따라서 이 다항식을 계속 반복한다고 수렴이 보장되는 것은 아니다. 다섯 번이라는 설정은 수렴 보장이 아니라 계산 예산에 따른 선택이다.

Muon의 한 업데이트에서는 기울기를 운동량에 반영하고, 운동량 행렬을 안정화 항 ε과 함께 프로베니우스 노름으로 정규화한 뒤 다항식 반복을 다섯 번 수행한다. 결과에 학습률을 곱해 가중치에서 뺀다. 선택적 가중치 감쇠 λ는 기울기 계산과 별도로 기존 가중치를 줄인다. 실제 구현은 현재 기울기를 다시 섞는 네스테로프 운동량, 행렬 모양에 따른 크기 조정, 다른 다항식 일정 등을 추가할 수 있다. 참조 구현은 이런 선택을 보여준다.

비교 결과가 말해주는 범위

같은 기울기 행렬을 적용하면 SGD는 기울기를 그대로 쓰고, AdamW는 원소별 기울기 제곱의 추정치로 각 위치를 따로 조정한다. Muon은 운동량 행렬의 특이 방향을 조정한다. 글의 시각화는 업데이트 규칙이 만드는 차이를 비교하며 학습 속도나 최종 손실을 측정하지 않는다.

좁은 이차 손실 골짜기를 사용한 예에서는 같은 움직임을 행렬의 여러 원소에 섞어 담으면 AdamW의 원소별 척도 추정이 가파른 방향과 완만한 방향의 기여를 함께 반영한다. 이 혼합 표현에서 설정한 학습률로는 Muon이 80단계 동안 골짜기를 더 내려갔지만, 두 방향을 별도 원소에 담으면 AdamW가 더 빨라졌다. 혼합 표현을 유지한 채 AdamW의 초기 학습률을 0.32로 높여도 Muon보다 낮은 손실에 도달했다. 따라서 이 예는 Muon의 보편적 우위를 입증하지 않는다. 또한 Muon은 한 단계에서 더 많은 계산을 하므로, 학습 단계 수만으로 실제 경과 시간을 비교할 수 없다.

예시의 손실에는 최소점이 하나뿐이며, 어떤 경로도 국소 최소점에서 탈출하는 능력을 보여주지 않는다. 기울기가 0이고 저장된 운동량도 없다면 Muon 역시 따라갈 방향이 없다. 운동량이 평탄하거나 방향이 바뀌는 구간을 지나게 할 수는 있지만, 특이값을 1로 바꾸는 것만으로 탈출이 보장되지는 않는다.

이 글의 유도는 손실의 국소 변화, 행렬 출력 변화의 제한, 기울기 평균, 극분해 근사를 연결해 Muon의 핵심 연산을 설명한다. 실제 모델에서 어떤 선택이 효과적인지는 별도의 실험 문제로 남는다.