Efficient Diffusion Distillation via Embedding Loss

📄 arXiv: 2604.22379v1 📥 PDF

作者: Jincheng Ying, Yitao Chen, Li Wenlin, Minghui Xu, Yinhao Xiao

分类: cs.CV

发布日期: 2026-04-24

💡 一句话要点

提出Embedding Loss，加速扩散模型蒸馏，提升生成质量，降低计算资源需求。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散模型蒸馏 Embedding Loss 生成模型 最大均值差异 特征嵌入 单步生成 模型加速

📋 核心要点

  1. 现有扩散模型蒸馏方法计算成本高昂，训练时间长，且依赖预生成数据或易训练不稳定，限制了其应用。
  2. 提出Embedding Loss (EL)，通过对齐学生模型和原始数据的特征分布，提升生成质量和训练速度。
  3. 实验表明，EL在CIFAR-10上取得了SOTA的FID，并在ImageNet等数据集上验证了其有效性，同时减少了训练迭代次数。

📝 摘要（中文）

本文提出了一种名为Embedding Loss (EL) 的新型辅助损失函数，旨在提升扩散模型蒸馏的效率和生成质量。现有的蒸馏方法通常需要大量的计算资源和长时间的训练，限制了资源受限的研究人员的使用。回归损失需要预先生成大型数据集，并将学生模型的性能限制在教师模型的水平，而基于GAN的损失则存在训练不稳定和需要仔细调整的问题。EL通过利用来自一组随机初始化的网络的特征嵌入，有效地对齐了精简的生成器和原始数据之间的特征分布。通过计算嵌入特征空间中的最大均值差异（MMD），EL确保了鲁棒的分布匹配，从而在蒸馏过程中保持了样本的保真度和多样性。在分布匹配蒸馏框架内，EL在一阶生成器上表现出强大的经验性能。在CIFAR-10数据集上，我们的方法在无条件生成和条件生成方面分别实现了1.475和1.380的最先进FID值。此外，我们还在ImageNet、AFHQ-v2和FFHQ数据集上，使用DMD、DI和CM蒸馏框架验证了EL，证明了其相对于现有单步蒸馏方法的一致改进。我们的方法还将训练迭代次数减少了高达80%，为在资源受限的环境中部署基于扩散的生成模型提供了一种更实用和可扩展的解决方案。

🔬 方法详解

问题定义：论文旨在解决扩散模型蒸馏过程中计算资源需求高、训练时间长以及现有辅助损失函数（如回归损失和GAN损失）存在的局限性问题。回归损失依赖于预先生成的大型数据集，限制了学生模型的性能上限，而GAN损失则容易出现训练不稳定，需要精细的参数调整。这些问题阻碍了扩散模型在资源受限环境中的部署和应用。

核心思路：论文的核心思路是利用从一组随机初始化的网络中提取的特征嵌入，来对齐蒸馏后的生成器和原始数据之间的特征分布。通过在嵌入特征空间中计算最大均值差异（MMD），可以实现鲁棒的分布匹配，从而在蒸馏过程中保持样本的保真度和多样性。这种方法避免了对预生成数据的依赖，并减少了训练不稳定的风险。

技术框架：整体框架是在现有的分布匹配蒸馏框架（如DMD、DI、CM）的基础上，引入Embedding Loss (EL) 作为辅助损失函数。训练过程中，首先通过随机初始化的网络提取生成器输出和真实数据的特征嵌入，然后在嵌入空间计算MMD损失，并将其与原有的蒸馏损失结合，共同优化学生模型。

关键创新：最重要的技术创新点在于Embedding Loss (EL) 的设计。EL通过利用随机初始化网络的特征提取能力，提供了一种无需预训练或对抗训练即可实现有效分布匹配的方法。与传统的回归损失和GAN损失相比，EL具有更强的鲁棒性和更低的计算成本。

关键设计：EL的关键设计包括：1) 使用多个随机初始化的网络来提取特征嵌入，以增加特征的多样性；2) 在嵌入空间中使用最大均值差异（MMD）作为分布距离度量，MMD能够有效地捕捉分布之间的差异；3) 将EL作为辅助损失函数，与原有的蒸馏损失进行加权组合，共同优化学生模型。具体的权重系数需要根据实验进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在CIFAR-10数据集上，使用EL的单步生成器在无条件生成和条件生成方面分别取得了1.475和1.380的FID值，达到了SOTA水平。此外，EL还在ImageNet、AFHQ-v2和FFHQ等数据集上验证了其有效性，并能够将训练迭代次数减少高达80%，显著提升了训练效率。

🎯 应用场景

该研究成果可广泛应用于图像生成、图像编辑、超分辨率等领域，尤其适用于资源受限的场景，例如移动设备上的图像处理、边缘计算等。通过降低扩散模型的计算成本和训练时间，可以加速其在各个领域的部署和应用，并推动生成式AI技术的发展。

📄 摘要（原文）

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.