面向改进稀疏专家混合模型的选择性 Sinkhorn 路由
机器学习
2025-11-13 v1
摘要
稀疏专家混合模型（SMoE）因其可扩展性和计算效率而日益受到青睐，使模型容量大幅提升而不增加推理成本。然而，现有 SMoE 模型常依赖额外损失（如 z 损失、负载平衡损失）和额外可训练参数（如噪声门控）以鼓励专家多样性，导致目标不一致和模型复杂度增加。此外，现有基于 Sinkhorn 的方法因高度依赖计算昂贵的 Sinkhorn 算法而存在显著训练开销。本文将 token 到专家分配问题形式化为最优输运问题，并引入约束确保专家利用的平衡。我们展示，引入最小程度的最优输运基于路由可提升 SMoE 性能，而无需额外平衡损失。与先前方法不同，本方法直接从输运图中导出门控得分，实现更有效的 token 到专家平衡，既有理论分析也有实证结果。基于这些洞见，我们提出选择性 Sinkhorn 路由（SSR），一种用轻量级 Sinkhorn 基于路由取代额外损失的路由机制。SSR 在保持专家选择灵活性的同时，促进平衡的 token 分配。跨语言建模和图像分类任务，SSR 实现更快的训练、更高的准确率和更强的对输入损坏的鲁棒性。
引用
@article{arxiv.2511.08972,
title = {Selective Sinkhorn Routing for Improved Sparse Mixture of Experts},
author = {Duc Anh Nguyen and Huu Binh Ta and Nhuan Le Duc and Tan M. Nguyen and Toan Tran},
journal= {arXiv preprint arXiv:2511.08972},
year = {2025}
}
备注
14 pages, 6 figures