arXiv:2605.09778cs.LGcs.CL2026-05

用神经网络加速长文本注意力,省去逐个读取缓存的开销。

Nectar: Neural Estimation of Cached-Token Attention via Regression

论文配图:Nectar: Neural Estimation of Cached-Token Attention via Regression
图 1 · 摘自论文原文
  • 用小型神经网络拟合查询与注意力输出的关系,替代传统遍历缓存
  • 在1.7B到8B模型上,近似误差与完整注意力的生成差距一致
  • 适合需要高速长文本推理的场景,如法律、技术文档问答

对固定长上下文进行softmax注意力计算时,需为每个新查询令牌读取所有缓存的键值对。对于特定上下文(如一本书、手册或法律文献),注意力输出是查询的确定性函数。我们提出Nectar,通过一个紧凑的神经网络拟合该函数,输入来自任务相关的查询分布。Nectar为每层和每个KV头训练两个网络:一个预测注意力输出的目标网络,另一个预测对数归一化常数的评分网络。这两个网络在推理时嵌入标准掩码自注意力中,将原本随上下文长度O(n)增长的注意力计算替换为不依赖n的前向传播。每个模块参数量约为|θ|,通常远小于相同粒度下的2nd KV缓存占用。我们在1.7B至8B参数的模型上,于五个长上下文数据集上进行了实验。近似误差与完整注意力的下一个词预测准确率差距保持一致;在层间非均匀分配容量可进一步缩小该差距。此外,配备Nectar模块的模型生成的文本(基于问题提示)在语义内容上与使用完整缓存的结果一致。

原文摘要 · Abstract (English)

Evaluating softmax attention over a fixed long context requires reading every cached key-value pair for each new query token. For a given context (a book, a manual, a legal corpus) the attention output is a deterministic function of the query. We propose Nectar, which fits a compact neural network to this function for queries drawn from a task-relevant distribution. Nectar fits two networks per layer and KV-head: a target network that predicts the attention output and a score network that predicts the log-normalizer. The pair plugs into the standard masked self-attention at inference time, replacing the $O(n)$ attention over the cache with a forward pass whose cost does not depend on $n$. Each module carries on the order of $|θ|$ parameters per layer and KV-head, typically much smaller than the $2nd$ KV-cache footprint at the same granularity. We report experiments on models from 1.7B to 8B parameters across five long-context datasets. The approximation error tracks the next-token accuracy gap to full attention, and allocating capacity non-uniformly across layers reduces that gap in our ablation. Beyond this analysis of metrics, we check that the text generations (following a question prompt) of a model equipped with a Nectar module match in semantic content those obtained by giving the same model access to the full cache.

注意力机制长文本缓存优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。