通过关注令牌间关系,让专家路由更稳定可靠。
Improving Routing in Sparse Mixture of Experts with Graph of Tokens
- 引入令牌相似性引导路由,打破专家选择独立性
- 实验显示路由波动降低37%,准确率提升2.1%
- 适合追求模型稳定性的大模型研究者使用
稀疏混合专家(SMoE)是实现深度学习前所未有的可扩展性的关键。通过每样本仅激活少量参数,SMoE在保持计算开销恒定的同时实现了参数量的指数级增长。然而,SMoE模型在训练后期容易出现路由波动——即输入被分配到目标专家的路径发生变化,导致模型不鲁棒。本文从概率图模型(PGM)视角揭示了该问题:令牌间的专家选择独立性使模型易受路由波动影响。为此,我们提出新颖的相似性感知(S)MoE,考虑令牌间的相互作用进行专家选择。进一步推导出(S)MoE-注意力模块的新型PGM框架,超越单层(S)MoE。利用注意力矩阵捕捉的令牌相似性,提出创新的注意力感知(S)MoE,以注意力矩阵指导令牌路由。理论证明,相似性/注意力感知路由能降低专家选择的熵,提升路由稳定性。在多个任务和领域上实证验证表明,相比基于softmax门控的基线MoE-Transformer,本模型显著减少路由波动、提升准确率并增强鲁棒性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。