Source-linked AI summary
Generalised Wasserstein Dice Score for Imbalanced Multi-class Segmentation using Holistic Convolutional Networks
Lucas Fidon, Wenqi Li, Luis C. Garcia-Peraza-Herrera, Jinendra Ekanayake, Neil Kitchen, Sebastien Ourselin, Tom Vercauteren
TL;DR
Multi-class brain tumour segmentation must handle severe class imbalance while accounting for semantic relationships between labels and spatial information at multiple scales. The paper combines a Wasserstein-based generalised Dice score with a holistic CNN using multi-scale predictions and deep supervision, and reports more semantically meaningful segmentations than mean Dice training and classic CNN architectures. The approach is evaluated on BraTS brain tumour segmentation data.
Problem
Existing multi-class Dice formulations do not fully use inter-class relationships or multi-scale information, despite these factors being relevant to imbalanced brain tumour segmentation.
Method
The paper combines a Wasserstein-based generalisation of the Dice score with a holistic CNN that uses multi-scale prediction and deep supervision.
Results
The joint use of the generalised Wasserstein Dice score and Holistic CNN achieves better generalisation and more semantically meaningful brain tumour segmentations than mean soft Dice training and classic CNN architectures.
Takeaways & Limitations
The study supports using inter-class relationships and multi-scale information for imbalanced multi-class brain tumour segmentation.
Takeaways & Limitations
The tree-based distance matrix was chosen heuristically, and learning it directly from data or improving its heuristics may yield further improvements.
Abstract
from arXiv · showhide
The Dice score is widely used for binary segmentation due to its robustness to class imbalance. Soft generalisations of the Dice score allow it to be used as a loss function for training convolutional neural networks (CNN). Although CNNs trained using mean-class Dice score achieve state-of-the-art results on multi-class segmentation, this loss function does neither take advantage of inter-class relationships nor multi-scale information. We argue that an improved loss function should balance misclassifications to favour predictions that are semantically meaningful. This paper investigates these issues in the context of multi-class brain tumour segmentation. Our contribution is threefold. 1) We propose a semantically-informed generalisation of the Dice score for multi-class segmentation based on the Wasserstein distance on the probabilistic label space. 2) We propose a holistic CNN that embeds spatial information at multiple scales with deep supervision. 3) We show that the joint use of holistic CNNs and generalised Wasserstein Dice scores achieves segmentations that are more semantically meaningful for brain tumour segmentation.
1 Introduction
Brain tumour segmentation faces severe class imbalance, while existing Dice-based training does not fully use label-space structure or multi-scale spatial information. The paper addresses both gaps with a Wasserstein-based Dice generalisation and a holistic CNN.
- Motivation: Severe class imbalance is a central challenge in automatic brain tumour segmentation.Weighted loss functions appear promising because fully convolutional training favors entire subject volumes over constrained patch sampling.
- Motivation: Soft Dice loss improves robustness to class imbalance by incorporating global spatial information.However, its multi-class formulation does not fully exploit relationships between labels or spatial information across scales.
- Contributions: The proposed generalised Wasserstein Dice score embeds semantic relationships between classes when comparing label probability vectors.It uses Wasserstein distance to incorporate prior structure in the label space.
- Contributions: The proposed holistic CNN embeds spatial information at different scales and introduces deep supervision during training.The joint approach achieves better generalisation than mean soft Dice training and classic CNN architectures for multi-class brain tumour segmentation.
2 A Wasserstein approach for multi-class soft Dice score
The paper extends soft Dice scoring to multi-class segmentation by comparing probabilistic label vectors with Wasserstein distance, allowing inter-class relationships to inform disagreement. The resulting Wasserstein Dice formulation incorporates label-space structure while retaining a loss-function form for training.
- Mean-class Dice scores compare probabilities associated with the same label and do not account for inter-class relationships.
- Wasserstein distance transforms a label-space distance matrix into a distance between probability vectors by minimizing the cost of transporting probability mass.
- The proposed Wasserstein Dice score uses this distance to generalize disagreement between predicted and ground-truth label probability vectors.
- Choosing class weights as the Wasserstein distance to background ensures that background voxels do not contribute to the generalized true-positive term.
- With a binary distance matrix, the Wasserstein distance becomes |pi − gi| and the proposed score reduces to the soft binary Dice score.
- For a crisp ground truth, a closed-form Wasserstein solution makes the loss computationally tractable; the loss is defined as 1 − DM.
3 Holistic convolutional networks for multi-scale fusion
The proposed holistic CNN combines predictions from multiple network scales and applies supervision both to the fused output and to scale-specific predictions. Its architecture uses pooling and dilated convolutions to expand receptive fields while retaining multi-scale spatial information.
- The HCNN produces multi-scale predictions and intermediate supervision, then combines the scale outputs through a fusion layer.
- As receptive fields grow across layers, predictions at different layers encode spatial information at different scales.
- Max pooling and dilated convolutions increase the receptive field and reduce redundancy between successive scale predictions.
- Class-specific fusion weights allow different scales to receive different importance for each class.
- The loss is applied to both the fused prediction and each scale-specific prediction, providing deep supervision during training.
4 Implementation details
The study evaluates HCNN models and Wasserstein Dice losses on the BraTS’15 brain-tumour segmentation dataset, using multimodal scans and specified training procedures.
- Dataset: BraTS’15 includes multimodal T1, T1c, T2, and Flair images from 274 glioma subjects.The dataset contains 220 high-grade and 54 low-grade gliomas.
- Dataset: The dataset is randomly split into 80% training, 10% validation, and 10% testing with matched high-grade and low-grade proportions.Images are skull stripped, resampled to a 1mm isotropic grid, co-registered, and histogram-standardised.
- Labels and evaluation: Segmentation uses five labels: background, necrotic core, edema, non-enhancing core, and enhancing tumour.Evaluation reports Dice scores for whole tumour, core tumour, and enhanced tumour.
- Evaluation: Table 1 evaluates different multi-class Dice scores for HCNN training and testing, including mean-Dice pre-training followed by LDMtree retraining.The pre-training lasts 4 epochs and retraining lasts 85 epochs.
- Training: Networks are trained with ADAM, early stopping, dropout, multimodal subject volumes, and class-maximising patch sampling.The implementation uses lr = 0.01, β1 = 0.9, β2 = 0.999, and dropout probability 0.6.
5 Results
The experiments compare semantically informed Wasserstein Dice losses with mean Dice training and examine multi-scale HCNN predictions. Pretraining with mean Dice followed by tree-based Wasserstein Dice improves semantic plausibility and true-positive performance, while scale-specific predictions expose resolution–receptive-field trade-offs.
- 5.1 Examples of distance metrics on BraTS label space: The study compares mean Dice with two semantically driven distance matrices, including a tumour-hierarchy-based Mtree, for training the HCNN.M0−1 encodes no inter-class relationship, whereas Mtree uses the hierarchical structure of whole, core, and enhancing tumour.
- 5.2 Evaluation and training with multi-class Dice score: Training with mean Dice followed by LDMtree produced the best results for all criteria and more semantically plausible segmentations than mean Dice alone.The network was pretrained for 4 epochs with mean Dice and then retrained with LDMtree; the approach also achieved a higher true-positive rate than either loss alone.
- 5.3 Impact of the Wasserstein Dice loss on class confusion: Pairwise class-confusion evaluation shows that Wasserstein losses balance misclassifications, whereas mean Dice primarily maximizes diagonal true positives.The proposed evaluation computes Dice scores for every ground-truth and predicted label pair.
- 5.3 Impact of the Wasserstein Dice loss on class confusion: Mean Dice training correctly segments most voxels but poorly differentiates edema from core tumour, producing semantically implausible confusions.The edema-related confusion is visible in the corresponding row of Table 2 and in the qualitative predictions.
- 5.3 Impact of the Wasserstein Dice loss on class confusion: LDMtree yields more meaningful confusion but cannot differentiate necrotic core from non-enhancing tumour when trained from random initialization.The authors attribute this difficulty to the additional constraints introduced by the tree-based loss.
- 5.4 Evaluation of deep supervision: Multi-scale results reveal class-specific trade-offs between receptive field and resolution, with fused predictions avoiding implausible regions while preserving performance across tumour regions.Scales 3 and 4 miss the thinnest regions at lower resolution, whereas scale 1 can produce holes; scales 2 to 4 and fused predictions achieve similar scores for several regions.
6 Conclusion and future work
The paper combines a Wasserstein-based, semantically informed Dice loss with a holistic multi-scale CNN for imbalanced multi-class brain tumour segmentation. It reports semantically meaningful segmentation and identifies heuristic distance-matrix selection and computational burden as future concerns.
- The proposed Wasserstein Dice score embeds prior knowledge about inter-class relationships through a distance matrix on the label space.
- The holistic convolutional network uses multi-scale predictions and deep supervision to exploit spatial information at multiple scales.
- Tree-based distance matrices lead to higher Dice scores than discrete distances for the imbalanced multi-class brain tumour segmentation task.
- The tree-based distance matrix was chosen heuristically, so improved heuristics or learning the matrix from data could further improve performance.
- Extending Wasserstein distance to embed spatial information directly in the loss remains computationally challenging.