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

TL;DR

MoE diffusion language models suffer from high expert activation per forward pass relative to the number of tokens actually produced. The authors introduce TEAM, a plug-in framework that leverages temporal-spatial consistency in block-diffusion decoding to reduce activations while increasing accepted tokens, via Delayed Caching for Decoded Tokens, Speculative Exploration for Hot Tokens, and Limited Activation for Cold Tokens. On SDAR 30B-A3B, TEAM achieves up to 2.2× speedups with negligible degradation, reducing activated experts per token and boosting tokens-per-forward-pass, with a practical caching strategy that eschews frequent refreshes. This approach enables more efficient MoE dLLMs suitable for latency-sensitive and edge deployments, while preserving model quality.

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.2x speedup over vanilla MoE dLLM, with negligible performance degradation. Code is released at https://github.com/PKU-SEC-Lab/TEAM-MoE-dLLM.

Figures (5)

• Figure 1: Activated experts vs. accepted tokens per forward pass in SDAR 30B-A3B. TEAM decodes more tokens with fewer experts activated in an iteration.
• Figure 2: Temporal-spatial characteristics of expert activation and decoding with the SDAR 30B-A3B model on a prompt from the GSM8K dataset. Results are shown for layers 0, 24, and 47 (of 47). (a) Number of activated experts across decoding iterations. (b) Distribution of experts activated by decoded and masked tokens at step 6 (of 11). (c) Token acceptance positions at each iteration, together with hidden state similarity relative to the subsequent iteration.
• Figure 3: Overview of our proposed TEAM. We apply differentiated expert activation and decoding strategies to tokens within each block. For decoded tokens, redundant computation is reduced through one-step delayed caching. For mask tokens (hot), we adopt aggressive multi-branch speculative exploration to exploit idle compute resources and increase the token acceptance rate. For mask tokens (cold), a double-round routing mechanism is introduced to constrain unnecessary expert activations.
• Figure 4: Expert activation with speculative exploration in SDAR for a response from the GSM8K dataset, measured at layer 24 (of 47).
• Figure 5: Ablation study on the Activated experts Per decoded Token (APT) and speedup compared to the vanilla model.