Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Encode, Think, Decode: Scaling test-time reasoning with recursive latent thoughts

Yeskendir Koishekenov, Aldo Lipani, Nicola Cancedda

TL;DR

Encode–Think–Decode (ETD) presents a mid-training technique that amplifies latent-space reasoning by iterating over a targeted subset of layers, forming a latent encoder–thinking block–latent decoder structure without adding parameters or data. By identifying reasoning-critical layers via angular-distance analysis and knee-point detection, ETD achieves substantial improvements across 17 reasoning benchmarks, including +28.4% on GSM8K and +36% on MATH for the OLMo-2 1B base, and supports adaptive test-time depth for token-specific computation. The work demonstrates that recursive latent reasoning can enhance reasoning performance on open-source LLMs with realistic training practices, offering a practical and scalable path to stronger LLM reasoning. The adaptive-depth variant further allocates compute where needed, balancing accuracy with efficiency and expanding the applicability of latent reasoning in real-world tasks.

Abstract

Most efforts to improve the reasoning capabilities of large language models (LLMs) involve either scaling the number of parameters and the size of training data, or scaling inference computation by letting models generate complex chains of thought. Motivated by interpretability studies showing that the crucial computation required for reasoning tasks is concentrated in a limited range of layers, we introduce Encode-Think-Decode (ETD), a method that enhances the reasoning capabilities of a base model by training it to iterate over a small subset of reasoning-relevant layers during the mid-training stage. ETD amplifies latent reasoning while preserving the original architecture, parameter count, hyperparameters, and training data composition. When iterating on the selected layers at inference time, ETD models yield substantial gains on 17 reasoning benchmarks, including +28.4% relative accuracy improvement on GSM8K and +36% on MATH with the OLMo-2 1B Base model. We also explore an adaptive depth strategy that adjusts the computation per input token. Our results show that recursive latent reasoning offers a simple and effective path to stronger LLM reasoning.

Encode, Think, Decode: Scaling test-time reasoning with recursive latent thoughts

TL;DR

Encode–Think–Decode (ETD) presents a mid-training technique that amplifies latent-space reasoning by iterating over a targeted subset of layers, forming a latent encoder–thinking block–latent decoder structure without adding parameters or data. By identifying reasoning-critical layers via angular-distance analysis and knee-point detection, ETD achieves substantial improvements across 17 reasoning benchmarks, including +28.4% on GSM8K and +36% on MATH for the OLMo-2 1B base, and supports adaptive test-time depth for token-specific computation. The work demonstrates that recursive latent reasoning can enhance reasoning performance on open-source LLMs with realistic training practices, offering a practical and scalable path to stronger LLM reasoning. The adaptive-depth variant further allocates compute where needed, balancing accuracy with efficiency and expanding the applicability of latent reasoning in real-world tasks.

Abstract

Most efforts to improve the reasoning capabilities of large language models (LLMs) involve either scaling the number of parameters and the size of training data, or scaling inference computation by letting models generate complex chains of thought. Motivated by interpretability studies showing that the crucial computation required for reasoning tasks is concentrated in a limited range of layers, we introduce Encode-Think-Decode (ETD), a method that enhances the reasoning capabilities of a base model by training it to iterate over a small subset of reasoning-relevant layers during the mid-training stage. ETD amplifies latent reasoning while preserving the original architecture, parameter count, hyperparameters, and training data composition. When iterating on the selected layers at inference time, ETD models yield substantial gains on 17 reasoning benchmarks, including +28.4% relative accuracy improvement on GSM8K and +36% on MATH with the OLMo-2 1B Base model. We also explore an adaptive depth strategy that adjusts the computation per input token. Our results show that recursive latent reasoning offers a simple and effective path to stronger LLM reasoning.

Paper Structure

This paper contains 27 sections, 7 equations, 3 figures, 6 tables.

Table of Contents

  1. Introduction
  2. On the Roles of Layers for Reasoning
  3. Choosing the optimal configuration for latent reasoning
  4. Experimental Setup
  5. Training Pipeline
  6. Evaluation Benchmarks
  7. Evaluating Recursive "Thinking" Blocks
  8. Performance Gains from Iterating over "Thinking" Blocks
  9. Scaling Behavior of Latent Reasoning
  10. Comparison with alternative recursive frameworks
  11. How does the choice of recursive layers change performance?
  12. Adaptive test-time scaling
  13. Methodology
  14. Results
  15. Related Work
  16. ...and 12 more sections

Figures (3)

  • Figure 1: Left: Illustration of the proposed architecture (Section \ref{['sec:choosing_config']}). The latent encoder (blue) maps inputs into latent space, the recursive “thinking” block (green) iteratively refines representations, and the latent decoder (red) maps them back to the output space. Each block consists of a different number of layers. Right: Angular distances $d(l, l+1)$ between consecutive layers for OLMo 2 1B base model. The plot highlights three groups of layers—latent encoder, recursive block, and latent decoder—corresponding to distinct trends in layer-to-layer evolution (Section \ref{['sec:choosing_config']}).
  • Figure 2: Results of the ETD method when varying the subset of layers in the recursive block. We report accuracy (Acc.) when increasing the size of the latent encoder $N_{E}$ from 1 to 11 in steps of 2, for each of 6 task categories (as defined in Sec. \ref{['sec:evaluation_metrics']}). The orange line marks selected configuration.
  • Figure 3: Results of fixed-depth ETD with varying numbers of recursive "thinking" iterations compared to adaptive-depth ETD. For fixed-depth ETD, we report accuracy (Acc.) at each iteration count. For adaptive-depth ETD, we report accuracy and the average number of iterations per task.