OrScale: 레이어별 신뢰 비율 스케일링을 이용한 직교화 최적화
요약
OrScale은 기존의 직교화 최적화 기법인 Muon을 확장한 새로운 스케일링 방법입니다. 이 방법은 각 레이어별로 업데이트 크기를 신뢰 비율에 기반하여 정밀하게 제어함으로써, 일반적인 행렬 레이어와 대규모 언어 모델(LLM) 모두에 적용 가능합니다. OrScale은 이론적으로 강력한 수렴 보장과 레이어 적응 하강 이득을 제공하며, 실험적으로 CIFAR-10 및 다양한 규모의 LLM 사전 학습에서 기존 최적화 기법 대비 성능 향상을 입증했습니다.
핵심 포인트
- OrScale은 Muon의 신뢰 비율 개념을 확장하여 레이어별 업데이트 크기를 정밀하게 스케일링합니다.
- 레이어별 분모를 실제 매개변수 공간 방향의 Frobenius 노름으로 측정하여 정확도를 높였습니다.
- 이론적으로 OrScale은 핵 노름 기준에서 $O(1/ ext{sqrt}(T))$의 비볼록 수렴 보장을 제공하며, 레이어 적응 하강 이득을 가집니다.
- 실험 결과, CIFAR-10 및 다양한 규모의 LLM 사전 학습 태스크에서 기존 최적화 기법 대비 성능 향상을 보여주었습니다.
Muon은 행렬 값 업데이트를 직교화하여 신경망 훈련을 개선하지만, 각 레이어의 업데이트 크기는 주로 전역 학습률에 의해 제어됩니다. 우리는 OrScale을 소개합니다. 이는 Muon의 신뢰 비율 확장(trust-ratio extension)으로, 간단한 규칙에 기반합니다: 레이어별 비율의 분모는 실제로 적용될 매개변수 공간 방향의 Frobenius 노름을 측정해야 합니다. 이를 통해 일반적인 행렬 레이어에 대한 OrScale과 언어 모델에 대한 OrScale-LM이 생성되며, 여기서는 Moonlight 형태 스케일링(shape scaling)을 한 번의 레이어별 보정(calibration)과 결합하여 모든 신뢰 비율이 1에서 시작하도록 합니다. 우리는 세 가지 자연스러운 Muon-LAMB 하이브리드가 모양 퇴화 분모(shape-degenerate denominators), 원시 모멘텀 클립 포화(raw-momentum clip saturation), 그리고 디커플링된 가중치 감쇠 폭주(decoupled weight-decay runaway)를 통해 실패하는 이유를 분석하고, 결합된 가중치 감쇠를 가진 실제 업데이트 방향 분모가 이러한 실패를 방지한다는 것을 보여줍니다. 이론적으로 OrScale은 핵 노름 기준(nuclear-norm criterion)에서 $O(1/ ext{sqrt}(T))$의 비볼록 수렴 보장(nonconvex convergence guarantee)을 가지며, 측정 가능한 레이어 이질성 하에서 엄격한 레이어 적응 하강 이득(layer-adaptive descent gain)을 갖고, 초기화 시 muP 스타일 학습률 전이(learning-rate transfer)를 유지하는 보정 속성을 가집니다. 경험적으로 OrScale은 세 개의 시드에 걸쳐 CIFAR-10/DavidNet에서 1위를 차지하며, Muon의 검증 Top-1 정확도를 93.70%에서 94.05%로 향상시키고, OrScale-LM은 125M부터 1.1B 매개변수까지 네 가지 규모 중 세 가지 규모에서 Muon+Moonlight 대비 FineWeb-Edu 사전 학습을 개선하며 모든 규모에서 AdamW를 능가합니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.CL의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기