Encode, Think, Decode: Scaling test-time reasoning with recursive latent thoughts
This addresses the need for more efficient reasoning in LLMs without scaling parameters or data, offering a domain-specific improvement for AI applications.
The paper tackled the problem of enhancing reasoning capabilities in large language models by introducing Encode-Think-Decode (ETD), a method that iterates over a small subset of reasoning-relevant layers during training and inference, resulting in substantial gains such as +28.4% relative accuracy improvement on GSM8K and +36% on MATH with the OLMo-2 1B Base model.
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.