[논문 리뷰] 3D U$^2$-Net: A 3D Universal U-Net for Multi-Domain Medical Image Segmentation

Chao Huang, Hu Han외 3명|arXiv (Cornell University)|2019. 09. 04.

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

한 줄 요약

이 논문은 분리 가능 컨볼루션을 사용하여 도메인 간 특징을 공유하면서 채널별 컨볼루션을 통해 도메인 특화 패턴을 학습하는 3D U²-Net을 제안한다. 이 모델은 작업별 전용 네트워크와 비슷한 분할 정확도를 달성하면서 파라미터 수의 1%만을 사용하며, 새로운 기관에 대해 약 0.3%의 추가 파라미터로 효과적으로 적응한다.

ABSTRACT

Fully convolutional neural networks like U-Net have been the state-of-the-art methods in medical image segmentation. Practically, a network is highly specialized and trained separately for each segmentation task. Instead of a collection of multiple models, it is highly desirable to learn a universal data representation for different tasks, ideally a single model with the addition of a minimal number of parameters steered to each task. Inspired by the recent success of multi-domain learning in image classification, for the first time we explore a promising universal architecture that handles multiple medical segmentation tasks and is extendable for new tasks, regardless of different organs and imaging modalities. Our 3D Universal U-Net (3D U$^2$-Net) is built upon separable convolution, assuming that {\it images from different domains have domain-specific spatial correlations which can be probed with channel-wise convolution while also share cross-channel correlations which can be modeled with pointwise convolution}. We evaluate the 3D U$^2$-Net on five organ segmentation datasets. Experimental results show that this universal network is capable of competing with traditional models in terms of segmentation accuracy, while requiring only about $1\%$ of the parameters. Additionally, we observe that the architecture can be easily and effectively adapted to a new domain without sacrificing performance in the domains used to learn the shared parameterization of the universal network. We put the code of 3D U$^2$-Net into public domain. \url{https://github.com/huangmozhilv/u2net_torch/}

연구 동기 및 목표

• 각 의료 영상 분할 작업에 대해 별도의 딥 러닝 모델을 훈련하는 데에 드는 비효율성과 확장의 어려움을 해결한다.
• 재학습 없이도 다양한 해부학적 구조와 영상 모odalities(예: CT, MRI) 간에 일반화할 수 있는 단일 신경망을 가능하게 한다.
• 다양한 도메인에서 높은 분할 정확도를 유지하면서 모델 크기를 최소화하는 파라미터 효율적인 아키텍처를 개발한다.
• 기존 작업의 성능 저하 없이 새로운 기관(예: 비장)에 대해 최소한의 추가 파라미터로 전이 학습을 가능하게 한다.
• 이미지 분류에서의 성공을 조밀한 예측 작업인 다중 도메인 분할에 확장함으로써, 의료 영상 분할 분야에서의 다중 도메인 학습 격차를 해소한다.

제안 방법

• 채널별 컨볼루션을 통해 도메인 특화된 공간 패턴을 캡처하고, 공유된 점 컨볼루션을 통해 채널 간 공유 상관관계를 모델링하는 분리 가능 컨볼루션을 사용한 3D U-Net 기반의 유니버설 아키텍처를 설계한다.
• 새로운 도메인에 적응할 때 공유된 구성 요소(예: 점 컨볼루션)를 동결하면서 도메인 특화 파라미터(예: 입력 레이어, 채널별 컨볼루션)를 도입한다.
• 최적화 과정에서 모든 도메인의 기여가 균형을 이루도록, 모든 도메인 간 라운드 로빈 샘플링 전략을 사용하여 유니버설 모델을 훈련시킨다.
• 학습 손실의 지수 이동 평균을 사용하여 학습률 감소 스케줄링을 적용하며, 손실이 정체될 경우 5배로 감소시킨다.
• 5개의 기관 분할 데이터셋에서 엔드 투 엔드 훈련을 위해 ADAM 옵timizer를 사용하고, 가중치 감소와 초기 학습률 3×10⁻⁴를 설정한다.
• 공유 레이어를 동결하고 도메인 특화 모듈(예: 새로운 입력 및 채널별 컨볼루션)만 미세조정함으로써, 새로운 도메인에 대해 제로샷 적응을 가능하게 한다.

실험 결과

연구 질문

• RQ1 다른 기관과 영상 모달리티를 가진 다양한 의료 영상 분할 도메인(예: 심장, 간, 전립선)에서 단일 딥 뉴럴 네트워크가 경쟁력 있는 분할 성능을 달성할 수 있는가?
• RQ2 공유 및 도메인 특화 구성 요소를 가진 유니버설 아키텍처가 독립 모델과 공유 전용 모델보다 분할 정확도와 파라미터 효율성 측면에서 뛰어나게 성능을 내는가?
• RQ3 기존 작업의 성능 저하 없이, 최소한의 추가 파라미터로 새로운 보기에(예: 비장) 일반화할 수 있는 정도는 어느 정도인가?
• RQ4 채널별로 도메인 특화 특징을 처리하고 점 컨볼루션을 통해 공유 특징을 모델링하는 분리 가능 컨볼루션의 사용이 다중 도메인 분할에서 모델 효율성과 성능을 어떻게 향상시키는가?
• RQ5 새로운 의료 영상 작업에 대해 유니버설 모델을 확장할 때, 모델 크기, 파라미터 공유, 분할 정확도 사이의 상충 관계는 어떻게 되는가?

주요 결과

• 3D U²-Net은 심장, 간, 해마, 전립선, 췌장 등 5개의 기본 도메인에서 평균 DICE 스코어 83.10%를 기록했으며, 독립 모델의 84.28%와 매우 유사한 성능을 달성했고, 파라미터 수는 1.7M(독립 모델의 126.7M 대비)에 불과했다.
• 유니버설 모델은 모든 기본 도메인에서 공유 모델(82.56% 평균 DICE)을 초월했으며, 전립선 분할의 외곽영역(PZ)과 전이영역(TZ)에서 특히 두드러진 성능 향상을 보였다(68.50% 대비 68.38%).
• 새로운 도메인(비장)에 적응한 결과, 유니버설 모델은 91.60%의 DICE 스코어를 기록했고, 추가 파라미터는 0.1M에 불과했으며, 공유 모델(90.67%)을 초월했고, 전체 독립 모델의 파라미터 수의 약 0.3%에 그쳤다.
• 유니버설 모델의 파라미터 수는 독립 모델의 약 1%(1.7M 대비 126.7M)에 불과하여 극도로 높은 파라미터 효율성을 입증했다.
• 새로운 기관에 적응한 후에도 모든 원래 도메인에서 뛰어난 성능을 유지하여, 공유 표현이 안정적이고 일반화 가능하다는 것을 확인했다.
• 공유 구성 요소를 동결하고 도메인 특화 모듈만 훈련함으로써, 모델은 효과적인 제로샷 적응을 가능하게 하였으며, 성능 저하 최소화와 높은 확장성 확보에 기여했다.