[논문 리뷰] Prior-aware Neural Network for Partially-Supervised Multi-Organ Segmentation

Yuyin Zhou, Zhe Li외 6명|arXiv (Cornell University)|2019. 04. 12.

Advanced Neural Network Applications참고 문헌 42인용 수 47

한 줄 요약

PaNN은 복부 장기 크기에 대한 해부학적 priors를 부분적으로 감독되는 다기관 세분화 프레임워크에 통합하고, min-max 학습 목표를 사용하여 출력 크기 분포를 경험적 priors toward에 정규화하며 MICCAI 2015에서 최첨단 성과를 달성한다.

ABSTRACT

Accurate multi-organ abdominal CT segmentation is essential to many clinical applications such as computer-aided intervention. As data annotation requires massive human labor from experienced radiologists, it is common that training data are partially labeled, e.g., pancreas datasets only have the pancreas labeled while leaving the rest marked as background. However, these background labels can be misleading in multi-organ segmentation since the "background" usually contains some other organs of interest. To address the background ambiguity in these partially-labeled datasets, we propose Prior-aware Neural Network (PaNN) via explicitly incorporating anatomical priors on abdominal organ sizes, guiding the training process with domain-specific knowledge. More specifically, PaNN assumes that the average organ size distributions in the abdomen should approximate their empirical distributions, a prior statistics obtained from the fully-labeled dataset. As our training objective is difficult to be directly optimized using stochastic gradient descent [20], we propose to reformulate it in a min-max form and optimize it via the stochastic primal-dual gradient algorithm. PaNN achieves state-of-the-art performance on the MICCAI2015 challenge "Multi-Atlas Labeling Beyond the Cranial Vault", a competition on organ segmentation in the abdomen. We report an average Dice score of 84.97%, surpassing the prior art by a large margin of 3.27%.

연구 동기 및 목표

• 복부 CT 스캔에서 라벨링 비용으로 주석이 부분적일 때 정확한 다기관 세분화를 촉진한다.
• 장기 크기에 대한 해부학적 priors를 도입하여 세분화 출력을 정규화한다.
• 확률적 경사 하강법에 적합한 min-max 최적화 스킴을 통해 priors를 통합하는 학습 목표를 개발한다.
• PaNN이 부분적으로 라벨링된 데이터를 사용한 표준 복부 세분화 벤치마크에서 이전 방법들을 능가함을 Demonstrate한다.

제안 방법

• 여러 복부 장기에 걸쳐 불완전한 감독이 이루어지는 partial supervision을 정의하다: 일부 데이터셋은 장기의 부분 집합만 라벨링한다.
• 네트워크 출력이 완전 라벨링 데이터에서 관측된 장기 크기 비율과 평균 출력 분포를 일치시키는 것을 권장하는 prior-aware 손실을 도입하여 경험적으로 관찰된 장기 크기 비율에 맞추도록 정규화한다.
• prior 손실을 장기에 대한 prior 분포와 네트워크의 평균 예측 분포 간의 KL 발산으로 형식화하고, 보조 이중 변수들을 이용해 objective를 min-max 문제로 변환하여 SGD 기반 최적화를 가능케 한다.
• 교대 최적화 전략을 채택한다: (a) 부분적으로 라벨링된 데이터의 의사 라벨을 추정하고 priors 항을 최대화하기 위해 이중 변수들을 업데이트한다; (b) 결합 손실을 최소화하도록 네트워크 매개변수를 업데이트한다.
• 두 단계 학습 절차를 제공한다: 먼저 완전 라벨링 데이터에서 초기화를 위해 학습한 뒤; 그다음 완전 라벨링 데이터와 부분적으로 라벨링 데이터를 함께 min-max 목적 함수로 학습한다.
• MICCAI 2015 데이터셋에 대해 2D 및 3D 세분화 모델을 교차 검증으로 평가하고 Dice 점수를 보고한다.

실험 결과

연구 질문

• RQ1 해부학적 priors가 부분적으로 라벨링된 학습 데이터일 때 다기관 세분화를 향상시킬 수 있는가?
• RQ2 세분화 목적 함수에 prior 정보를 확률적 최적화와 호환되도록 어떻게 통합할 수 있는가?
• RQ3 PaNN이 복부 장기 세분화 벤치마크에서 순수 부분 감독 또는 반지도 학습 방식보다 우수한가?
• RQ4 부분 감독하에서 서로 다른 백본 아키텍처(2D 대 3D)를 사용할 때 세분화 성능에 미치는 PaNN의 영향은 무엇인가?

주요 결과

• PaNN은 MICCAI 2015 Multi-Atlas Abdomen Labeling에서 평균 Dice 점수 84.97%로 최첨단 성능을 달성하며, 이전 연구보다 3.27% 포인트 개선했다.
• 부분 감독하에서 PaNN은 일반적으로 순수 부분 감독 방법을 능가하고 다양한 백본(예: ResNet-50, ResNet-101, 3D-UNet)에서 눈에 띄는 이득을 보인다.
• 방법은 간담췌, 간담췌처럼 작고 복잡한 장기를 포함한 다수의 장기에 대해 통계적으로 유의미한 개선을 보여주며, 여러 실험에서 개선을 보인다.
• 두 단계 학습 전략과 min-max 최적화는 추가 라벨링 예산 없이도 부분적으로 라벨링된 데이터셋의 효과적인 활용을 가능하게 한다.
• PaNN은 추가 데이터셋에서의 실험에서도 강한 일반화를 보이며, 경쟁 방법들과 비교해 완전 감독, 반감독, 부분 감독 성능에서 더 나은 성능을 달성한다.