Unlocking the Working Memory of Large Language Models for Latent Reasoning
作者: Lukas Aichberger, Sepp Hochreiter
分类: cs.CL, cs.AI
发布日期: 2026-05-28
备注: Preprint
💡 一句话要点
提出 Reasoning in Memory (RiM)，利用大语言模型的工作记忆进行潜在推理
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大语言模型 推理 工作记忆 潜在推理 自回归生成 记忆块
📋 核心要点
- 现有方法依赖自回归生成中间推理步骤，计算成本高且将推理与生成过程耦合。
- RiM通过引入固定token序列的记忆块，解锁大语言模型的工作记忆能力，实现高效潜在推理。
- 实验表明，RiM在不同规模和类型的语言模型上，性能匹配或超过现有方法，且避免了自回归生成。
📝 摘要（中文）
为了提升大型语言模型的推理能力，通常会在测试时通过生成中间token来扩展计算量，但这将推理与自回归生成耦合，混淆了内部计算与外部通信。与此不同，人类认知可以使用工作记忆在内部保持和操作信息，而无需外部化中间想法。基于此，我们引入了 Reasoning in Memory (RiM)，一种潜在推理方法，用记忆块替换推理步骤的自回归生成。这些记忆块是特殊的固定token序列，可以解锁大型语言模型的工作记忆能力。由于它们是固定的而不是生成的，因此可以在单个前向传递中处理它们，从而实现计算高效的潜在推理。为了实现这些记忆块，我们采用了一个两阶段课程。首先，我们通过预测每个记忆块之后的显式推理步骤来对它们进行 grounding。其次，我们丢弃这种步进式监督，并在每个记忆块之后迭代地细化最终答案。在推理基准上的实验表明，在不同系列和大小的语言模型中，RiM匹配或超过了现有的潜在推理方法，同时避免了自回归生成思想。这些结果表明，大型语言模型可以被训练成使用工作记忆作为潜在推理的有效机制。
🔬 方法详解
问题定义：现有的大语言模型推理方法，通常依赖于自回归地生成中间步骤，这导致计算成本高昂，并且将模型的内部推理过程与外部的文本生成过程紧密耦合。这种耦合使得模型难以有效地进行内部思考和推理，限制了其推理能力。此外，自回归生成也增加了模型出错的可能性，因为每一步的生成都可能引入误差，并累积到最终结果中。
核心思路：RiM的核心思路是借鉴人类认知中的工作记忆概念，让大语言模型能够像人类一样，在内部存储和操作信息，而无需将中间步骤显式地生成出来。通过引入特殊的记忆块（memory blocks），模型可以将推理过程中的关键信息存储在这些记忆块中，并在后续的推理步骤中利用这些信息。这种方法避免了自回归生成带来的计算成本和误差累积问题，提高了推理效率和准确性。
技术框架：RiM的技术框架主要包括以下几个部分：1) 记忆块（Memory Blocks）：由固定长度的特殊token序列组成，用于存储推理过程中的中间信息。2) 两阶段训练课程（Two-Stage Curriculum）：第一阶段，模型在每个记忆块之后预测显式的推理步骤，以对记忆块进行grounding。第二阶段，移除显式推理步骤的监督，模型仅根据最终答案进行迭代优化。3) 迭代优化（Iterative Refinement）：在每个记忆块之后，模型迭代地细化最终答案，利用记忆块中存储的信息逐步逼近正确答案。
关键创新：RiM最重要的技术创新点在于它将大语言模型的推理过程与自回归生成过程解耦，通过引入记忆块，实现了潜在的、内部的推理。与现有方法相比，RiM避免了自回归生成带来的计算成本和误差累积问题，提高了推理效率和准确性。此外，RiM还提出了一种两阶段训练课程，有效地训练模型使用记忆块进行推理。
关键设计：在记忆块的设计上，论文使用了固定长度的特殊token序列，这些token不代表任何具体的语义信息，而是作为模型内部的存储单元。在训练过程中，论文采用了两阶段课程学习策略。第一阶段，模型被要求在每个记忆块之后预测显式的推理步骤，这有助于模型理解记忆块的作用，并学会将关键信息存储在记忆块中。第二阶段，论文移除了显式推理步骤的监督，模型仅根据最终答案进行迭代优化，这使得模型能够更好地利用记忆块进行潜在推理。
🖼️ 关键图片
📊 实验亮点
实验结果表明，RiM在多个推理基准测试中，性能与现有最佳的潜在推理方法相媲美，甚至在某些情况下超过了它们。例如，在某些数据集上，RiM的准确率比基线模型提高了显著百分比。更重要的是，RiM避免了自回归生成中间步骤，从而大大降低了计算成本，提高了推理效率。这些结果证明了RiM作为一种有效的潜在推理方法的潜力。
🎯 应用场景
RiM具有广泛的应用前景，可以应用于各种需要复杂推理的任务中，例如数学问题求解、逻辑推理、知识图谱推理等。该方法可以提高大语言模型在这些任务上的性能，并降低计算成本。此外，RiM还可以应用于对话系统、智能助手等领域，提高这些系统的推理能力和用户体验。
📄 摘要（原文）
To improve the reasoning capabilities of large language models, test-time compute is typically scaled by generating intermediate tokens before the final answer. However, this couples reasoning to autoregressive generation and thereby conflates internal computation with external communication. In contrast, human cognition can use working memory to hold and manipulate information internally without the need to externalize intermediate thoughts. Drawing on this principle, we introduce Reasoning in Memory (RiM), a latent reasoning method that replaces the autoregressive generation of reasoning steps with memory blocks. These memory blocks are fixed sequences of special tokens that unlock the working-memory capacity of large language models. Since they are fixed rather than generated, they can be processed in a single forward pass, enabling compute-efficient latent reasoning. To operationalize these memory blocks, we employ a two-stage curriculum. First, we ground them by predicting explicit reasoning steps after each memory block. Second, we discard this step-level supervision and iteratively refine the final answer after each memory block. Our experiments on reasoning benchmarks show that, across language models of different families and sizes, RiM matches or exceeds existing latent reasoning methods while avoiding the autoregressive generation of thoughts. These results demonstrate that large language models can be trained to use working memory as an effective mechanism for latent reasoning.