High-Resolution Image Synthesis with Latent Diffusion Models

Paper:https://arxiv.org/abs/2112.10752 HackMD Link

1. Introductionof DMs (Diffusion Models)

1-1.Core Concept:

• generate picture from a simple distribtion(e.g.Gaussian distribution noise)

• This core concept is to simulate the process of thermodynamic diffusion from physics, breaking down data generation into a sequential series of denoising steps.

1-2.Forward and Reverse Processes

• Here is a detailed explanation of how Diffusion Models work:

1. Forward Process - Diffusion (Noising)

   • Nature: Fixed, non-learnable Markov Chain process.
   • Purpose: To define a “data corruption” path, degrading the original data ($x_0$) into noise.
   • Mechanism: Add a small, controlled amount of Gaussian white noise at each time step $t$.
   • Mathematical Definition:$$q(x_t∣x_{t−1})=\mathcal{N}(x_t;\space x_{t−1}\sqrt{1−\beta_t} ,\space I\beta_t)$$Where $\beta_t$ is the diffusion rate.
   • Result: The final state $x_T$ approximates a standard Gaussian distribution: $q(x_T) \approx \mathcal{N}(0,I)$.

2. Reverse Process - Denoising (Generation)

   • Nature: The learnable part of the model.
   • Purpose: The model learns how to reconstruct the data distribution from the noise.
   • Mechanism: The model ($\epsilon_{\theta}$) is trained to infer and gradually reduce the noise, reversing the forward process.
   • Mathematical Definition:$$p_{\theta}(x_{t−1}|x_t) \approx \mathcal{N}(x_{t−1};\space \mu_{\theta}(x_t,\space t),\space \Sigma_{\theta}(x_t,\space t))$$

1-3.Optimization(Training):

The model is trained to minimize the difference between the true reverse distribution and the model’s prediction. The standard objective used in practical DMs simplifies the variational lower bound (VLB) on the negative log-likelihood into a highly effective weighted form of the Mean Squared Error (MSE).

Training Objective: Noise Prediction

Instead of directly predicting the mean ($\mu$) and variance ($\Sigma$) of the reverse transition, DMs are typically trained to simply predict the noise component added at time $t$.Model Role: The neural network $\epsilon_{\theta}$ (usually a U-Net) is trained to estimate the noise $\epsilon$ that was added to $x_{t-1}$ to get $x_t$.Simplified Loss Function: The most common and effective loss function used for training Diffusion Models is:$$\text{Loss} = \mathbb{E}{x_0, \epsilon \sim \mathcal{N}(0,I), t} \left[ \mathbf{w}(t) \cdot ||\epsilon - \epsilon{\theta}(\sqrt{\bar{\alpha}_t}x_0 + \sqrt{1 - \bar{\alpha}t}\epsilon, t)||{2}^{2} \right]$$

Where:

• $\epsilon$: The ground-truth noise sampled from $\mathcal{N}(0,I)$.
• $\epsilon_{\theta}(x_t, t)$: The noise predicted by the model at time $t$ for the noisy input $x_t$.
• $x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon$: A simplified, non-sequential way to sample $x_t$ from $x_0$ (where $\bar{\alpha}_t$ is related to the product of $(1-\beta_s)$ from $s=1$ to $t$).
• $\mathbf{w}(t)$: A weighting function (often set to 1 or based on $1/\sqrt{1 - \bar{\alpha}_t}$) which balances the importance of different time steps.

Core Idea:

By minimizing this loss, the model $\epsilon_{\theta}$ learns to accurately estimate $\epsilon$. This estimated noise is then used during the sampling phase to derive the mean $\mu_{\theta}$ required for the reverse step, thus enabling the generation of high-quality images.

1-4.Sampling (Generation)

• Start: Sample pure noise $x_T \sim \mathcal{N}(0,I)$.
• Denoising: Iteratively apply the learned reverse process: $x_{t-1} \sim p_{\theta} (x_{t-1}|x_t)$.
• Result: Repeat until the clear data sample $x_0$ is obtained.

2. Problems of DMs:

Previously, DMs mostly used pixels as input data for computation. For high-resolution images, this required a huge amount of memory and time for training and inference, and in the end, it could only be applied to a limited range of applications. At the same time, not every pixel in an image contains important information; some are even classified as noise, which interferes with the output. To reduce computational costs, this work performs diffusion calculations in a compressed latent space. By adding cross-attention and using text embeddings as conditions, it allows image generation to be controlled with natural language.

3. Latent Diffusion Models (LDM)

ldms (e.g., Stable Diffusion) improve efficiency by applying the entire Denoising Diffusion process (including the Optimization phase) within a low-dimensional Latent Space, significantly reducing the computational cost compared to operating in the high-dimensional Pixel Space.

Operational structure and process:

This diagram illustrates the architecture and process of the Latent Diffusion Model (LDM). The model combines efficient perceptual compression with a powerful conditional diffusion model to achieve high-resolution image synthesis.

The operation of LDM can be divided into two independently trained stages, along with a core conditional generation mechanism.

Stage 1: Perceptual Compression and Reconstruction (The Autoencoder)

The goal of this stage is to compress high-dimensional image data into a lower-dimensional Latent Space while ensuring that sufficient visual information is preserved.

1. Encoder ($\mathcal{E}$)

• Input: The original image $x$ (in high-dimensional pixel space).
• Function: $\mathcal{E}$ compresses $x$ into a lower-dimensional latent representation $z$ ($z = \mathcal{E}(x)$). This latent space contains only the information most critical for visual perception, discarding minute, high-frequency details, which significantly reduces computational complexity.

2. Decoder ($\mathcal{D}$)

• Function: The role of $\mathcal{D}$ is to reconstruct the high-resolution image $\hat{x}$ from the latent representation $z$ ($\hat{x} = \mathcal{D}(z)$).
• Importance: The decoder ensures the quality and detail fidelity of the final output image during the last step of the generation process.

Stage 2: Generative Modeling in the Latent Space (Core Diffusion Process)

The core generative component of the model, the Diffusion Model (DM), is trained and operates within the low-dimensional latent space $z$.

3. Diffusion Model Backbone (UNet $\epsilon_{\theta}$)

• Input: The noisy latent representation $z_t$ (at time step $t$).
• Function: This is a time-conditioned U-Net network, $\epsilon_{\theta}$. Its task is to predict and remove the noise $\epsilon$ from $z_t$, thereby performing the reverse denoising process which ultimately recovers the clear latent representation $z$ from pure noise.
• Efficiency Advantage: Because the UNet operates in the lower-dimensional latent space, its training and inference computational cost is far lower than traditional DMs that operate directly in pixel space.

Stage 3: Conditioning Mechanism

One of LDM’s biggest advantages is its flexible conditioning capability, allowing users to guide the image generation through various types of inputs, such as text, semantic maps, or class labels.

4. Domain-Specific Encoder ($\tau_{\theta}$)

• Input: The conditional input $y$ (e.g., a text description).
• Function: The $\tau_{\theta}$ (e.g., a Transformer for text, or a Convolutional Net for image inputs) encodes the condition $y$ from different modalities into a Context Representation $C = \tau_{\theta}(y)$. $C$ contains the semantic information of the conditional input.

5. Cross-Attention

• Function: This is the key mechanism that connects the UNet backbone with the context $C$. Cross-Attention layers are inserted into the intermediate layers of the UNet.
• Operation:
  • It allows the intermediate features of the UNet (derived from the latent representation $z$) to serve as the Query (Q).
  • The context representation $C$ serves as the Key (K) and Value (V).
  • During each denoising step, the UNet can “query” the relevant semantic information within the context $C$ to dynamically guide the denoising process, ensuring that the generated latent representation $z$ (and subsequently the output image $\hat{x}$) adheres to the condition $y$‘s description.

Summary Flow (Conditional Generation)

1. Compression (Pre-trained): The original image $x$ is compressed to the latent $z$ via $\mathcal{E}$ (used during the training phase).

2. Encoding (Conditioning): The conditional input $y$ is encoded into the context $C$ via $\tau_{\theta}$.

3. Denoising (Generation): Starting from random noise, the UNet $\epsilon_{\theta}$ performs multiple denoising iterations in the latent space. In each iteration, the UNet references the context $C$ via Cross-Attention layers to guide the prediction and removal of noise.

4. Decoding (Reconstruction): Once the latent representation $z$ is near its noise-free state, it is passed through the $\mathcal{D}$ to output the final high-resolution image $\hat{x}$.

4. U-Net

U-Net construction:

This image illustrates the U-Net architecture, a highly successful Convolutional Neural Network (CNN) specifically designed for biomedical image segmentation. The name “U-Net” comes from its distinct “U” shape, composed of two symmetrical paths: a contracting path and an expanding path.

1. Contracting Path (The Left Side of the U)

This is the standard architecture for a convolutional network, used for feature extraction and context capturing.

• Operations:
  • Repeated $3\times3$ Convolutions (conv $3\times3$, ReLU, Blue Arrows): Applied twice, followed by a Rectified Linear Unit (ReLU) activation. This increases the number of feature channels (e.g., from 64 to 128) while slightly reducing the feature map size (e.g., $572\times572$ to $570\times570$) due to unpadded convolution.
  • Max Pooling ($max\ pool\ 2\times2$, Red Arrows): Used for downsampling at the end of each block. This halves the size of the feature map (e.g., $568\times568$ to $284\times284$).
• Purpose: As the path deepens, the number of feature channels increases (from 64 up to 1024), while the feature map size decreases. This allows the network to capture increasingly abstract, larger-scale contextual information.

2. Expanding Path (The Right Side of the U)

This path is symmetric to the contracting path and is used for precise localization and recovering spatial resolution.

• Operations:
  • Up-Convolution (up-conv $2\times2$, Green Arrows): A $2\times2$ Transposed Convolution (often called upsampling) is performed to double the feature map size (e.g., $52\times52$ to $104\times104$). This halves the number of feature channels.
  • Copy and Crop (Gray Arrows): This is the key innovation of U-Net. The high-resolution feature map (with fine spatial details) from the corresponding layer of the contracting path is cropped and concatenated with the upsampled feature map.
  • Repeated $3\times3$ Convolutions (conv $3\times3$, ReLU, Blue Arrows): Two more $3\times3$ convolutions are applied to process the combined features, which now contain both context and fine details.
• Purpose: The process progressively restores the spatial resolution of the feature maps, utilizing the detailed information brought in by the skip connections to achieve precise pixel-level localization of the segmentation target.

3. Output Layer

• Operation:
  • $1\times1$ Convolution (conv $1\times1$, Cyan Arrow): In the final layer, a $1\times1$ convolution maps the feature channels to the desired number of classes (2 in the diagram, typically foreground and background).
  • Purpose: To generate a class prediction for every pixel in the input image, resulting in the final output segmentation map.

Key Advantage: Skip Connections

The gray arrows (copy and crop) represent the Skip Connections, which are vital for U-Net’s success:

1. Mitigate Information Loss: Max pooling in the contracting path loses spatial details (localization information).
2. Combine Detail and Context: The skip connection transfers the high-resolution, high-detail feature maps from the contracting path directly to the expanding path.
3. Enable Accurate Segmentation: The output of the expanding path thus contains both the contextual information (knowing “what” the object is) extracted by the deeper layers and the localization information (knowing “where” it is) from the shallower layers, leading to highly accurate pixel-wise segmentation.

In summary, U-Net successfully balances context capture (via the deep network) and precise localization (via the skip connections), making it a preferred architecture for various image segmentation tasks, especially in biomedical imaging where data is often limited.