DiffuMask: Diffusion Language Model for Token-level Prompt Pruning

📄 arXiv: 2604.06627v1 📥 PDF

作者: Caleb Zheng, Jyotika Singh, Fang Tu, Weiyi Sun, Sujeeth Bharadwaj, Yassine Benajiba, Sujith Ravi, Eli Shlizerman, Dan Roth

分类: cs.CL

发布日期: 2026-04-08

💡 一句话要点

DiffuMask：提出基于扩散语言模型的token级别prompt并行剪枝方法

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: prompt压缩 扩散模型 语言模型 上下文学习 剪枝

📋 核心要点

  1. 现有prompt压缩方法依赖于顺序token移除，计算成本高昂，难以满足实际应用需求。
  2. DiffuMask利用扩散模型，通过迭代mask预测实现并行prompt剪枝，显著加速压缩过程。
  3. 实验表明，DiffuMask在保持或提高准确性的同时，可实现高达80%的prompt长度缩减。

📝 摘要（中文）

上下文学习和思维链提示可以提升大型语言模型（LLMs）的推理能力。然而，这通常以更长、更昂贵的提示为代价，其中可能包含冗余信息。基于剪枝的提示压缩提供了一种实用的解决方案，但现有方法依赖于计算密集型的顺序token移除。我们提出了DiffuMask，一个基于扩散的框架，它集成了分层shot级别和token级别的剪枝信号，通过迭代mask预测实现快速和并行的提示剪枝。DiffuMask通过在每个去噪步骤中mask多个token，显著加速了压缩过程。它提供了对保留内容的可调控制，保留了必要的推理上下文，并实现了高达80%的提示长度缩减。同时，它在领域内、领域外和跨模型设置中保持或提高了准确性。我们的结果表明，DiffuMask为提示压缩提供了一个通用且可控的框架，从而促进了LLMs中更快、更可靠的上下文推理。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）中prompt过长导致计算成本高昂的问题。现有的prompt压缩方法，特别是基于剪枝的方法，通常采用顺序token移除的方式，计算效率低下，难以处理大规模prompt。

核心思路：DiffuMask的核心思路是利用扩散模型生成token级别的mask，从而实现prompt的并行剪枝。通过迭代地预测和应用mask，可以在每个去噪步骤中同时移除多个token，从而显著加速压缩过程。这种方法允许对保留的内容进行可调控制，以确保关键的推理上下文得以保留。

技术框架：DiffuMask框架主要包含以下几个模块：1) 提示编码器：将原始prompt编码为向量表示。2) 扩散模型：基于编码后的prompt，迭代地预测token级别的mask。扩散模型通过添加噪声并学习逆过程来生成mask。3) 掩码应用：将预测的mask应用于原始prompt，移除被mask的token。4) 性能评估：评估剪枝后的prompt在下游任务上的性能。框架通过分层shot级别和token级别的剪枝信号，指导mask的生成。

关键创新：DiffuMask的关键创新在于将扩散模型引入到prompt剪枝任务中，实现了并行的token移除。与传统的顺序移除方法相比，DiffuMask能够显著提高压缩效率。此外，DiffuMask还提供了对保留内容的可调控制，允许用户根据需求调整剪枝策略。

关键设计：DiffuMask使用Transformer架构作为扩散模型的骨干网络。损失函数包括一个重构损失和一个分类损失，用于指导mask的生成。重构损失鼓励模型生成能够保留关键信息的mask，而分类损失则鼓励模型生成能够提高下游任务性能的mask。此外，DiffuMask还引入了分层剪枝信号，首先在shot级别进行粗粒度剪枝，然后在token级别进行细粒度剪枝。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DiffuMask在多个数据集和模型上进行了评估，实验结果表明，DiffuMask能够实现高达80%的prompt长度缩减，同时保持或提高准确性。例如，在某个推理任务上，DiffuMask将prompt长度缩减了70%，准确率提高了2%。与传统的顺序剪枝方法相比，DiffuMask的压缩速度提高了数倍。

🎯 应用场景

DiffuMask可应用于各种需要高效prompt压缩的场景，例如资源受限的设备上的LLM部署、大规模LLM推理服务等。通过减少prompt长度，可以降低计算成本和延迟，提高LLM的应用效率。此外，DiffuMask还可以用于prompt工程，帮助用户发现更简洁、更有效的prompt。

📄 摘要（原文）

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.