张量缓存：用于Transformer的基于驱逐条件的关联记忆
张量缓存是一种两层级缓存架构，结合滑动窗口注意力作为一级缓存（L1）和固定大小的外积快速权重记忆作为二级缓存（L2），L2由窗口驱逐的KV对填充。该方法通过线性注意力恒等式实现高效读取，并引入可学习的门控融合L1和L2输出。实验表明，张量缓存在记忆-质量边界上优于有状态基线。
文章情报
要点
- 提出张量缓存，一种两层级缓存机制，结合精确局部注意力和压缩记忆。
- 二级缓存使用外积快速权重记忆，仅由滑动窗口驱逐的KV对填充。
- 识别并修复了常见分块均值训练捷径中的虚假交叉标记外积问题。
- 在多个基准测试中，张量缓存改进了记忆-质量边界。
为什么重要
这条新闻值得关注，因为提出张量缓存，一种两层级缓存机制，结合精确局部注意力和压缩记忆。
技术影响
可能影响模型选型、推理成本、产品能力和评测基准。
近日，一篇题为“Tensor Cache: Eviction-conditioned Associative Memory for Transformers”的论文在arXiv上发布，提出了一种新颖的两层级缓存架构，旨在解决自回归Transformer模型中键值（KV）缓存随上下文长度线性增长的问题。传统的滑动窗口缓存虽然能够限制内存使用，但会完全丢弃被驱逐的令牌，导致窗口外的重要信息无法被访问，从而限制了模型处理长上下文的能力。
张量缓存（Tensor Cache）通过两层级设计巧妙地将精确的局部注意力与压缩的关联记忆结合起来。第一级缓存（L1）采用标准的滑动窗口softmax注意力，保留最近令牌的精确注意力计算。第二级缓存（L2）则是一个固定大小的外积快速权重记忆，其内容完全由从滑动窗口中驱逐的键值对填充。每当一个键值对被驱逐出窗口时，它会被压缩到每层的一个矩阵A中。未来查询可以通过一次矩阵乘法高效地读取这些压缩记忆，这得益于线性注意力恒等式q_t(k_i ⊗ v_i) = ⟨q_t, k_i⟩v_i。这种设计使得模型能够以最小的内存开销访问窗口外的历史信息。
论文的作者为Kabir Swain等五位研究人员，他们贡献了多个关键创新。首先，他们引入了可学习的标量门来融合L1和L2的输出，并且每个注意力头都有独立的衰减率和写入率参数，这些参数通过端到端训练得到。其次，他们识别并修正了常见分块均值训练捷径中存在的问题。该捷径使用A←λA+η(ķ⊗v)进行更新，但这样会在每个分块中引入C²-C个虚假的跨令牌外积，从而破坏了记忆的准确性。作者通过并行加权和扫描（等同于每令牌写入，误差在float32 epsilon内）来弥补这一差距，恢复了正确的记忆更新。
在系统扩展性、受控关联回忆、长上下文语言建模和记忆容量诊断等一系列实验中，张量缓存相比有状态基线显著改进了记忆-质量边界。实验结果表明，该方法能够有效地平衡内存效率与长上下文表现，为Transformer模型的推理优化提供了新思路。这项工作可能对需要长上下文理解的应用程序（如文档分析、对话系统和代码生成）产生重要影响。