Diffusion算法可以有多个角度进行理解,不同的理解方式只是对目标函数进行了不同的解释。其主体思想是不变的,可以归纳为:
- 训练时通过图片逐步添加噪声,变为一个纯噪声。然后学习每一步的噪声。
- 推理时给定一个随机噪声图片,然后通过学习到的噪声生成一个新的图片
目标
目标是已知上一步图像时,下一步图像的分布是什么。
每一步的图片用
x
0
,
x
1
,
.
.
.
,
x
T
x_0, x_1, ..., x_T
x0,x1,...,xT来表示,其中
x
0
x_0
x0是原图,
x
T
x_T
xT是纯噪声。它们的关系是:
x
t
=
α
t
x
t
−
1
+
1
−
α
t
ϵ
with
ϵ
∼
N
(
ϵ
;
0
,
I
)
\begin{align} \boldsymbol{x}_t = \sqrt{\alpha_t}\boldsymbol{x}_{t-1} + \sqrt{1 - \alpha_t}\boldsymbol{\epsilon} \quad \text{with } \boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{\epsilon}; \boldsymbol{0}, \textbf{I}) \end{align}
xt=αtxt−1+1−αtϵwith ϵ∼N(ϵ;0,I)
其中:
- ϵ \epsilon ϵ:是噪声,用一个均值为0,方差为I的高斯分布表示
- α t \alpha_t αt:是一个常数,只和t相关
为什么使用这个式子?可以看出,后一步的图片其实是前一步图片 x t − 1 x_{t-1} xt−1和另外一个噪声 ϵ \epsilon ϵ加权求和得到的。
我们需要把 x t − 1 x_{t-1} xt−1用 x t x_t xt表示,然后一步一步就可以推到 x 0 x_0 x0了:
这时候可能会有人想:上面那个式子不是有 x t − 1 x_{t-1} xt−1和 x t x_t xt吗?直接用上面那个式子不就可以了。
事实上, x t − 1 x_{t-1} xt−1和 x t x_t xt都是随机变量,可以进行恒等变换,但是算出来的仍然是一个随机变量。我们必须知道随机变量的分布才可以。
因此我们需要知道的其实是已知 x t x_t xt时 x t − 1 x_{t-1} xt−1的分布: q ( x t − 1 ∣ x t , x 0 ) q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) q(xt−1∣xt,x0),而这个值可以用贝叶斯公式变换为:
q
(
x
t
−
1
∣
x
t
,
x
0
)
=
q
(
x
t
∣
x
t
−
1
,
x
0
)
q
(
x
t
−
1
∣
x
0
)
q
(
x
t
∣
x
0
)
q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) }
q(xt−1∣xt,x0)=q(xt∣xt−1,x0)q(xt∣x0)q(xt−1∣x0)
这个式子中的值都是已知的,因为
x
t
x_{t}
xt和
x
t
−
1
x_{t-1}
xt−1的递推关系是已知的,因此可以不断地带入,然后使用
x
0
x_0
x0来表示
x
t
x_{t}
xt
具体如下:
x
t
=
α
ˉ
t
x
0
+
1
−
α
ˉ
t
ϵ
0
\begin{align} \boldsymbol{x}_t = \sqrt{\bar\alpha_t}\boldsymbol{x}_0 + \sqrt{1 - \bar\alpha_t}\boldsymbol{\boldsymbol{\epsilon}}_0 \\ \end{align}
xt=αˉtx0+1−αˉtϵ0
其中
α
t
ˉ
\bar{\alpha_t}
αtˉ指的是累乘:
α
1
⋅
α
2
⋅
.
.
.
⋅
α
t
\alpha_1\cdot \alpha_2\cdot ...\cdot\alpha_t
α1⋅α2⋅...⋅αt
上面的式子其实就是三个高斯分布相乘除,那么通过代入高斯分布的公式,然后经过一通计算以后,可以获得 q ( x t − 1 ∣ x t , x 0 ) q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) q(xt−1∣xt,x0),它的值如下:
q ( x t − 1 ∣ x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 ∣ x 0 ) q ( x t ∣ x 0 ) ∝ exp ( − 1 2 ( ( x t − α t x t − 1 ) 2 β t + ( x t − 1 − α ˉ t − 1 x 0 ) 2 1 − α ˉ t − 1 − ( x t − α ˉ t x 0 ) 2 1 − α ˉ t ) ) = exp ( − 1 2 ( x t 2 − 2 α t x t x t − 1 + α t x t − 1 2 β t + x t − 1 2 − 2 α ˉ t − 1 x 0 x t − 1 + α ˉ t − 1 x 0 2 1 − α ˉ t − 1 − ( x t − α ˉ t x 0 ) 2 1 − α ˉ t ) ) = exp ( − 1 2 ( ( α t β t + 1 1 − α ˉ t − 1 ) x t − 1 2 − ( 2 α t β t x t + 2 α ˉ t − 1 1 − α ˉ t − 1 x 0 ) x t − 1 + C ( x t , x 0 ) ) ) \begin{aligned} q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) &= q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) } \\ &\propto \exp \Big(-\frac{1}{2} \big(\frac{(\mathbf{x}_t - \sqrt{\alpha_t} \mathbf{x}_{t-1})^2}{\beta_t} + \frac{(\mathbf{x}_{t-1} - \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0)^2}{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big) \\ &= \exp \Big(-\frac{1}{2} \big(\frac{\mathbf{x}_t^2 - 2\sqrt{\alpha_t} \mathbf{x}_t \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \alpha_t} \color{red}{\mathbf{x}_{t-1}^2} }{\beta_t} + \frac{ \color{red}{\mathbf{x}_{t-1}^2} \color{black}{- 2 \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0} \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \bar{\alpha}_{t-1} \mathbf{x}_0^2} }{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big) \\ &= \exp\Big( -\frac{1}{2} \big( \color{red}{(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}})} \mathbf{x}_{t-1}^2 - \color{blue}{(\frac{2\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{2\sqrt{\bar{\alpha}_{t-1}}}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)} \mathbf{x}_{t-1} \color{black}{ + C(\mathbf{x}_t, \mathbf{x}_0) \big) \Big)} \end{aligned} q(xt−1∣xt,x0)=q(xt∣xt−1,x0)q(xt∣x0)q(xt−1∣x0)∝exp(−21(βt(xt−αtxt−1)2+1−αˉt−1(xt−1−αˉt−1x0)2−1−αˉt(xt−αˉtx0)2))=exp(−21(βtxt2−2αtxtxt−1+αtxt−12+1−αˉt−1xt−12−2αˉt−1x0xt−1+αˉt−1x02−1−αˉt(xt−αˉtx0)2))=exp(−21((βtαt+1−αˉt−11)xt−12−(βt2αtxt+1−αˉt−12αˉt−1x0)xt−1+C(xt,x0)))
上面这个高斯分布的均值和方差可以计算如下( β t = 1 − α t \beta_t=1-\alpha_t βt=1−αt):
β ~ t = 1 − α ˉ t − 1 1 − α ˉ t ⋅ β t μ ~ t ( x t , x 0 ) = α t ( 1 − α ˉ t − 1 ) 1 − α ˉ t x t + α ˉ t − 1 β t 1 − α ˉ t x 0 \begin{aligned} \tilde{\beta}_t &= \color{green}{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t} \\ \tilde{\boldsymbol{\mu}}_t (\mathbf{x}_t, \mathbf{x}_0) &= \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} \mathbf{x}_0\\ \end{aligned} β~tμ~t(xt,x0)=1−αˉt1−αˉt−1⋅βt=1−αˉtαt(1−αˉt−1)xt+1−αˉtαˉt−1βtx0
分析一下就可以知道,当 x t x_{t} xt已知时,其实这个 x t − 1 x_{t-1} xt−1的分布是已知的。有人问,那么均值中还有 x 0 x_0 x0怎么办呢,事实上可以通过上面那个递推公式的结果,使用 x t x_{t} xt把 x 0 x_0 x0表示出来,然后带入。带入后的结果如下:
μ q ( x t , x 0 ) = 1 α t x t − 1 − α t 1 − α ˉ t α t ϵ 0 β ~ t = 1 − α ˉ t − 1 1 − α ˉ t ⋅ β t \begin{align} \boldsymbol{\mu}_q(\boldsymbol{x}_t, \boldsymbol{x}_0) &= \frac{1}{\sqrt{\alpha_t}}\boldsymbol{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\boldsymbol{\epsilon}_0\\ \tilde{\beta}_t &= \color{green}{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t} \end{align} μq(xt,x0)β~t=αt1xt−1−αˉtαt1−αtϵ0=1−αˉt1−αˉt−1⋅βt
此时我们已经知道了 x t − 1 x_{t-1} xt−1的分布,只剩下一个是不知道的,就是噪声 ϵ 0 \epsilon_0 ϵ0。此时只需要用一个神经网络来估计每一步 t t t对应的 ϵ 0 \epsilon_0 ϵ0就可以了。
这也就是训练的过程:
Diffusion和VAE的关系:
VAE中引入了一个隐含的变量z,将p(x|y)看成了p(x|z)和q(z|y)这两个部分,然后获得了一个目标函数ELBO。下面的公式说明了ELBO是p(x)的下界,这个算法的目标就是最大化ELBO
log
p
(
x
)
=
log
p
(
x
)
∫
q
ϕ
(
z
∣
x
)
d
z
=
∫
q
ϕ
(
z
∣
x
)
log
p
(
x
)
d
z
=
E
q
ϕ
(
z
∣
x
)
[
log
p
(
x
)
]
=
E
q
ϕ
(
z
∣
x
)
[
log
p
(
x
,
z
)
p
(
z
∣
x
)
]
=
E
q
ϕ
(
z
∣
x
)
[
log
p
(
x
,
z
)
q
ϕ
(
z
∣
x
)
p
(
z
∣
x
)
q
ϕ
(
z
∣
x
)
]
=
E
q
ϕ
(
z
∣
x
)
[
log
p
(
x
,
z
)
q
ϕ
(
z
∣
x
)
]
+
E
q
ϕ
(
z
∣
x
)
[
log
q
ϕ
(
z
∣
x
)
p
(
z
∣
x
)
]
=
E
q
ϕ
(
z
∣
x
)
[
log
p
(
x
,
z
)
q
ϕ
(
z
∣
x
)
]
+
D
KL
(
q
ϕ
(
z
∣
x
)
∣
∣
p
(
z
∣
x
)
)
≥
E
q
ϕ
(
z
∣
x
)
[
log
p
(
x
,
z
)
q
ϕ
(
z
∣
x
)
]
\begin{align} \log p(\boldsymbol{x}) & = \log p(\boldsymbol{x}) \int q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})dz\\ & = \int q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})\log p(\boldsymbol{x})dz\\ & = \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\left[\log p(\boldsymbol{x})\right]\\ & = \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x}, \boldsymbol{z})}{p(\boldsymbol{z}\mid\boldsymbol{x})}\right]\\ & = \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x}, \boldsymbol{z})q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}{p(\boldsymbol{z}\mid\boldsymbol{x})q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\right]\\ & = \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\right] + \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\left[\log\frac{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}{p(\boldsymbol{z}\mid\boldsymbol{x})}\right]\\ & = \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\right] + \mathcal{D}_{\text{KL}}(q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x}) \mid\mid p(\boldsymbol{z}\mid\boldsymbol{x}))\\ & \geq \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\left[\log\frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z}\mid\boldsymbol{x})}\right] \end{align}
logp(x)=logp(x)∫qϕ(z∣x)dz=∫qϕ(z∣x)logp(x)dz=Eqϕ(z∣x)[logp(x)]=Eqϕ(z∣x)[logp(z∣x)p(x,z)]=Eqϕ(z∣x)[logp(z∣x)qϕ(z∣x)p(x,z)qϕ(z∣x)]=Eqϕ(z∣x)[logqϕ(z∣x)p(x,z)]+Eqϕ(z∣x)[logp(z∣x)qϕ(z∣x)]=Eqϕ(z∣x)[logqϕ(z∣x)p(x,z)]+DKL(qϕ(z∣x)∣∣p(z∣x))≥Eqϕ(z∣x)[logqϕ(z∣x)p(x,z)]
而VAE还有一个推广,就是Hierarchical VAE,表示中间的z不止一个,那么整个分布变成了p(x|z1), p(z1|z2), …, p(zt|y)。可以发现这个和扩散模型的思想是非常类似的。并且可以推导出来Hierarchical VAE的目标函数就是BLEO的形式是:
log
p
(
x
)
≥
E
q
(
x
1
:
T
∣
x
0
)
[
log
p
(
x
0
:
T
)
q
(
x
1
:
T
∣
x
0
)
]
=
[
t
]
E
q
(
x
1
∣
x
0
)
[
log
p
θ
(
x
0
∣
x
1
)
]
⏟
reconstruction term
−
D
KL
(
q
(
x
T
∣
x
0
)
∣
∣
p
(
x
T
)
)
⏟
prior matching term
−
∑
t
=
2
T
E
q
(
x
t
∣
x
0
)
[
D
KL
(
q
(
x
t
−
1
∣
x
t
,
x
0
)
∣
∣
p
θ
(
x
t
−
1
∣
x
t
)
)
]
⏟
denoising matching term
\begin{align} \log p(\boldsymbol{x}) &\geq \mathbb{E}_{q(\boldsymbol{x}_{1:T}\mid\boldsymbol{x}_0)}\left[\log \frac{p(\boldsymbol{x}_{0:T})}{q(\boldsymbol{x}_{1:T}\mid\boldsymbol{x}_0)}\right]\\ &= \begin{aligned}[t] \underbrace{\mathbb{E}_{q(\boldsymbol{x}_{1}\mid\boldsymbol{x}_0)}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_0\mid\boldsymbol{x}_1)\right]}_\text{reconstruction term} &- \underbrace{\mathcal{D}_{\text{KL}}(q(\boldsymbol{x}_T\mid\boldsymbol{x}_0) \mid\mid p(\boldsymbol{x}_T))}_\text{prior matching term} \\ &- \sum_{t=2}^{T} \underbrace{\mathbb{E}_{q(\boldsymbol{x}_{t}\mid\boldsymbol{x}_0)}\left[\mathcal{D}_{\text{KL}}(q(\boldsymbol{x}_{t-1}\mid\boldsymbol{x}_t, \boldsymbol{x}_0) \mid\mid p_{\boldsymbol{\theta}}(\boldsymbol{x}_{t-1}\mid\boldsymbol{x}_t))\right]}_\text{denoising matching term} \end{aligned} \end{align}
logp(x)≥Eq(x1:T∣x0)[logq(x1:T∣x0)p(x0:T)]=[t]reconstruction term
Eq(x1∣x0)[logpθ(x0∣x1)]−prior matching term
DKL(q(xT∣x0)∣∣p(xT))−t=2∑Tdenoising matching term
Eq(xt∣x0)[DKL(q(xt−1∣xt,x0)∣∣pθ(xt−1∣xt))]
然后扩散模型就选择了最后一项作为自己的目标函数。同时扩散模型假设了xt和xt-1之间的分布,然后把ELBO最后一项推呀推,推出最后需要学习一个噪声项。
总结一下VAE和Diffusion的区别:
- VAE的目标就是输入x,输出的y接近x的分布。做的方法是假设了一个中间变量z,然后问题变为计算两个条件概率:p(x|z)和p(z|y)。在传统VAE中这两个条件概率密度都是通过神经网络做的。
- Diffusion的目标和VAE挺类似的,但是没有用神经网络做,而是直接用一个线性的函数规定了z和x, y和z的关系(当然中间还有z1, z2, …)
- 对于VAE: 输入为x,输出为z的均值和方差: f ( x ) = ( σ , μ ) f(x)=(\sigma, \mu) f(x)=(σ,μ), f是一个神经网络
- 对于Diffusion: 规定了x和z的关系 z = α x + ( 1 − α ) ϵ z = \alpha x+(1-\alpha)\epsilon z=αx+(1−α)ϵ, ϵ \epsilon ϵ是一个高斯噪声,因此可以通过贝叶斯计算均值和方差。
- Diffusion的目标函数是VAE目标函数的一部分