[论文解读] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow

Xingchao Liu, Chengyue Gong+1 more|arXiv (Cornell University)|Sep 7, 2022

Model Reduction and Neural Networks被引用 75

一句话总结

修正流通过沿直线路径学习一个 ODE，在两个经验分布之间进行质量传输；通过一个简单的非线性最小二乘目标实现；递归修正产生几乎直线的流， enabling 生成与迁移任务的快速单步推断。

ABSTRACT

We present rectified flow, a surprisingly simple approach to learning (neural) ordinary differential equation (ODE) models to transport between two empirically observed distributions π_0 and π_1, hence providing a unified solution to generative modeling and domain transfer, among various other tasks involving distribution transport. The idea of rectified flow is to learn the ODE to follow the straight paths connecting the points drawn from π_0 and π_1 as much as possible. This is achieved by solving a straightforward nonlinear least squares optimization problem, which can be easily scaled to large models without introducing extra parameters beyond standard supervised learning. The straight paths are special and preferred because they are the shortest paths between two points, and can be simulated exactly without time discretization and hence yield computationally efficient models. We show that the procedure of learning a rectified flow from data, called rectification, turns an arbitrary coupling of π_0 and π_1 to a new deterministic coupling with provably non-increasing convex transport costs. In addition, recursively applying rectification allows us to obtain a sequence of flows with increasingly straight paths, which can be simulated accurately with coarse time discretization in the inference phase. In empirical studies, we show that rectified flow performs superbly on image generation, image-to-image translation, and domain adaptation. In particular, on image generation and translation, our method yields nearly straight flows that give high quality results even with a single Euler discretization step.

研究动机与目标

• 统一生成建模和领域迁移在一个传输映射框架下。
• 引入一种可扩展的基于 ODE 的方法，偏好分布之间的直线、最短路径流。
• 开发一个简单的最小二乘目标，用于学习引导传输的速度场。
• 展示理论性质：非相交流、边际保持，以及降低的凸传输成本。
• 在图像生成、图像到图像翻译和领域自适应方面展示经验有效性。

提出的方法

• 将修正流定义为一个 ODE dZ_t = v(Z_t, t) dt，通过沿接近直线方向的 X_1 - X_0（从耦合数据采样）将 pi_0 传输到 pi_1。
• 通过求解一个非线性最小二乘目标来学习速度场 v，该目标在 t ∈ [0,1] 内最小化 X_1 - X_0 与 v(t X_1 + (1 - t) X_0, t) 的不匹配。
• 通过重新布线交叉轨迹以获得非相交流，同时保持插值过程 X_t 的边际分布来确保因果性。
• 使用学习到的 v 对 pi_0 和 pi_1 进行前向与后向采样；可选地应用再流以获得越来越直的路径。
• 可选的蒸馏过程，用于训练神经网络 T 以模仿 k-rectified flow，实现快速一跳推断。
• 讨论将线性插值替换为一般时间可微的 X_t 的非线性扩展，并将其与概率流 ODE、DDIM 联系起来。

实验结果

研究问题

• RQ1 是否可以用简单的基于 ODE 的传输模型在两个经验分布之间移动质量，且计算开销极小？
• RQ2 修正流是否产生非相交、确定性的耦合，从而降低对于凸成本的传输成本？
• RQ3 递归修正是否产生越来越直的轨迹，使单步或少步推断变得更加准确？
• RQ4 修正流与概率流 ODE 和去噪扩散模型在生成与迁移任务中的关系与差异？
• RQ5 修正流在图像生成、图像到图像翻译和领域自适应方面是否有效？

主要发现

• 修正流提供了一种简单的基于优化的方式，在没有对抗训练或基于似然的推断的情况下学习经验分布之间的传输。
• 得到的流是非相交的，并在任意时刻保持线性插值过程的边际分布。
• 修正产生的耦合在所有凸成本 c 下比输入耦合具有更低的凸传输成本。
• 递归修正产生越来越直的流，使得在粗时间离散化下也能实现精确的模拟，且可能实现单步推断。
• 在经验上，修正流实现了高质量的图像生成和翻译结果，包括报道的最优状态之一的 CIFAR10 一步 FID 分数为 4.85，召回率 0.51。