论文

用于高效长上下文生成的上下文记忆

Context Memorization for Efficient Long Context Generation

模型推理KV Cache

摘要

现代 大语言模型 (LLM) 应用程序越来越依赖长条件前缀来控制推理时的模型行为。虽然前缀增强推理是有效的,但它带来了两个结构限制:i)前缀的影响随着生成的进行而减弱,ii)前缀上的注意力计算与其长度成线性比例。现有的方法要么在压缩时保持前缀的注意力,要么通过基于梯度的训练将其内化到模型参数中。前者仍然在推理时关注前缀,而后者是训练密集型的并且不适合前缀更新。为了解决这些问题,我们提出了注意力状态记忆,这是一种 无需训练 方法,它将前缀外部化为前缀和查询标记之间预先计算的注意力状态的轻量级、基于查找的记忆。在使用 LLaMA-3.1-8B 的 ManyICLBench 上,我们的方法在 1K-8K 内存预算下提高了上下文学习的准确性,同时在 8K 下将注意力延迟降低了 1.36 倍,并且仅使用 20% 的内存占用就超过了 NBA 基准测试上的全注意力 RAG 性能。