기술 요약: DynMuon: 무온 (Muon) 의 동적 스펙트럼 형성 관점

문제 제기

무온 (Muon) 은 최근 대규모 언어 모델 (LLM) 과 트랜스포머 학습을 위한 지배적인 최적화기로 부상했습니다. 그 핵심 메커니즘은 표준 경사 업데이트 행렬 $M = U\Sigma V^\top$을 그 극성 인자 (polar factor) $UV^\top$로 대체하여 특이값을 "평탄화"하되 특이 방향은 보존하는 것입니다. 이는 수렴성과 안정성을 향상시키지만, 해당 논문은 무온이 정적인 스펙트럼 변환 (구체적으로 스펙트럼 지수 $p=0$인 경우, 즉 $U\Sigma^p V^\top$ 계열) 을 사용한다는 치명적인 한계를 지적합니다.

이 연구가 다루는 핵심 미해결 질문은 단일 고정된 스펙트럼 형성 규칙이 전체 학습 과정에서 최적인지를 여부입니다. 기존 연구는 고정된 스펙트럼 형성 변형 (예: 고정된 멱법칙 형성 또는 지배적 부분 공간 억제) 을 탐구해 왔으나, 학습이 진행됨에 따라 서로 다른 스펙트럼 방향의 상대적 영향력이 어떻게 진화하는지에 대한 이론적 모델은 부재했습니다. 저자들은 최적의 스펙트럼 지수 $p$가 손실 함수의 변화하는 지형, 경사 노이즈의 성질, 그리고 학습 단계에 기반하여 동적으로 적응해야 한다고 가설을 세웠습니다.

방법론

1. 이론적 프레임워크: 노이즈 인지형 국소 모델링

저자들은 $D(p) := U\Sigma^p V^\top$로 정의된 스펙트럼 형성 연산의 학습 역학을 분석하기 위해 노이즈를 인지하는 국소 모델을 개발했습니다.

• 신호 - 노이즈 트레이드오프: 이 모델은 업데이트를 국소 곡률 방향 (모드) 으로 분해합니다. 이는 지수 $p$에 의해 제어되는 근본적인 트레이드오프를 드러냅니다.
  • 높은 $p$ (양수): 고곡률 (강한) 모드에서의 잔여 신호 수축을 가속화하지만, 평탄한 모드에서의 업데이트를 억제합니다.
  • 낮은 $p$ (음수): 저곡률 (평탄한) 모드에서의 수축을 증가시키지만, 동시에 해당 방향에서의 확률적 경사 노이즈를 증폭시킵니다.
• 단계 의존적 역학:
  • 초기 학습: 잔여 신호는 고곡률 방향에 집중됩니다. 양수 $p$는 이러한 방향을 우선시하여 제한된 노이즈 증폭으로 신호 수축을 가속화하므로 유익합니다.
  • 후기 학습: 학습이 진행됨에 따라 강한 모드에서의 잔여 신호가 감소하여 평탄한 모드로 집중도가 이동합니다. 이때는 약간 음수인 $p$가 유리해져, 여전히 유용한 신호를 포함하는 평탄한 방향으로 업데이트 강도를 재배분합니다. 다만, $p$가 너무 음수가 되면 노이즈 증폭으로 인해 최적화가 저하됩니다.

2. DynMuon 알고리즘

$p$가 양수에서 약간 음수로 전환되어야 한다는 이론적 통찰에 기반하여, 저자들은 DynMuon이라는 동적 스펙트럼 형성 알고리즘을 제안합니다.

• 스케줄링 전략: DynMuon 은 간단한 로지스틱 스케줄을 사용하여 스펙트럼 지수 $p$를 초기 단계의 양수 값 ($p_{max} = 1$) 에서 후기 단계의 약간 음수 값 ($p_{min} = -0.25$) 으로 단조 감소시킵니다.
• 효율적인 구현: 임의의 $p$에 대한 정확한 SVD 계산은 계산적으로 불가능합니다. DynMuon 은 무온에서 사용되는 뉴턴 - 슈르츠 (NS) 근사를 확장합니다.
  • $p \ge 0.25$인 경우, 원래 무온 업데이트 (또는 $p \in [0, 0.25)$인 경우 표준 NS 직교화) 를 사용합니다.
  • $p < 0$인 경우, 항등식 (identity) 주변에서 2 차 테일러 확장을 사용하여 분수 스펙트럼 형성 $A^{p/2}$ (여기서 $A = X_n X_n^\top$) 를 근사합니다. 이는 표준 NS 계산에 다항식 보정 항만 추가하여 무온과 동일한 점근적 복잡도를 유지합니다.
• 안정성 고정: $p$가 양수이고 클 때 안정성을 보장하기 위해, 알고리즘은 양수 영역을 안정적인 연산자에 고정하여 해당 영역에서 신뢰할 수 없는 테일러 근사를 피합니다.

주요 기여

1. 동적 스펙트럼 관점: 무온 스타일 최적화를 고정된 스펙트럼 연산에서 학습 역학에 따라 진화하는 스펙트럼 지수 $p$를 선택하는 적응형 문제로 재정의합니다.
2. 노이즈 인지형 국소 모델: 국소 곡률 방향 전반에 걸쳐 잔여 신호 감소와 확률적 노이즈 증폭 사이의 트레이드오프를 통해 $p$에 대한 단계 의존적 선호도를 설명하는 이론적 모델을 개발합니다.
3. 단계 의존적 영역 발견: 양수 지수가 초기 학습 (강한 모드 강조) 에 유리한 반면, 약간 음수인 지수 (이전에는 간과됨) 는 유용한 잔여 신호를 유지하는 평탄한 모드를 강조함으로써 후기 단계 최적화를 개선한다는 사실을 발견했습니다.
4. DynMuon 알고리즘: $p$를 양수에서 약간 음수로 스케줄링하는 효율적인 동적 스펙트럼 형성 알고리즘을 제안하여, 고정된 무온 대비 일관된 개선을 달성하고 계산 오버헤드는 무시할 수준으로 유지합니다.

실험 결과

저자들은 다양한 모델 크기 (1 억 2,700 만~11 억 파라미터), 아키텍처 (GPT 스타일 및 Qwen 스타일), 학습 설정 (FineWeb 및 FineWeb-Edu 데이터셋) 에서 DynMuon 을 평가했습니다.

• 성능: DynMuon 은 표준 무온 ($p=0$) 과 AdamW 베이스라인보다 일관되게 낮은 검증 손실을 달성합니다.
• 효율성: 동일한 목표 검증 손실 (학습 80% 지점에서 무온이 도달하는 손실로 정의됨) 에 도달하기 위해 DynMuon 은 10.6% 에서 26.5% 적은 학습 단계만 필요합니다.
• 오버헤드: DynMuon 의 단계당 시간 비용은 무온의 1.003~1.025 배에 불과하여 테일러 근사된 스펙트럼 형성의 효율성을 확인합니다.
• 강건성:
  • 아키텍처: 이득은 Qwen 스타일 아키텍처로 이전됩니다.
  • 학습 예산: 25 억에서 2,000 억 토큰 예산 전반에 걸쳐 개선 효과가 유지됩니다.
  • 학습률: DynMuon 은 무온에 비해 학습률 선택에 대한 민감도가 낮습니다.
  • 손실 목적 함수: CE-Brier 손실 보간 계열 전반에 걸쳐 약한 음수 형성의 이득이 지속되어, 이 메커니즘이 다양한 확률 공간 손실에 대해 강건함을 시사합니다.
• 절제 실험 (Ablations): 실험은 급격한 전환이나 고정된 음수 스케줄보다 부드러운 로지스틱 전환이 더 우수함을 확인합니다. 최적의 $p_{min}$은 $-0.25$로 발견되었으며, 더 공격적인 음수 값 (예:$-0.5$) 은 노이즈 증폭으로 인해 성능을 저하시킵니다.

중요성과 주장

이 논문은 DynMuon이 "최적"의 스펙트럼 편향이 정적이지 않고 학습 중에 진화함을 보여준다고 주장합니다. 고곡률 방향에서 평탄한 방향으로 강조를 동적으로 전환함으로써 DynMuon 은 무온과 같은 고정된 변환이 놓치는 최적화 역학을 포착합니다.

저자들은 이 작업을 더 적응적인 최적화기로 나아가는 한 걸음으로 위치시킵니다. 그들은 로지스틱 스케줄이 고정된 선택보다 우수하지만 이것이 최종 해답은 아닐 수 있음을 겸손하게 지적합니다. 최적의 $p$는 실시간 최적화 통계에 의존할 수 있으며, $p$의 유용한 범위가 상대적으로 작아 보이는 점을 고려할 때, 향후 연구는 덜 빈번하게 적용되는 더 비용이 많이 드는 신호를 사용하여 $p$의 온라인 추정에 집중할 수 있다고 제안합니다. 이 방법은 상당한 계산 비용 없이 효율성과 성능을 향상시키는 LLM 학습을 위한 범용 최적화 개선책으로 제시됩니다.