Distribution Matching Distillation Meets Reinforcement Learning

TL;DR

This paper introduces DMDR, a unified framework combining Distribution Matching Distillation and Reinforcement Learning to improve the quality and controllability of diffusion model distillation, surpassing multi-step teacher models.

Contribution

The work presents a novel joint optimization framework that unifies DMD and RL, enhancing diffusion model distillation with mutual benefits and new training strategies.

Findings

DMDR achieves state-of-the-art visual quality in few-step generation.

It surpasses the performance of multi-step teacher models.

Joint optimization improves controllability and reduces reward hacking.

Abstract

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.