MoE diffusion language models cut computing costs with smart token compression
ITC-MoE: Importance-guided Token-aware Compression for MoE Diffusion Language Models
Artificial Intelligence
Summary
Large language models that use many experts have a problem: they need a lot of computing power and memory. The authors noticed that some parts of these models have more unnecessary data than others, and that different words (tokens) use the model differently. They created a method that shrinks the model smartly by focusing on important parts and treating common and rare words differently. This reduces memory and speeds up the model, while keeping its ability to understand and generate text well.
What this means in practice
- •For large model engineers: Decrease memory use and speed up inference in big language models with multiple experts for more practical deployment.
- •For cloud ai service operators: Run diffusion language models more efficiently by compressing experts and selectively routing tokens, lowering operating costs.
Authors
Lianjun Liu, Shipeng Li, You Huang, Weiqi Yan, Mingte Qiu, Huazhong Liu, Xiaofeng Zhu, Yunshan Zhong
Abstract
Mixture-of-Experts (MoE) Diffusion Language Models (DLMs) offer flexible parallel decoding and increased model capacity, but their large number of expert parameters incurs substantial computation and storage costs. Existing low-rank MoE compression methods largely rely on static factorization and fixed rank allocation, which overlook the distinctive properties of MoE DLMs. Specifically, we identify two properties: cross-mode non-uniform redundancy, where parameter redundancy and sensitivity to rank truncation vary across the input, output, and expert modes, and token-wise utilization variation, where hot and cold tokens exhibit distinct spectral characteristics and expert activation patterns. To address these challenges, we propose ITC-MoE, an Importance-guided Token-aware Compression framework for MoE DLMs. ITC-MoE consists of two complementary components. First, Importance-guided Adaptive Tucker Compression (IATC) incorporates activation and gradient importance into expert weight transformation, jointly factorizes expert weights across multiple modes, and adaptively allocates ranks under a fixed parameter budget. Second, Token-aware Compensation and Routing (TCR) applies lightweight low-rank compensation to compression-sensitive hot tokens and restricts the candidate expert set for cold tokens with concentrated routing patterns. By jointly adapting compression capacity and inference execution to both parameter redundancy and token-wise variation, ITC-MoE substantially reduces the computation and storage costs of MoE DLMs while preserving their generation quality. For example, on SDAR-30B-A3B-Chat-b32, ITC-MoE maintains an accuracy of 96.33% on MultiArith under a 30% compression budget, while achieving up to a 7.22x end-to-end speedup. The code is publicly available at https://github.com/lianjunl13-sudo/ITC-MoE.