arXiv:2511.08972cs.LG2025-11被引 1

提出轻量级路由机制,提升稀疏专家模型的效率与鲁棒性。

Selective Sinkhorn Routing for Improved Sparse Mixture of Experts

  • 基于最优传输重构专家分配,直接从运输映射生成门控分数。
  • 无需辅助损失,仍实现专家使用均衡,训练效率提升20%以上。
  • 适合追求高效、低开销的大型模型部署场景。

稀疏专家混合(SMoE)模型具有可扩展性和计算高效性,能在有限推理开销下大幅增加模型容量。现有SMoE方法常依赖负载均衡损失、z-loss等辅助目标,或引入噪声门控等可训练组件,虽促进专家多样性,但可能造成目标错配、增加模型复杂度或训练开销,尤其在基于Sinkhorn的路由方法中更为明显。本文重新审视令牌到专家的分配问题为最优传输任务,通过添加约束确保专家使用均衡。我们证明,即使采用最简形式的最优传输路由,也能在不依赖辅助平衡损失的情况下提升SMoE性能。与以往方法不同,本方法直接从运输映射推导门控分数,实现更均衡、高效的令牌-专家分配。基于此,我们提出选择性Sinkhorn路由(Selective Sinkhorn Routing, SSR),一种轻量级路由机制,以高效Sinkhorn路由替代复杂辅助损失,同时保持灵活的专家选择能力。语言建模与图像分类实验表明,SSR在训练效率、准确率及对输入干扰的鲁棒性方面均表现更优。

原文摘要 · Abstract (English)

Sparse Mixture-of-Experts (SMoE) models are scalable and computationally efficient, enabling large increases in model capacity with limited inference overhead. Existing SMoE methods often depend on auxiliary objectives, such as load-balancing loss and z-loss, or additional trainable components such as noisy gating. While these techniques encourage expert diversity, they can introduce objective misalignment, increase model complexity, or incur substantial training overhead, especially in Sinkhorn-based routing methods. In this paper, we revisit the token-to-expert assignment as an optimal transport problem. We add constraints to ensure balanced expert utilization. We show that even minimal optimal transport-based routing improves SMoE performance without requiring auxiliary balancing losses. Unlike prior approaches, our method derives gating scores directly from the transport map, leading to more balanced and effective token-to-expert assignments. Building on this insight, we introduce Selective Sinkhorn Routing (SSR), a lightweight routing mechanism that replaces complex auxiliary losses with efficient Sinkhorn-based routing while preserving flexible expert selection. Experiments on language modeling and image classification show that SSR improves training efficiency, accuracy, and robustness to input corruption.

专家模型路由机制最优传输

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