文章

航海者Bernkastel的碎片之旅

航海者Bernkastel的碎片之旅

连续性方程

概念

连续性方程描述的是概率密度在空间和时间中的守恒关系。

设有一个随机过程 \(\mathbf{X}(t) \in \mathbb{R}^n\),其概率密度函数为 \(\rho(\mathbf{x}, t)\)

满足归一化条件:

\[\int_{\mathbb{R}^n} \rho(\mathbf{x}, t) \, d^n x = 1, \quad \forall t\]

假设该过程存在概率流密度矢量场(probability current density):

\[\mathbf{J}(\mathbf{x}, t) \in \mathbb{R}^n\]

它表示单位时间内通过单位面积的概率流量。

两种视角:

  • 拉格朗日视角(Lagrangian Perspective)跟踪粒子本身的运动轨迹和状态,而
  • 欧拉视角(Eulerian Perspective)则关注空间中固定点上的流体状态变化。

欧拉视角:局部概率守恒

考虑空间中任意体积元 \(V \subset \mathbb{R}^n\),其边界为 \(\partial V\)。那么:

  • 体积内总概率的变化率 = 流入该体积的净概率流
\[\frac{d}{dt} \int_V \rho(\mathbf{x}, t) \, d^n x = - \oint_{\partial V} \mathbf{J}(\mathbf{x}, t) \cdot d\mathbf{S}\]

说明:右侧的负号是因为 \(d\mathbf{S}\) 通常定义为指向外的法向量。若流矢量 \(\mathbf{J}\) 指向外(\(\mathbf{J} \cdot d\mathbf{S} > 0\)),则概率流出,导致体积 \(V\) 内的概率减少。

由于体积 \(V\) 是在空间中固定的(欧拉视角),不随时间变化,所以时间导数可以和空间积分交换顺序,即:\(\frac{d}{dt} \int_V \rho(\mathbf{x}, t) \, d^n x = \int_V \frac{\partial \rho(\mathbf{x}, t)}{\partial t} \, d^n x\)

根据散度定理(Gauss 定理) (\(\nabla \cdot\) 是 n 维空间中的散度算子):\(\oint_{\partial V} \mathbf{J} \cdot d\mathbf{S} = \int_V \nabla \cdot \mathbf{J} \, d^n x\)

所以:

\[\int_V \frac{\partial \rho}{\partial t} \, d^n x = - \int_V \nabla \cdot \mathbf{J} \, d^n x\]

由于体积 \(V\) 是任意选取的,被积函数必须处处相等:

\[\boxed{ \frac{\partial \rho}{\partial t} + \nabla \cdot \mathbf{J} = 0 }\]

流密度 \(\mathbf{J}\) 的具体形式

\(\mathbf{J}\) 的具体形式依赖于具体的动力学模型。

漂移-扩散过程(Fokker-Planck 方程)

若粒子服从朗之万方程:

\[d\mathbf{X}_t = \boldsymbol{\mu}(\mathbf{X}_t, t)\, dt + \boldsymbol{\sigma}(\mathbf{X}_t, t)\, d\mathbf{W}_t\]

\(d\mathbf{W}_t\) 是维纳过程(布朗运动)的无穷小增量,它有以下关键特性:

  • 它的期望(均值)是 0。
  • 它的方差是 \(dt\)。这意味着它的标准差,也就是它的大小尺度,是 \(\sqrt{dt}\)。
\[\frac{d\mathbf{X}_t}{dt} = \boldsymbol{\mu}(\mathbf{X}_t, t) + \boldsymbol{\sigma}(\mathbf{X}_t, t) \frac{d\mathbf{W}_t}{dt}\]

\(\frac{d\mathbf{W}_t}{dt}\) 这一项的尺度是 \(\frac{\sqrt{dt}}{dt} = \frac{1}{\sqrt{dt}}\)。

当 \(dt \to 0\) 时,\(\frac{1}{\sqrt{dt}} \to \infty\)!

这意味着:

  1. 单个粒子的瞬时速度是无穷大的、没有良好定义的
  2. 粒子的运动轨迹是连续的,但处处不可微。它的路径充满了无限微小的剧烈抖动。

所以,对于单个随机粒子,我们无法像在经典力学中那样谈论一个光滑的、良定义的瞬时速度。

上面讨论的 \(\frac{d\mathbf{X}_t}{dt}\)。它在数学上是病态的,无法用于描述一个平滑的“流场”。

宏观速度:概率流的速度 \(\mathbf{v} = \mathbf{J}/\rho\) 不是某一个粒子的速度,而是在空间中的一个点,概率的整体是如何流动的?

对应的 Fokker-Planck 方程为:

\[\frac{\partial \rho}{\partial t} = -\nabla \cdot (\boldsymbol{\mu} \rho) + \nabla \cdot (D \nabla \rho)\]

其中 \(D = \frac{1}{2} \boldsymbol{\sigma} \boldsymbol{\sigma}^\top\) 是扩散张量。

将其写成连续性方程形式:

\[\frac{\partial \rho}{\partial t} + \nabla \cdot \mathbf{J} = 0 \quad \text{其中} \quad \boxed{ \mathbf{J} = \boldsymbol{\mu} \rho - D \nabla \rho }\]

其中

  • 漂移流 (Drift Current): \(\mathbf{J}_{\text{drift}} = \boldsymbol{\mu} \rho\)。这一项描述了概率密度由于受到确定性力场(由漂移项 \(\boldsymbol{\mu}\) 描述)的驱动而产生的整体平移。概率“粒子”倾向于顺着场 \(\boldsymbol{\mu}\) 的方向移动。
  • 扩散流 (Diffusion Current): \(\mathbf{J}_{\text{diffusion}} = - D \nabla \rho\)。这一项是芬克第一定律 (Fick’s First Law) 的概率版本。它描述了概率从高密度区域流向低密度区域的趋势,其方向与概率密度梯度 \(\nabla \rho\) 相反。\(D\) 是扩散系数,衡量扩散的快慢。

确定性流(无扩散)

若系统完全由速度场 \(\mathbf{v}(\mathbf{x}, t)\) 驱动:

\[\frac{d\mathbf{x}}{dt} = \mathbf{v}(\mathbf{x}, t)\]

则概率流为:

\[\boxed{ \mathbf{J} = \mathbf{v} \rho }\]

连续性方程变为:

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

现代生成模型

核心思想:生成模型学习一个从简单先验分布(如高斯分布 \(p_0(\mathbf{z})\))到复杂数据分布 \(p_{data}(\mathbf{x})\) 的映射。我们可以将这个演化过程看作是概率“流体”从一个形状(高斯分布)变换到另一个形状(数据分布)的过程。

扩散模型 (Diffusion Models)

扩散模型通过两个过程来工作:一个前向过程将数据变成噪声,一个反向过程将噪声变回数据。

A. 前向过程:数据 → 噪声

对应于漂移-扩散过程。人为设计一个随机过程(SDE),逐渐地将任意数据点 \(\mathbf{X}_0 \sim p_{data}\) 转化为一个标准高斯噪声点 \(\mathbf{X}_T \sim \mathcal{N}(0, I)\)。一个常见的 SDE 形式是: \(d\mathbf{X}_t = -\frac{1}{2} \beta(t) \mathbf{X}_t dt + \sqrt{\beta(t)} d\mathbf{W}_t\)

  • \(\beta(t)\): 一个预先设定的噪声调度表。

  • 漂移项: \(\boldsymbol{\mu}(\mathbf{X}_t, t) = -\frac{1}{2} \beta(t) \mathbf{X}_t\)。这个力的方向与位移相反,因此是向心力,指向原点。
  • 扩散项: \(\boldsymbol{\sigma}(\mathbf{X}_t, t) = \sqrt{\beta(t)} \mathbf{I}\)。持续地注入高斯噪声。

B. 反向过程:噪声 → 数据

扩散式生成模型的核心是逆转上述过程。

一个数学结论:任何一个漂移-扩散过程,其时间反向过程也是一个漂移-扩散过程。反向过程的 SDE 为:

\[d\mathbf{X}_t = \left[ -\frac{1}{2} \beta(t) \mathbf{X}_t - \beta(t) \nabla_{\mathbf{x}} \log \rho(\mathbf{x}, t) \right] dt + \sqrt{\beta(t)} d\mathbf{W}_t\]

这里的 \(\nabla_{\mathbf{x}} \log \rho(\mathbf{x}, t)\) 被称为分数 (Score)

  • 模型目标:我们无法知道真实的 \(\rho(\mathbf{x}, t)\),因此也无法知道分数。扩散模型就是训练一个深度神经网络 \(\mathbf{s}_\theta(\mathbf{x}, t)\) 来近似这个分数函数: \(\mathbf{s}_\theta(\mathbf{x}, t) \approx \nabla_{\mathbf{x}} \log \rho(\mathbf{x}, t)\)。
  • 生成过程:一旦网络训练好了,从 \(\mathbf{z} \sim \mathcal{N}(0, I)\) 出发,使用学到的分数来模拟反向 SDE,从时间 T 积分到 0,最终得到的 \(\mathbf{X}_0\) 就是一个生成的样本。

流匹配 (Flow Matching) 与确定性流

https://github.com/helblazer811/Diffusion-Explorer

Diffusion-Explorer

tutorial video

tutorial video 2

Flow Matching 假设概率密度的演化是由一个确定性的速度场 \(\mathbf{v}(\mathbf{x}, t)\) 驱动的。

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

这个方程描述粒子的轨迹都遵循一个常微分方程 (ODE):

\[\frac{d\mathbf{x}}{dt} = \mathbf{v}(\mathbf{x}, t)\]

A. 核心思想 我们想学习一个速度场 \(\mathbf{v}_\theta(\mathbf{x}, t)\),使得如果我们从一个简单的先验分布 \(p_0(\mathbf{z})\) 中采样一个点 \(\mathbf{z}\),然后根据 \(\frac{d\mathbf{x}}{dt} = \mathbf{v}_\theta(\mathbf{x}, t)\) 从 t=0 积分到 t=1,最终得到的点 \(\mathbf{x}(1)\) 的分布恰好是数据分布 \(p_{data}\)。

B. “Matching” 的含义 直接学习这个理想的 \(\mathbf{v}\) 非常困难。Flow Matching 的贡献在于,它证明了我们不需要学习这个复杂的marginal速度场。我们可以定义一个更简单的、条件的(conditional)速度场。

例如,我们可以定义一个从噪声 \(\mathbf{x}_0\) 到数据 \(\mathbf{x}_1\) 的直线路径: \(\mathbf{x}(t) = (1-t) \mathbf{x}_0 + t \mathbf{x}_1\)

这个路径对应的瞬时速度是恒定的: \(\mathbf{u}(\mathbf{x}(t) | \mathbf{x}_0, \mathbf{x}_1) = \mathbf{x}_1 - \mathbf{x}_0\)

Flow Matching 证明了,我们只需要训练神经网络 \(\mathbf{v}_\theta(\mathbf{x}, t)\) 去匹配这个简单的条件速度场 \(\mathbf{u}\) 的期望即可。训练目标大致是:

\[\min_\theta \mathbb{E}_{t, \mathbf{x}_0, \mathbf{x}_1} \left\| \mathbf{v}_\theta((1-t) \mathbf{x}_0 + t \mathbf{x}_1, t) - (\mathbf{x}_1 - \mathbf{x}_0) \right\|^2\]

附 torch.autograd.grad (vjp: vector-Jacobian product)

torch.autograd.grad 是在计算 vector-Jacobian product(向量-雅可比积)

假设:

  • 输入:\(x ∈ R^N\)
  • 输出:\(y = f(x) ∈ R^M\)
  • 雅可比矩阵:\(J = ∂y/∂x ∈ R^{M×N}\),其中 \(J[i, j] = ∂y_i / ∂x_j\)
1
2
3
4
5
grads = torch.autograd.grad(
    outputs=y,           # shape (M,)
    inputs=x,            # shape (N,)
    grad_outputs=v       # shape (M,)
)[0]                     # 返回 shape (N,)

计算的是:

\[grads = v^T ⋅ J ∈ R^N\]

这正是 vector-Jacobian product —— 用向量 v 左乘雅可比矩阵 J

链式法则视角

在神经网络反向模式中通常计算:

dL/dx = (dL/dy) ⋅ (∂y/∂x) = v^T ⋅ J

其中:

  • dL/dy 是上游梯度(标量损失对输出的梯度)grad_outputs
  • ∂y/∂x 是当前函数的雅可比
  • dL/dx 是最终传播到输入的梯度

附 Time Reversal of Diffusions

反向过程 SDE

给定前向过程的 SDE:

\[d\mathbf{X}_t = \boldsymbol{\mu}(\mathbf{X}_t, t) dt + \boldsymbol{\sigma}(t) d\mathbf{W}_t\]

其对应的反向过程(从时间 T 到 0)由以下 SDE 描述:

\[\boxed{ d\mathbf{X}_t = \left[ -\boldsymbol{\mu}(\mathbf{X}_t, t) + \boldsymbol{\sigma}(t)^2 \nabla_{\mathbf{X}_t} \log p_t(\mathbf{X}_t) \right] dt + \boldsymbol{\sigma}(t) d\mathbf{\bar{W}}_t }\]

其中 \(d\mathbf{\bar{W}}_t\) 是一个反向的维纳过程。

前向过程 \(d\mathbf{X}_t = -\frac{1}{2} \beta(t) \mathbf{X}_t dt + \sqrt{\beta(t)} d\mathbf{W}_t\) 具体代入,可得反向过程:

\[d\mathbf{X}_t = \left[ \frac{1}{2} \beta(t) \mathbf{X}_t + \beta(t) \nabla_{\mathbf{X}_t} \log p_t(\mathbf{X}_t) \right] dt + \sqrt{\beta(t)} d\mathbf{\bar{W}}_t\]

方程的各个部分:

  1. 反向漂移项: \(\left[ \frac{1}{2} \beta(t) \mathbf{X}_t + \beta(t) \nabla_{\mathbf{X}_t} \log p_t(\mathbf{X}_t) \right]\)
    • \(\frac{1}{2} \beta(t) \mathbf{X}_t\): 方向与位移相同, 离心力将噪声点从原点推开。
    • \(\beta(t) \nabla_{\mathbf{X}_t} \log p_t(\mathbf{X}_t)\): 最关键的修正项
      • \(p_t(\mathbf{X}_t)\) 是在时间 \(t\) 时,所有数据点经过前向过程后形成的边缘概率密度。
      • \(\nabla_{\mathbf{X}_t} \log p_t(\mathbf{X}_t)\) 被称为分数函数 (Score Function)。该向量场指向 \(log p_t(\mathbf{X}_t)\) 增长最快的方向。
  2. 反向扩散项: \(\sqrt{\beta(t)} d\mathbf{\bar{W}}_t\)
    • 扩散项的幅度与前向过程相同,保证了过程的随机性量级是一致的。

如果没有分数修正项,反向过程就只是简单地将噪声从原点推开,无法恢复出任何有意义的结构。分数项 \(\nabla_{\mathbf{X}_t} \log p_t(\mathbf{X}_t)\) 包含了在每个时刻 \(t\) 恢复数据结构所需的所有信息。由于我们无法知道真实的 \(p_t(\mathbf{X}_t)\),扩散模型的训练目标就是用一个神经网络 \(\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{X}_t, t)\) 来近似这个分数函数

证明(基于离散时间的直观推导)

一个完全严格的证明需要用到随机过程理论中的 Fokker-Planck 方程或 Kolmogorov backward/forward 方程,过程比较复杂。这里我们提供一个在机器学习领域更常见、更直观的离散时间推导思路(源于 Anderson, 1982 的工作)。

考虑一个很小的时间步长 \(\Delta t\)。

前向转移概率

从 \(\mathbf{X}_t\) 到 \(\mathbf{X}_{t+\Delta t}\) 的转移可以写成:

\[\mathbf{X}_{t+\Delta t} \approx \mathbf{X}_t + \boldsymbol{\mu}(\mathbf{X}_t, t) \Delta t + \boldsymbol{\sigma}(t) \sqrt{\Delta t} \mathbf{z}\]

其中 \(\mathbf{z} \sim \mathcal{N}(0, \mathbf{I})\)。

这说明,给定 \(\mathbf{X}_t\),\(\mathbf{X}_{t+\Delta t}\) 的条件概率是一个高斯分布:

\[p(\mathbf{X}_{t+\Delta t} | \mathbf{X}_t) = \mathcal{N}(\mathbf{X}_{t+\Delta t}; \mathbf{X}_t + \boldsymbol{\mu}(\mathbf{X}_t, t)\Delta t, \boldsymbol{\sigma}(t)^2 \Delta t)\]

反向转移概率

我们想要求的是反向的转移概率 \(p(\mathbf{X}_t | \mathbf{X}_{t+\Delta t})\) 。根据贝叶斯定理:

\[p(\mathbf{X}_t | \mathbf{X}_{t+\Delta t}) = \frac{p(\mathbf{X}_{t+\Delta t} | \mathbf{X}_t) p_t(\mathbf{X}_t)}{p_{t+\Delta t}(\mathbf{X}_{t+\Delta t})}\]

取对数:

\[\log p(\mathbf{X}_t | \mathbf{X}_{t+\Delta t}) = \log p(\mathbf{X}_{t+\Delta t} | \mathbf{X}_t) + \log p_t(\mathbf{X}_t) - \log p_{t+\Delta t}(\mathbf{X}_{t+\Delta t})\]
  • 第一项: \(\log p(\mathbf{X}_{t+\Delta t} | \mathbf{X}_t) = C_1 - \frac{1}{2 \boldsymbol{\sigma}(t)^2 \Delta t} ||\mathbf{X}_{t+\Delta t} - (\mathbf{X}_t + \boldsymbol{\mu}(\mathbf{X}_t, t)\Delta t)||^2\)

  • 后两项: 我们对 \(\log p_{t+\Delta t}(\mathbf{X}_{t+\Delta t})\) 在 \((\mathbf{X}_t, t)\) 处一阶泰勒展开(忽略了高阶无穷小): \(\begin{aligned} \log p_{t+\Delta t}(\mathbf{X}_{t+\Delta t}) &\approx \log p_t(\mathbf{X}_t) + \frac{\partial \log p_t}{\partial t} \Delta t + (\mathbf{X}_{t+\Delta t} - \mathbf{X}_t)^T \nabla_{\mathbf{X}_t} \log p_t(\mathbf{X}_t) \\ \implies \log p_t(\mathbf{X}_t) - \log p_{t+\Delta t}(\mathbf{X}_{t+\Delta t}) &\approx - \frac{\partial \log p_t}{\partial t} \Delta t - (\mathbf{X}_{t+\Delta t} - \mathbf{X}_t)^T \nabla_{\mathbf{X}_t} \log p_t(\mathbf{X}_t) \end{aligned}\) 将展开式代入 \(\log p(\mathbf{X}_t | \mathbf{X}_{t+\Delta t})\)

我们只关心与 \(\mathbf{X}_t\) 相关的项,因为我们想得到一个关于 \(\mathbf{X}_t\) 的高斯分布。

\[\log p(\mathbf{X}_t | \mathbf{X}_{t+\Delta t}) \approx C - \frac{1}{2 \sigma^2 \Delta t} ||\mathbf{X}_{t+\Delta t} - \mathbf{X}_t - \boldsymbol{\mu} \Delta t||^2 - (\mathbf{X}_{t+\Delta t} - \mathbf{X}_t)^T \nabla \log p_t\]

展开二次项,保留与 \(\mathbf{X}_t\) 相关的主导项:

\[\approx C' - \frac{1}{2 \sigma^2 \Delta t} (\mathbf{X}_t^T\mathbf{X}_t - 2 \mathbf{X}_t^T(\mathbf{X}_{t+\Delta t} - \boldsymbol{\mu}\Delta t)) + \mathbf{X}_t^T \nabla \log p_t\]

我们对方程进行配方,使其形式为 \(-\frac{1}{2\text{Var}} ||\mathbf{X}_t - \text{Mean}||^2\)。 通过整理关于 \(\mathbf{X}_t\) 的一次项和二次项,可以发现这个分布是一个均值为 \(\mathbf{X}_{t+\Delta t} - \boldsymbol{\mu}_{rev} \Delta t\),方差为 \(\sigma^2 \Delta t\) 的高斯分布。

经过一些代数运算,可以推导出反向漂移项 \(\boldsymbol{\mu}_{rev}\) 满足:

\[\boldsymbol{\mu}_{rev}(\mathbf{X}_t, t) \approx \boldsymbol{\mu}(\mathbf{X}_t, t) - \boldsymbol{\sigma}(t)^2 \nabla_{\mathbf{X}_t} \log p_t(\mathbf{X}_t)\]

由于这个推导是针对从 \(t+\Delta t\) 到 \(t\) 的过程,时间在减少。如果我们定义反向 SDE 的时间是向前流逝的(即 \(dt\) 为正),那么漂移项需要取反。因此,反向过程的漂移项为:

\[\boldsymbol{\mu}_{reverse} = - \boldsymbol{\mu}_{rev} = -\boldsymbol{\mu}(\mathbf{X}_t, t) + \boldsymbol{\sigma}(t)^2 \nabla_{\mathbf{X}_t} \log p_t(\mathbf{X}_t)\]

warning

本文由作者按照 CC BY 4.0 进行授权