Improving the Training of Rectified Flows

Improving the Training of Rectified Flows

30 May 2024 | Sangyun Lee, Zinan Lin, Giulia Fanti
This paper addresses the challenge of sampling from state-of-the-art diffusion models, which require expensive numerical integration of a generative ODE. It proposes improved techniques for training rectified flows, a class of simulation-free flow models, to compete with knowledge distillation methods even in the low function evaluation (NFE) setting. The main insight is that under realistic settings, a single iteration of the Reflow algorithm is sufficient to learn nearly straight trajectories, eliminating the need for multiple Reflow iterations. The authors propose several techniques to enhance one-round training of rectified flows, including a U-shaped timestep distribution and an LPIPS-Huber premetric. These techniques significantly improve the performance of rectified flows, reducing the FID score by up to 72% on CIFAR-10. On ImageNet 64×64, the improved rectified flow outperforms state-of-the-art distillation methods such as consistency distillation and progressive distillation, and rivals the performance of improved consistency training (iCT) in terms of FID. The code for the proposed methods is available at <https://github.com/sangyun884/rfpp>.This paper addresses the challenge of sampling from state-of-the-art diffusion models, which require expensive numerical integration of a generative ODE. It proposes improved techniques for training rectified flows, a class of simulation-free flow models, to compete with knowledge distillation methods even in the low function evaluation (NFE) setting. The main insight is that under realistic settings, a single iteration of the Reflow algorithm is sufficient to learn nearly straight trajectories, eliminating the need for multiple Reflow iterations. The authors propose several techniques to enhance one-round training of rectified flows, including a U-shaped timestep distribution and an LPIPS-Huber premetric. These techniques significantly improve the performance of rectified flows, reducing the FID score by up to 72% on CIFAR-10. On ImageNet 64×64, the improved rectified flow outperforms state-of-the-art distillation methods such as consistency distillation and progressive distillation, and rivals the performance of improved consistency training (iCT) in terms of FID. The code for the proposed methods is available at <https://github.com/sangyun884/rfpp>.
Reach us at info@study.space
[slides] Improving the Training of Rectified Flows | StudySpace