14 KiB
title, draft, tags
| title | draft | tags | ||||
|---|---|---|---|---|---|---|
| 3-NF-CNF-标准化流与连续标准化流 | false |
|
Normalizing Flow 与 Continuous Normalizing Flow：从离散变换到微分方程的严格推导
一、数学基础与变量代换 (Change of Variables)
1.1 概率分布变换的基本定理
在讨论 Normalizing Flow 之前，必须严格建立从概率密度变换的基本定律。
设随机变量 z \in \mathbb{R}^D 服从已知分布 $p_z(z)$。我们通过一个双射（Bijunctive）映射 f: \mathbb{R}^D \rightarrow \mathbb{R}^D 将其变换为 $x = f(z)$。则 x 的概率密度 p_x(x) 与 z 的概率密度 p_z(z) 之间满足以下关系：
定理（Change of Variables）
p_x(x) = p_z(z) \left| \det \frac{\partial f^{-1}}{\partial x} \right| = p_z(z) \left| \det \frac{\partial f}{\partial z} \right|^{-1}
其中 $z = f^{-1}(x)$，而 \frac{\partial f^{-1}}{\partial x} 是逆函数在 x 处的 Jacobian 矩阵。
推导（严格形式）：
考虑一个无穷小的体积元 d z 在变换后的像 $d x$。由于 f 是双射，两者的体积通过 Jacobian 行列式关联：
d x = \left| \det \frac{\partial f}{\partial z} \right| d z \quad \Longrightarrow \quad d z = \left| \det \frac{\partial f^{-1}}{\partial x} \right| d x
概率质量在该体积元上保持不变：$p_x(x) d x = p_z(z) d z$。
将 d z 的表达式代入：
p_x(x) d x = p_z(z) \left| \det \frac{\partial f^{-1}}{\partial x} \right| d x
消除 d x 即得上述定理。
几何直觉：对于高维空间中的微小体积元，变换前后体积比率为 $|\det J|$，其中 $J = \frac{\partial f}{\partial z}$。概率守恒要求 $p_x(x) |\det J| = p_z(z)$，即 $p_x(x) = p_z(z) / |\det J|$。
补充：矩阵行列式与体积变换的关系
设 $J = \frac{\partial f}{\partial z}$，对于线性变换 $x = Jz$，D 维单位球被映射为椭球，椭球的体积为 $|\det J|$。这解释了为何 Jacobian 行列式出现在体积元变换中。
1.2 对数形式的变量代换公式
对两边取对数，得到实际计算中更常用的形式：
\log p_x(x) = \log p_z(f^{-1}(x)) + \log \left| \det \frac{\partial f^{-1}}{\partial x} \right|
\log p_x(x) = \log p_z(z) - \log \left| \det \frac{\partial f}{\partial z} \right|
这里的 \log \left| \det \frac{\partial f}{\partial z} \right| 被称为 log-determinant term，它是 NF 中梯度优化的核心组成部分。
1.3 Jacobian 行列式的计算复杂度诅咒
对于一般的矩阵 $J = \frac{\partial f}{\partial z} \in \mathbb{R}^{D \times D}$，计算其行列式的复杂度为 $O(D^3)$（通过 LU 分解或 SVD）。当维度 D 达到数千（如图像生成中常见的 $64 \times 64 \times 3$），这个计算代价是不可接受的。
这构成了 NF 设计中的核心约束：必须设计特殊结构的 $f$，使其 Jacobian 矩阵具有稀疏结构（如三角矩阵），从而将行列式计算降为 $O(D)$。 Affine Coupling Layer 正是针对这一约束的突破性设计。
二、离散流模型 (Discrete Flows) 的架构演进
2.1 耦合层 (Coupling Layer) 的数学构造
RealNVP 引入的 Affine Coupling Layer 是解决 Jacobian 计算问题的第一个工业级方案。
设输入 z \in \mathbb{R}^D 被沿通道维度划分为两部分：$z = [z_1, z_2]$，其中 $z_1 \in \mathbb{R}^{d}$，$z_2 \in \mathbb{R}^{D-d}$。
前向变换（从 z 到 $x$）定义为：
x_1 = z_1 \quad \text{（恒等变换）}
x_2 = z_2 \odot \exp(s(z_1)) + t(z_1)
其中 s(\cdot): \mathbb{R}^d \rightarrow \mathbb{R}^{D-d} 和 t(\cdot): \mathbb{R}^d \rightarrow \mathbb{R}^{D-d} 是神经网络输出的缩放和平移参数。\odot 表示逐元素乘法。
Jacobian 矩阵的结构分析：
\frac{\partial x}{\partial z} = \begin{bmatrix} I_d & 0 \\ \frac{\partial x_2}{\partial z_1} & \text{diag}(\exp(s(z_1))) \end{bmatrix}
这是一个下三角矩阵（假设 z_1 在前，z_2 在后）。其行列式为：
\det\left(\frac{\partial x}{\partial z}\right) = \prod_{i=1}^{D-d} \exp(s_i(z_1)) = \exp\left(\sum_{i=1}^{D-d} s_i(z_1)\right)
结论：行列式计算从 O(D^3) 降为 $O(D)$，且完全不涉及神经网络 s(\cdot) 和 t(\cdot) 的雅可比计算。
直觉解释：为什么下三角矩阵的行列式是对角线元素之积？
对于下三角矩阵 $A = \begin{bmatrix} A_{11} & 0 \ A_{21} & A_{22} \end{bmatrix}$，行列式为 $A_{11} \cdot A_{22}$。因为上三角部分全为 0 ，行列式展开式中所有涉及上三角位置元素的项都包含 0 因子。
在 Affine Coupling Layer 中，Jacobian 矩阵为：
\frac{\partial x}{\partial z} = \begin{bmatrix} I_d & 0 \\ \frac{\partial x_2}{\partial z_1} & \text{diag}(\exp(s(z_1))) \end{bmatrix}
对角块为 $I_d$（维度 $d$）和 $\text{diag}(\exp(s(z_1)))$（维度 $D-d$），故行列式为 $1 \cdot \prod_{i=1}^{D-d} \exp(s_i(z_1))$。
逆变换的简洁性：
z_1 = x_1
z_2 = (x_2 - t(x_1)) \odot \exp(-s(x_1))
完全可逆，且计算复杂度与前向传播相同。
2.2 Glow 中的 1 \times 1 可逆卷积
Glow 在 RealNVP 的基础上引入了 1 \times 1 可逆卷积作为通道混洗（Channel Permutation）的可学习替代。
设输入通道维度为 $C$，权重矩阵 W \in \mathbb{R}^{C \times C} 是可学习的置换矩阵（满足列归一化约束）。
前向与逆向分别为：
x = W \cdot z \quad \Longleftrightarrow \quad z = W^{-1} \cdot x
行列式为 $\det(W)$（通过矩阵分解预处理为 LU 形式以高效计算）。
通道混合的数学意义：标准的耦合层使通道 1 \sim d 与 d+1 \sim D 保持独立。1 \times 1 卷积允许任意通道间的线性混合，从而增强模型的表达能力，同时保持可逆性。
三、向连续流的跃迁：Neural ODEs 与 CNF
3.1 从离散叠加到连续流的极限推导
将有限个耦合层堆叠视为一条离散轨迹 $z_0 \rightarrow z_1 \rightarrow \cdots \rightarrow z_N$，每一步的变换记为 $z_{n+1} = f(z_n; \theta_n)$。
当我们无限细分这个轨迹，令层数 $N \rightarrow \infty$，步长 $\Delta t \rightarrow 0$，每一步的变换缩小至无穷小量。则离散路径收敛于一条连续时间流：
\frac{d x(t)}{d t} = f(x(t), t; \theta)
初始状态 $x(0) = z_0$，终止状态 $x(1) = x$。此时，f 被称为向量场 (Vector Field)。
3.2 核心推导：瞬时变量代换公式 (Instantaneous Change of Variables Formula)
定理：设 x(t) \in \mathbb{R}^D 满足上述 ODE，概率密度 p_t(x(t)) 对时间 t 的演化满足：
\frac{\partial \log p_t(x(t))}{\partial t} = -\text{tr}\left( \frac{\partial f(x(t), t; \theta)}{\partial x(t)} \right)
其中 \text{tr}(\cdot) 是矩阵的迹（Trace，即对角线元素之和）。
物理直觉：散度 $\text{tr}(\frac{\partial f}{\partial x}) = \nabla \cdot f$（仅在欧氏空间笛卡尔坐标系下成立）衡量向量场 f 在每一点的"源"与"汇"。若散度为正（膨胀），概率密度降低；若散度为负（收缩），概率密度升高。这与质量守恒的物理规律一致。
与离散流的关系：在离散 NF 中，第 n 层变换的 log-determinant 为 $\log |\det J_n|$。在连续情形下，这些离散的行列式对数被连续时间的 trace 积分所替代：
\log \det \frac{\partial x(1)}{\partial x(0)} = \int_0^1 \text{tr}\left( \frac{\partial f(x(t), t)}{\partial x(t)} \right) d t
这意味着连续标准化流将无穷多层离散的行列式计算压缩为单次迹积分。
推导（来自 Liouville 定理）：
考虑一个无穷小的体积元 V(t) 沿流线演化。概率密度演化严格遵循刘维尔方程，体积变化率由散度 (Divergence) 控制（无需不可压缩假设，流的微分同胚性质由 Lipschitz 条件保证）：
\frac{d}{d t} \log \det \frac{\partial x(t)}{\partial x(0)} = \text{tr}\left( \frac{\partial f}{\partial x} \right)
详细推导步骤：
设 $J(t) = \frac{\partial x(t)}{\partial x(0)}$，则 J(t) 满足矩阵微分方程：
\frac{d J(t)}{d t} = \frac{\partial f}{\partial x} J(t)
这是因为 $x(t) = x(0) + \int_0^t f(x(s), s) ds$，对初始条件求偏导：
\frac{\partial x_i(t)}{\partial x_j(0)} = \delta_{ij} + \int_0^t \sum_k \frac{\partial f_i}{\partial x_k} \frac{\partial x_k}{\partial x_j(0)} ds
对 t 求导得 $\frac{d}{dt} J = \frac{\partial f}{\partial x} J$。
对 J 的行列式取对数并求导，利用矩阵行列式的对数导数公式：
\frac{d}{dt} \log \det J(t) = \text{tr}\left( J(t)^{-1} \frac{d J(t)}{dt} \right) = \text{tr}\left( J(t)^{-1} \frac{\partial f}{\partial x} J(t) \right) = \text{tr}\left( \frac{\partial f}{\partial x} \right)
其中最后一步利用了 $\text{tr}(AB) = \text{tr}(BA)$。
概率质量在每个体积元上守恒：$p_t(x(t)) V(t) = p_0(x(0)) V(0)$。
取对数并整理：
\frac{\partial \log p_t}{\partial t} = -\frac{\partial}{\partial t} \log \det \frac{\partial x}{\partial x_0} = -\text{tr}\left( \frac{\partial f}{\partial x} \right)
3.3 迹算子为何取代行列式：从离散到连续的量级差异
在离散 NF 中，我们计算 $\log \left| \det \frac{\partial f_k}{\partial z_k} \right|$，这是一个标量。
在连续 NF 中，迹算子 \text{tr}\left( \frac{\partial f}{\partial x} \right) 是标量（Jacobian 矩阵的对角线元素之和），但它等价于无穷多层离散行列式的累积对数和：
\log \det \frac{\partial x(1)}{\partial x(0)} = \int_0^1 \text{tr}\left( \frac{\partial f(x(t), t)}{\partial x(t)} \right) d t
计算上的巨大进步：迹的计算复杂度是 $O(D)$（只需对角线元素），而行列式是 $O(D^3)$。这使得 CNF 能够处理任意维度、任意结构的向量场，无需担心 Jacobian 的稀疏结构约束。
3.4 Picard-Lindelöf 定理与可逆性保证
定理（Picard-Lindelöf）：若向量场 f(x, t; \theta) 关于 x 满足局部 Lipschitz 连续条件，则 ODE dx/dt = f(x, t) 存在唯一解，且解对初始条件连续依赖。（全局 Lipschitz 是更强的充分条件）
对可逆性的保证：该定理意味着给定初始点 $x(0)$，流 x(t; x(0)) 是唯一确定的。不存在两条不同的流线在 t > 0 时交汇。这从数学上保证了 CNF 的双射 (Diffeomorphism) 性质——只要 f 是 Lipschitz 的，积分过程就是可逆的。
逆向过程通过解伴随的 ODE（Reverse ODE）实现：
\frac{d x(t)}{d t} = -f(x(t), t; \theta)
四、训练与推理流程
4.1 Forward Path（从噪声到数据）
给定先验分布 $p_0(z) = \mathcal{N}(0, I_D)$（I_D 为 D 维单位矩阵），通过 ODE Solver 从 t=0 积分到 $t=1$：
x = z(T) = z(0) + \int_0^T f(z(t), t; \theta) d t
常用的数值方法包括：
Euler Method（简单但低阶）：
z_{n+1} = z_n + \Delta t \cdot f(z_n, t_n; \theta)
局部截断误差为 $O(\Delta t)$（全局积累误差为 $O(\Delta t)$）。
RK4（经典高阶方法）：
k_1 = f(z_n, t_n)
k_2 = f(z_n + \frac{\Delta t}{2} k_1, t_n + \frac{\Delta t}{2})
k_3 = f(z_n + \frac{\Delta t}{2} k_2, t_n + \frac{\Delta t}{2})
k_4 = f(z_n + \Delta t \cdot k_3, t_n + \Delta t)
z_{n+1} = z_n + \frac{\Delta t}{6}(k_1 + 2k_2 + 2k_3 + k_4)
局部截断误差为 $O(\Delta t^4)$（全局积累误差为 $O(\Delta t^3)$）。
4.2 训练目标：最大似然估计 (MLE)
CNF 的训练目标是最大化生成分布的对数似然。边缘似然为：
\log p_\theta(x) = \log p_0(z(0)) - \int_0^T \text{tr}\left( \frac{\partial f(x(t), t; \theta)}{\partial x(t)} \right) d t
其中 z(0) 是通过逆向 ODE 从 x 积分到 0 得到的噪声。
训练损失为负对数似然的期望：
\mathcal{L}(\theta) = -\mathbb{E}_{x \sim p_{data}} \left[ \log p_\theta(x) \right]
4.3 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，无需存储任何中间状态。
五、局限性与未来展望
5.1 CNF 的核心痛点：NFE (Number of Function Evaluations)
ODE Solver 的步数直接决定了计算量。对于高精度要求，NFE 可能达到数百甚至数千，远超同深度的离散流模型。
主要瓶颈：
- 数值稳定性：刚度 (Stiffness) 问题时需要自适应步长控制
- 精度与速度的权衡：高阶 Solver（RK45）精度高但慢；低阶（Euler）快但精度差
- 误差累积：长轨迹积分中误差可能累积
5.2 Flow Matching：超越 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) 可以通过插值路径（如线性插值或最优传输路径）预先构造。
优势：
- 训练时无需 ODE 求解（无 NFE 概念）
- 推理时可使用任何 ODE Solver
- 统一的框架可兼容扩散模型（DDPM）的噪声路径
六、总结
Normalizing Flow 从离散耦合层出发，逐步演化为连续标准化流（CNF），核心突破在于将行列式计算替换为迹算子，从而解放了对 Jacobian 稀疏结构的约束。然而 ODE Solver 的计算开销构成了实际应用中的主要瓶颈。Flow Matching 等新范式正在将研究重心从"如何解 ODE"转向"如何直接拟合向量场"，这标志着生成模型领域正在经历又一次根本性的数学统一。
延伸阅读：
- Rezende & Mohamed, "Variational Inference with Normalizing Flows" (2015)
- Grathwohl et al., "FFJORD: Free-Form Jacobian of Reversible Dynamics" (2019)
- Lipman et al., "Flow Matching for Generative Modeling" (2022)