[논문 리뷰] MixSearch: Searching for Domain Generalized Medical Image Segmentation Architectures

Luyan Liu, Zhiwei Wen외 7명|arXiv (Cornell University)|2021. 02. 26.

Advanced Neural Network Applications참고 문헌 50인용 수 5

한 줄 요약

MixSearch는 다중 도메인, 다중 작업 의료 영상의 혼합을 통해 복합 의료 데이터셋을 구성하고, 일반화된 U-형태 세그멘테이션 아키텍처를 훈련하기 위한 새로운 신경망 아키텍처 탐색(NAS) 프레임워크를 제안한다. 유연한 인코더-디코더 검색 공간과 학습 가능한 특징 융합을 활용함으로써, MixSearch는 사전에 수작업으로 설계된 모델들—Attention U-Net 및 U-Net 변종—을 초월하는 성능을 보이며, 예측 불가능한 데이터셋에서 강력한 도메인 일반화 성능을 입증한다. 이는 Dice 및 Jaccard 스코어에서 뚜렷한 향상을 보인다.

ABSTRACT

Considering the scarcity of medical data, most datasets in medical image analysis are an order of magnitude smaller than those of natural images. However, most Network Architecture Search (NAS) approaches in medical images focused on specific datasets and did not take into account the generalization ability of the learned architectures on unseen datasets as well as different domains. In this paper, we address this point by proposing to search for generalizable U-shape architectures on a composited dataset that mixes medical images from multiple segmentation tasks and domains creatively, which is named MixSearch. Specifically, we propose a novel approach to mix multiple small-scale datasets from multiple domains and segmentation tasks to produce a large-scale dataset. Then, a novel weaved encoder-decoder structure is designed to search for a generalized segmentation network in both cell-level and network-level. The network produced by the proposed MixSearch framework achieves state-of-the-art results compared with advanced encoder-decoder networks across various datasets.

연구 동기 및 목표

• 단일이고 소규모 데이터셋에서 훈련된 기존 의료 영상 세그멘테이션 모델의 일반화 부족 문제를 해결하기 위해.
• 의료 영상 분석에서의 데이터 부족 및 도메인 이동 문제를 극복하기 위해 대규모 복합 훈련 분포를 생성하기 위해.
• 다중 도메인 의료 영상 세그멘테이션을 위한 강력하고 일반화 가능한 U-형태 아키텍처를 탐색하는 엔드 투 엔드 NAS 프레임워크를 개발하기 위해.
• 다른 영상 모odal리티와 해부학적 영역에서 온 완전히 새로운 데이터셋에서 학습된 아키텍처의 일반화 능력을 평가하기 위해.

제안 방법

• 실제 및 증강된 복합 영상 포함, 여러 의료 영상 세그멘테이션 작업 및 도메인의 샘플을 혼합하여 복합 데이터셋을 구성한다.
• 셀 수준 및 네트워크 수준에서 동시 검색이 가능한 학습 가능한 수평 및 수직 특징 융합 모듈을 갖춘 새로운 웨이브드 인코더-디코더 아키텍처를 도입한다.
• 복합 데이터셋에서 아키텍처를 최적화하기 위해 DARTS 기반의 미분 가능 아키텍처 탐색 기법을 적용한다.
• 혼합 계수 μ=0.5를 사용한 MixUp과 같은 데이터 혼합 전략을 통해 도메인 다양성과 분포 내구성을 향상시킨다.
• 다중 해상도 감독을 위해 디코더 레이어의 여러 단계에 깊이 감독 헤드를 적용하여 특징 학습과 공간 정밀도를 향상시킨다.
• ISIC, LiTS, KiTS, ETIS를 병합한 복합 데이터셋에서 아키텍처 탐색을 수행하여 교차 도메인 일반화를 시뮬레이션한다.

실험 결과

연구 질문

• RQ1 다양한 도메인의 의료 영상 혼합을 통해 구성된 복합 데이터셋이 학습된 세그멘테이션 아키텍처의 일반화 능력을 향상시키는가?
• RQ2 학습 가능한 특징 융합을 갖춘 웨이브드 인코더-디코더 검색 공간이 다양한 의료 영상 도메인에서 더 높은 성능과 내구성을 제공하는가?
• RQ3 복합 데이터셋에서의 NAS가 사전에 수작업으로 설계된 모델보다 예측 불가능한 데이터셋에서 더 나은 일반화 성능을 보이는가?
• RQ4 제안된 MixSearch 프레임워크는 전용 및 유니온 기반 NAS 전략과 비교해 일반화 능력과 성능 면에서 어떻게 다른가?

주요 결과

• ETIS 데이터셋에서 MixSearch는 Attention U-Net보다 Dice 스코어에서 1.0% 향상되고 Jaccard 유사도에서 1.5% 향상되었다.
• LiTS 데이터셋에서 MixSearch는 MR U-Net(기존 U-Net++ 대비 0.01% Dice 스코어 우월)보다도 뚜렷한 격차로 앞서나갔다.
• KiTS 데이터셋에서 MixSearch는 R2T U-Net의 94.64% 대비 94.53%의 Dice 스코어를 기록하여 최신 기술 수준에 근접한 성능를 보였다.
• MixSearch가 학습한 아키텍처는 UnionSearch보다 더 우수한 일반화 성능를 보였으며, UnionSearch는 개별 소스 도메인에 과적합되어 기본 U-Net조차도 성능이 열 劣했다.
• MixSearch의 단계별 특징 융합은 가장자리 검출 및 마스크 품질 향상에 기여했으며, Jaccard 스코어는 X^{0,5}에서 X^{0,9}로 상승했다.
• 최종 출력 X^{0,7}은 높은 정확도를 보이며, 성능 손실 없이 계산량을 줄일 수 있는 프런팅 가능한 추론 헤드로 활용 가능하다.