UNet++: A Nested U-Net Architecture for Medical Image Segmentation

Summary

• 기존에 제안된 Unet구조보다 (IOU Score기준으로) 좋은 성능을 제시하였다고 제안
  기존의 skip-connection으로만 구성된 Unet구조와 다른 방식으로 Skip-path개념을 사용하여 모델의 capacity를 높임 ,
• Deep supervision 개념을 사용하여 전체 모델의 성능을 향상시켰다고 한다.

정리

기존의 Unet 구조는 위의 그림과 같이 진행된다. encoder로부터 온 skip connection을 사용하여 decoder의 Upsampling matrix에 concatenate을 진행한다. 압축된 latent vector는 segmentation하려는 feature를 추출하게 되고, skip connection의 정보를 추가로 입력받기 때문에, gradient vanishing문제 개선과 학습 성능이 더 좋아지는 장점을 가진다.

Unet++는 기존의 Skip connection으로 이루어진 Unet 구조에서 nested dense convolutional block을 연결한 encoder decoder 구조를 가진다.(이를 논문에서는 skip-pathway라고 부른다) 구조를 자세히 살펴보자

각 layer는 다음과 같은 구조로 연산을 진행한다. 여기서
$i$는 encoder에 따른 down sampling layer index
$j$는 skip pathway로 온 dense block의 convolution layer이다.
쉽게 말해 $i$는 encoder의 depth index , j는 skip-path의 depth를 의미한다.
$\mathcal{H}(\cdot)$는 Convolution 연산을 표현한 것이고, $\mathcal{U}(\cdot)$는 up-sampling layer를 의미한다.

• 식을 하나씩 보면, $j=0$일 경우(Unet에서 encoder부분)은 이전 layer에 Convolution 연산을 진행한 값이다. 그러나 $j>0$의 경우 $j-1$의 값과 upsampling된 값의 Convolution연산 값이다. 이를 그림으로 표한하면 밑의 그림이다.
  
  결국 Unet++에서 제안한 Skip-path는 기존 unet에서 skip connection의 부분을 모두 convolution연산이 진행되는 작은 Unet 구조를 제안한 것이다.

Deep supervision

Deep supervision은 Deeply-supervised nets에서 제안한 방법으로, 중간 layer의 출력을 추가적인 loss function으로 사용하여 전체 모델의 성능을 향상시키는 방법이라고 한다. layer가 깊어질 수록 생기는 gradient vanishing문제를 중간 layer의 loss를 계산하여 반영하므로, gradient vanishing을 해결할 수 있다고 한다.
이를 그림으로 보면 다음과 같다

중간 layer의 결과값을 사용하여 Loss를 계산하는 것을 볼 수 있다.

loss function

Unet++에서는 2개의 loss의 조화평균을 total loss로 사용하였다.
$Y_b$는 label을 $\hat{Y_b}$는 model predict를 의미하며 , total loss는

Cross entropy와

Dice coefficient를 결합하여 사용하였다.

평가

average IOU 성능이 Unet대비 2.8, 3.3point 좋아졌다고 한다.
재미있는 것은 model pruning부분인데,

Deep supervision을 사용한 loss를 각각 전체 loss로 가정하고 실험을 하였을 때 , Pruning된 만큼 성능은 감소하지만, Inference time이 급격하게 감소된다는 것을 보여준다.