按需注意力：语言模型知道何时回忆

推理和代理工作负载越来越需要高效的长上下文推理。然而，全注意力解码会在每一步读取不断增长的历史记录，而不管它对下一个预测的好处如何。我们表明，在全局读取之前，预训练模型的解码状态已经包含预测此好处的信息。基于这一发现，我们引入了按需注意力（ODA），这是一种本地优先解码方法，它使用轻量级召回头来选择性地调用全局注意力，因为其预测的收益在生成过程中发生变化。 ODA 仅训练召回头，保持预训练权重不变，并保留完整的历史 KV 缓存以供将来召回。我们进一步在 vLLM 中实现 GPU 端条件执行，将减少的全局读取转化为在长上下文长度下充分关注的实际解码加速。 Qwen 和 Gemma 模型（包括混合注意力主干）的实验表明，选择性召回恢复了局部注意力下损失的大部分性能，同时大幅减少了全局读取。这些发现支持长上下文推理，其中预训练的模型指导他们自己访问他们保留的信息。