Diffusion is one of the most used models for image / video generation. This note covers some basics of diffusion models.

Generation as Sampling

Unlike many other models, which generate the outcome in one forward pass, the diffusion models generate a single step towards the outcome, and it takes multiple steps to reach there. Generating a step is also known as sampling a step, because some random variable will affect the step direction.

In the context of image generation, a diffusion model typically starts with some random noise. It is able to denoise the image step by step, and reveal the meaningful image contents in the end. The model is called diffusion, but a better name would be inverse diffusion. Imagine we had this meaningful image in the beginning, and a naughty boy spilled the color palatte onto the canvas. The noise diffused into the canvas, but the model is able to repair it by reverting the process step by step.

Differential Equations

Mathematically, the process can be formulated as a Stochastic Differential Equation (SDE):

A few notes:

• $t \in [0, 1]$ are the normalized timestamps.
• We start at $X_0$ as random noise. We end at $X_1$ as meaningful outcome.
• The way we proceed from $t=0$ to $t=1$ is called simulation. We sample timestamps within the range and simulate $X_t$ throughout.
• $u_t$ is called drift coefficient, as it determines the main direction to proceed. It is also know as a vector field, as each vector in the field determines the direction to proceed in the next timestamp.
• $dW_t$ defines a stochastic process that adds randomness to the system. The process is called “Wiener process” or “Brownian motion”, which is defined as $W_{t+h} = W_t + \sqrt{h} \epsilon_t$, where $\epsilon_t \sim N(0, I)$.
• $\sigma_t$ is called diffusion coeeficient, as it determines the weight of the stochastic process in the system.

When $\sigma_t$ is 0, the system degrades to a Ordinary Differential Equation (ODE):

In both systems, we train a model $u_t^\theta$ to simulate the vector field $u_t$, which guides the process to navigate from random noise to meaningful outcome. When the diffusion term is included, we call it a “diffusion” model. Otherwise, we call it a “flow matching” model.

Conditional Probability Path

The way $X_0$ proceeds to $X_1$ forms a probabiliy path. We can add constraints to the system, so that the solution $u_t^\theta$ is tractable. Just like how Physics defines gravity to drop an apple from the tree to the ground so that we can revert the process to infer how an apple looks like when it was back on the tree, we are free to define our “Physics” in the system that is easy for models to learn.

A commonly used probability path is Gaussian probability path, and it has the following properties:

, where $\alpha_t$ is a sequence that goes from 0 to 1 and $\beta_t$ is a sequence that goes from 1 to 0. These two sequences define the noise schedule of the diffusion process. Model developers may choose one that works the best for their use cases. A linear schedule could be a good starting point:

$P(X_t|X_1)$ also known as a conditional probability path due to the condition on the data sample $X_1$. A quick path to generate $x_t$ given $x_1$ comes as follows:

Conditional Vector Field

A conditional vector field $u_t(x|z)$ is the solution to fulfilling a conditional probability path $p_t(x|z)$. ODE and SDE start to make a slight difference here. For ODE, it can be proved (but not in this article) that the following conditional vector field $u_t(x|z)$ can fulfill such conditional probability path:

For SDE, it has to include an additional score term:

$s_t$ is called a conditional score function according to classical statistics (wiki). For a Gaussian probability path, it has an analytical form:

Marginal Probability Path

However, a conditional probability path is not good enough because we have no such $X_1$ to condition on during inference time. Diffusion models are able to navigate to a meaningful outcome from random noise without seeing it in the first place. Instead, we need to find a marginal probability path to guide any $X_t$ to $X_1$ without seeing $X_1$:

It can be proved (but not in this article) that the following marginal vector field $u_t(x)$ can fulfill such marginal probability path:

The integral is intractible due to the complexity of the data distribution. However, a model $u_t^\theta$ can be trained to approximate this marginal conditional field.

Training Objective

An intuitive loss can be defined as the L2 loss between $u_t^\theta$ and $u_t$.

This loss is intractible, but we can find an upper bound that is tractible:

It can be proved (but not in this article) that

For ODE, we can just plug in $u_t^{ODE}$ for training. This algorithm is also known as “flow matching”.

for z in dataset:
  t = uniform(0, 1)
  xt = cond_prob_path(t, z)
  u_gt = cond_vec_field_ode(t, z, xt)
  u_pd = model(t, xt)
  loss = l2(u_pd, u_gt)
  loss.backward()

For SDE, we may encounter numerical issues because the score term is important but is way smaller than the ODE term.

However, we can rewrite $u_t^{SDE}$ in terms of $s_t$.

Instead of training a model to approximate $u_t$, we train $s_t^\theta$ to approximate $s_t$:

As $s_t(x|z)$ has an analytical form, we can just plug it in for training. This algorithm is also known as “score matching”.

for z in dataset:
  t = uniform(0, 1)
  xt = cond_prob_path(t, z)
  s_gt = cond_score(t, z, xt)
  s_pd = model(t, xt)
  loss = l2(s_pd, s_gt)
  loss.backward()

Simulation

Once the model is trained, we can sample it multiple times to generate meaningful contents. This process is called “simulation” in the context of differential equations. There is a whole literature on different methods, and a commonly used method is called Euler method.

For ODE, the simulation is straightforward:

xt = gauss(0, 1, shape)
for t in arange(0, 1, dt):
  ut = model(t, xt)
  xt += ut * dt

For SDE, we need to plug in the score term and handle the diffusion term properly:

xt = gauss(0, 1, shape)
for t in arange(0, 1, dt):
  alpha_t, beta_t, sigma_t = get_noise_schedule(t)
  ws = get_score_weight(alpha_t, beta_t)
  wx = get_base_weight(alpha_t, beta_t)
  st = model(t, xt)
  ut = ws * st + wx * xt
  dw = sqrt(dt) * gauss(0, 1, shape)
  xt = xt * ut + sigma_t * dw

Guidance

In practice, we also add guidance $y$ (e.g. label, text, image) as condition to the model, so as to generate the desirable contents. A commonly used technique is called classifier-free guidance (CFG). For ODE, the model becomes $u_t^\theta(x|y)$, and the loss function is extended to include a conditional dropout with rate $\eta$.

, where “…” includes the following:

Upon inference, the vector field $u_t$ is extended to include both the unguided and the guided term with a guidance scale $w > 1$:

For SDE, the model becomes $s_t^\theta(x|y)$, and the loss function is extended as follows:

, where “…” is the same as that in ODE. Upon inference, the score $s_t$ is extended to include both the unguided and the guided term with a guidance scale $w > 1$:

Acknowledgement

This note is heavily inspired by the MIT Class 6.S184.