LaDiR: Latent Diffusion Enhances LLMs for Text Reasoning
About
Large Language Models (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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Mathematical Reasoning | GSM8K | -- | 499 | |
| Code Generation | HumanEval+ | Pass@184.2 | 383 | |
| Code Generation | MBPP+ | Pass@159.5 | 216 | |
| Mathematical Reasoning | College | -- | 55 | |
| Code Generation | MBPP | Pass@166.8 | 41 | |
| Mathematical Reasoning | MATH | Pass@146.2 | 19 | |
| Mathematical Reasoning | Gaokao | Pass@135.4 | 18 | |
| Mathematical Reasoning | DM-Math | Pass@152.3 | 18 | |
| Mathematical Reasoning | TheoremQA | Pass@124.7 | 18 | |
| Mathematical Reasoning | Olympia | Pass@112.9 | 18 |