Mixture‑of‑Experts 모델의 전문가‑주의 할당 최적화: 확장 가능한 스케일링 법칙

초록

본 논문은 MoE 트랜스포머에서 전문가(Feed‑Forward)와 주의(attention) 서브레이어 간의 FLOPs 비율 $r$를 최적화하는 새로운 스케일링 법칙을 제시한다. 실험을 통해 $r^{*}$가 전체 학습 compute와 sparsity $S$에 따라 파워‑law 형태로 변한다는 것을 발견하고, 이를 기존 Chinchilla 법칙에 통합한 확장 모델을 제안한다. 결과적으로 고정된 compute 예산 하에서 효율적인 MoE 설계 지침을 제공한다.

상세 분석

이 논문은 MoE(전문가 혼합) 아키텍처가 대규모 언어 모델에서 파라미터 수는 크게 늘리면서도 토큰당 연산량을 거의 일정하게 유지한다는 점에 주목한다. 그러나 전문가와 주의 서브레이어 사이의 연산 배분은 기존 dense 트랜스포머에서는 거의 고려되지 않았던 새로운 설계 변수이다. 저자들은 이를 정량화하기 위해 FLOPs 비율 $r = C_E / C_A$ 를 정의하고, sparsity $S = (E - E_{\text{act}})/E$ 와 함께 두 변수의 상호작용을 분석한다.

이론적 동기는 “희소 전문가 활성화 하에서의 수확 체감”이다. 낮은 sparsity(많은 전문가가 활성화)일수록 추가적인 전문가 연산이 다양한 서브네트워크에 분산돼 marginal gain이 크지만, 높은 sparsity에서는 같은 전문가에 연산이 집중돼 포화 현상이 빨리 나타난다. 반면 주의 연산은 전역 토큰 상호작용을 담당하므로 sparsity에 크게 의존하지 않는다. 따라서 고정된 전체 compute $C = C_A + C_E$ 하에서 최적의 $r^{*}$는 $C$와 $S$에 따라 변해야 한다는 가설을 세운다.

실험 설계는 두 단계로 이루어진다. 첫째, 여러 모델 규모(수백만~수억 파라미터)와 다양한 sparsity(약 82%~98%)에 대해 $r$을 폭넓게 스윕하면서 동일한 토큰당 compute를 유지한다. 둘째, 각 설정에서 훈련 손실을 측정해 $r$에 대한 최소값을 찾는다. 결과는 손실 표면이 $r$ 축에 대해 명확한 볼트를 형성하며, $C$가 증가할수록 최적 $r^{}$가 꾸준히 오른다는 점을 보여준다. 특히 낮은 sparsity에서는 $r^{}$가 급격히 상승하고, 높은 sparsity에서는 완만하게 변한다.

정량적 분석에서는 $r^{}$가 파워‑law $r^{}= \alpha_r(S) , C^{\beta_r(S)}$ 로 모델링될 수 있음을 확인한다. $\alpha_r$와 $\beta_r$는 각각 $1-S$(활성화된 전문가 비율)의 파워‑law 형태를 보이며, 구체적으로 $\alpha_r = 6.7 \times 10^{-5} (1-S)^{-1.23}$, $\beta_r = 0.24 (1-S)^{0.21}$ 로 추정된다. 이는 sparsity가 증가할수록 $r^{*}$의 절대값은 작아지지만, compute가 커질 때의 성장률($\beta_r$)은 오히려 커지는 역설적인 현상을 설명한다.

이러한 경험적 법칙을 기존 Chinchilla 손실 스케일링식에 통합하기 위해 저자들은 두 개의 페널티 항을 추가한다. 첫 번째는 $r$이 $r^{}$에서 벗어날 때 발생하는 손실 증가를 $d \cdot \frac{r}{r^{}+1}$ 로, 두 번째는 과도한 전문가 할당을 $c \cdot e^{R(1-S)^\gamma}$ 로 모델링한다. 최종 확장식은

$$L = a N^{\alpha} + b D^{\beta} + c , e^{R(1-S)^\gamma} N^{\lambda} + d \frac{r}{r^{*}+1} + \tau$$

이며, 여기서 $N$은 파라미터 수, $D$는 토큰 수, $R$은 전체 compute, $r$은 실제 FLOPs 비율이다. 실험적으로 이 식은 훈련 손실을 높은 정확도로 예측하고, 훈련 중 손실 곡선까지도 잘 맞춘다.

실용적 시사점은 다음과 같다. (1) 고정된 compute 예산 하에서 전문가‑주의 비율을 고정하면 규모가 커질수록 비효율이 누적된다. (2) 설계자는 목표 compute와 선택한 sparsity에 따라 $r^{*}$를 계산해 전문가 레이어에 할당할 FLOPs를 사전에 결정할 수 있다. (3) 제안된 법칙은 MoE 모델을 데이터와 파라미터 규모뿐 아니라 내부 연산 배분까지 최적화하는 새로운 설계 프레임워크를 제공한다.