通过令牌图改善稀疏专家混合中的路由

稀疏专家混合（SMoE）已成为实现深度学习前所未有的可扩展性的关键。通过对每个样本仅激活少量参数，SMoE在参数数量上实现了指数级增长，同时保持计算开销恒定。然而，SMoE模型在模型训练后期易受路由波动影响——即输入路由到目标专家的路径发生变化——这导致模型不稳健。在本工作中，我们通过概率图模型（PGM）的视角揭示了SMoE的局限性。通过该PGM框架，我们强调了令牌专家选择中的独立性，这使模型容易受到路由波动和不稳健性的影响。为缓解这一独立性，我们提出了新颖的相似度感知（S）MoE，在专家选择时考虑令牌间的交互。随后我们推导出一个新的潜在PGM，涵盖一个（S）MoE-注意力模块，超越单个（S）MoE层。利用注意力矩阵捕获的令牌相似度，我们提出创新的注意力感知（S）MoE，利用注意力矩阵指导令牌在（S）MoE中路由至合适专家。我们理论证明，相似度/注意力感知路由有助于降低专家选择熵，从而实现更稳定的令牌路由机制。我们在多个任务和领域上对模型进行了实证验证，显示在减少路由波动、提升准确率和增强模型稳健性方面，相较基线通过softmax门控进行令牌路由的MoE-Transformer有显著改进。