Nagi-ovo

This article mainly organizes and elaborates on the teaching logic of the video, combined with the explanation 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, which achieves 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 the image 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 applying 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 denoise gradually? 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 proposes Denoising Diffusion Probabilistic Models (DDPM), significantly improving the generation quality and efficiency of diffusion models. By introducing a simple denoising network 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

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

2. Predict $x_0$ directly

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

3. Predict the added noise

   • Predicting the noise $ε$ added during the forward process.
   • This is mathematically equivalent to the first method (predicting the noise mean), just parameterized differently; they can be transformed 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 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:

We can see that the last few timesteps are close to complete noise, with very little information. Additionally, overall, the information is destroyed too quickly, so OpenAI used a Cosine Schedule to address these two issues:

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 using 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, which is achieved through sinusoidal positional encoding embeddings in the Transformer, which are projected into each Residual-Block. The model can also combine the Schedule 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 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 effect: $\mathcal{L}_{\text{reconstruction}} = \|X - \text{Decoder}(\text{Encoder}(X))\|^2$

Architecture Improvement#

OpenAI significantly improved the 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 certain 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 uses a neural network to output a denoised image $x_{t-1}$:

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, $I$ is the identity matrix, indicating that the covariance matrix is diagonal and independent across dimensions, increasing the amount of noise added as the timestep progresses.

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 (like 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 gradient calculations.

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 the normal distribution,

• Without using reparameterization:

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

• Using reparameterization:

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

In terms 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 calculate 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 calculated

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 whether you added too much sugar or too little.

Using reparameterization:

1. First, prepare a cup of standard concentration sugar water ($\epsilon$).
2. Then adjust the amount of sugar water ($\mu$) and dilution level ($\sigma$) to achieve the target sweetness.
3. If it doesn't taste good, you can clearly identify whether the sugar water amount or dilution level needs adjustment (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 independent of model parameters, during backpropagation, the gradient only needs to consider the corresponding parameters of $\sqrt{1 - \beta_t}$ and $\sqrt{\beta_t}$. This allows the model to be effectively optimized through gradient descent.

We simplify the notation with $\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 start analyzing 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:

Now 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}$, that is:

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 contents of the expression are 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 obtain 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 state, where these terms require estimating all samples, leading to high variance. For example, given the $x_t$ shown below, it is difficult to determine what the previous state looked like:

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

This way, with the noise-free image provided, the candidates for $x_{t-1}$ are reduced, and the variance will decrease:

Substituting back into the original expression:

Expanding the second sum term reveals that most terms cancel out:

Applying log rules to the last two terms simplifies some terms:

Moving the simplified first term to the front and combining into a single 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 expression for $x_0$ into the predicted mean formula $\tilde{\mu}_t$:

Now $\mu$ no longer depends on $x_0$. Continuing to simplify, 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 have found that neglecting the scaling term 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$, so it can be simplified 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.

In this case, 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, and then sample $x_{t-1}$ using the previously shown formula through reparameterization.

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

At $t = 1$, the formula is used to recover $x_0$ from $x_1$, 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 final step (i.e., without the $\sqrt{\beta_t} \epsilon$ term) can avoid introducing unnecessary randomness when generating the final image, thus maintaining the clarity and detail of the image.

Code Implementation#

Recommended Zhihu Sunrise's simplified MLP implementation.
If I have time later, I might consider doing a code implementation of Stable Diffusion, leaving a note...

References#

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