Source-linked AI summary
Marginal loss and exclusion loss for partially supervised multi-organ segmentation
Gonglei Shi, Li Xiao, Yang Chen, S. Kevin Zhou
TL;DR
Costly, incomplete organ annotations leave many multi-organ datasets only partially labeled, while learning from their union is a practical research problem. The paper combines marginal and exclusion losses in a single network trained on these datasets. Experiments on five benchmark datasets report significant performance gains and better results than existing partial-label frameworks.
Problem
Multi-organ annotation is costly and time-consuming, leaving datasets often small and only partially labeled, while the practical question is how to use them together.
Source: §1 · pp. 1–2Method
A single multi-organ network is trained on the union of partially labeled datasets using marginal loss for merged background labels and exclusion loss based on organ non-overlap.
Source: §1 · p. 2, §4.1 · p. 5Results
Experiments on five benchmark datasets show significant performance improvement from the proposed losses, and the method greatly outperforms existing frameworks for partially annotated data learning.
Source: §5 · pp. 9–10Takeaways & Limitations
Marginal and exclusion losses support learning a multi-class segmentation network from multiple datasets with partial organ labels.
Source: §5 · p. 9The method still has failure cases involving liver prediction near background with similar features and misjudgments of spleen and pancreas.
Source: §5 · p. 10
Abstract
from arXiv · showhide
Annotating multiple organs in medical images is both costly and time-consuming; therefore, existing multi-organ datasets with labels are often low in sample size and mostly partially labeled, that is, a dataset has a few organs labeled but not all organs. In this paper, we investigate how to learn a single multi-organ segmentation network from a union of such datasets. To this end, we propose two types of novel loss function, particularly designed for this scenario: (i) marginal loss and (ii) exclusion loss. Because the background label for a partially labeled image is, in fact, a `merged' label of all unlabelled organs and `true' background (in the sense of full labels), the probability of this `merged' background label is a marginal probability, summing the relevant probabilities before merging. This marginal probability can be plugged into any existing loss function (such as cross entropy loss, Dice loss, etc.) to form a marginal loss. Leveraging the fact that the organs are non-overlapping, we propose the exclusion loss to gauge the dissimilarity between labeled organs and the estimated segmentation of unlabelled organs. Experiments on a union of five benchmark datasets in multi-organ segmentation of liver, spleen, left and right kidneys, and pancreas demonstrate that using our newly proposed loss functions brings a conspicuous performance improvement for state-of-the-art methods without introducing any extra computation.
1. Introduction
Multi-organ segmentation is clinically useful, but fully annotated data are costly to produce and available datasets are often small and partially labeled. The paper proposes learning one multi-organ network from their union using marginal and exclusion losses.
1. Introduction
- Multi-organ segmentation is used in diagnostic interventions, treatment planning, and treatment delivery.Manual or semi-automated delineation is time-consuming in radiotherapy planning.Source: §1 · p. 1
- Partially labeled datasets are common because full annotation is costly and time-consuming, yet they are not usable by methods requiring full supervision.Existing labeled datasets are often small and cover only some organs.Source: §1 · p. 2
- The paper trains a single multi-organ network on the union of partially labeled datasets, without introducing extra computation.The aim is to use these data to improve segmentation accuracy and robustness.Source: §1 · p. 2
- Marginal loss treats a partially labeled image’s background as a merged label, summing probabilities for the unlabeled organs and true background before applying a loss.The resulting marginal probability can be used with losses such as cross entropy or Dice.Source: §1 · p. 2
- Exclusion loss uses organ non-overlap as prior knowledge to relate labeled ground truth to the estimated segmentation of unlabeled organs.The method uses this constraint alongside marginal loss in a single network that predicts all organ labels and background.Source: §1 · p. 2
- The authors report significant performance improvement from pooling partially labeled datasets and say their method outperforms two state-of-the-art approaches for partial-label learning.The approach uses datasets that are already labeled and available as auxiliary data.Source: §1 · p. 2
2. Related Work
Prior multi-organ segmentation methods include registration-based, statistical-model, and deep-learning approaches, but the surveyed methods rely on fully labeled images. Work on partial labels is limited, and approaches differ in their training setup and use of organ exclusiveness.
2.1. Multi-organ segmentation models
- Multi-organ segmentation approaches include multi-atlas registration, statistical models, deep-learning networks, and combinations of these methods.Multi-atlas performance is limited by image-registration quality; all methods discussed rely on fully labeled images.Source: §2.1 · p. 2
2.2. Multi-organ segmentation with partially annotated data learning
- For partial labels, prior work includes separate fully and partially labeled models, single-class-label learning, and unified multi-dataset strategies; this paper instead studies marginal losses and organ exclusiveness.Zhou et al. (2019) use separate models, Dmitriev and Kaufman (2019) support single-class-labeled datasets, and Fang and Yan (2020) focus on hierarchical architecture.Source: §2.2 · p. 3
2.3. Partially annotated data learning in other tasks
- Partial-label learning has also been studied for classification and object detection, including missing-label and semi-supervised settings.The cited examples address classification or detection rather than multi-organ segmentation.Source: §2.3 · p. 3
3. Method
The method trains one multi-class segmentation network using partially annotated data and introduces marginal and exclusion losses. These losses use merged-label probabilities and organ exclusivity to address partial labeling.
3. Method
- The method trains one multi-class network with partially annotated data alongside a few fully labeled data for baseline training.It uses label merging and organ exclusivity as prior information in the two proposed losses.Source: §3 · p. 3
3.1. Regular cross-entropy loss and regular Dice loss
- Regular cross-entropy and Dice losses provide the common objectives, with Dice loss defined as one minus the Dice score coefficient.The Dice score coefficient measures overlap between the segmentation map and ground truth.Source: §3.1 · p. 3
3.2. Marginal loss
- Marginal loss merges classes in partially labeled data, treating each merged class’s classification probability as a marginal probability.The method defines marginal cross-entropy and marginal Dice losses for the merged labels.Source: §3.2 · p. 4
3.3. Exclusion loss
- Exclusion loss encodes labels of classes mutually exclusive with a given class and penalizes overlap between that exclusion information and the network prediction.The exclusion label is an N-dimensional vector that need not be one-hot.Source: §3.3 · p. 4
4. Experiments and Results
Experiments pool a fully labeled base with four partially labeled datasets in a shared segmentation network. The proposed losses improve segmentation, and the final method outperforms PaNN and PIPO.
4.1. Problem setting and benchmark dataset
- The benchmark combines 30 fully labeled 13-organ CTs with partially labeled liver, spleen, pancreas, and kidney datasets.The kidney dataset was manually divided into left and right kidney labels; cancer labels were merged into their corresponding organ labels.Source: §4.1 · p. 5
4.2. Segmentation networks
- Ten segmentation configurations compare multiclass, binary, and ternary training across fully labeled, partially labeled, and pooled data.The configurations include a multiclass network trained on all five datasets and organ-specific networks trained on individual or combined datasets.Source: §4.2 · p. 5
4.3. Training procedure
- All networks use 3D nnUNet and a two-stage training procedure, first on the fully labeled dataset and then jointly on all five datasets.The second stage uses regular loss for fully labeled samples and marginal plus exclusion losses for partially labeled samples.Source: §4.3 · pp. 5–6
4.4. Ablation studies
- Pooling datasets generally improved results: mean Dice rose from .851 to .900 versus a single labeled dataset and from .874 to .900 versus training on F alone.The corresponding Hausdorff distance reductions were 37.5% and 28.7%, respectively.Source: §4.4 · p. 7
- Joint marginal and exclusion losses increased mean Dice by 3.1% (.931 vs .900); organ-specific gains averaged 4.0% for spleen, 5.05% for pancreas, and 4.45% for kidneys.The authors describe the boost as free because the pooled datasets already exist and are labeled.Source: §4.4 · p. 7
- With Dice loss, pooling reduced average Hausdorff distance from 9.33 to 4.43, and adding exclusion Dice reduced it further to 4.02.The authors attribute large Hausdorff outliers to small predicted regions appearing outside the corresponding organ.Source: §4.4 · pp. 7–8
- A 1:2 marginal-to-exclusion weight worked best on almost all metrics; exclusion alone barely helped pancreas or kidneys, Dice declined as annotations fell, and CE-versus-Dice superiority depended on setup.Combining CE and Dice was generally beneficial in the reported experiments.Source: §4.4 · pp. 7–8
4.5. Comparison with state-of-the-art
- Average Dice was 0.931 for our method versus 0.906 for PaNN and 0.907 for PIPO; mean Hausdorff distance fell 24.0% versus PaNN and 40.0% versus PIPO.The method also performed consistently better than PIPO across organs and datasets, with particular gains on smaller organs.Source: §4.5 · p. 9
5. Discussions and Conclusions
The proposed marginal and exclusion losses support multi-class segmentation from datasets with partial organ labels, and experiments across five benchmark datasets show a performance boost. The method also outperforms existing frameworks for partially annotated learning, though failure cases remain.
5. Discussions and Conclusions
- Marginal loss handles merged labels, while exclusion loss adds mutual exclusiveness as prior knowledge for multi-class segmentation with partial organ labels.The experiments use five benchmark datasets.Source: §5 · p. 9
- The method greatly outperforms existing frameworks for partially annotated data learning, with Table 7 comparing Dice coefficients and Hausdorff distances.The conclusion reports a significant performance boost from marginal and exclusion losses across five benchmark datasets.Source: §5 · pp. 9–10, Table 7 · p. 10
- Typical failures include incorrect liver prediction when background features resemble liver, and misjudgments involving spleen and pancreas.The authors propose incorporating more organ knowledge, such as a shape adversarial prior, in future work.Source: §5 · p. 10