# ROBUST-WT: Robust Uncertainty-aware Segmentation Transform via Whitening and Training Enhancements

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 36 references that are in the map.
- First posted on arXiv 2 Jun 2026.

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

- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [Statistical Aspects of Wasserstein Distances](https://mlsanity.com/paper/1806.05500/) (2018)
- [Comparing Kullback-Leibler Divergence and Mean Squared Error Loss in Knowledge Distillation](https://mlsanity.com/paper/2105.08919/) (2021)
- [Uncertainty Modeling for Out-of-Distribution Generalization](https://mlsanity.com/paper/2202.03958/) (2022)
- [Domain Generalization for Medical Imaging Classification with Linear-Dependency Regularization](https://mlsanity.com/paper/2009.12829/) (2020)
- [DoFE: Domain-oriented Feature Embedding for Generalizable Fundus Image Segmentation on Unseen Datasets](https://mlsanity.com/paper/2010.06208/) (2020)
- [Modeling Uncertainty with Hedged Instance Embedding](https://mlsanity.com/paper/1810.00319/) (2018)
- [Embracing the Dark Knowledge: Domain Generalization Using Regularized Knowledge Distillation](https://mlsanity.com/paper/2107.02629/) (2021)
- [Shape-Aware Complementary-Task Learning for Multi-Organ Segmentation](https://mlsanity.com/paper/1908.05099/) (2019)
- [Information Geometry Connecting Wasserstein Distance and Kullback-Leibler Divergence via the Entropy-Relaxed Transportation Problem](https://mlsanity.com/paper/1709.10219/) (2017)
- [Prior Attention Network for Multi-Lesion Segmentation in Medical Images](https://mlsanity.com/paper/2110.04735/) (2021)

## Abstract (arXiv)

Generalized segmentation of medical images prevents performance degradation when different imaging devices and clinical protocols are used across multiple domains. The Whitening Transform-based Probabilistic Shape Regularization Extractor (WT-PSE), published in IEEE Transactions on Medical Imaging in 2024, addresses this challenge by employing feature decorrelation and Wasserstein distance-based knowledge distillation to achieve robust cross-domain segmentation. This study systematically examines improvements to the WT-PSE learning framework. Four limitations in the original implementation are identified: limited training augmentations that fail to simulate real scanner variations, reliance on per-pixel binary cross-entropy loss that is sensitive to edge noise, the absence of a scheduled loss weighting strategy that may destabilize early training, and the lack of ablation switches for controlled scientific comparison. To address these issues, we propose four enhancements: (1) domain-adaptive augmentation including random erasing, gamma correction, and salt-and-pepper noise; (2) a hybrid BCE and Dice loss function for improved edge-aware segmentation under noisy conditions; (3) a curriculum-based Dice weight scheduling strategy; and (4) command-line control flags for systematic ablation studies. Experiments on the fundus optic disc segmentation benchmark demonstrate that the improved pipeline achieves a final epoch optic-disc Dice score of 0.956 and an ASD score of 13.31, outperforming the baseline epoch-5 Dice score of 0.939. These results indicate that training-level improvements can provide consistent performance gains without modifying the underlying WT-PSE architecture.