利用 Token 图改进稀疏混合专家模型中的路由
Improving Routing in Sparse Mixture of Experts with Graph of Tokens
AI总结：
针对SMoE模型训练后期的路由波动问题，提出相似性感知和注意力感知(S)MoE方法，利用Token相似性和注意力矩阵引导路由，有效降低专家选择熵，显著提升了模型稳定性与鲁棒性。
摘要：
稀疏混合专家模型（SMoE）已成为在深度学习中实现前所未有的可扩展性的关键。通过每个样本仅激活一小部分参数，SMoE在保持计算开销恒定的同时，实现了参数量的指数级增长。然而，SMoE模型在模型训练后期容易出现路由波动——即给定输入路由到其目标专家的路由发生变化——从而导致模型缺乏鲁棒性。在本研究中，我们通过概率图模型（PGM）的视角揭示了SMoE的局限性。通过这一PGM框架，我们强调了Token在专家选择中的独立性，这使得模型容易出现路由波动和缺乏鲁棒性。为了缓解这种独立性，我们提出了新颖的相似性感知（S）MoE，它在专家选择过程中考虑了Token之间的交互。随后，我们推导出一个支撑（S）MoE-Attention模块的新PGM，而不仅仅局限于单一的（S）MoE层。利用注意力矩阵捕获的Token相似性，我们提出了创新的注意力感知（S）MoE，它采用注意力矩阵来引导Token在（S）MoE中路由到合适的专家。我们从理论上证明了相似性/注意力感知路由有助于降低专家选择的熵，从而产生更稳定的Token路由机制。我们在各种任务和领域上实证验证了我们的模型，结果表明，与通过softmax门控进行Token路由的基线MoE-Transformer相比，在减少路由波动、提高准确性和增强模型鲁棒性方面有显著改善。
英文摘要：
Sparse Mixture of Experts (SMoE) has emerged as a key to achieving unprecedented scalability in deep learning. By activating only a small subset of parameters per sample, SMoE achieves an exponential increase in parameter counts while maintaining a constant computational overhead. However, SMoE models are susceptible to routing fluctuations--changes in the routing of a given input to its target expert--at the late stage of model training, leading to model non-robustness. In this work, we unveil the limitation of SMoE through the perspective of the probabilistic graphical model (PGM). Through this PGM framework, we highlight the independence in the expert-selection of tokens, which exposes the model to routing fluctuation and non-robustness. Alleviating this independence, we propose the novel Similarity-Aware (S)MoE, which considers interactions between tokens during expert selection. We then derive a new PGM underlying an (S)MoE-Attention block, going beyond just a single (S)MoE layer. Leveraging the token similarities captured by the attention matrix, we propose the innovative Attention-Aware (S)MoE, which employs the attention matrix to guide the routing of tokens to appropriate experts in (S)MoE. We theoretically prove that Similarity/Attention-Aware routing help reduce the entropy of expert selection, resulting in more stable token routing mechanisms. We empirically validate our models on various tasks and domains, showing significant improvements in reducing routing fluctuations, enhancing accuracy, and increasing model robustness over the baseline MoE-Transformer with token routing via softmax gating.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
