[최적화] Muon Optimizer
이 포스트는 Jia-Bin Huang님의 This Simple Optimizer Is Revolutionizing How We Train AI [Muon] 영상을 참고하였습니다.
기존 Adam optimizer의 한계
1. 메모리 효율성 문제
Adam의 업데이트 공식은 아래와 같으며, 각 파라미터마다 평균 $m$과 분산 $v$ 2개의 변수를 추가로 저장해야 한다.
\[\theta_{t+1}=\theta_t-\eta\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}\]예를 들어 Linear layer에서 $3\times2$ 크기의 가중치 행렬이 있다면, 그와 동일한 크기의 행렬을 저장해야 한다.
\[W=\begin{bmatrix} w_{11}&w_{12}\\ w_{21}&w_{22}\\ w_{31}&w_{32} \end{bmatrix} ~~,~~ M=\begin{bmatrix} m_{11}&m_{12}\\ m_{21}&m_{22}\\ m_{31}&m_{32} \end{bmatrix} ~~,~~ V=\begin{bmatrix} v_{11}&v_{12}\\ v_{21}&v_{22}\\ v_{31}&v_{32} \end{bmatrix}\]결과적으로, Optimizer State가 모델 파라미터 자체보다 2배 더 많은 메모리를 차지하게 된다.
2. Low-Rank 문제
Adam은 Vector-based Optimizer이다. 이러한 방식은 가중치 행렬을 행렬로 보지 않고, 단순히 하나의 긴 벡터로 취급하여 각 원소들을 독립적으로 업데이트한다.
예를 들어, $3\times2$ 가중치 행렬을 단순히 6개의 나열된 숫자로 취급한다.
독립적으로 업데이트하기 때문에 각 파라미터가 서로의 상관관계를 무시한 채 개별적인 보폭으로 업데이트되며, 이는 가중치 간의 스케일 불균형 문제로 이어진다.
실제 파라미터 업데이트 과정에서 특정 방향의 특이값만 커져서 모멘텀 행렬이 Low-rank가 되는 경향성이 나타나는데, 이는 소수의 방향만이 업데이트를 주도한다는 의미이다.
Muon optimizer
Muon은 MomentUm Orthogonalized by Newton-schulz의 줄임말이다.
Muon은 위와 같은 문제를 모멘텀 행렬 $M$을 직교화함으로써 해결한다.
$M$을 직교화함으로써 아래의 이점을 얻을 수 있다.
행렬의 모든 특이값이 1로 정규화
모든 방향에 대해 균일한 보폭이 유지되기 때문에 $v$를 별도로 저장할 필요도 없고, 특이값이 작아 학습에서 소외되었던 방향들도 활성화할 수 있다.
업데이트 방향이 서로 직교
각 파라미터가 중복되지 않는 정보를 학습할 수 있도록 유도할 수 있다.
수학적으로 모멘텀 행렬 $M$을 직교화하는 것은, $M$과 가장 가까운 직교 행렬 $O$를 찾는 것으로 볼 수 있다.
\[\begin{aligned} \text{Ortho}(M)=\underset{O}{\arg\min}\lVert O-M\rVert_F\\ \text{subject to}~OO^\top=O^\top O=I~~~~ \end{aligned}\]위 최적화 문제의 해 $O$는 SVD를 통해 간단히 구할 수 있다.
$M=USV^\top$라고 할 때, $S$의 모든 특이값을 $1$로 만든 행렬 $O=UV^\top$가 최적해가 된다.
그러나 SVD는 입력 행렬이 클수록 계산량이 기하급수적으로 증가하며, 연산 자체가 하드웨어 가속이 효율적이지 않기 때문에, 실제 학습에서 매 업데이트마다 SVD를 계산하는 것은 어렵다.
Newton-Schulz Iteration
이 문제를 해결하기 위해 Muon에서는 SVD를 직접 구하는 대신, 행렬 곱만으로 직교 행렬에 수렴시키는 Newton-Schulz iteration을 사용하였다.
Newton-Schulz Iteration을 간단히 말하면, 반복적인 다항식 연산을 통해 행렬의 특이값을 $1$로 수렴시키는 방식이다.
행렬 곱은 GPU와 같은 하드웨어 가속에 최적화되어 있기 때문에, SVD보다 연산 속도 측면에서 압도적으로 빠르다.
Muon은 직교화를 위해 홀수 차수 다항식 (Odd Polynomial) 형태의 행렬 함수 $\rho(\cdot)$를 사용한다. 이 함수는 행렬을 입력받아 아래와 같은 연산을 수행한다.
\[\rho(X)=aX+b(XX^\top)X\]위 함수 $\rho(\cdot)$는 행렬 전체를 복잡하게 변형시키는 것처럼 보이지만, 사실 SVD의 관점에서 보면 아래 연산과 동치이다.
\[\vphantom{\Big(} \rho(X)=U(aS+bS^3)V^\top\]다시 말해, 함수 $\rho(\cdot)$는 회전을 하는 $U$와 $V^\top$는 변화시키지 않고, $S$의 특이값들에만 개별적으로 적용되는 것으로 볼 수 있다.
\[\rho(X)=U\rho(S)V^\top\]공식 유도
SVD $X=USV^\top$ 대입
\[\rho(X)=aUSV^\top+b\left(USV^\top (USV^\top)^\top\right)USV^\top\]$U$와 $V$의 직교 행렬 성질을 이용해 두 번째 항 정리
\[\left(USV^\top VS^\top U^\top\right)USV^\top =USS^\top SV^\top=US^3V^\top\]정리
\[\rho(X)=aUSV^\top+bUS^3V^\top=U(aS+bS^3)V^\top\]
이러한 원리는 5차 다항식과 같은 모든 홀수 다항식에 적용할 수 있다.
\[\vphantom{\Big(} \rho(X)=aX+b(XX^\top)X+c(XX^\top)^2X~~\rightarrow~~ \rho(X)=U(aS+bS^3+cS^5)V^\top\]이 함수를 이용하면, SVD를 직접 계산하지 않고도 $a,b,c$에 적절한 값을 선택함으로써 특이값을 $1$에 가깝게 만들 수 있다.
예를 들어 3차식에서 $a=1.5,~b=-0.5$인 경우, 아래 애니메이션처럼 입력값 (노란 점)들이 반복 계산을 거쳐 $1$ 근처로 모여드는 것을 확인할 수 있다.
즉, 함수를 적용할수록 왼쪽 그림처럼 그래프가 계단 함수 형태로 변하며 $1$에 수렴한다.
또한 오른쪽 그림처럼 차수를 높이거나 계수를 조절함으로써, 훨씬 적은 반복 횟수만으로도 특이값들을 $1$로 빠르게 밀어낼 수 있다.
경험적으로, 모든 특이값을 완벽하게 $1$로 만들 필요가 없고, $0.7\sim1.3$과 같은 일정한 범위 안에만 머물게 해도 충분하다고 한다.
아래 그림은 특정 구간의 값들을 일정 범위 내로 진동하게 만드는 것을 보여준다.
Muon 알고리즘
Muon의 전체적인 동작 과정은 아래의 Pseudo-code로 요약할 수 있다. $G$는 그라디언트, $B$는 모멘텀 행렬을 의미한다.
초기에 모멘텀 행렬 $B$를 $0$으로 초기화한다.
\[\vphantom{\Big(}B_0=0\]각 학습 step $t=1\dots T$ 마다 아래의 과정을 반복한다.
손실 함수의 기울기를 계산하고 이를 $G$에 저장한다.
\[G_t=\nabla_\theta\mathcal{L}_t(\theta_{t-1})\]계산한 그라디언트를 이용해 모멘텀 행렬을 업데이트한다.
\[\vphantom{\Big(}B_t=\mu B_{t-1}+G_t\]이때, $\mu$는 모멘텀 계수를 의미한다.
모멘텀 행렬에 Newton-Schulze iteration을 5번 적용해 직교 행렬을 얻는다.
\[O_t=\text{NewtonSchulz5}(B_t)\]직교 행렬을 이용해 최종적으로 파라미터를 업데이트한다.
\[\vphantom{\Big(}\theta_t=\theta_{t-1}-\eta O_t\]
Muon Optimizer 적용
실제 모델을 학습할 때, 모든 파라미터에 Muon을 적용하지 않는다.
기본적으로 Muon은 행렬의 구조적 이점을 활용하는 최적화 도구이기 때문에, 2D 이상의 가중치 행렬 (Linear, Conv layer 등)에만 Muon을 사용하고, 1D 파라미터 (Bias, LayerNorm, Embedding 등)에는 여전히 AdamW를 병행하는 하이브리드 방식을 채택한다.
또한, 경험적으로 Linear layer에서 처음과 마지막 layer에서는 AdamW를 사용하고, hidden layer에만 Muon을 적용하는 것이 성능이 더 좋다고 한다.