arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2602.08404cs.CL
收藏

TEAM: 时空一致性引导的专家激活用于MoE扩散语言模型加速

TEAM: Temporal-Spatial Consistency Guided Expert Activation for MoE Diffusion Language Model Acceleration

Institute for Artificial Intelligence(人工智能研究院);School of Integrated Circuits(集成电路学院);Yuanpei College(元培学院);Beijing Advanced Innovation Center for Integrated Circuits(北京集成电路先进创新中心)

Linye Wei, Zixiang Luo, Pingzhi Tang, Meng Li

AI总结:

针对MoE扩散语言模型中专家激活过多导致推理开销大的问题,提出TEAM框架,利用路由决策的时空一致性,通过保守激活和激进推测解码实现加速,最高提速2.2倍且性能损失极小。

摘要:

扩散大语言模型(dLLM)因其天然支持并行解码而近期受到广泛关注。基于这一范式,具有自回归(AR)初始化的混合专家(MoE)dLLM进一步展示了与主流AR模型相媲美的强劲性能。然而,我们发现MoE架构与基于扩散的解码之间存在根本性不匹配。具体来说,每个去噪步骤中激活了大量专家,而最终只有一小部分token被接受,导致大量推理开销,限制了其在延迟敏感应用中的部署。在这项工作中,我们提出了TEAM,一个即插即用的框架,通过更少的激活专家实现更多被接受的token,从而加速MoE dLLM。TEAM的动机源于观察到专家路由决策在去噪层级间表现出强时间一致性,以及在token位置间表现出强空间一致性。利用这些特性,TEAM采用了三种互补的专家激活和解码策略,保守地选择解码和掩码token所需的专家,同时跨多个候选进行积极的推测性探索。实验结果表明,TEAM在性能下降可忽略的情况下,实现了相比原始MoE dLLM高达2.2倍的加速。代码已发布在https://github.com/PKU-SEC-Lab/TEAM-MoE-dLLM。

英文摘要:

Diffusion large language models (dLLMs) have recently gained significant attention due to their inherent support for parallel decoding. Building on this paradigm, Mixture-of-Experts (MoE) dLLMs with autoregressive (AR) initialization have further demonstrated strong performance competitive with mainstream AR models. However, we identify a fundamental mismatch between MoE architectures and diffusion-based decoding. Specifically, a large number of experts are activated at each denoising step, while only a small subset of tokens is ultimately accepted, resulting in substantial inference overhead and limiting their deployment in latency-sensitive applications. In this work, we propose TEAM, a plug-and-play framework that accelerates MoE dLLMs by enabling more accepted tokens with fewer activated experts. TEAM is motivated by the observation that expert routing decisions exhibit strong temporal consistency across denoising levels as well as spatial consistency across token positions. Leveraging these properties, TEAM employs three complementary expert activation and decoding strategies, conservatively selecting necessary experts for decoded and masked tokens and simultaneously performing aggressive speculative exploration across multiple candidates. Experimental results demonstrate that TEAM achieves up to 2.2x speedup over vanilla MoE dLLM, with negligible performance degradation. Code is released at https://github.com/PKU-SEC-Lab/TEAM-MoE-dLLM.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑