PaperlayerView exact revision on arXiv

Source-linked AI summary

ITC-MoE: Importance-guided Token-aware Compression for MoE Diffusion Language Models

Lianjun Liu, Shipeng Li, You Huang, Weiqi Yan, Mingte Qiu, Huazhong Liu, Xiaofeng Zhu, Yunshan Zhong

arXiv:2610.01296v1cs.AI

TL;DR

MoE DLMs incur substantial computation and storage costs, while existing low-rank compression strategies do not jointly adapt to cross-mode redundancy and token-wise variation. ITC-MoE combines adaptive Tucker compression with token-aware compensation and routing; across five benchmarks it reports higher speedups than TEAM on both models while maintaining comparable task performance.

  • Problem

    MoE DLMs have substantial expert-parameter costs, and existing low-rank MoE compression methods [Li et al., 2025, Gu et al., 2025, Chen et al., 2026, Mi et al., 2026, Wei et al., 2026a] do not jointly accommodate cross-mode redundancy and token-wise utilization variation.

    Source: §1 · p. 1, §3.2 · p. 5

  • Method

    ITC-MoE combines importance-guided adaptive Tucker factorization and rank allocation with low-rank compensation for hot tokens and candidate-expert restriction for cold tokens.

    Source: §1 · p. 2, §3.3.1 · p. 6

  • Results

    Across five benchmarks, ITC-MoE achieves higher speedups than TEAM on both models with comparable task performance; under 30% compression, SDAR-30B-A3B-Chat-b32 reports 96.33 accuracy on MultiArith and up to 7.22× speedup.

    Source: §1 · p. 2, §4.2 · p. 9

  • Takeaways & Limitations

    The reported results indicate that ITC-MoE improves MoE DLM inference efficiency while maintaining model fidelity.

    Source: §1 · p. 2

    ITC-MoE retains a performance gap relative to the full model, and hardware-aware implementations are still needed to translate cold-token routing savings more fully into end-to-end speedup.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

MoE DLMs enable parallel decoding and increased capacity, but their expert parameters impose substantial computation and storage costs. ITC-MoE adapts compression to cross-mode redundancy and token-wise variation, reporting efficient inference while retaining task performance.

1 Introduction

  • MoE DLMs have substantial expert-parameter computation and storage overhead, and existing low-rank compression strategies are not well suited to their distinctive properties.MoE DLMs also generate tokens at multiple positions in parallel within each denoising iteration.Source: §1 · p. 1
  • Redundancy and rank-truncation sensitivity vary across input, output, and expert modes, while hot and cold tokens differ in spectral characteristics and expert utilization.Hot tokens are more sensitive to rank truncation, whereas cold tokens exhibit more concentrated expert activation.Source: §1 · p. 2
  • ITC-MoE combines importance-guided Adaptive Tucker Compression (IATC) with token-aware Compensation and Routing (TCR).IATC transforms and factorizes expert weights with adaptive mode-wise ranks under a fixed parameter budget; TCR compensates hot tokens and restricts candidate experts for cold tokens.Source: §1 · p. 2
  • 96.33 accuracy on MultiArith and up to a 7.22× end-to-end inference speedup are reported for SDAR-30B-A3B-Chat-b32 under a 30% compression budget.The authors report that the combined IATC and TCR framework maintains model fidelity while improving inference efficiency.Source: §1 · p. 2

2 Related Work

Related work spans MoE routing and efficiency methods, including approaches to expert compression. Low-rank methods largely target autoregressive MoE models, while TEAM addresses MoE DLM inference overhead without fully compressing expert footprints.

2 Related Work

  • Diffusion MoE work explores routing adapted to timesteps, token or patch complexity, and token importance, alongside conversion of dense diffusion transformers to sparse MoE.Examples include Switch-DiT [Park et al., 2024], EC-DiT [Sun et al., 2025], Diff-MoE, Race-DiT [Yuan et al., 2025], and Dense2MoE [Zheng et al., 2025].Source: §2 · pp. 2–3
  • MoE efficiency methods include routing, quantization, expert pruning, parameter loading, and low-rank expert compression; TEAM [Wei et al., 2026a] reduces inference overhead but does not fully compress the model footprint.Low-rank approaches described include shared expert bases, novice modules, Tucker decomposition, low-rank compensators, and shared decomposed components.Source: §2 · p. 3
  • The low-rank methods discussed target autoregressive MoE LLMs and have not been adapted to MoE DLMs with bidirectional attention.TEAM [Wei et al., 2026a] addresses MoE DLM decoding by adapting expert activation to token states, but focuses on reducing inference overhead.Source: §2 · p. 3

3 Method

ITC-MoE combines importance-guided multi-mode compression with token-state-specific compensation and routing to address non-uniform redundancy and token-wise utilization in MoE DLMs.

3.1 Preliminaries

  • MoE layers route each input to selected experts; hot masked positions meet a confidence or neighborhood criterion, and remaining masked positions are cold.For position i, the hot-token criterion uses confidence ci,t and threshold τh or proximity to an accepted position; cold tokens are the remaining masked positions.Source: §3.1 · pp. 3–4

3.2 Motivation

  • Input and output modes have strong low-rank structure, expert weights contain cross-expert redundancy, and hot and cold tokens differ in rank sensitivity and expert use.Hot-token features have broader energy distributions and are more sensitive to truncation; cold tokens concentrate energy in fewer ranks and use a smaller set of experts.Source: §3.2 · pp. 4–5

3.3 Importance-guided Token-aware Compression for MoE DLMs

  • ITC-MoE combines Importance-guided Adaptive Tucker Compression (IATC) and Token-aware Compensation and Routing (TCR) to address both parameter redundancy and token-wise utilization variation.IATC adapts compression across modes, while TCR applies distinct strategies to hot and cold tokens.Source: §3.3 · p. 5

3.3.1 Importance-guided Adaptive Tucker Compression

  • IATC weights expert matrices using activation and gradient statistics, Tucker-factorizes them across expert, output, and input modes, and allocates ranks under a fixed parameter budget.It minimizes output reconstruction error to allocate more rank to truncation-sensitive modes; shared-factor execution reuses projections across experts and achieves up to a 2.18× speedup.Source: §3.3.1 · pp. 5–6

3.3.2 Token-aware Compensation and Routing

  • For hot tokens, TCR adds a trainable low-rank residual to the truncated Tucker reconstruction to recover truncation-induced information loss while anchoring outputs to initialization.The regularization term helps prevent overfitting and stabilize optimization.Source: §3.3.2 · pp. 6–7
  • At M = 48, cold-token Top-8 recall is 97.95% versus 88.34% for hot tokens, while candidate-restricted routing achieves up to a 13.0% speedup.Cold-token routing is restricted to a candidate set formed from routing probabilities for the cold-token centroid.Source: §3.3.2 · p. 7, Table 1 · p. 7

4 Experiments

ITC-MoE reports higher speedups than comparison methods while maintaining competitive task performance. Its ablations examine threshold sensitivity, importance-weighted transformation, hot-token compensation, and learned rank allocation.

4.2 Main Results

  • At 30% compression on MBPP, ITC-MoE scores 47.47 versus 34.24 for D2-MoE and 12.45 for TD-MoE.D2-MoE reaches only 0.80× speedup on that setting.Source: Table 2 · p. 8, §4.2 · p. 8

4.1 Experimental Setup

  • Zero-shot evaluation covers five benchmarks on SDAR-30B-A3B-Chat-b32 and LLaDA2.0-mini, using matched decoding settings on one NVIDIA A800 GPU.Comparisons include D2-MoE, TD-MoE, TEAM, and the uncompressed Original model.Source: §4.1 · pp. 7–8

4.2 Main Results

  • Across five benchmarks, ITC-MoE outpaces TEAM in speedup: 3.46×–4.61× versus 1.74×–2.08× on SDAR, and 2.66×–4.53× versus 1.26×–1.72× on LLaDA2.0-mini.ITC-MoE is set to 3% compression for this comparison because TEAM does not compress parameters.Source: §4.2 · pp. 8–9, Table 3 · p. 9

4.3 Ablation Studies

  • With Lh = 3, the best performance occurs at τh = 0.7; varying τh from 0.6 to 0.8 causes only small changes on MBPP and SingleQ.MBPP ranges from 56.81 to 59.92, while SingleQ remains within 95.41–97.25.Source: Figure 3 · p. 9, §4.3 · p. 9
  • At 10% compression, importance-weighted transformation raises MBPP from 45.91 to 59.92 and GSM8K from 83.87 to 87.87.At 20%, MBPP rises from 36.18 to 50.58 and GSM8K from 77.38 to 79.38.Source: Figure 3 · p. 9, §4.3 · pp. 9–10
  • Learned rank truncation outperforms fixed-proportion allocation on MBPP and GSM8K under both 10% and 20% compression ratios.Table 5 compares allocation across the three Tucker modes.Source: Table 5 · p. 10, §4.3 · p. 10
  • Enabling hot-token compensation improves MBPP from 51.75 to 59.92 and SingleQ from 94.49 to 97.25.The comparison uses a 10% compression ratio.Source: Table 4 · p. 10, §4.3 · p. 10

5 Limitations and Future Work

ITC-MoE retains a performance gap relative to the full model, and further hardware-aware implementation is needed to fully translate cold-token routing’s algorithmic reduction into end-to-end speedup.

5 Limitations and Future Work

  • ITC-MoE still trails the full model; cold-token routing reduces selection complexity from O(E) to O(M), but hardware-aware implementation is needed to realize that gain fully in end-to-end speedup.The performance gap may be alleviated by stronger compensation mechanisms.Source: §5 · p. 10

6 Conclusion

ITC-MoE targets cross-mode redundancy and token-wise variation in MoE DLMs through adaptive Tucker compression and token-aware compensation and routing. Experiments across DLMs and datasets report better performance and efficiency than existing methods.

6 Conclusion

  • Experiments across diverse DLMs and datasets report that ITC-MoE achieves superior performance and efficiency compared with existing methods.The framework combines importance-guided adaptive Tucker compression with token-aware compensation and routing.Source: §6 · p. 10
  • IATC incorporates importance into Tucker factorization and adaptively allocates ranks across modes under a fixed parameter budget.Source: §6 · p. 10
  • TCR applies lightweight low-rank compensation to compression-sensitive hot tokens and restricts candidate experts for cold tokens with concentrated routing patterns.Source: §6 · p. 10

AI use statement

The authors describe their use of generative AI for language polishing, code development, and debugging, and outline materials supplied to support reproduction of the main results.

AI use statement

  • Generative AI assisted with language polishing, code development, and debugging, but was not used to generate experimental results or unverified research conclusions.The authors say AI-assisted content was manually checked and verified, and they take responsibility for the paper's content.Source: AI use statement · pp. 10–11

Reproducibility statement

  • Source code, experimental configurations, method and implementation details, and supplementary results are provided to support reproduction of the main results.Source: Reproducibility statement · p. 11

A Appendix

The appendix describes expert similarity measurement, a task-aware weighted reconstruction objective for Tucker compression, and shared-factor execution strategies.

A.1 Centered Kernel Alignment

  • Linear CKA compares centered expert output features; scores lie in [0, 1], are invariant to isotropic scaling, and higher values indicate stronger geometric alignment.The compared feature matrices come from the same calibration samples.Source: §A.1 · pp. 14–15

A.2 Rationale for Importance-Weighted Transformation

  • Importance weighting makes reconstruction errors reflect input activation energy and output task sensitivity, while Tucker factorization finds shared low-rank structure across expert and feature modes.The weighted reconstruction error is a compression surrogate, not an exact expression for task-loss change.Source: §A.2 · p. 16

A.3 Shared-factor Execution Strategy.

  • Shared factors allow identical input projections to be computed once across activated experts and reduce K output projections to one.Source: §A.3 · p. 17

A.4 Rationale for Candidate-restricted Routing

The appendix relates cold-token hidden-state concentration to more consistent routing and motivates candidate-restricted routing using a representative centroid.

A.4 Rationale for Candidate-restricted Routing

  • For SDAR-30B-A3B-Chat-b32 calibration prompts, cold-token average pairwise cosine similarity is 0.4904, versus 0.2986 for hot tokens.The evaluation used GSM8K and MBPP calibration prompts.Source: §A.4 · p. 17
  • Higher pairwise cosine similarity implies lower directional dispersion among cold-token hidden states.Source: §A.4 · pp. 17–18
  • Cold tokens have more similar routing distributions and greater Top-K expert overlap, making their hidden-state centroid a representative basis for selecting candidate experts.Source: §A.4 · p. 18

B Backbone Model Details

ITC-MoE is implemented on TEAM, with model and decoding configurations and hyperparameters summarized in Tables 6–8. Compensation Rank is set to 11, while Lh follows TEAM’s default.

C Additional ablations

The ablations examine compensation regularization and candidate-set size, selecting λ = 0.1 and M = 48 as defaults based on reported performance and efficiency considerations.

Additional ablations

Related papers

Loading 2610.01296v1…