Technical Summary: Reasoning in Memory (RiM)

Problem Statement

Current methods for enhancing the reasoning capabilities of Large Language Models (LLMs) rely heavily on scaling test-time compute by generating intermediate tokens before the final answer, such as in Chain-of-Thought (CoT) reasoning. However, this approach couples internal computation directly to autoregressive text generation. Consequently, the model is forced to "think out loud," allocating a significant portion of its computational budget to generating grammatically fluent natural language rather than purely internal reasoning.

While recent "latent reasoning" approaches attempt to bypass natural language constraints by using continuous representations, they largely preserve the step-by-step autoregressive generation paradigm. In these methods, intermediate computations are still externalized sequentially before the next step can condition on them. This paper argues that human cognition offers a superior design principle: working memory. Humans can hold and manipulate task-relevant information internally without externalizing every intermediate thought. The authors posit that LLMs can be trained to utilize a similar internal workspace, decoupling reasoning from the sequential bottleneck of token generation.

Methodology: Reasoning in Memory (RiM)

The authors introduce Reasoning in Memory (RiM), a latent reasoning method that replaces the autoregressive generation of reasoning steps with fixed memory blocks.

Core Mechanism

Instead of generating a sequence of reasoning tokens, the model is provided with $K$ fixed memory blocks, $m_{1:K}$. Each block consists of a predefined sequence of special tokens:
$$m_k = [<b>, <m>, \dots, <m>, </b>]$$
where $<b>$ and $</b>$ delimit block boundaries, and $<m>$ tokens form the latent workspace.

• Fixed Input: These blocks are appended to the input question as fixed tokens.
• Single Forward Pass: Because the blocks are fixed inputs rather than generated outputs, the entire augmented sequence (question + memory blocks) is processed in a single forward pass. This eliminates the sequential decoding bottleneck associated with autoregressive generation.
• Learned Representations: While token identities and positions are fixed, their contextual representations are learned to be task-dependent, encoding intermediate reasoning steps.

Two-Stage Curriculum

To teach the model to use these fixed blocks as a working memory, RiM employs a two-stage training curriculum:

Stage 1: Reasoning Step Supervision

• Goal: Ground the memory blocks by assigning them a computational role.
• Mechanism: For a reasoning trace with $T$ steps, the model is trained with $T$ memory blocks. After each block $m_t$, the model must predict the next explicit reasoning step $r_{t+1}$.
• Constraint: A custom attention mask ensures that the readout for step $t+1$ can attend to the question and previous memory blocks, but not to previous explicit reasoning steps. This forces the model to encode the necessary intermediate computation within the latent workspace of the memory blocks.
• Supervision: Uses a weighted negative log-likelihood objective with an annealing schedule to gradually reduce supervision strength.

Stage 2: Final Answer Refinement

• Goal: Convert the grounded memory blocks into a latent workspace for final answer prediction.
• Mechanism: The explicit reasoning trace is discarded. The model uses a fixed number of $K$ memory blocks. After each block, the model is trained to predict the final answer directly.
• Objective: This shifts the supervision from predicting intermediate steps to progressively refining the final answer as more memory blocks (latent computation) become available.
• Inference: At test time, the model processes the question and memory blocks in one pass, and the final block's readout provides the answer.

Attention Masking

A custom causal attention mask (Figure 2 in the paper) is critical to the method. It prevents information leakage by ensuring that:

1. Memory blocks attend to the question and previous memory blocks.
2. Target readouts (reasoning steps in Stage 1, final answer in Stage 2) attend to the question and available memory blocks, but not to other target readouts.
   This allows all targets to be predicted in a single forward pass without the model "cheating" by reading future ground truth.

Key Contributions

1. Reasoning in Memory (RiM): A novel latent reasoning method that unlocks the working-memory capacity of LLMs by providing a dedicated, fixed latent workspace for intermediate computation, decoupling it from autoregressive generation.
2. Two-Stage Curriculum: A structured training approach that first grounds memory blocks via explicit step supervision and then refines them for final answer prediction, enabling the effective use of latent space without complex auxiliary objectives.
3. Efficiency and Performance: Empirical evidence showing that RiM matches or exceeds existing explicit and latent reasoning baselines while maintaining the inference speed of direct-answer models (no sequential generation overhead).

Experimental Results

The authors evaluated RiM on the GSM8K (in-distribution) and GSM-Hard (out-of-distribution) mathematical reasoning benchmarks using GPT-2 and Llama-3.2 models (1B and 3B parameters).

• Performance: RiM consistently outperformed the strongest latent reasoning baseline (Coconut) and direct-answer Supervised Fine-Tuning (SFT).
  • On GSM8K (Llama-3.2-1B), RiM achieved 42.1% greedy accuracy, compared to 36.9% for Coconut and 23.9% for SFT without CoT.
  • On GSM-Hard, RiM achieved 10.5%, surpassing Coconut (8.5%) and SFT without CoT (5.3%).
  • RiM also showed superior Pass@8 scores, indicating that correct solutions are distributed across the latent trajectory with high probability.
• Latency: RiM achieved inference speeds comparable to SFT without CoT (Time to First Token ~16ms for Llama-3.2-1B), whereas Coconut was ~7x slower and SFT with CoT was ~27x slower due to autoregressive generation.
• Representation Analysis: PCA projections of memory block representations showed that the model learns to create distinct, sample-dependent latent workspaces. The representations evolve systematically during training, confirming that the blocks are not merely fixed placeholders but active computational states.
• Robustness: Performance remained stable across varying inference-time memory budgets (number of blocks), suggesting the method is robust to deployment constraints.

Significance and Claims

The paper claims that LLMs can be trained to use working memory as an effective mechanism for latent reasoning.

The primary significance lies in the decoupling of computation from communication. By replacing the generation of intermediate thoughts with fixed memory blocks, RiM allows models to perform dense, parallelizable internal computation without the overhead of generating natural language or sequential continuous representations. This results in a method that is:

1. Computationally Efficient: It avoids the sequential bottleneck of autoregressive generation, enabling single-pass inference.
2. Effective: It achieves state-of-the-art performance among latent reasoning methods on mathematical benchmarks.
3. Scalable: It works across different model families and sizes without requiring complex recurrent architectures or multi-pathway training objectives.

The authors conclude that this approach demonstrates a viable path toward more efficient reasoning systems that mimic human cognitive efficiency by internalizing the reasoning process rather than externalizing it as text.