# Generalized Wasserstein Dice Loss, Test-time Augmentation, and Transformers for the BraTS 2021 challenge

arXiv: https://arxiv.org/abs/2112.13054 · HTML: https://mlsanity.com/paper/2112.13054/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 31 references that are in the map.
- First posted on arXiv 24 Dec 2021.

## Builds on (most cited first, 20 of 20)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation](https://mlsanity.com/paper/1606.06650/) (2016)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Instance Normalization: The Missing Ingredient for Fast Stylization](https://mlsanity.com/paper/1607.08022/) (2016)
- [Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations](https://mlsanity.com/paper/1707.03237/) (2017)
- [The RSNA-ASNR-MICCAI BraTS 2021 Benchmark on Brain Tumor Segmentation and Radiogenomic Classification](https://mlsanity.com/paper/2107.02314/) (2021)
- [Identifying the Best Machine Learning Algorithms for Brain Tumor Segmentation, Progression Assessment, and Overall Survival Prediction in the BRATS Challenge](https://mlsanity.com/paper/1811.02629/) (2018)
- [ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks](https://mlsanity.com/paper/2102.11600/) (2021)
- [Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks](https://mlsanity.com/paper/1807.07356/) (2018)
- [On the Compactness, Efficiency, and Representation of 3D Convolutional Networks: Brain Parcellation as a Pretext Task](https://mlsanity.com/paper/1707.01992/) (2017)
- [Are we using appropriate segmentation metrics? Identifying correlates of human expert perception for CNN training beyond rolling the DICE coefficient](https://mlsanity.com/paper/2103.06205/) (2021)
- [Generalised Wasserstein Dice Score for Imbalanced Multi-class Segmentation using Holistic Convolutional Networks](https://mlsanity.com/paper/1707.00478/) (2017)
- [Neural parameters estimation for brain tumor growth modeling](https://mlsanity.com/paper/1907.00973/) (2019)
- [Distributionally Robust Segmentation of Abnormal Fetal Brain 3D MRI](https://mlsanity.com/paper/2108.04175/) (2021)
- [Generalized Wasserstein Dice Score, Distributionally Robust Deep Learning, and Ranger for brain tumor segmentation: BraTS 2020 challenge](https://mlsanity.com/paper/2011.01614/) (2020)
- [Comparative study of deep learning methods for the automatic segmentation of lung, lesion and lesion type in CT scans of COVID-19 patients](https://mlsanity.com/paper/2007.15546/) (2020)
- [Partial supervision for the FeTA challenge 2021](https://mlsanity.com/paper/2111.02408/) (2021)

## Built on by (most cited first, 5 of 5)

- [Deep Quality Estimation: Creating Surrogate Models for Human Quality Ratings](https://mlsanity.com/paper/2205.10355/) (2022)
- [Average Calibration Error: A Differentiable Loss for Improved Reliability in Image Segmentation](https://mlsanity.com/paper/2403.06759/) (2024)
- [Average Calibration Losses for Reliable Uncertainty in Medical Image Segmentation](https://mlsanity.com/paper/2506.03942/) (2025)
- [Label tree semantic losses for rich multi-class medical image segmentation](https://mlsanity.com/paper/2507.15777/) (2025)
- [SegGuidedNet: Sub-Region-Aware Attention Supervision for Interpretable Brain Tumor Segmentation](https://mlsanity.com/paper/2605.22572/) (2026)

## Abstract (arXiv)

Brain tumor segmentation from multiple Magnetic Resonance Imaging (MRI) modalities is a challenging task in medical image computation. The main challenges lie in the generalizability to a variety of scanners and imaging protocols. In this paper, we explore strategies to increase model robustness without increasing inference time. Towards this aim, we explore finding a robust ensemble from models trained using different losses, optimizers, and train-validation data split. Importantly, we explore the inclusion of a transformer in the bottleneck of the U-Net architecture. While we find transformer in the bottleneck performs slightly worse than the baseline U-Net in average, the generalized Wasserstein Dice loss consistently produces superior results. Further, we adopt an efficient test time augmentation strategy for faster and robust inference. Our final ensemble of seven 3D U-Nets with test-time augmentation produces an average dice score of 89.4% and an average Hausdorff 95% distance of 10.0 mm when evaluated on the BraTS 2021 testing dataset. Our code and trained models are publicly available at https://github.com/LucasFidon/TRABIT_BraTS2021.