PromptDistill: Query-based Selective Token Retention in Intermediate Layers for Efficient Large Language Model Inference

📄 arXiv: 2503.23274v1 📥 PDF

作者: Weisheng Jin, Maojia Song, Tej Deep Pala, Yew Ken Chia, Amir Zadeh, Chuan Li, Soujanya Poria

分类: cs.CL

发布日期: 2025-03-30

💡 一句话要点

PromptDistill：一种基于查询的选择性token保留方法，用于高效的大语言模型推理。

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

关键词: 大语言模型 推理加速 token选择 注意力机制 长文本处理

📋 核心要点

  1. 大语言模型在处理长文本时，推理阶段的计算和内存开销巨大，成为性能瓶颈。
  2. PromptDistill通过在早期层选择性保留信息量大的token，减少后续层的计算负担，从而提高推理效率。
  3. 实验表明，PromptDistill在多个基准测试中，显著提升了推理效率，同时对模型输出质量的影响很小。

📝 摘要（中文）

为了解决大语言模型(LLMs)在处理复杂任务和长文档时推理计算和内存成本过高的问题，我们提出PromptDistill，一种新颖的、无需训练的方法，可在保持生成质量的同时提高推理效率。PromptDistill通过利用早期层的注意力交互来识别和保留信息量最大的token，保留它们的隐藏状态，同时减少后期层的计算负担。这使得模型能够专注于必要的上下文信息，而无需完全处理所有token。与仅在处理完整输入后进行压缩的H2O和SnapKV等先前方法，或选择固定比例的初始prompt而不考虑上下文依赖性的GemFilter不同，PromptDistill动态地将计算资源分配给最相关的token，同时保持对输入的全局感知。在LongBench、InfBench和Needle in a Haystack等基准测试中使用LLaMA 3.1 8B Instruct、Phi 3.5 Mini Instruct和Qwen2 7B Instruct等基础模型进行的实验表明，与原始模型相比，PromptDistill在显著提高效率的同时，对输出质量的影响极小。凭借单阶段选择策略，PromptDistill有效地平衡了性能和效率，由于其保留关键信息的能力更强，因此优于GemFilter、H2O和SnapKV等先前方法。具体而言，与GemFilter相比，PromptDistill实现了总体1%到5%的性能提升，同时还提供了更好的时间效率。此外，我们还探索了多阶段选择，这进一步提高了效率，同时保持了强大的生成性能。

🔬 方法详解

问题定义：论文旨在解决大语言模型在处理长文本输入时，推理阶段计算和内存成本过高的问题。现有方法，如H2O、SnapKV和GemFilter，要么在处理完整输入后进行压缩，要么仅关注初始prompt，无法动态地根据上下文信息的重要性来分配计算资源，导致效率提升有限或信息损失。

核心思路：PromptDistill的核心思路是，并非所有token对最终的生成都同等重要。通过在早期层识别并保留信息量最大的token，可以减少后续层的计算负担，同时保持模型的生成质量。这种方法允许模型专注于关键的上下文信息，避免对所有token进行完全处理。

技术框架：PromptDistill采用一种基于查询的选择性token保留框架。该框架主要包含以下阶段：1) Token重要性评估：利用早期Transformer层的注意力机制，计算每个token的重要性得分。2) Token选择：根据重要性得分，选择保留一部分token的隐藏状态。3) 后续层推理：在后续的Transformer层中，仅对保留的token进行计算，从而减少计算量。

关键创新：PromptDistill的关键创新在于其动态选择token的能力。与之前固定选择或压缩所有token的方法不同，PromptDistill能够根据token的上下文信息，自适应地分配计算资源。此外，PromptDistill是一种无需训练的方法，可以直接应用于现有的预训练模型，无需额外的训练成本。

关键设计：PromptDistill的关键设计包括：1) 注意力得分计算：使用Transformer早期层的注意力权重来衡量token的重要性。具体来说，可以将每个token的注意力权重进行聚合，得到一个代表其重要性的标量值。2) Token选择策略：可以采用不同的token选择策略，例如选择top-k个最重要token，或者设置一个阈值，保留重要性得分高于该阈值的token。3) 多阶段选择：可以采用多阶段选择策略，在不同的Transformer层中逐步减少token的数量，从而进一步提高效率。

🖼️ 关键图片

fig_0

📊 实验亮点

实验结果表明，PromptDistill在LongBench、InfBench和Needle in a Haystack等基准测试中，与原始模型相比，显著提高了推理效率，同时对输出质量的影响极小。与GemFilter相比，PromptDistill实现了总体1%到5%的性能提升，同时还提供了更好的时间效率。多阶段选择策略进一步提高了效率，同时保持了强大的生成性能。

🎯 应用场景

PromptDistill可应用于各种需要处理长文本的大语言模型应用场景，例如长文档摘要、问答系统、代码生成等。通过提高推理效率，PromptDistill可以降低部署成本，并支持在资源受限的设备上运行大型语言模型。该方法还有助于加速AI研究和开发，使研究人员能够更快地迭代和评估新的模型和算法。

📄 摘要（原文）

As large language models (LLMs) tackle increasingly complex tasks and longer documents, their computational and memory costs during inference become a major bottleneck. To address this, we propose PromptDistill, a novel, training-free method that improves inference efficiency while preserving generation quality. PromptDistill identifies and retains the most informative tokens by leveraging attention interactions in early layers, preserving their hidden states while reducing the computational burden in later layers. This allows the model to focus on essential contextual information without fully processing all tokens. Unlike previous methods such as H2O and SnapKV, which perform compression only after processing the entire input, or GemFilter, which selects a fixed portion of the initial prompt without considering contextual dependencies, PromptDistill dynamically allocates computational resources to the most relevant tokens while maintaining a global awareness of the input. Experiments using our method and baseline approaches with base models such as LLaMA 3.1 8B Instruct, Phi 3.5 Mini Instruct, and Qwen2 7B Instruct on benchmarks including LongBench, InfBench, and Needle in a Haystack demonstrate that PromptDistill significantly improves efficiency while having minimal impact on output quality compared to the original models. With a single-stage selection strategy, PromptDistill effectively balances performance and efficiency, outperforming prior methods like GemFilter, H2O, and SnapKV due to its superior ability to retain essential information. Specifically, compared to GemFilter, PromptDistill achieves an overall $1\%$ to $5\%$ performance improvement while also offering better time efficiency. Additionally, we explore multi-stage selection, which further improves efficiency while maintaining strong generation performance.