🤖 machine learning

Why Muon Outperforms Adam: A Curvature Perspective

이 논문은 Muon이 주로 더 낮은 정규화된 방향성 날카로움(NDS) 곡률 페널티를 통해 더 큰 1단계 손실 감소를 달성함으로써 대규모 언어 모델 학습에서 Adam보다 우수한 성능을 보인다는 것을 입증하며, 이러한 기하학적 이점은 데이터 불균형에 의해 증폭되고 레이어 내부의 곡률 감소에 의해 유지된다.

원저자: Shuche Wang, Fengzhuo Zhang, Jiaxiang Li, Dirk Bergemann, Zhuoran Yang

게시일 2026-06-04

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Shuche Wang, Fengzhuo Zhang, Jiaxiang Li, Dirk Bergemann, Zhuoran Yang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 광활하고 안개가 자욱한 산맥에서 가장 낮은 지점을 찾으려고 노력하고 있다고 상상해 보세요. 이것이 거대 AI 모델(LLM과 같은)을 훈련시키는 과정과 같습니다. 여기서 "산"은 오차율(error rate)이며, "가장 낮은 지점"은 완벽하고 가장 정확한 모델입니다. 그곳에 도달하기 위해서는 당신에게 어느 방향으로 발을 내디뎌야 할지 알려주는 가이드, 즉 옵티마이저(optimizer)가 필요합니다.

오랫동안 Adam은 표준적인 가이드 역할을 해왔습니다. 최근에는 Muon이라는 새로운 가이드가 등장하여 산의 밑바닥에 두 배 더 빠르게 도달하기 시작했습니다. 하지만 왜 그럴까요? 이 논문은 "곡률 관점"(지형의 모양을 살펴보는 것)을 사용하여 이 미스터리를 해결하는 탐정 역할을 합니다.

그 발견 내용을 이해하기 쉽게 설명하면 다음과 같습니다:

1. 두 가이드는 같은 보폭을 취하지만, Muon이 더 부드럽습니다

당신과 친구가 언덕을 내려가고 있다고 상상해 보세요. 두 사람 모두 정확히 같은 길이의 발걸음을 내딛기로 결정했습니다.

첫 번째 발걸음 (밀기): 당신과 친구 모두 동일한 힘으로 앞으로 밀어냅니다. 수학적으로 논문은 이를 "1차 이득(first-order gain)"이라고 부릅니다. 이 값은 같습니다.
울퉁불퉁함 (곡률 페널티): 지면은 완벽하게 평평하지 않습니다. 울퉁불퉁합니다. 만약 너무 세게 발을 내디디면, 약간 튕겨 올라오면서 에너지를 낭비하게 됩니다.
- Adam은 지형의 가장 "울퉁불퉁한" 부분으로 발을 내딛는 경향이 있습니다. 높은 둔덕을 치고, 다시 튕겨 올라오며 전진 동력을 잃습니다.
- Muon은 더 똑똑합니다. 지면의 모양을 살펴보고 더 부드럽고 평탄한 경로를 향해 발걸음을 조절합니다. 여전히 같은 길이의 발걸음을 내딛지만, 훨씬 더 작은 둔덕을 만납니다.

결과: Muon은 (Adam보다) 더 작은 "곡률 페널티"를 입기 때문에, 비록 동일한 초기 힘으로 밀었음에도 불구하고 한 번의 발걸음만으로도 Adam보다 실제로 더 멀리 내려갈 수 있습니다.

2. 비밀 병기: "정규화된 방향성 날카로움" (NDS)

논문은 NDS라는 멋진 용어를 도입했는데, 이는 본질적으로 "당신이 걷고 있는 방향으로 지면이 얼마나 날카로운가"를 측정하는 것입니다.

발견: Muon과 Adam은 같은 크기의 발걸음(동일한 "업데이트 노름/update norm")을 취하지만, Muon의 발걸음은 항상 지면이 덜 날카로운 방향을 향합니다.
비유: 키가 큰 옥수수 밭을 걷는다고 상상해 보세요.
- Adam은 가장 두껍고 날카로운 옥수수 줄기 사이를 곧장 뚫고 지나갑니다. (높은 NDS) 그러면 몸이 베이고, 옥수수가 반격하며 진행을 방해합니다.
- Muon은 특별한 나침반을 사용하여 옥수수 사이의 틈을 찾아냅니다. 같은 거리를 이동하지만, 옥수수의 밀도가 훨씬 낮습니다 (낮은 NDS). 저항 없이 미끄러지듯 지나갑니다.

3. 왜 Muon은 틈새를 더 잘 찾는가? (데이터 불균형)

논문은 데이터가 불균형할 때 지형이 더 울퉁불퉁해진다는 것을 발견했습니다.

시나리오: 도서관에 "고양이"에 관한 책이 90%, "강아지"에 관한 책이 10%뿐이라고 가정해 봅시다. 이것이 불균형한 데이터셋입니다.
효과: 이러한 불균형한 도서관에서는 지면이 극도로 들쭉날쭉하고 날카로운 스파이크로 가득 차게 됩니다.
승자: Adam은 "고양이" 스파이크에 심하게 갇히게 됩니다. 하지만 Muon은 이러한 들쭉날쭉하고 고르지 않은 풍경을 항해하는 데 훨씬 능숙합니다. 데이터가 더 불균형할수록, Muon의 부드러운 경로와 Adam의 울퉁불퉁한 경로 사이의 격차는 더 커집니다.

4. 마법은 어디에서 일어나는가? (레이어 내부)

AI 모델은 레이어들이 쌓인 구조(다층 케이크와 같은)로 이루어져 있습니다. 논문은 Muon의 이점이 전체 케이크에서 오는지, 아니면 특정 슬라이스에서 오는지 조사했습니다.

발견: 훈련이 진행됨에 따라 Muon의 이점은 변화합니다. Muon은 레이어 간의 상호작용(cross-layer)에 신경 쓰는 것을 멈추고, 각 개별 레이어 내부(within-layer)에서 발걸음을 부드럽게 만드는 데 완전히 집중합니다.
비유: 계주 경기를 생각해 보세요. 시작 단계에서 모든 선수는 주자 간의 바통 터치(cross-layer)를 걱정합니다. 하지만 경기가 격렬해지면, Muon은 속도의 비결이 다른 주자들의 소음은 무시한 채 자신의 구간을 완벽하게 부드럽게 달리는 것(within-layer)임을 깨닫습니다.

5. 이론적 증명: "이퀄라이저"

마지막으로, 저자들은 왜 이것이 작동하는지 증명하기 위해 간단한 수학적 모델("정형화된 이차 문제/stylized quadratic problem")을 구축했습니다.

문제: 산에 몇 개의 매우 가파르고 날카로운 절벽(높은 곡률)이 있고, 아주 완만한 경사지(낮은 곡률)가 넓게 펼쳐져 있다고 상상해 보세요.
Adam의 실수: "절벽"이 매우 가파르기 때문에, Adam은 그쪽으로 끌려갑니다. 가파른 절벽을 내려가기 위해 모든 에너지를 쏟아붓지만, 절벽이 너무 날카롭기 때문에 사방으로 튀어 오르며 요동칩니다.
Muon의 전략: Muon은 "스펙트럼 정규화(spectral normalization)" 기법을 사용합니다. 마치 마법의 이퀄라이저를 사용하여 가파른 절벽에 쏟는 에너지와 완만한 경사에 쏟는 에너지를 동일한 양으로 강제하는 것과 같습니다.
결과: 위험한 절벽에 과도하게 집중하지 않음으로써, Muon은 격렬하게 튀어 오르는 현상을 피합니다. 에너지를 고르게 분배함으로써, Adam은 따라올 수 없는 꾸준하고 효율적인 전진을 가능하게 합니다.

요약

Muon이 Adam을 이기는 이유는 더 큰 발걸음을 떼거나 더 세게 밀어서가 아니라, 더 나은 항해사이기 때문입니다. Muon은 AI를 튕겨 나가게 만드는 수학적 풍경의 "날카로운" 부분들을 피합니다. 특히 데이터가 지저분하거나 불균형할 때, 경로를 부드럽게 다듬음으로써 산의 밑바닥(최적의 모델)에 훨씬 더 빠르게 도달합니다.

기술 요약: 왜 Muon이 Adam보다 우수한가: 곡률 관점에서의 분석

문제 정의

Muon은 다양한 모델 규모에서 표준 Adam 옵티마이저에 비해 약 2배 빠른 학습 수렴 속도를 보여주며, 대규모 언어 모델(LLM) 사전 학습에서 매우 효율적인 옵티마이저로 부상했습니다. 최근 연구들은 이러한 이점을 연상 메모리(associative memory)나 데이터의 롱테일 특성(data long-tailedness)과 같은 메커니즘 덕분이라고 설명해 왔으나, Muon의 우수성을 뒷받는 국소적 기하학적 원인은 여전히 불분명합니다. 본 논문은 다음 두 가지 핵심 질문을 다룹니다:

어떤 구체적인 최적화 지형(optimization landscape)의 특성이 Adam 대비 Muon의 이점을 지원하는가?
사전 학습 요소들, 특히 학습 데이터 분포와 모델 구조는 이러한 특성에 어떠한 영향을 미치는가?

방법론

저자들은 LLM 학습 궤적에 대한 경험적 분석과 정형화된 이차 형식(quadratic) 최적화 문제에 대한 이론적 분석을 결과를 결합하여 사용합니다.

경험적 분석:

2차 테일러 전개(Second-Order Taylor Decomposition): 저자들은 테일러 전개를 사용하여 1차 이득(그래디언트 정렬)과 2차 곡률 페널티(curvature penalty)로 분해함으로써 1단계 손실 감소량( $\Delta D$ )을 분석합니다: $\Delta D \approx \langle G, Z \rangle - \frac{1}{2}\langle Z, H[Z] \rangle$ .
메트릭 정의: 업데이트 스케일로부터 업데이트 방향의 효과를 분리하기 위해, 저자들은 **정규화된 방향성 샤프니스(Normalized Directional Sharpness, NDS)**를 도입합니다. 이는 $SF(W; Z) = \langle Z, H[Z] \rangle / \|Z\|_F^2$ 로 정의됩니다. 이를 통해 곡률 페널티를 업데이트 노름(update norm)과 NDS로 분리할 수 있습니다.
실험 설정: 실험은 FineWeb 데이터셋으로 학습된 124M 파라미터 규모의 NanoGPT 모델을 대상으로 수행되었습니다. 공정성을 보장하기 위해 검증 손실(validation loss) 수준을 동일하게 맞춘 상태에서 Muon과 Adam을 비교하였습니다.
데이터 및 구조 변형:
- 데이터 불균형: 데이터 불균형 수준(Zipf 지수 $s$ )을 제어하는 Zipf-확률적 문맥 자유 문법(PCFG)을 사용하여 합성 데이터를 생성하여 헤비 테일(heavy-tailed) 분포의 영향을 테스트했습니다.
- 레이어 분해: 레이어 내(within-layer, 대각 헤시안 블록) 기여와 레이어 간(cross-layer, 비대각 상호작용) 기여로 NDS를 분해하여, 서로 다른 모델 구성 요소들이 곡률 이점에 어떻게 기여하는지 이해합니다.

이론적 분석:

정형화된 이차 모델: 저자들은 네 가지 경험적으로 검증된 가정(아래 참조)을 바탕으로 LLM 지형의 특성을 반영하는 이차 최적화 문제를 구축합니다:
1. 낮은 크로네커 랭크(Low Kronecker-Rankness): 헤시안이 낮은 랭크의 크로네커 근사를 허용함.
2. 동시 대각화(Simultaneous Diagonalization): 크로네커 인자들이 공통의 직교 기저를 공유함.
3. 곡률 이질성(Curvature Heterogeneity): 헤시안 스펙트럼이 강한 이질성(고윳값의 롱테일 분포)을 보임.
4. 그래디언트 정렬(Gradient Alignment): 그래디언트가 고곡률 고유 방향(high-curvature eigendirections)과 강하게 정렬됨.
비교: 이론적 증명에서는 스펙트럼 정규화된 업데이트인 Muon을 그래디언트 하강법(GD)과 비교합니다. Adam은 이 특정 정형 설정에서 GD와 유사하게 동작하므로, 저자들은 스펙트럼 정규화와 그래디언트 기반 업데이트 사이의 대조에 집중할 수 있습니다.

주요 기여 및 결과

1. Muon 우위의 원천: 감소된 곡률 페널티
분석 결과, Muon이 Adam보다 더 큰 1단계 손실 감소를 달하는 주된 이유는 더 큰 1차 그래디언트 이득을 얻기 때문이 아니라, 훨씬 더 작은 2차 곡률 페널티를 부담하기 때문입니다.

관찰 1: 검증 손실 수준이 동일할 때, Muon과 Adam은 유사한 1차 감소량을 보이지만, Muon의 곡률 페널티는 실질적으로 더 낮습니다.

2. 메커니즘: 낮은 정규화된 방향성 샤프니스(NDS)
감소된 곡률 페널티는 업데이트 스케일이 아닌 Muon의 업데이트 방향에 의해 유도됩니다.

관찰 2: Muon과 Adam은 유사한 업데이트 노름( $\|Z\|_F$ )을 가집니다. 그러나 Muon은 현저히 낮은 NDS를 나타냅니다(평균적으로 Adam보다 약 1.76배 낮음). 이는 Muon의 스펙트럼 정규화가 업데이트를 국소 곡률이 낮은 방향으로 유도함을 시사합니다.

3. 데이터 불균형의 역할
Muon의 NDS 이점은 데이터 불균형에 의해 증폭됩니다.

관찰 3: Zipf-PCFG 데이터를 사용한 결과, 데이터셋이 더 불균형해질수록(높은 Zipf 지수 $s$ ) 두 옵티마이저 모두 NDS가 증가하지만, Adam과 Muon 사이의 격차는 유의미하게 벌어집니다( $s=0$ 에서 $s=1$ 로 갈 때 1.8배 증가). Muon은 헤비 테일 데이터에 의해 유도되는 샤프니스 증가에 더 강건합니다.

4. 모델 구조의 역할
곡률 이점은 학습 과정에 따라 변화합니다.

관찰 4: 레이어 내(within-layer) 및 레이어 간(cross-layer) NDS 성분 모두 Muon에서 더 낮게 나타나지만, 학습 중기 및 후기로 갈수록 레이어 내 성분이 Muon의 이점에서 점점 더 지배적이 됩니다. Muon의 레이어 간 기여도는 급격히 떨어지는 반면, Adam은 레이어 내 및 레이어 간 샤프니스의 균형을 안정적으로 유지합니다.

5. 이론적 정당화
본 논문은 곡률이 이질적이고 그래디언트가 고곡률 모드에 정렬된 이차 문제에서 다음을 증명합니다:

Muon은 스펙트럼 정규화를 통해 업데이트 에너지를 곡률 그룹 전체에 균형 있게 배분하는 반면, GD는 에너지를 고곡률 모드에 집중시키기 때문에 평균적으로 더 작은 NDS를 달성합니다.
곡률 이질성이 충분히 강할 때, Muon은 동일한 스텝 수 이후에 GD보다 더 낮은 국소 이차 손실(local quadratic loss)을 달합니다.

의의 및 주장

본 논문은 Muon의 우월성을 국소 곡률 관점에서 규명하는 "첫 걸음"을 뗐다고 주장합니다. 핵심 결론은 Muon의 효율성이 스펙트럼 정규화된 업데이트를 수행하여 본질적으로 낮은 방향성 샤프니스를 갖게 함으로써, 최적화 단계의 2차 비용을 최소화하는 능력에서 비롯된다는 것입니다.

저자들은 이 연구가 기존의 "메모리 중심" 설명(예: Wang et al., 2025b)을 보완하는 작업임을 명시하며, 구별되는 기하학적 메커니즘을 제시합니다: 즉, 동일한 반복(iterate) 지점에서 Muon은 Adam보다 더 작은 방향성 샤프니스 페널티를 지불합니다. 본 연구는 인과적 LLM(causal LLMs)에 초점을 맞추고 국소적 메커니즘을 식별하는 데 집중하며, 전역적 수렴 보장이나 확산 모델(diffusion models)과 같은 모든 모델 클래스에 대한 적용성을 주장하지 않는다는 점에서 그 범위를 겸손하게 설정하고 있습니다.