[Paper Review] UNet++와 U-Net3+ 논문 리뷰

안녕하세요. 오늘은 segmentation에 많이 사용되는 architecture인 U-Net의 발전 architecture인 Nested U-Net과 U-Net3+ 논문에 대해 리뷰해보도록 하겠습니다.

참고 paper

• Nested U-Net : https://arxiv.org/abs/1807.10165
• U-Net3+ : https://arxiv.org/abs/2004.08790
  
  
  

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

해당 논문은 2018년에 나온 논문입니다.

[Abstract]

UNet++는 의료 영상에 대해 더 강력한 segmentation을 위한architecture  입니다. UNet++가 U-Net과 다른 점은 re-designed skip pathway를 설계하였다는 것입니다. 이는 encoder와 decoder feature map 사이에 semantic gap을 줄여 optimizer가 더 쉽고 빠르게 학습이 된다는 것 입니다. 이 논문에서는 다양한 의료 영상에 대해서 segmentation을 진행하였고, U-Net이나 wide U-Net보다 UNet++가 더 낫다는 것을 증명하였습니다.

[Introduction]

Segmentation에 사용되는 U-Net이나 FCN(Fully Connected Network)이 공통적으로 가지고 있는 것은 skip connection입니다. 이는 complex background에서 fine-grained detail을 얻기도 하고, 가려져 있는 물체를 segmentation하는 데도 탁월합니다.

그러나 이 논문에서 더 정교하게 segmentation을 할 수 있는 UNet++을 연구한 이유는 의료 영상의 segmentation 때문 입니다. 의료 영상에서는 작은 오류도 크게 작용 할 수 있고, 부정확한 segmentation는 추후의 진단에도 큰 변화를 가져올 수 있기 때문입니다. 그래서 더 fine detail을 얻을 수 있는 segmentation을 할 수 있는 architecture에 대해 연구하고자 하였습니다.

UNet++의 기저에 있는 가설은 encoder로부터의 high-resolution feature map이 점진적으로 semantically rich한 다른 feature map과 합쳐졌을 때 fine-grained detail을 더 잘 capture 할 수 있다는 것입니다. 그 이유는 점진적으로 feature map과 합쳐 졌을 때, encoder와 decoder 사이의 semantic gap은 줄어들 것이고, 그래서 더 학습하기 쉬운 task가 되기 때문입니다. 이는 plain skip connection을 사용해서 encoder와 decoder 사이의 semantically dissimilar한 U-Net과는 다른 것입니다. 다음 내용을 그림과 함께 더 자세히 설명하도록 하겠습니다.

[Proposed Network Architecture : UNet++]

그림1. UNet++ model 그림

UNet++의 특징 1. Re-Designed skip pathway

U-Net과 다른 UNet++의 특징은 skip connection이 dense convolution block으로 이루어져 있다는 것 입니다. Encoder와 Decoder 사이에는 여러 개의 dense convolution block이 존재합니다. 각각의 convolution layer는 같은 level에 있는 다른 block들과 lower level에 있는 block을 up-sampling 한 것을 합쳐서 얻습니다. 이를 수식으로 설명하면 다음과 같습니다.

식 1. UNet++의 Xi_j block 식

식1에서 H란 convolution을 뜻하고, U란 up-sampling을 뜻합니다. 예를 들어 설명하자면, X1_3 decoder의 경우에는 4개의 block의 fuse가 일어납니다. 같은 level에 있는 X1_0, X1_1, X1_2 block과 lower level에 있는 X2_2 block을 up-sampling 한 것을 모두 합쳐 convolution 하여서 얻습니다. 이는 그림1의 그림 (a)를 보면 쉽게 이해할 수 있습니다. 이런 dense convolution block을 이용하여 encoder와 decoder 사이의 semantic similarity를 높여줍니다. 이를 통해 이 논문의 가설에 따라 더  optimizer가 더 잘 학습할 수 있고, fine-grained detail을 얻을 수 있습니다.

UNet++의 특징 2. Deep Supervision

다음은 deep supervision에 관한 설명입니다. 이는 위의 그림1의 (c)에서 확인할 수 있습니다. UNet++는 dense convolution block 덕분에 multiple semantic level의 feature map을 얻을 수 있습니다.(X_0,1, X_0,2, X_0,3, X_0,4) 그래서 (1)accurate mode와 (2)fast mode로 작동할 수 있습니다.  Accurate mode는 모든 branch로부터 결과값의 평균값을 얻는 것입니다. Fast mode는 segmentation branch로부터 하나의 segmentation branch를 선택해서 model pruning과 speed gain을 얻도록 합니다.

식2. deep supervision 때 사용하는 Loss

deep supervision 때는 위의 Loss를 사용합니다. N은 batch size를 뜻하고, 위의 loss는 binary cross-entropy와 dice coefficient의 합입니다.

이렇게 UNet++는 re-designed skip connection과 deep supervision을 통해 같은 task에 대해서 U-Net이나 wide U-Net보다 더 높은 IoU Score를 도출해 낼 수 있었습니다.

UNet3+ : A FULL-SCALE CONNECTED UNET FOR MEDICAL IMAGE SEGMENTATION

해당 논문은 2020년에 나온 논문입니다.

[Abstract]

최근에 encoder-decoder architecture인 UNet이 의료 영상 segmentation에 널리 사용되고 있습니다. 이 때, multi-scale feature를 합치는 것이 중요한 요소로 생각되고 있고, UNet++는 dense convolution을 통해 이를 이뤘습니다. 하지만 이는 full scale로부터 충분한 정보를 얻기는 어려워서 full-scale의 skip connection과 Deep supervision을 가진 UNet3+에 대해 이 논문은 설명하고 있습니다. 여기에서 full-scale이란 low-level의 detail을 다른 scale에 있는 feature map으로부터 high-level semantic과 합치는 것을 뜻합니다. 이 방법은 특히 다양한 scale에서 나타나는 organ segmentation에서 유리합니다. 뿐만 아니라 UNet3+는 parameter의 수도 줄일 수 있습니다

[Introduction]

많은 segmentation 논문에서 보듯, 다른 scale에 있는 feature map들은 다른 특성, 정보를 탐색하고 있습니다. Low-level feature map들은 organ boundary와 같은 coarse feature를, high-level feature map은 organ 위치와 같은 fine feature의 정보를 가지고 있습니다. UNet++의 경우에는 dense convolution을 통해 low-level feature map과 high-level feature map을 합치고 있습니다. 그러나 이는 down-sampling과 up-sampling을 여러번 하다 보니깐 정보가 희석(dilute) 될 수 있다는 단점을 가지고 있습니다.

그래서 이 논문은 UNet3+라는 새로운 architecture를 제시하고 있습니다. 이는 encoder와 decoder 사이의 connection과 decoder 내부에서의 connection을 형성하고 있습니다. 이는 fine-grained detail과 coarse-grained semantic을 합쳐서 full-scale의 feature map을 형성할 수 있도록 합니다. 이런 UNet3+는 다양한 scale의 feature map을 인식 할 수 있을 뿐만 아니라, network parameter의 수도 줄여주고 있습니다. 그리고 false positive를 줄이기 위한 방법도 연구하였습니다.

즉, (1)더 적은 parameter로 full scale의 feature map을 형성하고 있고, (2)full scale의 feature map으로부터 deep supervision을 형성하고 있고, (3)false positive segmentation을 줄이기 위한 classification-guided module(4)마지막으로 이에 대한 실험과 결과 제시 에 대해서 이 논문은 제시하고 있습니다.

[Methods]

UNet3+는 full-scale skip connection과 full-scale deep supervision을 통해 multi-scale의 feature를 합치고 있습니다.

그림2. UNet3+ model 그림

UNet3+의 특징 1. Full-scale Skip Connections

UNet과 UNet++는 full-scale의 feature map을 사용하지 않아 충분한 정보를 얻지 못하고 있습니다. 이런 결점을 극복하기 위해 UNet3+는 encoder로부터 더 작거나 같은 scale의 feature map을 decoder로부터 더 큰 Scale의 Feature map을 합치고 있습니다.

그것을 식과 그림으로 표현하면 다음과 같습니다. (식3, 그림3)

식3. UNet3+의 Xi_De block 식

위의 식3에서 H는 convolution 후, batch normalization과 ReLU activation을 적용합니다. C는 convolution을 뜻하고, D와 U는 각각 down-sampling과 up-sampling을 뜻합니다.

그림3. UNet3+에서 X3_De 그림

X3_De를 살펴보자면 smaller-scale의encoder인 X1_En와 X2_En에 non-overlapping max pooling을 적용합니다. 그리고 high-scale의 decoder인 X4_De, X5_De를 bilinear interpolation을 적용합니다. 그리고 같은 scale의 encoder인 X3_En까지 5개의 layer를 concatenation해서 3*3 filter size, channel 320(각 layer는 64개의 channel로 구성)개로 만들어서 convolution을 적용 합니다.

이런 UNet3+는 concatenation 할 때 64로 channel수를 Reduction 하기 때문에 dense convolution을 사용하는 UNet++보다 현저하게 더 적은 수의 parameter를 사용하게 됩니다.

UNet3+의 특징 2.Full-scale Deep Supervision

UNet3+는 full-scale의 deep supervision을 제시하고 있습니다. Full-resolution Feature map을 deep supervision으로 사용하는 UNet++과 달리 UNet3+는 각각의 decoder stage로부터의 결과를 deep supervision으로 사용합니다. 이 때, 각 decoder는 concatenation 후에 3*3 convolution, bilinear up-sampling, sigmoid를 적용하고 있습니다.

이 때, UNet3+는 흐릿한 boundary에 더 큰 weight를 주기 위해서 multi-scale structural similarity index loss (MS-SSIM)를 사용하고 있습니다. 그래서 focal loss, MS-SSIM Loss, IoU Loss를 합친 hybrid loss를 사용하여 흐릿한 boundary 문제를 해결한 결과를 얻습니다.(식4)

Loss_seg = Loss_Focal + Loss_MS-SSIM + Loss_IOU

식4. UNet3+ Loss

UNet3+의 특징 3. Classification-guided Module (CGM)

많은 의료 영상 segmentation에서 나타나는 문제는 organ이 아닌 부분도 Segmentation false-positive가 많이 나타나는 것입니다. 이를 해결하기 위해 extra classification task를 합칩니다. 이 때, 기관이 있는지 없는지를 input image에서 예측하도록 설계 한 것입니다.

그림4. CGM

위의 그림4에서 볼 수 있듯이 X5_En에 dropout, convolution, max-pooling, sigmoid를 적용해서 2d tensor를 도출 합니다. 이는 classification probability로 organ이 있는지 없는지를 나타냅니다. 그래서 organ이 있으면 1, 없으면 0으로 만들어서 이를 deep supervision의 각각의 값에 곱해서 organ이 없다면 값이 통과하지 못하도록 만듭니다. 이를 통해 organ이 없는 부분에 대해서 over-segmentation이 되는 것을 방지합니다.

요약하자면 UNet3+는 full scale로 이루어진 skip connection을 가진 모델입니다. low-scale의 feature와 high-scale의 feature를 합쳐서 적은 parameter로 더 정확한 segmentation 결과를 얻었습니다.

UNet, UNet++, UNet3+ 성능 비교

세 모델을 의료 영상에 적용했을 때 결과를 비교해보겠습니다. 이는 UNet3+ 논문에서 결과를 발췌하였습니다.

1. Experiments

이 논문에서는 liver와 spleen의 CT data로 실험을 하였습니다. Liver data는 training에 103장, testing에 28장을 사용하였습니다. Spleen data는 training에 40장, testing에 9장을 사용하였습니다. Training 후에 dice coefficient를 이용하여 결과를 비교하였습니다.

2. Results

표1. Liver와 Spleen data에 대한 각 모델 적용 결과

각각의 model에 대한 dice score를 비교하면 위의 표와 같습니다. Resnet-101을 사용한 UNet, UNet++, UNet3+ without deep supervision에 대해서 비교해보자면 liver와 spleen 모두 UNet3+, UNet++, UNet 순으로 Dice Score가 높은 것을 알 수 있습니다. 그리고 모든 dice score를 average 했을 때에도 UNet3+, UNet++, UNet 순으로 성능이 좋은 것을 볼 수 있습니다. 이 때, varying scale의 CT slice를 사용하여 deep supervision을 포함한 UNet3+을 적용했을 때, dice score를 더 높일 수 있습니다,.