[논문 리뷰] A unified representation network for segmentation with missing modalities

Kenneth K. Lau, Jonas Adler외 1명|arXiv (Cornell University)|2019. 08. 19.

Advanced Neural Network Applications참고 문헌 15인용 수 17

한 줄 요약

이 논문은 입력 모odal이 누락될 경우 의료 영상에서 딥러닝 분할 성능을 햖을 수 있도록 통합 표현 네트워크(URN)와 모달리티 드롭아웃을 제안한다. URN은 다양한 입력 모달리티 간의 공통 표현을 학습하여, 데이터가 불완전한 경우에도 강건한 분할 성능을 달성하며, 표준 U-net과 모달리티 드롭아웃만을 사용한 경우보다 성능이 뛰어나며, 특히 BRATS와 HCP와 같은 다중모달 데이터셋에서 미리 훈련된 경우에 두드러진다.

ABSTRACT

Over the last few years machine learning has demonstrated groundbreaking results in many areas of medical image analysis, including segmentation. A key assumption, however, is that the train- and test distributions match. We study a realistic scenario where this assumption is clearly violated, namely segmentation with missing input modalities. We describe two neural network approaches that can handle a variable number of input modalities. The first is modality dropout: a simple but surprisingly effective modification of the training. The second is the unified representation network: a network architecture that maps a variable number of input modalities into a unified representation that can be used for downstream tasks such as segmentation. We demonstrate that modality dropout makes a standard segmentation network reasonably robust to missing modalities, but that the same network works even better if trained on the unified representation.

연구 동기 및 목표

• 시험 시 입력 모달리티가 누락될 경우 발생하는 실제 세계적 과제인 의료 영상 분할 문제를 해결하기 위해.
• 모든 가능한 입력 모달리티 조합에 대해 별도의 네트워크를 훈련하는 것을 방지하여 조합 폭발 문제를 피하기 위해.
• 비지도 사전 훈련과 통합 표현 학습을 통해 누락된 모달리티 상황에서의 분할 강건성과 성능 향상을 위해.
• 모달리티 드롭아웃과 통합 표현 네트워크가 다중모달 분할에서 일반화 능력을 향상시키고 누락된 데이터에 대한 민감도를 감소시키는지 평가하기 위해.

제안 방법

• 모달리티 드롭아웃은 훈련 중에 전체 입력 채널을 무작위로 0으로 설정하여, 네트워크가 모달리티에 관계없이 특징을 학습하도록 유도하고, 누락된 입력에 대한 강건성을 향상시킨다.
• 통합 표현 네트워크(URN)는 각 입력 모달리티에 대해 별도의 U-net 인코더를 사용하며, 학습 가능한 인코더를 통해 모든 모달리티를 공통 잠재 표현으로 매핑한다.
• URN은 BRATS와 HCP와 같이 부분적으로 겹치는 모달리티를 가진 데이터셋에서 비지도 방식으로 사전 훈련되며, 사전 훈련 중에는 분할 레이블이 필요하지 않다.
• 다음으로, 통합 표현을 기반으로 디코더 헤드를 훈련시켜 변수 입력 모달리티를 가진 종단 간 분할을 수행한다.
• 모달리티 드롭아웃 중 제거할 모달리티 수를 샘플링하기 위해 절단 기하분포를 사용한다: $ p_\theta(k) = \frac{(1-\theta)\theta^k}{1-\theta^{N+1}} $ for $ k = 0, \dots, N_{\text{max}} $.
• 사전 훈련은 HCP 데이터셋의 1108명의 건강한 피험자로부터의 T1 및 T2 MR 스캔을 대상으로 하며, 각 피험자 및 모달리티별로 강도 정규화와 비편향 필드 보정을 적용한다.

실험 결과

연구 질문

• RQ1 모달리티 드롭아웃은 표준 U-net의 누락된 입력 모달리티에 대한 강건성을 향상시키는가?
• RQ2 다양한 모달리티 간의 통합 표현을 학습하는 것이 모달리티 드롭아웃만 사용하는 것보다 더 나은 분할 성능을 낳는가?
• RQ3 부분적인 모달리티 겹침이 있는 여러 데이터셋에서의 비지도 사전 훈련이 후속 분할 성능 향상에 기여하는가?
• RQ4 불완전한 입력 조합에서 테스트했을 때, 통합 표현 네트워크가 표준 모델보다 더 효과적인가?

주요 결과

• 모달리티 드롭아웃은 표준 U-net의 강건성을 크게 향상시켜, 시험 시 모달리티가 누락되었을 때의 성능 저하를 줄였다.
• 통합 표현 네트워크(URN)는 모든 입력 조합에서 전암세포 분할의 최고 Dice 스코어를 기록했으며, 베이스라인 U-net과 모달리티 드롭아웃 변종보다 뛰어난 성능을 보였다.
• BRATS와 HCP 데이터를 결합하여 URN을 사전 훈련함으로써, 강화된 종양 및 종양 핵심 영역의 분할 성능이 향상되었으며, 전체 종양 성능은 사전 훈련되지 않은 URN과 유사했다.
• 모든 모달리티가 존재하는 상황에서도 모달리티 드롭아웃을 사용한 훈련은 베이스라인 U-net과 URN 양쪽 모두 성능 향상을 가져왔으며, 이는 모달리티 간 특징 공유를 장려하는 정규화 역할을 한다는 것을 시사한다.
• FLAIR는 BRATS에서 전체 종양 분할에 있어 가장 유용한 모달리티로 밝혀졌으며, 이 모달리티가 누락되었을 경우 베이스라인 모델에서 성능 저하가 가장 컸다.
• 정성적 결과에서는 URN에서 T1c만으로도 임상에서 강화된 종양에 주로 사용되는 바이어블이지만, 모든 종양 구조의 신뢰할 수 있는 분할을 생성할 수 있었다.