论文

超越轨迹匹配:边缘分布对齐回流

Beyond Trajectory Matching: Reflow with Marginal Distribution Alignment

摘要

扩散和连续流生成模型实现了高质量的生成,并且它们的确定性采样可以表示为求解学习的 ODE 动力学。然而,准确的 ODE 离散化通常需要许多步骤,这使得高效的少步生成成为一个关键挑战。在加速策略中,基于回流的 蒸馏 简化了教师 ODE 轨迹,以便学生模型可以用更少的步骤近似教师传输。我们确定了这种范式的理论局限性,即轨迹匹配可能无法确定学生模型引起的分布。特别是,两个学生模型可以获得相同的轨迹匹配损失,同时引入不同的端点边缘分布,这可能导致不同的生成质量。为了解决这个限制,我们引入了一个边际对齐正则化器,它可以惩罚每个 蒸馏 区间端点处学生引起的边际与相应教师边际之间的差异。正则化器是通过跟踪学生模型引起的 ODE 上的对数密度变化并评估冻结教师模型的分数来计算的,无需辅助可训练网络或对抗性优化。由此产生的框架统一适用于回流焊系列,包括普通回流焊和分段回流焊。我们进一步证明了一个伸缩总变异界限,表明局部边缘对齐控制了学生引起的分布和教师引起的分布之间的最终时间差异。对基准骨干网的实验证明了所提出的少步生成方法的有效性。