U-Net 논문 리뷰

U-Net: Convolutional Networks for Biomedical Image Segmentation

Stable Diffusion에서 알고 가야 하는 요소 중 하나이다. Diffusion도 공부해야 하는데... SDE도 좀 알아야 하고 할 게 많다 ㅠㅠ. 그러나 언제든 하나씩 천천히!

이번 리뷰에서는 Summary보다는 아키텍처(Network 구성, loss function 등)에 대해서 집중적으로 다루는 시간을 가지도록 하겠다. 논문을 보면서 느끼지만 왜 이렇게 했는지를 분석하면서 읽는 것이 많은 도움이 되는 것 같다.

---

Architecture

우선 왜 U자형의 네트워크를 만들었는지에 대해서 간단하게 역사를 알아보는 시간을 가지자. 본 논문은 생의학 이미지 분할 문제를 해결하기 위한 필요에서 나온 아키텍처임을 꼭 생각하자.

---

Problem Setting

생의학 이미지의 경우 레이블된 샘플이 극도로 적어 단순 CNN 기반 구조로는 훈련하기 어려우며, 더군다나 당시의 CNN 아키텍처는 하나의 이미지를 하나의 클래스에 대해 예측하는 테스크에만 최적화되어 있었다. 한편 biomedical image segmentation은 이미지 내의 모든 픽셀을 클래스 단위로 분할해야 한다는 테스크가 있었다.

즉 데이터는 적은데 픽셀 단위로 정밀하게 잘라야 하는 상황이었다고 보면 된다. 세포에 대한 정확한 위치와 경계 정보를 알아야 하는 동시에 픽셀이 세포의 내부에 존재하는지, 또는 배경에 속하는지에 대한 컨텍스트 정보를 알아야 하기 때문에,

• Down Sampling(Pooling, Stride Conv)의 경우 공간 해상도를 낮추어 컨텍스트 정보(객체 모양, 전반적 패턴 등)를 파악할 수 있지만, 세부적인 디테일은 손실된다는 문제점이 있다.
• Up Sampling(논문에서는 Up-Conv, Transposed Conv로 나옴)의 경우 공간 해상도를 키워 픽셀 단위 세부적인 정보를 이해할 수 있도록 모델을 학습하지만, global context가 부족해진다는 문제점이 생긴다.

따라서 본 논문은 U자형 인코더-디코더 모델을 제안함으로써 생의학 이미지 세그멘테이션 과제를 해결하고자 한다.

---

Network

• Contracting Path (왼쪽):
  2×(3×3 valid convolution → ReLU) → 2×2 max pooling (stride=2). 각 pooling 시 feature 채널 수 2배 증가.
• Expansive Path (오른쪽):
  “Up-convolution” (2×2 transposed convolution, stride=2)로 공간 크기 2배 증가, 채널 수 절반.
  Skip 연결 (cropping된 contracting path 피처맵)과 concatenation 후 2×(3×3 conv → ReLU).
• 최종 레이어: 1×1 convolution으로 64채널 → 클래스 수 매핑.

전체적으로 23개의 convolutional layer가 생기게 된다.

 

왜 U자인가? U를 거꾸로 뒤집은 형태는 안되는가?

하나하나 다 질문하면서 답해보자(내가 초짜라서 그렇다). 인코딩(contracting path)의 경우(일단 인코딩만 있고, 그다음에 분류라고 해보자) 해상도를 낮추지만, 문맥 정보를 압축하기 때문에 중요한 정보만을 추출하고자 모델이 학습될 것이다.

 

이후 축소된 정보는 다시 디코딩(expansive path)을 통해서 확장되는데, 이 때는 문맥을 바탕으로 디테일을 복원하는 과정을 거친다. 그렇다면 축소 과정에서도 문맥 정보만 압축하기보다는 디테일한 정보를 최대한 놓치지 않으면서 문맥을 복원하고자 할 것이고, 학습 가능한 Up-convolution에서 디테일한 정보를 복원하고자 학습될 것이다.

 

반면 U를 거꾸로 뒤집은 형태라면, 우선 모델이 원래의 input보다 큰 해상도를 인위적으로 만들기 때문에 interpolation noise만 늘어나도록 설계될 뿐이다. 따라서 downsampling 과정에서 단순히 늘였던 방법을 거꾸로 뒤집는 형태로만 연산이 작용하는, 즉 쓸데없이 확장하고 다시 축소하는 모습이 되는 것이다. 학습 자체가 잘 안될 것을 쉽게 알 수 있을 것이다.

Skip-Connection

Skip connection은 U-Net의 핵심적인 아이디어 중 하나라고 할 수 있다. 다운샘플링을 거치면서 네트워크는 점점 더 넓은 문맥 정보를 학습하지만, 그 과정에서 경계와 같은 세부 디테일이 손실된다. 반대로 업샘플링은 해상도를 높여 디테일을 복원하려 하지만, 문맥 정보가 부족해 경계가 흐릿해지는 문제가 생긴다. 따라서 U-Net은 얕은 층에서 얻은 디테일한 정보와 깊은 층에서 얻은 문맥 정보를 스킵커넥션으로 결합한다.

이 구조 덕분에 모델은 세포의 정확한 위치와 전체적인 구조를 동시에 파악할 수 있으며, 최종적으로 더 정밀한 segmentation 결과를 얻을 수 있다.

---

Loss Function

본 논문은 weighted pixel-wise cross-entropy를 제시한다:

여기서 $w(x)$는 다음과 같다:

$d_1(x)$, $d_2(x)$는 가장 가까운, 두 번째로 가까운 세포 경계까지의 거리이다.

어떻게 모델이 학습되는지 직관적으로 한 번 파악해보자.

• 해당 픽셀이 실제로 세포에 속하는 픽셀인 경우: 우선 해당 위치가 세포임을 예측한다면 $p_{\ell}(x)$의 값이 높아 $\log{p_{\ell}}(x)$의 값이 0에 가까워져 $w(x)$와 거의 무관하게 loss가 작은 모습을 보일 것이다. 즉 세포 안(경계 포함)에 속한다고만 예측하면 패널티 term이 거의 존재하지 않을 것이다.
• 해당 픽셀이 실제로 세포에 속하지 않지만, 세포와 위치가 가까운 경우: 해당 위치가 세포에 속한다고 모델이 잘못 예측하는 경우,  $\log{p_{\ell}}(x)$의 값이 explode하며 더군다나 $w(x)$의 값 또한 높을 것이다. 따라서 모델은 세포 경계면에 대해서 신중하게 판단하는 방식(경계 정보를 정확하게 인식하는 방향)으로 학습될 것이다.

왜 cross-entropy loss를 사용하지 않았을까? 만약 기본 cross-entropy loss를 사용하는 경우, 모델이 배경만 잘 맞추어도 loss가 작아지는 방향으로 적합될 수 있다. 즉 원래의 목적을 무시하고 학습하는 방향을 가지는 것이다.

 

한편 $w_c(x)$는 다음과 같은 식으로 정리된다.

$$
w_c(x) =
\begin{cases}
\frac{1}{\text{freq}(\text{background})}, & x \in \text{background} \\
\frac{1}{\text{freq}(\text{cell})}, & x \in \text{cell}
\end{cases}
$$

만약 $w_c$가 없다면 경계 값에 대해서는 모델이 신중하게 판단하겠지만, 여전히 세포 vs 배경에 대한 전체 성능이 다소 낮아질 수 있는 위험이 있다. 따라서 $w_c$까지 도입하여 클래스 불균형 문제를 해결하는 것이다.

---

Polygon 형태로 잘 예측하는 것을 확인할 수 있다!

 

논문이 가지고 있는 치명적인 문제는 이미지의 경계 쪽에 있는 데이터들이 소실되는 것인데, 사실 지금은 이미 그것을 trivial하게 해결하기 때문에 따로 다루지 않도록 하겠다. 끝!