LaDiR: Latent Diffusion Enhances LLMs for Text Reasoning

TL;DR

LaDiR keeps the LLM backbone but replaces token-level chain-of-thought with diffusion in a continuous latent reasoning space. A β-VAE (encoder fine-tuned from the LLM, frozen LLM as decoder) compresses each sentence of CoT into a block of latent “thought tokens”; a flow-matching model — same pretrained LLM, hybrid attention (bidirectional within a block, causal across blocks) — denoises these blocks one at a time, then switches to autoregressive text decoding to emit the final answer. The two-stage training (teacher-forcing on oracle latents, then rollout training on self-generated latents to fix exposure bias) closes the gap between training and inference. The headline pitch is iterative refinement at the semantic level rather than the token level, plus parallel diverse-trajectory exploration via increased initial noise and a Stein-style repulsion guidance during denoising.

Key claims

LaDiR (LLaMA-3.1 8B base) outperforms AR CoT SFT, iCoT, Pause Token, Coconut, and Discrete Latent on average across 7 math benchmarks (41.8 vs 40.4 / 31.9 / 40.8), with the biggest gains on out-of-domain TheoremQA (22.9 vs 21.2 CoT, 11.4 Coconut) and College-Math (48.6 vs 45.7 CoT) [Table 1].
On Countdown planning, LaDiR improves Pass@1 by >25 points over LLaMA 8B SFT on CD-4 (76.6 vs 46.7) and ~30 points on CD-5 (38.5 vs 8.9), with best-in-class Pass@100 (96.4 / 75.2) and diversity (7.3 / 8.9 unique solutions per 100 samples) [Table 2].
Stage-2 rollout training is load-bearing: removing it drops average math accuracy from 41.8 to 32.6, isolating the exposure-bias fix as the main quality lever beyond the architectural choices [Table 1, “–w/o Stage 2”].
Adaptive test-time compute holds: scaling denoising steps from 5 → 10 yields +11.7 avg, 10 → 30 adds +4.8, 10 → 50 adds +9.8 across the 7 math benchmarks — diffusion-steps act as a continuous knob for the accuracy/compute trade-off [§4.4, Fig. 5].
Iterative refinement is observable at the per-step level on Countdown: early decodes contain arithmetic errors that are progressively corrected, with the equation stabilizing several steps before the final denoising step — analogous to the self-correction behavior of long-CoT reasoning models, but executed in latent space [Table 3].
Parallel diversity via repulsion guidance + increased initial noise is a tunable Pareto curve: noise scale 2 + moderate repulsion improves both diversity and best-of-100 accuracy; excessive noise (scale 3) or strong repulsion harms accuracy despite higher diversity [§4.3, Table 4].
The flow-matching objective beats DDPM-style diffusion at the same architecture and budget — the authors adopt it as default after appendix ablations [§3.3, Appendix C].
VAE robustness augmentations (input-token substitution + latent-Gaussian-noise injection at σ ≈ 0.3) are required to make the latent space usable for downstream diffusion — without them the encoder overfits to lexical surface form [§3.2.1].

Method

The pipeline has two trained components on the same LLM backbone (LLaMA-3.1 8B). A β-VAE encodes each sentence of a chain-of-thought into a block of k continuous latent thought tokens (one sentence ↔ one block); its encoder is a full-parameter fine-tune of the base LLM with two linear heads producing (μ, σ), the decoder is the frozen base LLM conditioned on the sampled latent. To prevent the encoder from memorizing lexical form, training mixes input-token substitution (uniform vocabulary swap with probability p) and isotropic latent-noise injection at σ ≈ 0.3.

The reasoning model — initialized from the same LLM — denoises latent blocks under a flow-matching objective. The attention mask is the architectural trick: bidirectional within a block (the model “reasons” over the block as a whole, like a discrete diffusion LM but in continuous latents), strictly causal across blocks (later reasoning steps depend on earlier ones, preserving variable-length open-ended generation). Special tokens <BOT> / <EOT> / <SOA> delimit blocks and the transition from latent reasoning to text answering; a small binary classification head on <EOT> decides whether the next block continues reasoning or begins the answer. Stage 1 trains under teacher forcing on oracle VAE latents jointly with cross-entropy on the answer text and the special-token head. Stage 2 (“rollout training”) replaces oracle latents with self-generated ones (fewer denoising steps, gradients flowing through the trajectory) to close the train-inference gap that otherwise causes error accumulation. At inference, the model alternates iterative denoising on each latent block and the special-token head’s stop decision; once <SOA> fires, it switches to autoregressive text decoding for the final answer. Diversity is amplified by sampling the initial latent from a higher-variance Gaussian and adding a CFG-style repulsion term ∇ log Σ exp(−‖z − z'‖² / h) between latents in the same batch, with the guidance scale decayed linearly across denoising steps.

Results

On math (LLaMA-3.1 8B base, DART-MATH training set), LaDiR averages 41.8 pass@1 across MATH / GSM8K / Gaokao / DeepMind-Math / College / Olympiad / TheoremQA, vs 40.4 for AR CoT SFT, 40.8 for Discrete Latent VQ-VAE, 31.9 for Coconut, and 35.8 for LLaDA-8B CoT-SFT (a from-scratch diffusion LM) [Table 1]. The biggest absolute gains are out-of-domain (College +2.9, TheoremQA +1.7) — the authors argue this is because latent-space reasoning learns more abstract patterns than token-level CoT. On Countdown (combinatorial arithmetic puzzles, 500k training examples, 10% target numbers held out as OOD test), LaDiR reaches 76.6 Pass@1 / 96.4 Pass@100 on CD-4 vs 46.7 / 65.3 for LLaMA 8B SFT and 51.2 / 75.2 for LLaDA 8B SFT, and 38.5 / 75.2 on CD-5 vs 8.9 / 15.4 and 34.4 / 45.2 respectively [Table 2]. Diversity (unique valid solutions among 100 samples) is also the best across all baselines (7.3 / 8.9 vs ≤5.4 / ≤6.2). The Pass@k curve crosses the task-specific small-diffusion-model baseline (MGDM) at large k, indicating the diverse-trajectory mechanism is the load-bearing one for plan-style problems. Adaptive denoising-step scaling gives a clean accuracy-vs-compute knob (Fig. 5).

Why it’s interesting

LaDiR sits at the intersection of Diffusion Language Models and continuous-latent CoT — it’s the first paper on the wiki to push latent diffusion (not discrete masked diffusion) into LLM reasoning. The closest filed neighbors are Mercury: Ultra-Fast Language Models Based on Diffusion and Introspective Diffusion Language Models, which use discrete diffusion over token vocabulary for throughput; LaDiR uses continuous diffusion over a VAE latent for iterative semantic refinement, treating diffusion-steps as a test-time-compute knob rather than a parallel-decoding speedup. It complements Introspective Diffusion Language Models by occupying the opposite point on the substrate axis: ISD shows AR-base + discrete-diffusion can match AR quality with throughput wins; LaDiR shows AR-base + continuous-latent-diffusion can beat AR quality on reasoning at a comparable inference cost. The blockwise bidirectional/causal-across hybrid mask is structurally the same trick as in Block Diffusion (Arriola et al.) and shows up again in the dLLM line.

The Coconut comparison is the load-bearing one for the broader latent-reasoning thesis: prior latent-CoT (Coconut, Discrete Latent VQ-VAE) only sometimes beats AR CoT SFT; LaDiR’s gain over AR comes specifically from the diffusion objective on top of a VAE latent, not from the latent substrate alone. This is also relevant to Reasoning RL as a non-RL recipe for improving reasoning — the gains here come from architectural CoT (latent diffusion) rather than procedural CoT (GRPO on text rollouts). The Countdown diversity result connects to the planning-as-parallel-exploration thread in DiffThinker: Towards Generative Multimodal Reasoning with Diffusion Models: both papers report multiple-candidate-then-prune dynamics in the early denoising steps, but DiffThinker exploits it for visual-puzzle reasoning while LaDiR exploits it for text-space planning via explicit repulsion guidance.