通过动态保留关键词元,显著提升大模型推理效率
PromptDistill: Query-based Selective Token Retention in Intermediate Layers for Efficient Large Language Model Inference
- 基于早期层注意力机制,选择性保留最相关词元
- 相比基线方法,推理速度更快,性能损失仅1%-5%
- 适合需要高效推理的长文本任务场景
随着大语言模型处理更复杂任务和更长文档,其推理阶段的计算与内存开销成为主要瓶颈。为此,我们提出PromptDistill,一种无需训练的新方法,在保持生成质量的同时提升推理效率。该方法利用早期层中的注意力交互识别并保留最具信息量的词元,保留其隐状态,减少后续层的计算负担,使模型聚焦于核心上下文信息而不必完整处理所有词元。与仅在输入处理完毕后压缩的H2O、SnapKV,或固定选取初始提示部分的GemFilter不同,PromptDistill动态分配资源至最相关词元,同时维持对输入的全局感知。在LLaMA 3.1 8B Instruct、Phi 3.5 Mini Instruct和Qwen2 7B Instruct等基线模型上,结合LongBench、InfBench及Needle in a Haystack等基准测试,结果表明PromptDistill显著提升效率,对输出质量影响极小。采用单阶段选择策略即可在性能与效率间取得良好平衡,优于GemFilter、H2O和SnapKV。具体而言,相比GemFilter,性能提升1%至5%,时间效率更优。此外,探索多阶段选择进一步提升效率,同时保持强生成能力。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。