U-Net

Shrinking tells you what; growing tells you where

Shrink a picture and a single cell ends up standing in for a much wider area. Parts that were far apart end up folded into the same cell, which is exactly what makes a big-picture call — is this a cat or a car — easier. What gets lost is exactly where a boundary falls, down to the cell.

Growing works the opposite way. As the cells get finer, position can be pinned down precisely, but looking at only a narrow patch makes it hard to tell what the whole thing even is. Neither half can answer on its own.

U-Net folds both into a single structure. The front half keeps shrinking to work out what's there; the back half keeps growing to lay that judgment back down onto the original cells.

Sending the pre-shrink picture straight across

However hard the growing half works, it can't recreate fine detail that's already been thrown away. So U-Net saves the picture at each step going down, and hands it straight across to the matching step going up, at the same size. This shortcut is called a skip connection.

The growing half receives two things at once: the judgment about "what" coming up from below, and the original picture about "where things were" arriving from the side. It overlays the two and moves up to the next size.

Without this connection, boundaries blur and edges turn mushy. It also helps the learning signal reach the deep layers properly, so training holds up even when the network is stacked deep.

It comes out the same size it went in

An ordinary image classification model takes in one picture and puts out a single label. U-Net takes in one picture and puts out another picture at the same size — a shape built for jobs that need an answer for every single cell.

Tracing the boundary of a lesion in a medical scan was the original goal — every pixel needs a mark for whether it's part of that region or not, which requires matching sizes. Cutting a subject out of a background, or tracing roads in satellite imagery, are the same shape of problem.

What it does in image generation

These days, the place U-Net shows up most is image generation. Diffusion methods start from a smudge of noise and clear it away bit by bit to build up a picture. At every step, something has to report exactly where and how much noise still needs clearing from the current picture.

That job needs input and output at the same size, an understanding of the overall composition, and pixel-level position all at once — a fit tailor-made for U-Net. So a single finished image runs through a full loop of the U-shape dozens of times. That's one reason generating an image takes longer than generating text.

Is it still used?

Lately, structures that hold the same size all the way through instead of shrinking and growing in a U have also become widespread, driven by the fact that going bigger tends to keep improving performance. Even so, U-Net stays small, fast, and trains well on little data, which keeps it the default choice for tools that run on personal devices or for tracing boundaries.