Domain Generalizer: A Few-shot Meta Learning Framework for Domain Generalization in Medical Imaging

In biomedical imaging, deep learning models trained on one dataset are often hard to generalize to other related datasets. Generally, biomedical images can be represented as points on a high-dimensional non-linear manifold. Failure of segmentation and classification algorithms to generalize across imaging modalities, patients, image acquisition protocols, medical centers, healthy and diseased populations, age, etc., can be explained by significant differences in the statistical distributions of datasets on the image manifolds, known as covariate shift [2]. Addressing covariate shift by retraining deep learning models on each new data domain is impractical in most applications because of the scarcity of expert labeled data. Therefore, it is important to develop deep learning methods that generalize well to new related datasets not seen during training using few or no annotated examples from the new dataset. Domain adaptation [3] and domain generalization [4] paradigms aim at reducing the covariate shift between the training and test distributions by learning domain invariant features. Domain adaptation learns a feature representation that is invariant to the statistics of the source and target domains, and is discriminative enough for the actual learning task. Domain adaptation could either be unsupervised or semi-supervised. Domain generalization, a relatively less studied and harder problem, trains models using a variety of source domains to learn a generic feature representation which should perform well on unseen target domains. This flavor of transfer learning does not use any samples from the target distribution during training. Relatedly, few-shot learning is a paradigm which adapts a trained model to a completely new data distribution with very limited labeled training examples [5].

The biomedical imaging community has witnessed several applications of domain adaptation and few-shot learning. Adaptation across different medical centers is a known challenge in image segmentation [6], and has been achieved through both unsupervised [7], and supervised approaches [8]. Cross-modality domain adaptation methods between magnetic resonance (MR) and CT images have been proposed using variational autoencoders [9] for whole heart segmentation, and CycleGANs for segmentation of the prostate [10], [11]. Decision forests have been employed to adapt between in-vivo and in-vitro images [12] for intravascular ultrasound tissue segmentation. A few-shot network [13] was proposed to segment multiple organs in MR images. However, a priori knowledge of the unseen test domain is not always available, which hinders model generalizability as discussed above.

Very recently, some groups explored domain generalization for biomedical imaging. (1) A series of nine data augmentation techniques were applied to the training domains to mimic the test distribution [14] for heart ultrasound, heart and prostate MR image segmentation. (2) An episodic training-based meta-learning method [15] was applied to segment brain tissue in T1-weighted MRI across four medical centers. (3) A variational auto-encoder [16] was used to learn three latent subspaces to generalize across patients for a 2D cell segmentation task via domain disentanglement. These methods have certain limitations respectively: (1) In [14], the training data is very large and heavily augmented, which might not be the general case for many problems in medical imaging where the goal is to extract enough information from very limited data for domain generalization; the method is not tested on diseased populations which could vary significantly in anatomies and shapes from a generic healthy population. (2) In [15], again, the training set contains similar anatomies in both the train and test sets; the average performance is quite marginal than the compared baseline, with an improvement of only around 0.8% in Dice score (Baseline: 90.6% vs proposed: 91.4%); does not evaluate on cases with atrophied or irregular brain anatomies. (3) In [16], the method uses domain labels as an additional cue, which we argue is difficult to define precisely, due to its wide range of interpretation; the average performance is not significant with an improvement of 0.4% in average Dice score (Baseline: 95.4% vs proposed: 95.8%); the method evaluates only one 2D dataset, not extending to the 3D case; each patient is considered as different domain acquired at the same medical center, and hence the test and train sets might have similar statistical distributions.

Contributions. In the present work, we extend a gradient-based meta-learning domain generalization method (MLDG) [1] that has shown promise for image classification tasks to the context of biomedical image segmentation, termed MLDG-Seg. To evaluate this approach, we focus on the problem of vertebrae segmentation in CT images and utilize three publicly available databases. To address the above-mentioned drawbacks of existing domain generalization methods, we construct three domain generalization contexts: (a) generalization to new anatomies: the vertebrae are divided into four domains: lumbar, lower, middle, and upper thoracic regions. The model is trained on three domains, and then tested on the unseen fourth domain. (b) generalization to a diseased population with fractured vertebrae, dislocated discs: the model is trained on a healthy population, then tested on unseen data of a diseased population from another medical center. (c) generalization to unseen anatomies, surgical implants, different acquisition protocols, arbitrary orientations, and field of view (FoV): the model is tested on a very large dataset. Through these three contexts, we show that MLDG-Seg is able to learn generalized representations from very limited training examples. Finally, we show that the learned generalized representation can be quickly adapted with a few examples from the unseen target distribution in a k-shot learning setting to achieve additional performance gains.

Meta-learning, or learning to learn [17], aims at learning a variety of tasks, and then quickly adapting to new tasks in different settings. We adopt the optimization-based model-agnostic method proposed in [1], called Meta Learning Domain Generalization (MLDG). Here, we briefly describe the method.

Description. Let there be two distributions: source $\mathcal{S}$, and target $\mathcal{T}$. Both $\mathcal{S}$, and $\mathcal{T}$ share the same task, for example, segmentation or classification with the same label space. The goal of MLDG is to learn a single set of model parameters $\mathcal{\theta}$ via gradient descent and two meta-learners: meta-train and meta-test procedures. The model is trained on only the source $\mathcal{S}$ domains, and then tested on target $\mathcal{T}$ domain. The source domains $\mathcal{S}$ are split into two sets: meta-train domains $\hat{S}$, and meta-test domains $\bar{S}$ = $\mathcal{S}$ - $\hat{S}$. The goal of the two splits is to mimic the setting of domain shifts, and thereby make it easier for the model to generalize on an unseen target domain $\mathcal{T}$. We reproduce the learning procedure in Algorithm 1, and show the extended version in Fig. 1.

Explanation. Let’s consider a motivating example (carried out later in Experiment 1) from the image manifold of lumbar ($\mathcal{D}$1), lower thoracic ($\mathcal{D}$2) and middle thoracic regions ($\mathcal{D}$3), comprising the set of source domains $\mathcal{S}$; as well as upper ($\mathcal{D}$4) thoracic region, which constitutes the unseen test $\mathcal{T}$ domain. Refer to Algorithm 1. The MLDG method is supposed to learn a single model parameter $\mathcal{\theta}$ with the help of two optimization steps. At every iteration, the set of images in the source domains, here ($\mathcal{D}$1, $\mathcal{D}$2, and $\mathcal{D}$3) are randomly split into a meta-train (for example, consisting of images from $\mathcal{D}$1, $\mathcal{D}$2 set), and meta-test with the set $\mathcal{D}$3. Now, two losses are computed. The first loss $\mathcal{F}$ is computed using the training examples from meta-train set and the gradient is computed with respect to the model parameter $\theta$. The second loss $\mathcal{G}$ is computed on the meta-test set with the updated parameter $\theta^{{}^{\prime}}$ = $\theta-\alpha\nabla_{\theta}$. The key idea by introducing the second loss in the meta-test stage is that an improvement of the model’s performance on the meta-train set should also improve the model’s performance on the meta-test set. The final model parameter $\mathcal{\theta}$ is updated by taking the gradient of the weighted combination of the two losses $\mathcal{F}$, and $\mathcal{G}$. By doing so, the model is tuned in such a way that performance is improved in both meta-train, and meta-test domains. In other words, the model is regularized and does not overfit to one particular domain, by finding the best possible gradient direction due to the joint optimization of the two losses. Compare this to a “vanilla” setup, where a model is directly given images from the three domains $\mathcal{D}$1, $\mathcal{D}$2, and $\mathcal{D}$3 without any meta-learning setup, might overfit to a domain by minimizing its loss, and maximizing the loss for the other domains.

Experiment 1. Table 1 tabulates the results on different held-out test distributions. Each model is trained on three domains and then tested on the fourth unseen domain, except the oracle which is trained and tested on the same domain. MLDG-Seg consistently outperforms the baseline on both Dice and ASSD, and shows the desired low variance amongst the subjects. With a very few labeled examples from the test distribution, we see a further boost in performance. Fig. 2C shows that the MLDG-Seg is better able to segment the region of interest (ROI) by having significantly reduced background error than the baseline, in addition to the correct delineation of intervertebral discs (IVDs) when the held-out distribution is the middle thoracic region. A further improvement in performance is obtained using additional k-shot examples over MLDG-Seg. Ideally, the oracle should have the best performance than the rest of the procedures as the test domain distribution is similar to the training domain distribution. Here in this particular experimental setup, it is not surprising to see that the MLDG-Seg (and in some test domains, the baseline) performs better than the oracle. This might be due to the limited number of training (and validation) examples available to train the oracle. Furthermore, since all the procedures were trained for a fixed number of epochs, the oracle might have had a disadvantage of being trained for a lesser number of gradient steps than the MLDG-Seg procedure. An alternative approach would be to train a single oracle model which learns from all the domains simultaneously, and then test on each domain separately.

Experiment 2. Table 2 and the Spaghetti plots in Fig. 3 shows that the MLDG-Seg outperforms the baseline. Fig. 2B shows that MLDG-Seg is able to delineate the vertebrae by not segmenting the undesired spinal cord, as incorrectly segmented by the baseline. The k-shot setting further improves the segmentation. Here again, it is not necessarily surprising to see that MLDG-Seg performs better than the oracle. The oracle was trained with a smaller number of subjects than the baseline and MDG-Seg (see supplement).

Experiment 3. Table 3 shows the improved performance of MLDG-Seg over the baseline, and consistent performance boost in the few-shot learning regime. The k-shot procedure is either on par with the oracle or outperforms the same for k=4, 5 and 6. Fig. 2A depicts the superior performance of MLDG-Seg over the baseline on a compression fraction subject from a different distribution than the training set, where MLDG-Seg is able to segment more vertebrae completely than the baseline. The k-shot setting further improves the performance by segmenting the remaining vertebrae not segmented by baseline or MLDG-Seg. Here, the oracle performs better than most of the procedures, which is due to the fact that the training set for oracle is relatively larger than in experiments 1 and 2.

Conclusion In the present work, we benchmarked the performance of a gradient-based meta-learning domain generalization segmentation method in the context of biomedical image analysis, across a variety of training and test settings. The method was not only able to generalize across multiple medical sites and scanners, which is the most widely studied problem of generalization, but was also able to generalize to newly introduced settings of unseen complex vertebrae anatomies, surrounding inter-vertebral discs space, varying bone and soft tissue intensities distribution, diseased populations, different acquisition protocols, arbitrary orientations, and FoVs, thus resembling actual clinical settings. In future, we will evaluate the method on other modalities such as MRI, and US and compare with the recently proposed domain generalization methods. Our source code, scripts, and dataset-split files are available at:
https://github.com/Pulkit-Khandelwal/medical-mldg-seg.

Acknowledgments This work was supported by NIH grant R01 EB017255.

In the following set of figures, Spaghetti plots are shown for all the experiments. Y-axis shows the Dice (%), and the ASSD (mm) scores for the different procedures, shown on the X-axis. Each of the plotted lines denote a subject. Therefore, one can track the performance of a procedure for a given subject.