PromptDistill: Query-based Selective Token Retention in Intermediate Layers for Efficient LLM

提出了PromptDistill方法，通过中间层基于查询的选择性token保留，在保持生成质量的同时提高大语言模型推理效率，在多个基准测试中表现出色。

1. 主要内容
   • 研究背景：大语言模型在处理长文本时，推理的计算和内存成本高成为瓶颈。现有方法如MCSD、Mamba等在压缩内部状态时存在性能下降和可扩展性问题；H2O、SnapKV等虽无需额外训练，但压缩前需处理所有token，计算开销大；GemFilter仅处理部分层的所有token后选择重要token，可能导致信息丢失。
   • 方法介绍：PromptDistill是一种无需训练的方法。先像传统Transformer一样处理前几层所有token，然后基于最后一个提示token的查询与每个token的键之间的点积选择最重要的提示token，并保留其隐藏状态继续处理剩余层。同时引入缓存截断机制，清除未选中token的键值缓存，减少计算和内存成本。此外，该方法支持多阶段选择，在早期层选择较多token，后期层选择较少token，平衡效率和效果。
   • 实验结果：在LongBench、InfBench和Needle in a Haystack等基准测试中，对Llama 3.1 8B Instruct、Phi 3.5 Mini Instruct和Qwen2 7B Instruct等模型进行实验。结果表明，PromptDistill比全注意力模型效率更高，性能与全注意力模型接近，且在大多数情况下优于GemFilter、H2O和SnapKV等基线方法。多阶段选择可进一步提高效率，缓存截断不影响性能。