Improving Brain MRI Segmentation with Multi-Stage Deep Domain Unlearning
摘要
Inter-scanner variability of magnetic resonance imaging has an adverse impact on the diagnostic and prognostic quality of the scans and necessitates the development of models robust to domain shift inflicted by the unseen scanner data. Review of recent advances in domain adaptation showed that efficacy of strategies involving modifications or constraints on the latent space appears to be contingent upon the level and/or depth of supervision during model training. In this paper, we propose a domain adaptation technique based on multi-stage deep unlearning, with specialized training schedule, identified via an ablation study. Building upon the state-of-the-art segmentation framework nnUNet, we employ deep supervision at each convolutional block of the encoder using domain classifier unlearning, applied across multiple stages to suppress domain-related latent features. Validation on public multi-site scanner datasets for two independent tasks of brain lesion segmentation and tissue parcellation showed a notable and consistent increase in overall Dice score (+0.9% and +9.0%, resp.) and sensitivity (+5.5% and +8.6%, resp.) when the proposed model was applied on unseen scanner data, as compared to the baseline nnUNet segmentation model.