Complete-muE: MoE 모델을 위한 최적의 하이퍼파라미터 전이 및 스케일링

우리는 트랜스포머 (Transformer) 블록 내의 Dense FFN (Feed-Forward Network) 및 모든 Mixture-of-Experts (MoE) 설정 간의 하이퍼파라미터 전이 (Hyperparameter Transfer)를 목표로 하는 프레임워크인 Complete-muE를 제안합니다. $μ$P (고정된 아키텍처가 필요함) 또는 SDE (매 스텝당 고정된 토큰 수가 필요함)와 같은 기존 도구들은 Dense에서 MoE로의 전이나 MoE 전체 전문가 (Expert) 스케일링이 아키텍처와 전문가당 토큰 수를 모두 변화시키기 때문에, MoE 설정에서의 하이퍼파라미터 전이 문제를 직접적으로 해결할 수 없습니다. Complete-muE는 두 개의 브릿지 (Bridge) 시스템을 통해 이 과제를 해결합니다: Bridge I은 정규화된 라우터 스케일 (Router Scale)을 가진 활성 너비 (Active-width) $μ$P를 통해 Dense FFN과 Dense MoE 사이를 매핑합니다. Bridge II는 활성화된 전문가 스케일링 (Activated-expert scaling)을 통해 Dense MoE와 Sparse MoE 사이를 매핑하며, 여기서 1차 SDE LR/WD 보정은 상쇄되지만 유계된 잔차 $σ_0$ 이동이 남게 됩니다. 우리가 Complete-muE라고 명명한 결과적인 전이 규칙은 MoE 모델의 활성화된 전문가, 전체 용량 (Total Capacity), 입도 (Granularity), 그리고 공유/그룹 균형 하이브리드 (Shared/group-balanced hybrids)의 변화뿐만 아니라, 일반적인 트랜스포머 모델의 네트워크 너비/깊이, 배치 크기 (Batch Size), 그리고 지속 시간 (Duration)의 변화를 모두 포괄합니다. 광범위한 언어 모델 (Language Model) 및 확산 모델 (Diffusion Model) 사전 학습 (Pretraining) 실험을 통해, Complete-muE가 모델 아키텍처와 파라미터 수에 관계없이 비교적 안정적인 하이퍼파라미터 최적값 (Optima)을 산출함을 확인했습니다. 이는 Bridge II의 비엄격한 SDE 동작과 일치하는 미미한 드리프트 (Drift)만을 동반합니다. 실제로 이 드리프트는 매우 작아서, 단일 Dense 참조 모델에서 튜닝된 하이퍼파라미터가 모든 MoE 구성으로 거의 최적으로 전이됩니다. 즉, "Dense를 한 번만 튜닝하여 모두에게 전이한다 (Tune dense once, transfer to all)"가 Complete-muE의 핵심인 실질적인 레시피입니다. 이를 통해 MoE 모델은 비용이 많이 드는 하이퍼파라미터 탐색 없이도 모델 용량을 확장할 때 Dense 모델 대비 가속화된 수렴 속도를 달성할 수 있습니다.