LatentHarness: Learning Latent Actions for Memory and Reasoning via Counterfactual Policy Distillation

📄 arXiv: 2609.39740v1 📥 PDF

作者: Xiaoqiang Wang, Suyuchen Wang, Bang Liu

分类: cs.CL

发布日期: 2026-09-30

备注: Work in progress

💡 一句话要点

提出LatentHarness以解决长上下文推理中的记忆与推理瓶颈问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 长上下文推理 潜在动作选择 反事实策略蒸馏 记忆与推理 自然语言处理

📋 核心要点

  1. 长上下文推理面临记忆和计算效率的双重挑战，现有方法往往无法有效整合这两方面。
  2. LatentHarness通过顺序潜在动作选择将记忆访问与潜在推理结合，优化了推理过程中的决策。
  3. 在六个基准测试中，LatentHarness在性能和速度上均显著优于现有最强基线，展示了其有效性。

📝 摘要（中文）

长上下文推理面临两个互补的瓶颈：在长输入中保留证据和在多个推理步骤中维持计算。现有方法通常分别解决这两个问题，外部记忆扩展了对远程证据的访问，而潜在推理压缩了多步骤计算。本文提出LatentHarness，将记忆访问与潜在推理统一为顺序潜在动作选择。在每个内部步骤中，模型选择THINK进行进一步计算，RECALL从快速权重记忆中提取输入证据和中间推理状态，或EXIT以发出下一个标记。我们通过反事实策略蒸馏训练该策略，评估每个动作对发出标记的影响。这些改进教会策略在何时记忆比进一步推理更有用，同时通过反事实回忆的梯度教会哪些中间状态应保留以备将来使用。在六个通用和长上下文推理基准上，LatentHarness在1.4B参数下相较于最强基线提高了2.8%和10.0%的相对性能，并且运行速度比最强长上下文基线快5.9倍。

🔬 方法详解

问题定义：本文旨在解决长上下文推理中的记忆保留和多步骤计算效率低下的问题。现有方法往往分别处理这两个问题，导致整体性能受限。

核心思路：LatentHarness通过将记忆访问与潜在推理结合，采用顺序潜在动作选择的方式，允许模型在推理过程中动态选择进一步计算、回忆或退出的策略。

技术框架：该方法的整体架构包括三个主要模块：THINK（进行进一步计算）、RECALL（从快速权重记忆中提取信息）和EXIT（发出下一个标记）。模型在每个步骤中根据当前状态选择最优动作。

关键创新：LatentHarness的核心创新在于通过反事实策略蒸馏训练，评估每个动作对最终输出的影响，从而优化策略选择。这种方法与传统的单一推理或记忆方法有本质区别。

关键设计：在训练过程中，模型使用反事实回忆的梯度信息来指导哪些中间状态应被保留在记忆中，确保未来推理的有效性。模型参数设置为1.4B，确保在性能与计算效率之间取得良好平衡。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

在六个长上下文推理基准测试中，LatentHarness在1.4B参数下相较于最强基线提高了2.8%和10.0%的相对性能，同时运行速度比最强长上下文基线快5.9倍，显示出其在效率和效果上的显著优势。

🎯 应用场景

LatentHarness的研究成果在自然语言处理、智能问答系统和长文本理解等领域具有广泛的应用潜力。通过提高长上下文推理的效率和准确性，该方法能够显著提升智能系统的决策能力和用户体验，未来可能推动更复杂的推理任务的实现。

📄 摘要（原文）

Long-context reasoning faces two complementary bottlenecks: retaining evidence across long inputs and sustaining computation across many reasoning steps. Existing approaches largely address them separately, with external memory extending access to distant evidence and latent reasoning compressing multi-step computation. We introduce LatentHarness, which unifies memory access and latent reasoning as sequential latent action selection. At each internal step, the model chooses THINK for further computation, RECALL from a fast-weight memory of input evidence and intermediate reasoning states, or EXIT to emit the next token. We train this policy with counterfactual policy distillation, which branches every action for one step and scores its effect on the emitted token. These gains teach the policy when memory is more useful than further reasoning, while gradients through counterfactual recall teach which intermediate states should be retained in memory for future use. Across six general and long-context reasoning benchmarks, LatentHarness at 1.4B improves on the strongest baselines by 2.8% and 10.0% relative, respectively, and runs 5.9x faster than the strongest long-context baseline.