U-net leftover U-net leftover

modules

What I already understood

This is how I originally understood U-Net. A CNN shrinks the image and then restores it, the core information gets compressed along the way, and the model learns the rules it needs for segmentation through the feature map in the middle. Convolution, max pooling, padding, receptive field, segmentation: all known territory.

Two pieces were left over: upsampling and skip connections. This post picks up just those.

The heart of U-Net

The heart of U-Net is the side path that hands back the position information compression threw away.

While the encoder compresses the “what”, the skip connection carries the “where” to the decoder uncompressed.

U-Net overview: encoder, skip connections, decoder

The encoder on the left (the contracting path) shrinks the input while it squeezes out core features, and the decoder on the right (the expansive path) upsamples back to the original image size. At each stage, the high-resolution spatial information (“where”) extracted in the encoder crosses over horizontally through a skip connection and gets concatenated with the decoder’s upsampled feature map.

The position information max pooling throws away can’t be recovered by upsampling alone. That’s why the side path exists.

Upsampling

There are two main ways for the decoder to grow the size back.

  • Interpolation-based upsampling: grow the image by a fixed rule like bilinear or nearest. The interpolation itself learns nothing.
  • Transposed convolution: an upsampling layer with learnable weights. The up-convolution in the original U-Net is this kind.

Transposed convolution

Flatten the input and output into vectors and a convolution becomes multiplication by a sparse matrix CC: y=Cxy = Cx. The up-convolution in the early U-Net multiplies by the transpose of that same CC. The connectivity pattern stays and only the direction flips, so if the convolution took 16 values down to 4, C⊤C^\top takes 4 back up to 16.

Transposed convolution as the transpose of the convolution matrix

The entries of C⊤C^\top are learned weights, not a fixed interpolation rule. This is what makes up-convolution the mirror image of convolution, and it’s why U-Net ends up symmetric around the bottleneck, with matching channel counts on both sides (64–128–256–512–1024–512–256–128–64).

Checkerboard artifact

Use transposed convolution for a while, though, and a problem turns up: the checkerboard artifact. The kernel overlap in transposed convolution isn’t uniform. After upsampling, some pixels have been added into twice and some four times, which leaves a checker-patterned imbalance in the scale of the feature map.

So these days interpolation gets picked as the upsampling layer instead. Interpolation just fills the gaps with the average of the surrounding values. After interpolating, you concat with the skip connection and run a 3×3 conv.

Interpolation upsampling, concat with the skip connection, then a 3×3 conv

You’d think the expressiveness would drop. It didn’t. It turned out upsampling itself only grows spatial resolution, and the actual detail comes from the skip connection.

원래 내 이해

원래 나는 U-net을 이렇게 이해하고 있었다. CNN으로 이미지를 줄였다가 다시 복원하는 과정에서 핵심 정보가 압축되고, 가운데의 feature map을 통해 모델이 segmentation에 필요한 규칙을 학습한다. convolution, max pooling, padding, receptive field, segmentation까지는 아는 상태.

남아 있던 조각은 두 개 — upsampling, skip connection. 이 글은 그 leftover만 줍는다.

U-Net의 핵심

U-Net의 핵심은 압축이 아니라, 압축이 버린 위치 정보를 옆길로 되돌려 주는 것이다.

encoder가 “무엇(what)“을 압축하는 동안, skip connection은 “어디(where)“를 압축시키지 않고 그대로 decoder에 전달한다.

U-Net 전체 구조 — encoder, skip connection, decoder

왼쪽의 encoder(contracting path)는 입력 크기를 줄여가며 core feature를 압축·추출하고, 오른쪽의 decoder(expansive path)는 upsampling을 통해 원래의 이미지 크기로 복원한다. 이 과정에서 각 단계마다 encoder에서 추출된 고해상도 공간 위치 정보(“where”)가 skip connection을 통해 encoder에서 decoder 측으로 수평 전달되어, decoder의 upsampling된 feature map과 concat 된다.

max pooling이 버린 위치 정보는 upsampling만으로는 되살릴 수 없다. 그래서 옆길이 필요하다.

Upsampling

decoder가 크기를 다시 키우는 방법은 크게 두 가지다.

  • interpolation 기반 upsampling — bilinear·nearest 같은 고정된 규칙으로 크기를 키운다. interpolation 자체는 학습되지 않는다.
  • transposed convolution — 학습 가능한 upsampling layer. 초기 U-Net의 up-convolution이 이쪽이다.

Transposed convolution

convolution은 입력과 출력을 벡터로 펴면 sparse matrix CC와의 곱 y=Cxy = Cx로 쓸 수 있다. 초기 U-Net에서 쓰이던 up-convolution은 바로 이 CC의 transpose를 곱하는 연산이다 — 연결 패턴은 그대로 두고 방향만 뒤집어서, convolution이 16개 값을 4개로 줄였다면 C⊤C^\top는 4개를 16개로 되돌린다.

Transposed convolution — convolution 행렬의 transpose

C⊤C^\top의 값은 고정된 interpolation 규칙이 아니라 학습되는 weight다. 이 방식 덕에 up-convolution은 convolution과 대칭이 되고, U-Net은 가운데 bottleneck을 기준으로 양쪽의 채널 수가 같아지는 대칭 구조(64–128–256–512–1024–512–256–128–64)를 이룬다.

Checkerboard artifact

다만 transposed convolution을 사용하다 보면 checkerboard artifact 문제가 발견된다. transposed convolution은 kernel overlap이 일정하지 않다. upsampling 이후에 어떤 픽셀은 두 번 더해지고, 어떤 픽셀은 네 번 더해진다. 이로 인해 feature map 스케일에 체크무늬의 불균형이 생긴다.

그래서 지금은 interpolation을 upsampling layer로 선택한다. interpolation은 단순히 사이 갭을 주변 숫자의 평균으로 메운다. interpolation 이후에 skip connection과 concat 하고 3×3 conv를 한다.

Interpolation upsampling — skip connection concat과 3×3 conv

표현력이 떨어질 것 같지만 떨어지지 않았다. 즉 업샘플링 자체는 공간 해상도만 늘리는 연산이고, 실제 세부 정보는 skip connection에서 가져온다는 게 밝혀졌다.

← home← 랜딩으로