使用校正流生成并迁移数据.

paper：Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow

0. TL; DR

本文提出了校正流（rectified flow）方法，用于学习（神经）常微分方程（neural ordinary differential equation, ODE）模型，以在两个经验观测到的分布$\pi_0$和$\pi_1$之间进行传输。校正流的思想是，学习一个ODE，使其尽可能地遵循连接从$\pi_0$和$\pi_1$中抽取的点之间的直线路径。这通过求解一个直接的非线性最小二乘优化问题来实现。

作者表明，从数据中学习校正流的过程，称为“校正”（rectification），能将$\pi_0$和$\pi_1$的任意耦合转换为一个新的、确定性的耦合，并且该耦合被证明具有非递增的凸传输成本。此外，递归地应用校正过程（称为Reflow），可以获得一系列路径越来越直的流，这些流在推理阶段可以用粗糙的时间离散化进行精确模拟。

在实证研究中，作者展示了校正流在图像生成、图像到图像翻译和领域自适应方面表现出色。特别是在图像生成和翻译任务上，作者的方法产生了近乎笔直的流，即使只用一个欧拉离散化步骤也能得到高质量的结果。

1. 背景介绍

与监督学习相比，各种形式的无监督学习共同的困难在于缺乏成对的输入/输出数据。大多数无监督方法的核心，是以某种方式找到两个分布中点与点之间有意义的对应关系。这个统一的问题可以被描述为传输映射问题：

给定两个分布$\pi_0$和$\pi_1$的经验观测，找到一个传输映射$T$，使得如果$Z_0 \sim \pi_0$，那么$Z_1 = T(Z_0) \sim \pi_1$。

近年来，通过将传输映射表示为连续时间过程，该领域取得了显著进展，特别是基于ODE的流模型和基于SDE的扩散模型。其中，去噪扩散方法（denoising diffusion methods）在图像生成方面取得了巨大成功，其质量和多样性甚至超过了GAN。这些方法学习到的SDE可以通过概率流ODE（probability flow ODEs）转换为确定性的ODE模型，以实现更快的推理。

然而，与传统的单步模型（如GAN和VAE）相比，连续时间模型的一个关键缺点是推理成本高昂：生成一个样本需要数值求解器反复调用昂贵的神经网络。此外，现有的去噪扩散技术需要大量的超参数搜索，并且其经验和理论基础尚不完全清楚。

为了解决这些问题，作者提出了校正流（Rectified Flow），一个异常简单的、统一解决生成建模和领域迁移的方法。校正流是一个ODE模型，它通过尽可能遵循直线路径来传输分布。直线路径因其理论上的最短性和计算上的高效性（可被单步精确模拟）而备受青睐。

作者的方法通过一个简单的、可扩展的非线性最小二乘优化来训练，避免了GAN的不稳定性和MLE方法的难以处理性。更重要的是，作者引入了一个名为Reflow的递归过程，它能够迭代地“拉直”流的路径。经过一次Reflow后，流就变得近乎笔直，仅用一个欧拉步骤就能生成高质量的样本，极大地弥合了单步模型和连续时间模型之间的差距。

2. 方法介绍

2.1 校正流 (Rectified Flow)

给定从两个分布$\pi_0$和$\pi_1$中采样的点$X_0$和$X_1$，校正流旨在学习一个ODE模型

\[dZ_t = v(Z_t, t)dt\]

使其轨迹尽可能地跟随连接$X_0$和$X_1$的直线方向$(X_1 - X_0)$。这通过求解一个简单的最小二乘回归问题来实现：

\[\min_v \int_0^1 \mathbb{E}\left[ \left\| (X_1 - X_0) - v(X_t, t) \right\|^2 \right] dt, \quad X_t = tX_1 + (1-t)X_0\]

这里的$X_t$是$X_0$和$X_1$的线性插值。在实践中，作者用一个神经网络$v_\theta$来参数化$v$，并用随机梯度下降来优化这个目标。

训练完成后，就可以通过求解ODE来进行采样：

前向采样：从$Z_0 \sim \pi_0$开始，从$t=0$积分到$t=1$，得到$Z_1 \sim \pi_1$。
反向采样：从$Z_1 \sim \pi_1$开始，反向求解ODE，得到$Z_0 \sim \pi_0$。

校正流通过“重新布线”（rewiring）来避免线性插值路径可能出现的交叉，从而构建一个因果的、确定性的流。

2.2 校正流的关键性质

边际保持性 (Marginal Preserving Property)：由校正流生成的耦合$(Z_0, Z_1)$与输入耦合$(X_0, X_1)$具有相同的边际分布。即，如果$Z_0 \sim \pi_0$，那么$Z_1 \sim \pi_1$。这保证了校正流确实是在学习一个从$\pi_0$到$\pi_1$的有效传输。
降低传输成本 (Reducing Transport Costs)：对于任意凸代价函数$c$，校正流生成的耦合$(Z_0, Z_1)$的传输成本不高于输入耦合$(X_0, X_1)$的传输成本。这意味着校正过程天然地在寻找一个更“经济”的传输方案。

\[\mathbb{E}[c(Z_1 - Z_0)] \le \mathbb{E}[c(X_1 - X_0)]\]

2.3 Reflow：迭代地拉直路径

尽管1-rectified flow（即第一次校正得到的流）已经比扩散模型的路径更直，但它仍然是弯曲的。为了获得近乎完美的直线流，作者提出了Reflow程序。

Reflow是一个递归过程：

第一次校正：从一个任意的初始耦合（如独立耦合$(X_0, X_1) \sim \pi_0 \times \pi_1$）开始，训练一个1-rectified flow $Z^1$。
第k+1次校正：使用上一步得到的流$Z^k$生成一个新的数据对$(Z_0^k, Z_1^k)$。这个数据对不再是随机耦合的，而是由一个确定性的ODE连接。将这个新的耦合作为输入，训练一个新的(k+1)-rectified flow $Z^{k+1}$。

每进行一次Reflow，流的轨迹都变得更直。作者从理论上证明了这一点：

路径拉直效应 (Straightening Effect)：作者定义了一个衡量流“弯曲度”的指标$S(Z)$。Reflow过程保证了这个指标会以$O(1/K)$的速度收敛到零，其中$K$是Reflow的迭代次数。

这意味着，通过递归地应用Reflow，我们可以得到任意直的流。在实践中，作者发现2-rectified flow（即经过一次Reflow）的路径就已经近乎是完美的直线，这使得单步生成成为可能。

2.4 与概率流ODE (PF-ODEs) 的联系

作者进一步将校正流的框架推广到非线性插值路径，并由此澄清了其与扩散模型中的概率流ODE（如VP-ODE, sub-VP ODE）的关系。

作者证明，所有现有的PF-ODE变体都可以被看作是使用一个特定的非线性插值$X_t = \alpha_t X_1 + \beta_t \xi$（其中$\xi$是高斯噪声）的非线性校正流的特例。

通过这个统一的框架，作者指出现有PF-ODEs的设计存在两个问题：

路径非直：由于其$\alpha_t, \beta_t$的选择（源于SDE的推导），其路径是弯曲的，且无法通过Reflow拉直。
速度不均：其路径的行进速度不均匀，大部分变化集中在$t$接近1的后期阶段，这使得用大步长进行数值模拟时性能不佳。

相比之下，校正流采用的简单线性插值$X_t = (1-t)X_0 + tX_1$自然地保证了路径的直线性，并且速度是恒定的，这使其在少步数采样时具有天然的优势。

3. 实验分析

作者在一系列任务上对校正流及其Reflow程序进行了广泛的实证研究。

3.1 玩具示例

比较了线性校正流、VP-ODE和sub-VP-ODE在连接两个高斯分布时的轨迹。校正流的轨迹是直线且速度均匀。而VP-ODE和sub-VP-ODE的轨迹是弯曲的，且速度不均（早期慢，后期快），这导致它们在少步数采样时性能不佳。

3.2 无条件图像生成 (CIFAR-10)

当使用自适应步长的RK45求解器精确求解ODE时，1-rectified flow在FID（2.58）和Recall（0.57）指标上均优于所有基于ODE的扩散模型（如VP-ODE, sub-VP ODE），并与完全模拟的SDE模型性能相当。

在单步生成（NFE=1）场景下，Reflow和蒸馏（Distillation）起到了关键作用。2-rectified flow + Distill取得了4.85的FID，显著优于所有其他基于ODE的单步生成方法。

展示了不同NFE下的FID和Recall。Reflow次数越多，在低NFE下的性能越好。例如，2-rectified flow在NFE很小（如2-5步）时，性能就已经非常接近完全求解时的性能。

直观地展示了Reflow的拉直效果。2-rectified flow的路径（用像素值的轨迹表示）几乎是完美的直线。

3.3 高分辨率图像生成

展示了1-rectified flow在多个256x256高分辨率数据集（如LSUN Bedroom/Church, CelebA-HQ）上生成的图像，质量非常高，证明了该方法良好的可扩展性。

展示了一个简单的图像编辑应用。通过将一张“拼接”的猫的图像反向传输到噪声空间，然后对噪声进行简单的修改（使其更接近高斯分布），再正向生成，就可以得到一个更自然的融合图像。

3.4 图像到图像翻译

作者将校正流应用于无配对的图像到图像翻译任务，例如在猫、野生动物、人脸等不同领域之间进行风格转换。展示了在不同领域之间转换的高质量结果。例如，可以将猫脸转换为狐狸、狮子等多种野生动物的脸。2-rectified flow同样在单步生成时就能取得很好的效果。

3.5 领域自适应

作者还将校正流应用于领域自适应（domain adaptation）任务，即将在一个领域（source domain）训练的模型应用于另一个新的领域（target domain）。通过学习一个从目标域到源域的校正流，可以转换测试数据，以减轻领域偏移的影响。

在OfficeHome和DomainNet这两个基准数据集上，1-rectified flow取得了与当时最先进的方法（Deep CORAL）相当甚至更好的分类准确率。