Efficient Diffusion Distillation via Embedding Loss

TL;DR

This paper introduces Embedding Loss, a new method that improves diffusion model distillation by enhancing quality and reducing training time, especially for resource-limited settings.

Contribution

The paper proposes Embedding Loss, a novel feature embedding-based loss function that boosts diffusion model distillation efficiency and effectiveness across multiple datasets and frameworks.

Findings

Achieves state-of-the-art FID of 1.475 on CIFAR-10 for unconditional generation.

Reduces training iterations by up to 80%.

Demonstrates consistent improvements across datasets like ImageNet, AFHQ-v2, and FFHQ.

Abstract

Recent advances in distilling expensive diffusion models into efficient few-step generators show significant promise. However, these methods typically demand substantial computational resources and extended training periods, limiting accessibility for resource-constrained researchers, and existing supplementary loss functions have notable limitations. Regression loss requires pre-generating large datasets before training and limits the student model to the teacher's performance, while GAN-based losses suffer from training instability and require careful tuning. In this paper, we propose Embedding Loss (EL), a novel supplementary loss function that complements existing diffusion distillation methods to enhance generation quality and accelerate training with smaller batch sizes. Leveraging feature embeddings from a diverse set of randomly initialized networks, EL effectively aligns the feature distributions between the distilled few-step generator and the original data. By computing Maximum Mean Discrepancy (MMD) in the embedded feature space, EL ensures robust distribution matching, thereby preserving sample fidelity and diversity during distillation. Within distribution matching distillation frameworks, EL demonstrates strong empirical performance for one-step generators. On the CIFAR-10 dataset, our approach achieves state-of-the-art FID values of 1.475 for unconditional generation and 1.380 for conditional generation. Beyond CIFAR-10, we further validate EL across multiple benchmarks and distillation methods, including ImageNet, AFHQ-v2, and FFHQ datasets, using DMD, DI, and CM distillation frameworks, demonstrating consistent improvements over existing one-step distillation methods. Our method also reduces training iterations by up to 80%, offering a more practical and scalable solution for deploying diffusion-based generative models in resource-constrained environments.