U-Net

Every picture here is one of the sixteen by sixteen pictures from the convolutional networks page, a cross, a square outline, a filled disc or a diagonal bar drawn at a random place and size, brightness and background, with a little noise on every pixel. That page asked which of the four kinds each picture held.

This one asks, for every one of its 256 pixels, whether the pixel belongs to the shape. The answer is a picture of its own, one yes or no per pixel, and every picture in the collection is drawn with that answer attached, which is called its mask. The box at the top of the page shows the first held-out picture of each kind with its mask.

The last line is the trap in scoring this task by the share of pixels right. Only about one pixel in nine belongs to a shape, so a network that never calls anything a shape gets 0.889 of pixels right while finding nothing at all, and two networks that differ in whether they find the shape can differ by a few hundredths on that score. So the page reports the share of pixels right and always sets beside it a second score that a network answering nothing cannot win, which section 3 defines.

The obvious start is the network that already names these pictures, which gets 0.971 of the held-out kinds right. Its first four layers keep the arrangement of the picture, since a convolution answers at every position and a pooling halves the maps without shuffling them. Then the flattening lays the last maps out as a row of 128 numbers, the dense layers mix all of them into sixteen and then four, and from that point on no number stands for any place in the picture.

To see how little the kind alone says about where, I took the kind the classifier calls each held-out picture, averaged the masks of every training picture of that kind, and called a pixel shape wherever that average was above one half. The same was done with each picture’s true kind, to be sure the classifier’s few mistakes were not the problem.

Knowing the kind exactly is worth an overlap of 0.080, since a cross can be drawn anywhere and the average of crosses drawn everywhere is a faint blur in the middle. What the classifier learned to throw away, the position, is the whole of what this task asks for.

The second score counts only pixels that matter. Take the pixels a network calls shape and the pixels that truly are, count the pixels in both, and divide by the pixels in either. A network that calls nothing scores zero, a network that calls everything scores the shape’s share of the picture, and only an exact mask scores one. It is usually called the intersection over union, and this page calls it the overlap and averages it over pictures, so a thin square outline counts as much as a large disc.

overlap  =  pixels called shape and truly shape  ÷  pixels called shape or truly shape

union = actual shape + predicted shape − intersection
      = 28 + 23 − 19 = 32
IoU = 19 / 32 ≈ 0.594

correct pixels = 256 − 9 missed shape − 4 false shape = 243
pixel accuracy = 243 / 256 ≈ 0.949

The third score is the claim this whole architecture is built on. A pixel is on the edge if the pixel above, below, left or right of it has the other label, which counts both the shape’s outermost pixels and the ground’s innermost, and the frame of the picture does not count as a change. Of the 61,440 held-out pixels, 12,433 are on an edge and 49,007 are not, and on the cross above 52 are on an edge, of which the network got 39 right.

The left side of the U is the start of the classifier, kept as it was. A convolution with four filters reads the picture and answers four maps the size of the picture, a pooling halves them, a convolution with eight filters reads those and a second pooling halves again, leaving eight maps of four by four. The two convolutions hold 40 and 296 parameters, the same 336 as on the convolutional networks page, since the count depends on the filters and not on what comes after them.

parameters  =  filters × (channels read × 3 × 3 + 1)

What is different is what happens to the two convolutions’ answers. In the classifier each is read by the pooling above it and then forgotten. Here each is also kept, the four full-size maps and the eight half-size maps, because the way back up is going to need them.

At the bottom of the U one more convolution, eight filters of three by three, reads the eight maps of four by four and answers eight more of the same size. Each of its sixteen positions stands for a block of four by four pixels of the picture, and by the geometry the convolutional networks page worked out, a number here can depend on a patch eighteen pixels wide, wider than the picture itself.

So the middle is the one place where every number can take in a whole shape, and it is also the place with the least idea where, to finer than four pixels, anything is.

To answer at every pixel the maps have to grow back to sixteen by sixteen, and the cheapest way is to copy each value over a two by two block, which doubles both sides. It has nothing to learn. On the way back, each value was copied to four places, so it reaches the loss by four routes and its slope is the total of the four slopes that arrive there.

up[2i + a, 2j + b]  =  x[i, j]      for a and b each 0 or 1

slope at x[i, j]  =  sum over a, b of  slope at up[2i + a, 2j + b]

The paper used a learned upsampling instead, a two by two convolution run backwards that can learn how to spread each value. Repetition is used here because it has no weights to get wrong, and whatever detail it cannot put back has to come from somewhere else, which is the point of the next section.

At each depth on the way up, the decoder takes the maps that came up from below and the maps the encoder kept at the same size, and lays them side by side as one block with more channels. At half size eight repeated maps and the encoder’s eight make a block of sixteen, and at full size eight and four make twelve.

The convolution that follows reads every channel through its three by three windows, so at every position it sees both what the middle concluded and what the encoder saw there at full resolution, and its weights decide how to combine them.

…

The layers on this page are a chain, each reading only what the one before it answered, so the two joins are made outside the layers, by laying one block beside another between two of them, and the backward pass is walked by hand around them.

The joins cost 720 parameters, all of them in the two decoder convolutions, which read more channels. The U with them holds 2,521 parameters, close to the classifier’s 2,468, and the U without them 1,801, so the comparison in Part 4 is not between networks of equal size. It could not be, since the channels a skip adds are what the decoder spends weights on.

An answer pixel now has three routes back to the picture, and they see very different amounts of it. Through the middle it reaches the widest patch, at the coarsest spacing. Across the full-size skip it reaches only a few pixels around itself, at the finest. The half-size skip is in between. The rule for each step is the one from the convolutional networks page, where a window adds its width less one, times the spacing of what it reads, and a repetition halves the spacing without widening anything.

The last layer is a convolution with one filter of one by one, which gives each pixel a score from the four numbers the layer below answered there. The score is squashed into a probability of shape, and a pixel is called shape when that probability is above one half. Training makes small the negative logarithm of the probability each pixel gave its true answer, added over every pixel and divided by the number of pictures.

p  =  1 ÷ (1 + e^(−score))

loss  =  − (1 ÷ pictures) × sum over every pixel of  [ y log p + (1 − y) log (1 − p) ]

pixel loss = −ln(0.5208) ≈ 0.6524

Because the loss adds up 256 pixels for every picture, the slope on a shared kernel weight is a sum over 256 positions, and the step size has to be much smaller than the classifier’s. Both arrangements here are trained with a step of 0.003 against the classifier’s 0.1, for 20 passes over the same 240 training pictures in shuffled batches of 16.

The forward pass is the chain of layers with two extra moves. After each encoder convolution its answer is kept, and just before each decoder convolution the kept answer is laid alongside what came up from below, the repeated maps first and the encoder’s after them. The order matters, because the way back has to split the blame in the same order.

Walking back, each layer works out its own correction from the blame that arrives at its answer and hands down the blame for what it read. At a decoder convolution that blame covers the whole joined block, so it is split by channel in the order the block was built. The part for the repeated maps goes on down through the repetition to the layers below. The part for the encoder’s maps is held until the walk reaches that encoder convolution from above, and there the two are added.

blame at e2  =  blame down through the pooling  +  blame across the skip

A backward pass written by hand can be checked without trusting any of it. Nudge one weight up by a millionth, measure the loss, nudge it down by the same amount, measure again, and the difference divided by two millionths is the slope. I did that for twelve weights of every convolution in the untrained U with skips at seed 0, on the first held-out picture of each kind, and for twelve pixels of those pictures.

The walk agrees with the finite difference to within about two hundred-millionths everywhere, which is what the rounding in a nudge of a millionth allows. The last column is the mistake a hand-written join invites. Leaving out the addition does not touch the middle or the decoder, which are above the joins, and it puts the two encoder convolutions wrong by 6.838 and 4.054 on slopes no larger than 9.505.

The two parts added at each encoder convolution can be measured separately. I took the length of each block of blame, the square root of the sum of its squares, on the same four pictures, before training and after it.

Before training the first convolution hears mostly from the middle, 6.551 against 1.861 across the skip. After twenty passes the two routes carry about the same, 0.573 and 0.587, so the full-size skip has become half of what shapes the first filters. That is the second thing a skip connection does, beside carrying detail forward. It gives the layers nearest the picture a short route from the loss that passes through two convolutions, where the route through the middle passes through five.

To measure what the joins buy, the U is trained with them and without them from the same seed, which draws the same starting weights for the way down and the middle, so the only difference is whether the decoder can read the encoder’s maps. After every pass I scored both on the 240 held-out pictures.

…

At seed 0 the U with skips reaches an overlap of 0.909 after four passes and 0.987 after ten, and finishes at 0.992 with a loss of 0.0023 per pixel. The U without them does not get past 0.64 at any pass, and its curve jumps. At seed 2 it was at 0.556 after pass 10 and at exactly zero after pass 11, calling no pixel of any held-out picture shape, and at seed 1 it fell from 0.556 after pass 18 to 0.040 after pass 19.

A network this small ends somewhere different for each draw of its starting weights, so both arrangements were trained from three seeds, and the widgets draw seed 0 unless a button says otherwise.

With the skips the three seeds agree within 0.004 in overlap. Without them they spread from 0.163 to 0.623, and the worst of them gets 0.906 of pixels right, which is only 0.016 above the 0.889 that calling everything ground gets. Every seed of the U with skips beats every seed of the U without them by more than the spread of either.

The claim for skip connections is that they put the edges back, so I counted the wrong pixels separately on the edges and away from them.

Away from the edges both do nearly as well as each other, and the whole difference sits in the band a pixel wide on either side of each shape’s outline. That is what section 5 predicts, since the middle can place a shape to within about four pixels and no closer, and the full-size skip is what carries the one-pixel detail across.

Split by kind, the U without skips does worst on the square outline, which is one pixel thick and so is all edge, and best on the filled disc, which has an inside.

In the box at the top of the page, choose the square outline and switch the skips off. The probability map draws a soft ring roughly where the square is, too thick and with rounded corners, since a one-pixel line cannot be drawn from blocks of four. With the skips it is the square. The diagonal bar is the U with skips’ weakest kind at 0.977, since a stair of two-pixel steps puts every one of its pixels on an edge.

Every picture in this collection is a dim background with a brighter shape drawn on it, so the obvious way to find the shape is to call every pixel brighter than some level shape. I tried two versions. The first uses one level for every picture, the one of 41 levels from 0.20 to 0.60 that got most training pixels right, which was 0.41.

The second chooses a level for each picture from that picture alone, by Otsu’s rule of 1979, which tries every cut between two brightnesses and keeps the one that splits the picture into two groups whose means are furthest apart for their sizes.

…

The U with skips beats one fixed level, by about 0.01 in overlap and on the edges too, and a fixed level has no way to follow a picture whose background happens to be bright. It loses to a level chosen per picture, which gets one pixel wrong in 61,440. The U has 2,521 parameters, took twenty passes over 240 pictures and their masks to train, and is beaten by a rule with no parameters that has never seen a mask.

The threshold wins because the pictures were drawn so that it would. A shape pixel is its picture’s background plus a lift of at least 0.4, and the noise on each pixel has a spread of 0.05, so within one picture the shape and the ground are two groups of brightness with a gap between them, and splitting the gap is the whole task.

What the U-Net was invented for is the case where that is not true, where a membrane in an electron microscope picture is no brighter than the inside of the cell beside it and the only way to tell them apart is the shape of the surroundings. That case is not in this collection, and building it would need pictures drawn some other way.

So what this collection can show about a U-Net is the architecture itself, the gap between having the skips and not, and where that gap sits. It cannot show the architecture earning its place, since on this task a rule with no parameters already gets 61,439 of the 61,440 held-out pixels right.

The classifier on the convolutional networks page had no way to answer a ring, since it could only share each picture among the four kinds it knew, and it called most rings a square outline. A U-Net is asked only which pixels are shape, so a kind it never saw is a question it can still answer. I drew the same sixty rings that page used and scored both arrangements on them.

With the skips the rings come out almost as well as the kinds it trained on. Its errors are ring pixels called ground, 20, 70 and 11 of the 2,257 at the three seeds, and it calls at most 2 of the 2,046 pixels inside the rings’ holes shape, so it did not fill a ring in as if it were a disc.

Without the skips, at seed 0, it called 245 of those hole pixels shape and missed 908 ring pixels, which in the box at the top of the page looks like a smudge where a ring should be. The rings are the last four pictures there.

The classifier reads sixteen by sixteen and no other size, because its dense layer has one weight per number of the flattened row. The U has no dense layer. Every layer in it reads its input window by window, so the same trained weights can read a bigger picture, and I laid four held-out pictures out two by two as one picture of 32 by 32 and handed it to the U with skips at seed 0 without retraining anything.

…

It gets all 1,024 pixels right, and its call at every pixel is the same as when the four pictures were read apart, including the pixels next to the seams, where a window now reads the neighbouring picture instead of the zeros it saw at the frame during training. On a picture with nothing drawn on it, an even brightness of 0.15, it calls no pixel shape, and the largest probability of shape it gives any pixel is 0.0044, where the classifier called the same blank picture a diagonal bar.

A join lays two blocks of maps side by side, which means nothing unless they are the same height and width. The repeated maps come from halving the picture and doubling it back, and a side that does not halve evenly loses a row on the way down that repetition cannot restore, so the two blocks meeting at the join are of different sizes.

There are three ways out, and each costs something. Pad the picture up to the next side that halves evenly, which invents pixels at the border that the network then has to answer for. Crop the larger block down to the smaller, which is what the paper did throughout, since its unpadded convolutions shrank every map anyway, and which throws away the encoder’s outermost rows.

Or resize the picture, which changes every shape in it. Without the joins the question does not arise at the joins, but it moves to the answer, which for a side of 30 comes out 28 wide and cannot be compared with a mask of 30.

Each of these is a fact about answering every pixel with a shrinking and growing network, and each is a decision any implementation of it faces.