Tensor Cache: 基于驱逐条件的Transformer关联记忆

# 张量缓存：用于Transformer的基于驱逐条件的关联记忆
来源：https://arxiv.org/html/2605.22884

###### 摘要

自回归Transformer的KV缓存随上下文长度线性增长；滑动窗口缓存限制了内存但完全丢弃了被驱逐的token，因此窗口外的相关证据变得不可访问。我们引入了*张量缓存*，一种两级缓存，它将滑动窗口softmax注意力作为第一级缓存（L1），与一个固定大小的外积快速权重记忆作为第二级缓存（L2）配对，L2由从窗口驱逐的KV对提供。最近的token保留在精确的局部注意力中；被驱逐的对被压缩成一个每层矩阵$A$，并通过单个矩阵乘法供未来查询读取，利用线性注意力恒等式$q_t(k_i \otimes v_i)=\langle q_t,k_i\rangle v_i$。一个学习到的标量门控融合L1和L2的输出，并且每头的衰减和写入率参数是端到端训练的。外积记忆和读取恒等式是众所周知的；我们的贡献在于将它们用作专门由滑动窗口驱逐供给的L2缓存，并指出常见的块均值训练捷径$A \leftarrow \lambda A + \eta(\bar{k} \otimes \bar{v})$在每块中静默地引入了$C^2-C$个虚假的跨token外积，并通过一个与float32 epsilon内的逐token写入等效的并行加权和扫描来弥合差距。跨系统扩展、受控关联召回、长上下文语言建模和记忆容量诊断，张量缓存改善了与有界状态基线相比的记忆-质量前沿。

机器学习，ICML

参见图注图1：张量缓存。每层保留一个局部KV环形缓冲区（L1）和一个固定大小的记忆$A$（L2）。驱逐时，$(k,v)$通过外积更新写入$A$；查询读取两条路径并通过一个学习到的门控融合。

## 1 引言

自回归Transformer推理缓存每层的键和值（KV），以便前缀不会在每一步重复计算（Vaswani et al., 2017 (https://arxiv.org/html/2605.22884#bib.bib39); Shazeer,2019 (https://arxiv.org/html/2605.22884#bib.bib40)），但保留的KV状态随上下文长度、深度、头维度和服务并发度线性增长（Kwon et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib9)）。滑动窗口缓存通过仅保留最近的$W$个token来限制这一成本（Xiao et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib11)），但代价是*硬遗忘*：一旦token离开窗口，其键和值将完全从注意力中移除。

我们引入*张量缓存（TC）*，一种有界状态机制，将KV驱逐从删除事件转变为关联记忆写入。每层对近期窗口（第一级缓存，L1）保持精确的softmax注意力，当环形缓冲区覆盖旧条目时，被驱逐的对$(k_w,v_w)$被写入一个固定大小的外积矩阵$A$（第二级缓存，L2），这种矩阵用于快速权重编程器（Schmidhuber, 1992 (https://arxiv.org/html/2605.22884#bib.bib38)）和现代线性注意力变体（Schlag et al., 2021 (https://arxiv.org/html/2605.22884#bib.bib27); Katharopoulos et al., 2020 (https://arxiv.org/html/2605.22884#bib.bib47); Beck et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib34); Munkhdalai et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib26)）。未来查询通过下式读取此状态：

$$q_t A \approx \sum_i \alpha_i \langle q_t, k_i \rangle v_i,$$

即，注意力使用的相同查询-键对齐也用于寻址压缩的过去，但无需保留单个被驱逐的条目。一个学习到的标量门控融合L1和L2的输出。

外积记忆和读取恒等式$q_t(k_i \otimes v_i)=\langle q_t,k_i\rangle v_i$并非新事物。我们的贡献在于将它们用作专门由滑动窗口驱逐供给的L2缓存，而不是作为每个token的并行运行平均值（mLSTM、Infini-attention）或作为注意力的完全替代（RetNet、Mamba）（Sun et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib23); Gu and Dao,2023 (https://arxiv.org/html/2605.22884#bib.bib22)）。我们另外指出了一个先前被忽视的训练时间捷径：一个朴素的块扫描实现通过其均值总结每个块并写入$A \leftarrow \lambda A + \eta(\bar{k} \otimes \bar{v})$，这与逐token和相差$C^2-C$个跨token外积$k_i \otimes v_j\;(i\neq j)$，而这些在推理中从未出现。我们通过一个并行加权和扫描弥合了这一差距，该扫描在数学上等价于外规则的逐token写入，经验证与float32 epsilon一致。并行扫描公式本身在线性注意力文献中是标准的（Beck et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib34); Munkhdalai et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib26)）；我们使其对驱逐条件记忆的适用性明确，并量化了跳过它的代价。

TC并非完全注意力的无损替代。当负担得起时，完整KV仍然是精确内容基线；在超出训练块大小的评估中，完整KV的结果应解读为精确内容但位置外推（Su et al., 2021 (https://arxiv.org/html/2605.22884#bib.bib32); Chen et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib49); Peng et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib50)）。TC针对完整KV状态成本过高但纯滑动窗口解码遗忘过快的场景。

#### 贡献。

- • 一种两级缓存架构，结合了滑动窗口softmax注意力（L1）和一个专门由窗口驱逐的KV对供给的外积快速权重记忆（L2）。与之前处理每个token的外积记忆不同，我们的驱逐条件写入使矩阵特别成为局部注意力无法再触及内容的二级存储。
- • 识别并弥合了块均值训练捷径。该捷径在每个块中引入了$C^2-C$个虚假的跨token外积；并行加权和扫描以相同的计算成本消除了它们（在$10^{-7}$相对误差内验证）。
- • 跨系统扩展、真实文本长上下文语言建模、受控关联召回、原始记忆容量和消融实验的实证评估。张量缓存在使用比完整KV少72–84%推理状态的情况下，在合成召回中达到100%的匹配差距，并且在OpenWebText长上下文语言建模中，在从1024到32,768个token的每个评估上下文长度上，达到了所有测试方法（包括完整KV、窗口KV、StreamingLLM和Infini-attention）中最低的平均NLL，唯一的例外是$L=2048$，其中它与Infini-attention在统计上无法区分；在32倍于训练上下文的长度下，张量缓存达到NLL 5.14，而完整KV为6.00，且峰值GPU内存使用为2.4倍。

## 2 相关工作

#### KV缓存系统和驱逐。

KV缓存占用随上下文长度和服务并发度扩展。系统工作如PagedAttention（Kwon et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib9)）和CacheGen（Liu et al., 2023b (https://arxiv.org/html/2605.22884#bib.bib10)）改进了分配、压缩或移动，但逻辑保留状态仍随保留的上下文量扩展。另一条互补路线通过仅保留选定的token来减小KV大小：滑动窗口注意力（Longformer（Beltagy et al., 2020 (https://arxiv.org/html/2605.22884#bib.bib36)），带有注意力汇的StreamingLLM（Xiao et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib11)））和基于重要性的选择（H2O、SnapKV、CAOTE（Zhang et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib12); Li et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib13); Goel et al., 2025 (https://arxiv.org/html/2605.22884#bib.bib14)））。这些方法决定保留哪些离散条目。张量缓存则接受FIFO驱逐并将被驱逐的对转换为紧凑的关联记忆更新。

#### 记忆增强和循环Transformer。

Transformer-XL跨段重用隐藏状态（Dai et al., 2019 (https://arxiv.org/html/2605.22884#bib.bib18)），压缩Transformer压缩较早的激活（Rae et al., 2019 (https://arxiv.org/html/2605.22884#bib.bib19)），记忆Transformer使用外部kNN记忆（Wu et al., 2022 (https://arxiv.org/html/2605.22884#bib.bib20)），循环记忆Transformer通过记忆token携带信息（Bulatov et al., 2022 (https://arxiv.org/html/2605.22884#bib.bib28)）。与有界状态解码最相关的是，LESS结合了稀疏KV保留和一个低秩循环缓存（Dong et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib29)），以及Infini-attention结合了掩码局部注意力与一个通过归一化外积写入每个token更新的有界压缩记忆（Munkhdalai et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib26)）。张量缓存在写入语义上不同：一个键值对首先在精确的局部softmax注意力中使用，仅当环形缓冲区覆盖它时才将该对写入辅助状态。

#### 快速权重、外积记忆和线性注意力。

外积快速权重记忆可追溯到Schmidhuber的快速权重编程器（Schmidhuber, 1992 (https://arxiv.org/html/2605.22884#bib.bib38)）；Ba等人（2016 (https://arxiv.org/html/2605.22884#bib.bib21)）重新审视了作为短期记忆的快速权重。Widrow–Hoff（Widrow and Hoff, 1960 (https://arxiv.org/html/2605.22884#bib.bib37)）引入了delta规则背后的纠错原理。Schlag等人（2021 (https://arxiv.org/html/2605.22884#bib.bib27)）将线性注意力（Katharopoulos et al., 2020 (https://arxiv.org/html/2605.22884#bib.bib47)）重新构建为具有累积$k \otimes v$写入的快速权重编程器，形式化了我们利用的读取恒等式$q(k \otimes v)=\langle q,k\rangle v$。Yang等人（2024 (https://arxiv.org/html/2605.22884#bib.bib35)）开发了并行delta规则矩阵状态模型，而mLSTM（Beck et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib34)）重新审视了带有输入/遗忘门控和归一化向量的矩阵值循环记忆。张量缓存借用了矩阵原语和此工作线中的读取恒等式；其架构区别在于将矩阵用作*二级缓存*，由滑动窗口驱逐供给，而不是作为注意力替代（mLSTM、RetNet（Sun et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib23)）、Mamba（Gu and Dao,2023 (https://arxiv.org/html/2605.22884#bib.bib22)））或作为每个token的并行运行平均值（Infini-attention）。我们不需要Infini的归一化向量或特征映射，而是使用一个学习到的指数衰减。

#### 训练方面的考虑。

mLSTM和Infini-attention的段扫描在块训练期间保留了逐token写入语义。我们观察到块均值捷径$A \leftarrow \lambda A + \eta(\bar{k} \otimes \bar{v})$是一个诱人但有损的近似，它在每个块中引入了$C^2-C$个虚假的跨token外积，并且当关联片段适合单个块时降低检索性能。用标准的并行加权和扫描（在数学上等价于外规则的逐token写入）替换块均值写入，弥合了差距。扫描本身是标准线性循环材料；我们使其对驱逐条件外积记忆的适用性明确，并量化了跳过它的代价。

#### 其他长上下文方法。

环注意力（Liu et al., 2023a (https://arxiv.org/html/2605.22884#bib.bib16)）分布块状注意力；LongNet（Ding et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib17)）使用扩张注意力；FlashAttention-2（Dao,2023 (https://arxiv.org/html/2605.22884#bib.bib25)）提高吞吐量。这些在很大程度上是正交的：张量缓存通过一个固定大小的关联记忆路径来增强用于被驱逐KV条目的滑动窗口注意力堆栈，并且可以与此类系统级优化相结合。

## 3 架构

张量缓存（TC）在每个Transformer层中增加三个组件：用于精确局部注意力的固定容量KV环形缓冲区（L1），用于被驱逐信息的固定大小关联记忆状态（L2），以及一个融合两者的学习到的门控。

### 3.1 设置和局部KV窗口

设$B$为批量大小，$H$为头数，$D$为头维度。每层从单个输入投影产生$Q,K,V \in \mathbb{R}^{B \times H \times T \times D}$，可选地对$Q,K$应用RoPE。局部窗口

$$K_{\mathrm{win}}, V_{\mathrm{win}} \in \mathbb{R}^{B \times H \times W \times D} \quad (1)$$

保存最近的$W$个键/值对，并用于精确因果注意力$y_t^{\mathrm{local}} = \operatorname{Attn}(q_t, K_{\mathrm{win}}, V_{\mathrm{win}})$。RoPE在键进入缓冲区之前应用，因此稍后写入记忆的键是RoPE后的。

### 3.2 基于驱逐条件的记忆状态

每层维护一个快速权重矩阵$A \in \mathbb{R}^{B \times H \times D \times D}$，初始化为$A_0 = \mathbf{0}$。KV缓存是一个FIFO环形缓冲区；当槽位$\mathrm{pos}$被覆盖时，被替换的对$(k_{\mathrm{old}}, v_{\mathrm{old}})$成为TC写入源$(k_w, v_w)$（*写时驱逐*）。这是TC与先前外积记忆（如mLSTM（Beck et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib34)）和Infini-attention（Munkhdalai et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib26)））之间的架构区别，后者无条件写入每个token；TC仅写入局部窗口丢弃的内容。具体来说，每token写入会重复计算局部窗口已通过精确注意力服务到的token；将写入限制在驱逐上，可以清晰地划分token群体为“仍驻留在L1中”和“压缩到L2中”，因此每个先前的token恰好贡献给一个读取路径。

参见图注图2：张量缓存的一个流式步骤，从左到右显示。（1）局部窗口已满：局部KV环形缓冲区保存最近的$W$个键/值对，同时新对$(k_t, v_t)$到达。（2）驱逐最旧的：被替换的对$(k_{\mathrm{old}}, v_{\mathrm{old}})$从缓冲区内弹出以腾出空间。（3）更新张量缓存：被驱逐的对通过外积（或可选的delta规则）更新写入L2注意力记忆$A$：$A \leftarrow \lambda A + \eta \, \mathrm{update}(k_{\mathrm{old}}, v_{\mathrm{old}})$，压缩长期信息。（4）插入新token：$(k_t, v_t)$占据环形缓冲区的最新槽位。（5）读取并融合：当前查询$q_t$读取两个路径——局部窗口输出$y_t^{\mathrm{local}}$和记忆输出$m_t = q_t A$——通过一个学习到的标量门控组合：$y_t = y_t^{\mathrm{local}} + \sigma(g) \, m_t$。

### 3.3 记忆读取和写入

当前查询读取$r_t = q_t A \in \mathbb{R}^{B \times H \times D}$。在累积的外积写入$A \approx \sum_i \alpha_i (k_i$后，注意：原文在此处截断，但根据上下文，应该是$A \approx \sum_i \alpha_i (k_i \otimes v_i)$ 或类似。但原文只写到“$A \approx \sum_i \alpha_i (k_i$”，我们保持原文。实际上参考上文，应是$A \approx \sum_i \alpha_i (k_i \otimes v_i)$，但按原文保留。我们按原文翻译：

在累积的外积写入$A \approx \sum_i \alpha_i (k_i$之后，需要根据后续内容。但原文这样写，我们尊重原文。不过为了完整，我们可能需要在翻译时保持原样，但注意这是不完整的。检查原文：

原文："After accumulated outer-product writes $A \approx \sum_i \alpha_i (k_i$"

可能是笔误，应继续到 $k_i \otimes v_i$。但在翻译中我们保持原文的不完整状态，因为用户输入如此。我们按原文翻译。

所以译文：在累积的外积写入$A \approx \sum_i \alpha_i (k_i$之后，查询读取等。我们保持原样。# 张量缓存：用于Transformer的基于驱逐条件的关联记忆
来源: https://arxiv.org/html/2605.22884

###### 摘要

自回归Transformer的KV缓存随上下文长度线性增长；滑动窗口缓存限制了内存但完全丢弃了被驱逐的token，因此窗口外的相关证据变得不可访问。我们引入了*张量缓存*，一种两级缓存，它将滑动窗口softmax注意力作为第一级缓存（L1），与一个固定大小的外积快速权重记忆作为第二级缓存（L2）配对，L2由从窗口驱逐的KV对提供。最近的token保留在精确的局部注意力中；被驱逐的对被压缩成一个每层矩阵$A$，并通过单个矩阵乘法供未来查询读取，利用线性注意力恒等式$q_t(k_i \otimes v_i)=\langle q_t,k_i\rangle v_i$。一个学习到的标量门控融合L1和L2的输出，并且每头的衰减和写入率参数是端到端训练的。外积记忆和读取恒等式是众所周知的；我们的贡献在于将它们用作专门由滑动窗口驱逐供给的L2缓存，并指出常见的块均值训练捷径$A \leftarrow \lambda A + \eta(\bar{k} \otimes \bar{v})$在每块中静默地引入了$C^2-C$个虚假的跨token外积，并通过一个与float32 epsilon内的逐token写入等效的并行加权和扫描来弥合差距。跨系统扩展、受控关联召回、长上下文语言建模和记忆容量诊断，张量缓存改善了与有界状态基线相比的记忆-质量前沿。

机器学习，ICML

参见图注图1：张量缓存。每层保留一个局部KV环形缓冲区（L1）和一个固定大小的记忆$A$（L2）。驱逐时，$(k,v)$通过外积更新写入$A$；查询读取两条路径并通过一个学习到的门控融合。

## 1 引言

自回归Transformer推理缓存每层的键和值（KV），以便前缀不会在每一步重复计算（Vaswani et al., 2017 (https://arxiv.org/html/2605.22884#bib.bib39); Shazeer,2019 (https://arxiv.org/html/2605.22884#bib.bib40)），但保留的KV状态随上下文长度、深度、头维度和服务并发度线性增长（Kwon et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib9)）。滑动窗口缓存通过仅保留最近的$W$个token来限制这一成本（Xiao et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib11)），但代价是*硬遗忘*：一旦token离开窗口，其键和值将完全从注意力中移除。

我们引入*张量缓存（TC）*，一种有界状态机制，将KV驱逐从删除事件转变为关联记忆写入。每层对近期窗口（第一级缓存，L1）保持精确的softmax注意力，当环形缓冲区覆盖旧条目时，被驱逐的对$(k_w,v_w)$被写入一个固定大小的外积矩阵$A$（第二级缓存，L2），这种矩阵用于快速权重编程器（Schmidhuber, 1992 (https://arxiv.org/html/2605.22884#bib.bib38)）和现代线性注意力变体（Schlag et al., 2021 (https://arxiv.org/html/2605.22884#bib.bib27); Katharopoulos et al., 2020 (https://arxiv.org/html/2605.22884#bib.bib47); Beck et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib34); Munkhdalai et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib26)）。未来查询通过下式读取此状态：

$$q_t A \approx \sum_i \alpha_i \langle q_t, k_i \rangle v_i,$$

即，注意力使用的相同查询-键对齐也用于寻址压缩的过去，但无需保留单个被驱逐的条目。一个学习到的标量门控融合L1和L2的输出。

外积记忆和读取恒等式$q_t(k_i \otimes v_i)=\langle q_t,k_i\rangle v_i$并非新事物。我们的贡献在于将它们用作专门由滑动窗口驱逐供给的L2缓存，而不是作为每个token的并行运行平均值（mLSTM、Infini-attention）或作为注意力的完全替代（RetNet、Mamba）（Sun et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib23); Gu and Dao,2023 (https://arxiv.org/html/2605.22884#bib.bib22)）。我们另外指出了一个先前被忽视的训练时间捷径：一个朴素的块扫描实现通过其均值总结每个块并写入$A \leftarrow \lambda A + \eta(\bar{k} \otimes \bar{v})$，这与逐token和相差$C^2-C$个跨token外积$k_i \otimes v_j\;(i\neq j)$，而这些在推理中从未出现。我们通过一个并行加权和扫描弥合了这一差距，该扫描在数学上等价于外规则的逐token写入，经验证与float32 epsilon一致。并行扫描公式本身在线性注意力文献中是标准的（Beck et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib34); Munkhdalai et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib26)）；我们使其对驱逐条件记忆的适用性明确，并量化了跳过它的代价。

TC并非完全注意力的无损替代。当负担得起时，完整KV仍然是精确内容基线；在超出训练块大小的评估中，完整KV的结果应解读为精确内容但位置外推（Su et al., 2021 (https://arxiv.org/html/2605.22884#bib.bib32); Chen et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib49); Peng et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib50)）。TC针对完整KV状态成本过高但纯滑动窗口解码遗忘过快的场景。

#### 贡献。

- • 一种两级缓存架构，结合了滑动窗口softmax注意力（L1）和一个专门由窗口驱逐的KV对供给的外积快速权重记忆（L2）。与之前处理每个token的外积记忆不同，我们的驱逐条件写入使矩阵特别成为局部注意力无法再触及内容的二级存储。
- • 识别并弥合了块均值训练捷径。该捷径在每个块中引入了$C^2-C$个虚假的跨token外积；并行加权和扫描以相同的计算成本消除了它们（在$10^{-7}$相对误差内验证）。
- • 跨系统扩展、真实文本长上下文语言建模、受控关联召回、原始记忆容量和消融实验的实证评估。张量缓存在使用比完整KV少72–84%推理状态的情况下，在合成召回中达到100%的匹配差距，并且在OpenWebText长上下文语言建模中，在从1024到32,768个token的每个评估上下文长度上，达到了所有测试方法（包括完整KV、窗口KV、StreamingLLM和Infini-attention）中最低的平均NLL，唯一的例外是$L=2048$，其中它与Infini-attention在统计上无法区分；在32倍于训练上下文的长度下，张量缓存达到NLL 5.14，而完整KV为6.00，且峰值GPU内存使用为2.4倍。

## 2 相关工作

#### KV缓存系统和驱逐。

KV缓存占用随上下文长度和服务并发度扩展。系统工作如PagedAttention（Kwon et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib9)）和CacheGen（Liu et al., 2023b (https://arxiv.org/html/2605.22884#bib.bib10)）改进了分配、压缩或移动，但逻辑保留状态仍随保留的上下文量扩展。另一条互补路线通过仅保留选定的token来减小KV大小：滑动窗口注意力（Longformer（Beltagy et al., 2020 (https://arxiv.org/html/2605.22884#bib.bib36)），带有注意力汇的StreamingLLM（Xiao et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib11)））和基于重要性的选择（H2O、SnapKV、CAOTE（Zhang et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib12); Li et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib13); Goel et al., 2025 (https://arxiv.org/html/2605.22884#bib.bib14)））。这些方法决定保留哪些离散条目。张量缓存则接受FIFO驱逐并将被驱逐的对转换为紧凑的关联记忆更新。

#### 记忆增强和循环Transformer。

Transformer-XL跨段重用隐藏状态（Dai et al., 2019 (https://arxiv.org/html/2605.22884#bib.bib18)），压缩Transformer压缩较早的激活（Rae et al., 2019 (https://arxiv.org/html/2605.22884#bib.bib19)），记忆Transformer使用外部kNN记忆（Wu et al., 2022 (https://arxiv.org/html/2605.22884#bib.bib20)），循环记忆Transformer通过记忆token携带信息（Bulatov et al., 2022 (https://arxiv.org/html/2605.22884#bib.bib28)）。与有界状态解码最相关的是，LESS结合了稀疏KV保留和一个低秩循环缓存（Dong et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib29)），以及Infini-attention结合了掩码局部注意力与一个通过归一化外积写入每个token更新的有界压缩记忆（Munkhdalai et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib26)）。张量缓存在写入语义上不同：一个键值对首先在精确的局部softmax注意力中使用，仅当环形缓冲区覆盖它时才将该对写入辅助状态。

#### 快速权重、外积记忆和线性注意力。

外积快速权重记忆可追溯到Schmidhuber的快速权重编程器（Schmidhuber, 1992 (https://arxiv.org/html/2605.22884#bib.bib38)）；Ba等人（2016 (https://arxiv.org/html/2605.22884#bib.bib21)）重新审视了作为短期记忆的快速权重。Widrow–Hoff（Widrow and Hoff, 1960 (https://arxiv.org/html/2605.22884#bib.bib37)）引入了delta规则背后的纠错原理。Schlag等人（2021 (https://arxiv.org/html/2605.22884#bib.bib27)）将线性注意力（Katharopoulos et al., 2020 (https://arxiv.org/html/2605.22884#bib.bib47)）重新构建为具有累积$k \otimes v$写入的快速权重编程器，形式化了我们利用的读取恒等式$q(k \otimes v)=\langle q,k\rangle v$。Yang等人（2024 (https://arxiv.org/html/2605.22884#bib.bib35)）开发了并行delta规则矩阵状态模型，而mLSTM（Beck et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib34)）重新审视了带有输入/遗忘门控和归一化向量的矩阵值循环记忆。张量缓存借用了矩阵原语和此工作线中的读取恒等式；其架构区别在于将矩阵用作*二级缓存*，由滑动窗口驱逐供给，而不是作为注意力替代（mLSTM、RetNet（Sun et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib23)）、Mamba（Gu and Dao,2023 (https://arxiv.org/html/2605.22884#bib.bib22)））或作为每个token的并行运行平均值（Infini-attention）。我们不需要Infini的归一化向量或特征映射，而是使用一个学习到的指数衰减。

#### 训练方面的考虑。

mLSTM和Infini-attention的段扫描在块训练期间保留了逐token写入语义。我们观察到块均值捷径$A \leftarrow \lambda A + \eta(\bar{k} \otimes \bar{v})$是一个诱人但有损的近似，它在每个块中引入了$C^2-C$个虚假的跨token外积，并且当关联片段适合单个块时降低检索性能。用标准的并行加权和扫描（在数学上等价于外规则的逐token写入）替换块均值写入，弥合了差距。扫描本身是标准线性循环材料；我们使其对驱逐条件外积记忆的适用性明确，并量化了跳过它的代价。

#### 其他长上下文方法。

环注意力（Liu et al., 2023a (https://arxiv.org/html/2605.22884#bib.bib16)）分布块状注意力；LongNet（Ding et al., 2023 (https://arxiv.org/html/2605.22884#bib.bib17)）使用扩张注意力；FlashAttention-2（Dao,2023 (https://arxiv.org/html/2605.22884#bib.bib25)）提高吞吐量。这些在很大程度上是正交的：张量缓存通过一个固定大小的关联记忆路径来增强用于被驱逐KV条目的滑动窗口注意力堆栈，并且可以与此类系统级优化相结合。

## 3 架构

张量缓存（TC）在每个Transformer层中增加三个组件：用于精确局部注意力的固定容量KV环形缓冲区（L1），用于被驱逐信息的固定大小关联记忆状态（L2），以及一个融合两者的学习到的门控。

### 3.1 设置和局部KV窗口

设$B$为批量大小，$H$为头数，$D$为头维度。每层从单个输入投影产生$Q,K,V \in \mathbb{R}^{B \times H \times T \times D}$，可选地对$Q,K$应用RoPE。局部窗口

$$K_{\mathrm{win}}, V_{\mathrm{win}} \in \mathbb{R}^{B \times H \times W \times D} \quad (1)$$

保存最近的$W$个键/值对，并用于精确因果注意力$y_t^{\mathrm{local}} = \operatorname{Attn}(q_t, K_{\mathrm{win}}, V_{\mathrm{win}})$。RoPE在键进入缓冲区之前应用，因此稍后写入记忆的键是RoPE后的。

### 3.2 基于驱逐条件的记忆状态

每层维护一个快速权重矩阵$A \in \mathbb{R}^{B \times H \times D \times D}$，初始化为$A_0 = \mathbf{0}$。KV缓存是一个FIFO环形缓冲区；当槽位$\mathrm{pos}$被覆盖时，被替换的对$(k_{\mathrm{old}}, v_{\mathrm{old}})$成为TC写入源$(k_w, v_w)$（*写时驱逐*）。这是TC与先前外积记忆（如mLSTM（Beck et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib34)）和Infini-attention（Munkhdalai et al., 2024 (https://arxiv.org/html/2605.22884#bib.bib26)））之间的架构区别，后者无条件写入每个token；TC仅写入局部窗口丢弃的内容。具体来说，每token写入会重复计算局部窗口已通过精确注意力服务到的token；将写入限制在驱逐上，可以清晰地划分token群体为“仍驻留在L1中”和“压缩到L2中”，因此每个先前的token恰好贡献给一个读取路径。

参见图注图2：张量缓存的一个流式步骤，从左到右显示。（1）局部窗口已满：局部KV环形缓冲区保存最近的$W$个键/值对，同时新对$(k_t, v_t)$到达。（2）驱逐最旧的：被替换的对$(k_{\mathrm{old}}, v_{\mathrm{old}})$从缓冲区内弹出以腾出空间。（3）更新张量缓存：被驱逐的对通过外积（或可选的delta规则）更新写入L2注意力记忆$A$：$A \leftarrow \lambda A + \eta \, \mathrm{update}(k_{\mathrm{old}}, v_{\mathrm{old}})$，压缩长期信息。（4）插入新token：$(k_t, v_t)$占据环形缓冲区的最新槽位。（5）读取并融合：当前查询$q_t$读取两个路径——局部窗口输出$y_t^{\mathrm{local}}$和记忆输出$m_t = q_t A$——通过一个学习到的标量门控组合：$y_t = y_t^{\mathrm{local}} + \sigma(g) \, m_t$。

### 3.3 记忆读取和写入

当前查询读取$r_t =