Shape-aware Meta-learning for Generalizing Prostate MRI Segmentation to Unseen Domains

Quande Liu,Qi Dou,Pheng-Ann Heng
DOI: https://doi.org/10.48550/arXiv.2007.02035
2020-07-04
Abstract:Model generalization capacity at domain shift (e.g., various imaging protocols and scanners) is crucial for deep learning methods in real-world clinical deployment. This paper tackles the challenging problem of domain generalization, i.e., learning a model from multi-domain source data such that it can directly generalize to an unseen target domain. We present a novel shape-aware meta-learning scheme to improve the model generalization in prostate MRI segmentation. Our learning scheme roots in the gradient-based meta-learning, by explicitly simulating domain shift with virtual meta-train and meta-test during training. Importantly, considering the deficiencies encountered when applying a segmentation model to unseen domains (i.e., incomplete shape and ambiguous boundary of the prediction masks), we further introduce two complementary loss objectives to enhance the meta-optimization, by particularly encouraging the shape compactness and shape smoothness of the segmentations under simulated domain shift. We evaluate our method on prostate MRI data from six different institutions with distribution shifts acquired from public datasets. Experimental results show that our approach outperforms many state-of-the-art generalization methods consistently across all six settings of unseen domains.
Computer Vision and Pattern Recognition
What problem does this paper attempt to address?
This paper attempts to address the issue of insufficient model generalization in prostate MRI segmentation, especially when facing different domains (such as different imaging protocols and scanners). Specifically, the authors propose a novel Shape-Aware Meta-Learning (SAML) scheme aimed at improving the model's generalization ability in unseen target domains. The paper mainly focuses on two aspects: 1. **Shape Compactness**: By introducing a shape compactness loss function, the model ensures that the segmentation results generated in unseen domains have complete shapes. 2. **Boundary Smoothness**: By introducing a boundary smoothness loss function, the model ensures that the segmentation results generated in unseen domains have clear boundaries. These methods optimize the model during meta-training and meta-testing by simulating domain transfer, thereby improving the model's generalization performance. Experimental results show that this method outperforms several existing state-of-the-art methods across multiple datasets.