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 we chose every ingredient: the shrink, the noise and the schedule. Going the other way is the whole remaining problem, and the only part of this book that needs a network.

The machine we build does not look at a noisy apple and paint a clean one. It looks at a noisy apple and points at the noise, and then we do the subtraction ourselves, with arithmetic you can check by hand. That distinction sounds pedantic here, and everything in the chapter rests on it.

Why we ask for the noise

The forward process 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 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 ask: of the numbers in front of you, which part was the noise? That question has one right answer for every training example we will ever build, and we know the answer, because we drew the noise. That is the trick the whole field runs on.

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 says what comes out: noise, not an image. The subscript $\theta$, theta, stands for every adjustable number inside the network at once, several hundred million for a U-Net of the kind this chapter describes, and billions once the other trained pieces of a checkpoint are counted. If you have met the neuron, $\theta$ is every weight and every bias in the stack gathered into one symbol. The brackets hold its only inputs: $\mathbf{x}_t$, the noisy image, and $t$, the timestep, which is the noise level and 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, so the right answer is a faint noise field. High up it is almost all noise with a ghost of apple in it, so the right answer is nearly the whole picture. Telling the network which rung it is standing on 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, because of 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 way to manufacture training examples. 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$. The input is $(\mathbf{x}_t, t)$ and the correct answer is $\boldsymbol{\epsilon}$, which nobody had to label, because you drew it thirty milliseconds ago and still have it.

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:

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, and the smaller it is, the closer the guess was.

The double bars $\|\cdot\|$ are the length of a list of numbers, and 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 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. What changes the network is gradient descent, which this chapter does not re-derive, 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 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, and the arithmetic of updating any single one of them is what follows.

The training loop

Assembled, the pieces make the thing that runs for weeks on a lot of hardware. Five steps, 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, with $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. That is one multiplication and one addition per pixel, with no simulation and 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 step 1 with a different apple, a different timestep, and a different noise draw, millions of times. What is not in that list is a loop over the noise levels. The network never walks a whole ladder during training. It gets asked about single rungs picked at random, over and over, 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 algebra, and it is ours to do rather than 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, this would give back the true clean image. 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.

That amplification factor is the part to take seriously. The $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 the 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 becomes a large error in the apple. And at those levels the guess is poor for the reason we opened with. A heavily corrupted patch is compatible with many different apples, so the prediction hedges between all of them, and the $\hat{\mathbf{x}}_0$ you get from one leap is a soft, characterless blur.

The training objective is not at fault. The network is doing what we asked. The fix is in how we use it, and Chapter 5 builds it: 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

So far we have said what $\boldsymbol{\epsilon}_\theta$ must do without saying what it looks like inside. The usual answer is the U-Net, named for its shape when you draw it, which does look like the letter U. This section describes its structure and leaves its equations alone, because they are the ordinary layers of any image network.

The job pulls in two directions. 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 is delivered pixel by pixel, at full resolution, and wide-angle understanding and fine-grained output are hard to get from the same layers.

The U resolves that in three parts. Coming down the left arm, the U-Net repeatedly shrinks the image and widens it: fewer pixels, more numbers describing each surviving location. Detail is traded for scope, until at the bottom a single position's numbers are influenced by most of the original picture. That bottom is the bottleneck. Going up the right arm, the process reverses, expanding back toward full resolution and painting the coarse understanding onto a finer and finer grid. 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 detail it threw away on the way down. Without them the output is blurry. 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 it is working on.

See it move

Where you'll meet this

You meet this chapter first as 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$ and nothing else: the U-Net's hundreds of millions of numbers, billions once the encoder and the text encoder packed in beside it are counted, ground into shape by the training loop above. No images, no captions, no index. When people say one checkpoint "knows" fruit better than another, the difference is in those numbers, put there by a great many repetitions of draw-noise, predict-noise, nudge.

You meet it again if you ever watch a model train. The number scrolling past in the log is $L(\theta)$ from this chapter, averaged over a batch. It jumps around from one line to the next, because 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 does.

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 the loss on this page, and it tells you what the noise predictor's output means: static, to be subtracted, not a picture to be looked at.