[논문] Unet++ : A Nested U-Net Architecture for Medical Image Segmentation

https://arxiv.org/pdf/1807.10165.pdf

해당 논문을 공부하고 작성했습니다.

 

Abstract

이 논문에서는 UNet++라는 의료 분야에 image segmentation에 더 효과적인 새로운 구조를 제시합니다. Dense Net에서 사용한 dense skip pathway를 이용해 UNet++ 구조를 구현했습니다. 인코더와 디코더의 feature map의 차이를 줄이기 위해서 skip pathway를 새로 디자인했습니다.

 

1. Introduction

image segmentation에서 encoder, decoder 구조를 가지는 U-Net, Fully Convolution Network(FCN)을 사용하면서 좋은 성능을 보였는데

 

위 그림과 같이 skip connection이 존재해 좋은 성능을 보였습니다. skip connection을 이용하면 하위 decoder에서 output의 세부 정보를 복구하는데 효과적이기 때문입니다.

이러한 구조들은 자연 이미지에 대해서는 충분히 효과적이지만 더 정확한 결과를 요구하는 의학적 이미지에 대해서는 더 높은 정확도를 요구합니다. 그렇기 때문에 이 논문에서는 nested, dense skip connection을 이용한 새로운 구조인 UNet++를 제시했습니다. 이 UNet++는 일반적인 UNet과 wide UNet보다 더 좋은 성능을 보여줍니다.

 

2. Related Work

FCN와 UNet 모두 skip connection이라는 key idea를 가지고 있는데, FCN는 up-sample된 feature map을 encoder에서 얻은 feature map과 합치는 구조이고 Unet은 더하는 방식이 아닌 붙여서 각 up sampling 단계에서 convolution과 비선형성을 추가합니다. 

UNet++는 단순히 plain ckip connection을 nested dense skip pathway로 대채헤 Mask RCNN의 backbone 구조를 쉽게 배포할 수 있습니다.

 

3. Proposed Network Architecture : UNet++

위 그림은 UNet++ 구조를 보여줍니다. 그림에서 볼 수 있듯이 UNet++는 encoder sub-network(back bone)에서 시작해 decoder sub-network로 이어집니다. UNet 구조와 UNet++를 구분할 수 있는 차이는 새로 구현된 skip pathway부분입니다(위 그림에서 초록색과 파란색으로 표시).

 

3.1 Re-designed skip pathways

새로 구현된 skip pathways는 encoder와 decoder의 sub-networks를 연결을 변형했습니다. UNet에서는 encoder에서 나온 feature map들을 decoder에 바로 전달했습니다. 하지만 UNet++는 피라미드 레벨에 따라 convolution의 수가 달라지는 dense convolution block을 따라 전달됩니다. 예를 들어 노드 X0,0와 노드 X1,3 사이에 skip pathway는 각 convolution layer가 동일한 dense block 이전 convolution layer로부터의 출력을 하위 dense block의 대응하는 up sampling 된 출력과 concate하는 연결 레이어에 의해 선행되는 3개의 convolution layer를 가진 dense convolution block으로 구성된다. 기본적으로 dense convolution block은 encoder에서 나온 feature map을 decoder에서 대기하고 있는 feature map과 가깝게 만들어 줍니다.

 

위 식은 skip pathway의 흐름을 보여줍니다. x(i,j)은 노드 X(i, j)의 결과를 의미합니다. 여기서 i는 encoder에서 down sampling layer의 index를 의미하고 j는 skip pathway에 존재하는 dense block의 convolution layer의 index를 의미합니다. 위 식에서 H()라는 함수는 활성화 함수를 의미하고, U()는 up sampling을 의미합니다. 그리고 []는 concatenation layer를 의미합니다. j=0인 단계에서는 encoder의 이전 layer로부터 오직 하나의 input을 받습니다. j=1인 단계에서는 두 개의 input을 받는데, 연속된 level encoder의 sub-network에서 받습니다. j>1인 노드에서는 J+1 input을 받습니다. 여기서 j input은 같은 skip pathway에 있는 이전 j의 결과와 하위 skip pathway로부터 받은 up sampling 된 output을 받습니다.

 

3.2 Deep supervision

UNet++는 deep supervision을 이용해 두 가지 mode로 작동할 수 있습니다. 모든 세분화 분기의 결과가 평균인 정확한 mode, 마지막 분류 map이 하나의 세분화 분기를 선택해 속도를 향상시키는 mode로 두 가지가 있습니다.

 

위 그림은 fast mode에서 다양하고 복잡한 구조에서 세분화 분기를 가져오는 형태를 보여줍니다. nested skip pathway 때문에 UNet++는 여러 개의 level에서 feature map의 전부 복원할 수 있습니다. 네 가지의 의미 level 각각에 오차 함수로 이진 cross entropy와 주사위 계수를 추가했습니다.

 

식으로 표현하면 위와 같이 나타낼 수 있습니다. N은 batch size를 의미하고 위 loss는 binary cross entropy와 dice coefficient의 합입니다.

이렇게 UNet++는 re-designed skip connection과 deep supervision을 통해 같은 task에 대해서 U-Net이나 wide U-Net보다 더 높은 IoU(Intersection over Union) Score를 얻을 수 있습니다.

 

4. Experiments

Dataset: table1에서 보듯 4가지 의학 이미지 데이터셋을 사용해 model을 평가합니다.

Baseline models: U-Net과 wide U-Net 구조를 이용해 비교했는데, 이미지 세분화에서 주로 U-Net을 사용하기 때문에 이 두 구조를 이용했습니다. wide U-Net을 이 논문에서 제시한 구조와 비슷한 parameter를 갖도록 modeling 했습니다. 자세한 내용은 위 table 2를 통해 볼 수 있습니다.

Implementation details: dice coefficient와 Intersection over Union(IoU)를 이용해 측정했고 validation set에 early stop 메커니즘을 이용했습니다. 그리고 최적화 함수로 Adam을 이용했고 learning rate로 3e-4를 사용헀습니다. UNet++의 기본적인 구조는 위 UNet과 동일하게 만들고 skip pathway에 존재하는 모든 convolution layer에 3x3 kernel을 사용했습니다. deep supervision을 사용하기 위해 1x1 convolution layer를 활성화 함수인 sigmoid 뒤에 적용했습니다. 결과로 UNet++는 input 이미지로 주어지는 네 가지로 분류를 했는데 마지막 분류 map에서 평균적으로 가장 높은 값을 갖는 결과로 분류했습니다.

 

Result: 위 표를 보면 UNet과 wide UNet 그리고 UNet++를 비교했습니다. deep supervision을 사욯하지 않은 UNet++는 UNet과 wide UNet보다 IoU 2.8~3.3 만큼 향상된 결과를 보입니다. 특히 deep supervision을 사용한 UNet++는 deep supervision을 사용하지 않은 UNet++보다 평균적으로 0.6만큼 향상된 모습을 보입니다. 특히 간과 폐의 혹 세분화에 대해서 더 좋은 결과를 보입니다만, 세포 핵과 결장 용종에 대해서는 다른 결과를 보입니다. 용종과 간은 다양한 크기의 video frame을 나타내어 이러한 결과를 얻게 되었습니다.

 

5. Conclusion

의학 이미지 세분화에 대해 더 좋은 정확도를 얻기 위해 UNet++을 제시했습니다. re-designed skip pathway와 deep supervision을 이용해 구현되었습니다. re-designed skip pathway는 encoder의 feature map과 decoder의 feature map의 내용의 차이를 줄이기 위해 사용되었고 결과적으로 더 쉽게 최적화할 수 있게 되었습니다. 또한 Deep supervision을 통해 특히 대장 내시경 비디오에서 용종과 같이 여러 크기를 나타내는 병에 대해 정확한 세분화가 가능합니다.