Source-linked AI summary
Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach
Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Tom Goldstein
TL;DR
The paper addresses how to scale test-time reasoning without projecting all intermediate computation into verbalized tokens. It introduces a latent depth-recurrent transformer trained with variable recurrence counts, and reports dramatic reasoning-task improvements from additional test-time computation. The work remains a proof-of-concept with unresolved training and post-training questions.
Problem
The paper addresses how to scale test-time reasoning without projecting all intermediate computation into verbalized tokens.
Method
It introduces a transformer with a shared latent recurrent block that is repeatedly unrolled, with randomly sampled iteration counts during training.
Results
The proof-of-concept model can dramatically improve reasoning-task performance by expending additional test-time computation.
Takeaways & Limitations
Latent recurrent depth provides a way to scale model computation without specialized training data or verbalization at inference time.
Takeaways & Limitations
The models remain a proof-of-concept, with future work needed on optimized training and additional post-training schemes.
Abstract
from arXiv · showhide
We study a novel language model architecture that is capable of scaling test-time computation by implicitly reasoning in latent space. Our model works by iterating a recurrent block, thereby unrolling to arbitrary depth at test-time. This stands in contrast to mainstream reasoning models that scale up compute by producing more tokens. Unlike approaches based on chain-of-thought, our approach does not require any specialized training data, can work with small context windows, and can capture types of reasoning that are not easily represented in words. We scale a proof-of-concept model to 3.5 billion parameters and 800 billion tokens. We show that the resulting model can improve its performance on reasoning benchmarks, sometimes dramatically, up to a computation load equivalent to 50 billion parameters.
1. Scaling by Thinking in Continuous Space
The paper proposes scaling test-time computation through latent recurrent reasoning rather than longer verbalized chains of thought. A depth-recurrent model repeatedly updates hidden states, improving performance as additional computation is applied.
- Test-time computation can be increased by iterating a recurrent block in latent space instead of generating longer chains of thought.The model improves through latent reasoning without verbalizing intermediate calculations.
- 3.5B parameters: the proposed language model uses depth recurrence to scale computation during inference.
- Tasks requiring more reasoning, such as GSM8k, effectively use more recurrent computation than quicker-converging tasks such as OpenBookQA.
- Depth-recurrent language models can learn effectively, train efficiently, and improve significantly when test-time compute is scaled.
- The architecture scales to several billion parameters and over half a trillion pretraining tokens while supporting recurrent reasoning in latent space.
2. Why Train Models with Recurrent Depth?
Latent recurrent reasoning offers an alternative to verbalized long-context reasoning by performing additional computation in continuous hidden states. The paper presents this as a distinct axis for scaling model performance, with potential relevance to nonverbal reasoning.
- Recurrent layers let transformer models perform arbitrarily many computations before emitting each token.
- Latent reasoning can use standard training data with variable compute budgets, avoiding bespoke chain-of-thought demonstrations.
- Latent reasoning models require less memory for training and inference than chain-of-thought reasoning models.
- Latent reasoning may capture spatial thinking, physical intuition, and motor planning that are difficult to verbalize.
- Test-time latent computation is proposed as a third scaling axis alongside extended verbalized inference and pretraining parameter counts.
3. A scalable recurrent architecture
The architecture embeds inputs, repeatedly transforms latent states with a shared recurrent transformer block, and decodes the final state into token probabilities. Training samples variable recurrence depths so inference can unroll the model further, while design choices target stable recurrence at scale.
- 3.1. Macroscopic Design: The architecture comprises a prelude that embeds inputs, a recurrent core that updates latent states, and a coda that produces token probabilities.
- 3.1. Macroscopic Design: Given recurrence count r, the core repeatedly maps the previous latent state and embedded input to a new latent state.
- 3.1. Macroscopic Design: Injecting the embedded input at every recurrence step supports stable iterative operators and convergence independent of initialization.
- 3.2. Microscopic Design: At scale, the selected sandwich normalization is required to train the recurrence, whereas alternative normalization strategies work similarly at small scales.
- 3.2. Microscopic Design: The large model uses shape (2, 4, 2) with hidden size h = 5280; 32 recurrences unfold its 8 real layers to an effective depth of 132 layers.
- 3.3. Training Objective: Training randomly samples recurrence counts and optimizes expected loss over inputs and iteration counts, using a log-normal Poisson distribution with a heavy tail.
4. Training a large-scale recurrent-depth Language Model
The large-scale training run combines a recurrent-depth architecture with a code- and mathematics-heavy dataset, specialized initialization, and distributed training infrastructure. After early failed runs caused by representation collapse or ineffective recurrence, the main configuration trained smoothly for 750B tokens.
- A single medium-scale model was selected after small test models trained reliably up to 10B tokens.
- Pretraining Data: The pretraining mixture emphasized code and mathematical reasoning while retaining enough general webtext for standard language modeling.The data sources were publicly available, and instruction data was mixed directly into pretraining without an ablation of that choice.
- Architecture and Initialization: The scaled architecture used a (2, 4, 2) layer configuration, mean recurrence 32, hidden size 5280, and approximately 3.5B total parameters.The parameter allocation was about 1.5B in the non-recurrent prelude and head, 1.5B in the recurrent core, and 0.5B in tied input embeddings.
- Training Procedure: Locked-step depth sampling synchronized one recurrence depth per micro-batch, improving compute utilization without affecting convergence speed at small scale.The authors note that independent worker scheduling could better approximate the intended expectation at large batch sizes.
- Importance of Norms and Initializations at Scale: The first failed run suffered hidden-state representation collapse because recurrence increased token correlation until every token shared the same hidden state.A second run recovered from early correlation but learned to ignore the incoming recurrent state, so validation perplexity was unchanged from 1 to 32 recurrences.
- Importance of Norms and Initializations at Scale: The main run reverted to the sandwich block and lowered the peak learning rate to 4 × 10−5, then trained smoothly for another 750B tokens.The successful checkpoint is called the main model, Huginn-01254.
5. Benchmark Results
The 3.5B-parameter recurrent model performs competitively on standard benchmarks and improves with additional test-time recurrence, especially on harder reasoning tasks. Its training and compute budgets complicate direct comparisons with fixed-depth models.
- 800B tokens were used to train the recurrent model, compared with 180B tokens for its non-recurrent baseline.
- The model continuously improves up to test-time FLOP budgets equivalent to a standard 50B-parameter fixed-depth transformer.
- The model outperforms the older Pythia series, roughly matches first-generation OLMo-7B on most metrics, and trails later OLMo models trained on larger curated datasets.
- On mathematical reasoning, the model surpasses all evaluated models except the latest OLMo-2 on GSM8K and MATH.
- On coding benchmarks, the model beats other general-purpose open-source models but not dedicated code models trained for several trillion tokens.
- The recurrent model outperforms its matched non-recurrent baseline at 180B tokens, with an especially pronounced advantage on harder tasks such as ARC challenge.
- Test-time performance improves with recurrence, but saturation is task-dependent: easier tasks saturate sooner, while harder tasks benefit from more compute.
- With relevant facts supplied, the recurrent model improves substantially on open-book OBQA, nearly closing its gap to OLMo-2.
6. Recurrent Depth simplifies LLMs
Recurrent depth supports adaptive computation, cache sharing, and speculative decoding without requiring the same specialized adaptations as standard transformers. The model can vary computation across tokens and queries while preserving conversational performance in an initial test.
- Recurrent-depth models naturally support inference methods that standard transformers often require substantial tuning to implement.
- Zero-Shot Adaptive Compute at Test-Time: Zero-shot adaptive computation varies recurrence per token, allowing easy predictions to exit earlier while harder decisions receive more iterations.
- Zero-Shot Adaptive Compute at Test-Time: MMLU categories require different exit depths: high school mathematics exits earlier, while moral scenarios average 3.5 additional steps.
- Using the latest available KV states lets tokens at different recurrent depths attend without recomputing missing hidden states.
- The model can share KV-caches across recurrence iterations with minimal performance impact, reducing cache memory requirements.
- Warm-starting each token from the previous token’s final state reduces average convergence steps by 1–2 and shifts some philosophy exits earlier.
- Recurrent depth enables speculative decoding without a separate draft model by drafting with fewer iterations and later verifying with M > N iterations.
7. What Mechanisms Emerge at Scale in Recurrent-Depth Models
At scale, recurrent computation develops context-dependent latent trajectories rather than a single uniform convergence pattern. Observed dynamics include convergence, orbits, directional drifts, and initialization-independent behavior.
- Trajectory analysis examines latent states across tokens and recurrent iterations using PCA projections of high-dimensional trajectories.
- Convergence depends on context: key question tokens and early response tokens are deliberated more, and identical tokens can behave differently by position.
- Some latent trajectories converge to fixed points, while others trace non-monotonic orbits even though the training objective does not explicitly require them.
- The model uses directional slider-like trajectories for key tokens, which could implement iteration counting.
- The observed latent dynamics include convergent paths, orbits, and drifts that organize computation geometrically rather than as discrete verbal reasoning steps.
- Similar fixed points, orbital patterns, and directional drifts emerge from multiple initial states, indicating path-independent behavior.
8. Related Work Overview
The paper situates latent recurrent depth among recurrence, equilibrium, diffusion, energy-based, optimization, and hybrid architectures. Its distinguishing choices are scalable truncated unrolling and latent computation learned directly through recurrence.
- Depth recurrence has long been studied as a mechanism for adaptive computation and potentially universal machine-like computation.
- The approach relates to deep equilibrium and diffusion models, but differs in training objective: it uses truncated unrolling rather than direct equilibrium solving or a diffusion surrogate.
- Depth-recurrent architectures can also be viewed through connections to latent energy models, implicit optimization layers, and Kuramoto layers.
- Hybrid alternatives include latent subnetworks, LoRA adapters over weight-shared layers, and dynamic weight-tying of trained models.
- Related work differs in whether recurrence is pretrained from scratch or added by fine-tuning fixed-depth transformers, and whether chain-of-thought data is required.
- The paper also connects its motivation to approaches that prioritize reasoning and algorithm learning over memorizing simple patterns.
9. Future Work
The paper identifies several open directions for recurrent-depth models, including post-training methods, integration with other architecture improvements, and alternative memory mechanisms.
- Future work could fine-tune recurrence compression, use reinforcement learning with varying problem hardness, or internalize chain-of-thought reasoning into recurrence.
- The relationship between recurrent depth and modern architecture improvements remains insufficiently explored.
- Recurrent depth could let linear sequence-mixing operators repeat until necessary comparisons between sequence elements are computed.
- The architecture is compute-heavy and has more materialized parameters than actual parameters, paralleling mixture-of-expert models.
- The discussion also connects these considerations to other memory mechanisms proposed for large language models.
10. Conclusions
The paper presents latent recurrent depth as a proof-of-concept approach for scaling reasoning through continuous latent-space computation. It reports promising reasoning performance and emergent inference behaviors, while noting costs for optimization and oversight.
- The models are proof-of-concept systems, but latent reasoning can dramatically improve reasoning-task performance with additional test-time computation.
- The model exhibits context-dependent convergence speed, path independence, and various zero-shot abilities emerging from recurrent training.
- Latent-space reasoning operates without specialized training data or verbalization at inference time.
- The architecture offers promising reasoning performance and may benefit local deployment through its smaller size relative to parameter-scaled models.
- Moving reasoning into latent space may reduce oversight because its trajectories are less human-readable than verbalized chains of thought.
A.1. Classical Reasoning Problems
The appendix examines recurrent-depth reasoning on multi-operand arithmetic and inference-time behavior. Arithmetic performance depends on operand difficulty and improves with additional recurrences in some settings.
- A.1. Classical Reasoning Problems: The study evaluates whether recurrence can use increased test-time computation to solve verbalized addition problems of increasing difficulty.
- A.1. Classical Reasoning Problems: At 32 recurrences, the model reliably solves two-operand additions with operands up to 4 or 5 digits, but rarely solves four- or five-operand single-digit additions.
- A.1. Classical Reasoning Problems: The line charts sweep operand count from 1 to 64 recurrences while fixing the digit count.
- A.1. Classical Reasoning Problems: Table 6 reports first-turn MT-Bench scores for inference-time schemes native to the recurrent-depth model, with no statistically significant differences from the baseline.
A.2. Implementation Details
The appendix documents hardware and distributed-training choices alongside latent-space analyses of convergence, trajectories, and path independence across question types.
- A.2. Implementation Details: Training measures achievable MI250X performance on Frontier using the specified ROCm and PyTorch stack rather than nominal hardware throughput.
- A.2. Implementation Details: The compute-heavy model uses distributed data parallelism across nodes and zero-1 optimizer sharding within nodes with recurrent-step gradient checkpointing.
- A.2. Implementation Details: The implementation targets 512-node allocations and routes RCCL communication through Open Fabrics to achieve sufficient interconnect bandwidth.
- A.2. Implementation Details: Handwritten distributed data parallelism with exactly 64MB inter-node packages resolves hangs at 512 nodes, while achieved throughput varies across allocations.
- A.2. Implementation Details: Pretraining tracks optimizer, parameter, gradient, hidden-state, and validation-perplexity metrics across recurrent depths from 1 to 64 steps.
- A.2. Implementation Details: The latent-space appendix visualizes six PCA directions in grouped pairs and provides single-token analyses using the first 40 directions.
- A.2. Implementation Details: Latent trajectories separate system prompts and display question-dependent patterns, including rotation for a math-relevant number and convergence for trivia tokens.
- A.2. Implementation Details: Unsafe-question trajectories contain anomalous key tokens, while convergence plots examine distances to approximate steady states at r = 128 iterations.
C. Pretraining Data
The paper’s pretraining data are organized into standard sources and instruction data.
- Standard sources are listed as one part of the datasets used for model pre-training.
- Instruction data are listed as a separate part of the datasets used for model pre-training.