Activation-aware Probe-Query: Effective Key-Value Retrieval for Long-Context LLMs Inference

📄 arXiv: 2502.13542v1 📥 PDF

作者: Qingfa Xiao, Jiachuan Wang, Haoyang Li, Cheng Deng, Jiaqi Tang, Shuangyin Li, Yongqi Zhang, Jun Wang, Lei Chen

分类: cs.CL, cs.AI

发布日期: 2025-02-19

💡 一句话要点

提出ActQKV以解决长上下文LLMs推理效率问题

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

关键词: 长上下文推理 激活感知 键值对检索 大型语言模型 推理效率 动态查询 信息密度

📋 核心要点

  1. 现有方法在长上下文推理中面临效率低下和内存限制的问题，难以有效回忆相关的键值对。
  2. 本文提出ActQKV方法，通过动态选择探针查询，利用激活偏差来优化KV对的检索过程。
  3. 实验结果显示，ActQKV在多个基准测试中实现了最先进的性能，提升了推理质量和资源利用效率。

📝 摘要（中文）

近年来，大型语言模型（LLMs）在长上下文任务中表现出色，但在推理效率上面临显著挑战，尤其是在GPU内存有限的情况下。现有方法采用滑动窗口策略来积累历史键值对（KV），但由于长上下文中的稀疏注意力分布，识别和回忆相关的KV对变得困难。本文提出了一种新的激活感知方法ActQKV，通过动态确定探针查询，利用信息密度在解码阶段准确回忆相关KV对，从而提高推理效率。实验结果表明，ActQKV在Long-Bench和∞ Benchmarks上表现出色，具有竞争力的推理质量和资源效率。

🔬 方法详解

问题定义：本文旨在解决长上下文LLMs推理中的效率问题，现有方法在处理稀疏注意力分布时难以有效识别和回忆相关的KV对，导致推理性能下降。

核心思路：论文提出的ActQKV方法通过监测每个上下文窗口内的激活偏差，动态选择代表性令牌作为探针查询，从而有效表示整个上下文并优化KV对的检索。

技术框架：ActQKV的整体架构包括两个主要阶段：预填充阶段和解码阶段。在预填充阶段，利用激活偏差构建探针查询；在解码阶段，依据信息密度设计动态KV截断机制，以准确回忆相关KV对。

关键创新：ActQKV的核心创新在于其训练无关性和激活感知机制，能够动态选择探针查询，与传统方法相比，显著提高了KV对的检索效率和准确性。

关键设计：在设计中，采用了基于信息密度的动态KV截断机制，确保在解码阶段能够有效过滤掉无关的KV对，提升推理过程的资源利用效率。具体参数设置和损失函数设计在实验部分进行了详细说明。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，ActQKV在Long-Bench和∞ Benchmarks上实现了最先进的性能，相较于基线方法，推理质量提升显著，资源效率也得到了有效改善，展示了其在长上下文任务中的应用潜力。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等。通过提高长上下文推理的效率，ActQKV可以在资源受限的环境中实现更高效的模型推理，推动大型语言模型在实际应用中的普及与发展。

📄 摘要（原文）

Recent advances in large language models (LLMs) have showcased exceptional performance in long-context tasks, while facing significant inference efficiency challenges with limited GPU memory. Existing solutions first proposed the sliding-window approach to accumulate a set of historical \textbf{key-value} (KV) pairs for reuse, then further improvements selectively retain its subsets at each step. However, due to the sparse attention distribution across a long context, it is hard to identify and recall relevant KV pairs, as the attention is distracted by massive candidate pairs. Additionally, we found it promising to select representative tokens as probe-Query in each sliding window to effectively represent the entire context, which is an approach overlooked by existing methods. Thus, we propose \textbf{ActQKV}, a training-free, \textbf{Act}ivation-aware approach that dynamically determines probe-\textbf{Q}uery and leverages it to retrieve the relevant \textbf{KV} pairs for inference. Specifically, ActQKV monitors a token-level indicator, Activation Bias, within each context window, enabling the proper construction of probe-Query for retrieval at pre-filling stage. To accurately recall the relevant KV pairs and minimize the irrelevant ones, we design a dynamic KV cut-off mechanism guided by information density across layers at the decoding stage. Experiments on the Long-Bench and $\infty$ Benchmarks demonstrate its state-of-the-art performance with competitive inference quality and resource efficiency.