Teacher-Feature Drifting: One-Step Diffusion Distillation with Pretrained Diffusion Representations

2026-05-08

AI Summary

This work addresses the low sampling efficiency of pre-trained diffusion models and the reliance of existing distillation methods on multi-stage training or auxiliary networks. The authors propose an efficient single-step distillation framework that directly leverages intermediate hidden states from the teacher model as feature representations, eliminating the need for additional feature extraction networks. By integrating a single-step drift loss with a lightweight mode coverage loss, the method enables one-step forward generation while effectively mitigating mode collapse. Experimental results demonstrate that the approach achieves a Fréchet Inception Distance (FID) of 1.58 on ImageNet-64×64 and 18.4 on SDXL, attaining competitive generation quality and diversity while significantly simplifying the distillation pipeline.

Abstract

Sampling from pretrained diffusion and flow-matching models typically requires many forward passes to generate diverse and high-fidelity images. Existing distillation methods often rely on multiple auxiliary networks, carefully designed training stages, or complex optimization pipelines. In this work, we revisit the recently proposed Drifting Model objective and show that a single drifting loss can be directly used to simplify one step distillation. A key observation is that the pretrained diffusion teacher itself already provides a strong representation space. Unlike the original Drifting Model, which relies on an additional pretrained feature extractor, we use intermediate hidden states of the pretrained teacher model as the feature representation. This removes the need for training or introducing an extra representation network while preserving a semantically meaningful feature geometry for drifting. Furthermore, we introduce a lightweight mode coverage loss to mitigate mode collapse during distillation and encourage the student generator to cover diverse teacher-supported regions. Extensive experiments on ImageNet and SDXL demonstrate that our method achieves efficient one step generation with competitive image quality and diversity, achieving FID scores of 1.58 on ImageNet-64$\times$64 and 18.4 on SDXL, while substantially simplifying the overall distillation framework.

Problem

• diffusion distillation
• one-step generation
• teacher-feature drifting
• mode collapse
• pretrained diffusion models

Innovation

• diffusion distillation
• teacher-feature drifting
• one-step generation
• mode coverage loss
• pretrained diffusion representations