기술 요약: Muon은 Adam보다 더 견고하고 전이 가능한 특징을 학습한다

문제 정의

최근 대규모 언어 모델(LLM) 및 비전 분류기 사전 학습에서 뛰어난 효율성(종종 Adam이나 SGD보다 거의 두 배 빠르게 학습)을 입증한 Muon 옵티마이저(Newton-Schulz에 의한 Momentum Orthogonalized)의 구체적인 특징 학습 이점이 무엇인지는 여전히 불분명합니다. 기존 문헌은 주로 최단 경사 하강법(steepest gradient descent), 결합 메모리 역학(associative memory dynamics), 또는 헤비 테일(heavy-tailed) 데이터 처리와 같은 최적화 이론적 관점을 통해 Muon의 성공을 설명해 왔습니다. 그러나 Muon이 Adam이나 SGD보다 실제로 더 나은 특징을 학습하는지, 그리고 이러한 특징들이 모델 표현(representation)에서 어떻게 나타나는지에 대한 체계적인 조사는 부족했습니다. 본 논문은 다음 질문을 다룹니다: Muon은 Adam 및 SGD보다 더 견고하고 전이 가능한 특징을 학습하는가, 그리고 이러한 이점의 특성은 무엇인가?

방법론

저자들은 비전 및 언어 도메인에 걸친 광범위한 경험적 평가와 정형화된 분류 문제에 대한 이론적 분석을 결합한 이중 접근 방식을 채택했습니다.

경험적 평가

본 연구는 특징의 품질을 **견고성(robustness)**과 **전이 가능성(transferability)**이라는 두 가지 주요 지표를 통해 평가합니다.

• 견고성: 사전 학습된 모델(ResNet-18, ViT-S, GPT-2, GPT-2 Medium)을 변형된 입력값에 대해 테스트합니다. 비전 모델은 ImageNet-C(15가지 손상 유형)로 평가하며, 언어 모델은 문자 수준의 오타 공격(삽입, 삭제, 치환, 교체)을 적용한 새로운 벤치마크인 FineWeb10B-C로 평가합니다.
• 전이 가능성: 사전 학습된 특징을 다운스트림 태스크에 적응시킵니다. 비전의 경우, 고정된 백본(frozen backbone) 위에서 선형 분류기를 네 가지 데이터셋(EuroSAT, Food-101, Flowers-102, Stanford Cars)에 대해 학습시킵니다. 언어의 경우, 지시 이행(instruction-following) 데이터셋(Alpaca, Dolly, WizardLM)에 대해 전체 모델을 파인튜닝합니다.
• 은닉 상태 분석(Hidden-State Analysis): 이러한 성능 향상의 메커니즘을 이해하기 위해, 저자들은 다음을 사용하여 레이어별 은닉 상태를 분석합니다:
  • 로짓 마진(Logit Margins): 훈련된 선형 프로브(tuned-lens 프레임워크)를 통해 은닉 상태를 디코딩하여 정답 클래스/토큰과 가장 강력한 방해 요소(distractor) 사이의 마진을 계산합니다.
  • 스펙트럼 다양성(Spectral Diversity): 은닉 상태 행렬의 유효 랭크(effective rank)와 Top-k 에너지 분율(Top-k energy fraction)을 계산하여, 직교하는 방향들 사이에 정보가 얼마나 다양하게 분포되어 있는지 측정합니다.
• 실험 프로토콜: 모든 비교는 아키텍처, 데이터, 훈련 기간을 고정하되 하이퍼파라미터(학습률, 스케줄러)는 각 옵티마이저에 대해 독립적으로 튜닝하여, 표현력의 차이를 엔지니어링 노력의 차이로부터 격리하는 매치드 버젯(matched-budget) 프로토콜을 따릅니다.

이론적 분석

저자들은 **다중 구성 요소 특징(multi-component features)**을 가진 단일 레이어 선형 분류 모델을 사용하여 이론적 근거를 제공합니다. 데이터 모델은 특징이 공유 블록 수준 구성 요소(공통 클래스)와 클래스별 구성 요소(세밀한 클래스)로 이루어져 있다고 가정합니다. 분석은 무한소 단계 크기(infinitesimal step sizes)와 제로 모멘텀(zero momentum) 조건에서 GD, Adam, Muon의 훈련 역학을 비교하며, 각 옵티마이저가 특정 손실 임계값에 처음 도달하는 상태에 초점을 맞춥니다.

주요 기여 및 결과

1. 우수한 특징 품질에 대한 경험적 증거

본 논문은 Muon이 Adam 및 SGD가 학습하는 것보다 더 견고하고 전이 가능한 특징을 일관되게 학습한다는 것을 입증합니다:

• 견고성: Muon으로 사전 학습된 모델은 모든 손상 유형과 아키텍처에 대해 ImageNet-C에서 가장 높은 평균 정확도를 달-달하며, FineWeb10B-C에서 가장 낮은 퍼플렉시티(perplexity)를 기록했습니다. 특히, Muon은 비전 트랜스포머와 언어 모델에서 Adam보다 우수하며, 언어 모델에서는 SGD보다 우수합니다.
• 전이 가능성: Muon으로 사전 학습된 특징은 다운스트림 태스크로 더 효과적으로 전이됩니다. 비전 분야에서 Muon은 모든 다운스트림 벤치마크에서 Adam 및 SGD와 대등하거나 이를 능가합니다. 언어 분야에서 Muon으로 사전 학습된 모델은 Adam으로 사전 학습된 모델에 비해 지시 이행 튜닝 후 더 낮은 퍼플렉시티를 달성합니다.

2. 은닉 상태 메커니즘

경험적 이점은 학습된 은닉 표현의 특정 속성과 직접적으로 연결됩니다:

• 더 큰 로짓 마진: 레이어별 프로빙 결과, Muon은 대부분의 레이어에서 Adam 및 SGD에 비해 유의미하게 큰 로짓 마진을 생성합니다. 이는 정답 예측과 방해 요소 사이의 간격을 넓혀주며, 관찰된 견고성과 직접적인 상관관계를 갖습니다.
• 더 높은 유효 랭크: 스펙트럼 분석에 따르면, Muon은 레이어 전반에 걸쳐 더 높은 유효 랭크와 더 낮은 Top-10 에너지 분율을 달성합니다. 이는 Muon의 표현이 정보를 직교하는 방향들로 더 고르게 분산시켜, 다운스트림 태스크에 대한 적응을 용이하게 하는 더 풍부한 특징 서브스페이스(feature subspace)를 생성함을 시사합니다.

3. 이론적 보장

정형화된 다중 구성 요소 특징 설정에서, 저자들은 Muon의 스펙트럼 정규화가 그래디언트 행렬의 특이값 불균형(singular-value imbalances)을 제거함을 증명합니다.

• 마진 이점: Muon은 매칭된 훈련 손실(지수적으로 작은 터미널 손실 구간 제외)의 어떤 지점에서도 Adam 및 GD보다 엄격하게 더 큰 분류 마진을 확보합니다. 이는 Muon의 스펙트럼 정규화가 직교하는 특징 방향들의 기여도를 균형 있게 만드는 반면, Adam의 좌표별 정규화와 GD의 가공되지 않은 그래디언트 흐름은 특정 구성 요소에 대한 과도한 의존을 초래하기 때문입니다.
• 유효 랭크 이점: Muon은 매칭된 손실 조건에서 Adam 및 GD보다 엄격하게 더 높은 특징 유효 랭크를 달성함을 증명하며, 이는 더 다양한 표현을 생성한다는 경험적 발견을 뒷받침합니다.

의의 및 주장

본 논문은 초점을 최적화 효율성에서 특징 품질로 전환함으로써 Muon의 이점을 "규명(demystify)"한다고 주장합니다. 저자들은 Muon의 성공이 단순히 빠른 수렴의 결과가 아니라, 입력 섭동(perturbation)에 본질적으로 더 견고하고 새로운 태스크에 더 잘 적응할 수 있는 표현을 학습하는 능력에서 비롯된다고 주장합니다.

본 연구의 의의는 다음과 같습니다:

1. 최적화와 표현의 연결: 매트릭스 인지 업데이트 규칙(matrix-aware update rule)과 학습된 특징의 기하학적 속성(로짓 마진 및 유효 랭크) 사이의 구체적인 연결 고리를 확립했습니다.
2. 이론적 검증: 매칭된 손실 조건 하에서 Muon이 Adam 및 GD에 비해 체계적으로 다르고 우월한 특징 품질(더 큰 마진, 더 높은 유효 랭크)을 유도한다는 최초의 이론적 증명을 제공했습니다.
3. 실질적 시사점: 본 연구 결과는 분포 변화에 대한 견고성이나 다운스트림 태스크로의 강력한 전이성이 요구되는 응용 분야에서, 훈련 예산이 동일하게 맞춰졌을 때도 Muon이 Adam과 같은 표준 적응형 옵티마이저보다 우월한 선택임을 시사합니다.

저자들은 현재의 분석이 LLM과 비전 분류기에 집중되어 있으며, 다른 모델 제품군(예: 확산 모델)에서의 Muon 성능은 향후 연구 과제로 남겨둔다고 겸허히 언급했습니다.