Redesigning Skip Connections to Exploit (UNet++)

Redesigning Skip Connections to Exploit (UNet++)

0. Abstract

FCN 및 U-Net의 경우 2가지의 한계점이 존재합니다.

  • 데이터셋에 맞은 모델의 최적 깊이를 알 수가 없습니다. 그래서 비용을 들여서 이를 찾아내거나 여러 깊이의 모델들을 앙상블하는 비효율적인 작업이 필요합니다.
  • Skip Connection이 동일한 깊이를 가지는 인코더와 디코더만 연결되는 제한적인 구조를 가집니다.

이러한 2가지의 한계점을 극복하기 위해서 UNet++에서는 새로운 형태의 아키텍처를 제시합니다.

image-20210128000424488

  • 인코더를 공유하는 다양한 깊이의 U-Net을 만들어서 deep supervision을 이용해서 함께 학습하고 앙상블하는 형태를 제안합니다.
  • skip connection을 동일한 깊이에서의 특징맵들이 모두 결합하도록해서 유연한 특징맵을 만들어줍니다.
  • Pruning을 통해서 추론 속도를 올리는 방법을 제안합니다.

위의 과정을 통해서 만든 UNet++를 6개의 다른 이미지 데이터셋에 적용하였고 다음의 결과를 얻었습니다.

  • UNet++는 6개의 데이터셋에 대해서 일관성있는 높은 성능을 보입니다.
  • UNet++는 다양한 크기의 객체에 대해서 높은 성능의 세그멘테이션 품질을 보입니다.
  • Mask RCNN에 새로운 Skip Connection을 적용한 Mask RCNN++의 경우 Instance Segmentation에서 높은 성능을 보입니다.
  • Pruning을 적용한 UNet++는 높은 성능을 유지하면서 빠른 추론속도를 보입니다.

1. Introduction

Abstact에서 언급한 것처럼 기존의 encoder-decoder 형식의 모델들의 경우 2가지의 한계점이 존재합니다.

  • 첫째, 데이터 셋마다 최적의 깊이가 다릅니다. 그렇기에 이를 찾아주거나 다양한 깊이의 모델들을 각각 학습한 후에 결합하는 형식의 방법들이 제안되었습니다. 하지만, 이러한 접근은 인코더를 공유하지 않고 각각 돌아간다는 점에서 비효율적입니다. 특히, 이렇게 독립적으로 학습하게되면 multi-task learning의 장점이 없기도 합니다.
  • 둘째, Skip Connections의 디자인이 불필요하게 제한적입니다. 같은 크기의 특징맵을 가지는 경우의 인코더와 디코더가 결합하는 구조는 너무 약합니다.

image-20210128000424488

UNet++에서는 위의 두가지 한계를 극복하기 위해서 위의 그림에서와 같이 밀집된 Connection 형태를 제안합니다. 이렇게 구성할 경우에 몇가지 장점이 있습니다.

  • UNet++는 인코더를 공유하는 형태로 다양한 깊이에서의 학습을 공유합니다. Deep Supervision을 통해서 이미지의 표현(representation)을 공유하면서 같이 학습하고 깊이를 선택할 필요가 없습니다. 성능의 향상 뿐만아니라 Pruning을 통해서 추론 속도를 올릴 수 있다는 장점도 있습니다.
  • UNet++는 인코더와 디코더의 동일한 크기의 특징 맵만 결합되는 제한적인 연결을 사용하지 않습니다. 이를 Dense하게 Connection 함으로써 다양한 특징 맵의 특성들을 디코어의 특징 맵과 결합시켜줍니다.

그 결과 6개의 데이터셋에 대해서 높은 성능을 달성했고 모델의 주요 공헌을 정리하면 아래와 같습니다.

  1. UNet++에서 다양한 깊이의 U-Nets의 내부적인 앙상블을 도입하여 다양한 크기의 객체에 대해 성능 향상을 보였습니다.
  2. 우리는 UNet++에서 Skip Connection을 재설계하여 디코더의 유연한 특징 결합을 가능하게 합니다. 이는 동일한 깊이의 인코더와 디코더를 결합하는 U-Net에 비해 성능적으로 많은 향상을 가져옵니다.
  3. Pruning을 통해 성능은 유지하면서 추론 속도만 향상시키는 UNet++ 방법을 제안합니다.
  4. UNet++에 내장된 다양한 깊이의 U-Nets의 훈련은 U-Net 간의 협업 학습을 통해 개별 U-Net들을 학습하는 것 보다 나은 성능을 보입니다.
  5. UNet++의 경우 다양한 인코더 백본을 가져와서 학습이 가능하고 다양한 의료 영상 데이터에서 높은 성능을 보임으로서 확장성 및 활용 가능성을 증명합니다.

2. Proposed Network Architecture: UNET++

2.A Motivation behind the new architecture

U-Nets의 다양한 깊이에 따라서 모델의 성능이 어떤식으로 결정되고 결합했을때의 결과를 확인하기위해서 3개의 데이터셋에 대해 실험한 결과 두가지의 인사이트를 얻습니다.

image-20210128011758051

  • U-Net의 깊이를 깊게한다고 성능이 반드시 좋은 것은 아닙니다.
  • 모델의 최적 깊이는 데이터셋마다 다릅니다. (EM : L4, Cell : L3, Brain : L3)

위의 문제점을 해결하기 위해서 보통은 독립된 모델들을 학습시킨 후에 결과를 합칩니다. 하지만, UNet++에서는 하나의 네트워크에서 다양한 깊이를 학습해서 결합시키는 방법을 제안합니다. 그리고 Deep Supervision의 경우도 기존의 논문들과는 다르게 X4-j, j가 아닌 X0, j에 둠으로써 U-Net의 구조들을 앙상블한 형식으로 만들었습니다.

image-20210128014939622

이렇게하면 아래의 애니메이션 이미지처럼 노란색으로 칠해진 UNet이 계속 확장해가면서 앙상블을 이루는 것을 볼 수 있습니다. 즉, 4개의 U-Net이 결합된 형태입니다.

ezgif.com-gif-maker (6)

U-Nete 경우 UNet++와 마찬가지로 부분적으로 동일한 인코더를 사용해서 지식을 공유시키려는 이점을 제공합니다. 하지만, 이러한 네트워크는 아래의 두가지 단점을 가지게 됩니다.

image-20210128020619664

  • 디코더 X4-j, j가 분리되어 더 깊은 U-Nets는 앙상블에서 더 낮은 U-Net의 디코더에 신호를 제공하지 않습니다.
  • U-Nete의 디코더들의 경우 불필요하게 같은 크기의 특징 맵만 결합하기에 객체의 크기에 유연하지 못한 단점이 있습니다.

위의 이러한 한계들을 극복하기 위해서 기존의 Skip Connection을 제거하고 이웃한 노드간에 모두 Connection을 하는 UNet+라는 새로운 형태를 시도했습니다.

image-20210128000424488

UNet+는 이웃 노드간의 특징 맵도 결합하여 제한적인 Skip Connection을 완화하지만 아직 개선점이 존재합니다. 이를 극복하기 위해 같은 계층의 모든 노드들간의 조밀한 연결을 사용합니다.

2.B Technical details

image-20210128163559199

  • H : Convolution Operation
  • D : Down Sampling
  • U : Up Sampling
  • [] : Concatenation

image-20210128163815277

x0,1의 경우 x0,0과 x1,0을 Upsampling한 결과를 결합한 형태([x0,0, U(x1,0)])에 Convolution을 적용한 형태입니다. 디코더쪽으로 갈 수록 내부의 입력값들을 늘어나고 마지막 디코어 x0,4의 경우 모든 입력을 받습니다. j = 0인 경우에 대해서는 x1,0의 경우에는 H(D(x0,0))으로 인코더의 바로 윗 노드가 다운 샘플된 값에 Convolution을 수행한 결과가 됩니다.

image-20210128164428439

Deep supervision에서 X0,1, X0,2, X0,3, X0,4 노드에 1x1 Convolution + Sigmoid를 결합해서 출력을 내는 형태입니다. pixel wise cross-entropy와 soft dice-coefficient loss를 통해서 하이브리드 세그멘테이션을 수행하고 식은 위와 같습니다. 하이브리드 로스를 통해서 그래디언트 스무딩과 클래스의 불균형을 다룰 수 있습니다. 위의 수식에서 왼쪽이 크로스 엔트로피 로스이고 오른쪽의 수식이 다이스 로스입니다.

  • N : 하나의 배치내의 픽셀 수
  • nth : 배치내의 픽셀 번호
  • C : 클래스의 개수
  • yn,c : 타겟 레이블
  • pn,c : 예측 레이블

Deep supervision의 경우 pruning이 가능해서 전체 모델을 다 사용해서 평균을 내는 방법과 특정 깊이 이하의 값만을 이용해서 빠르게 추론하는 2가지 방법을 사용할 수 있습니다.

3. Experiments

3.A Datasets

  • Electron Microscopy (EM)
    • 30 images (512 x 512)
    • 2 classes
    • 96 x 96 패치를 적용하고 sliding window를 사용해서 패치의 절반이 겹치게 한 후 겹치는 부분은 aggregate를 적용
  • Cell
    • training 212 / validation 70 / test 72 이미지
    • 2 classes
  • Nuclei
    • training 335 / validation 134 / test 201 이미지
    • 96 x 96 패치를 적용하고 sliding window를 사용해서 32 pixel stride 적용
  • Brain Tumor
    • 256 x 256 의 30명의 환자들에 대한 데이터셋
  • Liver
    • training 100 / validation 15 / test 15 환자 데이터셋
  • Lung Nodule
    • training 510 / validation 100 / test 408 이미지
    • 64 x 64 x 64 crop 적용

3.B Baseline and implementation

  • Early Stop
  • Pixel-wise sensitivity, specificity, F1 and F2 scores
  • NVIDIA TITAN X

image-20210128162903488

4. Results

4.A Semantic Segmentation Results

image-20210128100917728

IoU를 보면 U-Net < wide U-Net < UNet+ < UNet++ 의 경향을 보입니다. 그리고 UNet++와 UNet의 점수를 비교해보면 여섯개의 데이터에 대해 neuronal structure (0.62±0.10, 0.55±0.01), cell (2.30±0.30, 2.12±0.09), nuclei (1.87±0.06, 1.71±0.06), brain tumor (2.00±0.87, 1.86±0.81), liver (2.62±0.09, 2.26±0.02), and lung nodule (5.06±1.42, 3.12±0.88) 의 상승을 보입니다. 특히 Deep Supervision을 이용할 경우 점수 상승의 폭이 꽤 큰 것을 볼 수 있습니다. 특히 EM과 Lung Nodule에서 점수의 상승이 큰데 그 이유는 다양한 크기의 객체들이 나오기 때문입니다.

image-20210128102040655

추가적으로 UNet++의 인코더를 vgg-19, resnet-152, densenet-201에 적용했을때의 결과는 아래와 같습니다.

image-20210128102620179

위의 결과를 보면 모든 인코더에 대해서 U-Net < UNet+ < UNet++의 성능을 일관성있게 보이고 있습니다. 20번의 반복된 실험을 통해서 나온 결과에 t-test를 수행해서 통계적으로 얼마나 유의미한 차이가 있는지 볼 수 있습니다.

4.B Instance Segmentation Results

Mask RCNN의 FPN 부분의 경우 왼쪽의 그림처럼 Conv한 결과와 2x UP한 결과를 결합하는 형태입니다. Mask RCNN++ 의 경우에는 Dense하게 이를 만들어줍니다. (코드로 구현되거나 그림이 없어서 정확하지는 않지만 개인적인 생각으로는 오른쪽 그림과 같이 구성되어있는 것으로 생각됩니다.)

image-20210128161305526

그 결과 아래의 테이블을 보면 Mask R-CNN의 FPN만 수정한 경우 93.28의 IoU가 95.10으로 높게 상승한 것을 볼 수 있습니다.

image-20210128153454860

5. Discussion

image-20210128171825949

  • UNet++가 U-Net 보다 다양한 크기에서 성능이 좋지만, 매우 큰 경우에 대해서는 서로 비슷합니다.

image-20210128171940135

  • Deep Supervision 사용하거나 깊은 디코더 계층으로 갈수록 세그멘테이션 결과가 뚜렷해집니다.

7. Conclusion

7.1 Advantages

image-20210128100720180

  • Dense하게 연결하는 방식은 Densenet에서도 이미 나왔던 개념인데 이를 잘 접목해서 활용한 것 같습니다.
  • wide U-Net을 도입해서 UNet++의 점수 상승 요인이 파라미터가 많아져서가 아닌 것을 보여주는게 인상깊습니다. 이 부분이 있기에 논리의 빈틈이 줄어든 것 같습니다.

7.2 Disadvantages

  • 처음에 읽었을때는 Same Scale의 특징 맵의 한계를 극복했다는게 다른 크기의 특징 맵도 결합해서 작은 객체와 큰 객체도 탐지했다는 의미인 줄 알았습니다. 하지만, 계속 읽어보니 인코더의 같은 계층의 노드만 사용한게 아니라 인코더와 디코더만 연결했다는 것을 같은 깊이의 노드를 전부 연결하고 업샘플링을 통해서 더 깊은 노드들의 정보도 활용한다는 의미 같습니다. 실제로 위와 같은 방법을 한계를 극복하기 위해서 UNet 3+라는 논문에서는 작은 특징맵과 큰 특징맵도 받아서 해결하려는 시도가 있었습니다.

image-20210128094341174

  • TABLE IV의 파라미터의 수를 보면 7.8M에서 9.0M으로 크게 증가하지는 않았지만 메모리 관점에서는 정보를 계속 저장해야하는 문제가 있습니다.

7.3 Appendix

댓글(0)

Designed by JB FACTORY