时间反演与概率流

前置

• 概率演化方程

设定与正则条件

考虑 $[0,T]$ 上的前向 Itô 扩散

$$\mathrm{d}{\mathbf{X}}_t=\mathbf{f}({\mathbf{X}}_t,t)\mathrm{d}t+\mathbf{G}({\mathbf{X}}_t,t)\mathrm{d}{\mathbf{W}}_t,$$

并记

$$\mathbf{a}(\mathbf{x},t):=\mathbf{G}(\mathbf{x},t)\mathbf{G}(\mathbf{x},t{)}^{\top }.$$

以下公式采用

$$(\nabla \cdot \mathbf{a}{)}_i:=\sum _j{\partial }_{x_j}{a}_{ij}$$

这一行散度约定。设过程非爆炸，$\mathbf{f}$、$\mathbf{a}$ 关于空间足够光滑，$\mathbf{a}$ 非退化或满足适用的扩散时间反演条件；并设每个 $t\in (0,T]$ 的边缘分布有严格正、足够光滑的密度 $p_t(\mathbf{x})$。还需密度和通量在无穷远充分衰减，或给定相容的边界条件，以保证 Fokker—Planck 与连续性方程解的唯一性。

前向边缘密度满足

$${\partial }_t{p}_t=-\nabla \cdot (\mathbf{f}{p}_t)+\frac{1}{2}\nabla \nabla :(\mathbf{a}{p}_t).$$

若密度在某处为零，$\nabla \log p_t$ 未定义；相应公式只能在 $p_t>0$ 的区域使用，或改写成不除以 $p_t$ 的通量形式。

一般扩散矩阵的反向 SDE

定义反向时钟过程

$${\mathbf{Y}}_s:={\mathbf{X}}_{T-s},0\le s\le T.$$

在上述正则条件下，${\mathbf{Y}}_s$ 满足反向过滤下的 SDE

$$\boxed{\begin{array}{rl}\mathrm{d}{\mathbf{Y}}_s=& [-\mathbf{f}({\mathbf{Y}}_s,T-s)+(\nabla \cdot \mathbf{a})({\mathbf{Y}}_s,T-s)\\ & +\mathbf{a}({\mathbf{Y}}_s,T-s)\nabla \log p_{T-s}({\mathbf{Y}}_s)]\mathrm{d}s\\ & +\mathbf{G}({\mathbf{Y}}_s,T-s)\mathrm{d}{\overline{\mathbf{W}}}_s.\end{array}}$$

扩散项只需具有协方差 $\mathbf{a}$，所以也可换成 $\mathbf{a}$ 的其他可测平方根。

等价地，若仍使用原时间变量 $t$，但从 $T$ 向 $0$ 积分（所以 $\mathrm{d}t<0$），则常写成

$$\boxed{\mathrm{d}{\mathbf{X}}_t=\left[\mathbf{f}-\nabla \cdot \mathbf{a}-\mathbf{a}\nabla \log p_t\right]({\mathbf{X}}_t,t)\mathrm{d}t+\mathbf{G}({\mathbf{X}}_t,t)\mathrm{d}{\overline{\mathbf{W}}}_t.}$$

这里的 $\overline{\mathbf{W}}$ 是相对于反向过滤的布朗运动，其增量方差对应 $|\mathrm{d}t|$；不能把这个式子误当作以 $t$ 递增求解的普通前向 SDE。

反向漂移可由密度方程核对。令

$$\tilde{\mathbf{f}}=-\mathbf{f}+\nabla \cdot \mathbf{a}+\mathbf{a}\nabla \log p.$$

由于

$$\tilde{\mathbf{f}}p=-\mathbf{f}p+(\nabla \cdot \mathbf{a})p+\mathbf{a}\nabla p=-\mathbf{f}p+\nabla \cdot (\mathbf{a}p),$$

反向过程在 $s$ 时钟下的 Fokker—Planck 右端为

$$\begin{array}{rl}-\nabla \cdot (\tilde{\mathbf{f}}p)+\frac{1}{2}\nabla \nabla :(\mathbf{a}p)& =\nabla \cdot (\mathbf{f}p)-\frac{1}{2}\nabla \nabla :(\mathbf{a}p)\\ & =-{\partial }_{t}p.\end{array}$$

这与 ${\partial }_s{p}_{T-s}=-{\partial }_t{p}_t$ 一致。严格的路径时间反演结论还依赖前述正则性；仅靠形式 PDE 计算不足以自行保证反向 SDE 的存在唯一性。

要真正复原前向过程的边缘路径，必须以

$${\mathbf{Y}}_0\sim p_T$$

作为终端分布，并沿 $s:0\to T$（等价于 $t:T\to 0$）求解。若终端初始化不是精确的 $p_T$，一般不会得到 $p_{T-s}$。

概率流 ODE

定义速度场

$$\boxed{\mathbf{v}(\mathbf{x},t)=\mathbf{f}(\mathbf{x},t)-\frac{1}{2}(\nabla \cdot \mathbf{a})(\mathbf{x},t)-\frac{1}{2}\mathbf{a}(\mathbf{x},t)\nabla \log p_t(\mathbf{x}).}$$

相应的概率流 ODE 为

$$\frac{\mathrm{d}{\mathbf{Z}}_t}{\mathrm{d}t}=\mathbf{v}({\mathbf{Z}}_t,t).$$

定理。 若 ODE 的流存在且唯一，连续性方程也在所给初边值类中唯一，则从 ${\mathbf{Z}}_0\sim p_0$ 出发时，

$$\mathcal{L}({\mathbf{Z}}_t)=p_t$$

对每个 $t$ 成立；也可从 ${\mathbf{Z}}_T\sim p_T$ 出发把同一 ODE 从 $T$ 反向积分到 $0$。

证明。 ODE 推动的密度 $q_t$ 满足连续性方程

$${\partial }_t{q}_t=-\nabla \cdot (\mathbf{v}{q}_t).$$

把 $q_t=p_t$ 和 $\mathbf{v}$ 代入：

$$\begin{array}{rl}-\nabla \cdot (\mathbf{v}p)=& -\nabla \cdot (\mathbf{f}p)\\ & +\frac{1}{2}\nabla \cdot \left[(\nabla \cdot \mathbf{a})p+\mathbf{a}\nabla \log pp\right]\\ =& -\nabla \cdot (\mathbf{f}p)+\frac{1}{2}\nabla \cdot \left[(\nabla \cdot \mathbf{a})p+\mathbf{a}\nabla p\right].\end{array}$$

逐指标有

$$\nabla \cdot \left[(\nabla \cdot \mathbf{a})p+\mathbf{a}\nabla p\right]=\nabla \nabla :(\mathbf{a}p),$$

故

$${\partial }_{t}p=-\nabla \cdot (\mathbf{f}p)+\frac{1}{2}\nabla \nabla :(\mathbf{a}p),$$

恰好是前向 SDE 的 Fokker—Planck 方程。由相同初值和密度方程的唯一性，ODE 与 SDE 在每个时刻都有同一边缘密度。$\square$

这里共享的只是单时刻边缘分布族 $(p_t{)}_{t\in [0,T]}$。非退化 SDE 的路径有噪声，而 ODE 路径由初值确定；二者通常不共享转移核、两时刻联合分布、有限维分布或路径测度。

SDE—ODE 凸组合下的边缘一致性

设前向 SDE 的漂移、扩散分别为 $\mathbf{f},\mathbf{G}$，并设 ODE 速度场 $\mathbf{v}$ 与它共享密度 $p_t$。对 $\lambda \in [0,1]$ 定义

$${\mathbf{f}}_{\lambda }:=\lambda \mathbf{f}+(1-\lambda )\mathbf{v},{\mathbf{G}}_{\lambda }:=\sqrt{\lambda }\mathbf{G},$$

以及新 SDE

$$\mathrm{d}{\mathbf{X}}_t={\mathbf{f}}_{\lambda }({\mathbf{X}}_t,t)\mathrm{d}t+{\mathbf{G}}_{\lambda }({\mathbf{X}}_t,t)\mathrm{d}{\mathbf{W}}_t.$$

定理。 在相同初始密度和前向方程唯一的条件下，每个 $\lambda \in [0,1]$ 构造出的 SDE 都具有同一组单时刻边缘密度 $p_t$。

证明。 令 $\mathbf{a}=\mathbf{G}{\mathbf{G}}^{\top }$，则

$${\mathbf{a}}_{\lambda }={\mathbf{G}}_{\lambda }{\mathbf{G}}_{\lambda }^{\top }=\lambda \mathbf{a}.$$

把 ${\mathbf{f}}_{\lambda },{\mathbf{a}}_{\lambda }$ 代入 Fokker—Planck 算子：

$$\begin{array}{rl}& -\nabla \cdot ({\mathbf{f}}_{\lambda }p)+\frac{1}{2}\nabla \nabla :({\mathbf{a}}_{\lambda }p)\\ =& -\nabla \cdot ([\lambda \mathbf{f}+(1-\lambda )\mathbf{v}]p)+\frac{\lambda }{2}\nabla \nabla :(\mathbf{a}p)\\ =& \lambda \left[-\nabla \cdot (\mathbf{f}p)+\frac{1}{2}\nabla \nabla :(\mathbf{a}p)\right]+(1-\lambda )\left[-\nabla \cdot (\mathbf{v}p)\right]\\ =& \lambda {\partial }_{t}p+(1-\lambda ){\partial }_{t}p={\partial }_{t}p.\end{array}$$

所以 $p_t$ 解新 SDE 的前向方程；由唯一性即得结论。$\square$

$\lambda =1$ 是原 SDE，$\lambda =0$ 是确定性概率流 ODE，中间值给出噪声强度连续变化的一族过程。这个插值仍只保持单时刻边缘，不保持路径规律。

反向扩散中的 Langevin 纠偏

考虑各向同性、状态无关扩散

$$\mathrm{d}{\mathbf{X}}_t=f(t){\mathbf{X}}_t\mathrm{d}t+g(t)\mathrm{d}{\mathbf{W}}_t.$$

此时 $\mathbf{a}=g(t{)}^{2}I$，空间散度 $\nabla \cdot \mathbf{a}=0$。从 $T$ 向 $0$ 写出的反向 SDE 是

$$\mathrm{d}{\mathbf{X}}_t=\left[f(t){\mathbf{X}}_t-g(t{)}^2\nabla \log p_t({\mathbf{X}}_t)\right]\mathrm{d}t+g(t)\mathrm{d}{\overline{\mathbf{W}}}_t.$$

把它分解为概率流 ODE 项和随机纠偏项：

$$\begin{array}{rl}\mathrm{d}{\mathbf{X}}_t=& \underset{\text{概率流 ODE 项}}{\underbrace{\left[f(t){\mathbf{X}}_t-\frac{1}{2}g(t{)}^2\nabla \log p_t({\mathbf{X}}_t)\right]\mathrm{d}t}}\\ & +\underset{\text{Langevin 纠偏项}}{\underbrace{-\frac{1}{2}g(t{)}^2\nabla \log p_t({\mathbf{X}}_t)\mathrm{d}t+g(t)\mathrm{d}{\overline{\mathbf{W}}}_t}}.\end{array}$$

固定一个采样时刻 $t$，冻结 $p_t$ 与 $g(t)$。由于反向采样中 $\mathrm{d}t<0$，引入正向虚拟时间 $\mathrm{d}s=-\mathrm{d}t$；布朗增量改变符号后分布不变，于是纠偏项对应

$$\mathrm{d}{\mathbf{Z}}_s=\frac{1}{2}g(t{)}^2\nabla \log p_t({\mathbf{Z}}_s)\mathrm{d}s+g(t)\mathrm{d}{\mathbf{B}}_s.$$

若 $q(\mathbf{x},s)$ 是虚拟动力学的密度，其 Fokker—Planck 方程为

$$\begin{array}{rl}{\partial }_{s}q& =-\nabla \cdot \left[\frac{1}{2}g(t{)}^2\nabla \log p_{t}q\right]+\frac{1}{2}g(t{)}^2\Delta q\\ & =\frac{1}{2}g(t{)}^2\left[-\nabla \cdot \left(\frac{\nabla p_t}{p_t}q\right)+\Delta q\right].\end{array}$$

当 $q=p_t$ 时，

$${\partial }_{s}q=\frac{1}{2}g(t{)}^2\left[-\nabla \cdot (\nabla p_t)+\Delta p_t\right]=0.$$

因此冻结时刻的 Langevin 动力学以 $p_t$ 为不变密度。这一计算只证明不变性，不自动证明从任意初值收敛到 $p_t$；收敛还需要归一化、尾部、不可约性或耗散性等遍历条件。实际离散采样中的有限步长、近似 score 和不精确终端分布也会破坏精确的边缘一致性。