단계적으로 확장하자: 대규모 전문가 혼합을 위한 연산 효율적 하이퍼파라미터 전이

초록

Mixture-of-Experts(MoE) 아키텍처는 계산 비용의 비례적 증가 없이 모델 용량을 크게 확장한다. 그러나 모델 크기와 토큰 예산 모두가 극단적인 규모에서 스윕(sweep)을 통한 하이퍼파라미터——특히 학습률——최적화는 계산적으로 불가능하다. 본 논문에서는 계산 효율적인 2단계 하이퍼파라미터 전이 프레임워크를 제안한다. 이 프레임워크는 모델 폭(width) 스케일링에 걸친 전이를 통해 대규모 MoE 모델 학습을 위한 최적 학습률을 추정하고, 이후 조 단위 토큰 수평선(horizon)으로 외삽(extrapolation)한다. 첫째, Multi-head Latent Attention(MLA)과 Muon 최적화기를 활용하는 MoE 아키텍처에 대한 Maximal Update Parameterization(μP) 적응 방식을 정식화하여, 최적 학습률이 폭 스케일링 모델 전반에 걸쳐 일관되게 전이됨을 입증한다. 둘째, 예측적 스케일링 법칙을 수립하여 토큰 차원을 따라 이러한 전이성을 확장한다. 제한된 예산 하에서 소형 프록시 모델로부터 도출된 최적값에 선형 회귀를 적용함으로써, 이상적 학습률을 대규모 학습 수평선(예: 10조 토큰)까지 높은 충실도(R²=0.95)로 성공적으로 외삽한다. 결과적으로 이는 대규모 MoE의 광범위한 학습을 위한 최적 학습률을 결정하는 데 소형 모델에 대한 프록시 학습으로 충분함을 시사한다. 제안된 방법론을 적용하여 파운데이션 모델(총 1,550억 개, 활성 파라미터 170억 개)을 처음부터 사전 학습하였으며, 안정적인 학습 및 평가 결과는 최소한의 제거(ablation) 비용으로 전체 규모 대상 모델의 최적 구성을 정확히 예측할 수 있음을 검증한다.

English

Mixture-of-Experts (MoE) architectures significantly expand model capacity without a proportional increase in computational cost. However, optimizing their hyperparameters---particularly the learning rate---at extreme scales of both model size and token budget via sweeping remains computationally prohibitive. In this paper, we propose a compute-efficient, two-step hyperparameter transfer framework that estimates optimal learning rates for training large MoE models by transferring them across scaling model widths, and subsequently extrapolating to trillion-token horizons. First, we formulate a Maximal Update Parameterization (μP) adaptation for MoE architectures utilizing Multi-head Latent Attention (MLA) and the Muon optimizer, demonstrating that optimal learning rates transfer consistently across width-scaled models. Second, we extend this transferability along the token dimension by establishing a predictive scaling law. By applying linear regression to the optimal values derived from small proxy models on limited budgets, we successfully extrapolate the ideal learning rate to massive training horizons (e.g., 10 trillion tokens) with high fidelity (R^2=0.95). Consequently, this indicates that proxy training on small models is sufficient to determine the optimal learning rate for the extensive training of large-scale MoEs. We apply the proposed methodology to pretrain our foundation model (155B total, 17B active parameters) from scratch, and the stable training and evaluation results validate that optimal configurations for full-scale target models can be accurately predicted with minimal ablation costs.