SoftMoE: LLM에서 Mixture-of-Experts를 위한 소프트하고 미분 가능한 라우팅

SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs

희소 Mixture-of-Experts (MoE) 아키텍처는 top-$k$ 라우팅을 통해 일부 전문가만 활성화하여 고정된 추론 예산 내에서 LLM의 파라미터를 확장할 수 있도록 합니다. 이는 인과 관계를 유지하고 순환 언어 모델에 적합하지만, 이산적인 top-$k$ 연산은 미분 불가능하며, 입력당 고정된 수의 전문가만 활성화되도록 하여 계산 효율성이 떨어집니다. 본 논문에서는 이 문제를 해결하기 위해, 이산적인 라우팅을 잘린 소프트 top-$k$ LapSum 완화 방식으로 대체하여, 전문가 라우팅에 대한 경사 기반 최적화를 가능하게 하는 SoftMoE를 제안합니다. 또한, 각 레이어에서 활성화되는 전문가의 평균 수를 파라미터화하고 전역 예산 제약을 적용하여 모델이 레이어 간에 전문가 용량을 할당하는 방법을 학습하도록 합니다. SoftMoE는 순환 모델링과 완벽하게 호환되며, 언어 모델링 및 후속 작업에서 희소 MoE와 동등하거나 더 나은 성능을 달성하면서도 훨씬 적은 수의 전문가를 활성화합니다. 주목할 만한 점은, 학습된 할당 방식이 매우 불균일하며, 후반 레이어에서 더 많은 전문가가 활성화된다는 것입니다. 소스 코드는 공개적으로 제공됩니다$^ extdagger$.

Sparse Mixture-of-Experts (MoE) architectures enable scaling LLM parameters under a fixed inference budget by activating only a small subset of experts via top-$k$ routing. While this preserves causality and suits autoregressive language models, the discrete top-$k$ operator is not differentiable, forcing a fixed number of active experts per input and resulting in inefficient use of computation. We propose SoftMoE, which replaces discrete routing with a truncated soft top-$k$ LapSum relaxation, allowing gradient-based optimization of expert routing. We further parameterize the mean number of active experts per layer and impose a global budget constraint, enabling the model to learn how to allocate expert capacity across layers. SoftMoE remains fully compatible with autoregressive modeling and achieves performance comparable to or better than sparse MoE on language modeling and downstream tasks, while activating significantly fewer experts. Notably, the learned allocation is highly non-uniform, with later layers activating more experts. The source code is publicly available$^\dagger$.