Your Model's MSE Is Lying to You III: Time Series Diffusion

In Part I of this series, we started from two forecasting models that had the same MSE to three decimal places and a very different level of risk, and we traced the problem back to the loss function itself: training with MSE is the same as assuming a Gaussian whose width never changes. The fix was a second output head for the variance, trained with the Gaussian negative log-likelihood (NLL). In Part II we found that this fix quietly breaks as soon as you forecast more than one step ahead, because feeding the predicted mean back into the model pretends that every previous prediction was perfect, and we repaired it by feeding back samples instead and running many rollouts.Both articles left one assumption untouched. A forecaster with a mean head and a variance head can tell you where the next value of the signal will be and how sure it is about that, but whatever it tells you, it tells you in the form of one symmetric bell curve. For a signal that drifts and jitters this is a perfectly reasonable description, whereas for a signal that can jump, or switch, or sit on a threshold and go either way, it is not, and the uncomfortable part is that the model can have exactly the right mean and exactly the right variance and still describe a future that never happens.So if Part I was about the size of the uncertainty and Part II was about how it travels through time, this part is about its shape. We will move from the simple single bell curve of Part I to a handful of bell curves, from a handful to infinitely many, and from there to a diffusion model, which we will finally plug into the sampled rollout of Part II without changing anything else.···Diffusion for time series, not for imagesAlmost every introduction to diffusion models I've read explains them with images, and actually for a good reason, since generating images is what made them famous. This article does not contain a single image of a cat. What our diffusion model generates is one number: the next value of a signal, given its past.I have to warn you, this is a long article, and deliberately so, because I did not want to skip a single step of the reasoning. You do not need to know anything about diffusion models to follow it. The theory comes first and the data comes last, and every figure can be reproduced with the companion notebook.Because the article is long, here is a map of it, and I would suggest looking at it for a moment before reading on and coming back to it whenever you feel lost in the details.Figure 1: A map of the article. We begin with a problem (stop 1), namely that a mean and a variance are not enough to describe the next value of a signal. We then try the obvious repair, several bell curves instead of one (stop 2), and push it to its limit, infinitely many bell curves (stop 3), which is powerful enough to express any shape but turns out to be very hard to train. Diffusion models are the way out of that difficulty, and the green stops explain them piece by piece: how noise is added (stop 4), how it is removed (stop 5), what the network has to learn in order to remove it (stop 6), why the resulting loss is MSE and why that is fine (stop 7), and how a forecast value is finally drawn (stop 8). Only then, with the theory complete, do we turn to a real signal and to the experiment (stop 9). Every section below opens with a line that tells you which stop you have reached. Image by author.Roadmap of the article in nine stops: the problem with a mean and a variance, mixtures of bell curves, infinitely many bell curves, adding noise, removing noise, the training loss, why MSE is enough, drawing samples, and the forecasting experiment. Each stop lists the question it answers and the answer in one line.···1. The problem: one situation, many possible futures1.1 What a probabilistic forecast actually claimsBefore we can talk about the shape of a forecast, we should be precise about what a probabilistic forecast is a statement about, because this is the point where the example that follows is most easily misunderstood.Suppose you could bring a physical system into exactly the same situation a thousand times, with the same history and the same recent readings, and each time let it run for one more step and write down the next value. You would not get the same number a thousand times, because there are always influences that the past readings do not determine, such as thermal noise, turbulence, or simply things the sensor does not see. What you would get is a collection of a thousand numbers. A probabilistic forecast is a description of that collection. When the two-head model of Part I outputs μ=0.5\mu = 0.5 and σ=0.01\sigma = 0.01, it is claiming that those thousand numbers would cluster tightly around 0.50.5, and when it outputs σ=2\sigma = 2 it is claiming that they would be scattered widely. The distribution does not say that several values occur together; it says that we do not know in advance which of them will occur, and it tells us how often each one would turn up if we could repeat the situation.1.2 A situation that can go two waysLet's now consider a system that sits on a threshold. A good picture is a ball balanced on the top of a small hill between two valleys, but you may equally think of a valve that will either open or stay shut, or of a detector that will either keep its lock or lose it. The ball will not stay on the hilltop, since the slightest disturbance makes it roll down, and it will roll either into the left valley, which I place at −1-1, or into the right valley, which I place at +1+1.If we repeat this situation a thousand times, then in roughly five hundred of the runs the next value yy ends up close to +1+1, and in the other five hundred it ends up close to −1-1, with a small scatter of about 0.10.1 around each of the two positions. In no single run is the ball in both valleys, and in practically no run is it still on the hilltop. The collection of outcomes, however, consists of two separate groups, and so its histogram has two narrow bumps with nothing in between.The forecast does not claim that the ball goes to both valleys; it claims that we cannot know which one, just as with a coin before it is flipped. The honest forecast for a coin is "heads or tails, fifty-fifty", and a forecast of "half-heads" would make no sense, even though it is the average of the two. In the same way, the honest forecast for the ball is "the next value will be either near +1+1 or near −1-1, and I cannot tell you which, because both are equally likely".1.3 What the two-head model reportsNow let's ask what the best possible Gaussian head would say about this situation. A Gaussian head can only ever answer in one format, "around μ\mu, give or take σ\sigma", and when it is trained with the Gaussian NLL, the best answer it can learn is always the same: μ\mu is the average of the outcomes, and σ\sigma is their typical distance from that average. A perfectly trained Gaussian head model therefore reports "around 0, give or take 1". Its single bump is centred exactly on the hilltop, which makes 0 its most likely value, while the real ball practically never ends up there.For the ball, both numbers are easy to work out. Half of the outcomes are near +1+1 and half near−1-1, so the average is 0, right on the hilltop. Every outcome lies about 1 away from 0, so σ\sigma is about 1. (More precisely, σ2=12+0.12=1.01\sigma^2 = 1^2 + 0.1^2 = 1.01, where the 0.120.1^2 comes from the small scatter around each valley, so σ≈1.005\sigma \approx 1.005.)Figure 2: The same situation, run again and again. In every run the ball rolls into exactly one valley, and the collection of outcomes builds up two narrow bumps. The best single Gaussian (blue) has the same mean and the same variance as this collection, and yet it peaks exactly where the ball never goes. Image by author.Animated thought experiment for probabilistic forecasting. The same situation is run again and again, and each time the ball rolls into exactly one of two valleys, so the outcomes form two narrow bumps at minus one and plus one. The best single Gaussian has the same mean and variance, yet it peaks on the hilltop and puts about 38% of its probability where the ball never goes.The mean is correct and the variance is correct, so nothing went wrong during training, and still the forecast is wrong in the only sense that matters, because the value it considers most probable is the one place where the ball will certainly not be. In Part I the two models differed in a number that MSE could not see, namely the variance, and here the forecast and the truth differ in something that even the mean and the variance together cannot see.1.4 Why this matters even more in a rolloutThe damage becomes worse once we remember Part II. If we sample from this Gaussian and feed the samples back into the model, which is exactly what a sampled rollout does, then roughly 38 % (12σ\frac{1}{2} \sigma) of our simulated futures take their very first step into the region between the valleys, −0.5<y<0.5-0.5 < y < 0.5, which the real system practically never visits. The model has never seen such inputs during training, so whatever it predicts next is extrapolation, and a wrong shape at one step has turned into an out-of-distribution input at the next step.A Gaussian head can get the mean right and the variance right and still get the future wrong!···2. First step beyond the bell curve: K bell curves2.1 The ideaIf one bell curve cannot describe two bumps, the most natural repair is to use two bell curves, or more generally KK of them, and to let the model say how much weight each one should get. This is called a mixture of Gaussians, and a forecaster that outputs its parameters is called a mixture density network. Its forecast for the next value is a weighted sum of bell curves:p(y∣context)⏟forecast = ∑k=1Kπk⏞weight N(y∣μk⏟centre, σk2⏟width)⏞bell curve k\underbrace{p(y \mid \text{context})}_{\text{forecast}} \;=\; \sum_{k=1}^{K} \overbrace{\textcolor{#E07B39}{\pi_k}}^{\text{weight}}\; \overbrace{\mathcal{N}\big(y \mid \underbrace{\textcolor{#2B6CB0}{\mu_k}}_{\text{centre}},\, \underbrace{\textcolor{#2F9E5B}{\sigma_k^2}}_{\text{width}}\big)}^{\text{bell curve } k}Every one of the KK components has three numbers. The weight πk\pi_k says how probable this component is, and the weights are positive and sum to one. The mean μk\mu_k says where the component is centred, and the variance σk2\sigma_k^2 says how wide it is. For the ball on the hilltop, K=2K = 2 components with π1=π2=12\pi_1 = \pi_2 = \tfrac12,μ1=−1\mu_1 = -1, μ2=+1\mu_2 = +1 and σ1=σ2=0.1\sigma_1 = \sigma_2 = 0.1 reproduce the truth exactly.2.2 How a mixture produces a sample, and the hidden variable inside itThe way a sample is drawn from a mixture deserves a close look, because it contains the seed of everything that follows. Drawing a sample takes two stages. In the first stage we choose one of the KK components at random, where component kk is chosen with probability πk\pi_k, and in the second stage we draw a value from the bell curve of the chosen component:k∼Categorical(π1,…,πK),y=μk+σk ε,ε∼N(0,1)k \sim \text{Categorical}(\pi_1, \dots, \pi_K), \qquad y = \mu_k + \sigma_k\,\varepsilon, \quad \varepsilon \sim \mathcal{N}(0, 1)For the ball, this is flip a coin to decide which valley, then add a little scatter. The index kk is a quantity that the model uses internally and that never appears in the data, since our sensor records the position of the ball and not a label that says which scenario it belongs to. Such a quantity is called a hidden variable (or latent variable), and with it the mixture can be written in a form that we will meet again and again, the probability of each scenario multiplied by a simple bell curve for that scenario:p(y)=∑k=1Kp(k)⏟how likely is scenario k p(y∣k)⏟a simple bell curvep(y) = \sum_{k=1}^{K} \underbrace{p(k)}_{\text{how likely is scenario } k} \;\underbrace{p(y \mid k)}_{\text{a simple bell curve}}This formula carries the central lesson of this section: a complicated distribution can be built from simple bell curves, provided that a hidden variable decides which bell curve is used. Each piece is as simple as the Gaussian head of Part I, and all the richness comes from the hidden choice.2.3 Where a handful of bell curves stops being enoughA mixture has three weaknesses, and they are worth stating clearly because the next section addresses precisely these.The first is that the number KK has to be chosen in advance and by hand. Two components are right for the ball on the hilltop, but, generally, for a realistic signal you do not know in advance how many components exist.The second is that many shapes are not "a few bumps" at all. A distribution with a long tail on one side, a ridge, or a bump whose position varies continuously can only be approximated by placing a great many narrow components side by side.The third is more subtle and concerns training. Since we never observe which component produced a given data point, the likelihood of that data point has to add up the contributions of all components, which is the sum in the formula above. With K=2K = 2 or K=10K = 10 this sum is easy to compute. Keep it in mind, however, because it is about to become the main obstacle.···3. Second step: infinitely many bell curves3.1 From a list of scenarios to a continuumIf the trouble with a mixture is that KK is a small number we have to choose, the boldest way out is to stop counting. Instead of a hidden index kk that takes one of KK values, we use a hidden variable zz that is a continuous number, and we draw it from the simplest distribution there is, a standard Gaussian. Instead of a list of KK means μ1,…,μK\mu_1, \dots, \mu_K, we use a function μ(z)\mu(z) that assigns a mean to every possible value of zz. With a finite list we could write down one centre per scenario, but there is no way to list a centre for every real number, so we need a rule that produces one, and that rule is the function μ(z)\mu(z). The sum over components then turns into an integral:p(y)⏟forecast = ∫p(z)⏞weight N(y∣μ(z)⏟centre, σ2⏟width)⏞bell curve for this z dz⏟add up one bell curve for every z,p(z)=N(z∣0,1)⏟plain bell curve\underbrace{p(y)}_{\text{forecast}} \;=\;\underbrace{\int\overbrace{\textcolor{#E07B39}{p(z)}}^{\text{weight}}\; \overbrace{\mathcal{N}\big(y \mid \underbrace{\textcolor{#2B6CB0}{\mu(z)}}_{\text{centre}},\, \underbrace{\textcolor{#2F9E5B}{\sigma^2}}_{\text{width}}\big)}^{\text{bell curve for this } z} \; dz}_{\text{add up one bell curve for every } z}, \qquad \textcolor{#E07B39}{p(z)} = \underbrace{\mathcal{N}(z \mid 0, 1)}_{\text{plain bell curve}}It helps to put the two formulas side by side and to see that every ingredient of the mixture has simply been replaced by its continuous counterpart:mixture of KK bell curvecontinuous mixturethe hidden variablean index k∈{1,…,K}k \in \{1, \dots, K\}a number zzhow likely each scenario isthe weight πk\pi_kthe density p(z)p(z)the centre of each bell curvean entry μk\mu_k of a lista value μ(z)\mu(z) of a functionhow the pieces are combineda sum over kkan integral over zznumber of bell curvesKKinfinitely manyDrawing a sample works exactly as before, in two stages, where we first draw the hidden variable and then draw from the bell curve it selects:z∼N(0,1),y=μ(z)+σ ε,ε∼N(0,1)z \sim \mathcal{N}(0, 1), \qquad y = \mu(z) + \sigma\,\varepsilon, \quad \varepsilon \sim \mathcal{N}(0, 1)3.2 An example you can compute by handTo see how much freedom this gives, take the function μ(z)=tanh⁡(8z)\mu(z) = \tanh(8z) and a small width σ=0.1\sigma = 0.1. The function tanh⁡\tanh is an S-shaped curve that is close to −1-1 for negative inputs and close to +1+1 for positive inputs, and the factor 8 makes the transition between the two very steep. Three values of zz show what happens: z=−0.5z = -0.5 receives the centre tanh⁡(−4)≈−1\tanh(-4) \approx -1, z=0.5z = 0.5 receives the centre +1+1, and only a value very close to zero, such as z=0.02z = 0.02, receives a centre in between (tanh⁡(0.16)≈0.16\tanh(0.16) \approx 0.16). A standard Gaussian zz is negative half of the time and positive half of the time and is rarely that close to zero, so about half of the samples land near −1-1 and the other half near +1+1.Figure 3: Gaussian noise in, any shape out. First, single values of zz rise to the curve μ(z)\mu(z), where their centre is read off, and land on the right with a little blur, where the forecast builds up. Then the function changes, from the S-curve to a straight line, a staircase and an exponential, and the forecast follows at once: a straight line gives back a bell curve, a staircase gives three bumps, and a curve that bends upwards gives a long tail. Image by author.Animated demonstration of a continuous mixture. Values drawn from a standard Gaussian are mapped through a function and blurred slightly, and the output histogram builds up on the right. When the function changes, the forecast distribution changes with it: a straight line gives a bell curve, an S-curve gives two bumps, a staircase gives three bumps, and an exponential gives a long tail.If μ\mu is a neural network μθ\mu_\theta with learnable weights, a model of this kind can in principle express any shape whatsoever, without anybody having to choose a number of components. This is the principle behind essentially all modern generative models, and it is the idea we will keep: Gaussian noise in, a learned nonlinear function, any shape out.3.3 The price: the likelihood can no longer be computedSo far we have only seen how such a model generates samples once the function μθ\mu_\theta is given. The hard part is to learn μθ\mu_\theta from data, and the way we have trained every model in this series is maximum likelihood, which means adjusting the weights so that the observed data becomes as probable as possible under the model. For that we need the probability of an observed value yy, which is the integral above.With a mixture of KK components, the corresponding quantity was a sum of KK terms, and we could simply compute it. Here the sum has become an integral over all possible values of the hidden variable, with a neural network inside it. Such an integral has no closed form, which means that there is no formula we could write down and evaluate exactly. The model is powerful enough, but in this form it is very hard to train. (Variational autoencoders train exactly this kind of model by learning a second network that guesses zz from yy; diffusion models take a different route, which is the subject of the rest of this article.)3.4 The way out: build the hidden variables ourselvesDiffusion models resolve this difficulty with two ideas that are both simple, and that I would like you to carry through the rest of the article.The first idea is to stop treating the hidden variable as a mystery. In the continuous mixture, the hidden variable was an abstract number whose relation to the data had to be discovered. A diffusion model instead constructs its hidden variables directly from the data, by a fixed and fully known recipe: it takes the data value and adds noise to it. The hidden variable is then nothing more than a noisy copy of the data, and for every training example we know exactly which hidden variables belong to it, because we made them ourselves.The second idea is to replace one large leap by many small steps. In the continuous mixture, a single function μ(z)\mu(z) had to turn pure noise into data in one go, deciding everything at once. A diffusion model instead adds noise gradually, creating a chain of copies that range from almost clean to pure noise, and learns to walk back one small step at a time. Think of a film of ink spreading in water: guessing the first frame from the last one is hopeless, but guessing frame 99 from frame 100 is easy, because the two frames are almost identical. Each small step back is so simple that one bell curve describes it well, which is exactly what the Gaussian head of Part I can do. And because we made the noisy copies ourselves (the first idea), we always know which copy belongs to which data point, so every step can be trained directly.Together they give us two paired processes. The forward process adds noise to a clean value step by step until only noise is left, and involves no learning at all (Section 4). The reverse process is a neural network that removes a little noise at each step (Section 5), and we will have to work out what it should be trained to do (Sections 6 and 7) and how it produces a forecast value (Section 8).Figure 4: Both processes in one picture. Forward (orange), the noising recipe (Section 4) melts the two bumps into a bell curve. In reverse (green), fifty small Gaussian steps split the bell curve into two bumps again. The reverse steps shown here use the exact best guess of the clean value, which is the function a perfectly trained network would learn (Section 7). Image by author.Animated view of both diffusion processes. Going forward, adding noise step by step melts the two-bump distribution into a plain bell curve. Going in reverse, fifty small Gaussian steps turn the bell curve back into two bumps. The reverse steps use the exact best guess of the clean value, with no neural network involved.3.5 A word on notationFrom here on there are two different kinds of "time" in play, and keeping them apart avoids most of the confusion around diffusion models for time series. In this series, tt has always been the time index of the signal and TT the context length, so I will keep those and use the letter nn for the noise level. Most diffusion papers call the noise level tt and the number of levels TT, which you should keep in mind if you read them alongside this article.Two remarks will keep the next sections light. Since we forecast one reading at a time, the quantity being noised and denoised is a single number, the next value of the signal (for vectors, such as images, every formula holds in the same form for each component). And everything in Sections 4 to 8 happens for a given context: the forecast distribution is always the distribution of the next value given the past, but I leave the context out of the formulas until Section 9, where it comes back as an extra input of the network.···4. The forward process: adding noiseGoal of this section. We want a recipe that gradually turns a clean value y0y_0 (during training) into pure Gaussian noise over NN steps, and we want to understand every symbol in it.4.1 Two things to know about GaussiansA Gaussian distribution N(μ,σ2)\mathcal{N}(\mu, \sigma^2) describes a random number that is probably close to its mean μ\mu, with a spread that is controlled by its variance σ2\sigma^2. The noise we will add is a draw from the standard Gaussian, ε∼N(0,1)\varepsilon \sim \mathcal{N}(0, 1), which has mean 0 and variance 1. There is a convenient way to draw from any Gaussian using only standard noise, which is to draw ε\varepsilon and then shift and scale it:X=μ+σ ε,ε∼N(0,1)X = \mu + \sigma\,\varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, 1)This is called reparameterisation, and it is exactly how the Gaussian head of Part I draws a sample. It looks like a small trick, but it is the engine of everything in this section, because it lets us write every noising step as an ordinary equation instead of as a probability distribution. The notation N(X∣μ,σ2)\mathcal{N}(X \mid \mu, \sigma^2) that appears below simply means "the density of a Gaussian with mean μ\mu and variance σ2\sigma^2, evaluated at XX".4.2 One small stepWe start with a clean value y0y_0 and define a noise schedule β1,β2,…,βN\beta_1, \beta_2, \dots, \beta_N, a list of small numbers that control how much noise is added at each step. At every step we take the value from the previous step, shrink it a little and add a little Gaussian noise. Written as a probability distribution, the rule isq(yn∣yn−1)⏟one noising step=N(yn ∣ 1−βn yn−1⏟centre: old value, shrunk, βn⏟width: noise added)\underbrace{q(y_n \mid y_{n-1})}_{\text{one noising step}}= \mathcal{N}\big(y_n \;\big|\; \underbrace{\textcolor{#2B6CB0}{\sqrt{1 - \beta_n}\;y_{n-1}}}_{\text{centre: old value, shrunk}},\; \underbrace{\textcolor{#2F9E5B}{\beta_n}}_{\text{width: noise added}}\big)and in reparameterised form, which is the form we actually use in code, yn=1−βn yn−1⏟shrink the old value a little + βn εn⏟add a little fresh noise,εn∼N(0,1) \boxed{\,y_n = \underbrace{\textcolor{#2B6CB0}{\sqrt{1 - \beta_n}\;y_{n-1}}}_{\text{shrink the old value a little}} \;+\; \underbrace{\textcolor{#2F9E5B}{\sqrt{\beta_n}\;\varepsilon_n}}_{\text{add a little fresh noise}}, \qquad \varepsilon_n \sim \mathcal{N}(0, 1)\,}The letterqq is used for this fixed forward process, and we will use pθp_\theta later for the learned model. Let's dissect each piece.What is βn\beta_n? It is a small number between 0 and 1, and you can think of it as a dial: βn=0\beta_n = 0 means that no noise is added at this step, and βn=1\beta_n = 1 means that the value is thrown away completely and replaced by noise. In the original paper on these models, βn\beta_n grows linearly from 0.00010.0001 to 0.020.02 over N=1000N = 1000 steps, so every individual step changes the value only slightly.4.3 Why 1−βn\sqrt{1 - \beta_n}, and not simply 1−βn1 - \beta_n?This is a design choice, but it is a principled one, and seeing where it comes from takes only one goal and two basic facts about random numbers.The goal: keep the variance equal to 1 at every step. Generation will start from a standard Gaussian, N(0,1)\mathcal{N}(0, 1), so the forward process must end exactly there, because otherwise the network would be trained on one kind of input and used on another. The simplest way to guarantee this is to standardise the data to variance 1 and to keep the variance at 1 at every single step, which also keeps the inputs of the network on the same scale at every noise level. Rule 1: scaling a random number squares the factor in its variance. If XX has variance σ2\sigma^2 and we multiply it by a constant aa, then Var(aX)=a2 Var(X)\mathrm{Var}(aX) = a^2\,\mathrm{Var}(X). If you stretch a distribution by a factor of 2, its standard deviation doubles and its variance becomes four times as large, because variance is measured in squared units.Rule 2: the variances of independent random numbers add up. If XX and YY are independent, then Var(X+Y)=Var(X)+Var(Y)\mathrm{Var}(X + Y) = \mathrm{Var}(X) + \mathrm{Var}(Y), with no cross term, because independent quantities do not vary together.Building the formula from the goal. Let's mix the previous value with fresh noise using two unknown constants aa and bb,yn=a yn−1+b εny_n = a\,y_{n-1} + b\,\varepsilon_nand compute the variance of the result. Since yn−1y_{n-1} and εn\varepsilon_n are independent, Rule 2 lets us add the variances of the two terms, and Rule 1 tells us how the constants enter:Var(yn)=a2 Var(yn−1)⏟= 1+b2 Var(εn)⏟= 1=a2+b2\mathrm{Var}(y_n) = a^2\,\underbrace{\mathrm{Var}(y_{n-1})}_{=\,1} + b^2\,\underbrace{\mathrm{Var}(\varepsilon_n)}_{=\,1} = a^2 + b^2Variance preservation therefore requires a2+b2=1a^2 + b^2 = 1. The noise schedule decides which share of this variance budget is handed to the fresh noise at step nn, namely b2=βnb^2 = \beta_n, and the rest must go to the old value, a2=1−βna^2 = 1 - \beta_n. Taking square roots givesa=1−βn,b=βna = \sqrt{1 - \beta_n}, \qquad b = \sqrt{\beta_n}so the square roots are not arbitrary at all, since they are the only positive solution of a2+b2=1a^2 + b^2 = 1 once βn\beta_n has fixed the share of the noise.Verifying that it works. Plugging back in, Var(yn)=(1−βn)2+(βn)2=(1−βn)+βn=1\mathrm{Var}(y_n) = (\sqrt{1 - \beta_n})^2 + (\sqrt{\beta_n})^2 = (1 - \beta_n) + \beta_n = 1. At every step the signal shrinks a little (it is scaled by 1−βn<1\sqrt{1 - \beta_n} < 1) while noise fills the gap (it is added with scale βn\sqrt{\beta_n}), and the total variance stays perfectly balanced.What goes wrong without the square roots? Suppose we had naively used yn=(1−βn) yn−1+βn εny_n = (1 - \beta_n)\,y_{n-1} + \beta_n\,\varepsilon_n. The variance would then follow the rule Var(yn)=(1−βn)2 Var(yn−1)+βn2\mathrm{Var}(y_n) = (1 - \beta_n)^2\,\mathrm{Var}(y_{n-1}) + \beta_n^2. With βn=0.1\beta_n = 0.1 and a starting variance of 1, the first step gives 0.81+0.01=0.820.81 + 0.01 = 0.82, and if we keep applying the rule the variance continues to fall until it settles at the value where it no longer changes, v=0.81 v+0.01v = 0.81\,v + 0.01, which is v≈0.053v \approx 0.053. The chain would end at a narrow Gaussian with variance 0.0530.053 and not at the standard Gaussian with which generation starts. Suppose instead we had used yn=yn−1+βn εny_n = y_{n-1} + \beta_n\,\varepsilon_n, adding noise without scaling the old value down. Then the variance would grow by βn2\beta_n^2 at every step, and, worse, the clean value y0y_0 would never be forgotten, because its coefficient would stay equal to 1 forever and the chain would never arrive at pure noise. 4.4 The jump formula: from the clean value to any noise level in one shotTraining will require noisy versions yny_n of our data at many different noise levels, millions of times. If we had to walk through all the individual steps each time, training would be hopelessly slow, so we need a shortcut, a direct formula that jumps from the clean value y0y_0 to any noise level nn in one go. The idea behind it is that many small independent Gaussian noises add up to one bigger Gaussian noise.Step 0: the Gaussian addition rule. Before doing any algebra, we need one fact that will do all the heavy lifting. If A∼N(0,σ12)A \sim \mathcal{N}(0, \sigma_1^2) and B∼N(0,σ22)B \sim \mathcal{N}(0, \sigma_2^2) are independent, thenA+B∼N(0, σ12+σ22)A + B \sim \mathcal{N}(0,\; \sigma_1^2 + \sigma_2^2)The variance part is Rule 2 from above. The additional statement, that the sum of two independent Gaussians is again a Gaussian, is a standard result of probability theory. Together they say that two independent Gaussian noises can always be replaced by a single Gaussian noise whose variance is the sum of the two variances.Step 1: shorter notation. Let αn=1−βn\alpha_n = 1 - \beta_n, so that the one-step rule becomesyn=αn yn−1+1−αn εny_n = \sqrt{\alpha_n}\;y_{n-1} + \sqrt{1 - \alpha_n}\;\varepsilon_nStep 2: expanding two steps. Let's write out two consecutive steps to see the pattern. The first step goes from y0y_0 to y1y_1:y1=α1 y0+1−α1 ε1y_1 = \sqrt{\alpha_1}\;y_0 + \sqrt{1 - \alpha_1}\;\varepsilon_1The second step goes from y1y_1 to y2y_2, and we substitute the expression for y1y_1:y2=α2 y1+1−α2 ε2=α2 (α1 y0+1−α1 ε1)+1−α2 ε2y_2 = \sqrt{\alpha_2}\;y_1 + \sqrt{1 - \alpha_2}\;\varepsilon_2= \sqrt{\alpha_2}\,\Big(\sqrt{\alpha_1}\;y_0 + \sqrt{1 - \alpha_1}\;\varepsilon_1\Big) + \sqrt{1 - \alpha_2}\;\varepsilon_2y2=α1α2 y0⏟signal + α2(1−α1) ε1+1−α2 ε2⏟two noise termsy_2 = \underbrace{\sqrt{\alpha_1 \alpha_2}\;y_0}_{\text{signal}} \;+\; \underbrace{\sqrt{\alpha_2 (1 - \alpha_1)}\;\varepsilon_1 + \sqrt{1 - \alpha_2}\;\varepsilon_2}_{\text{two noise terms}}We now have one clean signal term and two separate noise terms, and the next step is to merge the latter.Step 3: merging the two noise terms. This is where the addition rule of Step 0 does its work. The noisesε1\varepsilon_1andε2\varepsilon_2are independent standard Gaussians, so by Rule 1 the two scaled noise terms are Gaussians with variances α2(1−α1)\alpha_2(1 - \alpha_1) and 1−α21 - \alpha_2, and by the addition rule their sum is a single Gaussian whose variance isα2(1−α1)+(1−α2)=α2−α1α2+1−α2=1−α1α2\alpha_2 (1 - \alpha_1) + (1 - \alpha_2) = \alpha_2 - \alpha_1 \alpha_2 + 1 - \alpha_2 = 1 - \alpha_1 \alpha_2The two noise terms therefore collapse into one,α2(1−α1) ε1+1−α2 ε2 = 1−α1α2 εˉ,εˉ∼N(0,1)\sqrt{\alpha_2 (1 - \alpha_1)}\;\varepsilon_1 + \sqrt{1 - \alpha_2}\;\varepsilon_2 \;=\; \sqrt{1 - \alpha_1 \alpha_2}\;\bar\varepsilon, \qquad \bar\varepsilon \sim \mathcal{N}(0, 1)where the equality means that both sides have the same distribution, and substituting back givesy2=α1α2 y0+1−α1α2 εˉy_2 = \sqrt{\alpha_1 \alpha_2}\;y_0 + \sqrt{1 - \alpha_1 \alpha_2}\;\bar\varepsilonStep 4: the general formula. The pattern is now visible, and the same arithmetic extends to any number of steps: the signal is multiplied by the square root of the product of all the α\alpha's, and the combined noise has variance one minus that product. If we define the running productαˉn=α1 α2⋯αn=∏s=1nαs\bar\alpha_n = \alpha_1 \,\alpha_2 \cdots \alpha_n = \prod_{s=1}^{n} \alpha_sthen for any noise level nn all the intermediate noise terms collapse into a single one, and we obtain the jump formula: yn=αˉn y0⏟what is left of the clean value + 1−αˉn ε⏟all the noise of n steps, in one draw,ε∼N(0,1) \boxed{\,y_n = \underbrace{\textcolor{#2B6CB0}{\sqrt{\bar\alpha_n}\;y_0}}_{\text{what is left of the clean value}} \;+\; \underbrace{\textcolor{#2F9E5B}{\sqrt{1 - \bar\alpha_n}\;\varepsilon}}_{\text{all the noise of } n \text{ steps, in one draw}}, \qquad \varepsilon \sim \mathcal{N}(0, 1)\,}Written as a probability distribution, this says the same thing:q(yn∣y0)=N(yn∣αˉn y0⏟centre, 1−αˉn⏟width)q(y_n \mid y_0) = \mathcal{N}\big(y_n \mid \underbrace{\textcolor{#2B6CB0}{\sqrt{\bar\alpha_n}\;y_0}}_{\text{centre}},\; \underbrace{\textcolor{#2F9E5B}{1 - \bar\alpha_n}}_{\text{width}}\big)You can read it as a recipe, which says that we keep a fraction αˉn\sqrt{\bar\alpha_n} of the clean value and add noise with standard deviation1−αˉn\sqrt{1 - \bar\alpha_n}. Note that theε\varepsilonin this formula is the total noise accumulated since y0y_0, and not the noise of step nn alone.The payoff. We never have to loop through the noising steps during training. We pick a noise level, draw a single ε\varepsilon, and compute yny_n from y0y_0 in one line, which is what makes training a diffusion model computationally feasible.4.5 The noise scheduleWe still have to choose the numbers β1,…,βN\beta_1, \dots, \beta_N, and it is worth understanding first why the choice matters at all. The jump formula depends on the schedule only through αˉn\bar\alpha_n, the share of the signal that survives after nn steps, so choosing a schedule means choosing how this share falls from 1 (clean data at n=0n = 0) to practically 0 (pure noise at n=Nn = N).The schedule decides how our fixed number of steps is spent. If the signal disappears too slowly, the chain never reaches pure noise, and generation would then start from inputs the network has never seen. If it disappears too quickly, most of the steps are wasted on turning noise into more noise, while the interesting part, where the bumps merge, is squeezed into a few large steps, and large steps are exactly what the reverse process cannot handle (Section 5). A good schedule lets the signal fade at a steady pace, so that every step does a similar, small amount of work.One option is to choose theβn\beta_n directly, such as the linear schedule from 0.00010.0001 to 0.020.02 mentioned above. The other option is to design the curve αˉn\bar\alpha_n itself, as a smooth descent from 1 to 0, which guarantees an even pace for any number of steps, and to derive the individual steps from it. Since αˉn=αˉn−1 αn\bar\alpha_n = \bar\alpha_{n-1}\,\alpha_n by the definition of the running product, we haveβn=1−αn=1−αˉnαˉn−1\beta_n = 1 - \alpha_n = 1 - \frac{\bar\alpha_n}{\bar\alpha_{n-1}}A popular curve of this kind is the cosine schedule,αˉn=f(n)f(0),f(n)=cos⁡2 ⁣(π2⋅n/N+0.0081+0.008)\bar\alpha_n = \frac{f(n)}{f(0)}, \qquad f(n) = \cos^2\!\left(\frac{\pi}{2} \cdot \frac{n/N + 0.008}{1 + 0.008}\right)in which the division by f(0)f(0)makes sure that αˉ0=1\bar\alpha_0 = 1 exactly, the cosine reaches zero at n=Nn = N so that no signal is left at the end, and the small offset 0.0080.008 keeps the very first steps from being vanishingly small.The linear schedule, run over only 50 steps, still leaves 60 % of the signal at the end, so it fails the too slow test, and even with 1,000 steps it is rather fast, since about a third of its steps happen when less than 1 % of the signal is left. The cosine curve falls gently and evenly, which is why it was proposed and why the companion notebook uses it with N=50N = 50 steps. With so few steps, the first steps are small (βn\beta_n below0.010.01) but the last few are not (about 0.560.56, 0.750.75 and 0.900.90), which does little harm in practice because at that point almost no signal is left anyway, and which you can cure by increasing NN at the price of slower sampling.Figure 5: As the noise level nn grows, the share of signal αˉn\bar\alpha_n falls while the share of noise rises (left), and the two narrow bumps widen, slide together and merge into a standard Gaussian (right). Image by author.Animated diffusion forward process: a cosine noise schedule shrinks the signal and adds noise step by step, turning a two-peaked distribution into a standard Gaussian.···5. The reverse process: removing noise5.1 One big leap is hard, one small step is easySuppose somebody hands you pure noise yNy_N and asks what the clean value y0y_0 was. Pure noise carries no information about the data, so the honest answer is anything the data could be, which is the full two-bump distribution, and describing that with one Gaussian brings us straight back to Section 1. This is also the single large leap that the continuous mixture of Section 3 had to perform.Now suppose instead that you are handed yny_n and asked only what yn−1y_{n-1} was, one step earlier. The one-step rule tells us that yny_n is 1−βn yn−1\sqrt{1 - \beta_n}\;y_{n-1} plus a small amount of noise, so yn−1y_{n-1} must have been close to yn/1−βny_n / \sqrt{1 - \beta_n}, give or take about βn\sqrt{\beta_n}. The answer lives in a small window, and inside a small window a distribution cannot do anything dramatic, since it cannot contain two bumps that are far apart, which is why a single Gaussian describes it well. For the two-bump data, the distribution of "where was this value kk steps earlier" can be computed exactly, and Figure 6 shows how it changes as we look further and further back.Figure 6: One small step back is a bell curve, one big leap back is not. We hold a noisy value (the orange dot, at noise level 30) and ask where it was one step earlier, two steps earlier, and so on. One step back, the answer is a narrow bell curve, which the best single Gaussian (blue, dashed) covers almost perfectly. The further back we ask, the wider the answer becomes, until it splits into the two bumps of the data, of which a single bell curve covers less than a quarter. These distributions are exact; nothing is learned or simulated here. Image by author.Animated explanation of why diffusion models use many small steps. Starting from one noisy value at noise level 30, the animation shows where that value was one step earlier, two steps earlier, and so on. One step back the answer is a narrow bell curve that a single Gaussian covers almost completely; thirty steps back it has split into two bumps, of which a single Gaussian covers less than a quarter.This is the key observation of the whole method: when each forward step adds only a little noise, each reverse step is approximately Gaussian, even though the distribution of the data is not.5.2 Each reverse step is a Gaussian headWe therefore model every reverse step as a Gaussian whose mean is predicted by a neural network:pθ(yn−1∣yn)⏟one denoising step=N(yn−1∣μθ(yn,n)⏟centre: learned, σn2⏟width: fixed)\underbrace{p_\theta(y_{n-1} \mid y_n)}_{\text{one denoising step}}= \mathcal{N}\big(y_{n-1} \mid \underbrace{\textcolor{#2B6CB0}{\mu_\theta(y_n, n)}}_{\text{centre: learned}},\; \underbrace{\textcolor{#2F9E5B}{\sigma_n^2}}_{\text{width: fixed}}\big)If this looks familiar, it should, because it is the two-head model of Part I with two small changes:The variance is fixed, not learned. The network only predicts the mean; the width of each step is set by the noise schedule.It is used NN times in a row, not once. Each call removes a little noise, taking the noisy value yny_n and the noise level nn as extra inputs alongside the context.A single network serves all the steps, and it is told at which noise level it is working, so that it can behave differently when it is removing heavy noise and when it is polishing the last details.5.3 Where does the shape come from?Since every single step is Gaussian, you may wonder where the non-Gaussian shape comes from, and there are two ways of seeing it that tie this section to the earlier ones.The first way is to look at the last step of the chain. The final value y0y_0 is drawn from a bell curve whose centre μθ(y1,1)\mu_\theta(y_1, 1) depends on the previous value y1y_1, and y1y_1 is itself random, so the distribution of y0y_0 ispθ(y0)=∫pθ(y1) N(y0∣μθ(y1,1), σ12) dy1p_\theta(y_0) = \int p_\theta(y_1)\;\mathcal{N}\big(y_0 \mid \mu_\theta(y_1, 1),\; \sigma_1^2\big)\;dy_1If you compare this with the continuous mixture (Section 3), you will see that it is the same formula, with the noisy value y1y_1 in the role of the hidden variable zz. A diffusion model is an infinite mixture of bell curves, and in fact a whole tower of them, because the distribution of y1y_1 is in turn an infinite mixture over y2y_2, and so on up to pure noise.The second way is to look at the chain as a whole. Each step moves the sample by a small amount in a direction chosen by a nonlinear function μθ\mu_\theta. We have in fact already seen this mechanism at work in Part II without calling it by this name: a sampled rollout is also a chain of Gaussian steps in which each step depends on the previous sample, which is why the distribution of a rollout after sixteen steps is not a Gaussian even though every individual step is one. Part II composed Gaussian steps along the time axis of the signal, and a diffusion model composes them along a second, artificial axis, the noise level, so that even one single time step can have any shape.···6. Training: what should the network learn?Goal of this section. We want to train a model εθ(yn,n)\varepsilon_\theta(y_n, n) that, given a noisy value yny_n and its noise level nn, predicts the noise ε\varepsilon that was added to the clean value y0y_0.Step 1: why can't we just invert the forward process?The jump formula told us how to go from a clean value to a noisy one in one shot:yn=αˉn y0+1−αˉn εy_n = \sqrt{\bar\alpha_n}\;y_0 + \sqrt{1 - \bar\alpha_n}\;\varepsilonNaively, one might think that this is all we need, since we can simply rearrange it for y0y_0:y0⏟what we want=yn⏞what we hold−1−αˉn ε⏞unknown!αˉn\underbrace{y_0}_{\text{what we want}} = \frac{\overbrace{y_n}^{\text{what we hold}} - \sqrt{1 - \bar\alpha_n}\;\overbrace{\textcolor{#2F9E5B}{\varepsilon}}^{\text{unknown!}}}{\sqrt{\bar\alpha_n}}The problem is that ε\varepsilon is gone. When we ran the forward process we drew ε\varepsilon at random and mixed it into yny_n, and at generation time that particular ε\varepsilon is not available to us. In terms of the equation, we have one equation with two unknowns, y0y_0 and ε\varepsilon, and infinitely many pairs of a clean value and a noise could have produced the noisy value we are holding. (During training the situation is different, since there we chose y0y_0 and drew ε\varepsilon ourselves and therefore know both.)This is exactly why we need a neural network: to estimate ε\varepsilon from the noisy value alone. If we can make a good guess of what ε\varepsilon was, the rearranged formula gives us a guess of y0y_0.The derivation that follows is the hardest part of the article, so here is its destination and its main moves before the first equation. (1) We cannot compute how probable the data is under the model, because that would require every possible path of noisy values. (2) So we score the model only on paths that we generate ourselves, which gives a lower bound on that probability. (3) This score splits into one comparison per noise level, between the step of the network and an ideal step that knows the clean value. (4) Both steps are bell curves of the same width, so each comparison is a squared error, and after a change of variables it is a squared error on the noise.Step 2: what should the network actually optimise?The goal. As everywhere in this series, we want the model to give a high probability to the data we actually observed, so we want to maximise log⁡pθ(y0)\log p_\theta(y_0), averaged over the dataset. (We use logarithms because they turn products into sums and because they penalise near-zero probabilities heavily: if the model assigns a probability of 0.0010.001 to something that really happened, then log⁡0.001≈−6.9\log 0.001 \approx -6.9.)The problem. The model never produces y0y_0 directly. It walks a path, from pure noise yNy_N throughyN−1,…,y1y_{N-1}, \dots, y_1 toy0y_0, and what it defines directly is the probability of a complete path, which is the product of the probabilities of its steps:pθ(y0,y1,…,yN)=p(yN)∏n=1Npθ(yn−1∣yn)p_\theta(y_0, y_1, \dots, y_N) = p(y_N)\prod_{n=1}^{N} p_\theta(y_{n-1} \mid y_n)In words, this is the probability of starting from this particular noise (which is just a standard Gaussian) multiplied by the probability of each denoising step. To obtain the probability of y0y_0 alone, we would have to add up every path that ends there:pθ(y0)=∫pθ(y0,y1,…,yN) dy1⋯dyNp_\theta(y_0) = \int p_\theta(y_0, y_1, \dots, y_N)\;dy_1 \cdots dy_NThink of a city map. "How likely is this one route?" is an easy question, since you multiply the probabilities of each turn, whereas "how likely am I to end up at this address, by any route?" means summing over all routes. With N=50N = 50 steps, even a crude grid of 100 values per step would need 10050100^{50} evaluations of the network!The solution: score only the paths we make ourselves. We bring in our own forward process, q(y1:N∣y0)q(y_{1:N} \mid y_0), where y1:Ny_{1:N} is shorthand for all the noisy values y1,…,yNy_1, \dots, y_N together. We know this process exactly, we can sample from it as often as we like, and its paths are precisely the ones that belong to this particular y0y_0. Four short moves then give a quantity we can compute.(a) Multiply by 1. We multiply and divide by qq, which changes nothing:pθ(y0)=∫q(y1:N∣y0) pθ(y0,y1:N)q(y1:N∣y0) dy1:Np_\theta(y_0) = \int q(y_{1:N} \mid y_0)\;\frac{p_\theta(y_0, y_{1:N})}{q(y_{1:N} \mid y_0)}\;dy_{1:N}(b) Recognise an average. An integral in which something is weighted by a distribution qq is an average over samples from qq:pθ(y0)=Eq ⁣[pθ(y0,y1:N)q(y1:N∣y0)]p_\theta(y_0) = \mathbb{E}_q\!\left[\frac{p_\theta(y_0, y_{1:N})}{q(y_{1:N} \mid y_0)}\right](c) Take the logarithm of both sides, since the log-likelihood is what we are after.(d) Move the logarithm inside the average (Jensen's inequality). The logarithm bends downwards, so the logarithm of an average is always at least as large as the average of the logarithms. For the two numbers 1 and 100, for example, the logarithm of their average is log⁡50.5≈3.9\log 50.5 \approx 3.9, whereas the average of their logarithms is (0+4.6)/2=2.3(0 + 4.6)/2 = 2.3. Hencelog⁡pθ(y0)⏟what we want (cannot compute) ≥ Eq⏟average over noisepaths we generate ⁣[log⁡pθ(y0,y1:N)⏞model: denoise along the pathq(y1:N∣y0)⏟us: noise along the path]=ELBO\underbrace{\log p_\theta(y_0)}_{\text{what we want (cannot compute)}} \;\ge\; \underbrace{\mathbb{E}_q}_{\substack{\text{average over noise}\\\text{paths we generate}}}\!\left[\log \frac{\overbrace{p_\theta(y_0, y_{1:N})}^{\text{model: denoise along the path}}} {\underbrace{q(y_{1:N} \mid y_0)}_{\text{us: noise along the path}}}\right] = \text{ELBO}The right-hand side is called the evidence lower bound ("evidence" is another name for the probability of the data). Unlike the left-hand side, it can be computed: we generate noise paths ourselves and evaluate two things we know, qq (how we add noise) and pθp_\theta (how the network removes it). And because it is a floor under the log-likelihood, pushing the floor up during training pushes the real thing up with it.Step 3: how does the ELBO become MSE?A. One comparison per noise level. The ratio inside the ELBO is a product of NN backward steps (the model) divided by a product of NN forward steps (the noising). To compare them step by step, we turn every forward step around with Bayes' theorem, so that it points backwards too. We are allowed to condition everything on y0y_0 while doing so, because a forward step depends only on the value right before it and not on y0y_0. With only two steps you can see the whole trick at once:q(y1∣y0) q(y2∣y1) = q(y1∣y0)⋅q(y1∣y2,y0) q(y2∣y0)q(y1∣y0) = q(y2∣y0) q(y1∣y2,y0)q(y_1 \mid y_0)\;q(y_2 \mid y_1) \;=\; q(y_1 \mid y_0)\cdot\frac{q(y_1 \mid y_2, y_0)\;q(y_2 \mid y_0)}{q(y_1 \mid y_0)} \;=\; q(y_2 \mid y_0)\;q(y_1 \mid y_2, y_0)The factor q(y1∣y0)q(y_1 \mid y_0) cancels, and every remaining factor points backwards, just like the model. For more steps the same cancellation happens all along the chain. The logarithm then turns the products into sums, and what is left is−ELBO=∑n=2NEq[KL(q(yn−1∣yn,y0)⏟the cheat sheet ∥ pθ(yn−1∣yn)⏟the network)] + two end terms-\text{ELBO} = \sum_{n=2}^{N} \mathbb{E}_q\Big[\mathrm{KL}\big(\underbrace{q(y_{n-1} \mid y_n, y_0)}_{\text{the cheat sheet}}\;\big\|\;\underbrace{p_\theta(y_{n-1} \mid y_n)}_{\text{the network}}\big)\Big] \;+\; \text{two end terms}The symbol KL\mathrm{KL} stands for the Kullback-Leibler divergence, which measures how different two distributions are: it is zero when they agree and positive otherwise. Each term compares two answers to the same question, "where was the value one step earlier?". The first answer, q(yn−1∣yn,y0)q(y_{n-1} \mid y_n, y_0), is the ideal reverse step, and I will call it the cheat sheet, because it is the answer of somebody who knows the clean value y0y_0, and we can compute it exactly because we designed the noise ourselves. The second answer is that of the network, which sees only yny_n. Of the two end terms, the last-step term turns out to have the same form as the others, and the term without learnable weights compares the end of the forward chain with pure noise and can be ignored.Training means making the answer of the network agree with the cheat sheet, at every noise level.B. The cheat sheet is a bell curve. Two witnesses speak about the unknown yn−1y_{n-1}. The noisy value yny_n says that it was near yn/αny_n / \sqrt{\alpha_n}, since yny_n is αn yn−1\sqrt{\alpha_n}\;y_{n-1} plus a little noise. The clean value says that it was near αˉn−1 y0\sqrt{\bar\alpha_{n-1}}\;y_0, by the jump formula. Each of these statements is a bell curve over yn−1y_{n-1}, and multiplying two bell curves gives a bell curve again, whose centre is a weighted average of the two centres, with the more reliable witness counting more:q(yn−1∣yn,y0)=N(yn−1∣μ~n, σ~n2),μ~n=an yn⏟where we are now + cn y0⏟where we came fromq(y_{n-1} \mid y_n, y_0) = \mathcal{N}\big(y_{n-1} \mid \tilde\mu_n,\; \tilde\sigma_n^2\big), \qquad \tilde\mu_n = a_n\,\underbrace{y_n}_{\text{where we are now}} \;+\; c_n\,\underbrace{y_0}_{\text{where we came from}}The weights ana_n and cnc_n and the variance σ~n2\tilde\sigma_n^2 are fixed numbers that are determined by the schedule. Their exact values do not matter for what follows, and they are listed at the end of this section. If one witness is much more certain than the other, the combined centre sits close to it, and the combined bell curve is narrower than either, because two pieces of evidence together say more than either alone.C. Equal widths turn the KL divergence into a squared error. We make the step of the network a bell curve with the same width σ~n2\tilde\sigma_n^2 and the same form as the cheat sheet, except that the network has to supply its own guess y^0(yn,n)\hat y_0(y_n, n) of the clean value, which it cannot see:μθ=an yn+cn y^0\mu_\theta = a_n\, y_n + c_n\, \hat y_0For two bell curves of equal width, the KL divergence is simply the squared distance between their centres, divided by twice the variance. The parts anyna_n y_n are identical in both centres and cancel, so thatKL=(μ~n−μθ)22σ~n2=cn22σ~n2⏟fixed weight (y0−y^0)2⏟squared error on the clean value\mathrm{KL} = \frac{(\tilde\mu_n - \mu_\theta)^2}{2\tilde\sigma_n^2} = \underbrace{\frac{c_n^2}{2\tilde\sigma_n^2}}_{\text{fixed weight}}\; \underbrace{\big(y_0 - \hat y_0\big)^2}_{\text{squared error on the clean value}}Readers of Part I will recognise this moment, because it is the central statement of that article in a new place: under a Gaussian with a fixed width, maximum likelihood is a squared error.D. Clean value or noise: the same thing. By the jump formula, knowing yny_n and the noise ε\varepsilon is the same as knowing y0y_0. So if the network predicts the noise instead, with its guess written εθ(yn,n)\varepsilon_\theta(y_n, n), or εθ\varepsilon_\theta for short, its guess of the clean value follows from the same formula, and the two errors differ only by a known factor:y0=yn−1−αˉn εαˉn,y^0=yn−1−αˉn εθαˉn⟹y0−y^0=−1−αˉnαˉn (ε−εθ)y_0 = \frac{y_n - \sqrt{1-\bar\alpha_n}\;\varepsilon}{\sqrt{\bar\alpha_n}},\qquad \hat y_0 = \frac{y_n - \sqrt{1-\bar\alpha_n}\;\varepsilon_\theta}{\sqrt{\bar\alpha_n}} \qquad\Longrightarrow\qquad y_0 - \hat y_0 = -\sqrt{\frac{1-\bar\alpha_n}{\bar\alpha_n}}\;\big(\varepsilon - \varepsilon_\theta\big)Every term of the ELBO is therefore a weight, which depends only on the noise level, multiplied by (ε−εθ)2(\varepsilon - \varepsilon_\theta)^2. Picking the noise level at random then replaces the sum over nn by an average, and the training loss is Lsimple=E y0, n, ε[(ε⏟true noise−εθ(yn,n)⏟network’s guess)2],yn=αˉn y0+1−αˉn ε \boxed{\,L_{\text{simple}} = \mathbb{E}_{\,y_0,\; n,\; \varepsilon}\Big[\big(\underbrace{\textcolor{#2F9E5B}{\varepsilon}}_{\text{true noise}} - \underbrace{\varepsilon_\theta(y_n, n)}_{\text{network\textquoteright s guess}}\big)^2\Big], \qquad y_n = \textcolor{#2B6CB0}{\sqrt{\bar\alpha_n}\;y_0} + \textcolor{#2F9E5B}{\sqrt{1 - \bar\alpha_n}\;\varepsilon}\,}That is plain MSE. (Strictly speaking, once the weights are dropped it is no longer exactly the ELBO but a re-weighted version of it.)Step 4: the training loop, putting it all togetherNow that we have the loss function, training is straightforward, and one training step consists of the following eight actions:1. Sample a clean value y0y_0 from the dataset (in forecasting: a reading together with the context that preceded it).2. Sample a random noise level nn uniformly from {1,…,N}\{1, \dots, N\}.3. Sample a random noise ε∼N(0,1)\varepsilon \sim \mathcal{N}(0, 1).4. Compute the noisy value in one shot with the jump formula: yn=αˉn y0+1−αˉn εy_n = \sqrt{\bar\alpha_n}\;y_0 + \sqrt{1 - \bar\alpha_n}\;\varepsilon.5. Forward pass: feed (yn,n)(y_n, n) into the neural network. The noise level nn is turned into a vector by an embedding, so that the network knows how noisy its current input is. For images the network is typically a large U-Net, whereas for our single number a small multilayer perceptron is enough, and in the forecasting setting of Section 9 the network additionally receives a summary of the context.6. Network output: the predicted noise εθ(yn,n)\varepsilon_\theta(y_n, n).7. Compute the loss: (ε−εθ(yn,n))2\big(\varepsilon - \varepsilon_\theta(y_n, n)\big)^2.8. Backpropagate and update the weights of the network.This is repeated for many mini-batches, and the network gradually learns to denoise at every noise level. Notice that no chain is run during training and that there is no sampling loop, since every training step is one ordinary forward pass followed by one ordinary backward pass. Figure 7 shows this loop at work on the two-bump data.Figure 7: Training with plain MSE. Left: the network's best guess of the clean value as a function of the noisy value, at a high and at a low noise level (green), next to the exact answer (dashed). Right: what the network generates at that moment of training. Before training it knows nothing; after 2,000 steps of the loop above, its guesses have moved onto the dashed curves and its samples have split into the two bumps of the data. Image by author.Animated training run of a small denoising network on two-bump data. Left: the network's best guess of the clean value at a high and at a low noise level, next to the exact conditional mean shown dashed. Right: the samples the network generates at that moment. After 2,000 training steps with a mean squared error loss on the noise, the guesses match the exact curves and the samples form two bumps.···7. Wait, MSE again? Wasn't that the liar?We have arrived at a strange place. This series began by showing that MSE only learns the conditional mean and therefore cannot express uncertainty, and now we are training a model with MSE and claiming that it captures not only the uncertainty but its entire shape. Both statements are true, and reconciling them is the bridge between this article and the first one.7.1 What MSE learns, here as everywhereIf you predict a random quantity YY with a single number mm and you are penalised by E[(Y−m)2]\mathbb{E}[(Y - m)^2], the best you can do is m=E[Y]m = \mathbb{E}[Y], and with inputs the same holds for every input separately, so a network trained with MSE learns the conditional mean. Nothing about this has changed in the meantime, because our denoising network is trained with MSE to predict the noise from the pair (yn,n)(y_n, n), and the best it can possibly learn is the conditional mean of the noise, which through the jump formula is the same as learning the conditional mean of the clean value:y^0(yn,n)=E[ y0∣yn, n ]\hat{y}_0(y_n, n) = \mathbb{E}\big[\,y_0 \mid y_n,\, n\,\big]The dashed curves of Figure 7 are exactly this quantity. The diffusion network is therefore still a conditional-mean predictor, exactly like the MSE model of Part I. What has changed is the question it is asked. The model of Part I was asked one question, namely what the future is on average. The diffusion network is asked a whole family of questions, one for every noise level and every possible noisy value: "if the clean value had been noised this much and now looked like this, what was it on average?" A single mean does not determine a distribution, but the whole family of means across all noise levels does. Strictly, this holds in the limit of infinitely many, infinitely small steps; with a finite number of steps the shape is reproduced only approximately, and the approximation improves as the steps get smaller.7.2 Two cases we can compute by handThe two-bump example lets us see this with formulas. Suppose first that the clean value really is Gaussian,y0∼N(μ,σ2)y_0 \sim \mathcal{N}(\mu, \sigma^2). Then the best guess of the clean value isE[ y0∣yn ]=μ+αˉn σ2αˉn σ2+1−αˉn (yn−αˉn μ)\mathbb{E}[\,y_0 \mid y_n\,] = \mu + \frac{\sqrt{\bar\alpha_n}\;\sigma^2}{\bar\alpha_n\,\sigma^2 + 1 -\bar\alpha_n}\;\big(y_n - \sqrt{\bar\alpha_n}\;\mu\big)which is a straight line in the noisy value yny_n, and the intercept and the slope of that line are completely fixed by μ\mu and σ2\sigma^2, the two numbers that the heads of Part I output. Suppose now that the clean value is +1+1 or −1-1 with equal probability, as for the ball on the hilltop. Given a noisy value yny_n, Bayes' theorem tells us how probable each of the two origins is, and the average of the two origins weighted by those probabilities works out toE[ y0∣yn ]=tanh⁡ ⁣(αˉn yn1−αˉn)\mathbb{E}[\,y_0 \mid y_n\,] = \tanh\!\left(\frac{\sqrt{\bar\alpha_n}\;y_n}{1 - \bar\alpha_n}\right)which is an S-shaped curve, the same kind of curve that built the continuous mixture in Section 3. At high noise levels, where αˉn\bar\alpha_n is close to zero, this curve is flat at zero, so whatever the input, the best guess is the overall average of the data, and this is precisely the Part I answer, the mean that lands on the hilltop where the data never is. As the noise level decreases the curve becomes steeper, and at low noise it is almost a step that sends anything slightly positive to +1+1 and anything slightly negative to −1-1.7.3 Diffusion without a neural networkSince both best guesses are known in closed form, we can do something instructive and run the reverse process without any neural network, simply by plugging the formulas in. Figure 8 does this for two distributions that have the same mean and the same variance.Figure 8: Two reverse processes start from the same noise and use the same sampler, the same schedule and even the same fresh noise at every step. On the left the best guess is the straight line that belongs to Gaussian data, on the right it is the S-curve that belongs to two-bump data; the orange dots are sixty of the samples, sitting on the curve at their current position. The notebook produces a static version of this figure. Image by author.Animated comparison of two reverse diffusion processes that share the same starting noise, sampler and schedule. On the left the denoiser is the straight line that belongs to Gaussian data and produces one bump. On the right it is the S-curve that belongs to two-bump data and produces two bumps. It shows that a Gaussian head is a diffusion head whose denoiser is restricted to a straight line.Both runs aim at distributions with the same mean and (up to a small error that comes from using only fifty steps) the same variance. The only difference is whether the denoiser is a straight line or a curve, and that difference alone decides whether we end up with one bump or with two. This gives us a precise way of stating the connection between the two kinds of head:A Gaussian head is a diffusion head whose denoiser is restricted to be a straight line. Everything a diffusion head can do beyond a Gaussian head lives in the curvature of what it learns.MSE lied in Part I because it was asked a single question. In a diffusion model it is asked one question for every noise level, and together the answers describe the whole distribution.7.4 Do more steps make the model more accurate?Partly, and it is worth separating two sources of error. The first is the size of the steps: each reverse step is modelled as one bell curve, which is exact only for infinitely small steps, so this error shrinks when NN grows, even with a perfect network. The second is the network itself, which has to learn the conditional mean accurately at every noise level; more steps do not improve this, and small errors of the network accumulate along the chain. The first error can be measured in isolation with the no-network setup of the previous section, where the denoiser is exact.···8. Generating a forecast value: running the film backwardsOnce the network is trained, drawing one sample works as follows. We start from pure noise, yN∼N(0,1)y_N \sim \mathcal{N}(0, 1), and then repeat four small operations for n=N,N−1,…,1n = N, N-1, \dots, 1.First, we ask the network for the noise it believes is contained in the current value, ε^=εθ(yn,n)\hat\varepsilon = \varepsilon_\theta(y_n, n). Second, we turn that into a guess of the clean value with the rearranged jump formula (Section 6, Step 1):μ=an yn⏟where we are now + cn y^0⏟best guess of where we came from\mu = a_n\,\underbrace{y_n}_{\text{where we are now}} \;+\; c_n\,\underbrace{\hat{y}_0}_{\text{best guess of where we came from}}Third, we compute the centre of the reverse step with the formula of the cheat sheet, using our guess y^0\hat{y}_0 in place of the unknown clean value, exactly as in Section 6, Step 3:yn−1=μ⏟centre from the network + σ~n z⏟a little fresh noise,z∼N(0,1)y_{n-1} = \underbrace{\textcolor{#2B6CB0}{\mu}}_{\text{centre from the network}} \;+\; \underbrace{\textcolor{#2F9E5B}{\tilde\sigma_n\, z}}_{\text{a little fresh noise}}, \qquad z \sim \mathcal{N}(0, 1)And fourth, we take the step by adding a little fresh noise, yn−1=μ+σ~n zy_{n-1} = \mu + \tilde\sigma_n\, z with z∼N(0,1)z \sim \mathcal{N}(0, 1), except at the very last step, where we want the clean result and add nothing.Notice that the fourth operation is exactly the line mu + sigma * randn with which the Gaussian head of Part I draws a sample. Before attaching this to a forecaster, the notebook trains the smallest possible diffusion model, one that generates a single number with no context at all, on the two-bump data of Section 1. It takes a few seconds, and it is a good moment to convince yourself that the machinery of the last four sections works before anything else is added.···9. The experiment: a signal that really branches9.1 The dataWe need a signal on which the problem of Section 1 actually occurs, and ideally one for which we know the true answer, so that every model can be checked against it. I will use the ball from Section 1, now as a proper simulation that produces a time series. The ball rolls in a landscape with two wells, one at x=−1x = -1 and one at x=+1x = +1, separated by a small hill at x=0x = 0, which I will call the barrier. A force pushes the ball downhill towards the nearest well, and on top of this it receives small random kicks, which you can think of as thermal noise or turbulence. Most kicks only make it jiggle inside its well, but once in a while a lucky series of kicks carries it over the barrier into the other one. Our sensor records only one reading every hundred simulation steps, so a complete jump can happen between two consecutive readings. The equations of the landscape and of the simulation are in the notebook. Figure 9: (a) The landscape with its two wells and the barrier between them. (b) Four recorded series, where the solid part is the context the model is allowed to see and the dashed part is the future it has to forecast. (c) The histogram of all recorded values, which has two bumps and a thin region in between. Image by author.This dataset is a good test for two reasons. The future can really branch, because a ball that sits near the barrier can fall into either well. And we know the true answer, because the physics has no memory beyond the current position, so for any context we can restart the simulator from the last reading as often as we like. This is the "repeat the same situation a thousand times" (Section 1.1) made real, and it gives us samples of the true distribution of futures. As in the previous parts, the model sees a context of T=32T = 32 readings and has to forecast the next H=16H = 16 readings one step at a time.9.2 Plugging the diffusion head into the forecasterThe step that connects everything back to the previous articles is a small one. The forecaster of Parts I and II had an encoder that reads the context and compresses it into a summary vector h\mathbf{h}, and two heads that turn h\mathbf{h} into a mean and a variance. We keep the encoder and replace the heads:···10. ResultsWe train two forecasters on the double well that are identical in every respect except for the head, so that any difference between them is a difference in the shape they are allowed to express. For each of the 1,000 test contexts we then roll out M=100M = 100 trajectories of H=16H = 16 steps from each model, and we do the same with the true simulator.Figure 10: Left: the distribution of the very next reading for a context that ends at the barrier, according to the true physics (grey), the Gaussian head (blue) and the diffusion head (green). Middle and right: 25 sampled rollouts from each model for the same context, with the barrier region shaded. Image by author.Forecasts for a context that ends at the barrier. Left: the true distribution of the next reading has two bumps with a dip between them; the Gaussian head gives one broad bump centred on the dip, and the diffusion head reproduces both bumps. Middle and right: 25 sampled rollouts from the Gaussian head and from the diffusion head, with the barrier region shaded.Many of the Gaussian paths start in the barrier region, hesitate there for a few steps and only then drift to one side, whereas the diffusion paths leave the barrier quickly and commit to a well, as the real ball does. To put numbers on this over all test contexts, I use the Continuous Ranked Probability Score (CRPS), which for forecast samples XX and X′X' and an observed value yy is:CRPS=E ∣X−y∣−12 E ∣X−X′∣ \mathrm{CRPS} = \mathbb{E}\,\lvert X - y \rvert - \tfrac{1}{2}\,\mathbb{E}\,\lvert X - X' \rvert The first term rewards samples that land close to what actually happened, and the second gives credit back for honest spread, so that a model is not punished for being uncertain when the future really is uncertain. Lower is better, and since part of the future is genuinely random, even the true physics does not score zero.Near the barrierCRPS Share of forecasts on the barriertrue physics0.4850.4856.6%6.6 \%Gaussian head0.4940.49412.8%12.8 \%diffusion head0.4930.4935.8%5.8 \%By the CRPS, the two heads are practically tied: 0.4940.494 against 0.4930.493. Yet Figure 10 shows a difference that nobody could miss, and the second column puts a number on it, the barrier share, which is the fraction of forecast values that fall into the region ∣x∣<0.3\lvert x \rvert < 0.3, where the real ball spends less than7%7 \% of its time. The Gaussian head puts twice as many forecasts there, and the diffusion head matches the truth. A score designed to judge entire distributions simply does not see the difference in shape.There is a lesson here that goes beyond diffusion models. A single summary number hides a difference in risk in Part I, and here a far more sophisticated summary number hides a difference in shape, so whenever you can, look at the samples.···11. Costs and limitsNone of this is free, and it would not be in the spirit of this series to pretend otherwise. The costs:Sampling time. One trajectory of sixteen steps costs16×50=80016 \times 50 = 800 small network calls instead of 16. On a CPU, the rollouts of all test contexts took about a second for the Gaussian head and about a minute and a half for the diffusion head. Faster samplers exist, but they are a topic of their own.No formula. With two heads you can read off μ\mu and σ\sigma and compute any probability in closed form. With a diffusion head, every interval or threshold probability has to be estimated from samples, which is fortunately what Part II taught us to do anyway.Training time. Roughly five times longer for the diffusion head, well over a minute against less than twenty seconds.A small bias from few steps. As Section 7 explained, withN=50N = 50 even a perfect network produces bumps that are about a fifth too narrow, which is the price I accepted for a notebook that runs in minutes.And the limits of the experiment:Easy data. The data is one-dimensional and synthetic, and the system has no memory beyond its last reading, which makes the job of the encoder easy and is the reason why we know the true answer. For the same reason the notebook uses a small multilayer perceptron as the encoder instead of the transformer of Parts I and II; the transformer can be swapped back in without touching the heads.One run, no tuning. The results come from one random seed and from small networks without any tuning.Comparisons I left out. I did not train a mixture density network, which would handle this particular problem well as long as somebody tells it that there are two bumps, and I did not forecast the whole horizon in one shot, since both would have distracted from the single comparison this article is about.Where to read moreThis article took one path through diffusion models, the one that leads to forecasting the next value of a signal. The introductions below take other paths, nearly all of them through images, and they complement this one well. I list them by what you might be looking for, with the sections of this article they correspond to.If you want pictures and plain language first.Introduction to Diffusion Models for Machine Learning by Ryan O'Connor (AssemblyAI) explains the forward and the reverse process with clear diagrams and contrasts diffusion with earlier generative models.A Gentle Introduction to Diffusion by Brett Young (Weights & Biases) is a visual walk through noise schedules and noise prediction.Diffusion Models: A Practical Guide (Scale AI) is about using image generators in practice, and is the place to go if your interest is in generating pictures and not in the mechanism.If you want the full theory.What are Diffusion Models? by Lilian Weng is the standard reference. It derives the lower bound of our Section 6 for vectors, in full generality, and goes on to topics I left out, such as faster samplers and guidance.Understanding Diffusion Models: A Unified Perspective by Calvin Luo builds up to diffusion from variational autoencoders and shows that predicting the clean value, predicting the noise and predicting the "score" are three views of the same thing.What none of them does, as far as I know, is what this article is about: denoising a single future value of a time series, given its past, inside an autoregressive rollout, and checking the result against a known truth.···ReferencesJ. Ho, A. Jain, P. Abbeel, Denoising Diffusion Probabilistic Models, NeurIPS 2020.A. Nichol, P. Dhariwal, Improved Denoising Diffusion Probabilistic Models, ICML 2021 (the cosine schedule).C. M. Bishop, Mixture Density Networks, Technical Report, Aston University, 1994.D. P. Kingma, M. Welling, Auto-Encoding Variational Bayes, ICLR 2014 (continuous hidden variables and the ELBO).K. Rasul, C. Seward, I. Schuster, R. Vollgraf, Autoregressive Denoising Diffusion Models for Multivariate Probabilistic Time Series Forecasting (TimeGrad), ICML 2021.T. Li, Y. Tian, H. Li, M. Deng, K. He, Autoregressive Image Generation without Vector Quantization (MAR), NeurIPS 2024.C. M. Bishop, H. Bishop, Deep Learning: Foundations and Concepts, Springer, 2024, Chapters 15, 16 and 20.T. Gneiting, A. E. Raftery, Strictly Proper Scoring Rules, Prediction, and Estimation, Journal of the American Statistical Association, 2007 (CRPS).