第4章扩散模型入门

原课程: Robots That Learn, UC Berkeley CS294-291, Spring 2026 主讲: Angjoo Kanazawa (Guest Lecture L3) 本章基于 Lecture 3, 精读论文: Papamakarios et al. (Normalizing Flows, JMLR 2021) + Lipman et al. (Flow Matching, ICLR 2023)

一、本章概要

假设你有一张清晰的照片，然后你不停地往上面加一点噪声——每次加的噪声都是随机的、微小的，但累积起来，最终照片会变成一片完全无法辨认的雪花。扩散模型的核心洞察是：这个过程是可逆的。如果你精确地知道每一步加了多少噪声，你就可以学习一个神经网络来执行逆向操作——逐步从纯噪声中恢复出清晰的图像。

这个简单但深邃的思想，在2020年代彻底重塑了生成式人工智能的版图。然而，更深层的问题是：为什么要用扩散过程？为什么不像传统方法那样，直接学习一个从噪声到数据的确定性映射？

Angjoo Kanazawa 在 Lecture 3 中给出的答案是分层的。第一层：Denoising Diffusion Probabilistic Models (DDPM) 将生成建模归结为一个逐步去噪问题——这比 GAN 的对抗训练更稳定，比 VAE 的近似推断更灵活。第二层：Score-based 模型从 Langevin 动力学的视角揭示了扩散模型的本质——我们实际上是在学习数据分布的梯度场（score function），然后用它来引导随机采样。第三层：Normalizing Flows 提供了另一种视角——通过巧妙的变量替换（change of variables），我们可以直接从简单分布变换到复杂分布，而不需要任何随机性。第四层：Flow Matching 统一了扩散模型和 Normalizing Flows——它让我们直接设计概率路径（probability path），并大幅简化训练过程。

本章的核心使命是：从第一性原理出发完成四类模型的数学推导，使读者能够从空白 .py 文件出发，独立实现 DDPM、Normalizing Flow 和 Flow Matching。

具体而言，本章将覆盖以下内容：

DDPM 的完整推导链：从马尔可夫前向加噪过程开始，推导逆向去噪过程的变分下界 (ELBO)，通过重参数化将其化简为简单的均方误差损失，最终给出训练和采样算法。这条推导链的每一个等号都会被明确证明。
Score-based 模型的 Langevin 动力学视角：引入 Score Function（Stein score），推导 Denoising Score Matching (DSM) 损失，建立 score matching 与 denoising autoencoder 之间的等价关系，展示如何用 score 函数引导 Langevin 动力学采样。
Normalizing Flows 的数学基础：从变量替换公式出发，详细讲解自回归流（autoregressive flows）、耦合层（coupling layers）、残差流（residual flows）三类关键架构的设计原理和 Jacobian 行列式的计算方法。
Continuous Normalizing Flows (CNF) 与 ODE：推导连续时间流的瞬时变量替换公式，以及伴随方法（adjoint method）在梯度计算中的应用。
Flow Matching 的统一框架：推导 Conditional Flow Matching 定理（证明 CFM 与 FM 目标函数梯度的等价性），引入 Optimal Transport 条件概率路径，展示 OT 路径相对于扩散路径在训练和采样效率上的优势。
完整的 PyTorch 实现：DDPM（U-Net + 前向/逆向过程）、CNF（Neural ODE 求解器）、Flow Matching（OT 路径训练）。

在正式展开之前，请记住 Kanazawa 在这堂课中传递的核心方法论立场：扩散模型和 Normalizing Flows 不是竞争关系，而是互补的工具。理解它们的数学共性——都涉及将简单分布变换为复杂分布——比记住任何一个具体实现更有价值。

二、核心概念与数学基础

2.1 Denoising Diffusion Probabilistic Models (DDPM)

DDPM（Ho et al., 2020）定义了一个具有两条”时间线”的潜变量模型：前向过程（forward process）逐步向数据添加噪声，逆向过程（reverse process）学习去噪。

2.1.1 前向过程（Forward Process）

前向过程是一个固定的（无可训练参数）马尔可夫链，它在 \(T\) 个时间步内将数据分布 \(q(\mathbf{x}_0)\) 逐渐转化为各向同性的高斯噪声。给定真实数据样本 \(\mathbf{x}_0 \sim q(\mathbf{x}_0)\)，前向过程定义为：

\[q(\mathbf{x}_{1:T} \mid \mathbf{x}_0) = \prod_{t=1}^T q(\mathbf{x}_t \mid \mathbf{x}_{t-1})\]

其中每一步是一个高斯转移核：

\[q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = \mathcal{N}\left(\mathbf{x}_t; \sqrt{1 - \beta_t}\,\mathbf{x}_{t-1}, \beta_t \mathbf{I}\right) \tag{2.1}\]

这里 \(\beta_t \in (0, 1)\) 是预先定义的方差调度（variance schedule），它控制每步添加的噪声量。一个经典的选择是线性调度：\(\beta_1 < \beta_2 < \cdots < \beta_T\)（例如，\(\beta_1 = 10^{-4}\)，\(\beta_T = 0.02\)）。

重参数化技巧：令 \(\alpha_t = 1 - \beta_t\)，\(\bar{\alpha}_t = \prod_{s=1}^t \alpha_s\)。通过递推展开：

\[\begin{aligned} \mathbf{x}_t &= \sqrt{\alpha_t}\,\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\,\boldsymbol{\epsilon}_{t-1} \\ &= \sqrt{\alpha_t \alpha_{t-1}}\,\mathbf{x}_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}}\,\boldsymbol{\epsilon} \\ &\;\;\,\cdots \\ &= \sqrt{\bar{\alpha}_t}\,\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\,\boldsymbol{\epsilon} \end{aligned}\]

其中 \(\boldsymbol{\epsilon}_{t-1}, \boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})\)，第二步用到了两个高斯随机变量之和的性质：\(\sqrt{\alpha_t(1-\alpha_{t-1})}\boldsymbol{\epsilon}_{t-2} + \sqrt{1-\alpha_t}\boldsymbol{\epsilon}_{t-1}\) 等价于 \(\sqrt{1-\alpha_t \alpha_{t-1}}\boldsymbol{\epsilon}\)。

因此，我们可以直接从 \(\mathbf{x}_0\) 采样任意时刻的 \(\mathbf{x}_t\)：

\[\boxed{q(\mathbf{x}_t \mid \mathbf{x}_0) = \mathcal{N}\left(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\,\mathbf{x}_0, (1 - \bar{\alpha}_t) \mathbf{I}\right)} \tag{2.2}\]

关键设计约束：选择 \(\beta_t\) 使得 \(\bar{\alpha}_T \to 0\)，即 \(q(\mathbf{x}_T \mid \mathbf{x}_0) \approx \mathcal{N}(\mathbf{x}_T; 0, \mathbf{I})\)。这确保了前向过程的终点是纯噪声。

2.1.2 逆向过程（Reverse Process）

逆向过程也是一个马尔可夫链，但方向相反——从纯噪声 \(\mathbf{x}_T \sim \mathcal{N}(0, \mathbf{I})\) 开始，逐步去噪直到生成数据：

\[p_\theta(\mathbf{x}_{0:T}) = p(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)\] \[p(\mathbf{x}_T) = \mathcal{N}(\mathbf{x}_T; 0, \mathbf{I})\] \[p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t) = \mathcal{N}\left(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \sigma_t^2 \mathbf{I}\right) \tag{2.3}\]

这里 \(\boldsymbol{\mu}_\theta(\mathbf{x}_t, t)\) 是一个由神经网络参数化的函数，\(\sigma_t^2\) 通常固定为 \(\beta_t\) 或 \(\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t}\beta_t\)（后者是 \(q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)\) 的方差）。

核心问题：当我们只能观测到 \(\mathbf{x}_t\)，而无法获得 \(\mathbf{x}_0\) 时，真实的逆向分布 \(q(\mathbf{x}_{t-1} \mid \mathbf{x}_t)\) 是不可直接计算的。但如果我们同时以 \(\mathbf{x}_0\) 为条件，这个分布就变得可解了。

2.1.3 以 \(\mathbf{x}_0\) 为条件的逆向分布

利用贝叶斯定理和马尔可夫性质：

\[\begin{aligned} q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) &= \frac{q(\mathbf{x}_t \mid \mathbf{x}_{t-1}, \mathbf{x}_0) \cdot q(\mathbf{x}_{t-1} \mid \mathbf{x}_0)}{q(\mathbf{x}_t \mid \mathbf{x}_0)} \\ &= \frac{q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) \cdot q(\mathbf{x}_{t-1} \mid \mathbf{x}_0)}{q(\mathbf{x}_t \mid \mathbf{x}_0)} \end{aligned}\]

将三个高斯分布的密度函数代入：

\[\begin{aligned} q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) &\propto \exp\left(-\frac{\|\mathbf{x}_t - \sqrt{\alpha_t}\mathbf{x}_{t-1}\|^2}{2\beta_t} - \frac{\|\mathbf{x}_{t-1} - \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0\|^2}{2(1 - \bar{\alpha}_{t-1})} + \frac{\|\mathbf{x}_t - \sqrt{\bar{\alpha}_t}\mathbf{x}_0\|^2}{2(1 - \bar{\alpha}_t)}\right) \end{aligned}\]

将指数内的二次型关于 \(\mathbf{x}_{t-1}\) 配方（completing the square），得到高斯分布的形式。最终结果为：

\[\boxed{q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}\left(\mathbf{x}_{t-1}; \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0), \tilde{\beta}_t \mathbf{I}\right)} \tag{2.4}\]

其中：

\[\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) = \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t}\mathbf{x}_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t}\mathbf{x}_t \tag{2.5}\] \[\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t \tag{2.6}\]

直观理解：公式 (2.5) 告诉我们，\(\mathbf{x}_{t-1}\) 的均值是 \(\mathbf{x}_0\) 和 \(\mathbf{x}_t\) 的一个凸组合（加权平均）。\(\mathbf{x}_0\) 的权重是 \(\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\)，\(\mathbf{x}_t\) 的权重是 \(\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\)。两者的和恰好等于 \(\frac{\sqrt{\bar{\alpha}_{t-1}}(\beta_t + \alpha_t - \bar{\alpha}_t)}{1-\bar{\alpha}_t} = \frac{\sqrt{\bar{\alpha}_{t-1}}(1-\bar{\alpha}_t)}{1-\bar{\alpha}_t} = \sqrt{\bar{\alpha}_{t-1}}\) [UNSPECIFIED: 此和恒等式的代数验证基于 \(\beta_t + \alpha_t = 1\)]。

公式 (2.4) 是整个 DDPM 的理论核心——它给了我们一个训练信号：如果我们知道 \(\mathbf{x}_0\)，我们就知道在从 \(\mathbf{x}_t\) 去噪到 \(\mathbf{x}_{t-1}\) 时，应该”瞄准”什么样的分布均值。

2.1.4 变分下界（ELBO）推导

DDPM 通过最大化观测数据 \(\mathbf{x}_0\) 的对数似然来训练。由于精确似然不可计算，我们转而优化其变分下界（Evidence Lower BOund, ELBO）：

\[\begin{aligned} -\log p_\theta(\mathbf{x}_0) &\leq -\log p_\theta(\mathbf{x}_0) + D_{KL}\big(q(\mathbf{x}_{1:T} \mid \mathbf{x}_0) \;\big\|\; p_\theta(\mathbf{x}_{1:T} \mid \mathbf{x}_0)\big) \\ &= -\log p_\theta(\mathbf{x}_0) + \mathbb{E}_{q(\mathbf{x}_{1:T} \mid \mathbf{x}_0)}\left[\log\frac{q(\mathbf{x}_{1:T} \mid \mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T}) / p_\theta(\mathbf{x}_0)}\right] \\ &= \mathbb{E}_{q}\left[\log\frac{q(\mathbf{x}_{1:T} \mid \mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})}\right] =: \mathcal{L} \end{aligned}\]

展开 \(\mathcal{L}\)：

\[\begin{aligned} \mathcal{L} &= \mathbb{E}_{q}\left[\log\frac{\prod_{t=1}^T q(\mathbf{x}_t \mid \mathbf{x}_{t-1})}{p(\mathbf{x}_T)\prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)}\right] \\ &= \mathbb{E}_{q}\Bigg[-\log p(\mathbf{x}_T) + \sum_{t=1}^T \log\frac{q(\mathbf{x}_t \mid \mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)}\Bigg] \end{aligned}\]

利用马尔可夫性质 \(q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = q(\mathbf{x}_t \mid \mathbf{x}_{t-1}, \mathbf{x}_0)\)，以及贝叶斯重排：

\[q(\mathbf{x}_t \mid \mathbf{x}_{t-1}, \mathbf{x}_0) = \frac{q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) \cdot q(\mathbf{x}_t \mid \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \mid \mathbf{x}_0)}\]

代入并分离 \(t=1\) 项和求和：

\[\begin{aligned} \mathcal{L} &= \mathbb{E}_{q}\Bigg[-\log p(\mathbf{x}_T) + \sum_{t=2}^T \log\frac{q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) \cdot q(\mathbf{x}_t \mid \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t) \cdot q(\mathbf{x}_{t-1} \mid \mathbf{x}_0)} + \log\frac{q(\mathbf{x}_1 \mid \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \mid \mathbf{x}_1)}\Bigg] \\ &= \mathbb{E}_{q}\Bigg[\underbrace{D_{KL}\big(q(\mathbf{x}_T \mid \mathbf{x}_0) \;\big\|\; p(\mathbf{x}_T)\big)}_{L_T} \\ &\quad + \sum_{t=2}^T \underbrace{D_{KL}\big(q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) \;\big\|\; p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)\big)}_{L_{t-1}} \\ &\quad - \underbrace{\log p_\theta(\mathbf{x}_0 \mid \mathbf{x}_1)}_{L_0}\Bigg] \end{aligned}\]

这就是 DDPM 的完整变分下界分解。总结如下：

\[\boxed{\mathcal{L}_{\text{VLB}} = L_T + \sum_{t=2}^T L_{t-1} + L_0} \tag{2.7}\]

其中：

\(L_T = D_{KL}\big(q(\mathbf{x}_T \mid \mathbf{x}_0) \;\|\; p(\mathbf{x}_T)\big)\) — 最终噪声分布的匹配项（不含可训练参数，可忽略）
\(L_{t-1} = D_{KL}\big(q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) \;\|\; p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)\big)\) — 去噪步骤的 KL 散度
\(L_0 = -\log p_\theta(\mathbf{x}_0 \mid \mathbf{x}_1)\) — 最终重建项

2.1.5 简化损失函数（Simplified Training Objective）

\(L_{t-1}\) 是两个高斯分布之间的 KL 散度，可以解析计算。将 \(p_\theta\) 的均值参数化为与 \(q\) 相同的函数形式：

方法一：预测 \(\tilde{\boldsymbol{\mu}}\)

令 \(\boldsymbol{\mu}_\theta(\mathbf{x}_t, t) = \tilde{\boldsymbol{\mu}}_t\left(\mathbf{x}_t, \frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t - \sqrt{1-\bar{\alpha}_t}\,\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t))\right)\)。两个高斯分布之间的 KL 散度为：

\[L_{t-1} = \frac{1}{2\sigma_t^2} \left\|\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) - \boldsymbol{\mu}_\theta(\mathbf{x}_t, t)\right\|^2\]

利用公式 (2.2) 的逆——\(\mathbf{x}_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t - \sqrt{1-\bar{\alpha}_t}\,\boldsymbol{\epsilon})\) 以及公式 (2.5)，可以导出均值参数化的损失。

方法二：预测 \(\boldsymbol{\epsilon}\)（简化版，DDPM 采用的方案）

将公式 (2.5) 中的 \(\mathbf{x}_0\) 用公式 (2.2) 替代：

\[\mathbf{x}_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}\big(\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t}\,\boldsymbol{\epsilon}\big)\]

代入公式 (2.5)：

\[\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) = \frac{1}{\sqrt{\alpha_t}}\left(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\boldsymbol{\epsilon}\right)\]

由于 \(p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)\) 的均值也可以写成相同的形式（用神经网络 \(\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)\) 替代 \(\boldsymbol{\epsilon}\)），KL 散度化简为：

\[L_{t-1} = \frac{\beta_t^2}{2\sigma_t^2 \alpha_t (1 - \bar{\alpha}_t)} \left\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)\right\|^2\]

Ho et al. (2020) 发现，丢弃权重系数，直接使用简化损失在实践中效果更好：

\[\boxed{\mathcal{L}_{\text{simple}}(\theta) = \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol{\epsilon}}\left[\left\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta\big(\sqrt{\bar{\alpha}_t}\,\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\,\boldsymbol{\epsilon},\; t\big)\right\|^2\right]} \tag{2.8}\]

其中 \(t \sim \text{Uniform}(1, T)\)，\(\boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})\)。

这一简化的深层含义：DDPM 的训练归结为一个极其简单的任务——给定加噪数据 \(\mathbf{x}_t\) 和时间步 \(t\)，预测添加的噪声 \(\boldsymbol{\epsilon}\)。这与 denoising autoencoder 的训练目标高度相似，区别在于 DDPM 是在所有噪声水平上联合训练的。

2.1.6 训练与采样算法

训练算法（算法 2.1）：

重复直到收敛:
    1. 采样 x_0 ~ q(x_0)                    # 从数据集中取一个样本
    2. 采样 t ~ Uniform({1, ..., T})         # 随机选一个时间步
    3. 采样 ε ~ N(0, I)                      # 采样噪声
    4. 计算梯度并更新:
       ∇_θ ‖ε - ε_θ(√(ᾱ_t)·x_0 + √(1-ᾱ_t)·ε, t)‖²

采样算法（算法 2.2）：

输入: 训练好的模型 ε_θ
x_T ~ N(0, I)                            # 从纯噪声开始
for t = T, T-1, ..., 1:
    z ~ N(0, I) if t > 1, else z = 0     # 采样随机噪声（最后一步不加）
    x_{t-1} = (1/√α_t)·(x_t - (β_t/√(1-ᾱ_t))·ε_θ(x_t, t)) + σ_t·z
return x_0

为什么采样时需要加噪声？ 如果 \(p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)\) 是一个确定性映射，那么整个过程退化为一个简单的确定性的从噪声到数据的映射（类似 Normalizing Flow）。但 DDPM 在每一步重新注入随机噪声 \(\sigma_t \cdot \mathbf{z}\)，这赋予了采样过程的随机性——同一个噪声初始值不会总是生成同一个样本。从更深层次看，这个随机性对应于重参数化得分估计（reparameterized score estimation）中的 Langevin 动力学噪声项。

2.2 Score-based 生成模型

Diffusion 模型还有另一个等价的数学视角——通过得分函数（score function）来理解。

2.2.1 得分函数与 Stein Score

数据分布 \(p(\mathbf{x})\) 的得分函数（Stein score）定义为对数密度的梯度：

\[\boxed{\mathbf{s}(\mathbf{x}) = \nabla_{\mathbf{x}} \log p(\mathbf{x})} \tag{2.9}\]

得分函数是一个向量场：在每一点 \(\mathbf{x}\)，它指向对数密度增长最快的方向。如果你站在数据流形的”山坡”上，得分函数告诉你要往哪个方向走才能到达”山顶”（高概率密度区域）。

直观理解：得分函数 \(\neq\) 密度函数。我们不需要知道 \(p(\mathbf{x})\) 的绝对大小（归一化常数），只需要知道它的相对变化率（梯度）。这使得得分估计比密度估计更容易——因为我们可以忽略不可计算的归一化常数。

2.2.2 Langevin 动力学采样

给定得分函数 \(\mathbf{s}(\mathbf{x})\)，我们可以用 Langevin 动力学从 \(p(\mathbf{x})\) 中采样。Langevin 动力学是一个随机微分方程（SDE）：

\[d\mathbf{x} = \nabla_{\mathbf{x}} \log p(\mathbf{x}) \, dt + \sqrt{2} \, d\mathbf{w} \tag{2.10}\]

其中 \(d\mathbf{w}\) 是标准布朗运动（Wiener 过程）。

离散化版本（Euler-Maruyama 方法）：

\[\boxed{\mathbf{x}_{k+1} = \mathbf{x}_k + \eta \cdot \nabla_{\mathbf{x}} \log p(\mathbf{x}_k) + \sqrt{2\eta} \cdot \mathbf{z}_k,\quad \mathbf{z}_k \sim \mathcal{N}(0, \mathbf{I})} \tag{2.11}\]

这里 \(\eta\) 是步长。这个更新规则直观地分为两部分：

漂移项 \(\eta \cdot \nabla_{\mathbf{x}} \log p(\mathbf{x}_k)\)：向高概率密度区域移动
扩散项 \(\sqrt{2\eta} \cdot \mathbf{z}_k\)：随机扰动，防止卡在局部模式中

可以证明，当 \(\eta \to 0\) 且 \(k \to \infty\) 时，\(\mathbf{x}_k\) 的平稳分布恰好是 \(p(\mathbf{x})\)（Welling & Teh, 2011）。

2.2.3 得分匹配（Score Matching）

问题转化为：如何从数据样本中学习得分函数 \(\mathbf{s}_\theta(\mathbf{x})\)？

显式得分匹配（Explicit Score Matching, ESM）：

\[\mathcal{L}_{\text{ESM}}(\theta) = \frac{1}{2} \mathbb{E}_{p(\mathbf{x})}\left[\left\|\mathbf{s}_\theta(\mathbf{x}) - \nabla_{\mathbf{x}} \log p(\mathbf{x})\right\|^2\right] \tag{2.12}\]

这个目标函数不可用——因为我们不知道真实的数据密度 \(p(\mathbf{x})\) 的梯度。

隐式得分匹配（Implicit Score Matching, ISM）（Hyvarinen, 2005）：

通过分部积分，可以消除对 \(\nabla_{\mathbf{x}} \log p(\mathbf{x})\) 的依赖，得到等价的损失函数：

\[\mathcal{L}_{\text{ISM}}(\theta) = \mathbb{E}_{p(\mathbf{x})}\left[\frac{1}{2}\|\mathbf{s}_\theta(\mathbf{x})\|^2 + \text{tr}\big(\nabla_{\mathbf{x}} \mathbf{s}_\theta(\mathbf{x})\big)\right] \tag{2.13}\]

但这在高维空间中仍然代价高昂，因为需要计算 \(\mathbf{s}_\theta\) 的 Jacobian 的迹（\(\text{tr}(\nabla_{\mathbf{x}} \mathbf{s}_\theta)\)）。

2.2.4 去噪得分匹配（Denoising Score Matching, DSM）

Vincent (2011) 提出了一个更实用的方案：不去估计纯净数据的得分，而是估计加噪数据的得分。

给定数据样本 \(\mathbf{x}\)，构造加噪样本 \(\tilde{\mathbf{x}} = \mathbf{x} + \sigma \cdot \boldsymbol{\epsilon}\)，其中 \(\boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})\)。加噪分布的条件密度为 \(q_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) = \mathcal{N}(\tilde{\mathbf{x}}; \mathbf{x}, \sigma^2 \mathbf{I})\)。

去噪得分匹配的目标是：

\[\boxed{\mathcal{L}_{\text{DSM}}(\theta) = \frac{1}{2} \mathbb{E}_{q_\sigma(\tilde{\mathbf{x}}, \mathbf{x})}\left[\left\|\mathbf{s}_\theta(\tilde{\mathbf{x}}) - \nabla_{\tilde{\mathbf{x}}} \log q_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x})\right\|^2\right]} \tag{2.14}\]

由于 \(q_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x})\) 是已知的高斯分布，其得分函数可解析计算：

\[\nabla_{\tilde{\mathbf{x}}} \log q_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x}) = -\frac{\tilde{\mathbf{x}} - \mathbf{x}}{\sigma^2} = -\frac{\boldsymbol{\epsilon}}{\sigma}\]

代入得：

\[\mathcal{L}_{\text{DSM}}(\theta) = \frac{1}{2} \mathbb{E}_{\mathbf{x} \sim p(\mathbf{x}), \boldsymbol{\epsilon} \sim \mathcal{N}(0,\mathbf{I})}\left[\left\|\mathbf{s}_\theta(\mathbf{x} + \sigma\boldsymbol{\epsilon}) + \frac{\boldsymbol{\epsilon}}{\sigma}\right\|^2\right] \tag{2.15}\]

与 DDPM 的联系：令 \(\mathbf{s}_\theta(\mathbf{x}, \sigma) = -\boldsymbol{\epsilon}_\theta(\mathbf{x}, \sigma) / \sigma\)。则 DSM 损失变为：

\[\mathcal{L}_{\text{DSM}}(\theta) = \frac{1}{2\sigma^2} \mathbb{E}\left[\left\|\boldsymbol{\epsilon}_\theta(\mathbf{x} + \sigma\boldsymbol{\epsilon}, \sigma) - \boldsymbol{\epsilon}\right\|^2\right]\]

这实际上与 DDPM 的简化损失函数 (2.8) 是等价的！DDPM 和 DSM 本质上是同一个东西的两种语言——DDPM 从变分推断的角度描述，DSM 从得分估计的角度描述。

Vincent (2011) 的关键定理是：DSM 的最优解等于 ESM 的最优解，即 \(\arg\min_\theta \mathcal{L}_{\text{DSM}}(\theta) = \arg\min_\theta \mathcal{L}_{\text{ESM}}(\theta)\)。这意味着通过估计加噪数据的得分，我们间接地获得了原始数据分布的得分估计。

2.2.5 多噪声尺度与退火 Langevin 动力学

单一噪声尺度存在模态覆盖问题：如果噪声太小，训练样本太少无法覆盖低密度区域；如果噪声太大，得分估计过于平滑。解决方案是使用多个噪声尺度 \(\sigma_1 > \sigma_2 > \cdots > \sigma_L\)：

\[\mathcal{L}(\theta) = \frac{1}{L} \sum_{i=1}^L \lambda(\sigma_i) \cdot \mathcal{L}_{\text{DSM}}(\theta; \sigma_i) \tag{2.16}\]

采样时使用退火 Langevin 动力学：

从随机噪声 \(\mathbf{x}_0 \sim \mathcal{N}(0, \mathbf{I})\) 开始
对每个噪声级别 \(\sigma_i\)，运行 Langevin 动力学，以 \(\mathbf{s}_\theta(\cdot, \sigma_i)\) 为目标
当 \(i \to L\)（最小噪声），样本收敛到数据分布

这种多尺度方法直接催生了后来的 Score SDE 框架（Song et al., 2021），在连续时间极限下，退火 Langevin 动力学等价于求解反向时间 SDE。

2.2.6 Score SDE 的统一视角（简述）

Song et al. (2021) 将扩散过程和得分匹配统一到了 SDE 框架下。前向 SDE：

\[d\mathbf{x} = \mathbf{f}(\mathbf{x}, t)dt + g(t)d\mathbf{w}\]

其中 \(\mathbf{f}\) 是漂移系数，\(g(t)\) 是扩散系数。对应的时间反向 SDE 为：

\[d\mathbf{x} = \big[\mathbf{f}(\mathbf{x}, t) - g(t)^2 \nabla_{\mathbf{x}} \log p_t(\mathbf{x})\big]dt + g(t)d\bar{\mathbf{w}}\]

这里 \(\bar{\mathbf{w}}\) 是反向时间的 Wiener 过程，\(\nabla_{\mathbf{x}} \log p_t(\mathbf{x})\) 正是我们需要估计的得分函数。这个统一框架表明：

DDPM 的 VP-SDE：\(\mathbf{f}(\mathbf{x}, t) = -\frac{1}{2}\beta(t)\mathbf{x}\)，\(g(t) = \sqrt{\beta(t)}\)
Score SDE 的 VE-SDE：\(\mathbf{f}(\mathbf{x}, t) = 0\)，\(g(t) = \frac{d[\sigma^2(t)]}{dt}\)

2.3 Normalizing Flows：变量替换与可逆神经网络

Normalizing Flows（归一化流，Papamakarios et al., 2021）提供了一种完全不同但同样优雅的生成建模方法。

2.3.1 变量替换公式（Change of Variables）

Normalizing Flows 的核心是：将简单分布 \(p_{\mathbf{u}}(\mathbf{u})\)（如标准高斯）通过一个可逆且可微的变换 \(T\) 映射到复杂数据分布：

\[\mathbf{x} = T(\mathbf{u}), \quad \mathbf{u} \sim p_{\mathbf{u}}(\mathbf{u}) \tag{2.17}\]

根据概率密度的变量替换公式（change of variables formula）：

\[\boxed{p_{\mathbf{x}}(\mathbf{x}) = p_{\mathbf{u}}\big(T^{-1}(\mathbf{x})\big) \cdot \left\vert \det J_{T^{-1}}(\mathbf{x})\right\vert = p_{\mathbf{u}}(\mathbf{u}) \cdot \left\vert \det J_T(\mathbf{u})\right\vert ^{-1}} \tag{2.18}\]

其中 \(J_T(\mathbf{u}) \in \mathbb{R}^{D \times D}\) 是变换 \(T\) 在 \(\mathbf{u}\) 处的 Jacobian 矩阵：

\[[J_T(\mathbf{u})]_{ij} = \frac{\partial T_i}{\partial u_j}\]

几何直觉：\(\left\vert \det J_T(\mathbf{u})\right\vert\) 度量了 \(T\) 在 \(\mathbf{u}\) 附近对体积的局部缩放因子。如果 \(\mathbf{u}\) 附近的小邻域 vol\((d\mathbf{u})\) 被 \(T\) 映射为 \(\mathbf{x}\) 附近的小邻域 vol\((d\mathbf{x})\)，则 \(\left\vert \det J_T(\mathbf{u})\right\vert \approx \text{vol}(d\mathbf{x}) / \text{vol}(d\mathbf{u})\)。由于概率质量守恒，密度必须按体积变化的逆比例缩放。

可组合性（Composability）是 Normalizing Flows 最核心的性质。给定两个可逆可微变换 \(T_1\) 和 \(T_2\)，其复合 \(T = T_2 \circ T_1\) 也是可逆可微的，且：

\[T^{-1} = T_1^{-1} \circ T_2^{-1}\] \[\det J_{T_2 \circ T_1}(\mathbf{u}) = \det J_{T_2}(T_1(\mathbf{u})) \cdot \det J_{T_1}(\mathbf{u}) \tag{2.19}\]

这意味着我们可以通过堆叠多个简单变换来构建复杂的整体变换 \(T = T_K \circ T_{K-1} \circ \cdots \circ T_1\)，每一步 \(T_k\) 将中间表示 \(\mathbf{z}_{k-1}\) 映射到 \(\mathbf{z}_k\)（\(\mathbf{z}_0 = \mathbf{u}\)，\(\mathbf{z}_K = \mathbf{x}\)）：

\[\log p_{\mathbf{x}}(\mathbf{x}) = \log p_{\mathbf{u}}(\mathbf{u}) - \sum_{k=1}^K \log\left\vert \det J_{T_k}(\mathbf{z}_{k-1})\right\vert \tag{2.20}\]

这就是 “Normalizing Flow” 名称的由来：”Flow” 表示样本从简单分布逐步流动到复杂分布，”Normalizing” 表示反向过程将复杂数据”归一化”为标准分布。

2.3.2 自回归流（Autoregressive Flows）

自回归流直接实现了公式 (2.18) 的构造性证明中使用的三角映射。其形式为：

\[z'_i = \tau(z_i; \mathbf{h}_i), \quad \mathbf{h}_i = c_i(\mathbf{z}_{<i}) \tag{2.21}\]

其中：

\(\tau\) 是变换器（transformer），严格单调（因此可逆），由参数 \(\mathbf{h}_i\) 控制
\(c_i\) 是条件器（conditioner），仅以前面维度 \(\mathbf{z}_{<i} = (z_1, \ldots, z_{i-1})\) 为输入

Jacobian 结构：由于第 \(i\) 维的输出 \(z'_i\) 只依赖输入 \(z_1, \ldots, z_i\)，Jacobian 矩阵是下三角的：

\[J_f(\mathbf{z}) = \begin{bmatrix} \frac{\partial \tau}{\partial z_1}(z_1; \mathbf{h}_1) & 0 & \cdots & 0 \\ * & \frac{\partial \tau}{\partial z_2}(z_2; \mathbf{h}_2) & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ * & * & \cdots & \frac{\partial \tau}{\partial z_D}(z_D; \mathbf{h}_D) \end{bmatrix}\]

行列式为对角元素之积（\(O(D)\) 计算）：

\[\log\left\vert \det J_f(\mathbf{z})\right\vert = \sum_{i=1}^D \log\left\vert \frac{\partial \tau}{\partial z_i}(z_i; \mathbf{h}_i)\right\vert \tag{2.22}\]

变换器（Transformer）类型：

仿射变换器（Affine Transformer）：\(\tau(z_i; \alpha_i, \beta_i) = \alpha_i z_i + \beta_i\)，Jacobian 对数为 \(\sum_i \log\vert \alpha_i\vert\)。这是最简单、计算最快的选择，但表达能力有限——单个仿射自回归层只能将高斯分布变换为另一个高斯分布（条件分布意义上）。
基于组合的变换器（Combination-based Transformer）：\(\tau(z_i) = w_{i0} + \sum_{k=1}^K w_{ik} \sigma(\alpha_{ik} z_i + \beta_{ik})\)，满足 \(\alpha_{ik} > 0\) 和 \(w_{ik} > 0\) 以保证单调性。这是单调单层感知机的变体，具有任意逼近能力，但不可解析求逆（需要二分搜索）。
基于积分的变换器（Integration-based Transformer）：\(\tau(z_i) = \int_0^{z_i} g(t; \boldsymbol{\alpha}_i) dt + \beta_i\)，其中 \(g > 0\) 是一个神经网络。其导数直接等于 \(g(z_i)\)，避免了 Jacobian 计算中的额外开销。
基于样条的变换器（Spline-based Transformer）：将 \(\tau\) 实现为分段单调样条函数，如在每段上使用有理二次函数。样条变换器不仅灵活，而且解析可逆——这是其他非线性变换器难以同时满足的性质。

条件器（Conditioner）类型：

a. 循环条件器（Recurrent）：用 RNN 处理 \(\mathbf{z}\) 序列，输出逐维条件参数。顺序计算，对大维度数据慢。

b. 掩码条件器（Masked）：使用单个前馈网络，通过权重掩码确保自回归结构——第 \(i\) 个输出不依赖输入 \(\geq i\)。如 MADE (Germain et al., 2015)。评估快（单次前向传播），但逆向慢（需要 \(O(D)\) 次迭代）。

c. 耦合层（Coupling Layer）：将 \(\mathbf{z}\) 切分为两部分 \(\mathbf{z} = [\mathbf{z}_{\leq d}, \mathbf{z}_{>d}]\)，前半部分直接恒等映射，后半部分以前半为条件变换：

\[\begin{aligned} \mathbf{z}'_{\leq d} &= \mathbf{z}_{\leq d} \\ \mathbf{z}'_{>d} &= \tau(\mathbf{z}_{>d}; F(\mathbf{z}_{\leq d})) \end{aligned}\]

耦合层的 Jacobian 是分块下三角矩阵，行列式仅来自后半部分的对角元素。关键是：耦合层正逆计算代价相同。组成多层的耦合层并通过置换层穿插，可以获得足够的表达能力。被 Real NVP、Glow 等 SOTA 图像生成模型采用。

2.3.3 残差流（Residual Flows）

残差变换形式为 \(\mathbf{z}' = \mathbf{z} + g_\phi(\mathbf{z})\)，形如 ResNet 的残差块，但需要额外的约束来保证可逆性。

收缩残差流（Contractive Residual Flow）：要求 \(g_\phi\) 是压缩映射（Lipschitz 常数 \(L < 1\)）。由 Banach 不动点定理，存在唯一的逆，且可通过不动点迭代求得：

\[\mathbf{z}_{k+1} = \mathbf{z}' - g_\phi(\mathbf{z}_k), \quad \text{收敛到 } \mathbf{z}^* = f^{-1}(\mathbf{z}')\]

主要缺点是 Jacobian 行列式的计算代价高昂——这里没有三角结构可以利用。需要通过 Hutchinson 迹估计器和幂级数展开来近似。

基于矩阵行列式引理的残差流：

Planar Flow（Rezende & Mohamed, 2015）：

\[\mathbf{z}' = \mathbf{z} + \mathbf{v} \sigma(\mathbf{w}^T \mathbf{z} + b)\]

其中 \(\mathbf{v}, \mathbf{w} \in \mathbb{R}^D, b \in \mathbb{R}\)，\(\sigma\) 是激活函数。Jacobian 是 \(\mathbf{I} + \sigma'(\mathbf{w}^T\mathbf{z} + b) \mathbf{v}\mathbf{w}^T\)（单位阵 + 秩-1更新），利用矩阵行列式引理：

\[\det(\mathbf{I} + \sigma'(\mathbf{w}^T\mathbf{z} + b) \mathbf{v}\mathbf{w}^T) = 1 + \sigma'(\mathbf{w}^T\mathbf{z} + b) \cdot \mathbf{w}^T \mathbf{v}\]

计算复杂度 \(O(D)\)。

Sylvester Flow：推广到 \(M\) 个隐藏单元：

\[\mathbf{z}' = \mathbf{z} + \mathbf{V} \sigma(\mathbf{W}^T \mathbf{z} + \mathbf{b})\]

Jacobian 行列式可通过 \(\det(\mathbf{I}_M + \text{diag}(\sigma'(\mathbf{W}^T\mathbf{z} + \mathbf{b})) \cdot \mathbf{W}^T \mathbf{V})\) 计算，复杂度 \(O(M^3 + DM^2)\)。

2.3.4 连续时间流（Continuous Normalizing Flows, CNF）

如果我们将变换的堆叠推向极限——无限多步、步长趋于零——就得到了连续时间流：

\[\frac{d\mathbf{z}_t}{dt} = g_\phi(t, \mathbf{z}_t) \tag{2.23}\]

初始条件 \(\mathbf{z}_{t_0} = \mathbf{u}\)，最终状态 \(\mathbf{z}_{t_1} = \mathbf{x}\)。

瞬时变量替换公式（Instantaneous Change of Variables, Chen et al., 2018）：

\[\boxed{\frac{d \log p(\mathbf{z}_t)}{dt} = -\text{tr}\left(J_{g_\phi(t,\cdot)}(\mathbf{z}_t)\right) = -\nabla_{\mathbf{z}_t} \cdot g_\phi(t, \mathbf{z}_t)} \tag{2.24}\]

这是 Fokker-Planck 方程在零扩散极限下的特殊情况。通过联合积分 \(\mathbf{z}_t\) 和 \(\log p(\mathbf{z}_t)\)，可以同时计算正向变换和密度评估：

\[\begin{bmatrix} \mathbf{z}_{t_1} \\ \log p(\mathbf{x}) \end{bmatrix} = \begin{bmatrix} \mathbf{u} \\ \log p_{\mathbf{u}}(\mathbf{u}) \end{bmatrix} + \int_{t_0}^{t_1} \begin{bmatrix} g_\phi(t, \mathbf{z}_t) \\ -\text{tr}(J_{g_\phi})(\mathbf{z}_t) \end{bmatrix} dt \tag{2.25}\]

CNF 的关键优势：

不需要 Jacobian 行列式计算——只需要迹（trace），可以用 Hutchinson 迹估计器高效近似：\(\text{tr}(J_g) \approx \mathbf{v}^T J_g \mathbf{v}\)，\(\mathbf{v}\) 是零均值单位协方差的随机向量
对 \(g_\phi\) 的结构几乎无约束——只需满足一致 Lipschitz 条件
正向和逆向计算代价等同
可以使用自适应步长的 ODE 求解器（如 dopri5）

伴随方法（Adjoint Method）：CNF 的梯度通过求解伴随 ODE 来计算，而非回传整个求解器的计算图：

\[\frac{d}{dt}\frac{\partial \mathcal{L}}{\partial \mathbf{z}_t} = -\left(\frac{\partial \mathcal{L}}{\partial \mathbf{z}_t}\right)^T \frac{\partial g_\phi(t, \mathbf{z}_t)}{\partial \mathbf{z}_t}\]

这可以视为连续时间版本的反向传播。

2.4 Flow Matching：统一扩散模型与 CNF

Flow Matching（Lipman et al., 2023）提出了一种全新的 CNF 训练范式，它的核心贡献是：不需要模拟训练时的 ODE 轨迹（simulation-free training），即可训练任意概率路径的 CNF。

2.4.1 问题设置

我们有未知数据分布 \(q(\mathbf{x}_1)\) 的样本。目标：学习一个向量场 \(v_t: [0,1] \times \mathbb{R}^d \to \mathbb{R}^d\)（由参数 \(\theta\) 的神经网络 \(v_\theta\) 建模），使其对应的流 \(\phi_t\) 推动简单先验分布 \(p_0\) 到近似 \(q\) 的分布 \(p_1\)。

流 \(\phi_t\) 由 ODE 定义：

\[\frac{d}{dt}\phi_t(\mathbf{x}) = v_t(\phi_t(\mathbf{x})), \quad \phi_0(\mathbf{x}) = \mathbf{x} \tag{2.26}\]

概率路径 \(p_t\) 由 push-forward 操作定义：

\[p_t = [\phi_t]_* p_0 \tag{2.27}\]

连续性方程（Continuity Equation）提供了验证向量场 \(v_t\) 是否生成概率路径 \(p_t\) 的方法：

\[\frac{\partial p_t(\mathbf{x})}{\partial t} + \nabla_{\mathbf{x}} \cdot \big(p_t(\mathbf{x}) v_t(\mathbf{x})\big) = 0 \tag{2.28}\]

2.4.2 Flow Matching 目标函数

给定一个目标概率路径 \(p_t(\mathbf{x})\) 和其对应的向量场 \(u_t(\mathbf{x})\)（满足连续性方程），Flow Matching 目标函数为：

\[\boxed{\mathcal{L}_{\text{FM}}(\theta) = \mathbb{E}_{t, p_t(\mathbf{x})}\left[\|v_\theta(t, \mathbf{x}) - u_t(\mathbf{x})\|^2\right]} \tag{2.29}\]

其中 \(t \sim \mathcal{U}[0, 1]\)，\(\mathbf{x} \sim p_t(\mathbf{x})\)。

这个目标直观但不可直接使用，因为我们既不知道 \(p_t\) 也不知道 \(u_t\)。

2.4.3 条件概率路径与条件向量场

核心思想：定义每样本（per-sample）条件概率路径 \(p_t(\mathbf{x} \mid \mathbf{x}_1)\) 和条件向量场 \(u_t(\mathbf{x} \mid \mathbf{x}_1)\)，然后通过边缘化构建目标：

\[p_t(\mathbf{x}) = \int p_t(\mathbf{x} \mid \mathbf{x}_1) q(\mathbf{x}_1) d\mathbf{x}_1 \tag{2.30}\] \[u_t(\mathbf{x}) = \int u_t(\mathbf{x} \mid \mathbf{x}_1) \frac{p_t(\mathbf{x} \mid \mathbf{x}_1) q(\mathbf{x}_1)}{p_t(\mathbf{x})} d\mathbf{x}_1 \tag{2.31}\]

定理 1（Lipman et al., 2023）：如果每个 \(u_t(\cdot \mid \mathbf{x}_1)\) 生成对应的 \(p_t(\cdot \mid \mathbf{x}_1)\)，那么公式 (2.31) 定义的边缘向量场 \(u_t\) 生成边缘概率路径 \(p_t\)。

定理 2（条件 Flow Matching 的等价性）：令 \(\mathcal{L}_{\text{CFM}}\) 为条件目标函数：

\[\boxed{\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{t, q(\mathbf{x}_1), p_t(\mathbf{x} \mid \mathbf{x}_1)}\left[\|v_\theta(t, \mathbf{x}) - u_t(\mathbf{x} \mid \mathbf{x}_1)\|^2\right]} \tag{2.32}\]

则在期望意义上，\(\mathcal{L}_{\text{CFM}}\) 和 \(\mathcal{L}_{\text{FM}}\) 的梯度相等：\(\nabla_\theta \mathcal{L}_{\text{FM}}(\theta) = \nabla_\theta \mathcal{L}_{\text{CFM}}(\theta)\)。

定理二的证明概要：

\[\begin{aligned} \mathcal{L}_{\text{FM}}(\theta) &= \mathbb{E}_{t, p_t(\mathbf{x})}\left[\|v_\theta - u_t\|^2\right] \\ &= \mathbb{E}_{t, p_t(\mathbf{x})}\left[\|v_\theta\|^2 - 2\langle v_\theta, u_t \rangle + \|u_t\|^2\right] \end{aligned}\]

\(\|u_t\|^2\) 与 \(\theta\) 无关。对于交叉项：

\[\begin{aligned} \mathbb{E}_{p_t(\mathbf{x})}\langle v_\theta, u_t \rangle &= \int \langle v_\theta(\mathbf{x}), \int u_t(\mathbf{x} \mid \mathbf{x}_1) \frac{p_t(\mathbf{x} \mid \mathbf{x}_1) q(\mathbf{x}_1)}{p_t(\mathbf{x})} d\mathbf{x}_1 \rangle p_t(\mathbf{x}) d\mathbf{x} \\ &= \int \langle v_\theta(\mathbf{x}), u_t(\mathbf{x} \mid \mathbf{x}_1) \rangle p_t(\mathbf{x} \mid \mathbf{x}_1) q(\mathbf{x}_1) d\mathbf{x} d\mathbf{x}_1 \end{aligned}\]

对于 \(\|v_\theta\|^2\) 项，利用 \(p_t(\mathbf{x}) = \int p_t(\mathbf{x} \mid \mathbf{x}_1) q(\mathbf{x}_1) d\mathbf{x}_1\)，交换积分次序后得到同样的 \(\mathbb{E}_{q(\mathbf{x}_1), p_t(\mathbf{x} \mid \mathbf{x}_1)}[\|v_\theta\|^2]\)。因此：

\[\mathcal{L}_{\text{FM}}(\theta) = \mathbb{E}_{t, q(\mathbf{x}_1), p_t(\mathbf{x} \mid \mathbf{x}_1)}\left[\|v_\theta - u_t(\cdot \mid \mathbf{x}_1)\|^2\right] + \text{const}\]

2.4.4 高斯条件概率路径

选择条件概率路径为高斯分布：

\[p_t(\mathbf{x} \mid \mathbf{x}_1) = \mathcal{N}\big(\mathbf{x}; \mu_t(\mathbf{x}_1), \sigma_t(\mathbf{x}_1)^2 \mathbf{I}\big) \tag{2.33}\]

边界条件：\(t=0\) 时 \(p_0 = \mathcal{N}(0, \mathbf{I})\)，\(t=1\) 时 \(p_1(\mathbf{x} \mid \mathbf{x}_1) \approx \delta_{\mathbf{x}_1}\)（集中分布在 \(\mathbf{x}_1\) 附近）。

流映射为仿射变换：\(\psi_t(\mathbf{x}_0) = \sigma_t(\mathbf{x}_1) \mathbf{x}_0 + \mu_t(\mathbf{x}_1)\)，其中 \(\mathbf{x}_0 \sim \mathcal{N}(0, \mathbf{I})\)。

定理 3（高斯路径的向量场）：对于高斯条件概率路径，生成流 \(\psi_t\) 的唯一向量场为：

\[\boxed{u_t(\mathbf{x} \mid \mathbf{x}_1) = \frac{\sigma'_t(\mathbf{x}_1)}{\sigma_t(\mathbf{x}_1)}\big(\mathbf{x} - \mu_t(\mathbf{x}_1)\big) + \mu'_t(\mathbf{x}_1)} \tag{2.34}\]

CFM 损失简化为：

\[\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{t, q(\mathbf{x}_1), p(\mathbf{x}_0)}\left[\left\|v_\theta(t, \psi_t(\mathbf{x}_0)) - \frac{d}{dt}\psi_t(\mathbf{x}_0)\right\|^2\right] \tag{2.35}\]

2.4.5 Optimal Transport 条件路径

选择均值和标准差线性变化：

\[\mu_t(\mathbf{x}_1) = t \mathbf{x}_1, \quad \sigma_t(\mathbf{x}_1) = 1 - (1 - \sigma_{\min})t \tag{2.36}\]

代入定理 3，得到 OT 条件向量场：

\[u_t(\mathbf{x} \mid \mathbf{x}_1) = \frac{\mathbf{x}_1 - (1 - \sigma_{\min})\mathbf{x}}{1 - (1 - \sigma_{\min})t} \tag{2.37}\]

以及流：

\[\psi_t(\mathbf{x}_0) = (1 - (1 - \sigma_{\min})t)\mathbf{x}_0 + t \mathbf{x}_1 \tag{2.38}\]

CFM 训练损失：

\[\mathcal{L}_{\text{CFM-OT}}(\theta) = \mathbb{E}_{t, q(\mathbf{x}_1), p(\mathbf{x}_0)}\left[\left\|v_\theta(t, \psi_t(\mathbf{x}_0)) - (\mathbf{x}_1 - (1 - \sigma_{\min})\mathbf{x}_0)\right\|^2\right] \tag{2.39}\]

为什么 OT 路径更优？

直线轨迹：\(\psi_t(\mathbf{x}_0)\) 在 \(\mathbf{x}_0\) 和 \(\mathbf{x}_1\) 之间线性插值。粒子沿直线以恒定速度移动。相比之下，扩散路径（如 VP/VE）是弯曲的，常出现”overshoot”和”backtracking”。
更简单的回归目标：OT 向量场的方向与时间无关——\(u_t(\mathbf{x} \mid \mathbf{x}_1) = g(t) \cdot h(\mathbf{x}, \mathbf{x}_1)\)，简化了神经网络的拟合任务。
更快的训练和采样：实验表明，FM-OT 达到相同 FID 需要的训练迭代数显著更少；采样时可以使用更少的ODE求解器步数（NFE）。

OT 条件路径与扩散条件路径的对比：

属性	Diffusion (VP)	Flow Matching OT
\(\mu_t(\mathbf{x}_1)\)	\(\alpha_{1-t}\mathbf{x}_1\)	\(t\mathbf{x}_1\)
\(\sigma_t(\mathbf{x}_1)\)	\(\sqrt{1-\alpha_{1-t}^2}\)	\(1 - (1-\sigma_{\min})t\)
轨迹形状	弯曲	直线
\(t=0\) 是否精确为噪声	否（有限时间近似）	是
定义域	\(t \in (0, 1)\) 的近似	\(t \in [0, 1]\) 精确

三、论文精读

3.1 Papamakarios et al. (2021)：Normalizing Flows 的统一视角

这篇 64 页的 JMLR 综述文章建立了 Normalizing Flows 的完整数学框架和历史脉络。以下是对关键章节的逐段解读。

3.1.1 定义与基础（Section 2.1）

论文的核心定义已在 2.3.1 节中详细展开。值得强调的观点是：论文将 Normalizing Flows 明确为一种构造灵活概率分布的通用机制——仅需指定一个（简单的）基分布和一系列可逆变换。这种表述将 NF 从具体的技术细节中解放出来，将其提升为一种建模哲学。

论文中 Figure 1 展示了一个 4 步流将标准正态基分布变换为十字形目标分布的过程，这个可视化深刻地传达了 NF 的核心直觉：从简单到复杂，通过连续的空间”变形”。

3.1.2 表达力（Section 2.2）

NF 的表达力是通过一个构造性证明建立的：给定任意分布 \(p_{\mathbf{x}}(\mathbf{x}) > 0\)，利用概率链式法则分解 \(p_{\mathbf{x}}(\mathbf{x}) = \prod_{i=1}^D p_{\mathbf{x}}(x_i \mid \mathbf{x}_{<i})\)，然后定义变换 \(F_i(x_i, \mathbf{x}_{<i}) = \int_{-\infty}^{x_i} p_{\mathbf{x}}(x'_i \mid \mathbf{x}_{<i}) dx'_i\)（条件 CDF）。

\(F\) 的 Jacobian 是下三角的，对角线元素为 \(p_{\mathbf{x}}(x_i \mid \mathbf{x}_{<i})\)，因此 \(\det J_F(\mathbf{x}) = p_{\mathbf{x}}(\mathbf{x})\)。\(F\) 将 \(p_{\mathbf{x}}\) 变换为 \((0, 1)^D\) 上的均匀分布。通过 \(T = F^{-1} \circ G\)（\(G\) 类似地将基分布映射到均匀分布），可以得到从任意基分布到任意目标分布的可逆变换。

这个证明的重要性：它不仅保证了 NF 的理论完备性，而且直接启发了自回归流（autoregressive flows）的设计——自回归流正是这个构造性证明中三角映射的直接实现。

3.1.3 KL 散度训练（Section 2.3）

论文区分了两种训练模式：

前向 KL（\(D_{KL}[p_{\mathbf{x}}^* \| p_{\mathbf{x}}(\cdot; \theta)]\)）：等价于最大似然估计。需要 \(T^{-1}\) 和 Jacobian 行列式。
反向 KL（\(D_{KL}[p_{\mathbf{x}}(\cdot; \theta) \| p_{\mathbf{x}}^*]\)）：等价于变分推断。需要 \(T\)（正向变换）但不需要 \(T^{-1}\)。

两者关于 \(T\) 和 \(T^{-1}\) 的角色互换：前向 KL 等价于将训练数据通过 \(T^{-1}\) 拟合到基分布（反向 KL 意义下）。

3.1.4 有限流架构全景（Section 3）

论文将离散有限流分为三大类（Table 1），每一类都有 Jacobian 行列式计算的特定技巧：

自回归流：利用三角 Jacobian，行列式 = 对角线元素之积
线性流：通过结构化矩阵（三角、QR、PLU、Householder）参数化 \(\mathbf{W}\)
残差流：通过收缩映射或矩阵行列式引理控制可逆性

这个分类法的核心逻辑是按 Jacobian 行列式的计算方式分组——这反映了 NF 设计中最根本的约束：如何高效计算 \(\det J_T\)。

3.1.5 CNF 的连续时间表述（Section 4）

CNF 将离散步骤推向连续极限，得到了一个深刻的结论：连续流的设计约束远少于离散流。离散流需要精心选择架构以保证 Jacobian 行列式的高效计算；CNF 只需要 \(g_\phi\) 满足一致 Lipschitz 条件，Jacobian 的迹可以通过 Hutchinson 估计器高效近似。

Euler 离散化揭示了一个重要的联系：

\[\mathbf{z}_{t+\epsilon} = \mathbf{z}_t + \epsilon \cdot g_\phi(t, \mathbf{z}_t)\]

这恰好是一个残差流（公式 2.23 当 \(\epsilon < 1/L\) 时确实是一个收缩映射）。因此，CNF 可以视为残差流的连续极限。

3.1.6 关键评价

Papamakarios et al. 的综述在 NF 领域的价值相当于 Goodfellow et al. 在 DL 领域的价值——不是提出了新方法，而是提供了一种统一的思维框架。但其局限性在于：写于 CNF 大规模应用之前，未能覆盖 Flow Matching 等仿真无关训练方法。

3.2 Lipman et al. (2023)：Flow Matching 的突破

3.2.1 核心贡献识别

本文的核心贡献可以概括为一句话：提出 Conditional Flow Matching 定理，使得 CNF 可以无需模拟 ODE 即可训练，打破了此前 CNF 只能通过计算昂贵的最大似然或存在偏差梯度的仿真无关方法训练的瓶颈。

3.2.2 与去噪得分匹配的类比

FM 与 DSM 之间存在深刻的数学类比：

DSM	FM
目标：\(\| \mathbf{s}_\theta(\tilde{\mathbf{x}}) - \nabla \log q(\tilde{\mathbf{x}} \mid \mathbf{x})\|^2\)	目标：\(\| v_\theta(\mathbf{x}) - u_t(\mathbf{x} \mid \mathbf{x}_1)\|^2\)
条件目标等价于边缘目标	条件目标等价于边缘目标
限制于扩散路径	可用于任意概率路径

FM 的关键推广在于：DSM 只能用于扩散过程定义的路径，而 FM 可以用于任何满足边界条件的概率路径——尤其是 OT 路径。

3.2.3 OT 路径的实验验证

Table 1 中 FM-OT 在 ImageNet 各分辨率上一致地优于基于扩散的基线：

CIFAR-10: FM-OT 的 FID 为 6.35（DDPM 为 7.48），NFE 仅 142（DDPM 为 274）
ImageNet 64x64: FM-OT 的 FID 为 14.45（Score Matching 为 19.74），NFE 仅 138（Score Matching 为 441）

更关键的是训练效率：Figure 5 显示 FM-OT 在 ImageNet 64x64 上的 FID 下降速度远快于扩散基线。FM-OT 在 500k 迭代后即达到扩散模型需要数倍迭代才能达到的图像质量。

Figure 7 展示了采样效率：FM-OT 仅需扩散模型约 60% 的 NFE 即可达到相同的数值精度。

3.2.4 局限性

OT 路径的边缘向量场不一定是最优传输的边缘解——条件 OT 不等价于边缘 OT
论文未探索非高斯条件路径（如用于离散数据或结构化输出的路径）
FM 的高维似然评估仍需求解 ODE，虽然比 MLE 训练时的模拟要少

四、算法与代码实现

4.1 DDPM：完整 PyTorch 实现

以下代码实现了一个最小但可工作的 DDPM，包含 U-Net 去噪网络、前向加噪、逆向采样和训练循环。

"""
ddpm.py — Minimal DDPM Implementation
Based on: Ho et al., "Denoising Diffusion Probabilistic Models", NeurIPS 2020.
Ref: Eq. (2.1) forward process, Eq. (2.8) simplified loss, Algorithm 2.2 sampling.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# ============================================================
# U-Net building blocks — [UNSPECIFIED] channel dimensions
# ============================================================

class SinusoidalPositionEmbedding(nn.Module):
    """Sinusoidal time embedding as in Transformer/Vaswani et al."""
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        # t: (batch,) integer timesteps
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
        emb = t[:, None].float() * emb[None, :]          # (batch, half_dim)
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
        return emb  # (batch, dim)


class ResBlock(nn.Module):
    """Residual block with time embedding conditioning."""
    def __init__(self, in_ch: int, out_ch: int, time_emb_dim: int):
        super().__init__()
        self.norm1 = nn.GroupNorm(8, in_ch)
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.norm2 = nn.GroupNorm(8, out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.time_mlp = nn.Sequential(
            nn.SiLU(), nn.Linear(time_emb_dim, out_ch)
        )
        self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()

    def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
        # t_emb: (batch, time_emb_dim) -> (batch, out_ch, 1, 1)
        h = self.conv1(F.silu(self.norm1(x)))
        h = h + self.time_mlp(t_emb)[:, :, None, None]
        h = self.conv2(F.silu(self.norm2(h)))
        return h + self.skip(x)


class UNet(nn.Module):
    """Simple U-Net for DDPM noise prediction.
       [UNSPECIFIED]: channel multipliers and number of ResBlocks per level."""
    def __init__(self, in_ch: int = 3, base_ch: int = 64,
                 ch_mult: tuple = (1, 2, 4, 8), time_emb_dim: int = 256):
        super().__init__()
        self.time_mlp = nn.Sequential(
            SinusoidalPositionEmbedding(time_emb_dim),
            nn.Linear(time_emb_dim, time_emb_dim),
            nn.SiLU(),
            nn.Linear(time_emb_dim, time_emb_dim),
        )
        self.init_conv = nn.Conv2d(in_ch, base_ch, 3, padding=1)

        # Encoder
        self.downs = nn.ModuleList()
        chs = [base_ch]
        in_ch_curr = base_ch
        for mult in ch_mult:
            out_ch = base_ch * mult
            self.downs.append(nn.ModuleDict({
                'res1': ResBlock(in_ch_curr, out_ch, time_emb_dim),
                'res2': ResBlock(out_ch, out_ch, time_emb_dim),
                'down': nn.Conv2d(out_ch, out_ch, 3, stride=2, padding=1),
            }))
            chs.append(out_ch)
            in_ch_curr = out_ch

        # Bottleneck
        mid_ch = base_ch * ch_mult[-1]
        self.mid_res1 = ResBlock(mid_ch, mid_ch, time_emb_dim)
        self.mid_res2 = ResBlock(mid_ch, mid_ch, time_emb_dim)

        # Decoder
        self.ups = nn.ModuleList()
        for mult in reversed(ch_mult):
            out_ch = base_ch * mult
            self.ups.append(nn.ModuleDict({
                'res1': ResBlock(in_ch_curr + out_ch, out_ch, time_emb_dim),
                'res2': ResBlock(out_ch, out_ch, time_emb_dim),
                'up': nn.ConvTranspose2d(out_ch, out_ch, 4, stride=2, padding=1),
            }))
            in_ch_curr = out_ch

        self.out_norm = nn.GroupNorm(8, base_ch)
        self.out_conv = nn.Conv2d(base_ch, in_ch, 3, padding=1)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        t_emb = self.time_mlp(t)                        # (batch, time_emb_dim)
        h = self.init_conv(x)
        skips = [h]

        # Downsampling
        for layer in self.downs:
            h = layer['res1'](h, t_emb)
            h = layer['res2'](h, t_emb)
            skips.append(h)
            h = layer['down'](h)

        # Bottleneck
        h = self.mid_res1(h, t_emb)
        h = self.mid_res2(h, t_emb)

        # Upsampling with skip connections
        for layer in self.ups:
            skip = skips.pop()
            h = torch.cat([h, skip], dim=1)
            h = layer['res1'](h, t_emb)
            h = layer['res2'](h, t_emb)
            h = layer['up'](h)

        h = self.out_norm(h)
        h = F.silu(h)
        return self.out_conv(h)  # predicted noise ε_θ(x_t, t)


# ============================================================
# DDPM Scheduler — Eq. (2.1), (2.2), (2.4), Algorithm 2.2
# ============================================================

class DDPMScheduler:
    """Manages α_t, ᾱ_t, β_t and provides q-sample and p-sample steps."""
    def __init__(self, num_timesteps: int = 1000,
                 beta_start: float = 1e-4, beta_end: float = 0.02):
        self.num_timesteps = num_timesteps

        # Linear variance schedule [UNSPECIFIED: cosine schedule often better]
        betas = torch.linspace(beta_start, beta_end, num_timesteps)
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)    # ᾱ_t, Eq. (2.2)
        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)

        # Precompute quantities for Eq. (2.2) and Algorithm 2.2
        self.sqrt_alphas_cumprod = alphas_cumprod.sqrt()           # √ᾱ_t
        self.sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod).sqrt()  # √(1-ᾱ_t)

        # For sampling: Eq. (2.4) coefficients
        self.sqrt_recip_alphas = (1.0 / alphas).sqrt()             # 1/√α_t
        self.posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)

    def q_sample(self, x0: torch.Tensor, t: torch.Tensor,
                 noise: torch.Tensor) -> torch.Tensor:
        """Forward diffusion: x_t = √ᾱ_t · x_0 + √(1-ᾱ_t) · ε.
           Ref: Eq. (2.2)."""
        a = self._gather(self.sqrt_alphas_cumprod, t).to(x0.device)
        b = self._gather(self.sqrt_one_minus_alphas_cumprod, t).to(x0.device)
        return a * x0 + b * noise

    def p_sample(self, model: nn.Module, x: torch.Tensor,
                 t: torch.Tensor, t_index: int) -> torch.Tensor:
        """Reverse denoising step. Ref: Algorithm 2.2."""
        betas_t = self._gather(self.posterior_variance, t).to(x.device)
        sqrt_recip_alphas_t = self._gather(self.sqrt_recip_alphas, t).to(x.device)
        sqrt_one_minus_alphas_cumprod_t = self._gather(
            self.sqrt_one_minus_alphas_cumprod, t).to(x.device)

        # Eq. (2.5) rearranged to predict x_0 and then mean
        pred_noise = model(x, t)
        # x_0 estimate from Eq. (2.2) inverse
        x0_pred = sqrt_recip_alphas_t * (x - sqrt_one_minus_alphas_cumprod_t * pred_noise)

        # Clamp to [-1, 1] for stability (DDPM trick)
        x0_pred = torch.clamp(x0_pred, -1.0, 1.0)

        # Mean of q(x_{t-1} | x_t, x_0) — Eq. (2.5)
        # Rearranged: 1/√α_t · (x_t − β_t/√(1-ᾱ_t) · ε)
        mean = sqrt_recip_alphas_t * (
            x - betas_t * pred_noise / sqrt_one_minus_alphas_cumprod_t
        )

        if t_index == 0:
            return mean
        else:
            noise = torch.randn_like(x)
            return mean + torch.sqrt(betas_t) * noise

    def _gather(self, vals: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """Index into precomputed tensors by timestep."""
        return vals.gather(-1, t.cpu()).view(-1, 1, 1, 1)


# ============================================================
# Training Loop — Eq. (2.8) simplified loss
# ============================================================

def train_step(model: nn.Module, scheduler: DDPMScheduler,
               x0: torch.Tensor, optimizer: torch.optim.Optimizer,
               device: torch.device) -> float:
    """Single DDPM training step. Ref: Algorithm 2.1, Eq. (2.8)."""
    batch_size = x0.shape[0]
    x0 = x0.to(device)

    # Step 2: random timestep
    t = torch.randint(0, scheduler.num_timesteps, (batch_size,), device=device)

    # Step 3: noise
    noise = torch.randn_like(x0)

    # Forward diffusion (Eq. 2.2)
    xt = scheduler.q_sample(x0, t, noise)

    # Predict noise
    pred_noise = model(xt, t)

    # Simplified loss (Eq. 2.8)
    loss = F.mse_loss(pred_noise, noise)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return loss.item()


# ============================================================
# Sampling — Algorithm 2.2
# ============================================================

@torch.no_grad()
def sample(model: nn.Module, scheduler: DDPMScheduler,
           shape: tuple, device: torch.device) -> torch.Tensor:
    """Generate samples from the DDPM. Ref: Algorithm 2.2."""
    model.eval()
    x = torch.randn(shape, device=device)       # Step 1: x_T ~ N(0, I)

    for i in reversed(range(scheduler.num_timesteps)):
        t = torch.full((shape[0],), i, device=device, dtype=torch.long)
        x = scheduler.p_sample(model, x, t, i)   # Steps 2-4
    return x

代码关键注释：

前向扩散：q_sample 直接实现了公式 (2.2)，一次性从 \(\mathbf{x}_0\) 和 \(\boldsymbol{\epsilon}\) 生成任意时刻的 \(\mathbf{x}_t\)。这在训练中避免了 \(T\) 步循环。
逆向采样：p_sample 实现了算法 2.2 的单步去噪。它先通过预测的噪声 \(\boldsymbol{\epsilon}_\theta\) 反推 \(\hat{\mathbf{x}}_0\)（用于 clamp 以稳定采样），然后利用公式 (2.5) 计算 \(\tilde{\boldsymbol{\mu}}_t\)。
损失函数：F.mse_loss(pred_noise, noise) 直接对应公式 (2.8) 的简化损失。注意 DDPM 论文中的简化损失丢弃了时间步相关的权重系数——这是经验性的选择，在公式推导中没有理论必然性。

4.2 Continuous Normalizing Flow：PyTorch 实现

"""
cnf.py — Minimal Continuous Normalizing Flow with Neural ODE
Based on: Chen et al., "Neural Ordinary Differential Equations", NeurIPS 2018.
Ref: Eq. (2.23) ODE dynamics, Eq. (2.24) instantaneous change of variables.
"""
import torch
import torch.nn as nn
from torchdiffeq import odeint, odeint_adjoint


class CNF(nn.Module):
    """Continuous Normalizing Flow with Hutchinson trace estimator.

       The dynamics network g_φ(t, z_t) outputs both the velocity and
       an estimate of the log-density change via Hutchinson's estimator.
       Ref: Eq. (2.24), (2.25).
    """
    def __init__(self, dim: int, hidden_dim: int = 256,
                 num_layers: int = 3):
        super().__init__()
        self.dim = dim

        # g_φ(t, z): [dim+1] -> dim  (concatenate time)
        layers = []
        in_dim = dim + 1
        for i in range(num_layers):
            layers.append(nn.Linear(in_dim, hidden_dim))
            layers.append(nn.Softplus())
            in_dim = hidden_dim
        layers.append(nn.Linear(hidden_dim, dim))
        self.net = nn.Sequential(*layers)

        # For Hutchinson trace estimator [UNSPECIFIED: noise distribution,
        # Gaussian is standard; Rademacher can have lower variance]
        self.register_buffer('_hutch_noise', None)

    def forward(self, t: torch.Tensor, states: torch.Tensor) -> torch.Tensor:
        """Return dz/dt = g_φ(t, z)."""
        # states: (batch, dim+1) — last component is log p(z_t)
        z = states[:, :self.dim]
        t_expanded = t.expand(z.shape[0], 1)
        inp = torch.cat([z, t_expanded], dim=-1)

        # Check if we need trace estimation
        if states.shape[1] == self.dim + 1 and z.requires_grad:
            # Hutchinson trace estimator: v^T J_g v
            v = torch.randn_like(z)
            v = v / torch.sqrt(torch.sum(v ** 2, dim=-1, keepdim=True))
            dz = self.net(inp)

            # Jacobian-vector product for trace approximation
            vjp = torch.autograd.grad(
                dz, z, v, create_graph=True,
                retain_graph=True, only_inputs=True
            )[0]
            dlogp = -torch.sum(vjp * v, dim=-1, keepdim=True)  # Eq. (2.24)
            return torch.cat([dz, dlogp], dim=-1)
        else:
            # During sampling or evaluation without trace
            dz = self.net(inp)
            return dz

    def log_prob(self, x: torch.Tensor) -> torch.Tensor:
        """Compute log p(x) by integrating ODE backward.
           Eq. (2.25) — reverse integration from x to z_0."""
        # Augmented state: [z, log p(z)]
        logp0 = torch.zeros(x.shape[0], 1, device=x.device)
        augmented = torch.cat([x, logp0], dim=-1)

        # Solve ODE backward from t=1 to t=0
        t_span = torch.tensor([1.0, 0.0], device=x.device)
        # [UNSPECIFIED]: odeint tolerances (rtol=1e-5, atol=1e-5 are defaults)
        solution = odeint(
            self, augmented, t_span,
            rtol=1e-5, atol=1e-5, method='dopri5'
        )

        z0 = solution[-1, :, :self.dim]
        delta_logp = solution[-1, :, self.dim:]

        # Prior: standard Gaussian N(0, I)
        log_pz0 = -0.5 * (self.dim * math.log(2 * math.pi) + torch.sum(z0 ** 2, dim=-1, keepdim=True))

        return (log_pz0 + delta_logp).squeeze(-1)

    @torch.no_grad()
    def sample(self, num_samples: int, device: torch.device) -> torch.Tensor:
        """Sample from the CNF: z_0 ~ N(0,I) -> integrate to t=1."""
        z0 = torch.randn(num_samples, self.dim, device=device)
        t_span = torch.tensor([0.0, 1.0], device=device)
        # During sampling, we don't need trace — just integrate the ODE
        augmented = odeint(
            self, z0, t_span,
            rtol=1e-5, atol=1e-5, method='dopri5'
        )
        return augmented[-1]


import math  # for log_prob

代码关键注释：

Hutchinson 迹估计器：\(\text{tr}(J_g) \approx \mathbf{v}^T J_g \mathbf{v}\)，其中 \(\mathbf{v}\) 是随机向量（这里使用了标准化高斯向量）。这一步将 \(O(D)\) 的反向传播缩减为单次 Jacobian-vector 乘积。
增强状态：ODEsolver 同时积分 \(\mathbf{z}_t\) 和 \(\log p(\mathbf{z}_t)\)，如公式 (2.25) 所示。odeint 返回整个轨迹。
[UNSPECIFIED] 选择：Hutchinson 估计器的噪声分布（高斯 vs Rademacher）会影响方差——Rademacher 通常有更低的方差但在某些架构中可能导致数值问题。

4.3 Flow Matching：PyTorch 实现

"""
flow_matching.py — Minimal Flow Matching with OT Conditional Paths
Based on: Lipman et al., "Flow Matching for Generative Modeling", ICLR 2023.
Ref: Eq. (2.36) OT path, Eq. (2.38) flow ψ_t, Eq. (2.39) CFM-OT loss.
"""
import torch
import torch.nn as nn
from torchdiffeq import odeint


class FlowMatchingModel(nn.Module):
    """Vector field network v_θ(t, x) for Flow Matching.

       [UNSPECIFIED]: architecture — using a simple MLP; in practice
       images require a U-Net backbone similar to DDPM.
    """
    def __init__(self, dim: int, hidden_dim: int = 512, num_layers: int = 5):
        super().__init__()
        layers = []
        in_dim = dim + 1  # x concat t
        for i in range(num_layers):
            layers.append(nn.Linear(in_dim, hidden_dim))
            layers.append(nn.SiLU())
            in_dim = hidden_dim
        layers.append(nn.Linear(hidden_dim, dim))
        self.net = nn.Sequential(*layers)

    def forward(self, t: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
        """Predict vector field v_θ(t, x) ∈ R^d."""
        t_emb = t.unsqueeze(-1)
        inp = torch.cat([x, t_emb], dim=-1)
        return self.net(inp)


def ot_conditional_flow(x0: torch.Tensor, x1: torch.Tensor,
                        t: torch.Tensor, sigma_min: float = 0.001
                        ) -> tuple[torch.Tensor, torch.Tensor]:
    """OT conditional path: ψ_t(x_0) and target velocity.

       Ref: Eq. (2.38): ψ_t(x_0) = (1 - (1-σ_min)t) · x_0 + t · x_1
       Ref: Eq. (2.37): u_t = (x_1 - (1-σ_min)x) / (1 - (1-σ_min)t)

       Args:
           x0: samples from N(0, I), shape (batch, dim)
           x1: data samples, shape (batch, dim)
           t: time in [0, 1], shape (batch,)
           sigma_min: small constant for stability
       Returns:
           psi_t: (batch, dim) — point on OT path at time t
           ut:    (batch, dim) — target velocity (regression target)
    """
    # Eq. (2.36): σ_t = 1 - (1 - σ_min)t, μ_t = t·x_1
    sigma_t = 1.0 - (1.0 - sigma_min) * t
    mu_t = t * x1

    # ψ_t(x_0) = σ_t · x_0 + μ_t  (Eq. 2.38)
    psi_t = sigma_t * x0 + mu_t

    # Target velocity (Eq. 2.37)
    ut = (x1 - (1.0 - sigma_min) * psi_t) / (1.0 - (1.0 - sigma_min) * t + 1e-8)

    return psi_t, ut


def cfm_loss(model: FlowMatchingModel, x1: torch.Tensor,
             sigma_min: float = 0.001) -> torch.Tensor:
    """Conditional Flow Matching loss with OT paths.

       Ref: Eq. (2.39):
         L_CFM = E[‖v_θ(t, ψ_t(x_0)) - (x_1 - (1-σ_min)x_0)‖²]
    """
    batch_size, dim = x1.shape
    device = x1.device

    # Sample noise x0 ~ N(0, I)
    x0 = torch.randn(batch_size, dim, device=device)

    # Sample time t ~ U[0, 1]
    t = torch.rand(batch_size, device=device)

    # OT conditional flow
    psi_t, ut = ot_conditional_flow(x0, x1, t, sigma_min)

    # Predict vector field
    vt = model(t, psi_t)

    # MSE Loss (Eq. 2.39)
    return torch.mean((vt - ut) ** 2)


@torch.no_grad()
def fm_sample(model: FlowMatchingModel, dim: int, num_samples: int,
              device: torch.device, num_steps: int = 100) -> torch.Tensor:
    """Sample from Flow Matching model by solving ODE with Euler method.

       dφ_t(x_0)/dt = v_θ(t, φ_t(x_0)), φ_0(x_0) = x_0 ∼ N(0, I)
    """
    model.eval()
    x = torch.randn(num_samples, dim, device=device)  # z_0 ~ N(0, I)
    dt = 1.0 / num_steps

    for step in range(num_steps):
        t = torch.full((num_samples,), step * dt, device=device)
        vt = model(t, x)
        x = x + vt * dt  # Euler step

    return x


# ============================================================
# Training example usage
# ============================================================
def train_fm(model: FlowMatchingModel, dataloader, optimizer,
             num_epochs: int, device: torch.device):
    """Minimal FM training loop."""
    model.train()
    for epoch in range(num_epochs):
        for batch in dataloader:
            x1 = batch[0].to(device)  # Assuming (data, label) tuples
            x1 = x1.view(x1.shape[0], -1)  # Flatten to (batch, dim)

            loss = cfm_loss(model, x1)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        if epoch % 100 == 0:
            print(f"Epoch {epoch}: CFM Loss = {loss.item():.6f}")

代码关键注释：

OT 条件路径：ot_conditional_flow 直接实现了公式 (2.38) 和 (2.37)。\(\sigma_{\min}\) 设为一个小的正数（如 0.001）以使 \(p_1(\mathbf{x} \mid \mathbf{x}_1)\) 是紧密集中在 \(\mathbf{x}_1\) 周围的高斯分布，而非真正的 delta 函数。
CFM 损失：公式 (2.39) 直接编码为 MSE——将预测的向量场 \(v_\theta(t, \psi_t(\mathbf{x}_0))\) 与目标的 \(\mathbf{x}_1 - (1-\sigma_{\min})\mathbf{x}_0\) 比较。
采样：使用简单的 Euler 方法求解 ODE。对于高保真度样本，应该使用自适应步长求解器（如 dopri5）。num_steps 控制采样精度——步数越多越精确但也越慢，通常 50-100 步即可。
[UNSPECIFIED] 架构适配：对于图像数据，需要将 MLP 替换为 U-Net 架构（类似 DDPM 的网络），并处理 2D 空间结构。上述代码仅适用于低维数据演示。

五、核心 Takeaway 与延伸阅读

五个核心洞察

1. DDPM、Score Matching、Flow Matching 是同一图谱上的不同节点

这三者之间的统一关系可以用一张图来概括：DDPM 通过变分推断推导出的简化损失，等价于 DSM 中去噪得分匹配损失（令 \(\mathbf{s}_\theta = -\boldsymbol{\epsilon}_\theta / \sigma\)）。更进一步，扩散路径的 DSM 是 Flow Matching 框架的一个特例（当 \(p_t\) 被选择为扩散过程的边际分布时）。因此，这三个方向不是竞争关系，而是对同一个核心问题的三种语言描述——如何从噪声中恢复数据结构。

2. OT 路径优于扩散路径的根本原因

在扩散过程中，粒子的轨迹是弯曲的——它们先”overshoot”目标位置，然后”backtrack”回来（见 Lipman et al. 论文的 Figure 3）。这种冗余运动增加了训练和采样的计算开销。OT 路径将这种冗余运动消除为零：粒子严格沿直线以恒定速度移动。这不仅是一个数值效率问题——它使得回归目标（向量场）更简单、更平滑，因此神经网络可以更快地学习。

3. Jacobian 行列式计算是 NF 设计的核心约束

回顾本章所有 Normalizing Flow 架构，你会发现一个共同的线索：每一种架构设计都是在围绕如何高效计算 Jacobian 行列式（或其对数）来组织。自回归流利用三角 Jacobian；耦合层利用分块对角+下三角结构；残差流利用幂级数和 Hutchinson 估计器；CNF 则将行列式问题转化为迹估计问题。理解了这一点，你就理解了为什么这些架构看起来如此不同，却都属于同一个范式。

4. 连续时间流不需要架构约束

离散 NF 需要精心设计的架构来保证可逆性和高效的 Jacobian 行列式计算。而 CNF 的 ODE 公式（公式 2.23）几乎不对 \(g_\phi\) 施加架构约束——只需满足一致 Lipschitz 条件。这大大简化了模型设计，但代价是训练和评估时需要数值 ODE 求解。Flow Matching 进一步缓解了训练代价，使 CNF 真正可用于大规模生成建模。

5. “去噪”是生成建模的通用原理

DDPM 学习预测噪声，DSM 学习预测得分（噪声/标准差），FM 学习预测向量场（与 OT 路径的瞬时速度）。这三者的共同本质是：从被破坏的数据中重建干净数据的某种属性。这种”去噪即生成”的原理说明了一个深刻的观点：理解和反转破坏过程，等于学会了生成过程。

六、练习与思考

基础练习

6.1 DDPM 公式推导 从公式 (2.1) 出发，手动推导公式 (2.2)（直接采样 \(\mathbf{x}_t\) 的闭合形式）。提示：利用高斯随机变量之和的性质，对重参数化进行递推展开。

6.2 ELBO 分解验证 验证公式 (2.7) 中 ELBO 的分解。步骤提示：(a) 写出 \(\mathcal{L} = \mathbb{E}_q[\log \frac{q(\mathbf{x}_{1:T}\mid\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})}]\)；(b) 展开分子和分母的连乘；(c) 利用贝叶斯定理重组 \(t=1\) 项和 \(t>1\) 项；(d) 识别 KL 散度。

6.3 得分匹配等价性证明 证明去噪得分匹配 (DSM) 的最优解等于显式得分匹配 (ESM) 的最优解。即证明 \(\arg\min_\theta \mathbb{E}_{q_\sigma(\tilde{\mathbf{x}}, \mathbf{x})}[\|\mathbf{s}_\theta(\tilde{\mathbf{x}}) - \nabla \log q_\sigma(\tilde{\mathbf{x}} \mid \mathbf{x})\|^2] = \arg\min_\theta \mathbb{E}_{p(\mathbf{x})}[\|\mathbf{s}_\theta(\mathbf{x}) - \nabla \log p(\mathbf{x})\|^2]\)。提示：展开 DSM 损失的平方，并对 \(q_\sigma(\tilde{\mathbf{x}}, \mathbf{x})\) 求期望，利用条件高斯分布的得分函数的解析形式进行化简。

6.4 Planar Flow 的 Jacobian 行列式推导 手动推导 Planar Flow（公式 2.63）的 Jacobian 行列式公式 (2.65)。使用矩阵行列式引理：\(\det(\mathbf{I} + \mathbf{u}\mathbf{v}^T) = 1 + \mathbf{v}^T\mathbf{u}\)。

6.5 CFM 梯度等价性 在简化的 1 维情况下验证定理 2（CFM 和 FM 的梯度等价性）。取 \(q(x_1) = \mathcal{N}(5, 1)\)，\(p_t(x \mid x_1) = \mathcal{N}(x; t x_1, (1-t)^2)\)，手动计算 \(\mathcal{L}_{\text{FM}}\) 和 \(\mathcal{L}_{\text{CFM}}\) 的期望损失之差（应为常数）。[UNSPECIFIED：需要使用数值积分或蒙特卡洛采样来近似期望值]

编程练习

6.6 DDPM 训练与可视化 使用 4.1 节的代码，在 MNIST 数据集上训练一个 DDPM。任务：

(a) 可视化 \(t=0, 100, 300, 500, 700, 999\) 时刻的加噪图像，观察前向扩散过程
(b) 实现逆向过程的逐步可视化——从纯噪声开始，每 100 步保存一次中间图像
(c) 尝试不同的 \(\beta_t\) 调度（线性 vs 余弦），观察生成质量的变化

6.7 Flow Matching 2D 分布学习 使用 4.3 节的代码，在 2D 的 “双月”（two moons）或”棋盘”（checkerboard）分布上训练 FM-OT：

(a) 可视化训练后的向量场（用 quiver 绘制），确认粒子沿 OT 直线移动
(b) 对比 OT 路径和扩散路径（公式 2.17/2.19）的训练收敛速度
(c) 测量不同 ODE 求解步数（10, 20, 50, 100, 200）对样本质量的影响

6.8 CNF vs DDPM 对比实验 在 2D Swiss Roll 数据上分别实现 CNF（最大似然训练，通过 ODE 求解器）和 DDPM，对比两者的：

(a) 训练时间（wall clock time）
(b) 生成样本的密度估计质量（计算在测试点上的对数似然）
(c) 采样效率（NFE 需求）
讨论：哪种方法在什么场景下更优？

思考题

6.9 DDPM 的逆向方差选择 DDPM 将逆向过程的方差固定为 \(\sigma_t^2 = \beta_t\) 或 \(\tilde{\beta}_t\)，而非学习它。这是为什么？如果方差也是可学习的（如 Nichol & Dhariwal, 2021 的 improved DDPM），会产生什么改进？会引入什么问题？

6.10 Flow Matching 的极限 FM 可以在 \(t=1\) 时精确达到 \(\mathbf{x}_1\) 吗？考虑：当 \(\sigma_{\min} \to 0\) 时，\(p_1(\mathbf{x} \mid \mathbf{x}_1) \to \delta(\mathbf{x} - \mathbf{x}_1)\)，但公式 (2.37) 中向量场 \(u_t(\mathbf{x} \mid \mathbf{x}_1)\) 在 \(t \to 1\) 时是否存在奇异性？试分析 \(t \to 1\) 的极限行为，并讨论它与扩散模型中 \(t \to 1\) 时 \(p_T\) 仅近似等于噪声的区别。

6.11 统一视角 尝试用一个统一的数学框架来描述 DDPM、Score-based 模型和 Flow Matching。提示：考虑一个一般的连续性方程：

\[\frac{\partial p_t}{\partial t} + \nabla \cdot (p_t \mathbf{v}_t) = 0\]

对于 DDPM/Score SDE，\(\mathbf{v}_t\) 是什么？对于 Flow Matching OT，\(\mathbf{v}_t\) 是什么？两者在坐标系的选择上有何本质区别？

章末注：本章的代码实现是教学性的最小可行版本。生产级别的扩散模型（如 Stable Diffusion）需要更多工程细节——包括 latent space diffusion、classifier-free guidance、级联式超分辨率管道等——这些内容将在第12章（扩散策略）和第17章（VLA大模型）中继续深入讨论。但本章建立的数学基础（ELBO 推导、得分匹配等价性、CFM 梯度定理）将是后续所有章节的共同数学语言。

← 第03章机器人机构学：运动学与动力学

↑ 目录

第05章人手与灵巧操作 →

第4章 扩散模型入门

一、本章概要

二、核心概念与数学基础

2.1 Denoising Diffusion Probabilistic Models (DDPM)

2.1.1 前向过程（Forward Process）

2.1.2 逆向过程（Reverse Process）

2.1.3 以 \(\mathbf{x}_0\) 为条件的逆向分布

2.1.4 变分下界（ELBO）推导

2.1.5 简化损失函数（Simplified Training Objective）

2.1.6 训练与采样算法

2.2 Score-based 生成模型

2.2.1 得分函数与 Stein Score

2.2.2 Langevin 动力学采样

2.2.3 得分匹配（Score Matching）

2.2.4 去噪得分匹配（Denoising Score Matching, DSM）

2.2.5 多噪声尺度与退火 Langevin 动力学

2.2.6 Score SDE 的统一视角（简述）

2.3 Normalizing Flows：变量替换与可逆神经网络

2.3.1 变量替换公式（Change of Variables）

2.3.2 自回归流（Autoregressive Flows）

2.3.3 残差流（Residual Flows）

2.3.4 连续时间流（Continuous Normalizing Flows, CNF）

2.4 Flow Matching：统一扩散模型与 CNF

2.4.1 问题设置

2.4.2 Flow Matching 目标函数

2.4.3 条件概率路径与条件向量场

2.4.4 高斯条件概率路径

2.4.5 Optimal Transport 条件路径

三、论文精读

3.1 Papamakarios et al. (2021)：Normalizing Flows 的统一视角

3.1.1 定义与基础（Section 2.1）

3.1.2 表达力（Section 2.2）

3.1.3 KL 散度训练（Section 2.3）

3.1.4 有限流架构全景（Section 3）

3.1.5 CNF 的连续时间表述（Section 4）

3.1.6 关键评价

3.2 Lipman et al. (2023)：Flow Matching 的突破

3.2.1 核心贡献识别

3.2.2 与去噪得分匹配的类比

3.2.3 OT 路径的实验验证

3.2.4 局限性

四、算法与代码实现

4.1 DDPM：完整 PyTorch 实现

4.2 Continuous Normalizing Flow：PyTorch 实现

4.3 Flow Matching：PyTorch 实现

五、核心 Takeaway 与延伸阅读

五个核心洞察

延伸阅读

六、练习与思考

基础练习

编程练习

思考题

第4章扩散模型入门