[논문 리뷰] `Project & Excite' Modules for Segmentation of Volumetric Medical Scans

Anne-Marie Rickmann, Abhijit Guha Roy외 3명|arXiv (Cornell University)|2019. 06. 11.

Advanced Neural Network Applications참고 문헌 9인용 수 7

한 줄 요약

이 논문은 3D 컨volution 신경망(CNN) 재보정 블록인 'Project & Excite'(PE) 모듈을 제안한다. 기존의 전역 평균 풀링을 사용하는 전통적인 squeeze-and-excite(SE) 모듈과 달리, 주요 3개 축(세로, 가로, 깊이)을 따라 특징 맵을 투영한 후 자극을 주는 방식으로 공간 정보를 유지한다. 이 방법은 어려운 전체 뇌 및 전체 몸체 분할 작업에서 3D U-Net의 성능을 최대 5% 향상시키며, 모델 복잡도는 오직 2%만 증가시킨다.

ABSTRACT

Fully Convolutional Neural Networks (F-CNNs) achieve state-of-the-art performance for image segmentation in medical imaging. Recently, squeeze and excitation (SE) modules and variations thereof have been introduced to recalibrate feature maps channel- and spatial-wise, which can boost performance while only minimally increasing model complexity. So far, the development of SE has focused on 2D images. In this paper, we propose `Project & Excite' (PE) modules that base upon the ideas of SE and extend them to operating on 3D volumetric images. `Project & Excite' does not perform global average pooling, but squeezes feature maps along different slices of a tensor separately to retain more spatial information that is subsequently used in the excitation step. We demonstrate that PE modules can be easily integrated in 3D U-Net, boosting performance by 5% Dice points, while only increasing the model complexity by 2%. We evaluate the PE module on two challenging tasks, whole-brain segmentation of MRI scans and whole-body segmentation of CT scans. Code: https://github.com/ai-med/squeeze_and_excitation

연구 동기 및 목표

• 작은 해부학적 구조를 다룰 때 과적합과 제한된 훈련 데이터로 인해 3D 완전 컨volution 신경망(F-CNNs)이 약한 성능을 보이는 문제를 해결하기 위해.
• 전역 평균 풀링으로 인해 공간적 맥락이 손실되는 3D squeeze-and-excite(SE) 모듈의 문제를 해결하여 세밀한 분할에 필수적인 공간 정보를 유지하기 위해.
• 모델 복잡도를 크게 증가시키지 않으면서도 특징 맵 표현을 향상시키는 경량이고 효율적인 재보정 블록을 개발하기 위해.
• 3D 볼륨 의료 영상에 대한 3D U-Net 아키텍처에서 표준 3D SE 블록과 추가 컨볼루션 레이어에 비해 제안된 모듈의 우수성을 입증하기 위해.

제안 방법

• 3D 볼륨 의료 영상에 특화된 3D 확장된 squeeze-and-excite(SE) 블록으로서 'Project & Excite'(PE) 모듈을 제안한다.
• 3D SE에서 전역 평균 풀링을 대체하여 세로(H), 가로(W), 깊이(D)의 각 공간 축을 따라 채널별 평균을 계산하는 '투영' 연산을 적용함으로써 공간적 구조를 유지한다.
• 얻어진 세 개의 투영 벡터(각 축당 하나)를 자극(subnetwork)의 입력으로 사용하여, 이러한 공간적으로 정보가 반영된 기초 기술 간의 채널 간 의존성을 학습한다.
• 3D U-Net의 각 인코더 및 디코더 블록 뒤에 PE 블록을 통합하여, 모든 단계에서 전역 수신 영역를 유지하면서도 FLOPs나 파라미터 수를 크게 증가시키지 않는다.
• 표준 교차 엔트로피 손실을 사용하여 훈련하고, 평가 지표로 Dice 계수를 사용하며, PE를 기준 3D U-Net 및 3D 채널 SE(cSE) 모듈과 비교한다.
• 기존 3D F-CNN 아키텍처(예: 3D U-Net)와의 호환성을 확보하기 위해 PE 블록을 플러그-앤플레이스 방식으로 적용한다.

실험 결과

연구 질문

• RQ1 각 축을 따라 투영하는 방식으로 공간 정보를 유지하는 3D 재보정 블록이 전역 평균 풀링에 의존하는 표준 3D SE 블록보다 우수한 성능을 보일 수 있는가?
• RQ2 PE 모듈은 어려운 작은 해부학적 구조(예: 편도체 또는 기관지)의 3D 볼륨 스캔에서 분할 정확도를 향상시킬 수 있는가?
• RQ3 PE 모듈은 복잡도가 오직 2% 증가하는 수준에서 추가 컨볼루션 레이어를 추가한 것과 유사한 성능 향상을 달성할 수 있는가?
• RQ4 PE 모듈은 전체 뇌 MRI 및 전체 몸체 CT 분할과 같은 다양한 의료 영상 작업에서 효과적인가?
• RQ5 PE 모듈은 기준 모델 및 3D cSE 모듈에 비해 작은 구조에서의 성능 저하를 줄일 수 있는가?

주요 결과

• PE 모듈은 전체 뇌 MRI 데이터셋인 MALC에서 기준 3D U-Net 대비 평균 Dice 점수를 4% 향상시켰고, 전체 몸체 CT 데이터셋인 Visceral에서 3.5% 향상시켰다.
• 작은 구조인 하부 외측 측뇌실, 편도체, 보상핵에서는 PE 모듈이 Dice 점수를 3–36% 포인트 향상시켰지만, 3D cSE 모듈은 성능 저하를 보였다.
• 기관지와 경추에서 PE 모듈은 각각 0.845와 0.699의 Dice 점수를 기록했고, 3D cSE는 0.491과 0.427로 낮은 성능을 보여, 작은 복잡한 구조에서 뛰어난 성능을 입증했다.
• PE 모듈은 복잡도를 오직 2% 증가시켰지만, Dice 점수를 5% 향상시켰고, 복잡도를 40% 증가시키는 두 개의 추가 컨볼루션 레이어를 추가한 것보다 우수한 성능을 보였다.
• 시각화 결과, PE 모델이 기준 3D U-Net 및 3D cSE 모델이 놓친 왼쪽 하부 외측 측뇌실 및 기관지 분열부와 같은 작은 구조를 성공적으로 분할하는 것으로 확인되었다.
• PE 모듈은 전체 뇌 및 전체 몸체 분할 작업 모두에서 일관된 성능 향상을 보이며, 3D 의료 영상 분할 분야에 넓은 적용 가능성을 보였다.