Papers
High-Resolution Image Synthesis with Latent Diffusion Models
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, Björn Ommer · CVPR 2022, 2022
The problem: diffusion in pixel space is expensive to train and to sample from
Diffusion models (DMs) build an image by sequentially applying many denoising autoencoders and, at the time of the paper, already delivered state-of-the-art results in image synthesis. The problem the authors point out is cost: because these models operate directly in pixel space, training the most powerful ones usually consumes hundreds of GPU days (the paper cites 150 to 1000 V100 days in a prior work), and generating 50 thousand samples would take about 5 days on a single A100 in that same work. Inference is also expensive because the same network has to run sequentially for many steps (from 25 to 1000, in the cited work). The paper explains the root of the waste: because they are likelihood-based models, they tend to spend capacity modeling imperceptible details of the image, and, even though the loss function can be reweighted to ignore some of those details, the network and the gradients still have to be evaluated over all pixels. The authors see two consequences: training is restricted to those with enormous computing power, and using an already trained model is also expensive in time and memory. The stated goal is to reduce the computational complexity of training and sampling without hurting performance, which would make this class of models more accessible.
First stage: a perceptual autoencoder that compresses without destroying detail
The authors' idea is to separate compression from generation. They observe that learning a likelihood-based model has two stages: a perceptual compression, which removes high-frequency details but learns little semantic variation, and a semantic compression, in which the generative model learns the conceptual composition of the data. The first stage is handed to a separately trained autoencoder. Given an image x, the encoder E produces a latent z = E(x), with height and width reduced by a downsampling factor f, and the decoder D reconstructs the image from z. The autoencoder is trained by combining a perceptual loss with a patch-based adversarial objective, which keeps reconstructions locally realistic and avoids the blurriness caused by using only pixel losses such as L1 or L2. To prevent latents with arbitrarily high variance, the paper tests two regularizations: a slight KL penalty toward a standard normal (KL-reg.) and a vector quantization layer inside the decoder (VQ-reg.). Because the subsequent diffusion model exploits the two-dimensional structure of the latent, the compression can be mild, unlike earlier work that used a 1D ordering of the latent space and much more aggressive compression. This stage is trained only once and can be reused across several diffusion trainings.
Second stage: diffusing in the latent space
With E and D ready, the diffusion model now works in the latent space instead of on pixels. The process is the same as in an ordinary DM, a sequence of denoising autoencoders trained to predict a noise-free version of their input, but now the input is zt, the noisy version of z = E(x). The denoising network is a UNet conditioned on the time step, and the loss is the squared error between the true noise and the noise predicted by the network. Because the forward noising process is fixed, zt can be obtained efficiently from the encoder during training. At generation time, a sample from the latent is decoded into image space with a single pass through the decoder D. According to the authors, this space is better suited to likelihood-based models because they can focus on the important semantic bits and train in a lower-dimensional, computationally cheaper space. Because the latent keeps spatial structure, the model still takes advantage of the image-specific inductive biases of the UNet, which is built mainly from 2D convolutions. In contrast to prior work that used autoregressive transformers over a highly compressed discrete space, this avoids the need for aggressive compression. The paper also says that, because the model is convolutional, in densely conditioned tasks such as super-resolution and inpainting it can generate large, consistent images of about 1024 by 1024 pixels.
Cross-attention conditioning: from text to image
Diffusion models can, in principle, model conditional distributions, in which the denoising network also receives an input y, such as text, semantic maps or another image. The authors observe that combining DMs with conditioning beyond class labels or blurred versions of the input was an underexplored area. The proposed solution is to add cross-attention mechanisms to the UNet. A domain-specific encoder, called τθ, projects y to an intermediate representation. This representation supplies the keys and values of the attention, while the queries come from a flattened intermediate representation of the UNet itself, and the result is injected into intermediate layers of the network. The encoder τθ and the UNet are jointly optimized with the conditional loss. For text, the paper uses a BERT tokenizer and implements τθ as a transformer, training a KL-regularized model with 1.45 billion parameters on LAION-400M. The design is flexible: switching modality means swapping the encoder τθ, and the paper also trains, with the same mechanism, class-conditioned and layout-conditioned models. For some dense conditioning cases, such as semantic maps, the authors instead use concatenation of the inputs with the latent.
Compression trade-offs and limitations
The paper analyzes how the downsampling factor f affects the result, comparing LDM-f with f at 1, 2, 4, 8, 16, and 32, where LDM-1 is pixel-space diffusion. All models in this analysis use a single A100, the same number of steps, and the same number of parameters. The results show two bad extremes. With small f (LDM-1 and LDM-2), training progresses slowly, because much of the perceptual compression is left to the diffusion model. With very large f, fidelity stagnates after a few steps, because the first stage compresses too much, loses information, and limits the attainable quality. The values from LDM-4 to LDM-16 balance efficiency and fidelity; after 2 million steps, the paper reports a 38-point FID difference between LDM-1 and LDM-8. In tests with the DDIM sampler, LDM-4 and LDM-8 have much lower FID than pixel-space LDM-1 and, at the same time, higher sample throughput, and the paper concludes that LDM-4 and LDM-8 offer the best conditions for high-quality synthesis. Complex datasets such as ImageNet require lower compression rates. In the limitations, the authors acknowledge that sequential sampling is still slower than that of GANs and that the autoencoder's reconstruction capacity can become a bottleneck when fine pixel-level accuracy is required, and they suppose that their super-resolution models are already somewhat limited in this respect.
See the ideas at work
Practice the designs that use this paper in the System Design school.