기술 요약: 향상된 희소 Mixture of Experts를 위한 선택적 Sinkhorn 라우팅 (Selective Sinkhorn Routing)

문제 정의

희소 Mixture-of-Experts (SMoE) 모델은 토큰당 일부 전문가(expert)만을 활성화함으로써 계산 효율성을 유지하면서 모델 용량을 확장할 수 있는 확장 가능한 접근 방식을 제공합니다. 그러나 Softmax 기반 라우팅에 의존하는 기존의 SMoE 모델은 소수의 전문가에게 과도하게 집중되고 다른 전문가들은 제대로 사용되지 않는 "라우팅 붕괴(routing collapse)" 현상을 겪는 경우가 많습니다.

기존의 완화 전략은 일반적으로 다음을 활용합니다:

1. 보조 목적 함수 (Auxiliary Objectives): 불균형한 전문가 사용을 처벌하기 위한 로드 밸런싱 손실(load-balancing losses) 또는 z-loss. 이러한 방식은 목적 함수의 불일치(objective misalignment), 학습 불안정성, 그리고 복잡성 증가를 초래할 수 있습니다.
2. 노이즈 게이팅 (Noisy Gating): 탐색(exploration)을 장려하기 위해 노이즈를 주입하며, 이는 학습 가능한 파라미터나 하이퍼파라미터를 추가합니다.
3. Sinkhorn 기반 라우팅: 최적 운송(optimal transport, OT)을 위해 Sinkhorn 알고리즘을 사용하는 기존 방식들은 운송 맵(transport map)을 가중치 할당이 아닌 전문가 선택(top-k) 용도로만 사용합니다. 또한, 게이팅 행렬이 Softmax 방식과 동일하게 경사 기반 학습을 통해 직접 최적화되지 않거나, Sinkhorn 업데이트의 반복적인 특성으로 인해 상당한 학습 오버헤드를 발생시킬 수 있습니다.

방법론

본 논문은 **선택적 Sinkhorn 라우팅 (Selective Sinkhorn Routing, SSR)**을 제안합니다. 이는 엔트로피 정규화된 최적 운송 문제로 토큰-전문가 할당을 재구성한 가벼운 메커니즘입니다.

핵심 공식화

저자들은 $m$개의 토큰을 $n$개의 전문가에게 할당하는 것을 엔트로피 정규화된 최대 비용 최적 운송 문제로 모델링합니다. 목표는 특정 제약 조건을 만족하면서 호환성을 극대화하는 운송 계획 $\hat{\Pi}$를 찾는 것입니다:

• 양의 값 (Positivity): $\Pi > 0$.
• 행 합계 (Row Sum): 각 토큰은 자신의 전체 질량(mass)을 전문가에게 라우팅합니다 ($\Pi \mathbf{1}_n = \mathbf{1}_m$).
• 열 합계 (Column Sum): 각 전문가는 기대되는 총 부하(total load)를 균등하게 받습니다 ($\Pi^\top \mathbf{1}_m = \frac{m}{n} \mathbf{1}_n$).

비용 행렬 $C$는 게이팅 점수 $S$로부터 유도됩니다. 저자들은 두 가지 변형을 제안합니다:

• 선형 비용 (Linear Cost): $C = S$.
• Softmax 비용 (Softmax Cost): $C_{i,:} = \text{softmax}(S_{,:})$, Sinkhorn 반복 과정 중 수치적 오버플로를 방지하기 위해 사용됩니다.

이전의 OT 기반 방식들이 운송 계획을 오직 선택 용도로만 사용하는 것과 달리, SSR은 $\hat{\Pi}$의 값을 라우팅 가중치를 계산하는 데 직접 사용합니다. 각 토큰에 대해, $\hat{\Pi}$의 해당 행에서 가장 높은 항목들을 기준으로 상위-$k$ 전문가가 선택되며, 가중치는 이 항목들로부터 정규화됩니다.

선택적 적용 전략 및 노이즈 주입

핵심 과제는 $W_g$가 OT 목적 함수와 분리되어 있기 때문에 Sinkhorn 라우팅이 게이팅 가중치 행렬 $W_g$를 업데이트하지 않는다는 점입니다. 이를 해결하기 위해 SSR은 하이브리드 학습 전략을 채택합니다:

• 선택적 적용 (Selective Application): 학습 중에 각 MoE 블록은 작은 확률 $p$ (예: 0.001)로 Sinkhorn 라우팅을 사용하고, 그 외에는 표준 Softmax 게이팅을 사용합니다. 이를 통해 Softmax 게이팅이 $W_g$를 효과적으로 업데이트하는 동시에 Sinkhorn 라우팅이 균형 잡힌 전문가 활용 신호를 제공할 수 있습니다.
• 노이즈 주입 (Noise Injection): 전문가 붕괴를 방지하고 탐색을 장려하기 위해 비용 행렬에 가우시안 노이즈를 추가합니다 ($\tilde{C} = C + \alpha_{noise} \cdot \epsilon$). 이론적 분석(Proposition 4.3)은 이것이 모든 전문가가 선택될 비제로(non-zero) 확률을 보장함을 보여줍니다.

추론 동작 (Inference Behavior)

논문은 단일 입력에 대해 균형 잡힌 제약 조건(예: 동일한 열 합계)을 강제하는 것이 입력의 특정 호환성 점수와 상관없이 균일한 라우팅을 강요하므로 추론 시에는 역효등적이라고 주장합니다. 따라서 추론 시에는 Sinkhorn 라우팅과 노이즈 주입을 모두 비활성화하여, 일관되고 입력 의존적인 예측을 보장하는 결정론적 Softmax 라우팅으로 되돌립니다.

주요 기여

1. 새로운 라우팅 프레임워크: 보조 밸런싱 손실 없이 전문가 활용의 균형을 촉진하기 위해 엔트로피 정규화된 최적 운송과 확률적 노이즈 주입을 통합한 방법.
2. 이론적 통찰: Sinkhorn 기반 라우팅과 노이즈 주입이 탐색과 밸런싱을 도와 학습을 돕지만, 왜곡된 할당을 피하기 위해 추론 시에는 비활성화되어야 함을 입증하는 증명.
3. 효율성 및 성능: SSR이 상태 최첨단(SOTA) 베이스라인 모델에 비해 학습 효율성, 정확도 및 견고성을 개선함을 보여주는 광범적인 평가.

실험 결과

저자들은 언어 모델링 (WikiText-103, Enwik-8) 및 이미지 분류 (ImageNet-1K, ImageNet-A, ImageNet-O, ImageNet-R)에서 SSR을 평가했습니다.

• 언어 모델링:
  • WikiText-103에서 SSR 변형들은 Vanilla SMoE 및 로드 밸런싱 손실, z-loss 또는 노이즈를 사용하는 베이스라인들을 일관되게 능가했습니다.
  • SSR-L w/ noise는 테스트 퍼플렉시티(PPL) 34.367을 달성하여, Vanilla SMoE(35.550) 대비 1.183 감소하였으며, 차순위 베이스라인보다 상당한 격차로 앞섰습니다.
  • 결정적으로, SSR은 최소한의 학습 오버헤드(0.33% – 0.65%)를 발생시킨 반면, 전체 Sinkhorn 기반 SMoE는 **72.47%**의 오버헤드가 발생했습니다.
  • 모멘텀(Momentum) 설정에서 SSR은 Vanilla SMoE의 특징적인 노름 진화(norm evolution)를 유지하면서도 성능을 개선한 반면, 다른 밸런싱 방법들은 성능을 저하시켰습니다.
• 바이트 수준 모델링 (Enwik-8):
  • SSR-S w/ noise는 최고의 테스트 Bits-per-Character (BPC)인 1.128을 달성하여 Vanilla SMoE보다 0.010 BPC 개선되었습니다.
  • 전체 Sinkhorn 기반 SMoE보다 2.28배 더 빨랐습니다.
• 비전 작업:
  • ImageNet-1K에서 SSR-L w/ noise는 Top-1 정확도 77.420을 달성하여 Vanilla SMoE 대비 +2.368의 이득을 얻었습니다.
  • 또한 ImageNet-O 및 ImageNet-R에서 강력한 견고성을 보여주었습니다.
• 절제 연구 (Ablation Studies):
  • Sinkhorn 라우팅을 적용하는 확률 $p$가 매우 중요하며, 매우 작은 값(예: $10^{-3}$ 또는 $10^{-4}$)만으로도 라우터를 과도하게 제약하지 않으면서 충분한 밸런싱 신호를 제공할 수 있었습니다.
  • 노이즈 주입은 다양한 설정에서 일관되게 성능을 향상시켰습니다.
  • 추론 시 밸런싱 메커니즘을 비활성화하는 것이 최상의 결과를 냈으며, 이는 이론적 주장을 확인해 줍니다.

의의 및 주장

본 논문은 Selective Sinkhorn Routing이 SMoE 설계를 위한 실용적이고 효과적인 솔루션을 제공한다고 주장합니다. 복잡한 보조 손실을 효율적이고 간헐적인 Sinkhorn 기반 최적화로 대체함으로써, SSR은 다음을 달성합니다:

• 향상된 학습 안정성: 목적 함수 불일치 없이 내장된 전문가 밸런싱을 통해 구현.
• 높은 정확도: 텍스트 및 비전 벤치마크 전반에서 입증됨.
• 견고성: 입력 손상 및 적대적 예제에 대한 향상된 성능.
• 효율성: 보조 손실 방식 및 전체 Sinkhorn 라우팅과 비교하여 최소한의 계산 오버헤드 유지.

저자들은 이 방법이 가볍고, (노이즈 변형의 경우) 추가적인 학습 가능 파라미터가 필요 없으며, 학습 단계 동안 균형 잡힌 활용을 보장하면서도 전문가 선택의 유연성을 유지한다는 점을 강조합니다.