DiffuMask: Diffusion Language Model for Token-level Prompt Pruning

Caleb Zheng; Jyotika Singh; Fang Tu; Weiyi Sun; Sujeeth Bharadwaj; Yassine Benajiba; Sujith Ravi; Eli Shlizerman

cs.CLApr 8, 2026

DiffuMask: Diffusion Language Model for Token-level Prompt Pruning

Authors: Caleb Zheng, Jyotika Singh, Fang Tu, Weiyi Sun, Sujeeth Bharadwaj, Yassine Benajiba, Sujith Ravi, Eli Shlizerman,

Organizations: University of Washington · Oracle AI

Abstract

In-Context Learning and Chain-of-Thought prompting improve reasoning in large language models (LLMs). These typically come at the cost of longer, more expensive prompts that may contain redundant information. Prompt compression based on pruning offers a practical solution, yet existing methods rely on sequential token removal which is computationally intensive. We present DiffuMask, a diffusion-based framework integrating hierarchical shot-level and token-level pruning signals, that enables rapid and parallel prompt pruning via iterative mask prediction. DiffuMask substantially accelerates the compression process via masking multiple tokens in each denoising step. It offers tunable control over retained content, preserving essential reasoning context and achieving up to 80% prompt length reduction. Meanwhile, it maintains or improves accuracy across in-domain, out-of-domain, and cross-model settings. Our results show that DiffuMask provides a generalizable and controllable framework for prompt compression, facilitating faster and more reliable in-context reasoning in LLMs.

Figures & tables

Figure 1: Data generation pipeline for DiffuMask . The method applies hierarchical shot-level and token-level pruning to generate pruned prompt candidates, which are evaluated and filtered into training and validation sets. GSM8K is shown as a representative example; the pipeline is applicable to any dataset with long prompts.

Figure 2: Overview of the proposed DiffuMask framework compared with a Masked DLM baseline.

Methods	Full	Random	Llmlingua	Llmlingua2	CI	HP	DiffuMask
Token Counts $\downarrow$	5399.1	1061.3	1153.2	1155.2	1644.4	660.4	1113.8
EM Accuracy $\uparrow$	76.49	60.89	67.32	70.89	76.90	75.83	77.07
Pruning Inference(min) $\downarrow$	-	0	0.1	0.1	0.5	1440	0.75

Table 1: In-domain evaluation on the GSM8K benchmark under the 32-shot setting. Reported are average token counts, Exact Match (EM) accuracies and pruning inference of different methods.

Dataset	Shots	Token Counts		Accuracy
Dataset	Shots	Full	DiffuMask	Full	DiffuMask
Yelp5	1	877.7	285.4 (↓67.5%)	61.88	62.93 (+1.05)
	5	4800.7	149.6 (↓96.9%)	50.79	62.91 (+12.12)
Yahoo	1	533.3	224.8 (↓57.8%)	60.56	63.44 (+2.88)
	5	2484.4	759.8 (↓69.4%)	51.96	63.99 (+12.03)
AG’s News	1	227.0	224.3 (↓1.2%)	80.28	80.23 (-0.05)

Table 2: Out-of-domain evaluation on classification benchmarks. Relative token reduction (↓) and accuracy change (

\pm

) are shown with respect to the full prompt.

Model	Full	DiffuMask
GPT4o	90.57	92.39 (+1.8)
GPT4.1	90.93	92.76 (+1.8)
GPT4.1nan	83.95	83.48(-0.5)
Gem2.5Fl	97.03	96.40(-0.6)
Gem2.5Pro	98.12	98.28 (+0.2)
GROK 3	97.86	96.50(-1.4)

Table 3: Cross-model evaluation. Accuracy (%) on the GSM8K test set using full and DiffuMask -pruned prompts across eight LLMs. (GPT4.1nan = GPT-4.1-nano; Gem2.5Fl = Gemini-2.5-Flash; Gem2.5Pro = Gemini-2.5-Pro; Llam3.3-70 = Llama-3.3-70B-Instruct; Llam4-Mav = Llama-4-Maverick)

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

# Shots	Token Counts	Accuracy (%)
1	183	73.38
2	377	74.44
5	857	75.65
10	1704	75.64
15	2571	75.29
20	3412	75.05

Appendix

Table 4: Performance of top-32 random selection with 100 data points.

Method	Full		Fewer-shot
	Shots	Acc	Shots	Acc
RL Pruner (100)	32	74.42	10.3	75.32
Rand Var- $k$ (100)	32	74.42	10.3	75.33
Rand Fix- $k$ (600)	32	74.84	5	75.72

Appendix

Table 5: Shot-level pruning comparison on GSM8K.

Pruning Group	Improved	TA Traj	Total
RL Pruner	55	26	81
Rand Var- $k$	37	24	61
Rand Fix- $k$	61	-	61
Total	142	61	203

Appendix

Table 6: Summary of token-level pruning results.

Pruning Group	Train	Validation	Test
Trained Pruner	81	–	19
Random Var- $k$	37	24	–
Random Fix- $k$	50	11	–
Total	168	35	19

Appendix

Table 7: Dataset composition across training, validation, and testing sets.

Hyperparameter	Value
Batch size	1
Epochs	20
Learning rate	$1\times 10^{-4}$
Mask penalty coefficient	2.0
Balance ratio	0.8
Maximum sequence length	4096

Appendix

Table 8: Finetuning Hyperparameters for DiffuMask.

	Token Counts				Accuracy
	Threshold
Top-k	0.0001	0.001	0.01	0.1	0.0001	0.001	0.01	0.1
2	150.4	1061.1	2950.9	4219.7	75.31	80.76	80.95	79.75
3	132.1	1023.2	2943.3	4208.3	75.35	80.73	80.63	80.11
4	130.8	978.5	2891.4	4205.0	73.37	80.93	80.91	80.22

Appendix

Table 9: Effect of inference hyperparameters on GSM8K performance.

	One-shot		Five-shot
Top- $k$	Accuracy (%)	Tokens	Accuracy (%)	Tokens
topk10	56.88	591.0	62.78	1555.6
topk20	62.78	476.8	62.44	453.2
topk50	62.93	354.4	62.78	140.7
topk100	63.04	265.6	63.10	130.2

Appendix

Table 10: Hyperparameter search of DiffuMask on the Yelp dataset under one-shot and five-shot settings.

	One-shot		Five-shot
Top- $k$	Accuracy (%)	Tokens	Accuracy (%)	Tokens
topk10	59.26	353.4	63.94	1355.0
topk20	63.29	298.9	63.91	1082.9
topk50	63.53	206.0	64.09	666.4
topk100	63.19	170.0	63.43	400.5
topk200	62.77	127.3	62.50	178.1

Appendix

Table 11: Hyperparameter search of DiffuMask on the Yahoo dataset under one-shot and five-shot settings.

	One-shot		Five-shot		Twenty-shot
Top- $k$	Accuracy (%)	Tokens	Accuracy (%)	Tokens	Accuracy (%)	Tokens
topk2	80.48	220.9	67.58	1078.3	62.83	4287.7
topk5	80.64	219.2	73.93	991.8	79.40	3857.6
topk10	79.78	214.1	65.45	970.0	78.43	3613.7
topk20	71.22	203.3	76.58	897.2	76.38	3174.6
topk50	76.37	187.7	79.64	715.6	84.75	2121.9

Appendix

Table 12: Hyperparameter search of DiffuMask on the AG News dataset under one-shot, five-shot, and twenty-shot settings.

Panel of Figure 3 — Figure 3: Comparison of POS Tag Distribution (Original Prompts vs Removed Parts of Prompts)

ADJ:	adjective (e.g., big, old)
ADP:	adposition (e.g., in, to, during)
ADV:	adverb (e.g., very, tomorrow)
AUX:	auxiliary verb (e.g., is, has)
CCONJ:	coordinating conjunction (e.g., and, but)
DET:	determiner (e.g., the, a, some)
INTJ:	interjection (e.g., oh, wow)

Appendix

Table 13: Universal Part-of-Speech (POS) tags

Dataset	Primary Accuracy Measure	What the Accuracy is Checking
AGNews	Top-1 Accuracy	Whether the model correctly identified the single news topic (World, Sports, Business, or Sci/Tech).
Yelp5	Star Rating Accuracy	Whether the model predicted the exact star rating (1, 2, 3, 4, or 5) assigned by the user in their review.
Yahoo	Topic Classification Accuracy	Whether the model correctly assigned the text to one of 10 main categories (e.g., Health, Sports, Politics).

Appendix

Table 14: Accuracy-based evaluation used for out-of-domain datasets.

Prompt Compression

Prompt compression reduces the tokens sent to language models while trying to preserve the information needed for accurate, efficient responses.

See all 34 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
PC-SubMax: Efficient Prompt Compression via Regularized Submodular Maximization
Sep 26, 2026Ziyi Zhang, Shuang Cui, Haotian Zhang +1Prompt Compression Efficient Language Model Inference
Soochow University
OS-Pruner: Pruning Chains-of-Thought of Reasoning Models via Optimal Stopping
Jul 13, 2026Mohammed Ehab, Aymane El Gadarri, Vivek F. Farias +2Overthinking in Language Models Sequential Decision Making
MIT · Columbia Business School

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

Prompt Compression

Explore similar work

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

PC-SubMax: Efficient Prompt Compression via Regularized Submodular Maximization

OS-Pruner: Pruning Chains-of-Thought of Reasoning Models via Optimal Stopping