Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach

Tom Goldstein (University of Maryland) · Bhavya Kailkhura (Lawrence Livermore National Laboratory) · Brian Bartoldson (Lawrence Livermore National Laboratory) · Jonas Geiping (ELLIS Institute & MPI Intelligent Systems, Tübingen AI Center) · Sean McLeish (University of Maryland) · Neel Jain (University of Maryland, College Park) · John Kirchenbauer (University of Maryland, College Park) · Siddharth Singh (University of Maryland, College Park) · Abhinav Bhatele (University of Maryland, College Park)
adaptive computearbitrary depth unrollingchain-of-thought approachescompute cost reductionkv-cache sharinglanguage model architecturelatent space reasoningmainstream reasoning modelsnon-verbal reasoning representationparameter scalingproof-of-concept modelreasoning tasksrecurrent block iterationsmall context windowsspecialized training dataspeculative decoding

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 train a proof-of-concept model from scratch with 3.5 billion parameters and 800 billion tokens. We show that this model can effortlessly use varying levels of compute, significantly improving with additional compute especially on reasoning tasks, such as math and coding. Further, this architecture naturally reduces compute costs via zero-shot per-token adaptive compute, KV-cache sharing and speculative decoding.