DDPM (Denoising Diffusion Probabilistic Model)
Published:
Paper Reading: Denoising diffusion probabilistic models.
本文主要介绍 Denoising Diffusion Probabilistic Model (DDPM)， 以图片生成为例。
给定一张图片 $x_0$, 采用一个加噪的过程将其变为Guassian分布：
\[\begin{align*} p(x_{1:T} \mid x_0) = \prod_{t=1}^T p(x_t \mid x_{t-1}) , \quad x_t \mid x_{t-1} \sim \mathcal{N}( \sqrt{1 - \beta_t} x_{t-1}, \beta_t I). \end{align*}\]其中 $\beta_t$ 为噪声的调度。定义 $\alpha_t =1 - \beta_t$ 以及 $\bar \alpha_t = \prod_{s=1}^s \alpha_s$.
经过如下推导
\[\begin{align*} x_T \mid x_0 &\sim \sqrt{\alpha_T} x_{T-1} + \sqrt{1- \alpha_T} \epsilon_T, \quad \epsilon_T \sim \mathcal{N}(0,I) \\ &\sim \sqrt{\alpha_T} \left( \sqrt{\alpha_{T-1}} x_{T-2} + \sqrt{1- \alpha_{T-1}} \epsilon_{T-1} \right) + \sqrt{1- \alpha_T} \epsilon_T, \quad \epsilon_{T-1}, \epsilon_T \sim \mathcal{N}(0,I) \\ &\sim \sqrt{\alpha_T \alpha_{T-1}} x_{T-2} + \sqrt{1- \alpha_T \alpha_{T-1}} \epsilon_{T-1}, \quad \epsilon_{T-1} \sim \mathcal{N}(0,I) \\ & \sim \cdots \\ & \sim \sqrt{\bar \alpha_T} x_0 + \sqrt{1 - \bar \alpha_T} \epsilon_0, \quad \epsilon_0 \sim \mathcal{N}(0,I) \\ & \sim \mathcal{N}( \sqrt{\bar \alpha_T} x_0, (1- \bar \alpha_T) I ). \end{align*}\]令 $\bar \alpha_T \rightarrow 0$ 我们知道 $x_T \mid x_0 \sim \mathcal{N}(0,I)$. 得到了一个纯噪声。
利用Bayes公式，可以求解逆向过程的条件分布
\[\begin{align*} p(x_{t-1} \mid x_t, x_0) &= \frac{p (x_{t} \mid x_{t-1} , x_0 ) p (x_t \mid x_0)}{p (x_{t-1} \mid x_0)} \\ &= \exp\left( - \frac{\left(x_t - \sqrt{\alpha_t} x_{t-1}\right)^2}{2(1- \alpha_t)} \right) \exp\left( - \frac{\left(x_t - \sqrt{\bar \alpha_t} x_0 \right)^2}{2(1- \bar \alpha_t)}\right) \exp\left( - \frac{\left(x_{t-1} - \sqrt{\bar \alpha_{t-1}} x_0 \right)^2}{2(1- \bar \alpha_{t-1})}\right) \\ &\propto \exp\left( -\frac{1}{2} \left( \left( \frac{\alpha_t}{1 - \alpha_{t}} + \frac{1}{ 1- \bar \alpha_{t-1}} \right) x_{t-1}^2 - 2 \left( \frac{\sqrt{\alpha_t}}{1 - \alpha_t} x_t + \frac{\sqrt{\bar \alpha_t}}{1 - \bar \alpha_t} \right)x_{t-1} \right) \right) \\ &\propto \exp\left( - \frac{\left( x_{t-1} - \mu (x_t,x_0,t) \right)^2}{2 \sigma_t^2} \right). \end{align*}\]其中均值和方差项可以通过配方法得到，形如
\[\begin{align*} \mu(x_t,x_0,t) = \frac{\sqrt{\alpha_t} (1- \bar \alpha_{t-1})}{1 - \bar \alpha_t}x_t + \frac{\sqrt{\bar{\alpha}_{t-1}} (1- \alpha_t)}{1- \bar \alpha_t} x_0, \quad \sigma_t^2 = \frac{(1 - \alpha_t ) (1- \bar \alpha_{t-1})}{1 - \bar \alpha_t}. \end{align*}\]我们希望训练一个参数为 $ \theta$ 的模型预测逆向过程，如此便可以从噪声中恢复数据。
由于方差可以根据时间 $t$ 计算，自然地可以令模型预测均值，也即给定时间步 $t$ 求解
\[\begin{align*} \min \mathbb{E}_{x_0,x_t } \left \Vert \mu(x_t,x_0,t) - \mu_{\theta}(x_t,x_0,t) \right \Vert^2 \end{align*}\]原文 [1] 采用最大化变分下界的方式推导上述式子，但可以从预测均值的角度直观理解。
但上面的问题需要模型输入图片 $x_0$。采用重参数化，改为预测噪声，可以使得模型仅需要输入 $x_t,t$. 注意到
\[\begin{align*} x_t \sim \sqrt{\bar \alpha_t} x_0 + \sqrt{ 1- \bar \alpha_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0,I). \end{align*}\]可以将损失函数等价地写成
\[\begin{align*} &\quad\min \mathbb{E}_{x_t,\epsilon} \left \Vert \frac{\sqrt{\alpha_t} (1- \bar \alpha_{t-1})}{1 - \bar \alpha_t}x_t + \frac{\sqrt{\bar{\alpha}_{t-1}} (1- \alpha_t)}{(1- \bar \alpha_t) \sqrt{\bar \alpha_t}} \left( x_t - \sqrt{1-\bar \alpha_t} \epsilon \right) - \mu_{\theta}(x_t,x_0,t) \right \Vert^2 \\ &= \min \mathbb{E}_{x_t,\epsilon} \left \Vert \frac{1}{\sqrt{\alpha_t} } \left( x_t - \frac{1- \alpha_t}{\sqrt{1 - \bar \alpha_t}} \epsilon \right) - \mu_{\theta}(x_t,x_0,t) \right \Vert^2. \end{align*}\]令模型的预测为
\[\begin{align*} \mu_{\theta}(x_t,x_0, t) = \frac{1}{\sqrt{\alpha_t} } \left( x_t - \frac{1- \alpha_t}{\sqrt{1 - \bar \alpha_t}} \epsilon_{\theta}(x_t, t) \right) . \end{align*}\]此时模型已经不需要 $x_0$ 作为输入值，忽略前面的系数，上面的式子等价于
\[\begin{align*} &\quad \min \mathbb{E}_{x_t, \epsilon} \Vert \epsilon - \epsilon_{\theta}(x_t,t) \Vert^2 \\ &= \min \mathbb{E}_{x_0, \epsilon} \Vert \epsilon - \epsilon_{\theta}( \sqrt{\bar \alpha_t}x_0 + \sqrt{1 - \bar \alpha_t} \epsilon,t) \Vert^2. \end{align*}\]这就得到了DDPM的训练过程，也即采样一个图片数据 $x_0$ 和标准高斯随机变量 $\epsilon$, 然后用梯度下降对上面的损失函数求解最优的 $\theta$.
得到了训练完的模型后，需要生成数据仅需要首先采样 $x_T \sim \mathcal{N}(0,I)$ 然后使用网络进行如下的逆向过程即可
\[\begin{align*} x_{t-1} = \frac{1}{\sqrt{\alpha_t} } \left( x_t - \frac{1- \alpha_t}{\sqrt{1 - \bar \alpha_t}} \epsilon_{\theta}(x_t, t) \right) + \sigma_t z, \quad z \sim \mathcal{N}(0,I). \end{align*}\]迭代 $T$ 步就可以得到最终的数据 $x_0$.
Reference
[1] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020.