Fully Annotated Guide to "What are Diffusion Models?"

This article provides a comprehensive annotated guide to Lilian Weng’s “What are Diffusion Models?” ¹.

When writing this article, I had limited knowledge of image generation models and related mathematical background. If you have a similar background, this article should be suitable for you. Of course, I may have made some elementary mistakes due to this. Your corrections are welcome.

The structure of this article is largely consistent with the original:

• Key formulas and concepts in each section will be expanded with derivations or explanations.
• In addition, common symbol explanations are provided at the beginning of the article.
• And some background knowledge required to understand the original text is appended at the end.

Note: If you only want to learn the basics of the Diffusion model, the first section, What are Diffusion Models? is sufficient.

Notations

What are Diffusion Models?

The basic principle of diffusion models is to increase noise through forward diffusion to obtain samples following a pure Gaussian distribution. Train a model to mimic the reverse diffusion process so that it can recover real data samples from any Gaussian noise samples.

Forward diffusion process

These are two expressions of the forward diffusion process: the single-step diffusion process and the overall diffusion process.

In the single-step diffusion process:

• $\mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t\mathbf{I})$ indicates that $\mathbf{x}_t$ follows a Gaussian distribution with a mean of $\sqrt{1 - \beta_t} \mathbf{x}_{t-1}$ and a variance of $\beta_t\mathbf{I}$.
• $\beta_t$ is the noise variance schedule parameter. It corresponds to a variance schedule $\{\beta_t \in (0, 1)\}_{t=1}^T$, similar to the learning rate schedule.
• $\beta_t$ defines the variance at each time step in the diffusion process. Generally, $\beta_t$ increases gradually. As a result, the difference from the original data becomes larger ($\sqrt{1 - \beta_t}$ decreases), the data variability also increases ($\beta_t\mathbf{I}$ increases), and overall, more noise is added at each step.
• $\beta_t\mathbf{I}$ is the covariance matrix, which is also a diagonal matrix. All diagonal elements are $\beta_t$. The same intensity of noise is added to each dimension.

The overall diffusion process is a recurrence formula obtained by multiplying the single-step diffusion processes together according to the Markov property.

The overall diffusion process is necessary because it helps us quickly sample the final pure noise $\mathbf{x}_T$ from the real data distribution. However, it relies on the recurrence formula and is computationally expensive. Therefore, in practice, a simpler calculation method, i.e., the closed-form formula described below, is used.

A closed-form expression refers to a formula that can be directly written as a solution using a finite number of explicit mathematical expressions without iteration, numerical approximation, or solving equations ².

Based on the single-step diffusion process $q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t\mathbf{I})$ and the reparameterization trick $z = \mu + \sigma \cdot \epsilon$, we can rewrite the single-step diffusion process as:

This allows us to provide a more detailed derivation of the closed-form:

(*) Recall that when we merge two Gaussians with different variances, $\mathcal{N}(\mathbf{0}, \sigma_1^2\mathbf{I})$ and $\mathcal{N}(\mathbf{0}, \sigma_2^2\mathbf{I})$, the new distribution is $\mathcal{N}(\mathbf{0}, (\sigma_1^2 + \sigma_2^2)\mathbf{I})$. Here, the merged standard deviation is $\sqrt{\alpha_t (1-\alpha_{t-1}) + (1 - \alpha_t)} = \sqrt{1 - \alpha_t\alpha_{t-1}}$.

PS. It’s worth noting that $\mathbf{x}_t$ is an intermediate noisy sample that satisfies two conditions.

One is $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}$, which describes how the sample changes from the perspective of diffusion.

The other is its corresponding probability $q(\mathbf{x}_t \vert \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I})$, which indicates its probability density.

These two aspects are complementary.

Langevin dynamics is a statistical method in physics for simulating molecular motion. It describes the random disturbances (e.g., thermal noise) experienced by particles moving in a potential field and is often used to model the stochastic behavior of complex systems.

Stochastic Gradient Langevin Dynamics (SGLD) is a sampling method that combines Langevin dynamics with stochastic gradient descent (SGD) in machine learning. Its goal is to sample from a probability distribution $p(x)$ without knowing the specific form of this distribution, only requiring knowledge of its gradient.

The above sampling formula is an iterative formula. Its meaning is: “Move a little in the gradient direction and add some random disturbances to make the final sample distribution approximate the target distribution $p(x)$.” The meanings of the relevant symbols are as follows:

• $\mathbf{x}_t$: The sample at the $t$-th step.
• $\frac{\delta}{2} \nabla_\mathbf{x} \log p(\mathbf{x}_{t-1})$: The drift term. It moves according to the gradient of the target distribution, similar to being pulled by a force. It can also be analogized to $\sqrt{1 - \beta_t} \mathbf{x}_{t-1}$ in diffusion.
  • $\delta / 2$: The step size, controlling the magnitude of each update.
  • $p(x)$: The probability density function of the target distribution.
  • $\log p(x)$: The log-probability density, facilitating calculation and optimization.
  • $\nabla_\mathbf{x} \log p(\mathbf{x}_{t-1})$: The gradient of the log-probability density, also called the score function, representing the “uphill direction” at the current point.
• $\sqrt{\delta} \boldsymbol{\epsilon}_t$: The diffusion term, similar to the molecular collisions in Brownian motion. It can be analogized to $\sqrt{\beta_t} \boldsymbol{\epsilon}_{t-1}$ in diffusion.
  • $\sqrt{\delta}$: The step size, controlling the magnitude of each update.
  • $\epsilon_t \sim \mathcal{N}(0, I)$: Random noise following a standard normal distribution, adding randomness to avoid getting stuck in local optima.

Note: The $p(\cdot)$ mentioned here is a general target distribution and can be any distribution we want to sample from. It is different from the $q(\cdot)$ and $p_\theta(\cdot)$ we see in diffusion.

For the diffusion scenario, if we want to generate more realistic samples, we have:

Here, $q(\cdot)$ measures the authenticity of the samples. In each iteration, $\mathbf{x}_t$ is more realistic than $\mathbf{x}_{t-1}$. Meanwhile, $\boldsymbol{\epsilon}_t$ prevents the generated samples from getting stuck in local optima.

Reverse diffusion process

Let’s review the forward single-step diffusion formula:

When $\beta_t$ is small, it means that little noise is added at each step. Then:

• The relationship between $\mathbf{x}_t$ and $\mathbf{x}_{t-1}$ is very close to a linear transformation plus a small perturbation.
• A linear transformation of a Gaussian distribution still maintains a Gaussian form.
• This makes the reverse conditional distribution approximately Gaussian.

Therefore, we usually model the reverse process using a Gaussian distribution.

The above two formulas correspond to the likelihood formulas for the overall and single-step reverse diffusion processes, which are the form of the neural network we are going to build.

Since we model the reverse diffusion process as a Gaussian distribution, its learnable parameters are the mean and variance of the Gaussian distribution, $\boldsymbol{\mu}_\theta(\mathbf{x}_t, t)$ and $\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t)$.

Let’s take a look at the four important distributions involved in the training and inference of the diffusion model:

• The forward diffusion generates noise.
• The posterior generates the training target for the neural network, i.e., what the reverse diffusion neural network needs to learn at each step.
• The reverse diffusion likelihood is used to fit the posterior.
• According to the prior, we generate pure noise following a standard normal distribution and pass it to the fitted network for denoising to generate images.

Since we model the reverse diffusion process as a Gaussian distribution, we can first define the formula for the posterior as above. Then the problem is transformed into how to derive ${\tilde{\boldsymbol{\mu}}_t}(\mathbf{x}_t, \mathbf{x}_0)$ and $\tilde{\beta}_t$. We will derive them in two steps later, and their specific forms are:

First, we can use Bayes’ formula to transform the calculation of the posterior into the calculation of the prior, i.e., the calculation of the forward diffusion. This allows us to build on the previous derivations.

Second, we can expand the probabilities into calculations between Gaussian probability density functions to obtain a new Gaussian probability density form.

Let’s elaborate on the reasoning steps based on these two ideas:

(*) According to the Gaussian probability density function, we can expand it and perform linear simplification. Linear simplification is just a notational optimization because it does not affect our ability to obtain a new Gaussian distribution later. As long as the new Gaussian distribution has a similar form to the current one, it is acceptable.

(**) $C(\mathbf{x}_t,\mathbf{x}_0)$ does not contain $\mathbf{x}_{t-1}$. Since we are calculating $q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)$ and $\mathbf{x}_{t}$ and $\mathbf{x}_{0}$ are known, $C(\mathbf{x}_t,\mathbf{x}_0)$ is just a constant term. Moreover, since it is in the exponent, it can be factored out as a coefficient. Finally, when we use it to calculate the final loss, as it is inside the KL divergence and thus inside the logarithm, it can be factored out as a separate constant term. The gradient of a constant term is 0, so it can be ignored here.

Since

And according to the previous calculation:

We can obtain:

At this point:

This completes the proof.

(*) We can ignore the coefficient term here because, as mentioned earlier, the coefficient term will eventually become a constant term in the loss. Since the derivative of a constant is 0, it can be ignored.

The significance of this step is to make the calculation rely solely on noise rather than real data, enabling the direct recovery of real data from any noise.

The nice property mentioned here is the closed-form forward diffusion: $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}$.

Let’s provide a more detailed derivation:

This completes the proof.

We have now derived the formula for the posterior, which serves as the golden truth that the neural network needs to mimic. So how do we model the relationship between the golden truth and the model?

This is where variational inference (or ELBO, Evidence Lower Bound) comes in. The VLB can be understood as the objective function or the loss function. In short, we hope the VLB to be smaller so that the trained model can more easily generate realistic images.

Here is a brief introduction to the variational lower bound:

1. The log marginal likelihood, $\log p_\theta(\mathbf{x}_0)$, is what we expect from the model, i.e., we hope the model has a higher probability of generating real images.
2. However, it is difficult to calculate directly because we cannot directly integrate over all cases in the latent variable distribution.
3. Therefore, we need to find an alternative optimization lower bound. Optimizing this lower bound is equivalent to optimizing the log marginal likelihood.
4. If you want to fully understand the relevant concepts, I strongly recommend reading another article by Lilian Weng, “From Autoencoder to Beta-VAE” ³, specifically the section VAE: Variational Autoencoder.

Let’s re-derive the variational lower bound (VLB) with more details.

At this point, we have transformed the difficult-to-calculate log marginal likelihood into an expression involving $q(\mathbf{x}_{1:T} | \mathbf{x}_0)$ and $p_\theta(\mathbf{x}_{0:T})$. The former can be directly calculated using the posterior formula we derived earlier, and the latter is defined by the neural network and can be decomposed into conditional probabilities at each step.

For training, we cannot run the model on only one image $\mathbf{x}_0$. We also need to sample using the Monte Carlo method. Therefore, we have:

This completes the proof.

The derivation mainly uses marginalization in probability theory and Jensen’s inequality. It is easy to prove after understanding these two concepts.

Marginalizing over $\mathbf{x}_{1:T}$ gives:

Jensen’s inequality ⁴ states that if $\phi(\cdot)$ is a concave function ⁵ and $X$ is an integrable random variable, then the following inequality holds:

For example, $\log(\cdot)$ is a concave function.

Why expand the concise VLB into the final complex formula consisting of multiple KL divergences?

Recall the closed-form expression of the forward diffusion process we derived earlier, as well as the posterior and likelihood obtained when modeling the reverse diffusion process:

The $L_\text{VLB} = \mathbb{E}_{\mathbf{x}_{0:T}\sim q(\mathbf{x}_{0:T})}\Big[\log \frac{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})}{p_\theta(\mathbf{x}_{0:T})} \Big]$ we obtained just now does not use these existing formulas, so it cannot be calculated directly. Moreover, we do not want the neural network to fit the result in one step. Instead, we hope it can simulate the step-by-step denoising process. (Of course, later there emerged Consistency Models ⁶ that can achieve this in one step, which will be discussed later.)

Therefore, the purpose of expanding VLB is to transform the training objective from a log-likelihood function that is difficult to optimize directly into a set of computable KL divergence terms and reconstruction terms, thereby guiding the neural network to learn how to gradually recover the original data from noise.

Let’s write out the expansion process of VLB more completely:

For the sake of brevity, the expectation symbols are actually omitted here. Moreover, the range of the second term in the original formula is $2 \leq t \leq T$, not $1 \leq t \leq T-1$. Therefore, we write it more strictly as follows:

The following will explain how these three terms participate in the training of the neural network:

⭐ $L_T = \mathbb{E}_{\mathbf{x}_{0:T}\sim q(\mathbf{x}_{0:T})} D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T))$

As mentioned in the original blog post, this is just a constant with a backpropagation gradient of 0, so it can be ignored and does not participate in optimization.

⭐ $L_t = \mathbb{E}_{\mathbf{x}_{0:T}\sim q(\mathbf{x}_{0:T})} D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))$

This term is the most important and will be elaborated on later. It is temporarily skipped here.

⭐ $L_0 = - \mathbb{E}_{\mathbf{x}_{0:T}\sim q(\mathbf{x}_{0:T})} \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)$

For the convenience of calculation, this term is approximated as the denoising matching term at $t = 1$:

Parameterization of $L_t$ for Training Loss

The derivation still needs to use the closed-form expression of the forward diffusion process we derived earlier, as well as the posterior and likelihood obtained when modeling the reverse diffusion process:

In addition, we need to know the formula for calculating the KL divergence between two multivariate Gaussian distributions (read this post ⁷ to learn how to derive it). Given two Gaussian distributions, $\mathcal{N}(\boldsymbol{\mu_q},\,\Sigma_q)$ and $\mathcal{N}(\boldsymbol{\mu_p},\,\Sigma_p)$, with a data dimension of $k$, their KL divergence is:

Furthermore, we need to know an important assumption. The original DDPM paper ⁸ assumes that $\Sigma_\theta$ is a constant hyperparameter. If $\Sigma_\theta$ needs to be learned, it will lead to:

1. Gradient updates for the variance are required during training, which may cause divergence.
2. For each time step $t$, $\Sigma_\theta(\mathbf{x}_t, t)$ is high-dimensional (e.g., the image pixel dimension), resulting in a large training volume.
3. Experience shows that if only the mean is predicted, the model can already learn the reverse process well, and the quality of the generated samples is also high.

Based on these information, let’s write out the complete derivation of $L_t$:

Herein:

• $t, \mathbf{x}_0, \boldsymbol{\epsilon}_t$ are all obtained through Monte Carlo sampling, so they are inputs to the model.
• $\alpha_t, \bar{\alpha}_t$ both depend on the Noise variance schedule parameter $\beta_t$, which is a hyperparameter and a scheduling sequence, a value worth designing.
• $\boldsymbol{\Sigma}_{\theta,t}$ is the covariance of the reverse diffusion distribution at time step $t$, which can be designed as a constant or learnable parameters.
• $\boldsymbol{\epsilon}_\theta$ is a term that the neural network must learn, which needs no further explanation.

Next, the design strategies related to these three parameters will be explained one by one:

1. The $L_t^\text{simple}$ strategy where $\beta_t$ and $\boldsymbol{\Sigma}_{\theta,t}$ are simple constants.
2. The strategy where $\beta_t$ is a non-trivial schedule.
3. The strategy where $\boldsymbol{\Sigma}_{\theta,t}$ is a learnable parameter.

This mainly explains two things:

1. During training, Monte Carlo sampling is performed on any real image sample $\mathbf{x}_0$, any diffusion step $t$, and any noise $\boldsymbol{\epsilon}_t$.
2. The weight coefficient containing $\boldsymbol{\Sigma}_{\theta,t}$ is ignored during training because it is set to a constant in the original paper ⁸.

PS. This simplified formula also provides us with another perspective to observe $L_t$, i.e., it can also be regarded as an MSE loss instead of a KL divergence loss.

Training and Sampling in the DDPM Algorithm

• Algorithm 1: Training is to teach the model to “guess the noise in the image at a certain moment”. Thanks to the various closed-form formulas derived earlier, we do not need to perform step-by-step calculations.
• Algorithm 2: Sampling is to start from random noise and gradually denoise to generate new images. Therefore, it has to go through the entire diffusion process, so the generation speed is slow.

PS. In the field of image generation, sampling refers to using the trained model for inference.

Let’s recall the Stochastic Gradient Langevin Dynamics sampling formula mentioned earlier:

Here, $p(\cdot)$ is used to measure the authenticity of the generated samples. As stated in the original text, $p(\cdot)$ is defined as

Substituting this in, we obtain a more complete sampling formula:

Due to the existence of this closed-form formula, it means that we can continuously iterate this formula to obtain a sample with higher authenticity. So now the key question is to see if each term in it can be obtained.

In this formula, $\mathbf{x}_{t-1}$ is our input value; $\frac{\delta}{2}$ is a constant coefficient; $\sqrt{\delta} \boldsymbol{\epsilon}_t$ is randomly sampled; only the middle $\nabla_\mathbf{x} \log q(\mathbf{x}_{t-1})$ is the key term that the neural network needs to fit.

Before understanding how to model the middle term, let’s first derive the derivative of the logarithm of the Gaussian density function (this is a separate derivation and can be skipped first):

So, the conclusion is:

With the above formula, let’s derive our current ground truth:

There are two variables here, the time step $t$ and the real sample $\mathbf{x}_0$. So we have a supervision signal composed of the expectations of many ground truths:

Based on this, we can parameterize the neural network as:

• During training, we sample many sets of $(t, \mathbf{x}_0)$ to make the neural network fit the ground truth.

• During inference, we substitute the existing $\nabla_{\mathbf{x}_t} \log q(\mathbf{x}_t) \approx \mathbf{s}_\theta(\mathbf{x}_t, t) = - \frac{\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)}{\sqrt{1 - \bar{\alpha}_t}}$ into the Stochastic Gradient Langevin Dynamics sampling formula to get:

  $$ \mathbf{x}_t = \mathbf{x}_{t-1} - \frac{\delta \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)}{2\sqrt{1 - \bar{\alpha}_t}} + \sqrt{\delta} \boldsymbol{\epsilon}_t ,\quad\text{where } \boldsymbol{\epsilon}_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) $$

  At this point, we can iteratively sample to generate new images without relying on the real sample $\mathbf{x}_0$.

Parameterization of $\beta_t$

From trivial to non-trivial $\beta_t$ scheduling

The forward variances are set to be a sequence of linearly increasing constants in Ho et al. (2020), from $\beta_1=10^{-4}$ to $\beta_T=0.02$. They are relatively small compared to the normalized image pixel values between $[-1, 1]$. Diffusion models in their experiments showed high-quality samples but still could not achieve competitive model log-likelihood as other generative models.

Nichol & Dhariwal (2021) proposed several improvement techniques to help diffusion models to obtain lower NLL. One of the improvements is to use a cosine-based variance schedule. The choice of the scheduling function can be arbitrary, as long as it provides a near-linear drop in the middle of the training process and subtle changes around $t=0$ and $t=T$.

$$ \beta_t = \text{clip}(1-\frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}, 0.999) \quad\bar{\alpha}_t = \frac{f(t)}{f(0)}\quad\text{where }f(t)=\cos\Big(\frac{t/T+s}{1+s}\cdot\frac{\pi}{2}\Big)^2 $$

where the small offset $s$ is to prevent $\beta_t$ from being too small when close to $t=0$.

Let’s write out the scheduling methods from the original DDPM paper ⁸ and the scheduling formulas from Improved DDPM ⁹ in full:

The linear variance schedule in DDPM ⁸ is:

where:

• $\beta_{\text{min}}$ and $\beta_{\text{max}}$ are the preset minimum and maximum noise values (e.g., 0.0001 and 0.02).
• $T$ is the total number of diffusion steps (e.g., 1000).

The cosine-based variance schedule in Improved DDPM ⁹ is:

where:

• This schedule first defines $\bar{\alpha}_t$, then $\alpha_t$ and $\beta_t$.
• The small offset $s$ is used to prevent $\beta_t$ from becoming too small when close to $t = 0$. Improved DDPM ⁹ argues that having tiny amounts of noise at the beginning of the process makes it difficult for the network to predict accurately enough.
• The clip function ensures that $\beta_t$ does not exceed 0.999, avoiding numerical instability. Improved DDPM ⁹ suggests that this can prevent singularities at the end of the diffusion process near $t = T$.

Click the bottom-right corner of the following figure to enter an interactive interface and intuitively experience the two schedulers.

Parameterization of reverse process variance $\boldsymbol{\Sigma}_\theta$

Recall the posterior and likelihood in the reverse diffusion process we discussed earlier:

where $\tilde{\boldsymbol{\mu}}_t (\mathbf{x}_t, \mathbf{x}_0) = \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_t \Big)$ and $\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t$.

To make the training more convergent, in the original DDPM ⁸, the authors modeled the variance part as a constant, $\boldsymbol{\Sigma}_\theta \triangleq \tilde{\beta}_t \text{ or } \beta_t$, based on the formal consistency of the two equations. They only treated the noise $\boldsymbol{\epsilon}_t$ in the closed form of $\tilde{\boldsymbol{\mu}}(\mathbf{x}_t, \mathbf{x}_0)$ as a learnable parameter.

Why can it also be modeled as the noise schedule $\beta_t$ in the forward diffusion? Nichol from OpenAI provided an explanation in Improved DDPM ⁹. They are actually very close, especially in the later stages of the diffusion process.

They also found that improving this part to make it a learnable interpolation between $\beta_t$ and $\tilde{\beta}_t$ can lead to better results.

where:

• $\mathbf{I}$ (identity matrix), added here to ensure that the final $\boldsymbol{\Sigma}_\theta$ is a diagonal covariance matrix.
• $\beta_t$ (scalar) is the forward noise schedule.
• $\tilde{\beta}_t$ (scalar) is the posterior variance.
• $\mathbf{v}$ (vector) is the mixing coefficient output by the model, which is what the model needs to learn.
• $\boldsymbol{\Sigma}_\theta$ (vector) is the predicted variance for each dimension (diagonal covariance).

The overall meaning of this passage is relatively straightforward. The main goal is to combine the original $L_\text{simple}$ loss with the loss containing the learnable parameter $\boldsymbol{\Sigma}_\theta$ for joint optimization. However, two concepts mentioned here can be elaborated: “noisy gradient” and “time-averaged smoothed version of $L_\text{VLB}$ with importance sampling”.

Noisy gradient is a concept proposed in OpenAI’s paper An Empirical Model of Large-Batch Training ¹⁰.

In Stochastic Gradient Descent (SGD), instead of computing gradients using the entire dataset, we use a mini-batch. This introduces noise because gradients from different batches can vary significantly. The Gradient Noise Scale can measure this gradient volatility.

Its core idea is that if gradients vary greatly across different batches (high noise), we need a larger batch size to obtain more stable updates. Under simplified assumptions (e.g., the Hessian is a multiple of the identity matrix), the Gradient Noise Scale can be expressed as:

where:

• $\text{tr}(\Sigma)$: The trace of the gradient covariance matrix, which is the sum of the variances of all parameter gradients, representing the “volatility” of the gradients.
• $\|G\|_2^2$: The squared norm of the gradient (global gradient norm), which is the sum of the squares of all parameter gradients, representing the “average intensity” of the gradients.

$B_{\text{simple}}$ represents: the ratio of the gradient noise intensity to its average intensity

• If $B_{\text{simple}}$ is large, it indicates strong gradient noise, and a larger batch size is recommended.
• If it is small, it means the gradients are stable, and a smaller batch size can be used to speed up training.

In this passage, the “time-averaged smoothed version of $L_\text{VLB}$ with importance sampling” specifically refers to the following formula, a new loss function design proposed in Improved DDPM ⁹:

Let’s recall the original DDPM loss function:

We can see that the time step $t$ follows a uniform distribution, so its probability density function is $p'(t)=\frac{1}{T}$.

However, Improved DDPM ⁹ points out that the first few time steps contribute most of the loss values, so uniform sampling is inefficient.

Therefore, we can use another distribution to optimize sampling. The authors proposed a new distribution $p_t \propto \sqrt{\mathbb{E}[L_t^2]}$ where $\sum p_t = 1$. This distribution is proportional to the loss values, meaning we expect regions with high loss values (regions with small time steps) to be sampled more frequently. Using the importance sampling formula to adjust the expectation:

Conditioned Generation

Classifier Guided Diffusion

The overall idea is to combine a pre-trained diffusion model and a pre-trained image classifier model for conditional generation. By leveraging the Langevin dynamics technique, if we have a classifier, we can obtain the gradient information of class $l$. With this gradient information, we can perform sampling via Langevin dynamics and iteratively generate images that match class $l$:

The derivation ideas in the blog post are as follows:

• First, we convert the conditional probability into the gradients of the diffusion part and the classifier part.
• Second, we transform the theoretical formula into a form containing learnable parameters.
• Third, we perform a maybe unnecessary simplification.

Classifier-Free Guidance

The advantage of classifier-guided diffusion is that it can directly utilize two pre-trained models without additional operations. Another advantage is that we can control the strength of the “condition” through the weight $w$ mentioned above.

For classifier-free guidance, the simplest approach is to directly incorporate the condition information into the training of the diffusion model. However, this approach loses the feature of controlling the condition strength.

Therefore, another approach, as shown in the original blog post, is to train both conditional and unconditional cases on the same network architecture, only distinguishing them by the condition input. In this way, we can control the condition strength through “subtraction”.

Let’s write out the sampling formula more comprehensively:

FID and IS are important evaluation metrics for generative models.

FID (Fréchet Inception Distance) ¹¹ measures the distribution difference between generated images and real images in the feature space. It uses the Inception network to extract image features and then calculates the Fréchet distance between two high-dimensional Gaussian distributions.

where:

• $\mu_r, \Sigma_r$: Mean and covariance of real images.
• $\mu_g, \Sigma_g$: Mean and covariance of generated images.

A lower FID indicates that the generated images are closer to the real images.

IS (Inception Score) ¹² measures the “sharpness” and “diversity” of generated images. It uses the Inception network to predict the image category distribution and then calculates the KL divergence of the predicted distribution.

where:

• $p(y|x)$: Predicted distribution of the Inception network for generated images.
• $p(y)$: Average predicted distribution of all generated images.

A higher IS indicates that the images are sharp (low entropy of the predicted distribution) and diverse (high entropy of the average distribution).

GLIDE is a guided diffusion model. It tried two guiding strategies for image generation:

1. CLIP guidance: Leverages the image-text matching ability of the CLIP model to guide the image generation process.
2. Classifier-free guidance: Does not rely on an external classifier. Instead, it trains a single model to learn both conditional and unconditional image generation, achieving guidance in this way.

GLIDE’s experiments showed that classifier-free guidance outperforms CLIP guidance, as follows:

• Classifier-free guidance: The model learns to generate images independently without relying on external judgment.
• CLIP guidance: The model relies on CLIP scores but may “cheat” to deceive the CLIP model.
• GLIDE prefers the former because it is more natural and robust.

Speed up Diffusion Models

Fewer Sampling Steps & Distillation

Here is the standard DDPM sampling formula:

The naive strided sampling modifies the sampling formula as follows:

Of course, when taking multiple steps at once, the true mean and variance should be derived from the multi-step Gaussian combination formula. Therefore, this strided sampling strategy is only a rough acceleration method.

Naive Strided Sampling lacks mathematical rigor, which DDIM ¹³ addresses. Meanwhile, DDIM transforms the original random sampling process into a combination of random and deterministic sampling, making it more flexible.

Moreover, the main difference between DDIM and DDPM lies in how they model the likelihood to predict the posterior.

where:

• $\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t$
• $\sigma_t^2 = \eta \cdot \tilde{\beta}_t$. When $\eta = 0$, the variance is 0, making the sampling process deterministic (DDIM); when $\eta = 1$, the sampling process is random (DDPM).

Note: DDPM and DDIM are from different papers, so there may be math symbol conflicts. Here, we follow the original papers and do not unify them.

The derivation of the formula starts from the closed-form formula of the forward diffusion process. We attempt to separate a $\sigma_t^2$ term to the variance position using the following technique:

Progressive Distillation

We can provide a more detailed explanation of the Progressive Distillation algorithm on the right.

First, there are two nested loops: The outer loop controls the training of $K$ different student diffusion models, each capable of generating images with fewer sampling steps.

The inner loop trains the current student model:

1. $\text{Cat}[1,2,...,N]$ refers to the categorical distribution. We uniformly sample an integer $t$ from 1 to $N$, representing the current time step.
2. $t' = t - \frac{0.5}{N}, \quad t'' = t - \frac{1}{N}$ means splitting the time step $t$ into two smaller steps, $t'$ and $t''$.
3. $z_{t'} = \alpha_{t'} \tilde{x}_\eta(z_t) + \frac{\sigma_{t'}}{\sigma_t} \Big( z_t - \alpha_t \tilde{x}_\eta(z_t) \Big)$ is one of the update formulas for teacher DDIM sampling, indicating that $z_{t'}$ is obtained by taking one sampling step from $z_t$. For the specific derivation, please refer to the original paper ¹⁴.

The goal of the Consistency Model (CM) is to learn a direct mapping $f$ that can map any point $\mathbf{x}_t$ with a noise level $t > 0$ back to the “source” of the same generation trajectory (more precisely, a sample $\mathbf{x}_\epsilon$ at a small time point $\epsilon$ very close to 0).

Why map to $\mathbf{x}_\epsilon$ instead of $\mathbf{x}_0$ exactly?

• $t = 0$ is often numerically unstable (singular or has a poor condition number). Choosing a small $\epsilon > 0$ makes the training more stable and easier.
• $\mathbf{x}_\epsilon$ is already extremely close to $\mathbf{x}_0$. If necessary, we can add one or two refinement steps from $\epsilon$ to 0.

Why is it an identity mapping when $t = \epsilon$?

The paper parameterizes $f_\theta$ as a residual form with skip connections:

and specifically designs

Thus,

which is an identity mapping.

Latent Variable Space

This formula describes a key module for the cross-attention mechanism in the Latent Diffusion Model (LDM). It combines the attention mechanism in Transformer with the conditional control mechanism in diffusion models, serving as the core component for LDM to achieve text-to-image generation or other conditional generation tasks.

Scaled dot-product attention:

• $\mathbf{Q}$: Query, representing the information to be processed currently.
• $\mathbf{K}$: Key, representing the information that can be attended to.
• $\mathbf{V}$: Value, from which the actual content is extracted when $\mathbf{Q}$ and $\mathbf{K}$ match.
• The denominator $\sqrt{d}$ is a scaling factor to prevent the inner product from being too large, which could cause the softmax gradient to vanish.
• Softmax normalizes along the Key dimension to obtain attention weights.
• The final output is a weighted sum, indicating “which parts of the value are more important to extract based on the query”.

Sources of Query, Key, and Value in LDM:

• Query comes from the image (latent representation)
• Key and Value come from the text/conditional information

This means that every spatial position in the image “attends” to which text words are most relevant. For example, when generating an image of “a red dog running on the grass”, the “grass” area in the image will pay more attention to the word “grass” in the text.

This mechanism allows the model to precisely align semantic conditions with the spatial structure of the image.

Scale up Generation Resolution and Quality

This section discusses how to improve the image generation quality and resolution of diffusion models to a higher level through a series of techniques.

It mainly covers the Noise Conditioning Augmentation technique, the unCLIP model architecture, the Imagen model architecture, and some other improvements.

Here, we only pay attention to the Noise Conditioning Augmentation and the unCLIP model architecture.

When using a pipeline of multiple diffusion models, the input of each stage is the output of the previous stage, which can easily lead to error accumulation and damage to the final image.

Therefore, they introduced Noise Conditioning Augmentation to enable the model to generate clear images even under “slightly blurred or noisy” conditions.

Two types of noise:

• Low-resolution stage: Add Gaussian noise.
• High-resolution stage: Add Gaussian blur.
• These noises are only added during training and not during inference.

Two training methods:

1. Truncated Conditioning Augmentation: Terminate the diffusion process early in the low-resolution stage (e.g., stop at step $t$).
2. Non-Truncated Conditioning Augmentation: Complete the low-resolution diffusion process, then artificially add noise, and use the result as the input for the high-resolution model.

unCLIP

The two-stage diffusion model unCLIP (Ramesh et al. 2022) heavily utilizes the CLIP text encoder to produce text-guided images at high quality. Given a pretrained CLIP model $\mathbf{c}$ and paired training data for the diffusion model, $(\mathbf{x}, y)$, where $x$ is an image and $y$ is the corresponding caption, we can compute the CLIP text and image embedding, $\mathbf{c}^t(y)$ and $\mathbf{c}^i(\mathbf{x})$, respectively. The unCLIP learns two models in parallel:

• A prior model $P(\mathbf{c}^i \vert y)$: outputs CLIP image embedding $\mathbf{c}^i$ given the text $y$.
• A decoder $P(\mathbf{x} \vert \mathbf{c}^i, [y])$: generates the image $\mathbf{x}$ given CLIP image embedding $\mathbf{c}^i$ and optionally the original text $y$.

These two models enable conditional generation, because

$$ \underbrace{P(\mathbf{x} \vert y) = P(\mathbf{x}, \mathbf{c}^i \vert y)}_{\mathbf{c}^i\text{ is deterministic given }\mathbf{x}} = P(\mathbf{x} \vert \mathbf{c}^i, y)P(\mathbf{c}^i \vert y) $$

unCLIP follows a two-stage image generation process:

1. Given a text $y$, a CLIP model is first used to generate a text embedding $\mathbf{c}^t(y)$. Using CLIP latent space enables zero-shot image manipulation via text.
2. A diffusion or autoregressive prior $P(\mathbf{c}^i \vert y)$ processes this CLIP text embedding to construct an image prior and then a diffusion decoder $P(\mathbf{x} \vert \mathbf{c}^i, [y])$ generates an image, conditioned on the prior. This decoder can also generate image variations conditioned on an image input, preserving its style and semantics.

unCLIP is a two-stage text-to-image generation model that corelly leverages the text and image embeddings of CLIP:

1. Prior model: Input text → Output CLIP image embedding.
2. Decoder model: Input image embedding (and optional text) → Output image.

This design allows:

• Text-to-image generation.
• Generating variations of a given image while preserving its style and semantics.

Similar to unCLIP, the Imagen model is also a two-stage text-to-image generation model. The difference lies in the choice of the text encoder. unCLIP uses CLIP, while Imagen uses the large language model T5-XXL. Imagen also makes some optimizations to the U-Net architecture.

Model Architecture

Lilian Weng’s blog has already introduced these three models. Below are simplified PyTorch implementation examples. Some knowledge of CNN networks is required here. Refer to Wang’s CNN Explainer article ¹⁵ to grasp the basics.

U-Net mainly demonstrates the core ideas of downsampling, upsampling, and skip connections. It uses a simple double_conv block containing two convolutional layers, ReLU, and batch normalization.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

import torch
import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module):
    """(conv => BN => ReLU) * 2"""
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

class SimpleUnet(nn.Module):
    def __init__(self, n_channels, n_classes):
        super(SimpleUnet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes

        # Downsampling path
        self.inc = DoubleConv(n_channels, 64)
        self.down1 = nn.MaxPool2d(2)
        self.down2 = DoubleConv(64, 128)
        self.down3 = nn.MaxPool2d(2)
        self.down4 = DoubleConv(128, 256)
        self.down5 = nn.MaxPool2d(2)

        # Upsampling path
        self.up1 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.up2 = DoubleConv(256, 128)  # 128 from upconv + 128 from skip connection
        self.up3 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.up4 = DoubleConv(128, 64)   # 64 from upconv + 64 from skip connection
        self.outc = nn.Conv2d(64, n_classes, kernel_size=1)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down2(self.down1(x1))
        x3 = self.down4(self.down3(x2))

        x = self.up1(x3)
        x = torch.cat([x, x2], dim=1)
        x = self.up2(x)

        x = self.up3(x)
        x = torch.cat([x, x1], dim=1)
        x = self.up4(x)
        logits = self.outc(x)
        return logits

ControlNet is essentially a fine-tuning method for U-Net. Its core idea is to add a trainable copy to the frozen U-Net backbone and connect them through zero convolutions.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

import torch
import torch.nn as nn
import torch.nn.functional as F

class ZeroConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, *args, **kwargs):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, *args, **kwargs)
        self.conv.weight.data.fill_(0.)
        self.conv.bias.data.fill_(0.)

    def forward(self, x):
        return self.conv(x)

class ControlNetBlock(nn.Module):
    """A simplified example of a ControlNet block with a frozen backbone."""
    def __init__(self, in_channels, out_channels):
        super().__init__()
        # The frozen U-Net backbone block
        self.backbone_frozen = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)

        # A trainable copy of the backbone block
        self.backbone_trainable = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)

        # Initialize trainable part with backbone's weights
        self.backbone_trainable.load_state_dict(self.backbone_frozen.state_dict())

        # Freeze the original backbone
        for param in self.backbone_frozen.parameters():
            param.requires_grad = False

        # Zero convolutions
        self.zero_conv_in = ZeroConv2d(in_channels, in_channels, kernel_size=1)
        self.zero_conv_out = ZeroConv2d(out_channels, out_channels, kernel_size=1)

    def forward(self, x, conditioning_vector):
        # Apply the frozen backbone to the input
        output_frozen = self.backbone_frozen(x)

        # Apply the trainable copy to the sum of input and conditioning vector
        # The conditioning vector is first processed by a zero convolution
        conditioning_processed = self.zero_conv_in(conditioning_vector)
        trainable_input = x + conditioning_processed
        output_trainable = self.backbone_trainable(trainable_input)

        # Add the output of the trainable part (with zero convolution) to the frozen output
        final_output = output_frozen + self.zero_conv_out(output_trainable)
        return final_output

The core idea of DiT is to use a Transformer to process the latent representations of diffusion models. It “slices” the latent image representations into sequential patches. Here, we demonstrate the use of Adaptive Layer Normalization (adaLN) to inject time step and class information.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

import torch
import torch.nn as nn
import torch.nn.functional as F

class AdaLN(nn.Module):
    """
    Adaptive Layer Normalization (adaLN) for DiT.
    Generates scale and shift parameters from a conditioning vector.
    """
    def __init__(self, n_embd):
        super().__init__()
        self.n_embd = n_embd
        # We need a linear layer to predict gamma and beta
        self.linear = nn.Linear(n_embd, 2 * n_embd)
        self.norm = nn.LayerNorm(n_embd)

    def forward(self, x, c):
        # c is the conditioning vector (sum of time and class embeddings)
        gamma, beta = self.linear(c).chunk(2, dim=-1)
        x = self.norm(x)
        # Apply scaling and shifting
        return x * (1 + gamma.unsqueeze(1)) + beta.unsqueeze(1)

class DiTBlock(nn.Module):
    """A simplified DiT block, combining Attention and MLP with adaLN."""
    def __init__(self, n_embd, n_heads):
        super().__init__()
        self.adaLN1 = AdaLN(n_embd)
        self.attention = nn.MultiheadAttention(n_embd, n_heads, batch_first=True)
        self.adaLN2 = AdaLN(n_embd)
        self.mlp = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd)
        )

    def forward(self, x, c):
        # x is the sequence of patches, c is the conditioning vector
        x = x + self.attention(self.adaLN1(x, c), self.adaLN1(x, c), self.adaLN1(x, c))[0]
        x = x + self.mlp(self.adaLN2(x, c))
        return x

class SimpleDiT(nn.Module):
    """Simplified DiT model for latent diffusion."""
    def __init__(self, img_size, patch_size, n_embd, n_heads, n_layers, num_classes=1000):
        super().__init__()
        self.patch_size = patch_size
        num_patches = (img_size // patch_size) ** 2

        # Patch embedding and positional embedding
        self.patch_embed = nn.Linear(3 * patch_size**2, n_embd)
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, n_embd))

        # Time and class embeddings
        self.time_embed = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd)
        )
        self.class_embed = nn.Embedding(num_classes, n_embd)

        # Transformer blocks
        self.blocks = nn.ModuleList([DiTBlock(n_embd, n_heads) for _ in range(n_layers)])

        # Final output layer
        self.final_layer = nn.Linear(n_embd, 3 * patch_size**2)

    def forward(self, x, t, c):
        # 1. Patchify and embed
        x = x.unfold(2, self.patch_size, self.patch_size).unfold(3, self.patch_size, self.patch_size)
        x = x.contiguous().view(x.size(0), -1, 3 * self.patch_size**2)
        x = self.patch_embed(x) + self.pos_embed

        # 2. Time and class conditioning
        t_emb = self.time_embed(t)
        c_emb = self.class_embed(c)
        cond = t_emb + c_emb

        # 3. Pass through DiT blocks
        for block in self.blocks:
            x = block(x, cond)

        # 4. Final layer to predict noise
        output = self.final_layer(x)
        return output

Quick Summary

In the design of generative models, “tractability” and “flexibility” are often conflicting goals:

• Tractability: This refers to models with clear mathematical structures that facilitate inference and optimization. For example, Gaussian distributions, Laplace distributions, etc., allow for direct likelihood calculation, sampling, and training.
• Flexibility: This refers to models that can fit complex, high-dimensional, and non-linear data structures, such as images, audio, and text. However, these models are often difficult to train, sample from, or evaluate.

Traditional models include:

• VAE (Variational Autoencoder):
  • Principal: The input image is compressed into latent space variables by an encoder and then reconstructed by a decoder.
  • Pro/Con: Tractable but sacrifices generation quality.
• GAN (Generative Adversarial Network):
  • Principal: Two networks are trained, one for generating images and the other for distinguishing real from fake images.
  • Pro/Con: Flexible but suffers from unstable training and lacks an explicit likelihood.
• Flow-based models:
  • Principal: Images are generated through flow matching, which combines function-level processes.
  • Pro/Con: Reversible but have limited architectures.

The uniqueness of diffusion models lies in the fact that they maintain the tractability of high-dimensional Gaussian modeling in theory by gradually adding noise (forward process) and learning to denoise (reverse process). Meanwhile, they achieve flexibility by using deep neural networks to learn complex data structures.

Moreover, nowadays, diffusion models have become the de facto standard for generative models, especially in the field of image generation. They excel in generation quality, sample diversity, and training stability, and can be easily scaled to larger and more complex tasks.

Appendix

Relevant Deep Learning Concepts

From AE to VAE and Then to VQ-VAE

• AE (Autoencoder): A neural network with an encoder and a decoder. Its goal is to learn a compressed representation (latent variables) of the input data, mainly used for dimensionality reduction and feature extraction.
• VAE (Variational Autoencoder): Based on the AE, it introduces probabilistic modeling and variational inference, enabling the model to have generative capabilities.
• VQ-VAE (Vector Quantized Variational Autoencoder): It further introduces a discrete latent space, replacing the continuous latent distribution with vector quantization. It is more suitable for modeling discrete structures (e.g., symbolic features in speech and images) and provides discrete token representations for subsequent generative models (e.g., Transformer decoders).

Reparameterization Trick

Definition: Reparameterization is a method that decouples random variables from non-differentiable sampling operations, allowing the sampling operations to participate in gradient descent optimization.

Principle: It does not eliminate random sampling but minimizes the impact of random sampling on gradient propagation.

Example: If you have a random variable $z \sim \mathcal{N}(\mu, \sigma^2)$, directly sampling from this distribution prevents gradients from propagating through $\mu$ and $\sigma$. You can convert the random sampling of $z$ to the random sampling of $\epsilon$ using the following formula:

PyTorch Code Example:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

import torch

# Assume the mean and standard deviation output by the encoder
mu = torch.tensor([0.0, 1.0], requires_grad=True)
log_sigma = torch.tensor([0.0, 0.0], requires_grad=True)  # Usually output log(sigma) to avoid negative values
sigma = torch.exp(log_sigma)

# Reparameterization sampling
epsilon = torch.randn_like(mu)  # Sample from the standard normal distribution
z = mu + sigma * epsilon  # z is differentiable

# Assume a simple loss function
loss = (z**2).sum()
loss.backward()

print("grad mu:", mu.grad)
print("grad log_sigma:", log_sigma.grad)

Importance Sampling Trick

Suppose you want to compute the expectation of a function $f(x)$ with respect to a probability distribution $p(x)$:

However, it might be difficult to sample directly from $p(x)$, or most samples of $f(x)$ under $p(x)$ contribute little, with only a few regions making significant contributions (e.g., tail events). In such cases, direct Monte Carlo estimation can be inefficient.

If we can sample from another distribution $q(x)$ that is easier to sample from and then correct the bias through weighting, the estimation becomes more convenient. This is known as importance sampling.

Importance sampling rewrites the original expectation as:

Where:

• $q(x)$ is the proposal distribution, and $q(x) > 0$ when $p(x) > 0$ (i.e., the support of $q$ contains the support of $p$).
• $w(x_i) = \frac{p(x_i)}{q(x_i)}$ is called the importance weight.

At this point, we can sample $x_i \sim q(x)$ from $q(x)$ and then estimate:

Let’s take an example. Assume:

• $p(x) = \mathcal{N}(0, 1)$
• $f(x) = x^2 \cdot \mathbb{1}[x > 2]$
• If we sample directly from $p$, most samples satisfy $x < 2$, contributing zero to the result, which is inefficient.
• If we switch to sampling from $q(x) = \mathcal{N}(2.5, 1)$, we are more likely to sample points in the region $x > 2$.

PyTorch Code Example:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

import numpy as np

N = 10000
# Sample from q(x) ~ N(2.5, 1)
x = np.random.normal(2.5, 1, N)

# Compute unnormalized densities (ignoring constants)
log_p = -0.5 * x**2           # log N(0,1)
log_q = -0.5 * (x - 2.5)**2   # log N(2.5,1)

# Importance weights (unnormalized)
w = np.exp(log_p - log_q)

# Function values
f = x**2 * (x > 2)

# Normalized importance sampling estimate
mu_hat = np.sum(f * w) / np.sum(w)
print("Importance Sampling estimate:", mu_hat)

Probability Theory

Gaussian Distribution

Univariate Gaussian Distribution $x \sim \mathcal{N}(\mu, \sigma^2)$ has the following probability density function (PDF):

• Mean $\mu$: The central location of the distribution.
• Variance $\sigma^2$ (Standard Deviation $\sigma$): The width (uncertainty) of the distribution.

Multivariate Gaussian Distribution $\mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma})$ has the following probability density function (pdf):

Where:

• $\mathbf{x} \in \mathbb{R}^d$ is a $d$-dimensional vector.
• Mean vector $\boldsymbol{\mu} \in \mathbb{R}^d$: The central location of the distribution.
• Covariance matrix $\boldsymbol{\Sigma} \in \mathbb{R}^{d\times d}$: Describes the variance and correlation between different dimensions.
• $|\boldsymbol{\Sigma}|$ = The determinant of the covariance matrix, representing “volume scaling”.
• $\boldsymbol{\Sigma}^{-1}$ = The inverse of the covariance matrix, defining the “ellipsoidal” isodensity curves.

Decomposing $\boldsymbol{\Sigma}$ can reveal the geometric properties of the distribution:

• Diagonal elements: The variance of each dimension (the numerical value equals the “width” along that axis).
• Off-diagonal elements: The correlation between different dimensions, determining whether the distribution is a “rotated ellipse/ellipsoid”.

Examples:

• If $\boldsymbol{\Sigma} = \sigma^2 I$, it is an isotropic “circular/spherical” distribution.
• If $\boldsymbol{\Sigma}$ is not a diagonal matrix, there is correlation, and the isodensity lines are “tilted ellipses/ellipsoids”.

The following two illustrations more intuitively show the effects of $\mu$ and $\Sigma$ on the shape of the PDF ¹⁶:

Joint, Marginal, and Conditional Distributions

• Joint distribution $P(A, B)$: A panoramic map (containing the probabilities of all combinations).
• Marginal distribution $P(B)$: The projection of the panoramic map onto a single axis.
• Conditional distribution $P(A \vert B)$: Cutting a line from the panoramic map (given the value of another variable) and observing the probability distribution along this line. The formula is: $P(A \vert B) = \frac{P(A, B)}{P(B)}$.

Prior, Likelihood, and Posterior

The relationship among them is:

Bayes’ Rule

Bayes’ Rule is a core principle in probability theory, used to update the probability of an event given certain conditions. Its basic form is:

Where:

• $P(A)$: The prior probability of event A (the belief in A before observing B).
• $P(B \vert A)$: The likelihood of observing B given that A has occurred.
• $P(B)$: The marginal probability of event B (the probability of B occurring under all possible scenarios).
• $P(A \vert B)$: The posterior probability of event A after observing B (the updated belief).

Law of Total Probability

The Law of Total Probability is used to calculate the probability of an event when the probability of this event is not easy to calculate directly. It can be solved by decomposing it into the sum of conditional probabilities under several mutually exclusive scenarios.

Suppose events $B_1, B_2, \dots, B_n$ form a complete set of events in the sample space (i.e., they are mutually exclusive and their union is the entire space), that is:

• $B_i \cap B_j = \emptyset$ (Mutually exclusive).
• $\bigcup_{i=1}^n B_i = \Omega$ (Exhaustive).

Then for any event $A$, we have:

Marginalization

Marginalization refers to obtaining the probability distribution of another variable from a joint probability distribution by “summing (or integrating) over certain variables”.

For example, given the joint distribution $P(X, Y)$ of two random variables $X$ and $Y$, if we want to obtain the distribution $P(X)$ of $X$, we need to “marginalize” over all possible values of $Y$:

Moreover, we can also view the Law of Total Probability as an application of marginalization:

Expectation

Definition of expectation:

• Discrete random variable: $\mathbb{E}_{x \sim p}[g(x)] = \sum_{x} g(x) \cdot p(x)$
• Continuous random variable: $\mathbb{E}_{x \sim p}[g(x)] = \int_{-\infty}^{\infty} g(x) \cdot p(x) dx$

Three important properties:

1. Linearity property: $\mathbb{E}_{x \sim p, y \sim q}[a \cdot f(x, y) + b \cdot g(x, y)] = a \cdot \mathbb{E}_{x \sim p, y \sim q}[f(x, y)] + b \cdot \mathbb{E}_{x \sim p, y \sim q}[g(x, y)]$
2. Law of Iterated Expectation / Tower Property: $\mathbb{E}_{(x, y) \sim p(x, y)}[f(x, y)] = \mathbb{E}_{y \sim p(y)}\left[ \mathbb{E}_{x \sim p(x|y)}[f(x, y)] \right]$
3. Law of Total Expectation: $\mathbb{E}_{x \sim p(x)}[f(x)] = \mathbb{E}_{y \sim p(y)}\left[ \mathbb{E}_{x \sim p(x|y)}[f(x)] \right]$

Stochastic Process

A Stochastic Process is a family of random variables that evolve over time (or space). A stochastic process can be written as:

• $t$: The index set, which can be discrete (e.g., integer time points $t = 0, 1, 2, \dots$) or continuous (e.g., real time $t \geq 0$).
• $X_t$: A random variable at each time point $t$.
• The entire process is a family of random variables, reflecting the randomness of the system as it evolves with $t$.

Intuitive understanding:

• A random variable is a “random quantity at a certain moment”.
• A stochastic process is a “sequence of random quantities that change over time”.

Markov Property

A stochastic process satisfies the Markov property if:

That is, the future only depends on the present and is independent of the past.

Score Function

In probability theory and statistics, the original definition of the score function is:

• This is the gradient of the log-likelihood function with respect to the parameter θ.
• In classical statistics, it is used for maximum likelihood estimation or Fisher information calculation.
• Intuitively, it tells you “how to adjust the model parameter θ to make the observed data x more likely”.

However, in diffusion models / score-based generative models, the score function is extended to:

• Here, it is the gradient of the log-density of the data itself x.
• It tells you the ascending direction of the data distribution, that is, “where the data is more likely to appear”.

Information Theory

Let’s use a scenario to explain the four most important formulas in information theory: You need to send a series of messages (e.g., letters), and each letter has a different probability of occurrence. To save bandwidth, you need to send them using the shortest binary codes…

Information Content $I(x) = -\log p(x)$

• Intuition: The amount of information contained in an event is related to how surprising it is. The less likely an event is to occur, the more information it carries when it does happen.
  • “The sun rises in the east” (probability ~1): Carries almost no information, with extremely little information content.
  • “It will rain tomorrow” (probability 0.3): Carries some information.
  • “It rained diamonds in Beijing tomorrow” (probability almost 0): If it happens, the information content is explosive.
• Why is it $-\log p(x)$?
  • From the encoding perspective: To optimally encode an event, we should assign short codewords to events with high probabilities and long codewords to events with low probabilities. This minimizes the average code length.
  • Optimal code length: The optimal code length for an event $x$ is $-\log p(x)$. The smaller the probability $p(x)$, the larger the code length $-\log p(x)$, which perfectly corresponds to the “surprise factor”.
  • Why use logarithm: Using logarithms can convert the multiplicative relationship of probabilities (joint probability of independent events) into an additive relationship of code lengths, which is very convenient.
• Conclusion: The information content $I(x)$ is the optimal code length for event $x$.

Shannon Entropy $H(p) = -\sum_i p(x_i) \log p(x_i)$

• Intuition: For a probability distribution $p$, when we encode the events in it, what is the shortest average code length required per event? This “average minimum cost” represents the uncertainty of the distribution.
  • The more uniform a distribution is (e.g., a fair die), the harder it is to predict the next outcome, the higher its uncertainty, and the longer the average code length.
  • The more concentrated a distribution is (e.g., a biased die), the easier it is to predict the next outcome, the lower its uncertainty, and the shorter the average code length.
• Why is it $-\sum_i p(x_i) \log p(x_i)$?
  • This is the expected value of the information content: $\mathbb{E}_{x \sim p}[I(x)] = \mathbb{E}_{x \sim p}[-\log p(x)]$.
  • Each event $x_i$ has its own optimal code length $-\log p(x_i)$. We use the probability $p(x_i)$ of this event as the weight and average all possible code lengths to obtain the average optimal code length for the entire distribution.
• Conclusion: The Shannon entropy $H(p)$ is the average code length required when using the optimal encoding scheme of the distribution $p$ itself. It measures the inherent uncertainty of the distribution $p$.

Cross-Entropy $H(p, q) = -\sum_i p(x_i) \log q(x_i)$

• Intuition: Now the situation has changed. The true distribution of the data is $p$, but you mistakenly assume the distribution is $q$ and adopt the optimal encoding scheme designed for $q$ (i.e., assign a codeword of length $-\log q(x_i)$ to event $x_i$). How long, on average, are the codewords needed to encode the real data using this incorrect scheme?
• Why is it $-\sum_i p(x_i) \log q(x_i)$?
  • In the real data, the probability of event $x_i$ occurring is $p(x_i)$.
  • The codeword length you assign to it is $-\log q(x_i)$.
  • Therefore, the overall average code length is $\mathbb{E}_{x \sim p}[-\log q(x)]$.
• Conclusion: The cross-entropy $H(p, q)$ is the average code length required when using the optimal encoding designed for $q$ to represent data from $p$.

Relative Entropy (KL Divergence) $D_{KL}(p\|q) = \sum_i p(x_i) \log \frac{p(x_i)}{q(x_i)}$

• Intuition: Continuing from the concept of cross-entropy. Since encoding with the incorrect scheme $q$ is longer than encoding with the correct scheme $p$, on average, how much extra code length do we waste per sample? This “extra wasted length” represents the difference between the two distributions.
• Why is it $\sum_i p(x_i) \log \frac{p(x_i)}{q(x_i)}$?
  • For event $x_i$, the length using the incorrect scheme is $-\log q(x_i)$, and the length using the correct scheme is $-\log p(x_i)$.
  • The extra wasted length is: $(-\log q(x_i)) - (-\log p(x_i)) = \log \frac{p(x_i)}{q(x_i)}$.
  • We average this “wastage” using the true probability $p(x_i)$ to obtain the average wasted length: $\mathbb{E}_{x \sim p}[\log \frac{p(x)}{q(x)}]$.
• Conclusion: The KL divergence $D_{KL}(p\|q)$ measures the additional average code length incurred by using the approximate distribution $q$ instead of the true distribution $p$. It directly measures the difference between the two distributions.

Relationship $D_{KL}(p\|q) = H(p, q) - H(p) \geq 0$

Let’s conduct a complete derivation:

Citation

@article{shichaosong2025fullyannotatedg,
  title = {Fully Annotated Guide to "What are Diffusion Models?"},
  author = {Shichao Song},
  journal = {The Kiseki Log},
  year = {2025},
  month = {September},
  url = {https://ki-seki.github.io/posts/250902-diffusion-annotated/}
}

References