Unlocking the Working Memory of Large Language Models for Latent Reasoning

2026-05-28

AI Summary

This work addresses the computational inefficiency and cognitive implausibility of existing large language models, which rely on autoregressive generation of intermediate reasoning steps. To overcome these limitations, the authors propose Reasoning in Memory (RiM), a novel approach that incorporates human-like working memory mechanisms into large language models for the first time. RiM employs fixed-length memory blocks to perform implicit reasoning within a single forward pass, effectively decoupling internal computation from external output generation. The method is trained via a two-stage curriculum: an initial pretraining phase where memory blocks predict explicit reasoning steps, followed by a fine-tuning phase that removes intermediate supervision and optimizes only for the final answer. Experiments demonstrate that RiM matches or surpasses state-of-the-art implicit reasoning methods across diverse model architectures and scales while substantially improving inference efficiency.

Abstract

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.

Problem

• working memory
• latent reasoning
• autoregressive generation
• large language models
• intermediate reasoning

Innovation

• working memory
• latent reasoning
• memory blocks
• non-autoregressive reasoning
• Reasoning in Memory