import Comments from '../../components/Comments'

Introduction

In previous discussion, we have shown power of pixel-based idiffusion models on a variety of dataset and tasks such as image synthesis and sampling. These models achieved state-of-the-art synthesis quality. In this post, we are going to discuss some recent works on conditional generation which means guide the generation process with additional conditions. A naive solution is to train a diffusion model specific on certain dataset and generate samples with it. However, more commonly, we want to generate samples conditioned on class labels or a piece of descriptive text. Building on this, an more sophisticated method is add label $y$ into input and, therefore, the diffusion process will take class information into consideration. However, due to performance reason, there are multiple algorithm have been proposed to achieve higher generation quality.

Classifier Guided Diffusion

In order to explicit utilize class label information to guide the diffusion process, Dhariwal and Nichol applying the gradient of a trained classifier to guide the diffusion sampling process.

There are three important components in this approach: classifier training, incorporate label information into diffusion model training and classifier-guided sample generation.

Classifier training: A classifier $p (y ∣ x)$ can be exploited to improve a diffusion generator by providing gradient $\nabla_{x} p (y ∣ x)$ to the sampling process. Since the generated images at intermediate steps are noisy, the trained classifier should be able to adapt to these noises. Therefore, the classifier $p_{ϕ} (y ∣ x_{t}, t)$ is trained on noisy images $x_{t}$ and then use gradients $\nabla_{x_{t}} lo g p_{ϕ} (y ∣ x_{t}, t)$ to guide the diffusion sampling process towards an arbitrary class label y.
Adaptive group normalization: The paper incorporated adaptive group normalization layer $AdaGN (h, y) = y_{s} GroupNorm (h) + y_{b}$ into the neural network, where $h$ is the output of previous hidden layer and $y = [y_{s}, y_{b}]$ is obtained from a linear projection of the timestep and class embedding.
Conditional reverse noising process: The paper proved that the reverse transition distribution can be written in the form as $p_{θ, ϕ} (x_{t} ∣ x_{t + 1}, y) = Z p_{θ} (x_{t} ∣ x_{t + 1}) p_{ϕ} (y ∣ x_{t})$ .\ This can be observed from the following relationship: $q (x_{t} ∣ x_{t + 1}, y) = = = = = \frac{q ( x _{t} , x _{t + 1} , y )}{q ( x _{t + 1} , y )} q (y ∣ x_{t}, x_{t + 1}) \frac{q ( x _{t} , x _{t + 1} )}{q ( x _{t + 1} , y )} (\frac{q ( x _{t + 1} ∣ x _{t} , y ) q ( x _{t} , y )}{q ( x _{t} , x _{t + 1} )}) \frac{q ( x _{t} ∣ x _{t + 1} )}{q ( y ∣ x _{t + 1} )} (\frac{q ( x _{t + 1} ∣ x _{t} ) q ( y ∣ x _{t} )}{q ( x _{t + 1} ∣ x _{t} )}) \frac{q ( x _{t} ∣ x _{t + 1} )}{q ( y ∣ x _{t + 1} )} \frac{q ( x _{t} ∣ x _{t + 1} ) q ( y ∣ x _{t} )}{q ( y ∣ x _{t + 1} )},$ where $q (y ∣ x_{t + 1})$ can be viewed as a constant since it does not contain $x_{t}$ .

We can write the reverse process(step 3) in DDIM's language. Recall that $\nabla_{x_{t}} lo g p_{θ} (x_{t}) = - \frac{1}{1 - α ˉ _{t}} ϵ_{θ} (x_{t}, t)$ and we can write the score function for the joint distribution of $(x_{t}, y)$ as follows,

\nabla_{x_{t}} lo g q (x_{t}, y) = \nabla_{x_{t}} lo g q (x_{t}) + \nabla_{x_{t}} lo g q (y ∣ x_{t}) \approx - \frac{1}{1 - α ˉ _{t}} ϵ_{θ} (x_{t}, t) + \nabla_{x_{t}} lo g p_{ϕ} (y ∣ x_{t}) = - \frac{1}{1 - α ˉ _{t}} (ϵ_{θ} (x_{t}, t) - 1 - \overset{α}{ˉ}_{t} \nabla_{x_{t}} lo g p_{ϕ} (y ∣ x_{t})) .

Therefore, we obtained the new noise prediction $\overset{ϵ}{^} (x_{t})$ as

\overset{ϵ}{^} (x_{t}) := ϵ_{θ} (x_{t}) - 1 - \overset{α}{ˉ}_{t} \nabla_{x_{t}} lo g p_{ϕ} (y ∣ x_{t}) .

The paper provided detailed algorithms based on DDPM and DDIM.

Fig.1 - DDPM and DDIM based classifier-guided diffusion model.

Classifier-Free Guidance Diffusion model

Since training an independent classifier $p_{ϕ} (y ∣ x)$ involved extra effort, Ho and Salimans proposed algorithm to run conditional diffusion steps without an independent classifier. The paper incorporated the scores from a conditional and an unconditional diffusion model. The method includes two components:

Replace the previously trained classifier with the implicit classifier according to Bayesian Rule (opens in a new tab).

\nabla_{x_{t}} lo g p (y ∣ x_{t}) = \nabla_{x_{t}} lo g p (x_{t} ∣ y) - \nabla_{x_{t}} lo g p (x_{t}) = - \frac{1}{1 - α ˉ _{t}} (ϵ_{θ} (x_{t}, t, y) - ϵ_{θ} (x_{t}, t))

Use a single neural network to function as two noise generators-- a conditional one and a unconditional one. It can be done by let $ϵ_{θ} (x_{t}, t) = ϵ_{θ} (x_{t}, t, y = \emptyset)$ for unconditional generation and $ϵ_{θ} (x_{t}, t, y)$ for conditional generation towards class label $y$ . Therefore, the new noise generation function can be deduced as follows.

\overline{ϵ}_{θ} (x_{t}, t, y) = ϵ_{θ} (x_{t}, t, y) - 1 - \overset{α}{ˉ}_{t} w \nabla_{x_{t}} lo g p (y ∣ x_{t}) = ϵ_{θ} (x_{t}, t, y) + w (ϵ_{θ} (x_{t}, t, y) - ϵ_{θ} (x_{t}, t)) = (w + 1) ϵ_{θ} (x_{t}, t, y) - w ϵ_{θ} (x_{t}, t)

Latent Diffusion Models

Operating on pixel space is exceptional costful. For algorithms like diffusion models, it is even more demanding, since the recursive updates amplified this cost. A common solution in ML to deal with high dimensionality is embedding data into lower dimensional latent space. It is observed in Rombach et al. that most bits of an image contribute to perceptual details and the semantic and conceptual composition remains intact after undergoing aggressive compression. This motivates Rombach et al. to first embed the image into latent space, with models like VAE, then train a diffusion model in latent space. Moreover, it loosely decomposes the perceptual compression (removing high-frequency details) and semantic compression (semantic and conceptual composition of the data). In practice, a VAE can be used first to trimming off pixel-level redundancy and an U-Net backboned diffusion process can be used to learn to manipulate semantic concepts.

Fig.1 - The plot shows the [rate-distortion trade-off](https://en.wikipedia.org/wiki/Rate%E2%80%93distortion_theory). Notice that semantic compression and perceptual compression happens at different stages in the plot. The graph is taken from Rombach et al. Bits per dim and RMSE measures different aspect for generative model. Details can be found in Theis et al. [Simple explaination.](https://blog.csdn.net/fanyue1997/article/details/109703025)

Methods

The perception compression process is depended on an autoencoder model. And encoder $E$ encodes an image $x \in R^{H \times W \times 3}$ in RGB space into a latent representation $z = E (x)$ , and an decoder $D$ reconstructs the image from its latent $\tilde{x} = D (z) = D (E (x))$ . In contrary to other previous work, the paper use a two dimensional latent space to better suit the follow up diffusion model. The paper explored two types of regularization in autoencoder to avoid arbitrarily high-variance in the latent space.

KL-reg: A small KL penalty towards a standard normal.
VQ-reg: Uses a vector quantization layer within the decoder, like VQVAE but the quantization layer absorbed by the decoder.

The semantic compression stage happens in the latent space. After the autoencoder, the paper construct a diffusion model in latent space with U-Net being the backbone neural network. Denote the backbone neural network as $ϵ_{θ} (\circ, t)$ and the loss function is

L_{L D M} := E_{E (x), ϵ \sim N (0, 1), t} [∥ ϵ - ϵ_{θ} (z_{t}, t) ∥_{2}^{2}] .

As in many other generative models and the topic of this blog, conditional mechanisms can be applied to this framework and, to be more specific, in the latent space. The paper implemented this by adding the additional inputs $y$ to the denoising autoencoder as $ϵ_{θ} (z_{t}, t, y)$ . The additional inputs $y$ can be text, semantics maps or other "embedible information" like images and it aims to controll the synthesis process.

Due to the various modalities of the inputs, the paper first project the inputs $y$ to an "intermediate representation"(embedding) $τ_{θ} (y) \in R^{M \times d_{τ}}$ .
Cross-attention layer is used to apply controlling signal $τ_{θ} (y)$ to the diffusion process through U-Net backbone. To be more specific, $Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d})$ , with

Q = W_{Q}^{(i)} \cdot φ_{i} (z_{t}), K = W_{K}^{(i)} \cdot τ_{θ} (y), V = W_{V}^{(i)} \cdot τ_{θ} (y) .

Based on image-conditioning pairs, we then learn the conditional LDM via the loss

L_{LDM} := E_{E (x), y, ϵ \sim N (0, 1), t} [∥ ϵ - ϵ_{θ} (z_{t}, t, τ_{θ} (y)) ∥_{2}^{2}] .

Experiments

The paper examine the model performance in two aspects:

Enerated samples' perceptual quality and training efficiency
Sampling efficiency

Perceptual Compression Tradeoffs

Fig.3 Image samples generated from LDM-8(KL) on LAION dataset with 200 DDIM steps. Perceptually, the LDM is capable of generating high fidelity images under the prompt constraint. Image source: Rombach et al.

In this experiment, the paper compared FID/Inception scores under different downsample rate at different training steps. The experiment observed that LDM, with proper downsample rate, achieved siginificant better FID scores comparing with pixel-based diffusion model.

Fig.4 FID and Inception score under different setup. LDM-k where k means downsample coefficient. As training steps increase, the both scores improves under all circumstance. Meanwhile, LDM-8 achieved siginificant better quality than LDM-1(pixel-based diffusion) and achieved faster. Image source: Rombach et al.

Sampling Efficiency

Fig.5 Log FID vs throughput. Left CelebA-HQ and right ImageNet. As we expected, updating reverse diffusion steps in lower dimensional embedding space is less costful and have higher throughput. Moreover, LDM-32 is not only generating samples much faster(visually, at x20) than LDM-1 and LDM-2, it also generate samples with higher quality(in terms of FID). Image source: Rombach et al.

The LDM also demonstrated better sampling efficiency. Moreover, it generates samples faster and at a higher quality.

References

(Citations removed)

Diffusion Models for 3D Shape Generation Notes on Diffusion Models (I) -- DDPM and DDIM