LLM2D
改进修正流的训练方法
Improving the Training of Rectified Flows
作者: Sangyun Lee, Zinan Lin, Giulia Fanti
发布日期: 10/10/2024
arXiv ID: oai:arXiv.org:2405.20320v2

摘要

扩散模型在图像和视频生成方面展现出巨大的潜力,但从最先进的模型中采样需要对生成性常微分方程进行代价高昂的数值积分。解决这个问题的一种方法是校正流,它迭代地学习不易受截断误差影响的平滑常微分方程路径。然而,校正流仍然需要相对大量的函数评估次数 (NFE)。在这项工作中,我们提出了改进的校正流训练技术,使其即使在低 NFE 设置下也能与知识蒸馏方法竞争。我们的主要见解是,在现实环境下,用于训练校正流的 Reflow 算法的一次迭代就足以学习近乎直线的轨迹;因此,目前使用多次 Reflow 迭代的做法是不必要的。因此,我们提出了改进校正流单轮训练的技术,包括 U 形时间步长分布和 LPIPS-Huber 预度量。利用这些技术,我们在 CIFAR-10 的 1 NFE 设置下,将之前的 2-校正流的 FID 提高了高达 75%。在 ImageNet 64×64 上,我们改进的校正流在单步和两步设置中都优于最先进的蒸馏方法,如一致性蒸馏和渐进式蒸馏,并在 FID 上与改进的一致性训练 (iCT) 的性能相媲美。代码可在 https://github.com/sangyun884/rfpp 获取。