【论文阅读】LatentConsistencyModels(LDMs)、LCM-LoRa

文章目录 IntroductionPreliminariesDiffusion ModelsConsistency Models Latent Consistency ModelsConsistency Distillation in the Latent SpaceOne-Stage Guided Distillation by Solving Augmented PF-ODEAccelerating Distillation with Skipping Time StepsLatent Consistency Fine-tuing for Customized Dataset ExperimentLCM-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σt​g2(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σt​g2(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)(αt​z−σ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+1​tn​​(f(t)zt​+2σt​g2(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​​αtn​​​ztn+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，收敛更快，训练消耗显著降低.