Beyond Mode Elicitation: Diversity-Preserving Reinforcement Learning via Latent Diffusion Reasoner

TL;DR

The paper tackles diversity collapse in reinforcement learning for large language model reasoning by shifting exploration from discrete token spaces to continuous latent spaces. It introduces LaDi-RL, which combines Latent Diffusion Reasoning (LaDiR) with Group Relative Policy Optimization (GRPO) to explore semantic-level reasoning trajectories via a diffusion process, while decoupling latent exploration from final text generation with a complementary text policy. A repulsion-based diversity guidance and multi-step denoising preserve multiple solution modes, yielding substantial improvements in pass@1 and pass@k on code generation and math reasoning benchmarks, including absolute gains of +$9.4 ext{ extendash}5.7$ percentage points. The approach also enhances efficiency through latent compression, reducing inference time by about 33% and enabling more scalable reasoning in practical settings. Overall, diffusion-based latent RL provides a principled alternative to token-level RL for robust, diverse, and accurate LLM reasoning.

Abstract

Recent reinforcement learning (RL) methods improve LLM reasoning by optimizing discrete Chain-of-Thought (CoT) generation; however, exploration in token space often suffers from diversity collapse as policy entropy decreases due to mode elicitation behavior in discrete RL. To mitigate this issue, we propose Latent Diffusion Reasoning with Reinforcement Learning (LaDi-RL), a framework that conducts exploration directly in a continuous latent space, where latent variables encode semantic-level reasoning trajectories. By modeling exploration via guided diffusion, multi-step denoising distributes stochasticity and preserves multiple coexisting solution modes without mutual suppression. Furthermore, by decoupling latent-space exploration from text-space generation, we show that latent diffusion-based optimization is more effective than text-space policy optimization alone, while a complementary text policy provides additional gains when combined with latent exploration. Experiments on code generation and mathematical reasoning benchmarks demonstrate consistent improvements in both pass@1 and pass@k over discrete RL baselines, with absolute pass@1 gains of +9.4% on code generation and +5.7% on mathematical reasoning, highlighting diffusion-based latent RL as a principled alternative to discrete token-level RL for reasoning.

Figures (7)

• Figure 1: Average pass@$k$ performance over all tested benchmarks. LaDi-RL preserves solution diversity and converts LaDiR’s pass@$k$ gains into strong pass@$1$ improvements, whereas vanilla GRPO exhibits diversity collapse at large $k$.
• Figure 2: Comparison between standard discrete CoT reasoning and latent continuous CoT reasoning. Left: conventional autoregressive generation produces reasoning directly in token space. Right: reasoning is represented as a fixed-size latent block, enabling continuous modeling and decoupling reasoning from surface text generation.
• Figure 3: Overview of LaDi-RL training pipeline. The latent diffusion policy samples $N$ latent blocks $z^0$, which condition the Text Policy to generate $M$ candidate answers $x_{i,j}$. Outcome rewards evaluated on text are aggregated to compute group-relative advantages for jointly updating both the diffusion and text policies.
• Figure 4: Latent diffusion–based exploration with diversity guidance. Multi-step denoising maps a noisy latent to diverse solutions. Repulsion force pushes trajectories apart in latent space to encourage exploration of diverse reasoning paths.
• Figure 5: pass@$k$ performance on code generation and math reasoning benchmarks across different $k$.