Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference — part 2/4
Cleaned-text code points [15389, 31306). Bibliography omitted here; source.md includes it.
Utilizing image latent space in large-scale diffusion models like Stable Diffusion (SD) ( Rombach et al., 2022 ) has effectively enhanced image generation quality and reduced computational load. In SD, an autoencoder ( ℰ , 𝒟 \mathcal{E},\mathcal{D} ) is first trained to compress high-dim image data into low-dim latent vector z = ℰ ( x ) z=\mathcal{E}(x) , which is then decoded to reconstruct the image as x ^ = 𝒟 ( z ) \hat{x}=\mathcal{D}(z) . Training diffusion models in the latent space greatly reduces the computation costs compared to pixel-based models and speeds up the inference process; LDMs make it possible to generate high-resolution images on laptop GPUs. For LCMs, we leverage the advantage of the latent space for consistency distillation, contrasting with the pixel space used in CMs ( Song et al., 2023 ) . This approach, termed Latent Consistency Distillation (LCD) is applied to pre-trained SD, allowing the synthesis of high-resolution (e.g., 768 × \times 768) images in 1 ∼ \sim 4 steps. We focus on conditional generation. Recall that the PF-ODE of the reverse diffusion process ( Song et al., 2020b ; Lu et al., 2022a ) is d 𝒛 t d t = f ( t ) 𝒛 t + g 2 ( t ) 2 σ t ϵ θ ( 𝒛 t , 𝒄 , t ) , 𝒛 T ∼ 𝒩 ( 𝟎 , σ ~ 2 𝑰 ) , \frac{\mathrm{d}\bm{{\bm{z}}}_{t}}{\mathrm{~d}t}=f(t)\bm{{\bm{z}}}_{t}+\frac{g^{2}(t)}{2\sigma_{t}}\bm{\epsilon}_{\theta}\left(\bm{{\bm{z}}}_{t},\bm{c},t\right),\quad\bm{{\bm{z}}}_{T}\sim\mathcal{N}\left(\mathbf{0},\tilde{\sigma}^{2}\bm{I}\right), (8) where 𝒛 t \bm{{\bm{z}}}_{t} are image latents, ϵ θ ( 𝒛 t , 𝒄 , t ) \bm{\epsilon}_{\theta}\left(\bm{{\bm{z}}}_{t},\bm{c},t\right) is the noise prediction model, and 𝒄 \bm{c} is the given condition (e.g text). Samples can be drawn by solving the PF-ODE from T T to 0 0 . To perform LCD , we introduce the consistency function 𝒇 𝜽 : ( 𝒛 𝒕 , 𝒄 , t ) ↦ 𝒛 𝟎 \bm{f_{\theta}}:(\bm{z_{t}},\bm{c},t)\mapsto\bm{z_{0}} to directly predict the solution of PF-ODE (Eq. 8 ) for t = 0 t=0 . We parameterize 𝒇 𝜽 \bm{f_{\theta}} by the noise prediction model ϵ ^ θ \hat{\bm{\epsilon}}_{\theta} , as follows: 𝒇 𝜽 ( 𝒛 , 𝒄 , t ) = c skip ( t ) 𝒛 + c out ( t ) ( 𝒛 − σ t ϵ ^ θ ( 𝒛 , 𝒄 , t ) α t ) , ( ϵ -Prediction ) {\bm{f}}_{\bm{\theta}}({\bm{z}},{\bm{c}},t)=c_{\text{skip}}(t){\bm{z}}+c_{\text{out}}(t)\left(\frac{{\bm{z}}-\sigma_{t}\hat{\bm{\epsilon}}_{\theta}({\bm{z}},{\bm{c}},t)}{\alpha_{t}}\right),\quad(\bm{\epsilon}\text{-Prediction}) (9) where c skip ( 0 ) = 1 , c out ( 0 ) = 0 c_{\text{skip}}(0)=1,c_{\text{out}}(0)=0 and ϵ ^ θ ( 𝒛 , 𝒄 , t ) \hat{\bm{\epsilon}}_{\theta}({\bm{z}},{\bm{c}},t) is a noise prediction model that initializes with the same parameters as the teacher diffusion model. Notably, 𝒇 𝜽 \bm{f_{\theta}} can be parameterized in various ways, depending on the teacher diffusion model parameterizations of predictions (e.g., 𝒙 \bm{x} , ϵ \bm{\epsilon} ( Ho et al., 2020 ) , 𝒗 \bm{v} ( Salimans & Ho, 2022 ) ). We discuss other possible parameterizations in Appendix D .
We assume that an efficient ODE solver Ψ ( 𝒛 t , t , s , 𝒄 ) \Psi({\bm{z}}_{t},t,s,{\bm{c}}) is available for approximating the integration of the right-hand side of Eq equation 8 from time t t to s s . In practice, we can use DDIM ( Song et al., 2020a ) , DPM-Solver ( Lu et al., 2022a ) or DPM-Solver++ ( Lu et al., 2022b ) as Ψ ( ⋅ , ⋅ , ⋅ , ⋅ ) \Psi(\cdot,\cdot,\cdot,\cdot) . Note that we only use these solvers in training/distillation, not in inference. We will discuss these solvers further when we introduce the skipping-step technique in Sec. 4.3 . LCM aims to predict the solution of the PF-ODE by minimizing the consistency distillation loss ( Song et al., 2023 ) : ℒ 𝒞 𝒟 ( 𝜽 , 𝜽 − , Ψ ) = 𝔼 𝒛 , 𝒄 , n [ d ( 𝒇 𝜽 ( 𝒛 t n + 1 , 𝒄 , t n + 1 ) , 𝒇 𝜽 − ( 𝒛 ^ t n Ψ , 𝒄 , t n ) ) ] . \mathcal{L_{CD}}\left(\bm{\theta},\bm{\theta^{-}};\Psi\right)=\mathbb{E}_{{\bm{z}},\bm{c},n}\left[d\left(\bm{f_{\theta}}(\bm{{\bm{z}}}_{t_{n+1}},\bm{c},t_{n+1})\ ,\bm{f_{\theta^{-}}}(\bm{\hat{{\bm{z}}}}^{\Psi}_{t_{n}},\bm{c},t_{n})\right)\right].\vskip-3.61371pt (10) Here, 𝒛 ^ t n Ψ \bm{\hat{{\bm{z}}}}^{\Psi}_{t_{n}} is an estimation of the evolution of the PF-ODE from t n + 1 → t n t_{n+1}\rightarrow t_{n} using ODE solver Ψ \Psi : 𝒛 ^ t n Ψ − 𝒛 t n + 1 = ∫ t n + 1 t n ( f ( t ) 𝒛 t + g 2 ( t ) 2 σ t ϵ θ ( 𝒛 t , 𝒄 , t ) ) 𝑑 t ≈ Ψ ( 𝒛 t n + 1 , t n + 1 , t n , 𝒄 ) , \displaystyle\hat{{\bm{z}}}^{\Psi}_{t_{n}}-\bm{{\bm{z}}}_{t_{n+1}}=\int_{t_{n+1}}^{t_{n}}\left(f(t)\bm{{\bm{z}}}_{t}+\frac{g^{2}(t)}{2\sigma_{t}}\bm{\epsilon}_{\theta}\left(\bm{{\bm{z}}}_{t},\bm{c},t\right)\right)\mathrm{d}t\approx\Psi({\bm{z}}_{t_{n+1}},t_{n+1},t_{n},{\bm{c}}), (11) where the solver Ψ ( ⋅ , ⋅ , ⋅ , ⋅ ) \Psi(\cdot,\cdot,\cdot,\cdot) is used to approximate the integration from t n + 1 → t n t_{n+1}\rightarrow t_{n} .
### 4.2 One-Stage Guided Distillation by solving augmented PF-ODE
Classifier-free guidance (CFG) ( Ho & Salimans, 2022 ) is crucial for synthesizing high-quality text-aligned images in SD, typically needing a CFG scale ω \omega over 6 6 . Thus, integrating CFG into a distillation method becomes indispensable. Previous method Guided-Distill ( Meng et al., 2023 ) introduces a two-stage distillation to support few-step sampling from a guided diffusion model. However, it is computationally intensive (e.g. at least 45 A100 GPUs Days for 2-step inference, estimated in ( Liu et al., 2023 ) ). An LCM demands merely 32 A100 GPUs Hours training for 2-step inference, as depicted in Figure 1 . Furthermore, the two-stage guided distillation might result in accumulated error, leading to suboptimal performance. In contrast, LCMs adopt efficient one-stage guided distillation by solving an augmented PF-ODE. Recall the CFG used in reverse diffusion process: ϵ ~ θ ( 𝒛 t , ω , 𝒄 , t ) := ( 1 + ω ) ϵ θ ( 𝒛 t , 𝒄 , t ) − ω ϵ θ ( 𝒛 t , ∅ , t ) , \bm{\tilde{\epsilon}}_{\theta}\left(\bm{{\bm{z}}}_{t},\omega,\bm{c},t\right):=(1+\omega)\bm{\epsilon}_{\theta}\left(\bm{{\bm{z}}}_{t},\bm{c},t\right)-\omega\bm{\epsilon}_{\theta}\left(\bm{{\bm{z}}}_{t},\varnothing,t\right), (12) where the original noise prediction is replaced by the linear combination of conditional and unconditional noise and ω \omega is called the guidance scale . To sample from the guided reverse process, we need to solve the following augmented PF-ODE : (i.e., augmented with the terms related to ω \omega ) d 𝒛 t d t = f ( t ) 𝒛 t + g 2 ( t ) 2 σ t ϵ ~ θ ( 𝒛 t , ω , 𝒄 , t ) , 𝒛 T ∼ 𝒩 ( 𝟎 , σ ~ 2 𝑰 ) . \frac{\mathrm{d}\bm{{\bm{z}}}_{t}}{\mathrm{~d}t}=f(t)\bm{{\bm{z}}}_{t}+\frac{g^{2}(t)}{2\sigma_{t}}\bm{\tilde{\epsilon}}_{\theta}\left(\bm{{\bm{z}}}_{t},\omega,\bm{c},t\right),\quad\bm{{\bm{z}}}_{T}\sim\mathcal{N}\left(\mathbf{0},\tilde{\sigma}^{2}\bm{I}\right). (13) To efficiently perform one-stage guided distillation, we introduce an augmented consistency function 𝒇 𝜽 : ( 𝒛 𝒕 , ω , 𝒄 , t ) ↦ 𝒛 𝟎 \bm{f_{\theta}}:(\bm{z_{t}},\omega,\bm{c},t)\mapsto\bm{z_{0}} to directly predict the solution of augmented PF-ODE (Eq. 13 ) for t = 0 t=0 . We parameterize the 𝒇 𝜽 \bm{f_{\theta}} in the same way as in Eq. 9 , except that ϵ ^ θ ( 𝒛 , 𝒄 , t ) \hat{\bm{\epsilon}}_{\theta}({\bm{z}},{\bm{c}},t) is replaced by ϵ ^ θ ( 𝒛 , ω , 𝒄 , t ) \hat{\bm{\epsilon}}_{\theta}({\bm{z}},\omega,{\bm{c}},t) , which is a noise prediction model initializing with the same parameters as the teacher diffusion model, but also contains additional trainable parameters for conditioning on ω \omega . The consistency loss is the same as Eq. 10 except that we use augmented consistency function 𝒇 𝜽 ( 𝒛 𝒕 , ω , 𝒄 , t ) \bm{f_{\theta}}(\bm{z_{t}},\omega,\bm{c},t) . ℒ 𝒞 𝒟 ( 𝜽 , 𝜽 − , Ψ ) = 𝔼 𝒛 , 𝒄 , ω , n [ d ( 𝒇 𝜽 ( 𝒛 t n + 1 , ω , 𝒄 , t n + 1 ) , 𝒇 𝜽 − ( 𝒛 ^ t n Ψ , ω , ω , 𝒄 , t n ) ) ] \mathcal{L_{CD}}\left(\bm{\theta},\bm{\theta^{-}};\Psi\right)=\mathbb{E}_{\bm{{\bm{z}},c},\omega,n}\left[d\left(\bm{f_{\theta}}(\bm{{\bm{z}}}_{t_{n+1}},\omega,\bm{c},t_{n+1})\ ,\bm{f_{\theta^{-}}}(\bm{\hat{{\bm{z}}}}^{\Psi,\omega}_{t_{n}},\omega,\bm{c},t_{n})\right)\right] (14) In Eq 14 , ω \omega and n n are uniformly sampled from interval [ ω min , ω max ] [\omega_{\text{min}},\omega_{\text{max}}] and { 1 , … , N − 1 } \{1,\ldots,N-1\} respectively. 𝒛 ^ t n Ψ , ω \bm{\hat{{\bm{z}}}}^{\Psi,\omega}_{t_{n}} is estimated using the new noise model ϵ ~ θ ( 𝒛 t , ω , 𝒄 , t ) \bm{\tilde{\epsilon}}_{\theta}\left(\bm{{\bm{z}}}_{t},\omega,\bm{c},t\right) , as follows: 𝒛 ^ t n Ψ , ω − 𝒛 t n + 1 = \displaystyle\hat{{\bm{z}}}^{\Psi,\omega}_{t_{n}}-\bm{{\bm{z}}}_{t_{n+1}}= ∫ t n + 1 t n ( f ( t ) 𝒛 t + g 2 ( t ) 2 σ t ϵ ~ θ ( 𝒛 t , ω , 𝒄 , t ) ) 𝑑 t \displaystyle\int_{t_{n+1}}^{t_{n}}\left(f(t)\bm{{\bm{z}}}_{t}+\frac{g^{2}(t)}{2\sigma_{t}}\bm{\tilde{\epsilon}}_{\theta}\left(\bm{{\bm{z}}}_{t},\omega,\bm{c},t\right)\right)\mathrm{d}t (15) = \displaystyle= ( 1 + ω ) ∫ t n + 1 t n ( f ( t ) 𝒛 t + g 2 ( t ) 2 σ t ϵ θ ( 𝒛 t , 𝒄 , t ) ) 𝑑 t − ω ∫ t n + 1 t n ( f ( t ) 𝒛 t + g 2 ( t ) 2 σ t ϵ θ ( 𝒛 t , ∅ , t ) ) 𝑑 t \displaystyle(1+\omega)\int_{t_{n+1}}^{t_{n}}\left(f(t)\bm{{\bm{z}}}_{t}+\frac{g^{2}(t)}{2\sigma_{t}}\bm{\epsilon}_{\theta}\left(\bm{{\bm{z}}}_{t},\bm{c},t\right)\right)\mathrm{d}t-\omega\int_{t_{n+1}}^{t_{n}}\left(f(t)\bm{{\bm{z}}}_{t}+\frac{g^{2}(t)}{2\sigma_{t}}\bm{\epsilon}_{\theta}\left(\bm{{\bm{z}}}_{t},\varnothing,t\right)\right)\mathrm{d}t ≈ \displaystyle\approx ( 1 + ω ) Ψ ( 𝒛 t n + 1 , t n + 1 , t n , 𝒄 ) − ω Ψ ( 𝒛 t n + 1 , t n + 1 , t n , ∅ ) . \displaystyle(1+\omega)\Psi({\bm{z}}_{t_{n+1}},t_{n+1},t_{n},{\bm{c}})-\omega\Psi({\bm{z}}_{t_{n+1}},t_{n+1},t_{n},\varnothing). Again, we can use DDIM ( Song et al., 2020a ) , DPM-Solver ( Lu et al., 2022a ) or DPM-Solver++ ( Lu et al., 2022b ) as the PF-ODE solver Ψ ( ⋅ , ⋅ , ⋅ , ⋅ ) \Psi(\cdot,\cdot,\cdot,\cdot) .
### 4.3 Accelerating Distillation with Skipping Time Steps
Discrete diffusion models ( Ho et al., 2020 ; Song & Ermon, 2019 ) typically train noise prediction models with a long time-step schedule { t i } i \{t_{i}\}_{i} (also called discretization schedule or time schedule) to achieve high quality generation results. For instance, Stable Diffusion (SD) has a time schedule of length 1,000. However, directly applying Latent Consistency Distillation ( LCD ) to SD with such an extended schedule can be problematic. The model needs to sample across all 1,000 time steps, and the consistency loss attempts to aligns the prediction of LCM model 𝒇 𝜽 ( 𝒛 t n + 1 , 𝒄 , t n + 1 ) \bm{f_{\theta}}(\bm{{\bm{z}}}_{t_{n+1}},\bm{c},t_{n+1}) with the prediction 𝒇 𝜽 ( 𝒛 t n , 𝒄 , t n ) \bm{f_{\theta}}(\bm{{\bm{z}}}_{t_{n}},\bm{c},t_{n}) at the subsequent step along the same trajectory. Since t n − t n + 1 t_{n}-t_{n+1} is tiny, 𝒛 t n \bm{{\bm{z}}}_{t_{n}} and 𝒛 t n + 1 \bm{{\bm{z}}}_{t_{n+1}} (and thus 𝒇 𝜽 ( 𝒛 t n + 1 , 𝒄 , t n + 1 ) \bm{f_{\theta}}(\bm{{\bm{z}}}_{t_{n+1}},\bm{c},t_{n+1}) and 𝒇 𝜽 ( 𝒛 t n , 𝒄 , t n ) \bm{f_{\theta}}(\bm{{\bm{z}}}_{t_{n}},\bm{c},t_{n}) ) are already close to each other, incurring small consistency loss and hence leading to slow convergence. To address this issues, we introduce the skipping-step method to considerably shorten the length of time schedule (from thousands to dozens) to achieve fast convergence while preserving generation quality.
Consistency Models (CMs) ( Song et al., 2023 ) use the EDM ( Karras et al., 2022 ) continuous time schedule, and the Euler, or Heun Solver as the numerical continuous PF-ODE solver. For LCMs, in order to adapt to the discrete-time schedule in Stable Diffusion, we utilize DDIM ( Song et al., 2020a ) , DPM-Solver ( Lu et al., 2022a ) , or DPM-Solver++ ( Lu et al., 2022b ) as the ODE solver. ( Lu et al., 2022a ) shows that these advanced solvers can solve the PF-ODE efficiently in Eq. 8 . Now, we introduce the Skipping-Step method in Latent Consistency Distillation (LCD). Instead of ensuring consistency between adjacent time steps t n + 1 → t n t_{n+1}\rightarrow t_{n} , LCMs aim to ensure consistency between the current time step and k k -step away, t n + k → t n t_{n+k}\rightarrow t_{n} . Note that setting k k =1 reduces to the original schedule in ( Song et al., 2023 ) , leading to slow convergence, and very large k k may incur large approximation errors of the ODE solvers. In our main experiments, we set k k =20, drastically reducing the length of time schedule from thousands to dozens. Results in Sec. 5.2 show the effect of various k k values and reveal that the skipping-step method is crucial in accelerating the LCD process. Specifically, consistency distillation loss in Eq. 14 is modified to ensure consistency from t n + k t_{n+k} to t n t_{n} : ℒ 𝒞 𝒟 ( 𝜽 , 𝜽 − , Ψ ) = 𝔼 𝒛 , 𝒄 , ω , n [ d ( 𝒇 𝜽 ( 𝒛 t n + k , ω , 𝒄 , t n + k ) , 𝒇 𝜽 − ( 𝒛 ^ t n Ψ , ω , ω , 𝒄 , t n ) ) ] , \mathcal{L_{CD}}\left(\bm{\theta},\bm{\theta^{-}};\Psi\right)=\mathbb{E}_{\bm{{\bm{z}},c},\omega,n}\left[d\left(\bm{f_{\theta}}(\bm{{\bm{z}}}_{t_{n+k}},\omega,\bm{c},t_{n+k})\ ,\bm{f_{\theta^{-}}}(\bm{\hat{{\bm{z}}}}^{\Psi,\omega}_{t_{n}},\omega,\bm{c},t_{n})\right)\right], (16) with 𝒛 ^ t n Ψ , ω \bm{\hat{{\bm{z}}}}^{\Psi,\omega}_{t_{n}} being an estimate of 𝒛 t n {\bm{z}}_{t_{n}} using numerical augmented PF-ODE solver Ψ \Psi : 𝒛 ^ t n Ψ , ω ⟵ 𝒛 t n + k + ( 1 + ω ) Ψ ( 𝒛 t n + k , t n + k , t n , 𝒄 ) − ω Ψ ( 𝒛 t n + k , t n + k , t n , ∅ ) . \hat{{\bm{z}}}^{\Psi,\omega}_{t_{n}}\longleftarrow{\bm{z}}_{t_{n+k}}+(1+\omega)\Psi({\bm{z}}_{t_{n+k}},t_{n+k},t_{n},{\bm{c}})-\omega\Psi({\bm{z}}_{t_{n+k}},t_{n+k},t_{n},\varnothing). (17) The above derivation is similar to Eq. 15 . For LCM, we use three possible ODE solvers here: DDIM ( Song et al., 2020a ) , DPM-Solver ( Lu et al., 2022a ) , DPM-Solver++ ( Lu et al., 2022b ) , and we compare their performance in Sec 5.2 . In fact, DDIM ( Song et al., 2020a ) is the first-order discretization approximation of the DPM-Solver (Proven in ( Lu et al., 2022a ) ). Here we provide the detailed formula of the DDIM PF-ODE solver Ψ DDIM \Psi_{\text{DDIM}} from t n + k t_{n+k} to t n t_{n} . The formulas of the other two solver Ψ DPM-Solver \Psi_{\text{DPM-Solver}} , Ψ DPM-Solver++ \Psi_{\text{DPM-Solver++}} are provided in Appendix E . Ψ DDIM ( 𝒛 t n + k , t n + k , t n , 𝒄 ) = α t n α t n + k 𝒛 t n + k − σ t n ( σ t n + k ⋅ α t n α t n + k ⋅ σ t n − 1 ) ϵ ^ θ ( 𝒛 t n + k , 𝒄 , t n + k ) ⏟ DDIM Estimated 𝒛 t n − 𝒛 t n + k \displaystyle\Psi_{\text{DDIM}}({\bm{z}}_{t_{n+k}},t_{n+k},t_{n},{\bm{c}})=\underbrace{\frac{\alpha_{t_{n}}}{\alpha_{t_{n+k}}}{\bm{z}}_{t_{n+k}}-\sigma_{t_{n}}\left(\frac{\sigma_{t_{n+k}}\cdot\alpha_{t_{n}}}{\alpha_{t_{n+k}}\cdot\sigma_{t_{n}}}-1\right)\hat{\bm{\epsilon}}_{\theta}({\bm{z}}_{t_{n+k}},{\bm{c}},t_{n+k})}_{\text{DDIM Estimated}\ {\bm{z}}_{t_{n}}}-{\bm{z}}_{t_{n+k}} (18) We present the pseudo-code for LCD with CFG and skipping-step techniques in Algorithm 1 The modifications from the original Consistency Distillation (CD) algorithm in Song et al. (2023) are highlighted in blue. Also, the LCM sampling algorithm 3 is provided in Appendix B .
### 4.4 Latent Consistency Fine-tuning for customized dataset
Foundation generative models like Stable Diffusion excel in diverse text-to-image generation tasks but often require fine-tuning on customized datasets to meet the requirements of downstream tasks. We propose Latent Consistency Fine-tuning (LCF), a fine-tuning method for pretrained LCM. Inspired by Consistency Training (CT) ( Song et al., 2023 ) , LCF enables efficient few-step inference on customized datasets without relying on a teacher diffusion model trained on such data. This approach presents a viable alternative to traditional fine-tuning methods for diffusion models. The pseudo-code for LCF is provided in Algorithm 4 , with a more detailed illustration in Appendix C .
## 5 Experiment