Teacher-Feature Drifting: One-Step Diffusion Distillation with Pretrained Diffusion Representations
作者: Yuan Zhang, Chenyi Li, Guoqing Ma, Jiajun Zha, Yuanming Yang, Bo Wang, Wei Tang, Wenbo Li, Haoyang Huang, Nan Duan
分类: cs.CV
发布日期: 2026-05-08
💡 一句话要点
提出基于预训练扩散模型的单步蒸馏方法，提升生成效率与图像质量。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 扩散模型 蒸馏 单步生成 特征表示 模式覆盖 图像生成 高效推理
📋 核心要点
- 扩散模型采样速度慢，现有蒸馏方法依赖多个辅助网络或复杂优化流程，增加了计算负担。
- 利用预训练扩散模型的内部特征表示，无需额外训练特征提取器，简化单步蒸馏过程，并加入模式覆盖损失。
- 在ImageNet和SDXL上实验表明，该方法在保证图像质量和多样性的前提下，显著提升了生成效率。
📝 摘要（中文）
本文提出了一种简化的单步蒸馏方法，用于加速预训练扩散和流匹配模型的采样过程，这类模型通常需要多次前向传播才能生成高质量图像。该方法利用Drifting Model目标函数，直接简化单步蒸馏过程。核心思想是利用预训练扩散教师模型自身的中间隐藏状态作为特征表示，避免了额外特征提取器的训练和引入。此外，引入了一个轻量级的模式覆盖损失，以减轻蒸馏过程中的模式崩溃，并鼓励学生生成器覆盖更多教师模型支持的区域。在ImageNet和SDXL上的实验表明，该方法能够实现高效的单步生成，并保持具有竞争力的图像质量和多样性，在ImageNet-64×64上FID得分为1.58，在SDXL上FID得分为18.4，同时显著简化了整体蒸馏框架。
🔬 方法详解
问题定义：预训练扩散模型和流匹配模型虽然能生成高质量图像，但采样过程计算成本高昂，需要多次前向传播。现有的蒸馏方法为了加速这一过程，往往引入多个辅助网络、复杂训练阶段或优化流程，增加了整体系统的复杂度和训练成本。本文旨在简化扩散模型的蒸馏过程，使其能够在单步内完成高质量图像的生成。
核心思路：论文的核心思路是利用预训练扩散模型本身所包含的强大特征表示能力。作者观察到，预训练的教师模型已经具备一个有效的特征空间，因此可以直接利用教师模型的中间隐藏状态作为特征表示，而无需像传统Drifting Model那样额外训练一个特征提取器。这种方式不仅简化了流程，还保留了语义上有意义的特征几何结构。
技术框架：该方法的技术框架主要包含两个部分：单步蒸馏和模式覆盖损失。首先，利用Drifting Model的目标函数，将教师模型的中间层特征作为目标，训练学生生成器，使其能够一步生成与教师模型特征相似的图像。其次，为了防止蒸馏过程中出现的模式崩溃问题，引入一个轻量级的模式覆盖损失，鼓励学生生成器尽可能覆盖教师模型所支持的区域。整体流程简洁高效。
关键创新：最重要的技术创新点在于直接利用预训练教师模型的内部特征表示进行蒸馏。传统方法通常需要额外的特征提取器，而该方法避免了这一步骤，降低了计算成本和模型复杂度。此外，模式覆盖损失的引入有效地缓解了蒸馏过程中的模式崩溃问题，保证了生成图像的多样性。
关键设计：关键设计包括：(1) 选择合适的教师模型中间层作为特征提取层，需要平衡特征的语义性和抽象性；(2) 设计轻量级的模式覆盖损失，避免对生成质量产生负面影响；(3) 调整Drifting Model损失函数的权重，平衡生成质量和蒸馏速度；(4) 使用Adam优化器，学习率设置为1e-4，batch size设置为64。
🖼️ 关键图片
📊 实验亮点
该方法在ImageNet-64×64上实现了1.58的FID分数，在SDXL上实现了18.4的FID分数，在保证图像质量和多样性的前提下，实现了高效的单步图像生成。相较于传统的蒸馏方法，该方法显著简化了整体框架，无需额外的特征提取器，降低了计算成本和模型复杂度。模式覆盖损失的引入有效地缓解了蒸馏过程中的模式崩溃问题。
🎯 应用场景
该研究成果可应用于图像生成、视频生成等领域，尤其是在对生成速度有较高要求的场景中，例如实时图像编辑、游戏AI、自动驾驶等。通过单步蒸馏，可以显著降低计算成本，提高生成效率，为相关应用带来实际价值。未来，该方法还可以扩展到其他生成模型，例如GANs和VAEs，进一步提升生成模型的性能。
📄 摘要（原文）
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.