On the Width Scaling of Neural Optimizers Under Matrix Operator Norms I: Row/Column Normalization and Hyperparameter Transfer

이 논문은 행렬 연산자 노름의 관점에서 신경망 옵티마이저의 폭 확장성을 분석하고, 층별 합성 가능한 평균 정규화 노름을 도입하여 폭에 무관한 학습률 전이를 가능하게 하는 새로운 옵티마이저 MOGA 를 제안하며, 대규모 사전 학습에서 Muon 과 경쟁력 있으면서도 더 빠른 성능을 입증합니다.

핵심

🏗️ 1. 문제 상황: "집을 키우면 문이 좁아지는 이유"

상상해 보세요. 여러분이 작은 오두막을 짓고 있습니다. 이때 문 (학습률) 은 1 미터 정도면 충분합니다. 하지만 이 집을 100 층짜리 마천루로 키우려 할 때, 문이 그대로라면 어떻게 될까요?

• 기존의 문제: 기존 최적화 알고리즘 (AdamW, Muon 등) 은 건물의 크기 (모델의 너비, $w$) 가 커질수록 문이 상대적으로 너무 좁아지거나 너무 넓어집니다.
  • 문이 너무 좁으면 (학습률이 너무 작으면) 사람들이 (데이터가) 천천히 지나가서 공사 (학습) 가 매우 느려집니다.
  • 문이 너무 넓으면 (학습률이 너무 크면) 사람들이 서로 부딪혀 넘어지고 (발산), 건물이 무너집니다.
• 현실: 그래서 연구자들은 건물을 키울 때마다 문 크기를 일일이 다시 재고 조정해야 했습니다. 이는 매우 비효율적이고 비용이 많이 듭니다.

🔍 2. 새로운 관점: "기하학적 눈으로 보기"

저자들은 이 문제를 **기하학 (Geometry)**의 관점에서 바라봤습니다.
"우리가 모델을 업데이트할 때, 어떤 '규칙'을 따라 움직이는가?"

• 기존 방식: 대부분의 알고리즘은 "가장 가파른 경사"를 따라 내려가는 방식 (최강 하강법) 을 사용하지만, 이때 사용하는 '거리 측정 도구' (노름, Norm) 가 건물의 크기에 따라 왜곡되었습니다.
• 저자의 통찰: 건물의 층이 높아질수록, 층과 층 사이의 연결 고리가 늘어나는데, 기존의 측정 도구는 이 연결 고리에서 발생하는 '왜곡'을 잡아내지 못했습니다. 마치 1 층에서는 평평한 땅으로 보이지만, 100 층으로 올라가면 땅이 기울어져 보이는 것과 같습니다.

💡 3. 해결책: "평균화된 자 (Mean-Normalized Ruler)"

저자들은 새로운 측정 도구, 즉 **'평균화된 자 (Mean-Normalized Norm)'**를 개발했습니다.

• 비유: 기존의 자는 건물이 커질수록 눈금이 늘어나서 (너비가 $w$일 때 $w$배 커짐) 거리를 잘못 재게 했습니다.
• 새로운 자: 이 새로운 자는 건물이 커질수록 눈금을 자동으로 보정합니다. "아, 건물이 2 배 커졌구나? 그럼 내 눈금도 2 배로 줄여서 실제 거리는 똑같이 재자!"라고 말합니다.
• 결과: 이 자를 사용하면 건물이 작든 크든, 문 (학습률) 의 크기를 똑같이 유지해도 됩니다. 100 층 건물이든 1000 층 건물이든, 처음에 정한 문 크기가 그대로 통합니다.

🚀 4. 제안된 방법: "MOGA (모가)"

이 이론을 바탕으로 저자들은 **MOGA(Matrix Operator Geometry Aware)**라는 새로운 최적화 기법을 제안했습니다.

• 핵심 아이디어: 모델의 각 층에서 행렬 (데이터) 의 **행 (Row)**이나 **열 (Column)**을 평균화하여 정규화하는 것입니다.
• 특히 '행 정규화 (Row Normalization)'가 강력합니다:
  • 기존에 인기 있던 Muon이라는 방법은 건물이 커질수록 바닥이 점점 울퉁불퉁해져서 (수학적 용어로 '스무스함'이 떨어짐) 학습이 불안정해질 수 있다는 이론적 한계가 있었습니다.
  • 반면, MOGA(행 정규화) 는 건물이 커져도 바닥이 항상 평평하게 유지되도록 설계되었습니다.

🧪 5. 실험 결과: "실제 건물에서 증명되다"

저자들은 GPT-2 와 LLaMA 같은 거대 언어 모델을 이용해 실험했습니다.

1. 학습률 이동 (Transfer): 작은 모델 (GPT-2 Small) 에서 찾은 최적의 문 크기 (학습률) 를 그대로 큰 모델 (GPT-XL) 에 적용했습니다. 결과는? 완벽하게 작동했습니다. 다시 튜닝할 필요가 없었습니다.
2. 속도와 안정성:
   • Muon 과 비교: Muon 과 비슷한 성능을 내면서도, 특히 학습이 끝날 때쯤 (손실 값이 낮아질 때) 더 빠르고 안정적이었습니다.
   • AdamW 와 비교: 기존 표준인 AdamW 보다 훨씬 빠르게 수렴했습니다.

📝 요약: 왜 이것이 중요한가요?

1. 시간과 돈 절약: 모델을 키울 때마다 학습률을 다시 찾아야 하는 번거로움이 사라집니다. 작은 모델에서 실험한 결과를 큰 모델에 바로 적용할 수 있습니다.
2. 더 큰 모델, 더 빠른 학습: 건물이 커질수록 학습이 불안정해지는 문제를 해결하여, 더 큰 규모의 AI 모델을 더 효율적으로 훈련할 수 있게 합니다.
3. 이론적 근거: 단순히 "실험해보니까 잘 됐다"가 아니라, "왜 잘 되는지"에 대한 수학적이고 확실한 이유 (기하학적 안정성) 를 제시했습니다.

한 줄 결론:
이 논문은 **"AI 모델의 크기가 커져도 학습 속도를 조절할 필요가 없게 만드는, 더 똑똑하고 안정적인 새로운 나침반 (MOGA)"**을 개발했다고 말할 수 있습니다.

기술 요약

1. 문제 정의 (Problem Statement)

현대 딥러닝의 핵심 과제 중 하나는 모델의 규모 (scale) 가 커짐에 따라 최적화 알고리즘이 어떻게 작동해야 하는지입니다.

• 현황: AdamW, Muon 과 같은 널리 사용되는 옵티마이저들은 네트워크 너비 ($w$) 가 증가함에 따라 최적의 학습률이 크게 변합니다. 예를 들어, 512 은닉 유닛으로 튜닝된 학습률은 2048 유닛으로 확장될 때 발산하거나 수렴 속도가 현저히 느려집니다.
• 근본 원인: 기존 옵티마이저들은 네트워크의 기하학적 구조 (forward map) 와 최적화 기하학 (optimizer geometry) 간의 불일치로 인해, 너비가 커질 때 Lipschitz 상수나 $L$-smoothness(기울기의 매끄러움) 상수가 왜곡됩니다.
• 목표: 네트워크 너비가 변해도 최적의 학습률이 안정적으로 전이 (transfer) 될 수 있는, 너비 불변 (width-independent) 인 최적화 방법론을 설계하는 것입니다.

2. 방법론 (Methodology)

2.1. 최적화자의 통합적 관점: 행렬 연산자 노름

저자들은 다양한 옵티마이저 (SignSGD, AdamW, Muon 등) 를 행렬 연산자 노름 (Matrix Operator Norms) 하에서의 **가장 가파른 하강 (Steepest Descent)**으로 해석합니다.

• 벡터 공간의 $\ell_p$ 노름을 행렬 공간으로 확장하여, $\ell_p \to \ell_q$ 연산자 노름을 도입합니다.
• 예를 들어, AdamW 는 $\ell_1 \to \ell_\infty$ 노름 하의 가파른 하강으로, Muon 은 $\ell_2 \to \ell_2$ (스펙트럼 노름) 하의 가파른 하강으로 해석됩니다.

2.2. 기존 방법의 한계: 층간 호환성 부족

기존의 표준 $p \to q$ 연산자 노름 ($p \le q$) 은 여러 층을 거치면서 너비에 의존적인 Lipschitz 상수를 생성합니다.

• 이유: 인접한 층 사이의 노름 불일치 (mismatch) 로 인해, 작은 섭동도 층을 거치며 증폭됩니다. 특히 $p < q$인 경우, $\|I\|_{q \to p} = n^{1/p - 1/q} > 1$이 되어 너비가 커질수록 안정성이 떨어집니다.

2.3. 핵심 제안: 평균 정규화 연산자 노름 (Mean-Normalized Operator Norms)

이 문제를 해결하기 위해 저자들은 **평균 정규화 노름 (Mean-Normalized Norm)**을 도입합니다.

• 정의: $\|\mathbf{x}\|_{(p, \text{mean})} := n^{-1/p} \|\mathbf{x}\|_p$.
• 효과: 이 정규화는 차원 ($n$) 에 따른 스케일링을 상쇄하여, 인접 층 간의 호환성 조건 ($\|\cdot\|_{\text{in}} \le \|\cdot\|_{\text{out}}$) 을 만족시킵니다.
• 결과: $(p, \text{mean}) \to (q, \text{mean})$ 기하학 하에서는 네트워크의 Lipschitz 상수가 너비와 무관하게 유지됩니다.

2.4. $L$-smoothness 분석 및 MOGA 옵티마이저

• Smoothness 분석: Lipschitz 제어뿐만 아니라 $L$-smoothness(기울기의 변화율) 도 너비와 무관해야 합니다.
  • Muon ($\ell_2 \to \ell_2$): Worst-case 에서 $L$-smoothness 상수가 $O(\sqrt{w})$로 증가하여 너비가 커질수록 최적화가 불안정해질 수 있음을 보였습니다.
  • 새로운 기하학: $(1, \text{mean}) \to (q, \text{mean})$ ($q \ge 2$) 또는 $(p, \text{mean}) \to \infty$ 기하학은 너비와 무관한 $L$-smoothness 를 보장합니다.
• MOGA (Matrix Operator Geometry Aware) 옵티마이저:
  • 위 이론에 기반하여 학습률에 너비 인식 스케일링 규칙을 적용한 새로운 옵티마이저를 제안합니다.
  • 행 정규화 (Row Normalization): $(p, \text{mean}) \to \infty$ 기하학에 기반하며, 학습률 스케일링은 $d_{\text{in}}^{-1/p}$에 비례합니다.
  • 열 정규화 (Column Normalization): $(1, \text{mean}) \to (q, \text{mean})$ 기하학에 기반합니다.
  • $\mu P$와의 관계: Adam/SignSGD 의 경우 MOGA 스케일링은 기존의 $\mu P$ (Maximal Update Parametrization) 스케일링과 일치하지만, MOGA 는 스펙트럼 조건이 아닌 최적화 기하학 (Lipschitz 및 Smoothness) 관점에서 이를 유도하여 더 넓은 범위의 옵티마이저에 적용 가능합니다.

3. 주요 기여 (Key Contributions)

1. 이론적 통찰: 신경망 옵티마이저를 행렬 연산자 노름 하의 가파른 하강으로 통일하여 해석하고, 기존 $p \to q$ 노름이 너비 확장 시 왜곡되는 이유를 층간 노름 불일치로 규명했습니다.
2. 새로운 기하학 제안: 너비 불변의 Lipschitz 및 Smoothness 상수를 보장하는 평균 정규화 연산자 노름을 도입했습니다.
3. Muon 의 한계 규명: Muon 옵티마이저가 너비가 커질수록 Worst-case $L$-smoothness 가 $O(\sqrt{w})$로 증가할 수 있음을 이론적으로 증명하여, 대규모 모델에서의 잠재적 불안정성을 지적했습니다.
4. MOGA 옵티마이저 개발: 행/열 정규화를 기반으로 한 너비 인식 학습률 스케일링 규칙을 포함한 새로운 옵티마이저를 제안했습니다. 이는 $\mu P$를 일반화하면서도 더 넓은 클래스의 옵티마이저에 적용 가능합니다.

4. 실험 결과 (Results)

GPT-2 와 LLaMA 아키텍처를 사용한 대규모 사전 학습 실험을 통해 검증되었습니다.

• 학습률 전이 (Learning Rate Transfer):

  • MOGA(행 정규화) 를 사용하여 GPT-2 Small(1.24 억 파라미터) 에서 XL(15 억 파라미터) 까지 다양한 크기의 모델을 학습시켰습니다.
  • 결과: 모델 크기가 크게 달라져도 최적의 피크 학습률 (peak learning rate) 이 거의 동일하게 유지되었습니다. 이는 작은 모델에서 튜닝된 학습률을 큰 모델에 그대로 적용할 수 있음을 의미합니다.
  • 특히 $p=3$과 같이 $\mu P$의 스펙트럼 가정을 만족하지 않는 경우에도 안정적인 전이가 이루어져, 제안된 프레임워크의 보편성을 입증했습니다.

• 학습 효율성 (Training Efficiency):

  • 표준 토큰 예산: GPT-2 Small 및 LLaMA-130M 에서 AdamW 보다 빠르고 Muon 과 유사한 성능을 보였습니다.
  • 대규모 토큰 예산 (Large Token Budget): Chinchilla 최적 토큰 수의 약 8 배에 해당하는 데이터로 학습 시, MOGA(행 정규화) 는 학습 후반부 (low-loss regime) 에서 Muon 보다 더 빠른 수렴 속도와 더 낮은 손실 값을 기록했습니다. 이는 최적화 안정성이 중요한 대규모 모델 배포 시나리오에서 MOGA 가 우월함을 보여줍니다.

5. 의의 및 결론 (Significance & Conclusion)

• 하이퍼파라미터 전이의 이론적 기반: 이 연구는 단순히 경험적 규칙이 아닌, **최적화 기하학 (Lipschitz 및 Smoothness)**을 통해 학습률 전이의 원리를 설명합니다.
• 실용적 가치: 대규모 언어 모델 (LLM) 의 사전 학습 비용이 매우 높은 상황에서, 작은 모델에서 튜닝된 학습률을 큰 모델에 그대로 적용하여 튜닝 비용을 획기적으로 줄일 수 있습니다.
• Muon 대안: Muon 이 가진 잠재적인 너비 의존성 문제를 해결하면서도, 행 정규화를 통해 표현력 (representational capacity) 과 최적화 안정성 사이의 균형을 더 잘 잡은 새로운 옵티마이저 (MOGA) 를 제시했습니다.

요약하자면, 이 논문은 신경망의 너비 확장에 따른 최적화 불안정성을 행렬 연산자 노름의 기하학적 재정의로 해결하고, 이를 통해 학습률 전이와 대규모 모델 학습 효율성을 동시에 개선하는 실용적인 방법론을 제시했습니다.