Stable Diffusion / LDM — High-Resolution Image Synthesis with Latent Diffusion Models

By decomposing the image formation process into a sequential application of denoising autoencoders, diffusion models (DMs) achieve state-of-the-art synthesis results on image data and beyond. However, since these models typically operate directly in pixel space, optimization of powerful DMs often consumes hundreds of GPU days and inference is expensive due to sequential evaluations. To enable DM training on limited computational resources while retaining their quality and flexibility, we apply them in the latent space of powerful pretrained autoencoders.

In contrast to previous work, training diffusion models on such a representation allows for the first time to reach a near-optimal point between complexity reduction and detail preservation, greatly boosting visual fidelity. By introducing cross-attention layers into the model architecture, we turn diffusion models into powerful and flexible generators for general conditioning inputs such as text or bounding boxes and high-resolution synthesis becomes possible in a convolutional manner. Our latent diffusion models (LDMs) achieve new state-of-the-art scores for image inpainting and class-conditional image synthesis and highly competitive performance on various tasks, while significantly reducing computational requirements compared to pixel-based DMs.

Image synthesis is one of the computer vision fields with the most spectacular recent development, but also among those with the greatest computational demands. Recently, diffusion models, which are built from a hierarchy of denoising autoencoders, have shown to achieve impressive results in image synthesis and define the state-of-the-art in class-conditional image synthesis and super-resolution. Being likelihood-based models, they do not exhibit mode-collapse and training instabilities as GANs and, by heavily exploiting parameter sharing, they can model highly complex distributions without involving billions of parameters as in AR models.

DMs belong to the class of likelihood-based models, whose mode-covering behavior makes them prone to spend excessive amounts of capacity (and thus compute) on modeling imperceptible details of the data. Training the most powerful DMs often takes hundreds of GPU days (e.g. 150–1000 V100 days), and repeated evaluations on a noisy version of the input space render also inference expensive, so that producing 50k samples takes approximately 5 days on a single A100 GPU. This has two consequences: training requires massive resources available only to a small fraction of the field and leaves a huge carbon footprint; and evaluating an already trained model is also expensive in time and memory.

Our approach starts with the analysis of already trained diffusion models in pixel space: learning can be roughly divided into two stages. First is a perceptual compression stage which removes high-frequency details but still learns little semantic variation. In the second stage, the actual generative model learns the semantic and conceptual composition of the data (semantic compression). We thus aim to first find a perceptually equivalent, but computationally more suitable space, in which we will train diffusion models for high-resolution image synthesis.

We separate training into two distinct phases: first, we train an autoencoder which provides a lower-dimensional representational space which is perceptually equivalent to the data space. Importantly, we do not need to rely on excessive spatial compression, as we train DMs in the learned latent space, which exhibits better scaling properties with respect to spatial dimensionality. We dub the resulting model class Latent Diffusion Models (LDMs). A notable advantage is that we need to train the universal autoencoding stage only once and can therefore reuse it for multiple DM trainings or for different tasks.

In sum, our work makes the following contributions: (i) Our method scales more gracefully to higher dimensional data and can work on a compression level which provides more faithful and detailed reconstructions than previous work, and can be efficiently applied to high-resolution synthesis of megapixel images. (ii) We achieve competitive performance on multiple tasks while significantly lowering computational costs. (iii) Our approach does not require a delicate weighting of reconstruction and generative abilities, ensuring faithful reconstructions with very little regularization of the latent space. (iv) For densely conditioned tasks, our model can be applied in a convolutional fashion and render large, consistent images of ~1024² px. (v) We design a general-purpose conditioning mechanism based on cross-attention, enabling multi-modal training (class-conditional, text-to-image, layout-to-image). (vi) We release pretrained latent diffusion and autoencoding models.

[摘要合并] Generative Models: GANs sample high-res images efficiently but are hard to optimize and struggle to capture the full distribution; VAEs/flows enable efficient synthesis but sample quality lags GANs; ARMs achieve strong density estimation but their sequential sampling limits them to low resolution. Diffusion Models recently achieved SOTA in density estimation and sample quality with a UNet backbone, but evaluating/optimizing in pixel space has low inference speed and very high training cost.

Two-Stage Synthesis: VQ-VAEs use ARMs over a discretized latent; VQGANs employ adversarial + perceptual first stage to scale AR transformers. However, the high compression rates required for feasible ARM training limit overall performance. Our LDMs scale more gently to higher-dimensional latent spaces due to their convolutional backbone, so we are free to choose the compression level that optimally mediates between a powerful first stage and the generative diffusion model.

3.1 Perceptual Image Compression. Our perceptual compression model is based on an autoencoder trained by combination of a perceptual loss and a patch-based adversarial objective. Given an image $x \in \mathbb{R}^{H\times W\times 3}$, the encoder $\mathcal{E}$ encodes $x$ into a latent $z=\mathcal{E}(x)$, and the decoder $\mathcal{D}$ reconstructs $\tilde{x}=\mathcal{D}(z)=\mathcal{D}(\mathcal{E}(x))$, where $z\in\mathbb{R}^{h\times w\times c}$. The encoder downsamples by a factor $f=H/h=W/w$, and we investigate $f=2^m$.

To avoid arbitrarily high-variance latent spaces, we experiment with two regularizations: KL-reg. imposes a slight KL-penalty towards a standard normal (similar to a VAE), whereas VQ-reg. uses a vector quantization layer within the decoder. Because our subsequent DM works with the two-dimensional structure of $z$, we can use relatively mild compression rates and achieve very good reconstructions — in contrast to previous works which relied on an arbitrary 1D ordering of $z$.

3.2 Latent Diffusion Models. Diffusion models learn a data distribution $p(x)$ by gradually denoising a normally distributed variable, corresponding to learning the reverse of a fixed Markov Chain of length $T$. The objective simplifies to

$$L_{DM} = \mathbb{E}_{x,\,\epsilon\sim\mathcal{N}(0,1),\,t}\Big[\,\lVert \epsilon - \epsilon_\theta(x_t, t)\rVert_2^2\,\Big].$$

With our trained perceptual model $(\mathcal{E},\mathcal{D})$, we now have an efficient, low-dimensional latent space. The neural backbone $\epsilon_\theta(\circ,t)$ is realized as a time-conditional UNet. Since the forward process is fixed, $z_t$ can be efficiently obtained from $\mathcal{E}$ during training, and samples from $p(z)$ can be decoded with a single pass through $\mathcal{D}$. The corresponding objective becomes

$$L_{LDM} := \mathbb{E}_{\mathcal{E}(x),\,\epsilon\sim\mathcal{N}(0,1),\,t}\Big[\,\lVert \epsilon - \epsilon_\theta(z_t, t)\rVert_2^2\,\Big].$$

3.3 Conditioning Mechanisms. We turn DMs into more flexible conditional generators by augmenting the UNet with the cross-attention mechanism. To pre-process $y$ from various modalities (e.g. language prompts), we introduce a domain specific encoder $\tau_\theta$ that projects $y$ to an intermediate representation $\tau_\theta(y)\in\mathbb{R}^{M\times d_\tau}$, then mapped into intermediate UNet layers via cross-attention:

$$\text{Attention}(Q,K,V)=\text{softmax}\!\left(\tfrac{QK^\top}{\sqrt{d}}\right)\!\cdot V,$$ with $Q=W_Q^{(i)}\varphi_i(z_t),\; K=W_K^{(i)}\tau_\theta(y),\; V=W_V^{(i)}\tau_\theta(y)$.

Based on image-conditioning pairs, we learn the conditional LDM via $$L_{LDM}:=\mathbb{E}_{\mathcal{E}(x),y,\epsilon,t}\Big[\lVert\epsilon-\epsilon_\theta(z_t,t,\tau_\theta(y))\rVert_2^2\Big],$$ where both $\tau_\theta$ and $\epsilon_\theta$ are jointly optimized.

4.1 On Perceptual Compression Tradeoffs. We analyze LDMs with different downsampling factors $f\in\{1,2,4,8,16,32\}$ (LDM-1 = pixel-based DM). Fig. 6 shows sample quality vs training progress for class-conditional models on ImageNet. We see that i) small downsampling for LDM-{1,2} results in slow training, whereas ii) overly large $f$ causes stagnating fidelity after few steps. We attribute this to leaving most of perceptual compression to the diffusion model (small $f$), vs too strong first-stage compression losing information (large $f$). LDM-{4-16} strike a good balance, manifesting in a significant FID gap of 38 between LDM-1 and LDM-8 after 2M steps.

4.2 Image Generation with Latent Diffusion. We train unconditional models on CelebA-HQ, FFHQ, LSUN-Churches/-Bedrooms and evaluate sample quality (FID) and coverage (Precision-Recall). On CelebA-HQ, we report a new state-of-the-art FID of 5.11, outperforming previous likelihood-based models and GANs, and also outperforming LSGM where a latent DM is trained jointly with the first stage. We outperform prior diffusion approaches on all but LSUN-Bedrooms, where our score is close to ADM despite using half its parameters and 4× less train resources.

4.3 Conditional LDM — Text-to-Image. We train a 1.45B-parameter KL-regularized LDM conditioned on language prompts on LAION-400M, using BERT-tokenizer and a transformer as $\tau_\theta$. On MS-COCO, applying classifier-free guidance greatly boosts sample quality, such that the guided LDM-KL-8-G is on par with recent SOTA AR and diffusion models for text-to-image while substantially reducing parameter count. We also train layout-to-image models and class-conditional ImageNet models (Tab. 3), outperforming ADM while reducing compute.

4.3.2 – 4.5 Convolutional Sampling, Super-Resolution & Inpainting. By concatenating spatially aligned conditioning to the input of $\epsilon_\theta$, LDMs serve as general-purpose image-to-image models. A model trained at 256² generalizes to larger resolutions (e.g. 512×1024) when applied convolutionally, used for semantic synthesis, super-resolution and inpainting to render images between 512² and 1024². For super-resolution we condition on low-resolution images via concatenation (following SR3); for inpainting LDMs achieve new SOTA. Their sequential sampling process is still slower than GANs, and reconstruction can become a bottleneck for tasks requiring fine-grained pixel accuracy.

Limitations. While LDMs significantly reduce computational requirements compared to pixel-based approaches, their sequential sampling process is still slower than that of GANs. Moreover, the use of LDMs can be questionable when high precision is required: although the loss of image quality is very small in our $f=4$ autoencoding models, their reconstruction capability can become a bottleneck for tasks that require fine-grained accuracy in pixel space.

Conclusion. We have presented latent diffusion models, a simple and efficient way to significantly improve both the training and sampling efficiency of denoising diffusion models without degrading their quality. Based on this and our cross-attention conditioning mechanism, our experiments could demonstrate favorable results compared to state-of-the-art methods across a wide range of conditional image synthesis tasks without task-specific architectures.