Distribution Matching Distillation Meets Reinforcement Learning

Dengyang Jiang; Dongyang Liu; Zanyi Wang; Qilong Wu; Liuzhuozheng Li; Hengzhuang Li; Xin Jin; David Liu; Changsheng Lu; Zhen Li; Bo Zhang; Mengmeng Wang; Steven Hoi; Peng Gao; Harry Yang

Source-linked AI summary

Distribution Matching Distillation Meets Reinforcement Learning

Dengyang Jiang, Dongyang Liu, Zanyi Wang, Qilong Wu, Liuzhuozheng Li, Hengzhuang Li, Xin Jin, David Liu, Changsheng Lu, Zhen Li, Bo Zhang, Mengmeng Wang, Steven Hoi, Peng Gao, Harry Yang

arXiv:2511.13649v5Nov 17, 2025cs.CV

TL;DR

Few-step diffusion distillation and reinforcement-learning alignment are usually treated as separate stages, leaving their combination underexplored. DMDR unifies them through reward-tilted and dynamic initial training followed by joint DMD-RL optimization, achieving state-of-the-art few-step performance and surpassing multi-step teachers on visual quality and prompt adherence.

Problem
Existing studies typically treat efficient DMD distillation and RL-based preference alignment as separate stages, leaving their systematic unification unexplored.
Source: §1 · p. 2, §2.2 · p. 3
Method
DMDR combines Reward-Tilted Distribution Matching and two dynamic distillation strategies initially, then jointly optimizes DMD and RL.
Source: §1 · p. 2
Results
DMDR achieves state-of-the-art few-step performance across diverse architectures and consistently surpasses its multi-step teachers in visual quality and prompt adherence.
Source: §1 · p. 2, §3.2 · p. 8, §6 · p. 11
Takeaways & Limitations
Joint DMD-RL training makes few-step models more preference-aware while using DMD regularization to mitigate reward hacking.
Source: §1 · p. 2, §6 · p. 11
Reward optimization and distillation reduce diversity, although tuning λrl allows users to trade diversity against preference gains.
Source: §4 · pp. 9–10

Abstract

from arXiv · show

Distribution Matching Distillation (DMD) facilitates efficient inference by distilling multi-step diffusion models into few-step variants. Concurrently, Reinforcement Learning (RL) has emerged as a vital tool for aligning generative models with human preferences. While both represent critical post-training stages for large-scale diffusion models, existing studies typically treat them as independent, sequential processes, leaving a systematic framework for their unification largely unexplored. In this work, we demonstrate that jointly optimizing these two objectives yields mutual benefits: RL enables more preference-aware and controllable distillation rather than uniformly compressing the full data distribution, while DMD serves as an effective regularizer to mitigate reward hacking during RL training. Building on these insights, we propose DMDR, a unified framework that incorporates Reward-Tilted Distribution Matching optimization alongside two dynamic distillation training strategies in the initial stage, followed by the joint DMD and RL optimization in the second stage. Extensive experiments demonstrate that DMDR achieves state-of-the-art visual quality and prompt adherence among few-step generation methods, even surpassing the performance of its multi-step teacher model.

1. Introduction

Diffusion models need efficient few-step sampling and stronger alignment with human preferences, but distillation and reinforcement learning have largely remained separate. DMDR unifies them, using joint optimization and dynamic training strategies to achieve state-of-the-art few-step generation across models and RL algorithms.

1. Introduction

Existing work generally studies step-distillation and preference alignment in isolation, leaving their systematic unification and potential synergy underexplored.Earlier attempts mainly applied RL after a model had already been distilled.Source: §1 · p. 2
DMDR unifies Distribution Matching Distillation and Reinforcement Learning in a synergistic training pipeline rather than treating them as separate post-training stages.The framework is designed to combine efficient sampling with preference alignment.Source: §1 · p. 2
RL steers distillation toward high-reward regions, while DMD regularizes RL to mitigate reward hacking and maintain proximity to the teacher manifold.This joint optimization provides preference-aware distillation and distributional regularization.Source: §1 · p. 2
DMDR begins with Reward-Tilted Distribution Matching plus Dynamic Distribution Guidance and Dynamic Renoise Sampling to stabilize initial distillation.These strategies address the early training phase before joint DMD and RL optimization.Source: §1 · p. 2
The method achieves state-of-the-art few-step generative performance across different model types and RL algorithms.The reported compatibility includes flow-based and denoising-based models, as well as ReFL, DPO, and GRPO.Source: §1 · p. 2

2. Method

DMDR unifies DMD and RL through a two-stage pipeline that first integrates preference signals during distillation, then jointly optimizes reward and distribution matching. Dynamic strategies address early cold-start instability, while DMD regularization constrains RL-driven drift and reward hacking.

2.2. Problem Formulation and Pipeline Overview.

DMDR combines reward-aware distillation with joint RL+DMD optimization in two stages, integrating preference signals early and directly maximizing rewards later.Stage 1 uses Reward-Tilted Distribution Matching with dynamic training; Stage 2 applies joint RL and DMD.Source: Figure 2 · p. 3, §2.2 · p. 3, §2.4 · p. 5

2.3. Bring DMD and RL Together with Benefit

Joint DMD+RL improves over DMD alone while avoiding the reward hacking and semantic collapse observed when RL is applied directly to a distilled model.The joint approach keeps improving DPG and aesthetic quality as reward increases.Source: Figure 4 · p. 4, §2.3 · p. 5
DMD projects reward-driven updates toward the teacher’s data manifold, while RL steers distillation toward high-reward regions beyond uniform teacher imitation.This combination is intended to preserve semantic integrity while overcoming the performance ceiling of pure DMD.Source: §2.3 · p. 4

2.4. Reward-Tilted Distribution Matching for Integrating RL into the Early Phase of DMD

RT-DM weights DMD gradients by reward scores, enabling preference-aware warm starts without backpropagating through the full multi-step denoising trajectory.RT-DM accelerates early preference alignment, but alone saturates because it remains constrained by the teacher’s manifold; joint DMD+RL breaks this plateau.Source: §2.4 · pp. 5–6

2.5. Annealed Distributional Overlap Maximization for Better Initial Distillation

DynaDG and DynaRS increase early distributional overlap and provide more reliable optimization signals, accelerating global structure formation during initial distillation.High-noise-biased renoising is progressively annealed toward uniform sampling as the student improves, supporting later fine-grained detail learning.Source: §2.5 · pp. 6–7, Figure 7 · p. 6

3. Experiments

DMDR is evaluated across architectures, model sizes, paradigms, benchmarks, and RL configurations, achieving strong few-step generation and surpassing multi-step teachers on DPG Bench and GenEval. Ablations show that reward choice, RL compatibility, and DMD regularization govern preference alignment and general capability.

3.1. Experimental Setup

Experiments use SDXL-Base, SD3-Medium, and SD3.5-Large for distillation, with ShareGPT-4o-Image prompts and CLIP, Aesthetic, Pick, and Human Preference scores for evaluation.Additional comparisons use DPG Bench and GenEval against the corresponding teachers.Source: §3.1 · pp. 7–8

3.2. Main Results

DMDR achieves state-of-the-art prompt coherence and high-quality, aesthetically pleasing images across UNet and Transformer architectures, 2B and 8B models, and flow-based and denoising-based paradigms.Qualitative comparisons also report superior quality and prompt coherence.Source: §3.2 · p. 8, Figure 8 · p. 8
DMDR’s four-step models consistently outperform their multi-step teachers on DPG Bench and GenEval across all base models.These benchmarks were not directly optimized with RL, indicating gains beyond the training reward objectives.Source: §3.2 · p. 8, Table 3 · p. 8, Table 4 · p. 8

3.3. Ablation Study

Reward-model ablations show that optimizing a specific reward improves its corresponding metric, while combining CLIP and HPS yields more balanced overall performance.The results support customizing the distillation preference signal for downstream tasks.Source: §3.3 · p. 9, Table 5 · p. 9
ReFL, DPO, and GRPO all substantially improve every evaluation dimension over the Stage 1 initialization, with ReFL achieving the best overall performance.In few-step models, ReFL gradients can return directly to the initial state because each step predicts a clean image.Source: §3.3 · p. 9
The λrl ablation shows that too little RL leaves the student constrained by the teacher, whereas excessive RL reduces general capacity despite high reward scores.DMD therefore acts as a structural regularizer that balances preference gains with distributional stability and mitigates reward hacking.Source: §3.3 · p. 9

4. Limitation and Discussion

The discussion identifies a quality–diversity trade-off: CFG and RL sharpen probability mass toward high-reward, high-confidence regions, while DMD regularization improves but does not eliminate the diversity cost. Adjusting λrl provides a tunable balance between preference gains and diversity.

4. Limitation and Discussion

CFG and RL act as density sharpeners in class-conditional ImageNet and text-to-image experiments, increasing quality while reducing diversity.The ImageNet experiments use a SiT backbone for one-step generation and a DINOv2-based reward model for classification accuracy.Source: §4 · p. 9, Table 6 · p. 10
The diversity reduction is presented as a characteristic of score-based distillation and preference alignment, which steer generation toward high-reward, high-confidence regions.Additional LPIPS evaluations report similar degradation for ImageNet and other methods.Source: §4 · p. 10, Table 7 · p. 10
DMD regularization mitigates reward hacking relative to direct or sequential RL but cannot completely remove the inherent diversity cost of reward optimization.The trade-off remains adjustable through λrl.Source: §4 · p. 10, Table 8 · p. 10
A smaller λrl improves diversity but reduces preference gains, allowing users to choose the desired quality–alignment balance.Adaptive reward weighting and diversity-preserving RL objectives are identified as future directions.Source: §4 · p. 10, Table 8 · p. 10

5. Related Work

Prior work includes distribution-matching, trajectory-based, adversarial, and reward-model approaches for efficient or preference-aligned diffusion generation. DMDR differs by integrating few-step distribution matching with reinforcement learning rather than treating these goals independently or applying RL only after distillation.