提出可微概率路由，让专家模型训练更高效稳定。
ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts
- 将专家选择建模为概率分布，通过边际概率近似梯度。
- Exact-$k$路由在多个基准上表现优于基线，专家利用率更高。
- 动态$ k $路由能自适应分配专家，激活数量更少但性能相当。
Mixture-of-Experts（MoE）模型通过每标记仅激活少量专家实现规模扩展。然而，由于top-$k$路由是离散且不可微的，专家选择的梯度估计设计仍是核心挑战。我们提出ProbMoE，一种将专家选择建模为基数约束专家子集分布的概率路由框架，并将其路由问题转化为该离散子集空间中的概率推断。我们首先提出ProbMoE Exact-$k$路由，在前向传播中采样$k$-专家子集，反向传播中使用每个专家的精确边缘概率作为真实梯度的可处理代理。ProbMoE自然推广至动态$ k $路由设置，训练与推理均限制路由基数在相同预定义范围内，支持每标记自适应专家分配。在多个基准和模型主干上，ProbMoE Exact-$k$相较于竞争基线表现出强劲性能，具有更高的专家利用率和路由多样性；ProbMoE Dynamic-$k$在激活专家更少的情况下仍保持相当性能。
原文摘要 · Abstract (English)
Mixture-of-Experts (MoE) models scale by activating only a small subset of experts per token. However, training such models remains challenging because top-$k$ routing is discrete and non-differentiable, requiring gradient estimators for expert selection whose design remains a central open problem. We introduce ProbMoE, a probabilistic routing framework that models expert selection as a distribution over cardinality-constrained expert subsets and formulates routing as probabilistic inference in this discrete subset space. We first propose ProbMoE Exact-$k$ routing, which samples $k$-expert subsets in the forward pass, and the backward pass uses gradients through each expert's exact marginal probability as a tractable surrogate for the true gradient. ProbMoE naturally generalizes to a dynamic-$k$ routing setting, where both training and inference constrain the routing cardinality to the same predefined range, allowing adaptive expert allocation per token. Across benchmarks and model backbones, ProbMoE Exact-$k$ achieves strong performance compared to competitive baselines, with improved expert utilization and routing diversity; ProbMoE Dynamic-$k$ achieves comparable performance with fewer activated experts.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。