InfoMem: 使用答案条件信息增益训练长上下文记忆代理

# InfoMem: 训练长上下文记忆代理：基于答案条件的信息增益  
**来源：** https://arxiv.org/html/2606.03329  

Tiancheng Han¹,², Yong Li¹, Wuzhou Yu¹, Qiaosheng Zhang²,³,†, Wenqi Shao²,³,†  
¹同济大学  
²上海创新研究院  
³上海人工智能实验室  
[email protected]  
[email protected]  

###### 摘要  

长上下文任务要求 LLM 从大规模上下文中识别并保留与答案相关的信息。逐块记忆代理通过顺序读取文档块、更新紧凑记忆，并从累积的记忆中生成最终答案来解决这一问题。然而，现有的基于强化学习的逐块代理要么依赖稀疏的最终答案奖励，要么使用词汇级中间奖励来监督记忆和检索行为。这些信号监督任务成功与否或局部词汇重叠，但不能直接评估最终记忆是否支持真实答案。我们提出 InfoMem，一种用于训练逐块记忆代理的奖励机制，该机制利用基于答案条件的信息增益来评估最终记忆的效用。InfoMem 衡量最终记忆在多大程度上增加了模型对真实答案的逐 token 对数似然。为稳定强化学习优化，InfoMem 仅对成功轨迹应用此信号，并在奖励组合前进行归一化。在相同的 GRPO 框架和训练预算下，InfoMem 在长上下文记忆代理性能上优于可比较的记忆代理强化学习基线。分析表明，有效的最终记忆奖励应作用于成功轨迹，在奖励组合前进行归一化，并基于答案而非查询进行条件化。我们的代码可在 https://github.com/GenSouKa1/InfoMem 获取。  

**InfoMem: 训练长上下文记忆代理：基于答案条件的信息增益**  

Tiancheng Han¹,², Yong Li¹, Wuzhou Yu¹, Qiaosheng Zhang²,³,†, Wenqi Shao²,³,†  
¹同济大学  
²上海创新研究院  
³上海人工智能实验室  
[email protected]  
[email protected]  

{NoHyper}††footnotetext:†通讯作者。  

## 1 引言  

长上下文理解已成为大型语言模型（LLM）的核心能力，其应用涵盖长文档问答到语料级证据聚合等领域（OpenAI, 2025 (https://arxiv.org/html/2606.03329#bib.bib15); Wu 等, (https://arxiv.org/html/2606.03329#bib.bib17); Hsieh 等, (https://arxiv.org/html/2606.03329#bib.bib13); Lu 等, 2026 (https://arxiv.org/html/2606.03329#bib.bib16)）。先前的工作已通过扩展上下文窗口（Shen 等, 2025 (https://arxiv.org/html/2606.03329#bib.bib19)）、注意力或位置修改（Munkhdalai 等, 2024 (https://arxiv.org/html/2606.03329#bib.bib20); Press 等, (https://arxiv.org/html/2606.03329#bib.bib21)）、检索增强（Zhao 等, 2024 (https://arxiv.org/html/2606.03329#bib.bib4)）以及基于记忆或智能体的流水线（Packer 等, 2024 (https://arxiv.org/html/2606.03329#bib.bib22); Zhou 等, 2025 (https://arxiv.org/html/2606.03329#bib.bib23)）来改进长上下文处理。然而，当相关证据稀疏、分布在遥远的片段中，或需要在漫长的阅读过程中加以保留时，有效利用长输入仍然具有挑战性。  

在这些方法中，逐块记忆代理为长上下文推理提供了一种简单而有效的范式。模型不是一次性处理整个文档，而是顺序读取较短的块，更新一个紧凑的记忆状态，并从累积的记忆中生成最终答案。这一范式出现在免训练阅读代理（Lee 等, 2024 (https://arxiv.org/html/2606.03329#bib.bib24)）、递归或记忆增强模型架构（Dai 等, 2019 (https://arxiv.org/html/2606.03329#bib.bib25)）以及后训练方法（Yu 等, 2025 (https://arxiv.org/html/2606.03329#bib.bib1); Shi 等, 2026 (https://arxiv.org/html/2606.03329#bib.bib32)）中。至关重要的是，其显式的记忆状态使记忆质量成为一个可观察的优化目标。  

尽管有效，现有的逐块长上下文系统仍然缺乏一种可扩展的方法来监督记忆的形成。免训练方法通常依赖人工设计的记忆更新提示、摘要启发式或固定的遍历流程（Lee 等, 2024 (https://arxiv.org/html/2606.03329#bib.bib24)）。架构级方法虽可能从更根本上改善长上下文能力（Dai 等, 2019 (https://arxiv.org/html/2606.03329#bib.bib25)），但通常需要昂贵的预训练。基于强化学习的记忆代理也可以通过任务反馈改进长上下文行为，但现有方法主要依赖稀疏的答案奖励（Yu 等, 2025 (https://arxiv.org/html/2606.03329#bib.bib1)）或用于记忆和检索动作的词汇级中间奖励（Shi 等, 2026 (https://arxiv.org/html/2606.03329#bib.bib32)）。这些奖励监督任务成功与否或局部词汇重叠，但不直接评估最终记忆是否在语义上支持真实答案。  

这种局限性在成功轨迹中尤为突出：稀疏的结果奖励无法区分最终记忆是包含聚焦于答案的支持证据，还是包含冗余的干扰信息；而词汇奖励可能无法捕捉对最终答案的语义支持。这激发了针对逐块长上下文强化学习的特定记忆奖励信号的需求。  

我们提出 InfoMem，一种基于答案条件信息增益的奖励塑形方法，用于逐块长上下文记忆代理。其核心直觉是：一个有用的最终记忆应增加模型对真实答案的支持。InfoMem 不估计分布级互信息，而是通过比较模型在有无最终记忆条件下对真实答案的逐 token 对数似然，使用逐点信息增益代理。InfoMem 进一步通过仅对成功轨迹应用此信号并在奖励组合前进行归一化来提升训练稳定性。  

实验表明，InfoMem 在长上下文记忆代理性能上持续优于仅使用结果的 GRPO 和可比较的记忆代理强化学习基线 ReMemR1（Shi 等, 2026 (https://arxiv.org/html/2606.03329#bib.bib32)）。进一步的分析表明，有效的最终记忆奖励应作用于成功轨迹，在奖励组合前进行归一化，并基于真实答案而非仅基于查询进行条件化。这些发现表明，基于答案条件的信息增益为逐块长上下文强化学习中的最终记忆监督提供了一个原则性的框架。  

我们的贡献有三方面：(1) 我们从信息论视角阐述了最终记忆的效用，即有用的记忆应降低模型对真实答案的不确定性。(2) 我们引入了 InfoMem，一种基于答案条件的信息增益奖励，用于在成功轨迹上直接塑造最终记忆。(3) 我们发现 InfoMem 在逐块长上下文记忆代理上持续优于可比较的记忆代理强化学习基线，并进一步识别了有效最终记忆奖励的三个关键属性：成功侧监督、组合前归一化和基于答案的条件化。  

## 2 相关工作  

### 2.1 长上下文 LLM 与逐块记忆代理  

长上下文 LLM 研究已通过上下文扩展（Shen 等, 2025 (https://arxiv.org/html/2606.03329#bib.bib19)）、注意力或位置修改（Munkhdalai 等, 2024 (https://arxiv.org/html/2606.03329#bib.bib20); Press 等, (https://arxiv.org/html/2606.03329#bib.bib21)）以及检索增强（Zhao 等, 2024 (https://arxiv.org/html/2606.03329#bib.bib4)）改进了长输入处理。与这些方法互补，逐块记忆代理在顺序长文档处理过程中维护一个显式的记忆状态。该范式包括免训练阅读流程（Lee 等, 2024 (https://arxiv.org/html/2606.03329#bib.bib24)）、片段级递归架构（Dai 等, 2019 (https://arxiv.org/html/2606.03329#bib.bib25); Hutchins 等, 2022 (https://arxiv.org/html/2606.03329#bib.bib26); Ding 等, 2021 (https://arxiv.org/html/2606.03329#bib.bib27)）以及后训练记忆代理（Yu 等, 2025 (https://arxiv.org/html/2606.03329#bib.bib1); Shi 等, 2026 (https://arxiv.org/html/2606.03329#bib.bib32)）。现有方法依赖人工设计的流程、昂贵的递归架构，或基于稀疏最终答案奖励和中间记忆启发式的强化学习目标。因此，对基于答案条件的最终记忆效用的直接监督仍相对未被充分探索。  

### 2.2 用于长上下文问答的强化学习  

DeepSeek-R1（Guo 等, 2025 (https://arxiv.org/html/2606.03329#bib.bib28)）表明，强化学习可以在专业领域大幅提升 LLM 能力。近期工作进一步显示，强化学习可以改善长上下文问答和推理。早期方法主要优化可验证的端任务结果，如最终答案正确性或基于验证器的响应质量（Shen 等, 2025 (https://arxiv.org/html/2606.03329#bib.bib19); Yu 等, 2025 (https://arxiv.org/html/2606.03329#bib.bib1)）。最近的研究引入了更密集的监督信号，用于基础事实定位、证据提取和上下文推理（Chen 等, (https://arxiv.org/html/2606.03329#bib.bib29); Guan 等, 2026 (https://arxiv.org/html/2606.03329#bib.bib30); Ping 等, 2026 (https://arxiv.org/html/2606.03329#bib.bib31); Shi 等, 2026 (https://arxiv.org/html/2606.03329#bib.bib32)）。这些结果共同表明，基于奖励的优化是改进长上下文推理的一个有前景的方向。  

尽管取得了这些进展，现有奖励主要监督定位质量、证据选择、阅读效用或中间记忆更新行为。ReMemR1（Shi 等, 2026 (https://arxiv.org/html/2606.03329#bib.bib32)）引入了用于记忆和回调动作的信息型奖励，但这些信号基于词级召回率，而非基于答案条件的最终记忆效用。对最终记忆表示本身的直接监督仍相对未被充分探索。特别是，现有方法很少评估最终记忆是否直接支持真实答案，而这正是我们工作的重点。  

## 3 问题设定与动机  

### 3.1 逐块长上下文记忆代理  

我们将逐块记忆代理作为长上下文设置中强化学习的一种实用范式。形式上，令 \(x\) 表示查询，\(D\) 表示长文档，\(y^*\) 表示真实答案。给定预定义的块大小 \(C\)，文档被划分为 \(K\) 个块：  

\[
D = \{c_1, c_2, \ldots, c_K\}, \tag{1}
\]

基于查询，模型顺序读取块并维护一个记忆状态：  

\[
M_t = \pi_\theta(M_{t-1}, c_t, x), \quad t = 1, \ldots, K. \tag{2}
\]

处理完所有块后，模型获得最终记忆 \(M_K\)，并基于查询和该记忆生成最终答案：  

\[
\hat{y} = \pi_\theta(x, M_K). \tag{3}
\]

### 3.2 为何结果奖励对记忆学习不足  

结果奖励直接监督最终答案的正确性，但对记忆效用仅提供稀疏且间接的监督。在逐块记忆代理中，最终预测基于最终记忆 \(M_K\) 生成，该记忆应保留支持正确答案所需的信息。然而，二元结果奖励仅评估生成答案是否匹配 \(y^*\)，而不直接区分不同最终记忆的质量。  

这种局限性在成功轨迹中尤为明显。多次 rollout 可能生成相同的正确答案，因此获得相同的结果奖励，但其最终记忆的效用可能差异显著。某些记忆可能仅保留关键支持证据，而其他记忆可能保留相同证据但附带大量冗余或干扰信息，却仍能产生正确的预测。因此，仅靠结果奖励无法区分成功轨迹内记忆的效用，这激发了对逐块长上下文强化学习中的最终记忆效用进行直接奖励信号的需求。  

### 3.3 从互信息到模型诱导的逐点代理  

长上下文问答可以视为从大型上下文中提取与答案相关的信息。在逐块记忆代理的框架下，最终记忆应降低在查询条件下答案的不确定性。理想情况下，这种效用可以通过条件互信息 \(I(M; Y \mid X)\) 来刻画，该指标衡量给定查询 \(X\) 后，记忆 \(M\) 关于答案 \(Y\) 提供了多少额外信息。  

然而，分布级互信息依赖于查询、记忆和答案的完整联合分布，并且在 LLM 的高维语义空间中难以可靠估计（Qian 等, 2025 (https://arxiv.org/html/2606.03329#bib.bib5)）。这激发了第 4 节（https://arxiv.org/html/2606.03329#S4）中引入的 InfoMem 奖励，其中我们将互信息目标实例化为一个单样本逐点代理。该代理衡量当前实例上最终记忆是否增加了模型对真实答案的支持。  

## 4 方法：InfoMem  

参见说明  

**图 1：** InfoMem 用于逐块长上下文强化学习的概述。InfoMem 通过比较教师强制下真实答案 \(y^*\) 的逐 token 平均对数似然（在有和没有最终记忆的情况下），利用基于答案条件的信息增益来衡量最终记忆的效用。在 GRPO 训练期间，信息增益监督仅应用于成功轨迹，在成功 rollout 间进行归一化，并与稀疏结果奖励结合用于策略优化。  

### 4.1 奖励定义  

我们使用组相对策略优化（GRPO）（Shao 等, 2024 (https://arxiv.org/html/2606.03329#bib.bib11)）训练逐块记忆代理。对于每个提示，GRPO 采样一组 \(n\) 个 rollout，  

\[
\mathcal{G} = \{1, 2, \ldots, n\}. \tag{4}
\]

每个 rollout \(i \in \mathcal{G}\) 产生一个最终记忆 \(M_i\) 和一个最终答案 \(\hat{y}_i\)。基本结果奖励定义为：  

\[
R_{\mathrm{outcome}, i} = \mathbb{1}[\hat{y}_i = y^*], \tag{5}
\]

在我们的主要实验中，答案正确性通过归一化字符串匹配来评估。给定查询 \(x\)、最终记忆 \(M\) 和真实答案 \(y^*\)，我们将基于答案条件的信息增益奖励 \(r_{\mathrm{gain}}\) 定义为：  

\[
r_{\mathrm{gain}}(x, M, y^*) = \frac{1}{|y^*|} \log P_\theta(y^* \mid x, M) - \frac{1}{|y^*|} \log P_\theta(y^* \mid x, \emptyset), \tag{6}
\]

其中 \(\emptyset\) 表示空记忆。LLM 在真实 token 上的序列似然 \(P_\theta\) 在教师强制下计算：  

\[
\log P_\theta(y^* \mid x, M) = \sum_{j=1}^{|y^*|} \log p_\theta(y_j^* \mid y_0^{<j}, x, M), \quad y_0^{<j} = \{y_1^*, \ldots, y_{j-1}^*\}.
\]

该奖励直接量化最终记忆相对于空记忆所提供的信息增益：正值表示支持真实答案的条件对数似然增加，负值则相反。  

### 4.2 训练流程  

**基于信息增益的记忆奖励。** 对于每个 rollout \(i \in \mathcal{G}\)，我们首先计算结果奖励 \(R_{\mathrm{outcome}, i}\)。成功轨迹定义为 \(\mathcal{S} = \{i \in \mathcal{G} \mid R_{\mathrm{outcome}, i} = 1\}\)。对于每个成功 rollout \(i \in \mathcal{S}\)，我们根据公式 (6) 计算信息增益分数 \(r_{\mathrm{gain}, i}\)。为增强训练稳定性，我们将这些分数在 \(\mathcal{S}\) 内进行归一化：  

\[
\tilde{r}_{\mathrm{gain}, i} = \frac{r_{\mathrm{gain}, i} - \mu_{\mathrm{gain}}}{\sigma_{\mathrm{gain}}}, \quad i \in \mathcal{S},
\]

其中 \(\mu_{\mathrm{gain}}\) 和 \(\sigma_{\mathrm{gain}}\) 是集合 \(\{r_{\mathrm{gain}, i} \mid i \in \mathcal{S}\}\) 的均值和标准差。若 \(|\mathcal{S}| = 1\)，我们跳过归一化并直接使用原始分数。若 \(|\mathcal{S}| = 0\)，则不应用信息增益奖励。  

**奖励组合。** 我们组合结果奖励和归一化的信息增益奖励：  

\[
R_{\mathrm{final}, i} = R_{\mathrm{outcome}, i} + \beta \cdot \mathbb{1}[i \in \mathcal{S}] \cdot \tilde{r}_{\mathrm{gain}, i},
\]

其中 \(\beta\) 是一个超参数，控制信息增益奖励的权重。  

**训练算法。** 整体训练过程如算法 1 所示。注意我们仅区分成功和非成功轨迹；非成功 rollout 不会收到信息增益信号，但参与策略优化以降低其概率。  

**推理。** 在推理时，我们直接使用训练后的智能体进行块级阅读和最终答案生成，无需计算信息增益分数。  

### 算法 1 InfoMem 训练流程  

**输入：** 策略 \(\pi_\theta\)，参考策略 \(\pi_{\mathrm{ref}}\)，组大小 \(n\)，KL 系数 \(\lambda\)，信息增益系数 \(\beta\)，查询 \(x\)，文档 \(D\)，真实答案 \(y^*\)，块大小 \(C\)  
**输出：** 更新后的策略 \(\pi_\theta\)  

1: 将 \(D\) 按块大小 \(C\) 划分为块 \(\{c_1, \ldots, c_K\}\)  
2: 初始化记忆 \(M_0\)（空）  
3: **for** \(i = 1\) **to** \(n\) **do**  
4:   **for** \(t = 1\) **to** \(K\) **do**  
5:     从策略中采样阅读和记忆操作，更新记忆为 \(M_t\)  
6:   **end for**  
7:     生成最终答案 \(\hat{y}_i \sim \pi_\theta(\cdot \mid x, M_K)\)  
8:     计算结果奖励 \(R_{\mathrm{outcome}, i} = \mathbb{1}[\hat{y}_i = y^*]\)  
9: **end for**  
10: \(\mathcal{S} \leftarrow \{i \in [n] \mid R_{\mathrm{outcome}, i} = 1\}\)  
11: **if** \(|\mathcal{S}| > 0\) **then**  
12:   **for** \(i \in \mathcal{S}\) **do**  
13:     \(r_i \leftarrow r_{\mathrm{gain}}(x, M_i, y^*)\) ▷ 教师强制评分  
14:   **end for**  
15:   **if** \(|\mathcal{S}| = 1\) **then**  
16:     \(\tilde{r}_i \leftarrow r_i\) 对于唯一的 \(i \in \mathcal{S}\)  
17:   **else**  
18:     计算 \(\mu_{\mathrm{gain}} = \frac{1}{|\mathcal{S}|} \sum_{i \in \mathcal{S}} r_i\)，\(\sigma_{\mathrm{gain}} = \sqrt{\frac{1}{|\mathcal{S}|} \sum_{i \in \mathcal{S}} (r_i - \mu_{\mathrm{gain}})^2}\)  
19:     **for** \(i \in \mathcal{S}\) **do**  
20:       \(\tilde{r}_i \leftarrow (r_i - \mu_{\mathrm{gain}}) / \sigma_{\mathrm{gain}}\)  
21:     **end for**  
22:   **end if**  
23:   **for** \(i \in \mathcal{S}\) **do**  
24:     \(R_{\mathrm{final}, i} \leftarrow R_{\mathrm{outcome}, i} + \beta \cdot \tilde{r}_i\)  
25:   **end for**  
26: **end if**  
27: **for** \(i \notin \mathcal{S}\) **do**  
28:   \(R_{\mathrm{final}, i} \leftarrow R_{\mathrm{outcome}, i}\)  
29: **end for**  
30: 使用 GRPO 更新 \(\pi_\theta\)：最大化 \(\mathbb{E}_{i \sim \mathcal{G}} \left[ \frac{1}{n} \sum_{i=1}^n \left( \frac{\pi_\theta(y_i | x, D)}{\pi_{\mathrm{ref}}(y_i | x, D)} R_{\mathrm{final}, i} - \lambda \mathrm{KL}(\pi_\theta \| \pi_{\mathrm{ref}}) \right) \right]\)  
31: **返回** \(\pi_\theta\)  

（注：公式和算法中语言已翻译，数学符号保持原样。）