[논문리뷰] Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts

링크: 논문 PDF로 바로 열기

메타데이터

저자: Nayeon Kim, Hojin Lee, Yunju Bak, Jaesun Park, Boseop Kim

1. Key Terms & Definitions (핵심 용어 및 정의)

• MoE (Mixture-of-Experts): 전체 파라미터 수는 크지만, 입력마다 일부 전문가(Expert)만 활성화하여 계산 효율성과 모델 용량을 동시에 확보하는 아키텍처입니다.
• μP (Maximal Update Parameterization): 모델의 크기(Width)가 변하더라도 최적의 하이퍼파라미터(특히 Learning Rate)를 유지하거나 예측할 수 있도록 하는 파라미터 초기화 및 스케일링 기법입니다.
• WSD (Warmup-Stable-Decay): 훈련 과정에서 Learning Rate를 Warmup, Stable, Decay 순으로 조정하는 스케줄러로, 학습의 안정성을 높입니다.
• EMA (Exponential Moving Average): 모델 파라미터의 과거 이력을 지수 이동 평균으로 산출하여 노이즈를 줄이고 훈련 성능을 안정화하는 기법입니다.
• Muon Optimizer: 대규모 모델 훈련에서 AdamW보다 우수한 수렴 속도와 성능을 보여주는 최적화 알고리즘입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 대규모 MoE 모델 훈련 시 필수적인 하이퍼파라미터, 특히 Learning Rate 최적화의 높은 비용과 복잡성 문제를 해결합니다. 기존의 2D 하이퍼파라미터 스윕(모델 규모 및 토큰 수 탐색)은 모델이 커질수록 계산 비용이 기하급수적으로 증가하여 실행 불가능한 수준입니다 [Figure 1]. 또한, 기존의 μP 기반 전이 학습은 주로 밀집(Dense) 모델의 Width 확장에 집중되어 있어, MoE의 핵심 요소인 희소성(Sparsity) 및 다양한 토큰 버짓 환경에서의 전이 가능성이 명확하지 않았습니다. 이에 저자들은 모델 규모 확장에 따른 μP 적용과 토큰 수 증가에 따른 Learning Rate 외삽(Extrapolation)을 결합한 효율적인 프레임워크를 제안합니다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

본 논문은 모델 Width 확장에 최적화된 μP를 MoE에 맞게 조정하고, 토큰 버짓에 따른 Learning Rate Scaling Law를 유도하는 2단계 프레임워크를 제안합니다. 첫 번째 단계에서는 MLA(Multi-head Latent Attention) 및 Muon 옵티마이저를 활용한 MoE 모델에서 μP를 적용하여, 모델 규모가 커져도 최적의 Learning Rate가 일정하게 전이됨을 확인하였습니다 [Figure 3]. 두 번째 단계에서는 짧은 토큰 버짓에서 추정된 최적 Learning Rate 데이터를 바탕으로 Linear Regression을 수행하여 장기 훈련 토큰 버짓(예: 10T 토큰)에서의 최적 값을 예측합니다 [Figure 5]. 이 방법론을 통해 155B 파라미터(17B Active) 규모의 모델을 10T 토큰으로 훈련한 결과, 손실 스파이크 없는 안정적인 학습을 달성하였습니다 [Figure 6]. 정량적으로, 제안된 방식은 $R^{2}=0.95$의 높은 정확도로 최적의 Learning Rate를 예측하며, 기존 모델 대비 유사한 연산량 내에서 더 높은 MMLU-Pro 정확도를 기록하였습니다 [Figure 7, Figure 8].

4. Conclusion & Impact (결론 및 시사점)

본 연구는 고비용의 하이퍼파라미터 스윕 없이도 대규모 MoE 모델의 최적 훈련 설정을 예측할 수 있는 실용적인 2단계 전이 프레임워크를 정립하였습니다. 이 프레임워크는 소규모 프록시 모델 훈련만으로 초대규모 모델의 훈련 환경을 설정할 수 있게 하여, 모델 훈련의 컴퓨팅 자원 효율성을 극대화합니다. 제안된 방법론은 향후 초거대 언어 모델(LLM)을 개발하는 학계 및 산업계에서 훈련 비용을 획기적으로 절감하고 개발 주기를 단축하는 데 기여할 것으로 기대됩니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.