[Paper] 행렬 연산자 노름 하에서 신경 최적화기의 폭 스케일링에 관하여 I: 행/열 정규화와 하이퍼파라미터 전이

Source: arXiv - 2603.09952v1

개요

이 논문은 많은 인기 있는 신경망 최적화기(예: AdamW, Muon)가 모델의 폭 (w)가 증가함에 따라 불안정해지는 이유를 조사하고, 이를 폭‑인식으로 만드는 원칙적인 방법을 제안한다. 이러한 최적화기들을 특수하게 설계된 행렬 연산자 노름 아래에서 최급강하 단계로 재구성함으로써, 저자들은 폭에 관계없이 안정적인 학습률 스케일링 규칙을 도출하고, 작은 모델에서 큰 모델로 하이퍼파라미터를 원활하게 전이할 수 있게 한다.

주요 기여

• Operator‑norm 재해석: AdamW, Muon 및 관련 옵티마이저를 행렬 노름 하에서 최급강하법(steepest‑descent)으로 바라봄.
• 표준 (p!\to!q) 노름에서 발생하는 composability 문제를 식별하여, 깊은 네트워크에서 폭(width)에 독립적인 보장을 방해함.
• mean‑normalized operator norms (\pmean!\to!\qmean) 도입: 층별로 합성 가능하며 폭에 독립적인 smoothness 경계를 제공함.
• learning‑rate scaling rules 도출: μ‑parameterization (μP)을 특수 사례로 복원하고, 다양한 옵티마이저 계열에 대해 폭 간 하이퍼파라미터 전이(cross‑width hyper‑parameter transfer)를 지원함.
• 이론적 분석을 통해 Muon이 smoothness에서 (\mathcal{O}(\sqrt{w})) 급증을 겪을 수 있음을 보이며, 새로운 row‑normalized 옵티마이저는 이를 회피함.
• MOGA (Matrix Operator Geometry Aware) 제안: 행/열 정규화만을 기반으로 하는 실용적인 옵티마이저.
• GPT‑2 및 LLaMA 사전학습에 대한 실증 검증: MOGA(특히 row‑normalized 변형)가 Muon의 성능을 맞추거나 능가하면서, 대규모 토큰·저손실 환경에서 더 빠르게 동작함.

방법론

1. 옵티마이저에 대한 기하학적 관점 – 저자들은 각 옵티마이저 단계가 매개변수 변화가 네트워크 출력에 미치는 영향을 측정하는 행렬 노름에 대한 지역 최급강하 문제를 푸는 것으로 본다.
2. 연산자 노름 분석 – 먼저 고전적인 (p!\to!q) 노름(예: 스펙트럴, Frobenius)을 살펴보고, 이러한 노름이 층을 가로질러 잘 합성되지 않아 너비에 의존하는 부드러움 상수를 초래함을 증명한다.
3. 평균 정규화 노름 – 합성 가능성을 해결하기 위해, 행이나 열을 평균낸 뒤 (p) 또는 (q) 거듭제곱을 취하는 (\pmean!\to!\qmean) 노름을 정의한다. 이 구성은 신경망의 “층별” 곱 구조를 보존하여 너비에 따라 증가하지 않는 깔끔한 경계를 가능하게 한다.
4. 스케일링 규칙 도출 – 새로운 노름을 최급강하 공식에 대입함으로써, 학습률 스케일링 계수를 명시적으로 얻는다. 이 계수는 노름의 파라미터에만 의존하고 (w)에는 의존하지 않는다. 평균 노름 파라미터를 적절히 설정하면 자연스럽게 μP 스케일링으로 귀결된다.
5. 옵티마이저 설계 (MOGA) – 도출된 노름을 사용하여, 행별 또는 열별 정규화를 그래디언트에 적용한 뒤 기존 Adam 스타일의 모멘트 업데이트를 수행하는 옵티마이저 군을 구현한다. 표준 AdamW 파라미터 외에 추가 하이퍼파라미터는 필요하지 않다.
6. 실험적 평가 – 대규모 언어 모델 사전 학습 실험에서 MOGA 변형을 Muon 및 일반 AdamW와 비교한다. 실험은 다양한 너비, 토큰 예산, 손실 구간을 아우른다.

Source: …

결과 및 발견

• 이론적 보장: 행 정규화 최적화기는 폭에 독립적인 부드러움 상수를 달성하는 반면, Muon은 (\mathcal{O}(\sqrt{w})) 증가를 일으킬 수 있다.
• 속도: MOGA는 Muon의 추가 파라미터당 스케일링 비용을 피하기 때문에, 대형 모델 환경에서 훈련 단계당 실제 시간(벽시계 시간)이 약 5‑10 % 감소한다.
• 하이퍼파라미터 전이: 좁은 모델에서 조정한 학습률이 8배 넓은 모델에서도 바로 적용 가능하여, 도출된 스케일링 규칙을 확인한다.

Practical Implications

• Stable scaling of models: 개발자는 학습률을 다시 조정하지 않고도 더 넓은 트랜스포머를 훈련할 수 있어 실험 주기를 줄일 수 있습니다.
• Faster large‑scale pre‑training: MOGA의 가벼운 행/열 정규화는 거의 오버헤드가 없으며, 대규모 언어 모델 파이프라인에 매력적입니다.
• Unified optimizer family: 기존 AdamW 코드베이스는 간단한 사전 그래디언트 정규화 단계만 교체하면 MOGA로 업그레이드할 수 있어 옵티마이저 핵심을 다시 작성할 필요가 없습니다.
• Better theoretical footing for optimizer design: 평균 정규화 연산자 노름 프레임워크는 깊은 네트워크의 기하학을 고려한 새로운 옵티마이저 개발을 안내할 수 있으며, 이는 다른 분야(예: 비전, 강화 학습)에서 수렴성을 향상시킬 가능성이 있습니다.
• Cross‑project reproducibility: 팀은 서로 다른 폭의 모델들 간에 단일 하이퍼파라미터 세트를 공유할 수 있어 재현성과 배포가 간소화됩니다.

제한 사항 및 향후 연구

• 완전 연결 층에 대한 가정: 합성 가능성 증명은 행렬 곱셈 층에 의존합니다; 이론을 컨볼루션이나 어텐션 스타일 연산자에 확장하려면 추가 작업이 필요합니다.
• 실험 범위: 실험은 언어 모델(GPT‑2, LLaMA)에 초점을 맞추었습니다. 비전 트랜스포머, 디퓨전 모델, 그래프 신경망에 대한 검증은 아직 진행 중입니다.
• 다른 기법과의 상호 작용: 본 논문은 MOGA가 학습률 워밍업, 그래디언트 클리핑, 혼합 정밀도 훈련과 어떻게 상호 작용하는지 탐구하지 않았으며, 이는 실제 성능에 영향을 줄 수 있는 영역입니다.
• 적응형 노름 파라미터 가능성: 향후 연구에서는 훈련 중에 (\pmean)와 (\qmean) 지수를 동적으로 조정하여 변화하는 곡률을 포착하는 방안을 조사할 수 있습니다.

전반적으로, 이 연구는 수학적으로 기반을 둔 실용적인 폭 인식 최적화 경로를 제공하여 현대 딥러닝 시스템의 보다 원활한 확장을 약속합니다.

저자

• Ruihan Xu
• Jiajin Li
• Yiping Lu

논문 정보

• arXiv ID: 2603.09952v1
• 분류: cs.LG, eess.SY, math.NA, math.OC, stat.ML
• 출판일: 2026년 3월 10일
• PDF: Download PDF