기술 요약: Muown - Muon 최적화를 위한 행 노름 제어

문제 제기

Muon 은 트랜스포머 선형 계층의 행렬 구조를 활용하여 AdamW 를 종종 능가하는 대규모 언어 모델 (LLM) 사전 학습용 경쟁력 있는 최적화기로 부상했습니다. 그러나 최근 관찰 결과에 따르면, Muon 의 대규모 확장 시 행동은 가중치 감쇠 (weight decay) 에 민감하며, bfloat16 오버플로우와 폭발적인 어텐션 로짓 (attention logits) 과 같은 수치적 불안정성에 취약한 것으로 나타났습니다.

식별된 핵심 문제는 Muon 을 분해된 가중치 감쇠 없이 사용할 때 훈련 중 가중치 행렬의 스펙트럼 노름 ($\|W_t\|_{S_\infty}$) 이 지속적으로 상승하는 경향입니다. 이 상승은 안정적인 특징 학습에 필요한 스펙트럼 조건 ($\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$) 을 위반합니다. 분해된 가중치 감쇠는 가중치 행렬을 균일하게 축소함으로써 이를 완화할 수 있지만, 이는 성장의 원인이 되는 특정 방향과 무관하게 전체 스펙트럼을 축소하는 비표적 개입이므로 훈련 초기 단계에서 수렴 속도를 저하시킬 수 있습니다. 본 논문은 이러한 스펙트럼 노름 상승의 구체적인 경험적 원인을 규명하고, 보다 표적화된 제어 메커니즘을 제안하고자 합니다.

방법론: Muown

저자들은 집합적 제약이 아닌 재매개변수화 전략을 통해 스펙트럼 노름 상승을 해결하는 Muon 의 대체재인 Muown을 소개합니다.

1. 스펙트럼 노름 분해

핵심 통찰은 스펙트럼 노름을 두 가지 요인으로 분해하는 것입니다 (명제 1):
$$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$$
여기서:

• $g_t = \|W_t\|_{row}$는 행 크기 (행의 유클리드 노름) 의 벡터입니다.
• $D_t$는 행 정규화 행렬입니다 ($W_t = \text{Diag}(g_t)D_t$).
• $C_t = D_t D_t^\top$는 단위 노름 행들의 그람 행렬입니다.
• $P_t$는 정규화된 행 크기의 대각 행렬입니다.
• $\lambda_{\max}(P_t C_t P_t)$는 행 일관성 (row-coherence) 요인을 나타냅니다.

5 억 파라미터 모델에 대한 경험적 분석 결과, 행 크기 요인 ($\|g_t\|_\infty$) 이 스펙트럼 노름 상승의 주요 원동력이며 $\|W_t\|_{S_\infty}$와 동기적으로 움직이는 것으로 나타났습니다. 반면 행 일관성 요인은 훈련 내내 잘 제어되고 유계 (bounded) 상태를 유지합니다.

2. 재매개변수화 및 최적화

이 진단에 기반하여 Muown 은 가중치 행렬 $W$를 다음과 같이 재매개변수화합니다:
$$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$$
여기서 $g \in \mathbb{R}^m$은 명시적인 행 크기 벡터를 나타내고, $R \in \mathbb{R}^{m \times n}$은 방향 성분을 나타냅니다.

• 방향 성분 ($R$): Muon 을 $R$에 변경 없이 적용하여 스펙트럼 노름 기하학 (최대 경사 스펙트럼 하강) 하에서 처리합니다.
• 크기 성분 ($g$): 행 크기 $g$는 명시적인 최적화 변수로 승격됩니다. 대각 행렬의 연산자 노름에 대한 자연스러운 쌍대 노름인 $\ell_\infty$ 기하학 하에서 업데이트됩니다.
• 업데이트 규칙: $g$의 업데이트는 $\ell_\infty$ 기하학 하에서 평활화된 SignSGD 로 작용하는 Adam 최적화기를 사용하여 확률적 노이즈를 처리하며, $R$의 업데이트는 뉴턴 - 슈르츠 직교화 (Newton-Schulz orthogonalization) 를 따르는 표준 Muon 절차를 따릅니다.

이 접근법은 최적화기 상태 내에서 암시적으로 구현되어 모델 아키텍처의 순전파 (forward pass) 를 변경하지 않습니다.

주요 기여

1. 스펙트럼 상승 진단: 본 논문은 스펙트럼 노름을 분해하여 Muon 하에서 스펙트럼 노름 성장의 경험적 원동력인 최대 행 크기 ($\|g_t\|_\infty$) 를 식별하고, 이를 안정적인 행 일관성 요인과 구별합니다.
2. Muown 알고리즘: 행 크기를 $\ell_\infty$ 기하학 하에서 업데이트되는 명시적 학습 가능한 변수로 취급하면서 Muon 의 방향 업데이트를 유지하는 새로운 최적화기입니다. 이는 가중치 감쇠의 무분별한 축소와 스펙트럼 재매개변수화의 직교성 사전 (priors) 을 피합니다.
3. 수렴 보장: 저자들은 Muown 이 결정론적 및 확률적 regimes 모두에서 최적의 비볼록 수렴 속도를 달성함을 증명합니다. 분석은 업데이트 기하학에 정렬된 혼합 쌍대 노름 (크기에 대한 $\ell_1$, 방향에 대한 Schatten-1) 에서 수행됩니다. 이론적 경계는 Muon 에 비해 유효 확률적 노름 계수의 감소를 시사합니다.
4. 경험적 검증: 1 억 2 천 4 백만에서 27 억 파라미터까지의 다양한 모델 크기에 걸친 FineWeb-Edu 에 대한 광범위한 실험은 일관된 개선을 입증합니다.

실험 결과

• 퍼플렉시티 개선: Muown 은 모든 테스트된 모델 규모 (1 억 2 천 4 백만에서 27 억) 와 학습률에서 Muon, SOAP, AdamW, Lion 보다 일관되게 낮은 퍼플렉시티를 달성합니다.
• 하이퍼파라미터에 대한 강건성:
  • 학습률: Muown 은 다양한 모델 너비에 걸쳐 근사 최적 학습률의 플래토를 넓혀 규모별 학습률 튜닝에 대한 민감도를 낮춥니다.
  • 가중치 감쇠: Muown 은 가중치 감쇠 튜닝의 필요성을 완화합니다. 가중치 감쇠 없이 ($\lambda=0$) 도 잘 수행되는 반면, Muon 은 발산이나 비최적 성능을 방지하기 위해 신중한 튜닝 (최적 $\lambda=0.1$) 이 필요합니다.
• 안정성: Muown 은 Muon 에서 관찰된 체계적인 스펙트럼 노름 상승을 방지합니다. 높은 학습률 영역 (예: $4 \times 10^{-3}$) 에서 Muon 은 종종 발산하는 반면, Muown 은 안정적이고 근사 최적 상태를 유지합니다.
• 경사 특성: 분석 결과, Muown 은 훈련 내내 Muon 에 비해 더 높은 정규화된 유효 랭크의 미니배치 경사를 유지하고 더 낮은 확률적 노름 계수를 보임을 보여줍니다.
• 계산 오버헤드: 이 방법은 단계 시간 (step-time) 오버헤드를 무시할 수준으로 도입합니다. 최적화기가 랭크 간에 샤딩 (분산 설정) 될 때 Muown 과 Muon 의 단계 시간은 구별할 수 없게 됩니다. 메모리 오버헤드는 최소 (행 수에 선형) 로 레이어당 몇 메가바이트만 추가됩니다.

중요성 및 주장

본 논문은 Muon 에서 관찰된 스펙트럼 노름 상승이 최적화기의 불가피한 속성이 아니라 통제되지 않은 행 크기 성장의 결과라고 주장합니다. 이 원동력을 격리하고 명시적인 최적화 변수로 승격시킴으로써 Muown 은 다음과 같은 성과를 달성합니다:

• 개선된 수렴: 공격적인 가중치 감쇠 없이도 더 나은 퍼플렉시티와 안정성을 달성합니다.
• 이론적 건전성: 업데이트의 혼합 성격을 명시적으로 고려하는 기하학 인식 프레임워크에서 최적의 수렴 속도를 제공합니다.
• 실용적 효율성: 분산 설정에서 Muon 과 동일한 실행 시간을 유지하면서 하이퍼파라미터 선택에 대한 우수한 강건성을 제공하는 대체재입니다.

저자들은 Muown 을 특징 학습의 이론적 안정성 요구사항과 대규모 LLM 훈련의 실용적 필요성 사이의 간극을 메우는 해결책으로 위치시키며, 분해된 가중치 감쇠에 대한 보다 표적화된 대안을 제시합니다.