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 images

Almost 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.

···

1. The problem: one situation, many possible futures

1.1 What a probabilistic forecast actually claims

Before 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 $mu = 0.5$

1.2 A situation that can go two ways

Let’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$, or into the right valley, which I place at $+1$.

If we repeat this situation a thousand times, then in roughly five hundred of the runs the next value $y$ ends up close to $+1$, and in the other five hundred it ends up close to $-1$, with a small scatter of about $0.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$ or near $-1$, and I cannot tell you which, because both are equally likely”.

1.3 What the two-head model reports

Now 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$ and half near$-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, $sigma^2 = 1^2 + 0.1^2 = 1.01$

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 rollout

The 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 % ($frac{1}{2} sigma$

···

2. First step beyond the bell curve: K bell curves

2.1 The idea

If one bell curve cannot describe two bumps, the most natural repair is to use two bell curves, or more generally $K$ 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:

$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 $K$ components has three numbers. The weight $pi_k$

2.2 How a mixture produces a sample, and the hidden variable inside it

The 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 $K$ components at random, where component $k$ is chosen with probability $pi_k$

$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 $k$ 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) = 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 enough

A mixture has three weaknesses, and they are worth stating clearly because the next section addresses precisely these.

• The first is that the number $K$ 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 = 2$

···

3. Second step: infinitely many bell curves

3.1 From a list of scenarios to a continuum

If the trouble with a mixture is that $K$ is a small number we have to choose, the boldest way out is to stop counting. Instead of a hidden index $k$ that takes one of $K$ values, we use a hidden variable $z$ that is a continuous number, and we draw it from the simplest distribution there is, a standard Gaussian. Instead of a list of $K$ means $mu_1, dots, mu_K$

$underbrace{p(y)}_{text{forecast}} ;=;underbrace{intoverbrace{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:

Drawing 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 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 hand

To see how much freedom this gives, take the function $mu(z) = tanh(8z)$

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$

3.3 The price: the likelihood can no longer be computed

So far we have only seen how such a model generates samples once the function $mu_theta$

With a mixture of $K$ components, the corresponding quantity was a sum of $K$ 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 $z$ from $y$; 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 ourselves

Diffusion 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 $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).

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 notation

From 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, $t$ has always been the time index of the signal and $T$ the context length, so I will keep those and use the letter $n$ for the noise level. Most diffusion papers call the noise level $t$ and the number of levels $T$, 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 noise

Goal of this section. We want a recipe that gradually turns a clean value $y_0$

4.1 Two things to know about Gaussians

A Gaussian distribution $mathcal{N}(mu, sigma^2)$

$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 $mathcal{N}(X mid mu, sigma^2)$

4.2 One small step

We start with a clean value $y_0$

$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,

$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 letter$q$ is used for this fixed forward process, and we will use $p_theta$

What is $beta_n$

4.3 Why $sqrt{1 – 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, $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 $X$ has variance $sigma^2$ and we multiply it by a constant $a$, then $mathrm{Var}(aX) = a^2,mathrm{Var}(X)$

• Rule 2: the variances of independent random numbers add up. If $X$ and $Y$ are independent, then $mathrm{Var}(X + Y) = mathrm{Var}(X) + mathrm{Var}(Y)$

Building the formula from the goal. Let’s mix the previous value with fresh noise using two unknown constants $a$ and $b$,

$y_n = a,y_{n-1} + b,varepsilon_n$

and compute the variance of the result. Since $y_{n-1}$

$mathrm{Var}(y_n) = a^2,underbrace{mathrm{Var}(y_{n-1})}_{=,1} + b^2,underbrace{mathrm{Var}(varepsilon_n)}_{=,1} = a^2 + b^2$

Variance preservation therefore requires $a^2 + b^2 = 1$

$a = 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 $a^2 + b^2 = 1$

Verifying that it works. Plugging back in, $mathrm{Var}(y_n) = (sqrt{1 – beta_n})^2 + (sqrt{beta_n})^2 = (1 – beta_n) + beta_n = 1$

What goes wrong without the square roots? Suppose we had naively used $y_n = (1 – beta_n),y_{n-1} + beta_n,varepsilon_n$

4.4 The jump formula: from the clean value to any noise level in one shot

Training will require noisy versions $y_n$

Step 0: the Gaussian addition rule. Before doing any algebra, we need one fact that will do all the heavy lifting. If $A sim mathcal{N}(0, sigma_1^2)$

$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 $alpha_n = 1 – beta_n$

$y_n = sqrt{alpha_n};y_{n-1} + sqrt{1 – alpha_n};varepsilon_n$

Step 2: expanding two steps. Let’s write out two consecutive steps to see the pattern. The first step goes from $y_0$

$y_1 = sqrt{alpha_1};y_0 + sqrt{1 – alpha_1};varepsilon_1$

The second step goes from $y_1$

$y_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_1Big) + sqrt{1 – alpha_2};varepsilon_2$

$y_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$varepsilon_1$

$alpha_2 (1 – alpha_1) + (1 – alpha_2) = alpha_2 – alpha_1 alpha_2 + 1 – alpha_2 = 1 – alpha_1 alpha_2$

The two noise terms therefore collapse into one,

$sqrt{alpha_2 (1 – alpha_1)};varepsilon_1 + sqrt{1 – alpha_2};varepsilon_2 ;=; sqrt{1 – alpha_1 alpha_2};barvarepsilon, qquad barvarepsilon sim mathcal{N}(0, 1)$

where the equality means that both sides have the same distribution, and substituting back gives

$y_2 = sqrt{alpha_1 alpha_2};y_0 + sqrt{1 – alpha_1 alpha_2};barvarepsilon$

Step 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

$baralpha_n = alpha_1 ,alpha_2 cdots alpha_n = prod_{s=1}^{n} alpha_s$

then for any noise level $n$ all the intermediate noise terms collapse into a single one, and we obtain the jump formula:

$boxed{,y_n = underbrace{textcolor{#2B6CB0}{sqrt{baralpha_n};y_0}}_{text{what is left of the clean value}} ;+; underbrace{textcolor{#2F9E5B}{sqrt{1 – baralpha_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(y_n mid y_0) = mathcal{N}big(y_n mid underbrace{textcolor{#2B6CB0}{sqrt{baralpha_n};y_0}}_{text{centre}},; underbrace{textcolor{#2F9E5B}{1 – baralpha_n}}_{text{width}}big)$

You can read it as a recipe, which says that we keep a fraction $sqrt{baralpha_n}$

The payoff. We never have to loop through the noising steps during training. We pick a noise level, draw a single $varepsilon$, and compute $y_n$

4.5 The noise schedule

We still have to choose the numbers $beta_1, dots, beta_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$beta_n$

$beta_n = 1 – alpha_n = 1 – frac{baralpha_n}{baralpha_{n-1}}$

A popular curve of this kind is the cosine schedule,

$baralpha_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)$makes sure that $baralpha_0 = 1$

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 = 50$

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 noise

5.1 One big leap is hard, one small step is easy

Suppose somebody hands you pure noise $y_N$

Now suppose instead that you are handed $y_n$

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 head

We therefore model every reverse step as a Gaussian whose mean is predicted by a neural network:

$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 $N$ times in a row, not once. Each call removes a little noise, taking the noisy value $y_n$

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 $y_0$

$p_theta(y_0) = int p_theta(y_1);mathcal{N}big(y_0 mid mu_theta(y_1, 1),; sigma_1^2big);dy_1$

If you compare this with the continuous mixture (Section 3), you will see that it is the same formula, with the noisy value $y_1$

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$

···

6. Training: what should the network learn?

Goal of this section. We want to train a model $varepsilon_theta(y_n, n)$

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:

$y_n = sqrt{baralpha_n};y_0 + sqrt{1 – baralpha_n};varepsilon$

Naively, one might think that this is all we need, since we can simply rearrange it for $y_0$

$underbrace{y_0}_{text{what we want}} = frac{overbrace{y_n}^{text{what we hold}} – sqrt{1 – baralpha_n};overbrace{textcolor{#2F9E5B}{varepsilon}}^{text{unknown!}}}{sqrt{baralpha_n}}$

The problem is that $varepsilon$ is gone. When we ran the forward process we drew $varepsilon$ at random and mixed it into $y_n$

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 $y_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_theta(y_0)$

The problem. The model never produces $y_0$

$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 $y_0$

$p_theta(y_0) = int p_theta(y_0, y_1, dots, y_N);dy_1 cdots dy_N$

Think 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 = 50$

The solution: score only the paths we make ourselves. We bring in our own forward process, $q(y_{1:N} mid y_0)$

(a) Multiply by 1. We multiply and divide by $q$, which changes nothing:

$p_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 $q$ is an average over samples from $q$:

$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 approx 3.9$

$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, $q$ (how we add noise) and $p_theta$

Step 3: how does the ELBO become MSE?

A. One comparison per noise level. The ratio inside the ELBO is a product of $N$ backward steps (the model) divided by a product of $N$ 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 $y_0$

$q(y_1 mid y_0);q(y_2 mid y_1) ;=; q(y_1 mid y_0)cdotfrac{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(y_1 mid y_0)$

$-text{ELBO} = sum_{n=2}^{N} mathbb{E}_qBig[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 $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(y_{n-1} mid y_n, y_0)$

B. The cheat sheet is a bell curve. Two witnesses speak about the unknown $y_{n-1}$

$q(y_{n-1} mid y_n, y_0) = mathcal{N}big(y_{n-1} mid tildemu_n,; tildesigma_n^2big), qquad tildemu_n = a_n,underbrace{y_n}_{text{where we are now}} ;+; c_n,underbrace{y_0}_{text{where we came from}}$

The weights $a_n$

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 $tildesigma_n^2$

$mu_theta = a_n, y_n + c_n, hat y_0$

For two bell curves of equal width, the KL divergence is simply the squared distance between their centres, divided by twice the variance. The parts $a_n y_n$

$mathrm{KL} = frac{(tildemu_n – mu_theta)^2}{2tildesigma_n^2} = underbrace{frac{c_n^2}{2tildesigma_n^2}}_{text{fixed weight}}; underbrace{big(y_0 – hat y_0big)^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 $y_n$

$y_0 = frac{y_n – sqrt{1-baralpha_n};varepsilon}{sqrt{baralpha_n}},qquad hat y_0 = frac{y_n – sqrt{1-baralpha_n};varepsilon_theta}{sqrt{baralpha_n}} qquadLongrightarrowqquad y_0 – hat y_0 = -sqrt{frac{1-baralpha_n}{baralpha_n}};big(varepsilon – varepsilon_thetabig)$

Every term of the ELBO is therefore a weight, which depends only on the noise level, multiplied by $(varepsilon – varepsilon_theta)^2$

$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{networktextquoteright s guess}}big)^2Big], qquad y_n = textcolor{#2B6CB0}{sqrt{baralpha_n};y_0} + textcolor{#2F9E5B}{sqrt{1 – baralpha_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 together

Now that we have the loss function, training is straightforward, and one training step consists of the following eight actions:

1. Sample a clean value $y_0$

2. Sample a random noise level $n$ uniformly from ${1, dots, N}$

3. Sample a random noise $varepsilon sim mathcal{N}(0, 1)$

4. Compute the noisy value in one shot with the jump formula: $y_n = sqrt{baralpha_n};y_0 + sqrt{1 – baralpha_n};varepsilon$

5. Forward pass: feed $(y_n, n)$

6. Network output: the predicted noise $varepsilon_theta(y_n, n)$

7. Compute the loss: $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.

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 everywhere

If you predict a random quantity $Y$ with a single number $m$ and you are penalised by $mathbb{E}[(Y – m)^2]$

$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 hand

The two-bump example lets us see this with formulas. Suppose first that the clean value really is Gaussian,$y_0 sim mathcal{N}(mu, sigma^2)$

$mathbb{E}[,y_0 mid y_n,] = mu + frac{sqrt{baralpha_n};sigma^2}{baralpha_n,sigma^2 + 1 -baralpha_n};big(y_n – sqrt{baralpha_n};mubig)$

which is a straight line in the noisy value $y_n$

$mathbb{E}[,y_0 mid y_n,] = tanh!left(frac{sqrt{baralpha_n};y_n}{1 – baralpha_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 $baralpha_n$

7.3 Diffusion without a neural network

Since 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.

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 $N$ 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 backwards

Once the network is trained, drawing one sample works as follows. We start from pure noise, $y_N sim mathcal{N}(0, 1)$

First, we ask the network for the noise it believes is contained in the current value, $hatvarepsilon = varepsilon_theta(y_n, n)$

$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 $hat{y}_0$

$y_{n-1} = underbrace{textcolor{#2B6CB0}{mu}}_{text{centre from the network}} ;+; underbrace{textcolor{#2F9E5B}{tildesigma_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, $y_{n-1} = mu + tildesigma_n, z$

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 branches

9.1 The data

We 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 = -1$

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 = 32$

9.2 Plugging the diffusion head into the forecaster

The 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 $mathbf{h}$, and two heads that turn $mathbf{h}$ into a mean and a variance. We keep the encoder and replace the heads:

context x_1 ... x_T --> ENCODER --> h --> HEAD --> a sample of the next value Gaussian head (Parts I and II): mu(h), sigma(h) one draw: mu + sigma * noise Diffusion head (this part) : eps(y_n, n, h) N = 50 small draws in a row

···

10. Results

We 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 = 100$

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 $X$ and $X’$ and an observed value $y$ is:

$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.

By the CRPS, the two heads are practically tied: $0.494$ against $0.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 $lvert x rvert < 0.3$

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 limits

None 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 costs$16 times 50 = 800$

• 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, with$N = 50$

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 more

This 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.

···

References

• J. 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).