利用令牌图改进稀疏专家混合模型的路由

稀疏专家混合（Sparse Mixture of Experts, SMoE）已成为深度学习实现空前可扩展性的关键。通过每个样本仅激活一小部分参数，SMoE 在保持恒定计算开销的同时，实现了参数量的指数级增长。然而，SMoE 模型在训练后期容易受到路由波动的影响——即给定输入到其目标专家的路由发生变化——导致模型缺乏鲁棒性。在本工作中，我们从概率图模型（Probabilistic Graphical Model, PGM）的视角揭示了 SMoE 的局限性。通过此 PGM 框架，我们强调了令牌专家选择中的独立性，这使模型暴露于路由波动和非鲁棒性之下。为缓解这种独立性，我们提出了新颖的相似性感知（Similarity-Aware, S）MoE，它在专家选择过程中考虑了令牌之间的交互。随后，我们推导出一个超越单一 (S)MoE 层的、构成 (S)MoE-注意力模块基础的新的 PGM。利用注意力矩阵捕获的令牌相似性，我们提出了创新的注意力感知（Attention-Aware, S）MoE，它使用注意力矩阵来指导 (S)MoE 中令牌到合适专家的路由。我们从理论上证明了相似性/注意力感知路由有助于降低专家选择的熵，从而产生更稳定的令牌路由机制。我们在各种任务和领域上对我们的模型进行了实证验证，结果表明，与通过 softmax 门控进行令牌路由的基线 MoE-Transformer 相比，我们的模型在减少路由波动、提高准确性和增强模型鲁棒性方面均有显著改进。