🤖 machine learning

Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization

This paper introduces Distribution Matching Policy Optimization (DMPO), a novel reinforcement learning framework that significantly enhances the reasoning capabilities of diffusion large language models by aligning their policy distribution with an optimal reward-tilted target, achieving substantial accuracy improvements over existing baselines without supervised fine-tuning.

Original authors: Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

Published 2026-07-14

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

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 have a super-smart robot that writes stories, solves math problems, and plays games. Most robots today work like a person reading a book one word at a time, from left to right. They can't skip ahead or look back; they just have to guess the next word based on the ones before it. This is slow, like trying to fill a giant puzzle by only looking at the piece right in front of you.

Enter Diffusion Large Language Models (dLLMs). These are the new kids on the block. Instead of writing word-by-word, they start with a messy, scrambled mess of "mystery tokens" (like a puzzle where every piece is covered in fog) and slowly clean it up, revealing the answer all at once or in any order they want. This makes them potentially much faster at thinking.

But here's the catch: while these diffusion robots are fast, they aren't always the smartest at tricky reasoning tasks like advanced math or logic puzzles. To make them smarter, scientists usually use a technique called Reinforcement Learning (RL). Think of this as a video game where the robot gets points for a good answer and loses points for a bad one.

The Problem: The "Mode-Seeking" Trap

The old way of training these robots (using methods like GRPO) was a bit like a student who only studies the one answer they think is right. If the robot finds a way to get points that feels "safe," it stops exploring. It gets stuck in a rut, ignoring other clever solutions that might be just as good. In the paper, the authors call this "mode-seeking." It's like a robot that only learns to solve a math problem one specific way, and if that way gets blocked, it panics. It also tends to ignore the "messy" but correct answers that look different from the first one it found.

The paper argues that this old way of just chasing the highest score is flawed for diffusion models. It suggests that by only focusing on the single best path, we miss out on the robot's unique ability to explore many different paths at once.

The Solution: DMPO (The "Map-Matching" Robot)

The authors propose a new method called Distribution Matching Policy Optimization (DMPO).

Instead of telling the robot, "Just find the highest score," DMPO says, "Here is the entire map of all the good answers you could possibly give. Your job is to learn to match that whole map."

Imagine you are trying to teach a dog to fetch.

The Old Way: You throw a ball, and the dog runs to the one spot where the ball landed. If the ball lands in a bush, the dog learns to only fetch from bushes.
The DMPO Way: You show the dog a map of every possible place the ball could land (bushes, grass, sand, water) and say, "Learn to fetch from any of these places, weighted by how good the spot is." The dog learns to be flexible and explore the whole yard, not just one corner.

In technical terms, DMPO uses a special math trick called Weighted Denoising Cross-Entropy (WDCE). This allows the robot to learn from its past attempts (even the ones that weren't perfect) without needing to re-generate everything from scratch every time. It's like having a "replay buffer" where the robot can study its old games over and over, learning from the good and the bad moves without getting confused.

The Secret Sauce: The "Weight Baseline"

The authors found a tricky problem when training with small groups of examples (small "batch sizes"). Sometimes, the robot would get confused and start rewarding bad answers just because they were the only ones it saw.

To fix this, they invented a clever trick called Weight Baseline Subtraction.
Think of it like a teacher grading a test. If a student gets a question right, they get a gold star. But if the teacher only saw that one question, they might think everything the student did was great. The "baseline" is like a "standard of average." The teacher subtracts the "average" score from the student's score.

If the student did better than average, they get a big gold star.
If they did worse than average, they get a "penalty" (a negative weight), even if they technically got some points.

This ensures the robot doesn't get too excited about mediocre answers and keeps pushing for the truly excellent ones.

The Results: How Much Better?

The authors tested this new method on some very hard reasoning benchmarks. They applied DMPO to pretrained models (like LLaDA-Instruct and Dream-Instruct) in two different ways:

Direct Application: They applied DMPO directly to the base models without any prior "homework" (Supervised Fine-Tuning or SFT) on reasoning datasets. This "R1-Zero-like" approach was used to clearly demonstrate the raw potential of DMPO.
Enhanced Application: They also applied DMPO to models that had already undergone SFT to show it works as a powerful upgrade for models that are already trained.

The results were quite impressive across both setups:

On a math puzzle called GSM8K, the model applied directly to the pretrained base (without SFT) showed massive improvements. Specifically, DMPO achieved up to a 39.63 percentage-point accuracy improvement over prior non-DMPO RL baselines, and a staggering 67.97 percentage-point improvement over the base model itself.
On a logic puzzle called Sudoku, they saw huge gains, with one version of their model jumping from a 16.41% success rate to 80.86% (a difference of over 64 points!).
Crucially, DMPO continued to deliver significant performance gains even when applied to models that had already undergone SFT, proving it is a robust method that works whether you start from scratch or build on existing training.

The paper suggests that DMPO is not just a small tweak but a fundamental shift. It allows the robot to be off-policy, meaning it can learn from old data efficiently, and forward-only, meaning it doesn't need to do expensive backward calculations that slow down other robots.

The Bottom Line

The authors are confident that this method works well for the specific models they tested (like LLaDA and Dream) on these specific reasoning tasks. They measured these improvements through rigorous experiments on standard datasets, showing that DMPO can boost performance whether applied to a model that has just been pretrained or one that has already undergone supervised fine-tuning. However, they admit that they haven't tested this on every possible type of model or task yet, so while the results are strong, the full potential is still being explored.

In short, DMPO teaches diffusion robots to stop fixating on a single "perfect" answer and instead learn to appreciate the whole landscape of good solutions, making them smarter, faster, and more creative problem solvers.

Technical Summary: Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization

Problem Statement

Diffusion Large Language Models (dLLMs) offer a promising alternative to autoregressive LLMs (AR-LLMs) by enabling any-order generation and potentially higher inference throughput. However, scaling dLLMs to match the reasoning capabilities of AR-LLMs remains a significant challenge. Existing Reinforcement Learning (RL) fine-tuning methods, such as Group Relative Policy Optimization (GRPO), are ill-suited for dLLMs due to two primary factors:

Computational Mismatch: dLLMs rely on bidirectional processing, making the estimation of sequence log-probabilities (required for standard policy gradient methods) computationally expensive and often intractable compared to AR models.
Algorithmic Limitations: Current RL approaches for dLLMs often rely on heuristic approximations (e.g., mean-field decomposition) that ignore correlations between token positions. Furthermore, they typically focus on reward maximization, which tends to optimize for a single reward mode. This "mode-seeking" behavior can lead to reward hacking and a failure to explore the diverse, high-quality reasoning paths that dLLMs are theoretically capable of generating.

Methodology: Distribution Matching Policy Optimization (DMPO)

The paper proposes Distribution Matching Policy Optimization (DMPO), a principled RL fine-tuning framework grounded in Stochastic Optimal Control (SOC). Unlike traditional methods that maximize expected reward, DMPO aims to match the model's policy distribution to an optimal, reward-tilted target distribution.

Core Framework

The objective is to learn a policy $\pi_\theta$ such that the induced sequence distribution $p_\theta(o|q)$ matches the optimal distribution $p^*(o|q)$ , defined as:
$p^*(o|q) = \frac{1}{Z(q)} p_{ref}(o|q) e^{r(q,o)/\alpha}$
where $p_{ref}$ is the pre-trained policy, $r$ is the reward function, and $\alpha$ is a temperature parameter.

To achieve this, the authors formulate the problem as minimizing the Forward Kullback-Leibler (KL) divergence (equivalent to Cross-Entropy) between the target and the current policy:
$\min_\theta \mathbb{E}_{q \sim D} \mathbb{E}_{p^*(o|q)} \left[ \log \frac{p^*(o|q)}{p_\theta(o|q)} \right]$
This "mass-covering" objective encourages the model to explore diverse high-reward trajectories rather than collapsing to a single mode.

Key Technical Components

Weighted Denoising Cross-Entropy (WDCE):
Since $p^*$ is intractable (due to the unknown partition function $Z(q)$ and the inability to sample directly from it), DMPO employs Importance Sampling. It treats samples from a reference policy $p_v$ (a stop-gradient copy of the model) as weighted samples from $p^*$ .
The resulting loss function, WDCE, leverages the unique forward (noising) process of dLLMs. It computes the loss using only the final clean samples and their associated importance weights, discarding the intermediate rollout trajectories. This makes WDCE an off-policy and forward-only objective.
Weight Baseline Subtraction:
A critical implementation challenge identified is the instability of WDCE when using small training batch sizes. In small batches, positive weights for "bad" responses (which still have non-zero probability under $p^*$ ) can dominate, leading to incorrect gradient updates.
To address this, the authors propose Weight Baseline Subtraction, where a baseline weight $w_{base}$ is subtracted from the importance weight:
$w_{real} = w - w_{base}$
Three strategies for $w_{base}$ are proposed:
- Group Weight Baseline: A constant value (1) to encourage convergence.
- Individual Weight Baseline: Based on a negative reward tilt.
- Model Weight Baseline: Based on the current model's own importance weights.
  This technique effectively injects negative gradients for low-reward samples, preventing the model from promoting bad responses in low-batch scenarios.
Weighted Direct Discriminative Optimization (WDDO):
As an alternative to cross-entropy, the paper explores a GAN-inspired objective that naturally incorporates negative gradients for bad samples, further enhancing distribution matching.

Key Contributions

Novel RL Framework: The proposal of DMPO, a distribution-matching framework specifically designed for dLLMs, shifting the paradigm from reward maximization to matching the entire reward-tilted policy distribution.
Theoretical Grounding: Derivation of the WDCE loss using importance sampling and Stochastic Optimal Control, providing a theoretically sound method for off-policy training without requiring expensive sequence log-probability estimates.
Practical Solutions: Identification of the small-batch failure mode in forward KL objectives and the introduction of Weight Baseline Subtraction and WDDO to stabilize training and improve sample efficiency.
Efficiency: Demonstration that WDCE enables training via the forward process alone, allowing for the reuse of rollouts (off-policy) and compatibility with fast inference techniques (e.g., KV-cache), significantly accelerating training compared to trajectory-based methods.

Experimental Results

The authors evaluated DMPO on the LLaDA-8B-Instruct and Dream-7B models across four reasoning benchmarks: GSM8K, MATH500, Countdown, and Sudoku.

Performance Gains: DMPO consistently outperformed base models and existing RL baselines (including diffu-GRPO, cGRPO, and d1).
- On Sudoku, DMPO-LLaDA achieved a 39.63 percentage-point improvement over the best non-DMPO RL baseline and a 67.97 percentage-point improvement over the base model.
- On Countdown, improvements of up to 16.40 percentage points were observed over strong baselines.
No SFT Required: The method achieved these results without Supervised Fine-Tuning (SFT) on curated datasets, demonstrating its ability to scale reasoning capabilities directly from pre-trained models.
Training Efficiency: Due to its off-policy nature and forward-only computation, DMPO achieved 2–3x higher sample efficiency and up to 8x faster rollout sampling (when combined with fast-dLLM inference) compared to GRPO-type methods. It reached the reward levels of d1 with only 31% of the training compute budget.
Stability: The method remained stable even with highly stale data (up to 24x reuse of the same rollout batch), a property attributed to the stability of the sequence-level importance weights used in DMPO.

Significance and Claims

The paper claims that DMPO represents a significant step forward in scaling dLLMs for reasoning tasks. By moving away from the reward-maximization paradigm and the computational bottlenecks of trajectory-based RL, DMPO fully leverages the unique characteristics of diffusion models:

Diversity: The forward KL objective prevents mode collapse, allowing dLLMs to explore diverse reasoning paths.
Efficiency: The forward-only, off-policy nature of WDCE aligns with the high-throughput inference capabilities of dLLMs, offering a path to scalable RL fine-tuning that is computationally feasible for large models.
Generalizability: The framework is presented as a principled approach that can be applied to various dLLM backbones, as evidenced by its success on both LLaDA and Dream architectures.

The authors conclude that while the current work focuses on specific models and benchmarks, the distribution-matching framework opens new directions for improving generative models and designing more effective weight-baseline techniques for future research.