Distribution Matching Distillation Meets Reinforcement Learning

TL;DR

DMDR presents a unified framework that synergistically combines Distribution Matching Distillation (DMD) with reinforcement learning (RL) to train few-step diffusion generators that can outperform their multi-step teachers. By integrating RL during distillation, it guides mode coverage beyond imitation while using DMD as a regularizer to curb reward hacking. Two dynamic cold-start strategies, DynaDG and DynaRS, accelerate early training, enabling rapid alignment of distributions and global structure learning. Empirical results across multiple model families and RL algorithms show state-of-the-art performance for four-step generation and even surpassing the teacher in several settings, with strong generalization and stability. The work highlights the practical potential of image-free, jointly optimized distillation and RL for fast, high-quality diffusion synthesis.

Abstract

Distribution Matching Distillation (DMD) distills a pre-trained multi-step diffusion model to a few-step one to improve inference efficiency. However, the performance of the latter is often capped by the former. To circumvent this dilemma, we propose DMDR, a novel framework that combines Reinforcement Learning (RL) techniques into the distillation process. We show that for the RL of the few-step generator, the DMD loss itself is a more effective regularization compared to the traditional ones. In turn, RL can help to guide the mode coverage process in DMD more effectively. These allow us to unlock the capacity of the few-step generator by conducting distillation and RL simultaneously. Meanwhile, we design the dynamic distribution guidance and dynamic renoise sampling training strategies to improve the initial distillation process. The experiments demonstrate that DMDR can achieve leading visual quality, prompt coherence among few-step methods, and even exhibit performance that exceeds the multi-step teacher.

Figures (12)

• Figure 1: More Images generated by Z-Image-Turbo zimage distilled through our DMDR.
• Figure 2: Overview of DMDR, which contains three key elements: (1) A DMD framework to optimize the generator by utilizing the gradient derived from an implicit distribution matching objective; (2) A RL branch to concurrently incorporate reward feedback from the reward model; (3) Two dynamic training strategies to facilitate more efficient and effective distillation during the initial phase.
• Figure 3: Illustration for mode seeking process.
• Figure 4: Visualization of distribution matching directions.
• Figure 5: Visual comparison between the teachers, selected competing methods dmd2hypersdladd, and ours. All images are generated using identical noise. Our model produces images with superior quality and prompt coherence. More comparisons are available in Appendix \ref{['app:vis_comp']}.