Unsupervised domain adaptation for brain structure segmentation via mutual information maximization alignment

Abstract

Brain structure segmentation in Magnetic Resonance image (MRI) plays an important role in the diagnosis of various neuropsychiatric diseases. Recently, the high cost of manual annotation has motivated research in medical image analysis to focus on unsupervised learning methods. As a promising solution, unsupervised domain adaptation (UDA) has received increasing attention in various medical image segmentation tasks. In this work, we propose a novel UDA framework for brain structure segmentation from MRI volumes. Specifically, we design a self-training-based network architecture to learn cross-domain representations, which is achieved by transferring essential knowledge from a well-trained cumbersome teacher network to a non-trained compact student network using distillation learning technology. The pseudo-labels generated by the pre-trained student model are used to supervise the domain adaptation process of the teacher model. Furthermore, we develop a mutual information maximization alignment module to learn domain-invariant features and cross-domain segmentation knowledge, in which contrastive learning techniques are utilized to maximize the correlation between samples. Finally, we build a unified distillation framework by integrating mutual information distillation and output map distillation, correcting pseudo-labels during the self-training stage, and continuously refining the final predictions of compact networks. Extensive experiments conducted on two datasets with vastly different styles and settings demonstrate that our method is effective in improving segmentation performance on unlabeled target images, and outperforms the state-of-the-art domain adaptation approaches by a large margin. The source code is available at https://github.com/huqian999/UDA-MIMA.