연관 기억 학습에서 Muon 옵티마이저의 동역학과 스케일링 법칙

초록

본 논문은 선형 연관 기억 모델에 Muon 옵티마이저를 적용해, 주파수 계층 구조와 라벨 노이즈가 존재할 때의 학습 동역학을 분석한다. Gradient Descent는 저주파 성분 학습이 느려 전체 수렴을 병목시키는 반면, Muon은 행렬 부호 연산을 통해 주파수 불균형을 완화하고, 무노이즈 상황에서는 지수적 가속, 노이즈가 있는 경우에는 𝑂̃(1/T²) 수준의 스케일링을 달성한다. 또한 Muon을 암묵적 행렬 전처리기로 해석하고, SignGD와의 관계를 이론적으로 규명한다. 실험은 합성 장기 분포 데이터와 LLaMA‑style 사전학습에서 이론을 검증한다.

상세 분석

논문은 먼저 대규모 언어 모델 학습에서 매트릭스 파라미터를 업데이트할 때 행렬 부호 연산을 이용하는 Muon 옵티마이저가 SGD·Adam보다 뛰어난 실험적 성능을 보인다는 사실을 상기한다. 그러나 기존 이론은 수렴 속도와 일반적인 경계만 제공해, Muon이 왜 빠른 수렴을 보이는지, 특히 장기 분포(Long‑tail) 데이터에서 어떤 메커니즘으로 이득을 얻는지 설명하지 못한다. 이를 해소하기 위해 저자들은 연관 기억(Associative Memory) 모델을 선택한다. 이 모델은 K개의 지식 항목을 정규 직교 임베딩으로 표현하고, 각 항목은 계층적 주파수 r₁>r₂>…>r_M을 갖는 M개의 그룹으로 나뉜다. 각 그룹은 C개의 항목을 포함하며, C·M=K이다. 이렇게 구성된 데이터는 실제 LLM 사전학습에서 흔히 관찰되는 ‘head‑tail’ 구조를 수학적으로 단순화한다.

학습 목표는 소프트맥스 기반 교차 엔트로피 손실을 최소화하는 것이며, 두 옵티마이저의 업데이트 규칙을 명시한다. GD는 전통적인 경사 하강법이며, Muon은 행렬 부호 연산 msgn(∇L) = U sgn(Σ) Vᵀ (SVD 기반) 를 사용한다. 초기 가중치는 0이며, 학습률 η는 일정하게 유지한다.

무노이즈(α=0) 상황에서는 저자는 GD가 각 그룹의 손실 L_{GD}^{(j)}(t) ≈ 1/(p_j t) 로 감소함을 증명한다. 즉, 고주파(높은 p_j) 그룹은 빠르게 학습되지만 저주파(낮은 p_j) 그룹은 t에 비례해 매우 느리게 감소한다. 전체 손실은 L_{GD}(t) ≈ K/t 로, 전체 수렴 속도가 저주파에 의해 제한된다. 반면 Muon은 행렬 부호 연산이 그래디언트의 스펙트럼을 균등화시켜, 모든 그룹에 동일한 지수적 감쇠 L_{Muon}^{(j)}(t) ≈ exp(−c t) 를 보인다. 결과적으로 전체 손실은 L_{Muon}(t) ≈ K·exp(−c t) 로, GD 대비 지수적 가속을 달성한다. 이 현상은 Muon이 ‘블록 대칭’ 구조와 ‘작업 정렬(task alignment)’을 자동으로 구현해, 저주파 성분의 효과적인 전처리를 수행한다는 해석으로 설명된다.

노이즈(0<α<1) 상황에서는 라벨이 α 확률로 무작위로 섞이면서 각 그룹의 실제 학습 신호가 약화된다. 저자는 Muon의 학습 과정을 세 단계(초기 급증, 중간 완화, 최종 수렴)로 구분하고, 각 단계에서 손실 감소율을 정량화한다. 특히, 그룹 크기 C와 전체 그룹 수 M에 비례하는 Ω(p C^q) 수준의 속도 향상을 보이며, 이는 ‘지식 그룹 크기’가 클수록 Muon의 이점이 커짐을 의미한다.

스케일링 법칙에서는 주파수 분포가 파워 법칙 p_i ∝ i^{−β} (β>1) 로 가정된다. 이때 GD의 손실 하한은 Ω̃(1/T^{1−1/β}) 로, β가 클수록 수렴이 느려진다. 반면 Muon은 ℴ̃(1/T²) 로, β에 무관하게 2차 스케일링을 달성한다. 이는 Muon이 고주파와 저주파를 동일하게 ‘전처리’함으로써, 전체 스펙트럼에 대해 동일한 효과적인 학습률을 적용한다는 의미다.

전처리 관점에서는 Muon을 ‘암묵적 행렬 전처리기’로 해석한다. 행렬 부호 연산은 그래디언트의 특이값을 1 혹은 −1 로 정규화하고, 좌·우 특이벡터(U,V)를 그대로 유지한다. 따라서 업데이트는 W_{t+1}=W_t−η U sgn(Σ) Vᵀ 로, 이는 실제 작업 행렬 W* = γ EEᵀ 와의 정렬을 촉진한다. 저자는 좌표별 SignGD가 동일한 효과를 얻으려면 작업 행렬의 특이벡터를 사전에 알 수 있는 ‘오라클’이 필요하다고 지적한다. 실제 상황에서는 이러한 정보가 없으므로, Muon이 제공하는 자동 정렬 메커니즘이 실질적인 이점을 만든다.

실험에서는 (1) 합성 장기 분포 데이터셋에서 그룹 크기 C와 β를 다양하게 바꾸어 손실 감소 곡선을 확인하고, (2) LLaMA‑style 사전학습(7B 파라미터)에서 Muon이 동일한 FLOP 대비 1.8× 빠른 수렴과 낮은 퍼플렉시티를 달성함을 보고한다. 실험 결과는 이론적 예측(지수 가속, 1/T² 스케일링)과 일치한다.

종합하면, Muon은 행렬 부호 연산을 통한 스펙트럼 정규화가 저주파 병목을 해소하고, 전반적인 학습 효율을 크게 향상시킨다. 이는 대규모 LLM 사전학습에서 특히 장기 분포 데이터를 다룰 때 유용한 메커니즘임을 입증한다.