[논문 리뷰] SAUNet: Shape Attentive U-Net for Interpretable Medical Image Segmentation

Jesse Sun, Fatemeh Darbehani외 2명|arXiv (Cornell University)|2020. 01. 21.

Advanced Neural Network Applications참고 문헌 39인용 수 22

한 줄 요약

SAUNet는 전통적인 텍스처 기반 특징 외에 병렬으로 구조적 특징을 학습하기 위해 형태 인식 스트림을 통합한 형태 인식 U-Net 아키텍처를 제안한다. 이는 의료 영상 분할의 해석 가능성과 내구성을 향상시킨다. 모델은 SUN09 및 AC17 심장 MRI 데이터셋에서 최신 기술 수준의 성능을 달성하면서 추론 중에 다중 해상도 열화상 맵을 생성하여 후행적 기반의 기울기 기반 해석 방법이 필요 없도록 한다.

ABSTRACT

Medical image segmentation is a difficult but important task for many clinical operations such as cardiac bi-ventricular volume estimation. More recently, there has been a shift to utilizing deep learning and fully convolutional neural networks (CNNs) to perform image segmentation that has yielded state-of-the-art results in many public benchmark datasets. Despite the progress of deep learning in medical image segmentation, standard CNNs are still not fully adopted in clinical settings as they lack robustness and interpretability. Shapes are generally more meaningful features than solely textures of images, which are features regular CNNs learn, causing a lack of robustness. Likewise, previous works surrounding model interpretability have been focused on post hoc gradient-based saliency methods. However, gradient-based saliency methods typically require additional computations post hoc and have been shown to be unreliable for interpretability. Thus, we present a new architecture called Shape Attentive U-Net (SAUNet) which focuses on model interpretability and robustness. The proposed architecture attempts to address these limitations by the use of a secondary shape stream that captures rich shape-dependent information in parallel with the regular texture stream. Furthermore, we suggest multi-resolution saliency maps can be learned using our dual-attention decoder module which allows for multi-level interpretability and mitigates the need for additional computations post hoc. Our method also achieves state-of-the-art results on the two large public cardiac MRI image segmentation datasets of SUN09 and AC17.

연구 동기 및 목표

• 딥 러닝 기반 의료 영상 분할 모델의 해석 가능성과 내구성 부족 문제를 해결하기 위해.
• 텍스처 기반 특징보다 더 내구성이 높은 형태 특징을 학습함으로써 모델의 일반화 능력을 향상시키기 위해.
• SmoothGrad와 같은 계산 비용이 큰 후행적 열화상 방법에 의존하지 않고 모델 내부에 다중 해상도 해석 가능성을 직접 통합함으로써 이를 제거하기 위해.
• 대규모 공개 심장 MRI 분할 벤치마크(SUN09 및 AC17)에서 최신 기술 수준의 성능을 달성하기 위해.
• 진단 적용 분야에서 투명성과 신뢰성을 지원하는 임상에서 활용 가능한 분할 프레임워크를 개발하기 위해.

제안 방법

• 주 텍스처 스트림과 보조 게이트형 형태 스트림을 갖춘 이중 스트림 인코더-디코더 아키텍처를 도입하여 형태에 의존하는 특징을 병렬로 학습한다.
• 추론 중에 다중 해상도 공간 및 형태 주의 맵을 생성하는 이중 주의 디코더 모듈을 사용하여 내장형 해석 가능성을 실현한다.
• 분할 정확도와 형태 특징 학습을 동시에 최적화하는 이중 작업 손실 목표를 적용하여 내구성을 향상시킨다.
• 스케일 간 특징 융합을 가능하게 하면서도 공간적 맥락을 유지하기 위해 인코더 및 디코더 경로 간 스킵 연결을 적용한다.
• SmoothGrad와 같은 방법이 요구하는 추가 정방향-역방향 전파를 피하기 위해, 디코더의 다양한 수준에서 주의 맵을 통해 열화상 맵을 생성한다.
• SUN09 및 AC17 데이터셋에서 학습하기 위해 RAdam 옵timizer를 사용하고 학습률은 5e-4, 배치 크기는 4로 설정한다.

실험 결과

연구 질문

• RQ1 형태 특징을 명시적으로 학습하는 이중 스트림 U-Net 아키텍처가 표준 CNN보다 분할 내구성을 향상시킬 수 있는가?
• RQ2 형태 스트림을 통합함으로써 SUN09 훈련 데이터에서 다른 테스트 세트(AC17)로 일반화할 때 성능 저하가 감소하는가?
• RQ3 내장형 주의 맵이 SmoothGrad와 같은 후행적 기울기 기반 방법에 의존하지 않고 다중 해상도 해석 가능성을 제공할 수 있는가?
• RQ4 형태 인식 특징이 이중심실 심장 MRI 분할 작업에서 분할 정확도를 얼마나 향상시키는가?
• RQ5 제안된 방법이 내재된 해석 가능성을 유지하면서 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• 형태 스트림을 포함한 모델은 SUN09 훈련 세트에서 90.84%의 mIoU를 달성하여 형태 스트림이 없는 베이스라인(90.20%)을 초월했다.
• AC17 검증 세트에서 형태 스트림을 포함한 모델은 81.73%의 mIoU를 기록하여, 베이스라인(79.33%) 대비 1.76%의 더 작은 성능 저하를 보였다.
• 형태 스트림은 일반화 능력을 향상시켰으며, SUN09에서 AC17로 이행할 때 mIoU 감소 폭이 1.76% 더 작아져 더 높은 내구성을 나타냈다.
• 제안된 방법은 전체 AC17 검증 세트에 대해 다중 해상도 열화상 맵을 20초 내에 생성했으며, SmoothGrad 대비 24분이 소요되는 것과 비교해 뚜렷한 계산 효율성을 입증했다.
• 형태 스트림의 주의 맵은 좌심실과 우심실의 해부학적 형태에 명확하게 집중되어 있어 모델이 의미 있는 구조적 특징을 학습했다는 것을 확인했다.
• 이중 주의 디코더는 계층적 해석 가능성을 가능하게 했으며, 고해상도 블록은 세부 사항에 집중하고 저해상도 블록은 전반적인 형태 사전 지식을 포착했다.