17 天前

提升修正流的训练效果

Sangyun Lee, Zinan Lin, Giulia Fanti
提升修正流的训练效果
摘要

扩散模型在图像与视频生成任务中展现出巨大潜力,但当前最先进的模型在采样时需要对生成性常微分方程(ODE)进行昂贵的数值积分。为应对这一挑战,一种有效方法是修正流(rectified flows),其通过迭代学习平滑的ODE轨迹,从而降低截断误差的影响。然而,现有修正流方法仍需较多的函数求值次数(NFE)。本文提出改进的修正流训练技术,使其在极低NFE设置下也能与知识蒸馏(knowledge distillation)类方法相媲美。我们的核心洞察在于:在实际应用场景中,仅需一次Reflow算法迭代即可学习到近乎直线的轨迹,因此当前普遍采用多轮Reflow迭代的做法实属冗余。为此,我们提出一系列优化单轮训练的技术,包括U型时间步分布策略以及基于LPIPS-Huber的预度量(premetric)。实验表明,在CIFAR-10数据集上,采用上述方法后,单次采样(1 NFE)下的FID指标相比先前的两轮修正流方法提升高达75%。在ImageNet 64×64数据集上,我们改进后的修正流在一步和两步采样设置下均超越了当前最优的知识蒸馏方法(如一致性蒸馏与渐进式蒸馏),并接近改进型一致性训练(iCT)在FID指标上的表现。代码已开源,地址为:https://github.com/sangyun884/rfpp。