Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics

기술 요약: 행렬 분해에서 느린 안장점-안장점 역학 없이 균형 잡힌 솔루션을 학습하는 Muon

문제 정의
본 논문은 행렬 분해 $L(P, Q) = \frac{1}{2}\|M^* - PQ^\top\|_F^2$를 최소화하는 맥락에서 Muon 옵티마이저의 학습 역학을 조사한다. 행렬 분해는 비선형 파라미터 역학 및 표현 학습을 연구하기 위한 최소 모델로서, 저계수 편향(low-rank bias) 및 표현 정렬(representation alignment)과 같은 딥러닝 관련 동작을 보여준다. 저자들은 특히 Muon이 표준 경사 하강법(GD)과 비교하여 어떻게 고유한 업데이트 규칙을 통해 최적화 궤적을 변화시키는지, 즉 초기화 시의 수렴 속도, 안정성, 그리고 작은 무작위 초기화로부터 찾아내는 솔루션의 성질을 이해하고자 한다.

방법론
본 연구는 그래디언트를 학습률에 적용하기 전에 직교화(orthogonalize)하는 Muon 업데이트 규칙에 초점을 맞춘다. 가중치 행렬 $W$에 대해, 업데이트는 $\Delta W = -\eta \cdot \widehat{\nabla_W L}$이며, 여기서 $\widehat{G} = UV^\top$는 그래디언트 $G=USV^\top$의 극분해(polar decomposition, 직교화)이다. 이는 모든 업데이트가 $\eta$로 스케일링된 (반)-직교 행렬임을 보장한다.

저자들은 다음 세 가지 관점을 통해 역학을 분석한다:

1. 스칼라 감소(Scalar Reduction): 가중치가 정렬되어 있다고 가정할 때, 문제는 잔차 $r = m^* - pq$가 업데이트를 결정하는 독립적인 스칼라 문제들로 분리된다.
2. 정렬 역학(Alignment Dynamics): 저자들은 서로 다른 파라미터 체제(과소 매개변수화, 정확한 매개변수화, 과잉 매개변수화) 하에서 $P$와 $Q$의 특이 벡터들이 서로 어떻게 정렬되는지(내부 정렬), 그리고 타겟 행렬 $M^*$과 어떻게 정렬되는지(타겟 정렬)에 대한 이론적 비율을 도출한다.
3. 보존량(Conserved Quantities): 논문은 Muon과 GD의 흐름(flow) 동안 보존되는 특정 행렬량들을 식별한다.

주요 기여 및 결과

1. 느린 안장점-안장점(Saddle-to-Saddle) 역학의 회피
작은 초기값으로부터 시작하는 표준 경사 하강법은 "안장점-안장점" 역학을 보인다. 특이값들이 곱셈적으로 성장하기 때문에, 작은 모드(mode)는 큰 모드보다 느리게 성장한다. 이로 인해 옵티마이저는 저계수 안장점 근처에 머물며 모드들을 순차적으로 학습한다(단계적 학습).
반면, Muon은 **가법적 역학(additive dynamics)**을 보인다. 일단 정렬이 이루어지면, 가중치 행렬의 모든 특이값은 현재 크기와 관계없이 일정한 비율($\eta$ per step)로 성장한다. 결과적으로, 작은 모드가 큰 모드보다 먼저 수렴하며, 옵티마이저는 중간 안장점 없이 $M^*$의 모든 상위 모드를 동시에 학습한다.

2. 임계 학습률을 넘어서는 안정성
경사 하강법은 국소 손실 날카로움(특히 $2/\lambda_{\max}(\nabla^2 L)$)에 의해 결정되는 임계 학습률 임계값에 의해 제약을 받는다. 이를 초과하면 발산하며, 수렴 속도는 문제의 조건수(condition number)에 의해 제한된다.
Muon은 학습률이 이 임계 임계값을 초과하더라도 안정성을 유지한다. 업데이트 크기가 직교화 단계에 의해 고정되어 있기 때문에 반복값(iterates)이 발산할 수 없다. 대신, 큰 학습률을 사용할 경우 최적화는 솔루션 매니폴드로 빠르게 진행되어 그곳에서 진동한다. 이는 학습률을 조건수로부터 분리하여, 지수적(기하적) 학습률 어닐링을 통한 빠른 수렴을 가능하게 한다.

3. 구별되는 보존량 및 균형 잡힌 솔루션
두 옵티마이저 모두 소멸하는 초기값으로부터 "균형 잡힌 솔루션"($P^\top P = Q^\top Q$)을 찾지만, 서로 다른 양을 보존한다:

• 그래디언트 흐름(Gradient Flow): 행렬량 $\Delta_2 = P^\top P - Q^\top Q$를 보존한다. 파라미터 공간에서의 흐름선은 쌍곡선(hyperbolic) 형태를 띤다.
• Muon 흐름(Muon Flow): 행렬량 $\Delta_1 = \sqrt{P^\top P} - \sqrt{Q^\top Q}$를 보존한다. 흐름선은 아핀(affine, 평행한 광선) 형태를 띤다.
  저자들은 Muon의 흐름이 해밀토니안 $H(P, Q) = \text{Tr}(\Delta_1) = \|P\|_* - \|Q\|_*$ (핵 노름의 차이)에 의해 생성된다고 도출하였다. 이 보존 법칙은 가중치가 균형 잡힌 상태로 시작되면 균형 잡힌 상태를 유지하도록 보장하며, 균형 잡힌 솔루션으로 이끈다.

4. 정렬 역학 및 비율
논문은 "침묵 정렬(silent alignment)" 단계(손실이 평탄한 동안 특이 벡터들이 정렬되는 단계)에 대한 정렬 비율을 도출한다:

• 과소 매개변수화 ($d < n$): 내부 정렬은 자명하며, 타겟 정렬은 상대적 스펙트럼 갭 $s$에 대해 $(t/\alpha_{\text{eff}})^{-2(1-s^{-1})}$로 다항식적으로 감소한다.
• 과잉 매개변수화 ($d > n$): 타겟 정렬은 자명하며, 내부 불일치는 $(t/\alpha_{\text{eff}})^{-4}$로 빠르게 감소한다.
• 일반적 상황: 정렬은 더 작은 초기화 스케일에서 더 일찍 발생한다. 저자들은 이러한 이론적 비율이 일반적인 설정에서의 경험적 정렬을 예측함을 보여준다.

5. 스파이크형 학습률 스케줄(Spiked Learning Rate Schedule)
정렬 역학에 대한 이해를 활용하여, 저자들은 정확히 두 번의 최적화 단계 내에 거의 완벽한 정렬을 달계하는 특정 학습률 스케줄을 구성한다 (for $d \ge n$):

1. 1단계: $\eta^{(0)} = \alpha$ (초기화 스케일)로 설정하여 $P$를 $Q$와 정렬시킨다.
2. 2단계: 매우 큰 스파이크 $\eta^{(1)} \gg \eta^{(0)}$를 설정하여 $P$를 $M^*$과 동시에 정렬시키고 가중치를 솔루션 매니폴드로 몰아넣는다.
3. 어닐링: 선형 수렴을 달성하기 위해 학습률을 기하학적으로 어닐링한다.
   이 스케줄은 흐름에 대한 이론적 통찰이 어떻게 실질적인 최적화 전략으로 변환될 수 있는지를 보여주지만, 저자들은 이것이 특정 체제에 대한 개념 증명(proof-of-concept)임을 언급한다.

의의 및 주장
본 논문은 Muon이 곱셈적이고 안장점이 많은 역학을 균등한 가법적 성장으로 대체함으로써 행렬 분해의 최적화 지형을 근본적으로 변화시킨다고 주장한다. 이를 통해 Muon은 다음과 같은 이점을 가진다:

• 타겟 행렬의 모든 모드를 동일한 속도로 학습하여, GD의 느린 순차적 학습을 피한다.
• GD를 제약하는 임계 임계값을 훨씬 초과하는 학습률에서도 안정적으로 작동하여, 학습률을 문제의 조건수로부터 자유롭게 한다.
• 가중치 균형을 유지하기 위해 GD의 쌍곡선 기하학을 요구하지 않는 별도의 보존 법칙($\sqrt{P^\top P} - \sqrt{Q^\top Q}$)을 통해 균형 잡힌 솔루션에 도달한다.

저자들은 본 분석이 수정되지 않은 스펙트럴 그래디언트 디센트(모멘텀, 가중치 감쇠, 뉴턴-슐츠 반복을 무시함)와 단순 행렬 분해에 국한되어 있다는 점을 강조하면서도, 정렬 및 균등 성장에 대한 기계적 이해가 더 크고 복잡한 모델에서 Muon의 성공을 이해하는 토대를 제공한다고 강조한다. 이 연구는 왜 Muon이 조건이 나쁘거나 이질적인 스펙트럼 설정에서 훈련을 가속화하는지에 대한 이론적 설명을 제공한다.