DDIL: Diversity Enhancing Diffusion Distillation With Imitation Learning

TL;DR

This paper introduces DDIL, a novel diffusion distillation method using imitation learning to improve sample diversity and quality while reducing inference passes in generative models.

Contribution

The paper formulates diffusion distillation within an imitation learning framework, addressing covariate shift and enhancing training with data and student distributions.

Findings

DDIL improves sample diversity and quality.

Enhanced training stability across distillation methods.

Outperforms baseline algorithms like PD, LCM, and DMD2.

Abstract

Diffusion models excel at generative modeling (e.g., text-to-image) but sampling requires multiple denoising network passes, limiting practicality. Efforts such as progressive distillation or consistency distillation have shown promise by reducing the number of passes at the expense of quality of the generated samples. In this work we identify co-variate shift as one of reason for poor performance of multi-step distilled models from compounding error at inference time. To address co-variate shift, we formulate diffusion distillation within imitation learning (DDIL) framework and enhance training distribution for distilling diffusion models on both data distribution (forward diffusion) and student induced distributions (backward diffusion). Training on data distribution helps to diversify the generations by preserving marginal data distribution and training on student distribution addresses compounding error by correcting covariate shift. In addition, we adopt reflected diffusion formulation for distillation and demonstrate improved performance, stable training across different distillation methods. We show that DDIL consistency improves on baseline algorithms of progressive distillation (PD), Latent consistency models (LCM) and Distribution Matching Distillation (DMD2).