[Paper Review] Prior-aware Neural Network for Partially-Supervised Multi-Organ Segmentation

Yuyin Zhou, Zhe Li+6 more|arXiv (Cornell University)|Apr 12, 2019

Advanced Neural Network Applications42 references47 citations

TL;DR

PaNN integrates anatomical priors on abdominal organ sizes into a partially-supervised multi-organ segmentation framework, using a min-max training objective to regularize output size distributions toward empirical priors and achieve state-of-the-art results on 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%.

Motivation & Objective

• Motivate accurate multi-organ segmentation when annotations are partial due to labeling costs in abdominal CT scans.
• Introduce anatomical priors on organ sizes to regularize segmentation outputs.
• Develop a training objective that incorporates priors via a min-max optimization scheme suitable for stochastic gradient descent.
• Demonstrate that PaNN surpasses previous methods using partially labeled data on a standard abdominal segmentation benchmark.

Proposed method

• Define partial supervision across multiple abdominal organs where some datasets label only a subset of organs.
• Introduce a prior-aware loss that encourages the network output to match empirically observed organ-size proportions from fully-labeled data, by matching average output distributions to prior statistics.
• Formulate the prior loss as a KL divergence between the prior distribution over organs and the network’s average predicted distribution, and convert the objective to a min-max problem using auxiliary dual variables to enable SGD-based optimization.
• Adopt an alternating optimization strategy: (a) estimate pseudo-labels for partially-labeled data and update dual variables to maximize the prior term; (b) update network parameters to minimize the combined loss.
• Provide a two-stage training procedure: first train on fully-labeled data for initialization; then jointly train on fully- and partially-labeled data with the min-max objective.
• Evaluate using 2D and 3D segmentation models on MICCAI 2015 dataset with cross-validation and report Dice scores.

Experimental results

Research questions

• RQ1 Can anatomical priors on organ sizes improve multi-organ segmentation when training data is partially labeled?
• RQ2 How can prior information be integrated into a segmentation objective in a way that is compatible with stochastic optimization?
• RQ3 Does PaNN outperform naive partial supervision or semi-supervised approaches on abdominal organ segmentation benchmarks?
• RQ4 What is the impact of using PaNN with different backbone architectures (2D vs 3D) on segmentation performance under partial supervision?

Key findings

• PaNN achieves state-of-the-art performance on MICCAI 2015 Multi-Atlas Abdomen Labeling with an average Dice score of 84.97%, surpassing prior art by 3.27%.
• Under partial supervision, PaNN generally outperforms naive partially-supervised methods and shows notable gains across various backbones (e.g., ResNet-50, ResNet-101, 3D-UNet).
• The method demonstrates statistically significant improvements for many organs, including small or complex ones such as the gallbladder, esophagus, pancreas, and adrenal glands, across multiple experiments.
• The two-stage training strategy and the min-max optimization enable effective utilization of partially-labeled datasets without requiring additional labeling budget.
• PaNN shows strong generalization in experiments on additional datasets, achieving better fully-, semi-, and partial-supervision performance compared to competing approaches.