Reverse-Time SDE

$$$$

问题引入

考虑如下一维 SDE： $$\mathrm{d}{X}_t=f(X_t,t)\mathrm{d}t+g(X_t,t)\mathrm{d}{W}_t$$ 其 Fokker-Planck 方程 为： $$\begin{array}{}\text{(1)}& \frac{\partial p(x,t)}{\partial t}=-\frac{\partial }{\partial x}[f(x,t)p(x,t)]+\frac{1}{2}\frac{{\partial }^2}{\partial x^2}[g^2(x,t)p(x,t)]\end{array}$$ FP 方程刻画了分布 $p(x,t)$ 在 $t$ 时刻的瞬时变化情况。如果两个 SDE 有着相同的 FP 方程和相同的初始分布 $p(x,0)$，那么它们在任意时刻 $t$ 的分布 $p(x,t)$ 也都相同——在扩散模型中，这个结论被用于寻找对应于给定 SDE 的 Probability Flow ODE（因为 ODE 可以视为扩散项为零的 SDE，此时 FP 方程退化为连续性方程）。

本文关心一个不同但类似的问题——假设系统的总运行时间为 $T$，能否找到一个 reverse-time SDE，其在 $T-t$ 时刻的分布与原 SDE 在 $t$ 时刻的分布相同？如果能够做到这一点，那么 reverse-time SDE 就是我们在扩散模型中要寻求的生成模型。

推导过程

变量代换

视 $\text{(1)}$ 式中的时间参数 $t$ 为系统已运行的时间，考虑将其代换为距离结束还剩下的时间 $s:=T-t$，有： $$\frac{\partial p(x,T-s)}{\partial (T-s)}=-\frac{\partial }{\partial x}[f(x,T-s)p(x,T-s)]+\frac{1}{2}\frac{{\partial }^2}{\partial x^2}[g^2(x,T-s)p(x,T-s)]$$ 为书写方便起见，引入记号： $$\overline{f}(x,s):=f(x,T-s),\overline{g}(x,s):=g(x,T-s),q(x,s):=p(x,T-s)$$ 则上式改写作： $$-\frac{\partial q(x,s)}{\partial s}=-\frac{\partial }{\partial x}[\overline{f}(x,s)q(x,s)]+\frac{1}{2}\frac{{\partial }^2}{\partial x^2}[{\overline{g}}^2(x,s)q(x,s)]$$ 取个负号得： $$\begin{array}{}\text{(2)}& \frac{\partial q(x,s)}{\partial s}=\frac{\partial }{\partial x}[\overline{f}(x,s)q(x,s)]-\frac{1}{2}\frac{{\partial }^2}{\partial x^2}[{\overline{g}}^2(x,s)q(x,s)]\end{array}$$ 注意，截至目前我们所做的只是一堆变量代换而已，还没有涉及任何与 reverse-time SDE 相关的东西。

大胆猜想

注意到 $\text{(2)}$ 式与 FP 方程的形式还挺像的，于是我们有了一个大胆的猜想：是否真的存在一个以 $s$ 为时间参数的 SDE，其 FP 方程正好就是 $\text{(2)}$ 式呢？具体而言，设有如下 SDE： $$\mathrm{d}{Y}_s=b(Y_s,s)\mathrm{d}s+\sigma (Y_s,s)\mathrm{d}{W}_s$$ 那么它的 FP 方程为： $$\begin{array}{}\text{(3)}& \frac{\partial q(x,s)}{\partial s}=-\frac{\partial }{\partial x}[b(x,s)q(x,s)]+\frac{1}{2}\frac{{\partial }^2}{\partial x^2}[{\sigma }^2(x,s)q(x,s)]\end{array}$$

如果猜想成立，那么 $\text{(3)}$ 式应与 $\text{(2)}$ 式相同，联立便可解出 $b(x,s)$ 与 $\sigma (x,s)$，即找到了我们想要的 SDE.

小心求解

联立 $$ 式与 $$ 式，得： $$ 合并化简得： $$

左右乘上任一光滑紧支的测试函数 $$ 并对 $$ 积分： $$ 左右各自做一次分部积分，得： $$ 这里假设了 $$ 在无穷远处为零。根据 $$ 的任意性，应有： $$ 将右式按求导乘积法则展开，并将左边的 $$ 除到右边，整理得： $$ 因此我们要找的 SDE 为： $$ 注意这是以 $$ 为自由变量的一族 SDE，而不是一个特定的 SDE. 这是因为 SDE 与 FP 方程是多对一的关系，不同的 SDE 可以具有相同的 FP 方程。下面我们考察两个特解。

Reverse-Time SDE

当取 $$ 时，代入 $$ 式得到一个特解： $$ 将 $$ 代换回原本的记号 $$，有： $$ 再将 $$ 代换回 $$，有： $$

引入记号： $$ 则有： $$ 这就是 reverse-time SDE.

Probability Flow ODE

当取 $$ 时，代入 $$ 式得到另一个特解： $$ 将 $$ 代换回原本的记号 $$，有： $$ 再将 $$ 代换回 $$，有： $$ 引入记号： $$ 则有： $$ 这就是 Probability Flow ODE. 由于 ODE 本来就是可逆的，所以上式与 forward-time 的 PF ODE 完全一致。

多维情形

上文推导的是一维情形，其结论可以推广到多维情形下。设有如下 $$ 空间中的 SDE： $$ 其中 $$，$$，则 reverse-time SDE 为： $$ 其中对于矩阵值函数 $$，符号 $$.

特别地，倘若 $$，即原 SDE 为： $$ 则 reverse-time SDE 简化为： $$ 这就是扩散模型中常用到的形式。

参考资料

1. Ji-Ha Kim. Deriving Reverse-Time Stochastic Differential Equations (SDEs). https://jiha-kim.github.io/posts/deriving-reverse-time-stochastic-differential-equations-sdes/ ↩︎
2. Song, Yang, et al. Score-Based Generative Modeling through Stochastic Differential Equations. International Conference on Learning Representations. ↩︎