文章目录
- Introduction
- Preliminaries
-
- [Diffusion Models](#Diffusion Models)
- [Consistency Models](#Consistency Models)
- [Latent Consistency Models](#Latent Consistency Models)
-
- [Consistency Distillation in the Latent Space](#Consistency Distillation in the Latent Space)
- [One-Stage Guided Distillation by Solving Augmented PF-ODE](#One-Stage Guided Distillation by Solving Augmented PF-ODE)
- [Accelerating Distillation with Skipping Time Steps](#Accelerating Distillation with Skipping Time Steps)
- [Latent Consistency Fine-tuing for Customized Dataset](#Latent Consistency Fine-tuing for Customized Dataset)
- Experiment
- 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的训练过程用如下算法所示：
多步采样算法如下：
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，收敛更快，训练消耗显著降低.