Preserving Mathematical Reasoning in Compressed Diffusion Language Models via Trajectory-Aware Low-Rank Approximation

Tian Liang; Zishan Shao; Yiran Chen

cs.AIOct 2, 2026

Preserving Mathematical Reasoning in Compressed Diffusion Language Models via Trajectory-Aware Low-Rank Approximation

Authors: Tian Liang, Zishan Shao, Yiran Chen

Organizations: Department of Electrical and Computer Engineering, Duke University

Abstract

Diffusion language model (dLLM) compression faces a known challenge because calibration is typically performed on clean, fully visible activations, whereas inference traverses partially masked intermediate states. For low-rank compression, this raises two questions. First, can low-rank optimality still be characterized when approximation quality is measured over trajectory-distributed states, and second, does the choice of calibration states affect mathematical reasoning preservation under compression? We address these questions by formulating a trajectory-aware low-rank objective over corruption levels and masking realizations. To estimate this objective efficiently, we propose Traj-MC, which estimates the trajectory second moment through Monte Carlo sampling and yields exact sampled-state optimality and population consistency. Under matched compression budgets, trajectory-aware calibration improves reconstruction over the generation trajectory and preserves substantially more mathematical reasoning than clean calibration on mathematical reasoning benchmarks. Our results connect trajectory-aware low-rank optimality to the reasoning capability retained after dLLM compression. Our code is available at: https://github.com/Zishan-Shao/traj-mc.git.

Figures & tables

Figure 1: Impact of trajectory-aware calibration. (a) Traj-MC reduces second-moment mismatch to generation states (details in Appendix. B ), (b) lowers held-out reconstruction error

E_{\text{gen}}

, and (c) preserves substantially higher mathematical reasoning performance at 20% compression with similar understanding accuracy.

Model	Setting	Method	Math & Reasoning				General Understanding
Model	Setting	Method	GSM8K	MATH-500	SVAMP	ARC-C	MMLU	HellaSwag	PIQA
	Dense	BF16	80.2	30.2	85.7	87.9	64.3	76.6	82.8
		SVD	12.3	4.8	28.8	42.3	39.7	38.2	55.9
		SVD-LLM	42.2	9.0	52.6	77.8	57.2	65.6	78.7
	20%	+Traj-MC	56.7	11.0	68.0	79.9	57.3	65.4	77.7
		SVD	0.0	0.2	0.0	16.7	0.0	0.1	0.4

Table 1: Effect of calibration distribution on low-rank dLLM compression under matched settings, where bold marks Traj-MC improvements over SVD-LLM (higher is better).

Figure 2: Empirical validation of Traj-MC. (a) Uniform sampling achieves lower generation-state error

E_{\text{gen}}

than any fixed-

t

baseline. (b) Estimation error

e_{\Sigma}(N)

follows the theoretical

O(N^{-1/2})

convergence rate. (c) Covariance error directly tracks the resulting low-rank subspace deviation

d_{\text{sub}}

Diffusion Model Compression

See all 27 papers

Prompt Compression in Diffusion Large Language Models: Evaluating LLMLingua-2 on LLaDA
May 18, 2026Sterling Huang, Abigayle Brown, Jiyoo Noh +4LLM Evaluation LLM Compression
University of Toronto Toronto, Canada · King Mongkut’s University of Technology Thonburi Bangkok, Thailand
What does FFN compression change downstream? Same-state causal restoration in diffusion language models
Sep 16, 2026Shaurya OmarLLM Compression Diffusion Model Compression
Indian Institute of Technology, Roorkee
Trace-Based On-Policy Distillation for Masked Diffusion Language Models
Jul 18, 2026Haolin Ren, Ziyang Huang, Chenhao Yuan +2Masked Diffusion Models Diffusion Model Distillation
Institute of Automation, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China

May 18, 2026J/K move · Enter open · S save

Prompt Compression in Diffusion Large Language Models: Evaluating LLMLingua-2 on LLaDA

Sterling Huang, Abigayle Brown, Jiyoo Noh, Jiakang Xu+3

University of Toronto Toronto, Canada · King Mongkut’s University of Technology Thonburi Bangkok, Thailand

LLM Evaluation LLM Compression Diffusion Language Models Prompt Compression

Open paper AskPDF arXiv

Prompt compression reduces inference cost and context length in large language models, but prior evaluations focus mainly on autoregressive architectures. This study examines whether LLMLingua-2 transfers effectively to diffusion large language models (DLLMs), specifically LLaDA-8B-Instruct. We evaluate GSM8K, DUC2004, and ShareGPT using 250 prompts per dataset at an approximate 50% compression ratio, covering mathematical reasoning, prompt reconstruction, and summarization. Outputs from original, compressed, and reconstructed prompts are compared using exact-match accuracy, BLEU, ROUGE, and BERTScore. Results show that high semantic preservation does not necessarily ensure stable downstream behavior in diffusion models. Summarization remains relatively robust, while mathematical reasoning degrades substantially despite high semantic similarity. Reconstruction further shows that semantically similar prompts may omit reasoning-critical information needed for stable denoising. Overall, compression failures are mainly driven by information omission rather than semantic drift, suggesting that autoregressive prompt compression methods may not transfer uniformly to DLLMs. These findings motivate diffusion-aware compression strategies.

Abstract

Diffusion Model Compression

Explore similar work

Prompt Compression in Diffusion Large Language Models: Evaluating LLMLingua-2 on LLaDA

What does FFN compression change downstream? Same-state causal restoration in diffusion language models

Trace-Based On-Policy Distillation for Masked Diffusion Language Models