Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ReDiF: Reinforced Distillation for Few Step Diffusion

Amirhossein Tighkhorshid, Zahra Dehghanian, Gholamali Aminian, Chengchun Shi, Hamid R. Rabiee

TL;DR

Diffusion models deliver high-fidelity samples but suffer from slow sampling. ReDiF recasts distillation as policy optimization, training a few-step student via reward signals that quantify alignment with a high-step teacher. The framework is model- and data-free, compatible with existing acceleration strategies, and can encode downstream task preferences through rewards. Empirical results on LAION-5B and COCO demonstrate superior fidelity and diversity with far fewer steps and lower compute than prior distillation methods, highlighting RL as a flexible, general mechanism for diffusion acceleration.

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.

ReDiF: Reinforced Distillation for Few Step Diffusion

TL;DR

Diffusion models deliver high-fidelity samples but suffer from slow sampling. ReDiF recasts distillation as policy optimization, training a few-step student via reward signals that quantify alignment with a high-step teacher. The framework is model- and data-free, compatible with existing acceleration strategies, and can encode downstream task preferences through rewards. Empirical results on LAION-5B and COCO demonstrate superior fidelity and diversity with far fewer steps and lower compute than prior distillation methods, highlighting RL as a flexible, general mechanism for diffusion acceleration.

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.
Paper Structure (24 sections, 7 equations, 5 figures, 5 tables, 2 algorithms)

This paper contains 24 sections, 7 equations, 5 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Training Free Methods
  4. Distillation Based Methods
  5. Reward Based Methods
  6. ReDiF Algorithm
  7. Preliminary
  8. Methodology
  9. Reward Functions
  10. Low Step Distillation Issues
  11. Markov Property of the Low Step Student
  12. Gaussian Assumption in Reverse Process
  13. Experiments
  14. Reward Function
  15. RL Algorithm Ablation
  16. ...and 9 more sections

Figures (5)

  • Figure 1: Qualitative comparison of ReDiF and the best diffusion distillation methods.
  • Figure 2: Trajectories of the teacher and student models during the denoising process. A well-trained student model should be capable of approximating multiple denoising steps of the teacher model within a single, longer step.
  • Figure 3: Overview of the ReDiF framework setup: the prompt is given to both Teacher and Student diffusion models, their outputs are compared via CLIP/DINOv3 image encoders to compute a cosine similarity-based reward, which is used by PPO or GRPO to update the Student model. The Teacher model is frozen while the Student is trained.
  • Figure 4: Qualitative comparison of ReDiF and the best diffusion models acceleration methods.
  • Figure 5: Qualitative evolution of generated samples during ReDiF Framework. Each column corresponds to a training epoch, and each row shows samples generated from a fixed text prompt. The first column shows the teacher reference images.