扩散模型（四）：PF-SDE

扩散模型（四）：PF-SDE

从离散到连续

本文我们尝试利用随机微分方程SDE将两种扩散模型DDPM和NCSN统一起来。首先我们给出SDE的定义，随机微分方程SDE描述了随机过程的动态行为，其中包括确定性部分和随机部分，一般形式为：

\[dX_t=f(x_t,t)dt+g(X_t,t)dW_t \]

\(X_t\) 表示随机过程， \(f(x_t,t)\) 是漂移系数，表示确定性的变化率， \(g(X_t,t)\) 是扩散系数，表示随机性的强度， \(W_t\) 是标准布朗运动，代表随机噪声， \(W_t\) 是期望为 \(0\) ，方差为 \(t\) 的正态分布： \(W_t\sim\mathcal{N}(0,t) \Rightarrow W_{t+\Delta t}-W_t\sim\mathcal{N}(0,\Delta t)\) 当 \(\Delta t \to 0\) 时， \(dw=\sqrt{dt}w\) 。

下面我们尝试把DDPM和NSCN的前向过程从离散形式推广到连续的SDE形式。首先分析DDPM，给出DDPM的前向过程：

\[x_t=\sqrt{1-\beta_i}x_{i-1}+\sqrt{\beta_i}\epsilon_i,\, i=1,...,N \]

令 \(\overline{\beta}_i=N\beta_i\) ，那么有 \(x_i=\sqrt{1-\frac{\overline{\beta}_i}{N}}x_{i-1}+\sqrt{\frac{\overline{\beta}_i}{N}}\epsilon_i\)， 我们进一步令 \(\Delta t=\frac{1}{N}\) ，那么当 \(\Delta t\to \infty\) 时， \(\{\overline{\beta}_i\}_{i=1}^N\) 成为了一个关于时间 \(t\) 的连续函数 \(\beta(t),\,t\in\left[0,1\right]\) ，此时，前向过程转化为如下形式：

\[\begin{align*} x(t+\Delta t) &= \sqrt{1-\beta(t+\Delta t)\Delta t}x(t)+\sqrt{\beta(t+\Delta t)\Delta t}\epsilon \\ &\approx x(t)-\frac{1}{2}\beta(t+\Delta t)\Delta tx(t)+\sqrt{\beta(t+\Delta t)\Delta t}\epsilon \\ &\approx x(t)-\frac{1}{2}\beta(t)\Delta tx(t)+\sqrt{\beta(t)\Delta t}\epsilon \end{align*} \]

上述推导使用了一阶泰勒展开和 \(\Delta t \to 0\) 的性质。当 \(\Delta t \to 0\) 时，我们得到DDPM前向过程的SDE形式：

\[dx=-\frac{1}{2}\beta_tdt+\sqrt{\beta_t}dw \]

由于DDPM的前向过程中 \(x_t=\sqrt{\overline{\alpha}_t}x_0+\sqrt{1-\overline{\alpha}_t}\overline{\epsilon_t}\) 方差始终有界，我们将DDPM对应的SDE称为Variance Preserving方差缩紧SDE，简称为VP-SDE。

NCSN的前向过程为：

\[\begin{align*} x_{x+\Delta t} &= x_t+\sqrt{\sigma^2_{t+\Delta t}-\sigma^2_t}\epsilon \\ &= x_t+\sqrt{\frac{\sigma^2_{t+\Delta t}-\sigma^2_t}{\Delta t}}\sqrt{\Delta t}\epsilon \\ &= x_t+\sqrt{\frac{\Delta\sigma^2_t}{\Delta t}}\sqrt{\Delta t}\epsilon \end{align*} \]

当 \(\Delta t \to 0\) 时，我们得到NCSN前向过程的SDE形式：

\[dx=\sqrt{\frac{d\sigma^2_t}{dt}}dw \]

分析NSCN的前向过程 \(x_t=x_t+\sigma_t\epsilon\) ，由于 \(\sigma_t\) 随着 \(t\) 的增大而增大，当 \(t\to\infty\) 时，方差会愈发增大，故我们将NCSN对应的SDE称为Variance Exploding方差爆炸SDE，简称为VE-SDE。

因此，我们可以使用下述SDE描述扩散模型的前向过程：

\[dx=f_t(x)dt+g_tdw \]

当 \(f_t(x)=-\frac{1}{2}\beta_t,\,g_t=\sqrt{\beta_t}\) 时，该SDE对应DDPM的前向过程；当 \(f_t(x)=0,\,g_t=\sqrt{\frac{d\sigma^2_t}{dt}}\) 时，该SDE对应NCSN的前向过程。

逆向过程SDE

得到扩散模型的前向SDE后，我们需要还需推导一个逆向的SDE方程用于还原真实数据分布。

首先，我们把前向SDE写成离散形式：

\[x_{t+\Delta t}-x_t=f_t(x_t)\Delta t+g_t\sqrt{\Delta t}\epsilon,\,\epsilon\sim\mathcal{N}(0,\mathbf{I}) \]

我们将其转化为条件概率形式：

\[\begin{align*} p(x_{t+\Delta t}|x_t) &= \mathcal{N}(x_{t+\Delta t};x_t+f_t(x_t)\Delta t,g_t^2\Delta t\mathbf{I}) \\ &\propto exp\left(-\frac{||x_{t+\Delta t}-x_t-f_t(x_t)\Delta t||^2}{2g_t^2\Delta t}\right) \end{align*} \]

我们尝试使用贝叶斯定理推导 \(p(x_t|x_{t+\Delta t})\) 的形式：

\[\begin{align*} p(x_t|x_{t+\Delta t}) &= \frac{p(x_{t+\Delta t}|x_t)p(x_t)}{p(x_{t+\Delta t})} = p(x_{t+\Delta t}|x_t)exp(\log{p(x_t)}-\log{p(x_{t+\Delta t})}) \\ &\propto exp\left(-\frac{||x_{t+\Delta t}-x_t-f_t(x_t)\Delta t||^2}{2g_t^2\Delta t}+ \log{p(x_t)}-\log{p(x_{t+\Delta t})}\right) \end{align*} \]

我们对关于变量 \((x_t,t)\) 的函数 \(\log{p(x_{t+\Delta t})}\) 作泰勒展开：

\[\log{p(x_{t+\Delta T})} \approx \log{p(x_t)}+(x_{t+\Delta T}-x_t)\cdot\nabla_{x_t}\log{p(x_t)}+\Delta t\frac{\partial}{\partial t}\log{p(x_t)} \]

代入上式后，我们配方得到如下式子：

\[p(x_t|x_{t+\Delta t}) \propto exp\left(-\frac{||x_{t+\Delta t}-x_t-[f_t(x_t)-g_t^2\nabla_{x_t}\log{p(x_t)}]\Delta t||^2}{2g_t^2\Delta t} +\mathcal{O}(\Delta t) \right) \]

当 \(\Delta t\to0\) 时， \(\mathcal{O}(\Delta t)\to0\) 因此该项不起作用：

\[\begin{align*} p(x_t|x_{t+\Delta t}) &\propto exp\left(-\frac{||x_{t+\Delta t}-x_t-[f_t(x_t)-g_t^2\nabla_{x_t}\log{p(x_t)}]\Delta t||^2}{2g_t^2\Delta t} \right) \\ &\approx exp\left(-\frac{||x_{t}-x_{t+\Delta t}+[f_t(x_{t+\Delta t})-g_{t+\Delta t}^2\nabla_{x_{t+\Delta t}}\log{p(x_{t+\Delta t})}]\Delta t||^2}{2g_{t+\Delta t}^2\Delta t} \right) \end{align*} \]

因此，我们认为 \(p(x_t|x_{t+\Delta t})\) 近似于正态分布： \(\mathcal{N}(x_t;x_{t+\Delta t}-[f_t(x_{t+\Delta t})-g_{t+\Delta t}^2\nabla_{x_{t+\Delta t}}\log{p(x_{t+\Delta t})}]\Delta t,g_{t+\Delta t}^2\Delta t\mathbf{I})\) ，即：

\[x_t-x_{t+\Delta t}=-[f_t(x_{t+\Delta t})-g_{t+\Delta t}^2\nabla_{x_{t+\Delta t}}\log{p(x_{t+\Delta t})}]\Delta t-g_{t+\Delta t}\sqrt{\Delta t}\epsilon,\,\epsilon\sim(0,\mathbf{I}) \]

取 \(\Delta t\to0\) 的极限，我们得到逆向过程对应的SDE：

\[dx=\left[f_t(x)-g_t^2\nabla_x\log{p_t(x)} \right]dt+g_tdw \]

式中的 \(\nabla_x\log{p_t(x)}\) 便是上篇文章所提到的score function，我们可以利用上文所提到的条件得分匹配方法，设计神经网络估计score function \(s_{\theta}(x_t,t)\) ，并通过以下损失函数进行优化：

\[\mathbb{E}_{x_0,x_t\sim p_0(x_0)p(x_t|x_0)}\left[||\nabla_{x_t}\log{p(x_t|x_0)}-s_\theta(x_t,t)||^2\right] \]

与DDPM的联系

在DDPM中，加噪图像与预测噪声有如下关系：

\[x_t = \sqrt{\overline{\alpha}_t}x_0+\sqrt{1-\overline{\alpha}_t}\epsilon_{\theta}(x_t,t) \\ \epsilon_{\theta}(x_t,t) =\frac{x_t-\sqrt{\overline{\alpha}_t}x_0}{\sqrt{1-\overline{\alpha}_t}} \]

而由于 \(x_t\sim\mathcal{N}(x_t;\sqrt{\overline{\alpha}_t}x_0,(1-\overline{\alpha}_t)\mathbf{I})\) ，我们可以对score项作以下化简：

\[\begin{align*} s_{\theta}(x_t,t) &= \nabla_{x_t}\log{p_t(x)} \\ &\propto -\nabla_{x_t}\frac{(x_t-\sqrt{\overline{\alpha}_t}x_0)^2}{2(1-\overline{\alpha}_t)} \\ &= -\frac{x_t-\sqrt{\overline{\alpha}_t}x_0}{1-\overline{\alpha}_t} \\ &=-\epsilon_{\theta}(x_t,t) \end{align*} \]

因此，DDPM噪声模型训练得到的噪声方向刚好是score的反方向，因此二者的预训练模型可以互相使用。