LaDi-RL: Latent Diffusion Reasoning Prevents Entropy Collapse in Reinforcement Learning

TL;DR

LaDi-RL introduces a diffusion-based latent reasoning approach in reinforcement learning to improve global reasoning capabilities in language models, addressing the limitations of token-level optimization.

Contribution

It proposes a hierarchical latent-text rollout method that enhances reward estimation for latent diffusion policies in RL, leading to improved reasoning performance.

Findings

Outperforms token-level RL by 9.4% on code generation

Achieves 5.7% improvement on math reasoning

Surpasses base model pass@k performance

Abstract

Reinforcement learning has become a central paradigm for improving LLM reasoning, but most existing methods optimize policies over discrete token sequences. This creates a mismatch between the optimization space and the structure of reasoning: many important decisions are semantic, global, and trajectory-level rather than local token choices. Continuous latent-space RL offers a promising alternative by allowing policies to explore higher-level reasoning representations. However, simply moving to latent space is not sufficient. The resulting policy must model a complex, multi-modal distribution over valid reasoning trajectories. We therefore propose Latent Diffusion Reasoning with Reinforcement Learning (LaDi-RL), where a diffusion model generates latent reasoning trajectories through iterative denoising. This formulation enables structured exploration and expressive distribution modeling, but also introduces a fundamental credit-assignment challenge: the policy acts in latent space, while rewards are observed only after the latent is decoded into text. A naive rollout strategy therefore entangles latent reasoning quality with text decoding quality, making it unclear whether an incorrect answer results from a poor latent trajectory or from an imperfect textual realization. To address this, we introduce hierarchical latent-text rollouts. We sample multiple text completions for each latent trajectory and aggregate their rewards to obtain a decoder-marginalized estimate of latent utility. This provides a cleaner and lower-variance reward signal for optimizing the diffusion policy. Empirically, LaDi-RL outperforms token-level RL by 9.4% on code generation and 5.7% on math reasoning in pass@1, and even surpasses the base model's pass@k performance.