为什么将残差流限制在层而非令牌？用于连续潜在推理的持久记忆

# 为什么将残差流限制为层而不是标记？ 用于连续潜在推理的持久内存 来源：https://arxiv.org/html/2606.07720 ###### 摘要 大型语言模型 \(LLMs\) 在数学和多跳规划任务上展现了卓越的推理能力。CoCoNuT（连续思维链）范式 (Hao et al., 2024 (https://arxiv.org/html/2606.07720#bib.bib9)) 通过使模型能够在潜在空间中进行推理，同时探索多条推理路径而非早期就承诺单一链条，扩展了这种能力。然而，我们识别出一个我们称之为**概念瓶颈**的局限性。在每次推理过程中，中间隐藏状态被覆盖，导致模型随着推理深度的增加而丢失早期步骤中计算出的关键事实。我们在经验上观察到了这一点。在 HotpotQA 上，原始 CoCoNuT (10.4% EM) 未能超越 CoT 基线 (11.0% EM)，并且在 GSM8K 上，性能随着课程深度而下降。为了解决这个问题，我们提出了 **AGCLR** (自适应门控连续潜在推理)，它通过一个**门控概念流**增强了 CoCoNuT。这是一个在所有推理过程中保持的持久残差记忆，由三个学习得到的门控制：一个**写入门**，将中间事实提交到记忆；一个**读取门**，检索相关的先前状态；以及一个**遗忘门**，修剪不相关的上下文。在 GSM8K、HotpotQA 和 ProsQA 上使用 GPT-2 作为基础模型进行评估，AGCLR 在所有类型的数据集上都取得了一致的改进。随着课程深度的增加，性能差距放大，直接解决了概念瓶颈。代码可在 https://anonymous.4open.science/r/JJJJ/README.md 获取。 机器学习, ICML, 潜在推理, 思维链, 记忆 ## 1 引言 多步推理仍然是大型语言模型能力中最具挑战性的方面之一。Wei 等人 (2022 (https://arxiv.org/html/2606.07720#bib.bib14)) 表明，使用中间推理步骤提示 LLM 能显著提高数学和逻辑基准上的性能。然而，链式思维 (CoT) 推理被限制在单次前向传播中。每个生成的标记成为下一个的输入，迫使模型早期就承诺一个推理路径，并阻止探索替代路径。此外，显式的推理轨迹往往不完整或与底层计算不一致 (Su et al., 2026 (https://arxiv.org/html/2606.07720#bib.bib11); Swaroop et al., 2025 (https://arxiv.org/html/2606.07720#bib.bib12))，这激励了直接在潜在空间中进行的推理。 GSM8K HotpotQA ProsQA 0 50 50 100 100 40.6 40.6 1.0 1.0 55.0 55.0 31.4 31.4 10.4 10.4 92.0 92.0 34.0 34.0 14.0 14.0 96.0 96.0 数据集 准确率 / 精确匹配 (%) CoT 原始 CoCoNuT AGCLR (我们的) 图 1: AGCLR 擅长多跳推理。跨 GSM8K (数学)、HotpotQA (多跳 QA) 和 ProsQA (规划) 的性能。AGCLR 的持久内存使其在多跳任务 (HotpotQA: +3.6%, ProsQA: +4.0%) 上取得了显著提升，而 CoT 在单步数学推理上仍然更优。 更近期的工作探索了将这些推理链内化。Deng 等人 (2024 (https://arxiv.org/html/2606.07720#bib.bib3)) 提出了 iCoT，它在训练过程中逐步去除推理链的前缀，直到模型在没有任何显式链条的情况下预测答案。Goyal 等人 (2023 (https://arxiv.org/html/2606.07720#bib.bib6)) 引入了暂停标记，即在问题和答案之间插入的固定嵌入的特殊标记，以提供额外的计算时间。这两种方法都在语言空间中操作，无法在推理步骤间保持持久状态。最雄心勃勃的扩展是 CoCoNuT (Hao et al., 2024 (https://arxiv.org/html/2606.07720#bib.bib9))，它用连续的潜在思维替换了离散的推理标记。模型最后的隐藏状态直接作为下一个输入嵌入被反馈，从而实现在无约束潜在空间中的推理，并支持对推理路径的隐式广度优先搜索。CoCoNuT 通过一个多阶段课程进行训练，该课程逐步用潜在标记替换显式的推理步骤，每个阶段一步。 尽管有其前景，但原始 CoCoNuT 遭受**概念瓶颈**：中间推理状态在多轮推理过程中逐渐丢失，因为每个新的潜在标记都会覆盖早期轮次的信息，而没有持久记忆。这在需要更长链条的多跳推理中变得严重。我们在图 1 (https://arxiv.org/html/2606.07720#S1.F1) 中通过 GSM8K (算术)、HotpotQA (多跳 QA) 和 ProsQA (规划) 的经验证据展示了这一点。为了解决这个问题，我们提出了 **AGCLR** (自适应门控连续潜在推理)，它通过一个门控概念流增强了 CoCoNuT，该概念流在多个轮次间保留中间推理状态。虽然门控机制可以追溯到用于顺序状态更新的 LSTM (Hochreiter & Schmidhuber, 1997 (https://arxiv.org/html/2606.07720#bib.bib10))，但我们的门控操作于连续潜在推理中**跨轮次的持久内存**：每个轮次都优化相同的表示，而不是处理新的顺序输入，并且记忆在迭代推理周期中积累事实，而不是在每个时间步丢弃它们。 参见标题 图 2: AGCLR 架构。在每个潜在标记位置，三个学习得到的门 (读取、遗忘、写入) 控制当前隐藏状态 \(h_t\) 和持久概念流 \(c_t\) 之间的信息流。读取门从 \(c_{t-1}\) 检索相关的先前事实，遗忘门从 \(h_t\) 修剪不相关的上下文，写入门将门控后的隐藏状态 \(h'_t\) 提交到残差流中，直接解决了原始 CoCoNuT 中的概念瓶颈。 AGCLR 通过一个**门控概念流**增强了 CoCoNuT：一个在所有推理过程中保持的持久残差记忆向量 \(c_t \in \mathbb{R}^d\)。在每个潜在标记位置，三个学习得到的 sigmoid 门控制信息流：一个**写入门**将当前隐藏状态中相关的中间事实提交到记忆；一个**读取门**将先前的记忆检索到当前推理状态中；一个**遗忘门**从隐藏状态中修剪不相关的上下文。图 2 (https://arxiv.org/html/2606.07720#S1.F2) 展示了该架构。我们做出以下贡献： - • 我们识别并实证展示了原始 CoCoNuT 在三种不同类型的推理数据集上的**概念瓶颈**。 - • 我们提出了 **AGCLR**，一种门控残差记忆机制，它通过仅增加 GPT-2 的 1.41% 参数就解决了概念瓶颈。 - • AGCLR 在 GSM8K (算术)、HotpotQA (多跳 QA) 和 ProsQA (图规划) 上持续优于原始 CoCoNuT，并且优势随着课程深度的增加而放大。 ## 2 相关工作 用于控制信息流的门控机制可以追溯到长短期记忆网络 (Hochreiter & Schmidhuber, 1997 (https://arxiv.org/html/2606.07720#bib.bib10))，它引入了遗忘门以选择性地在循环隐藏状态中保留或丢弃信息。然而，LSTM 在时间步上对顺序输入进行门控，而我们的门控则作用于在相同潜在表示的迭代推理过程中跨轮次的持久内存。 (Deng et al., 2024 (https://arxiv.org/html/2606.07720#bib.bib3)) 提出了 iCoT，它在训练过程中逐步去除显式的推理前缀标记；虽然 iCoT 将推理压缩到前向传播中，但它缺乏任何在推理步骤间保留信息的机制，并且不适用于多轮次潜在推理设置。 (Hao et al., 2024 (https://arxiv.org/html/2606.07720#bib.bib9)) 引入了 CoCoNuT，它通过递归地将模型的隐藏状态作为下一个输入嵌入反馈来实现连续潜在推理，允许对推理路径进行隐式广度优先搜索。CoCoNuT 作为我们在所有三个数据集上的直接基线，但在每次轮次丢弃所有先前的隐藏状态，并且缺乏持久内存，导致了我们识别并解决的概念瓶颈。 (Wang et al., 2024 (https://arxiv.org/html/2606.07720#bib.bib13)) 提出了一种并行的后训练方法，使用一个固定标量 \(\alpha\) 在推理时混合连续的隐藏状态；与我们的方法不同，它们的门控不是端到端学习的，仅在连续状态上操作而非持久残差流，并且作为后处理以无需训练的方式应用。相关地，摊销潜在引导学习一种低成本干预，替代测试时的潜在优化 (Egbuna et al., 2025 (https://arxiv.org/html/2606.07720#bib.bib4))；然而，与固定 \(\alpha\) 混合类似，它应用一个瞬态引导信号，而不是在轮次之间维护一个持久的、门控的残差流。 记忆增强架构，如神经图灵机 (Graves et al., 2014 (https://arxiv.org/html/2606.07720#bib.bib7)) 和可微神经计算机 (Graves et al., 2016 (https://arxiv.org/html/2606.07720#bib.bib8))，已经探索了用于顺序推理的外部记忆，但它们通过跨序列块的读/写操作来增强模型，而不是像我们一样在多轮次潜在推理中维护持久的内部状态。 ## 3 方法: AGCLR ### 3.1 门控概念流 我们用一个持久概念流 \(c_t \in \mathbb{R}^d\) 增强 CoCoNuT，该流在每个前向调用开始时初始化为零，并在每个潜在标记位置更新。在轮次 \(t\)，给定潜在标记位置的隐藏状态 \(h_t\)： \[ \begin{aligned} \hat{h}_t &= \text{LayerNorm}(h_t), \tag{1} \\ r_t &= \sigma(W_r \hat{h}_t), \quad f_t = \sigma(W_f \hat{h}_t), \quad w_t = \sigma(W_w \hat{h}_t), \tag{2} \\ h'_t &= (1 - f_t) \odot h_t + r_t \odot c_{t-1}, \tag{3} \\ c_t &= \text{LayerNorm}(c_{t-1} + w_t \odot h'_t), \tag{4} \end{aligned} \] 其中 \(r_t, f_t, w_t \in [0,1]^d\) 是读取门、遗忘门和写入门，\(W_r, W_f, W_w \in \mathbb{R}^{d \times d}\) 是学习得到的权重矩阵。门控隐藏状态 \(h'_t\) 替换 \(h_t\) 作为下一个潜在标记位置的输入嵌入。**读取门** \(r_t\) 控制从概念流 \(c_{t-1}\) 中检索多少信息到当前隐藏状态中，允许轮次 \(t\) 访问来自所有早期轮次的事实。**遗忘门** \(f_t\) 控制当前隐藏状态被保留多少，被检索记忆替换多少，从而能够选择性修剪不相关的上下文。**写入门** \(w_t\) 控制门控隐藏状态 \(h'_t\) 中有多少被提交到概念流中，防止低置信状态污染残差记忆。 ## 4 训练协议 ### 4.1 多阶段课程 我们利用语言链式思维 (CoT) 数据通过实施一个受 Hao 等人 (2024 (https://arxiv.org/html/2606.07720#bib.bib9)) 启发的多阶段训练课程来监督连续潜在推理。在初始阶段 (阶段 0)，模型在具有显式推理步骤的常规 CoT 实例上进行训练。在后续阶段，我们逐步用连续潜在思维替换推理步骤。在阶段 \(k\)，CoT 中的前 \(k\) 个推理步骤被替换为 \(k \times c\) 个潜在标记，其中 \(c\) 是一个超参数，控制替换单个语言推理步骤的潜在思维数量。我们插入 <BOT> (思维开始) 和 <EOT> (思维结束) 标记来封装连续思维。遵循 Hao 等人 (2024 (https://arxiv.org/html/2606.07720#bib.bib9))，我们在训练阶段之间转换时重置优化器状态。 ### 4.2 实现细节 我们使用预训练的 GPT-2 基础模型 (117M 参数)，学习率为 \(1 \times 10^{-4}\)，有效批量大小为 128。我们在三个多跳推理基准上进行训练：GSM8K (Cobbe et al., 2021 (https://arxiv.org/html/2606.07720#bib.bib2))、HotpotQA (Yang et al., 2018 (https://arxiv.org/html/2606.07720#bib.bib15)) 和 ProsQA (Hao et al., 2024 (https://arxiv.org/html/2606.07720#bib.bib9))。遵循原始 CoCoNuT (Hao et al., 2024 (https://arxiv.org/html/2606.07720#bib.bib9)) 的课程结构，我们在第 1-9 个 epoch 期间经历阶段 0-2 (部分潜在推理)，逐步用潜在标记替换推理步骤。从第 10 个 epoch 开始，我们停留在阶段 3，其中所有推理都是潜在的，在 GSM8K 和 HotpotQA 上总共训练 15 个 epoch，在 ProsQA 上训练 20 个 epoch (该数据集包含多达 6 步的更复杂推理链)。对于 HotpotQA，我们格式化实例以包含问题、支持段落、中间推理步骤和答案跨度，以鼓励 CoT 阶段的多跳推理。使用最后阶段中验证准确率最佳的检查点进行评估。 ## 5 结果 ### 5.1 主要结果 表 1 (https://arxiv.org/html/2606.07720#S5.T1) 显示 AGCLR 在所有三个数据集上持续优于原始 CoCoNuT。 表 1: 在三个数据集：GSM8K、HotpotQA 和 ProsQA 上的结果。更高的准确率表示更强的推理能力。† 结果来自 Deng 等人 (2024 (https://arxiv.org/html/2606.07720#bib.bib3))，使用相同的 GPT-2 架构，如 Hao 等人 (2024 (https://arxiv.org/html/2606.07720#bib.bib9)) 所报告。HotpotQA 在先前的未评估。ProsQA 在阶段 6 (所有推理步骤都为潜在) 进行评估以进行公平比较。 ### 5.2 缓解概念瓶颈 原始 CoCoNuT 和 AGCLR 在早期课程阶段表现相当，但 AGCLR 的优势随着推理深度的增加而放大。在 ProsQA 上，原始 CoCoNuT 在阶段 5 达到 95% 的准确率，但在阶段 6 (最后一个课程阶段，所有推理步骤都被潜在标记替换) **下降**到 92%。而 AGCLR 则持续改进，在同一检查点达到 96%。这种下降与改进的模式展示了**概念瓶颈**：当模型过渡到完全潜在推理时，中间计算状态在没有保留机制的情况下逐渐丢失。 **跨推理轮次的记忆保持。** 为了理解门控如何解决这个瓶颈，我们分析了跨推理轮次的隐藏状态演化。图 3 (https://arxiv.org/html/2606.07720#S5.F3) 显示了在第 15 个 epoch 的 100 个验证样本上，第一轮隐藏状态与后续轮次之间的余弦相似度。原始 CoCoNuT 表现出单调的记忆衰减：相似度从 1.0 下降到第 6 轮的 0.126，随着中间推理步骤被逐步覆盖，信息损失达 87%。AGCLR 缓解了这种衰减。虽然相似度最初下降 (第 1 轮 → 第 2 轮)，但在第 3-6 轮时稳定在约 0.22，在最终生成时保留了比原始 CoCoNuT 多 71% 的信息 (0.216 vs 0.126)。门控概念流充当了一个持久记忆缓冲区，在轮次间保留关键推理状态。这种记忆保持直接解释了 AGCLR 在 HotpotQA 上 +3.6% EM 的提升。 第1轮 第2轮 第3轮 第4轮 第5轮 第6轮 0.0 0.2 0.2 0.4 0.4 0.6 0.6 0.8 0.8 1.0 0.216 0.126 +0.090 AGCLR 稳定 推理轮次 与第1轮的余弦相似度 原始 CoCoNuT AGCLR (我们的) 图 3: 隐藏状态记忆保持。第 1 轮与后续轮次之间的余弦相似度 (100 个样本，第 15 个 epoch)。原始 CoCoNuT 表现出单调衰减 (1.0 → 0.126)，而 AGCLR 在第 3 轮后稳定。阴影区域：±1 标准差。AGCLR 保留了多 71% 的信息 (+0.090)。