🤖 machine learning

Convergence Bound and Critical Batch Size of Muon Optimizer

이 논문은 다양한 설정에 걸쳐 Muon 옵티마이저에 대한 이론적 수렴 증명을 제공하고, 가중치 감쇠(weight decay)가 유계 그레디언트 가정 없이도 파라미터 및 그레디언트 노름을 유계로 보장함을 입증하며, 훈련 효율성을 결정하는 임계 배치 크기에 대한 하이퍼파라미터 의존적 하한선을 도출한다.

원저자: Naoki Sato, Hiroki Naganuma, Hideaki Iiduka

게시일 2026-06-09

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Naoki Sato, Hiroki Naganuma, Hideaki Iiduka

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

당신이 거대하고 복잡한 로봇(신경망)에게 사진 속 고양이를 인식하거나 이야기를 쓰는 법을 가르치려 한다고 상상해 보세요. 이를 위해 당신은 로봇의 내부 설정을 어떻게 조정하여 더 나아지게 할지 알려주는 "코치"(최적화 도구, optimizer)가 필요합니다. 오랫동안 표준적인 코치는 매우 신뢰할 수 있는 훈련사였던 AdamW였습니다. 하지만 최근, Muon이라는 새로운 코치가 등장하여 놀라운 결과를 보여주고 있습니다.

이 논문은 마치 탐정 보고서와 같습니다. 왜 Muon이 그렇게 뛰어난지 설명하고, 이를 가장 효율적으로 사용하는 규칙을 제공하려 합니다. 다음은 쉬운 용어로 정리한 내용입니다.

1. 핵심 아이디어: 목록이 아닌 형태를 보는 것

대부분의 코치는 로봇의 두뇌를 거대하고 지저nd한 숫자 목록처럼 취급합니다. 그들은 오차를 보고 이렇게 말합니다. "이 숫자는 높이고, 저 숫자는 낮춰라."

Muon은 다릅니다. Muon은 로봇의 두뇌를 **형태(행렬, matrices)**의 집합으로 바라봅니다.

비유: 구겨진 종이를 펴려고 노력한다고 상상해 보세요.
- 기존 코치들 (AdamW): 종이 위의 개별 점들을 무작위로 밀어서 종이를 평평하게 만들려고 합니다. 효과는 있지만, 다소 무질서합니다.
- Muon: 종이 전체를 봅니다. 종이가 특정한 "접힘"이나 구조를 가지고 있다는 것을 깨닫습니다. 단순히 미는 대신, 종이를 자연스러운 선을 따라 완벽하게 펴기 위해 특별한 기동(직교화, orthogonalization)을 수행합니다. 에러 신호가 얼마나 크거나 노이즈가 심하든 상관없이, 가장 "깨끗한" 방향을 찾아냅니다.

2. "안정성"의 비밀: 브레이크 페달

이 논문은 가중치 감쇠(Weight Decay)(로봇이 너무 날뛰지 않도록 제어하는 브레이크 역할을 함)라는 기능과 함께 Muon을 사용하는 데 있어 결정적인 규칙을 발견했습니다.

발견: 로봇을 밀어붙이는 강도(학습률, Learning Rate)와 브레이크를 밟는 강도(가중치 감쇠, Weight Decay) 사이에는 엄격한 관계가 있습니다.
규칙: 브레이크 없이 너무 세게 밀기만 하면 로봇은 통제 불능 상태가 됩니다. 이 논문은 수학적으로 "미는 힘"이 결코 "브레이크"의 용량을 초과해서는 안 된다는 것을 증명합니다. 구체적으로, 학습률은 1 / 가중치 감사보다 작아야 합니다.
결과: 이 규칙을 따르면 로봇의 설정값이 안전하고 예측 가능한 범위 내에 머물게 됩니다. 로봇이 폭주하거나 미쳐 날뛰지 않습니다. 이는 오차가 작다고 가정할 필요가 없다는 것을 의미하므로 매우 큰 장점입니다. 수학적으로 로봇이 스스로 안정성을 유지하도록 보장하기 때문입니다.

3. "스윗 스팟" 배치 크기 (임계 배치 크기)

이러한 로봇을 훈련할 때, 사진을 한 번에 한 장씩 보여줄 수도 있고, 사진 뭉치를 한꺼번에 보여줄 수도 있습니다(이를 배치, Batch라고 합니다).

문제: 사진을 너무 적게 보여주면 로봇의 학습이 느려집니다. 반대로 너무 많이 보여주면 컴퓨터는 뭉치를 처리하느라 바쁘지만, 로봇의 학습 속도는 크게 빨라지지 않습니다. 가장 적은 컴퓨터 작업량으로 최대의 학습 효과를 얻을 수 있는 '골디락스(딱 적당한)' 지점이 존재합니다. 이를 **임계 배치 크기(Critical Batch Size)**라고 합니다.
Muon의 비법: 이 논문은 Muon을 위한 이 "골디락스" 지점에 대한 공식을 유도해 냈습니다.
- 모멘텀 요소: Muon은 "모멘텀(momentum)" 설정(로봇이 같은 방향으로 계속 움직이게 하는 무거운 플라이휠 같은 것)을 사용합니다. 논문에 따르면 모멘텀을 높이면(플라이휠을 더 무겁게 만들면), "골디락스" 배치 크기가 작아집니다. 즉, 더 작은 배치로도 효율적으로 학습할 수 있습니다.
- 브레이크 요소: 브레이크(가중치 감쇠)를 더 강하게 사용하면, "골디락스" 배치 크기는 커집니다. 최고의 효율을 내기 위해서는 한 번에 더 많은 사진을 보여줘야 합니다.

4. 이것이 왜 중요한가 ( "랭크"의 이점)

이 논문은 Muon이 효율적인 이유가 로봇의 두뇌가 실제로는 보이는 것만큼 복잡하지 않다는 것을 이해하기 때문이라고 설명합니다.

비유: 100x100 격자 형태의 불빛(10,000개의 불빛)을 상상해 보세요. 대부분의 시간 동안, 실제로 켜지는 것은 몇 가지 패턴뿐입니다. 나머지는 그저 노이즈일 뿐입니다.
Muon의 우위: Muon은 "실제" 정보가 저차원(low-rank, 단순한 패턴)이라는 것을 깨닫습니다. 노이즈를 무시하는 것입니다. 이 덕분에 Muon은 올바른 방향을 찾아내는 데 거대한 배치 크기가 필요하지 않습니다. 다른 방법들에 비해 더 작은 배치로도 효과적으로 학습할 수 있어 시간과 에너지를 절약합니다.

요약된 결과

효과가 입증됨: Muon은 네 가지 설정(모멘텀 유무, 가중치 감쇠 유무)에서 성공적으로 수렴(학습)한다는 것이 수학적으로 증명되었습니다.
안정적임: 적절한 학습률과 함께 가중치 감쇠를 사용하면, 오차가 작다고 가정하지 않더라도 로봇이 폭주하지 않음을 보장합니다.
효율적임: 이 논문은 완벽한 배치 크기를 찾기 위한 공식을 제공합니다.
- 높은 모멘텀 = 더 작은 배치를 사용할 수 있음.
- 높은 가중치 감쇠 = 더 큰 배치를 사용해야 함.
실제 검증: 저자들은 이미지 인식(개/고양이)과 언어 모델(텍성 작성)에서 이를 테스트했습니다. 실험 결과, 수학적 모델이 현실과 일치함을 확인했습니다. 즉, 새로운 규칙에 따라 배치 크기를 조절할 때 Muon은 기존 표준인 AdamW보다 더 빠르고 효율적입니다.

요약하자면, 이 논문은 실무에서 잘 작동하던 "블랙박스" 최적화 도구를 열어 그 내부 메커니즘을 설명하고, 그것이 안전하게 사용될 수 있음을 증명하며, 최대 속도를 내기 위해 어떻게 튜닝해야 하는지 정확히 알려줍니다.

기술 요약: Muon 옵티마이저의 수렴 경계 및 임계 배치 크기

문제 정의

Muon 옵티마이저는 AdamW와 같은 표준 옵티마이저의 잠재적 후계자로서 강력한 경험적 성능을 입증해 왔으나, 그 이론적 토대는 여전히 미비한 상태이다. 기존 문헌에는 특히 네스테로프 모멘텀(Nesterov momentum)과 가중치 감쇠(weight decay)가 포함된 실질적인 설정에서 Muon에 대한 엄밀한 수렴 증명이 부족하다. 또한, Muon의 하이퍼파라미터와 임계 배치 크기(학습의 확률적 1차 오라클(SFO) 복잡도를 최소화하는 배치 크기) 사이의 관계는 이론적으로 규명되지 않았다. 본 논문은 엄밀한 수렴 분석을 제공하고 Muon의 임계 배치 크기에 대한 하한선을 도출함으로써 이러한 격차를 해소하고자 한다.

방법론

저자들은 다음 네 가지 실용적인 구성에 대해 Muon을 이론적으로 분석한다:

네스테로프 모멘텀 없이, 가중치 감쇠 없이
네스테로프 모멘텀 포함, 가중치 감쇠 없이
네스테로프 모멘텀 없이, 가중치 감쇠 포함
네스테로프 모멘텀 포함, 가중치 감쇠 포함

이론적 프레임워크:

수렴 분석: 저자들은 평균 기대 핵 노름(nuclear norm)의 상한인 $\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}[\|\nabla f(W_t)\|_*]$ 를 유도한다. 이를 위해 $L$ -매끄러움(smoothness) 및 유계된 그래디언트 분산이라는 표준 가정을 활용한다. 핵심적인 혁신은 모멘텀 행렬 $C_t$ 를 직교화하여 얻은 Muon의 탐색 방향을 분석하는 것이다. 저자들은 전체 파라미터 차원이 아닌, 모멘텀 오차 행렬( $r_1$ )과 탐색 방향( $r_2$ )의 랭크(rank)를 고려한다.
Newton-Schulz 근사: Muon은 실제적으로 정확한 SVD 대신 Newton-Schulz 반복법을 통해 직교화를 근사하므로, 저자들은 이 근사 오차( $\delta(k)$ )가 수렴에 미치는 영향을 분석하며, 이것이 무시할 만한 항만을 도입함을 보여준다.
임계 배치 크기(CBS) 도출: 목표 정밀도 $\epsilon$ 에 도달하는 데 필요한 학습 단계 수 $T$ 를 배치 크기 $b$ 의 함수로 모델링한다. SFO 복잡도( $T \times b$ )를 최소화함으로써, 임계 배치 크기 $b^*$ 에 대한 하한선을 도출한다.

실험적 검증:
저자들은 두 단계 전략을 통해 이론적 예측을 검증한다:

제어된 Full-Muon 실험: 이론적 가정을 엄격히 준수하며 모든 파라미터를 Muon으로 업데이트하는 합성 Teacher-Student 태스크.
실제 규모의 하이브리드 실험: CIFAR 데이터셋에 대해 ResNet-18/VGG-16을 학습시키고, C4 코퍼스에 대해 Llama3.1 (320M)을 학습시킨다. 이 설정에서는 Muon이 행렬 형태의 파라미터에 적용되고 나머지는 AdamW가 처리하는 일반적인 배포 방식을 따른다.

주요 기여

1. 네 가지 변형에 대한 수렴 증명

본 논문은 네 가지 Muon 변형 모두에 대한 수렴 보장을 확립한다. 유도된 평균 기대 그래디언트 노름의 상한은 다음과 같다:
$O\left(\frac{1}{\eta T} + \sqrt{\frac{(1-\beta)r_1}{b}} + \bar{r}\eta\right)$
(네스테로프 및 가중치 감쇠 항에 따라 약간의 변동 있음).

의의: 벡터 기반 옵티마이저에서 분산이 전체 파라미터 수에 따라 스케일링되는 것과 달리, Muon의 경계는 모멘텀 오차 행렬의 랭크 $r_1$ 에 따라 스케일링된다. 경험적 증거는 과잉 매개변수화(over-parameterized)된 영역에서 $r_1 \ll n$ 임을 시사하며, 이는 더 타이트한 이론적 경계를 의미한다.
핵 노름(Nuclear Norm): 분석은 많은 벡터 기반 분석에서 사용되는 Frobenius 노름보다 더 엄격한 기준인 핵 노름을 사용하며, 이는 Muon이 신경망 파라미터의 내재적 행렬 구조를 존중함을 시사한다.

2. 가중치 감쇠의 역할과 안정성

저자들은 가중치 감쇠를 포함하는 것이 유계된 그래디언트라는 제한적인 가정에 의존하지 않고도 파라미터 노름( $\|W_t\|_F$ )과 그래디언트 노름( $\|\nabla f(W_t)\|_F$ )의 **거의 확실한 유계성(almost-sure boundedness)**을 보장함을 증명한다.

안정성 조건: 저자들은 안정성을 위한 필요 조건을 도출한다: 학습률 $\eta$ 와 가중치 감쇠 계수 $\lambda$ 는 $\eta \leq 1/\lambda$ 를 만족해야 한다. 실험 결과는 $\eta$ 가 이 임계값을 초과할 때 학습이 불안정해짐을 확인시켜 준다.
암묵적 정규화: 가중치 감쇠가 점근적 수렴 속도의 차수를 개선하지는 않지만, 노름을 제한함으로써 암묵적 정규화를 제공하며, 이것이 우수한 경험적 성능의 이유임을 설명한다.

3. 임계 배치 크기(CBS) 하한선

본 논문은 하이퍼파라미터에 따른 의존성을 보여주는 임계 배치 크기 $b^*$ 의 하한선을 도출한다:
$b^*_{\text{Muon}} > \frac{9(1-\beta)(1+\sqrt{2\beta})^2 r_1 \sigma^2}{(1-\lambda)^2 \epsilon^2}$

하이퍼파라미터 의존성: 공식은 다음을 드러낸다:
- 모멘텀 ( $\beta$ ): $\beta \to 1$ 로 갈수록 항 $(1-\beta)$ 가 소멸하여, 임계 배치 크기가 감소함을 시사한다.
- 가중치 감쇠 ( $\lambda$ ): $1/(1-\lambda)^2$ 인자는 가중치 감쇠를 높이면 임계 배치 크기가 증가함을 나타낸다.
- 랭크 ( $r_1$ ): 경계는 모멘텀 오차의 랭크에 따라 스케일링되며, 이를 통해 CBS를 그래디언트의 저계수(low-rank) 구조와 연결한다.
질적 스케일링: 저자들은 이 공식이 문제별 상수( $\sigma^2, r_1, \epsilon$ )에 의존하기 때문에 절대적인 예측보다는 질적인 스케일링 법칙을 제공한다는 점을 강조한다. 그러나 서로 다른 Muon 변형 간의 CBS 비율은 $\beta$ 와 $\lambda$ 에 의해서만 결정된다.

결과

이론적 결과:

수렴 속도는 표준 비볼록(non-convex) SGD 유형의 방법( $특정 하이퍼파라미터 스케줄 하에서$ O(T^{-1/4})$)과 일치하지만, 랭크 의존성 덕분에 더 타이트한 분산 항을 갖는다.
Newton-Schulz 근사 오차는 수렴에 제한적인 영향만을 미치며, 반복 횟수에 따라 지수적으로 감소한다.

실험적 결과:

안정성: ResNet-18/CIFAR-10에 대한 실험은 안정성 조건 $\eta \leq 1/\lambda$ 를 검증한다. 그래디언트 노름과 학습 손실은 이 경계 근처에서 최소화된다.
수렴 속도: 네스테로프 모멘텀과 가중치 감쇠를 포함한 Muon은 AdamW, Momentum SGD 및 다른 Muon 변형들에 비해 지속적으로 가장 빠른 손실 및 그래디언트 노름 감소를 달로한다.
임계 배치 크기 검증:
- 하이퍼파라미터 효과: 실험은 네스테로프 모멘텀을 추가하면 CBS가 더 큰 값으로 이동하고, 가중치 감쇠를 높이면 CBS가 더욱 증가함을 확인한다. 예를 들어, $\lambda$ 를 0에서 0.5로 높였을 때 CBS가 $\approx 1024$ 에서 $\geq 4096$ 으로 이동하였으며, 이는 예측된 $1/(1-\lambda)^2$ 스케일링(비율 4.0)과 일치한다.
- 랭크 스케일링: 모델 너비를 변화시킴에 따라 유효 랭크 $r_1$ 이 너비와 함께 스케일링됨을 보였으나, $r_1 \cdot \sigma^2$ 의 곱은 비교적 일정하게 유지되었다. 이는 왜 경험적 CBS 값이 이론적 랭크 의존성에도 불구하고 너비에 따라 평탄하게 나타났는지를 설명한다.
- 효율성: Muon은 특히 작거나 중간 정도의 배치 크기에서 AdamW보다 적은 단계와 낮은 SFO 복잡도를 요구한다.

의의 및 주장

본 논문은 등장하는 옵티마이저인 Muon에 대해 네스테로프 모멘텀과 가중치 감쇠를 동시에 포함하는 최초의 엄밀한 수렴 분석을 제공한다고 주장하며, 이를 통해 이론적 이해의 공백을 메운다.

이론적 통찰: 가중치 감쇠가 단순한 정규화 도구가 아니라, 유계된 그래디언트 가정 없이도 유계성과 안정성을 보장하는 메커니즘임을 입증한다.
실무적 가이드라인: 도출된 CBS 하한선은 실무자들에게 실행 가능한 가이드를 제공한다. 이는 $\beta$ 와 $\lambda$ 를 튜닝하는 것이 최적의 배치 크기에 어떻게 질적으로 영향을 미치는지 명확히 하며, 높은 모멘텀 설정은 작은 배치로도 효율적으로 학습할 수 있는 반면, 높은 가중치 감쇠는 효율성을 유지하기 위해 더 큰 배치를 필요로 함을 시사한다.
검증: 연구는 비전 및 언어 모델링 워크로드 전반에서 이러한 예측을 검증함으로써 이론과 실제 사이의 간극을 메우며, 이론적 스케일링 법칙이 Muon이 행렬 형태의 파라미터에만 적용되는 하이브리드 최적화 설정에서도 유효함을 확인한다.

저자들은 단일 행렬 분석이 레이어별 이질성(heterogeneity)을 명시적으로 모델링하지 않는다는 점을 겸허히 인정하면서도, 그래디언트 노이즈 스케일의 높은 상관관계가 레이어 간에 존재하므로 연구 결과의 일반성을 뒷받침한다고 주장한다. 또한, 실험에서 Muon이 Shampoo보다 우수한 성능을 보였으나, Shampoo에 대한 유사한 경계를 도출하는 것은 여전히 미해결 과제로 남아 있음을 언급한다.