LaDiR: Latent Diffusion Enhances LLMs for Text Reasoning
TL;DRQuick summary
LaDiR models intermediate reasoning with latent diffusion. A learned latent space represents reasoning steps, and a diffusion model generates them through joint, iterative updates within each block. This gives us control over both refinement and exploration at inference time. The crucial training lesson: these generated states must be useful for reaching the answer, so we train the model on its own latent trajectories.
Why diffusion for reasoning?
The appeal of diffusion is the kind of generation process it gives us. A sample starts as noise and takes shape through repeated updates. The model can coordinate several parts of that sample at once, and the sampling process offers ways to control how much computation and exploration go into producing it.
For LaDiR, we wanted to bring these properties to intermediate reasoning. What would it mean to generate a thought through diffusion?
That question requires two choices. We need a space in which to represent reasoning, and a model of how to generate states in that space. Continuous vectors provide a representation; diffusion supplies a learned process that turns noise into a candidate reasoning step. The combination lets us study refinement, diverse sampling, and inference-time compute within the same framework.
The interesting challenge is making that process useful for solving a problem. A plausible intermediate state still has to support the steps that follow and ultimately lead to a correct answer.
A latent space that diffusion can work with
In LaDiR, a variational autoencoder maps a written reasoning step into a small block of latent tokens. A diffusion reasoner generates each block through repeated denoising, updating its tokens together. Completed blocks provide context for later blocks; the final answer is generated in text. The decoder lets us inspect a block as language. Method.
The block is the key compromise. It gives the model a local workspace where parts of a tentative step can influence one another, while preserving an ordered sequence of reasoning steps. Earlier completed blocks stay fixed, so the scope of revision is deliberately local.
Why use a learned latent space? Consider changing the quantity an equation refers to. In text, that change may require coordinated edits to a variable, an operation, and its explanation. We want a representation in which those pieces can change together. That is the design motivation; continuous vectors alone do not guarantee meaningful edits.
This makes the latent space a demanding interface. It must retain useful information, be predictable by the reasoner, and remain useful after prediction errors. Reconstructing a clean training example tests only part of that contract.
The model has to live with its own thoughts
That last requirement is easy to underestimate. During teacher-forced training, the reasoner sees latent states encoded from correct solutions. At inference, it sees states it generated itself. Even small deviations can change what a later step assumes.
LaDiR therefore includes rollout training: the model generates its own latent trajectory, and the answer loss backpropagates through denoising while the flow-matching objective remains active. In the LLaMA 3.1 8B math experiments, removing this stage lowers average Pass@1 across seven benchmarks from 43.5% to 27.9%. Training and ablation.
Our reading of this result is that a reasoning representation must be evaluated through its downstream use. A latent can encode a good explanation and still be a poor state for the model to continue from. Rollout training puts the producer and consumer of that state into the same loop: the model has to answer using what it actually produced.
For us, this is one of the most useful lessons from the project. When designing latent reasoning systems, reconstruction quality deserves attention, but the decisive question is what happens after the model takes over generation.
Denoising gives us another way to spend compute
Once a thought can be updated, extra inference compute has another destination. We can spend more denoising steps refining a block, alongside spending compute on longer trajectories or more candidates. This creates an allocation problem: when is an uncertain step worth more refinement, and when is another attempt more valuable?
Exploration introduces a related tension. LaDiR uses diversity guidance to push sampled latent trajectories apart. In our Countdown ablations, moderate guidance balances accuracy and diversity; excessive guidance reduces accuracy. Ablation.
We interpret this as a constraint on useful search. Making candidates different is easy if we allow them to become invalid. The valuable diversity is disagreement among plausible ways of solving the problem. A good inference procedure must create room for alternatives while keeping them connected to what the model knows how to solve.
For us, LaDiR points toward treating the reasoning process itself as a generative modeling problem. The latent space defines what the model works with; diffusion defines how it generates and improves candidate states; answer supervision connects those states to the task. The next question is how to allocate denoising and exploration where they make the largest difference to solving the problem.