技术摘要：通过关联递归记忆扩展 LLM 上下文

问题陈述

大语言模型（LLMs）正日益被要求处理跨越数十万甚至数百万个 token 的输入，用于技术报告分析、软件开发和多文档推理等任务。然而，标准的 Transformer 架构面临着一个根本性的瓶颈：自注意力机制的计算和内存成本随序列长度呈二次方增长（$O(N^2)$）。此外，即使在模型的标称窗口内，随着上下文长度的增加，性能也往往会下降。虽然递归架构（如 Mamba、RWKV）提供了线性缩放能力，但它们通常需要从头开始训练，这阻碍了利用现有预训练 LLM 的可能性，并且与 Transformer 相比，它们在处理复杂算法任务和指令遵循方面往往表现较弱。

方法论

作者提出了关联递归记忆 Transformer（ARMT），作为一种在保持恒定内存缩放和效率的同时扩展上下文长度的实用解决方案。ARMT 作为预训练基础 LLM 的一个封装器（wrapper），能够实现对输入的分段处理。

核心架构

ARMT 将长上下文输入划分为固定长度的非重叠分段。在每个分段内，模型利用全自注意力机制（短期/工作记忆）。至关重要的是，它引入了一个层级关联记忆模块，用于在分段之间传播信息（长期记忆）。该机制分为三个阶段运行：

1. 记忆提取： 每个 Transformer 层将输入分段压缩为记忆嵌入（memory embeddings）。
2. 记忆整合： 这些嵌入被整合为作为键值对（key-value pairs）的每层关联矩阵。
3. 关联： 后续分段中的嵌入被转化为查询向量（query vectors），并通过与关联矩阵相乘，从之前的分段中检索相关信息。

训练方案

论文概述了一套全面的训练策略，用于将预训练的 LLM 适配为 ARMT：

• 持续预训练： 未初始化的关联记忆参数通过在长上下文（例如使用 FineWeb-Edu 的 19B 个 token）上的无监督语言建模进行初始化，以便在进行特定任务微调之前学习有效的记忆传播。
• 课程学习： 为了解决从头学习长程依赖关系的难度，模型通过逐步增加分段数量（例如从 2 到 4 再到 8）并退火学习率来进行微调。
• 合成数据生成： 为了克服长上下文场景下的数据稀缺问题，作者通过将长文档中的短段落拼接在一起，并为每个段落生成问答对，从而生成合成训练实例，创建了具有特定上下文长度的大型数据桶。
• 层剪枝与选择： 研究调查了是否需要在每一层都使用关联记忆。它提出了一种策略，仅在部分层（例如特定的中间层和最终层）保留关联记忆，从而在不显著损失性能的情况下减少可训练参数和计算成本。

核心贡献

1. 领域特定数据集： 构建了两个新数据集：用于代码中变量类型预测的 ManyTypes-long (MT) 和用于长文档问答的 GovReport-long (GR)，旨在评估现实的、窄领域的负载。
2. 训练方案： 一个利用 ARMT 扩展 LLM 上下文的新颖框架，结合了持续预训练、合成数据生成、课程学习和选择性层集成。
3. 实证验证： 通过对 Gemma-3-1B-IT 和 SmolLM-2-360M-IT 主干网络的广泛实验研究表明，ARMT 增强的模型：
   • 在处理超出原始上下文限制（高达 64k tokens）的输入时，相对于其限内基准模型，性能不会出现退化。
   • 与基础模型相比，在分布外（OOD）上下文长度方面表现出更优越的泛化能力。
   • 在保持原始上下文窗口内性能的同时，实现了约 30% 的 FLOPs 减少。

实验结果

作者使用 Gemma-3-1B-IT 和 SmolLM-2-360M-IT 作为主干网络，在 MT 和 GR 数据集上对 ARMT 进行了评估。

• 性能： ARMT 模型在高达 65k tokens 的上下文长度下保持了稳定的性能。相比之下，基础模型（即使经过微调）在其原生上下文窗口（例如 8k 或 32k）之外会出现剧烈的性能下降。ARMT 在长 OOD 机制（32k–65k）下显著优于基础模型。
• 效率： 无论上下文长度如何变化，ARMT 均表现出恒定的 GPU 显存占用，而基础模型的显存使用量随之线性增长。对于 32k token 的序列，在相同的显存预算下，ARMT 允许将 Batch Size 提高 4 倍（32 对比 8）。
• FLOPs 减少： 通过理论分析和经验推理时间测量确认，全局注意力 FLOPs 减少了 $T/S$ 倍（序列长度除以分段大小），导致与全注意力模型相比，总 FLOPs 降低了约 30%。
• 消融实验：
  • 层剪枝： 仅在约 20% 的层（特别是预选的中间层和最终层）中使用关联记忆的模型，达到了与完整 ARMT 模型相当或更好的性能。
  • 预训练： 发现持续预训练对于初始化关联记忆至关重要，能显著提高其在领域内及 OOD 的性能。
  • 基准测试： ARMT 在长上下文泛化方面优于其他长上下文基准模型（包括 Mamba-2、DeltaNet 和 xLSTM），且相比于从头训练的模型，所需的预训练量更少。

重要性与主张

论文将 ARMT 定位为一种实用的、计算高效的方法，用于实现中小规模 LLM（最高达 1B 参数）的长上下文处理。作者声称，这种方法弥合了 Transformer 强大的短上下文性能与递归模型的线性缩放能力之间的差距。

关于其重要性的关键主张包括：

• 隐私与本地部署： 通过使小型模型能够高效处理长上下文，ARMT 为不依赖远程大规模 API 驱动 LLM 的隐私保护应用提供了便利。
• 可扩展性： 该方法允许以恒定的内存处理任意长的上下文，这是现实世界文档分析的关键要求。
• 效率： 恒定内存缩放与减少 FLOPs 的结合，使 ARMT 成为资源受限环境下的可行替代方案。
• 泛化： 所提出的训练方案有效地解决了标准 Transformer 在处理长序列时出现的“迷失在中部（lost in the middle）”现象以及性能退化问题。

作者承认了局限性，指出实验仅限于 1B 参数以下的模型以及特定的任务集（代码和文档问答），并且关联记忆与 Transformer 表示之间的底层交互机制在一定程度上仍不完全明确。然而，他们认为，在小型模型中展示出的扩展上下文的能力，是迈向实用的、本地化长上下文 AI 的重要一步。