🤖 machine learning

Tensor Cache: Eviction-conditioned Associative Memory for Transformers

本文介绍了张量缓存，这是一种两级关联记忆系统，它将滑动窗口注意力与固定大小的外积快速权重记忆相结合，以保留和压缩被驱逐的令牌，从而在改进长上下文 Transformer 的内存质量前沿的同时，纠正了一种引入虚假跨令牌交互的常见训练捷径。

原作者： Kabir Swain, Sijie Han, Daniel Karl I. Weidele, Mauro Martino, Antonio Torralba

发布于 2026-05-25

📖 1 分钟阅读☕ 轻松阅读

原作者： Kabir Swain, Sijie Han, Daniel Karl I. Weidele, Mauro Martino, Antonio Torralba

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在向朋友讲述一个非常长的故事，同时试图记住它。为了做到这一点，你的大脑（即人工智能）会保留一个“草稿纸”，记录你最近说过的话，以便随时回看。这就是当今标准人工智能模型的工作原理。

然而，这里存在一个问题：

完整记忆问题：如果你在草稿纸上保留你说过的每一个字，它最终会变得太大而无法容纳。你的大脑会不堪重负，处理速度也会变慢。
“滑动窗口”问题：为了解决大小问题，一些模型使用“滑动窗口”。它们只保留最近（例如）1,000 个字。一旦某个字滑出这个窗口，它就会被永远丢弃。如果当前问题的答案是在 5,000 个字之前提到的，模型就完全不知道那是什么。这就像对几分钟之前的事情患上了失忆症。

Tensor Cache 是一项新发明，它通过为模型提供两部分记忆系统来解决这个问题，就像一张书桌和一个文件柜。

两部分系统

1. 书桌（一级缓存）：
这就是“滑动窗口”。模型将最近说过的字（token）直接保留在书桌上。它可以瞬间且完美地查看它们。这是为了应对即时的过去。

2. 文件柜（二级缓存）：
这是神奇的部分。当一个字从书桌上滑落，按常理会被扔进垃圾桶时，Tensor Cache 并不会将其丢弃。相反，它提取该字，并将一条摘要笔记写入一个固定大小的文件柜中。

工作原理：它不保存整个字，而是保存一个“压缩指纹”（该字的键和值的数学组合）。
检索过程：当模型稍后提出问题时，它会先查看书桌。如果答案不在那里，它就会询问文件柜：“你有关于这个主题的笔记吗？”文件柜利用一种特殊的数学技巧，快速扫描所有摘要笔记，并回答：“是的，我在很久以前有一条关于那个的提示。”

“智能”归档技巧

该论文强调了一种模型学习如何填充这个文件柜的巧妙方法。通常，当你试图一次性总结一整页文本时，你可能会不小心混淆细节（例如，因为平均了他们的话语，而误以为两个不同的人说了同样的话）。

作者发现了一种特定的数学捷径，可以防止这种混淆。他们开发了一种方法，将这些笔记逐个写入文件柜（甚至在训练期间），从而确保模型永远不会搞混哪条笔记属于哪个字。这保证了当模型回看时，“指纹”是准确的。

为什么这更好？

该论文将此方法与其他方法进行了测试，发现：

内存效率：与保留完整历史记录相比，它使用的计算机内存要少得多。即使故事变得非常长，它也能保持小巧和快速。
更好的召回率：与那些会忘记窗口外一切内容的“滑动窗口”模型不同，Tensor Cache 仍然能够记住很久以前的事情。在测试中，它能够以近乎完美的准确率回忆起数百个字之前的具体细节，而其他“小内存”模型则失败了。
速度：它比尝试保留完整历史记录要快，并且通常比其他“压缩内存”方法更快。

核心结论

可以将 Tensor Cache 想象成一位聪明的图书管理员。

旧方法：图书管理员把每一本书都放在书架上（太重），或者只保留最后几本书并烧毁其余的（你会失去整个故事）。
Tensor Cache 方法：图书管理员将最后几本书放在书桌上以便随时取用。当一本书从书桌上移走时，图书管理员会为其写下一张完美且压缩的索引卡片，并将其放入一个小型的固定大小盒子里。当你提出问题时，图书管理员会先检查书桌，如果答案不在那里，他们会快速扫描盒子里的索引卡片以找到答案。

这使得人工智能能够拥有一个“有界”（有限）的内存大小，不会无限增长，同时仍能记住非常长对话中的重要部分。

技术摘要：张量缓存（Tensor Cache）

问题陈述
自回归 Transformer 推理依赖于缓存键值（KV）对以避免重复计算前缀。然而，保留的 KV 状态随上下文长度、深度和并发量线性增长，从而造成内存瓶颈。滑动窗口缓存通过仅保留最近的 $W$ 个令牌来限制内存使用，但这种方法存在“硬性遗忘”问题：一旦令牌被逐出窗口，其信息即被完全丢弃，导致窗口外的相关证据无法获取。现有的有界状态方法要么完全丢弃被逐出的令牌，要么尝试将每个令牌压缩为运行平均值，这可能会引入噪声，或无法有效区分近期与远期上下文。

方法：张量缓存（TC）
张量缓存引入了一种双层有界状态架构，将 KV 逐出从删除事件转化为关联记忆写入操作。

双层架构：
- L1（第一级缓存）： 一个固定大小的 KV 环形缓冲区，持有最近的 $W$ 个令牌。这支持局部上下文的精确 Softmax 注意力。
- L2（第二级缓存）： 一个固定大小的外积矩阵 $A$ （快速权重记忆），专门存储从 L1 窗口中被逐出的令牌的信息。
- 写入机制： 当环形缓冲区覆盖某个条目时，被置换的 $(k, v)$ 对不会被丢弃。相反，它通过外积更新写入 L2 矩阵 $A$ ： $A \leftarrow \lambda A + \eta (k \otimes v)$ 。其中， $\lambda$ （衰减）和 $\eta$ （写入率）是每个头学习的标量。
- 读取机制： 未来的查询 $q_t$ 通过计算 $r_t = q_t A$ 来访问压缩后的过去。利用线性注意力恒等式 $q(k \otimes v) = \langle q, k \rangle v$ ，该操作有效地从所有被逐出的令牌中检索信息，而无需单独存储它们。
- 融合： 一个学习的标量门控 $g$ 融合局部注意力（L1）的输出和记忆读取（L2）的输出： $y_t = y_{local} + \sigma(g) m_t$ 。
训练优化（并行扫描）：
该论文指出了标准分块训练捷径中的一个关键问题。一种朴素的方法通过其均值（ $\bar{k} \otimes \bar{v}$ ）来总结一个分块，这会引入 $C^2 - C$ 个虚假的跨令牌外积（ $k_i \otimes v_j$ ，其中 $i \neq j$ ），而这些在流式推理中并不存在。
- 解决方案： 作者实现了一种并行加权求和扫描，在数学上等价于逐个令牌的顺序写入。这将训练与推理之间的差距缩小至 float32 epsilon 范围内（相对误差 $< 10^{-7}$ ），确保模型学习正确的关联动态，而不会受到跨令牌干扰噪声的影响。

主要贡献

逐出条件关联记忆： 与以往摄入每个令牌的外积记忆（如 Infini-attention、mLSTM）不同，TC 将外积矩阵严格用作仅由滑动窗口逐出事件驱动的 L2 缓存。这清晰地将令牌划分为“驻留在 L1"和“压缩在 L2"。
训练伪影的修正： 该论文识别并解决了一个此前被忽视的训练捷径（分块均值），该捷径引入了虚假的跨令牌乘积。所提出的并行加权求和扫描以相同的计算成本消除了这种不匹配。
有界状态效率： TC 提供了一种机制，使旧上下文能够影响预测，而无需保留单个被逐出的条目，为完全 KV 保留提供了一种实用的替代方案。

实验结果
作者在系统扩展、合成关联回忆以及真实文本长上下文语言建模（OpenWebText 和莎士比亚语料）方面评估了 TC。

合成关联回忆： 在需要检索滑动窗口之外信息的任务中，TC 实现了 100% 的准确率，与完整 KV 性能持平。与完整 KV 相比，它减少了 72–84% 的保留推理状态，同时显著优于 Infini-attention（随着间隔增加其性能下降）和滑动窗口基线（其表现仅相当于随机猜测）。
长上下文语言建模： 在 OpenWebText 上，TC 在从 1,024 到 32,768 个令牌的每个评估上下文长度下均实现了最低的负对数似然（NLL）均值，唯一的例外是 $L=2,048$ ，在此处其表现与 Infini-attention 在统计上无显著差异。在 $L=32,768$ 时，TC 的 NLL 达到 5.14，优于完整 KV（6.00）和 Infini-attention（5.42），同时使用的峰值 GPU 内存比完整 KV 少 2.4 倍。
吞吐量： 与纯滑动窗口方法相比，TC 产生了约 30% 的吞吐量开销（由于 L2 读取/更新），但在所有测试的上下文长度下，其速度始终快于 Infini-attention。
容量诊断： 关于原始关联记忆容量的实验表明，随着更多不相关的关联被压缩到固定状态中，检索性能会下降，这证实了该记忆的有界性质。

意义与主张
该论文声称，张量缓存成功弥合了滑动窗口注意力的内存效率与全上下文注意力的回忆能力之间的差距。通过将 KV 逐出视为一种结构化记忆操作而非破坏性删除，TC 使 Transformer 能够在不付出完整 KV 缓存的线性内存成本的情况下，访问即时窗口之外的相关证据。作者强调，其主要贡献并非发明外积记忆本身，而是将其具体应用于逐出条件 L2 缓存，并严格修正了此前阻碍此类方法的训练伪影。TC 被提出作为一种实用的有界状态替代方案，提升了长上下文推理的内存 - 质量前沿。

两部分系统

“智能”归档技巧

为什么这更好？

核心结论

技术摘要：张量缓存（Tensor Cache）

类似论文