多对话优化的挑战。然而，由于MemAgent方法的特性，它会为单个查询生成多个上下文独立的对话，如图2所示。因此，策略优化不能像在多轮工具调用优化中那样简单地应用注意力掩码来实现。

将每次对话视为独立优化目标。为了解决这个问题，我们将每次对话视为一个独立的优化目标，如图3所示。假设对于一组中的一个样本$(q_i, a_i)$，生成了$n_i$次对话$(o_{i,1}, o_{i,2}, ..., o_{i,n_i})$。$o_{i,j}$进一步分解为token级别的输出$(o_{i,j,1}, o_{i,j,2}, ..., o_{i,j,|o_{i,j}|})$。我们通过包含最终答案的最后一次对话为每个样本计算一个结果奖励$R_i$，并将组归一化的优势分配给所有相关的对话。

Multi-Conv DAPO的优势和损失计算。方程4和5展示了在我们的MemAgent算法中，如何为上下文独立的多对话rollouts计算优势和损失。优势值从包含最终答案的对话中得出，然后均匀地应用于源自同一 لm 样本的所有对话。我们的损失函数类似于DAPO【索引35，DAPO: An open-source llm reinforcement learning system at scale，2025，arXiv】中使用的函数，该函数包含了一个token级别的平均损失。此外，我们将损失计算的维度从传统的（组，token）结构扩展到（组，对话，token）。遵循DrGRPO【索引58，Understanding r1-zero-like training: A critical perspective，2025，arXiv】的做法，我们不通过奖励的标准差来归一化优势。

3.3 奖励建模

基于规则验证器的最终结果奖励。遵循RLVR配方【索引33, Search-r1: Training llms to reason and leverage search engines with reinforcement learning, 2025, arXiv; 索引35, Dapo: An open-source llm reinforcement learning system at scale, 2025, arXiv; 索引51, Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, 2025, arXiv】，我们使用由基于规则的验证器计算出的最终结果奖励来训练模型。

多标准答案任务的奖励函数。在RULER【索引1，Ruler: What’s the real context size of your long-context language models?，2024，arXiv】和其他数据集中，问题可能有多个真实答案。对于某些任务，例如问答，这些真实答案被认为是等价的。给定一组多个真实答案$Y = \{y_1, y_2, . . . , y_n\}$，奖励分数定义为：

其中$\hat{y}$是预测的答案，$I(·)$表示指示函数。

需要包含所有答案任务的奖励函数。对于其他任务，最终输出中需要包含所有真实答案。一个例子是“大海捞针”中的多值任务，问题可能是：“XXX的所有特殊魔法数字是什么？”在这种情况下，奖励函数定义为：

其中$| · |$表示集合的基数。

3.4 从自回归建模视角重新思考MemAgent

标准自回归模型的局限性。为了更深入地理解MemAgent的设计，我们建议从以下方式重新思考语言模型的因式分解。一个标准的自回归LLM将序列$x_{1:N}$的联合似然分解为$p(x_{1:N}) = \prod^N_{n=1} p(x_n | x_{1:n-1})$，这隐含地假设每个过去的token（或至少其隐藏状态）都必须保留在活动上下文中。这正是使二次注意力成为长上下文瓶颈的原因。

MemAgent的建模方式。MemAgent用一个固定长度的记忆$m \in V_M$取代了无限的历史记录，如图4所示。输入文本以$K$个连续的块$c_1, . . . , c_K$（每个块的长度$≤ C$）流式传输给模型。在读取块$k$之后，模型用一个新的向量$m_k$覆盖面板，该向量总结了迄今为止所见的所有证据。因为$|m_k| = M$是恒定的，所以每一步的计算和内存都是$O(C + M)$，从而得到一个整体的线性复杂度$O(N)$。

分解后的似然函数。引入潜在序列$m_{1:K-1}$将原始似然分解为：

其中基准情况为$m_0 = \emptyset$。在每个块内部，我们仍然运行一个普通的Transformer解码器，但条件是一个恒定的上下文窗口$(c_k, m_{k-1})$。读取路径逐个token地因式分解，$p(c_k | m_{k-1}) = \prod^{kC}_{i=(k-1)C+1}p(x_i | x_{1:i-1}, m_{k-1})$，而写入路径以相同的自回归方式生成下一个记忆。

MemAgent的优势与概念。MemAgent实现了上下文的token级压缩，而局部-全局或线性注意力模型则在特征空间中压缩长上下文；它们的摘要是隐式的和不透明的。MemAgent的摘要存在于token空间中，因此每个中间记忆都是人类可读的，可以被检查甚至编辑——我们在设计RL奖励时利用了这一特性（§3.3）。从概念上讲，方程8将Transformer变成了一个状态大小可由用户控制的循环网络。

强化学习的必要性。因为记忆token是潜在的，并且通过离散的覆盖规则进行更新，仅靠反向传播无法教会模型保留什么和丢弃什么。我们的多对话GRPO算法（§A）将每个读-写-读循环视为一个RL转换，直接奖励那些能导向正确最终答案的记忆。这弥合了显式监督（答案）和隐式结构（好的记忆）之间的差距，完成了前面介绍的训练流程。

总结。最终的MemAgent架构保留了原生解码器的训练配方，不需要奇特的注意力核心，并满足了长上下文的三个难题：任意长度、无损外推和线性成本。

A3 实验环境

• 数据集:

  • 主要使用数据集: RULER，一个包含多种可控上下文长度的合成任务基准。具体使用了其问答（QA）子集，该子集通过将包含正确答案的黄金段落嵌入到从同一数据集中采样的大量干扰内容中，来适配长上下文评估。
  • 训练数据合成: 使用 HotpotQA 数据集的方法合成训练样本，共包含200篇文章，token长度约为28K。
  • 数据清洗与筛选: 从HotpotQA训练集中处理了80,000个样本，过滤掉了在没有上下文的情况下，Qwen2.5-7B-Base或Qwen2.5-7B-Instruct模型Best-Of-2得分为100%的问题（这些问题可能属于模型的内部知识）。约50%的数据被过滤，最终选取了前32,768个样本用于训练。
  • 测试数据合成: 使用HotpotQA验证集合成了128个样本。为研究模型性能随长度的变化，合成了不同上下文长度的测试集，文章数量从50到6400不等，对应上下文长度约为7K到3.5M tokens。

• 模型架构:

  • 基础模型: Qwen2.5-7B-Instruct 和 Qwen2.5-14B-Instruct。
  • 基线模型: DeepSeek-R1-Distill-Qwen 系列、Qwen-2.5-Instruct-1M 和 QwenLong-L1。
  • 基线模型配置: Qwen2.5-Instruct-1M 使用DCA外推至1M上下文；DeepSeek-R1-Distill-Qwen 和 QwenLong 上下文长度设为128K。

• 硬件配置:

  • 论文未明确提及具体的硬件配置（如GPU型号/数量、CPU等）。

• 软件与训练配置:

  • 代码实现: 基于 verl 框架实现多对话独立上下文。
  • 训练细节: 训练时，模型上下文窗口限制在 8K。窗口分配为：1024 tokens用于查询，5000 tokens用于上下文块，1024 tokens用于记忆，1024 tokens用于输出，剩余部分为聊天模板。这使得模型处理整个上下文通常需要5到7个对话轮次。
  • 超参数: 使用 GRPO 算法，KL因子为1e-3，禁用熵损失。优化器为 AdamW，学习率为1e-6，采用带线性预热的恒定学习率调度。7B和14B模型的rollout批大小分别为128和256，组大小为16。

A4 实验结果

主要结果

实验结果如表2所示，比较了不同模型在7K到896K上下文长度下的性能，并对MemAgent模型额外评估了其在1.75M和3.5M超长上下文上的外推能力。

• 实验内容：在合成的HotpotQA数据集上，评估MemAgent与多个基线模型（包括推理模型、长上下文后训练模型）在不同上下文长度下的问答准确率。

• 实验结果：

  • MemAgent表现出卓越的长度外推能力，随着输入上下文长度的增加，性能衰减非常小。RL-MEMAGENT-14B在7K长度下准确率为83.59%，在3.5M长度下仍能达到78.12%。
  • 基线模型即使在其声称的上下文窗口内也表现出明显的失效模式：
    • 推理模型（DS-Distill-Qwen系列）性能迅速下降。
    • QwenLong-L1在其训练长度60K内表现尚可，但之后性能大幅下降。
    • Qwen2.5-Instruct-1M系列在112K tokens内性能尚可，但在896K时性能降至零，远未达到其理论的1M token容量。

• 分析结论：MemAgent通过记忆机制与强化学习的结合，能有效处理超长上下文场景。而基线模型尽管拥有扩展的上下文窗口，但在超长上下文中有效利用信息的能力有限。

Table 2: 主要实验结果，比较了不同模型在各种上下文长度下的性能。所有数值均为准确率（%）。

消融研究

1. 强化学习（RL）训练的影响
- 实验内容：在RULER-HotpotQA任务上，比较了三种模型：1）原始的Qwen2.5-Instruct模型；2）配备了记忆机制但未经RL训练的MemAgent模型；3）经过RL训练的MemAgent模型。
- 实验结果（如图5所示）：
- 原始模型随着上下文长度增加性能严重下降，超过112K后因上下文窗口限制输入被截断。
- 无RL训练的MemAgent模型性能优于原始模型，且能在超过上下文长度的任务上保持一定性能，但整体性能随输入长度增加而下降。
- 经过RL训练的MemAgent模型在所有上下文长度上均保持了持续的高性能，衰减极小。

• 分析结论：记忆机制为处理长上下文提供了结构支持，但强化学习是教会模型如何正确利用记忆的关键。

  

2. 分布外（Out-of-Distribution, OOD）任务的泛化能力
- 实验内容：在RULER基准的OOD任务上进行评估，包括“大海捞针”变体、变量跟踪、高频词提取，以及由SQuAD合成的问答任务。上下文长度从8K到512K tokens。
- 实验结果（如图6所示）：
- MemAgent在各种任务类型上都保持了持续的优越性能。MemAgent-14B在8K到512K的RULER平均任务上准确率超过95%。
- MemAgent-7B/14B在基于SQuAD的QA任务上性能稳定，表明记忆能力可以泛化到训练数据之外。
- 基线模型在所有任务类别上，超过128K tokens后性能均显著下降。

• 分析结论：MemAgent的记忆机制能够有效泛化到各种长上下文场景，而不是过拟合于特定格式。

  

  

案例研究

• 实验内容：分析MemAgent-14B在一个2跳（2-hop）问答任务上的生成轨迹。问题是：“浪漫喜剧‘Big Stone Gap’的导演常驻在纽约哪个城市？”

• 轨迹分析：

  1. 第一轮：模型读到一个可能相关的条目（一个同样位于纽约的制作团队Ghost），并选择保留此信息。
  2. 第二轮：没有提供相关上下文，但模型保持其记忆状态，表现出抗干扰能力。
  3. 第三轮：提供了两个相关的条目（电影信息和导演Adriana Trigiani的介绍），模型正确识别关键信息并相应更新记忆，得到正确答案“Greenwich Village, New York City”。
  4. 后续轮次：推理过程完成，模型记忆保持不变，并最终生成答案。

• 分析结论：该案例展示了记忆智能体学习到的核心能力：

  • 根据查询关键词预估潜在相关内容并抢先存储。
  • 在遇到与查询匹配的上下文时立即更新记忆。
  • 不受无关信息的影响。
    这些行为并非来自架构上的注意力机制，而是通过强化学习增强的文本生成能力。

A5 结论

本文提出了一种通过引入潜在变量记忆来建模长上下文任务的新方法。该方法将连续的自回归生成过程分解为一系列从记忆中顺序生成上下文的步骤。基于现有的密集注意力Transformer，我们的方法可以在不改变生成范式或引入额外模型架构的情况下，以O(N)的计算复杂度处理无限长的输入文本。

我们引入了MemAgent来实现这种建模方法，为LLM配备了一个经过强化学习训练的记忆，使模型能够学会记录相关信息和忽略无关细节的能力。实验表明，当在32K长度的数据上使用8K上下文（包括一个1024-token的记忆和每步处理5000 tokens的输入）进行训练时，模型在测试中可以外推到3.5M的文本长度，且性能几乎没有损失。

消融研究证明了将记忆本身作为长上下文处理机制的有效性，以及在此基础上进行强化学习训练带来的好处。在领域内和领域外任务上的结果表明，MemAgent超越了长上下文后训练模型、推理模型和其他基线模型，在长上下文任务上达到了最先进的性能。

A6 附录

A 计算复杂度

采用的浮点运算（FLOP）估算器。我们采用了verl【索引61，Hybridflow: A flexible and efficient rlhf framework，2024，arXiv】中针对Qwen2Model的浮点运算（FLOP）估算器来计算基线模型和我们所提方法的FLOP成本。结果如图7所示。基线模型表现出$O(n^2)$的复杂度，而MemAgent实现了$O(n)$的复杂度。

基线模型的FLOP计算。对于基线模型，需要处理的token数量为$q + c + o$，其中$q$代表问题的长度，$c$是上下文长度，$o$代表输出长度。

MemAgent的FLOP计算。对于MemAgent，总FLOP成本是所有阶段FLOP的总和。涉及的详细阶段如下：

• 初始化：在第一阶段，模型处理一个由$q + 200 + o$组成的输入，其中200代表一个常数，用于提示模型遵循MemAgent工作流。
• 记忆更新：重复次数由$k = \lceil \frac{c}{N} \rceil$确定，其中$c$是输入的可变部分。每次重复需要一个长度为$q + 200 + N + o$的输入。
• 最终回答：最后阶段处理一个长度为$q + 100 + o$的输入，其中包括了前面步骤累积的输出。

计算参数。我们设置$q = 1024$，$o = 1024$，$N = 5000$，并且$c$的范围从8K到4M，以计算最终结果。

B 完整的领域外任务结果

RULER基准测试概述。RULER基准测试包含四大类任务：检索、多跳追踪、聚合和问答。每项任务都旨在通过基于可配置输入参数（定义序列长度和复杂性）自动生成的示例来评估长上下文建模能力的特定方面。在这个受限的评估框架内，任务复杂性可以被概念化为目标输出token数量和上下文中信噪比的函数。

大海捞针（NIAH）任务。大海捞针（Needle-in-a-Haystack, NIAH）范式通过在大型干扰内容（“草堆”）中插入键值对（“针”）来评估长上下文场景中的检索性能。它包括四个变体：
- 单针NIAH：需要检索单个针。
- 多键NIAH：涉及在多个干扰项中检索一个针。
- 多值NIAH：必须提取与一个键关联的所有值。
- 多查询NIAH：需要检索具有不同键的多个针。

其他任务。变量追踪（Variable Tracking, VT）通过在扩展序列中追踪实体链来测试多跳推理。高频词提取（Frequent Words Extraction, FWE）则挑战模型识别幂律分布中的高频词，评估其分析语言数据中词频的能力。

方法细节中的引用汇总