Towards Training One-Step Diffusion Models Without Distillation

📄 arXiv: 2502.08005v3 📥 PDF

作者: Mingtian Zhang, Wenlin Chen, Jiajun He, Zijing Ou, José Miguel Hernández-Lobato, Bernhard Schölkopf, David Barber

分类: cs.LG, cs.CV

发布日期: 2025-02-11 (更新: 2025-05-27)

备注: 21 pages, 8 figures, 3 tables, 2 algorithms


💡 一句话要点

提出无需教师模型监督的单步扩散模型训练方法,性能超越蒸馏方法。

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

关键词: 单步扩散模型 蒸馏训练 无监督学习 图像生成 模型初始化

📋 核心要点

  1. 现有单步扩散模型训练依赖于教师模型的监督和初始化,流程繁琐且依赖性强。
  2. 论文提出直接训练单步扩散模型的新方法,无需教师模型的得分函数监督,简化训练流程。
  3. 实验表明,新方法在没有教师监督的情况下,性能超越了大多数教师引导的蒸馏方法。

📝 摘要(中文)

近年来,训练单步扩散模型的先进方法通常遵循两阶段流程:首先训练一个教师扩散模型,然后将其蒸馏成一个单步学生模型。这个过程通常依赖于教师模型的得分函数进行监督,并使用其权重初始化学生模型。本文探讨了是否可以直接训练单步扩散模型,而无需这种蒸馏过程。我们提出了一系列新的训练方法,完全放弃了教师得分监督,但性能优于大多数教师引导的蒸馏方法。这表明得分监督对于有效训练单步扩散模型并非必不可少。然而,我们发现用教师模型的权重初始化学生模型仍然至关重要。令人惊讶的是,教师初始化的关键优势并非在于更好的潜在空间到输出空间的映射,而是教师扩散模型提供的跨不同噪声水平的丰富特征表示。这些见解使我们朝着无需蒸馏训练单步扩散模型更进一步,并更好地理解了教师监督和初始化在蒸馏过程中的作用。

🔬 方法详解

问题定义:论文旨在解决单步扩散模型训练过程中对教师模型依赖的问题。现有的单步扩散模型训练方法通常需要先训练一个教师模型,然后通过蒸馏的方式将知识迁移到单步学生模型。这种方法存在两个主要痛点:一是训练流程复杂,需要训练两个模型;二是学生模型的性能受到教师模型的限制。

核心思路:论文的核心思路是探索是否可以直接训练单步扩散模型,而无需教师模型的得分函数监督。作者认为,教师模型的初始化权重可能比得分函数监督更重要。因此,他们提出了一系列新的训练方法,放弃了教师得分监督,但保留了使用教师模型权重初始化学生模型的做法。

技术框架:论文提出的方法主要包括以下几个步骤:首先,训练一个标准的扩散模型作为教师模型。然后,使用教师模型的权重初始化一个单步扩散模型(学生模型)。最后,使用一种新的损失函数直接训练学生模型,而无需教师模型的得分函数监督。整体框架仍然是扩散模型的标准流程,但训练方式发生了改变。

关键创新:论文最重要的技术创新点在于证明了教师模型的得分函数监督对于单步扩散模型的训练并非必不可少。通过实验,作者发现,即使没有教师模型的得分函数监督,只要使用教师模型的权重初始化学生模型,仍然可以获得很好的性能。这表明教师模型的初始化权重包含了有用的信息,可以帮助学生模型更快地收敛。

关键设计:论文的关键设计包括:1)使用教师模型的权重初始化学生模型;2)设计一种新的损失函数,用于直接训练学生模型。具体的损失函数细节在论文中进行了详细描述。此外,论文还对教师模型的选择、学生模型的网络结构等进行了实验分析。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,论文提出的方法在没有教师得分监督的情况下,性能优于大多数教师引导的蒸馏方法。这证明了教师得分监督并非单步扩散模型训练的必要条件。此外,实验还发现,教师模型的初始化权重对于学生模型的性能至关重要,即使没有教师得分监督,使用教师模型的权重初始化学生模型仍然可以获得很好的性能。

🎯 应用场景

该研究成果可应用于图像生成、图像编辑、音频合成等领域。通过简化单步扩散模型的训练流程,降低了训练成本,使得扩散模型更容易被应用到各种实际场景中。未来,该研究有望推动扩散模型在低资源设备上的部署,并促进其在更多领域的应用。

📄 摘要(原文)

Recent advances in training one-step diffusion models typically follow a two-stage pipeline: first training a teacher diffusion model and then distilling it into a one-step student model. This process often depends on both the teacher's score function for supervision and its weights for initializing the student model. In this paper, we explore whether one-step diffusion models can be trained directly without this distillation procedure. We introduce a family of new training methods that entirely forgo teacher score supervision, yet outperforms most teacher-guided distillation approaches. This suggests that score supervision is not essential for effective training of one-step diffusion models. However, we find that initializing the student model with the teacher's weights remains critical. Surprisingly, the key advantage of teacher initialization is not due to better latent-to-output mappings, but rather the rich set of feature representations across different noise levels that the teacher diffusion model provides. These insights take us one step closer towards training one-step diffusion models without distillation and provide a better understanding of the roles of teacher supervision and initialization in the distillation process.