AIGC拾遗：PF-ODE

前言

本文试图从SDE过渡到PF-ODE，实现推理过程中的ODE采样。

等价SDE

假设前向SDE过程为

\[dx_{t}=f(t)x_{t}dt+g(t)dw_{t} \tag{1} \label{eq.1} \]

其Fokker-Planck方程为

\[\frac{\partial}{\partial t}p(x_{t}, t)=-\nabla·[f(t)x_{t}p(x_{t}, t)]+\frac{1}{2}g^{2}(t)\nabla·\nabla p(x_{t}, t) \]

进一步化简可得

\[\begin{aligned} \frac{\partial}{\partial t}p(x_{t}, t) &=-\nabla·[f(t)x_{t}p(x_{t}, t)-\frac{1}{2}g^{2}(t)\nabla{p(x_{t}, t)}] \\ &=-\nabla·[f(t)x_{t}-\frac{1}{2}g^{2}(t)\nabla{\log{p(x_{t}, t)}}]p(x_{t}, t) \\ &=-\nabla·[f(t)x_{t}-\frac{1}{2}(g^{2}(t)-\sigma^{2}(t))\nabla{\log{p(x_{t}, t)}}]p(x_{t}, t)+\frac{1}{2}\sigma^{2}(t)\nabla·\nabla{p(x_{t}, t)} \end{aligned} \]

其对应的SDE过程为

\[dx_{t}=[f(t)x_{t}-\frac{1}{2}(g^{2}(t)-\sigma^{2}(t))\nabla_{x_{t}}{\log{p_{t}(x_{t})}}]dt+\sigma(t)dw_{t} \tag{2}\label{eq.2} \]

因此，公式\eqref{eq.1}和公式\eqref{eq.2}具有相同的边缘概率密度\(p_{t}(x_{t})\)，求解公式\eqref{eq.2}的逆向SDE过程等价于求解公式\eqref{eq.1}的逆向SDE过程。

PF-ODE

当公式\eqref{eq.2}中的\(\sigma_{t}(t)=0\)时，该式化简为

\[dx_{t}=[f(t)x_{t}-\frac{1}{2}g^{2}(t)\nabla_{x_{t}}\log{p_{t}(x_{t})}]dt \tag{3} \label{eq.3} \]

公式\eqref{eq.3}是一个ODE过程，称为PF-ODE(probability flow ODE)，因此，我们可以采用SDE进行加噪，而采用PF-ODE进行去噪。

Velocity prediction

上述分析可知，对于前向加噪过程

\[x_{t}=\alpha_{t}x_{0}+\sigma_{t}\epsilon \]

存在对应的PF-ODE

\[dx_{t}=[f(t)x_{t}-\frac{1}{2}g^{2}(t)\nabla_{x_{t}}\log{p_{x_{t}}(x_{t}, t)}]dt \]

其中，

\[f(t)=\frac{d(ln\alpha_{t})}{dt}, ~~~g^{2}(t)=\frac{d(\sigma_{t}^{2})}{dt}-2\sigma_{t}^{2}\frac{d(ln\alpha_{t})}{dt} \]

因此，我们可以采用一个神经网络\(s_{\theta}(x_{t}, t)\)拟合\(\nabla_{x_{t}}\log{p_{t}(x_{t})}=-\frac{x_{t}-\alpha_{t}x_{0}}{\sigma_{t}^{2}}\)，然后通过ODE求解器生成样本。

然而，我们也可以直接拟合\(\frac{dx_{t}}{dt}\)，令

\[v_{\theta}(x_{t}, t)=f(t)x_{t}-\frac{1}{2}g^{2}(t)\nabla_{x_{t}}\log{p_{x_{t}}(x_{t}, t)}=\frac{dx_{t}}{dt}=\frac{d\alpha_{t}}{dt}x_{0}+\frac{d\sigma_{t}}{dt}\epsilon \]

当\(\alpha_{t}=t\)，\(\sigma_{t}=1-t\)时，\(v_{\theta}(x_{t}, t)=x_{0}-\epsilon\)，即为ReFlow的目标。