🤖 machine learning

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

The paper proposes LaDi-RL, a reinforcement learning framework that utilizes latent diffusion models to generate structured reasoning trajectories and employs hierarchical latent-text rollouts to decouple latent planning quality from text decoding errors, thereby preventing entropy collapse and significantly outperforming traditional token-level RL on code and math reasoning tasks.

Original authors: Haoqiang Kang, Yizhe Zhang, Nikki Lijing Kuang, Yi-An Ma, Lianhui Qin

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haoqiang Kang, Yizhe Zhang, Nikki Lijing Kuang, Yi-An Ma, Lianhui Qin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very smart but slightly stubborn student (a Large Language Model) how to solve complex puzzles, like writing computer code or solving advanced math problems.

For a long time, the best way to teach this student has been Reinforcement Learning (RL). Think of this like a game where the student tries to solve a problem, and if they get it right, they get a gold star. If they get it wrong, they get a "try again."

However, the old way of teaching has a big flaw. It forces the student to think one word at a time. It's like asking a master chef to plan a whole banquet by deciding only the very next ingredient to chop.

The Problem: The student gets stuck focusing on tiny details (like using the word "therefore" instead of "so") and forgets the big picture. They start repeating the same few strategies over and over, ignoring other clever ways to solve the problem. In the paper, they call this "Entropy Collapse." It's like a student who only ever learns to solve a math problem by adding numbers, even when subtraction or multiplication would be faster. Eventually, they stop trying new things, and their creativity dies.

The New Idea: LaDi-RL (The "Dreaming" Student)

The authors of this paper, LaDi-RL, propose a different way to teach the student. Instead of making them decide word-by-word, they let the student "dream" the whole solution first in a hidden, abstract space, and then just write it down.

Here is how it works, using a simple analogy:

1. The "Blueprint" vs. The "House"

Old Way (Token-Level): The student builds a house brick by brick, deciding the color of every single brick as they go. If they make a mistake early on, the whole house might be crooked.
New Way (Latent Diffusion): The student first draws a blueprint (a "latent trajectory") in their mind. This blueprint represents the logic and the strategy of the solution, not the specific words.
- Think of this blueprint as a "thought cloud." It's a continuous, fluid representation of the idea.
- The student uses a special tool called a Diffusion Model. Imagine this tool as a sculptor who starts with a block of clay (random noise) and slowly chips away the excess until a perfect statue (the solution strategy) emerges. This allows the student to explore many different types of strategies (e.g., "let's try geometry" vs. "let's try algebra") before committing to any specific words.

2. The "Translator" Problem

There's a catch. The student's "dream" (the blueprint) is in a secret language they understand, but the final answer must be written in plain English (or code).

The Risk: Sometimes the student has a brilliant blueprint, but the "translator" (the part that turns the blueprint into text) messes up the translation. If the student gets a bad grade, how do we know if the idea was bad, or just the translation?
The Solution (Hierarchical Rollouts): The paper introduces a clever fix. Instead of judging the blueprint based on just one translation, the student generates multiple translations for the same blueprint.
- Analogy: Imagine a chef has a great recipe (the blueprint). Instead of cooking it once and judging the dish, they cook it five times. If four out of five dishes are delicious, we know the recipe is good, even if one dish was burnt by accident. This gives the teacher a much clearer signal on whether the student's strategy was actually smart.

3. Keeping the Party Diverse

To make sure the student doesn't get lazy and stick to just one blueprint, the authors add a "Repulsion Force."

Analogy: Imagine a group of explorers trying to find treasure. If they all follow the same path, they might miss a hidden cave. The "Repulsion Force" is like a gentle nudge that says, "Hey, you're too close to your friend! Go explore a different direction!"
This forces the student to generate blueprints that are structurally different from each other, ensuring they explore a wide variety of solutions.

What Did They Find?

The paper tested this new method on two tough challenges: writing code and solving math problems.

Better Accuracy: The "Dreaming Student" (LaDi-RL) got significantly more problems right on the first try compared to the old "Brick-by-Brick" student.
- In coding, they improved accuracy by 9.4%.
- In math, they improved by 5.7%.
More Creativity: The old student eventually stopped trying new things (Entropy Collapse). The new student kept finding diverse, creative solutions. Even when asked to generate 100 different answers, the new student's answers were all unique and correct, whereas the old student's answers started looking the same and getting worse.
Efficiency: The new student didn't need to write thousands of words of "thinking out loud" to get the answer. They compressed their thinking into a short, efficient "blueprint," saving time and computing power.

The Bottom Line

LaDi-RL changes how AI learns to reason. Instead of forcing the AI to think in tiny, rigid steps (word-by-word), it lets the AI explore the landscape of ideas first (using a diffusion process) and then translate those ideas into words. This prevents the AI from getting stuck in a rut, keeps its solutions diverse, and helps it solve harder problems more accurately.

Technical Summary: LaDi-RL

Problem Statement

Reinforcement learning (RL) has become a dominant paradigm for enhancing the reasoning capabilities of Large Language Models (LLMs), particularly through the optimization of reasoning trajectories. However, existing methods predominantly optimize policies over discrete token sequences. This approach creates a fundamental mismatch: critical reasoning decisions are often semantic, global, and trajectory-level, whereas token-level optimization relies on local perturbations.

This mismatch leads to diversity collapse, a phenomenon where the model's policy concentrates probability mass on a narrow set of reasoning patterns, suppressing alternative strategies. While token-level exploration can generate lexical variants (e.g., replacing "therefore" with "so"), it struggles to discover semantically distinct strategies (e.g., switching from algebraic to geometric reasoning). Consequently, RL-tuned models often improve pass@1 (single-sample accuracy) but fail to maintain or improve pass@k performance at large $k$ , indicating a loss of diverse correct trajectories.

Furthermore, moving exploration to continuous latent spaces introduces a credit-assignment challenge. In latent-space RL, the policy acts in a continuous latent space, but rewards are only observed after the latent trajectory is decoded into text. A naive rollout strategy entangles the quality of the latent reasoning with the quality of the textual decoding. A low reward could stem from a flawed latent trajectory or merely an imperfect textual realization of a sound latent, creating a noisy and misleading learning signal.

Methodology: LaDi-RL

The authors propose Latent Diffusion Reasoning with Reinforcement Learning (LaDi-RL), a framework that reframes exploration in continuous latent space using diffusion models to prevent entropy collapse.

1. Latent Diffusion Policy

Instead of generating reasoning token-by-token, LaDi-RL encodes an entire reasoning trajectory into a compact, continuous latent block (e.g., 64 tokens) using a Variational Autoencoder (VAE). The policy is parameterized as a diffusion model over these trajectory representations.

Mechanism: Starting from noise, the policy iteratively denoises toward a latent reasoning trajectory. This process allows for structured exploration and expressive distribution modeling, naturally supporting multi-modal reasoning strategies.
Diversity Guidance: To further prevent collapse, the authors introduce a repulsion-based guidance mechanism. During the denoising process, a repulsive force is applied to nearby latent trajectories, encouraging them to separate in the latent space. This acts as a local geometric regularization, explicitly preserving diversity without relying solely on policy entropy.

2. Hierarchical Latent-Text Rollouts

To address the credit-assignment challenge, LaDi-RL employs a hierarchical $N \times M$ rollout strategy:

Latent Sampling ( $N$ ): For a given query, the diffusion policy samples $N$ distinct latent reasoning trajectories.
Text Decoding ( $M$ ): For each sampled latent trajectory, the text policy generates $M$ candidate answer texts.
Decoder-Marginalized Reward: The rewards of the $M$ text completions are aggregated (averaged) to estimate the utility of the underlying latent trajectory. This "decoder-marginalized" estimate isolates the quality of the latent reasoning from the stochasticity of text decoding, providing a cleaner, lower-variance reward signal for optimizing the diffusion policy.

3. Joint Optimization

The framework optimizes two heterogeneous policies simultaneously using Group Relative Policy Optimization (GRPO):

Latent Policy: Optimized using the average reward of the $M$ texts associated with each latent block. The objective encourages the diffusion model to explore diverse, high-quality reasoning modes.
Text Policy: Optimized using the relative differences among the $M$ texts generated from the same latent block. This forces the text policy to find the optimal textual realization for a given reasoning strategy, decoupling it from the exploration of new strategies.
Loss Function: A weighted sum of the latent and text losses ensures the latent policy explores diverse chains of thought while the text policy remains a faithful and effective generator.

Key Results

The authors evaluated LaDi-RL on code generation and mathematical reasoning benchmarks, comparing it against token-level RL baselines (e.g., Vanilla GRPO, GRPO with entropy advantages) and other latent reasoning methods.

Performance Gains: LaDi-RL achieved absolute pass@1 improvements of +9.4% on code generation and +5.7% on math reasoning compared to token-level RL baselines.
Diversity Preservation: Unlike vanilla GRPO, which exhibits diversity collapse (where pass@k drops below the base model at large $k$ ), LaDi-RL consistently surpassed the base model's pass@k performance across all benchmarks. For instance, on AIME 2025, it achieved an absolute improvement of +13.4 in pass@128.
Efficiency: By compressing reasoning into a latent block (64 tokens) before decoding, LaDi-RL reduced the number of generated tokens significantly (e.g., from ~3,447 tokens in standard GRPO to ~422 tokens in LaDi-RL for code tasks), resulting in a 33% reduction in end-to-end inference time.
Ablation Studies:
- Trajectory vs. Token Level: Exploration in latent trajectory space yielded significantly higher semantic diversity (0.52 vs. 0.18–0.22 for token-level) and stronger performance.
- Diffusion vs. VAE: The diffusion-based policy outperformed a single-step VAE sampler, demonstrating the importance of iterative denoising for learning complex latent reasoning distributions.
- Hierarchical Rollouts: The $N \times M$ structure was critical; increasing $N$ (latent diversity) drove semantic diversity, while increasing $M$ (text averaging) provided cleaner reward signals.

Significance and Claims

The paper claims that LaDi-RL provides a principled alternative to token-level RL for improving LLM reasoning. Its primary significance lies in:

Mitigating Diversity Collapse: By shifting exploration from local token perturbations to global trajectory-level diffusion, the method maintains a diverse set of solution strategies, preventing the policy from converging to suboptimal local minima.
Solving Credit Assignment: The hierarchical rollout and decoder-marginalized reward estimation effectively disentangle reasoning quality from decoding noise, enabling stable optimization of latent policies.
Efficiency and Quality: The approach achieves state-of-the-art performance in both accuracy (pass@1) and diversity (pass@k) while significantly reducing computational overhead through latent compression.

The authors conclude that latent diffusion offers a robust mechanism for structured exploration, enabling models to discover semantically distinct reasoning strategies that token-level methods miss, thereby unlocking higher reasoning boundaries without sacrificing solution diversity.