Nagi-ovo

This article mainly follows the teaching logic of the video, organizing and explaining the content. If there are any errors, please feel free to correct them in the comments!

Intuitive Part#

Theoretical Support#

Deep Unsupervised Learning using Nonequilibrium Thermodynamics

This paper lays the theoretical foundation for diffusion models. The authors propose a generative model based on nonequilibrium thermodynamics, achieving data generation by gradually adding and removing noise. This provides important theoretical support for subsequent diffusion model research.

We apply a large amount of noise to images and then use a neural network to denoise. If this neural network learns well, it can start from completely random noise and ultimately obtain images from our training data.

Forward Diffusion Process:#

Iteratively apply noise to the image; when the steps are sufficient, the image will completely turn into noise, using a normal distribution as the noise source:

Reverse Diffusion Process#

From pure noise to image, involving a neural network that learns to denoise step by step.

Why is it gradual denoising? The authors mention in the paper that the result of "denoising in one step" is very poor.

So what does this network look like? What does it need to predict?

Algorithm Improvement#

Denoising Diffusion Probabilistic Models

This paper introduces Denoising Diffusion Probabilistic Models (DDPM), significantly improving the generation quality and efficiency of diffusion models. By introducing simple denoising networks and optimized training strategies, DDPM has become an important milestone in the field of diffusion models.

The authors discuss three targets that the neural network can predict:

1. Predict the mean of the noise at each timestep (predict the mean of the noise at each timestep)

   • That is, predict the mean of the conditional distribution $p(x_{t-1}|x_t)$
   • The variance is fixed and not learnable

2. Predict the original image (predict $x_0$ directly)

   • Directly predict the original, unpolluted image
   • Experiments show that this method performs poorly

3. Predict the added noise (predict the added noise)

   • Predict the noise $ε$ added during the forward process
   • This is mathematically equivalent to the first method (predicting the mean of the noise), just parameterized differently; they can be converted into each other through simple transformations.

The paper ultimately chooses predicting noise (the third method) as the main approach, as this method is more stable in training and performs better.

The amount of noise added at each step is not fixed, controlled by a Linear Schedule to prevent instability during training.

It looks something like this:

You can see that the last few timesteps are close to complete noise, with very little information; moreover, overall, the information is destroyed too quickly. Therefore, OpenAI used a Cosine Schedule to solve these two problems:

Model Architecture#

U-Net#

Along with this paper, a model architecture called U-Net was published:

This model has a Bottleneck (a layer with fewer parameters) in the middle, using Downsample-Block and Resnet-Block to project the input image to a lower resolution, and Upsample-Block to project it back to the original size during output.

At certain resolutions, the authors added Attention-Blocks and made Skip-Connections between layers in the same resolution space. The model is designed for each timestep, achieved through sinusoidal position encoding embeddings in Transformers, which are projected into each Residual-Block. The model can also combine schedules to remove different amounts of noise at different timesteps to enhance generation effects, which will be discussed in detail later.

Bottleneck and Autoencoder#

The concept of Bottleneck (bottleneck layer) was initially proposed and widely used in the unsupervised learning method "Autoencoder." As the lowest-dimensional hidden layer in the autoencoder architecture, it is located between the encoder and decoder, forming the narrowest part of the network, forcing the network to learn a compressed representation of the data, minimizing reconstruction error and acting as a regularization term: $\mathcal{L}_{\text{reconstruction}} = \|X - \text{Decoder}(\text{Encoder}(X))\|^2$

Architecture Improvement#

OpenAI significantly improved overall performance by enhancing the architecture in their second paper Diffusion Models Beat GANs on Image Synthesis:

1. Increase network depth (more layers) and reduce width (number of channels per layer)
2. Increase the number of Attention-Blocks
3. Expand the number of heads in each Attention Block
4. Introduce BigGAN-style Residual Blocks for upsampling and downsampling
5. Introduce Adaptive Group Normalization (AdaGN) to dynamically adjust normalization parameters based on conditional information (such as timestep)
6. Use Separate Classifier Guidance to help the model generate specific types of images

Mathematical Part#

Symbol Table#

• $X_t$ represents the image at timestep $t$, i.e., $X_0$ is the original image. It can be simplified as $t$ being smaller means less noise:

• The final image of noise is an isotropic (uniform in all directions) complete noise, denoted as $X_T$. In the initial studies, $T=1000$, and subsequent work has reduced this significantly:

• Forward Process: $q(x_t|x_{t-1})$, inputting the $x_{t-1}$ image outputs a noisier image $X_t$:

• Backward Process: $p(x_{t-1}|x_t)$, inputting the $x_t$ image outputs a denoised image $x_{t-1}$ using a neural network:

Forward Process#

Where,

• $\sqrt{1 - \beta_t} x_{t-1}$ is the mean of the distribution; $\beta_t$ is the noise schedule parameter, ranging from 0 to 1, combined with $\sqrt{1 - \beta_t}$ to scale the noise, decreasing as the timestep increases, representing the retained portion of the original signal.

• $\beta_t I$ is the covariance matrix of the distribution, where $I$ is the identity matrix, indicating that the covariance matrix is diagonal and independent in each dimension, increasing the amount of noise added as the timestep increases.

Now we just need to iterate this step continuously to get the result after 1000 steps, but actually, these can be completed in one step.

Reparameterization Trick#

The reparameterization trick is very important in diffusion models and other generative models (such as Variational Autoencoders, VAE). Its core idea is to transform the sampling process of a random variable into a deterministic function plus a standardized random variable. This transformation allows the model to be optimized through gradient descent because it eliminates the randomness in the sampling process from affecting the gradient calculation.

Here is a simple example to explain its significance.

You can implement a dice roll in two ways:

• The first way has randomness inside the function:

# 1. Directly rolling the dice (random sampling)
def roll_dice():
    return random.randint(1, 6)

result = roll_dice()

• The second way separates randomness outside the function, making the function itself deterministic:

# 2. Separating randomness
random_number = random.random()  # Generate a random number between 0 and 1

def transformed_dice(random_number):
    # Map the random number from 0-1 to 1-6
    return math.floor(random_number * 6) + 1

result = transformed_dice(random_number)

In probability theory, we learned that if $X$ is a random variable, and $X ∼ 𝒩(0,1)$, then: $aX + b ∼ 𝒩(b, a²)$

Thus, for a normal distribution $\mathcal{N}(\mu, \sigma^2)$, samples can be generated as follows:

where $\epsilon \sim \mathcal{N}(0, 1)$ is a standard normal distribution.

Similarly, in a normal distribution,

• Without using reparameterization:

# Directly sampling from the target normal distribution
x = np.random.normal(mu, sigma)

• Using reparameterization:

# First sample from the standard normal distribution
epsilon = np.random.normal(0, 1)
# Then obtain the target distribution through a deterministic transformation
x = mu + sigma * epsilon

In the context of gradient calculations involved in model training,

Without using reparameterization:

def sample_direct(mu, sigma):
    return np.random.normal(mu, sigma)

# In this case, it's difficult to compute the gradient with respect to mu and sigma
# because the random sampling operation blocks gradient propagation

Using reparameterization:

def sample_reparameterized(mu, sigma):
    epsilon = np.random.normal(0, 1)  # Gradient does not need to propagate through here
    return mu + sigma * epsilon        # Gradients for mu and sigma can be easily computed

Taking VAE (Variational Autoencoder) as an example:

class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.encoder = Encoder()  # Outputs mu and sigma
        self.decoder = Decoder()

    def reparameterize(self, mu, sigma):
        # Reparameterization trick
        epsilon = torch.randn_like(mu)  # Sample from standard normal distribution
        z = mu + sigma * epsilon        # Deterministic transformation
        return z

    def forward(self, x):
        # Encoder outputs mu and sigma
        mu, sigma = self.encoder(x)
        
        # Use reparameterization to sample
        z = self.reparameterize(mu, sigma)
        
        # Decoder reconstructs the input
        reconstruction = self.decoder(z)
        return reconstruction

Reparameterization from a Foodie's Perspective#

Imagine you are making milk tea:

Without using reparameterization:

• Directly making a cup of milk tea with a specific sweetness
• If it doesn't taste good, you don't know if you added too much or too little sugar

Using reparameterization:

1. First, prepare a cup of standard sugar water ($\epsilon$)
2. Then adjust the amount of sugar ($\mu$) and dilution level ($\sigma$) to achieve the desired sweetness
3. If it doesn't taste good, you can clearly identify whether to adjust the sugar water amount or the dilution level (parameters can be optimized)

In summary, through reparameterization:

• Gradients can propagate through deterministic transformations
• Parameters can be optimized through gradient descent
• Randomness is isolated, not affecting gradient calculations

Forward Mathematical Derivation#

Transition from $x_{t-1}$ to $x_t$:

• Given $x_{t-1}$, we want to generate $x_t$.
• Using the reparameterization trick on $q(x_t \mid x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t I)$,

We can express $x_t$ as a deterministic transformation of $x_{t-1}$ plus a noise term:

• Here, $\sqrt{1 - \beta_t} x_{t-1}$ is the mean part, and $\sqrt{\beta_t} \epsilon$ is the noise part. Since $\epsilon$ is a sample from the standard normal distribution and is independent of model parameters, during backpropagation, the gradient only needs to consider the parameters corresponding to $\sqrt{1 - \beta_t}$ and $\sqrt{\beta_t}$. This allows the model to be effectively optimized through gradient descent.

We simplify notation using $\alpha_t$ and record its cumulative product:

We obtain:

Calculating the two-step transition: from $x_{t-2}$ to $x_t$

Since $\epsilon_{t-1}$ and $\epsilon_t$ are independent standard normal distributions, we can combine the noise parts into a new noise term $\epsilon \sim \mathcal{N}(0, I)$:

Similarly:

The complete derivation process is as follows:

Reverse Mathematical Derivation#

Since the variance is fixed and does not need to be learned (see 1.3), we only need the neural network to predict the mean:

Our ultimate goal is to predict the noise in two timesteps. Now let's analyze from the loss function:

However, the probability of $x_0$ in this negative log-likelihood depends on all other previous timesteps. We can learn a model that approximates these conditional probabilities as a solution, which requires using the Variational Lower Bound to obtain a more computable formula.

Variational Lower Bound#

Assuming we have an uncomputable function $f(x)$, in our scenario, it is the negative log-likelihood, we can find a computable function $g(x)$ that always satisfies $g(x) \leq f(x)$. Then optimizing $g(x)$ can also increase $f(x)$:

We ensure this by subtracting the KL divergence, which is a measure of the similarity between two distributions and is always non-negative:

Subtracting a term that is always non-negative ensures that the result is always less than the original function. Here, we use "+" because we want to minimize the loss, so adding it guarantees that the original negative log-likelihood remains large:

In this form, since the negative log-likelihood is still present, the lower bound remains uncomputable. We need to obtain a better expression. First, rewrite the KL divergence as the log ratio of two terms:

Next, apply Bayes' theorem to the denominator:

The numerator $p_\theta(x_0 \mid x_{1:T}) p_\theta(x_{1:T})$ is actually the joint probability $p_\theta(x_0, x_{1:T})$ because:

Typically, $p_\theta(x_{0:T})$ represents the joint probability of $x_0$ and all intermediate steps $x_{1:T}$, i.e.:

Substituting gives:

Thus, following the flow in the diagram below, we arrive at the final form:

Suddenly, the annoying two terms cancel out:

This gives us a lower bound that can be minimized, and the content in the equation is all known:

• The numerator is the joint probability distribution of the forward process: $q(x_{1:T} \mid x_0)=\prod_{t=1}^T q(x_t \mid x_{t-1})$;
• The denominator is the joint probability distribution of the reverse process: $p_\theta (x_{0:T})=p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} \mid x_t)$

To have an analytical solution, a few additional reorganization steps are needed:

Rewriting the numerator of the sum term using Bayes' theorem: $q(x_t \mid x_{t-1})=\frac{q(x_{t-1}\mid x_t)q(x_t)}{q(x_{t-1})}$

But this returns to the previous point; these terms require estimating all samples, leading to high variance. For example, given the $x_t$ shown in the diagram below, it is difficult to determine what the previous state looked like:

The improvement idea is to condition directly on the original data $x_0$:

This way, providing the noise-free image reduces the candidates for $x_{t-1}$, thus reducing variance:

Substituting back into the original expression:

Expanding the second sum term, we find that most terms simplify:

Applying log rules to the last two terms can simplify some terms:

Then moving the simplified first term to the front and merging into one log gives the final analytical form:

The first term in this form can be ignored because $q$ has no learnable parameters; it is just the forward process of adding noise, which will converge to a normal distribution, while $p(x_T)$ is just noise sampled from a Gaussian distribution. Therefore, it can be determined that this KL divergence term will be very small.

The remaining two terms' derivation results are as follows (process omitted, see Lilian's Blog)

$\beta$ is fixed, so we focus on the form of $\mu$:

The closed form generated by the forward process $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$ can be rewritten in terms of $x_0$:

Substituting the above expression for $x_0$ into the predicted mean formula $\tilde{\mu}_t$:

Now $\mu$ no longer depends on $x_0$. Continuing to simplify, we first expand the second term:

Combining the $x_t$ terms:

Further combining and simplifying the coefficients of $x_t$, we ultimately obtain:

This indicates that we are essentially just subtracting the randomly scaled noise generated by $x_t$, which is what the neural network needs to predict.

The loss function $L_t$ defined after substitution is a mean squared error:

The final form is the mean squared error between the actual noise at timestep $t$ and the noise predicted by the neural network. Researchers found that ignoring the scaling terms in front yields better sampling quality and is easier to implement.

Returning to the original formula

The authors decided not to add extra random noise in the final sampling step to make the generation process more stable:

The final form is:

• $\mathbb{E}_{t, \mathbf{x}_0,\boldsymbol{\epsilon}}$ represents the expectation over timestep $t$, original data $\mathbf{x}_0$, and noise $\boldsymbol{\epsilon}$
• $\boldsymbol{\epsilon}$ is the actual added random noise
• $\boldsymbol{\epsilon}_\theta$ is the noise predicted by the neural network
• $\sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}$ is the closed-form solution of the forward process, representing the noisy data at timestep $t$, thus simplifying to:
  • $\mathbf{x}_t$ directly represents the noisy data at timestep $t$, i.e., $\sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}$
  • The entire loss function essentially measures the mean squared error between the predicted noise and the actual noise.

Where the timestep $t$ is typically sampled from a uniform distribution (i.e., $t∼Uniform(1,T)$, where T is the total number of timesteps). This choice ensures that during training, each timestep has an equal probability of being selected, allowing the model to effectively learn the denoising process across all timesteps.

Training#

First, we sample some images from the dataset, then sample $t$ and noise from a normal distribution, and optimize the objective through gradient descent.

Sampling#

First, sample $x_t$ from a normal distribution, then use the previously shown formula to sample $x_{t-1}$ through reparameterization.

Note that at $t=1$, no noise is added. According to the formula

At $t = 1$, this formula is used to recover from $x_1$ to $x_0$, which is the last step of the denoising process. At this point, we want to reconstruct the original image as accurately as possible. Not adding noise in the last step (i.e., without the $\sqrt{\beta_t} \epsilon$ term) avoids introducing unnecessary randomness when generating the final image, thus preserving the clarity and detail of the image.

Code Implementation#

Recommended Zhihu Sunrise's simplified MLP implementation
If there is time later, consider doing a code implementation of Stable Diffusion, leaving a pit...

References#

• Diffusion Models | Paper Explanation | Math Explained
• Lilian Weng: From Autoencoder to Beta-VAE
• Lilian Weng: What are Diffusion Models?