# DiffuMask: Diffusion Language Model for Token-level Prompt Pruning

arXiv: https://arxiv.org/abs/2604.06627 · HTML: https://mlsanity.com/paper/2604.06627/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 27 of its 39 references that are in the map.
- First posted on arXiv 8 Apr 2026.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)
- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)
- Efficient LLM inference and serving / [Prompt and context compression](https://mlsanity.com/inference/prompt-and-context-compression/)

## Builds on (most cited first, 27 of 27)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Large Language Models are Zero-Shot Reasoners](https://mlsanity.com/paper/2205.11916/) (2022)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Dream 7B: Diffusion Large Language Models](https://mlsanity.com/paper/2508.15487/) (2025)
- [DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models](https://mlsanity.com/paper/2210.08933/) (2022)
- [Learning to Compress Prompts with Gist Tokens](https://mlsanity.com/paper/2304.08467/) (2023)
- [Scaling Diffusion Language Models via Adaptation from Autoregressive Models](https://mlsanity.com/paper/2410.17891/) (2024)
- [Scaling up Masked Diffusion Models on Text](https://mlsanity.com/paper/2410.18514/) (2024)
- [Walking Down the Memory Maze: Beyond Context Limit through Interactive Reading](https://mlsanity.com/paper/2310.05029/) (2023)
- [Prompt Compression and Contrastive Conditioning for Controllability and Toxicity Reduction in Language Models](https://mlsanity.com/paper/2210.03162/) (2022)
- [An Information-theoretic Approach to Prompt Engineering Without Ground Truth Labels](https://mlsanity.com/paper/2203.11364/) (2022)
- [Yelp Dataset Challenge: Review Rating Prediction](https://mlsanity.com/paper/1605.05362/) (2016)
- [Prompt Compression for Large Language Models: A Survey](https://mlsanity.com/paper/2410.12388/) (2024)
- [CompAct: Compressing Retrieved Documents Actively for Question Answering](https://mlsanity.com/paper/2407.09014/) (2024)
- [Diffusion Beats Autoregressive in Data-Constrained Settings](https://mlsanity.com/paper/2507.15857/) (2025)
- [OpenICL: An Open-Source Framework for In-context Learning](https://mlsanity.com/paper/2303.02913/) (2023)
- [Discrete Prompt Compression with Reinforcement Learning](https://mlsanity.com/paper/2308.08758/) (2023)
- [TACO-RL: Task Aware Prompt Compression Optimization with Reinforcement Learning](https://mlsanity.com/paper/2409.13035/) (2024)
- [Evolving Prompts In-Context: An Open-ended, Self-replicating Perspective](https://mlsanity.com/paper/2506.17930/) (2025)
- [Can LLMs Narrate Tabular Data? An Evaluation Framework for Natural Language Representations of Text-to-SQL System Outputs](https://mlsanity.com/paper/2510.23854/) (2025)
- [Barriers to Discrete Reasoning with Transformers: A Survey Across Depth, Exactness, and Bandwidth](https://mlsanity.com/paper/2602.11175/) (2026)

## Built on by (most cited first, 1 of 1)

- [JTPRO: A Joint Tool-Prompt Reflective Optimization Framework for Language Agents](https://mlsanity.com/paper/2604.19821/) (2026)

## Abstract (arXiv)

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.