【论文阅读】Latent Consistency Models (LDMs)、LCM-LoRa
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)}{2\sigma_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 ) \large\hat{\epsilon_{\theta}}(z_t,\omega,c,t)=(1+\omega)\epsilon_{\theta}(z_t,c,t)-\omega\epsilon_{\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}(\mu\theta^-+(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)}{2\sigma_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_t\hat{\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)}{2\sigma_t}\epsilon_{\theta}(z_t,c,t)\right)\mathrm{d}t\approx\Psi(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}}-1\right)\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的训练过程用如下算法所示：
多步采样算法如下：
![[图片]](https://img-blog.csdnimg.cn/direct/e7c78c490a114ab388fac40ebbf76bf2.png#pic_center)
Latent Consistency Fine-tuing for Customized Dataset
全量微调算法：
![[图片]](https://img-blog.csdnimg.cn/direct/ab73130a357d4b848282a0957dce22e7.png#pic_center)
Experiment
测试数据集使用 LAION-Aesthetic-6+ 和 LAION-Aesthetic-6.5+，teacher model 是 Stable Diffusion-v2.1.
![[图片]](https://img-blog.csdnimg.cn/direct/bd50bfd5c664484caa70b81ceddfec28.png#pic_center)
LCM的推理步数在
1
1
1到
4
4
4步的时候效果会比其他 baseline 方法好. 因为DPM和DPM++算实践中很常用的 ODE Solver，正常使用时推理步数在
20
20
20以上. 所以综合速度和质量，LCM表现不错.
训练时间 32 A100 GPU Hours
LCM-LoRA
![[图片]](https://img-blog.csdnimg.cn/direct/9d98e3d88dc24b109ca9f61971b381f1.png#pic_center)
原理：在原本的 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
免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！