ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

기술 요약: ProbMoE: Mixture-of-Experts를 위한 미분 가능한 확률적 라우팅

문제 정의
Mixture-of-Experts (MoE) 모델은 토큰당 일부 전문가(expert) 하위 집합만을 활성화함으로써 효율적으로 규모를 확장합니다. 그러나 이러한 모델의 학습은 표준 라우팅이 이산적이고 미분 불가능한 top-$k$ 연산자에 의존하기 때문에 여전히 까다롭습니다. 기존의 학습 절차는 라우터 로짓(logits)에 대한 의존성을 무시하고 선택된 전문가의 소프트맥스 확률을 통해서만 그래디언트를 전파함으로써 이를 우회합니다. 이러한 근사는 대안적인 전문가 하위 집합에 대한 학습 신호를 제한하여, 라우팅 분포가 특정 전문가에 집중되게 만들고, 전문가 활용도를 저하시키며, 학습 역학을 불안정하게 만드는 원인이 됩니다. 이를 완화하려는 기존의 시도들(예: 밀집 역전파 또는 휴리스틱 그래디언트 추정법)은 종종 가능한 전문가 하위 집합에 대한 분포를 명시적으로 모델링하지 못하거나, 학습 과정에서 대안적인 조합을 탐색하는 메커니즘을 제한적으로 제공합니다.

방법론
저자들은 ProbMoE를 제안하며, 이는 전문가 라우팅을 카디널리티(cardinality)가 제약된 전문가 하위 집합에 대한 확률적 추론 문제로 재정의합니다. 라우팅을 결정론적 연산자로 취급하는 대신, ProbMoE는 전문가 선택을 이산적 하위 집합에 대한 분포로 모델링합니다.

1. 확률적 정식화:

   • Exact-$k$ 라우팅: 이 프레임워크는 $|S|=k$인 하위 집합 $S$에 대한 분포를 정의합니다. (라우터 로짓으로부터 유도된 시그모이드 기반의) 독립적인 베르누이 선택 확률을 정확히 $k$개의 전문가가 선택되어야 한다는 제약 조건에 따라 조건화합니다. 이 분포의 정규화 상수 $Z_k$는 동적 계획법(dynamic programming)을 사용하여 정확하게 계산되므로, 이 정식화는 트랙터블(tractable)합니다.
   • Dynamic-$k$ 라우팅: 이 프레임워크는 하위 집합 크기 $|S|$가 $[k_{min}, k_{max}]$ 범위 내에 있는 범위 제약 설정으로 일반화됩니다. 이 분포는 범위에 대해 조건화되어, 모델이 각 토큰에 대해 전문가의 정체성과 라우팅 카디널리티를 공동으로 추론할 수 있게 합니다.

2. 순전파 (Forward Pass):

   • 학습 중, ProbMoE는 정의된 분포(exact-$k$ 또는 range-constrained)로부터 전문가 하위 집합 $S$를 샘플링합니다.
   • MoE 레이어는 샘플링된 전문가만을 사용하여 출력을 계산하며, 이는 표준 MoE 모델의 희소 실행 패턴을 유지합니다.
   • 추론 시에는 최대 사후 확률(Maximum-A-Posteriori, MAP) 집합이 선택되어, 추론 비용이 기존 MoE 라우팅과 일치하도록 합니다.

3. 역전파 (Backward Pass):

   • 그래디언트 기반 최적화를 가능하게 하기 위해, ProbMoE는 SIMPLE 그래디언트 추정기(Ahmed et al., 2023)를 채택합니다.
   • 이 역전파 과정은 이산 샘플을 직접 미분하는 대신, 제약된 분포 하에서 각 전문가가 선택될 **한계 확률(marginal probabilities)**을 사용합니다.
   • 라우팅 가중치는 다음과 같은 스트레이트-스루 추정기(straight-through estimator)를 사용하여 구성됩니다: $w = (\text{stopgrad}(z - m) + m) \odot \pi$, 여기서 $z$는 샘플링된 마스크, $m$은 정확한 한계 확률 벡터, $\pi$는 소프트맥스 라우팅 확률입니다.
   • 이를 통해 그래디언트가 한계 확률을 통해 흐를 수 있게 하여, 선택 확률의 변화가 기대 손실(expected loss)에 어떻게 영향을 미치는지 반영하는 정보가 풍부한 업데이트를 라우터에 제공하는 동시에, 희소한 전문가 실행을 유지합니다.

주요 기여

1. 확률적 추론 정식화: 본 논문은 MoE 라우팅을 휴리스틱한 그래디언트 근사를 넘어 카디널리티가 제약된 전문가 하위 집합에 대한 확률적 추론으로 규정합니다.
2. 트랙터블한 한계 기반 최적화: SIMPLE 추정기를 활용함으로써, 저자들은 확률적 하위 집합 샘플링과 정확한 한계 기반 라우터 그래디언트를 결-합하여, 밀집된 전문가 측 계산 없이도 미분 가능한 학습을 가능하게 합니다.
3. 통합된 Dynamic-$k$ 프레임워크: 이 방법은 휴리스틱한 게이팅 규칙에 의존하지 않고, 모델이 사전 정의된 범위 내에서 전문가의 정체성과 토큰당 활성 전문가 수를 함께 추론할 수 있는 다이내믹 라우팅으로 자연스럽게 확장됩니다.
4. 실증적 개선:
   • Exact-$k$: ProbMoE는 OLMoE 및 Qwen 백본 상에서 GSM8K, Law, Translation, MBPP, MMLU 등의 벤치마크에서 강력한 성능을 보였으며, 종종 DenseMixer나 기존의 Top-$k$와 같은 경쟁적인 베이스라인과 대등하거나 이를 능가하는 성능을 보여줍니다.
   • Dynamic-$k$: ProbMoE Dynamic-$k$는 평균적으로 더 적은 수의 전문가를 활성화하면서도 Exact-$k$와 대등한 성능을 달야내며, 효과적인 적응형 할당을 입증합니다.
5. 라우팅 다양성: 분석 결과, ProbMoE는 기존 베이스라인에 비해 더 넓은 라우팅 분포와 더 균형 잡힌 전문가 활용도를 유도하여 전문가 붕괴(expert collapse)를 완화함을 보여줍니다.

결과 및 의의
본 논문은 ProbMoE가 전문가 활용도와 라우팅 다양성을 개선함을 입증합니다. 절제 연구(ablation studies)를 통해, 확률적 샘플링과 한계 기반 최적화의 결합이 성능 향상의 주요 동력이며, 두 구성 요소 중 하나만 사용하는 설정보다 우수함이 밝혀졌습니다.

저자들은 ProbMoE가 전문가 선택을 위한 원칙적인 방법을 제공하여 학습 안정성과 계산 효율성을 개선한다고 주장합니다. 라우팅을 확률적으로 모델링함으로써, 이 프레임워크는 학습 과정에서 대안적인 전문가 조합을 탐색할 수 있게 하여 더 잘 보정된 라우팅 분포를 이끌어냅니다. Dynamic-$k$ 변형 모델은 적응형 카디널리티가 토큰 수준의 복잡도에 맞춰 학습될 수 있음을 보여주며, 작업 성능을 희생하지 않으면서 실질적인 라우팅 비용을 줄입니다. 본 연구는 향-후 방향으로 dynamic 설정에서의 희소성을 완전히 활용하기 위한 ProbMoE의 사전 학습(pre-training) 확장 및 시스템 레벨 최적화를 제시합니다.