Improving cross-domain generalizability of medical image segmentation using uncertainty and shape-aware continual test-time domain adaptation
Jiayi Zhu, Bart Bolsterlee, Yang Song, Erik Meijering
Abstract
Continual test-time adaptation (CTTA) aims to continuously adapt a source-trained model to a target domain with minimal performance loss while assuming no access to the source data. Typically, source models are trained with empirical risk minimization (ERM) and assumed to perform reasonably on the target domain to allow for further adaptation. However, ERM-trained models often fail to perform adequately on a severely drifted target domain, resulting in unsatisfactory adaptation results. To tackle this issue, we propose a generalizable CTTA framework. First, we incorporate domain-invariant shape modeling into the model and train it using domain-generalization (DG) techniques, promoting target-domain adaptability regardless of the severity of the domain shift. Then, an uncertainty and shape-aware mean teacher network performs adaptation with uncertainty-weighted pseudo-labels and shape information. As part of this process, a novel uncertainty-ranked cross-task regularization scheme is proposed to impose consistency between segmentation maps and their corresponding shape representations, both produced by the student model, at the patch and global levels to enhance performance further. Lastly, small portions of the model’s weights are stochastically reset to the initial domain-generalized state at each adaptation step, preventing the model from ‘diving too deep’ into any specific test samples. The proposed method demonstrates strong continual adaptability and outperforms its peers on five cross-domain segmentation tasks, showcasing its effectiveness and generalizability. • Deep segmentation models trained with ERM can be difficult to adapt to drifted target domains. • Existing adaptation methods struggle with moving target domains due to catastrophic forgetting. • We use a domain-robust stochastic weight reset scheme to alleviate catastrophic forgetting. • Shape & uncertainty awareness and multi-level cross-task loss are used for online adaptation. • Experiments on five cross-domain segmentation datasets showcase the efficacy of our method.