Introduction
- 提出
Latent Consistency Models (LCMs)
,图像生成速度更快、质量更好. - 提出一种简单高效的 one-stage guided consistency distillation 方法,用极少的采样步数蒸馏
Stable Diffusion
,进一步提出skipping-step
技术加快收敛过程. - 介绍针对
LCMs
的微调方法.
Preliminaries
Diffusion Models
使用 empirical PF-ODE 表示模型的逆扩散过程:
d
x
t
d
t
=
f
(
t
)
x
t
+
g
2
(
t
)
2
σ
t
ϵ
θ
(
x
t
,
t
)
large frac{mathrm{d}x_t}{mathrm{d}t}=f(t)x_t+frac{g^2(t)}{2sigma_t}epsilon_{theta}(x_t,t)
dtdxt=f(t)xt+2σtg2(t)ϵθ(xt,t)
对于 class-conditioned 扩散模型,Classifier-Free Guidance (CFG)
有效地提高了生成样本的质量,用
ω
omega
ω表示 CFG
系数,原始的噪声预测模型可以被替换为:
ϵ
θ
^
(
z
t
,
ω
,
c
,
t
)
=
(
1
+
ω
)
ϵ
θ
(
z
t
,
c
,
t
)
−
ω
ϵ
θ
(
z
t
,
∅
,
t
)
largehat{epsilon_{theta}}(z_t,omega,c,t)=(1+omega)epsilon_{theta}(z_t,c,t)-omegaepsilon_{theta}(z_t,varnothing,t)
ϵθ^(zt,ω,c,t)=(1+ω)ϵθ(zt,c,t)−ωϵθ(zt,∅,t)
Consistency Models
令
F
θ
(
x
,
t
)
F_{theta}(mathrm{x}, t)
Fθ(x,t)表示任意形式的神经网络,使用 sikp connection 可以将模型表示为:
f
θ
(
x
,
t
)
=
c
s
k
i
p
(
t
)
x
+
c
o
u
t
(
t
)
F
θ
(
x
,
t
)
large f_{theta}(mathrm{x}, t)=c_{skip}(t)mathrm{x}+c_{out}(t)F_{theta}(mathrm{x},t)
fθ(x,t)=cskip(t)x+cout(t)Fθ(x,t)
其中边界条件为
c
s
k
i
p
(
ϵ
)
=
1
c_{skip}(epsilon)=1
cskip(ϵ)=1,
c
o
u
t
(
ϵ
)
=
0
c_{out}(epsilon)=0
cout(ϵ)=0.
损失函数为:
L
C
D
N
(
θ
,
θ
−
;
ϕ
)
=
E
[
λ
(
t
n
)
d
(
f
θ
(
x
t
n
+
1
,
t
n
+
1
)
,
f
θ
−
(
x
^
t
n
ϕ
,
t
n
)
]
large mathcal{L}_{CD}^{N}(theta, theta^-;phi)=mathbb{E}left[lambda(t_n)d(f_{theta}(mathrm{x}_{t_{n+1}},t_{n+1}),f_{theta^-}(hat{mathrm{x}}_{t_n}^{phi}, t_n) right]
LCDN(θ,θ−;ϕ)=E[λ(tn)d(fθ(xtn+1,tn+1),fθ−(x^tnϕ,tn)]
θ
−
theta^-
θ−使用 EMA 更新,计算公式如下:
θ
−
←
s
t
o
p
g
a
r
d
(
μ
θ
−
+
(
1
−
μ
)
θ
)
large theta^- leftarrow mathrm{stopgard}(mutheta^-+(1-mu)theta)
θ−←stopgard(μθ−+(1−μ)θ)
x
^
t
n
ϕ
hat{mathrm{x}}_{t_n}^{phi}
x^tnϕ是从
x
t
n
+
1
mathrm{x}_{t_{n+1}}
xtn+1到
x
t
n
mathrm{x}_{t_{n}}
xtn的估计,计算公式如下:
x
^
t
n
ϕ
=
x
t
n
+
1
+
(
t
n
−
t
n
+
1
)
Φ
(
x
t
n
+
1
,
t
n
+
1
;
ϕ
)
large hat{mathrm{x}}_{t_n}^{phi}=mathrm{x}_{t_{n+1}} + (t_n-t_{n+1})Phi(mathrm{x}_{t_{n+1}}, t_{n+1};phi)
x^tnϕ=xtn+1+(tn−tn+1)Φ(xtn+1,tn+1;ϕ)
Latent Consistency Models
Consistency Distillation in the Latent Space
针对类似 Stable Diffusion
的隐空间上的条件扩散模型,其 PF- ODE 逆过程可以表示为:
d
z
t
d
t
=
f
(
t
)
z
t
+
g
2
(
t
)
2
σ
t
ϵ
θ
(
z
t
,
c
,
t
)
large frac{mathrm{d}z_t}{mathrm{d}t}=f(t)z_t+frac{g^2(t)}{2sigma_t}epsilon_{theta}(z_t,c,t)
dtdzt=f(t)zt+2σtg2(t)ϵθ(zt,c,t)
其中
z
t
z_t
zt是图像隐向量,
c
c
c是给定的条件. 类似CM
中的做法,引入
f
θ
:
(
z
t
,
c
,
t
)
↦
z
0
f_{theta}:(z_t,c,t)mapsto z_0
fθ:(zt,c,t)↦z0,将其参数化为:
f
θ
(
z
,
c
,
t
)
=
c
s
k
i
p
(
t
)
z
+
c
o
u
t
(
t
)
(
z
−
σ
t
ϵ
^
θ
(
z
,
c
,
t
)
α
t
)
large f_{theta}(z,c,t)=c_{skip}(t)z+c_{out}(t)left(frac{z-sigma_that{epsilon}_{theta}(z,c,t)}{alpha_{t}} right)
fθ(z,c,t)=cskip(t)z+cout(t)(αtz−σtϵ^θ(z,c,t))
具体的参数化形式由被蒸馏的扩散模型决定.
损失函数表示为:
L
C
D
(
θ
,
θ
−
;
Ψ
)
=
E
z
,
c
,
n
[
d
(
f
θ
(
z
t
n
+
1
,
c
,
t
n
+
1
)
,
f
θ
−
(
z
^
t
n
Ψ
,
c
,
t
n
)
]
large mathcal{L}_{CD}(theta,theta^-;Psi)=mathbb{E}_{z,c,n}left[d(f_{theta}(z_{t_{n+1}},c,t_{n+1}),f_{theta^-}(hat{z}_{t_n}^{Psi},c,t_n) right]
LCD(θ,θ−;Ψ)=Ez,c,n[d(fθ(ztn+1,c,tn+1),fθ−(z^tnΨ,c,tn)]
z
^
t
n
Ψ
hat{z}_{t_n}^{Psi}
z^tnΨ为
z
t
n
+
1
z_{t_{n+1}}
ztn+1到
z
t
n
z_{t_{n}}
ztn的估计,计算方法如下:
z
^
t
n
Ψ
−
z
t
n
+
1
=
∫
t
n
+
1
t
n
(
f
(
t
)
z
t
+
g
2
(
t
)
2
σ
t
ϵ
θ
(
z
t
,
c
,
t
)
)
d
t
≈
Ψ
(
z
t
n
+
1
,
t
n
+
1
,
t
n
,
c
)
large hat{z}_{t_n}^{Psi}-z_{t_{n+1}}=int_{t_{n+1}}^{t_n}left(f(t)z_t+frac{g^2(t)}{2sigma_t}epsilon_{theta}(z_t,c,t)right)mathrm{d}tapproxPsi(z_{t_{n+1}}, t_{n+1}, t_n, c)
z^tnΨ−ztn+1=∫tn+1tn(f(t)zt+2σtg2(t)ϵθ(zt,c,t))dt≈Ψ(ztn+1,tn+1,tn,c)
One-Stage Guided Distillation by Solving Augmented PF-ODE
使用CFG
,损失函数可以表示为:
L
C
D
(
θ
,
θ
−
;
Ψ
)
=
E
z
,
c
,
n
[
d
(
f
θ
(
z
t
n
+
1
,
ω
,
c
,
t
n
+
1
)
,
f
θ
−
(
z
^
t
n
Ψ
,
ω
,
c
,
t
n
)
]
large mathcal{L}_{CD}(theta,theta^-;Psi)=mathbb{E}_{z,c,n}left[d(f_{theta}(z_{t_{n+1}},omega,c,t_{n+1}),f_{theta^-}(hat{z}_{t_n}^{Psi},omega,c,t_n) right]
LCD(θ,θ−;Ψ)=Ez,c,n[d(fθ(ztn+1,ω,c,tn+1),fθ−(z^tnΨ,ω,c,tn)]
z
^
t
n
Ψ
hat{z}_{t_n}^{Psi}
z^tnΨ的计算方法更新为
z
^
t
n
Ψ
−
z
t
n
+
1
≈
(
1
+
ω
)
Ψ
(
z
t
n
+
1
,
t
n
+
1
,
t
n
,
c
)
−
Ψ
(
z
t
n
+
1
,
t
n
+
1
,
t
n
,
∅
)
large hat{z}_{t_n}^{Psi}-z_{t_{n+1}}approx(1+omega)Psi(z_{t_{n+1}}, t_{n+1}, t_n, c)-Psi(z_{t_{n+1}}, t_{n+1}, t_n, varnothing)
z^tnΨ−ztn+1≈(1+ω)Ψ(ztn+1,tn+1,tn,c)−Ψ(ztn+1,tn+1,tn,∅)
Accelerating Distillation with Skipping Time Steps
扩散模型例如Stable Diffusion
的总时间步长有
1000
1000
1000步,LCM在训练的采样需要覆盖这
1000
1000
1000步,既然相邻时间步之间的差值小,那么
f
θ
(
z
t
n
+
1
,
c
,
t
n
+
1
)
f_{theta}(z_{t_{n+1}},c,t_{n+1})
fθ(ztn+1,c,tn+1)和
f
θ
(
z
t
n
,
c
,
t
n
)
f_{theta}(z_{t_{n}},c,t_{n})
fθ(ztn,c,tn)之间的差距也小,这导致计算出来的损失小、收敛慢.
作者介绍了skipping-step
方法,原来度量时间步
t
n
+
1
t_{n+1}
tn+1和
t
n
t_n
tn间的差距,改为度量
t
n
+
k
t_{n+k}
tn+k和
t
n
t_n
tn间的差距. 至此,LCM训练的损失函数为
L
C
D
(
θ
,
θ
−
;
Ψ
)
=
E
z
,
c
,
n
[
d
(
f
θ
(
z
t
n
+
k
,
ω
,
c
,
t
n
+
k
)
,
f
θ
−
(
z
^
t
n
Ψ
,
ω
,
c
,
t
n
)
]
large mathcal{L}_{CD}(theta,theta^-;Psi)=mathbb{E}_{z,c,n}left[d(f_{theta}(z_{t_{n+k}},omega,c,t_{n+k}),f_{theta^-}(hat{z}_{t_n}^{Psi},omega,c,t_n) right]
LCD(θ,θ−;Ψ)=Ez,c,n[d(fθ(ztn+k,ω,c,tn+k),fθ−(z^tnΨ,ω,c,tn)]
z
^
t
n
Ψ
hat{z}_{t_n}^{Psi}
z^tnΨ中
Ψ
(
⋅
,
⋅
,
⋅
,
⋅
)
Psi(·,·,·,·)
Ψ(⋅,⋅,⋅,⋅)的计算方法对应跨
k
k
k步,作者分别使用了DDIM
、DPM-Solver
、DPM-Solver++
作为 PF-ODE solver,以DDIM
为例,其对应的
Ψ
(
⋅
,
⋅
,
⋅
,
⋅
)
Psi(·,·,·,·)
Ψ(⋅,⋅,⋅,⋅)计算方法为
Ψ
(
z
t
n
+
k
,
t
n
+
k
,
t
n
,
c
)
=
α
t
n
α
t
n
+
k
z
t
n
+
k
−
σ
t
n
(
σ
t
n
+
k
α
t
n
α
t
n
+
k
σ
t
n
−
1
)
ϵ
^
θ
(
z
t
n
+
k
,
c
,
t
n
+
k
)
−
z
t
n
+
k
large Psi(z_{t_{n+k}}, t_{n+k}, t_n, c)=frac{alpha_{t_n}}{alpha_{t_{n+k}}}z_{t_{n+k}}-sigma_{t_n}left(frac{sigma_{t_{n+k}}alpha_{t_n}}{alpha_{t_{n+k}}sigma_{t_n}}-1right)hat{epsilon}_{theta}(z_{t_{n+k}},c,t_{n+k})-z_{t_{n+k}}
Ψ(ztn+k,tn+k,tn,c)=αtn+kαtnztn+k−σtn(αtn+kσtnσtn+kαtn−1)ϵ^θ(ztn+k,c,tn+k)−ztn+k
再加入CFG
和skipping-step
之后,LCM的训练过程用如下算法所示:
多步采样算法如下:
Latent Consistency Fine-tuing for Customized Dataset
全量微调算法:
Experiment
测试数据集使用 LAION-Aesthetic-6+ 和 LAION-Aesthetic-6.5+,teacher model 是 Stable Diffusion-v2.1.
LCM的推理步数在
1
1
1到
4
4
4步的时候效果会比其他 baseline 方法好. 因为DPM
和DPM++
算实践中很常用的 ODE Solver,正常使用时推理步数在
20
20
20以上. 所以综合速度和质量,LCM表现不错.
训练时间 32 A100 GPU Hours
LCM-LoRA
原理:在原本的 Latent Diffusion Model (LDM)
中,可以使用 LoRa
训练一个额外结构附加到模型的 TextEncoder 和 Unet 中,做到模型的风格迁移. 即图中所示的
τ
′
mathbb{tau}’
τ′,它是原模型微调后额外结构的参数向量. LCM
的 backbone 和被它蒸馏模型的 backbone 结构是一致的,所以LCD
过程也可以视作对原模型的微调,所以也可以利用 LoRa
,在初始化 student Unet 之后,整个蒸馏过程只训练 LoRa
引入的额外结构,也就是获得
τ
L
C
M
mathbb{tau}_{mathrm{LCM}}
τLCM. 理论上可以结合
τ
′
mathbb{tau}’
τ′,最终做到既能加速生成,又能风格迁移.
LCD
过程仅微调 LoRa
,收敛更快,训练消耗显著降低.
原文地址:https://blog.csdn.net/weixin_43843835/article/details/135600936
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。
如若转载,请注明出处:http://www.7code.cn/show_58668.html
如若内容造成侵权/违法违规/事实不符,请联系代码007邮箱:suwngjj01@126.com进行投诉反馈,一经查实,立即删除!