精读笔记：Infini-attention——无限上下文 Transformer

论文全名：Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention
arXiv 编号：2404.07143
发表机构：Google（作者：Tsendsuren Munkhdalai, Manaal Faruqui, Sasha Golovnev）
发表时间：2024 年 4 月
本地 PDF：Infini-attention 无限上下文 [2404.07143].pdf

阅读地图

本笔记按以下顺序精读：

前置问题：标准注意力的 KV 缓存为什么会"撑爆"显存？
Abstract：论文用一段话说了什么？
Introduction：为什么现有方案不够好？Infini-attention 的核心思路是什么？
方法核心（全文最重要部分）：
- 标准缩放点积注意力回顾
- 压缩记忆（Compressive Memory）是什么
- 记忆检索（Memory Retrieval）
- 记忆更新（Memory Update）——Linear 与 Linear+Delta 两种策略
- 长期上下文注入（Long-term Context Injection）——门控机制
- 多头压缩记忆（Multi-head Compressive Memory）
- 内存复杂度对比
实验：三个关键任务的结论
与 StreamingLLM / Mamba 的思路对比
总结与局限

前置问题：固定大小的"笔记本"能装下无限的历史吗？

在正式进入论文之前，先把最核心的工程矛盾讲清楚。

标准注意力（Standard Self-Attention）的 KV 缓存问题

当 Transformer 处理一段文本时，它对每一个词都要计算 Key（键）和 Value（值）向量，并把它们存下来，以便后续的词可以"回头看"前面所有词——这就是 KV Cache（KV 缓存）。

术语：KV Cache（KV 缓存）
在自回归生成中，为了避免重复计算，模型把历史 token 的 Key 和 Value 向量全部存进一块显存缓冲区。每新来一个 token，缓冲区就增长一格。

问题在于：KV 缓存的大小与序列长度成正比。论文给出了一个具体数字：一个 500B 参数的模型，批大小 512、上下文 2048，KV 缓存就已经占据了 3 TB 显存。如果上下文扩展到 1M（百万）tokens，显存需求将彻底不可承受。

核心矛盾

能不能用一个固定大小的记忆，装下近乎无限长的历史？

这个问题的答案，正是 Infini-attention 试图给出的。

类比：固定页数的随身笔记本

想象你在边读一本百万字的长篇小说，边用一本只有固定页数的随身笔记本做记录。你不可能把每句话都抄进去，但你可以不断把新读到的内容压缩、浓缩，覆盖到已有的笔记上。这本笔记因此永远保持固定页数，却隐约"记住"了你读过的全部内容——只是记忆是有损的、浓缩的，不是完美的复现。

Infini-attention 做的正是这件事：在每个注意力头里放一个固定大小的"压缩记忆矩阵"，每处理完一段文字，就把这段的信息"压缩累加"进去，然后继续处理下一段。

1. Abstract（摘要）

原文（关键句）

"This work introduces an efficient method to scale Transformer-based Large Language Models (LLMs) to infinitely long inputs with bounded memory and computation. A key component in our proposed approach is a new attention technique dubbed Infini-attention. The Infini-attention incorporates a compressive memory into the vanilla attention mechanism and builds in both masked local attention and long-term linear attention mechanisms in a single Transformer block."

"We demonstrate the effectiveness of our approach on long-context language modeling benchmarks, 1M sequence length passkey context block retrieval and 500K length book summarization tasks with 1B and 8B LLMs. Our approach introduces minimal bounded memory parameters and enables fast streaming inference for LLMs."

翻译

本文提出了一种高效方法，可将基于 Transformer 的大型语言模型（LLM）扩展到无限长的输入，同时保持有界的内存和计算量。我们方法的核心是一种名为 Infini-attention 的新型注意力机制。Infini-attention 将压缩记忆（compressive memory）融入标准注意力机制，在单个 Transformer 块内同时构建了带掩码的局部注意力和长期线性注意力两种机制。

我们在长上下文语言建模基准测试、百万 token（1M）长度的密钥检索（passkey retrieval）以及 500K token 长度的书籍摘要任务上，使用 1B 和 8B 参数的 LLM 验证了方法的有效性。我们的方法仅引入极少量固定的记忆参数，并支持 LLM 的快速流式推理。

新手讲解

摘要用三句话说了三件事：

做了什么：提出 Infini-attention，让 Transformer 能处理"无限长"输入，但内存是固定的。
怎么做到的：在注意力层里塞一个"压缩记忆"，同时保留局部注意力（看眼前这段）和长期线性注意力（查历史压缩记忆）。
效果如何：在 1M token 的密钥检索和 500K token 的书籍摘要上都能正常工作。

术语：Linear Attention（线性注意力）
标准注意力的计算量是 O(N²)（N 是序列长度），因为每个词都要和所有其他词做点积。线性注意力通过数学变换，把复杂度降到 O(N)，代价是无法精确地使用 softmax。它常被用来实现高效的循环式记忆更新。

2. Introduction（引言）

原文（关键段落 1）

"Memory serves as a cornerstone of intelligence, as it enables efficient computations tailored to specific contexts. However, Transformers and Transformer-based LLMs have a constrained context-dependent memory, due to the nature of the attention mechanism."

"The attention mechanism in Transformers exhibits quadratic complexity in both memory footprint and computation time. For example, the attention Key-Value (KV) states have 3TB memory footprint for a 500B model with batch size 512 and context length 2048. Indeed, scaling LLMs to longer sequences (i.e. 1M tokens) is challenging with the standard Transformer architectures and serving longer and longer context models becomes costly financially."

翻译

记忆是智能的基石，因为它使得针对特定上下文的高效计算成为可能。然而，由于注意力机制的固有特性，Transformer 及其衍生的大型语言模型的上下文记忆是受约束的。

Transformer 中的注意力机制在内存占用和计算时间上均表现出平方级复杂度。以一个 500B 参数模型为例，批大小 512、上下文长度 2048 时，注意力的 KV 状态就已占用 3 TB 显存。实际上，使用标准 Transformer 架构将 LLM 扩展到更长序列（如 100 万 token）面临极大挑战，服务越来越长的上下文模型的经济成本也越来越高。

新手讲解

这两段确立了问题的严峻性。标准 Transformer 注意力是"平方级"复杂度：序列长度翻倍，显存和计算量就变成 4 倍。100 万 token 的序列，计算量是 1000 token 序列的 100 万倍。

原文（关键段落 2）

"Compressive memory systems promise to be more scalable and efficient than the attention mechanism for extremely long sequences. Instead of using an array that grows with the input sequence length, a compressive memory primarily maintains a fixed number of parameters to store and recall information with a bounded storage and computation costs. In the compressive memory, new information is added to the memory by changing its parameters with an objective that this information can be recovered back later on. However, the LLMs in their current state have yet to see an effective, practical compressive memory technique that balances simplicity along with quality."

翻译

对于极长序列，压缩记忆系统有望比注意力机制更具可扩展性和效率。压缩记忆不是使用随输入序列长度增长的数组，而是维护固定数量的参数来存储和检索信息，具有有界的存储和计算成本。在压缩记忆中，新信息通过修改参数的方式被写入，目标是使这些信息日后可以被恢复。然而，当前的 LLM 还没有见到一种既简单又高质量的有效压缩记忆技术。

新手讲解

这段是铺垫：已有"压缩记忆"的想法，但实际做好的方案还没出现，于是本文来填这个空。

术语：Compressive Memory（压缩记忆）
一种固定大小的记忆系统。不把所有历史信息原样存储（那会无限增长），而是把信息"压缩"进固定大小的参数矩阵里。就像把一本书浓缩成一页摘要——占空间固定，但内容是有损的。

原文（关键段落 3：核心思路）

"Our Infini-attention reuses all the key, value and query states of the standard attention computation for long-term memory consolidation and retrieval. We store old KV states of the attention in the compressive memory, instead of discarding them like in the standard attention mechanism. We then retrieve the values from the memory by using the attention query states when processing subsequent sequences. To compute the final contextual output, the Infini-attention aggregates the long-term memory-retrieved values and the local attention contexts."

翻译

我们的 Infini-attention 复用了标准注意力计算中的所有键（Key）、值（Value）和查询（Query）状态，用于长期记忆的整合与检索。我们将注意力的旧 KV 状态存入压缩记忆，而不是像标准注意力那样直接丢弃。在处理后续序列时，我们通过注意力的查询状态从记忆中检索相关值。最终，Infini-attention 将长期记忆检索到的值与局部注意力上下文聚合，得到最终输出。

新手讲解

这段是整篇论文最关键的设计思路，用三步描述了 Infini-attention 的工作流程：

存：处理完每一段文字后，不丢弃 KV 状态，而是把它"压缩累加"进一个固定大小的记忆矩阵。
取：处理下一段时，用当前的 Query 向量去查询这个压缩记忆，取出历史相关信息。
融合：把"从压缩记忆取出的历史信息"和"当前段的局部注意力输出"用一个门控机制混合，作为最终输出。

原文（贡献总结）

"1. We introduce a practical and yet powerful attention mechanism – Infini-attention with long-term compressive memory and local causal attention for efficiently modeling both long and short-range contextual dependencies.
2. Infini-attention introduces minimal change to the standard scaled dot-product attention and supports plug-and-play continual pre-training and long-context adaptation by design.
3. Our approach enables Transformer LLMs to scale to infinitely long context with a bounded memory and compute resource by processing extremely long inputs in a streaming fashion."

翻译

提出了 Infini-attention：将长期压缩记忆与局部因果注意力结合，高效建模长程与短程依赖。
对标准注意力的改动极小，支持即插即用（plug-and-play）的持续预训练和长上下文适配。
通过流式（streaming）处理极长输入，使 LLM 在有界内存和计算资源下扩展到无限上下文。

3. 方法：Infini-attention 机制（核心，全部段落精读）

这是本文最重要的部分，请逐段仔细阅读。

3.1 标准缩放点积注意力回顾

原文

"The multi-head scaled dot-product attention, specially its self-attention variant, has been the main building block in LLMs. The MHA's strong capability to model context-dependent dynamic computation and its conveniences of temporal masking have been leveraged extensively in the autoregressive generative models.

A single head in the vanilla MHA computes its attention context from sequence of input segments as follows. First, it computes attention query, key, and value states:

K = XW_K, V = XW_V, Q = XW_Q

Here, W_K, W_V and W_Q are trainable projection matrices. Then, the attention context is calculated as a weighted average of all other values as

A_dot = softmax(QKᵀ / √d_model) V

For MHA, we compute H number of attention context vectors for each sequence element in parallel, concatenate them along the second dimension and then finally project the concatenated vector to the model space to obtain the attention output."

翻译

多头缩放点积注意力（尤其是自注意力变体）一直是 LLM 的核心构建模块。MHA 强大的上下文相关动态计算能力及其对时序掩码的便捷支持，在自回归生成模型中得到了广泛应用。

标准 MHA 中的单个注意力头按如下方式从输入片段序列中计算注意力上下文：首先计算注意力的查询、键和值状态：

K = XW_K，V = XW_V，Q = XW_Q

其中 W_K、W_V、W_Q 是可训练的投影矩阵。然后，注意力上下文计算为所有其他值的加权平均：

A_dot = softmax(QKᵀ / √d_model) · V

对于 MHA，我们并行计算 H 个注意力上下文向量，沿第二维度拼接，最后投影到模型空间得到注意力输出。

新手讲解

这是在复习标准注意力机制，作为对比基线。

核心操作：每个 token 生成三个向量——Query（我想找什么？）、Key（我是什么？）、Value（我的内容是什么？）。Query 和 Key 做点积得到"相关性分数"，经过 softmax 归一化后，加权平均 Value，得到这个 token 的上下文表示。

问题：Q 要和序列中所有 K 做点积，计算量是 O(N²)，且所有 KV 都要存在显存里——这正是 KV 缓存随序列长度无限增长的根本原因。

3.2 压缩记忆（Compressive Memory）

原文

"In Infini-attention, instead of computing new memory entries for compressive memory, we reuse the query, key and value states (Q, K and V) from the dot-product attention computation. The state sharing and reusing between the dot-product attention and compressive memory not only enables efficient plug-in-play long-context adaptation but also speeds up training and inference."

"While there are different forms of compressive memory proposed in the literature, for simplicity and computational efficiency, in this work we parameterize the memory with an associative matrix. This approach further allows us to cast the memory update and retrieval process as linear attention mechanism and to leverage stable training techniques from the related methods."

翻译

在 Infini-attention 中，我们不为压缩记忆单独计算新的记忆条目，而是复用点积注意力计算中的查询、键和值状态（Q、K、V）。点积注意力与压缩记忆之间的状态共享与复用，不仅支持高效的即插即用长上下文适配，还能加速训练和推理。

在现有文献中有多种形式的压缩记忆方案，为了简单和计算效率，本文使用关联矩阵（associative matrix）来参数化记忆。这种方法使得记忆的更新和检索过程可以被视为线性注意力机制，并能借鉴相关方法中的稳定训练技术。

新手讲解

这两段说了两件很聪明的设计决策：

决策 1：复用 Q、K、V，不另起炉灶

标准注意力已经计算了 Q、K、V。Infini-attention 不增加额外的 Q'/K'/V' 专门服务压缩记忆，而是直接用同一套 Q、K、V。这样做的好处：改动极小，几乎可以"零成本"地插入到已有模型中。

决策 2：用关联矩阵（一个固定大小的矩阵 M）来表示记忆

关联矩阵就像一本固定页数的"知识字典"。Key 是索引，Value 是内容，矩阵 M 存储了 Key→Value 的映射关系。关键是，这个矩阵 M 的大小是固定的（d_key × d_value），不随序列长度增长。

为什么用关联矩阵能对应线性注意力？

线性注意力的本质是把 softmax(QKᵀ)V 近似为 φ(Q)(φ(K)ᵀV)——通过结合律，可以先算 φ(K)ᵀV（一个固定大小的矩阵），再用 φ(Q) 去查询它。这个 φ(K)ᵀV 的形式，正是关联矩阵的更新方式。

3.3 记忆检索（Memory Retrieval）

原文

"Memory retrieval. In Infini-attention, we retrieve new content from the memory by using the query as:

A_mem = σ(Q) M_{s-1} / (σ(Q) z_{s-1})

Here, σ and z_{s-1} are a nonlinear activation function and a normalization term, respectively. As the choice of the non-linearity and the norm method is crucial for training stability, following prior work we record a sum over all keys as the normalization term z_{s-1} and use element-wise ELU + 1 as the activation function."

翻译

记忆检索：在 Infini-attention 中，我们使用查询向量从记忆中检索内容：

A_mem = σ(Q) · M_{s-1} / (σ(Q) · z_{s-1})

其中 σ 是非线性激活函数，z_{s-1} 是归一化项。由于非线性的选择和归一化方法对训练稳定性至关重要，我们参考先前工作，以所有键的累加和作为归一化项 z_{s-1}，并使用逐元素的 ELU+1 作为激活函数。

新手讲解

用类比来理解这个公式：

把压缩记忆矩阵 M 想象成一本"浓缩词典"。检索时的操作是：

把当前的 Query 向量通过激活函数 σ（ELU+1）变换一下，得到"查询特征"。
用这个查询特征去乘以记忆矩阵 M（矩阵乘法），相当于"在词典里查找"。
除以归一化项 z，保证数值稳定（不会爆炸或消失）。

为什么用 ELU+1 而不是 softmax？

Softmax 需要看到所有 Key 才能归一化，这意味着必须把所有历史 Key 都存着——那就又回到了 KV 缓存无限增长的问题。ELU+1（确保输出为正数）配合累加归一化项 z，可以在不保存所有 Key 的情况下，用固定大小的参数做到近似检索。

M_{s-1} 中的下标 s-1 表示处理到第 s 段文字时，使用的是前 s-1 段更新过的记忆矩阵。

3.4 记忆更新（Memory Update）

原文（Linear 更新规则）

"Memory update. Once the retrieval is done, we update the memory and the normalization term with the new KV entries and obtain the next states as

M_s ← M_{s-1} + σ(K)ᵀ V
z_s ← z_{s-1} + Σ_{t=1}^{N} σ(K_t)

The new memory states M_s and z_s are then passed to the next segment S+1, building in a recurrence in each attention layer. The right side term σ(K)ᵀ V in the above equation is known as an associative binding operator."

翻译

记忆更新：检索完成后，用新的 KV 条目更新记忆和归一化项，得到下一个状态：

M_s ← M_{s-1} + σ(K)ᵀ · V
z_s ← z_{s-1} + Σ_{t=1}^{N} σ(K_t)

新的记忆状态 M_s 和 z_s 被传递给下一段 S+1，在每个注意力层中建立循环（recurrence）机制。上式中右侧项 σ(K)ᵀ · V 被称为关联绑定算子（associative binding operator）。

新手讲解

这就是"把新内容压缩进笔记本"的操作。

更新规则非常直观：

M_s = M_{s-1} + σ(K)ᵀ · V：新的记忆矩阵 = 旧的记忆矩阵 + 当前段的 KV 绑定信息。
每处理完一段，σ(K)ᵀ · V 就把这段的键值对"压缩求和"后累加进 M。

为什么叫"关联绑定"？

在关联记忆（associative memory）的理论框架里，"绑定"（binding）是指把一个 Key 和一个 Value 关联存储。σ(K)ᵀ · V 做的正是这件事：把当前段所有 token 的 Key-Value 对压缩成一个矩阵，加进全局记忆。

关键洞察：这引入了循环（Recurrence）

注意 M_s 依赖 M_{s-1}，M_{s-1} 依赖 M_{s-2}，……这形成了一个跨段的隐式循环结构。这就像 RNN 里的隐状态，但这里的"隐状态"是一个矩阵而不是向量，且是在每个注意力头里独立维护的。

原文（Linear+Delta 改进规则）

"Inspired by the success of delta rule, we have also incorporated it into our Infini-attention. The delta rule attempts a slightly improved memory update by first retrieving existing value entries and subtracting them from the new values before applying the associative bindings as new update.

M_s ← M_{s-1} + σ(K)ᵀ (V - σ(K) M_{s-1} / σ(K) z_{s-1})

This update rule (Linear+Delta) leaves the associative matrix unmodified if the KV binding already exists in the memory while still tracking the same normalization term as the former one (Linear) for numerical stability."

翻译

受 Delta 规则成功经验的启发，我们也将其融入 Infini-attention。Delta 规则在应用关联绑定之前，先从记忆中检索出已有的值条目，再用新值减去它，以此实现略有改进的记忆更新：

M_s ← M_{s-1} + σ(K)ᵀ · (V - σ(K)M_{s-1} / σ(K)z_{s-1})

这种更新规则（Linear+Delta）在 KV 绑定已存在于记忆中时，不修改关联矩阵；同时仍跟踪与前一规则（Linear）相同的归一化项，以保证数值稳定性。

新手讲解

Delta 规则是一种"只写新信息"的策略。

普通的 Linear 更新会无条件地把新 KV 加进去，即使记忆里已经有类似的信息，也会重复叠加，可能造成"过度强化旧信息"。

Delta 规则的思路：

先用当前的 Key 去查记忆，看看记忆里已经"认为"这个 Key 对应的 Value 是什么（即检索出"旧值"）。
计算"新值 - 旧值"，即记忆需要被"纠正"的量（delta）。
只把这个"差值"写进记忆。

效果：如果记忆里已经完美记住了这个 KV 对，差值为 0，记忆不变。如果是全新信息，差值就是新 Value 本身，和 Linear 规则等价。这样能更高效地利用有限的记忆容量。

实验表明 Linear+Delta 规则的效果通常优于普通 Linear 规则。

3.5 长期上下文注入——门控融合（Long-term Context Injection / Gating）

原文

"Long-term context injection. We aggregate the local attention state A_dot and memory retrieved content A_mem via a learned gating scalar β:

A = sigmoid(β) ⊙ A_mem + (1 - sigmoid(β)) ⊙ A_dot

This adds only a single scalar value as training parameter per head while allowing a learnable trade-off between the long-term and local information flows in the model."

翻译

长期上下文注入：我们通过一个可学习的门控标量 β，将局部注意力状态 A_dot 和记忆检索内容 A_mem 聚合：

A = sigmoid(β) ⊙ A_mem + (1 - sigmoid(β)) ⊙ A_dot

这在每个注意力头里只增加一个标量训练参数，同时允许模型学习长期信息流与局部信息流之间的权衡。

新手讲解

这是整个机制的"混合器"，也是最优雅的设计之一。

术语：Gating（门控）
门控机制是神经网络中常见的动态加权策略。LSTM 里的遗忘门/输入门、GRU 里的更新门，都是门控机制的例子。门控让模型自己学习"该相信哪个来源的信息多一些"。

公式的直观理解：

A_mem：从压缩记忆里取出的历史信息（"我从笔记本里查到的"）。
A_dot：当前段的标准注意力输出（"我正在阅读的这页"）。
sigmoid(β)：一个介于 0 和 1 之间的动态权重（"我该多信任笔记本？"）。

当 sigmoid(β) 接近 1：模型更多依赖压缩记忆（更关注历史）。
当 sigmoid(β) 接近 0：模型更多依赖局部注意力（更关注当前段）。

只有一个参数 β，却很强大

整个门控机制每个注意力头只增加一个标量参数 β。这是极其轻量级的设计。更重要的是，实验发现不同注意力头会自动学习出不同的 β 值：

有些头的 β ≈ 1（专职负责长期历史，"全靠笔记本"）
有些头的 β ≈ 0（专职处理局部信息，"只看眼前"）
有些头的 β ≈ 0.5（混合型，"笔记本和眼前各参考一半"）

这种自发的"专业化分工"是模型主动学习出来的，而不是人为指定的。

3.6 多头压缩记忆（Multi-head Compressive Memory）

原文

"Similar to the standard MHA, for the multi-head Infini-attention we compute H number of context states in parallel, and concatenate and project them for the final attention output:

O = [A¹; …; Aᴴ] W_O

where W_O is trainable weights."

"Similar to multi-head attention (MHA), it maintains H number of parallel compressive memory per attention layer."

翻译

与标准 MHA 类似，对于多头 Infini-attention，我们并行计算 H 个上下文状态，然后将它们拼接并投影得到最终注意力输出：

O = [A¹; …; Aᴴ] · W_O

其中 W_O 是可训练权重。每个注意力层维护 H 个并行的压缩记忆。

新手讲解

这部分说明 Infini-attention 完全兼容多头注意力（MHA）框架：

标准 MHA 有 H 个注意力头，每个头独立计算 A_dot。
Infini-attention 同样有 H 个头，每个头同时维护一个独立的压缩记忆矩阵 M。
H 个头的输出 A 拼接后，通过 W_O 投影，得到最终输出——和标准 MHA 完全一致的接口。

这就是"即插即用"的含义：从外部接口看，Infini-attention 和标准注意力一模一样，可以直接替换现有模型中的注意力层。

3.7 内存复杂度对比

原文（Table 1 关键内容）

"Infini-Transformer has a constant memory complexity of d_key × d_value + d_key for storing compressed context in M_s and z_s for each head in single layer."

Table 1 对比：Transformer-XL 的内存复杂度为 (d_key + d_value) × H × N × l，而 Infini-Transformer 为 d_key × (d_value + 1) × H × l。

翻译

Infini-Transformer 存储每层每个头的压缩上下文（M_s 和 z_s）的内存复杂度为常数级：d_key × d_value + d_key，与序列长度 N 完全无关。

方案	内存复杂度
Transformer-XL（缓存上一段）	(d_key + d_value) × H × N × l
Memorizing Transformers（存所有 KV）	(d_key + d_value) × H × N × l
Infini-Transformer（本文）	d_key × (d_value + 1) × H × l（与 N 无关）

其中 N 是段长（缓存的 token 数），l 是层数，H 是头数。

新手讲解

这张表是 Infini-attention 最硬核的优势所在：

Transformer-XL：只缓存上一个段的 KV，内存与段长 N 成正比。换段时旧的 KV 就丢掉了——等于只有一段的"短期记忆"。
Memorizing Transformers：把所有 KV 都存着，内存随序列总长度线性增长，长文章会撑爆。
Infini-Transformer：内存与 N 完全无关，无论处理多长的文本，压缩记忆矩阵的大小永远是 d_key × d_value。

这 114 倍的压缩比怎么来的？

实验中，Infini-Transformer 使用的序列总长是 32768 token（段长 2048），与 Memorizing Transformers 相比，内存节省了 114 倍（32768/2048 × 某些参数差异 ≈ 114）。

代价是什么？

压缩必然有损。把 32768 个 token 的 KV 信息塞进一个固定大小的矩阵，细节必然丢失。但实验表明，在困惑度（perplexity）和检索准确率上，损失是可以接受的，甚至在某些任务上优于只看有限窗口的方案。

3.8 Infini-attention 完整流程图解（作者总结）

用一段伪代码表示整个前向传播流程，方便整体理解：

对每个输入段 X_s（当前段文本）：

  1. 用 X_s 计算 Q, K, V（与标准注意力相同）

  2. [局部注意力]
     A_dot = softmax(QKᵀ / √d) · V
     （只看当前段内部的注意力，带因果掩码）

  3. [从压缩记忆检索历史]
     A_mem = σ(Q) · M_{s-1} / (σ(Q) · z_{s-1})
     （用 Q 查询上一段存下的压缩记忆矩阵 M）

  4. [门控融合]
     A = sigmoid(β) ⊙ A_mem + (1 - sigmoid(β)) ⊙ A_dot

  5. [更新压缩记忆，供下一段使用]
     M_s ← M_{s-1} + σ(K)ᵀ · (V - A_mem_from_K)  [Linear+Delta]
     z_s ← z_{s-1} + Σ σ(K_t)

  6. 输出 A（送入后续网络层）
     同时把 M_s, z_s 传递给下一段

4. 实验（关键结果精读）

4.1 长上下文语言建模

原文

"Infini-Transformer (Linear + Delta) achieved 9.67 perplexity on PG19 and 2.23 perplexity on Arxiv-math with 114x better compression than Memorizing Transformers."

"The model achieves even better perplexity when trained with 100K sequence length."

翻译

Infini-Transformer（Linear+Delta 规则）在 PG19 数据集上取得了 9.67 的困惑度（perplexity），在 Arxiv-math 上取得了 2.23，在内存节省 114 倍的情况下优于或持平于 Memorizing Transformers。

在使用 100K 序列长度进行训练时，模型可以取得更好的困惑度。

新手讲解

术语：Perplexity（困惑度）
语言模型的评估指标，越低越好。它衡量模型对文本的"惊讶程度"——如果模型总是能准确预测下一个词，困惑度就低。

实验设置：使用 12 层 Transformer，在 PG19（英文小说数据集）和 Arxiv-math（数学论文数据集）上进行语言建模。训练序列长度 32768 token，段长 2048。

关键结论：Infini-Transformer 在内存只有 Memorizing Transformers 1/114 的情况下，语言建模性能相当甚至更好。这说明压缩记忆虽然有损，但损失是可控的，且节省的内存空间可以用来处理更长的训练序列（100K），进一步提升性能。

4.2 密钥检索任务（Passkey Retrieval）

原文

"Infini-Transformers solved the passkey task with up to 1M context length when fine-tuned on 5K length inputs. Results show 100% accuracy at 1M length after 400 fine-tuning steps, tested at positions: start/middle/end."

翻译

Infini-Transformer 在仅用 5K 长度样本微调的情况下，成功解决了百万（1M）token 上下文长度的密钥检索任务，在 400 步微调后，对开头/中间/结尾各位置均达到接近 100% 的准确率。

新手讲解

术语：Passkey Retrieval（密钥检索 / 大海捞针任务）
一个测试长上下文能力的极端基准：在一段很长的随机文字中，随机某处藏着一个 5 位数字（密钥）；模型需要找到并复述这个数字。这测试的是模型是否真的"读完"了全文，而不是只看了开头结尾。

为什么这个结果令人震惊？

模型在 5K 长度（5000 个 token）的样本上微调，然后直接被测试在 1M token（100 万个 token）的文本中找密钥——相当于在 200 倍长的文本上泛化。这在标准注意力或 Transformer-XL 上是不可能实现的。

为什么 Infini-attention 能做到？

因为压缩记忆的更新是流式的、无限延续的：不管文本有多长，每处理一段就更新一次记忆矩阵。只要密钥在某一段被"压缩写入"了记忆矩阵，后续的 Query 就能（大概率）把它检索出来。5K 到 1M 的外推，靠的是这种流式记忆的结构性泛化。

4.3 长文档摘要（BookSum）

原文

"A 8B model with Infini-attention reaches a new SOTA result on a 500K length book summarization task after continual pre-training and task fine-tuning."

"Infini-Transformers (Linear + Delta) achieved: Rouge-1: 40.0, Rouge-2: 8.8, Rouge-L: 17.9, Overall: 18.5"

翻译

一个 8B 参数的 Infini-attention 模型，经过持续预训练和任务微调后，在 500K token 长度的书籍摘要任务（BookSum）上取得了新的最优（SOTA）结果：ROUGE-1 为 40.0，ROUGE-2 为 8.8，ROUGE-L 为 17.9，综合分 18.5。

新手讲解

BookSum 是一个需要对整本书（通常超过 10 万词）生成摘要的数据集，输入极长。普通 LLM 受限于上下文窗口，只能"截断"输入，无法处理完整的书。

Infini-attention 的 8B 模型可以流式处理完整的 500K token 输入（约等于一本中等长度的英文小说），生成的摘要质量达到了当时的最优水平。

4.4 门控可视化（Gating Visualization）

原文

"Figure 3 visualization reveals: specialized heads with gating score near 0 or 1 and mixer heads with score close to 0.5."

翻译

论文图 3 的可视化显示：部分注意力头的门控分数接近 0 或 1（专职头），另一些接近 0.5（混合头）。

新手讲解

这是一个非常有意思的"涌现"现象：模型自发地将不同注意力头分配了不同职能——

门控 ≈ 1 的头：几乎完全依赖压缩记忆，专门负责"查历史"。
门控 ≈ 0 的头：几乎完全依赖局部注意力，专门负责"看当前段"。
门控 ≈ 0.5 的头：混合使用，同时关注历史和当前。

这种职能分化不是人为设计的，而是通过端到端训练自然涌现出来的。这说明单个 β 参数虽然轻量，但足够让模型学会合理的信息分配策略。

5. 与 StreamingLLM、Mamba 的思路对比

这是很多读者关心的问题：Infini-attention 和其他处理长上下文的方案有什么本质区别？

对比表格

方案	如何处理历史信息	内存	是否有损	适合场景
标准注意力	完整保存所有 KV	O(N)，随长度增长	无损	短上下文
Transformer-XL	只保留上一段 KV	O(段长)，固定	有损（丢弃更早历史）	中等长度
StreamingLLM	保留最近 KV + Attention Sink	O(窗口大小)，固定	有损（只记得最近和开头）	流式生成
Mamba（状态空间模型）	用固定大小隐状态做选择性压缩	O(状态维度)，固定	有损	各种长度
Infini-attention（本文）	把历史 KV 累积压缩进固定矩阵 M	O(d_key × d_value)，固定	有损	超长上下文

核心思路差异

StreamingLLM（丢弃中间）

StreamingLLM 的策略是：只保留"Attention Sink"（最开始的几个 token，因为它们接收最多注意力）和最近的一个滑动窗口内的 KV。中间的历史信息直接丢掉，不做任何压缩保存。

类比：读书时只记第一章和最近几页，中间全扔掉。

这对于流式生成场景（如聊天机器人）很有效，因为对话通常只需要关注最近上下文。但对于需要引用远端信息的任务（如"回忆第三章提到的密钥"），StreamingLLM 会彻底失败。

Mamba（状态空间模型）

Mamba 是一种完全不同的架构，不基于注意力机制。它使用选择性状态空间模型（Selective SSM），通过一个固定大小的"状态向量"来传递序列信息。每一步根据当前输入，有选择地更新状态（更新多少、遗忘多少）。

类比：脑中有一块固定大小的"黑板"，每读一个词就用海绵擦掉一点、写上新内容，擦多写多还是擦少写少由当前词决定。

Mamba 的状态是向量（一维），而 Infini-attention 的记忆是矩阵（二维），理论上能存储更多的关联信息。此外，Mamba 是从零设计的新架构，而 Infini-attention 是对现有 Transformer 的增量改造——后者在利用已有预训练权重方面有天然优势。

Infini-attention（压缩累积）

Infini-attention 的独特之处在于：它不丢弃任何历史信息，而是把所有历史 KV 通过线性注意力的方式累积压缩进固定大小的矩阵 M。新信息不断以"叠加"的方式写入，旧信息以"浓缩"的形式保留。

类比：每读一段都把要点浓缩写进随身笔记本，笔记本页数固定，越到后面写的越密——但所有读过的内容都有痕迹，不是完全丢失。

代价是：压缩必然有信息损失，早期信息可能被后期信息"覆盖"或"稀释"。但实验表明，这种有损压缩在实际任务中效果很好，尤其是能泛化到训练时从未见过的极长序列长度。

6. 相关工作与附录（一句话带过）

相关工作：论文对比了 Transformer-XL、Memorizing Transformers、Longformer、BigBird、线性 Transformer 等方案，Infini-attention 的核心优势在于用固定内存实现真正无限上下文，且改动极小可即插即用。

附录：包含更多消融实验（验证 Linear vs Linear+Delta、不同段长的影响）、门控可视化图、以及 BookSum 任务的详细超参数配置，感兴趣可参阅原文附录。

7. 总结

核心贡献（一句话版）

Infini-attention 在每个标准注意力头里塞了一个固定大小的压缩记忆矩阵，用线性注意力机制做读写操作，用一个可学习的门控标量融合局部注意力和历史记忆，使 Transformer 可以以有界内存处理无限长序列。

方法优缺点

优点	局限
内存复杂度与序列长度无关	压缩有损，精确信息可能丢失
对现有 Transformer 改动极小	压缩记忆矩阵大小需要提前选定
支持即插即用持续预训练	理论上早期信息可能被后期信息稀释
流式推理，支持无限长输入	线性注意力与 softmax 注意力存在近似误差
实验在 1M token 上验证有效	目前还没有在超大规模（如 GPT-4 量级）上验证

给新手的最终类比总结

Infini-attention 就像给 Transformer 配备了一本"固定页数的随身笔记本"：

读每一段时，把要点浓缩写进笔记本（压缩记忆更新）。
读下一段时，先翻笔记本找相关历史信息（记忆检索），再结合当前段的内容（局部注意力）。
用一个可学习的"判断力"决定该多信任笔记本还是多信任眼前的内容（门控融合）。
笔记本只有固定页数，但可以一直读下去，永远不需要更大的笔记本。

关键代价：笔记本越写越密，早期内容的细节会逐渐"模糊"（有损压缩）。但模型依然能在 100 万 token 的文本中找到一个隐藏的密钥——说明即便有损，关键信息的痕迹依然留存。

精读完成。本笔记涵盖：Abstract 全文、Introduction 核心段落、Infini-attention 方法章节全部段落（含 3.1-3.8 共 8 个子节）、三项关键实验、与 StreamingLLM/Mamba 的对比分析。