Harmonization and Targeted Feature Dropout for Generalized Segmentation: Application to Multi-site Traumatic Brain Injury Images

Abstract

While learning based methods have brought extremely promising results in medical imaging, a major bottleneck is the lack of generalizability. Medical images are often collected from multiple sites and/or protocols for increasing statistical power, while CNN trained on one site typically cannot be well-transferred to others. Further, expert-defined manual labels for medical images are typically rare, making training a dedicated CNN for each site unpractical, so it is important to make best use of the limited labeled source data. To address this problem, we harmonize the target data using adversarial learning, and propose targeted feature dropout (TFD) to enhance the robustness of the model to variations in target images. Specifically, TFD is guided by attention to stochastically remove some of the most discriminative features. Essentially, this technique combines the benefits of attention mechanism and dropout, while it does not increase parameters and computational costs, making it well-suited for small neuroimaging datasets. We evaluated our method on a challenging Traumatic Brain Injury (TBI) dataset collected from 13 sites, using labeled source data of only 14 healthy subjects. Experimental results confirmed the feasibility of using the Cycle-consistent adversarial network for harmonizing multi-site MR images, and demonstrated that TFD further improved the generalization of the vanilla segmentation model on TBI data, reaching comparable accuracy with that of the supervised learning. The code is available at https://github.com/YilinLiu97/Targeted-Feature-Dropout.git.