Notes/机器学习/深度学习模型/04-扩散模型/6-修正流/1-Rectified-Flow.md

title, draft, tags

title	draft	tags
Rectified Flow：路径直线化与少步生成	false	Rectified-Flow 扩散模型 生成模型 深度学习

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

---

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

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

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

1. 预先定义一条从噪声分布 p_0 到数据分布 p_1 的概率路径 p_t
2. 通过条件向量场 u_t(x_t | x_1) 的回归来训练向量场 v_\theta

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

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

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

• 起点 x_0 \sim \mathcal{N}(-5, 1)
• 终点 x_1 \sim \mathcal{N}(5, 1)
• 路径非线性：x_t = \tanh((1-t) \cdot \text{arctanh}(x_0) + t \cdot \text{arctanh}(x_1))

当用 Euler 方法求解时：

x_{t+\Delta t} = x_t + \Delta t \cdot v_\theta(x_t, t)

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

数学上：Euler 方法的局部截断误差为 $O(\Delta t^2)$，全局误差为 $O(\Delta t)$。当 \Delta t 增大时（减少步数），误差急剧增加。

1.3 Rectified Flow 的核心思想

Rectified Flow 的核心观察是：

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

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

x_t = (1-t)x_0 + t x_1

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

\frac{d x_t}{d t} = x_1 - x_0 = \text{const}

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

x_1 = x_0 + (x_1 - x_0)

其中 $v_\theta(x_0) = x_1 - x_0$。

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

---

二、Rectified Flow 的数学框架

2.1 问题形式化

目标：学习一个向量场 $v_\theta: \mathbb{R}^D \times [0,1] \to \mathbb{R}^D$，使得从 x_0 \sim p_0 出发， ODE \frac{d x(t)}{d t} = v_\theta(x(t), t) 的解 x(1) 分布在 p_1 上。

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

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

2.2 线性插值路径

最简单的直线路径定义：

x_t = (1-t) x_0 + t x_1, \quad t \in [0,1]

验证：对两边求导：

\frac{d x_t}{d t} = -x_0 + x_1 = x_1 - x_0

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

2.3 边际分布与曲率

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

答案：是的！因为我们不知道 \left(x_0, x_1\right) 的配对关系。训练数据通常是 \left\{ x_0^i, x_1^j \right\} 独立采样的，我们需要学习一个从噪声到数据的映射。

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

对于一条参数化路径 $x(t)$，标量曲率定义为：

\kappa(t) = \frac{\|x''(t)\|}{\|x'(t)\|^3}

对于直线路径 $x(t) = (1-t)x_0 + t x_1$，有：

x'(t) = x_1 - x_0, \quad x''(t) = 0 \quad \Rightarrow \quad \kappa(t) = 0

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

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

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

设 $x_0 \sim p_0$，x_1 \sim p_1 独立采样。考虑两个起点 x_0^a 和 $x_0^b$，以及它们对应的直线轨迹：

x_t^a = (1-t)x_0^a + t x_1^a x_t^b = (1-t)x_0^b + t x_1^b

如果 x_1^a \approx x_1^b 且 $x_0^a \neq x_0^b$，这两条轨迹会在某个中间点相交。

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

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

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

最优传输理论提出，将 \left(x_0, x_1\right) 按某种最优方式配对。Monge 问题：

\min_{\pi \in \Pi(p_0, p_1)} \int c(x_0, x_1) d\pi(x_0, x_1)

其中 c(x_0, x_1) = \| x_0 - x_1 \|^2 是代价函数，\Pi 是边际为 p_0, p_1 的联合分布集合。

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

---

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

3.1 固定点方程

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

定理（Rectified Flow 固定点）：

设向量场 v^* 满足：

v^*(x_t) = \mathbb{E}_{x_1 \sim p_1 | x_t} \left[ \frac{x_1 - x_t}{1-t} \right]

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

则 ODE 轨迹 \frac{d x}{d t} = v^*(x(t)) 产生直线运动。

证明：

从 x_0 出发的轨迹满足：

x(t) = x_0 + t \cdot \mathbb{E}[x_1 - x_0 | x_0] = (1-t)x_0 + t \cdot \mathbb{E}[x_1 | x_0]

如果配对是最优传输映射，即 $x_1 = T(x_0)$，则：

x(t) = (1-t)x_0 + t \cdot T(x_0)

这正是一条直线路径！

3.2 Re-flow 算法

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

Algorithm: Re-flow

输入：数据分布 $p_{data}$，迭代次数 $K$，神经网络 v_\theta

初始化：

1. 采样配对数据 $\left{ (x_0^i, x_1^i) \right}$（初始时独立采样）
2. 定义线性路径：x_t^i = (1-t)x_0^i + t x_1^i

for k = 1 to K do:

1. 训练当前向量场：

   \theta^{(k)} = \arg\min_\theta \mathbb{E}_{t, i} \left[ \left\| v_\theta^{(k-1)}(x_t^i, t) - (x_1^i - x_0^i) \right\|^2 \right]

   其中 v_\theta^{(k-1)} 是上一轮固定的向量场

2. 生成轨迹样本（如果 $k < K$）：

   • 使用当前向量场 v_{\theta^{(k)}} 解 ODE
   • 从 x_0 积分到 $t=1$，得到轨迹 \left\{ x(t)^i \right\}
   • 用生成样本的轨迹终点作为新终点：x_1'^i = x(1)^i

3. 重新定义路径：

   • 令 $x_t^i \leftarrow (1-t)x_0^i + t \cdot x_1'^i$（新的直线插值）
   • 此步骤用生成样本重构造无偏插值路径，修正随机独立配对带来的分布偏移

输出：训练好的向量场 v_{\theta^{(K)}}

3.3 收敛性分析

直线路径与误差的关系：

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

\text{Error} \approx \frac{1}{2} \kappa \cdot (\Delta t)^2

关键观察：

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

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

考虑 Euler 一步更新：

x_{t+\Delta t} = x_t + \Delta t \cdot v(x_t)

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

---

四、Rectified Flow 的目标函数

4.1 基础目标函数

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

\mathcal{L}_{RF}(\theta) = \mathbb{E}_{t \sim \mathcal{U}(0,1), x_0 \sim p_0, x_1 \sim p_1} \left[ \| v_\theta(x_t, t) - (x_1 - x_0) \|^2 \right]

其中 $x_t = (1-t)x_0 + t x_1$。

问题：这个目标要求我们事先知道配对 $\left(x_0, x_1\right)$！

4.2 无配对版本

实际训练中的配对策略：

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

   \mathcal{L} = \mathbb{E}_{t, x_0 \sim p_0, x_1 \sim p_1} \left[ \| v_\theta(x_t, t) - (x_1 - x_0) \|^2 \right]

   优点：简单直接 缺点：可能学到"平凡解"（如 $v_\theta \approx 0$）

2. 最优传输配对：

   • 先通过 Sinkhorn 算法计算最优配对 \pi^*
   • 令 x_1 = T(x_0) 是 \pi^* 下的配对样本
   • 然后使用上述目标

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

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

4.3 损失权重

实践中常用的加权损失：

\mathcal{L}_\lambda = \mathbb{E}_{t \sim \text{LogitNormal}(\mu, \sigma^2)} \left[ \lambda(t) \cdot \| v_\theta(x_t, t) - (x_1 - x_0) \|^2 \right]

Logit-Normal 采样：

t = \frac{1}{1 + e^{-z}}, \quad z \sim \mathcal{N}(\mu, \sigma^2)

常用配置：

• $\mu = 0, \sigma = 1.0$（使 t 集中在 0.5 附近）
• $\lambda(t) = 1$（均匀权重）
• $\lambda(t) = 1 - \bar{\alpha}_t$（与信噪比相关）

---

五、Rectified Flow 与少步采样

5.1 采样步数的数学分析

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

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

\| x(T) - x_N \| \leq \frac{T L}{2N} \left( e^{LT} - 1 \right)

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

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

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

5.2 步长选择策略

固定步长：

步数 N	步长 \Delta t = 1/N	适用场景
1	1.0	路径完全直线化后
2-4	0.5-0.25	Re-flow 充分收敛
10-20	0.1-0.05	部分直线化
100+	<0.01	未直线化（原始 FM）

自适应步长：

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

5.3 与 DDPM/SDPM 的对比

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

---

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

6.1 平凡解问题

问题：当 v_\theta(x_t, t) \approx 0 时，损失函数取得零值，但这对应的是"不移动"的平凡解。

原因：

• x_1 - x_0 和 v_\theta 的期望都可能接近零
• 网络可能学到恒为零的输出

解决方案：

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

6.2 路径交叉问题

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

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

数学表现：

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

解决方案：

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

6.3 训练不稳定性

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

原因：

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

解决方案：

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

6.4 方差估计问题

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

对于高维数据：

\mathbb{E}[|x_1 - x_0|^2] = \mathbb{E}[|x_1|^2] + \mathbb{E}[|x_0|^2] - 2 \mathbb{E}[x_1]^T \mathbb{E}[x_0]

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

解决方案：

1. 使用 batch normalization
2. 归一化目标：\| v_\theta - (x_1 - x_0) \|^2 / D
3. 使用更大的 batch size

---

七、Rectified Flow 的扩展与变体

7.1 单次迭代与多次迭代 RF

单次迭代 RF：

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

多次迭代 RF：

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

7.2 条件 Rectified Flow

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

条件 RF：

v_\theta(x_t, t | c) = \mathbb{E}[x_1 - x_0 | x_0, c]

Classifier-Free Guidance 在 RF 中的应用：

\tilde{v}_\theta(x_t, t) = (1 + w) \cdot v_\theta(x_t, t | c) - w \cdot v_\theta(x_t, t | \emptyset)

其中 w 是 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 核心公式

线性插值路径：

x_t = (1-t)x_0 + t x_1

速度场（直线轨迹）：

u_t(x_t | x_1) = x_1 - x_0

CFM 损失（RF 版本）：

\mathcal{L}_{RF} = \mathbb{E}_{t, x_0, x_1} \left[ \| v_\theta(x_t, t) - (x_1 - x_0) \|^2 \right]

Re-flow 更新：

x_1^{(k+1)} = x(1)^{(k)} \quad (trajectory endpoint)

最优传输目标：

\min_\pi \int \| T(x_0) - x_0 \|^2 d\pi(x_0, x_1)

9.2 误差分析

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

\| x(T) - x_N \| \leq \frac{T L}{2N} \left( e^{LT} - 1 \right)

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

---

十、总结

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