RectifiedFlow

Rectified Flow (整流)

核心直觉

Rectified Flow (RF) 是一种基于 ODE 的生成模型。 它的核心理念非常简单:两点之间直线最短

  • 我们希望数据从噪声分布 $Z_0$ (Noise) 演化到数据分布 $Z_1$ (Data) 的路径是尽可能直的
  • 路径越直,采样时的数值误差越小,需要的步数就越少 (Few-step sampling)。

1. 理论基础: 从 ODE 出发

Rectified Flow 属于 CNF (Continuous Normalizing Flow) 的一种简化与特例。 我们定义一个随时间变化的变量 $Z_t$,其中 $t \in [0, 1]$。

  • $Z_0 \sim \pi_0$ (通常是标准高斯噪声 $\mathcal{N}(0, I)$)
  • $Z_1 \sim \pi_1$ (真实数据分布)

常微分方程 (ODE) 描述了 $Z_t$ 的演化:

$$ dZ_t = v(Z_t, t) dt $$

其中 $v(Z, t)$ 是速度场 (Velocity Field),也就是我们需要训练的神经网络模型 $v_\theta(Z, t)$。

1.1 为什么是直线?

在没有任何约束的情况下,从 $Z_0$ 到 $Z_1$ 可以有无数种路径。 但在“传输成本”最小化的角度(Optimal Transport),直线路径是最优的。 如果我们假设每一对 $(Z_0, Z_1)$ 都是通过直线连接的,那么中间时刻 $Z_t$ 可以显式写出:

$$ Z_t = t Z_1 + (1-t) Z_0 $$

这条路径的速度(导数)就是:

$$ \frac{dZ_t}{dt} = Z_1 - Z_0 $$

这是一个常数速度(虽然对不同的样本对来说常数不同,但在这一条路径上是常数)。

2. 训练目标: Flow Matching

既然我们想要这样的直线路径,我们就通过Flow Matching来训练网络去拟合在这个直线假设下对应的速度场。

2.1 损失函数

我们要让网络预测的速度 $v_\theta(Z_t, t)$ 尽可能接近真实的速度 $Z_1 - Z_0$。 目标函数 (Loss Function) 就是简单的均方误差 (MSE):

$$ \mathcal{L}(\theta) = \mathbb{E}_{Z_0, Z_1, t} \left[ || v_\theta(Z_t, t) - (Z_1 - Z_0) ||^2 \right] $$

训练流程

  1. 随机采样噪声 $Z_0 \sim \mathcal{N}(0, I)$。
  2. 随机采样数据 $Z_1 \sim Data$。
  3. 随机采样时间 $t \sim U[0, 1]$。
  4. 构造中间状态 $Z_t = t Z_1 + (1-t) Z_0$。
  5. 计算 Loss: $|| v_\theta(Z_t, t) - (Z_1 - Z_0) ||^2$ 并更新网络。
与 DDPM 的区别
  • DDPM: 拟合的是噪声 $\epsilon$ (Score Matching)。
  • Rectified Flow: 拟合的是速度 $v = Z_1 - Z_0$。 在 Flux / SD3 的代码中,通常能看到 model output 直接就是 x_1 - x_0 的形式。

3. Reflow (重流) 过程

虽然我们训练时构造了直线插值,但训练好的神经网络 $v_\theta$ 生成的 ODE 轨迹不一定完全是直的。因为不同的 $(Z_0, Z_1)$ 对可能会在空间中交叉,导致学到的平均速度场产生弯曲。

Reflow 是一种递归优化技术,用来“拉直”这些轨迹。

  1. 1-Rectified Flow: 用原始的一队 $(Z_0, Z_1)$ 训练出模型 $v_\theta^1$。
  2. 生成伪数据对: 用 $v_\theta^1$ 从 $Z_0$ 采样生成对应的 $Z_1'$。这时 $(Z_0, Z_1')$ 这对数据是根据 ODE 轨迹耦合好的 (Coupled)。
  3. 2-Rectified Flow: 用这些新的对 $(Z_0, Z_1')$ 重新训练一个模型 $v_\theta^2$。

结果:经过 Reflow 后的模型,其 ODE 轨迹会变得更直,从而允许我们在采样时使用更大的步长(甚至 1-step 生成)。

实际上,在 Stable Diffusion 3 和 Flux 等现代模型中,通常只进行 1-Rectified Flow (即标准的 Flow Matching 训练) 就已经足够好了,很少再做昂贵的 Reflow。

4. 现代应用 (SD3 & Flux)

Rectified Flow 已经成为新一代文生图模型的标准范式。

  • Stable Diffusion 3 (SD3): 采用了 Rectified Flow 的公式,抛弃了原来 DDPM/LDM 的 $\epsilon$-prediction。
  • Flux: 同样基于 Flow Matching。
  • InstaFlow: 专门利用 Reflow 技术将 SD 蒸馏成少步数模型。

关键优势

  1. 训练简单: 无需复杂的噪声调度 (Noise Schedule),$t$ 均匀采样即可。
  2. 采样高效: 轨迹接近直线,使用 Euler Solver 走几步效果就很好。
  3. 零信噪比问题 (Zero-SNR): 天然解决了 DDPM 在即使 $t=1$ 时也无法生成纯噪声的问题(RF 定义 $t=0$ 就是纯噪声,$t=1$ 是纯数据,边界清晰)。

相关链接:

  • CNF (Rectified Flow 是一种高效训练 CNF 的方法)
  • DDPM (上一代主流范式)
  • Euler (最常用的 RF 采样器)