解锁大型语言模型的工作记忆以实现潜在推理
Unlocking the Working Memory of Large Language Models for Latent Reasoning
AI总结:
提出一种名为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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。