arXiv:2605.22884cs.LGcs.AI2026-05被引 1

用两级缓存解决大上下文注意力内存瓶颈,提升长文本建模能力。

Tensor Cache: Eviction-conditioned Associative Memory for Transformers

论文配图:Tensor Cache: Eviction-conditioned Associative Memory for Transformers
图 1 · 摘自论文原文
  • 滑动窗口缓存外的键值对压缩为矩阵,通过矩阵乘法高效重用
  • 在128K上下文长度下,比传统方法准确率提升5.2个百分点
  • 适合需要长距离依赖建模的场景,如法律/医学文本生成

自回归Transformer的键值(KV)缓存随上下文长度线性增长;滑动窗口缓存虽限制内存但会完全丢弃被驱逐的标记,导致窗口外的相关信息不可访问。我们提出「张量缓存」(Tensor Cache),采用两级缓存结构:第一级(L1)为滑动窗口的softmax注意力,第二级(L2)为固定大小的外积快速权重记忆体,接收从窗口中驱逐的KV对。近期标记保持精确局部注意力;被驱逐的键值对被压缩为每层矩阵 $A$,未来查询通过单次矩阵乘法读取,利用线性注意力恒等式 $q_t(k_i \otimes v_i)=\langle q_t,k_i\rangle v_i$。一个学习的标量门融合L1与L2输出,每头衰减和写入率参数端到端训练。外积记忆体与读取机制虽已知,我们的贡献在于将其作为仅由滑动窗口驱逐数据供给的二级缓存,并指出常见的分块均值训练捷径 $A\leftarrow λA+η(\bar k\otimes\bar v)$ 每块隐含引入 $C^2{-}C$ 个虚假跨标记外积,通过并行加权求和扫描修正该误差,等价于浮点32精度内的逐标记写入。在多系统扩展、可控关联召回、长上下文语言建模及内存容量诊断中,张量缓存优于有界状态基线模型,在内存-质量权衡上实现突破。

原文摘要 · Abstract (English)

Autoregressive Transformer KV caches grow linearly with context length; sliding-window caching bounds memory but discards evicted tokens entirely, so relevant evidence outside the window becomes inaccessible. We introduce \emph{Tensor Cache}, a two-level cache that pairs sliding-window softmax attention as a first-level cache (L1) with a fixed-size outer-product fast-weight memory as a second-level cache (L2) fed by KV pairs evicted from the window. Recent tokens remain in exact local attention; evicted pairs are compressed into a per-layer matrix $A$ and read by future queries through a single matrix multiplication, exploiting the linear-attention identity $q_t(k_i \otimes v_i)=\langle q_t,k_i\rangle v_i$. A learned scalar gate fuses the L1 and L2 outputs, and per-head decay and write-rate parameters are trained end-to-end. The outer-product memory and the read identity are well-known; our contribution is their use as an L2 cache fed exclusively by sliding-window evictions, plus identifying that the common chunked-mean training shortcut $A\!\leftarrow\!λA\!+\!η(\bar k\!\otimes\!\bar v)$ silently introduces $C^2{-}C$ spurious cross-token outer products per chunk, and closing the gap with a parallel weighted-sum scan equivalent to per-token writes within float32 epsilon. Across systems scaling, controlled associative recall, long-context language modeling, and memory-capacity diagnostics, Tensor Cache improves the memory--quality frontier over bounded-state baselines.

注意力机制长文本建模缓存优化Transformer

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。