选择性Sinkhorn路由以提高稀疏专家混合模型
Selective Sinkhorn Routing for Improved Sparse Mixture of Experts
AI总结：
本文提出了一种选择性Sinkhorn路由方法，通过将token到专家的分配问题转化为最优传输问题，并引入约束以确保专家利用率均衡，从而在不依赖辅助平衡损失的情况下提升稀疏专家混合模型的性能。
摘要：
稀疏专家混合模型（SMoE）模型具有可扩展性和计算效率，能够在有限的推理开销下实现模型容量的大幅增加。现有的SMoE方法通常依赖于辅助目标，如负载均衡损失和z损失，或额外的可训练组件如噪声门控。虽然这些技术鼓励专家多样性，但可能会引入目标不一致、增加模型复杂性或带来显著的训练开销，尤其是在基于Sinkhorn的路由方法中。在本文中，我们重新审视token到专家的分配问题作为最优传输问题。我们添加约束以确保专家利用率的平衡。我们证明，即使是最小的基于最优传输的路由也能在不需辅助平衡损失的情况下提升SMoE性能。与以往方法不同，我们的方法直接从传输图中推导出门控分数，从而实现更平衡和有效的token到专家分配。基于这一见解，我们引入了选择性Sinkhorn路由（SSR），一种轻量级的路由机制，它用高效的Sinkhorn路由替代了复杂的辅助损失，同时保持灵活的专家选择。在语言建模和图像分类实验中，SSR在训练效率、准确性和对输入损坏的鲁棒性方面均有所提升。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
