LaDiR: Latent Diffusion Enhances LLMs for Text Reasoning (2510.04573v2)

Abstract: LLMs demonstrate their reasoning ability through chain-of-thought (CoT) generation. However, LLM's autoregressive decoding may limit the ability to revisit and refine earlier tokens in a holistic manner, which can also lead to inefficient exploration for diverse solutions. In this paper, we propose LaDiR (Latent Diffusion Reasoner), a novel reasoning framework that unifies the expressiveness of continuous latent representation with the iterative refinement capabilities of latent diffusion models for an existing LLM. We first construct a structured latent reasoning space using a Variational Autoencoder (VAE) that encodes text reasoning steps into blocks of thought tokens, preserving semantic information and interpretability while offering compact but expressive representations. Subsequently, we utilize a latent diffusion model that learns to denoise a block of latent thought tokens with a blockwise bidirectional attention mask, enabling longer horizon and iterative refinement with adaptive test-time compute. This design allows efficient parallel generation of diverse reasoning trajectories, allowing the model to plan and revise the reasoning process holistically. We conduct evaluations on a suite of mathematical reasoning and planning benchmarks. Empirical results show that LaDiR consistently improves accuracy, diversity, and interpretability over existing autoregressive, diffusion-based, and latent reasoning methods, revealing a new paradigm for text reasoning with latent diffusion.

• The paper introduces LaDiR, a novel framework that enhances LLM reasoning by leveraging latent diffusion to improve both accuracy and solution diversity.
• It employs block-wise latent reasoning with a VAE and flow-matching objective to enable semantic-level iterative refinement of thought tokens.
• LaDiR achieves superior performance on mathematical reasoning benchmarks and planning tasks through parallel exploration and transparent interpretability of intermediate steps.

Latent Diffusion Reasoning for Enhanced Text Reasoning in LLMs

Motivation and Paradigm Shift

The paper introduces LaDiR (Latent Diffusion Reasoner), a framework that augments the reasoning capabilities of LLMs by leveraging latent diffusion in a structured semantic space. Traditional autoregressive (AR) LLMs generate reasoning chains sequentially, which constrains their ability to revise earlier steps and limits exploration of diverse solutions. Discrete diffusion LLMs, while enabling parallel refinement, operate at the token level and lack semantic abstraction. LaDiR addresses these limitations by reasoning over continuous latent thought tokens, enabling iterative semantic refinement and parallel exploration of diverse reasoning trajectories.

Figure 1: Comparison of reasoning paradigms: AR CoT and latent CoT generate tokens sequentially; diffusion LMs refine masked tokens in parallel; LaDiR reasons via latent diffusion over thought tokens, enabling semantic-level refinement and diverse solution exploration.

Block-wise Latent Reasoning Framework

LaDiR constructs a structured latent reasoning space using a Variational Autoencoder (VAE) that encodes text reasoning steps into blocks of latent thought tokens. Each block corresponds to a semantically meaningful reasoning step, typically a sentence, and is delimited by special tokens (<BOT>, <EOT>). The VAE encoder is initialized from a pretrained LLM and fine-tuned to produce mean and variance vectors for each block, from which latent tokens are sampled. The decoder, a frozen LLM, reconstructs the original text from these latent tokens, ensuring semantic alignment and interpretability.

Figure 2: Block-wise latent reasoning: a question $Q$ is input as condition to generate latent blocks; each block is iteratively denoised with bidirectional attention inside the block and causal attention across blocks.

The reasoning model, also based on a pretrained LLM, predicts and refines latent blocks via latent diffusion. Within each block, bidirectional attention enables holistic reasoning over the block's horizon, while causal attention across blocks maintains autoregressive dependencies. The reasoning process terminates upon emission of a <SOA> token, after which the answer is generated autoregressively conditioned on the denoised latent sequence.

Training and Inference Procedures

LaDiR training proceeds in two stages:

1. VAE Training: The VAE is trained to encode reasoning steps into robust latent representations. Robustness is enhanced via latent Gaussian noise injection and input token substitution, promoting invariance to semantic perturbations and lexical noise.
2. Reasoning Model Training: The reasoning model is trained to denoise latent blocks using a flow-matching objective, which learns the continuous velocity field of the denoising process. Answer tokens are predicted autoregressively conditioned on the latent blocks, with cross-entropy supervision. A special token classification head controls block transitions.

Stage 1 employs teacher-forcing with oracle latents; Stage 2 introduces rollout training, where the model generates its own latents from noise, mitigating error accumulation and improving generalization.

Inference involves iterative denoising of latent blocks, with diversity promoted via increased initial noise scale and a repulsion gradient that disperses latent trajectories within a batch. This enables parallel generation of diverse reasoning paths, enhancing coverage of alternative solutions.

Figure 3: Effect of initial noise scale $\tilde{\sigma}$ on diversity and accuracy during inference.

VAE Architecture and Interpretability

The VAE encoder integrates original text embeddings and learnable latent embeddings, projecting hidden states to mean and variance vectors for latent sampling. The decoder reconstructs text from latent tokens, preserving interpretability and semantic fidelity.

Figure 4: VAE architecture: encoder (finetuned LLM) produces mean/variance vectors from text and latent embeddings; decoder (frozen LLM) reconstructs text from sampled latent thought tokens.

This design ensures that each latent block can be decoded to human-readable text, facilitating transparent inspection of the reasoning process and enabling semantic-level analysis of intermediate steps.

Empirical Results and Analysis

LaDiR is evaluated on seven mathematical reasoning benchmarks (GSM8K, MATH, College-Math, DeepMind-Math, OlympiaBench-Math, TheoremQA, Fresh-Gaokao-Math-2023) and the Countdown planning task. Key findings include:

• Mathematical Reasoning: LaDiR achieves the highest average pass@1 accuracy, outperforming AR CoT, latent, and diffusion baselines. Notably, it improves accuracy on challenging benchmarks (DM-Math, College-Math) where token-level reasoning struggles with long-horizon consistency. The latent diffusion objective yields stronger generalization to out-of-domain tasks.
• Planning (Countdown): LaDiR delivers over 30% absolute improvement in Pass@1 and Pass@100 on Countdown compared to AR and diffusion baselines, with the highest diversity of valid solutions. The pass@k curve rises steeply, indicating effective exploration of reasoning trajectories.
• Iterative Refinement: The flow-matching model progressively corrects reasoning errors across denoising steps, stabilizing partial results and converging to correct solutions.
• Adaptive Test-Time Compute: Increasing the number of denoising steps at inference consistently improves accuracy, enabling flexible trade-offs between compute and performance.

Ablation and Architectural Insights

Ablation studies reveal:

• Latent Prediction Objective: Flow-matching yields superior accuracy compared to MSE, $x_0$, $\epsilon$, and $v$-prediction objectives.
• Block Size: Optimal reasoning accuracy is achieved with 6 latent tokens per block; excessive tokens introduce redundancy and degrade performance.
• Diversity Guidance: Moderate repulsion scales maximize diversity without sacrificing accuracy; excessive repulsion destabilizes reasoning.

Interpretability and Semantic Reasoning

LaDiR's latent blocks are decodable to interpretable text, enabling inspection of intermediate reasoning steps. The denoising process operates at the semantic level, refining logical relations and causal dependencies rather than merely correcting lexical errors. This contrasts with SFT baselines, which may produce syntactically fluent but semantically brittle reasoning chains.

Implications and Future Directions

LaDiR establishes a new paradigm for text reasoning in LLMs by unifying continuous latent representation with iterative refinement via diffusion. The framework offers:

• Flexible accuracy–compute trade-off: Iterative denoising enables adaptive allocation of compute to harder queries.
• Parallel and diverse exploration: Batch-level diversity guidance facilitates coverage of multiple solution strategies.
• Enhanced interpretability: Structured latent space allows transparent analysis of reasoning steps.

Potential future developments include integration with reinforcement learning for post-training, dynamic policies for adaptive compute allocation, and extension to multimodal reasoning tasks. The semantic-level reasoning enabled by LaDiR may inform the design of more robust, generalizable, and interpretable AI systems.

Conclusion

LaDiR demonstrates that latent diffusion over structured thought tokens substantially enhances the reasoning capabilities of LLMs, outperforming AR and diffusion baselines in both accuracy and diversity. The framework's semantic abstraction, iterative refinement, and interpretability mark a significant advance in the modeling of reasoning processes, with broad implications for the development of next-generation AI systems.