arXiv:2602.09001cs.LG2026-02被引 2

用狄利克雷分布重构专家路由,让大模型更高效地选择并分配专家任务。

DirMoE: Dirichlet-routed Mixture of Experts

  • 将专家选择与贡献分配解耦,分别用伯努利和狄利克雷分布建模
  • 训练时可精确控制激活专家数量,实现端到端可微优化
  • 适合追求高效率、强专家分工的大规模语言模型研究者

混合专家(MoE)模型在大规模语言模型中表现出色。现有路由机制通常依赖非可微的Top-k+Softmax,限制了性能与可扩展性。我们指出,标准方法将激活哪些专家和如何分配贡献这两个决策混淆在一起。为此提出狄利克雷路由的混合专家(DirMoE),基于狄利克雷变分自编码器框架,从根本上解耦核心路由问题:专家选择由伯努利成分建模,所选专家间的贡献分配由狄利克雷成分处理。整个前向传播通过Gumbel-Sigmoid松弛实现专家选择的可微性,以及狄利克雷分布的隐式重参数化。训练目标为变分ELBO,包含直接的稀疏性惩罚项,精确控制预期激活专家数,并配合关键超参数调度,引导模型从探索态过渡到确定性路由状态。实验表明,该方法在性能上匹配或超越其他方法,同时提升了专家专业化程度。

原文摘要 · Abstract (English)

Mixture-of-Experts (MoE) models have demonstrated exceptional performance in large-scale language models. Existing routers typically rely on non-differentiable Top-$k$+Softmax, limiting their performance and scalability. We argue that two distinct decisions, which experts to activate and how to distribute expert contributions among them, are conflated in standard Top-$k$+Softmax. We introduce Dirichlet-Routed MoE (DirMoE), a novel end-to-end differentiable routing mechanism built on a Dirichlet variational autoencoder framework. This design fundamentally disentangles the core routing problems: expert selection, modeled by a Bernoulli component, and expert contribution among chosen experts, handled by a Dirichlet component. The entire forward pass remains fully differentiable through the use of Gumbel-Sigmoid relaxation for the expert selection and implicit reparameterization for the Dirichlet distribution. Our training objective, a variational ELBO, includes a direct sparsity penalty that precisely controls the number of active experts in expectation, alongside a schedule for key hyperparameters that guides the model from an exploratory to a definitive routing state. Moreover, our DirMoE router matches or exceeds other methods while improving expert specialization.

混合专家可微路由稀疏模型语言模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。