提出新型嵌入损失,让高效生成模型更快更准地复现扩散模型效果。
Efficient Diffusion Distillation via Embedding Loss

- 用随机网络的特征嵌入计算最大均值差异,对齐学生与数据分布。
- 在CIFAR-10上实现1.475(无条件)和1.380(有条件)的FID最优表现。
- 训练迭代减少80%,适合算力有限的研究者快速部署生成模型。
近期将昂贵的扩散模型蒸馏为高效少步生成器的方法展现出巨大潜力。然而,这些方法通常需要大量计算资源和长时间训练,限制了资源受限研究者的使用,且现有辅助损失函数存在明显缺陷:回归损失需预先生成大规模数据集,且使学生模型受限于教师性能;GAN-based损失则存在训练不稳、调参复杂等问题。本文提出嵌入损失(Embedding Loss, EL),一种新型辅助损失函数,可与现有扩散蒸馏方法结合,提升生成质量并加速训练,支持更小批量。通过利用一组随机初始化网络的特征嵌入,在嵌入空间中计算最大均值差异(MMD),EL有效对齐蒸馏后的少步生成器与原始数据的特征分布,确保样本保真度与多样性。在分布匹配蒸馏框架下,EL在单步生成器上表现出色。在CIFAR-10上,该方法实现无条件生成1.475、有条件生成1.380的当前最优FID值。此外,我们在ImageNet、AFHQ-v2和FFHQ等多个数据集及DMD、DI、CM等蒸馏框架上验证了其有效性,结果一致优于现有单步蒸馏方法。本方法还可将训练迭代减少高达80%,为资源受限环境下的扩散生成模型部署提供了更实用、可扩展的解决方案。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。