353 lines
15 KiB
Markdown
353 lines
15 KiB
Markdown
---
|
||
title: 4-Neural ODE + FFJORD：从离散到连续的标准化流
|
||
draft: false
|
||
tags:
|
||
- Neural-ODE
|
||
- FFJORD
|
||
- CNF
|
||
- 标准化流
|
||
- 生成模型
|
||
- 深度学习
|
||
---
|
||
|
||
# Neural ODE + FFJORD：从离散到连续标准化流的数学基础
|
||
|
||
---
|
||
|
||
## 一、从离散流到连续流的哲学跃迁
|
||
|
||
### 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 精确迹的计算成本与 FFJORD 的真实动机
|
||
|
||
精确计算迹 $\text{tr}(\frac{\partial f}{\partial x})$ 本身复杂度仅 $O(D)$（只需 $D$ 个对角线偏导数），而非 $O(D^2)$。传统 CNF 的实际瓶颈是**完整雅可比行列式**$O(D^3)$。
|
||
|
||
FFJORD 使用 Hutchinson 估计的真正目的：省去逐个计算 $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$，有 $v^T A v = \sum_{i,j} v_i A_{ij} v_j$。由随机向量正交性 $\mathbb{E}[v_i v_j] = 0\ (i \neq j)$，$\mathbb{E}[v_i^2] = 1$，取期望得 $\mathbb{E}[v^T A v] = \sum_i A_{ii} = \mathrm{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. **数值精度**：积分过程中的数值误差会累积，影响似然估计的准确性。
|
||
|
||
---
|
||
|
||
## 五、训练与推理流程
|
||
|
||
### 5.1 Forward Path（从噪声到数据）
|
||
|
||
给定先验分布 $p_0(z) = \mathcal{N}(0, I)$，通过 ODE Solver 从 $t=0$ 积分到 $t=1$：
|
||
|
||
$$x = z(T) = z(0) + \int_0^T f(z(t), t; \theta) d t$$
|
||
|
||
### 5.2 Adjoint Method：常数显存的反向传播
|
||
|
||
传统的反向传播需要存储所有中间激活值 $z(t)$，内存复杂度为 $O(N \cdot D)$（$N$ 为步数）。这在深层 ODE 中是不可接受的。
|
||
|
||
**Adjoint Method** 通过引入**伴随状态 (Adjoint State)** $\lambda(t) = \frac{\partial \mathcal{L}}{\partial z(t)}$，将梯度计算转化为另一个 ODE 的求解，从而实现**单次前向 + 无需存储中间激活值**的反向传播。
|
||
|
||
定义增广损失：
|
||
|
||
$$\bar{\mathcal{L}} = \mathcal{L}(z(T)) + \int_0^T \lambda(t)^T (f(z(t), t) - \dot{z}(t)) d t$$
|
||
|
||
对 $\theta$ 求导，利用欧拉-拉格朗日方程，可推导出伴随 ODE：
|
||
|
||
$$\frac{d \lambda(t)}{d t} = -\lambda(t)^T \frac{\partial f}{\partial z}$$
|
||
|
||
最终梯度的计算只需在反向时间积分这个 ODE，无需存储任何中间状态。
|
||
|
||
---
|
||
|
||
## 六、与 Flow Matching 的关系
|
||
|
||
### 6.1 从 ODE 约束中解放
|
||
|
||
**Flow Matching** 提出了一个根本性问题：为什么训练时必须"解 ODE"？能否直接拟合向量场 $f$ 而不经过数值积分？
|
||
|
||
**核心思想**：将向量场拟合定义为一个简单的回归任务。设目标向量场为 $u_t(x)$，训练目标为：
|
||
|
||
$$\mathcal{L}_{FM} = \mathbb{E}_{t, x, x_1} \left[ \| f_t(x) - u_t(x) \|^2 \right]$$
|
||
|
||
其中 $u_t(x)$ 可以通过**插值路径**（如线性插值或最优传输路径）预先构造。
|
||
|
||
### 6.2 FFJORD 与 Flow Matching 的核心区别
|
||
|
||
| 特性 | FFJORD | Flow Matching |
|
||
|------|--------|---------------|
|
||
| **训练目标** | 最大似然（需要 ODE 求解） | 向量场回归（simulation-free） |
|
||
| **是否需要 ODE 求解** | 是（反向 ODE） | 否 |
|
||
| **损失函数** | 负对数似然 | MSE 回归 |
|
||
| **推理** | ODE Solver | 任意 ODE Solver |
|
||
|
||
---
|
||
|
||
## 七、常见问题与解决方案
|
||
|
||
### 7.1 NFE（Number of Function Evaluations）过高
|
||
|
||
**问题**：CNF/FFJORD 训练时需要 ODE 逆向求解，NFE 可能达到数百甚至数千。
|
||
|
||
**解决方案**：
|
||
- 使用自适应步长求解器（Dormand-Prince）
|
||
- 采用 Checkpointing 策略减少内存
|
||
- 切换到 Flow Matching 等 simulation-free 方法
|
||
|
||
### 7.2 方差问题（FFJORD）
|
||
|
||
**问题**：随机迹估计引入方差，$M$ 越大越稳定但越慢。
|
||
|
||
**解决方案**：
|
||
- 使用更大的 batch size 抵消方差
|
||
- 采用 antithetic variates（负相关采样）
|
||
- 多次估计取平均
|
||
|
||
### 7.3 刚度问题（Stiffness）
|
||
|
||
**问题**：某些向量场导致 ODE 刚度增加，求解器需要极小步长。
|
||
|
||
**解决方案**：
|
||
- 使用专门处理刚度的求解器（如 BDF 方法）
|
||
- 调整向量场架构避免刚度
|
||
- 使用半隐式方法
|
||
|
||
---
|
||
|
||
## 八、公式速查
|
||
|
||
**Neural ODE 核心方程**：
|
||
$$\frac{d x(t)}{d t} = f_\theta(x(t), t)$$
|
||
|
||
**瞬时变量代换公式**：
|
||
$$\frac{\partial \log p_t}{\partial t} = -\text{tr}\left( \frac{\partial f}{\partial x} \right)$$
|
||
|
||
**对数似然积分形式**：
|
||
$$\log p_T(x) = \log p_0(z(0)) - \int_0^T \text{tr}\left( \frac{\partial f}{\partial x} \right) d t$$
|
||
|
||
**Hutchinson 迹估计**：
|
||
$$\text{tr}(A) = \mathbb{E}_v \left[ v^T A v \right]$$
|
||
|
||
**维度与复杂度**：
|
||
|
||
| 操作 | 维度 | 复杂度 |
|
||
|------|------|--------|
|
||
| 行列式计算 | $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 将离散流模型的思想推广到连续时间极限，用 ODE 取代了离散的层叠加。核心突破在于**瞬时变量代换公式**，将行列式计算替换为迹算子，从而解放了对 Jacobian 稀疏结构的约束。
|
||
|
||
FFJORD 进一步使用 **Hutchinson 迹估计器**，通过随机投影将迹计算从 $O(D^2)$ 降为 $O(D)$，实现了真正意义上的无 Jacobian 计算的连续标准化流训练。
|
||
|
||
然而，FFJORD 仍然需要 ODE 求解来计算损失函数，这与 Flow Matching 提出的"直接拟合向量场"的思想形成了鲜明对比。这一思想转变为后续 Flow Matching 和 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. Rezende & Mohamed, "Variational Inference with Normalizing Flows" (ICML 2015) |