论文

ProbMoE:混合专家的可微概率路由

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

摘要

专家混合 (MoE) 模型通过每个词元仅激活一小部分专家来进行扩展。然而,训练此类模型仍然具有挑战性,因为 top-$k$ 路由是离散且不可微的,需要梯度估计器来进行专家选择,其设计仍然是一个核心开放问题。我们引入了 ProbMoE,一种概率路由框架,它将专家选择建模为基数约束的专家子集上的分布,并将路由制定为该离散子集空间中的概率推理。我们首先提出 ProbMoE Exact-$k$ 路由,它在前向传递中对 $k$-专家子集进行采样,而后向传递则使用通过每个专家的精确边际概率的梯度作为真实梯度的易处理替代。 ProbMoE 自然地推广到动态 $k$ 路由设置,其中训练和推理将路由基数限制在相同的预定义范围内,从而允许每个词元进行自适应专家分配。在基准和模型主干中,与竞争基准相比,ProbMoE Exact-$k$ 实现了强大的性能,并提高了专家利用率和路由多样性; ProbMoE Dynamic-$k$ 用更少的激活专家即可达到相当的性能。