616 lines
25 KiB
Markdown
616 lines
25 KiB
Markdown
---
|
||
title: Neural ODE + FFJORD:从离散到连续的标准化流
|
||
draft: false
|
||
tags:
|
||
- Neural-ODE
|
||
- FFJORD
|
||
- CNF
|
||
- 标准化流
|
||
- 生成模型
|
||
- 深度学习
|
||
---
|
||
|
||
# Neural ODE + FFJORD + Score-Based SDE:连续标准化流的数学基础与统一理论
|
||
|
||
---
|
||
|
||
## 一、从离散流到连续流的哲学跃迁
|
||
|
||
### 1.1 离散标准化流的核心瓶颈
|
||
|
||
在讨论 Neural ODE 之前,我们需要回顾离散标准化流(Discrete Normalizing Flow)的根本限制。
|
||
|
||
设我们有一个由 $K$ 个耦合层组成的变换链:
|
||
$$x = f_K \circ f_{K-1} \circ \cdots \circ f_1(z_0)$$
|
||
|
||
其中每个变换 $f_k$ 都是可逆的。由变量代换公式,完整变换的对数概率密度为:
|
||
|
||
$$\log p_X(x) = \log p_{Z_0}(z_0) - \sum_{k=1}^{K} \log \left| \det \frac{\partial f_k}{\partial z_{k-1}} \right|$$
|
||
|
||
**关键问题**:计算 $\det \frac{\partial f_k}{\partial z_{k-1}}$ 的复杂度为 $O(D^3)$($D$ 为维度)。即使 Coupling Layer 将其降为 $O(D)$,整个流程仍然需要 $K$ 次这样的计算。
|
||
|
||
### 1.2 连续化的动机:无穷层叠加
|
||
|
||
**核心洞察**:当层数 $K \to \infty$ 且每层的步长 $\Delta t \to 0$ 时,有限次离散的 Jacobian 行列式计算会连续地累积为一次积分。
|
||
|
||
设第 $k$ 层变换为 $z_{k+1} = z_k + f(z_k, t_k; \theta) \Delta t$,其中 $t_k = k \Delta t$。
|
||
|
||
当 $\Delta t \to 0$ 时,记 $z(t) = z_k$(连续时间极限),则:
|
||
|
||
$$\frac{d z(t)}{d t} = f(z(t), t; \theta)$$
|
||
|
||
这正是 **Neural ODE** 的核心方程。
|
||
|
||
---
|
||
|
||
## 二、Neural ODE:连续时间神经网络的数学框架
|
||
|
||
### 2.1 形式化定义
|
||
|
||
**Neural ODE**(Chen et al., 2018)将数据演化建模为一个一阶常微分方程:
|
||
|
||
$$\frac{d x(t)}{d t} = f_\theta(x(t), t)$$
|
||
|
||
其中:
|
||
- $x(t) \in \mathbb{R}^D$ 是时刻 $t$ 的隐变量状态
|
||
- $f_\theta: \mathbb{R}^D \times \mathbb{R} \to \mathbb{R}^D$ 是由神经网络参数化的向量场(velocity field)
|
||
- $\theta$ 是待学习的参数
|
||
- 初始条件:$x(0) = z_0 \sim p_0$(先验噪声分布)
|
||
- 终止条件:$x(1) = x$(目标数据分布)
|
||
|
||
**物理直觉**:$f_\theta(x(t), t)$ 描述了在 $t$ 时刻、位置 $x(t)$ 处,速度场赋予粒子的速度向量。ODE 的解轨迹 $x(t)_{t \in [0,1]}$ 描述了粒子从噪声运动到数据的路径。
|
||
|
||
### 2.2 ODE 求解的数值方法
|
||
|
||
#### Euler 方法(一阶)
|
||
|
||
$$x_{n+1} = x_n + \Delta t \cdot f_\theta(x_n, t_n)$$
|
||
|
||
误差:$O(\Delta t)$。当 $\Delta t$ 足够小时可用,但精度较低。
|
||
|
||
#### Runge-Kutta 4 阶(RK4)
|
||
|
||
$$k_1 = f_\theta(x_n, t_n)$$
|
||
$$k_2 = f_\theta(x_n + \frac{\Delta t}{2} k_1, t_n + \frac{\Delta t}{2})$$
|
||
$$k_3 = f_\theta(x_n + \frac{\Delta t}{2} k_2, t_n + \frac{\Delta t}{2})$$
|
||
$$k_4 = f_\theta(x_n + \Delta t \cdot k_3, t_n + \Delta t)$$
|
||
$$x_{n+1} = x_n + \frac{\Delta t}{6}(k_1 + 2k_2 + 2k_3 + k_4)$$
|
||
|
||
误差:$O(\Delta t^4)$。标准的高精度方法。
|
||
|
||
#### 自适应步长方法(Runge-Kutta-Fehlberg 45)
|
||
|
||
通过比较 4 阶和 5 阶 RK 估计的差异来自动调节步长,在精度和效率间取得平衡。
|
||
|
||
### 2.3 可逆性保证:Picard-Lindelöf 定理
|
||
|
||
**定理(Picard-Lindelöf)**:若向量场 $f_\theta(x, t)$ 关于 $x$ 满足 **Lipschitz 连续条件**:
|
||
$$\| f_\theta(x, t) - f_\theta(y, t) \| \leq L \| x - y \|, \quad \forall x, y \in \mathbb{R}^D, t \in [0,1]$$
|
||
|
||
则 ODE $\frac{dx}{dt} = f_\theta(x, t)$ 存在唯一解,且解对初始条件连续依赖。
|
||
|
||
**对可逆性的保证**:
|
||
- 唯一性 $\Rightarrow$ 不存在两条不同的流线在 $t > 0$ 后交汇
|
||
- 连续依赖 $\Rightarrow$ 初始条件的微小变化只会导致解的微小变化
|
||
|
||
这从数学上保证了 **双射(Diffeomorphism)** 性质:流 $x(t; x_0)$ 是可逆的。逆向过程通过解**伴随 ODE** 实现:
|
||
$$\frac{d x(t)}{d t} = -f_\theta(x(t), t; \theta)$$
|
||
|
||
---
|
||
|
||
## 三、连续标准化流(CNF):核心推导
|
||
|
||
### 3.1 瞬时变量代换公式(Instantaneous Change of Variables)
|
||
|
||
**定理**:设 $x(t) \in \mathbb{R}^D$ 服从 ODE $\frac{dx(t)}{dt} = f(x(t), t)$,概率密度 $p_t(x(t))$ 对时间的演化满足:
|
||
|
||
$$\frac{\partial \log p_t(x(t))}{\partial t} = -\text{tr}\left( \frac{\partial f(x(t), t)}{\partial x(t)} \right)$$
|
||
|
||
其中 $\text{tr}(\cdot)$ 是矩阵的迹(Trace)。
|
||
|
||
**物理直觉**:散度 $\text{tr}(\frac{\partial f}{\partial x}) = \nabla \cdot f$ 衡量向量场在每一点的"源"与"汇"。若散度为正(膨胀),概率密度降低;若散度为负(收缩),概率密度升高。这与质量守恒一致。
|
||
|
||
**严格推导**:
|
||
|
||
设 $J(t) = \frac{\partial x(t)}{\partial x(0)}$,则 $J(t)$ 满足矩阵微分方程(对 ODE 初始条件求偏导):
|
||
$$\frac{d J(t)}{d t} = \frac{\partial f}{\partial x} J(t)$$
|
||
|
||
对 $J(t)$ 的行列式取对数并求导,利用矩阵行列式的对数导数公式 $\frac{d}{dt}\log\det J = \text{tr}(J^{-1}\frac{dJ}{dt})$:
|
||
|
||
$$\frac{d}{dt} \log \det J(t) = \text{tr}\left( J^{-1} \frac{d J}{dt} \right) = \text{tr}\left( J^{-1} \frac{\partial f}{\partial x} J \right) = \text{tr}\left( \frac{\partial f}{\partial x} \right)$$
|
||
|
||
其中最后一步利用了 $\text{tr}(AB) = \text{tr}(BA)$。
|
||
|
||
概率质量守恒:$p_t(x(t)) d x(t) = p_0(x(0)) d x(0)$
|
||
|
||
取体积元比例:
|
||
$$\frac{p_t(x(t))}{p_0(x(0))} = \frac{d x(0)}{d x(t)} = \frac{1}{\det J(t)}$$
|
||
|
||
取对数并对 $t$ 求导:
|
||
|
||
$$\frac{\partial \log p_t}{\partial t} = -\frac{d}{dt} \log \det J(t) = -\text{tr}\left( \frac{\partial f}{\partial x} \right)$$
|
||
|
||
### 3.2 对数似然的积分形式
|
||
|
||
从 $t=0$ 积分到 $t=T$:
|
||
|
||
$$\log p_T(x(T)) = \log p_0(x(0)) - \int_0^T \text{tr}\left( \frac{\partial f(x(t), t)}{\partial x(t)} \right) d t$$
|
||
|
||
**解释**:
|
||
- $\log p_0(x(0))$ 是先验分布的对数概率(已知)
|
||
- 积分项是 trace 的累积,描述了整个路径上的体积变化
|
||
|
||
**与离散流的关系**:在离散 NF 中,第 $k$ 层变换的 log-determinant 为 $\log |\det J_k|$。在连续情形下,这些离散的行列式对数被连续时间的 trace 积分所替代:
|
||
|
||
$$\log \det \frac{\partial x(T)}{\partial x(0)} = \int_0^T \text{tr}\left( \frac{\partial f(x(t), t)}{\partial x(t)} \right) d t$$
|
||
|
||
### 3.3 迹算子为何取代行列式
|
||
|
||
| 操作 | 离散 NF | 连续 NF (CNF) |
|
||
|------|----------|----------------|
|
||
| Jacobian 结构 | 需精心设计的稀疏结构(如三角矩阵) | 任意结构的 $f_\theta$ |
|
||
| 计算复杂度 | $O(D^3)$(通用)$/ O(D)$(三角结构) | $O(D)$(只需对角线元素) |
|
||
| 层数依赖 | 固定层数 $K$ | 连续时间,无穷层 |
|
||
|
||
**核心突破**:迹 $\text{tr}\left( \frac{\partial f}{\partial x} \right)$ 只需计算 Jacobian 矩阵的对角线元素(偏导数 $\partial f_i / \partial x_i$ 之和),而行列式需要计算所有特征值的乘积。
|
||
|
||
---
|
||
|
||
## 四、FFJORD:Hutchinson 迹估计器与无似然训练
|
||
|
||
### 4.1 核心问题:迹的精确计算仍是瓶颈
|
||
|
||
即使 CNF 将行列式 $O(D^3)$ 降为迹 $O(D)$,对于高维数据(如图像 $D = 64 \times 64 \times 3 = 12288$),精确计算 $\text{tr}(\frac{\partial f}{\partial x})$ 仍然代价高昂。
|
||
|
||
**精确迹的计算**:需要计算所有 $D$ 个偏导数 $\partial f_i / \partial x_i$。
|
||
|
||
### 4.2 Hutchinson 迹估计器:随机正交投影
|
||
|
||
**随机化技巧**(Hutchinson, 1990):迹可以通过随机投影来估计。
|
||
|
||
**定理**:对于任意矩阵 $A \in \mathbb{R}^{D \times D}$,有:
|
||
$$\text{tr}(A) = \mathbb{E}_{p} \left[ v^T A v \right]$$
|
||
|
||
其中 $v \in \mathbb{R}^D$ 是任意满足 $\mathbb{E}[v] = 0$ 且 $\mathbb{E}[v v^T] = I$ 的随机向量。
|
||
|
||
**常用选择**:Rademacher 分布($v_i = \pm 1$ with equal probability)或标准正态分布。
|
||
|
||
**验证**:取 $v$ 为标准基向量 $e_i$ 时,$v^T A v = A_{ii}$,故 $\mathbb{E}[v^T A v] = \sum_i A_{ii} = \text{tr}(A)$。
|
||
|
||
**对角线估计公式**:
|
||
$$\text{tr}\left( \frac{\partial f_\theta}{\partial x} \right) \approx \frac{1}{M} \sum_{j=1}^{M} v_j^T \frac{\partial f_\theta}{\partial x} v_j = \frac{1}{M} \sum_{j=1}^{M} \sum_{i=1}^D \frac{\partial f_i}{\partial x_i} v_{j,i}^2$$
|
||
|
||
其中 $v_j$ 是第 $j$ 个随机向量。
|
||
|
||
**复杂度优势**:计算 $v^T \frac{\partial f}{\partial x} v$ 可以通过一次反向传播完成(对于每个随机向量 $v$,先计算 $\frac{\partial f}{\partial x} v$,再与 $v$ 做点积)。
|
||
|
||
### 4.3 FFJORD 的训练目标
|
||
|
||
**FFJORD**(Grathohl et al., 2019)将 Hutchinson 迹估计器引入 CNF,实现了完全无需精确 Jacobian 计算的连续标准化流训练。
|
||
|
||
**原始损失(精确迹)**:
|
||
$$\mathcal{L}(\theta) = -\mathbb{E}_{x \sim p_{data}} \left[ \log p_\theta(x) \right]$$
|
||
|
||
其中 $\log p_\theta(x) = \log p_0(z(0)) - \int_0^T \text{tr}\left( \frac{\partial f}{\partial x} \right) d t$
|
||
|
||
**FFJORD 损失(随机迹)**:
|
||
$$\mathcal{L}_{FFJORD}(\theta) = -\mathbb{E}_{x \sim p_{data}, \{v_j\}_{j=1}^M} \left[ \log p_0(z(0)) - \int_0^T \frac{1}{M} \sum_{j=1}^M v_j^T \frac{\partial f}{\partial x} v_j \, d t \right]$$
|
||
|
||
**实现细节**:
|
||
1. 从数据 $x$ 逆向积分到 $z(0)$(ODE 逆向求解)
|
||
2. 对每个随机向量 $v_j$,计算迹估计
|
||
3. 累积得到无偏估计的似然梯度
|
||
|
||
### 4.4 训练流程与潜在问题
|
||
|
||
**Algorithm: FFJORD Training**
|
||
|
||
**输入**:数据分布 $p_{data}$,向量场网络 $f_\theta$,迹估计随机向量数 $M$,ODE 步数 $N$
|
||
|
||
**for** each batch **do**:
|
||
1. 采样 $x \sim p_{data}$
|
||
2. 逆向 ODE 求解:从 $x$ 逆向积分到 $z(0)$
|
||
3. 计算先验对数概率 $\log p_0(z(0))$
|
||
4. **for** $j = 1$ to $M$ **do**:
|
||
- 采样随机向量 $v_j$
|
||
- 计算迹估计项 $\hat{tr}_j = v_j^T \frac{\partial f}{\partial x} v_j$
|
||
5. 累积积分 $\hat{\mathcal{T}} = \frac{1}{M} \sum_{j=1}^M \int_0^T \hat{tr}_j \, d t$
|
||
6. 更新损失 $\mathcal{L} = -(\log p_0(z(0)) - \hat{\mathcal{T}})$
|
||
7. 反向传播更新 $\theta$
|
||
|
||
**潜在问题**:
|
||
|
||
1. **方差问题**:随机迹估计引入方差。$M$ 越大估计越稳定,但计算成本也越高。通常 $M=1$ 足够(单次随机估计),但可能需要更大的 batch 来抵消方差。
|
||
|
||
2. **NFE 问题**:训练时需要逆向 ODE 求解,NFE 仍然可能达到数百。自适应求解器(如 Dormand-Prince)可以缓解但不能根本解决。
|
||
|
||
3. **刚度(Stiffness)**:某些向量场会导致 ODE 刚度增加,需要更小的步长或特殊求解器。
|
||
|
||
4. **数值精度**:积分过程中的数值误差会累积,影响似然估计的准确性。
|
||
|
||
---
|
||
|
||
## 五、Score-Based Generative Modeling through SDEs:扩散模型的统一理论
|
||
|
||
### 5.1 从 ODE 到 SDE:引入随机性的必然性
|
||
|
||
在 Neural ODE / CNF 中,变换是完全确定性的:
|
||
$$\frac{d x(t)}{d t} = f_\theta(x(t), t)$$
|
||
|
||
然而,扩散模型的前向过程是**随机的**(每一步都注入独立的高斯噪声)。为了统一这两个世界,我们需要将 ODE 扩展为 **随机微分方程(SDE)**。
|
||
|
||
### 5.2 前向 SDE(Forward Process)
|
||
|
||
设前向过程满足 Itô SDE:
|
||
$$d x(t) = f_t(x(t)) d t + g_t(x(t)) d W(t)$$
|
||
|
||
其中:
|
||
- $W(t)$ 是维纳过程(Wiener Process,即标准布朗运动)
|
||
- $f_t: \mathbb{R}^D \to \mathbb{R}^D$ 是漂移系数(drift coefficient)
|
||
- $g_t: \mathbb{R}^D \to \mathbb{R}^{D \times D}$ 是扩散系数(diffusion coefficient)
|
||
|
||
**标准前向 SDE**(DDPM 类):
|
||
$$d x(t) = -\frac{1}{2} \beta(t) x(t) d t + \sqrt{\beta(t)} d W(t)$$
|
||
|
||
其中 $\beta(t)$ 是噪声调度函数(通常为线性或余弦调度)。
|
||
|
||
**对应的离散化形式**(Euler-Maruyama 方法):
|
||
$$x_{t+1} = x_t - \frac{1}{2} \beta_t x_t + \sqrt{\beta_t} \epsilon_t, \quad \epsilon_t \sim \mathcal{N}(0, I)$$
|
||
|
||
### 5.3 Reverse-Time SDE:反向生成的数学推导
|
||
|
||
**核心定理**:对于任意前向 Itô SDE:
|
||
$$d x(t) = f(x(t), t) d t + g(x(t), t) d W(t)$$
|
||
|
||
对应的**反向时间 SDE** 为:
|
||
$$d x(t) = \left[ f(x(t), t) - g(x(t), t) g^T(x(t), t) \nabla_x \log p_t(x(t)) \right] d t + g(x(t), t) d \bar{W}(t)$$
|
||
|
||
其中:
|
||
- $\bar{W}(t)$ 是反向时间的维纳过程
|
||
- $\nabla_x \log p_t(x(t))$ 是 **score function**(对数概率密度的梯度)
|
||
- $g g^T \nabla_x \log p_t$ 是 "逆风"项,抵消前向扩散的效应
|
||
|
||
**推导要点**:
|
||
|
||
设前向过程从 $x(0) \sim p_0$ 演化到 $x(T) \sim p_T$。我们需要找到一个反向过程,从 $p_T$ 回到 $p_0$。
|
||
|
||
由贝叶斯定理,在反向时间我们有:
|
||
$$p_{T-t}(x(t)) \propto p_T(x(t)) \cdot \frac{p_{data}(x(0))}{p_{data}(x(T))}$$
|
||
|
||
对 $t$ 求导并利用前向 SDE 的 Fokker-Planck 方程,可以推导出反向 SDE 的形式。
|
||
|
||
**物理直觉**:score function $\nabla_x \log p_t$ 指向概率密度增长的方向。前向 SDE 添加噪声(增加熵),反向 SDE 需要减去熵(沿 score 方向去噪)。
|
||
|
||
### 5.4 Probability Flow ODE:确定性等价的数学基础
|
||
|
||
**关键发现**:对于任意 SDE,存在一个**等价的确定性 ODE**,它产生相同的边际分布 $\{p_t\}$。
|
||
|
||
**Probability Flow ODE**(Song et al., 2021):
|
||
$$d x(t) = \left[ f(x(t), t) - \frac{1}{2} g(x(t), t) g^T(x(t), t) \nabla_x \log p_t(x(t)) \right] d t$$
|
||
|
||
**对比**:
|
||
- Reverse-Time SDE:$d x(t) = [\cdots] d t + g d \bar{W}(t)$(随机)
|
||
- Probability Flow ODE:$d x(t) = [\cdots] d t$(确定性)
|
||
|
||
**为什么它们等价?**
|
||
|
||
两者对应的 Fokker-Planck 方程给出相同的解 $p_t$。概率流 ODE 通过将随机项替换为"drift correction"项来消除随机性。
|
||
|
||
**与 DDPM 的联系**:
|
||
|
||
对于 DDPM 风格的 SDE $d x = -\frac{1}{2} \beta(t) x d t + \sqrt{\beta(t)} d W$:
|
||
|
||
Probability Flow ODE 为:
|
||
$$d x = -\frac{1}{2} \beta(t) \left[ x + \nabla_x \log p_t(x) \right] d t$$
|
||
|
||
这正是 DDPM 反向过程在连续时间的形式!
|
||
|
||
### 5.5 Score Matching:训练目标的推导
|
||
|
||
**Score Function**:$\nabla_x \log p_t(x)$ 是对数概率密度关于 $x$ 的梯度,指向概率密度增长最快的方向。
|
||
|
||
**三种等价的 score matching 训练目标**:
|
||
|
||
#### 5.5.1 切片得分匹配(Slice Score Matching)
|
||
$$\mathcal{L}_{SSM} = \mathbb{E}_{t, x \sim p_t, v \sim \mathcal{N}(0,I)} \left[ v^T \nabla_x \log p_t(x) - \frac{1}{2} \| v \|^2 \right]^2$$
|
||
|
||
#### 5.5.2 切片得分匹配 denoising 版本
|
||
$$\mathcal{L}_{DSM} = \mathbb{E}_{t, x_0 \sim p_{data}, \epsilon \sim \mathcal{N}(0,I)} \left[ \| s_\theta(x_t, t) - \nabla_{x_t} \log p_{0|t}(x_t | x_0) \|^2 \right]$$
|
||
|
||
其中 $s_\theta(x_t, t)$ 是网络预测的 score function,$\nabla_{x_t} \log p_{0|t}(x_t | x_0) = -\frac{\epsilon}{\sigma_t}$。
|
||
|
||
#### 5.5.3 随机变分推断(Score VI)
|
||
$$\mathcal{L}_{SVI} = \mathbb{E}_{t} \left[ \lambda(t) \cdot \mathbb{E}_{x_0, x_t} \left[ \| s_\theta(x_t, t) + \frac{x_t - \sqrt{\bar{\alpha}_t} x_0}{\sigma_t^2} \|^2 \right] \right]$$
|
||
|
||
### 5.6 采样过程
|
||
|
||
**从 Score 到生成**:给定训练好的 score network $s_\theta(x, t) \approx \nabla_x \log p_t(x)$,我们有三种采样范式:
|
||
|
||
#### 5.6.1 Reverse-Time SDE 采样
|
||
使用数值 SDE 求解器(如 Euler-Maruyama)沿反向时间 ODE + 随机项求解。
|
||
|
||
#### 5.6.2 Probability Flow ODE 采样
|
||
使用确定性 ODE 求解器(可使用更少步数,因为没有随机性)。
|
||
|
||
#### 5.6.3 predictor-corrector 方法
|
||
结合两者:先用 ODE predictor 预测,再用 SDE corrector 校正。
|
||
|
||
**Algorithm: Score-Based SDE Sampling**
|
||
|
||
1. **初始化**:$x_T \sim p_T$(纯噪声)
|
||
2. **反向时间迭代**:对 $t = T, T-\Delta, \dots, 0$:
|
||
- **Predictor**:使用 ODE step 更新 $x$
|
||
- **Corrector**:使用 score 修正(例:Langevin Monte Carlo)
|
||
3. **输出**:$x_0$ 即为生成样本
|
||
|
||
---
|
||
|
||
## 六、Flow Matching:更一般的训练视角
|
||
|
||
### 6.1 从 SDE/ODE 约束中解放
|
||
|
||
无论是 Neural ODE 还是 Score-Based SDE,训练过程都受到微分方程数值解的约束:
|
||
|
||
- **Neural ODE**:需要前向/反向 ODE 求解
|
||
- **Score-Based SDE**:需要 predictor-corrector 采样
|
||
|
||
**Flow Matching**(Lipman et al., 2022)提出一个根本性问题:
|
||
|
||
> 能否直接拟合驱动数据分布演化的**向量场**,而不经过任何微分方程求解?
|
||
|
||
### 6.2 概率路径(Probability Path)的定义
|
||
|
||
设 $p_0$ 为噪声分布(通常 $\mathcal{N}(0, I)$),$p_1$ 为数据分布。
|
||
|
||
Flow Matching 引入一个随时间 $t \in [0,1]$ 演化的**边际分布族** $p_t$,满足:
|
||
$$p_{t=0} = p_0, \quad p_{t=1} = p_1$$
|
||
|
||
**直觉**:$p_t$ 描述了从噪声到数据的演化过程中,每个时刻 $t$ 所有样本的概率密度。
|
||
|
||
### 6.3 条件向量场与边际向量场
|
||
|
||
**条件向量场(Conditional Vector Field)**:$u_t(x | x_1)$
|
||
|
||
对于从 $x_0 \sim p_0$ 演化到 $x_1 \sim p_1$ 的特定粒子,其在时刻 $t$ 的速度向量定义为 $u_t(x | x_1)$。
|
||
|
||
**边际向量场(Marginal Vector Field)**:$v_t(x)$
|
||
|
||
边际向量场是所有条件向量场的**期望**:
|
||
$$v_t(x) = \mathbb{E}_{x_1 \sim p_1} \left[ u_t(x | x_1) \cdot w(x, x_1) \right]$$
|
||
|
||
其中 $w$ 是与路径定义相关的权重函数。
|
||
|
||
**核心目标**:找到 $v_\theta$ 使得 $v_\theta \approx v_t$。
|
||
|
||
### 6.4 Conditional Flow Matching(CFM)损失函数
|
||
|
||
**直接监督边际向量场是困难的**——因为 $p_t$ 的解析形式未知。
|
||
|
||
**关键洞察**:条件匹配等价于边际匹配!
|
||
|
||
**CFM 损失函数**:
|
||
$$\mathcal{L}_{CFM}(\theta) = \mathbb{E}_{t \sim \mathcal{U}(0,1), \, x_1 \sim p_{data}, \, x_t \sim p_t(\cdot | x_1)} \left[ \| v_\theta(x_t, t) - u_t(x_t | x_1) \|^2 \right]$$
|
||
|
||
其中 $x_t$ 是根据预设路径 $p_t$ 从 $x_0$ 采样得到的中间状态。
|
||
|
||
### 6.5 关键定理:条件匹配等价于边际匹配
|
||
|
||
**定理(Flow Matching Replacement Theorem)**:
|
||
|
||
设 $u_t(x | x_1)$ 为条件向量场,$v_t(x)$ 为对应的边际向量场。则:
|
||
$$\mathbb{E}_{t, x_t} \| v_\theta(x_t, t) - v_t(x_t) \|^2 = \mathbb{E}_{t, x_t, x_1} \| v_\theta(x_t, t) - u_t(x_t | x_1) \|^2$$
|
||
|
||
**证明**:
|
||
|
||
对于固定时刻 $t$ 和 $x_t$,边际向量场可以写成条件向量场的条件期望:
|
||
$$v_t(x_t) = \mathbb{E}_{x_1} \left[ u_t(x_t | x_1) | x_t \right]$$
|
||
|
||
展开误差平方:
|
||
$$\| v_\theta - v_t \|^2 = \mathbb{E}_{x_1} \left[ \| v_\theta - u_t(x_t | x_1) \|^2 | x_t \right] - \| \mathbb{E}_{x_1}[u_t | x_t] - v_\theta \|^2$$
|
||
|
||
第二项是常数(非负),因此最小化左侧等价于最小化右侧的期望。
|
||
|
||
**直觉**:这就像是说,"所有人对正确答案的误差平方的平均"最小化,等价于"每个人对正确答案的误差平方"最小化。
|
||
|
||
### 6.6 扩散路径作为特例
|
||
|
||
标准的 DDPM 前向过程:
|
||
$$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)$$
|
||
|
||
可以写成 Flow Matching 的路径形式:
|
||
$$x_t = \alpha(t) x_0 + \beta(t) \epsilon$$
|
||
|
||
其中 $\alpha(t) = \sqrt{\bar{\alpha}_t}$,$\beta(t) = \sqrt{1 - \bar{\alpha}_t}$。
|
||
|
||
**关键结论**:Flow Matching 不是"另一套完全不同的东西",而是扩散模型的**更一般的训练视角**。
|
||
|
||
| 特性 | DDPM | Flow Matching |
|
||
|------|------|---------------|
|
||
| **路径定义** | 噪声调度 $\bar{\alpha}_t$ 预设 | 任意可设计 |
|
||
| **训练目标** | 预测噪声 $\epsilon$ 或 score | 预测速度场 $v_\theta$ |
|
||
| **采样步数** | 50-1000 步 | 可 1-10 步(路径直线化) |
|
||
| **数学框架** | 变分推断 + ELBO | 向量场回归 |
|
||
|
||
---
|
||
|
||
## 七、Rectified Flow:路径直线化与最优传输
|
||
|
||
### 7.1 线性插值路径
|
||
|
||
**Rectified Flow** 起源于一个最简单的直觉:**直线是最短的路径**。
|
||
|
||
定义线性插值路径:
|
||
$$x_t = (1-t)x_0 + t x_1, \quad t \in [0,1]$$
|
||
|
||
其中 $x_0 \sim p_0$(噪声),$x_1 \sim p_1$(数据)。
|
||
|
||
**速度场**:
|
||
$$\frac{d x_t}{d t} = x_1 - x_0$$
|
||
|
||
**关键发现**:$u_t(x_t | x_1) = x_1 - x_0$ 是一个**与 $t$ 和 $x_t$ 无关**的常数向量!
|
||
|
||
这意味着我们可以用一个与时间无关的向量场来驱动样本从噪声演化为数据。
|
||
|
||
### 7.2 1-Step 生成公式
|
||
|
||
路径完全直线化后,ODE 求解退化为**单步仿射变换**:
|
||
$$x_1 = x_0 + v_\theta(x_0)$$
|
||
|
||
这相当于在噪声 $x_0$ 上直接加一个"位移向量" $v_\theta(x_0)$。
|
||
|
||
### 7.3 Re-flow 流程:迭代掰直路径
|
||
|
||
**问题**:线性插值生成的路径未必是最优的——可能经过低密度区域("路径交叉"问题)。
|
||
|
||
**Re-flow** 是一种迭代训练策略,用于逐步"掰直"流线:
|
||
|
||
**Algorithm: Re-flow**
|
||
|
||
1. **初始化**:使用线性路径训练初始向量场 $v_\theta^{(0)}$
|
||
2. **迭代优化**:
|
||
- 给定当前向量场 $v_\theta^{(k)}$,解 ODE 生成样本轨迹
|
||
- 估计轨迹的"曲率"(Curvature)
|
||
- **重新定义路径**:令新路径为上一条轨迹的直线插值
|
||
- 在新路径上训练更新的向量场 $v_\theta^{(k+1)}$
|
||
3. **收敛**:当路径曲率足够小时停止
|
||
|
||
### 7.4 最优传输视角
|
||
|
||
路径交叉问题本质上是**非最优传输(Non-optimal Transport)** 性质。精确求解 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$ 是代价函数。
|
||
|
||
**实际挑战**:精确求解 OT 需要 Sinkhorn 等近似算法,复杂度 $O(N^2)$。
|
||
|
||
---
|
||
|
||
## 八、训练过程与潜在问题汇总
|
||
|
||
### 8.1 训练算法对比
|
||
|
||
| 模型 | 训练目标 | 是否需要 ODE/SDE 求解 | 关键技巧 |
|
||
|------|----------|----------------------|----------|
|
||
| **NF** | 最大似然(log-determinant) | 否 | Jacobian 结构设计 |
|
||
| **CNF** | 最大似然(trace 积分) | 是(反向) | Adjoint Method |
|
||
| **FFJORD** | 最大似然(随机 trace) | 是(反向) | Hutchinson 估计器 |
|
||
| **Score SDE** | Score Matching | 否 | 去噪 score matching |
|
||
| **Flow Matching** | 向量场回归 | 否 | 路径设计 |
|
||
|
||
### 8.2 常见问题与解决方案
|
||
|
||
#### 8.2.1 NFE(Number of Function Evaluations)过高
|
||
|
||
**问题**:CNF/FFJORD 训练时需要 ODE 逆向求解,NFE 可能达到数百甚至数千。
|
||
|
||
**解决方案**:
|
||
- 使用自适应步长求解器(Dormand-Prince)
|
||
- 采用 Checkpointing 策略减少内存
|
||
- 切换到 Flow Matching 等 simulation-free 方法
|
||
|
||
#### 8.2.2 方差问题(FFJORD)
|
||
|
||
**问题**:随机迹估计引入方差,$M$ 越大越稳定但越慢。
|
||
|
||
**解决方案**:
|
||
- 使用更大的 batch size 抵消方差
|
||
- 采用 antithetic variates(负相关采样)
|
||
- 多次估计取平均(对于关键步骤)
|
||
|
||
#### 8.2.3 刚度问题(Stiffness)
|
||
|
||
**问题**:某些向量场导致 ODE 刚度增加,求解器需要极小步长。
|
||
|
||
**解决方案**:
|
||
- 使用专门处理刚度的求解器(如 BDF 方法)
|
||
- 调整向量场架构避免刚度
|
||
- 使用半隐式方法
|
||
|
||
#### 8.2.4 路径交叉问题(Flow Matching)
|
||
|
||
**问题**:独立采样的 $x_0$ 和 $x_1$ 可能导致路径交叉。
|
||
|
||
**解决方案**:
|
||
- Re-flow 迭代训练
|
||
- 最优传输(OT)配对
|
||
- Unbalanced Flow Matching
|
||
|
||
#### 8.2.5 数值精度累积
|
||
|
||
**问题**:长轨迹积分中数值误差累积,影响似然估计。
|
||
|
||
**解决方案**:
|
||
- 使用高精度求解器
|
||
- 积分路径两端使用更小步长
|
||
- 定期重置数值方法
|
||
|
||
---
|
||
|
||
## 九、数学推导速查表
|
||
|
||
### 9.1 核心公式汇总
|
||
|
||
**变量代换(离散)**:
|
||
$$p_X(x) = p_Z(z) \left| \det \frac{\partial f^{-1}}{\partial x} \right|$$
|
||
|
||
**变量代换(连续/瞬时)**:
|
||
$$\frac{\partial \log p_t}{\partial t} = -\text{tr}\left( \frac{\partial f}{\partial x} \right)$$
|
||
|
||
**Hutchinson 迹估计**:
|
||
$$\text{tr}(A) = \mathbb{E}_v \left[ v^T A v \right], \quad v \sim \mathcal{N}(0,I)$$
|
||
|
||
**前向 SDE**:
|
||
$$d x(t) = f(x(t), t) d t + g(x(t), t) d W(t)$$
|
||
|
||
**反向时间 SDE**:
|
||
$$d x(t) = \left[ f - g g^T \nabla_x \log p_t \right] d t + g d \bar{W}(t)$$
|
||
|
||
**Probability Flow ODE**:
|
||
$$d x(t) = \left[ f - \frac{1}{2} g g^T \nabla_x \log p_t \right] d t$$
|
||
|
||
**Flow Matching CFM 损失**:
|
||
$$\mathcal{L}_{CFM} = \mathbb{E}_{t, x_1, x_t} \| v_\theta(x_t, t) - u_t(x_t | x_1) \|^2$$
|
||
|
||
### 9.2 维度与复杂度
|
||
|
||
| 操作 | 维度 | 复杂度 |
|
||
|------|------|--------|
|
||
| 行列式计算 | $D \times D$ | $O(D^3)$ |
|
||
| 迹计算(精确) | $D \times D$ | $O(D^2)$ |
|
||
| 迹估计(Hutchinson,$M=1$) | $D \times D$ | $O(D)$ |
|
||
| ODE 求解($N$ 步) | $D$ | $O(N \cdot D)$ |
|
||
| 完整 CNF 训练 | $D$ | $O(N \cdot D^2)$(精确)/ $O(N \cdot D)$(FFJORD) |
|
||
|
||
---
|
||
|
||
## 十、总结
|
||
|
||
从 Neural ODE 到 FFJORD,再到 Score-Based SDE 和 Flow Matching,我们看到了一条清晰的演进路径:
|
||
|
||
1. **Neural ODE**:用 ODE 取代离散层叠加,实现连续标准化流
|
||
2. **FFJORD**:用随机迹估计器消除 Jacobian 计算瓶颈
|
||
3. **Score-Based SDE**:将前向/反向过程统一到 SDE 框架,给出严格的反向时间 SDE 和概率流 ODE
|
||
4. **Flow Matching**:从"解 ODE"转向"拟合向量场",实现 simulation-free 训练
|
||
|
||
**核心数学工具**:
|
||
- 瞬时变量代换公式(trace 代替 determinant)
|
||
- Hutchinson 迹估计器(随机化降低复杂度)
|
||
- Itô Calculus(SDE 框架)
|
||
- 向量场回归(Flow Matching)
|
||
|
||
**实际权衡**:
|
||
- 精度 vs 速度:RK45 vs Euler
|
||
- 准确性 vs 方差:精确迹 vs 随机估计
|
||
- 采样质量 vs 步数:多步 vs 单步(Rectified Flow)
|
||
|
||
---
|
||
|
||
**延伸阅读**:
|
||
|
||
1. Chen et al., "Neural Ordinary Differential Equations" (NeurIPS 2018)
|
||
2. Grathwohl et al., "FFJORD: Free-Form Jacobian of Dynamics Reversible" (ICLR 2019)
|
||
3. Song et al., "Score-Based Generative Modeling through Stochastic Differential Equations" (ICLR 2021)
|
||
4. Lipman et al., "Flow Matching for Generative Modeling" (NeurIPS 2022)
|
||
5. Liu et al., "Rectified Flow: A Marginal Preserving Approach to Optimal Transport" (ICML 2023)
|
||
6. Albergo & Vanden-Eijnden, "Building Normalizing Flows with Stochastic Interpolants" (ICLR 2023) |