Domain Generalization with Small Data

Existing DG methods can be generally categorized into three different streams, namely data augmentation/generation (Yue et al., 2019; Graves, 2011; Zhou et al., 2021), meta-learning (Li et al., 2018; Kim et al., 2021) and feature representation learning (Li et al., 2018; Gong et al., 2019; Xiao et al., 2021). Among these methods, feature representation learning, which aims to explore invariant feature information that can be shared across domains, demonstrates to be a widely adopted method for the problem of DG. For the feature representation learning-based DG method, (Li et al., 2018) proposed to conduct multi-domain alignment in latent space via a multi-domain MMD distance. Gong et al. (2019) leveraged adversarial training to eliminate the domain discrepancy such that domain-invariant representation can be learned in a manifold space. Due to the varieties of imaging protocol (e.g., the choice of image solution for MRI image), device vendors (e.g., Philips or Siemens CT scanners), and patient populations (the race and age group), the acquired imaging data from different medical sites may exist significant domain shift problem (Liu et al., 2021). Dou et al. (2019) proposed a meta-learning framework to perform local and global semantic alignment for medical image classification. A similar design is also adopted by Li et al. (2022) for tissue image classification. Qi et al. (2020) utilized the curriculum learning scheme to transfer the knowledge for histopathological images classification. Li et al. (2020) combined the data augmentation and domain alignment to achieve decent performance on multiple medical data classification tasks. However, these methods may not focus on learning domain-invariant representation on insufficient samples from source domains.

Compared with deterministic models, probabilistic neural networks turns to learn a distribution over model parameters, which can integrate the uncertainty in predictive modeling (Kingma et al., 2015; Gal & Ghahramani, 2016). When the data is insufficient, probabilistic models usually can achieve better generalized performance due to its probabilistic property (as an implicit regularization) (Blundell et al., 2015). In the context of insufficient samples, Bayesian neural network (Neal, 2012) (BNN) with variational inference, a representative probabilistic model, not only can improve predictive accuracy as a classifier (Wilson & Izmailov, 2020), but also can build up the quality of low-dimensional embeddings of insufficient data (Mallick et al., 2021), which is a crucial motivation for this paper. Meanwhile, modern analytical approximation techniques (e.g., Variational inference (Blei et al., 2017), empirical Bayes (Krishnan et al., 2020)) can efficiently infer the posterior distribution of model parameters with stochastic gradient descent method, which can integrate BNN with deterministic DNN conveniently.

In Xiao et al. (2021), the authors proposed to consider the uncertainty of a generalizable model based on BNN, where the distances of positive probabilistic embedding pairs and class distribution are minimized via KL measure. Despite the effectiveness, the dissimilar pairs (i.e., negative pairs) are ignored, which may not benefit feature representation learning. Moreover, they only focused on sample similarity while the distribution information is ignored. Instead, our proposed method comprehensively considers both positive and negative probabilistic embedding pairs via a novel distribution-based contrastive semantic loss. Last but not the least, our proposed method highlights the benefit of the BNN for building up the quality of latent embeddings under insufficient sample scenarios.

Compared with deterministic point embeddings, probabilistic embeddings aim to characterize the data with a distribution. Due to its high robustness and effective representation (Nguyen et al., 2017), probabilistic representation has been applied to several fields, such as video representation learning (Park et al., 2022), image representation learning (Oh et al., 2018), face recognition (Shi & Jain, 2019; Chang et al., 2020), speaker diarization (Silnova et al., 2020), and human pose estimation (Sun et al., 2020). Recently, some researchers further leveraged the probabilistic embeddings to bridge the gap between data modalities (Chun et al., 2021; Neculai et al., 2022; Chun, 2023). For example, Chun et al. (2021) found that the probabilistic representation can lead to a richer embedding space for the challengeable relation reasoning between the images and their captions. These probabilistic embedding-based approaches either inherit the inherent distribution property of data (e.g., the multiple frames of a video) or tackle the one-to-many correspondences through distributional representation. Instead, our proposed method imposes the Bayesian neural network to generate probabilistic embedding. As such, the representative capacity of the data in the small-data regime can be enhanced. Moreover, we devise a novel probabilistic MMD to measure the discrepancy between mixture probabilistic embeddings for domain-invariant learning.

In this section, we introduce an approach to learning domain-invariant representation from a global perspective by minimizing the discrepancy among domains. Among various distribution distance metrics, Maximum Mean Discrepancy (MMD) is widely adopted (Long et al., 2017; Li et al., 2018) which aims to measure the distance between two probability distributions in a non-parametric manner. Specifically, assume that latent embeddings \(\textbf{Z}_{l}\)=\(\{\textbf{z}_{l_{1}},\ldots ,\textbf{z}_{l_{n_{l}}}\}\) and \(\textbf{Z}_{t}\)=\(\{\textbf{z}_{t_{1}},\ldots ,\textbf{z}_{t_{n_{t}}}\}\) are drawn from two unknown distributions \(\mathbb {P}_{l}\) and \(\mathbb {P}_{t}\). The probability measure \(\mathbb {P}\) can be mapped into a reproducing kernel Hilbert space (RKHS) \(\mathcal {H}\) as a element by setting,where a reproducing kernel \(k: \mathcal {X} \times \mathcal {X} \rightarrow \mathbb {R}\) and corresponding feature map \(\phi : \mathcal {X} \rightarrow \mathcal {H}\) are defined. Let the kernel k is characteristic such that the map \(\mu :\mathbb {P}\rightarrow \mu _{\mathbb {P}}\) is injective. In this case the MMD can be defined as the distance \(\Vert \mu _{\mathbb {P}_l}-\mu _{\mathbb {P}_k}\Vert _{\mathcal {H}}\) in \(\mathcal {H}\) between mean embeddings and it can be used as a measure of distance between the distributions \(\mathbb {P}_{l}\) and \(\mathbb {P}_{t}\) (Borgwardt et al., 2006; Gretton et al., 2012). The explicit computation of MMD can be derived by unbiased empirical estimation of mean map (Gretton et al., 2012), i.e.,The idea of using MMD for domain generalization has been explored in several works (e.g., (Li et al., 2018; Hu et al., 2020)).

In the probabilistic framework, instead of the individual latent embeddings \(\textbf{z}_{l_{1},\ldots }\), we have latent probabilistic embeddings \(\Pi _{l_{1}}:= p(\textbf{z}|\textbf{x}_{l_{1}},\textbf{W}),\ldots \). For a source domain \(D_l\), we have the associated distribution over distributions \(\mathbb {P}_{l} = \{\Pi _{l_{1}},\ldots ,\Pi _{l_{n_{l}}}\}\). For this scenario, we propose to extend the existing point-based empirical MMD estimate to a distribution-based empirical probability MMD (P-MMD) estimate. P-MMD utilizes empirical estimation by kernels on distributions to measure the discrepancy between mixture distributions \(\mathbb {P}_{l}\) and \(\mathbb {P}_{t}\) under the probabilistic framework.

Specifically, we first represent latent probabilistic embeddings as elements in RKHS \(\mathcal {H}_k\) using the kernel k, that we coin a level-1 kernel in the sequel, e.g., \(\mu _{\Pi _{l_{1}}}:= \mathbb {E}_{\textbf{z}\sim \Pi _{l_{1}}}[\phi (\textbf{z})]=\mathbb {E}_{\textbf{z}\sim \Pi _{l_{1}}}[k(\textbf{z},\cdot )]\), which is an analog to the Eq. (1). The kernel mean embedding \(\mu _{\Pi _{l_{1}}}\) can be regarded as a new feature map for a variety of tasks (Yoshikawa et al., 2014). Here, to enable non-linear learning on distributions, we introduce a level-2 kernel K (Muandet et al., 2012). Consider a level-1 kernel \(\kappa \) on \(\mathcal {H}\) and its reproducing kernel Hilbert space (RKHS) \(\mathcal {H}_{\kappa }\). Define K aswhere K and its explicit form on kernel mean embeddings \(\kappa \) are p.d. kernels (Berlinet & Thomas-Agnan, 2011). We define a novel probabilistic MMD (P-MMD) empirical estimation method using the level-2 kernel K:In this work, the level-1 and level-2 kernels, k and K, are both Gaussian RBF kernel due to its impressive performance on a limited amount of distribution data (Muandet et al., 2012). Namely, \(K = K_{Gau}(\Pi _{l_{i}}, \Pi _{t_{j}}) = \kappa (\mu _{\Pi _{l_{i}}}, \mu _{\Pi _{t_{j}}})\) can be represented aswhere \(m_{l}\) and \(m_{t}\) are determined by sampling times T. The kernel mean embedding using the level-1 kernel k creates distributions \(\mu (\mathbb {P}_1),\dots ,\mu (\mathbb {P}_N)\) represented by the samples \(\{\mu _{\Pi _{l_1}},\ldots ,\mu _{\Pi _{l_n}}\}\) for \(l=1,\ldots ,N\) respectively in the RKHS \(\mathcal {H}_k\). The underlying strategy of P-MMD is to apply the classic MMD to these distributions (with respect to the kernel \(\kappa \)). To access the effect that the minimization of P-MMD has on the original latent probability distributions across different domains, we recall the following:

In our setting Theorem 1 and Corollary 3.1 along with the fact that k is a characteristic kernel imply the following:

Following Corollary 3.2 we define the following loss function:Corollary 3.2 implies that as Eq. 6 tends to 0 so does the distance between the distributions of means, variances and higher moments of the distributions \(\Pi _l\) associated to points of different domains.

Although we focus on the scenario of insufficient samples, the computational consumption from Eqs. (4) and (5) may be still prohibitive as the calculation of MMD distance between distributions can scale at least quadratically with the increasing of sample size (especially for image segmentation task), i.e., \(O(n^{2})\) in a domain. Here, by following the linear statistic theory of MMD, the unbiased estimate can be derived by drawing pairs from two domains with replacement, i.e., \({\text {P-MMD}}(\mathbb {P}_{l}, \mathbb {P}_{t})^{2} \approx \frac{2}{n_{l}}\sum _{i=1}^{\frac{2}{n_{l}}}[K(\Pi _{l_{2i}}, \Pi _{l_{2i+1}^{\prime }})+K(\Pi _{t_{2i}}, \Pi _{t_{2i+1}^{\prime }})-K(\Pi _{l_{2i}}, \Pi _{t_{2i+1}})-K(\Pi _{l_{2i+1}}, \Pi _{t_{2i}})\), where assuming \(n_{l}=n_{t}\) for simplicity. Borgwardt et al. (2006) gives proof about the unbiased property of the linear statistic of MMD and shows that statistic power does not be sacrificed too much.

To learn domain-invariant representation from a local perspective, a popular idea is to encourage positive pairs with same label closer together, while pulling other negative ones with different labels further apart (Motiian et al., 2017; Dou et al., 2019). These methods usually measure the Euclidean distance between samples in the embedding space. However, this scheme may not satisfy our probabilistic framework due to its probabilistic embeddings.

To this end, we propose a probabilistic contrastive semantic alignment (P-CSA) loss that can utilize the empirical MMD to measure the discrepancy between probabilistic embeddings. The proposed P-CSA loss \(\mathcal {L}_{local}\) consists of two components, including the positive probabilistic contrastive loss and negative probabilistic contrastive loss. The former aims to minimize the distance between the intra-class distributions from different domains, i.e.,where \(M_{\Theta }(\cdot )\) denotes the embedding network of metric learning, which will contribute to learn the distance between features better (Dou et al., 2019). Note that \(\textbf{y}_{n} = \textbf{y}_{q}\) needs to be satisfied. Then, the negative probabilistic contrastive loss is denoted bywhere \(\xi \) is a distance margin that can guarantee an appropriate repulsion range. Note that \(\textbf{y}_{n} \ne \textbf{y}_{q}\) needs to be satisfied.

Model Training Our proposed framework consists of three modules, a BNN-based probabilistic extractor \(Q_{\phi }\), a BNN-based classifier \(C_{\omega }\), and a metric network \(M_{\Theta }(\cdot )\). For the \(Q_{\phi }\), we only add a Bayesian layer with ReLU layer on the bottom of a pretrained deterministic model (e.g., ResNet18 by removing fully-connected layers) by following Xiao et al. (2021). For the \(C_{\omega }\), a Bayesian layer is also introduced to adapt the classification on insufficient sample better. More implement details of BNN can be found in “Appendix A.1”. The structure of \(M_{\Theta }\) is the same as Dou et al. (2019). The images \(\mathcal {X}=\{\textbf{x}_{l_{i}}\}\) conduct T stochastic forward passes on the \(Q_{\phi }\) and \(C_{\omega }\) by MC sampling to obtain probabilistic predicts \(\{\hat{y}_{l_{i}}^{j}\}_{j=1}^{T}\), where the outputs (i.e., probabilistic embeddings) of \(Q_{\phi }\) serve as the inputs for the calculations of \(\mathcal {L}_{global}\) and \(\mathcal {L}_{local}\). The final predicts \(\{\hat{y}_{l_{i}}^{j}\}\) are the expectation of \(\{\hat{y}_{l_{i}}^{j}\}_{j=1}^{T}\). The total objectives can be summarized as follows,

Discussion The rationale that our proposed method can benefit DG performance on small-data scenario can come from two aspects. First, BNN can be adaptive to insufficient data well compared with deterministic models (Graves, 2011; Mallick et al., 2021). For our proposed method, BNN is introduced to the DG problem in the context of insufficient data, where BNN-based feature extractor and classification layers can take both consistent improvements (see 2nd and 3rd columns in Table 4). More importantly, domain-invariant representation learning under this probabilistic framework from global and local perspectives contributes to more robust cross-domain performance (see 4th and 5th columns in Table 4; Fig. 3 in Sect. 4.5 for the effectiveness of P-MMD).

Epithelium stroma classification is a fundamental step for the prognostic analysis of the tumor. The public histopathological image datasets for binary classification (epithelium or stroma) are collected from three healthcare centers with different staining types and tissue densities¹: IHC, NKI, and VGH. After the patching operation, IHC, NKI and VGH datasets respectively have 1342,1230, and 1376 patches, which means that the insufficient sample problem exist in all source domains compared with large-scale natural images. We randomly split the data of each source domain into a training set (80%) and a test set (20%) and adopt the leave-one-domain-out strategy for evaluation. The pretrained ResNet18 is introduced as the backbone. The structure of Bayesian layer in \(Q_{\phi }\) is a fully-connected-based BNN with \(512 \times 512\). The structure of Bayesian layer in \(C_{\omega }\) is also a fully-connected-based BNN with \(512 \times 2\). We utilize Adam optimizer with learning rate as \(5 \times 10^{-5}\) for training. The batch size is 32 for each source domain with 4000 iterations. The hyperparameters are selected in a wide range on the validation set, where the \(\beta _{1}\) and \(\beta _{2}\) are 0.1 and 0.7 for the \(\mathcal {L}_{local}\) and the \(\mathcal {L}_{global}\), respectively. For the P-MMD, level-1 and level-2 kernels are the Gaussian RBF kernels (the kernel bandwidth is empirically set to 1 for all kernels) by following Muandet et al. (2012). For the P-CSA loss, the distance margin \(\xi \) is set to 1. By balancing the performance and computational efficiency, the number of MC sampling in each Bayesian layer (a.k.a., T), is set to 10. We also discuss the influence of using different T in the Sect. 4.5. We report the results based on average value and standard deviation in each target domain by running the experiment for five different times.

Results Table 1 shows the epithelium stroma classification results on different target domains. We compare with five different approaches. All methods are constructed via a SWAD-based (Cha et al., 2021) framework (where a pretrained ResNet18 as the backbone) with the same training schedule. DeepAll is a deterministic version of our proposed method that is trained on all source domains without any DG strategy in the sequel. Some observations can be summarized as following. First, both our proposed method and BDIL (Xiao et al., 2021) achieve promising results on both IHC and VGH tasks, which may benefit from the positive impact of probabilistic framework on insufficient samples. However, BDIL has an obvious performance drop on the more challenging NKI domain (which has a lower average accuracy compared with other target domains) compared with our proposed method (i.e., 71.89% v.s. 76.71%). This may be due to the introduction of global alignment (i.e., P-MMD, as used by ours), which is more powerful for learning domain-invariant representations than only using local negative pairs (as used by BDIL). Second, compared with contrastive semantic alignment-based method (i.e., MSFA (Dou et al., 2019)), our proposed method achieves a significantly better performance (80.45% v.s. 88.82% on IHC domain), due to a more reliable distribution-based contrastive learning manner on insufficient samples from all source domains. Finally, our proposed method achieves the best average performance with a clear margin compared with other approaches (i.e., LDDG (Li et al., 2020), KDDG (Wang et al., 2021), DNA (Chu et al., 2022), MIRO (Cha et al., 2022), and DSU (Chu et al., 2022)).

Seven public skin lesion datasets² for seven classes of lesions are collected from various institutes using different dermatoscope types: HAM10000 with 10015 images, UDA with 601 images, SON with 9251 images, DMF with 1212 images, MSK with 3551 images, D7P with 1926 images, and PH2 with 200 images. We can observe that the insufficient sample problem exists in some source domains, especially in PH2 and UDA domains. Following previous work (Li et al., 2020), each domain is randomly split into a 50% training set, 30% test set, and 20% validate set, respectively. As adopted in Li et al. (2020), one domain from DMF, D7P, MSK, PH2, SON and UDA is as target domain and the remaining domains together with HAM10000 as source domains. The pretrained ResNet18 is introduced as the backbone. The structure of Bayesian layer in \(Q_{\phi }\) is a fully-connected-based BNN with \(512 \times 512\). The structure of Bayesian layer in \(C_{\omega }\) is also a fully-connected-based BNN with \(512 \times 7\). The Adam optimizer is employed with learning rate of \(5 \times 10^{-5}\) for 2000 iterations. The hyperparameters are selected in a wide range on the validation set, where the \(\beta _{1}\) and \(\beta _{2}\) are 0.1 and 0.7 for the \(\mathcal {L}_{local}\) and the \(\mathcal {L}_{global}\), respectively. The batch size is 32 for each source domain. The remaining settings are the same with epithelium stroma classification. For the results, the average value and standard deviation are reported by running five times.

Results Table 2 shows the skin lesion classification accuracies on different target domains. Six different methods are utilized for comparison. All approaches are implemented using the SWAD-based framework and a pretrained ResNet18 model as the backbone. One has some observations as following. First, compared with contrastive semantic alignment-based method (e.g., MASF), our proposed method provides more reliable distribution-based pairs, leading to a better performance on insufficient samples from some source domains (MASF:0.2692 v.s. Ours: 0.3781 on DMF domain, where PH2 and UDA as the parts of source domain). Second, although BDIL slightly outperforms our proposed method on D7P domain (0.6204 v.s. 0.6120), our proposed method has a significantly better performance on the challenging DMF domain (0.2985 v.s. 0.3781) and the average results (0.7049 v.s. 0.7328) and other domains. Third, it seems that other baseline methods (e.g., SWAD, DNS, and DSU) impose respective schemes to relieve the impact of insufficient samples from some source domains. For example, DSU achieves the best performance on DMF domain. This may be due to the positive impact of straightforward domain randomization. Yet, it is difficult for DSU to realize consistently better results in multiple domains compared with our proposed method, owing to the lack of explicit domain alignment.

The spinal cord gray matter (GM) segmentation (pixel-level classification) is an emergent task for predicting disability via evaluating the atrophy of GM area. The acquired magnetic resonance imaging (MRI) data are collected from four healthcare centers³: “site1” with 30 slices, “site2” with 113 slices, “site3” with 246 slices, and “site4” with 122 slices. One can observe that the insufficient sample problem exists in some source domains, especially in “site1” domain. By following previous work (Li et al., 2020), we randomly split the data of each source domain into a training set (80%) and a test set (20%) and adopt the leave-one-domain-out strategy for evaluation. The hyperparameters are selected in a wide range on the validation set, where the \(\beta _{1}\) and \(\beta _{2}\) are 0.01 and 0.001 for the \(\mathcal {L}_{local}\) and the \(\mathcal {L}_{global}\), respectively. The 2D-Unet (Ronneberger et al., 2015) is leveraged as the backbone for all methods. The structures of \(Q_{\phi }\) and \(C_{\omega }\) as well as more experimental details can be found in “Appendix B”. By following Li et al. (2020), the average results in each target domain are reported by running three times.

Results Table 3 shows the spinal cord GM segmentation results on different target domains. Dice Similarity Coefficient (DSC), Jaccard Index (JI), and Conformity Coefficient (CC) are used to measure the accuracy of obtained segmentation results. True Positive Rate (TPR) and Average Surface Distance (ASD) are introduced from statistical and distance-based perspectives. We compare with five different methods (which have effective segmentation performance). Some observations can be found as follows. First, our proposed method outperforms all baseline methods in terms of average results for five quantitative metrics. Second, compared with the contrastive semantic alignment-based method (e.g., MASF), reliable pixel-level pairs constructed by our proposed method also contribute to improving performance in scenarios with insufficient samples from some source domains. Third, as we can see, our proposed method and DSU achieve the best and second-best performance compared with other baselines, which may be reasonable as they can benefit from the modeling of the data uncertainty in the small-data regime using the BNN or the distribution modeling.

The spinal cord gray matter segmentation task is utilized to explore the effectiveness of each component for our proposed method, due to its various quantitative metrics. The results can be shown in Table 4. First, we observe that better performance can be achieved by introducing a probabilistic layer compared with the results that using Unet, which reflects the superiority of probabilistic models. Secondly, we observe that by either introducing local or global alignment for domain-invariant information learning, better performance can be achieved compared with the results of only using the probabilistic layer, which shows the effectiveness of the introduced probabilistic feature regularization term. Last but not least, by imposing domain-invariant learning with both local and global views, the performances are further improved, which justifies the effectiveness of our proposed method by jointly considering local and global alignment.

Effectiveness of domain-invariant loss We are also interested in the impacts of domain-invariant losses on different tasks. Thus, we visualize the learning curves of different task in Fig. 3. As we can observe, for the skin lesion (on DMF) and epithelium-stroma (on IHC) classification tasks, the loss curves with iterations reflect the global discrepancy converges faster than the local discrepancy, while the more challenging cross-domain task converges more slowly on global alignment.

Results on Different Fractions of Training Data We are interested in how different fractions of training samples influence the final performance based on small-data scenario. To this end, we adopt skin lesion classification task for evaluation since the number of samples in each domain turns out to be imbalance (some domains such as HAM10000 contain sufficient data while the number of samples in some other domains such as PH2 and UDA is insufficient). As such, we can better simulate the scenario that the issue of insufficient samples either exists in all source domains or in partial source domains. We choose MSK dataset with the second-highest number of samples as the target domain and the remaining datasets as the source domains (including PH2 and UDA). Table 5 shows the accuracies. As we can see, compared with DNA (with second-best performance on 100% of samples), our proposed method and BDIL have better average results and lower performance attenuation as the decrease of sample number, due to their probabilistic gain under small data scenarios. Compared with BDIL, our proposed method shows a more robust performance on this challenging small data scenario, due to the integration of local (P-CSA) and global (P-MMD) alignments.

The results can be found in Table 6. As we can see, our proposed method achieved the best performance among all settings compared with all baseline methods. Meanwhile, it seems that the Bayesian-based DG approaches (e.g., our proposed method and BIDL) have better performance compared with other methods, which is reasonable as the BNN can be adaptive to the small data scenario well. Especially, our proposed method has around 5% improvements compared with the second-best method when T is set to smaller, i.e., 20.

Kernel Mean Embedding (i.e., P-MMD) v.s. Mean Embedding We explore the effect of different schemes for probabilistic embeddings. A straightforward method is to represent probabilistic embeddings with the expectation, which is called the “Mean Embedding”. Then, a probabilistic embedding can be considered as a latent point and the MMD can be used to measure the discrepancy between distributions consisting of latent points.

For the mean embedding-based \(\mathcal {L}_{global}\), the computational process of this scheme for MMD distance can be formulated asEqation (10) can be further constructed a global alignment loss \(\mathcal {L}_{global}\). For the local alignment loss \(\mathcal {L}_{local}\), the Euclidean distance can be used to compute the distance between latent points, which is similar to the original CAS loss in Motiian et al. (2017). For the positive pairs with the same label, the mean embedding-based positive contrastive loss can be represented as\(M_{\Theta }(\cdot )\) denotes the embedding network of metric learning. For the negative pairs with different labels, the negative contrastive loss \(\mathcal {L}_{local}^{neg}\) is denoted byAs a result, a mean embedding-based contrastive loss with the view of local alignment can be calculated asInstead, we can observe from Fig. 2 that our proposed method induces a level-2 kernel-based MMD with empirical estimation for probabilistic embeddings. Specifically, our proposed scheme can preserve higher moments of a probabilistic embedding via nonlinear level-1 kernel (see the fourth component in Fig. 2). Moreover, by introducing a level-2 kernel, the similarities between probabilistic embeddings also can be measured based on their own moment information (see the last component in Fig. 2). Benefiting from these virtues, the proposed probabilistic MMD can accurately capture the discrepancy between mixture distributions via an extended empirical MMD fashion.

Here, we validate the effectiveness of different schemes on the NKI task of Epithelium Stroma classification in each aligned view. The experimental settings are similar for different methods. The experimental results can be found in Fig. 4. As we can see, our proposed method achieves consistent improvements in each alignment method with different Monte Carlo samples, which may be reasonable as the kernel mean representation can preserve many statistical components due to the injective property. Second, when the number of MC samples is 10, we can observe an obvious margin in global alignment, which refers to the computation between mixture distributions.

Influence of Different Number of MC Samples It is much important to balance the number of Monte Carlos samples and the computational efficiency. On the one hand, the property of probabilistic embeddings can be affected by the Monte Carlos sampling. On the other hand, too many Monte Carlos samples may suffer from the heavy computational cost. Xiao et al. (2021) suggested that the distributional property and computational cost are both acceptable for the computation of the KL divergence when the number of Monte Carlos samples is chosen appropriately, the practical performance of our proposed method needs to be explored. We conduct the experiments on the NKI task of Epithelium-Stromal classification with different Monte Carlo samples T.

The results are shown in Fig. 5. As we can see, if the number of Monte Carlo samples is too small, it is difficult to capture the property of distribution for probabilistic embeddings. As the increase of T, there is an obvious improvement for our proposed method. Interestingly, the performance is gradually saturated. As a result, by balancing the number of Monte Carlos samples and the computational efficiency, the number of Monte Carlos samples T in each Bayesian layer is set to 10.

Here, we introduce three DG benchmarks, namely PACS (Art: 2048 images, Cartoon: 2344 images, Photo: 1670 images, Sketch: 3929 images), OfficeHome (has 15588 samples with 65 classes from four domains) and VLCS (has 10729 samples with 5 classes from four domains), for comparison. Compared with some large-scale benchmarks (e.g., DomainNet and Wilds), these three datasets are more appropriate to explore the effectiveness of different DG models under the scenario of insufficient samples. performance. We adopt pretrained ResNet50 as the backbone for all benchmarks. The structure of the overall framework is similar with the model mentioned in lesion skin classification. Our proposed method as well as baseline methods are all based on DomainBed, where the holdout fraction (the proportion of validation set) rate for DomainBed is set to 0.2 for all methods. A domain is the target domain and the remaining domains are the source domain for training. The testing is on the overall data of a target domain.

Here, our proposed method is optimized by Adam optimizer with learning rate as \(5\times 10^{-5}\). The batch size for each source domain is 32. The training steps are set to 20000 for PACS and OfficeHome, and 2000 for VLCS. By following the SWAD framework, the training process will be stopped for our proposed method when the validation loss increases significantly. The hyperparameters are selected in a wide range on the validation set. For the \(\mathcal {L}_{local}\) and \(\mathcal {L}_{global}\), the \(\beta _{1}\) and the \(\beta _{2}\) are set to 0.1 and 1 for all benchmark datasets. Other hyperparameters such as kernel function, kernel bandwidth and distance margin are similar to the settings mentioned before.

The experimental results on PACS, OfficeHome and VLCS can be shown in Tables 7, 8, and 9. As we can see, compared with domain-invariant-based approaches (e.g., DANN), our proposed method has a significant improvement due to the introduction of a probabilistic framework. Our proposed method also outperforms the data augmentation-based approach (e.g., Mixstyle, Manifold Mixup, and CutMix). Although data generation methods (e.g., MixStyle) can effectively tackle the insufficient sample problem via additional generative samples, the lack of effective domain-invariant learning may hamper the improvement of the performance. Our proposed method achieves better performances compared with feature disentanglement-based (e.g., pAdaI). Last but not least, compared with the state-of-the-art domain generalization method (e.g., MIRO), our proposed method achieves comparable performance on small-scale benchmark datasets, which demonstrates the scalability of our proposed method in a small-data regime.