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] $$训练流程:
- 随机采样噪声 $Z_0 \sim \mathcal{N}(0, I)$。
- 随机采样数据 $Z_1 \sim Data$。
- 随机采样时间 $t \sim U[0, 1]$。
- 构造中间状态 $Z_t = t Z_1 + (1-t) Z_0$。
- 计算 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-Rectified Flow: 用原始的一队 $(Z_0, Z_1)$ 训练出模型 $v_\theta^1$。
- 生成伪数据对: 用 $v_\theta^1$ 从 $Z_0$ 采样生成对应的 $Z_1'$。这时 $(Z_0, Z_1')$ 这对数据是根据 ODE 轨迹耦合好的 (Coupled)。
- 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 蒸馏成少步数模型。
关键优势
- 训练简单: 无需复杂的噪声调度 (Noise Schedule),$t$ 均匀采样即可。
- 采样高效: 轨迹接近直线,使用 Euler Solver 走几步效果就很好。
- 零信噪比问题 (Zero-SNR): 天然解决了 DDPM 在即使 $t=1$ 时也无法生成纯噪声的问题(RF 定义 $t=0$ 就是纯噪声,$t=1$ 是纯数据,边界清晰)。
相关链接: