ReDiF: Reinforced Distillation for Few Step Diffusion

TL;DR

This paper introduces ReDiF, a reinforcement learning-based distillation method for diffusion models that optimizes the denoising process, enabling high-quality generation with fewer steps and less computation.

Contribution

ReDiF is the first RL-driven distillation framework for diffusion models that dynamically guides the denoising process, improving efficiency and performance.

Findings

Achieves superior performance with fewer inference steps.

Reduces computational resources compared to existing methods.

Model-agnostic and applicable to various diffusion models.

Abstract

Distillation addresses the slow sampling problem in diffusion models by creating models with smaller size or fewer steps that approximate the behavior of high-step teachers. In this work, we propose a reinforcement learning based distillation framework for diffusion models. Instead of relying on fixed reconstruction or consistency losses, we treat the distillation process as a policy optimization problem, where the student is trained using a reward signal derived from alignment with the teacher's outputs. This RL driven approach dynamically guides the student to explore multiple denoising paths, allowing it to take longer, optimized steps toward high-probability regions of the data distribution, rather than relying on incremental refinements. Our framework utilizes the inherent ability of diffusion models to handle larger steps and effectively manage the generative process. Experimental results show that our method achieves superior performance with significantly fewer inference steps and computational resources compared to existing distillation techniques. Additionally, the framework is model agnostic, applicable to any type of diffusion models with suitable reward functions, providing a general optimization paradigm for efficient diffusion learning.