Distribution Matching Distillation Meets Reinforcement Learning
This work addresses efficiency and quality trade-offs in diffusion models for image generation, representing an incremental improvement over existing distillation methods.
The paper tackles the performance limitation of few-step diffusion models distilled from multi-step teachers by integrating reinforcement learning into the distillation process, achieving leading visual quality and even surpassing the teacher model in some cases.
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.