[논문 리뷰] V-NAS: Neural Architecture Search for Volumetric Medical Image Segmentation

Zhuotun Zhu, Chenxi Liu외 3명|arXiv (Cornell University)|2019. 06. 06.

Advanced Neural Network Applications참고 문헌 29인용 수 12

한 줄 요약

이 논문은 체적 의료 영상 분할을 위한 2D, 3D 및 허위 3D(P3D) 컨볼루션의 최적 조합을 자동으로 학습하는 연속적이고 미분 가능한 신경망 아키텍처 탐색 프레임워크인 V-NAS를 제안한다. 아키텍처 탐색을 연속적이고 미분 가능한 최적화 문제로 공식화함으로써, V-NAS는 세 가지 공개 데이터셋에서 정상 및 비정상 장기 분할 작업 전반에서 최신 기술 수준의 성능을 달성하며, 검색된 아키텍처는 데이터셋 간에 잘 일반화됨을 보였다.

ABSTRACT

Deep learning algorithms, in particular 2D and 3D fully convolutional neural networks (FCNs), have rapidly become the mainstream methodology for volumetric medical image segmentation. However, 2D convolutions cannot fully leverage the rich spatial information along the third axis, while 3D convolutions suffer from the demanding computation and high GPU memory consumption. In this paper, we propose to automatically search the network architecture tailoring to volumetric medical image segmentation problem. Concretely, we formulate the structure learning as differentiable neural architecture search, and let the network itself choose between 2D, 3D or Pseudo-3D (P3D) convolutions at each layer. We evaluate our method on 3 public datasets, i.e., the NIH Pancreas dataset, the Lung and Pancreas dataset from the Medical Segmentation Decathlon (MSD) Challenge. Our method, named V-NAS, consistently outperforms other state-of-the-arts on the segmentation task of both normal organ (NIH Pancreas) and abnormal organs (MSD Lung tumors and MSD Pancreas tumors), which shows the power of chosen architecture. Moreover, the searched architecture on one dataset can be well generalized to other datasets, which demonstrates the robustness and practical use of our proposed method.

연구 동기 및 목표

• 체적 의료 영상 분할에서 고정된 2D, 3D 또는 P3D 컨볼루션 네트워크의 한계를 해결하기 위해, 3D 공간적 맥락을 忽시하거나 높은 계산 비용으로 인해 성능이 떨어지는 문제를 해결한다.
• 체적 의료 데이터의 특성에 맞게 각 레이어에서 2D, 3D 및 P3D 컨볼루션 간의 최적 아키텍처 구성(선택)을 자동으로 학습하는 방법을 개발한다.
• 검색된 아키텍처가 다양한 의료 영상 데이터셋 간에 일반화되는지 평가한다. 특히 정상 및 비정상 장기 분할에 대해 검토한다.
• 미분 가능한 신경망 아키텍처 탐색이 수작업으로 설계된 아키텍처보다 분할 정확도를 높이면서도 계산 효율성을 유지할 수 있음을 입증한다.

제안 방법

• 체적 분할을 위한 신경망 아키텍처 탐색을 연속적이고 미분 가능한 최적화 문제로 공식화함으로써, 아키텍처 파rameter의 엔드 투 엔드 학습을 가능하게 한다.
• 각 레이어가 2D, 3D 및 허위 3D(P3D) 컨볼루션 중에서 동적으로 선택할 수 있는 탐색 공간을 도입하며, 여기서 P3D는 3D 커널을 두 개의 연속된 2D 컨볼루션(k×k×1 이후 1×1×k)으로 대체한다.
• DARTS 방식과 유사한 연속적 근사 기반으로 아키텍처 탐색 공간을 미분 가능하게 변환함으로써, 각 레이어에서 컨볼루션 유형 선택에 대한 기울기 기반 최적화를 가능하게 한다.
• 사전 훈련 없이 단일 단계의 미분 가능한 탐색 프로세스를 사용하여 네트워크를 훈련시키며, 아키텍처와 가중치를 동시에 최적화한다.
• 단일 경로 탐색 전략을 사용하고 공유된 가중치를 하나의 세트로 사용하여, 여러 데이터셋에 걸쳐 아키텍처 공간을 효율적으로 탐색한다.
• NIH 췌장, MSD 기관지 폐 종양 및 MSD 췌장 종양 데이터셋의 보류된 테스트 세트에서 최종 검색된 아키텍처의 성능과 일반화 능력을 평가한다.

실험 결과

연구 질문

• RQ1 미분 가능한 신경망 아키텍처 탐색 프레임워크는 체적 의료 영상 분할을 위한 2D, 3D 및 허위 3D 컨볼루션의 최적 조합을 자동으로 발견할 수 있는가?
• RQ2 제안된 V-NAS 방법은 컨볼루션 유형을 고정한 수작업 아키텍처(예: 순수 2D, 3D 또는 P3D)보다 분할 정확도에서 뛰어나게 성능을 높일 수 있는가?
• RQ3 한 데이터셋에서 검색된 아키텍처는 해부학적 구조와 병변이 다른 다른 의료 영상 데이터셋으로 잘 일반화되는가?
• RQ4 V-NAS의 성능는 정상 장기 및 비정상 종양 분할 작업 전반에서 최신 기술 수준의 방법들과 비교해 어떻게 되는가?

주요 결과

• NIH 췌장 데이터셋에서 V-NAS는 평균 DSC 85.15%를 달성하여, 최고의 수작업 구성인 P3D/P3D(84.75%)를 초월했다.
• MSD 기관지 폐 종양 데이터셋에서 V-NAS는 평균 DSC 55.27%를 기록하여, 최고의 수작업 구성인 3D/3D(53.74%)보다 1.53% 높은 성능을 보였다.
• MSD 췌장 종양 데이터셋에서 V-NAS는 최신 기술 수준의 성능을 달성하였으며, 정상 및 비정상 장기 분할 작업 전반에서 일관된 우수성을 입증했다.
• 한 데이터셋에서 검색된 아키텍처는 다른 데이터셋으로 잘 일반화되어, 실제 임상 적용에 있어 강건성과 실용성을 보여주었다.
• 제거 실험을 통해 NAS를 통한 컨볼루션 유형 자동 선택이 고정된 수작업 설정보다 뛰어나다는 것을 확인하였으며, 다양한 해부학적 및 병리적 변화를 다루는 아키텍처의 유연성의 중요성을 입증하였다.
• 이 방법은 계산적으로 효율적이며 사전 훈련이 필요 없어, 아키텍처와 가중치를 동시에 엔드 투 엔드 최적화할 수 있다.