Diffusion Models
Generation by denoising
Diffusion models are generative models built from two parts: a forward process which gradually adds noise to the data, and a reverse process that removes the noise and restores the data. The forward process is a fixed rule; only the reverse process is learned (Figure 1). They are used in state-of-the-art image generation such as DALL-E, Stable Diffusion and Midjourney.
Assumed background: the normal distribution, expectation and variance, conditional probability, and the basics of neural networks. Markov chains and Bayes' theorem are explained in the article to the extent needed.
The model this article treats mathematically is its basic form, DDPM (Denoising Diffusion Probabilistic Models). Throughout, DDPM refers to this basic form and its standard settings.
Basic Concepts of Diffusion Models
Mathematical Formulation
Forward Process
Noise-adding process
Markov chain adds Gaussian noise gradually. Here $q(x_t | x_{t-1})$ is “the conditional probability distribution of the next state $x_t$ given the previous state $x_{t-1}$”:
\[ q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t} x_{t-1}, \beta_t I) \]The notation \(\mathcal{N}(x;\, \mu,\, \Sigma)\) on the right-hand side means “the density of a normal distribution in the variable \(x\), with mean \(\mu\) and covariance matrix \(\Sigma\)”. The semicolon separates the variable from the parameters. In this equation the correspondence is as follows:
- variable \(x \to x_t\): which variable the density is over. The left-hand side \(q(x_t | x_{t-1})\) is a distribution over \(x_t\), so it is that one
- mean \(\mu \to \sqrt{1-\beta_t}\, x_{t-1}\): the previous state shrunk by a factor of \(\sqrt{1-\beta_t}\) becomes the centre of the distribution
- covariance matrix \(\Sigma \to \beta_t I\): \(I\) is the identity matrix, so the components are independent and every variance equals \(\beta_t\) (isotropic)
The third argument is a covariance matrix rather than a plain variance because \(x_t\) is an image, that is, a vector. For a one-dimensional normal distribution the third argument would just be the variance \(\sigma^2\), but in several dimensions the correlations between components must also be specified, so it becomes a matrix. In this forward process, however, \(\Sigma = \beta_t I\) is a constant multiple of the identity, so it is the simplest case: “every component has variance \(\beta_t\) and they are mutually independent”. This is why the variance computation below only needs to follow a single component.
In words: “scatter noise of variance \(\beta_t\) around a slightly shrunk copy of \(x_{t-1}\)”, which is the same content as the update rule \(x_t = \sqrt{1-\beta_t}\, x_{t-1} + \sqrt{\beta_t}\, \epsilon_t\) introduced below.
Note that a form with two arguments and no semicolon, such as \(\epsilon_t \sim \mathcal{N}(0, I)\), denotes not a density but the distribution itself, and its arguments are (mean, covariance matrix). Since the variable is dropped, the argument positions shift by one compared with \(\mathcal{N}(x;\, \mu,\, \Sigma)\).
Here \(\beta_t\) is the noise schedule (typically \(0.0001 \sim 0.02\)).
Why multiply by \(\sqrt{1-\beta_t}\) — the variance computation
Written as a single-step update, this forward process is \[x_t = \sqrt{1-\beta_t}\, x_{t-1} + \sqrt{\beta_t}\, \epsilon_t, \quad \epsilon_t \sim \mathcal{N}(0, I)\] where \(\epsilon_t\) is drawn afresh at every step, independently of \(x_{t-1}\). The factor \(\sqrt{1-\beta_t}\) is a coefficient that shrinks the signal, and its role becomes clear once the variance is tracked. Below we follow only the \(i\)-th component of \(x_t\) and write its variance as \[v_t := \mathrm{Var}\big(x_t^{(i)}\big) = \mathbb{E}\Big[\big(x_t^{(i)} - \mathbb{E}[x_t^{(i)}]\big)^2\Big]\] Since the added noise is isotropic (\(\Sigma = \beta_t I\)) and the coefficients are scalars, the recurrence has the same form for every component. That said, the value is also independent of the component only when the initial variances agree across components, in which case the covariance matrix can be written as \(v_t I\). Even when they do not agree, the spread between components shrinks by a factor of \(\bar{\alpha}_t\) and vanishes, by the same argument as in (iii). From here on we also abbreviate this as \(\mathrm{Var}(x_t)\). Note \(\mathrm{Var}(\epsilon_t) = 1\).
(i) How the variance changes in one step
For independent random variables \(X, Y\) and constants \(a, b\) we have \(\mathrm{Var}(aX + bY) = a^2\mathrm{Var}(X) + b^2\mathrm{Var}(Y)\). Setting \(a = \sqrt{1-\beta_t}\), \(b = \sqrt{\beta_t}\), \(X = x_{t-1}\) and \(Y = \epsilon_t\) gives \[ v_t = \left(\sqrt{1-\beta_t}\right)^2 v_{t-1} + \left(\sqrt{\beta_t}\right)^2 \cdot 1 = (1-\beta_t)\, v_{t-1} + \beta_t \tag{1} \] The square roots on the coefficients are there because variance scales with the square of a coefficient. If one multiplies by \(\sqrt{1-\beta_t}\) and \(\sqrt{\beta_t}\), their squares \((1-\beta_t)\) and \(\beta_t\) add up to exactly 1.
(ii) \(v = 1\) is a fixed point
Substituting \(v_{t-1} = 1\) into equation (1) gives \[ v_t = (1-\beta_t)\cdot 1 + \beta_t = 1 - \beta_t + \beta_t = 1 \] Hence if \(\mathrm{Var}(x_0) = 1\), induction shows that \(\mathrm{Var}(x_t) = 1\) is preserved for every \(t\) (variance preserving).
Real data need not have variance 1, however. Normalising images to \([-1, 1]\) does not make the variance 1: measured on MNIST normalised to \([-1,1]\), the variance is \(0.3797\) and the mean is \(-0.7387\) (the black background dominates). Scaling to \([-1,1]\) is done mainly to put the values in a range that neural networks handle comfortably. Why a variance away from 1 is acceptable is answered by (iii) below.
(iii) Even a mismatched initial variance is pulled towards 1
Real data does not have variance exactly 1. Subtracting 1 from both sides of equation (1) and writing \(\alpha_t = 1 - \beta_t\), \[ v_t - 1 = (1-\beta_t)\, v_{t-1} + \beta_t - 1 = (1-\beta_t)(v_{t-1} - 1) = \alpha_t\,(v_{t-1} - 1) \] and iterating this \(t\) times gives \[ v_t - 1 = \bar{\alpha}_t\,(v_0 - 1), \qquad \bar{\alpha}_t = \prod_{s=1}^{t}(1-\beta_s) \tag{2} \] In other words, the deviation from 1 shrinks by a factor of \(\bar{\alpha}_t\). Under the standard DDPM setting (\(\beta_t\) linear from \(10^{-4}\) to \(0.02\), \(T = 1000\)) we have \(\bar{\alpha}_T \approx 4.0 \times 10^{-5}\), so even a large mismatch such as \(\mathrm{Var}(x_0) = 4\) gives \(v_T = 1 + 4.0\times10^{-5}\times 3 \approx 1.00012\); the deviation is absorbed automatically. Putting in the MNIST value \(v_0 = 0.3797\) measured above gives \(v_T = 0.999975\). The practical meaning of (iii) is that there is no need to rescale the data to variance 1.
(iv) Without the coefficient the variance diverges
If instead \(x_t = x_{t-1} + \sqrt{\beta_t}\,\epsilon_t\), then in place of equation (1) we get \(v_t = v_{t-1} + \beta_t\), so \[ v_T = v_0 + \sum_{s=1}^{T}\beta_s \] grows monotonically. Under the standard setting above \(\sum_{s}\beta_s = 1000 \times \frac{10^{-4} + 0.02}{2} = 10.05\), so with \(v_0 = 1\) it swells to \(v_T = 11.05\) (standard deviation \(3.32\)). Worse, this value depends on the schedule and on \(v_0\), so the distribution the reverse process must start from is not fixed in advance. The factor \(\sqrt{1-\beta_t}\) is the brake that stops this divergence by shrinking the signal exactly as much as noise is added.
(v) The mean collapses to 0
Preserving the variance at 1 is not by itself enough to reach \(\mathcal{N}(0, I)\). Looking at the mean as well, since \(\mathbb{E}[\epsilon_t] = 0\), \[ \mathbb{E}[x_t] = \sqrt{1-\beta_t}\;\mathbb{E}[x_{t-1}] = \sqrt{\bar{\alpha}_t}\;\mathbb{E}[x_0] \] Since \(\bar{\alpha}_T \approx 4.0\times10^{-5}\), that is \(\sqrt{\bar{\alpha}_T} \approx 0.0064\), the mean shrinks to 0.64% of its original value.
Note that one cannot argue “the mean is 0 and the variance is 1, therefore the distribution is \(\mathcal{N}(0, I)\)”. A distribution with mean 0 and variance 1 need not be normal — a uniform distribution can have both. Mean and variance only show the behaviour required for approaching a normal distribution; on their own they do not determine the shape of the distribution.
That \(x_T\) really does approach the standard normal follows directly from the closed form \(x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon\) in (vi) below. As \(\bar{\alpha}_t \to 0\) the first term vanishes and the coefficient of the second tends to 1, so \[x_T \;\approx\; \epsilon \;\sim\; \mathcal{N}(0, I)\] In other words \(x_T\) is essentially the added noise itself, and it is normal because the noise is normal. This is what justifies the “start from standard normal noise” step of the reverse process.
(vi) Consistency with the closed form
The same conclusion can be checked from \(q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t) I)\) shown next. Computing the variance of \(x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon\) by the same rule gives \[ \mathrm{Var}(x_t) = \bar{\alpha}_t \,\mathrm{Var}(x_0) + (1-\bar{\alpha}_t) \;\overset{\mathrm{Var}(x_0)=1}{=}\; \bar{\alpha}_t + 1 - \bar{\alpha}_t = 1 \] which agrees with applying equation (1) \(t\) times. The fact that the signal weight \(\bar{\alpha}_t\) and the noise weight \(1-\bar{\alpha}_t\) always sum to 1 is what variance preservation really is (Figure 3). Because of this property, a diffusion process of this form is called variance preserving (VP) . For a unified treatment together with the other family that does not preserve the variance (variance exploding, VE), see the SDE formulation in The Mathematics of Stable Diffusion .
With \(\alpha_t = 1 - \beta_t\) and \(\bar{\alpha}_t = \displaystyle\prod_{s=1}^t \alpha_s\):
\[ q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t) I) \]This makes it possible to jump directly to any step \(t\):
\[ x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0, I) \]Reverse Process
A map of this section: the algebra below gets a little long, but the conclusion it reaches is simple.
- To step back from \(x_t\), all that is needed is how much noise it contains.
- That noise \(\epsilon\) is unknown, so a neural network \(\epsilon_\theta(x_t, t)\) is asked to estimate it.
- Training can therefore be run as supervised learning: guess the known noise you added yourself.
The Bayes' theorem and completing-the-square work below is the business of confirming that these three lines really do hold. If you lose the thread, come back here.
The reverse process is parameterised by a neural network \(\theta\):
\[ p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t)) \]Note that what appears before the semicolon here is not \(x_t\) but \(x_{t-1}\) . This is because it describes the distribution of \(x_{t-1}\), one step earlier, conditioned on \(x_t\); the fact that time runs opposite to the forward process shows up in the positions of the variables.
The subscript \(\theta\) is a symbol standing for all the weights of the neural network, not an argument. \(\mu_\theta\) means “a function whose shape is determined by \(\theta\)”: \(x_t, t\) are inputs that vary from sample to sample, whereas \(\theta\) is fixed during inference and changes only during training. \(p_\theta, \mu_\theta, \Sigma_\theta\) all share the same \(\theta\), that is, they come from a single network. The forward process \(q\) carries no subscript because it is determined by \(\beta_t\) alone and is not learned.
Here \(\epsilon_\theta\) is the noise-prediction network. It takes the noisy \(x_t\) and the time step \(t\), and returns a tensor of the same shape as \(x_t\) (for an image, the same width, height and number of channels): \[\epsilon_\theta:\ (x_t,\ t)\ \longmapsto\ \hat{\epsilon}\ \in \mathbb{R}^{\dim(x_t)}\] What that output represents is the noise that was added in producing \(x_t\) from \(x_0\), that is, an estimate of the \(\epsilon\) in the closed form \(x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon\). Concretely it is the U-Net described below (Figure 5).
The time step \(t\) is a single integer, but it is not fed in as a raw number; it is converted into a vector by a sinusoidal embedding and injected into every layer.
This is a fixed transformation, with no learnable parameters, that maps \(t\) to a vector of \(\sin\) and \(\cos\) values at different wavelengths,
\[\big(\sin(\omega_1 t),\ \cos(\omega_1 t),\ \sin(\omega_2 t),\ \cos(\omega_2 t),\ \dots\big),
\qquad \omega_k = 1/10000^{2k/d}\]
(the SinusoidalPositionEmbeddings in the implementation below is exactly this).
Feeding in the raw integer would give an awkwardly wide range of values, whereas in this form every component lies in \([-1,1]\)
and each \(t\) produces a different combination of waveforms, making it easy for the network to tell the time steps apart.
It is the same mechanism as the positional encoding used to represent word positions in a Transformer,
with “time step \(t\)” in place of “position \(pos\)”.
The meaning of the formula, and the derivation showing that relative positions can be expressed by a rotation matrix, are covered in the positional-encoding section of
the Attention Mechanism .
In principle, however, \(\epsilon\) cannot be pinned down from \(x_t\) alone, because many pairs \((x_0, \epsilon)\) can produce the same \(x_t\). The best prediction in the least-squares sense is the conditional expectation \[\epsilon^{\star}(x_t, t) := \mathbb{E}\big[\,\epsilon \ \big|\ x_t,\, t\,\big]\] and this is what training aims at. It is not about hitting each individual \(\epsilon\).
The superscript \(\star\) marks the ideal optimum, and its difference from \(\epsilon_\theta\) is as follows.
- \(\epsilon_\theta(x_t, t)\): the value actually produced by the network with the weights \(\theta\) currently at hand
- \(\epsilon^{\star}(x_t, t)\): the ideal function determined by the forward process and the data distribution alone. It can be defined without any network and does not depend on \(\theta\) (which is why it carries no subscript \(\theta\))
How close \(\epsilon_\theta\) can get to \(\epsilon^{\star}\) depends on the expressive power of the network and on whether optimisation succeeds. There is no guarantee that the two coincide, which is why the notation distinguishes them.
The only thing the network outputs is this \(\epsilon_\theta\); the mean is then computed from it by a fixed formula: \[\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\epsilon_\theta(x_t, t)\right)\] In plain DDPM the covariance matrix \(\Sigma_\theta\) is usually not learned but fixed to \(\beta_t I\) (the subscript \(\theta\) is a remnant of the general form; there are improved variants that do learn it). The implementation below also uses this formula and this fixed variance as they stand.
Where \(\mu_\theta\) comes from — derivation from the posterior
The formula for \(\mu_\theta\) above was not pulled out of thin air; it comes from the true posterior \(q(x_{t-1} | x_t, x_0)\), which is determined by the forward process alone. The reverse \(q(x_{t-1} | x_t)\) on its own is intractable, but adding \(x_0\) to the conditioning makes it Gaussian and writable in closed form.
(i) Decompose into known distributions with Bayes' theorem
The forward process is a Markov chain, so \(q(x_t | x_{t-1}, x_0) = q(x_t | x_{t-1})\), and \[q(x_{t-1} | x_t, x_0) = \frac{q(x_t | x_{t-1})\; q(x_{t-1} | x_0)}{q(x_t | x_0)}\] All three factors on the right are normal distributions already introduced:
- \(q(x_t | x_{t-1}) = \mathcal{N}(x_t;\ \sqrt{\alpha_t}\,x_{t-1},\ \beta_t I)\) (the definition of the forward process itself)
- \(q(x_{t-1} | x_0) = \mathcal{N}(x_{t-1};\ \sqrt{\bar{\alpha}_{t-1}}\,x_0,\ (1-\bar{\alpha}_{t-1}) I)\)
- \(q(x_t | x_0) = \mathcal{N}(x_t;\ \sqrt{\bar{\alpha}_t}\,x_0,\ (1-\bar{\alpha}_t) I)\)
(ii) Complete the square in the exponent
The following is again written for a single component (the computation is identical for each). First, the denominator can be dropped. \(q(x_t | x_0)\) does not contain \(x_{t-1}\), so as a function of \(x_{t-1}\) it is a constant and is absorbed into the normalising constant: \[q(x_{t-1} | x_t, x_0) \;\propto\; q(x_t | x_{t-1})\; q(x_{t-1} | x_0)\] Writing out the two remaining exponents and collecting the terms that do not depend on \(x_{t-1}\) into \(\mathrm{const}\), \[\log q(x_{t-1} | x_t, x_0) = -\frac{1}{2}\left[\frac{\big(x_t - \sqrt{\alpha_t}\,x_{t-1}\big)^2}{\beta_t} + \frac{\big(x_{t-1} - \sqrt{\bar{\alpha}_{t-1}}\,x_0\big)^2}{1-\bar{\alpha}_{t-1}}\right] + \mathrm{const}\]
Expand the bracket and collect in \(x_{t-1}\). Writing the coefficient of the square as \(A\) and that of the linear term as \(-2C\), \[A = \frac{\alpha_t}{\beta_t} + \frac{1}{1-\bar{\alpha}_{t-1}}, \qquad C = \frac{\sqrt{\alpha_t}}{\beta_t}\,x_t + \frac{\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_{t-1}}\,x_0\] (terms in \(x_t^2\) or \(x_0^2\) do not contain \(x_{t-1}\) and go into \(\mathrm{const}\)), so \[\log q = -\frac{1}{2}\Big[A\,x_{t-1}^2 - 2C\,x_{t-1}\Big] + \mathrm{const} = -\frac{A}{2}\left(x_{t-1} - \frac{C}{A}\right)^{2} + \mathrm{const}\] completes the square.
Comparing this with the exponent \(-\dfrac{(x_{t-1}-\tilde{\mu}_t)^2}{2\tilde{\beta}_t}\) of a normal distribution \(\mathcal{N}(x_{t-1};\ \tilde{\mu}_t,\ \tilde{\beta}_t)\), matching coefficients immediately gives \[\tilde{\beta}_t = \frac{1}{A}, \qquad \tilde{\mu}_t = \frac{C}{A} = \tilde{\beta}_t\, C\] All that remains is to compute \(A\). Putting it over a common denominator, \[A = \frac{\alpha_t(1-\bar{\alpha}_{t-1}) + \beta_t}{\beta_t\,(1-\bar{\alpha}_{t-1})}\] and the numerator folds up as \[\alpha_t(1-\bar{\alpha}_{t-1}) + \beta_t = \underbrace{\alpha_t + \beta_t}_{=\,1} - \underbrace{\alpha_t\bar{\alpha}_{t-1}}_{=\,\bar{\alpha}_t} = 1 - \bar{\alpha}_t\] Therefore \[\tilde{\beta}_t = \frac{1}{A} = \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\,\beta_t\] and substituting into \(\tilde{\mu}_t = \tilde{\beta}_t C\) and collecting the coefficients of \(x_t\) and \(x_0\) yields \[q(x_{t-1} | x_t, x_0) = \mathcal{N}\big(x_{t-1};\ \tilde{\mu}_t(x_t, x_0),\ \tilde{\beta}_t I\big) \tag{3}\] \[\tilde{\mu}_t(x_t, x_0) = \frac{\sqrt{\bar{\alpha}_{t-1}}\,\beta_t}{1-\bar{\alpha}_t}\,x_0 + \frac{\sqrt{\alpha_t}\,(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\,x_t\] Note that the mean is a linear combination of \(x_0\) and \(x_t\).
The two identities used here, \(\alpha_t + \beta_t = 1\) and \(\alpha_t\bar{\alpha}_{t-1} = \bar{\alpha}_t\), do the same work again in (iii) below when the coefficient of \(x_t\) folds up into \(1/\sqrt{\alpha_t}\). It is the choice of definitions \(\alpha_t = 1-\beta_t\) and \(\bar{\alpha}_t = \prod_s \alpha_s\) that keeps the algebra clean in both places.
(iii) Eliminate \(x_0\)
Equation (3) cannot be used as it stands, because at generation time \(x_0\) — the very image we are trying to produce — is unknown. So we solve the closed form \(x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon\) for \(x_0\), \[x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}\left(x_t - \sqrt{1-\bar{\alpha}_t}\,\epsilon\right)\] and substitute. Since \(\bar{\alpha}_t = \bar{\alpha}_{t-1}\alpha_t\) gives \(\sqrt{\bar{\alpha}_{t-1}}/\sqrt{\bar{\alpha}_t} = 1/\sqrt{\alpha_t}\), the coefficient of \(x_t\) collects into \[\frac{1}{\sqrt{\alpha_t}\,(1-\bar{\alpha}_t)}\Big[\beta_t + \alpha_t(1-\bar{\alpha}_{t-1})\Big] = \frac{1}{\sqrt{\alpha_t}\,(1-\bar{\alpha}_t)}\Big[\underbrace{\beta_t + \alpha_t}_{=\,1} - \underbrace{\alpha_t\bar{\alpha}_{t-1}}_{=\,\bar{\alpha}_t}\Big] = \frac{1}{\sqrt{\alpha_t}}\] In the end \(x_0\) disappears and \[\tilde{\mu}_t = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\epsilon\right) \tag{4}\] The heart of this derivation is that the definition \(\beta_t + \alpha_t = 1\) is exactly what makes it fold up so neatly.
(iv) Replace the unknown \(\epsilon\) by the network
The only unknown left in equation (4) is the noise \(\epsilon\), and nothing else, the noise that was added when \(x_t\) was made from \(x_0\). Replacing it by the prediction \(\epsilon_\theta(x_t, t)\) is precisely the \(\mu_\theta(x_t, t)\) given above.
This is why the network is asked to predict the noise. Since \(\epsilon\) is all that is needed to assemble the mean of the reverse process, the training objective need only be “guess the noise that was added”. That is why the loss \(L_{simple}\) in the next section takes the simple form \(\|\epsilon - \epsilon_\theta(x_t, t)\|^2\).
Incidentally, it makes almost no difference whether the variance is fixed to \(\tilde{\beta}_t\) or to \(\beta_t\). Their ratio \((1-\bar{\alpha}_{t-1})/(1-\bar{\alpha}_t)\) differs only while \(t\) is small (\(0.743\) at \(t=5\)) and is essentially 1 for large \(t\) (\(1.0000\) at \(t=999\)).
Training Objective
Maximising the log-likelihood directly is hard, so a lower bound on it (the variational lower bound) is maximised instead. Simplifying the loss that follows from it gives the form below (the derivation is in The Mathematics of Stable Diffusion ):
\[ L_{simple} = \mathbb{E}_{t, x_0, \epsilon}\left[ \| \epsilon - \epsilon_\theta(x_t, t) \|^2 \right] \]This is the intuitive goal “predict the added noise \(\epsilon\) accurately”.
U-Net Architecture
Improved Methods
DDPM vs DDIM
DDIM (Denoising Diffusion Implicit Models) speeds up generation by using the same trained network as DDPM and swapping out only the sampling rule. There is no need to retrain.
More precisely, DDIM is not an approximation of DDPM's sampling. It constructs a family of non-Markovian diffusion processes that give the same training objective as DDPM, and includes a deterministic sampler among them. Because the training objective is the same, an \(\epsilon_\theta\) trained for DDPM can be reused as it is.
Classifier-Free Guidance(CFG)
To strengthen generation conditioned on \(c\) (a text prompt, for example), the conditional and unconditional predictions are combined:
\[ \tilde{\epsilon}_\theta(x_t, t, c) = \epsilon_\theta(x_t, t, \varnothing) + w \cdot (\epsilon_\theta(x_t, t, c) - \epsilon_\theta(x_t, t, \varnothing)) \]- \(w\): guidance scale (earlier Stable Diffusion models commonly used fairly large values, around 7, but the appropriate value differs greatly between models and samplers)
- \(w = 1\): ordinary conditional generation
- \(w > 1\): stronger adherence to the condition (higher quality, lower diversity)
What \(w\) actually does — extrapolation along the conditioning direction
The formula above has the shape of an extrapolation. The difference \(\Delta = \epsilon_\theta(x_t,t,c) - \epsilon_\theta(x_t,t,\varnothing)\) is “how much the prediction changes when the condition \(c\) is supplied”, that is, the direction that makes a generic image look like \(c\), and \(w\) is how far to travel in that direction.
- \(w = 0\): purely unconditional. The condition is ignored completely
- \(w = 1\): the \(\varnothing\) terms cancel in \(\epsilon_\theta(\varnothing) + \big(\epsilon_\theta(c) - \epsilon_\theta(\varnothing)\big) = \epsilon_\theta(c)\), leaving the conditional prediction itself
- \(w > 1\): extrapolates past the conditional prediction. In other words it manufactures a prediction the network would never produce on its own
By exaggerating, the trajectory at each step moves towards regions of high density under the conditional distribution. Landing on more typical specimens raises quality, and at the same time distinctive outputs away from the typical become less likely. “Higher quality” and “lower diversity” are not separate phenomena; they are two sides of the same operation.
The price is computation. When \(w \neq 1\), both \(\epsilon_\theta(c)\) and \(\epsilon_\theta(\varnothing)\) are needed, so two forward passes are required and the cost per step roughly doubles.
Latent Diffusion (Stable Diffusion)
Implementation in PyTorch
Everything above is now collected into the smallest code that actually runs. The target is MNIST (28×28 handwritten digits), at a scale where training and generation fit on a CPU alone. Class conditioning and Classifier-Free Guidance are built in, so this single file covers everything from unconditional generation to strongly conditioned generation.
import torch
import torch.nn as nn
import torch.nn.functional as F
class SinusoidalPositionEmbeddings(nn.Module):
"""Positional encoding for the time step"""
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, time):
device = time.device
half_dim = self.dim // 2
embeddings = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
embeddings = time[:, None] * embeddings[None, :]
embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
return embeddings
class ResBlock(nn.Module):
"""The basic block that makes up the U-Net.
Runs GroupNorm -> SiLU -> convolution twice, adding the time and class embedding in
between, and finally adds the input back as a residual. The normalisation is placed
*before* the convolution (pre-activation), so gradients flow well even when many
blocks are stacked.
"""
def __init__(self, in_ch, out_ch, emb_dim, groups=8):
super().__init__()
self.norm1 = nn.GroupNorm(min(groups, in_ch), in_ch)
self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
self.emb_proj = nn.Linear(emb_dim, out_ch) # corresponds to the yellow arrows in Figure 5
self.norm2 = nn.GroupNorm(min(groups, out_ch), out_ch)
self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
# match the residual path with a 1x1 convolution only when the channel counts differ
self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
def forward(self, x, emb):
h = self.conv1(F.silu(self.norm1(x)))
# add the embedding as a per-channel bias (uniform over space)
h = h + self.emb_proj(F.silu(emb))[:, :, None, None]
h = self.conv2(F.silu(self.norm2(h)))
return h + self.skip(x)
class UNet(nn.Module):
"""Class-conditional noise predictor (about 4.7M parameters).
The label y is a class in 0..num_classes-1, and num_classes denotes "no condition".
Keeping one extra embedding slot for the null condition lets the same weights learn
both conditional and unconditional prediction, which is what makes Classifier-Free
Guidance possible at generation time.
"""
def __init__(self, base=64, emb_dim=256, num_classes=10, n_res=2):
super().__init__()
self.num_classes = num_classes
self.time_mlp = nn.Sequential(
SinusoidalPositionEmbeddings(emb_dim),
nn.Linear(emb_dim, emb_dim),
nn.SiLU(),
nn.Linear(emb_dim, emb_dim),
)
# the +1 is the slot for the null condition
self.label_emb = nn.Embedding(num_classes + 1, emb_dim)
c1, c2, c3 = base, base * 2, base * 2 # channel counts at 28x28 / 14x14 / 7x7
self.stem = nn.Conv2d(1, c1, 3, padding=1)
# encoder: n_res residual blocks at each resolution
self.down1 = nn.ModuleList([ResBlock(c1, c1, emb_dim) for _ in range(n_res)])
self.pool1 = nn.Conv2d(c1, c1, 3, stride=2, padding=1) # 28 -> 14
self.down2 = nn.ModuleList([ResBlock(c1 if i == 0 else c2, c2, emb_dim)
for i in range(n_res)])
self.pool2 = nn.Conv2d(c2, c2, 3, stride=2, padding=1) # 14 -> 7
self.down3 = nn.ModuleList([ResBlock(c2 if i == 0 else c3, c3, emb_dim)
for i in range(n_res)])
# bottleneck
self.mid1 = ResBlock(c3, c3, emb_dim)
self.mid2 = ResBlock(c3, c3, emb_dim)
# decoder: skip connections are concatenated, so only the first block sees twice the input
self.up3 = nn.ModuleList([ResBlock(c3 * 2 if i == 0 else c3, c3, emb_dim)
for i in range(n_res)])
self.upsample2 = nn.ConvTranspose2d(c3, c2, 4, stride=2, padding=1) # 7 -> 14
self.up2 = nn.ModuleList([ResBlock(c2 * 2 if i == 0 else c2, c2, emb_dim)
for i in range(n_res)])
self.upsample1 = nn.ConvTranspose2d(c2, c1, 4, stride=2, padding=1) # 14 -> 28
self.up1 = nn.ModuleList([ResBlock(c1 * 2 if i == 0 else c1, c1, emb_dim)
for i in range(n_res)])
self.out_norm = nn.GroupNorm(8, c1)
self.out_conv = nn.Conv2d(c1, 1, 3, padding=1)
# start the output layer at zero: training then begins from "predict no noise", which is stable
nn.init.zeros_(self.out_conv.weight)
nn.init.zeros_(self.out_conv.bias)
def forward(self, x, t, y):
# add time and class into a single embedding, reused by every block (yellow arrows, Figure 5)
emb = self.time_mlp(t) + self.label_emb(y)
h = self.stem(x)
for block in self.down1:
h = block(h, emb)
skip1 = h
h = self.pool1(h)
for block in self.down2:
h = block(h, emb)
skip2 = h
h = self.pool2(h)
for block in self.down3:
h = block(h, emb)
skip3 = h
h = self.mid2(self.mid1(h, emb), emb)
# skip connections are concatenated, not added (red arrows, Figure 5)
h = torch.cat([h, skip3], dim=1)
for block in self.up3:
h = block(h, emb)
h = self.upsample2(h)
h = torch.cat([h, skip2], dim=1)
for block in self.up2:
h = block(h, emb)
h = self.upsample1(h)
h = torch.cat([h, skip1], dim=1)
for block in self.up1:
h = block(h, emb)
return self.out_conv(F.silu(self.out_norm(h)))
# Diffusion scheduler
class DiffusionScheduler:
def __init__(self, num_timesteps=1000, beta_start=1e-4, beta_end=0.02):
self.num_timesteps = num_timesteps
# noise schedule
self.betas = torch.linspace(beta_start, beta_end, num_timesteps)
self.alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
def add_noise(self, x_0, t, noise=None):
"""Forward process: add noise to x_0 to produce x_t"""
if noise is None:
noise = torch.randn_like(x_0)
sqrt_alpha_cumprod = self.alphas_cumprod[t].sqrt()
sqrt_one_minus_alpha_cumprod = (1 - self.alphas_cumprod[t]).sqrt()
# reshape for broadcasting
sqrt_alpha_cumprod = sqrt_alpha_cumprod[:, None, None, None]
sqrt_one_minus_alpha_cumprod = sqrt_one_minus_alpha_cumprod[:, None, None, None]
return sqrt_alpha_cumprod * x_0 + sqrt_one_minus_alpha_cumprod * noise
@torch.no_grad()
def sample(self, model, shape, device, labels=None, guidance_scale=1.0):
"""Reverse process: generate images from noise
labels : the classes to generate. None means unconditional generation
with the null condition everywhere
guidance_scale : the w of Classifier-Free Guidance
w = 0 -> ignore the condition, w = 1 -> the conditional prediction itself,
w > 1 -> extrapolate along the condition (higher quality, lower diversity)
"""
n = shape[0]
null = torch.full((n,), model.num_classes, device=device, dtype=torch.long) # ∅
if labels is None:
labels, guidance_scale = null, 1.0
x = torch.randn(shape, device=device)
for t in reversed(range(self.num_timesteps)):
t_batch = torch.full((n,), t, device=device, dtype=torch.long)
# noise prediction (for w=1 a single conditional pass is enough)
if guidance_scale == 1.0:
predicted_noise = model(x, t_batch, labels)
else:
eps_cond = model(x, t_batch, labels)
eps_uncond = model(x, t_batch, null)
predicted_noise = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
# compute x_{t-1}
alpha = self.alphas[t]
alpha_cumprod = self.alphas_cumprod[t]
beta = self.betas[t]
if t > 0:
noise = torch.randn_like(x)
else:
noise = 0
x = (1 / alpha.sqrt()) * (x - (beta / (1 - alpha_cumprod).sqrt()) * predicted_noise)
x = x + (beta.sqrt() * noise)
return x
# Training loop
def train_diffusion(model, scheduler, dataloader, epochs, device, p_uncond=0.15):
"""Drop the label to the null condition with probability p_uncond, so that one set of
weights learns both the conditional and the unconditional model. Without this, eps(null)
cannot be produced at generation time and CFG is unavailable."""
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
for epoch in range(epochs):
for batch in dataloader:
x_0 = batch[0].to(device)
y = batch[1].to(device)
batch_size = x_0.size(0)
# replace the label with the null condition at a fixed probability
drop = torch.rand(batch_size, device=device) < p_uncond
y = torch.where(drop, torch.full_like(y, model.num_classes), y)
# sample a random time step
t = torch.randint(0, scheduler.num_timesteps, (batch_size,), device=device)
# add noise
noise = torch.randn_like(x_0)
x_t = scheduler.add_noise(x_0, t, noise)
# predict the noise
predicted_noise = model(x_t, t, y)
# compute the loss
loss = F.mse_loss(predicted_noise, noise)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Example usage
if __name__ == "__main__":
model = UNet(base=64, num_classes=10)
scheduler = DiffusionScheduler()
# train_diffusion(model, scheduler, mnist_loader, epochs=20, device="cpu")
# generate ten "3"s with strong guidance
labels = torch.full((10,), 3, dtype=torch.long)
imgs = scheduler.sample(model, (10, 1, 28, 28), "cpu",
labels=labels,
guidance_scale=5.0) # <- the w of the formula; larger than 1 follows the condition more strongly
# passing no labels gives unconditional generation (same as w = 0)
imgs_uncond = scheduler.sample(model, (10, 1, 28, 28), "cpu")
Why the model is built this way
Including normalisation layers, placing several convolutions at each resolution, and securing a reasonable parameter count all bear directly on quality. Keeping everything else identical and removing the normalisation layers, cutting each resolution down to a single convolution, and dropping to about 2 million parameters makes the final loss worse, $0.0219 \to 0.0320$, and the accuracy without guidance ($w = 1$) falls from the $0.920$ reported below to $0.460$. The version above catches up with 20 epochs of the stripped-down version by its second epoch。
Results from actually running it
The code above was trained for 20 epochs on MNIST (60,000 handwritten digits, each pixel normalised to \([-1,1]\)) and then sampled with the \(T = 1000\)-step reverse process. The model has 4,702,593 parameters,
and the training loss after 20 epochs is $0.0219$, essentially flat beyond 10 epochs.
On CPU alone, training takes about 2 hours 45 minutes and generation about 3–6 minutes per 100 images
(for \(w \neq 1\) both the conditional and the unconditional pass are needed, so it roughly doubles).
Because labels are dropped to \(\varnothing\) with probability p_uncond during training,
no retraining is required: changing \(w\) alone switches between unconditional generation and strongly conditioned generation.
Whether the requested digit actually appeared was judged with a separately trained classifier (test accuracy 0.979), giving the following. “Diversity” is the mean pairwise \(L_2\) distance in pixel space among the ten images requested for the same digit; larger means less alike. The Real MNIST row was measured by the same procedure (ten images per class drawn from the training set, put through the same formula).
| Setting | Accuracy | Mean confidence | Diversity |
|---|---|---|---|
| Real MNIST | 0.983 | 0.983 | 18.13 |
| $w = 0$ (condition ignored) | 0.150 | 0.911 | 21.14 |
| $w = 1$ (no guidance) | 0.920 | 0.970 | 17.71 |
| $w = 2$ | 1.000 | 1.000 | 16.66 |
| $w = 3$ | 1.000 | 1.000 | 16.26 |
| $w = 5$ | 1.000 | 1.000 | 15.89 |
$w = 0$ is the setting that ignores the condition entirely. The accuracy of 0.150 is indistinguishable from the chance level $0.10$ for 10 classes (with 100 images the chance fluctuation is about $\pm 0.03$, and $0.150$ falls within twice that). This is what the top-left panel of Figure 8 shows, where the rows do not match the digits. The advantage of CFG is that a single model moves continuously from $w = 0$ to $w = 5$, so no separate model is needed for unconditional generation.
Changing $w$ from $0$ to $1$ alone raises the accuracy from $0.150$ to $0.920$. Since $w = 1$ uses the conditional prediction $\epsilon_\theta(x_t, t, c)$ as it is, without any extrapolation, this gain comes from the conditioning itself, not from guidance. Re-measured over four different initial noises, the mean confidence is $0.976 \pm 0.010$, far larger than the spread. Once the class condition has been learned, the model draws the requested digit without relying on guidance.
Raising $w$ further, the accuracy and mean confidence reach $1.000$ at $w = 2$ and do not move after that. Quality saturates。
Diversity, on the other hand, keeps falling past $w = 2$. Re-measured over four initial noises it is $16.87 \pm 0.15$ at $w = 2$ and $15.89 \pm 0.04$ at $w = 5$; the difference of $-0.98$ is more than three times the spread ($0.30$) and cannot be explained by measurement error. In other words, raising $w$ beyond $2$ means that quality no longer improves while diversity alone keeps shrinking. $w$ is not a quantity that is better when larger; there is a region where only the cost remains.
Setting this beside the diversity of real MNIST, $18.13 \pm 0.27$, makes the meaning of $w$ clearer still. Without guidance ($w = 1$) the value $17.71$ is the closest to the real thing, and raising $w$ moves monotonically away from it towards “too alike”. Note that $21.14$ at $w = 0$ exceeds the real value not because the diversity is better: with the condition inactive, the ten images requested for one digit contain a mixture of different digits, which inflates the distances.
$1.000$ does not mean “better than real”
At $w \geq 2$ both accuracy and mean confidence are $1.000$, exceeding real MNIST ($0.983$). This does not mean the quality has surpassed the real thing. It means that the shapes have become so “typical” that the classifier never hesitates. Real handwritten digits include idiosyncratic shapes, so their confidence does not reach 1.000. The supporting evidence is that the diversity falls monotonically, $21.14 \to 15.89$.
A metric such as confidence measures typicality, not fidelity to the distribution. Distribution-level distances such as FID are used to evaluate generative models precisely in order to avoid this limitation.
The pixel values of the generated images slightly exceed the training range $[-1,1]$. Nothing in the reverse update rule constrains the range to $[-1,1]$, so this is only to be expected; implementations clip just before saving or displaying (as Figure 8 does).
That excess, however, barely depends on $w$. Already at $w = 0$, which uses no extrapolation at all, the range is $[-1.06,\ 1.08]$, and even at $w = 5$ it only reaches $[-1.09,\ 1.10]$. In other words, extrapolation widens the range by only about $0.03$, and the side effect of the form “pushing so hard towards the condition that the pixel values break down” was not observed in this implementation. The price of raising $w$ shows up not in the range but in reduced diversity.
How far these numbers may be generalised
The table above is measured from a single initial noise, 100 images per setting. The accuracy is judged on 100 images, so its chance fluctuation is about $\pm 0.03$; the spread across seeds is about $\pm 0.01$ for the mean confidence and about $\pm 0.15$ for the diversity. No difference smaller than these margins should be read into the table. The three places where the text claims a difference ($w=0$ versus $w=1$, the diversity at $w=2$ versus $w=5$, and the comparison against the stripped-down model) have all been checked to exceed these margins.
Also, the fact that saturation occurred at $w = 2$ cannot be carried over to other subjects. Real image-generation models often use much larger guidance scales (earlier Stable Diffusion models used around 7; the appropriate value differs greatly between models). MNIST has only ten classes and little variation in shape, which is presumably why saturation comes early, but that is an unverified conjecture.
Summary
- Diffusion models learn noise addition (forward process) and noise removal (reverse process)
- The training objective is simple: predict the added noise accurately
- U-Net and a time embedding make up the noise predictor
- DDIM gives fast sampling, and CFG strengthens conditional generation
- Latent Diffusion: efficiency through computation in a latent space
- Already in production use: Stable Diffusion, DALL-E 3, Midjourney and others
The Mathematics of Stable Diffusion develops, rigorously, the unified SDE formulation, the equivalence of the score function and noise prediction, the ODE interpretation of DDIM, and the Bayesian derivation of Classifier-Free Guidance.
Frequently Asked Questions (FAQ)
Q1. What is a diffusion model?
A generative model that first adds noise to data step by step until it becomes white noise (forward process), then learns a network that predicts and removes that noise (reverse process). DDPM (Denoising Diffusion Probabilistic Models) is the basic form, and it underpins Stable Diffusion and similar systems.
Q2. Why do diffusion models generate higher-quality samples than GANs?
They train stably through likelihood-based optimisation without adversarial training, so mode collapse is rare. They are theoretically equivalent to score matching and can generate diverse samples. The drawback is that inference requires many denoising steps and is computationally expensive.
Q3. What is DDIM?
Short for Denoising Diffusion Implicit Models. It is not an approximation of DDPM's stochastic sampling: it constructs non-Markovian diffusion processes that give the same training objective as DDPM and includes a deterministic sampler among them. A network trained for DDPM can be reused as it is, and because no random noise is injected the trajectory survives skipping time steps, cutting the number of inference steps drastically (typically 1000 to 10–50). This made generation practical at usable speeds.
Q4. Why does the forward process multiply by \(\sqrt{1-\beta_t}\)?
To keep the variance at 1. Variance scales with the square of a coefficient, so the update \(x_t = \sqrt{1-\beta_t}\,x_{t-1} + \sqrt{\beta_t}\,\epsilon_t\) has variance \(\mathrm{Var}(x_t) = (1-\beta_t)\mathrm{Var}(x_{t-1}) + \beta_t\); if \(\mathrm{Var}(x_{t-1}) = 1\) then \((1-\beta_t) + \beta_t = 1\) preserves it exactly (variance preserving). Without this coefficient the variance grows as \(\mathrm{Var}(x_0) + \sum_s \beta_s\), diverging to 11.05 under the standard DDPM schedule, and the reverse process could no longer start from standard normal noise.