# A Bit More Bayesian: Domain-Invariant Learning with Uncertainty

Zehao Xiao  Jiayi Shen  Xiantong Zhen  Ling Shao  Cees G. M. Snoek

###### Abstract

Domain generalization is challenging due to the domain shift and the
uncertainty caused by the inaccessibility of target domain data. In this
paper, we address both challenges with a probabilistic framework based on
variational Bayesian inference, by incorporating uncertainty into neural
network weights. We couple domain invariance in a probabilistic formula with
the variational Bayesian inference. This enables us to explore domain-
invariant learning in a principled way. Specifically, we derive domain-
invariant representations and classifiers, which are jointly established in a
two-layer Bayesian neural network. We empirically demonstrate the
effectiveness of our proposal on four widely used cross-domain visual
recognition benchmarks. Ablation studies validate the synergistic benefits of
our Bayesian treatment when jointly learning domain-invariant representations
and classifiers for domain generalization. Further, our method consistently
delivers state-of-the-art mean accuracy on all benchmarks.

Machine Learning, ICML

##  1 Introduction

Learning to improve the generalization of deep neural networks to data out of
their training distribution remains a fundamental yet challenging problem for
machine learning (Wang et al., 2018; Bengio et al., 2019; Krueger et al.,
2020). Domain generalization (Muandet et al., 2013) aims to train a model on
several source domains and have it generalize well to unseen target domains.
The main challenge stems from the large shift of distributions between the
source and target domains, which is further complicated by the prediction
uncertainty (Malinin & Gales, 2018) introduced by the inaccessibility to data
from the target domains during training. Established approaches learn domain-
invariant features by dedicated loss functions (Muandet et al., 2013; Li et
al., 2018a) or specific architectures (Li et al., 2017; D’Innocente & Caputo,
2018). The state-of-the-art relies on advanced deep neural network backbones,
known to degenerate 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).

Domain generalization is susceptible to uncertainty as the domain shift from
the source to the target domain is unknown a priori. Hence, uncertainty should
be taken into account during domain-invariant learning. As deep neural
networks are commonly trained by maximum likelihood estimation, they fail to
effectively capture model uncertainty. This tends to make the models
overconfident in their predictions, especially on out-of-distribution data
(Daxberger & Hernández-Lobato, 2019). As a possible solution, approximate
Bayesian inference offers a natural framework to represent prediction
uncertainty (Kristiadi et al., 2020; MacKay, 1992). It 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). Moreover, the prediction uncertainty can be
improved, even when Bayesian approximation is only applied to the last network
layer (Kristiadi et al., 2020; Atanov et al., 2019). These properties make it
appealing to introduce Bayesian learning into the challenging and, as of yet,
unexplored scenario of domain generalization.

In this paper, we address domain generalization under a probabilistic
framework111Source code is publicly available at
https://github.com/zzzx1224/A-Bit-More-Bayesian.git.. To better explore
domain-invariant learning, we introduce weight uncertainty to the model by
leveraging variational Bayesian inference. To this end, we introduce the
principle of domain invariance in a probabilistic formulation and incorporate
it into the variational Bayesian inference framework. This enables us to
explore domain invariance in a principled way to achieve domain-invariant
feature representations and classifiers jointly. To better handle the domain
shifts between seen and unseen domains, we explore our method under the meta-
learning framework (Du et al., 2020; Balaji et al., 2018; Li & Malik, 2017).
The meta-learning setting is utilized to expose the model to domain shift and
mimic the generalization process by randomly splitting source domains into
several meta-source domains and a meta-target domain in each iteration. We
evaluate our method on four widely-used benchmarks for cross-domain object
classification. Our ablation studies demonstrate the benefit of domain-
invariant learning in a probabilistic framework through its synergy with
variational Bayesian inference, as well as the advantage of jointly learning
domain-invariant feature extractors and classifiers for domain generalization.
Our method achieves state-of-the-art mean accuracy on all four benchmarks.

##  2 Methodology

###  2.1 Preliminaries

In domain generalization, we have
$\mathcal{D}{=}{\left\\{D_{i}\right\\}}^{|\mathcal{D}|}_{i=1}{=}\mathcal{S}\cup\mathcal{T}$
as a set of domains, where $\mathcal{S}$ and $\mathcal{T}$ denote the source
and target domains, respectively. $\mathcal{S}$ and $\mathcal{T}$ do not have
any overlap besides sharing the same label space. Data from the target domains
$\mathcal{T}$ is never seen during training. For each domain
$D_{i}\in\mathcal{D}$, we define a joint distribution
$p(\mathbf{x},\mathbf{y})$ on $\mathcal{X}\times\mathcal{Y}$, where
$\mathcal{X}$ and $\mathcal{Y}$ denote the input space and output space,
respectively. We aim to learn a model $f:\mathcal{X}{\rightarrow}\mathcal{Y}$
in the source domains $\mathcal{S}$ that generalizes well to the target
domains $\mathcal{T}$.

We address domain generalization in a probabilistic framework of Bayesian
inference by introducing weight uncertainty into neural networks. To be
specific, we adopt variational Bayesian inference to learn a neural network
that is assumed to be parameterized by weights $\boldsymbol{\theta}$ with a
prior distribution $p(\boldsymbol{\theta})$ and the posterior distribution
$p(\boldsymbol{\theta}|\mathbf{x},\mathbf{y})$, where
$(\mathbf{x},\mathbf{y})$ are samples from the source domain $\mathcal{S}$. To
learn the model, we seek for a variational distribution
$q(\boldsymbol{\theta})$ to approximate
$p(\boldsymbol{\theta}|\mathbf{x},\mathbf{y})$ by minimizing the Kullback-
Leibler divergence
$\mathbb{D}_{\rm{KL}}\big{[}q(\boldsymbol{\theta})||p(\boldsymbol{\theta}|\mathbf{x},\mathbf{y})\big{]}$
between them. This amounts to minimizing the following objective:

$\displaystyle\mathcal{L}_{\rm{Bayes}}=-\mathbb{E}_{q(\boldsymbol{\theta})}[\log
p(\mathbf{y}|\mathbf{x},\boldsymbol{\theta})]+\mathbb{D}_{\rm{KL}}[q(\boldsymbol{\theta})||p(\boldsymbol{\theta})],$
(1)

which is also known as the negative value of the evidence lower bound (ELBO)
(Blei et al., 2017). We learn the model on the source domain $\mathcal{S}$ in
the hope that it will generalize well to the target domain $\mathcal{T}$.

Kristiadi et al. (2020) show that applying Bayesian approximation to the last
layer of a neural network effectively captures model uncertainty. This is also
appealing when dealing with uncertainty in domain generalization since
applying Monte Carlo sampling to the full Bayesian neural network can be
computationally infeasible. In this paper, to obtain a model that generalizes
well across domains in an efficient way, we incorporate variational Bayesian
approximation into the last two network layers to jointly establish domain-
invariant feature representations and classifiers. More analyses are provided
in the experiments and supplementary. To achieve better domain invariance, we
introduce a domain-invariant principle in a probabilistic form under the
Bayesian framework. We use $\boldsymbol{\phi}$ and $\boldsymbol{\psi}$ to
denote the parameters of the feature extractor and classifier. We incorporate
the domain-invariant principle into the inference of posteriors over
$\boldsymbol{\psi}$ and $\boldsymbol{\phi}$, which results in our two-layer
Bayesian network. Next we detail the Bayesian domain-invariant learning of the
network and the objective for optimization.

Figure 1: Illustrative contrast between deterministic invariance (top) and
probabilistic invariance (bottom), where colors indicate domains. The
deterministic invariance tends to minimize the distance between two
deterministic samples ($a\rightarrow b$). The samples come from their
respective distributions, but with limited distributional awareness. This
means that the samples cannot represent the complete distributions. Thus, the
deterministic invariance can lead to small distances between samples, but
large gaps between distributions, as shown in figure (b). In contrast, the
probabilistic invariance directly minimizes the distance between different
distributions ($c\rightarrow d$), which yields a better domain invariance as
most of the samples in the distributions are taken into account.

###  2.2 Bayesian Domain-Invariant Learning

Existing domain generalizations try to achieve domain invariance by minimizing
the distance between intra-class samples from different domains in the hope of
reducing the domain gaps (Motiian et al., 2017). However, individual samples
cannot be assumed to be representative of the distributions of samples from a
certain domain. Therefore, it is preferable to minimize the distributional
distance of intra-class samples from different domains, which can directly
narrow the domain gap. A contrastive illustration is provided in Figure 1.
Motivated by this, we propose to address domain generalization under the
probabilistic modeling by incorporating weight uncertainty into invariant
learning in the variational Bayesian inference framework.

We first introduce the definition of domain invariance in a probabilistic
form, which we adopt to learn domain-invariant feature representations and
classifiers in a unified way. We define a continuous domain space
$\mathfrak{D}$ containing a set of domains
$\left\\{D_{i}\right\\}^{|\mathcal{D}|}_{i=1}$ in $\mathcal{D}$ and a domain-
transform function $g_{\zeta}(\cdot)$ with parameters $\zeta$ in the domain
space. $\left\\{D_{i}\right\\}^{|\mathcal{D}|}_{i=1}$ are discrete samples in
$\mathfrak{D}$. Function $g_{\zeta}(\cdot)$ transforms samples $\mathbf{x}$
from a reference domain to a different domain $D_{\zeta}$ with respect to
$\zeta$, where $\zeta{\thicksim}q(\zeta)$, and different samples $\zeta$ lead
to different sample domains $D_{\zeta}{\thicksim}{\mathfrak{D}}$.

As a concrete example, consider the rotation domains in Rotated MNIST (Ghifary
et al., 2015). $\mathfrak{D}$ is the rotation domain space, containing images
rotated by continuous angles from $0$ to $2\pi$. A sample $\mathbf{x}$ from
any domain in $\mathfrak{D}$ can be transferred into a different domain
$D_{\zeta}$ by:

$g_{\zeta}(\mathbf{x})=\left[\begin{matrix}cos(\zeta)&-sin(\zeta)\\\
sin(\zeta)&-cos(\zeta)\end{matrix}\right]\left[\begin{matrix}x_{1}\\\
x_{2}\end{matrix}\right],$ (2)

where $\zeta\thicksim$ Uniform$(0,2\pi)$. $\zeta$ is the parameter of
$g_{\zeta}(\cdot)$ and different $\zeta$ lead to different rotation angles of
the original sample $\mathbf{x}$, which form another domain $D_{\zeta}$ in
$\mathfrak{D}$. In practice, transformations between domains are more
complicated and the exact forms of $g_{\zeta}(\cdot)$ and $q(\zeta)$ are not
explicitly known.

Based on the above assumptions about $\mathfrak{D}$, $g_{\zeta}(\cdot)$ and
$q(\zeta)$, we introduce our definition of the probabilistic form of domain
invariance as follows.

######  Definition 2.1 (Domain Invariance)

Let $\mathbf{x}_{i}$ be a given sample from domain $D_{i}$ in the domain space
$\mathfrak{D}$, and $\mathbf{x}_{\zeta}{=}g_{\zeta}(\mathbf{x}_{i})$ be a
transformation of $\mathbf{x}_{i}$ in another domain $D_{\zeta}$ from the same
domain space, where $\zeta{\thicksim}q(\zeta)$.
$p_{\boldsymbol{\theta}}(\mathbf{y}|\mathbf{x})$ denotes the output
distribution of input $\mathbf{x}$ with model $\boldsymbol{\theta}$. Model
$\boldsymbol{\theta}$ is domain-invariant in $\mathfrak{D}$ if

$p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})=p_{\boldsymbol{\theta}}(\mathbf{y}_{\zeta}|\mathbf{x}_{\zeta}),\qquad\forall\zeta\thicksim
q(\zeta).$ (3)

Here, we use $\mathbf{y}$ to represent the output from a neural layer with
input $\mathbf{x}$, which can either be the prediction vector from the last
layer or the feature vector from the last convolutional layer of a deep neural
network.

To adopt the domain-invariant principle in the variational Bayesian inference
framework, we reformulate (3) to an expectation form with respect to
$q_{\zeta}$, following Nalisnick & Smyth (2018):

$p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})=\mathbb{E}_{q_{\zeta}}[p_{\boldsymbol{\theta}}(\mathbf{y}_{\zeta}|\mathbf{x}_{\zeta})].$
(4)

According to Definition 2.1, we use the Kullback-Leibler divergence between
the two terms in (4),
$\mathbb{D}_{\rm{KL}}\big{[}p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})||\mathbb{E}_{q_{\zeta}}[p_{\boldsymbol{\theta}}(\mathbf{y}_{\zeta}|\mathbf{x}_{\zeta})]\big{]}$,
to quantify the domain invariance of the model, which will be zero when the
model is domain invariant in the domain space $\mathfrak{D}$. To further
facilitate the computation, we derive the upper bound of the KL divergence:

$\displaystyle\mathbb{D}_{\rm{KL}}$
$\displaystyle\big{[}p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})||\mathbb{E}_{q_{\zeta}}[p_{\boldsymbol{\theta}}(\mathbf{y}_{\zeta}|\mathbf{x}_{\zeta})]\big{]}$
(5)
$\displaystyle\leq\mathbb{E}_{q_{\zeta}}\Big{[}\mathbb{D}_{\rm{KL}}\big{[}p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})||p_{\boldsymbol{\theta}}(\mathbf{y}_{\zeta}|\mathbf{x}_{\zeta})\big{]}\Big{]},$

which can be approximated by Monte Carlo sampling as we usually have access to
samples from different domains. The complete derivation of (5) is provided in
the supplementary material.

We now adopt the probabilistic domain invariance principle in the variational
Bayesian approximation of the last two layers with parameters of
$\boldsymbol{\phi}$ and $\boldsymbol{\psi}$, respectively. The introduced
weight uncertainty results in distributional representations
$p_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$ and predictions
$p_{\boldsymbol{\psi}}(\mathbf{y}|\mathbf{z})$, based on which we derive
domain-invariant learning.

#### Invariant Classifier.

For the classifier with parameters $\boldsymbol{\psi}$ and input features
$\mathbf{z}$, the probabilistic domain-invariant property in (5) is
represented as
$\mathbb{E}_{q_{\zeta}}\big{[}\mathbb{D}_{\rm{KL}}[p_{\boldsymbol{\psi}}(\mathbf{y}_{i}|\mathbf{z}_{i})||p_{\boldsymbol{\psi}}(\mathbf{y}_{\zeta}|\mathbf{z}_{\zeta})]\big{]}$.
Under the Bayesian framework, the predictive distribution
$p_{\boldsymbol{\psi}}(\mathbf{y}|\mathbf{z})$ is obtained by taking the
expectation over the distribution of parameter $\boldsymbol{\theta}$, i.e.,
$p_{\boldsymbol{\psi}}(\mathbf{y}|\mathbf{z}){=}\mathbb{E}_{q(\boldsymbol{\psi})}[p(\mathbf{y}|\mathbf{z},\boldsymbol{\psi})]$.
As the KL divergence is a convex function (Nalisnick & Smyth, 2018), we
further extend it to the upper bound:

$\displaystyle\mathbb{E}_{q_{\zeta}}$
$\displaystyle\Big{[}\mathbb{D}_{\rm{KL}}\big{[}p_{\boldsymbol{\psi}}(\mathbf{y}_{i}|\mathbf{z}_{i})||p_{\boldsymbol{\psi}}(\mathbf{y}_{\zeta}|\mathbf{z}_{\zeta})\big{]}\Big{]}$
(6)
$\displaystyle=\mathbb{E}_{q_{\zeta}}\Big{[}\mathbb{D}_{\rm{KL}}\big{[}\mathbb{E}_{q(\boldsymbol{\psi})}[p(\mathbf{y}_{i}|\mathbf{z}_{i},\boldsymbol{\psi})]||\mathbb{E}_{q(\boldsymbol{\psi})}[p(\mathbf{y}_{\zeta}|\mathbf{z}_{\zeta},\boldsymbol{\psi})]\big{]}\Big{]}$
$\displaystyle\leq\mathbb{E}_{q_{\zeta}}\Big{[}\mathbb{E}_{q(\boldsymbol{\psi})}\mathbb{D}_{\rm{KL}}\big{[}p(\mathbf{y}_{i}|\mathbf{z}_{i},\boldsymbol{\psi})||p(\mathbf{y}_{\zeta}|\mathbf{z}_{\zeta},\boldsymbol{\psi})\big{]}\Big{]},$

which is tractable with unbiased approximation by using Monte Carlo sampling.

In practice, we estimate the expectation over $q(\zeta)$ in an empirical way.
Specifically, in each iteration, we choose one domain from the source domains
$\mathcal{S}$ as the meta-target domain $D_{t}$ and the rest are used as the
meta-source domains ${\left\\{D_{s}\right\\}}^{S}_{s=1}$, where
$S{=}|\mathcal{S}|-1$. Then we use a batch of samples $\mathbf{x}_{s}$ from
each meta-source domain in the same category as $\mathbf{x}_{t}$ to
approximate $\mathbf{x}_{\zeta}{=}g_{\zeta}(\mathbf{x}_{t})$.

Thereby, the domain-invariant classifier is established by minimizing:

$\displaystyle\frac{1}{SN}\sum_{s=1}^{S}\sum_{i=1}^{N}\mathbb{E}_{q(\boldsymbol{\psi})}\Big{[}\mathbb{D}_{\rm{KL}}\big{[}p(\mathbf{y}_{t}|\mathbf{z}_{t},\boldsymbol{\psi})||p(\mathbf{y}_{s}^{i}|\mathbf{z}_{s}^{i},\boldsymbol{\psi})\big{]}\Big{]},$
(7)

where ${\left\\{\mathbf{z}_{s}^{i}\right\\}}^{N}_{i=1}$ are representations of
samples $\mathbf{x}_{s}$ from $D_{s}$, which are in the same category as
$\mathbf{x}_{t}$.

The variational Bayesian inference enables us to flexibly incorporate our
domain invariant principle into different layers of the neural network. To
enhance the domain invariance, we obtain domain-invariant representations by
adopting the principle in the penultimate layer.

#### Invariant Representations.

To obtain domain-invariant representations $p(\mathbf{z}|\mathbf{x})$ with
feature extractor $\boldsymbol{\phi}$, we extend the quantification of our
probabilistic domain invariance in (5) to
$\mathbb{E}_{q_{\zeta}}\big{[}\mathbb{D}_{\rm{KL}}[p_{\boldsymbol{\phi}}(\mathbf{z}_{i}|\mathbf{x}_{i})||p_{\boldsymbol{\phi}}(\mathbf{z}_{\zeta}|\mathbf{x}_{\zeta})]\big{]}$.
Based on the Bayesian layer, the probabilistic distribution of features
$p_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$ will be a factorized Gaussian
distribution if the posterior of $\boldsymbol{\phi}$ is as well. We illustrate
this as follows. Let $\boldsymbol{\phi}$ be the last Bayesian layer in the
feature extractor with a factorized Gaussian posterior and $\mathbf{x}$ be the
input feature of $\boldsymbol{\phi}$. The posterior of the activation
$\mathbf{z}$ is also a factorized Gaussian (Kingma et al., 2015):

$\displaystyle
q(\phi_{i,j})\thicksim\mathcal{N}(\mu_{i,j},\sigma^{2}_{i,j})~{}~{}~{}\forall\phi_{i,j}\in\boldsymbol{\phi}$
(8) $\displaystyle\Rightarrow
p(z_{j}|\mathbf{x},\boldsymbol{\phi})\thicksim{\mathcal{N}(\gamma_{j},\delta_{j}^{2})},$
$\displaystyle\gamma_{j}=\sum_{i=1}^{N}x_{i}\mu_{i,j},~{}~{}~{}~{}\text{and}~{}~{}~{}~{}\delta_{j}^{2}=\sum_{i=1}^{N}x_{i}^{2}\sigma_{i,j}^{2},$

where $z_{j}$ denotes the $j$-th element in $\mathbf{z}$, likewise for
$x_{i}$, and $\phi_{i,j}$ denotes the element at position $(i,j)$ in
$\boldsymbol{\phi}$. Thus, we assume the posterior of the last Bayesian layer
in the feature extractor has a factorized Gaussian distribution. Then, it is
easy to obtain the Bayesian domain invariance of the feature extractor.

In a similar way to (7), we estimate the expectation over $q(\zeta)$
empirically. Therefore, the domain-invariant representations are established
by minimizing:

$\frac{1}{SN}\sum_{s=1}^{S}\sum_{i=1}^{N}\Big{[}\mathbb{D}_{\rm{KL}}\big{[}p(\mathbf{z}_{t}|\mathbf{x}_{t},\boldsymbol{\phi})||p(\mathbf{z}_{s}^{i}|\mathbf{x}_{s}^{i},\boldsymbol{\phi})\big{]}\Big{]},$
(9)

where ${\left\\{\mathbf{x}_{s}^{i}\right\\}}^{N}_{i=1}$ are from $D_{s}$, and
denote the samples in the same category as $\mathbf{x}_{t}$. More details and
an illustration of the Bayesian domain-invariant learning are provided in the
supplementary material.

###  2.3 Objective Function

Having the Bayesian treatment for the last two layers with respect to
parameters $\boldsymbol{\phi}$ and $\boldsymbol{\psi}$, the
$\mathcal{L}_{\rm{Bayes}}$ in (1) is instantiated as the objective with
respect to $\boldsymbol{\psi}$ and $\boldsymbol{\phi}$ as follows:

$\displaystyle{\mathcal{L}}_{\rm{Bayes}}$
$\displaystyle=-\mathbb{E}_{q(\boldsymbol{\psi})}\big{[}\mathbb{E}_{q(\boldsymbol{\phi})}[\log
p(\mathbf{y}|\mathbf{x},\boldsymbol{\psi},\boldsymbol{\phi})]\big{]}$ (10)
$\displaystyle+\mathbb{D}_{\rm{KL}}[q(\boldsymbol{\psi})||p(\boldsymbol{\psi})]+\mathbb{D}_{\rm{KL}}[q(\boldsymbol{\phi})||p(\boldsymbol{\phi})].$

The detailed derivation of (10) is provided in the supplementary material.

By integrating (7) and (9) into (10), we obtain the Bayesian domain-invariant
learning objective as follows:

$\displaystyle{\mathcal{L}}_{\rm{BIL}}$
$\displaystyle=\frac{1}{L}\sum_{\ell=1}^{L}\Big{[}\frac{1}{M}\sum_{m}^{M}[-\log
p(\mathbf{y}_{t}|\mathbf{x}_{t},\boldsymbol{\psi}^{(\ell)},\boldsymbol{\phi}^{(m)})]$
(11)
$\displaystyle+\frac{1}{SN}\sum_{s=1}^{S}\sum_{i=1}^{N}\big{[}\lambda_{\boldsymbol{\psi}}\mathbb{D}_{\rm{KL}}[p(\mathbf{y}_{t}|\mathbf{z}_{t},\boldsymbol{\psi}^{(\ell)})||p(\mathbf{y}^{i}_{s}|\mathbf{z}^{i}_{s},\boldsymbol{\psi}^{(\ell)})]$
$\displaystyle+\lambda_{\boldsymbol{\phi}}\mathbb{D}_{\rm{KL}}[p(\mathbf{z}_{t}|\mathbf{x}_{t},\boldsymbol{\phi})||p(\mathbf{z}^{i}_{s}|\mathbf{x}^{i}_{s},\boldsymbol{\phi})]\big{]}$
$\displaystyle+\mathbb{D}_{\rm{KL}}[q(\boldsymbol{\psi})||p(\boldsymbol{\psi})]+\mathbb{D}_{\rm{KL}}[q(\boldsymbol{\phi})||p(\boldsymbol{\phi})]\Big{]},$

where $\lambda_{\boldsymbol{\psi}}$ and $\lambda_{\boldsymbol{\phi}}$ are
hyperparameters that control the domain-invariant terms. $\mathbf{x}_{t}$ and
$\mathbf{z}_{t}$ denote the input and its feature from $D_{t}$, and
$\mathbf{x}^{i}_{s}$ and $\mathbf{z}^{i}_{s}$ are from $D_{s}$. The posteriors
are set to factorized Gaussian distributions, i.e.,
$q(\boldsymbol{\psi}){=}{\mathcal{N}(\boldsymbol{\mu}_{\boldsymbol{\psi}},\boldsymbol{\sigma}_{\boldsymbol{\psi}}^{2})}$
and
$q(\boldsymbol{\phi}){=}{\mathcal{N}(\boldsymbol{\mu}_{\boldsymbol{\phi}},\boldsymbol{\sigma}_{\boldsymbol{\phi}}^{2})}$.
We adopt the reparameterization trick to draw Monte Carlo samples (Kingma &
Welling, 2014) as
$\boldsymbol{\psi}^{(\ell)}{=}\boldsymbol{\mu}_{\boldsymbol{\psi}}+\epsilon^{(\ell)}*\boldsymbol{\sigma}_{\boldsymbol{\psi}}$,
where $\epsilon^{(\ell)}\thicksim{\mathcal{N}(0,I)}$. Likewise, we draw the
Monte Carlo samples $\boldsymbol{\phi}^{(m)}$ for $q(\boldsymbol{\phi})$.

When implementing our Bayesian invariant learning, to increase the flexibility
of the prior distribution in our Bayesian layers, we place a scale mixture of
two Gaussian distributions as the priors $p(\boldsymbol{\psi})$ and
$p(\boldsymbol{\phi})$ (Blundell et al., 2015):

${\pi\mathcal{N}(0,\boldsymbol{\sigma}_{1}^{2})}+{(1-\pi)\mathcal{N}(0,\boldsymbol{\sigma}_{2}^{2})},$
(12)

where $\boldsymbol{\sigma}_{1}$, $\boldsymbol{\sigma}_{2}$ are set according
to (Blundell et al., 2015) and $\pi$ is chosen by cross-validation. More
experiments and analyses on the hyperparameters $\lambda_{\boldsymbol{\psi}}$,
$\lambda_{\boldsymbol{\phi}}$ and $\pi$ are provided in the supplementary
material.

##  3 Related Work

To address the domain shifts, domain adaptation (Saenko et al., 2010; Long et
al., 2015; Wang et al., 2020) and domain generalization (Muandet et al., 2013;
Li et al., 2017) are the two main settings. Domain adaptation has a key
assumption that the data from the target domain is accessible during training,
which is often invalid in realistic applications. By contrast, no target
domain data is available during training in domain generalization, which
introduces more prediction uncertainty into the problem and makes it more
challenging. Our method is developed specifically for domain generalization.

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) augment the
data by perturbing the input images with adversarial gradients generated by an
auxiliary classifier. Qiao et al. (2020) propose an even more challenging
scenario of domain generalization named single domain generalization, which
only has one source domain, and introduce an adversarial domain augmentation
method that creates “fictitious” yet “challenging” data. Recently, Zhou et al.
(2020) employ a generator to synthesize data from pseudo-novel domains to
augment the source domains, maximizing the distance between the source and
pseudo-novel domains as measured by optimal transport (Peyré & Cuturi, 2019).

Another solution for domain generalization involves learning domain-invariant
features (D’Innocente & Caputo, 2018; Li et al., 2018b, 2017). Muandet et al.
(2013) propose domain-invariant component analysis to learn invariant
transformations by minimizing the dissimilarity across domains. Louizos et al.
(2015) learn invariant representations by a variational auto-encoder (Kingma &
Welling, 2014), introducing Bayesian inference into invariant feature
learning. Both Dou et al. (2019) and Seo et al. (2019) achieve a similar goal
by introducing two complementary losses and employing multiple normalizations.
Li et al. (2019) propose an episodic training algorithm to obtain both a
domain-invariant feature extractor and classifier. Zhao et al. (2020) propose
entropy-regularization to learn domain-invariant features. Seo et al. (2019)
and Zhou et al. (2021) design normalizations to combine different feature
statistics and embed the samples into a domain-invariant feature space.

Meta-learning has also been considered for domain generalization. Li et al.
(2018a) introduce a gradient-based method, i.e., model agnostic meta-learning
(Finn et al., 2017), for domain generalization. Balaji et al. (2018) meta-
learn a regularization function, making their model robust to domain shifts.
Du et al. (2020) propose the meta-variational information bottleneck to learn
domain-invariant representations through episodic training.

Gulrajani & Lopez-Paz (2020) find that with careful implementation, empirical
risk minimization methods outperform many state-of-the-art models in domain
generalization. They claim that model selection is non-trivial for domain
generalization and algorithms for this task should specify their own model
selection criteria, which is important for the completeness and comparability
of the method. In this paper, we implement our method based on the pretrained
ResNet-18, as done for our baseline and most methods we compare against in
Section 4.

Motiian et al. (2017) previously considered representation alignment across
domains in the same class. They propose a classification and contrastive
semantic alignment loss based on the L2 distance between deterministic
features. Different from them, we exploit Bayesian neural networks to learn
domain-invariant representations by minimizing the distance between
probabilistic distributions.

Bayesian neural networks have not yet been explored for domain generalization.
Our method introduces variational Bayesian approximation to both the feature
extractor and classifier of the neural network in conjunction with the newly
introduced domain-invariant principle for domain generalization. The resultant
Bayesian domain-invariant learning combines the representational power of deep
neural networks and variational Bayesian inference.

##  4 Experiments

###  4.1 Datasets and Settings

We conduct our experiments on four widely used benchmarks for domain
generalization and report the mean classification accuracy on target domains.

PACS222https://domaingeneralization.github.io/ (Li et al., 2017) consists of
9,991 images of seven classes from four domains – photo, art-painting, cartoon
and sketch. We follow the “leave-one-out” protocol from (Li et al., 2017,
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.

Office-Home333https://www.hemanthdv.org/officeHomeDataset.html (Venkateswara
et al., 2017) also has four domains: art, clipart, product and real-world.
There are about 15,500 images of 65 categories for object recognition in
office and home environments. We use the same experimental protocol as for
PACS.

Rotated MNIST444http://yann.lecun.com/exdb/mnist/ and Fashion-
MNIST555https://www.kaggle.com/zalando-research/fashionmnist are introduced in
(Piratla et al., 2020) for evaluating domain generalization. For 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 are
considered to have been rotated by $0^{\circ}$. The subset of images is then
rotated by $15^{\circ}$ through $75^{\circ}$ in intervals of $15^{\circ}$,
creating five source domains. The target domains are created by rotations of
$0^{\circ}$ and $90^{\circ}$. These datasets allow us to demonstrate the
generalizability of our model by comparing its performance on in-distribution
and out-of-distribution data.

Table 1: Ablation study on PACS. The “$\checkmark$” and “$\times$” in the “Bayesian” column indicate whether the classifier $\boldsymbol{\psi}$ and feature extractor $\boldsymbol{\phi}$ are Bayesian layers or deterministic layers. In the “Invariant” column, they indicate whether the domain-invariant learning is introduced into the classifier and the feature extractor. Note that with “$\checkmark$” in the Bayesian column, the invariant column denotes Bayesian domain-invariant learning. Otherwise it denotes the deterministic one. The results show that both Bayesian and domain-invariant learning benefit domain generalization, but our Bayesian domain-invariant learning is better. We obtain the best performance with Bayesian domain-invariant learning in both the classifier and feature extractor. We also provide the visualizations of the features for each of the settings in Figure 2.  | Classifier $\boldsymbol{\psi}$ | Feature extractor $\boldsymbol{\phi}$ | PACS
---|---|---|---
ID | Bayesian | Invariant | Bayesian | Invariant | Photo | Art-painting | Cartoon | Sketch | Mean
(a) | $\times$ | $\times$ | $\times$ | $\times$ | 92.85 $\pm$0.21 | 75.12 $\pm$0.48 | 77.44 $\pm$0.26 | 75.72 $\pm$0.47 | 80.28 $\pm$0.42
(b) | $\checkmark$ | $\times$ | $\times$ | $\times$ | 93.89 $\pm$0.29 | 77.88 $\pm$0.53 | 78.20 $\pm$0.39 | 77.75 $\pm$0.75 | 81.93 $\pm$0.22
(c) | $\times$ | $\checkmark$ | $\times$ | $\times$ | 93.95 $\pm$0.51 | 80.03 $\pm$0.72 | 78.03 $\pm$0.77 | 77.83 $\pm$0.52 | 82.46 $\pm$0.67
(d) | $\checkmark$ | $\checkmark$ | $\times$ | $\times$ | 95.21 $\pm$0.26 | 81.25 $\pm$0.76 | 80.67 $\pm$0.73 | 79.31 $\pm$0.94 | 84.11 $\pm$0.39
(e) | $\times$ | $\times$ | $\checkmark$ | $\times$ | 92.81 $\pm$0.35 | 78.66 $\pm$0.56 | 77.90 $\pm$0.40 | 78.72 $\pm$0.86 | 82.02 $\pm$0.26
(f) | $\times$ | $\times$ | $\times$ | $\checkmark$ | 94.17 $\pm$0.35 | 79.75 $\pm$0.68 | 79.51 $\pm$0.98 | 78.31 $\pm$1.11 | 82.94 $\pm$0.53
(g) | $\times$ | $\times$ | $\checkmark$ | $\checkmark$ | 95.15 $\pm$0.26 | 80.96 $\pm$0.69 | 79.57 $\pm$0.85 | 79.15 $\pm$0.98 | 83.71 $\pm$0.65
(h) | $\checkmark$ | $\times$ | $\checkmark$ | $\times$ | 93.83 $\pm$0.19 | 82.13 $\pm$0.41 | 79.18 $\pm$0.48 | 79.03 $\pm$0.78 | 83.54 $\pm$0.34
(i) | $\times$ | $\checkmark$ | $\times$ | $\checkmark$ | 94.12 $\pm$0.22 | 80.52 $\pm$0.61 | 80.39 $\pm$0.81 | 78.53 $\pm$0.95 | 83.39 $\pm$0.52
(j) | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$ |  95.97 $\pm$0.24 |  83.92 $\pm$0.71 |  81.61 $\pm$0.59 |  80.31 $\pm$0.91 |  85.45 $\pm$0.24
Figure 2: Illustration of the benefit of Bayesian domain-invariant learning by
visualization of the feature representations. Colors denote domains, where the
target domain “art-painting” is in violet, and shapes indicate classes. Our
Bayesian domain-invariant learning (shown in the right column), especially
when employed in both the classifier and feature extractor ((j)), achieves
better results than the other cases shown in the ten subfigures, which
correspond to the ten settings in Table 1 (identified by ID).

Settings. For all four benchmarks, we employ a ResNet-18 (He et al., 2016)
pretrained on ImageNet (Deng et al., 2009) as the backbone. During training we
use Adam optimization (Kingma & Ba, 2014) with a learning rate of 0.0001, and
train for 10,000 iterations. In each iteration we choose one source domain as
the meta-target domain. The batch size is 128. To fit the memory footprint, we
choose a maximum number of samples per category and target domain to implement
the domain-invariant learning, i.e., sixteen for PACS, Rotated MNIST and
Fashion-MNIST, and four for Office-Home. We select
$\lambda_{\boldsymbol{\phi}}$ and $\lambda_{\boldsymbol{\psi}}$ based on
validation set performance and summarize their influence in the supplementary
material. The optimal values of $\lambda_{\boldsymbol{\phi}}$ and
$\lambda_{\boldsymbol{\psi}}$ are 0.1 and 100. Parameters $\sigma_{1}$ and
$\sigma_{2}$ in (12) are set to 0.1 and 1.5. The model with the highest
validation set accuracy is employed for evaluation on the target domain.

###  4.2 Results

We first conduct an ablation study on PACS to investigate the effectiveness of
our Bayesian invariant layer for domain generalization. Since the major
contributions of this work are the Bayesian treatment and the probabilistic
domain-invariant principle, we evaluate their effect by individually
incorporating them into the classifier – the last layer – $\boldsymbol{\psi}$
and the feature extractor – the penultimate layer – $\boldsymbol{\phi}$. The
results are shown in Table 1, with a corresponding t-SNE visualization of the
features learned by various settings in Figure 2, following Du et al. (2020).
We also provide a comparison with the state-of-the-art methods on four widely
used benchmarks. More results and visualizations are provided in the
supplementary material.

Benefits of Bayesian Invariant Classifier. In Table 1, rows (a) to (d)
demonstrate the benefits of the Bayesian invariant classifier. Row (a) serves
as our baseline model, which is a vanilla deep convolutional network without
any Bayesian treatment or domain-invariant loss. The backbone is also a
ResNet-18 pretrained on ImageNet. Rows (b), (c) and (d) show the performance
with a Bayesian classifier, a deterministic domain-invariant classifier and
our Bayesian-invariant classifier. By comparing (b) and (a), it is clear that
the Bayesian treatment for the classifier improves the performance, especially
in the “art-painting” and “sketch” domains. The deterministic invariant
property, as shown in row (c), also benefits the performance. Nevertheless,
our Bayesian invariant learning based on the probabilistic framework performs
better. The results demonstrate that the Bayesian invariant learning enhances
the robustness of Bayesian neural networks on out-of-distribution data and
better leverages domain-invariance than the deterministic invariant model.

The three subfigures in the first row of Figure 2 also demonstrate the
benefits of the Bayesian invariant classifier. The Bayesian treatment enlarges
the inter-class distance in all domains, as shown in Figure 2 (b). The
Bayesian layer incorporates weight uncertainty into the predictions and
improves their diversity, which enhances the classification performance on the
source domains as well as the generalization to the target domain. The
deterministic invariant classifier tends to minimize the distance between
samples with the same label. However, the effect is mainly on the source
domain, without obvious impact on the target domain (pink samples), as shown
in Figure 2 (c). By introducing uncertainty into both the classification and
domain-invariant procedure, our Bayesian domain-invariant classifier further
enlarges the inter-class distance in all domains and better generalizes the
domain-invariant property to the target domain (compare Figure 2 (d) to (b)
and (c)).

Benefits of Bayesian Invariant Feature Extractor. The benefit of our Bayesian
domain-invariant principle for the feature extractor is demonstrated in rows
(e), (f), and (g) of Table 1. Similar to the classifier, introducing Bayesian
inference into the feature extractor (comparing row (e) with row (a)) brings a
good accuracy improvement. By comparing (g) to (e) and (f), we observe that
our Bayesian invariant layer achieves consistently better accuracy than the
deterministic invariant model and regular Bayesian layer.

The subfigures in the second row of Figure 2 further demonstrate the
effectiveness of the Bayesian domain-invariant feature extractor. Similar to
the classifier, introducing the Bayesian framework into the feature extractor
also enlarges the inter-class distance of all domains by introducing model
uncertainty, as shown in (e). The deterministic invariant feature extractor
can also minimize the intra-class distance between different domains. However,
while the effect is more obvious than the deterministic domain-invariant
classifier, it is still poor on the target domain (Figure 2 (f)). Compared to
(e) and (f), the Bayesian invariant feature extractor in (g) tends to further
maximize the inter-class distance while minimizing the intra-class distance
between different domains, which obviously enlarges the inter-class distance
between samples from the target domain. Since the Bayesian invariant feature
extractor achieves domain invariance in a probabilistic way, more
uncertainties are taken into account, enabling the model to achieve better
generalization on the target domain. Note that although the Bayesian invariant
classifier in (d) and feature extractor in (g) achieve domain-invariant
learning from different directions of the feature space, they both improve the
performance on the target domain and do not conflict with each other.

Table 2: Performance of different priors on PACS. “Standard” denotes the standard Gaussian prior while “Mixture” denotes the scale mixture prior.  Prior | Photo | Art-painting | Cartoon | Sketch | Mean
---|---|---|---|---|---
Standard | 95.17 $\pm$0.25 | 81.95 $\pm$0.42 | 79.41 $\pm$0.82 | 77.88 $\pm$0.57 | 83.61 $\pm$0.12
Mixture | 95.97 $\pm$0.24 | 83.92 $\pm$0.71 | 81.61 $\pm$0.59 | 80.31 $\pm$0.91 | 85.45 $\pm$0.24
Table 3: Effect of more Bayesian layers on PACS. The “Bayesian” and “Invariant” columns indicate whether the penultimate layer $\boldsymbol{\phi}^{\prime}$ in the feature extractor has a Bayesian and/or domain-invariant property. More Bayesian layers benefit the performance while excessive domain-invariant learning is harmful.  $\boldsymbol{\phi}^{\prime}$ | PACS
---|---
Bayesian | Invariant | Photo | Art-painting | Cartoon | Sketch | Mean
$\times$ | $\times$ | 95.97 $\pm$0.24 | 83.92 $\pm$0.71 | 81.61 $\pm$0.59 | 80.31 $\pm$0.91 | 85.45 $\pm$0.24
$\checkmark$ | $\times$ | 95.69 $\pm$0.20 | 83.28 $\pm$0.74 |  82.06 $\pm$0.25 |  81.00 $\pm$0.55 |  85.51 $\pm$0.13
$\checkmark$ | $\checkmark$ | 95.72 $\pm$0.29 | 82.33 $\pm$0.47 | 81.10 $\pm$0.73 | 80.67 $\pm$1.01 | 84.96 $\pm$0.34
Table 4: Comparison on PACS. Our method achieves the best performance on the “Cartoon” domain, is competitive on the other three domains and obtains the best overall mean accuracy. | Photo | Art-painting | Cartoon | Sketch | Mean
---|---|---|---|---|---
Baseline | 92.85 | 75.12 | 77.44 | 75.72 | 80.28
Carlucci et al. (2019) | 96.03 | 79.42 | 75.25 | 71.35 | 80.51
Dou et al. (2019) | 94.99 | 80.29 | 77.17 | 71.69 | 81.04
Zhao et al. (2020) | 96.65 | 80.70 | 76.40 | 71.77 | 81.38
Piratla et al. (2020) | 94.10 | 78.90 | 75.80 | 76.70 | 81.40
Chattopadhyay et al. (2020) | 93.35 | 76.90 | 80.38 | 75.21 | 81.46
Li et al. (2019) | 93.90 | 82.10 | 77.00 | 73.00 | 81.50
Balaji et al. (2018) | 95.50 | 83.70 | 77.20 | 70.30 | 81.68
Zhou et al. (2020) | 96.20 | 83.30 | 78.20 | 73.60 | 82.83
Seo et al. (2019) | 95.87 | 84.67 | 77.65 | 82.23 | 85.11
Huang et al. (2020) | 95.99 | 83.43 | 80.31 | 80.85 | 85.15
This paper | 95.97 $\pm$0.24 | 83.92 $\pm$0.71 | 81.61 $\pm$0.59 | 80.31 $\pm$0.91 | 85.45 $\pm$0.24

Benefits of Synergistic Bayesian Invariant Learning. The last three rows of
Table 1 show the performance when introducing Bayesian and invariant learning
into both the classifier and feature extractor. Both the Bayesian learning
(row (h)) and deterministic invariant learning (row (i)) in two layers perform
better than introducing the corresponding properties into only one layer of
the model (compared with (b), (c) and (e), (f)) on most domains, as well as in
terms of mean performance. Overall, employing our Bayesian invariant layer in
both the classifier and feature extractor achieves the best performance, as
shown in row (j).

The last row in Figure 2 further demonstrates the benefits of synergistic
Bayesian invariant learning. Incorporating the Bayesian classifier and the
Bayesian feature extractor further maximizes the inter-class distance of all
domains, as comparing (h) to (b) and (e). Introducing deterministic invariance
into both the classifier and feature extractor also improves the domain-
invariant property of features and the generalization to the target domain to
some extent, as shown in subfigure (i). Moreover, by utilizing both the
Bayesian domain-invariant classifier and Bayesian domain-invariant feature
extractor, our method combines their benefits shown in (d) and (g) and
achieves the best performance, as shown in (j). This indicates the benefits of
the synergy between Bayesian inference and probabilistic domain-invariant
learning for domain generalization.

Effect of Scale Mixture Priors. Instead of using an uninformative diagonal
Gaussian prior, we adopt the scale mixture prior (Blundell et al., 2015). To
show the effect of the prior, we conduct experiments with both the standard
Gaussian prior and scale mixture Gaussian prior as in (12). The results are
reported in Table 2. The scale mixture prior achieves better performance on
all domains. We have also done an ablation on the scaling mixture prior by
changing the value of $\pi$ in (12). The prior with $\pi{=}0.5$ achieves the
best performance on the “cartoon” domain of PACS (Figure C.2 in
supplementary), and thereby we use this value in all other experiments.

Effect of More Bayesian Layers. We also experiment with more Bayesian layers
in the feature extractor, as shown in Table 3. The settings of the model in
the first row are the same as row (j) in Table 1. When introducing another
Bayesian layer $\boldsymbol{\phi}^{\prime}$ without domain invariance into the
model, as shown in the second row, the average performance improves slightly.
However, if we introduce the Bayesian domain-invariant learning into
$\boldsymbol{\phi}^{\prime}$ (third row), the overall performance deteriorates
slightly. This may due to the loss of information in the features caused by
the excessive use of domain-invariant learning. In addition, due to the
Bayesian inference and Monte Carlo sampling, more Bayesian layers leads to
higher memory usage and more computations (more detailed discussion in
supplementary material). As such, we prefer to apply the Bayesian invariant
learning only in the last feature extraction layer and the classifier.

Table 5: Comparison on Office-Home. Our method achieves the best performance on the “Art” and “Clipart” domains, while being competitive on the “Product” and “Real” domains. Again we report the best overall mean accuracy. | Art | Clipart | Product | Real | Mean
---|---|---|---|---|---
Baseline | 54.84 | 49.85 | 72.40 | 73.14 | 62.55
Carlucci et al. (2019) | 53.04 | 47.51 | 71.47 | 72.79 | 61.20
Li et al. (2018b) | 56.50 | 47.30 | 72.10 | 74.80 | 62.68
Seo et al. (2019) | 59.37 | 45.70 | 71.84 | 74.68 | 62.90
Huang et al. (2020) | 58.42 | 47.90 | 71.63 | 74.54 | 63.12
Zhou et al. (2020) | 60.60 | 50.10 | 74.80 | 77.00 | 65.63
This paper | 61.81 $\pm$0.36 | 53.27 $\pm$0.37 | 74.27 $\pm$0.35 | 76.31 $\pm$0.24 | 66.42 $\pm$0.18
Table 6: Comparison on Rotated MNIST and Fashion-MNIST. In-distribution accuracy is evaluated on the test sets of MNIST and Fashion-MNIST with rotation angles of $15^{\circ}$, $30^{\circ}$, $45^{\circ}$, $60^{\circ}$ and $75^{\circ}$, while out-of-distribution accuracy is evaluated on test sets with angles of $0^{\circ}$ and $90^{\circ}$. Our method achieves best performance on both the in-distribution and out-of-distribution test sets. | MNIST | Fashion-MNIST
---|---|---
| In-distribution | Out-of-distribution | In-distribution | Out-of-distribution
Baseline | 98.4 | 93.5 | 89.6 | 76.9
Dou et al. (2019) | 98.2 | 93.2 | 86.9 | 72.4
Piratla et al. (2020) | 98.4 | 94.7 | 89.7 | 78.0
This paper | 99.0 $\pm$0.02 | 96.5 $\pm$0.08 | 91.5 $\pm$0.10 | 83.5 $\pm$0.63

#### State-of-the-Art Comparisons.

We compare our method with the state-of-the-art on four datasets. For
comprehensive comparison, we also include a vanilla deep convolutional
ResNet-18 network as a baseline, without any Bayesian treatment, as done in
row (a) of Table 1. The results for each dataset are reported in Tables 4, 5
and 6.

On PACS, in Table 4, our method achieves the best mean accuracy. For each
individual domain, we are competitive with the state-of-the-art and even
exceed all other methods on the “cartoon” domain. On Office-Home, in Table 5,
we again achieve the best mean accuracy. It is worth mentioning that on the
most challenging “art” and “clipart” domains, we also deliver the highest
accuracy, with a good improvement over previous methods. However, Zhou et al.
(2020) and Seo et al. (2019) outperform the proposed model on some domains of
PACS and Office-Home. In (Zhou et al., 2020), the source domains are augmented
by a generator that synthesizes data from pseudo-novel domains, which often
have similar characteristics with the source data. This pays off when the
target data also has similar characteristics to the source domains, as the
pseudo domains are more likely to cover the target domain, as can be seen for
“product” and “real world” in Office-Home and “photo” in PACS. When the test
domain is different from all of the training domains the performance suffers,
e.g., “clipart” in Office-Home and “sketch” in PACS. We highlight that our
method generates domain-invariant representations and classifiers, resulting
in competitive results across all domains and overall. In addition, Seo et al.
(2019) combine batch and instance normalization for domain generalization.
This tactic is effective on PACS, but less so on Office-Home. We attribute
this to the larger number of categories in Office-Home, where instance
normalization is known to make features less discriminative with respect to
object categories (Seo et al., 2019). In contrast, our Bayesian domain-
invariant learning establishes domain-invariant features and predictions in a
probabilistic form by introducing uncertainty into the model, resulting in
good performance on both PACS and Office-Home.

On Rotated MNIST and Fashion-MNIST, 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 6, our method achieves the best
performance on both sets of the two datasets. In particular, our method
improves the classification performance on the out-of-distribution sets,
demonstrating its strong generalizability to unseen domains, which is also
consistent with the findings in Figure 2.

##  5 Conclusion

In this work, we propose a variational Bayesian learning framework for domain
generalization. We introduce Bayesian neural networks into the model, which
are 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 classifier. Our method combines the representational
power of deep neural networks and uncertainty modeling ability of Bayesian
learning, demonstrating effectiveness for domain generalization. Ablation
studies further validate these benefits. Our Bayesian invariant learning sets
a new state-of-the-art on four domain generalization benchmarks.

## Acknowledgements

This work is financially supported by the Inception Institute of Artificial
Intelligence, the University of Amsterdam and the allowance Top consortia for
Knowledge and Innovation (TKIs) from the Netherlands Ministry of Economic
Affairs and Climate Policy.

## References

  * Atanov et al. (2019) Atanov, A., Ashukha, A., Struminsky, K., Vetrov, D., and Welling, M.  The deep weight prior.  In _International Conference on Learning Representations_ , 2019.
  * Balaji et al. (2018) Balaji, Y., Sankaranarayanan, S., and Chellappa, R.  Metareg: Towards domain generalization using meta-regularization.  In _Advances in Neural Information Processing Systems_ , pp. 998–1008, 2018.
  * Bengio et al. (2019) Bengio, Y., Deleu, T., Rahaman, N., Ke, N. R., Lachapelle, S., Bilaniuk, O., Goyal, A., and Pal, C.  A meta-transfer objective for learning to disentangle causal mechanisms.  In _International Conference on Learning Representations_ , 2019.
  * Blei et al. (2017) Blei, D. M., Kucukelbir, A., and McAuliffe, J. D.  Variational inference: A review for statisticians.  _Journal of the American Statistical Association_ , 112(518):859–877, 2017.
  * Blundell et al. (2015) Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D.  Weight uncertainty in neural networks.  _ArXiv_ , abs/1505.05424, 2015.
  * Carlucci et al. (2019) Carlucci, F. M., D’Innocente, A., Bucci, S., Caputo, B., and Tommasi, T.  Domain generalization by solving jigsaw puzzles.  In _IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 2229–2238, 2019.
  * Chattopadhyay et al. (2020) Chattopadhyay, P., Balaji, Y., and Hoffman, J.  Learning to balance specificity and invariance for in and out of domain generalization.  _arXiv preprint arXiv:2008.12839_ , 2020.
  * Daxberger & Hernández-Lobato (2019) Daxberger, E. and Hernández-Lobato, J. M.  Bayesian variational autoencoders for unsupervised out-of-distribution detection.  _arXiv preprint arXiv:1912.05651_ , 2019.
  * Deng et al. (2009) Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L.  Imagenet: A large-scale hierarchical image database.  In _IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 248–255, 2009.
  * Dou et al. (2019) Dou, Q., de Castro, D. C., Kamnitsas, K., and Glocker, B.  Domain generalization via model-agnostic learning of semantic features.  In _Advances in Neural Information Processing Systems_ , pp. 6450–6461, 2019.
  * Du et al. (2020) Du, Y., Xu, J., Xiong, H., Qiu, Q., Zhen, X., Snoek, C. G. M., and Shao, L.  Learning to learn with variational information bottleneck for domain generalization.  In _European Conference on Computer Vision_ , 2020.
  * D’Innocente & Caputo (2018) D’Innocente, A. and Caputo, B.  Domain generalization with domain-specific aggregation modules.  In _German Conference on Pattern Recognition_ , pp. 187–198. Springer, 2018.
  * Finn et al. (2017) Finn, C., Abbeel, P., and Levine, S.  Model-agnostic meta-learning for fast adaptation of deep networks.  In _International Conference on Machine Learning_ , pp. 1126–1135, 2017.
  * Ghifary et al. (2015) Ghifary, M., Bastiaan Kleijn, W., Zhang, M., and Balduzzi, D.  Domain generalization for object recognition with multi-task autoencoders.  In _IEEE International Conference on Computer Vision_ , pp. 2551–2559, 2015.
  * Gulrajani & Lopez-Paz (2020) Gulrajani, I. and Lopez-Paz, D.  In search of lost domain generalization.  _arXiv preprint arXiv:2007.01434_ , 2020.
  * Guo et al. (2017) Guo, C., Pleiss, G., Sun, Y., and Weinberger, K. Q.  On calibration of modern neural networks.  In _International Conference on Machine Learning_ , pp. 1321–1330, 2017.
  * He et al. (2016) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 770–778, 2016.
  * Huang et al. (2020) Huang, Z., Wang, H., Xing, E. P., and Huang, D.  Self-challenging improves cross-domain generalization.  _arXiv preprint arXiv:2007.02454_ , 2020.
  * Ilse et al. (2019) Ilse, M., Tomczak, J. M., Louizos, C., and Welling, M.  Diva: Domain invariant variational autoencoders.  _arXiv preprint arXiv:1905.10427_ , 2019.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kingma & Welling (2014) Kingma, D. P. and Welling, M.  Auto-encoding variational bayes.  In _International Conference on Learning Representations_ , 2014.
  * Kingma et al. (2015) Kingma, D. P., Salimans, T., and Welling, M.  Variational dropout and the local reparameterization trick.  In _Advances in neural information processing systems_ , pp. 2575–2583, 2015.
  * Kristiadi et al. (2020) Kristiadi, A., Hein, M., and Hennig, P.  Being Bayesian, even just a bit, fixes overconfidence in relu networks.  _arXiv preprint arXiv:2002.10118_ , 2020.
  * Krueger et al. (2020) Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Binas, J., Priol, R. L., and Courville, A.  Out-of-distribution generalization via risk extrapolation (rex).  _arXiv preprint arXiv:2003.00688_ , 2020.
  * Li et al. (2017) Li, D., Yang, Y., Song, Y.-Z., and Hospedales, T. M.  Deeper, broader and artier domain generalization.  In _IEEE International Conference on Computer Vision_ , pp. 5542–5550, 2017.
  * Li et al. (2018a) Li, D., Yang, Y., Song, Y.-Z., and Hospedales, T. M.  Learning to generalize: Meta-learning for domain generalization.  In _AAAI Conference on Artificial Intelligence_ , 2018a.
  * Li et al. (2019) Li, D., Zhang, J., Yang, Y., Liu, C., Song, Y.-Z., and Hospedales, T. M.  Episodic training for domain generalization.  In _IEEE/CVF International Conference on Computer Vision_ , pp. 1446–1455, 2019.
  * Li et al. (2018b) Li, H., Jialin Pan, S., Wang, S., and Kot, A. C.  Domain generalization with adversarial feature learning.  In _IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 5400–5409, 2018b.
  * Li & Malik (2017) Li, K. and Malik, J.  Learning to optimize.  In _International Conference on Learning Representations_ , 2017.
  * Long et al. (2015) Long, M., Cao, Y., Wang, J., and Jordan, M.  Learning transferable features with deep adaptation networks.  In _International Conference on Machine Learning_ , pp. 97–105. PMLR, 2015.
  * Louizos & Welling (2017) Louizos, C. and Welling, M.  Multiplicative normalizing flows for variational bayesian neural networks.  In _International Conference on Machine Learning_ , pp. 2218–2227, 2017.
  * Louizos et al. (2015) Louizos, C., Swersky, K., Li, Y., Welling, M., and Zemel, R.  The variational fair autoencoder.  _arXiv preprint arXiv:1511.00830_ , 2015.
  * MacKay (1992) MacKay, D. J.  The evidence framework applied to classification networks.  _Neural computation_ , 4(5):720–736, 1992.
  * Malinin & Gales (2018) Malinin, A. and Gales, M.  Predictive uncertainty estimation via prior networks.  In _Advances in Neural Information Processing Systems_ , pp. 7047–7058, 2018.
  * Motiian et al. (2017) Motiian, S., Piccirilli, M., Adjeroh, D. A., and Doretto, G.  Unified deep supervised domain adaptation and generalization.  In _IEEE International Conference on Computer Vision_ , pp. 5715–5725, 2017.
  * Muandet et al. (2013) Muandet, K., Balduzzi, D., and Schölkopf, B.  Domain generalization via invariant feature representation.  In _International Conference on Machine Learning_ , pp. 10–18, 2013\.
  * Nalisnick & Smyth (2018) Nalisnick, E. and Smyth, P.  Learning priors for invariance.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 366–375, 2018.
  * Nguyen et al. (2015) Nguyen, A., Yosinski, J., and Clune, J.  Deep neural networks are easily fooled: High confidence predictions for unrecognizable images.  In _IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 427–436, 2015.
  * Nguyen et al. (2018) Nguyen, C. V., Li, Y., Bui, T. D., and Turner, R. E.  Variational continual learning.  In _International Conference on Learning Representations_ , 2018.
  * Peyré & Cuturi (2019) Peyré, G. and Cuturi, M.  Computational optimal transport: With applications to data science.  _Foundations and Trends® in Machine Learning_ , 11(5-6):355–607, 2019.
  * Piratla et al. (2020) Piratla, V., Netrapalli, P., and Sarawagi, S.  Efficient domain generalization via common-specific low-rank decomposition.  _arXiv preprint arXiv:2003.12815_ , 2020.
  * Qiao et al. (2020) Qiao, F., Zhao, L., and Peng, X.  Learning to learn single domain generalization.  In _IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 12556–12565, 2020.
  * Saenko et al. (2010) Saenko, K., Kulis, B., Fritz, M., and Darrell, T.  Adapting visual category models to new domains.  In _European Conference on Computer Vision_ , pp. 213–226, 2010\.
  * Seo et al. (2019) Seo, S., Suh, Y., Kim, D., Han, J., and Han, B.  Learning to optimize domain specific normalization for domain generalization.  _arXiv preprint arXiv:1907.04275_ , 2019.
  * Shankar et al. (2018) Shankar, S., Piratla, V., Chakrabarti, S., Chaudhuri, S., Jyothi, P., and Sarawagi, S.  Generalizing across domains via cross-gradient training.  _arXiv preprint arXiv:1804.10745_ , 2018.
  * Venkateswara et al. (2017) Venkateswara, H., Eusebio, J., Chakraborty, S., and Panchanathan, S.  Deep hashing network for unsupervised domain adaptation.  In _IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 5018–5027, 2017.
  * Volpi et al. (2018) Volpi, R., Namkoong, H., Sener, O., Duchi, J. C., Murino, V., and Savarese, S.  Generalizing to unseen domains via adversarial data augmentation.  In _Advances in Neural Information Processing Systems_ , pp. 5334–5344, 2018.
  * Wang et al. (2018) Wang, H., He, Z., Lipton, Z. C., and Xing, E. P.  Learning robust representations by projecting superficial statistics out.  In _International Conference on Learning Representations_ , 2018.
  * Wang et al. (2020) Wang, Z., Cheng, X., Sapiro, G., and Qiu, Q.  A dictionary approach to domain-invariant learning in deep networks.  In _Advances in Neural Information Processing Systems_ , 2020.
  * Zhao et al. (2020) Zhao, S., Gong, M., Liu, T., Fu, H., and Tao, D.  Domain generalization via entropy regularization.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * Zhou et al. (2020) Zhou, K., Yang, Y., Hospedales, T., and Xiang, T.  Learning to generate novel domains for domain generalization.  _arXiv preprint arXiv:2007.03304_ , 2020.
  * Zhou et al. (2021) Zhou, K., Yang, Y., Qiao, Y., and Xiang, T.  Domain generalization with mixstyle.  In _International Conference on Learning Representations_ , 2021.

##  Appendix A Derivation

###  A.1 Derivation of the Upper Bounds of Probabilistic Domain-invariant
Learning

To implement
$\mathbb{E}_{q_{\zeta}}\big{[}p_{\boldsymbol{\theta}}(\mathbf{y}_{\zeta}|\mathbf{x}_{\zeta})\big{]}$
in a tractable way, we derive the upper bound in Section 2.2, which is
achieved via Jensen’s inequality:

$\displaystyle\mathbb{D}_{\rm{KL}}\big{[}p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})||\mathbb{E}_{q_{\zeta}}[p_{\boldsymbol{\theta}}(\mathbf{y}_{\zeta}|\mathbf{x}_{\zeta})]\big{]}$
$\displaystyle=\mathbb{E}_{p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})}\big{[}\log
p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})-\log\mathbb{E}_{q_{\zeta}}[p_{\boldsymbol{\theta}}(\mathbf{y}_{\zeta}|\mathbf{x}_{\zeta})]\big{]}$
(13)
$\displaystyle\leq\mathbb{E}_{p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})}\big{[}\log
p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})-\mathbb{E}_{q_{\zeta}}[\log
p_{\boldsymbol{\theta}}(\mathbf{y}_{\zeta}|\mathbf{x}_{\zeta})]\big{]}$
$\displaystyle=\mathbb{E}_{q_{\zeta}}\Big{[}\mathbb{D}_{\rm{KL}}\big{[}p_{\boldsymbol{\theta}}(\mathbf{y}_{i}|\mathbf{x}_{i})||p_{\boldsymbol{\theta}}(\mathbf{y}_{\zeta}|\mathbf{x}_{\zeta})\big{]}\Big{]}.$

###  A.2 Derivation of Variational Bayesian Approximation for Feature
Extractor ($\boldsymbol{\phi}$) and Classifier ($\psi$) layers

We consider the model with two Bayesian layers $\boldsymbol{\phi}$ and
$\boldsymbol{\psi}$ as the last layer of the feature extractor and the
classifier. The prior distribution of the model is
$p(\boldsymbol{\phi},\boldsymbol{\psi})$, and the true posterior distribution
is $p(\boldsymbol{\phi},\boldsymbol{\psi}|\mathbf{x},\mathbf{y})$. Following
the settings in Section 2.1, we need to learn a variational distribution
$q(\boldsymbol{\phi},\boldsymbol{\psi})$ to approximate the true posterior by
minimizing the KL divergence from $q(\boldsymbol{\phi},\boldsymbol{\psi})$ to
$p(\boldsymbol{\phi},\boldsymbol{\psi}|\mathbf{x},\mathbf{y})$:

$\displaystyle\boldsymbol{\phi}^{*},\boldsymbol{\psi}^{*}=\mathop{\arg\min}\limits_{\boldsymbol{\phi},\boldsymbol{\psi}}\mathbb{D}_{\rm{KL}}\big{[}q(\boldsymbol{\phi},\boldsymbol{\psi})||p(\boldsymbol{\phi},\boldsymbol{\psi}|\mathbf{x},\mathbf{y})\big{]}.$
(14)

By applying the Bayesian rule
$p(\boldsymbol{\phi},\boldsymbol{\psi}|\mathbf{x},\mathbf{y})\propto
p(\mathbf{y}|\mathbf{x},\boldsymbol{\phi},\boldsymbol{\psi})p(\boldsymbol{\phi},\boldsymbol{\psi})$,
the optimization is equivalent to minimizing:

$\displaystyle{\mathcal{L}}_{\rm{Bayes}}=\int
q(\boldsymbol{\phi},\boldsymbol{\psi})\log\frac{q(\boldsymbol{\phi},\boldsymbol{\psi})}{p(\boldsymbol{\phi},\boldsymbol{\psi})p(\mathbf{y}|\mathbf{x},\boldsymbol{\phi},\boldsymbol{\psi})}\mathrm{d}\boldsymbol{\phi}\mathrm{d}\boldsymbol{\psi}=\mathbb{D}_{\rm{KL}}\big{[}q(\boldsymbol{\phi},\boldsymbol{\psi})||p(\boldsymbol{\phi},\boldsymbol{\psi})\big{]}-\mathbb{E}_{q(\boldsymbol{\phi},\boldsymbol{\psi})}\big{[}\log
p(\mathbf{y}|\mathbf{x},\boldsymbol{\phi},\boldsymbol{\psi})\big{]}.$ (15)

With $\boldsymbol{\phi}$ and $\boldsymbol{\psi}$ being independent,

$\displaystyle{\mathcal{L}}_{\rm{Bayes}}$
$\displaystyle=-\mathbb{E}_{q(\boldsymbol{\psi})}\mathbb{E}_{q(\boldsymbol{\phi})}[\log
p(\mathbf{y}|\mathbf{x},\boldsymbol{\psi},\boldsymbol{\phi})]+\mathbb{D}_{\rm{KL}}[q(\boldsymbol{\psi})||p(\boldsymbol{\psi})]+\mathbb{D}_{\rm{KL}}[q(\boldsymbol{\phi})||p(\boldsymbol{\phi})].$
(16)

##  Appendix B Details of Bayesian Domain-invariant Training

In domain generalization, the given domains
$\mathcal{D}{=}{\left\\{D_{i}\right\\}}^{|\mathcal{D}|}_{i=1}$ are divided
into the source domains $\mathcal{S}$ and the target domains $\mathcal{T}$.
During training, as the data from $\mathcal{T}$ is never seen, only the source
domains $\mathcal{S}$ are accessible. In each iteration, as shown in Figure 3,
we randomly choose one source domain from $\mathcal{S}$ as the meta-target
domain $D_{t}$, and the rest of the source domains
${\left\\{D_{s}\right\\}}^{S}_{s=1}$ are treated as the meta-source domains,
where $S{=}|\mathcal{S}|-1$. We randomly select a batch of samples
$\mathbf{x}_{t}$ from $D_{t}$. We also select $N$ samples
${\left\\{\mathbf{x}^{i}_{s}\right\\}}^{N}_{i=1}$, which are in the same
category as $\mathbf{x}_{t}$, from each of the meta-source domains $D_{s}$.
All these samples are sent to the Bayesian invariant feature extractor
$\boldsymbol{\phi}$ after a ResNet18 backbone to get the representations
$\mathbf{z}_{t}$ and ${\left\\{\mathbf{z}^{i}_{s}\right\\}}^{N}_{i=1}$, which
are then sent to the Bayesian invariant classifier $\boldsymbol{\psi}$ to get
the predictions $\mathbf{y}_{t}$ and
${\left\\{\mathbf{y}^{i}_{s}\right\\}}^{N}_{i=1}$. We obtain the Bayesian
invariant objective function for the feature extractor
$\mathcal{L}_{I}(\boldsymbol{\phi})$ by calculating the mean of the KL
divergence of $p(\mathbf{z}_{t}|\mathbf{x}_{t},\boldsymbol{\phi})$ and each
$p(\mathbf{z}^{i}_{s}|\mathbf{x}^{i}_{s},\boldsymbol{\phi})$ as (9) in the
main paper. The Bayesian invariant objective function for feature classifier
$\mathcal{L}_{I}(\boldsymbol{\psi})$ is calculated in a similar way on
$p(\mathbf{y}_{t}|\mathbf{z}_{t},\boldsymbol{\psi})$ and
${\left\\{p(\mathbf{y}^{i}_{s}|z^{i}_{s},\boldsymbol{\psi})\right\\}}^{N}_{i=1}$
as (7) in the main paper. In addition to the Bayesian invariant objective
functions, there are also a cross-entropy loss $\mathcal{L}_{CE}$ on
$p(\mathbf{y}_{t}|\mathbf{x}_{t},\boldsymbol{\psi},\boldsymbol{\phi})$ and two
Kullback-Leibler terms between the posteriors and priors of
$\boldsymbol{\psi}$ and $\boldsymbol{\phi}$ as detailed in (11) in the main
paper.

Figure 3: Illustration of the training phase of our Bayesian domain-invariant
learning. $\mathcal{S}$ denotes the source domains, $\mathcal{T}$ denotes the
target domains, and $\mathcal{D}=\mathcal{S}\cup\mathcal{T}$. $\mathbf{x}$,
$\mathbf{z}$ and $\mathbf{y}$ denote inputs, features and outputs of samples
in each domain. $\mathcal{L}_{I}(\boldsymbol{\psi})$ and
$\mathcal{L}_{I}(\boldsymbol{\phi})$ denote the domain-invariant objective
functions for the classifier and feature representations as detailed in (7)
and (9) in the main paper.

##  Appendix C Ablation Study for Hyperparamters

We also ablate the hyperparameters $\lambda_{\boldsymbol{\phi}}$,
$\lambda_{\boldsymbol{\psi}}$ and $\pi$ to show their effects. The experiments
are conducted on PACS by using cartoon as the target domain. The results are
shown in Figure 4. Specifically, we produce Figure 4 (a) by fixing
$\lambda_{\boldsymbol{\psi}}$ as 100 and adjusting
$\lambda_{\boldsymbol{\phi}}$, Figure 4 (b) by fixing
$\lambda_{\boldsymbol{\phi}}$ as 1 and adjusting $\lambda_{\boldsymbol{\psi}}$
and Figure 4 (c) by adjusting $\pi$ while fixing other settings as in Section
4.1. $\lambda_{\boldsymbol{\phi}}$ and $\lambda_{\boldsymbol{\psi}}$ balance
the influence of the Bayesian learning and domain-invariant learning, and
their optimal values are 1 and 100, respectively. If the values are too small,
the model tends to overfit to source domains as the performance on target data
drops more obviously than on validation data. By contrast, too large values of
harm the overall performance of the model as there are obvious decrease of
accuracy on both validation data and target data. Moreover, $\pi$ 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 $\pi$ is 0.5 according to Figure 4
(c), which demonstrates the effectiveness of using two components in the scale
mixture prior.

Figure 4: Performance on “cartoon” domain in PACS with different
hyperparameters $\lambda_{\phi}$, $\lambda_{\psi}$ and $\pi$. The red line
denotes the accuracy on validation data while the blue line denotes accuracy
on target data. The optimal value of $\lambda_{\phi}$, $\lambda_{\psi}$ and
$\pi$ are 1, 100 and 0.5.

##  Appendix D Additional Ablations on PACS

To further demonstrate the effectiveness of the Bayesian inference and
Bayesian invariant learning, we conduct some more supplementary experiments
with other settings on PACS as shown in Table 7. For more comprehensive
comparisons, we add four more settings with IDs (k), (l), (m), and (n) to the
settings in Table 1.

Comparing (k) and (l) with (d), we find that when employing the Bayesian
invariant classifier $\boldsymbol{\psi}$ in the model, introducing Bayesian or
deterministic invariant learning into the last layer of the feature extractor
$\boldsymbol{\phi}$ both further improves the overall performance. Moreover,
introducing Bayesian invariant learning into $\boldsymbol{\phi}$ achieves the
best improvements as shown in row (j). A similar phenomenon occurred on the
classifier as comparing (m), (n) and (j) with (g). When there is a Bayesian
invariant layer in the last layer of the feature extractor, introducing
Bayesian learning or deterministic invariant learning both improves the
performance, while Bayesian invariant learning achieves the best result. This
further indicates the benefits of introducing uncertainty into the model and
the invariant learning under the Bayesian framework.

Table 7: More detailed ablation study on PACS. Compared to Table 1 we add four more settings with IDs (k), (l), (m) and (n). Bayesian inference benefits domain generalization. Our Bayesian invariant learning achieves better performance than the deterministic ones.  | Classifier $\boldsymbol{\psi}$ | Feature extractor $\boldsymbol{\phi}$ | PACS
---|---|---|---
ID | Bayesian | Invariant | Bayesian | Invariant | Photo | Art-painting | Cartoon | Sketch | Mean
(a) | $\times$ | $\times$ | $\times$ | $\times$ | 92.85 $\pm$0.21 | 75.12 $\pm$0.48 | 77.44 $\pm$0.26 | 75.72 $\pm$0.47 | 80.28 $\pm$0.42
(b) | $\checkmark$ | $\times$ | $\times$ | $\times$ | 93.89 $\pm$0.29 | 77.88 $\pm$0.53 | 78.20 $\pm$0.39 | 77.75 $\pm$0.75 | 81.93 $\pm$0.22
(c) | $\times$ | $\checkmark$ | $\times$ | $\times$ | 93.95 $\pm$0.51 | 80.03 $\pm$0.72 | 78.03 $\pm$0.77 | 77.83 $\pm$0.52 | 82.46 $\pm$0.67
(d) | $\checkmark$ | $\checkmark$ | $\times$ | $\times$ | 95.21 $\pm$0.26 | 81.25 $\pm$0.76 | 80.67 $\pm$0.73 | 79.31 $\pm$0.94 | 84.11 $\pm$0.39
(e) | $\times$ | $\times$ | $\checkmark$ | $\times$ | 92.81 $\pm$0.35 | 78.66 $\pm$0.56 | 77.90 $\pm$0.40 | 78.72 $\pm$0.86 | 82.02 $\pm$0.26
(f) | $\times$ | $\times$ | $\times$ | $\checkmark$ | 94.17 $\pm$0.35 | 79.75 $\pm$0.68 | 79.51 $\pm$0.98 | 78.31 $\pm$1.11 | 82.94 $\pm$0.53
(g) | $\times$ | $\times$ | $\checkmark$ | $\checkmark$ | 95.15 $\pm$0.26 | 80.96 $\pm$0.69 | 79.57 $\pm$0.85 | 79.15 $\pm$0.98 | 83.71 $\pm$0.65
(m) | $\checkmark$ | $\times$ | $\checkmark$ | $\checkmark$ | 95.87 $\pm$0.48 | 81.15 $\pm$0.49 | 79.39 $\pm$0.33 | 80.15 $\pm$1.16 | 84.14 $\pm$0.41
(n) | $\times$ | $\checkmark$ | $\checkmark$ | $\checkmark$ | 95.39 $\pm$0.12 | 82.32 $\pm$0.37 | 80.27 $\pm$0.58 | 79.61 $\pm$0.93 | 84.40 $\pm$0.45
(h) | $\checkmark$ | $\times$ | $\checkmark$ | $\times$ | 93.83 $\pm$0.19 | 82.13 $\pm$0.41 | 79.18 $\pm$0.48 | 79.03 $\pm$0.78 | 83.54 $\pm$0.34
(i) | $\times$ | $\checkmark$ | $\times$ | $\checkmark$ | 94.12 $\pm$0.22 | 80.52 $\pm$0.61 | 80.39 $\pm$0.81 | 78.53 $\pm$0.95 | 83.39 $\pm$0.52
(j) | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$ |  95.97 $\pm$0.24 |  83.92 $\pm$0.71 |  81.61 $\pm$0.59 |  80.31 $\pm$0.91 |  85.45 $\pm$0.24

##  Appendix E Extra Visualizations

Figure 5: Visualization of feature representations of the target domain. The
subfigures have the same experimental settings as the experiments in Table 1
and Figure 2 in the main paper. To obtain more obvious contrast between the
subfigures, we use different colors to denote different categories. The target
domain is “art-painting”, the same as in Figure 2. We obtain a similar
conclusion to Section 4.2, where by introducing uncertainty into both the
model and the domain-invariant learning, our Bayesian invariant learning
achieves better performance and generalization to the target domain.

To further observe and analyze the benefits of the Bayesian invariant
learning, we conduct more detailed visualizations of the features in Figure 5
and Figure 6. Visualizations in Figure 5 show the features of all categories
from the target domain only, while visualizations in Figure 6 show features of
only one category from all domains. Similar to Figure 2 in Section 4.2 of the
main paper, the visualization is conducted on the PACS dataset and the target
domain is “art-painting”. The chosen category in Figure 6 is “horse”. We also
visualize the features on the rotated MNIST and Fashion MNIST datasets in
Figure 7. The results further demonstrate the effectiveness of our Bayesian
invariant learning on both in-distribution data and out-of-distribution data.

###  E.1 Visualizations on the Target Domain of PACS

Figure. 5 provides a more intuitive observation of the benefits of Bayesian
domain-invariant learning on the target domain.

  * •

Bayesian learning benefits classification on the target domain as shown in the
“Bayesian” column of subfigures. Comparing (b), (c) to (a), it is obvious that
by introducing uncertainty into the model, both the Bayesian learning on the
classifier and feature extractor enlarge the inter-class distance on the
target domain. Samples from different categories tend to align on a thinner
line, which makes them easier to classify. This phenomenon is more obvious
when introducing Bayesian invariant learning into both the classifier and the
feature extractor, as shown in (h).

  * •

Bayesian invariant learning is better than deterministic invariant learning.
Comparing the subfigures in the “Bayesian invariant” column to those in the
“Deterministic invariant” column demonstrates the benefits of our Bayesian
invariant learning. The Bayesian invariant classifier achieves larger inter-
class distances than the deterministic invariant one, as shown in (d) and (c).
The Bayesian invariant feature extractor pushes the features in all categories
away from the center as shown in (g), which also results in larger inter-class
distance. Moreover, employing the Bayesian invariant learning to both the
classifier and the feature extractor conducts more obvious improvements, as
comparing (j) to (i).

  * •

Bayesian invariant learning benefits generalization to the target domain by
introducing uncertainty into both model and domain-invariant learning. It is
interesting to look at subfigures (b), (d), and (h). The visualizations in (d)
are similar to (h), which also occurs in Figure 2 in the main paper. Compared
to (b), (h) introduces more uncertainty by employing Bayesian inference in the
feature extractor, while (d) applies the Bayesian invariant learning based on
the Bayesian classifier. The similar observation in comparing (d) to (h)
indicates that the benefits of our Bayesian invariant learning can be
attributed to the uncertainty introduced into the domain-invariant learning.
The objective function in (6) in the main paper also indicates this
conclusion. Note that although both introduce uncertainty into the domain-
invariant learning, visualization in (g) is different from (d). This is due to
the uncertainty in the domain-invariant classifier and domain-invariant
features having different effects on the feature space. As shown in (d) and
(g), both the Bayesian invariant classifier and feature extractor enlarge the
inter-class distance on the target domain, though from different directions.

###  E.2 Visualizations of Samples in One Class from Different Domains on PACS

Figure 6 provides a deeper insight into the intra-class feature distributions
of the same category from different domains.

  * •

Introducing uncertainty into the model gathers features from different domains
to the same manifold. By comparing (b), (e), and (h) to (a), we find that
introducing Bayesian inference into the model tends to gather the features to
a certain manifold. This is even more obvious with more Bayesian layers as
shown in (h). The same phenomenon can be found in subfigure (d), where
uncertainty is introduced in both the classifier and the domain-invariant
learning. Gathering the features into a certain manifold is beneficial to
domain generalization as it facilitates classification, both on the source
domains and the target domain.

  * •

Bayesian invariant learning enables better generalization to the target domain
than the deterministic counterpart. The deterministic invariant learning
minimizes the distance between deterministic samples from different source
domains. That tends to be overfitting on the source domains. As shown in (c),
(f), and (i), the samples from source domains are clustered well, while some
samples from the target domains tend to be out of the cluster. There is even
an obvious gap when applying deterministic invariant learning on both
classifier and feature extractor, as shown in (i). By contrast, our Bayesian
invariant learning introduces uncertainty into domain-invariant learning.
Thus, it achieves better mixture of features from the source domains and the
target domain as shown in (d), (g), and (j). This indicates the Bayesian
invariant learning enables better generalization to the target domain than the
deterministic one.

###  E.3 Visualizations on Rotated MNIST and Fashion-MNIST

Figure 6: Visualization of feature representations of one category. All
samples are from the “horse” category with colors denoting different domains.
The target domain is “art-painting” (violet). The “Bayesian” column shows
Bayesian inference benefits domain generalization by gathering features from
different domains to the same manifold. Comparing the subfigures in the
“Bayesian invariant” column to those in the “Deterministic invariant” column
indicates that our Bayesian invariant learning has better generalization to
the target domain than the deterministic one.

To further demonstrate the effectiveness of our Bayesian invariant learning,
we also visualize the features on rotated MNIST and Fashion MNIST, as shown in
Figure 7. 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 as well
as the out-of-distribution set, and clusters the out-of-distribution samples
of the same categories better, especially in the rotated Fashion-MNIST
dataset.

Figure 7: Visualization of feature representations in rotated MNIST and
rotated Fashion-MNIST datasets. Samples from the in-distribution and out-of-
distribution sets are in red and blue, respectively. Different shapes denote
different categories. Compared to the baseline, our Bayesian invariant
learning achieves better performance on both the in-distribution and out-of-
distribution sets in each dataset, and especially on the out-of-distribution
set from the Fashion-MNIST benchmark.

##  Appendix F Computational Cost Analysis of Different Number of Bayesian
Layers

Table 8: Computational cost in FLOPs, parameters and GPU memory usage for different number of Bayesian layers in the model. Number of Bayesian layers | Monte Carlo samples | Extra FLOPs (M) | Extra parameters (M) | Memory usage (G)
---|---|---|---|---
0 | 0 | 0 | 0 | 19.4
1 | 10 | 0.04 | 0.004 | 19.4
2 | 10 | 2.98 | 0.530 | 19.5
3 | 10 | 32.42 | 1.060 | 22.6
1 | 15 | 0.05 | 0.004 | 19.4
2 | 15 | 4.74 | 0.530 | 19.8
3 | 15 | 75.01 | 1.060 | 25.5
1 | 20 | 0.07 | 0.004 | 19.4
2 | 20 | 6.68 | 0.530 | 20.4
3 | 20 | 138.77 | 1.060 |  $>$ 31.7

To quantify the effect of the Bayesian framework on the computations and
parameters of the model, we show the FLOPs and parameters for different
numbers of Bayesian layers in Table 8. The parameters of Bayesian layers are
irrelevant to the batch size and the number of Monte Carlo samples. As the
number of Bayesian layers increases, the extra parameters brought by the
Bayesian layers grow smoothly, which will not cause much trouble. However, the
extra FLOPs will have a significant increase with more Bayesian layers in the
model. In addition, the FLOPs are related to the number of Monte Carlo samples
and batch size. As shown in the third column, a bigger number of Monte Carlo
samples leads to much higher FLOPs, especially with more Bayesian layers in
the model. The extra FLOPs in the table are based on one image with size (224,
224, 3). When the batchsize becomes bigger, the FLOPs will also increase.

To be more intuitive, we also show the GPU memory usage of different numbers
of Bayesian layers. The experiments are conducted on the PACS dataset based on
a single Tesla V100 GPU. The batch size is 128, and the number of samples per
category and meta-source domain is 32. The GPU memory increases slightly when
employing one or two Bayesian layers into the model, especially when the
number of Monte Carlo samples is small. However, when introducing three
Bayesian layers into the model, there will be a significant increase in GPU
memory usage. When the number of Monte Carlo samples is bigger, such as 20,
the model with three Bayesian layers will even run out of the GPU memory
during training, as shown in the last row.

Based on Table 8, we finally apply the Bayesian invariant learning only in the
last feature extraction layer and the classifier. The batch size in our
experiment is set to 128 and the number of Monte-Carlo sampling in each
Bayesian layer is 10.

##  Appendix G Failure Cases

Figure 8: Some failure cases of our Bayesian invariant learning on PACS. The
numbers associated with each image are the top two prediction probabilities of
our method, with ground truth labels in red. Our method fails to make the
correct predictions, but provides high probabilities for the true label,
indicating the effectiveness of introducing uncertainty.

Finally, we consider several failure cases of our method. As shown in Figure
8, a complex scene makes it challenging to make correct predictions, e.g.,
(b), (d) in the “photo” domain and (d) in the “cartoon” domain. The failure
cases in the “art painting” domain show the generalization performance of the
method still needs further improvements, especially for samples with a very
special and specific domain appearance. Our method also gets confused when
samples have objects of different categories in the same image, e.g., (a) and
(c) in the “photo” domain, or cluster the characteristics of different
categories to the same sample, e.g., (a) and (b) in the “cartoon” domain and
“sketch” domain. Although our Bayesian invariant learning fails to make the
correct predictions in these challenging cases, it does provide reasonable
probabilities for possible categories, which we attribute to the introduced
model uncertainty and domain-invariant learning.