Stable Diffusion: Generate Quality Images with Diffusion Model on Latent Space

May 4, 2024, 7:49 p.m. · 10 min read · 🌐︎ en

deep learning paper review multimodal generative

Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

Introduction

As generative computer vision models generate increasingly high-quality images, the computational demands have also risen. Initially, Generative Adversarial Networks (GANs) have shown promising results, but its distinctive, unstable training method hampers the ability to scale to model a more complex, multi-modal distribution 1.
Diffusion Models (DMs), in contrast, does not exhibit such properties. Highly benefited from parameter sharing, it is able to model complex distributions with a smaller model.
However, being a likelihood-based model, DM allocates too much resources on modelling inperceptible (and thus, meaningless) details of images. Also, repeated function evaluation included in the inference process makes the computation cost too expensive.
Therefore, researchers have been finding methods to reduce the training/sampling complexity of DMs without reducing its performance. Proposed in this paper, Latent Diffusion Model(LDM) is the very solution to such an issue. The authors were inspired by the finding that the learning process of likelihood-based models can be rougly divided into two steps:

  1. The perceptual compression stage, which removes high-frequency details,
  2. And the semantic compression stage, where the model learns the semantic components of the data.

The authors decided to explicitly divide these two stages, so that each of them can be done in separate models.

  1. An autoencoder is employed to compress the image into a lower-dimensional
    representation. This eliminates much information from the original data, but the vast majority of them are semantically meaningless details.

    • This autoencoder is trained only once, and is reused for each training of DMs throughout the paper.
  2. Then, the diffusion model is run on the latent representation compressed by the autoencoder.

Through this, the diffusion model is no longer run on the complex pixel space, but on the latent space, which has much better scaling property over spatial dimension. As a result, the reduced computation complexity enables the LDMs to generate images more efficiently.
The authors summarize the contribution of the paper in five. They claim that LDMs

Previous Work

Generative Model for Image Synthesis

The majority of prior work can be divided into the classes: (1) GAN methods, (2) likelihood-based methods including VAE(Variational Autoencoder) and flow model, and (3) autoregressive models (ARM). As previously discussed, each has pros and cons:

Diffusion Models, empowered by the inductive bias of its underlying UNet-like structure, have appeared as a powerful alternative for these methods. Among them, DDPM, in short of Denoising Diffusion Probabilistic Model, have achieved the most promising performance by engaging in the so-called reweighted objective. However, being trained directly on the pixel space, these models also suffer from very high training cost and low speed. This is why the authors introduce the compressed latent space on which diffusion model is run.

Two-Stage Image Synthesis

As described above, LDM can be classified as a two-stage method, exploiting separate models (autoencoder and DM) for each of the steps. This section studies previously suggested two-step approaches applied to image synthesis task.

As opposed to these models, LDMs are free from such trade-offs. The compression level that both guarantee a high quality and good efficiency can be freely chosen. As for the design choice whether to train the two models jointly or separately, the authors explain that the former requires difficult weighting between the reconstruction and generative capability.

Method

As explained previously in the introduction, the authors separate the compressive phase from the generative learning phase, exploiting a separate model for it: an autoencoder. This enables the DM to train on a latent space which is perceptually equivalent with the original image, but where the redundant and inperceptible details are removed. There are several advantages with such architecture:

Step 1: Perceptual Image Compression

The perceptual image compression model is based on the autoencoder employed in the
VQGAN, where a combination of (1) perceptual loss and (2) patch-based adversarial loss is utilized. Mathematically, speaking, the architecture comprises an encoder $\mathcal{E}$ and a decoder $\mathcal{D}$. Here, the encoder is configured to downsample the input by a factor $f=H/h=W/w$, where $f$ is varied in the experiments. The decoder $\mathcal{D}$ is trained to
restore the image from the latent, yielding $\tilde{x} = \mathcal{D}(\mathcal{E}(x))$.

To control the variance of the latent space, two different kinds of regularizations are used:

Step 2: Latent Diffusion Models

Briefly speaking, diffusion models are the models that utilizes a sequential denoising process to generate an image. They approximate the reverse diffusion process, which calculates the denoised image $x_t$ from the noised image $x_{t+1}$. Among the DMs, the most successful ones including DDPM engage in the reweighted loss:

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

Since the latent diffusion model is trained on the latent space, not pixel space, the above formula must be modified as follows:

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

Here, $\epsilon_\theta(\cdot, t)$ is implemented using time-conditional UNet.

Conditioning Mechanism

The above figure outlines the overall architecture of the conditional LDM demonstrated in the paper. Here, the crucial part to control the generated image is the conditioning mechanism, which is implemented by cross-attention. The encoder with condition can be denoted as $\mathcal{E}_\theta(z_t, t, y)$ where $y$ is the condition, which can either be text, semantic maps, or other image inputs.

As shown in the figure, the condition $y$ first passes through a network $\tau_\theta$ to be mapped to the latent space. The resulting $\tau_\theta(y)$ is fed into the attention modules that composes up the denoising UNet. In detail, each attention layer in the UNet uses $\tau_\theta(y)$ as the key($K$) and the value($V$), where the intermediate outputs of the UNet are utilized as the query($Q$), that is,
$$Q = W_Q^{(i)}\cdot \phi_i(z_t),$$

$$K = W_K^{(i)}\cdot \tau_\theta(y),$$

$$V = W_V^{(i)}\cdot \tau_\theta(y)$$
Considering that the projection of a vector $v$ onto the span of column vectors of $W$ (where the column vectors of W are unit vectors) is formulated by
$$\text{proj}_W(\mathbf{u}) = (\mathbf{u}^T W)W$$

This is roughly equivalent to the projection of UNet representations to the space determined by the input condition.

Putting this all together, the training objective can be expressed as
$$L_{LDM} = \mathbb{E}_{z\sim \mathcal{E}(x), y, \epsilon \sim\mathcal{N}(0, I), t}\left[\lVert \epsilon - \epsilon_\theta(z_t, t, \tau_\theta(y))\rVert_2^2\right]$$
where $\tau_\theta$ and $\epsilon_\theta$ are jointly optimized.

Experiments

The authors have conducted a number of experiments to study the adaptation of LDM on various tasks.

Downsampling Factor

First, the authors conducted an experiment varying the value of downsampling factor $f\in{1, 2, 4, 8, 16, 32}$, each of which the authors refer to as LDM-$f$. The paper reports that there is the trade-off between the training time and the fidelity of generated image over $f$.
Too small value of $f$ is bound to make the training process slow, while too large values can incur an excessive information loss, resulting in the low quality. LDM-4 and LDM-8 are considered as a reasonable choice between this trade-off.

Image Generation

The authors then evaluated the (1) sample quality and (2) their coverage of the data manifold. For the latter, FID(Frechet Inception Distance) and Precision-and-Recall were used. Evaluating the metrics on four different datasets, LDM showed the best performance in the majority of cases, including CelebA-HQ, where the model achieved a new SOTA FID. Especially, LDM consistently showed the precision and recall better than GAN models, indicating that it has a wider mode coverage.

Note: How do we evaluate the precision and recall of a generative model?
The precision and the recall of the generative vision model represents how realistic and how diversified the generated images are, respectively. One way to assess these (Kynkäänniemi et al., 2019) is to embed the dataset images using a pre-trained VGGNet, approximate the manifold of real images using k-nearest neighbor (k-NN), and evaluate the precision and the recall by determining if the generated images are included in the learned manifold.

Conditional Image Generation

There are two ways in controlling the generation process of LDM: one using the cross-attention as described above, and the other held by concatenating the conditioning image to the input of denoising UNet.

Cross-Attention Based Conditioning

The authors apply cross-attention based conditioning on three types of tasks: (1) text-to-image, (2) semantic layout, and (3) ImageNet class condition.
Firstly, the authors trained a KL-reg LDM conditioned on the transformer output which embeds the text input to the latent space LDM uses.

As depicted in the above table, the guided LDM-KL-8-G was on par with the SOTA autoregressive model, whie employing a significantly less number of parameters.

Secondly, the LDM was trained to synthesize images based on semantic layouts on COCO, performing layout-to-image synthesis. Below are the samples of generated images.

Finally, the authors conditioned LDM to synthesize images that correspond to each class of ImageNet. The results showed that LDM outperforms the SOTA LDMs can be conditioned by feeding the conditioning information together with the input of denoising UNet. Note that the image must be spatially aligned to be utilized as a condition.

Concatenation of Conditioning Image

Another way of controlling the image generation is to concatenating the spatially aligned conditioning image to the input of denoising UNet $\epsilon_\theta$. LDMs can thereby serve as image-to-image translation models. The authors apply this method into three tasks: (1) semantic synthesis, (2) super-resolution, and (3) inpainting.

Firstly, in the semantic synthesis, the images paired with corresponding semantic map were used to train the LDM. The downsampled version ($f=4$) of semantic map was fed to the diffusion model to generate the image.

While experimenting for semantic synthesis, the authors discovered that the model generalizes well to a larger dimension. This is enabled by the convolution-based structure of the denoising UNet, which is insensitive to the input dimension, although it was trained on $256^2$. This feature was further applied to the super resolution and inpainting task, which are discussed in the continuing sections.

Super-Resolution

LDMs could effectively handle the super-resolution task by generating a high resolution image, conditioned on the low resolution version of it. The authors trained the model to restore the original image based on the image degraded via bicubic interpolation.

Although LDM did not succeed to outperform existing models in quantitative metrics, these metrics are known not to align with human-perceived quality. Human evaluation results show that LDM has a better performance in super-resolution.

Inpainting

Inpainting is a task where the model is trained to fill the masked region of the image. The authors have compared the LDM’s performance on the task with the SOTA approaches.

As the above figures indicate, LDMs showed a surprisingly good performance in inpainting.

The quantitative results also suggests that LDM improves FID by a factor of at least 1.6x, while also reducing the time consumption.

Conclusion

In this paper, the authors have suggested a new methodology for generating images, stable diffusion. By introducing an autoencoder model to map the data to the latent space, on which the diffusion models are run, the authors could drastically reduce the computational demand while maintaining the image quality. The model allows various, multi-modal conditioning mechanisms to be applied, including text, class, and image. Through this general conditioning, the authors could apply the model into various downstream tasks, yielding the performance better than or at least on par of task-specific SOTA models.


  1. Here, a distribution being multimodal means it has multiple peaks. This has nothing to do with multi-modal models such as CLIP). In general, GANs are known to be prone to mode collapse, which means it does not generate diversified images. 

  2. This includes super-resolution, inpainting and semantic. synthesis.