Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts

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.

한 줄 요약

대규모 MoE 모델의 학습률 최적화를 위해 계산 효율적인 2단계 하이퍼파라미터 전이 프레임워크를 제안한다.

핵심 기여도

MoE 아키텍처에 대한 μP(최대 업데이트 매개변수화) 적응을 제안하여 학습률이 모델 너비 확장에 따라 일관되게 전이됨을 보인다.
토큰 수 증가에 따른 학습률 예측을 위한 선형 스케일링 법칙을 도출하여, 10조 토큰 규모에서도 R²=0.95의 높은 정확도로 학습률을 추정한다.
155B 총 파라미터, 17B 활성 파라미터를 가진 대규모 MoE 기반 모델을 10조 토큰으로부터 학습하여 제안된 프레임워크의 실용성을 검증한다.

핵심 아이디어

기존의 MoE 모델 학습률 최적화는 모델 크기와 토큰 수가 커질수록 계산 비용이 급증하는 문제를 해결하기 위해, μP와 MLA, Muon 옵티마이저를 결합한 새로운 접근법을 제안한다. μP는 모델 너비 확장에 따라 학습률이 일관되게 전이됨을 보장하며, MLA와 Muon은 MoE의 스파스성과 병렬 처리를 효과적으로 활용한다. 이에 따라, 소규모 프록시 모델에서 얻은 학습률을 대규모 모델로 전이할 수 있다.

두 번째 핵심 아이디어는 토큰 수 증가에 따른 학습률 예측을 위한 선형 스케일링 법칙 도입이다. 소규모 모델에서 얻은 학습률 최적값을 기반으로 선형 회귀를 적용하여, 10조 토큰 규모에서도 높은 정확도로 학습률을 예측한다. 이는 기존의 2D 스윕 방식에 비해 계산 비용을 크게 절감한다.

기술적 접근법

**μP 적응**: MoE 아키텍처에 MLA와 Muon 옵티마이저를 결합한 μP를 적용하여, 모델 너비 확장에 따른 학습률 전이 가능성을 검증한다.
**토큰 수 스케일링**: 소규모 프록시 모델에서 얻은 학습률 최적값을 선형 회귀를 통해 10조 토큰 규모로 확장한다.
**대규모 학습**: 155B 총 파라미터, 17B 활성 파라미터를 가진 MoE 기반 모델을 10조 토큰으로부터 학습하여 프레임워크의 실용성을 검증한다.
**하이퍼파라미터**: Muon 옵티마이저, MLA, μP, R²=0.95의 스케일링 정확도, 3.85×10⁻⁴의 최적 학습률 등이 사용된다.

주요 결과

**학습률 전이 정확도**: μP를 통해 모델 너비 확장에 따른 학습률 전이가 일관되게 이루어짐을 확인.
**토큰 수 스케일링 성능**: 10조 토큰 규모에서도 R²=0.95의 높은 정확도로 학습률을 예측.
**대규모 모델 학습 성과**: 155B 총 파라미터, 17B 활성 파라미터를 가진 MoE 모델을 10조 토큰으로부터 학습한 결과, 안정적인 손실 트레이젝토리와 경쟁력 있는 벤치마크 점수를 달성.

의의 및 한계

본 연구는 대규모 MoE 모델 학습에서 하이퍼파라미터 최적화의 계산 비용을 획기적으로 줄이는 기법을 제시하며, μP와 MLA, Muon의 결합을 통해 모델 너비와 토큰 수 확장에 따른 학습률 전이 가능성을 입증한다. 특히, 소규모 프록시 모델에서의 학습률 최적값을 기반으로 대규모 학습에 필요한 하이퍼파라미터를 예측함으로써, 기존의 2D 스윕 방식에 비해 훨씬 효율적인 최적화가 가능하다는 점에서 실용적 의의가 있다.

한편, 현재 연구는 스파스성 확장에 따른 μP 전이 가능성에 대한 명확한 분석이 부족하며, 전문가별 학습률 조정이나 다양한 MoE 구조 및 옵티마이저에 대한 일반화 가능성도 추가 연구가 필요하다.

실용적 활용

본 연구의 프레임워크는 대규모 MoE 모델의 학습률 최적화에 있어 계산 비용을 줄이며 정확한 예측을 가능하게 하므로, LLM 연구 및 산업적 MoE 모델 개발에 유용하게 활용될 수 있다. 특히, 10조 토큰 이상의 대규모 학습 환경에서 하이퍼파라미터 최적화를 효율적으로 수행할 수 있는 기반을 제공한다.