기술 요약: MONA – 네스테로프 가속을 갖춘 뮤온 옵티마이저

1. 문제 제기

대형 언어 모델 (LLM) 의 학습은 점점 더 우수한 샘플 효율성을 가진 옵티마이저를 요구합니다. AdamW가 10 년 이상 표준으로 자리 잡았지만, 최근 Muon과 같은 대안들이 등장했습니다. Muon 은 가중치 행렬을 단일 기하학적 단위로 취급하고 모멘텀 버퍼에 행렬 직교화 (뉴턴 - 슈르츠 반복법을 통해) 를 적용함으로써 1 차 방법을 개선합니다. 이 접근법은 스펙트럼 노름 하에서 가장 가파른 하강 (steepest descent) 과 정렬되어, 대규모 전문가 혼합 (MoE) 학습에서 AdamW 대비 약 $2\times$의 계산 효율성을 제공합니다.

그러나 Muon 은 모든 1 차 경사 방법과 근본적인 한계를 공유합니다: 국소 경사 정보 이상의 손실 지형을 탐색할 명시적 메커니즘이 부족합니다. 결과적으로 Muon 은 특히 대규모 배치 학습 시나리오에서 **날카로운 국소 최소값 (sharp local minima)**에 갇힐 수 있습니다. 날카로운 최소값은 낮은 일반화 능력과 연관되는 반면, 더 평탄한 최소값은 더 나은 하류 성능을 제공하는 경향이 있습니다. ALTO 옵티마이저는 경사 차이의 지수 이동 평균 (EMA) 을 기반으로 한 가속 항을 도입하여 날카로운 최소값을 탈출하려 했지만, Adam 과 LAMB 를 위한 범용 어댑터로 설계된 특성상 Muon 의 비선형성이 매우 높은 직교화 과정과의 통합에 어려움을 겪습니다.

2. 방법론: MONA

저자들은 Muon 의 직교화 프레임워크에 곡률 인식 가속 항을 매끄럽게 통합한 MONA(Nesterov 가속을 갖춘 Muon 옵티마이저) 를 제안합니다.

핵심 알고리즘

MONA 는 모멘텀 누적 및 직교화 이전에 가속 항을 도입함으로써 표준 Muon 파이프라인을 수정합니다. $k$번째 반복에서의 행렬 매개변수 $W$에 대한 업데이트 과정은 다음과 같습니다:

1. 경사 차이 계산: 현재 확률적 경사 $G_k$와 이전 경사 $G_{k-1}$ 사이의 차이를 계산합니다:
   $$D_k = G_k - G_{k-1}$$
2. 가속 버퍼 업데이트: 이러한 차이의 지수 이동 평균을 유지합니다:
   $$A_k = \beta_a A_{k-1} + (1 - \beta_a) D_k$$
   여기서 $\beta_a$는 가속의 기억력을 조절합니다.
3. 곡률 인식 경사: 원시 경사에 스케일링된 가속 항을 추가하여 확장된 경사 $\tilde{G}_k$를 구성합니다:
   $$\tilde{G}_k = G_k + \alpha A_k$$
   계수 $\alpha$는 가속을 스케일링합니다. 이론적 동기는 $G_k - G_{k-1}$이 $H_k(\theta_k - \theta_{k-1})$를 근사한다고 제안하며, 여기서 $H_k$는 헤시안입니다. 따라서 이 항은 경사 변화가 큰 날카로운 최소값에서 멀어지도록 지시하고, 평탄한 영역에서는 작게 유지됩니다.
4. Muon 직교화: 확장된 경사는 표준 Muon 모멘텀 및 직교화 단계로 입력됩니다:
   $$M_k = \mu M_{k-1} + \tilde{G}_k$$
   $$O_k = \text{Newton-Schulz}(M_k)$$
   $$W_{k+1} = W_k - \eta (\gamma O_k + \lambda W_k)$$

설계 근거

직교화 이전에 가속 항을 적용하는 것이 결정적입니다. 만약 그 이후에 적용된다면, Muon 이 스펙트럼 노름 정규화를 위해 의존하는 직교 구조가 교란될 것입니다. 원시 경사를 먼저 곡률 인식 방향으로 변환함으로써, 모멘텀 버퍼는 이러한 풍부한 방향들을 포착하고, 이는 이후 Muon 의 표준 파이프라인을 통해 처리됩니다.

벡터 값 매개변수 (예: 임베딩, 편향) 의 경우, MONA 는 Muon 의 기존 관례를 따르며 AdamW 를 기본으로 사용합니다.

3. 이론적 분석

이 논문은 표준 가정 (L-부드러움, 편향되지 않은 경사 및 유계 분산, 기대 방향 정렬) 하에서 MONA 에 대한 수렴 분석을 제공합니다.

• 수렴 속도: 저자들은 MONA 가 비볼록 환경에서 SGD 와 AdamW 의 점근적 속도와 일치하는 $O(K^{-1/2})$의 수렴 속도를 달성함을 증명합니다. 가속 항은 경계 내의 상수에 영향을 미치지만 점근적 속도에는 영향을 주지 않습니다.
• 날카로운 최소값 탈출: 이 분석은 가속 항이 어떻게 날카로운 최소값을 선택적으로 저항하는지 형식화합니다. 정상점 근처에서 기대 가속 $E[A_k]$는 과거 방향들의 헤시안 가중 합에 비례합니다. 큰 고유값 (날카로운 최소값) 을 가진 방향의 경우 가속 크기가 커져 옵티마이저를 밀어냅니다. 작은 고유값 (평탄한 방향) 의 경우 가속은 작아 안정적인 수렴을 허용합니다. 이 메커니즘은 곡률 인식을 추가하면서도 Muon 의 스펙트럼 노름 정규화를 유지합니다.

4. 실증 결과

MONA 는 세 가지 규모의 MoE 사전 학습 모델 (1B, 6B, 68B 매개변수) 과 지도형 미세 조정 (SFT) 단계를 통해 평가되었습니다.

사전 학습 성능

• 모델: MOE-1B-A0d2B(400B 토큰), MOE-6B-A0d5B(1.2T 토큰), MOE-68B-A3B(700B 토큰).
• 지표: 코드, 수학적 추론, 일반 영어, 중국어 학술 텍스트 도메인에서 검증 손실이 모니터링되었습니다.
• 결과: MONA 는 세 가지 모델 규모와 모든 도메인에서 Muon 과 AdamW 모두보다 일관되게 낮은 검증 손실을 달성했습니다. 성능 차이는 특히 더 큰 68B 모델에서 두드러졌으며, 여기서 MONA 는 더 안정적인 수렴을 보여주었습니다.

하류 벤치마크 (MOE-68B-A3B)

MONA 로 학습된 68B 모델은 포괄적인 벤치마크 세트를 통해 평가되었습니다:

• 일반 능력: MONA 는 MMLU, MMLU-Pro, CMMLU, CEVAL, BBH 를 포함한 벤치마크에서 평균 점수 (0.4557) 가 가장 높았으며, Muon(0.4478) 과 AdamW(0.4382) 를 능가했습니다.
• 수학적 추론 및 코드 생성: MONA 는 GSM8K, GPQA, 그리고 다양한 코드 생성 벤치마크 (HumanEval+, MBPP+, BigCodeBench, LiveCodeBench) 에서 우수한 성능을 보여주었습니다. 특히 MONA 는 이전이 가능한 코드 표현이 필요한 작업에서 상당한 개선을 달성했습니다.
• 지도형 미세 조정 (SFT): MONA 로 사전 학습된 체크포인트가 코드 데이터에 대해 SFT 를 거칠 때, Muon 으로 사전 학습된 기준선과 비교하여 전체 과정에서 더 낮은 학습 손실을 유지했습니다. BigCode 세트를 통한 SFT 후 평가에서 MONA 로 초기화된 모델은 Muon 으로 초기화된 모델보다 7 개 작업 중 6 개에서 더 높은 점수를 달성했습니다.

효율성 및 배포

이전 경사 ($G_{k-1}$) 와 가속 버퍼 ($A_k$) 저장에 따른 메모리 오버헤드를 해결하기 위해 저자들은 MONA-Lite를 도입했습니다:

• 양자화: 보조 상태를 float32 대신 bfloat16(BF16) 으로 저장.
• 스트리밍 계산: 전용 $G_{k-1}$ 버퍼의 필요성을 제거하기 위해 경사 차이를 제자리에서 계산.
• 결과: 이러한 전략은 학습 품질을 희생하지 않으면서 추가 메모리 오버헤드를 약 75% 감소시킵니다. 계산 오버헤드 분석에 따르면 MONA 는 옵티마이저 단계 수준에서 약 1% 만 느리며, Muon 과 비교하여 엔드투엔드 반복 시간에는 실질적인 차이가 없습니다.

5. 중요성 및 주장

이 논문은 MONA 가 Muon 의 기하학적 효율성과 곡률 인식 가속 필요성 사이의 간극을 성공적으로 연결한다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

1. 알고리즘 통합: 가속 항이 기하학적 속성을 교란하지 않고 행렬 직교화 프레임워크에 통합될 수 있음을 입증.
2. 확장성: 다양한 모델 크기 (1B~68B 매개변수) 와 토큰 수 (최대 1 조 토큰) 에 걸쳐 방법을 검증하여 AdamW 와 Muon 모두에 대한 일관된 개선을 보여줌.
3. 일반화: 사전 학습 중 날카로운 최소값을 탈출하는 것이 복잡한 추론 및 코드 생성과 같은 하류 작업에서 더 나은 일반화로 이어진다는 실증적 증거 제공.
4. 실용성: 리소스 제약 환경에서 옵티마이저를 실행 가능하게 만드는 메모리 효율적인 변형 (MONA-Lite) 제공.

저자들은 MONA 가 우수한 수렴과 하류 성능을 달성하면서도 Muon 프레임워크의 계산 효율성을 유지하는 대규모 LLM 학습을 위한 최첨단 옵티마이저를 대표한다고 결론지었습니다.