Huajishe Notes

2-Rectified Flow：路径直线化与少步生成

20 min read

Rectified Flow：路径直线化与少步生成的数学理论

一、从 Flow Matching 到 Rectified Flow：问题的引入

1.1 Flow Matching 的核心成就与遗留问题

在上一份笔记（Flow Matching）中，我们已经看到 Flow Matching 如何将连续流模型的训练从”解 ODE”转变为”向量场回归”。其核心思想是：

  1. 预先定义一条从噪声分布 到数据分布 的概率路径
  2. 通过条件向量场 的回归来训练向量场

遗留问题：即使训练目标简化，推理时我们仍然需要通过 ODE Solver 来生成样本。对于非线性路径，ODE 求解可能需要数十甚至数百步才能保证精度。

1.2 核心观察：为什么路径”弯曲”会导致采样效率低下？

考虑一个简单的一维例子。假设：

  • 起点
  • 终点
  • 路径非线性：

当用 Euler 方法求解时：

如果路径弯曲，相邻的两个点 和 之间的割线与真实路径切线存在较大偏差。这个偏差在每一步都会累积。

数学上：Euler 方法的局部截断误差为 ，全局误差为 。当 增大时（减少步数），误差急剧增加。

1.3 Rectified Flow 的核心思想

Rectified Flow 的核心观察是：

如果我们能够让粒子从 到 的轨迹沿着直线路径运动，那么 ODE 求解就可以退化为单步仿射变换，完全不需要迭代。

一个沿直线运动的粒子满足：

对应的单轨迹内速度场是常数：

此时，Euler 方法的单步更新就是精确解（针对该固定起点终点对）：

其中 。

关键澄清：以上结论仅对固定的一组起点终点对成立。全局层面，不同样本对对应不同的位移向量，因此学习到的全局向量场依旧是关于状态的非线性函数，全局 Lipschitz 常数 。直线路径无法做到理论无条件零误差单步生成，仅能大幅降低积分误差、放宽步长限制。

二、Rectified Flow 的数学框架

2.1 问题形式化

目标：学习一个向量场 ，使得从 出发， ODE 的解 分布在 上。

传统方法的局限：即使 足够强大，可以精确拟合目标分布，推理时的 ODE 求解仍然可能不精确。

Rectified Flow 的洞见：与其训练一个”弯曲”的向量场然后祈祷 ODE Solver 能拟合它，不如直接让向量场沿着直线路径运动。

2.2 线性插值路径

最简单的直线路径定义：

验证：对两边求导：

这确实是一个常数向量场！无论 当前在什么位置，速度永远指向 的方向，大小等于 。

2.3 边际分布与曲率

问题：如果我们强制所有粒子走直线，那还需要学习吗？

答案：是的！因为我们不知道 的配对关系。训练数据通常是 独立采样的，我们需要学习一个从噪声到数据的映射。

定义轨迹曲率（标量形式）：

对于一条参数化路径 ，标量曲率定义为：

对于直线路径 ，有：

结论：直线轨迹的曲率恒为零。

2.4 最优传输视角（可选优化）

为什么独立采样的路径可能造成分布偏移？

设 ， 独立采样。考虑两个起点 和 ，以及它们对应的直线轨迹：

如果 且 ，这两条轨迹会在某个中间点相交。

澄清：不同初始条件的轨迹相交并不违反 ODE 唯一性定理。Picard-Lindelöf 唯一性要求的是同一个初始条件仅对应唯一一条演化轨迹，而非禁止不同起点轨迹相交。高维生成流中，多条轨迹交汇是正常现象。

Re-flow 真正要解决的问题是流折叠（Path Crossing）导致边际分布偏移：当多个起点收敛到同一个终点时，边际分布 会偏离目标分布。Re-flow 通过**重边缘化（Re-marginalization）**保持演化全程的边际一致性。

最优传输（Optimal Transport）配对（可选的高阶优化）：

最优传输理论提出，将 按某种最优方式配对。Monge 问题：

其中 是代价函数， 是边际为 的联合分布集合。

重要说明：原生 RF 最核心的优势是无需任何精确 OT 配对，随机独立采样配对是工业界主流用法。OT 配对仅作为高阶优化方案，用于精细化训练，而非基础必备条件。Flux、SD3 等大模型均未使用 Sinkhorn OT 配对。

三、Re-flow 流程：迭代路径直线化

3.1 固定点方程

关键问题：什么样的向量场能够产生直线轨迹？

定理（Rectified Flow 固定点）：

设向量场 满足：

其中条件分布 是在给定当前位置 时，终点 的后验分布。

则 ODE 轨迹 产生直线运动。

证明：

从 出发的轨迹满足：

如果配对是最优传输映射，即 ，则：

这正是一条直线路径！

3.2 Re-flow 算法

Re-flow 全称 Re-marginalized Flow（重边缘化流），其核心目标是保证演化全程边际分布始终匹配目标分布。

Algorithm: Re-flow

输入：数据分布 ，迭代次数 ，神经网络

初始化：

  1. 采样配对数据 （初始时独立采样）
  2. 定义线性路径：

for to do:

  1. 训练当前向量场： 其中 是上一轮固定的向量场

  2. 生成轨迹样本（如果 ）：

    • 使用当前向量场 解 ODE
    • 从 积分到 ，得到轨迹
    • 用生成样本的轨迹终点作为新终点：

  3. 重新定义路径：

    • 令 （新的直线插值）
    • 此步骤用生成样本重构造无偏插值路径，修正随机独立配对带来的分布偏移

输出：训练好的向量场

3.3 收敛性分析

直线路径与误差的关系：

对于 Euler 方法，当路径曲率 较小时，局部误差约为：

关键观察：

  • 当 （直线路径）时，误差
  • 即使 很大（如 ），误差仍然很小
  • 这意味着我们可以用 10 步而不是 1000 步完成采样

3.4 为什么路径变直后可以用粗步长？

考虑 Euler 一步更新：

局部截断误差为 ，全局误差为 。当路径充分直线化后，误差主要来源于 的表达能力而非数值离散化，因此可以使用更大的步长。

四、Rectified Flow 的目标函数

4.1 基础目标函数

最直接的目标（线性路径假设）：

其中 。

问题：这个目标要求我们事先知道配对 ！

4.2 无配对版本

实际训练中的配对策略：

  1. 独立采样（最常用）：

    优点：简单直接 缺点：可能学到”平凡解”（如 ）

  2. 最优传输配对：

    • 先通过 Sinkhorn 算法计算最优配对
    • 令 是 下的配对样本
    • 然后使用上述目标

  3. 动态配对（Re-flow）：

    • 每轮迭代更新配对
    • 逐步改善轨迹

4.3 损失权重

实践中常用的加权损失：

Logit-Normal 采样：

常用配置：

  • （使 集中在 0.5 附近）
  • （均匀权重）
  • （与信噪比相关）

五、Rectified Flow 与少步采样

5.1 采样步数的数学分析

Euler 方法的全局误差界（标准 Lipschitz 条件）：

对于 Lipschitz 向量场 ，Euler 方法的全局误差为：

其中 是步数， 是终点时间， 是向量场的全局 Lipschitz 常数。

澄清：直线路径仅保证单轨迹内速度为常数。全局层面，不同样本对应不同起点终点，因此学习到的向量场 依旧是关于状态 的非线性函数，全局 。

实际意义：路径直线化后， 更容易学习到近似常数值向量场，使得 Lipschitz 常数大幅降低，从而允许使用更大的步长，而非理论零误差。

5.2 步长选择策略

固定步长：

步数 步长 适用场景
11.0路径完全直线化后
2-40.5-0.25Re-flow 充分收敛
10-200.1-0.05部分直线化
100+<0.01未直线化（原始 FM）

自适应步长：

可以使用标准自适应 ODE 求解器（如 Dormand-Prince RK45）来自动选择步长。在直线化路径上，求解器会自动增大步长。

5.3 与 DDPM/SDPM 的对比

特性DDPMScore SDERectified Flow
采样步数50-100010-10001-10
步长限制必须小（稳定性）必须小（噪声项）可以大（直线化）
理论基础变分推断SDE最优传输
训练目标噪声预测 Score function速度场

六、实际训练中的问题与解决方案

6.1 平凡解问题

问题：当 时，损失函数取得零值，但这对应的是”不移动”的平凡解。

原因：

  • 和 的期望都可能接近零
  • 网络可能学到恒为零的输出

解决方案：

  1. 使用非对称权重初始化
  2. 添加正则项 防止平凡解
  3. 使用噪声感知的训练（如 noise schedule）

6.2 路径交叉问题

问题：独立采样的 可能导致流折叠（Path Crossing），即多个起点可能汇聚到同一个人终点。

核心成因：随机配对下，插值路径分布与真实流分布存在偏移，这是分布层面的问题，而非单纯网络初始化问题。

数学表现：

  • 学习到的映射 不是双射
  • 某些 可能被映射到多个不同的
  • 导致边际分布 偏离目标分布

解决方案：

  1. Re-flow 迭代：逐步改善配对
  2. 最优传输正则：在损失中添加 OT 项
  3. Unbalanced Flow Matching：允许边际分布不完全匹配

6.3 训练不稳定性

问题：某些配置下，训练可能不稳定，特别是在 时。

原因：

  • 噪声端 通常是简单的高斯分布
  • 速度场的 Lipschitz 常数可能在噪声端较大

解决方案：

  1. 使用学习率 warmup
  2. 对极端 值使用更小的学习率权重
  3. 梯度裁剪

6.4 方差估计问题

问题： 的方差随数据维度增加。

对于高维数据：

当维度 很大时，即使每个维度方差很小，总方差也可能很大。

解决方案：

  1. 使用 batch normalization
  2. 归一化目标：
  3. 使用更大的 batch size

七、Rectified Flow 的扩展与变体

7.1 单次迭代与多次迭代 RF

单次迭代 RF：

  • 只进行一次 Re-flow 迭代
  • 路径：从独立采样的直线插值开始，然后训练
  • 通常需要 4-10 步采样

多次迭代 RF：

  • 进行两次或以上 Re-flow 迭代
  • 第二次迭代使用第一次生成的轨迹终点作为新的
  • 通常可以用 2-4 步采样

7.2 条件 Rectified Flow

无条件生成的局限：在实际应用中，我们通常需要条件生成（给定文本描述等条件 ）。

条件 RF：

Classifier-Free Guidance 在 RF 中的应用：

其中 是 guidance 权重。

八、与大模型的结合

8.1 为什么 RF 适合大模型？

理论原因：

  1. 少步采样 = 少梯度步数：大模型的梯度计算代价极高，少步采样意味着少梯度步数
  2. 常数路径 = 简单优化 landscape：直线轨迹对应的损失 landscape 更平滑
  3. 可扩展性：路径直线化与模型规模无关

实证支持：

  • Stable Diffusion 3（12B 参数）使用 RF
  • Flux.1（12B 参数）使用 RF
  • 两者都实现了 4-8 步高质量生成

8.2 与 DiT 的结合

**DiT（Diffusion Transformer）**架构与 Rectified Flow 的结合：

  1. 噪声到潜空间的映射：VAE 编码器将图像压缩到低维潜空间
  2. Transformer 处理条件：文本 embedding 和时间步 embedding
  3. RF 训练：在潜空间中进行向量场回归
  4. 少步采样：解码器从潜空间恢复图像

优势：

  • Transformer 的并行计算能力适合大模型
  • RF 的少步特性降低推理成本
  • 两者结合实现质量与效率的平衡

九、数学推导速查

9.1 核心公式

线性插值路径：

速度场（直线轨迹）：

CFM 损失（RF 版本）：

Re-flow 更新：

最优传输目标：

9.2 误差分析

Euler 方法全局误差界（标准 Lipschitz 条件）：

路径直线化的实际意义：直线化后 更容易学习到近似常数值向量场，使得 Lipschitz 常数 大幅降低，从而允许使用更大步长，而非理论零误差。

十、总结

Rectified Flow 提供了一个朴素而有力的洞察：直线是最短的路径，也是最高效的生成路径。

核心贡献：

  1. 路径直线化：通过 Re-flow 迭代，逐步将路径直线化，使 ODE 求解退化为少步仿射变换
  2. 少步采样：直线化后可用 1-10 步完成高质量生成
  3. 大模型友好：少步采样 = 少梯度步数，适合大参数模型

数学本质：

  • 将 ODE 数值求解的精度问题转化为路径设计问题
  • 从”如何精确求解弯曲的 ODE”转变为”如何让路径更容易求解”
  • 最优传输理论为路径设计提供了理论基础

与 FM 的关系：

  • FM 提供了通用的向量场回归框架
  • RF 是 FM 的一种特殊路径设计选择
  • RF 的线性路径假设简化了 FM 的损失函数

延伸阅读：

  1. Liu et al., “Rectified Flow: A Marginal Preserving Approach to Optimal Transport” (ICML 2023)
  2. Albergo & Vanden-Eijnden, “Building Normalizing Flows with Stochastic Interpolants” (ICLR 2023)
  3. Lipman et al., “Flow Matching for Generative Modeling” (NeurIPS 2022)
  4. Esser et al., “Theketing Diffusion Models: Rectified Flow at Scale” (2024)
  5. Flux.1 Technical Report

Graph View