Python 全栈体系【四阶】（五十）

第五章 深度学习

十一、扩散模型

4. 附录：Diffusion的数学推导过程

4.2 Diffusion正向扩散过程推导

设初始数据$x_0$符合分布$q(x_0)$，即训练集分布，然后不断向其中添加高斯噪声，高斯噪声本身是不可训练参数，或者说均值和方差是固定的，通过方差系数${\beta }_1,\cdots ,{\beta }_n$来控制添加噪声的强度，它们是0~1之间的小数，一般会越来越大. 另外，这个过程被固定为马尔科夫链，每步的条件转移分布为$q(x_t|x_{t-1})$，整体后验分布表示为$q(x_{1:T}|x_0)$，也就是连乘的形式.

$$\begin{gathered} q(x_t|x_{t-1})=N(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I) \\ q(x_{1:T}|x_0)=\prod _{t=1}^{T}q(x_t|x_{t-1}) \end{gathered}$$

正向过程的特点在于，可以根据系数$\beta$和$x_0$直接求出任意时刻的转移分布$q(x_t|x_0)$，如下所示：

$$q(x_t|x_0)=N(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I)$$

其中，${\alpha }_t=1-{\beta }_t$，${\bar{\alpha }}_t={\prod }_{i=1}^t{\alpha }_i$.

4.2.1 正向扩散过程的推导

利用重参数化技巧，可将$x_t$表示为$x_{t-1}$加上一个噪声值

$$x_t=\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1};where{ϵ}_{t-1},{ϵ}_{t-2},\cdots \sim N(0,1)$$

$t_{x-1}$又可以表示为$x_{t-1}$加上一个噪声值：

$$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t}(\sqrt{t-1}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{ϵ}_{t-2})+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{ϵ}_{t-2}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\end{array}$$

上式第二项、第三项为两个高斯噪声相加. 两个均值为-的高斯值相加，均值仍为0；方程相加${\sigma }_1^2+{\sigma }_2^2={\alpha }_t(1-{\alpha }_{t-1})+1-{\alpha }_t=1-{\alpha }_t{\alpha }_{t-1}$

所以，第二项、第三项合并后可表示为：

$$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{ϵ}}_{t-2};where{\overline{ϵ}}_{t-2}\text{merges two Gaussians(*)}\\ & =...\\ & =\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-\bar{\alpha }}ϵ\end{array}$$

这样，就能根据初始分布和时间t直接求出t时刻的分布.

4.2.2 逆向去噪过程推导

如果把正向扩散过程比作墨水在水中扩散的过程，逆向过程就相当于从水中提取出墨水的过程. 为了简化分析，也把它假定为马尔科夫链，转移分布也是高斯的，这样就变成了一个参数估计问题，用神经网络来学习转移分布.

$$p_{\theta }(x_{0:T})=p(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t)=N(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$$

其中，网络的输入是$x_t$和$t$，转移分布为$\mu$，协方差为$\Sigma$，$\theta$为模型参数（要求的目标），转移概率$p_{\theta }$为未知的. 逆向过程比正向过程要难（这就好比把墨水融到水中容易，把墨水从水中提取出来更困难），Diffusion模型的做法是，通过公式推导，把逆向过程的转移分布$p_{\theta }(x_{t-1}|x_t)$用正向扩散过程的后验分布$q(x_{t-1}|x_t,x_0)$来逼近：

$$q(x_{t-1}|x_t,x_0)=N(x_{t-1};\tilde{\mu }(x_t,x_0),\widetilde{{\beta }_t}I)$$

根据贝叶斯定理：

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

将上面的公式写成指数形式，并省略掉前面的系数，展开，凑成一元二次方程标准形式：

$$\begin{array}{rl}& q(x_t|x_{t-1},x_0)\frac{q(x_{t-1}|x_0)}{q(x_t|x_0)}\\ & \propto exp(-\frac{1}{2}(\frac{(x_t-\sqrt{\alpha }{x}_{t-1}{)}^2}{{\beta }_t}+\frac{(x_{t-1}-\sqrt{{\bar{\alpha }}_{t-1}}{x}_0{)}^2}{1-{\bar{\alpha }}_{t-1}}-\frac{(x_t-\sqrt{{\bar{\alpha }}_t}{x}_0{)}^2}{1-{\bar{\alpha }}_t}\left)\right)\\ & =exp(-\frac{1}{2}(\frac{(x_t^2-2{\sqrt{\alpha }}_t{x}_t{x}_{t-1}+{\alpha }_t{x}_{t-1}^2}{{\beta }_t}+\frac{(x_{t-1}^2-2\sqrt{{\bar{\alpha }}_{t-1}}{x}_0{x}_{t-1}+{\bar{\alpha }}_{t-1}{x}_0^2)}{1-{\bar{\alpha }}_{t-1}}-\frac{(x_t-\sqrt{{\bar{\alpha }}_t}{x}_0{)}^2}{1-{\bar{\alpha }}_t}\left)\right)\\ & =exp(-\frac{1}{2}((\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})x_{t-1}^2-(\frac{2{\sqrt{\alpha }}_t}{{\beta }_t}{x}_t+\frac{2\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{x}_0)x_{t-1}+C(x_t,x_0)\left)\right)\end{array}$$

根据一元二次方程方程公式$a(x-\frac{b}{2a}{)}^2+(\frac{4ab-b^2}{4a})=0$和高斯概率密度函数$f(x)=\frac{1}{\sqrt{2\pi \sigma }}exp(-\frac{(x-\mu {)}^2}{2{\sigma }^2})$可知，均值为 $-\frac{b}{2a}$，方差为$\frac{1}{a}$，带入公式计算出均值和方差（以下公式省略掉了常数）：

$$\begin{array}{rl}\tilde{\mu }(x_t,x_0)& =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{x}_0)/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})\\ & =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{x}_0)\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}.{\beta }_t\\ & =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0\\ \widetilde{{\beta }_t}& =1/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})=1/(\frac{{\alpha }_t-\overline{{\alpha }_t}+{\beta }_t}{{\beta }_t(1-{\bar{\alpha }}_{t-1})})=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}.{\beta }_t\end{array}$$

这里使用了之前的定义:

$$\begin{gathered} {\alpha }_t=1-{\beta }_t \\ {\bar{\alpha }}_t=\prod _{i=1}^T{\alpha }_t \end{gathered}$$

这样，就得到扩散过程后验分布$q(x_{t-1}|x_t,x_0)$的解析式，它是一个高斯分布. 其中，均值是关于$\alpha ,\bar{\alpha },{\beta }_t$以及$x_0,x_t$的表达式， 而方差完全是个常数，跟x没有关系. 进一步根据前面正向扩散过程中，应用重参数技巧推导的$x_0$与$x_t$之间的关系，可以得到：

$$x_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-\overline{{\alpha }_t}}{ϵ}_t)$$

带入上面均值表达式，替换掉$x_0$，于是${\tilde{\mu }}_t$就等于：

$$\begin{array}{rl}{\tilde{\mu }}_t& =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t)\\ & =\frac{1}{{\sqrt{\alpha }}_t}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_t)\end{array}$$

其中，$ϵ$是t时刻从标准正态分布中采样得到的随机值. 到此为止，就完成扩散过程的后验概率分布$q(x_{t-1}|x_t,x_0)$， 它依旧是一个高斯分布，均值只和$x_t$及标准正态分布的噪声相关，方差只跟常数$\alpha$和$\beta$相关. 扩散模型的重要意义在于提供了一种全新的生成模型范式，可以更好地描述数据的演化过程. 总体上来说，正向扩散和逆向扩散都是马尔科夫链，其中正向过程是确定性的、可控的，通过不断调整系数${\beta }_t$逐步添加噪声；转移分布是高斯的；逆向过程虽然复杂，但是转移分布$p_{\theta }(x_{t-1}|x_t)$也可以假设为高斯，用神经网络来逼近求解. 由于直接求解缺少有效数据，因此先推导了更容易求得，有解析式的正向过程后验条件分布$q(x_{t-1}|x_t,x_0)$，用来逼近$p_{\theta }(x_{t-1}|x_t)$. 这三个分布某种意义上刻画了Diffusion模型的全部演化过程，在后面求损失函数过程中用到.

4.2.3 损失函数变分推导

先求数据的负对数似然函数，直接不好求解，从而变通一下，寻找它的上界，其方法就是加上一个KL散度，因为KL散度是非负数.

$$-log{p}_{\theta }(x_0)\le -log{p}_{\theta }(x_0)+D_{KL}(q(x_{1:T}|x_0)||p_{\theta }(x_{1:T}|x_0))$$

最小化化负对数似然，等价于最小化他的上界（上式中右边的部分）. 将右边部分进行变形，先利用贝叶斯定理展开，然后消掉最后一个无关项：

$$\begin{array}{rl}-log{p}_{\theta }(x_0)& \le -log{p}_{\theta }(x_0)+D_{KL}(q(x_{1:T}|x_0)||p_{\theta }(x_{1:T}|x_0))\\ & =-log{p}_{\theta }(x_0)+{\mathbb{E}}_{x_{1:T}\sim q(x_{1:T}|x_0)}\left[log\frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})/p_{\theta }(x_0)}\right]\\ & =-log{p}_{\theta }(x_0)+{\mathbb{E}}_q[log\frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}+log{p}_{\theta }(x_0)]\\ & ={\mathbb{E}}_q\left[log\frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right]\end{array}$$

不等式两边都加上一个对$q(x_0)$的期望，左侧就是交叉熵，右边就变成对$q(x_{0:T})$求期望：

$${\mathbb{E}}_q(x_0)log{p}_{\theta }(x_0)\le {\mathbb{E}}_{q(x_{0:T})}\left[log\frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right]$$

最小化交叉熵就等价于最小化它的上界，右侧部分称为证据下界（Evidence Lower Bound），也就是变分推断中的ELBO，只不过前面加了符号，因此最大化对数似然，变成了最小化负对数似然，右侧的下界也变成了上界. Diffusion模型选取的损失函数就是目标数据的交叉熵，然后通过变分先找到上界，然后持续化简上界表达式，因为分子部分就是正向过程的条件概率分布，分布部分是逆向过程的联合分布. 接下来，就是对右侧一顿暴推，推导出迭代形式：

$$\begin{array}{rl}{L}_{VLB}& ={\mathbb{E}}_{q(x_{0:T})}\left[log\frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right]\\ & ={\mathbb{E}}_q\left[log\frac{{\prod }_{t=1}^{T}q(x_t|x_{t-1})}{p_{\theta }(x_T){\prod }_{t=1}^T{p}_{\theta }(x_{t-1}|x_t)}\right]\\ & ={\mathbb{E}}_q[-log{p}_{\theta }(x_T)+\sum _{t=1}^{T}log\frac{q(x_t|x_{t-1})}{p_{\theta }(x_{t-1}|x_t)}]\\ & ={\mathbb{E}}_q[-log{p}_{\theta }(x_T)+\sum _{t=2}^{T}log\frac{q(x_t|x_{t-1})}{p_{\theta }(x_{t-1}|x_t)}+log\frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}]\\ & ={\mathbb{E}}_q[-log{p}_{\theta }(x_T)+\sum _{t=2}^{T}log(\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}.\frac{q(x_t|x_0)}{q(x_{t-1}|x_0)})+log\frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}]\\ & ={\mathbb{E}}_q[-log{p}_{\theta }(x_T)+\sum _{t=2}^{T}log\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}+\sum _{t=2}^T\frac{q(x_t|x_0)}{q(x_{t-1}|x_0)}+log\frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}]\\ & ={\mathbb{E}}_q[-log{p}_{\theta }(x_T)+\sum _{t=2}^{T}log\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}+log\frac{q(x_T|x_0)}{q(x_1|x_0)}+log\frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}]\\ & ={\mathbb{E}}_q[log\frac{q(x_T|x_0)}{p_{\theta }(x_T)}+\sum _{t=2}^{T}log\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}-log{p}_{\theta }(x_0|x_1)]\\ & ={\mathbb{E}}_q[\underset{L_T}{\underbrace{D_{KL}(q(x_T|x_0)||p_{\theta }(x_T))}}+\sum _{t=2}^T\underset{L_{T-1}}{\underbrace{D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))}}-\underset{L_0}{\underbrace{log{p}_{\theta }(x_0|x_1)}}]\end{array}$$

4.2.4 损失函数的参数化

两个高斯分布p和q的KL散度其实是可以直接求解的，只和它们的均值和方差有关：

$$KL(p,q)=log\frac{{\sigma }_2}{{\sigma }_1}+\frac{{\sigma }^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}-\frac{1}{2}$$

在推导得到的KL散度损失函数中，两个高斯分布的方差都是常数，因此对最优化没有贡献，可以忽略掉，只剩下含有两个均值的部分，得到：

$$L_t={\mathbb{E}}_{x_0,ϵ}[\frac{1}{2||{\Sigma }_{\theta }(x_t,t)|{|}_2^2}||{\tilde{\mu }}_t(x_t,x_0)-{\mu }_{\theta }(x_t,t)|{|}^2]$$

其中$\tilde{\mu }(x_t,x_0)$是$q(x_{t-1}|x_t,x_0)$的均值，${\mu }_{\theta }(x_t,t)$是$p_{\theta }(x_{t-1}|x_t)$的均值，优化的目标就是后面逆向过程的均值${\mu }_{\theta }$要尽量毕竟前面正向过程后验分布的均值$\tilde{\mu }$，或者说训练的目标就是让${\mu }_{\theta }$来预测$\tilde{\mu }$. 前面这个均值我们刚才已经求出来了，有具体解析式.

$$\begin{array}{rl}{\tilde{\mu }}_t& =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t)\\ & =\frac{1}{{\sqrt{\alpha }}_t}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_t)\end{array}$$

因为$x_t$在训练过程中是已知的，因此后面的均值${\mu }_{\theta }$也可以通过重参数化技巧写成$x_t$和一个含参的高斯噪声${ϵ}_{\theta }$的形式（下式中第二步），整理合并同类项，就消掉了$x_t$，只剩下两个$ϵ$之间的差值. 再根据正向过程的重参数化推导，把$x_t$替换成$x_0$和${ϵ}_t$的形式.

$$\begin{array}{rl}{L}_t& ={\mathbb{E}}_{x_0,ϵ}[\frac{1}{2||{\Sigma }_{\theta }(x_t,t)|{|}_2^2}||{\tilde{\mu }}_t(x_t,x_0)-{\mu }_{\theta }(x_t,t)|{|}^2]\\ & ={\mathbb{E}}_{x_0,ϵ}[\frac{1}{2||{\Sigma }_{\theta }|{|}_2^2}||\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_t)-\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)|{|}^2]\\ & ={\mathbb{E}}_{x_0,ϵ}[\frac{(1-{\alpha }_t{)}^2}{2{\alpha }_t(1-{\bar{\alpha }}_t)||{\Sigma }_{\theta }|{|}_2^2}||{ϵ}_t-{ϵ}_{\theta }(x_t,t)|{|}^2]\\ & ={\mathbb{E}}_{x_0,ϵ}[\frac{(1-{\alpha }_t{)}^2}{2{\alpha }_t(1-{\bar{\alpha }}_t)||{\Sigma }_{\theta }|{|}_2^2}||{ϵ}_t-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t,t)|{|}^2]\end{array}$$

上面的式子，表达的含义是：有一个神经网络，输入$x_0,{ϵ}_t$和时间戳$t$，输出是预测的${ϵ}_{\theta }$，用来逼近扩散过程噪声${ϵ}_t$. 这样就实现了对负对数似然的优化. 论文作者进一步发现前面一项系数可以丢掉，并不影响结果，而且还更稳定，所以可以进一步简化为：

$$\begin{array}{rl}{L}_t^{simple}& ={\mathbb{E}}_{t\sim [1,T],x_0,{ϵ}_t}[||{ϵ}_t-{ϵ}_{\theta }(x_t,t)|{|}^2]\\ & ={\mathbb{E}}_{t\sim [1,T],x_0,{ϵ}_t}||{ϵ}_t-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t,t)|{|}^2]\end{array}$$

其中，${ϵ}_{\theta }$是由一个含有参数的神经网络预测得到的值，${ϵ}_t$就是当前时刻的噪声随机量. 到此，就完成了损失函数部分的推导，过程较为复杂，但结果却十分简洁.