기술 요약: AMUSE (Anytime Muon with Stable Gradient Evaluation)

문제 제기

현대 딥러닝 최적화는 일반적으로 학습률 스케줄 (예: 코사인 감쇠) 과 결합된 AdamW 옵티마이저에 의존합니다. 최근 발전들은 이러한 구성 요소들을 개별적으로 도전했습니다: Schedule-Free (SF) 최적화는 현재 값과 평균화된 반복 값 사이의 보간점에서 그래디언트를 평가함으로써 명시적인 스케줄을 제거하고, Muon은 모멘텀을 직교화하여 행렬 값 파라미터의 업데이트 기하학을 개선합니다.

대규모 언어 모델 (LLM) 사전 학습에서 Muon 의 경험적 성공에도 불구하고, 그 근본 메커니즘은 완전히 이해되지 않았습니다. 해당 논문은 Muon 에서 특정 불안정성을 식별합니다: Muon 의 직교화는 손실 지형의 "평탄한" 방향 (즉, "강" 부분) 을 따라 진행을 가속화하지만, 우연히 "가파른" 고곡률 방향 (즉, "계곡 벽") 에서의 노이즈를 증폭시켜 진동하는 궤적을 초래합니다. 반면, Schedule-Free 방법은 궤적을 강 근처에 유지함으로써 궤적을 안정화하지만, Muon 의 직교화가 제공하는 빠른 대량 진행을 결여할 수 있습니다. 핵심 문제는 수동 학습률 스케줄 없이 Muon 의 빠른 수렴과 Schedule-Free 평균화의 안정성을 결합한 옵티마이저를 설계하는 것입니다.

방법론: AMUSE

저자들은 AMUSE(Anytime Muun with Stable gradient Evaluation) 를 제안합니다. 이는 Muon 의 직교화된 행렬 업데이트를 시간 가변 보간 계수를 사용하여 Schedule-Free 프레임워크와 통합합니다.

이론적 기반: 강 - 계곡 손실 지형

분석은 강 - 계곡 손실 지형 관점에 기반합니다:

• 강: 유용한 학습 진행이 발생하는 고차원, 저곡률 대역 부분 공간.
• 계곡 벽: 진동을 유발하는 저차원, 고곡률 지배적 부분 공간.
• Muon 의 메커니즘: 모멘텀 행렬을 직교화하면 업데이트의 상대적 대역 성분이 증가하여 강을 따라 진행이 가속화됩니다. 그러나 지배적 성분의 감쇠를 방지하여 지배적 방향의 작은 노이즈가 계곡 벽을 가로지르는 진동으로 증폭되게 합니다.
• Schedule-Free 메커니즘: 평균화된 궤적 (강에 더 가까움) 에서 그래디언트를 평가함으로써 업데이트 단계 전에 고곡률 변동을 필터링합니다.

알고리즘 설계

AMUSE 는 빠른 기본 시퀀스 $Z_t$(Muon 업데이트) 와 안정적인 평균화된 시퀀스 $X_t$ 사이의 그래디언트 평가점 $Y_t$를 결정하기 위해 시간 가변 보간 계수 $\beta_t$를 사용합니다:
$$Y_t = (1 - \beta_t)Z_t + \beta_t X_t$$

$\beta_t$의 스케줄은 초기 적응과 후기 안정성을 균형 있게 조정하도록 설계되었습니다:

1. 워밍업 단계 ($t \le T_0$): $\beta_t$는 작은 값 $\beta_1$로 고정됩니다. 이는 옵티마이저가 빠른 Muon 시퀀스 ($Z_t$) 근처에서 그래디언트를 평가하여 초기 적응을 빠르게 할 수 있게 합니다.
2. 워밍업 이후 ($t > T_0$): $\beta_t$는 다음에 따라 1 을 향해 점진적으로 증가합니다:
   $$\beta_t = 1 - \left( \frac{T_0 - 1}{t - 1} \right)^\rho (1 - \beta_1)$$
   여기서 $\rho \in [0, 1]$은 전환 속도를 제어합니다. $\beta_t$가 증가함에 따라 그래디언트 평가점은 평균화된 시퀀스 $X_t$ 쪽으로 이동합니다. 이는 Muon 의 직교화를 거치기 전에 그래디언트의 고곡률 (지배적) 성분을 억제하여 궤적을 안정화합니다.

중요하게도, 이 스케줄은 전체 학습 기간과 무관하여 언제나 가능한 학습(어떤 시점에서든 중단해도 성능 저하 없이 학습 가능) 을 가능하게 합니다. AMUSE 는 학습률 감쇠 스케줄이 필요하지 않습니다.

주요 기여

1. Muon 의 기하학적 분석: 논문은 Muon 의 효율성과 불안정성에 대한 기하학적 설명을 제공합니다. 직교화가 대역 지향적 업데이트를 촉진하여 강 진행을 가속화하지만, 지배적 방향의 노이즈를 증폭시켜 계곡 벽 진동을 유발함을 보여줍니다.
2. 보간을 통한 안정화: 저자들은 평균화된 궤적에 더 가까운 곳에서 (Schedule-Free 를 통해) 그래디언트를 평가함으로써 직교화 전에 지배적 성분을 줄여 진동을 완화함을 보여줍니다.
3. AMUSE 알고리즘: 그래디언트 평가점을 빠른 Muon 궤적에서 안정적인 평균화된 궤적으로 동적으로 전환하는 새로운 옵티마이저입니다. 이는 Muon 의 빠른 수렴과 Schedule-Free 방법의 안정성을 결합합니다.
4. 언제나 가능한 학습 지원: 총 단계 수와 무관한 시간 가변 스케줄을 사용하고 학습률 감쇠를 피함으로써, AMUSE 는 유연한 학습 기간을 지원합니다.

실험 결과

AMUSE 는 이미지 분류 (CIFAR-10/100, SVHN, ImageNet-1k), 이미지 분할 (ISIC 2018), LLM 사전 학습 (FineWeb 에서 124M, 720M, 1B 파라미터 Llama 스타일 모델) 을 포함한 다양한 벤치마크에서 평가되었습니다.

• 성능 - 반복 Pareto 프론티어: AMUSE 는 모든 작업에서 (Schedule-Free) AdamW 및 Muon 일관되게 능가합니다.
• 효율성 향상:
  • 720M LLM 사전 학습에서 AMUSE 는 Muon 의 최종 성능에 도달하는 데 1.51 배 적은 단계가 필요합니다.
  • ResNet-50 을 사용한 ImageNet 에서 동일한 성능을 달성하는 데 1.12 배 적은 단계가 필요합니다.
  • ViT 미세 조정을 사용한 ImageNet 에서 3.08 배 적은 단계가 필요합니다.
• 안정성: 사후 지수 가중 평균 (EWA) 또는 학습률 감쇠로부터 상당한 이점을 얻는 일반 Muun 과 달리, AMUSE 는 학습 전반에 걸쳐 안정적인 궤적을 유지하며 이러한 사후 기법을 적용했을 때 최소한의 개선만 보입니다.
• 확장성: AMUSE 는 대규모 배치 크기 (최대 2048) 와 대규모 모델 규모 (1B 파라미터) 에서 효과적으로 확장되며, 다른 Schedule-Free 변형들이 때때로 어려움을 겪는 영역에서도 작동합니다.

중요성과 주장

이 논문은 AMUSE 가 학습 역학과 손실 지형 기하학에 대한 체계적인 이해에서 도출된 옵티마이저 설계의 원칙적인 접근법임을 주장합니다. 빠른 대역 진행과 계곡 벽 안정성 사이의 트레이드오프를 명시적으로 해결함으로써, AMUSE 는 수동 학습률 스케줄링 없이도 우수한 성능을 달성합니다.

저자들은 겸손하게도 AMUSE 가 기존 옵티마이저보다 성능이 뛰어나지만, 기본 시퀀스 $Z_t$를 위한 하나의 추가 상태 복사본으로 인해 일반 Muon 에 비해 약간의 메모리 오버헤드가 발생한다고 지적합니다. 다만 이는 AdamW 및 Schedule-Free AdamW 와 비교할 만한 수준입니다. 그들은 향후 연구가 이 추가 메모리 비용 없이 유사한 안정성과 성능을 달성하는 데 초점을 맞출 수 있음을 시사합니다. 논문은 "강 - 계곡" 관점이 효율적이고 안정적인 옵티마이저를 설계하기 위한 견고한 프레임워크를 제공함을 강조합니다.