TEAM: Temporal–Spatial Consistency Guided Expert Activation for MoE Diffusion Language Model Acceleration
Abstract
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.2× speedup over vanilla MoE dLLM, with negligible performance degradation. Code is released at https://github.com/PKU-SEC-Lab/TEAM-MoE-dLLM.
Lay Summary
Large language models are becoming more powerful, but they are also increasingly expensive and slow to run. A recent type of language model can generate several parts of an answer in parallel instead of writing strictly one word after another, which has the potential to make text generation faster. However, when these models are combined with a design that selects different specialized components for different words, they often do too much work: many components are used repeatedly, while only a small number of words are finally kept in each step. In this paper, we propose TEAM, a method that makes this process more efficient. The key idea is that the model’s choices are often stable over time and similar across nearby words, so we can avoid repeatedly using unnecessary components. TEAM carefully reuses useful information, limits extra computation for less important words, and explores several promising word choices at the same time. Experiments show that TEAM can make these models up to 2.2 times faster while keeping their output quality nearly unchanged.