Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference — part 1/4
Cleaned-text code points [0, 15389). Bibliography omitted here; source.md includes it.
# Latent Consistency Models: Synthesizing High-Resolution Images with Few-step Inference
###### Abstract
Latent Diffusion models (LDMs) have achieved remarkable results in synthesizing high-resolution images. However, the iterative sampling process is computationally intensive and leads to slow generation. Inspired by Consistency Models ( Song et al., 2023 ) , we propose Latent Consistency Models ( LCMs ), enabling swift inference with minimal steps on any pre-trained LDMs, including Stable Diffusion ( Rombach et al., 2022 ) . Viewing the guided reverse diffusion process as solving an augmented probability flow ODE (PF-ODE), LCMs are designed to directly predict the solution of such ODE in latent space, mitigating the need for numerous iterations and allowing rapid, high-fidelity sampling. Efficiently distilled from pre-trained classifier-free guided diffusion models, a high-quality 768 × \times 768 2 ∼ \sim 4-step LCM takes only 32 A100 GPU hours for training. Furthermore, we introduce Latent Consistency Fine-tuning (LCF), a novel method that is tailored for fine-tuning LCMs on customized image datasets. Evaluation on the LAION-5B-Aesthetics dataset demonstrates that LCMs achieve state-of-the-art text-to-image generation performance with few-step inference. Project Page: https://latent-consistency-models.github.io/
## 1 Introduction
Diffusion models have emerged as powerful generative models that have gained significant attention and achieved remarkable results in various domains ( Ho et al., 2020 ; Song et al., 2020a ; Nichol & Dhariwal, 2021 ; Ramesh et al., 2022 ; Song & Ermon, 2019 ; Song et al., 2021 ) . In particular, latent diffusion models (LDMs) (e.g., Stable Diffusion ( Rombach et al., 2022 ) ) have demonstrated exceptional performance, especially in high-resolution text-to-image synthesis tasks. LDMs can generate high-quality images conditioned on textual descriptions by utilizing an iterative reverse sampling process that performs gradual denoising of samples. However, diffusion models suffer from a notable drawback: the iterative reverse sampling process leads to slow generation speed, limiting their real-time applicability. To overcome this drawback, researchers have proposed several methods to improve the sampling speed, which involves accelerating the denoising process by enhancing ODE solvers ( Ho et al., 2020 ; Lu et al., 2022a ; Lu et al., 2022b ) , which can generate images within 10 ∼ \sim 20 sampling steps. Another approach is to distill a pre-trained diffusion model into models that enable few-step inference Salimans & Ho (2022) ; Meng et al. (2023) . In particular, Meng et al. (2023) proposed a two-stage distillation approach to improving the sampling efficiency of classifier-free guided models. Recently, Song et al. (2023) proposed consistency models as a promising alternative aimed at speeding up the generation process. By learning consistency mappings that maintain point consistency on ODE-trajectory, these models allow for single-step generation, eliminating the need for computation-intensive iterations. However, Song et al. (2023) is constrained to pixel space image generation tasks, making it unsuitable for synthesizing high-resolution images. Moreover, the applications to the conditional diffusion model and the incorporation of classifier-free guidance have not been explored, rendering their methods unsuitable for text-to-image generation synthesis.
In this paper, we introduce Latent Consistency Models (LCMs) for fast, high-resolution image generation. Mirroring LDMs, we employ consistency models in the image latent space of a pre-trained auto-encoder from Stable Diffusion ( Rombach et al., 2022 ) . We propose a one-stage guided distillation method to efficiently convert a pre-trained guided diffusion model into a latent consistency model by solving an augmented PF-ODE. Additionally, we propose Latent Consistency Fine-tuning, which allows fine-tuning a pre-trained LCM to support few-step inference on customized image datasets. Our main contributions are summarized as follows:
• We propose Latent Consistency Models (LCMs) for fast, high-resolution image generation. LCMs employ consistency models in the image latent space, enabling fast few-step or even one-step high-fidelity sampling on pre-trained latent diffusion models (e.g., Stable Diffusion (SD)).
• We provide a simple and efficient one-stage guided consistency distillation method to distill SD for few-step (2 ∼ \sim 4) or even 1-step sampling. We propose the Skipping-Step technique to further accelerate the convergence. For 2- and 4-step inference, our method costs only 32 A100 GPU hours for training and achieves state-of-the-art performance on the LAION-5B-Aesthetics dataset.
• We introduce a new fine-tuning method for LCMs, named Latent Consistency Fine-tuning, enabling efficient adaptation of a pre-trained LCM to customized datasets while preserving the ability of fast inference.
## 2 Related Work
Diffusion Models have achieved great success in image generation ( Ho et al., 2020 ; Song et al., 2020a ; Nichol & Dhariwal, 2021 ; Ramesh et al., 2022 ; Rombach et al., 2022 ; Song & Ermon, 2019 ) . They are trained to denoise the noise-corrupted data to estimate the score of data distribution. During inference, samples are drawn by running the reverse diffusion process to gradually denoise the data point. Compared to VAEs ( Kingma & Welling, 2013 ; Sohn et al., 2015 ) and GANs ( Goodfellow et al., 2020 ) , diffusion models enjoy the benefit of training stability and better likelihood estimation.
Accelerating DMs. However, diffusion models are bottlenecked by their slow generation speed. Various approaches have been proposed, including training-free methods such as ODE solvers ( Song et al., 2020a ; Lu et al., 2022a ; Lu et al., 2022b ) , adaptive step size solvers ( Jolicoeur-Martineau et al., 2021 ) , predictor-corrector methods ( Song et al., 2020b ) . Training-based approaches include optimized discretization ( Watson et al., 2021 ) , truncated diffusion ( Lyu et al., 2022 ; Zheng et al., 2022 ) , neural operator ( Zheng et al., 2023 ) and distillation ( Salimans & Ho, 2022 ; Meng et al., 2023 ) . More recently, new generative models for faster sampling have also been proposed ( Liu et al., 2022 ; Liu et al., 2023 ) .
Latent Diffusion Models (LDMs) ( Rombach et al., 2022 ) excel in synthesizing high-resolution text-to-images. For example, Stable Diffusion (SD) performs forward and reverse diffusion processes in the data latent space, resulting in more efficient computation.
Consistency Models (CMs) ( Song et al., 2023 ) have shown great potential as a new type of generative model for faster sampling while preserving generation quality. CMs adopt consistency mapping to directly map any point in ODE trajectory to its origin, enabling fast one-step generation. CMs can be trained by distilling pre-trained diffusion models or as standalone generative models. Details of CMs are elaborated in the following section.
## 3 Preliminaries
In this section, we briefly review diffusion and consistency models and define relevant notations.
Diffusion Models: Diffusion models, or score-based generative models Ho et al. (2020) ; Song et al. (2020a) is a family of generative models that progressively inject Gaussian noises into the data, and then generate samples from noise via a reverse denoising process. In particular, diffusion models define a forward process transitioning the origin data distribution p d a t a ( x ) p_{data}(x) to marginal distribution q t ( 𝒙 t ) q_{t}(\bm{x}_{t}) , via transition kernel: q 0 t ( 𝒙 t ∣ 𝒙 0 ) = 𝒩 ( 𝒙 t ∣ α ( t ) 𝒙 0 , σ 2 ( t ) 𝑰 ) q_{0t}\left(\bm{x}_{t}\mid\bm{x}_{0}\right)=\mathcal{N}\left(\bm{x}_{t}\mid\alpha(t)\bm{x}_{0},\sigma^{2}(t)\bm{I}\right) , where α ( t ) , σ ( t ) \alpha(t),\sigma(t) specify the noise schedule. In continuous time perspective, the forward process can be described by a stochastic differential equation (SDE) Song et al. (2020b) ; Lu et al. (2022a) ; Karras et al. (2022) for t ∈ [ 0 , T ] t\in[0,T] : d 𝒙 t = f ( t ) 𝒙 t d t + g ( t ) d 𝒘 t , 𝒙 0 ∼ p d a t a ( 𝒙 0 ) , \mathrm{d}\bm{x}_{t}=f(t)\bm{x}_{t}\mathrm{~d}t+g(t)\mathrm{d}\bm{w}_{t},\,\,\bm{x}_{0}\sim p_{data}\left(\bm{x}_{0}\right), where 𝒘 t \bm{w}_{t} is the standard Brownian motion, and f ( t ) = d log α ( t ) d t , g 2 ( t ) = d σ 2 ( t ) d t − 2 d log α ( t ) d t σ 2 ( t ) . f(t)=\frac{\mathrm{d}\log\alpha(t)}{\mathrm{~d}t},\quad g^{2}(t)=\frac{\mathrm{d}\sigma^{2}(t)}{\mathrm{~d}t}-2\frac{\mathrm{d}\log\alpha(t)}{\mathrm{~d}t}\sigma^{2}(t). (1) By considering the reverse time SDE (see Appendix A for more details), one can show that the marginal distribution q t ( 𝒙 ) q_{t}(\bm{x}) satisfies the following ordinary differential equation, called the Probability Flow ODE (PF-ODE) ( Song et al., 2020b ; Lu et al., 2022a ) : d 𝒙 t d t = f ( t ) 𝒙 t − 1 2 g 2 ( t ) ∇ 𝒙 log q t ( 𝒙 t ) , 𝒙 T ∼ q T ( 𝒙 T ) . \frac{\mathrm{d}\bm{x}_{t}}{\mathrm{~d}t}=f(t)\bm{x}_{t}-\frac{1}{2}g^{2}(t)\nabla_{\bm{x}}\log q_{t}\left(\bm{x}_{t}\right),\,\,\bm{x}_{T}\sim q_{T}\left(\bm{x}_{T}\right). (2) In diffusion models, we train the noise prediction model ϵ θ ( 𝒙 t , t ) \bm{\epsilon}_{\theta}(\bm{x}_{t},t) to fit − ∇ log q t ( 𝒙 t ) -\nabla\log q_{t}(\bm{x}_{t}) (called the score function ). Approximating the score function by the noise prediction model in 21 , one can obtain the following empirical PF-ODE for sampling: d 𝒙 t d t = f ( t ) 𝒙 t + g 2 ( t ) 2 σ t ϵ θ ( 𝒙 t , t ) , 𝒙 T ∼ 𝒩 ( 𝟎 , σ ~ 2 𝑰 ) . \frac{\mathrm{d}\bm{x}_{t}}{\mathrm{~d}t}=f(t)\bm{x}_{t}+\frac{g^{2}(t)}{2\sigma_{t}}\bm{\epsilon}_{\theta}\left(\bm{x}_{t},t\right),\quad\bm{x}_{T}\sim\mathcal{N}\left(\mathbf{0},\tilde{\sigma}^{2}\bm{I}\right). (3) For class-conditioned diffusion models, Classifier-Free Guidance (CFG) ( Ho & Salimans, 2022 ) is an effective technique to significantly improve the quality of generated samples and has been widely used in several large-scale diffusion models including GLIDE Nichol et al. (2021) , Stable Diffusion ( Rombach et al., 2022 ) , DALL·E 2 ( Ramesh et al., 2022 ) and Imagen ( Saharia et al., 2022 ) . Given a CFG scale ω \omega , the original noise prediction is replaced by a linear combination of conditional and unconditional noise prediction, i.e., ϵ ~ 𝜽 ( 𝒛 t , ω , 𝒄 , t ) = ( 1 + ω ) ϵ 𝜽 ( 𝒛 t , 𝒄 , t ) − ω ϵ 𝜽 ( 𝒛 , ∅ , t ) . \bm{\tilde{\epsilon}_{\theta}}({\bm{z}}_{t},\omega,{\bm{c}},t)=(1+\omega)\bm{\epsilon_{\theta}}({\bm{z}}_{t},{\bm{c}},t)-\omega\bm{\epsilon_{\theta}}({\bm{z}},\varnothing,t).
Consistency Models: The Consistency Model (CM) ( Song et al., 2023 ) is a new family of generative models that enables one-step or few-step generation. The core idea of the CM is to learn the function that maps any points on a trajectory of the PF-ODE to that trajectory’s origin (i.e., the solution of the PF-ODE). More formally, the consistency function is defined as 𝒇 : ( 𝒙 t , t ) ⟼ 𝒙 ϵ , {\bm{f}}:({\bm{x}}_{t},t)\longmapsto{\bm{x}}_{\epsilon}, where ϵ \epsilon is a fixed small positive number. One important observation is that the consistency function should satisfy the self-consistency property : 𝒇 ( 𝒙 t , t ) = 𝒇 ( 𝒙 t ′ , t ′ ) , ∀ t , t ′ ∈ [ ϵ , T ] . {\bm{f}}({\bm{x}}_{t},t)={\bm{f}}({\bm{x}}_{t^{\prime}},t^{\prime}),\forall t,t^{\prime}\in[\epsilon,T]. (4) The key idea in ( Song et al., 2023 ) for learning a consistency model 𝒇 𝜽 {\bm{f}}_{\bm{\theta}} is to learn a consistency function from data by effectively enforcing the self-consistency property in Eq. 4 . To ensure that 𝒇 𝜽 ( 𝒙 , ϵ ) = 𝒙 {\bm{f}}_{\bm{\theta}}({\bm{x}},\epsilon)={\bm{x}} , the consistency model 𝒇 𝜽 {\bm{f}}_{\bm{\theta}} is parameterized as: 𝒇 𝜽 ( 𝒙 , t ) = c skip ( t ) 𝒙 + c out ( t ) 𝑭 𝜽 ( 𝒙 , t ) , {\bm{f}}_{\bm{\theta}}({\bm{x}},t)=c_{\text{skip}}(t){\bm{x}}+c_{\text{out}}(t)\bm{F}_{\bm{\theta}}({\bm{x}},t), (5) where c skip ( t ) c_{\text{skip}}(t) and c out ( t ) c_{\text{out}}(t) are differentiable functions with c skip ( ϵ ) = 1 c_{\text{skip}}(\epsilon)=1 and c out ( ϵ ) = 0 c_{\text{out}}(\epsilon)=0 , and 𝑭 𝜽 ( 𝒙 , t ) \bm{F}_{\bm{\theta}}({\bm{x}},t) is a deep neural network. A CM can be either distilled from a pre-trained diffusion model or trained from scratch. The former is known as Consistency Distillation . To enforce the self-consistency property, we maintain a target model 𝜽 − {\bm{\theta}}^{-} , updated with exponential moving average (EMA) of the parameter 𝜽 {\bm{\theta}} we intend to learn, i.e., 𝜽 − ← μ 𝜽 − + ( 1 − μ ) 𝜽 {\bm{\theta}}^{-}\leftarrow\mu{\bm{\theta}}^{-}+(1-\mu){\bm{\theta}} , and define the consistency loss as follows: ℒ ( 𝜽 , 𝜽 − , Φ ) = 𝔼 𝒙 , t [ d ( 𝒇 𝜽 ( 𝒙 t n + 1 , t n + 1 ) , 𝒇 𝜽 − ( 𝒙 ^ t n ϕ , t n ) ) ] , \mathcal{L}({\bm{\theta}},{\bm{\theta}}^{-};\Phi)=\mathbb{E}_{\bm{x},t}\left[d\left({\bm{f}}_{\bm{\theta}}({\bm{x}}_{t_{n+1}},t_{n+1}),{\bm{f}}_{{\bm{\theta}}^{-}}(\hat{{\bm{x}}}^{\phi}_{t_{n}},t_{n})\right)\right], (6) where d ( ⋅ , ⋅ ) d(\cdot,\cdot) is a chosen metric function for measuring the distance between two samples, e.g., the squared ℓ 2 \ell_{2} distance d ( 𝒙 , 𝒚 ) = ‖ 𝒙 − 𝒚 ‖ 2 2 d({\bm{x}},{\bm{y}})=||{\bm{x}}-{\bm{y}}||_{2}^{2} . 𝒙 ^ t n ϕ \hat{{\bm{x}}}_{t_{n}}^{\phi} is a one-step estimation of 𝒙 t n \bm{x}_{t_{n}} from 𝒙 t n + 1 \bm{x}_{t_{n+1}} as: 𝒙 ^ t n ϕ ← 𝒙 t n + 1 + ( t n − t n + 1 ) Φ ( 𝒙 t n + 1 , t n + 1 , ϕ ) . \hat{{\bm{x}}}_{t_{n}}^{\phi}\leftarrow{\bm{x}}_{t_{n+1}}+(t_{n}-t_{n+1})\Phi({\bm{x}}_{t_{n+1}},t_{n+1};\phi). (7) where Φ \Phi denotes the one-step ODE solver applied to PF-ODE in Eq. 24 . ( Song et al., 2023 ) used Euler ( Song et al., 2020b ) or Heun solver ( Karras et al., 2022 ) as the numerical ODE solver. More details and the pseudo-code for consistency distillation (Algorithm 2 ) are provided in Appendix A .
## 4 Latent Consistency Models
Consistency Models (CMs) ( Song et al., 2023 ) only focused on image generation tasks on ImageNet 64 × \times 64 ( Deng et al., 2009 ) and LSUN 256 × \times 256 ( Yu et al., 2015 ) . The potential of CMs to generate higher-resolution text-to-image tasks remains unexplored. In this paper, we introduce Latent Consistency Models (LCMs) in Sec 4.1 to tackle these more challenging tasks, unleashing the potential of CMs. Similar to LDMs, our LCMs adopt a consistency model in the image latent space. We choose the powerful Stable Diffusion (SD) as the underlying diffusion model to distill from. We aim to achieve few-step (2 ∼ \sim 4) and even one-step inference on SD without compromising image quality. The classifier-free guidance (CFG) ( Ho & Salimans, 2022 ) is an effective technique to further improve sample quality and is widely used in SD. However, its application in CMs remains unexplored. We propose a simple one-stage guided distillation method in Sec 4.2 that solves an augmented PF-ODE , integrating CFG into LCM effectively. We propose Skipping-Step technique to accelerate the convergence of LCMs in Sec. 4.3 . Finally, we propose Latent Consistency Fine-tuning to finetune a pre-trained LCM for few-step inference on a customized dataset in Sec 4.4 .
### 4.1 Consistency Distillation in the Latent Space