Robust White Matter Hyperintensity Segmentation on Unseen Domain

TL;DR

This paper tackles the challenge of domain generalization in medical imaging, specifically for white matter hyperintensity segmentation, by combining domain adversarial learning and mix-up techniques to improve performance on unseen datasets.

Contribution

It introduces a novel approach that synergizes domain adversarial learning and mix-up for robust WMH segmentation across unseen domains.

Findings

Significant performance improvements on unseen target domains.

Effective combination of domain adversarial learning and mix-up.

Enhanced generalization in multi-site medical imaging datasets.

Abstract

Typical machine learning frameworks heavily rely on an underlying assumption that training and test data follow the same distribution. In medical imaging which increasingly begun acquiring datasets from multiple sites or scanners, this identical distribution assumption often fails to hold due to systematic variability induced by site or scanner dependent factors. Therefore, we cannot simply expect a model trained on a given dataset to consistently work well, or generalize, on a dataset from another distribution. In this work, we address this problem, investigating the application of machine learning models to unseen medical imaging data. Specifically, we consider the challenging case of Domain Generalization (DG) where we train a model without any knowledge about the testing distribution. That is, we train on samples from a set of distributions (sources) and test on samples from a new, unseen distribution (target). We focus on the task of white matter hyperintensity (WMH) prediction using the multi-site WMH Segmentation Challenge dataset and our local in-house dataset. We identify how two mechanically distinct DG approaches, namely domain adversarial learning and mix-up, have theoretical synergy. Then, we show drastic improvements of WMH prediction on an unseen target domain.