Inference-Time Diffusion Model Distillation

TL;DR

Distillation++ enhances inference-time diffusion model distillation by incorporating teacher-guided refinement during sampling, significantly reducing performance gaps without extra data or fine-tuning.

Contribution

We introduce Distillation++, a novel framework that integrates teacher guidance into the sampling process to improve diffusion model distillation.

Findings

Substantial improvements over state-of-the-art baselines.

Enhanced early-stage sampling quality.

Robust guided sampling without additional data.

Abstract

Diffusion distillation models effectively accelerate reverse sampling by compressing the process into fewer steps. However, these models still exhibit a performance gap compared to their pre-trained diffusion model counterparts, exacerbated by distribution shifts and accumulated errors during multi-step sampling. To address this, we introduce Distillation++, a novel inference-time distillation framework that reduces this gap by incorporating teacher-guided refinement during sampling. Inspired by recent advances in conditional sampling, our approach recasts student model sampling as a proximal optimization problem with a score distillation sampling loss (SDS). To this end, we integrate distillation optimization during reverse sampling, which can be viewed as teacher guidance that drives student sampling trajectory towards the clean manifold using pre-trained diffusion models. Thus, Distillation++ improves the denoising process in real-time without additional source data or fine-tuning. Distillation++ demonstrates substantial improvements over state-of-the-art distillation baselines, particularly in early sampling stages, positioning itself as a robust guided sampling process crafted for diffusion distillation models. Code: https://github.com/geonyeong-park/inference_distillation.