Mathematical Principle of DDPM | Hexo

DDPM数学推导

前向过程和反向过程

约定$x_0$为原始输入干净图片，$x_T$为经过T步加噪声得到的纯噪声图，前向过程希望向原始图片中加入噪声，反向过程尝试通过UNet预测每一步加入的噪声并将其从图片中去除之

前向过程$x_{t-1}\to x_t$

借助一个无参数的条件概率$q(x_{t}|x_{t-1})$生成加噪后的图像

采样随机噪声$\epsilon_t|\mathcal N(0,I)$加入到原始图片中

$$x_t = a_t x_{t-1} + b_t \epsilon_t$$

系数满足$a_t^2 + b_t^2 = 1$，进一步拆分$x_{t-1}$，得到

$$x_t = a_t(a_{t-1} x_{t-2} + b_{t-1}\epsilon_{t-1}) + b_t \epsilon_t = \prod_{i=1}^t a_i x_0 + \prod_{i=2}^t a_i b_1 \epsilon_1 + \prod_{i=3}^t a_i b_2 \epsilon_2 +\cdots b_t \epsilon_t$$

RHS第二项到最后一项为多个高斯分布的加权和，对应也是高斯分布，方差为系数平方和，所有系数的平方和满足

$$(\prod_{i=1}^t a_i)^2+(\prod_{i=2}^t a_i b_1)^2 + (\prod_{i=3}^t a_i b_2)^2 + \cdots + b_t^2 =1$$

这样加噪过程记作

$$x_t = \sqrt {\overline\alpha_t} x_0 + \sqrt {1-\overline\alpha_t } \overline\epsilon_t,\overline \alpha_t = \prod_{i=1}^t a_i$$

令$a_t = \sqrt {\alpha_t}$，则一步加噪写成

$$x_t = \sqrt {\alpha_t} x_{t-1} + \sqrt {1- \alpha_t} \epsilon_t$$

显然，前向过程的条件分布写成

$$P(x_t|x_{t-1}) = \mathcal N(\sqrt {\alpha_t} x_{t-1},(1-\alpha_t)I)$$

是一个正态分布

逆向过程$x_T\to x_0$

借助一个有参数的模型$p_\theta(x_{t-1}|x_t)$进行去噪

前向过程得到的多层加噪$x_1,x_2,\cdots,x_T$相当于编码，逆向过程对其进行解码，这里我们希望估计真实图片$x_0$的概率$p_\theta(x_0)$并采取maximize likelihood的方式学习参数，将其视为一个马尔可夫过程，则概率估算为

$$p_{\theta}(x_0) = \int p_\theta(x_0,x_1,\cdots,x_T) d x_1 d x_2 \cdots d x_T$$

引入latent variable在输入$x_0$下的联合后验分布$q(x_{1:T}|x_0)$，则上述概率写成

$$p_\theta(x_0) = \int \frac{p_\theta(x_{0:T})q(x_{1:T}|x_0)}{q(x_{1:T}|x_0)} dx_1 dx_2\cdots d x_T = E_{x_{1:T}|q(x_{1:T}|x_0)}[\frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)}]$$

写成log形式，根据$E_{p(x)}[\log f(x)]\leq \log E[f(x)]$，有

$$\log E_{x_{1:T}|q(x_{1:T}|x_0)}[\frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)}]\geq E_{q(x_{1:T}|x_0)}[\log \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)}]$$

注意到$q(x_{1:T}|x_0)$和$p_\theta(x_{0:T})$实际上可以拆分为T步生成，写成

$$p_\theta(x_{0:T} )= p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}|x_t)\\ q(x_{1:T}|x_0) = \prod_{t=1}^T q(x_t|x_{t-1})$$

于是有

$$LHS = E_{q(x_{1:T}|x_0)}[\log \frac{p(x_T)\prod_{t=1}^T p_\theta(x_{t-1}|x_t)}{\prod_{t=1}^T q(x_t|x_{t-1})}]$$

分子分母拆分出第一项

$$RHS = E_{q(x_{1:T}|x_0)}[\log \frac{p(x_T)p_\theta(x_0|x_1)}{q(x_T|x_{T-1})}] + E_{q(x_{1:T}|x_0)}[\log \prod_{t=1}^{T-1} \frac{p_\theta(x_t|x_{t+1})}{q(x_t|x_{t-1})}]$$

对于RHS第一项，考虑分别优化

1. $\frac{p(x_T)}{q(x_T|x_{T-1})}$
2. $p_\theta(x_0|x_1)$

写成

$$E_{q(x_1|x_0)}[\log p_\theta(x_0|x_1)] + E_{q(x_{T-1},x_T|x_0)}[\log \frac{p(x_T)}{q(x_T|x_{T-1})}]$$

进一步将（13）中的第二项拆分为$E_{q(x_{T-1}|x_0) q(x_T|x_{T-1})}$，最终优化目标包含三项

1. $E_{q(x_1|x_0)}[\log p_\theta(x_0|x_1)]$，希望从加噪的$x_1$中重建$x_0$
2. $E_{q(x_{T-1}|x_0)}E_{ q(x_T|x_{T-1})}[\log \frac{p(x_T)}{q(x_T|x_{T-1})}] =- E_{q(x_{T-1}|x_0)}[D_{KL}(q(x_T|x_{T-1})\parallel p(x_T))]$，希望通过T步扩散得到的结果和先验概率（纯噪声高斯分布）尽量接近
3. 和2类似$-\sum_{t=1}^{T-1} E_{q(x_{t-1},x_{t+1}|x_0)}[D_{KL}(q(x_t|x_{t-1})\parallel p_\theta(x_t|x_{t+1}))]$

因为前向过程满足马尔可夫性，因此有

$$q(x_t|x_{t-1},x_0) = q(x_t|x_{t-1})$$

并且

$$q(x_t|x_{t-1},x_0) =\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)}$$

拆分

$$\prod_{t=1}^T q(x_t|x_{t-1}) = q(x_1|x_0)\prod_{t=2}^T q(x_t|x_{t-1},x_0)$$

得到

$$\begin{aligned} \log p_\theta(x_0) &\geq E_{q(x_{1:T}|x_0)}[\log \frac{p(x_T)p_\theta(x_0|x_1)}{q(x_1|x_0)} + \log \prod_{t=2}^T \frac{p_\theta(x_{t-1}|x_t)}{q(x_t|x_{t-1},x_0)}] \\ &=E_{q(x_{1:T}|x_0)}[\log \frac{p(x_T)p_\theta(x_0|x_1)}{q(x_1|x_0)} + \sum_{t=2}^T\log \frac{p_\theta(x_{t-1}|x_t)q(x_{t-1}|x_0)}{q(x_{t-1}|x_t,x_0)q(x_t|x_0)} ]\\ &=E_{q(x_{1:T}|x_0)}[\log \frac{p(x_T)p_\theta(x_0|x_1)}{q(x_1|x_0)} + \sum_{t=2}^T (\log \frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t,x_0)} + \log \frac{q(x_{t-1}|x_0)}{q(x_{t}|x_0)})]\\ &=E_{q(x_{1:T}|x_0)}{[\log \frac{p(x_T)p_\theta(x_0|x_1)}{q(x_1|x_0)} + \sum_{t=2}^T\log \frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t,x_0)}+\log \frac{q(x_1|x_0)}{q(x_T|x_0)}(和\theta优化无关)}\\ &=E_{q(x_{1:T}|x_0)}[\log \frac{p(x_T)p_\theta(x_0|x_1)}{q(x_T|x_0)} +\sum_{t=2}^T\log \frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t,x_0)}] \end{aligned}$$

进一步拆分为

1. $E_{q(x_1|x_0)}[\log p_\theta(x_0|x_1)]$，对应最后一步逆扩散的重建损失
2. $E_{q(x_{1:T}|x_0)}[\frac{p(x_T)}{q(x_T|x_0)}] = -D_{KL}(q(x_T|x_0)\parallel p(x_T))$，最终扩散出的图片应该完全是噪声
3. $E_{q(x_t,x_{t-1}|x_0)}[\log \frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t,x_0)}] = E_{q(x_{t-1}|x_t)}E_{q(x_{t}|x_0)}[\log \frac{p_\theta(x_{t-1}|x_r)}{q(x_{t-1}|x_t,x_0)} ]$，这里用$p_\theta(x_{t-1}|x_t)$建模$q(x_{t-1}|x_{t})$，得到

$$-E_{q(x_t|x_0)}[D_{KL}(q(x_{t-1}|x_t,x_0)\parallel p_\theta(x_{t-1}|x_t))]$$

训练

第一项用重建损失代替，第二项和$\theta$无关，关注第三项如何计算$q(x_{t-1}|x_t,x_0)$和$p_\theta(x_{t-1}|x_t)$之间的KL散度，这两个概率写成

1. 利用前向过程建模$q(x_{t-1}|x_t,x_0)$，写成（已知加噪马氏链上游和下游值，推断中间值）

$$\begin{aligned} q(x_{t-1}|x_t,x_0) &= \frac{q(x_t|x_{t-1},x_0)q(x_{t-1}|x_0)}{q(x_t|x_0)}\\ \end{aligned}$$

对应三个前向过程

• $x_{t-1}\to x_t,x_{t} = \sqrt {\alpha_t} x_{t-1} + \sqrt{1-\alpha_t}\epsilon_t,q(x_{t}|x_{t-1},x_0) = \mathcal N(\sqrt{\alpha_t} x_{t-1},(1-\alpha_t)I)$
• $x_0\to x_{t-1},x_{t-1} = \sqrt{\overline \alpha_{t-1}} x_0 + \sqrt{1- \overline\alpha_{t-1}}\epsilon,q(x_{t-1}|x_0) =\mathcal N(\sqrt{\overline{\alpha}_{t-1}}x_{0},(1-\overline{\alpha}_{t-1})I)$
• $x_0\to x_t$同理，$q(x_t|x_0) = \mathcal N(\sqrt{\alpha}_t x_0,(1-\overline \alpha_t)I)$

得到$q(x_{t-1}|x_t,x_0)$实际上也是一个高斯分布，均值和方差记作$\mu_q(x_t,x_0),\sum_q(t)I$

1. 利用后向过程建模单步去噪$p_\theta(x_{t-1}|x_{t})$，对应一个参数化的高斯分布$\mathcal N(\mu_\theta,\sum_\theta)$

优化两个概率分布之间的距离写成

$$\arg\min_\theta \frac{1}{2\sigma_q^2(t)}[\parallel \mu_\theta - \mu_q\parallel_2^2]\label{distance}$$

$\mu_q$是$x_0$和$x_t$的函数

$$\mu_q = \frac{\overline {\alpha_{t-1}}\beta_t}{1 -\overline \alpha_t}x_0 + \frac{\sqrt{\alpha_t}(1 -\overline \alpha_{t-1})}{1 -\overline \alpha_t}x_t$$

在扩散过程中，真值根据纯噪声$x_t$经过t步逆扩散得到，将其视为一个参数为$\theta$的神经网络，记作

$$x_0 = f_\theta(x_t,t)$$

因此KL散度写成

$$\arg\min_\theta \parallel f_\theta(x_t,t) - x_0\parallel_2^2$$

$\mu_q$的另一种形式为

$$\mu_q = \frac{1}{\sqrt{\alpha_t}} x_t -\frac{1 -\alpha_t}{\sqrt{1-\overline \alpha_t}\sqrt {\alpha_t}} \epsilon_t$$

第二项对应的是噪声，因此得到

$$\arg\min_\theta \parallel f_\theta(x_t,t) -\epsilon_t\parallel$$

逆扩散实际上学习到的是t step加入的噪声