DiffuMask: Diffusion Language Model for Token-level Prompt Pruning

TL;DR

DiffuMask is a diffusion-based prompt pruning framework that rapidly compresses prompts by masking multiple tokens simultaneously, reducing length up to 80% while maintaining reasoning accuracy.

Contribution

It introduces a novel diffusion-based method for parallel prompt pruning, significantly speeding up compression and improving control over retained content.

Findings

Achieves up to 80% prompt length reduction.

Maintains or improves reasoning accuracy across various settings.

Accelerates prompt compression via parallel token masking.

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.