Diffusion models compressed effectively from large to small with guidance folding

GFD-OPD: Guidance-Folded On-Policy Distillation of Diffusion Models Across Scales

Machine LearningArtificial Intelligence

Summary

Diffusion models help computers create complex images and language by generating data step-by-step. Scientists found it hard to shrink large diffusion models into smaller ones without losing quality. The authors discovered why the usual training methods fail and created a new technique called GFD-OPD that fixes those problems by carefully managing how smaller models learn from bigger ones. This new method helps smaller models perform better and learn faster while avoiding errors that often occur when shrinking these models.

What this means in practice

  • For machine learning engineers: Create smaller yet accurate diffusion models for deployment on devices with limited resources using the improved distillation method presented here.
  • For software product teams: Integrate efficient, compact diffusion models for image or language generation features to reduce application size and improve responsiveness.

Authors

Zhenxing Zhang, Jiayan Teng, Wenxu Wu, Zhuoyi Yang, Jiazheng Xu, Wendi Zheng, Jie Tang, Dan Guo, Meng Wang

Abstract

On-policy distillation (OPD) has demonstrated two important capabilities in language models: compressing large teachers into smaller students and merging expert models into a single model. Existing diffusion OPD, however, mostly focus on the latter, with teachers and students sharing the same backbone and scale. We investigate large-to-small diffusion opd from large teachers to a small student and find that the standard recipe fails. To find the underlying cause, we propose Fixed-State KL, an effective and fair way to measure the distribution gap between student and teacher during OPD training for diffusion models. We are the first to clarify why large-to-small OPD is challenging for diffusion models: a smaller student struggles to perfectly match the distribution of a larger teacher, while classifier-free guidance can accumulate and amplify the distributional discrepancies between the student's conditional and unconditional branches and those of the teacher. To solve this problem, we propose GFD-OPD, a simple yet effective method that reduces the student-teacher gap while avoiding the error amplification of the CFG composition. Across numerous experiments, GFD outperforms previous baselines in both training efficiency and final performance, achieving state-of-the-art results on all benchmarks.