arXiv:2605.30343cs.CLcs.AI2026-05

用固定记忆块模拟工作记忆,让大模型高效进行隐式推理

Unlocking the Working Memory of Large Language Models for Latent Reasoning

论文配图:Unlocking the Working Memory of Large Language Models for Latent Reasoning
图 1 · 摘自论文原文
  • 用固定特殊标记序列替代生成中间思考过程
  • 单次前向传播完成推理,计算效率显著提升
  • 适配多种模型架构,无需依赖自回归生成

为提升大语言模型的推理能力,现有方法通常通过生成中间标记来扩展推理时间计算。但这种方式将内部计算与外部输出耦合,难以区分思维过程与表达行为。人类认知则可通过工作记忆在不外显的情况下处理信息。受此启发,我们提出一种名为「记忆中的推理」(RiM)的隐式推理方法,用固定序列的特殊标记块替代自回归生成的推理步骤。这些记忆块可直接利用模型的工作记忆能力,在单次前向传播中完成处理,实现高效率的隐式推理。为实现该机制,我们采用两阶段课程学习:首先通过预测每个记忆块后的显式推理步骤进行对齐;随后移除步骤级监督,仅通过迭代优化最终答案完成精炼。在多个推理基准测试中,RiM在不同家族和规模的模型上均达到或超越现有隐式推理方法性能,且避免了中间思考的自回归生成。结果表明,大语言模型可通过训练具备使用工作记忆进行有效隐式推理的能力。

原文摘要 · Abstract (English)

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.

隐式推理工作记忆大模型高效计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。