Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention | Tsendsuren Munkhdalai, Manaal Faruqui, Siddharth Gopal | 2024-04-10

不遗漏任何上下文：用 Infini-attention 实现高效无限上下文 Transformer

论文提出 Infini-attention，把压缩记忆与局部掩码注意力、长期线性注意力融合进单个 Transformer 块，使 LLM 能以有界内存与计算处理无限长输入。

背景

Transformer 及基于 Transformer 的大语言模型受注意力机制本身限制，上下文相关的记忆容量有限。注意力机制的内存占用与计算时间均为二次复杂度；论文引用的例子指出，500B 模型在批大小 512、上下文长度 2048 时，注意力 KV 状态的内存占用可达 3TB。把 LLM 扩展到更长序列（如 100 万 token）在标准 Transformer 架构下十分困难，服务长上下文模型的成本也很高。压缩记忆系统被认为比注意力机制更具可扩展性与效率：它不维护随输入长度增长的数组，而是用固定数量的参数存储与召回信息，新信息通过修改参数写入，之后可被恢复。但论文认为，当前 LLM 尚未出现一种兼顾简单性与质量的实用压缩记忆技术，本文即针对这一空白提出新方法。

创新点

  1. 提出 Infini-attention：一种带有长期压缩记忆与局部因果注意力的实用注意力机制，可高效建模长程与短程上下文依赖。
  2. Infini-attention 对标准缩放点积注意力的改动极小，从设计上支持即插即用的持续预训练与长上下文适配。
  3. 该方法以流式方式处理极长输入，使 Transformer LLM 在有界内存与计算资源下扩展到无限长上下文。

方法

Infini-Transformer 与 Transformer-XL 类似，按片段序列处理输入：在每个片段内计算标准因果点积注意力，因此该点积注意力是局部的，只覆盖当前片段 S 的 N 个 token。区别在于，Infini-Transformer 不丢弃旧片段的 KV 注意力状态，而是复用它们，通过压缩记忆维护整个上下文历史，使每个注意力层同时拥有全局压缩状态与局部细粒度状态。Infini-attention 是一种循环注意力机制：与多头注意力一样，每层维护 H 个并行压缩记忆（H 为注意力头数），并像 RNN 与 MNM 那样维护循环记忆状态，形式为 O_s, M_s = infini-attention(X_s, M_{s-1})。压缩记忆不另算新条目，而是复用点积注意力的 Q、K、V 状态；记忆用关联矩阵参数化，从而把记忆更新与检索转化为线性注意力，并采用 Katharopoulos 等人的更新与检索机制。检索时用查询 Q 从 M_{s-1} 读取内容：A_mem = σ(Q)M_{s-1} / σ(Q)z_{s-1}，其中 σ 为非线性激活、z 为归一化项；为训练稳定，z 记录所有键之和，激活用逐元素 ELU+1。更新规则为 M_s ← M_{s-1} + σ(K)^T V，z_s ← z_{s-1} + Σσ(K_t)，新状态传给下一片段，从而在每层建立循环。论文还引入 delta 规则：先检索已有值并从新值中减去再加关联绑定，即 M_s ← M_{s-1} + σ(K)^T(V − σ(K)M_{s-1}/(σ(K)z_{s-1}))；若 KV 绑定已存在则不改动关联矩阵，同时沿用同一归一化项保证数值稳定。最终输出通过可学习门控标量 β 聚合：A = sigmoid(β)⊙A_mem + (1−sigmoid(β))⊙A_dot，每头仅增加一个标量参数，实现长期与局部信息流的可学习权衡；多头结果拼接后经 W_O 投影得到输出。实现上，输入整段前向传播，在每个 Infini-attention 层内做片段切分并逐段处理再拼接，对现有 Transformer 实现改动最小；各层用随时间反向传播（BPTT）训练压缩记忆状态，并做梯度检查点以省内存；压缩记忆的键与查询不使用位置嵌入，仅存全局上下文信息，位置嵌入只在记忆读写之后施加于 QK 向量。表 1 的对比显示，Infini-Transformer 每层每头的记忆复杂度恒为 d_key×d_value+d_key，而 Transformer-XL、Compressive Transformer、Memorizing Transformers 依赖缓存大小，RMT 与 AutoCompressors 依赖软提示大小，均随序列增长；因此 Infini-Transformer 在有界内存下获得无界上下文窗口。

实验

论文在极长输入序列的基准上评估 Infini-Transformer：长上下文语言建模、100 万长度 passkey 上下文块检索、50 万长度书籍摘要。语言建模从头训练，passkey 与书籍摘要则持续预训练现有 LLM，以突出即插即用的长上下文适配能力。实现细节：片段切分在每个 Infini-attention 层内完成；用 BPTT 训练并做梯度检查点；压缩记忆的 K/Q 不用位置嵌入。长上下文语言建模在 PG19 与 Arxiv-math 上进行，设置贴近 Memorizing Transformers：12 层、8 个注意力头、每头维度 128、FFN 隐层 4096；Infini-attention 片段长度 N=2048，训练输入长度 32768，使记忆状态展开 16 步。表 2 中，Infini-Transformer（Linear）记忆 1.6M（114x 压缩），PG19 困惑度 9.65、Arxiv-math 2.24；Linear+Delta 为 9.67 与 2.23；Transformer-XL 为 11.88 与 2.42，Memorizing Transformers（183M，1x）为 11.37 与 2.26，RMT 为 13.27 与 2.55。即模型以比 65K 长度向量检索 KV 记忆的 Memorizing Transformer 少 114 倍的记忆参数取得更优结果。把训练序列长度从 32K 增至 100K 后，Linear 与 Linear+Delta 在 Arxiv-math 上的困惑度进一步降至 2.21 与 2.20。门控分数可视化显示训练后出现两类头：门控接近 0 或 1 的专门化头（或走局部注意力、或从压缩记忆检索）与接近 0.5 的混合头；每层至少有一个短程头，保证输入信号可前向传播到输出层，且前向计算中长短期内容检索交替出现。1M passkey 检索：将 1B LLM 的普通 MHA 替换为 Infini-attention，在 4K 长度输入上继续预训练 30K 步、批大小 64，再在 passkey 任务上微调；此前工作显示 8B LLaMA 用 32K 输入与位置插值微调可解到 32K，本文仅在 5K 长度输入上微调以测试 1M 长度。表 3 报告 32K 至 1M 输入、passkey 位于开头/中间/结尾的 token 级准确率：零样本下 Linear 在 1M 为 8/6/98，Linear+Delta 为 7/6/97；微调 400 步后 Linear+Delta 在 32K、128K、256K、512K、1M 上分别达到 100/100/100、100/100/99、100/100/99、100/100/100、100/100/100。50 万长度书籍摘要（BookSum）：用 8K 输入长度对 8B LLM 持续预训练 30K 步，再在 BookSum 上微调，微调输入 32K、评估增至 500K，生成温度 0.5、top_p=0.95、解码 1024 步。表 4 中 Infini-Transformer（Linear+Delta）Rouge-1/2/L 为 40.0/8.8/17.9、Overall 18.5，优于 BART（36.4/7.6/15.3/16.2）、BART+Unlimiformer（16.9）、PRIMERA（16.3）、PRIMERA+Unlimiformer（17.2），在 BookSum 上取得新 SOTA；验证集曲线显示输入书籍文本越多，摘要指标越好。

局限

  • 压缩记忆采用关联矩阵与线性注意力近似，其检索是粗粒度的；论文未报告该方法在需要精确长程细节回忆的任务上的表现，因此长上下文质量的边界仍待检验。（解读者观点）
  • passkey 与书籍摘要的结果依赖持续预训练与任务微调（如 30K 步持续预训练、400 步微调），零样本 1M 检索准确率明显低于微调后，说明无限上下文能力并非开箱即用。（解读者观点）
  • 论文指出压缩记忆系统的参数通过更新规则修改来存储信息、再经读取机制召回；相较于随长度增长的 KV 数组，其容量固定，信息写入与覆盖方式可能影响久远内容的保真度。（解读者观点）

AI 依据论文正文生成，经人工审核；数字与结论以原文为准。