Denoising diffusion probabilistic models (DDPMs)从马尔科夫链中采样生成样本,需要迭代多次,速度较慢。Denoising diffusion implicit models (DDIMs)的提出是为了在复用DDPM训练的网络的前提下,加速采样过程。
加速采样的基本思路是,原本的生成过程是从
[
T
,
⋯
,
1
]
[T,\cdots,1]
[T,⋯,1]的序列逐步采样,加速时考虑从子序列
{
τ
1
,
…
,
τ
S
}
,
τ
1
>
τ
2
>
⋯
>
τ
S
∈
[
1
,
T
]
\{\tau_1, \dots, \tau_S\},\tau_1 > \tau_2 > \dots > \tau_S \in [1, T]
{τ1,…,τS},τ1>τ2>⋯>τS∈[1,T]采样,通过跳步的方式减少采样的步数。比如DDPM网络原始训练包含1000步,但是采样时可以只从1000步中均匀的选出50步,用这50步采样出图像。
DDPM和DDIM都可以跳步采样,这个作者在实验中也进行了证明。DDIM的贡献主要是在复用DDPM训练的网络的前提下,提出了一种可以调节方差的生成形式,在步数较少的时候使用小的方差生成效果好。
DDIM论文中的符号和DDPM论文不同,本笔记中采用DDPM论文的符号。
引子
DDPM的优化目标是
L
VLB
=
E
q
(
x
0
:
T
)
[
log
q
(
x
1
:
T
∣
x
0
)
p
θ
(
x
0
:
T
)
]
=
E
q
[
D
KL
(
q
(
x
T
∣
x
0
)
∥
p
θ
(
x
T
)
)
⏟
L
T
+
∑
t
=
2
T
D
KL
(
q
(
x
t
−
1
∣
x
t
,
x
0
)
∥
p
θ
(
x
t
−
1
∣
x
t
)
)
⏟
L
t
−
1
−
log
p
θ
(
x
0
∣
x
1
)
⏟
L
0
]
\begin{aligned} L_\text{VLB} &= \mathbb{E}_{q(\mathbf{x}_{0:T})} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \\ &= \mathbb{E}_q [\underbrace{D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T))}_{L_T} + \sum_{t=2}^T \underbrace{D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))}_{L_{t-1}} \underbrace{- \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)}_{L_0} ] \end{aligned}
LVLB=Eq(x0:T)[logpθ(x0:T)q(x1:T∣x0)]=Eq[LT
DKL(q(xT∣x0)∥pθ(xT))+t=2∑TLt−1
DKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))L0
−logpθ(x0∣x1)]
其中
L
t
−
1
=
E
q
[
D
KL
(
q
(
x
t
−
1
∣
x
t
,
x
0
)
∥
p
θ
(
x
t
−
1
∣
x
t
)
)
]
L_{t-1}= \mathbb{E}_q [D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))]
Lt−1=Eq[DKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))]是DDPM网络优化的目标项。
注意
L
t
−
1
L_{t-1}
Lt−1只和边际分布
q
(
x
t
∣
x
0
)
q(\pmb{x}_t|\pmb{x}_0)
q(xt∣x0)有关,而不是联合分布
q
(
x
1
:
T
∣
x
0
)
q(\pmb{x}_{1:T}|\pmb{x}_0)
q(x1:T∣x0)。所以,我们可以定义一个更为灵活的推理过程,只要它的边际分布
q
(
x
t
∣
x
0
)
q(\pmb{x}_t|\pmb{x}_0)
q(xt∣x0)和DDPM一致,就可以复用DDPM优化的网络。
非马尔科夫的前向过程
DDPM中推理分布 q ( x 1 : T ∣ x 0 ) q(\mathbf x_{1:T}|\mathbf x_0) q(x1:T∣x0)(推理分布是从可观测变量 x 0 \mathbf x_0 x0推理隐变量 x 1 : T \mathbf x_{1:T} x1:T的分布)是固定的马尔科夫链(DDPM中要求 q ( x t ∣ x t − 1 , x 0 ) = q ( x t ∣ x t − 1 ) : = N ( 1 − β t x t − 1 , β t I ) q(\pmb{x}_t|\pmb{x}_{t-1},\pmb{x}_0) = q(\pmb{x}_t|\pmb{x}_{t-1}) := \mathcal{N}(\sqrt{1 - \beta_t}\pmb{x}_{t-1}, \beta_t \pmb{I}) q(xt∣xt−1,x0)=q(xt∣xt−1):=N(1−βtxt−1,βtI))。现在放宽限制,不要求前向过程是马尔科夫的,也就是对 q ( x t ∣ x t − 1 ) q(\pmb{x}_t|\pmb{x}_{t-1}) q(xt∣xt−1)不做形式要求。
作者定义由实向量
σ
∈
R
≥
0
T
\sigma \in \mathbb{R}^T_{\ge 0}
σ∈R≥0T索引的推理分布族Q:
q
σ
(
x
1
:
T
∣
x
0
)
:
=
q
σ
(
x
T
∣
x
0
)
∏
t
=
2
T
q
σ
(
x
t
−
1
∣
x
t
,
x
0
)
\begin{split} q_\sigma (\pmb{x}_{1:T}|\pmb{x}_0) := q_\sigma(\pmb{x}_T|\pmb{x}_0)\prod_{t=2}^Tq_\sigma(\pmb{x}_{t-1}|\pmb{x}_t, \pmb{x}_0) \end{split}
qσ(x1:T∣x0):=qσ(xT∣x0)t=2∏Tqσ(xt−1∣xt,x0)只要求满足边际分布和DDPM一样,即
q
σ
(
x
t
∣
x
0
)
=
N
(
α
ˉ
t
x
0
,
(
1
−
α
ˉ
t
)
I
)
q_\sigma(\pmb{x}_t|\pmb{x}_0) = \mathcal{N}(\sqrt{\bar{\alpha}_t}\pmb{x}_0,(1-\bar{\alpha}_t)\pmb{I})
qσ(xt∣x0)=N(αˉtx0,(1−αˉt)I)。
通过待定系数法(参考[1])可以得到形式更自由的生成过程:
q
σ
(
x
t
−
1
∣
x
t
,
x
0
)
:
=
N
(
α
ˉ
t
−
1
x
0
+
1
−
α
ˉ
t
−
1
−
σ
t
2
⋅
x
t
−
α
ˉ
t
x
0
1
−
α
ˉ
t
,
σ
t
2
I
)
q_\sigma(\pmb{x}_{t-1}|\pmb{x}_t, \pmb{x}_0) := \mathcal{N}(\sqrt{\bar{\alpha}_{t-1}}\pmb{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1}-\sigma_t^2}\cdot \frac{\pmb{x}_t - \sqrt{\bar{\alpha}_t}\pmb{x}_0}{\sqrt{1 - \bar{\alpha}_t}},\sigma_t^2 \pmb{I})
qσ(xt−1∣xt,x0):=N(αˉt−1x0+1−αˉt−1−σt2⋅1−αˉtxt−αˉtx0,σt2I)
对应的前向过程也是高斯分布,但前向过程变成了非马尔科夫的,因为每一步都依赖
x
0
\mathbf x_0
x0
q
σ
(
x
t
∣
x
t
−
1
,
x
0
)
=
q
σ
(
x
t
−
1
∣
x
t
,
x
0
)
q
σ
(
x
t
∣
x
0
)
q
σ
(
x
t
−
1
∣
x
0
)
q_\sigma(\pmb{x}_t|\pmb{x}_{t-1}, \pmb{x}_0) = \frac{q_\sigma(\pmb{x}_{t-1}|\pmb{x}_t, \pmb{x}_0)q_\sigma(\pmb{x}_t|\pmb{x}_0)}{q_\sigma(\pmb{x}_{t-1}|\pmb{x}_0)}
qσ(xt∣xt−1,x0)=qσ(xt−1∣x0)qσ(xt−1∣xt,x0)qσ(xt∣x0)
如下图所示,DDPM的推理过程是非马尔科夫的。
注意,DDIM构造的推理分布和DDPM不同,但和DDPM优化相同的优化目标。
反向生成过程
根据上面的推理过程,定义需要学习的生成过程
p
θ
(
x
0
:
T
)
p_\theta(\mathbf{x}_{0:T})
pθ(x0:T),该过程利用
q
σ
(
x
t
−
1
∣
x
t
,
x
0
)
q_\sigma(\pmb{x}_{t-1}|\pmb{x}_t,\pmb{x}_0)
qσ(xt−1∣xt,x0)。
直观地说,给定
x
t
\pmb{x}_t
xt,我们首先预测对应的
x
0
\pmb{x}_0
x0,然后使用我们定义的反向条件分布
q
σ
(
x
t
−
1
∣
x
t
,
x
0
)
q_\sigma(\pmb{x}_{t-1}|\pmb{x}_t, \pmb{x}_0)
qσ(xt−1∣xt,x0)获得
x
t
−
1
\pmb{x}_{t-1}
xt−1。
预测对应的
x
0
\pmb{x}_0
x0如下:
x
^
0
=
f
θ
(
t
)
(
x
t
)
:
=
x
t
−
1
−
α
ˉ
t
ϵ
θ
(
t
)
(
x
t
)
α
ˉ
t
\hat{\pmb{x}}_0 = f_\theta^{(t)}(\pmb{x}_t) := \frac{\pmb{x}_t - \sqrt{1-\bar{\alpha}_t} \pmb{\epsilon}_\theta^{(t)}(\pmb{x}_t)}{\sqrt{\bar{\alpha}_t}}
x^0=fθ(t)(xt):=αˉtxt−1−αˉtϵθ(t)(xt)使用预测的
x
^
0
\hat{\pmb{x}}_0
x^0通过
q
σ
(
x
t
−
1
∣
x
t
,
x
0
)
q_\sigma(\pmb{x}_{t-1}|\pmb{x}_t, \pmb{x}_0)
qσ(xt−1∣xt,x0)获得
x
t
−
1
\pmb{x}_{t-1}
xt−1如下:
x
^
t
−
1
=
α
ˉ
t
−
1
x
^
0
+
1
−
α
ˉ
t
−
1
−
σ
t
2
⋅
x
t
−
α
ˉ
t
x
^
0
1
−
α
ˉ
t
+
σ
t
z
,
z
∼
N
(
0
,
I
)
\hat{\pmb{x}}_{t-1} = \sqrt{\bar{\alpha}_{t-1}}\hat{\pmb{x}}_0 + \sqrt{1- \bar{\alpha}_{t-1}-\sigma_t^2}\cdot \frac{\pmb{x}_t - \sqrt{\bar{\alpha}_t}\hat{\pmb{x}}_0}{\sqrt{1 - \bar{\alpha}_t}} + \sigma_t \pmb{z}, \pmb{z} \sim \mathcal{N}(\pmb{0}, \pmb{I})
x^t−1=αˉt−1x^0+1−αˉt−1−σt2⋅1−αˉtxt−αˉtx^0+σtz,z∼N(0,I)写成使用
ϵ
θ
(
t
)
(
x
t
)
\pmb{\epsilon}_\theta^{(t)}(\pmb{x}_t)
ϵθ(t)(xt)的形式:
x
^
t
−
1
=
α
ˉ
t
−
1
x
t
−
1
−
α
ˉ
t
ϵ
θ
(
t
)
(
x
t
)
α
ˉ
t
+
1
−
α
ˉ
t
−
1
−
σ
t
2
⋅
ϵ
θ
(
t
)
(
x
t
)
+
σ
t
z
,
z
∼
N
(
0
,
I
)
\hat{\pmb{x}}_{t-1} = \sqrt{\bar{\alpha}_{t-1}} \frac{\pmb{x}_t -\sqrt{1-\bar{\alpha}_t}\pmb{\epsilon}_\theta^{(t)}(\pmb{x}_t)}{\sqrt{\bar{\alpha}_t}} + \sqrt{1- \bar{\alpha}_{t-1}-\sigma_t^2}\cdot \pmb{\epsilon}_\theta^{(t)}(\pmb{x}_t) + \sigma_t \pmb{z}, \pmb{z} \sim \mathcal{N}(\pmb{0}, \pmb{I})
x^t−1=αˉt−1αˉtxt−1−αˉtϵθ(t)(xt)+1−αˉt−1−σt2⋅ϵθ(t)(xt)+σtz,z∼N(0,I)
选择不同的
σ
\sigma
σ值会导致不同的生成过程,但它们使用相同的
ϵ
θ
\epsilon_{\theta}
ϵθ模型。
DDPM和DDIM对比
DDPM : q ( x t − 1 ∣ x t , x 0 ) = N ( α ˉ t − 1 β t 1 − α ˉ t x 0 + α t ( 1 − α ˉ t − 1 ) 1 − α ˉ t x t , β t ( 1 − α ˉ t − 1 ) ( 1 − α ˉ t ) I ) DDIM : q σ ( x t − 1 ∣ x t , x 0 ) : = N ( α ˉ t − 1 x 0 + 1 − α ˉ t − 1 − σ t 2 ⋅ x t − α ˉ t x 0 1 − α ˉ t , σ t 2 I ) \begin{split} \text{DDPM}:\ &q(\pmb{x}_{t-1}|\pmb{x}_t, \pmb{x}_0) = \mathcal{N}(\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\pmb{x}_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 -\bar{\alpha}_t}\pmb{x}_t, \frac{\beta_t(1-\bar{\alpha}_{t-1})}{(1-\bar{\alpha}_t)}\pmb{I})\\ \text{DDIM}:\ &q_\sigma(\pmb{x}_{t-1}|\pmb{x}_t, \pmb{x}_0) := \mathcal{N}(\sqrt{\bar{\alpha}_{t-1}}\pmb{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1}-\sigma_t^2}\cdot \frac{\pmb{x}_t - \sqrt{\bar{\alpha}_t}\pmb{x}_0}{\sqrt{1 - \bar{\alpha}_t}},\sigma_t^2 \pmb{I}) \end{split} DDPM: DDIM: q(xt−1∣xt,x0)=N(1−αˉtαˉt−1βtx0+1−αˉtαt(1−αˉt−1)xt,(1−αˉt)βt(1−αˉt−1)I)qσ(xt−1∣xt,x0):=N(αˉt−1x0+1−αˉt−1−σt2⋅1−αˉtxt−αˉtx0,σt2I)当 σ t = ( 1 − α ˉ t − 1 ) / ( 1 − α ˉ t ) ( 1 − α ˉ t / α ˉ t − 1 ) = β t ( 1 − α ˉ t − 1 ) ( 1 − α ˉ t ) \sigma_t = \sqrt{(1-\bar\alpha_{t-1})/(1-\bar\alpha_{t})}\sqrt{(1-\bar\alpha_{t}/\bar\alpha_{t-1})} = \sqrt{\frac{\beta_t(1-\bar{\alpha}_{t-1})}{(1-\bar{\alpha}_t)}} σt=(1−αˉt−1)/(1−αˉt)(1−αˉt/αˉt−1)=(1−αˉt)βt(1−αˉt−1)时,DDPM的 q ( x t − 1 ∣ x t , x 0 ) q(\pmb{x}_{t-1}|\pmb{x}_t, \pmb{x}_0) q(xt−1∣xt,x0)和DDIM的 q σ ( x t − 1 ∣ x t , x 0 ) q_\sigma(\pmb{x}_{t-1}|\pmb{x}_t,\pmb{x}_0) qσ(xt−1∣xt,x0)是一样的,生成过程变成和DDPM是一样的。另外,DDIM的前向过程变成了马尔科夫的。
当
σ
t
=
0
\sigma_t=0
σt=0时,随机噪声前的系数是0,
x
0
\mathbf x_0
x0和
x
T
\mathbf x_T
xT之间的关系是固定的,这属于隐概率模型(implicit probabilistic model)。
这时生成过程的每一步变为
x
^
t
−
1
=
α
ˉ
t
−
1
x
t
−
1
−
α
ˉ
t
ϵ
θ
(
t
)
(
x
t
)
α
ˉ
t
+
1
−
α
ˉ
t
−
1
⋅
ϵ
θ
(
t
)
(
x
t
)
=
1
α
t
(
x
t
−
(
1
−
α
ˉ
t
−
α
t
1
−
α
ˉ
t
−
1
)
ϵ
θ
(
t
)
(
x
t
)
)
\hat{\pmb{x}}_{t-1} = \sqrt{\bar{\alpha}_{t-1}} \frac{\pmb{x}_t - \sqrt{1-\bar{\alpha}_t}\pmb{\epsilon}_\theta^{(t)}(\pmb{x}_t)}{\sqrt{\bar{\alpha}_t}} + \sqrt{1- \bar{\alpha}_{t-1}}\cdot \pmb{\epsilon}_\theta^{(t)}(\pmb{x}_t)\\ =\frac{1}{\sqrt{\alpha_{t}}}(\pmb{x}_t - (\sqrt{1-\bar{\alpha}_t} - \sqrt{\alpha_{t}}\sqrt{1-\bar{\alpha}_{t-1}}) \pmb{\epsilon}_\theta^{(t)}(\pmb{x}_t))
x^t−1=αˉt−1αˉtxt−1−αˉtϵθ(t)(xt)+1−αˉt−1⋅ϵθ(t)(xt)=αt1(xt−(1−αˉt−αt1−αˉt−1)ϵθ(t)(xt))对比DDPM生成过程的每一步
x
^
t
−
1
=
1
α
t
(
x
t
−
1
−
α
t
1
−
α
ˉ
t
ϵ
θ
(
t
)
(
x
t
)
)
+
σ
t
z
\hat{\pmb{x}}_{t-1} = \frac{1}{\sqrt{\alpha_t}}(\pmb{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}}\pmb{\epsilon}_\theta^{(t)}(\pmb{x}_t)) + \sigma_t\pmb{z}
x^t−1=αt1(xt−1−αˉt1−αtϵθ(t)(xt))+σtz
[2]中给出了下面的分解:
x
t
−
1
=
α
ˉ
t
−
1
x
0
+
1
−
α
ˉ
t
−
1
ϵ
t
−
1
=
α
ˉ
t
−
1
x
0
+
1
−
α
ˉ
t
−
1
−
σ
t
2
ϵ
t
+
σ
t
ϵ
=
α
ˉ
t
−
1
x
0
+
1
−
α
ˉ
t
−
1
−
σ
t
2
x
t
−
α
ˉ
t
x
0
1
−
α
ˉ
t
+
σ
t
ϵ
\begin{aligned} \mathbf{x}_{t-1} &= \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1}}\boldsymbol{\epsilon}_{t-1} \\ &= \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \boldsymbol{\epsilon}_t + \sigma_t\boldsymbol{\epsilon} \\ &= \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \frac{\mathbf{x}_t - \sqrt{\bar{\alpha}_t}\mathbf{x}_0}{\sqrt{1 - \bar{\alpha}_t}} + \sigma_t\boldsymbol{\epsilon} \end{aligned}
xt−1=αˉt−1x0+1−αˉt−1ϵt−1=αˉt−1x0+1−αˉt−1−σt2ϵt+σtϵ=αˉt−1x0+1−αˉt−1−σt21−αˉtxt−αˉtx0+σtϵ分解的依据是不相关高斯随机变量
N
(
0
,
σ
1
2
I
)
\mathcal{N}(\mathbf{0}, \sigma_1^2\mathbf{I})
N(0,σ12I)和
N
(
0
,
σ
2
2
I
)
\mathcal{N}(\mathbf{0}, \sigma_2^2\mathbf{I})
N(0,σ22I)之和的分布是
N
(
0
,
(
σ
1
2
+
σ
2
2
)
I
)
\mathcal{N}(\mathbf{0}, (\sigma_1^2 + \sigma_2^2)\mathbf{I})
N(0,(σ12+σ22)I)。
在DDPM的逆向生成过程的每一步中,也可以认为是先估计
x
^
0
\hat{\pmb{x}}_0
x^0,再求
x
^
t
−
1
\hat{\pmb{x}}_{t-1}
x^t−1。
DDIM与DDPM的主要区别是DDIM构造了一种更自由的过程,通过
σ
\sigma
σ改变了方差的大小,同时改变均值,使其依然符合DDPM的边际分布
q
(
x
t
∣
x
0
)
q(\pmb{x}_t|\pmb{x}_0)
q(xt∣x0)。
参考资料
[1] https://kxz18.github.io/2022/06/21/DDIM/
[2] https://lilianweng.github.io/posts/2021-07-11-diffusion-models/