LaDiR: Latent Diffusion Enhances LLMs for Text Reasoning

2025-10-06

AI Summary

Current large language models (LLMs) rely on autoregressive decoding for chain-of-thought (CoT) reasoning, limiting holistic evaluation and optimization of early reasoning steps and hindering efficient exploration of diverse solution paths. To address this, we propose LaDiR—a latent-space iterative reasoning framework grounded in continuous representations. First, a variational autoencoder (VAE) encodes reasoning steps into structured latent blocks. Second, a block-level bidirectional attention-masked latent diffusion model enables parallel denoising and multi-path generation, facilitating holistic reasoning planning and test-time adaptive refinement. Experiments demonstrate that LaDiR significantly improves accuracy, reasoning diversity, and interpretability over state-of-the-art autoregressive, diffusion-based, and latent reasoning methods on mathematical reasoning and planning benchmarks.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: (Large) Language Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

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.

Problem

• Enhancing LLMs' holistic reasoning refinement capabilities
• Improving efficiency of diverse solution exploration in text
• Unifying latent representation with iterative diffusion refinement

Innovation

• Latent diffusion enables iterative refinement of reasoning
• VAE encodes reasoning steps into compact latent tokens
• Blockwise bidirectional attention allows holistic reasoning revision