Learning to Undo It

Now we learn to walk back

We can destroy an apple at any noise level in one shot. Now we learn to walk back. In destroying an apple we built a one-way machine: hand it our apple and a timestep, and it hands back that apple part-dissolved into static. It never needed to learn anything, because every ingredient — the shrink, the noise, the schedule — was something we chose. Going the other way is the entire remaining problem, and the only part of this book that requires a network at all.

Here is the shape of what we are about to do. We will not build a machine that looks at a noisy apple and paints a clean one. We will build a machine that looks at a noisy apple and points at the noise — and then we will do the subtraction ourselves, with arithmetic you can check by hand. That distinction sounds pedantic in this paragraph and it is the most important sentence in the chapter. By the end you will have computed a loss, taken a gradient step with real numbers, and pulled a pixel back from $1.4725$ to close to the $0.8$ it started as.

Why we ask for the noise

Start with the honest difficulty. The forward process — adding noise on purpose — throws information away, and throwing information away is not reversible. Look at a heavily speckled patch and ask "which apple was this?" There is no single right answer: a red apple with a bright highlight and a slightly larger green apple in dimmer light could both have landed on that exact patch of static, given different luck with the noise draw. A network trained on a question with many answers learns to output the blurry average of all of them.

So we change the question. We do not ask for the apple. We ask: of the numbers in front of you, which part was the noise? That question has exactly one right answer for every training example we will ever build, and — this is the trick that makes the whole field work — we know that answer, because we are the ones who drew the noise.

The thing that answers it is the noise predictor, written $\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)$ and read "epsilon-theta". Take that symbol apart. The $\boldsymbol{\epsilon}$ — epsilon — is Chapter 2's noise blob, so the letter itself is a promise about what comes out: noise, not an image. The subscript $\theta$ — theta — stands for every adjustable number inside the network at once — for a U-Net of the kind this chapter describes, several hundred million of them, and billions once the other trained pieces of a checkpoint are counted alongside it; if you have met the neuron, $\theta$ is every weight and every bias in the whole stack, gathered into one symbol. And the two things in the brackets are its only two inputs: $\mathbf{x}_t$, the noisy image, and $t$, the timestep — the noise level, the same $t$ that indexes Chapter 2's schedule. Nothing else goes in. The network does not know what we want, does not know it is looking at an apple, and has never been told the word "apple".

Why does the timestep need to go in at all? Because the same speckle means different things at different noise levels. Low down the ladder the picture is almost all apple with a faint dusting on top, and the right answer is a faint noise field; high up it is almost all noise with a ghost of apple in it, and the right answer is nearly the whole picture. Telling the network which rung it is standing on is what lets one network handle all of them.

The answer key we make ourselves

Most machine learning is limited by labeled data: somebody has to say what the right answer is, and people are slow and expensive. Diffusion has no such problem, and the reason is the closed form we derived last chapter:

In words: the noisy image at timestep $t$ is the clean image faded by alpha-bar-t — the cumulative fraction of the original signal still surviving — plus one blob of noise scaled by whatever is left over.

Read that equation as a recipe for manufacturing homework. Take any real apple photograph and call it $\mathbf{x}_0$. Pick a timestep $t$ at random. Draw a noise blob $\boldsymbol{\epsilon}$ — one number per pixel, each pulled independently from the bell curve centered at zero with spread one. Combine them with the two coefficients, and out comes $\mathbf{x}_t$. You now hold a training example whose input is $(\mathbf{x}_t, t)$ and whose correct answer is $\boldsymbol{\epsilon}$ — not because anyone labeled it, but because you drew it thirty milliseconds ago and still have it in your hand.

The loss: one number for how wrong

To train anything we need a single number that says how bad the current guess was, and it has to get smaller as the guess gets better. Here it is:

In words: subtract the network's guessed noise from the true noise, one number at a time, square each of those differences, and add the squares up. That total is the loss — the smaller it is, the closer the guess was.

The double bars $\|\cdot\|$ are the length of a list of numbers; squaring that length means "skip the square root and keep the sum of squares". This is the mean-squared-error idea from measuring wrongness, applied to a noise field instead of a prediction about text. Squaring does two jobs: it makes every disagreement positive, so overshoots and undershoots cannot cancel out, and it punishes one big miss far harder than several small ones.

One step downhill, by hand

A loss on its own changes nothing; it is a thermometer, not a treatment. The treatment is gradient descent, which we are not going to re-derive here because rolling downhill already does it properly and blame flows backward explains how the blame reaches a weight buried deep inside a network. The one-sentence version: work out which way each adjustable number would have to move to make the loss smaller, then move it a little way in that direction.

To see that happen with real numbers, shrink the network until it has exactly one adjustable number in it. Our toy noise predictor is a single multiplication — it takes the noisy value and scales it by one weight $w$:

In words: this pretend network's entire opinion is "the noise is some fixed fraction of what I am looking at", and the fraction $w$ is the only thing it can learn. A real U-Net has hundreds of millions of numbers where this has one; the arithmetic of updating any single one of them is exactly what follows.

The training loop

Now assemble the pieces into the thing that actually runs, for weeks, on a lot of hardware. Five moves, repeated:

1. Take a real apple. Draw one image $\mathbf{x}_0$ from the training set — one of the many varied apples from the big picture, a point sitting inside the data distribution.
2. Pick a noise level. Choose the timestep $t$ uniformly at random from $1$ to $T$, the length of the schedule — illustratively $T = 1000$ in a model of this kind. Uniformly means every rung of the ladder is equally likely, so the network gets equal practice at barely-speckled apples and at near-total static.
3. Draw the noise. Sample $\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ — the tilde reads "drawn from", and the rest reads "the standard bell curve, mean zero, spread one", one independent draw per pixel. Keep it. This is the answer key.
4. Build the noisy image. Compute $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\,\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t}\,\boldsymbol{\epsilon}$ with the two coefficients looked up from the schedule. One multiplication and one addition per pixel — no simulation, no chain of a thousand small corruptions, because Chapter 2's closed form jumps straight there.
5. Score and nudge. Run $\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)$, compute $L(\theta) = \|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)\|^2$ against the answer key from move 3, and take one gradient step on every number in $\theta$ at once.

Then go back to move 1 with a different apple, a different timestep, and a different noise draw. Millions of times. Notice what is not in that list: there is no loop over the noise levels. Each visit to the loop touches exactly one randomly chosen $t$. The network never walks a whole ladder during training; it just gets asked, over and over, about single rungs picked at random, until it is good at all of them.

Subtracting the noise

We now have a network that points at the noise. Turning that into a cleaner apple is pure algebra, and it is ours to do, not the network's — this is the first half of the reverse process, also called the denoising process. Start from the closed form and isolate the clean image. Move the noise term across:

In words: the faded clean image is whatever is left of the noisy image once the scaled noise has been taken away.

Then divide both sides by the fading factor to undo the fade:

In words: subtract the scaled noise, then divide by alpha-bar-t's square root to restore the original brightness. If we knew the true noise exactly, this would give back the true clean image exactly — the closed form is an equation, not an approximation.

We do not know the true noise. We have a guess. So swap $\boldsymbol{\epsilon}$ for the network's output and mark the result with a hat, which is this book's notation for "an estimate rather than the real thing":

In words: x-zero-hat is our best guess at the clean image — take the noisy image, subtract the noise the network thinks is in there, and undo the fade. Every symbol on the right is something we hold: the noisy image, two numbers from the schedule, and one network output.

Here is the part to take seriously. The factor $0.6237$ came from $\bar{\alpha}_2 = 0.72$ — a mild noise level, where most of the apple was still there. Push to a high timestep where $\bar{\alpha}_t$ is small and the picture is nearly all static, and that ratio grows sharply: the divisor $\sqrt{\bar{\alpha}_t}$ shrinks toward zero while the numerator's coefficient grows toward one, so a small error in the noise guess is blown up into a large error in the apple. And at those levels the guess is necessarily poor, for the reason we opened with: a heavily corrupted patch is compatible with many different apples, so the honest prediction hedges between all of them, and the $\hat{\mathbf{x}}_0$ you get from one leap is a soft, characterless blur.

That is not a flaw in the training objective — the network is doing precisely what we asked. It is an argument about how to use it, and the resolution, which Chapter 5 builds, is to never take the leap: move a short way down the ladder, look again at a slightly less noisy picture, and let the estimate sharpen as the ambiguity drains away.

The shape of the network

We have said what $\boldsymbol{\epsilon}_\theta$ must do without once saying what it looks like inside. The usual answer is the U-Net — named for its shape when you draw it, which really does look like the letter U. We will describe it architecturally and leave its equations alone; they are the ordinary machinery of any image network.

The job has a genuine tension in it. To decide whether a smear of static is noise or the edge of an apple, the noise predictor needs to see the whole picture. But its answer has to be delivered pixel by pixel, at full resolution, with fine detail intact — and wide-angle understanding and fine-grained output are hard to get from the same layers.

The U resolves it in three moves. Coming down the left arm, the U-Net repeatedly shrinks the image and widens it: fewer and fewer pixels, more and more numbers describing each surviving location. Detail is traded away for scope, until at the bottom a single position's numbers are influenced by essentially the entire original picture. That bottom is the bottleneck. Going up the right arm, the process reverses: expand back toward full resolution, step by step, painting the coarse understanding back onto a finer and finer grid. And across the middle run the skip connections — direct wires from each level of the left arm to the matching level of the right arm, handing the upward path the sharp, un-shrunk detail it threw away on the way down. Without them the output is a blurry approximation of the right answer; with them it is crisp. The timestep $t$ is fed in too, converted into a small set of numbers and mixed into every block, so every layer knows which rung of the ladder it is working on.

See it move

Where you'll meet this

The most concrete place you meet this chapter is a file. When a workflow graph loads a checkpoint — a multi-gigabyte .safetensors file with a name like a small brand — what is inside is $\theta$: the U-Net's hundreds of millions of numbers — billions, once the encoder and the text encoder packed in beside it are counted — that the training loop above ground into shape, and nothing else. No images, no captions, no index. When people compare two checkpoints and say one "knows" fruit better, the difference is entirely in those numbers, put there by a great many repetitions of draw-noise, predict-noise, nudge.

You meet it a second time if you ever watch a model train. The number scrolling past in the log, the one everybody stares at, is $L(\theta)$ from this chapter, averaged over a batch of examples. It is noisy from one line to the next — of course it is: each line used a different apple at a different randomly chosen timestep, and an easy example at a low timestep scores much better than a hard one at a high timestep. A rising loss for a few lines means nothing. Only the trend over thousands of lines means anything.

And you meet it a third time in the vocabulary. When documentation for a model says it was trained with an "epsilon-prediction objective", that phrase is exactly the loss on this page, and it is telling you what the noise predictor's output means: static, to be subtracted, not a picture to be looked at.