论文

Nectar:通过回归对缓存词元注意力进行神经估计

Nectar: Neural Estimation of Cached-Token Attention via Regression

模型推理KV Cache

摘要

在固定的长上下文上评估 softmax 注意力需要读取每个新查询标记的每个缓存的键值对。对于给定的上下文(一本书、一本手册、一个法律语料库),注意力输出是查询的确定性函数。我们提出了 Nectar,它将一个紧凑的神经网络拟合到该函数中,用于从任务相关分布中提取查询。 Nectar 每层和 KV-head 适合两个网络:一个预测注意力输出的目标网络和一个预测对数归一化器的评分网络。该对在推理时插入标准屏蔽自注意力,用成本不依赖于 $n$ 的前向传递替换缓存上的 $O(n)$ 注意力。每个模块都按每层和 KV 头的 $|θ|$ 参数顺序进行,通常比相同粒度下的 $2nd$ KV 缓存占用空间小得多。我们报告了跨五个长上下文数据集从 1.7B 到 8B 参数的模型的实验。近似误差跟踪下一个词元的准确度差距,以充分注意,并且跨层不均匀地分配容量可以减少我们消融中的差距。除了对指标的分析之外,我们还检查配备 Nectar 模块的模型的文本生成(在问题提示之后)在语义内容上是否与通过给予相同模型访问完整缓存而获得的语义内容相匹配。