番外篇Diffusion&Stable Diffusion扩散模型与稳定扩散模型

Diffusion&Stable Diffusion扩散模型与稳定扩散模型

摘要

本篇文章为阅读笔记，，主要内容围绕扩散模型和稳定扩散模型展开，介绍了kl loss、vae模型的损失函数以及变分下限作为扩展部分。扩散模型是一种生成模型，定义了一个逐渐扩散的马尔科夫链，逐渐项数据添加噪声，然后学习逆扩散过程，从噪声中构建所需的数据样本。稳定扩散模型在其基础上添加了编码器用以降维训练数据、降低训练成本，该模型亦添加了额外的文本嵌入向量，通过该向量模型得以根据文本生成图片。

主要内容参考deephub,“Diffusion 和Stable Diffusion的数学和工作原理详细解释”,https://zhuanlan.zhihu.com/p/597924053

Abstract

This article is a reading note, the main content is around the diffusion model and stable diffusion model, and introduces the loss function of VAE model, KL loss and the variational lower bound as an extended part. Diffusion model is a generative model that defines a progressively diffused Markov chain, gradually adding noise to the data, and then learning the reverse diffusion process to construct the required data samples from the noise. On the basis of this, an encoder is added to the stable diffusion model to reduce the dimensionality of the training data data and reduce the training cost. The model also adds an additional text embedding vector, through which images can be generated from the text.

Diffusion Model扩散模型

训练

• 正向扩散$\to$在图中添加噪声
• 反向扩散$\to$去除图中噪声

Forward Diffusion Process正向扩散过程

噪声图像的分布

$q(x_t|x_{t-1})=N(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)$

• t：time step时间步长
• $x_0$:真实数据分布中抽取的样本数据
• ${\beta }_t$:方差进度刻($0\le {\beta }_t\le 1$,${\beta }_0$小$\to {\beta }_T$大)
• $I$:单位阵
• $x_t$:输出output，亦是经过t次添加噪声的数据
• $\sqrt{1-{\beta }_t}{x}_{t-1}$:均值mean
• ${\beta }_{t}I$:方差variance
• 注：在正态分布中，各项参数依次是输出、均值、方差

逐步向图像中添加高斯噪声，从而迭代的产生样本$x_1,...,x_T$

当$T\to \infty$，将有完全高斯噪声图像

采样方法：可使用一个封闭公式在特定时间步长对图像直接取样

封闭公式

可通过重新参数化技巧得到

若z时高斯分布$N(\mu ,{\sigma }^2)$的取样，则$z=\mu +\sigma ϵ$

• $epsilon$是$N(0,1)$的取样

又$\mu =\sqrt{1-{\beta }_t}{x}_{t-1},{\sigma }^2={\beta }_{t}I$

令${\alpha }_t=1-{\beta }_t,{\overline{\alpha }}_{t={\prod }_{i-1}^t}{\alpha }_i,{ϵ}_i\sim N(0,I),{\overline{ϵ}}_i\sim N(0,1)$

$\begin{array}{rl}{x}_t& =\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}{ϵ}_{t-1}\\ & =\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{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{(1-{\alpha }_{t-1}){\alpha }_t}{ϵ}_{t-2}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\end{array}$

$X=\sqrt{(1-{\alpha }_{t-1}){\alpha }_t}{ϵ}_{t-2},Y=\sqrt{1-{\alpha }_t}{ϵ}_{t-1}$

$\because {ϵ}_{t-2},{ϵ}_{t-1}\sim N(0,1)$

$\therefore \sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{ϵ}_{t-2}\to X\sim N(0,{\alpha }_t(1-{\alpha }_{t-1})I,\sqrt{1-{\alpha }_{t-1}}{ϵ}_{t-1}\to Y\sim N(0,\sqrt{1-{\alpha }_t}I)$

令$Z=X+Y$，则有$Z\sim N(0,(1-{\alpha }_t{\alpha }_{t-1})I)$

$X+Y=\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{ϵ}}_{t-2},{\overline{ϵ}}_{t-2}\sim N(0,1)$

$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{(1-{\alpha }_{t-1}){\alpha }_t}{ϵ}_{t-2}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+X+Y\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{ϵ}}_{t-2}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}\dots {\alpha }_1}{x}_0+\sqrt{1-{\alpha }_t{\alpha }_{t-1}\dots {\alpha }_1}ϵ\\ & =\sqrt{{\overline{\alpha }}_t{x}_0}+\sqrt{1-{\overline{\alpha }}_t}ϵ\end{array}$

令${\prod }_{i=1}^t{\alpha }_i={\overline{\alpha }}_t$，则有①式$x_t=\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ$

注：各$ϵ$是相互独立但取自相同分布的标准正态随机变量

①式可用于在任何时间步骤直接对$x_t$采样

Reverse Diffusion Process反向扩散过程

该过程无法通过反转噪声实现，而需要训练神经网络approximated distribution$p_{\theta }(x_{t-1}|x_t)$来近似target distribution$q(x_{t-1}|x_t)$，其服从正态分布

目标分布target distribution：$q(x_{t-1}|x_t)=N(x_{t-1};\overline{\mu }(x_t,x_0),{\overline{\beta }}_{t}I)$

近似分布approximated distribution：$p_{\theta }(x_{t-1}|x_t)=N(x_{t-1};{\mu }_{\theta }(x_t,t),{\sum }_{\theta }(x_t,t))$

$\{\begin{array}{l}\mu (x_t,t):={\overline{\mu }}_t(x_t,x_0)\\ {\sum }_{\theta }(x_t,t):={\beta }_{t}I\end{array}$

loss function损失函数

$LOSS=-\log (p_{\theta }(x_0))$

其中$p_{\theta }(x_0)$取决于$x_1,x_2,\dots ,x_T$，较为棘手，因此采用与VAE中的类似的方法，通过优化下阶，简介优化不可处理的损失函数，关于KL LOSS与VAE loss function详见VAE模型中的损失函数

Loss Function of VAE model VAE模型的损失函数

$LOSS=MSE(X,X^{\prime })+KL(N(\mu ,{\sigma }^2),N(0,1))$

首项$MSE(X,X^{\prime })$称reconstruct loss 重构损失，反应生成结果与输入之间的差异

次项$KL(N(\mu ,{\sigma }^2),N(0,1))$称KL lossKL散度正则项：在去除reconstruct loss的vae中其编码空间会呈现标准正态分布，故称其作用：使得编码器生成的隐藏变量尽可能符合标准正态分布

已知kl loss非负，故

$-\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)\le E_q[\frac{\log q(x_{1:T}|x_0)}{p_{\theta }(x_{1:T})}]$tip:此处引入变分下界，详见变分推断

展开有$-\log p_{\theta }(x_0)\le E_q[D_{KL}(q(x_T|x_0)||p_{\theta }(x_T))+{\sum }_{t=2}^T{D}_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))]-\log p_{\theta }(x_0|x_1)$

1. 常数项$E_q[D_{KL}(q(x_T|x_0)||p_{\theta }(x_T))$：q无可学习参数，q为高斯噪声概率，故本项为可忽略的常数项
2. 逐步去噪项${\sum }_{t=2}^T{D}_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))]$
   1. 真实分布$q(x_{t-1}|x_t,x_0)=N(x_{t-1};\tilde{\mu }(x_t,x_0),{\overline{\beta }}_{t}I)$
      • ${\tilde{\mu }}_t(x_t,x_0)=\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{x}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}{x}_0$
        • 又$x_0=\frac{1}{\sqrt{{\overline{\alpha }}_t}}(x_t-\sqrt{1-\overline{\alpha })t}{ϵ}_t)$
        • 故${\tilde{\mu }}_t(x_t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_t)$
      • ${\overline{\beta }}_t=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}\beta )t$
   2. 预测分布$p_{\theta }(x_{t-1}|x_t))=N(x_{t-1};{\mu }_{\theta }(x_t,t),{\beta }_{t}I)$
      • ${\mu }_{\theta }(x_t,t)$是由神经网络训练的部分
      • 近似均值${\mu }_{\theta }$设置为与真实分布的均值${\overline{\mu }}_t$（即目标均值）相同的形式
        • target mean ${\overline{\mu }}_t(x_t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_t)$
        • approximated mean${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }(x_t,t))$
      • 二者的比较可以使用MSE进行
      • 
      • 
3. 重构项$\log p_{\theta }(x_0|x_1)$
   • 本项可以忽略，由于可用$L_{t-1}$中相同神经网络进行拟合，且忽略使得样本质量更好实施
     • 

kl loss KL散度正则项

关于KL散度的解释

1. 离散的kl loss
   • $D_{KL}(P||Q)={\sum }_{i=1}^N[p(x_i)\log p(x_i)-p(x_i)\log q(x_i)]$
   • P为真实时间的概率分布，Q为理论拟合出的该事件的概率分布
   • 两个概率分布之间的KL散度可以表达两个分布的相似性，例子如下
   • 
2. 连续的kl loss
   • $D_{KL}(P||Q)={\int }_{-\infty }^{+\infty }p(x)\log (\frac{p(x)}{q(x)})dx$
   • 基本性质：非负、仿射变换不变性、非对易性、一定条件下$D_{KL}(P||Q)$
     • 仿射变换不变性
     • $\because y=ax+b\therefore D_{KL}(P(x)||Q(x))=D_{KL}(P(x)||Q(y))$
3. 多元正态分布的kl loss
   • 设$P_1(x)=N({\mu }_1,{\sigma }_1^2),P_2(x)=N({\mu }_2,{\sigma }_2^2)$
   • 则有$KL(N(\mu ,{\sigma }_1^2),N({\mu }_2,{\sigma }_2^2))=-\frac{1}{2}+\log (\frac{{\sigma }_2}{{\sigma }_1})+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}$
   • 若修改P2为$P_2(x)=N(0,1)$，则有$KL(N(\mu ,{\sigma }_1^2),N(0,1))=-\frac{1}{2}({\mu }_1^2+{\sigma }_1^2-2\log ({\sigma }_1)-1)$

Variational Lower Bound变分下界

this part is mainly about variational inference.

变分推断可以针对贝叶斯统计中对于后验概率的计算，MCMC马尔科夫链蒙特卡洛方法对于大量数据计算较慢，而变分推断提供了更快更简单的近似推断方法

假定观测量$x=x_{1:n}$，隐藏变量$z=z_{1:m}$，推断问题需要推测后验条件概率分布$p(z|x)$。变分法的基本思想是将该问题转化为优化问题。首先引入一族关于隐藏变量的近似概率分布Q，需要从中寻找一个与真实分布的KL Divergence的最小的分布，即$q^{\ast }(z)=argmi{n}_{q(z\in Q)}KL(q(z)||p(z|x))$。可用$q^{\ast }(z)$近似代替真实后验分布。将变分推断转换为求极值，可通过选择更合适的Q使得密度分布足够灵活，同时足够简单适于优化

注：

• 信息熵公式$H(x)=-{\sum }_{i=1}^{n}p(x_i)\log (p(x_i))$
  • 其中$p(x_i)$代表随机事件的概率
• 对于同一个随机变量的两种不同分布P(x)和Q(x)可以使用KL Divergence来描述差异
  • $D_{KL}(P||Q)=E_{xP}[\log \frac{P(x)}{Q(x)}]=E_{xP}[\log P(x)-\log Q(x)]$

变分推断详述

1. from KL to ELBO
   • $\begin{array}{rl}KL(q(z)||p(z|x))& =E[\log q(z)]-E[\log q(z|x)]\\ & =E[\log q(z)]-E[\log p(z,x)]+\log p(x)\end{array}$
   • 最后一项为$\log p(x)$，属于边缘分布问题难以求解。由于无法直接计算kl，所以改变优化目标维与KL前两项相关量ELBO(evidence lower bound)
   • KaTeX parse error: Expected 'EOF', got '&' at position 8: ELBO(q)&̲=E[\log{p(\vec{…
   • 上式为ELBO的公式，可以看作$\log p(z)-KL$
2. ELBO各项直观解释
   • $\begin{array}{rl}ELBO(q)& =E[\log (z)]+E[\log p(x|z)]-E[\log q(z)]\\ & =E[\log p(x|z)]-KL(q(z)||p(z))\end{array}$
   • 可将ELBO展开为首行的形式，根据首行可以转换为次行的形式
   • 通过次行可以更直观的理解ELBO的作用。
     • 首项为期望项，促使模型将隐藏变量集中于可解释的观察量配置上
     • 次项为隐藏变量变分分布与先验分布的KL divergence的相反数，促使变分分布接近于先验分布
3. ELBO名称来源
   • $\because ELBO(q)=\log p(x)-KL(q(z)||p(z|x))$
   • $\therefore \log p(x)=KL(q(z)||p(z|x))+ELBO(q)$
   • $\because KL(q(z)||p(z|x))\ge 0$
   • $\therefore \log p(x)\ge ELBO(q)$
   • ELBO提供了$\log p(x)$的下限。当$q(z)=p(z|x)$时，有$\log p(x)=ELBO(q)$

近似概率分布的选择

• 可以选择平均场变分族，其假设隐藏变量间相互独立
• $q(z)={\prod }_{j=1}^m{q}_j(z_j)$

根据ELBO及平均场假设，可用coordinate ascent variational inference（坐标上升变分推断CAVI）

利用条件概率分布的链式法则有$p(z_{1:m},x_{1:n})=p(x_{1:n}){\prod }_{j=1}^{m}p(z_j|z_{1:(j-1)},x_{1:n})$

变分分布的期望为$E[\log q(z_{1:m})]={\sum }_{j=1}^m{E}_j[\log q(z_j)]$

此时，代入ELBO有

$\begin{array}{rl}ELBO& =E[\log (z)]+E[\log p(x|z)]-E[\log q(z)]\\ & =\log p(x_{1:n})+\sum _{j=1}^m{E}_j[\log q(z_j)]-E_j[\log q(z_j)]\end{array}$

对$z_k$求导并令其为零有$\frac{dELBO}{dq(z_k)}=E_{-k}[\log p(z_k|z_{-k},x)]-\log q(z_k)-1=0$

由CAVI有更新形式$q^{\ast }(z_k)$ ∝$E[\log (conditional)]$

UNet

每个训练轮次中各样本随机取步长t，各步长组合为嵌入向量，噪声依据步长生成

训练步骤

1. 随机选取一个步长样本并编码
2. 根据步长合成原图与高斯噪声为噪声图像
3. 输入噪声图像和嵌入向量，经UNet处理后输出预测图像，将预测噪声与实际噪声比较

训练过程

• 重复以下直至契合
  • x0为原始数据
  • t为从1-T中抽取的一个数字
  • $ϵ$为正态分布N(0,1)
  • 做梯度下降，对象为预测值与实际值方差

反向扩散过程

• xT为由N(0,1)生成噪声图像
• for t from T to 1:
  • z由N(0,1)生成
  • 若t>1，否则z置零
  • $x_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }(x_t,t))+{\sigma }_{t}z$
    • 处理前图像减去噪声图像生成去噪图像
    • $x_{t-1}$去噪图像
    • $x_t$处理前图像
    • ${ϵ}_{\theta }(x_t,t)$噪声图像
• 返回x0
  • x0:根据学习的均值${\mu }_{\theta }(x,1)$生成的预测去噪图像

扩散模型的速度问题

扩散采样会迭代的项UNet提供完整尺寸图象获得结果，总扩散步数T和图像尺寸较大时较慢，由此引入稳定扩散stable diffusion

Stable Diffusion稳定扩散

潜在空间

训练自编码器E，学习将图像数据降维成潜在数据，然后用解码器D，将潜在数据解码成图像。降噪过程将对潜在空间进行

相当于由处理三维空间，转换为处理三维空间的照片

潜在空间的扩散

正向扩散、反向扩散过程的操作内容同扩散模型

条件作用/调节

该模型在训练后可将文字编码后作为嵌入文本向量（额外输入），再输入随机潜在图像和步长嵌入向量，经扩散后解码成与文本有关的图像

通过使用交叉注意机制增强其去噪UNet，将内部扩散器转变为适应性(condition)图像生成器，通过开关调节输入，控制功能切换

• 非空间对齐输入（文本）：用语言模型转换为嵌入向量，然后映射至UNet
• 空间对齐输入：使用concat链接

训练

$L_{LDM}=E_{t,z_0,ϵ,y}[||ϵ-{ϵ}_{\theta }(z_t,t,{\tau }_{\theta }(y))|{|}^2]$

• ${\tau }_{\theta }(y)$，适应性输入
• $z_t=\sqrt{(}{\overline{\alpha }}_t)z_0+\sqrt{(}1-{\overline{\alpha }}_t)ϵ$
  • $z_0=E(x_0)$

架构比较

扩散模型框架

稳定扩散模型框架

与扩散模型相比，

• 输入为潜在数据，而采用潜在数据进行训练，训练速度更快
• 增加适应性输入，更富功能性

参考文章

1. deephub,“Diffusion 和Stable Diffusion的数学和工作原理详细解释”,https://zhuanlan.zhihu.com/p/597924053
2. 川陀学者,“变分推断——深度学习第十九章”,https://zhuanlan.zhihu.com/p/49401976
3. David M. Blei, Alp Kucukelbir, Jon D. McAuliffe,“Variational Inference: A Review for Statisticians”，arXiv:1601.00670
4. 捡到一束光,“关于KL散度（Kullback-Leibler Divergence）的笔记”,https://zhuanlan.zhihu.com/p/438129018
5. 哇哦,“VAE模型解析（loss函数，调参…）”,https://zhuanlan.zhihu.com/p/578619659