Generalised Wasserstein Dice Score for Imbalanced Multi-class Segmentation using Holistic Convolutional Networks (1707.00478v4)

Abstract: 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.

• The paper introduces a generalised Wasserstein Dice Score that leverages inter-class relationships to enhance multi-class segmentation.
• It employs a holistic CNN with deep supervision and multi-scale predictions to capture both fine boundary details and contextual features.
• Experimental results on the BraTS'15 dataset show reduced semantically implausible misclassifications compared to traditional Dice loss.

Generalised Wasserstein Dice Score for Imbalanced Multi-class Segmentation using Holistic Convolutional Networks

In this paper, the authors introduce an innovative approach to address the specific challenges of imbalanced multi-class segmentation in convolutional neural networks (CNNs), particularly focusing on automatic brain tumor segmentation. Traditional methods, while effective in binary segmentation, fall short when it comes to leveraging inter-class relationships and multi-scale information necessary for a nuanced understanding of complex anatomical structures such as brain tumors. The paper presents a framework that combines a semantically-informed loss function with a novel holistic CNN architecture to address these challenges comprehensively.

The main contributions of the paper are twofold. Firstly, the introduction of a generalised Wasserstein Dice Score, which builds on the classical Dice loss used for binary segmentation but extends it to multi-class segmentation through a semantic reasoning approach. By utilizing the Wasserstein distance, this new metric incorporates inter-class relationships by introducing a distance matrix that contextualizes the costs associated with misclassifying a voxel from one class to another. This is particularly relevant in medical imaging where certain classes share more anatomical features or proximity, and therefore, misclassifying them incurs different penalties compared to others.

Secondly, the authors introduce a holistic CNN architecture which addresses the limitations of traditional CNNs in capturing multi-scale spatial information. The proposed architecture leverages deep supervision and multi-scale prediction to enhance the segmentation performance. By combining predictions from different scales, the model not only captures high-resolution boundary details but also encodes contextual information crucial for semantic labeling. This is achieved by incorporating residual connections, batch normalization, and exponentially linear units (ELUs) to bolster model robustness and convergence speed.

The robustness of this approach is empirically validated using the BraTS'15 dataset focused on brain tumor segmentation, which is known for high-grade and low-grade gliomas. By comparing their Wasserstein Dice loss with the mean Dice score, the authors demonstrate significant improvements in capturing the underlying spatial and hierarchical relationships among tumor classes. Notably, the evaluation using class confusion matrices indicates a reduction in semantically implausible misclassifications, showcasing the enhanced capabilities of the proposed method in handling intricate multi-class segmentation tasks.

Furthermore, the paper discusses the impact of training the holistic CNN with different choices of the distance matrix needed for the Wasserstein calculation. The results suggest that encoding semantic relationships through hierarchical structures (as in $M_{tree}$) provides superior segmentation outcomes compared to a naive implementation without inter-class considerations ($M_{0-1}$).

In terms of future implications, this framework deviates from purely class-count based metrics, setting a new direction for employing more nuanced, contextually aware metrics in medical imaging and beyond. The research explores the feasible transition towards end-to-end training on full volumetric datasets as GPU capacities continue to increase, potentially fracturing the limitations imposed by current patch-based methods.

Moreover, the discussion points towards the sustainability of refining loss functions by incorporating learned distance matrices, seeks to address the computational overheads involved in more complex ground distance matrices. This has potential ramifications for broadening the scope of applications into other domains featuring class imbalance and intricate multi-label requirements.

Overall, the paper provides a constructive methodological shift towards more intelligent and context-aware automated segmentation tools, which could cascade into other fields requiring intricate contextual understanding and segmentation accuracy in imbalanced data scenarios.