U-Net Architecture

Most neural nets you have met so far do one job: take a big thing and produce a small thing. A classifier eats a 3×224×224 image and emits a 1000-way probability. A regressor eats an image and emits a single number. One convolution after another, one pooling layer after the next, the spatial grid gets crushed down until only a vector survives. The whole machine is a funnel.

Denoising is weirder. Given a noisy image x_t of shape 3×H×W, the network has to predict the noise ε at every single pixel — an output of shape 3×H×W. Same resolution in, same resolution out. A plain funnel cannot do this. You would have to bolt a second, reversed funnel onto the first — compress, then decompress. An encoder taped to a decoder. Shape-wise, that works.

And then the output looks like soup. Here is why: by the time your encoder has crushed H×W down to H/16×W/16, it has thrown the crisp edges away. The decoder only sees that tiny, compressed summary. It knows roughly where things live; it does not know where the exact pixel boundary is. A cell membrane becomes a smudge. Whiskers become a suggestion. The output is blurry because the information that made it sharp was last seen four downsamples ago, and nothing in the decoder ever gets to look at it again. You cannot re-invent pixel-accurate detail from a feature map the size of a postage stamp.

Ronneberger and co-authors hit this wall in 2015 trying to segment cells in microscopy images. Their fix was the one fix that everything after has basically just reused: don't make the decoder re-invent the detail — hand it across. Staple wires from each encoder level directly to its mirror decoder level, so the high-resolution features skip the bottleneck entirely. They called the thing a U-Net, because when you draw the resolutions on a page, it looks like the letter U. Six years later, the diffusion crowd realized the same shape was the perfect denoiser, bolted on a timestep, and now it is the backbone of Stable Diffusion, DALL·E 2, Imagen, and every image-diffusion model you can name. Two careers, one hourglass.

Before any math, stare at the shape. The name is not a marketing choice — it is literally what the data flow looks like when you draw the feature-map resolutions on a page. A funnel on the left, a mirror funnel on the right, a thin neck in the middle, and wires running straight across.

Two things to notice. First: the input and output have identical spatial dimensions. H and W at the top, H and W at the bottom. That constraint is the whole reason the hourglass exists. Second: every encoder level has a wire running across to its mirror decoder level. Those are the skip connections, and they are the only thing separating a real U-Net from the naive encoder-taped-to-decoder that returns soup.

Here is the reveal. The encoder funnel and the decoder funnel are not doing the same job. The encoder's job is to get abstract — to pack more and more global context into fewer and fewer spatial positions until the neck holds a single view of the whole image. The decoder's job is the opposite: to get sharp, to reinflate that global view into a pixel-exact output. Those two jobs need different information. Abstraction wants low-frequency summaries; sharpness wants high-frequency edges.

The bottleneck has the abstraction. It threw the edges away at each downsample — that was the point. Ask the decoder to reconstruct whiskers from a H/16×W/16 feature map and it is guessing. The skip wires are the fix: the encoder keeps a copy of its high-resolution features before it crushes them, and hands that copy directly across the hourglass to the decoder at the matching resolution. The decoder concatenates the skip onto its upsampled blur, runs a conv to mix them, and the whiskers snap back. The wire does nothing in the forward pass except exist. And that is exactly why it works — no information had to survive the squeeze.

The channel-doubling trick is not accidental. As resolution halves, you have four times fewer spatial positions, so you can afford twice as many channels before compute per block goes up. Roughly constant FLOPs per level of the funnel, more abstract features as you descend toward the neck. It is the same pattern you saw in skip connections and ResNet stages, for the same reason.

Here is the hourglass drawn interactively. Click any level — encoder side or mirror decoder side — to see the exact feature map shape at that point for a 256×256 input. Pay attention to the memory wires: they always connect matching resolutions, never crossing levels. And notice how each decoder block's channel count stays higher than you'd expect right after the upsample — that is because the skip was just concatenated onto the upsampled tensor, and the conv inside the block halves it back down.

Why this shape is perfect for diffusion. Step back and think about what diffusion actually asks of a network. You hand it a noisy image and say: tell me, at every pixel, what noise you added. The output has to be pixel-exact (3×H×W out, same as in). The answer at any given pixel depends on the whole scene — the model can't know what noise lives on a cat's ear unless it knows there is a cat. Global context in, pixel-precise output. Soup not acceptable.

That is the hourglass's entire design brief. The funnel down aggregates global context into the neck, where every position can see everything. The funnel up reinflates that decision back to the original grid. The memory wires keep the output from going soupy by carrying fine detail across the pinch. Input shape preserved, global reasoning available at the neck, crisp edges delivered at the mirror top — every requirement diffusion has, the U-Net already met in 2015 for a completely different problem. The architecture did not need to be invented for diffusion. It needed to be noticed.

One extra wrinkle. The same network has to denoise at every noise level, from faint haze (t=1) to pure Gaussian static (t=1000). If it does not know which level it is looking at, it cannot know how aggressively to clean. So we inject a timestep signal into every block of the hourglass — encoder funnel, neck, decoder funnel, every single one. That is the one thing the diffusion people added on top of the 2015 shape.

The timestep is a single integer. The network needs a smooth, expressive encoding of it that every block can read. The standard recipe is the same one the Transformer used for token positions: map the scalar t into a high-dimensional sinusoidal vector, push it through a tiny MLP, then add it into the activations inside every residual block.

The sinusoidal part is recycled wisdom. Same trick the original Transformer used for token positions, for the same reason: a smooth, periodic, infinitely-distinguishable encoding of a scalar that the network does not have to learn from scratch. Slow sines give you “rough region of t”; fast sines give you “exactly which step”. The downstream MLP picks whichever granularity it needs.

The injection point is the real design choice. Every residual block in the hourglass — top of the funnel down, every level of the encoder, the neck itself, every mirror level of the decoder funnel — adds t_emb as a per-channel bias on its hidden activations. Not once at the input. Every block. That is how the model gets a strong, repeated reminder of “this is noise level t, adjust accordingly” at every scale of processing.

Here is the timestep embedding and its injection, made visible. Drag the t slider and watch two things. First, the sinusoidal vector changes smoothly — slow components evolve slowly, fast components ripple. No hard jumps. Second, the colored bias bars alongside each block of the hourglass change too: that is the per-channel shift being added to the activations inside every residual block, telling that block what noise level to denoise.