Learning to Denoise: The U-Net and the Sampler

The denoiser's job, restated

Let's pick up exactly where Guide 2 left off. We trained a single neural network, written \epsilon_\theta, whose entire job is humble but powerful: hand it a noisy image \mathbf{x}_t together with a number t telling it how noisy that image is, and it predicts the noise hiding inside. Not the clean picture — just the static that was mixed in. The subscript \theta is shorthand for 'all the trainable weights', so \epsilon_\theta literally means 'the noise guess produced by our trained weights'.

But spotting noise is only half a magic trick. The real goal of diffusion is to generate a brand-new picture starting from nothing but pure static \mathbf{x}_T — a screen of random fuzz with no image in it at all. So the question this guide answers is: now that we own a good noise predictor, how do we actually use it to walk that sky of static, step by step, down to a sharp, never-before-seen image?

The answer comes in two parts, and we'll take them in order. Part one is the architecture that makes \epsilon_\theta a good predictor in the first place — a particular network shape called the denoising U-Net. Part two is the recipe that turns a stack of these predictions into a finished image — the reverse denoising process, also called the sampler. Architecture first, then sampling: a great brain, then a great procedure for using it.

Inside the denoising U-Net

Why a special shape at all? Because removing noise needs two skills that pull in opposite directions. To know what the picture should be — a face, a cat, a landscape — the network must see the whole image at once and reason about big, coarse structure. But to put back crisp edges and fine texture, it must work at full resolution, pixel by pixel. The denoising U-Net is an elegant shape that does both. Picture an artist who first steps far back from the canvas to plan the overall composition, then leans right in with a fine brush to restore the detail. The U-Net's two halves are exactly that stepping-back and leaning-in.

The left half is the encoder (the 'stepping back'). It repeatedly shrinks the image — halving the width and height while letting the channel count grow — so that after a few stages a 256×256 picture has become a tiny grid of, say, 16×16 feature vectors. Shrinking forces the network to throw away exact pixel positions and instead summarise what is present and roughly where: this is a cat, the body is here, the head is up-left. This coarse, global understanding is exactly what you need to decide what the de-noised image ought to look like.

The right half is the decoder (the 'leaning in'). It does the reverse, upsampling the tiny summary back up to full 256×256 resolution. But on its own, upsampling from a 16×16 grid can only produce something blurry — the fine detail was thrown away on the way down. The fix is the U-Net's signature trick: skip connections. At every level, the decoder is handed a direct copy of the matching encoder feature map (the horizontal bridges in the figure). So the decoder gets the best of both worlds — the coarse 'what and where' rebuilt from the bottom, plus the sharp, high-resolution edges piped straight across from the encoder. That is why U-Net outputs keep their crisp lines instead of melting into mush.

One more essential input: the timestep t. The same network has to behave very differently at high noise (where it should commit to bold, coarse strokes) versus low noise (where it should make tiny, careful touch-ups). We tell it which regime it's in with a time embedding: we turn the integer t into a vector (using sinusoids, much like positional encodings) and inject that vector into every block of the network. Concretely it scales and shifts the features, so a single set of weights can act like 1000 slightly different denoisers, one for each noise level — without us training 1000 separate networks.

Finally, modern U-Nets sprinkle self-attention layers at the low-resolution stages (where the grid is small enough to afford it). Attention lets distant parts of the image talk to each other and agree — so the two eyes of a face end up matching, or the left and right ends of a horizon line up. Hold on to this detail, because it is the scaffold for Guide 4: that same backbone will host cross-attention, where the network attends not to other image patches but to the words of a text prompt. Build the mental model now — encoder, decoder, skips, time embedding, attention — and text-to-image will slot right in.

The reverse step, one rung at a time

Now the sampler. Generation is a ladder climbed downwards: we start at the top in pure static \mathbf{x}_T and descend one rung at a time, \mathbf{x}_T \to \mathbf{x}_{T-1} \to \dots \to \mathbf{x}_0, getting a little cleaner each step. The reverse denoising process is just the rule for taking one rung down: given the current noisy image \mathbf{x}_t, how do we compute the slightly cleaner \mathbf{x}_{t-1}? This is the DDPM ancestral-sampling update, and here it is in full.

One DDPM reverse step: un-shrink, subtract the predicted noise, then add a controlled splash of fresh randomness.

Let's dissect every piece — and recall from the noise schedule in Guide 2 that \alpha_t is the 'keep fraction' at step t and \bar{\alpha}_t = \alpha_1\alpha_2\cdots\alpha_t is the cumulative keep fraction from the start. (1) \epsilon_\theta(\mathbf{x}_t,t) is our denoising U-Net's guess of the noise inside \mathbf{x}_t. (2) The factor \tfrac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}} is the weight on that guess: it scales the predicted noise to exactly the right size to remove for this single step, not all the noise at once. (3) We subtract that scaled guess from \mathbf{x}_t, stepping toward a cleaner image. (4) The whole bracket is then multiplied by \tfrac{1}{\sqrt{\alpha_t}}, which 'un-shrinks' the image — remember the forward step multiplied by \sqrt{\alpha_t}<1 to shrink the signal, so dividing by it here restores the proper scale. (5) Finally \sigma_t\,\mathbf{z} adds a little fresh randomness, where \mathbf{z} is a brand-new sample of standard Gaussian noise (mean 0, variance 1 in every pixel) and \sigma_t is a small step-size set by the schedule.

Why on earth add noise back when we are trying to remove it? Here's the intuition. The U-Net's noise guess is only an average — at high noise there are many plausible clean images consistent with the same static. If we greedily subtracted all the noise and never added any back, we'd march straight to the single 'safest average' answer, and averages of many faces look like a blurry non-face. Instead we estimate the noise, remove most of it to step toward cleaner, then add back a touch of randomness so we keep exploring the space of real images rather than collapsing to one mushy mean. The \sigma_t\,\mathbf{z} term is that exploratory nudge — and on the very last step (t=1) we set \mathbf{z}=0, because we want the final output crisp, not freshly speckled. And reassuringly, none of this is invented: every coefficient drops out of algebraically inverting the forward equations you already met in Guide 2.

import torch

@torch.no_grad()
def ddpm_sample(model, shape, T, alpha, alpha_bar, sigma):
    # Start from pure Gaussian static x_T
    x = torch.randn(shape)
    # Walk DOWN the ladder: t = T, T-1, ..., 1
    for t in reversed(range(1, T + 1)):
        # Network's guess of the noise inside x at this noise level
        eps = model(x, t)
        # Weight that removes just this step's worth of noise
        coef = (1 - alpha[t]) / (1 - alpha_bar[t]).sqrt()
        # Step toward cleaner, then un-shrink by 1/sqrt(alpha_t)
        mean = (x - coef * eps) / alpha[t].sqrt()
        # Fresh randomness on every step EXCEPT the last (t == 1)
        z = torch.randn(shape) if t > 1 else torch.zeros(shape)
        x = mean + sigma[t] * z
    return x  # x_0: a finished image

The score-based view: climbing toward real images

Here is the deeper idea that ties everything together — and it's worth the master-level badge even though we'll keep it intuitive. Imagine all possible images as points in an enormous space (a million-dimensional space for a megapixel image). Real photos aren't scattered everywhere; they cluster on a thin, structured region — the 'land' where realistic images live. At any point in this space we can ask: which way should I nudge this image to make it more like real data? That direction, defined at every point, is called the score. Picture it as an arrow planted at every location, all the arrows pointing 'uphill' toward where the believable images are.

The key bridge: the noise predictor IS a score estimator, up to a known scale and a sign flip.

Let's read this carefully. On the left, p(\mathbf{x}_t) is the probability density of noisy images at level t — big where realistic (noisy) images live, tiny in implausible regions. The symbol \nabla_{\mathbf{x}} means 'gradient with respect to the image', i.e. the vector that points in the direction of steepest increase. So \nabla_{\mathbf{x}}\log p(\mathbf{x}_t) is exactly the 'uphill arrow' — the direction in image-space that most rapidly increases the (log) likelihood of being real. The equation says this arrow is, up to a known constant, just the negative of the U-Net's noise prediction. The minus sign is the heart of it: noise points away from the clean data, so to climb toward real images you walk in the opposite direction of the predicted noise. The divisor \sqrt{1-\bar{\alpha}_t} is the standard deviation of the total noise added by step t (straight from the Guide 2 forward equation), and it rescales the raw noise guess into a true gradient — at high noise the denominator is large, so the same noise guess corresponds to a gentler slope.

Once you see the score as a compass, the reverse denoising process becomes a single clean idea: gradient ascent on the data density. Each step nudges the image a little way uphill, toward where real images live, and repeats until you arrive on the land of believable pictures. Doing exactly this — repeatedly stepping along the score and adding a dash of noise so you don't get stuck — is a classic recipe called Langevin dynamics, and it's the engine of score-based generative models. The punchline: DDPM (predict-and-subtract noise) and score-based models (climb the log-density) are not rivals — they are two descriptions of the very same process, which is why the field treats them as one.

Faster sampling with DDIM

There's a painful catch with the DDPM sampler from Section 3: it typically uses around 1000 rungs, and every single rung is one full call to the U-Net. That's ~1000 forward passes of a large network per image — far too slow for an interactive tool where you want a picture in a second or two. DDIM sampling is the fix. It reformulates the reverse process to be deterministic (it drops the random \sigma_t\,\mathbf{z} term entirely) and, crucially, lets you skip timesteps — taking maybe 20 to 50 steps instead of 1000 with surprisingly little loss of quality.

The key idea, in words: at each step, instead of only nudging the image one small notch cleaner, DDIM first uses the noise prediction to leap to a guess of the fully clean image \hat{\mathbf{x}}_0, then re-noises that guess down to whatever next timestep we've chosen to visit. Here is that clean-image guess — and notice it's just the Guide 2 forward equation solved for the clean image:

The network's current best guess of the clean image, recovered by stripping the predicted noise out of x_t.

Reading it: the Guide 2 forward rule said a noisy image is \mathbf{x}_t = \sqrt{\bar{\alpha}_t}\,\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon — signal shrunk by \sqrt{\bar{\alpha}_t} plus noise scaled by \sqrt{1-\bar{\alpha}_t}. We don't know the true noise \epsilon, but the U-Net hands us its estimate \epsilon_\theta(\mathbf{x}_t,t). So we subtract that estimated noise from \mathbf{x}_t and divide by \sqrt{\bar{\alpha}_t} to undo the shrink — and out pops \hat{\mathbf{x}}_0, the network's best guess of the finished picture as seen from where we stand right now. Early on this guess is rough and dreamlike; near the end it's sharp. It is the same denoiser doing the work, just rearranged to point at the destination instead of the next rung.

DDIM's deterministic step: take the clean-image guess and re-noise it exactly to the next chosen timestep — no random term.

This update has a beautifully simple shape — it's the very same forward formula again: take the clean-image guess \hat{\mathbf{x}}_0, shrink it by \sqrt{\bar{\alpha}_{t-1}}, and add back noise scaled by \sqrt{1-\bar{\alpha}_{t-1}} — but using the same predicted-noise direction \epsilon_\theta(\mathbf{x}_t,t) rather than fresh random noise. Because there is no random \mathbf{z} term, the path is a smooth, deterministic curve from static to image, and we can land on any chosen timestep t-1, not just the immediate neighbour. That is exactly why we can skip: the destination \hat{\mathbf{x}}_0 is re-estimated every step, so even a big jump (say 1000 → 950 → ... in steps of 50) stays on a coherent trajectory. From 1000 steps down to 20 is a 50× speed-up — the difference between waiting half a minute and getting an image now.