跳到主内容
人工智能

面向长上下文Transformer推理的高效提示词压缩方法:基于评估头技术

Efficient Prompt Compression with Evaluator Heads for Long-Context Transformer Inference

Weizhi Fei, Xueyan Niu, Guoqing Xie, Yi Liu, Bo Bob Bai, Wei Na Han

neural information processing systems · 2025-01-01

开放获取大模型与语言被引前 10%影响因子 ≈3.8

摘要

尽管涉及长上下文输入的应用对于充分发挥大型语言模型的价值至关重要,但这类应用也会带来计算成本上升与性能下降的问题。为应对这一挑战,我们提出了一种无需训练的高效提示词压缩方法,能够在压缩后的提示词中保留关键信息。我们发现基于Transformer的大型语言模型中存在特定的注意力头,将其命名为评估头,这类注意力头能够筛选出长输入中对推理最为重要的词元。基于这一发现,我们开发了EHPC——一种基于评估头的提示词压缩方法,该方法在预填充阶段仅利用包含评估头的前几层即可让大型语言模型快速“浏览”输入提示词,随后仅将重要词元送入模型进行推理。EHPC在提示词压缩与长上下文推理加速两项主流基准测试中均取得了最优结果,从而有效降低了调用商业API时的复杂性与成本。我们还进一步证明,EHPC的性能可与基于键值缓存的加速方法相媲美,这凸显了其提升大型语言模型处理长上下文任务效率的潜力。

本摘要由 AI(Hy-MT2-Pro)翻译生成,仅供参考,请以原文为准。

AI 综述

基于全文生成
主要观点
该文针对长上下文输入给基于 Transformer 的大语言模型带来的计算成本上升与性能下降问题,提出一种无需训练、在压缩提示中保留关键信息的提示压缩方法。其核心发现是 LLM 中存在一组特定注意力头,作者称为 evaluator heads,可在长输入中选取对推理最显著的 token。基于此构建 EHPC:在预填充阶段仅用带 evaluator heads 的前几层快速“略读”输入提示,再把重要 token 送入模型推理。EHPC 在提示压缩与长上下文推理加速两个主流基准上取得 SOTA,并可降低商业 API 调用复杂度与成本;与基于 KV cache 的加速方法相比也具竞争力,在问答数据集上较直接推理最高提升 40%。
研究方法
研究先用合成的“Needle-in-a-Haystack”式长文本基准做 pilot experiment,计算证据子序列在各层各注意力头上的累积注意力分数,形成 evidence score 矩阵,并选取最高累计分数层的 top-k 头作为 evaluator heads。随后采用 token deletion 策略,依据 evaluator heads 的注意力分数经 pooling 操作删除非关键 token,按原序构造压缩提示。实验设置包括 Extended Model Inference(EMI,用不同模型或 API 推理压缩提示)和 Native Model Inference(NMI,同一模型压缩并推理),并给出预填充与解码复杂度分析。评估覆盖 LongBench 与 ZeroSCROLLS 的提示压缩基准,以及 LongBench 中单文档、多文档、摘要、FewShot、合成、代码等任务;还通过对比同层其他头、任务感知头与 QA 头来检验 evaluator heads 的存在性、实用性与鲁棒性。
学术贡献
该文识别出一类可评估输入 token 显著性的注意力头(evaluator heads),将 LLM 内在注意力机制用于训练-free 的提示压缩,属于方法创新与机制发现。EHPC 在 LongBench 和 ZeroSCROLLS 上刷新提示压缩性能,并在压缩长度 2048 token 约束下以 0.88 延迟取得 49.6 平均性能,显示其效率优势。它证明压缩后的提示可迁移至 API 黑盒模型,降低商业 API 成本,也可在本地模型 NMI 下减少内存并加速长文本推理。对长上下文 LLM 推理、提示压缩、KV cache 压缩加速以及商业 API 成本控制等子领域有直接价值;材料还显示 EHPC 在问答数据集上较直接推理最高提升 40%。
存在局限
材料中未明确设置专门的局限性讨论。方法上,evaluator heads 的识别依赖合成的 Needle-in-a-Haystack pilot 数据,虽然作者报告其任务无关性,但 top-k 头与最高分层的选择规则在材料中未展示对合成数据分布、证据位置和长度的敏感性分析。实验样本以 LongBench、ZeroSCROLLS 及部分多跳、代码任务为主,材料中未明确提及多语言、多模态或极长上下文场景的覆盖。部署上,EHPC 需要本地可访问 LLM 的前几层与注意力分数来完成压缩,EMI 虽可用于黑盒 API,但压缩阶段仍依赖本地模型,材料未量化这种部署约束。论证链条上,材料比较了 evaluator heads 与同层其他头,但未明确给出不同层数、头数和 pooling 核大小 r 的系统消融,因而对压缩后信息丢失与语义连贯性的边界说明不足。
进一步研究方向
第一,可在法律文档、文学小说、指令手册和代码文档等真实长文本上分别运行 EMI 与 NMI,测量 API 成本、延迟、内存和下游准确率,以检验 evaluator heads 的跨领域泛化。第二,可构造多跳推理、代码补全、摘要等不同 pilot 合成数据,自动搜索层与头组合,并与 QA 识别出的 task-agnostic heads 对照,验证 evaluator heads 的稳定性与可迁移性。第三,可将 evaluator heads 与 KV cache 压缩方法联合,在 NMI 下同时压缩提示 token 和 KV cache,测量 LongBench、ZeroSCROLLS 性能与内存-延迟曲线,验证增益是否叠加。第四,可对压缩率 κ2、压缩层比例 κ1、pooling 核大小 r 和保留 token 数做系统消融,分析性能-成本权衡及压缩提示的可读性与连贯性。第五,可在不同规模压缩模型(如 Llama-3.1-8B 及其他长上下文模型)上比较 EMI 对商业 API 输出质量的影响,并测试任务指令变化时 evaluator heads 是否需重新识别。

本综述由 AI(deepseek-flash)阅读开放获取全文自动生成,仅供参考,不代表原文作者观点,请以原文为准。

参考文献

相关论文

话题:Topic Modeling