혼합 전문가 레이어의 하이퍼파라미터 전이

초록

본 논문은 Mixture‑of‑Experts(MoE) 레이어를 포함한 트랜스포머 모델을 확장할 때, 폭·깊이·전문가 수·전문가 크기 등 네 가지 스케일 차원을 동시에 조정하면서도 학습 하이퍼파라미터를 작은 모델에서 큰 모델로 안정적으로 전이할 수 있는 새로운 파라미터화 방식을 제안한다. 동적 평균장 이론(DMFT)으로 이론적 근거를 제시하고, 1 B 토큰 예산 하에서 51 M ~ 2 B 파라미터 규모의 모델에 대해 초기 가중치 표준편차와 학습률을 규칙적으로 스케일링하면 학습 안정성과 성능이 유지됨을 실험적으로 검증한다. 또한 짧은 토큰 수로 튜닝한 하이퍼파라미터를 장기간 학습에 그대로 적용해도 경쟁력 있는 결과를 얻는다.

상세 분석

이 연구는 기존의 “max‑update parameterization(µP)”과 “Complete P”가 dense 트랜스포머에 적용된 한계를 넘어, sparse MoE 구조에 특화된 스케일링 법칙을 도출한다는 점에서 의미가 크다. MoE 레이어는 라우터 가중치와 전문가(Expert) 파라미터라는 두 종류의 새로운 파라미터 집합을 도입한다. 라우터는 토큰을 상위 k 개의 전문가에 할당하는 역할을 하며, 전문가 자체는 일반적인 FFN과 동일하게 up‑projection, 비선형, down‑projection으로 구성된다. 논문은 (1) 라우터 가중치 초기화 표준편차를 n_embd^(-γ) 형태로, (2) 전문가 bias는 0, (3) 전문가 up/down 가중치는 각각 n_embd^(-1/2)·α_fnn^(-1)·f_n 형태로 스케일링한다. 여기서 α_ffn은 전문가 내부 hidden 차원의 배수이며, f_n은 폭 스케일링에 따른 보정 인자다. 이러한 규칙은 Adam 옵티마이저의 업데이트 크기가 스케일에 무관하게 Θ(1) 수준을 유지하도록 설계되었다.

동적 평균장 이론(DMFT) 분석은 무한 폭·깊이·전문가 수·전문가 크기 한계에서 네트워크 요약 통계(예: 레이어별 커널, 그라디언트 코릴레이션)가 일정하게 유지된다는 것을 증명한다. 특히, residual stream이 전문가 출력의 평균장에, 전문가 출력이 개별 뉴런의 평균장에 의존하는 3‑level 계층 구조를 밝혀냈다. 이 구조는 전문가 수와 크기를 동시에 늘려도 전체 학습 역학이 변하지 않음을 보장한다.

실험에서는 FineWeb 데이터셋을 1 B 토큰(≈2000 step) 예산으로 사용해, 폭, 깊이, 전문가 수, 전문가 hidden 배수 네 차원을 각각 독립적으로 변형하면서도 동일한 초기 학습률·초기화 표준편차를 적용했다. 결과는 전반적으로 loss 곡선이 겹치며, 최종 검증 손실 차이가 미미함을 보여준다. 특히, 38 M 활성 파라미터(≈51 M 전체) 모델에서 찾은 최적 LR·σ를 2 B 파라미터 모델에 그대로 적용했을 때, dense GPT‑2와 비교해 동일하거나 더 낮은 손실을 기록했다. 라우터 바이어스 기반의 부하 균형 기법은 별도 auxiliary loss 없이도 Load_i≈κ(=n_act/n_exp) 를 유지했으며, 전문가 수가 늘어나도 부하 불균형이 발생하지 않았다.

또한, 고정된 sparsity κ를 유지하면서 전문가 수와 크기를 조절했을 때, 동일 파라미터 예산 하에서는 전문가 수를 늘리는 것이 성능 향상에 더 유리함을 확인했다. 이는 기존 연구(Krajevski et al., 2024)와 일치하며, 전문가 수와 크기 간의 트레이드오프를 명시적으로 보여준다.

요약하면, 이 논문은 (1) MoE 모델에 대한 이론적 스케일링 프레임워크를 제시하고, (2) 실험적으로 작은 모델에서 찾은 하이퍼파라미터를 큰 모델에 그대로 전이해도 학습 안정성과 성능이 보장됨을 증명했으며, (3) 전문가 수 vs. 크기 선택에 대한 실용적인 가이드를 제공한다는 점에서 MoE 기반 대규모 언어 모델 개발에 실질적인 가치를 제공한다.