V ARIATIONAL I NVARIANT L EARNING FOR B AYESIAN D OMAIN G ENERALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Domain generalization addresses the out-of-distribution problem, which is challenging due to the domain shift and the uncertainty caused by the inaccessibility to data from the target domains.
In this paper, we propose variational invariant learning , a probabilistic inference framework that jointly models domain invariance and uncertainty.
We introduce variational Bayesian approximation into both the feature representation and classiﬁer layers to facilitate invariant learning for better generalization across domains.
In the probabilistic modeling framework, we introduce a domain-invariant principle to explore invariance across domains in a uniﬁed way.
We incorporate the principle into the variational Bayesian layers in neural networks, achieving domain-invariant representations and classiﬁer.
We empirically demonstrate the effectiveness of our proposal on four widely used cross-domain visual recognition benchmarks.
Ablation studies demonstrate the beneﬁts of our proposal and on all benchmarks our variational invariant learning consistently delivers state-of-the-art performance.
I NTRODUCTION
Domain generalization (Muandet et al., 2013), as an out-of-distribution problem, aims to train a model on several source domains and have it generalize well to unseen target domains.
The major challenge stems from the large distribution shift between the source and target domains, which is further complicated by the prediction uncertainty (Malinin & Gales, 2018) introduced by the inaccessibility to data from target domains during training.
Previous approaches focus on learning domain-invariant features using novel loss functions (Muandet et al., 2013; Li et al., 2018a) or speciﬁc architectures (Li et al., 2017a; D’Innocente & Caputo, 2018).
Meta-learning based methods were proposed to achieve similar goals by leveraging an episodic training strategy (Li et al., 2017b;
Balaji et al., 2018; Du et al., 2020).
Most of these methods are based on deep neural network backbones (Krizhevsky et al., 2012;
He et al., 2016).
However, while deep neural networks have achieved remarkable success in various vision tasks, their performance is known to degenerate considerably when the test samples are out of the training data distribution (Nguyen et al., 2015;
Ilse et al., 2019), due to their poorly calibrated behavior (Guo et al., 2017; Kristiadi et al., 2020).
As an attractive solution, Bayesian learning naturally represents prediction uncertainty (Kristiadi et al., 2020; MacKay, 1992), possesses better generalizability to out-of-distribution examples (Louizos & Welling, 2017) and provides an elegant formulation to transfer knowledge across different datasets (Nguyen et al., 2018).
Further, approximate Bayesian inference has been demonstrated to be able to improve prediction uncertainty (Blundell et al., 2015; Louizos & Welling, 2017;
Atanov et al., 2019), even when only applied to the last network layer (Kristiadi et al., 2020).
These properties make it appealing to introduce Bayesian learning into the challenging and unexplored scenario of domain generalization.
In this paper, we propose variational invariant learning (VIL), a Bayesian inference framework that jointly models domain invariance and uncertainty for domain generalization.
We apply variational
Bayesian approximation to the last two network layers for both the representations and classiﬁer by placing prior distributions over their weights, which facilitates generalization.
We adopt Bayesian neural networks to domain generalization, which enjoys the representational power of deep neural networks while facilitating better generalization.
To further improve the robustness to domain shifts, we introduce the domain-invariant principle under the Bayesian inference framework, which enables us to explore domain invariance for both feature representations and the classiﬁer in a uniﬁed way.
We evaluate our method on four widely-used benchmarks for cross-domain visual object classiﬁcation.
Our ablation studies demonstrate the effectiveness of the variational Bayesian domain-invariant features and classiﬁer for domain generalization.
Results further show that our method achieves the best performance on all of the four benchmarks.
M ETHODOLOGY
We explore Bayesian inference for domain generalization.
In this task, the samples from the target domains are never seen during training, and are usually out of the data distribution of the source domains.
This leads to uncertainty when making predictions on the target domains.
Bayesian inference offers a principled way to represent the predictive uncertainty in neural networks (MacKay, 1992; Kristiadi et al., 2020).
We brieﬂy introduce approximate Bayesian inference, under which we will introduce our variational invariant learning for domain generalization.
2. A PPROXIMATE B AYESIAN INFERENCE
Given a dataset { x ( i ) , y ( i ) } Ni =1 of N input-output pairs and a model parameterized by weights θ with a prior distribution p ( θ ) , Bayesian neural networks aim to infer the true posterior distribution p ( θ | x , y ) .
As the exact inference of the true posterior is computationally intractable, Hinton & Camp (1993) and Graves (2011) recommended learning a variational distribution q ( θ ) to approximate p ( θ | x , y ) by minimizing the Kullback-Leibler (KL) divergence between them:
[Equation]
The above optimization is equivalent to minimizing the loss function:
[Equation]
which is also known as the negative value of the evidence lower bound (ELBO) (Blei et al., 2017).
2. V ARIATIONAL DOMAIN - INVARIANT LEARNING
In domain generalization, let D =
{ D i } |D| i =1 = S ∪ T be a set of domains, where S and T denote source domains and target domains respectively.
S and T do not have any overlap with each other but share the same label space.
For each domain D i ∈ D , we can deﬁne a joint distribution p ( x i , y i ) in the input space X and the output space Y .
We aim to learn a model f : X → Y in the source domains S that can generalize well to the target domains T .
The fundamental problem in domain generalization is to achieve robustness to domain shift between source and target domains, that is, we aim to learn a model invariant to the distributional shift between the source and target domains.
In this work, we mainly focus on the invariant property across domains instead of exploring general invariance properties (Nalisnick & Smyth, 2018).
Therefore, we introduce a formal deﬁnition of domain invariance, which is easily incorporated as criteria into the Bayesian framework to achieve domain-invariant learning.
Provided that all domains in D are in the same domain space, then for any input sample x s in domain
D s , we assume that there exists a domain-transform function g ζ ( · ) which is deﬁned as a mapping function that is able to project x s to other different domains D ζ with respect to the parameter ζ , where ζ ∼ q ( ζ ) , and a different ζ lead to different post-transformation domains D ζ . Usually the exact form of g ζ ( · ) is not necessarily known.
Under this assumption, we introduce the deﬁnition of domain invariance, which we will incorporate into the Bayesian layers of neural networks for domain-invariant learning.
Deﬁnition 2.1 (Domain Invariance) Let x s be a given sample from domain D s ∈ D , and x ζ = g ζ ( x s ) be a transformation of x s in another domain D ζ , where ζ ∼ q ( ζ ) .
p θ ( y | x ) denotes the output distribution of input x with model θ .
The model θ is domain-invariant if,
[Equation]
) , ∀
[Equation] [Equation]
Here, we use y to represent the output from a neural layer with input x , which can either be the prediction vector from the last layer or the feature vector from the convolutional layers.
To make the domain-invariant principle easier to implement, we then extend the Eq.
(3) to an expectation form:
[Equation]
Based on this deﬁnition, we use the Kullback-Leibler divergence between the two terms in Eq.
(4), D
KL (cid:2) p θ ( y s | x s )
||
E q ζ [ p θ ( y ζ | x ζ )] (cid:3) , to quantify the domain invariance of the model, which will be zero when the model is domain invariant.
As in most cases, there is no analytical form of the domaintransform function and only a few samples from D ζ are available, which makes E q ζ [ p θ ( y ζ | x ζ )] intractable.
Thus, we derive the following upper bound of the divergence:
[Equation]
which can be approximated by Monte Carlo sampling.
We deﬁne the complete objective function of our variational invariant learning by combining Eq.
(5) with Eq.
However, in Bayesian inference, the likelihood is obtained by taking the expectation over the distribution of parameter θ , i.e., p θ ( y | x ) =
E q ( θ ) [ p ( y | x , θ )] , which is also intractable in
Eq.
As the KL divergence is a convex function (Nalisnick & Smyth, 2018), we further extend
Eq.
(5) to an upper bound:
[Equation]
which is tractable with the unbiased Monte Carlo approximation.
The complete derivations of
Eq.
(5) and Eq.
(6) are provided in Appendix A.
In addition, it is worth noting that the domain-transformation distribution q ( ζ ) is implicit and inexpressible in reality and there are only a limited number of domains available in practice.
This problem is exacerbated because the target domain is unseen during training, which further limits the number of available domains.
Moreover, in most of the domain generalization databases, for a certain sample x s from domain D s , there is no transformation corresponding to x s in other domains.
This prevents the expectation with respect to q ζ from being directly tractable in general.
Thus, we resort to use an empirically tractable implementation and adopt an episodic setting as in (Li et al., 2019).
In each episode, we choose one domain from the source domains S as the metasource domain D s and the rest are used as the meta-target domains { D t } Tt =1 .
To achieve variational invariant learning in the Bayesian framework, we use samples from meta-target domains in the same category as x s to approximate the samples of g ζ ( x s ) .
Then we obtain a general loss function for domain-invariant learning:
[Equation]
where (cid:8) x it (cid:9) N i =1 are from D t , denoting the samples in the same category as x s .
More details and an illustration of the domain-invariant loss function can be found in Appendix B.
With the aforementioned loss functions, we develop the loss function of variational invariant learning for domain generalization:
[Equation]
Our variational invariant learning combines the Bayesian framework, which is able to introduce uncertainty into the network and is beneﬁcial for out-of-distribution problems (Daxberger &
Hern ´ andez-Lobato, 2019), and a domain-invariant loss function L I , which is designed based on predictive distributions to make the model generalize better to the unseen target domains.
For Bayesian learning, it has been demonstrated that being just “a bit” Bayesian in the last layer of the neural network can well represent the uncertainty in predictions (Kristiadi et al., 2020).
This indicates that applying the Bayesian treatment only to the last layer already brings sufﬁcient beneﬁts of Bayesian inference.
Although adding Bayesian inference to more layers improves the performance, it also increases the computational cost.
Further, from the perspective of domain invariance, making both the classiﬁer and feature extractor more robust to the domain shifts also leads to better performance (Li et al., 2019).
Thus, there is a trade-off between the beneﬁts of variational Bayesian domain invariance and computational efﬁciency.
Instead of applying the Bayesian principle to all the layers of the neural network, in this work we explore domain invariance by applying it to only the classiﬁer layer ψ and the last feature extraction layer φ .
In this case, the L Bayes in Eq.
(2) becomes the ELBO with respect to ψ and φ jointly.
As they are independent, the L Bayes is expressed as:
[Equation]
The above variational inference objective allows us to explore domain-invariant representations and classiﬁer in a uniﬁed way.
The detailed derivation of Eq.
(9) is provided in Appendix A.
Domain-Invariant Classiﬁer To establish the domain-invariant classiﬁer, we directly incorporate the proposed domain-invariant principle into the last layer of the network, which gives rise to
[Equation]
where z denotes the feature representations of input x , and the subscripts s and t indicate the metasource domain and the meta-target domains as in Eq.
Since p ( y | z , ψ ) is a Bernoulli distribution, we can conveniently calculate the KL divergence in Eq.
Domain-Invariant Representations To also make the representations domain invariant, we have
[Equation]
where φ are the parameters of the feature extractor.
Since the feature extractor is also a Bayesian layer, the distribution of p ( z | x , φ ) will be a factorized Gaussian if the posterior of φ is as well.
We illustrate this as follows.
Let x be the input feature of a Bayesian layer φ , which has a factorized Gaussian posterior, the posterior of the activation z of the Bayesian layer is also a factorized
Gaussian (Kingma et al., 2015):
[Equation]
where z j denotes the j -th element in z , likewise for x i , and φ i,j denotes the element at the position ( i, j ) in φ .
Based on this property of the Bayesian framework, we assume that the posterior of our variational invariant feature extractor has a factorized Gaussian distribution, which leads to an easier calculation of the KL divergence in Eq. Note that with the domain-invariant representations, z in Eq.
(10) corresponds to samples of the feature representation distributions: z s ∼ p
( z s | x s , φ ) and z t ∼ p ( z t | x t , φ ) .
2. O BJECTIVE FUNCTION
The objective function of our variational invariant learning is deﬁned as:
[Equation]
13) where λ ψ and λ φ are hyperparameters to control the domain-invariant terms.
We adopt Monte Carlo sampling and obtain the empirical objective function for variational invariant learning as follows: [Equation]
[Equation]
where x s and z s denote the input and its feature from D s , respectively, and x it and
z it are from D t as in Eq.
The posteriors are set to factorized Gaussian distributions, i.e., q ( ψ ) =
N ( µ ψ , σ 2 ψ ) and q ( φ ) =
N ( µ φ , σ 2 φ ) .
We adopt the reparameterization trick to draw Monte Carlo samples (Kingma & Welling, 2014) as ψ ( (cid:96) )
= µ ψ + (cid:15) ( (cid:96) ) ∗ σ ψ , where (cid:15) ( (cid:96) )
∼ N (0 , I ) .
We draw the samples for φ ( m ) in a similar way.
In the implementation of our variational invariant learning, to increase the ﬂexibility of the prior distribution in our Bayesian layers, we choose to place a scale mixture of two Gaussian distributions as the priors p ( ψ ) and p ( φ ) (Blundell et al., 2015):
[Equation]
where σ 1 , σ 2 and π are hyperparameters chosen by cross-validation.
R ELATED WORK
One solution for domain generalization is to generate more source domain data to increase the probability of covering the data in the target domains (Shankar et al., 2018; Volpi et al., 2018).
Shankar et al.
(2018) augmented the data by perturbing the input images with adversarial gradients generated by an auxiliary classiﬁer.
Qiao et al.
(2020) proposed a more challenging scenario of domain generalization named single domain generalization, which only has one source domain, and they designed an adversarial domain augmentation method to create “ﬁctitious” yet “challenging” data.
Recently, Zhou et al.
(2020) employed a generator to synthesize data from pseudo-novel domains to augment the source domains, maximizing the distance between source and pseudo-novel domains as measured by optimal transport (Peyr ´ e et al., 2019).
Another solution for domain generalization is based on learning domain-invariant features (D’Innocente & Caputo, 2018; Li et al., 2018b; 2017a).
Muandet et al.
(2013) proposed domain-invariant component analysis to learn invariant transformantions by minimizing the dissimilarity across domains.
Louizos et al.
(2015) learned invariant representations by variational autoencoder (Kingma & Welling, 2014), which introduced Bayesian inference into invariant feature learning.
Dou et al.
and Seo et al.
(2019) tried to achieve a similar goal by introducing two complementary losses, global class alignment loss and local sample clustering loss, and employing multiple normalization methods.
Li et al.
(2019) proposed an episodic training algorithm to obtain both domain-invariant feature extractor and classiﬁer.
Recently, meta-learning based techniques have been considered to solve domain generalization problems.
Li et al.
(2018a) proposed a meta-learning domain generalization method which introduced the gradient-based method,
i.e., model agnostic meta-learning
(Finn et al., 2017), for domain generalization.
Balaji et al.
(2018) address the domain generalization problem by learning a regularization function in a meta-learning framework, making the model robust to domain shifts.
Du et al.
(2020) propose the meta variational information bottleneck to learn domain-invariant representations through episodic training.
To the best of our knowledge, Bayesian neural networks have not yet been explored in domain generalization.
Our method introduces variational Bayesian approximation to both the feature extractor and classiﬁer of the neural network in conjunction with the newly introduced domain-invariant principle for domain generalization.
The resultant variational invariant learning combines the representational power of deep neural networks and variational Bayesian inference.
Similar to our proposal, CCSA by Motiian et al.
(2017) also aligns representations across domains in the same class.
Speciﬁcally, CCSA utilizes an L2 distance between deterministic features while we exploit Bayesian neural networks to learn domain-invariant representations by minimizing the distance between domain distributions.
Theoretically, minimizing the distance between distributions incorporates larger inter-class variance than minimizing distance of deterministic features. Moreover, we apply our variational invariant learning to both the feature extractor and the classiﬁer, while CCSA only considers an alignment loss on the feature representations.  E XPERIMENTS 4. D ATASETS AND SETTINGS
We conduct our experiments on four widely used benchmarks in domain generalization:
PACS (Li et al., 2017a) consists of 9,991 images of seven classes from four domains - photo , artpainting , cartoon and sketch .
We follow the “leave-one-out” protocol in (Li et al., 2017a; 2018b;
Carlucci et al., 2019), where the model is trained on any three of the four domains, which we call source domains, and tested on the last (target) domain.
Ofﬁce-Home (Venkateswara et al., 2017) also has four domains: art , clipart , product and realworld .
There are about 15,500 images of 65 categories for object recognition in ofﬁce and home environments.
We use the same experimental protocol as for PACS.
Rotated MNIST and Fashion-MNIST were introduced for evaluating domain generalization in (Piratla et al., 2020).
For the fair comparison, we follow their recommended settings and randomly select a subset of 2,000 images from MNIST and 10,000 images from Fashion-MNIST, which is considered to have been rotated by 0 ◦ .
The subset of images is then rotated by 15 ◦ through 75 ◦ in intervals of 15 ◦ , creating ﬁve source domains.
The target domains are created by rotation angles of 0 ◦ and 90 ◦ .
We use these two datasets to demonstrate the generalizability by comparing the performance on in-distribution and out-of-distribution data.
For all four benchmarks, we use ResNet-18 (He et al., 2016) pretrained on ImageNet (Deng et al., 2009) as the CNN backbone.
During training, we use Adam optimization (Kingma & Ba, 2014) with the learning rate set to 0.0001, and train for 10,000 iterations.
In each iteration we choose one source domain as the meta-source domain.
The batch size is 128.
To ﬁt memory footprint, we choose a maximum number of samples per category per target domain to implement the domaininvariant learning, i.e. sixteen for PACS, Rotated MNIST and Fashion-MNIST datasets, and four for the Ofﬁce-Home dataset.
We choose λ φ and λ ψ based on the performance on the validation set and their inﬂuence is summarized in the new Fig 3 in Appendix C.
The optimal values of λ φ and λ ψ are 0.1 and 100 respectively.
Parameters σ 1 and σ 2 in Eq.
(15) are set to 0.1 and 1.5.
The model with the highest validation set accuracy is employed for evaluation on the target domain.
All code will be made publicly available.
4. A BLATION STUDY
We conduct an ablation study to investigate the effectiveness of our variational invariant learning for domain generalization.
The experiments are performed on the PACS dataset.
Since the major contributions of this work are the Bayesian treatment and the domain-invariant principle, we evaluate their effect by individually incorporating them into the classiﬁer - the last layer - ψ and the feature extractor - the penultimate layer - φ .
The results are shown in Table 1.
The “ (cid:88) ” and “ × ” in the “Bayesian” column denote whether the classiﬁer ψ and feature extractor φ are Bayesian layers or deterministic layers.
In the “Invariant” column they denote whether the domain-invariant loss is introduced into the classiﬁer and the feature extractor.
Note that the predictive distribution is a
Bernoulli distribution, which also admits the domain-invariant loss, we therefore include this case for a comprehensive comparison.
In Table 1, the ﬁrst four rows demonstrate the beneﬁt of the Bayesian treatment.
The ﬁrst row (a) serves as a baseline model, which is a vanilla deep convolutional network without any Bayesian treatment and domain-invariant loss.
The backbone is also a ResNet-18 pretrained on ImageNet. It is clear the Bayesian treatment, either for the classiﬁer (b) or the feature extractor (c), improves the performance, especially in the “Art-painting” and “Sketch” domains, and this is further demonstrated in (d) where we employ the Bayesian classiﬁer and feature extractor simultaneously.
The beneﬁt of the domain-invariant principle for classiﬁers is demonstrated by comparing (e) to (a) and (f) to (b).
The settings with domain invariance consistently perform better than those without it.
A similar trend is also observed when applying the domain-invariant principle to the feature extractor, as shown by comparing (g) to (c).
Overall, our variational invariant learning (h) achieves the best performance compared to other variants, demonstrating its effectiveness for domain generalization.
Note that the feature distributions p ( z | x ) are unknown without Bayesian formalism, leading to an intractable L I ( φ ) .
Therefore, we do not conduct the experiment with only the domain-invariant loss on both the classiﬁer and the feature extractor.
To further demonstrate the domain-invariant property of our method, we visualize the features learned by different settings of our method in Table 1.
We use t-SNE (Maaten & Hinton, 2008) [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
invariant principle enlarges the inter-class distance in the target domain by reducing the intra-class distances between the source and target domains.
to reduce the feature dimension into a two-dimensional subspace, following Du et al.
The visualization is shown in Fig.
More visualization results are provided in Appendix E.
Figs.
1 (a), (b), (c) and (d) show the baseline, Bayesian classiﬁer, Bayesian representations and both Bayesian classiﬁer and representations , demonstrating the beneﬁts of Bayesian inference for learning domain-invariant features.
The Bayesian treatment on both the classiﬁer and the feature extractor enlarges the inter-class distance in all domains, which beneﬁts the classiﬁcation performance on the target domain, as shown in Fig.
1 (d).
Figs.
1 (e), (f), (g), (h) are visualizations of feature representations after introducing the domain-invariant principle, compared to the upper row.
Comparing the two ﬁgures in each column indicates that our domain-invariant principle imposed on either the representation or the classiﬁer further enlarges the inter-class distances.
At the same time, it reduces the distance between samples of the same class from different domains.
This is even more apparent in the intra-class distance between samples from source and target domains.
As a result, the inter-class distances in the target domain become larger, therefore improving performance.
It is worth noting that the domain-invariant principle on the classiﬁer in Fig.
1 (f) and on the feature extractor in Fig.
1 (g))
both improve the domain-invariant features.
Our variational invariant learning in Fig.
1 (h) therefore has better performance by combining their beneﬁts.
[Table]
We also experiments with more layers in the feature extractor, see Table 2.
“Bayesian φ (cid:48) ” and “Invariant φ (cid:48) ” denote whether the additional feature extraction layer φ (cid:48) has the Bayesian property and domain-invariant property.
The classiﬁers have both properties in all cases in Table 2. The ﬁrst row is the setting with only one variational invariant layer in the feature extractor. When introducing another Bayesian learning layer φ (cid:48) without the domain-invariant property into the model, as shown in the second row in Table 2, the average performance improves slightly. If we introduce both the Bayesian learning and domain-invariant learning into φ (cid:48) , as shown in the third row, the overall performance declines a bit.
One reason might be the information loss in feature representations caused by the excessive use of domain-invariant learning.
In addition, due to the Bayesian inference and Monte-Carlo sampling, more variational-invariant layers leads to higher memory usage and more computations, which is also one reason for us to apply the variational invariant learning only to the last feature extraction layer and the classiﬁer.
[Table]
painting Cartoon Sketch Mean
[Table]
[Table]
4. S TATE - OF - THE - ART COMPARISON
In this section, we compare our method with several state-of-the-art methods on four datasets.
The results are reported in Tables 3-5.
The baseline on PACS (Table 3), Ofﬁce-Home (Table 4), and rotated MNIST and Fashion-MNIST (Table 5) are all based on the same vanilla deep convolutional
ResNet-18 network, without any Bayesian treatment, the same as row (a) in Table
On PACS, as shown in Table 3, our variational invariant learning method achieves the best overall performance.
On each domain, our performance is competitive with the state-of-the-art and we exceed all other methods on the “Cartoon” domain.
On Ofﬁce-Home, as shown in Table 4, we again achieve the best recognition accuracy.
It is worth mentioning that on the most challenging “Art” and “Clipart” domains, our variational invariant learning also delivers the highest performance, with a good improvement over previous methods.
L2A-OT and DSON outperform the proposed model on some domains of PACS and Ofﬁce-Home. L2A-OT learns a generator to synthesize data from pseudo-novel domains to augment the source domains. The pseudo-novel domains often have similar characteristics with the source data. Thus, when the target data also have similar characteristics with the source domains this pays off as the
[Table]
[Table]
[Table]
[Table]
pseudo domains are more likely to cover the target domain, such as “Product” and “Real World” in Ofﬁce-Home and “Photo” in PACS.
When the test domain is different from all of the training domains the performance suffers, e.g., “Clipart” in Ofﬁce-Home and “Sketch” in PACS.
Our method generates domain-invariant representations and classiﬁers, resulting in competitive results across all domains and overall.
DSON mixtures batch and instance normalization for domain generalization.
This tactic is effective on PACS, but less competitive on Ofﬁce-Home.
We attribute this to the larger number of categories on Ofﬁce-Home, where instance normalization is known to make features less discriminative with respect to object categories (Seo et al., 2019).
Our domain-invariant network makes feature distributions and predictive distributions similar across domains, resulting in good performance on both PACS and Ofﬁce-Home.
On the Rotated MNIST and Fashion-MNIST datasets, following the experimental settings in (Piratla et al., 2020), we evaluate our method on the in-distribution and out-of-distribution sets.
As shown in Table 5, our VIL achieves the best performance on both sets of the two datasets, surpassing other methods.
Moreover, our method especially improves the classiﬁcation performance on the out-of-distribution sets, demonstrating its strong generalizability to unseen domains, which is also consistent with the ﬁndings in Fig.
C ONCLUSION
In this work, we propose variational invariant learning (VIL), a variational Bayesian learning framework for domain generalization.
We introduce Bayesian neural networks into the model, which is able to better represent uncertainty and enhance the generalization to out-of-distribution data.
To handle the domain shift between source and target domains, we propose a domain-invariant principle under the variational inference framework, which is incorporated by establishing a domain- invariant feature extractor and classiﬁer.
Our variational invariant learning combines the representational power of deep neural networks and uncertainty modeling ability of Bayesian learning, showing great effectiveness for domain generalization.
Extensive ablation studies demonstrate the beneﬁts of the Bayesian inference and domain-invariant principle for domain generalization.
Our variational invariant learning sets a new state-of-the-art on four domain generalization benchmarks.
A D ERIVATION
D ERIVATION OF THE UPPER BOUNDS OF DOMAIN - INVARIANT LEARNING
As in most cases the E q ζ [ p θ ( y ζ | x ζ )] is intractable, we derive the upper bound in Eq.
(5), which is achieved via Jensen’s inequality:
[Equation]
In Bayesian inference
[Equation]
computing the likelihood p θ ( y | x ) = (cid:82) θ p ( y | x , θ )d θ = E q ( θ )
[ p ( y | x , θ )] is notoriously difﬁcult.
Thus, as the fact that KL divergence is a convex function
[Equation]
we obtain the upper bound in Eq.
(6) achieved via Jensen’s inequality similar to Eq.
[Equation]
A. D ERIVATION OF VARIATIONAL B AYESIAN APPROXIMATION FOR REPRESENTATION ( φ ) AND CLASSIFIER ( ψ ) LAYERS .
We consider the model with two Bayesian layers φ and ψ as the last layer of feature extractor and the classiﬁer respectively.
The prior distribution of the model is p ( φ , ψ ) , and the true posterior distribution is p ( φ , ψ | x , y ) .
Following the settings in Section 2.1, we need to learn a variational distribution q ( φ , ψ ) to approximate the true posterior by minimizing the KL divergence from q ( φ , ψ ) to p ( φ , ψ | x , y ) :
[Equation]
[Figure]
[Equation] [Figure] [Figure]
𝝓 𝝍 𝝓
[Figure]
𝓛 𝑰 (𝝓) 𝓛 𝑰 (𝝍)
[Figure]
𝑖
[Figure]
𝑁
[Figure] y [Figure] By applying the Bayesian rule p ( φ , ψ | x , y ) ∝ p ( y | x , φ , ψ ) p ( φ , ψ ) , the optimization is equivalent to minimizing:
[Equation]
With φ and ψ are independent,
[Equation]
B D ETAILS OF DOMAIN - INVARIANT LOSS IN VIL TRAINING
We split the training phase of VIL into several episodes.
In each episode, as shown in Fig. 2, we randomly choose a source domain as the meta-source domain D s , and the rest of the source domains { D t } Tt =1 are treated as the meta-target domains.
From D s , we randomly select a batch of samples x s .
For each x s , we then select N samples (cid:8) x it (cid:9) N i =1 , which are in the same category as x s , from each of the meta-target domains D t .
All of these samples are sent to the variational invariant feature extractor φ to get the representations z s and (cid:8)
z it (cid:9) N i =1 , which are then sent to the variational invariant classiﬁer ψ to obtain the predictions y s and (cid:8)
y it (cid:9) N i =1 .
We obtain the domain-invariant loss for feature extractor L I ( φ ) by calculating the mean of the KL divergence of z s and each z it as
Eq.(11).
The domain-invariant loss for feature classiﬁer L I ( ψ ) is calculated in a similar way on y s and (cid:8)
y it (cid:9) N i =1 as Eq.(10).
C A BLATION STUDY FOR HYPERPARAMTERS
We also ablate the hyperparameters λ φ , λ ψ and π on PACS with cartoon as the target domain.
Results are shown in Fig 3 (a), (b) and (c).
We obtain Fig 3 (a) by ﬁxing λ ψ as 100 and adjusting λ φ ,
Fig 3 (b) by ﬁxing λ φ as 1 and adjusting λ ψ and Fig 3 (c) by adjusting π while ﬁxing other settings as in Section 4.1.
λ φ and λ ψ balance the inﬂuence of the Bayesian learning and domain-invariant learning, and their optimal values are 1 and 100.
If the values are too small, the model tends to overﬁt to source domains as the performance on target data drops more obviously than on validation data.
In contrast, too large values of them harm the overall performance of the model as there are obvious decrease of accuracy on both validation data and target data.
Moreover, π balances the two components of the scale mixture prior of our Bayesian model.
According to Blundell et al. (2015), the two components cause a prior density with heavier tail while many weights tightly concentrate around zero.
Both of them are important.
The performance is the best when π is 0.5 according to
Fig 3 (c), which demonstrates the effectiveness of the two components in the scale mixture prior.
[Figure] [Table] D D ETAILED ABLATION STUDY ON PACS In addition to the aforementioned experiments, we conduct some supplementary experiments with other settings on PACS to further demonstrate the effectiveness of the Bayesian inference and domain-invariant loss as shown in Table 6.
The evaluated components are the same as in Table 1.
For better comparison, we show the contents of Table 1 again in Table 6, and add three other settings with IDs (i), (j) and (k).
Note that as the distribution of features z is unknown without a Bayesian feature extractor φ , the settings with L I ( φ ) and a non-Bayesian feature extractor is intractable.
Comparing (i) with (d), we ﬁnd that employing Bayesian inference to the last layer of the feature extractor improves the overall performance and the classiﬁcation accuracy on three of the four domains.
Moreover, comparing (j) and (k) to (g) shows the beneﬁts of introducing variational Bayes and the domain-invariant loss to the classiﬁer on most of the domains and the average of them.
E E XTRA VISUALIZATION RESULTS
To further observe and analyze the beneﬁts of the individual components of VIL for domain-invariant learning, we visualize the features of all categories from the target domain only in Fig. 4, and features of only one category from all domains in Fig.
The same as Fig.
1, the visualization is conducted on the PACS dataset and the target domain is “art-painting”.
The chosen category in Fig.
5 is “horse”.
Fig.
4 provides a more intuitive observation of the beneﬁts of the Bayesian framework and domaininvariant learning in our method for enlarging the inter-class distance in the target domain.
The conclusion is similar as in Fig.
From the ﬁgures in the ﬁrst row, it is clear that the Bayesian framework whether in the classiﬁer ((b)) or the feature extractor ((c)) increases the inter-class distance compared with the baseline method ((a)).
With both of them ((d)), the performance becomes
[Figure]
[Figure]
Fig.
5 provides a deeper insight into the intra-class feature distributions of the same category from different domains.
By introducing the Bayesian inference into the model, the features demonstrate the manifold of the category as shown in the ﬁrst row ((b), (c) and (d)).
This makes recognition easier.
Indeed, the visualization of features from multiple categories has similar properties as shown in Fig.
As shown in each column, introducing the domain-invariant learning into the model leads to a better mixture of features from different domains.
The resultant domain-invariant representation makes the model more generalizable to unseen domains.
We also visualize the features in rotated MNIST and Fashion MNIST datasets, as shown in Fig.
Different shapes denote different categories.
Red samples denote features from the in-distribution set and blue samples denote features from the out-of-distribution set.
Compared with the baseline, our method reduces the intra-class distance between samples from the in-distribution set and the out-ofdistribution set and clusters the out-of-distribution samples of the same categories better, especially in the rotated Fashion-MNIST dataset.
[Figure]