通过激活值动态选关键查询,提升长文本推理效率
Activation-aware Probe-Query: Effective Key-Value Retrieval for Long-Context LLMs Inference
- 用激活值引导选择关键查询,精准定位相关上下文
- 在Long-Bench和∞Bench上实现顶尖性能,节省显存
- 无需训练,适合部署在资源受限的长文本场景
大型语言模型在长上下文任务中表现优异,但受限于GPU显存,推理效率面临挑战。现有方法采用滑动窗口累积历史键值对以复用,再通过选择性保留子集优化。然而,由于长上下文中注意力分布稀疏,大量候选键值对导致注意力分散,难以有效识别相关项。我们发现,在每个滑动窗口中选取代表性标记作为探针查询(probe-Query),可有效表征整体上下文,而此思路未被现有方法重视。为此,我们提出无训练的激活感知方法ActQKV:在预填充阶段,通过监测每窗口的令牌级激活偏置(Activation Bias),动态构建探针查询;在解码阶段,设计基于层间信息密度的动态键值截断机制,精确召回相关键值对并减少无关项。在Long-Bench和∞Bench上的实验表明,该方法在保持竞争性推理质量的同时,达到领先的资源效率。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。