💬 NLP

LaDiR: Latent Diffusion Enhances LLMs for Text Reasoning

The paper proposes LaDiR, a novel reasoning framework that enhances Large Language Models by encoding text reasoning into a structured latent space and utilizing a latent diffusion model with bidirectional attention to enable holistic, iterative refinement and parallel generation of diverse reasoning trajectories, thereby outperforming existing autoregressive and diffusion-based methods on mathematical and planning benchmarks.

Original authors: Haoqiang Kang, Yizhe Zhang, Nikki Lijing Kuang, Nicklas Majamaki, Navdeep Jaitly, Yi-An Ma, Lianhui Qin

Published 2026-03-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haoqiang Kang, Yizhe Zhang, Nikki Lijing Kuang, Nicklas Majamaki, Navdeep Jaitly, Yi-An Ma, Lianhui Qin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

🧠 The Problem: The "One-Way Street" of Current AI

Imagine a Large Language Model (LLM) like a student taking a math test. Currently, these students use a method called Chain-of-Thought (CoT). They write down their thinking step-by-step, from left to right, just like you write a sentence.

The Catch: Once the student writes a word, they can't go back and erase it easily. If they make a mistake in step 3, they have to keep writing step 4, 5, and 6 based on that wrong idea. It's like driving a car on a one-way street where you can't turn around. If you take a wrong turn, you have to drive all the way to the end, realize you're lost, and start the whole trip over.

This makes it hard for AI to:

Fix mistakes once they happen.
Explore different solutions (they tend to get stuck in the same "rut").

🚀 The Solution: LaDiR (The "Sculptor" Approach)

The authors propose a new method called LaDiR (Latent Diffusion Reasoner). Instead of writing words one by one, LaDiR thinks in concepts and shapes first, then turns them into words.

Think of it like sculpting a statue from clay instead of typing a story.

1. The Clay (Latent Space)

First, the AI compresses its thoughts into "clay." In the paper, this is called Latent Tokens.

Old Way: Writing "The cat sat on the mat." (Discrete words).
LaDiR Way: Holding a lump of clay that feels like "a cat sitting on a mat." It's a smooth, continuous blob of meaning, not a string of letters.

2. The Sculpting Process (Diffusion)

This is where the magic happens. The AI starts with a lump of random noise (like a messy pile of clay). It then slowly "denoises" it, refining the shape step-by-step until it looks like a perfect statue.

The Analogy: Imagine you are trying to draw a perfect circle.
- Old AI (Autoregressive): You draw one dot, then another, then another. If you mess up the first dot, the whole circle is crooked. You can't fix the first dot without erasing the whole page.
- LaDiR (Diffusion): You start with a rough, messy scribble. You look at the whole scribble and say, "Hmm, that part is too high, let's push it down. That part is too low, let's pull it up." You can fix any part of the drawing at any time while you are working on it.

3. The "Block" Strategy

To make this work for long reasoning tasks, LaDiR breaks the problem into blocks (like paragraphs).

It refines one "thought block" at a time.
Inside the block, it can look back and forth (bidirectional) to fix errors.
Once a block is perfect, it moves to the next one.

🎨 Why is this better? (The Superpowers)

1. The "Do-Over" Button (Iterative Refinement)

Because LaDiR works on the "clay" (concepts) rather than the "words," it can fix a mistake in the middle of a thought without deleting everything that came after.

Analogy: If you are writing a recipe and realize you forgot the salt, you can just add it. If you are using LaDiR, you can "smooth out" the concept of "salty soup" in your mind before you even write the word "salt."

2. The "Multiverse" Explorer (Diversity)

Current AI models often give you the exact same answer every time because they follow the most obvious path.

LaDiR's Trick: The paper uses a "repulsion force." Imagine you are in a room with 10 other people trying to find an exit. If everyone clumps together, you might all get stuck. LaDiR pushes the different "thought paths" apart, forcing them to explore different corners of the room.
Result: It generates many different, diverse solutions to the same problem, increasing the chance that one of them is the correct answer.

3. The "Translator" (Interpretability)

Since the AI thinks in "clay" (latent space) but needs to talk to humans, it uses a VAE (Variational Autoencoder) as a translator.

It takes the smooth clay shape, translates it back into human-readable text, and shows you the steps. This makes the AI's "thinking process" visible and understandable, unlike some "black box" methods.

🏆 The Results: Did it work?

The authors tested LaDiR on three tough challenges:

Math Problems: It solved more complex math problems than standard AI, especially when allowed to "think longer" (more refinement steps).
Coding: It wrote better code, fixing logic errors that other models missed.
Puzzle Planning (Countdown): In a game where you have to combine numbers to reach a target, LaDiR was 30% better at finding solutions and found many more different ways to solve the same puzzle.

📝 The Bottom Line

LaDiR changes how AI thinks. Instead of rushing to write the next word (like a frantic typist), it pauses to sculpt its ideas (like a thoughtful artist). It allows the AI to look at the whole picture, fix mistakes in the middle, and explore many different paths to the solution, leading to smarter, more creative, and more accurate reasoning.

1. Problem Statement

Large Language Models (LLMs) currently rely on autoregressive (AR) decoding for Chain-of-Thought (CoT) reasoning. This paradigm has two fundamental limitations:

Inability to Self-Refine: AR models generate tokens sequentially from left to right. Once a token is generated, it cannot be revised in a holistic manner. This makes self-correction inefficient and difficult, often leading to error propagation.
Limited Diversity: AR sampling tends to collapse into repetitive or similar reasoning trajectories, restricting the exploration of diverse valid solutions.

While Diffusion Models offer iterative refinement and global coherence, applying them directly to text (discrete tokens) is challenging. Existing latent diffusion approaches for text generation often focus on fluency rather than the complex causal dependencies required for multi-step reasoning.

2. Methodology: LaDiR (Latent Diffusion Reasoner)

LaDiR proposes a novel framework that unifies the expressiveness of continuous latent representations with the iterative refinement capabilities of latent diffusion models. It separates the reasoning process from the final answer generation.

A. Core Architecture

Latent Space Construction (VAE):
- A Variational Autoencoder (VAE) is used to encode intermediate reasoning steps (CoT) into a continuous latent space.
- Blockization: The reasoning chain is split into sentences, where each sentence is treated as a "block" of latent tokens ( $Z^{(b)}$ ).
- Encoder: Initialized from a pretrained LLM, it maps text blocks to latent distributions ( $\mu, \sigma$ ).
- Decoder: A frozen pretrained LLM that reconstructs text from latent blocks.
- Robustness: The VAE is trained with augmentations (Gaussian noise injection and token substitution) to ensure the latent space is robust to perturbations and captures semantic content rather than exact lexical forms.
Reasoning Model (Latent Diffusion):
- Instead of generating text tokens directly, the model generates latent thought tokens using a Flow Matching objective (a variant of diffusion).
- Blockwise Diffusion: Diffusion occurs at the block level. Within a block, tokens attend bidirectionally (allowing internal refinement); across blocks, attention is strictly causal (maintaining the order of reasoning steps).
- Hybrid Attention: The model uses a special token <BOT> (Begin of Thought) and <EOT> (End of Thought) to delimit blocks, and <SOA> (Start of Answer) to signal the end of reasoning.
Training Pipeline:
- Stage 1 (Teacher-Forcing): The model is trained to predict latent blocks given the "oracle" latents provided by the VAE encoder. It optimizes a joint loss: Flow Matching (for latents) + Cross-Entropy (for final answer) + Special Token Classification (to predict <SOA> or <BOT>).
- Stage 2 (Rollout Training): To address the train-test mismatch (where the model must use its own generated latents), the model generates latents from noise with fewer denoising steps. Gradients are backpropagated through the trajectory, allowing the answer supervision to directly shape latent predictions and mitigate error accumulation.
Inference & Diversity:
- Iterative Denoising: The model starts from Gaussian noise and iteratively refines latent blocks until the <SOA> token is predicted.
- Diversity Guidance: To explore multiple solutions, the framework employs:
  - Increased Initial Noise: Broadens the starting distribution.
  - Repulsion Force: A gradient-based repulsion term pushes latent trajectories apart within a batch, preventing mode collapse and encouraging exploration of distinct reasoning paths.

3. Key Contributions

Semantic-Level Reasoning: LaDiR shifts reasoning from the discrete token level to a continuous semantic latent space, enabling the model to refine entire reasoning steps (e.g., correcting an arithmetic logic) rather than just individual words.
Iterative Self-Correction: Unlike AR models, LaDiR allows for holistic refinement of reasoning chains during the denoising process, significantly improving accuracy on complex tasks.
Adaptive Compute: The framework allows for a flexible trade-off between accuracy and test-time compute. Adding more denoising steps consistently improves performance, offering a new paradigm for "test-time scaling."
Enhanced Diversity: The repulsion mechanism enables the generation of diverse, valid reasoning trajectories in parallel, outperforming AR models that often collapse to similar solutions.
Interpretability: By decoding latent blocks back to text via the VAE, the reasoning process remains transparent and human-readable, unlike "black box" continuous reasoning methods.

4. Experimental Results

The authors evaluated LaDiR on mathematical reasoning, code generation, and puzzle planning benchmarks, comparing it against AR baselines, masked diffusion models, and other latent reasoning methods.

Mathematical Reasoning (7 Benchmarks):
- Using LLaMA 3.1 8B, LaDiR achieved the highest average accuracy, surpassing the best prior latent reasoning method (TaH+) by 1.5% in Pass@1.
- It showed a 6.1% absolute gain in Pass@100 over AR CoT SFT, demonstrating superior diversity.
- It significantly outperformed specialized latent methods like Coconut and CODI, particularly on out-of-domain benchmarks (e.g., College-Math, DM-Math).
Code Generation (4 Benchmarks):
- On HumanEval+, LaDiR surpassed the AR SFT baseline by nearly 8% and the training-free Soft Thinking method by a significant margin.
- It established new state-of-the-art results on HumanEval and HumanEval+, outperforming specialized code models like Ouro 2.6B by over 10%.
Puzzle Planning (Countdown):
- On the Countdown task, LaDiR showed over 30% absolute improvement in Pass@1 and Pass@100 compared to AR baselines.
- It demonstrated the ability to explore diverse global planning paths, with Pass@k curves rising steeply as $k$ increased.
Ablation Studies:
- Flow Matching was found superior to standard DDIM objectives ( $\epsilon$ , $x_0$ , $v$ prediction).
- Block Size: 1 sentence per block with 4 latent tokens offered the best balance.
- Diversity: Moderate repulsion strength ( $\gamma_{max} \approx 0.3-0.5$ ) provided the best trade-off between diversity and accuracy.

5. Significance

LaDiR represents a paradigm shift in LLM reasoning. By decoupling reasoning from token generation and utilizing latent diffusion, it overcomes the sequential constraints of autoregressive models.

Practical Impact: It offers a principled way to balance accuracy and diversity, crucial for applications requiring robust problem-solving (e.g., math, coding, planning).
Theoretical Insight: The results suggest that modeling reasoning at the semantic level (via continuous latents) yields more faithful intermediate steps than token-level manipulation, leading to stronger final answers.
Future Direction: The ability to flexibly allocate test-time compute via denoising steps opens new avenues for adaptive inference policies, potentially making LLMs more capable of "System 2" (deliberate) thinking.