Combo Loss: Handling Input and Output Imbalance in Multi-Organ Segmentation

Saeid Asgari Taghanaki; Yefeng Zheng; S. Kevin Zhou; Bogdan Georgescu; Puneet Sharma; Daguang Xu; Dorin Comaniciu; Ghassan Hamarneh

Source-linked AI summary

Combo Loss: Handling Input and Output Imbalance in Multi-Organ Segmentation

Saeid Asgari Taghanaki, Yefeng Zheng, S. Kevin Zhou, Bogdan Georgescu, Puneet Sharma, Daguang Xu, Dorin Comaniciu, Ghassan Hamarneh

arXiv:1805.02798v6May 8, 2018cs.CV

TL;DR

Multi-organ segmentation must address input class imbalance and the false-positive/false-negative trade-off in model outputs. The paper introduces a curriculum-based loss combining Dice with modified cross entropy and evaluates it across PET, MRI, and ultrasound segmentation tasks. The proposed loss improves segmentation performance across multi-organ and single-organ settings and can reduce result variance.

Problem
Medical segmentation faces input imbalance from small foregrounds or differently sized organs and output imbalance between false positives and false negatives.
Method
The paper introduces a curriculum-learning loss that combines Dice with weighted cross entropy to address both imbalance types.
Results
The proposed loss improves multi-organ and single-organ segmentation across different modalities, with PET Dice scores ranging from 0.58 to 0.91.
Takeaways & Limitations
The loss can improve competing segmentation models while controlling the trade-off between false positives and false negatives.

Abstract

from arXiv · show

Simultaneous segmentation of multiple organs from different medical imaging modalities is a crucial task as it can be utilized for computer-aided diagnosis, computer-assisted surgery, and therapy planning. Thanks to the recent advances in deep learning, several deep neural networks for medical image segmentation have been introduced successfully for this purpose. In this paper, we focus on learning a deep multi-organ segmentation network that labels voxels. In particular, we examine the critical choice of a loss function in order to handle the notorious imbalance problem that plagues both the input and output of a learning model. The input imbalance refers to the class-imbalance in the input training samples (i.e., small foreground objects embedded in an abundance of background voxels, as well as organs of varying sizes). The output imbalance refers to the imbalance between the false positives and false negatives of the inference model. In order to tackle both types of imbalance during training and inference, we introduce a new curriculum learning based loss function. Specifically, we leverage Dice similarity coefficient to deter model parameters from being held at bad local minima and at the same time gradually learn better model parameters by penalizing for false positives/negatives using a cross entropy term. We evaluated the proposed loss function on three datasets: whole body positron emission tomography (PET) scans with 5 target organs, magnetic resonance imaging (MRI) prostate scans, and ultrasound echocardigraphy images with a single target organ i.e., left ventricular. We show that a simple network architecture with the proposed integrative loss function can outperform state-of-the-art methods and results of the competing methods can be improved when our proposed loss is used.

I. INTRODUCTION

Medical image segmentation must address both training-time input imbalance and inference-time output imbalance. The paper motivates a curriculum-based loss combining Dice and cross entropy because existing approaches do not explicitly handle both types together.

Motivation: Deep encoder-decoder models have become widely used for binary and multi-class medical image segmentation across imaging modalities.Examples include 2D U-Net, 3D U-Net, 3D V-Net, and 2D SegNet.
Input and output imbalance: Input imbalance arises when foreground voxels or smaller organs are overwhelmed by abundant background voxels and larger classes.Classes with more observations can overshadow minority classes during training.
Input and output imbalance: Output imbalance concerns the trade-off between false positives and false negatives, whose relative importance depends on the medical application.False positives may need stronger penalties for missing organs or normal active regions in PET images.
Limitations of existing losses: Cross entropy can control output imbalance but performs suboptimally with highly input-imbalanced images, whereas Dice handles input imbalance without explicitly controlling the FP/FN trade-off.The paper identifies these complementary limitations as motivation for combining the two objectives.
Proposed contribution: The proposed curriculum-learning loss combines Dice with cross entropy to handle input and output imbalance during segmentation training.The contribution is positioned as an algorithmic-level loss-function improvement for existing segmentation models.

II. METHOD

The method combines Dice and modified cross-entropy losses to address input class imbalance and control false-positive/false-negative trade-offs in multi-organ segmentation. It applies this loss with a simple deep architecture, volume handling, and backpropagation-based optimization.

Loss design: Dice-based optimization cannot explicitly control the trade-off between false positives and false negatives, motivating the modified cross-entropy component.Generalized Dice additionally requires whole-volume statistics to produce meaningful weights, which conflicts with sub-volume processing.
Optimization: Cross entropy avoids the sigmoid-saturation gradient slowdown described for Dice-based losses, so parameter gradients depend on output, target, and preceding-layer input.Model parameters are optimized by error backpropagation and the chain rule.
Combo Loss: Combo Loss is a weighted sum of Dice loss and modified cross entropy, combining input-imbalance handling with false-positive/false-negative control.The cross-entropy term encodes curriculum learning and supports smoother training.
Loss design: β controls error penalization: β < 0.5 weights false positives more heavily, whereas larger β shifts emphasis toward false negatives.Add-one smoothing is used to prevent division by zero.
Loss design: The loss is computed after sigmoid normalization over flattened W × H × D × C predicted and one-hot ground-truth volumes, extending directly to multi-class segmentation.This formulation controls false positives and false negatives across the entire volume.
Architecture and training: The study uses a simple architecture to test whether improvements arise from the loss rather than elaborate network design, with random sub-volumes for PET and entire volumes for binary datasets.PET testing uses a volumetric sliding-window procedure because whole-body volumes are too large for memory.

III. IMPLEMENTATION DETAILS

The implementation uses modality-specific volume handling and standard training choices, with memory constraints shaping PET processing more strongly than MRI or ultrasound.

Training and inference: PET training extracts randomly positioned sub-volumes per organ plus negative background sub-volumes, while MRI and ultrasound use entire volumes.The PET setup samples 100 80×80×80-voxel sub-volumes per organ and another 100 background sub-volumes.
Memory management: MRI and ultrasound images are slightly resampled to fit whole volumes into memory without substantial information loss.PET requires much greater resizing to fit memory, causing a considerable accuracy drop.
Optimization: Models use ADADELTA optimization, batch normalization after each convolution, and grid search for the loss weights α and β.ADADELTA uses learning rate 1, ρ = 0.95, ϵ = 1e −08, and decay = 0.
Evaluation protocol: All models train for a fixed number of epochs, with results reported from the best validation-set epoch.Competing methods retain their authors’ hyperparameters; replacing their losses with Combo loss leaves network hyperparameters unchanged.

IV. DATASETS

The study evaluates segmentation across whole-body PET, prostate MRI, and left-ventricular ultrasound datasets, using representative multimodal views.

Datasets: The evaluation includes 58 whole-body PET scans, 958 MRI prostate scans, and ultrasound echocardiography images for left ventricular myocardial segmentation.The PET data use 10 test volumes and 48 training volumes; the MRI data use 700 training volumes and a separate test set; ultrasound uses 430 training and 20 test images.
Dataset examples: Figure 2 shows left ventricular myocardium in ultrasound, prostate in MRI, and coronal whole-body PET views.The MRI and ultrasound rows show axial and coronal views.

V. RESULTS

The results compare the proposed method with competing loss-function and network variants on PET, MRI, and ultrasound segmentation tasks.

Evaluation design: The evaluation has two parts: PET multi-organ segmentation followed by single-organ segmentation using MRI and ultrasound.The study compares PCE, PWCE, PD, PDCE, and PCombo with the proposed architecture.
PET evaluation: PET experiments compare competing methods and variants of the proposed loss both qualitatively and quantitatively.DCE integrates Dice with traditional cross entropy, whereas Combo combines Dice with weighted cross entropy.

A. Performance of the proposed vs. competing methods on multi-organ PET segmentation

On whole-body PET, the proposed architecture with Combo loss outperforms competing methods and improves several baseline networks, while reducing organ misclassification and false positives.

Overall PET performance: 57% ± 24% Jaccard, 38% ± 18% Dice, and 86% ± 5% FPR are reported for the proposed architecture with Combo loss.The multi-class case is treated as binary by flattening one-hot predicted and ground-truth volumes before computing Combo loss.
Baseline improvements: 3D V-Net improves by 5.8% Jaccard, 4.5% Dice, and 28% FPR with Combo loss.The reported improvement is based on the comparison of rows in section b of Table III.
Baseline improvements: 3D SegNet improves by 34.1% Jaccard, 23.2% Dice, 44% FPR, and 12.5% FNR with Combo loss.Comparing PCE with PWCE shows that weighted cross entropy helps, while comparing PD with PCombo shows further improvement from Combo loss.
Qualitative and organ-wise results: The proposed method reduces organ misclassification and false positives, and its organ-wise Dice scores range from 0.58 to 0.91.It also handles multiple missing organs to a high extent, although some missing PET organs may result from insufficient radiotracer uptake.

B. Performance of the proposed vs. competing methods on single organ segmentation from MRI and ultrasound

The Combo loss improved competing 3D U-Net and 3D V-Net models on MRI prostate and ultrasound left-ventricular segmentation, with gains in Dice and false-negative rates. It also outperformed Focal loss and reduced result variability.

4.6% and 1.13% Dice improvements were achieved for 3D U-Net and 3D V-Net, respectively, on MRI prostate segmentation.
8.23% and 3.4% Dice improvements were achieved for 3D U-Net and 3D V-Net, respectively, on ultrasound left-ventricular myocardial segmentation.
43.8% and 16.7% reductions in FNR were reported for 3D U-Net and 3D V-Net on MRI prostate segmentation.
33.3% and 16.7% reductions in FNR were reported for 3D U-Net and 3D V-Net on ultrasound left-ventricular myocardial segmentation.
The loss reduced segmentation-result variance and produced smaller Dice and Hausdorff-distance ranges with fewer outliers after application to competing methods.
Combo loss outperformed Focal loss by 2.3% in Dice and 47.4% in FNR for MRI, and by 10.8% in Dice for ultrasound.
Replacing U-Net cross entropy with Combo increased Dice from 0.87 to 0.91 for MRI and from 0.85 to 0.92 for ultrasound.

VI. CONCLUSION

The paper proposes Combo loss to address input/class imbalance and output imbalance by controlling the false-positive/false-negative trade-off. Across modalities and organ settings, it reports improved segmentation performance and existing-model results.

Combo loss is a curriculum-learning-based function designed to handle input/class imbalance and output imbalance by enforcing a false-positive/false-negative trade-off.
The loss sharply decreases FPR or FNR when models are penalized for false positives or false negatives, respectively, across PET, MRI, and ultrasound settings.
The proposed loss improved multi-organ and single-organ segmentation across different imaging modalities, including higher Dice and lower false-positive and false-negative rates.
The authors state that Combo loss can also be leveraged for other segmentation tasks beyond the multi-organ setting studied here.