# How Does Distribution Matching Help Domain Generalization: An Information-
theoretic Analysis

Yuxin Dong, Tieliang Gong, Hong Chen, Shuangyong Song, Weizhan Zhang, Chen Li
Corresponding author: Tieliang Gong.Y. Dong (yxdong9805@gmail.com), T. Gong
(adidasgtl@gmail.com), W. Zhang<EMAIL_ADDRESS>and C. Li
<EMAIL_ADDRESS>are with the School of Computer Science and Technology,
Xi’an Jiaotong University, Xi’an 710049, China.H. Chen
<EMAIL_ADDRESS>is with the College of Science, Huazhong Agriculture
University, Wuhan 430070, China.S. Song<EMAIL_ADDRESS>is with the
China Telecom Corporation, Beijing 100033, China.

###### Abstract

Domain generalization aims to learn invariance across multiple training
domains, thereby enhancing generalization against out-of-distribution data.
While gradient or representation matching algorithms have achieved remarkable
success, these methods generally lack generalization guarantees or depend on
strong assumptions, leaving a gap in understanding the underlying mechanism of
distribution matching. In this work, we formulate domain generalization from a
novel probabilistic perspective, ensuring robustness while avoiding overly
conservative solutions. Through comprehensive information-theoretic analysis,
we provide key insights into the roles of gradient and representation matching
in promoting generalization. Our results reveal the complementary relationship
between these two components, indicating that existing works focusing solely
on either gradient or representation alignment are insufficient to solve the
domain generalization problem. In light of these theoretical findings, we
introduce IDM to simultaneously align the inter-domain gradients and
representations. Integrated with the proposed PDM method for complex
distribution matching, IDM achieves superior performance over various baseline
methods.

###### Index Terms:

Information Theory, Domain Generalization, Out-of-distribution Generalization,
Generalization Analysis.

##  I Introduction

Distribution shifts are prevalent in various real-world learning contexts,
often leading to machine learning systems overfitting environment-specific
correlations that may negatively impact performance when facing out-of-
distribution (OOD) data [1, 2, 3, 4]. Domain generalization (DG) is then
introduced to address this challenge: By assuming the training data
constitutes multiple domains that share some invariant underlying
correlations, DG algorithms then attempt to learn this invariance so that
domain-specific variations do not affect the model’s performance. To this end,
various DG approaches have been proposed, including invariant representation
learning [5, 6], adversarial learning [7, 8], causal inference [9, 10],
gradient manipulation [11, 12, 13], and robust optimization [14, 15, 16].

DG is typically formulated as an average-case [17, 18] or worst-case [9, 14]
optimization problem, which however either lacks robustness against OOD data
[9, 19] or leads to overly conservative solutions [16]. In this paper, we
introduce a novel probabilistic formulation that aims to minimize the gap
between training and test-domain population risks with high probability. Our
comprehensive generalization analysis then reveals that the input-output
mutual information and the representation space covariate shift are pivotal in
controlling this domain-level generalization gap, which could be achieved by
aligning inter-domain gradients and representations, respectively.

Although distribution matching techniques are already well-explored in
existing DG literature, these methods generally lack generalization guarantees
or depend on strong assumptions, e.g. controllable invariant features [11],
quadratic bowl loss landscape [13] or Lipschitz-continuous gradients [20]. In
contrast, we derive instructive generalization bounds by leveraging a relaxed
i.i.d domain assumption [16], which is easily satisfied in practice. Our
results indicate that combining gradient and representation matching
effectively minimizes the domain-level generalization gap. Crucially, we
reveal the complementary nature of these two components, highlighting that
neither of them alone is sufficient to solve the DG problem.

In light of these theoretical findings, we propose inter-domain distribution
matching (IDM) for high-probability DG by simultaneously aligning gradients
and representations across training domains. Furthermore, we point out the
limitations of traditional distribution alignment techniques, especially for
high-dimensional and complex probability distributions. To circumvent these
issues, we further propose per-sample distribution matching (PDM) by slicing
and aligning individual sorted data points. IDM jointly working with PDM
achieves superior performance on the Colored MNIST dataset [9] and the
DomainBed benchmark [21]. Our primary contributions can be summarized as
follows:

  * •

Probabilistic formulation: We introduce a probabilistic perspective for
evaluating DG algorithms, focusing on their ability to minimize the domain-
level generalization gap with high probability. Our approach leverages milder
assumptions about the domains and enables generalization analysis with
information-theoretic tools.

  * •

Information-theoretic insights: Our analysis comprehensively elucidates the
role of gradient and representation matching in promoting domain
generalization. Most importantly, we reveal the complementary relationship
between these two components, indicating that neither of them alone is
sufficient to solve the DG problem.

  * •

Novel algorithms: We propose IDM for high-probability DG by simultaneously
aligning inter-domain gradients and representations, and PDM for complex
distribution matching by slicing and aligning individual sorted data points.
IDM jointly working with PDM achieves superior performance over various
baseline methods.

##  II Problem Setting

We denote random variables by capitalized letters ($X$), their realizations by
lower-case letters ($x$), and the corresponding spaces by calligraphic letters
($\mathcal{X}$). Let $\mathcal{Z}=\mathcal{X}\times\mathcal{Y}$ be the
instance space of interest, where $\mathcal{X}$ and $\mathcal{Y}$ are the
input space and the label space, respectively. Let $\mathcal{W}$ be the
hypothesis space, each $w\in\mathcal{W}$ characterizes a predictor $f_{w}$:
$\mathcal{X}\mapsto\mathcal{Y}$, comprised of an encoder $f_{\phi}$:
$\mathcal{X}\mapsto\mathcal{R}$ and a classifier $f_{\psi}$:
$\mathcal{R}\mapsto\mathcal{Y}$ with the assist of the representation space
$\mathcal{R}$.

Following [16], we assume that there exists a distribution $\nu$ over the
space of possible environments $\mathcal{D}$, where each domain
$d\in\mathcal{D}$ corresponds to a specific data-generating distribution
$\mu_{d}=P_{Z|D=d}$. The source $D_{s}=\\{D_{i}\\}_{i=1}^{m}$ and target
$D_{t}=\\{D_{k}\\}_{k=1}^{m^{\prime}}$ domains are both random variables
sampled from $\nu$. Let $S=\\{S_{i}\\}_{i=1}^{m}$ denote the training dataset,
with each subset $S_{i}=\\{Z_{j}^{i}\\}_{j=1}^{n}$ containing $n$ i.i.d data
sampled from $\mu_{D_{i}}$. The task is to design algorithm
$\mathcal{A}:\mathcal{D}^{m}\mapsto\mathcal{W}$, taking $D_{s}$ as the input
(with proxy $S$) and providing possibly randomized hypothesis
$W=\mathcal{A}(D_{s})$. Given the loss function
$\ell:\mathcal{Y}\times\mathcal{Y}\mapsto\mathbb{R}^{+}$, the ability of some
hypothesis $w\in\mathcal{W}$ to generalize in average is evaluated by the
global population risk:

$L(w)=\mathbb{E}_{D\sim\nu}[L_{D}(w)]=\mathbb{E}_{Z\sim\mu}[\ell(f_{w}(X),Y)],$

where $L_{d}(w)=\mathbb{E}_{Z\sim\mu_{d}}[\ell(f_{w}(X),Y)]$ is the domain-
level population risk. Since $\nu$ is unknown, only the source and target-
domain population risks are tractable in practice:

$L_{s}(w)=\frac{1}{m}\sum_{i=1}^{m}L_{D_{i}}(w),L_{t}(w)=\frac{1}{m^{\prime}}\sum_{k=1}^{m^{\prime}}L_{D_{k}}(w).$

Main Assumptions. We list the assumptions considered in our theoretical
analysis as follows:

######  Assumption II.1.

$D_{t}$ is independent of $D_{s}$.

######  Assumption II.2.

$\ell(\cdot,\cdot)$ is bounded in $[0,M]$.

######  Assumption II.3.

$\ell(f_{w}(X),Y)$ is $\sigma$-subgaussian w.r.t $Z\sim\mu$ for any
$w\in\mathcal{W}$.

######  Assumption II.4.

$\ell(\cdot,\cdot)$ is symmetric and satisfies the triangle inequality, i.e.
for any $y_{1},y_{2},y_{3}\in\mathcal{Y}$,
$\ell(y_{1},y_{2})=\ell(y_{2},y_{1})$ and
$\ell(y_{1},y_{2})\leq\ell(y_{1},y_{3})+\ell(y_{3},y_{2})$.

######  Assumption II.5.

$\ell(f_{w}(X),Y)$ is $\beta$-Lipschitz w.r.t a metric $c$ on $\mathcal{Z}$
for any $w\in\mathcal{W}$, i.e. for any $z_{1},z_{2}\in\mathcal{Z}$,
$\lvert\ell(f_{w}(x_{1}),y_{1})+\ell(f_{w}(x_{2}),y_{2})\rvert\leq\beta
c(z_{1},z_{2})$.

Subgaussianity (Assumption II.3) is one of the most common assumptions for
information-theoretic generalization analysis [22, 23, 24, 25]. Notably,
Assumption II.2 is a strengthened version of Assumption II.3, since any
$[0,M]$-bounded random variable is always $M/2$-subgaussian. Lipschitzness
(Assumption II.5) is a crucial prerequisite for stability analysis and has
also been utilized in deriving Wasserstein distance generalization bounds [26,
27, 28, 29, 30, 31]. Assumption II.4 is fulfilled when distance functions,
such as mean absolute error (MAE) and 0-1 loss, are used as loss functions.
This assumption has also been examined in previous studies [32, 33, 34].

High-Probability DG. The classical empirical risk minimization (ERM)
technique, which minimizes the average-case risk: $\min_{w}L(w)$, is found
ineffective in achieving invariance across different environments [9, 19]. To
overcome this limitation, recent works [15, 35, 12, 13, 36, 37] have cast DG
as a worst-case optimization problem: $\min_{w}\max_{d}L_{d}(w)$. However,
this approach is generally impractical without strong assumptions made in the
literature [38, 16], e.g. linearity of the underlying causal mechanism [9, 15,
35], or strictly separable spurious and invariant features [37]. On the
contrary, we propose the following high-probability objective by leveraging
the mild Assumption II.1:

######  Problem II.6.

(High-Probability DG)

$\min_{\mathcal{A}}\mathbb{E}[L_{s}(W)],\ \textrm{ s.t. }\ \mathbb{P}\\{\lvert
L_{t}(W)-L_{s}(W)\rvert\geq\epsilon\\}\leq\delta.$

This formulation is directly motivated by the intuition that the optimal
algorithm $\mathcal{A}$ should be chosen in consideration of minimizing the
generalization gap between training-domain $L_{s}(W)$ and test-domain
$L_{t}(W)$ population risks. The probability is taken over both the sampled
domains ($D_{s}$ and $D_{t}$) and the learning algorithm ($W$). Notably, our
Assumption II.1 is significantly weaker than the previously adopted i.i.d
domain assumption [16] by allowing correlations between source (or target)
domains, and should be trivially satisfied in practice.

##  III Generalization Analysis

The primary goal of DG is to tackle the distribution shift problem, raised by
the variation in the data-generating distribution $\mu_{d}$ for different
environments $d$. This inconsistency can be quantified by the mutual
information $I(Z;D)$ between the data pair $Z$ and the environment identifier
$D$, which can be further decomposed into:

$I(Z;D)\textrm{ (distribution shift)}=I(X;D)\textrm{ (covariate
shift)}+I(Y;D|X)\textrm{ (concept shift)}.$ (1)

While $D$ is binary to distinguish training and test samples in [39, 40], we
extend this concept to any discrete or continuous space, provided that each
$d\in\mathcal{D}$ corresponds to a distinct data distribution $\mu_{d}$. The
right hand side (RHS) characterizes the changes in the marginal input
distribution $P_{X}$ (covariate shift) as well as the predictive distribution
$P_{Y|X}$ (concept shift). We first show that the achievable level of average-
case risk $L(w)$ is constrained by the degree of concept shift as following:

######  Proposition III.1.

For any predictor $Q_{Y|X}$, we have

$\mathrm{KL}(P_{Y|X,D}\,\|\,Q_{Y|X})\geq I(Y;D|X).$

When $\ell$ represents the cross-entropy loss, the population risk of
predictor $Q$ on domain $d$ can be represented as the KL divergence
$\mathrm{KL}(P_{Y|X,D=d}\,\|\,Q_{Y|X})$, provided that $H(Y|X,D)=0$ (i.e. the
label can be entirely inferred from $X$ and $D$). This implies that any model
fitting well in training domains will suffer from strictly positive risks in
test domains once concept shift is induced. This implication verifies the
trade-off between optimization and generalization as we characterized in
Problem II.6, and highlights the inherent difficulty of solving the DG
problem.

###  III-A Decomposing the Generalization Gap

We further demonstrate that by connecting source and target-domain population
risks via the average-case risk $L(W)$, one can decompose the constraint of
Problem II.6 into training and test-domain generalization gaps. To be
specific, since the predictor $W$ is trained on the source domains $D_{s}$, it
is reasonable to assume that $W$ achieves lower population risks on $D_{s}$
than on average, i.e. $L_{s}(W)\leq L(W)$. Moreover, since the sampling
process of test domains is independent of the hypothesis, the test-domain
population risk $L_{t}(W)$ is an unbiased estimate of $L(W)$. Combining these
two observations, it is natural to assume that $L_{s}(W)\leq L(W)\approx
L_{t}(W)$, implying that the average-case risk $L(W)$ acts as a natural bridge
between the two. For any constant $\lambda\in(0,1)$, one can prove that:

$\mathbb{P}\\{\lvert
L_{s}(W)-L_{t}(W)\rvert\geq\epsilon\\}\leq\mathbb{P}\\{\lvert
L_{s}(W)-L(W)\rvert\geq\lambda\epsilon\\}+\mathbb{P}\\{\lvert
L_{t}(W)-L(W)\rvert\geq(1-\lambda)\epsilon\\}.$

While the first event heavily correlates with the hypothesis $W$, the second
event is instead hypothesis-independent. This observation inspires us to
explore both hypothesis-based and hypothesis-independent bounds to address
source and target-domain generalization errors, respectively.

###  III-B Source-domain Generalization

We first provide a sufficient condition for source-domain generalization. Our
results are motivated by recent advancements in generalization analysis within
the information-theoretic framework [41, 42]. Specialized to our problem, we
quantify the changes in the hypothesis once the training domains are observed
through the input-output mutual information $I(W;D_{i})$:

######  Theorem III.2.

If Assumption II.2 holds, then

$\mathbb{P}\left\\{\left\lvert
L_{s}(W)-L(W)\right\rvert\geq\epsilon\right\\}\leq\frac{M}{m\epsilon\sqrt{2}}\sum_{i=1}^{m}\sqrt{I(W,D_{i})}+\frac{1}{\epsilon}\mathbb{E}_{W,D}\lvert
L_{D}(W)-L(W)\rvert,$

where $D\sim\nu$ is independent of $W$.

Intuitively, extracting correlations between $X$ and $Y$ that are invariant
across training domains enhances the generalization ability of machine
learning models. The mutual information $I(W;D_{i})$ approaches zero when the
correlations that a model learns from a specific training domain $D_{i}$ are
also present in other training environments. This does not imply that the
model learns nothing from $D_{s}$: by further assuming the independence of
these domains, the summation of $I(W,D_{i})$ can be relaxed to $I(W;D_{s})$,
which measures the actual amount of information learned by the model. By
minimizing each $I(W;D_{i})$ and $L_{s}(W)$ simultaneously, learning
algorithms are encouraged to discard domain-specific correlations while
preserving invariant ones and thus achieve high generalization performance.

Interestingly, the second term at the RHS of Theorem III.2 is highly relevant
to the target-domain generalization gap, so we will postpone related analysis
to Section III-C.

Next, we demonstrate that the minimization of $I(W;D_{i})$ can be achieved by
matching the conditional distributions of inter-domain gradients. To see this,
we assume that $W$ is optimized by some noisy and iterative learning
algorithms, e.g. stochastic gradient descent (SGD). Then the rule of updating
$W$ at step $t$ through ERM can be formulated as:

$W_{t}=W_{t-1}-\eta_{t}\sum_{i=1}^{m}g(W_{t-1},B_{t}^{i}),\quad\mathrm{where}\quad
g(w,B_{t}^{i})=\frac{1}{m\lvert B_{t}^{i}\rvert}\sum_{z\in
B_{t}^{i}}\nabla_{w}\ell(f_{w}(x),y),$

providing $W_{0}$ as the initial guess. Here, $\eta_{t}$ is the learning rate,
and $B_{t}^{i}$ is the batch of data points randomly drawn from training
environment $D_{i}$ to compute the gradient. Suppose that algorithm
$\mathcal{A}$ finishes in $T$ steps, we then have:

######  Theorem III.3.

Let $G_{t}=-\eta_{t}\sum_{i=1}^{m}g(W_{t-1},B_{t}^{i})$, then

$I(W_{T};D_{i})\leq\sum_{t=1}^{T}I(G_{t};D_{i}|W_{t-1}).$

Although our analysis is derived from the update rule of SGD, the same
conclusion applies to a variety of iterative and noisy learning algorithms,
e.g. SGLD and AdaGrad. Theorem III.3 suggests that minimizing
$I(G_{t};D_{i}|W_{t-1})$ in each update penalizes $I(W_{T};D_{i})$ and thus
leads to training-domain generalization. Notably, this conditional mutual
information $I(G_{t};D_{i}|W_{t-1})$ can be rewritten as the KL divergence
$\mathrm{KL}(P_{G_{t}|D_{i},W_{t-1}}\,\|\,P_{G_{t}|W_{t-1}})$, which directly
motivates matching the distributions of inter-domain gradients:  Gradient
matching promotes training-domain generalization when $\mathbb{E}_{W,D}\lvert
L_{D}(W)-L(W)\rvert$ is minimized. Intuitively, gradient alignment enforces
the model to learn common correlations shared across training domains, thus
preventing overfitting to spurious features and promoting invariance [12, 13].

We further present an alternative approach by assuming Lipschitzness instead
of Subgaussianity, which usually leads to tighter bounds beyond information-
theoretic measures:

######  Theorem III.4.

If $\ell(f_{w}(X),Y)$ is $\beta^{\prime}$-Lipschitz w.r.t $w$, then

$\left\lvert\mathbb{E}_{W,D_{s}}[L_{s}(W)]-\mathbb{E}_{W}[L(W)]\right\rvert\leq\frac{\beta^{\prime}}{m}\sum_{i=1}^{m}\mathbb{E}_{D_{i}}[\mathbb{W}(P_{W|D_{i}},P_{W})].$

Besides the elegant symmetry compared to KL divergence metrics, Wasserstein
distance bounds are generally considered to be tighter improvements over
information-theoretic bounds. To see this, we assume that the adopted metric
$c$ is discrete, which leads to the following reductions:

$\mathbb{E}_{D_{i}}[\mathbb{W}(P_{W|D_{i}},P_{W})]=\mathbb{E}_{D_{i}}[\mathrm{TV}(P_{W|D_{i}},P_{W})]\leq\mathbb{E}_{D_{i}}\sqrt{\frac{1}{2}\mathrm{KL}(P_{W|D_{i}}\,\|\,P_{W})}\leq\sqrt{\frac{1}{2}I(W;D_{i})},$
(2)

where $\mathrm{TV}$ is the total variation. These reductions confirm that the
RHS of Theorem III.2 also upper bounds other alternative measures of domain
differences i.e. total variation and Wasserstein distance. This observation
encourages us to directly penalize the mutual information $I(W;D_{i})$, which
is not only more stable for optimization [43, 34] but also enables
simultaneous minimization of these alternative metrics.

###  III-C Target-domain Generalization

We then investigate sufficient conditions for target-domain generalization.
Since the training process is independent of the test domains, the predictor
could be considered as some constant hypothesis $w\in\mathcal{W}$. It is
straightforward to verify that $\mathbb{E}_{D_{t}}[L_{t}(w)]=L(w)$ due to the
identical domain distribution $\nu$. We then establish the following bound for
the target-domain generalization gap:

######  Theorem III.5.

If Assumption II.3 holds, then $\forall w\in\mathcal{W}$,

$\mathbb{P}\left\\{\lvert
L_{t}(w)-L(w)\rvert\geq\epsilon\right\\}\leq\frac{\sigma}{\epsilon}\sqrt{2I(Z;D)}.$

The result above can be interpreted from two perspectives. Firstly, evaluating
the predictor $w$ on randomly sampled test environments reflects its ability
to generalize on average, since $L_{t}(w)$ is an unbiased estimate of $L(w)$.
Secondly, knowledge about $L(w)$ can be used to predict the ability of $w$ to
generalize on unseen domains, which complements Theorem III.2 in solving
Problem II.6.

In Theorem III.5, the probability of generalization is mainly controlled by
the extent of distribution shift $I(Z;D)$. Notably, $I(Z;D)$ is an intrinsic
property of the data collection procedure, and thus cannot be penalized from
the perspective of learning algorithms. Fortunately, the encoder $\phi$ can be
considered as part of the data preprocessing procedure, enabling learning
algorithms to minimize the representation space distribution shift. Under the
same conditions as Theorem III.5, we have that for any classifier $\psi$:

$\mathbb{P}\left\\{\lvert
L_{t}(\psi)-L(\psi)\rvert\geq\epsilon\right\\}\leq\frac{\sigma}{\epsilon}\sqrt{2I(R,Y;D)}.$

Let $P_{R,Y}$ be the joint distribution by pushing forward $P_{Z}$ via the
encoder as $R=f_{\phi}(X)$, the representation space distribution shift can
then be decomposed into:

$I(R,Y;D)\textrm{ (distribution shift)}=I(R;D)\textrm{ (covariate
shift)}+I(Y;D|R)\textrm{ (concept shift)}.$

This motivates us to simultaneously minimize the representation space
covariate shift and concept shift to achieve test-domain generalization. We
further demonstrate that bounding the covariate shift $I(R;D)$ solely is
sufficient for target-domain generalization with Assumption II.4:

######  Theorem III.6.

If Assumptions II.2 and II.4 hold, then $\forall\psi$,

$\displaystyle\mathbb{P}\left\\{L_{t}(\psi)-L(\psi)\geq\epsilon\right\\}\leq\frac{M}{\epsilon\sqrt{2}}\sqrt{I(R;D)}+\frac{2}{\epsilon}L^{*},$

where $L^{*}=\min_{f^{*}:\mathcal{R}\mapsto\mathcal{Y}}[L(f^{*})]$.

Similarly, we further refine these test-domain generalization bounds by
incorporating the more stringent Assumption II.5 in Appendix Further
Discussions. Theorem III.6 indicates that test-domain generalization is mainly
controlled by the amount of covariate shift. Notably, the optimal classifier
$f^{*}$ is chosen from the entire space of functions mapping from
$\mathcal{R}$ to $\mathcal{Y}$. In the noiseless case where there exists a
ground-truth labeling function $h^{*}$ such that $Y=h^{*}(R)$, we have
$L^{*}=0$. Therefore, $L^{*}$ serves as an indicator for the level of label
noise in the data-generating distribution $\mu$. Moreover, the representation
space covariate shift $I(R;D)$ is equivalent to the KL divergence
$\mathrm{KL}(P_{R|D}\,\|\,P_{R})$, which directly motivates matching the
distributions of inter-domain representations:  Representation matching
promotes test-domain generalization when the level of label noise $L^{*}$ is
low.

A byproduct of the proof of Theorem III.6 is an upper bound on the expected
absolute target-domain generalization error:

$\mathbb{E}_{W,D}\lvert
L_{D}(W)-L(W)\rvert\leq\frac{M}{\sqrt{2}}\sqrt{I(R;D)}+2L^{*}.$

This result complements the training-domain generalization bound in Theorem
III.2, confirming that penalizing $I(W;D_{i})$ and $I(R;D)$ simultaneously is
sufficient to minimize the training-domain generalization gap.

While minimizing $I(R;D)$ guarantees test-domain generalization, this
operation requires knowledge about test-domain samples, which is not available
during the entire training process. A walkaround is to match the training-
domain representations instead, utilizing $I(R_{i};D_{i})$ for each $D_{i}\in
D_{s}$ as a proxy to penalize $I(R;D)$. The following proposition verifies the
feasibility of this approach:

######  Proposition III.7.

Under mild conditions, we have

$\mathrm{SKL}(P_{R,D}\,\|\,P_{R_{i},D_{i}})=O\left\lparen\sqrt{I(W;D_{i})}\right\rparen,$

where $\mathrm{SKL}(P\,\|\,Q)=\mathrm{KL}(P\,\|\,Q)+\mathrm{KL}(Q\,\|\,P)$.

Interestingly, the discrepancy between the test-domain joint distribution of
$P_{R,D}$ and its training-domain counterpart $P_{R_{i},D_{i}}$ can be upper
bounded by the input-output mutual information $I(W;D_{i})$. By simultaneously
minimizing $I(W;D_{i})$, one can use $I(R_{i};D_{i})$ as a proxy to penalize
$I(R;D)$ and achieve test-domain generalization.

While our analysis does not necessitate the independence condition between
source domains or target domains, such a condition is also naturally satisfied
in most learning scenarios and leads to tighter generalization bounds.
Specifically, Theorem III.5 and III.6 can be further tightened by a factor of
$\frac{1}{m^{\prime}}$ when test domains are i.i.d. We refer the readers to
Appendix Omitted Proofs in Section III and Omitted Proofs in Section IV for
the proof of these results.

##  IV Inter-domain Distribution Matching

Motivated by our theoretical analysis in Section III, we propose inter-domain
distribution matching (IDM) to achieve high-probability DG (Problem II.6).
Recall that the average-case risk $L(W)$ serves as a natural bridge to connect
$L_{s}(W)$ and $L_{t}(W)$, the regularization in Problem II.6 directly
indicates an objective for optimization by combining the high-probability
concentration bounds in Theorem III.2 and III.5. Specifically, for any
$\lambda\in(0,1)$, we have:

$\mathbb{P}\\{\lvert
L_{t}(W)-L_{s}(W)\rvert\geq\epsilon\\}\leq\frac{M}{m\epsilon\lambda\sqrt{2}}\sum_{i=1}^{m}\sqrt{I(W,D_{i})}+\frac{1}{\epsilon\lambda(1-\lambda)}\left\lparen\frac{M}{\sqrt{2}}\sqrt{I(R;D)}+2L^{*}\right\rparen.$
(3)

This observation directly motivates aligning inter-domain distributions of the
gradients and representations simultaneously. While the idea of distribution
matching is not new, we are the first to explore the complementary
relationship between gradient and representation matching:  Gradient and
representation matching together minimize $\mathbb{P}\\{\lvert
L_{t}(W)-L_{s}(W)\rvert\geq\epsilon\\}$ in Problem II.6. Specifically,
training-domain generalization requires the minimization of the test-domain
generalization gap (Theorem III.2), and test-domain generalization requires
the minimization of the input-output mutual information (Proposition III.7).
Therefore, existing works focusing exclusively on either gradient or
representation alignment are insufficient to fully address the domain-level
generalization gap. To our best knowledge, this is the first exploration in
the literature where gradient and representation matching are combined to
yield a sufficient solution for the DG problem.

###  IV-A Per-sample Distribution Matching

While various distribution matching methods have been proposed in the
literature, these techniques are generally either ineffective or insufficient
for high-dimensional and complex distributions. Typically, learning algorithms
have no knowledge about the underlying distribution of either the
representation or the gradient, and the only available way is to align them
across batched data points. We first provide an impossibility theorem for
high-dimensional distribution matching in the cases of limited number of
samples:

######  Theorem IV.1.

(Informal) Let $n$ and $b$ be the dimension and the number of data points
respectively. If $n>b+1$, then there exist infinite environments whose
conditional probabilities given an arbitrarily sampled group of data points
are indistinguishable. If $n>2b+1$, then there exist infinite environments
whose conditional probabilities cannot distinguish two arbitrarily sampled
groups of data points.

We refer the readers to Appendix Omitted Proofs in Section IV for a formal
statement of Theorem IV.1. In real-world scenarios, the dimensionality of the
feature or the gradient easily exceeds that of the batch size, making
algorithms that aim to align the entire distribution (e.g. CORAL [5] and MMD
[6]) generally ineffective since distribution alignment is basically
impossible given such few data points. This observation is also verified by
[13] that aligning the entire covariance matrix achieves no better performance
than aligning the diagonal elements only. Furthermore, prior distribution
alignment techniques mainly focus on aligning the directions [20, 44, 12] or
low-order moments [5, 11, 13], which are insufficient for complex probability
distributions. For example, while the standard Gaussian distribution $N(0,1)$
and the uniform distribution $U(-\sqrt{3},\sqrt{3})$ share the same
expectation and variance, they are fundamentally different to one another. To
address these issues, we propose the per-sample distribution matching (PDM)
technique that aligns distributions in a per-dimension manner, by minimizing
an upper bound of the KL divergence between probability density estimators.

Let $\\{x_{i}^{1}\\}_{i=1}^{b}$ and $\\{x_{i}^{2}\\}_{i=1}^{b}$ be two groups
of 1-dimensional data points drawn from probability distributions $P$ and $Q$
respectively. Let $p_{i}$ denote the density of Gaussian distribution with
expectation $x_{i}^{1}$ and variance $\sigma^{2}$, then the kernel density
estimator $\bar{P}$ for $P$ can be written as
$\bar{p}(x)=\frac{1}{b}\sum_{i}p_{i}(x)$ (respectively for $q_{i}$, $\bar{Q}$
and $\bar{q}$). The following theorem suggests a computable upper bound for
the KL divergence (Wasserstein distance) between probability density
estimators:

######  Theorem IV.2.

Let $f$ be a bijection: $[1,b]\leftrightarrow[1,b]$ and $P_{i}$ be the
probability measure defined by $p_{i}$ (respectively for $Q_{i}$ and $q_{i}$),
then
$\mathrm{KL}(\bar{P}\,\|\,\bar{Q})\leq\frac{1}{b}\sum_{i=1}^{b}\mathrm{KL}(P_{i}\,\|\,Q_{f(i)})$,
and
$\mathbb{W}(\bar{P},\bar{Q})\leq\frac{1}{b}\sum_{i=1}^{b}\mathbb{W}(P_{i},Q_{f(i)})$.

Hence, distribution matching can be achieved by minimizing the KL divergence
or Wasserstein distances between point Gaussian densities, which can be
achieved by aligning individual data points. The following theorem suggests an
optimal bijection for choosing the order of alignment:

######  Theorem IV.3.

Suppose that $\\{x_{i}^{1}\\}_{i=1}^{b}$ and $\\{x_{i}^{2}\\}_{i=1}^{b}$ are
both sorted in the same order, then $f(j)=j$ is the minimizer of both
$\sum_{i=1}^{b}\mathrm{KL}(P_{i}\,\|\,Q_{f(i)})$ and
$\sum_{i=1}^{b}\mathbb{W}(P_{i},Q_{f(i)})$.

To summarize, the procedure of PDM is to slice the data points into separate
dimensions, sort the data points in ascending (or descending) order for each
dimension, and then match the sorted data points across different training
domains. PDM improves over previous distribution matching techniques by
simultaneously capturing multiple orders of moments, avoiding ineffective
high-dimensional distribution matching, as well as enabling straightforward
implementation and efficient computation. We provide pseudo-codes for both PDM
and IDM in Appendix Experiment Details.

###  IV-B Algorithm Design

Combining the methods discussed above, we finally propose the IDM algorithm
for high-probability DG by simultaneously aligning inter-domain gradients and
representations. Recall that Problem II.6 incorporates an additional
regularization based on ERM, we adopt the following Lagrange multipliers to
optimize the IDM objective:

$\mathcal{L}_{\mathrm{IDM}}=\mathcal{L}_{\mathrm{E}}+\lambda_{1}\mathcal{L}_{\mathrm{G}}+\lambda_{2}\mathcal{L}_{\mathrm{R}}=\frac{1}{m}\sum_{i=1}^{m}\left[L_{D_{i}}(W)+\lambda_{1}\mathrm{PDM}(G_{i})+\lambda_{2}\mathrm{PDM}(R_{i})\right].$
(4)

Here $\mathcal{L}_{\mathrm{E}}$ is the risk of ERM, $\mathcal{L}_{\mathrm{G}}$
and $\mathcal{L}_{\mathrm{R}}$ denote the penalty of distribution matching for
the gradients and representations respectively, implemented with the proposed
PDM method. To cooperate representation alignment which regards the classifier
$\psi$ as the true predictor and also for memory and time concerns, we only
apply gradient alignment for the classifier $\psi$ as in [13]. Furthermore,
$\lambda_{1}$ and $\lambda_{2}$ should be adaptively chosen according to the
extent of covariate and concept shifts respectively: Firstly, $I(R;D)$ is
upper bounded by $I(X;D)$ by the Markov chain $D\rightarrow X\rightarrow R$,
so the representations are naturally aligned when $I(X;D)=0$. Secondly,
gradient alignment is not required when $I(Y;D|X)=0$, since the distribution
shift can then be minimized by aligning the representations solely. Therefore,
$\lambda_{1}$ and $\lambda_{2}$ should scale with the amount of the covariate
and concept shifts, respectively.

Figure 1: Learning dynamics of IDM.

TABLE I: The Colored MNIST task.

Method Train Acc Test Acc Gray Acc ERM 86.4 $\pm$ 0.2 14.0 $\pm$ 0.7 71.0
$\pm$ 0.7 IRM 71.0 $\pm$ 0.5 65.6 $\pm$ 1.8 66.1 $\pm$ 0.2 V-REx 71.7 $\pm$
1.5 67.2 $\pm$ 1.5 68.6 $\pm$ 2.2 IGA 68.9 $\pm$ 3.0 67.7 $\pm$ 2.9 67.5 $\pm$
2.7 Fishr 69.6 $\pm$ 0.9 71.2 $\pm$ 1.1 70.2 $\pm$ 0.7 IDM 70.2 $\pm$ 1.4 70.6
$\pm$ 0.9 70.5 $\pm$ 0.7

##  V Related Works

In the literature, various approaches have been proposed by incorporating
external domain information to achieve OOD generalization. Most recent works
achieve invariance by employing additional regularization criteria based on
ERM. These methods differ in the choice of the statistics used to match across
training domains and can be categorized by the corresponding objective of 1)
gradient, 2) representation, and 3) predictor, as follows:

Invariant Gradients. Gradient alignment enforces batched data points from
different domains to cooperate and promotes OOD generalization by finding loss
minima shared across training domains. Specifically, IGA [11] aligns the
empirical expectations, Fish [12] maximizes the dot-product of inter-domain
gradients, AND-mask [20] and SAND-mask [44] only update weights when the
gradients share the same direction, and Fishr [13] matches the gradient
variance. These gradient-based objectives are generally restricted to aligning
the directions or low-order moments, resulting in substantial information loss
in more granular statistics. Besides, these works either lack generalization
guarantees or rely on strong assumptions including the existence of invariant
and controllable features (IGA), the shape of loss landscape around the local
minima (Fishr), Lipschitz continuous gradients and co-diagonalizable Hessian
matrix (AND-mask). In contrast, our analysis successfully connects gradient
alignment and training-domain generalization by leveraging the mild assumption
of identical domain distributions.

Invariant Representations. Extracting domain-invariant features has been
extensively studied to solve both DG and domain adaptation (DA) problems. DANN
[7] and CDANN [8] align inter-domain representations via adversarial learning,
MMD [6] uses kernel methods for distribution alignment, and CORAL [5] matches
low-order moments of the representations. Still, these methods are
insufficient for complex probability distributions [45], ineffective for high-
dimensional distributions (Theorem IV.1), and incapable of addressing the
concept shift. Besides, the viability of minimizing the representation shift
$I(R;D)$ through the training-domain proxy $I(R_{i};D_{i})$ remains
questionable without gradient matching (Proposition III.1). Our analysis sheds
light on understanding how representation alignment enhances test-domain
generalization by minimizing the variance of target-domain risks.

Invariant Predictors. A recent line of works proposes to explore the
connection between invariance and causality. IRM [9] and subsequent works [37,
36] learn an invariant classifier that is simultaneously optimal for all
training domains. However, later works have shown that IRM may fail on non-
linear data and lead to sub-optimal predictors [46, 35]. Parallel works
include: V-REx [15] which minimizes the variance of training-domain risks,
GroupDRO [14] which minimizes the worst-domain training risk, and QRM [16]
which optimizes a quantile of the risk distribution. As shown in the next
section, IDM also promotes domain-invariant predictors and ensures optimality
across different training domains.

TABLE II: The DomainBed benchmark. We format best, second best and worse than
ERM results.

Algorithm Accuracy ($\uparrow$) Ranking ($\downarrow$) CMNIST RMNIST VLCS PACS
OffHome TerraInc DomNet Avg Mean Median Worst ERM 57.8 $\pm$ 0.2 97.8 $\pm$
0.1 77.6 $\pm$ 0.3 86.7 $\pm$ 0.3 66.4 $\pm$ 0.5 53.0 $\pm$ 0.3 41.3 $\pm$ 0.1
68.7 12.3 11 20 IRM 67.7 $\pm$ 1.2 97.5 $\pm$ 0.2 76.9 $\pm$ 0.6 84.5 $\pm$
1.1 63.0 $\pm$ 2.7 50.5 $\pm$ 0.7 28.0 $\pm$ 5.1 66.9 18.3 20 22 GroupDRO 61.1
$\pm$ 0.9 97.9 $\pm$ 0.1 77.4 $\pm$ 0.5 87.1 $\pm$ 0.1 66.2 $\pm$ 0.6 52.4
$\pm$ 0.1 33.4 $\pm$ 0.3 67.9 11.7 10 19 Mixup 58.4 $\pm$ 0.2 98.0 $\pm$ 0.1
78.1 $\pm$ 0.3 86.8 $\pm$ 0.3 68.0 $\pm$ 0.2 54.4 $\pm$ 0.3 39.6 $\pm$ 0.1
69.0 7.3 6 15 MLDG 58.2 $\pm$ 0.4 97.8 $\pm$ 0.1 77.5 $\pm$ 0.1 86.8 $\pm$ 0.4
66.6 $\pm$ 0.3 52.0 $\pm$ 0.1 41.6 $\pm$ 0.1 68.7 12.6 13 18 CORAL 58.6 $\pm$
0.5 98.0 $\pm$ 0.0 77.7 $\pm$ 0.2 87.1 $\pm$ 0.5 68.4 $\pm$ 0.2 52.8 $\pm$ 0.2
41.8 $\pm$ 0.1 69.2 6.4 5 14 MMD 63.3 $\pm$ 1.3 98.0 $\pm$ 0.1 77.9 $\pm$ 0.1
87.2 $\pm$ 0.1 66.2 $\pm$ 0.3 52.0 $\pm$ 0.4 23.5 $\pm$ 9.4 66.9 10.0 10 22
DANN 57.0 $\pm$ 1.0 97.9 $\pm$ 0.1 79.7 $\pm$ 0.5 85.2 $\pm$ 0.2 65.3 $\pm$
0.8 50.6 $\pm$ 0.4 38.3 $\pm$ 0.1 67.7 15.0 18 22 CDANN 59.5 $\pm$ 2.0 97.9
$\pm$ 0.0 79.9 $\pm$ 0.2 85.8 $\pm$ 0.8 65.3 $\pm$ 0.5 50.8 $\pm$ 0.6 38.5
$\pm$ 0.2 68.2 12.4 14 18 MTL 57.6 $\pm$ 0.3 97.9 $\pm$ 0.1 77.7 $\pm$ 0.5
86.7 $\pm$ 0.2 66.5 $\pm$ 0.4 52.2 $\pm$ 0.4 40.8 $\pm$ 0.1 68.5 11.7 10 21
SagNet 58.2 $\pm$ 0.3 97.9 $\pm$ 0.0 77.6 $\pm$ 0.1 86.4 $\pm$ 0.4 67.5 $\pm$
0.2 52.5 $\pm$ 0.4 40.8 $\pm$ 0.2 68.7 11.3 9 17 ARM 63.2 $\pm$ 0.7 98.1 $\pm$
0.1 77.8 $\pm$ 0.3 85.8 $\pm$ 0.2 64.8 $\pm$ 0.4 51.2 $\pm$ 0.5 36.0 $\pm$ 0.2
68.1 13.0 16 21 VREx 67.0 $\pm$ 1.3 97.9 $\pm$ 0.1 78.1 $\pm$ 0.2 87.2 $\pm$
0.6 65.7 $\pm$ 0.3 51.4 $\pm$ 0.5 30.1 $\pm$ 3.7 68.2 10.6 8 20 RSC 58.5 $\pm$
0.5 97.6 $\pm$ 0.1 77.8 $\pm$ 0.6 86.2 $\pm$ 0.5 66.5 $\pm$ 0.6 52.1 $\pm$ 0.2
38.9 $\pm$ 0.6 68.2 13.4 13 19 AND-mask 58.6 $\pm$ 0.4 97.5 $\pm$ 0.0 76.4
$\pm$ 0.4 86.4 $\pm$ 0.4 66.1 $\pm$ 0.2 49.8 $\pm$ 0.4 37.9 $\pm$ 0.6 67.5
17.0 16 22 SAND-mask 62.3 $\pm$ 1.0 97.4 $\pm$ 0.1 76.2 $\pm$ 0.5 85.9 $\pm$
0.4 65.9 $\pm$ 0.5 50.2 $\pm$ 0.1 32.2 $\pm$ 0.6 67.2 17.9 19 22 Fish 61.8
$\pm$ 0.8 97.9 $\pm$ 0.1 77.8 $\pm$ 0.6 85.8 $\pm$ 0.6 66.0 $\pm$ 2.9 50.8
$\pm$ 0.4 43.4 $\pm$ 0.3 69.1 11.3 11 18 Fishr 68.8 $\pm$ 1.4 97.8 $\pm$ 0.1
78.2 $\pm$ 0.2 86.9 $\pm$ 0.2 68.2 $\pm$ 0.2 53.6 $\pm$ 0.4 41.8 $\pm$ 0.2
70.8 5.4 3 16 SelfReg 58.0 $\pm$ 0.7 98.1 $\pm$ 0.7 78.2 $\pm$ 0.1 87.7 $\pm$
0.1 68.1 $\pm$ 0.3 52.8 $\pm$ 0.9 43.1 $\pm$ 0.1 69.4 5.0 3 19 CausIRLCORAL
58.4 $\pm$ 0.3 98.0 $\pm$ 0.1 78.2 $\pm$ 0.1 87.6 $\pm$ 0.1 67.7 $\pm$ 0.2
53.4 $\pm$ 0.4 42.1 $\pm$ 0.1 69.4 5.0 3 15 CausIRLMMD 63.7 $\pm$ 0.8 97.9
$\pm$ 0.1 78.1 $\pm$ 0.1 86.6 $\pm$ 0.7 65.2 $\pm$ 0.6 52.2 $\pm$ 0.3 40.6
$\pm$ 0.2 69.2 10.4 10 20 IDM 72.0 $\pm$ 1.0 98.0 $\pm$ 0.1 78.1 $\pm$ 0.4
87.6 $\pm$ 0.3 68.3 $\pm$ 0.2 52.8 $\pm$ 0.5 41.8 $\pm$ 0.2 71.2 3.3 3 6

##  VI Experimental Results

In this section, we evaluate the proposed IDM algorithm on the Colored MNIST
task [9] and the DomainBed benchmark [21] to demonstrate its capability of
generalizing against various types of distribution shifts111The source code is
available at https://github.com/Yuxin-Dong/IDM.. Detailed settings of these
experiments and further empirical results including ablation studies are
reported in Appendix Experiment Details and Additional Experimental Results.

###  VI-A Colored MNIST

The Colored MNIST task [9] is carefully designed to create high correlations
between image colors and the true labels, leading to spurious features that
possess superior predictive power ($90\%$ and $80\%$ accuracy) over the actual
digits ($75\%$). However, this correlation is reversed in the test domain
($10\%$), causing any learning algorithm that solely minimizes training errors
to overfit the color information and fail when testing. As such, Colored MNIST
is an ideal task to evaluate the capability of learning algorithms to achieve
invariance across source domains.

Following the settings of [9], we adopt a two-stage training technique, where
the penalty strength $\lambda$ is set low initially and higher afterward. We
visualize the learning dynamics of relevant DG penalties, including IRM,
V-Rex, IGA, and Fishr, using the IDM objective for optimization in Figure 1.
The penalty values are normalized for better clarity. This visualization
confirms Theorem III.2 that IDM promotes source-domain generalization by
minimizing the gap between training risks, thus ensuring the optimality of the
predictor across different training domains. Moreover, it verifies the
superiority of PDM by showing that penalizing the IDM objective solely is
sufficient to minimize other types of invariance penalties.

Table 1 presents the performance comparison on Colored MNIST across $10$
independent runs. Following the hyper-parameter tuning technique as [9], we
select the best model by $\max_{w}\min(L_{s}(w),L_{t}(w))$. As can be seen,
IDM achieves the best trade-off between training and test-domain accuracies
(70.2%), and near-optimal gray-scale accuracy (70.5%) compared to the Oracle
predictor (71.0%, ERM trained with gray-scale images).

###  VI-B DomainBed Benchmark

The DomainBed Benchmark [21] comprises multiple synthetic and real-world
datasets for assessing the performance of both DA and DG algorithms. To ensure
a fair comparison, DomainBed limits the number of attempts for hyper-parameter
tuning to $20$, and the results are averaged over $3$ independent trials.
Therefore, DomainBed serves as a rigorous and comprehensive benchmark to
evaluate different DG strategies. We compare the performance of our method
with $20$ baselines in total for a thorough evaluation. Table II summarizes
the results using test-domain model selection, which is a common choice for
validation purposes [13, 15] and highly motivated by our discussion in
Appendix -U.

As can be seen, IDM achieves top-1 accuracy (72.0%) on CMNIST which is
competitive with the Oracle (75.0%), outperforming all previous distribution
alignment techniques by aligning the directions (AND-mask, SAND-mask, Fish) or
low-order moments (Fishr). This verifies the superiority of the proposed PDM
method as well as the complementary relationship between gradient and
representation alignment. On the contrary, algorithms that only align the
representations (CORAL, MMD, DANN, CDANN) are incapable of addressing the
concept shift, thus performing poorly on CMNIST. Moreover, IDM achieves the
highest accuracy among all distribution matching algorithms on RMNIST / PACS,
competitive performances to the best algorithm on RMNIST (98.0% v.s. 98.1%),
PACS (87.6% v.s. 87.7%), OfficeHome (68.3% v.s. 68.4%), the highest average
accuracy (71.2%) and best rankings (mean, median and worst rankings on $7$
datasets) among all baseline methods. IDM also enables efficient computation,
such that the running-time overhead is only $5\%$ compared to ERM on the
largest DomainNet dataset, and negligible for other smaller datasets. Notably,
IDM is the only algorithm that consistently achieves top rankings (Top 6 of
22), while any other method failed to outperform most of the competitors on at
least $1$ dataset.

While the overall performance is promising, we notice that IDM is not very
effective on TerraIncognita. There are several possible reasons: Firstly, the
number of hyper-parameters in IDM exceeds most competing methods, which is
critical to model selection since the number of tuning attempts is limited in
DomainBed. Recall that the value of $\lambda_{1}$ and $\lambda_{2}$ should
adapt to the amount of covariate and concept shifts respectively: While CMNIST
manually induces high concept shift, covariate shift is instead dominant in
other datasets, raising extra challenges for hyper-parameter tuning. Secondly,
representation space distribution alignment may not always help since
$L_{t}(w)\leq L(w)$ is possible by the randomized nature of target domains.
These factors together result in sub-optimal hyper-parameter selection
results.

##  VII Conclusion

In this work, we explore a novel perspective for DG by minimizing the domain-
level generalization gap with high probability, which facilitates information-
theoretic analysis for the generalization behavior of learning algorithms. Our
analysis sheds light on understanding how gradient or representation matching
enhances generalization and unveils the complementary relationship between
these two elements. These theoretical insights inspire us to design the IDM
algorithm by simultaneously aligning inter-domain gradients and
representations, which then achieves superior performance on the DomainBed
benchmark.

## Acknowledgments

This work has been supported by the Key Research and Development Program of
China under Grant 2021ZD0110700; The National Natural Science Foundation of
China under Grant 62106191;

## Prerequisite Definitions and Lemmas

######  Definition .1.

(Subgaussian) A random variable $X$ is $\sigma$-subgaussian if for any
$\rho\in\mathbb{R}$,
$\mathbb{E}[\exp(\rho(X-\mathbb{E}[X]))]\leq\exp(\rho^{2}\sigma^{2}/2)$.

######  Definition .2.

(Kullback-Leibler Divergence) Let $P$ and $Q$ be probability measures on the
same space $\mathcal{X}$, the KL divergence from $P$ to $Q$ is defined as
$\mathrm{KL}(P\,\|\,Q)\triangleq\int_{\mathcal{X}}p(x)\log(p(x)/q(x))\mathop{}\\!\mathrm{d}x$.

######  Definition .3.

(Mutual Information) Let $(X,Y)$ be a pair of random variables with values
over the space $\mathcal{X}\times\mathcal{Y}$. Let their joint distribution be
$P_{X,Y}$ and the marginal distributions be $P_{X}$ and $P_{Y}$ respectively,
the mutual information between $X$ and $Y$ is defined as
$I(X;Y)=\mathrm{KL}(P_{X,Y}\,\|\,P_{X}P_{Y})$.

######  Definition .4.

(Wasserstein Distance) Let $c(\cdot,\cdot)$ be a metric and let $P$ and $Q$ be
probability measures on $\mathcal{X}$. Denote $\Gamma(P,Q)$ as the set of all
couplings of $P$ and $Q$ (i.e. the set of all joint distributions on
$\mathcal{X}\times\mathcal{X}$ with two marginals being $P$ and $Q$), then the
Wasserstein distance of order $p$ between $P$ and $Q$ is defined as
$\mathbb{W}_{p}(P,Q)\triangleq\left\lparen\inf_{\gamma\in\Gamma(P,Q)}\int_{\mathcal{X}\times\mathcal{X}}c(x,x^{\prime})^{p}\mathop{}\\!\mathrm{d}\gamma(x,x^{\prime})\right\rparen^{1/p}$.

Unless otherwise noted, we use $\log$ to denote the logarithmic function with
base $e$, and use $\mathbb{W}(\cdot,\cdot)$ to denote the Wasserstein distance
of order $1$.

######  Definition .5.

(Total Variation) The total variation between two probability measures $P$ and
$Q$ is $\mathrm{TV}(P,Q)\triangleq\sup_{E}\lvert P(E)-Q(E)\rvert$, where the
supremum is over all measurable set $E$.

######  Lemma .6.

(Lemma 1 in [42]) Let $(X,Y)$ be a pair of random variables with joint
distribution $P_{X,Y}$ and let $\bar{Y}$ be an independent copy of $Y$. If
$f(x,y)$ is a measurable function such that $E_{X,Y}[f(X,Y)]$ exists and
$f(X,\bar{Y})$ is $\sigma$-subgaussian, then

$\left\lvert\mathbb{E}_{X,Y}[f(X,Y)]-\mathbb{E}_{X,\bar{Y}}[f(X,\bar{Y})]\right\rvert\leq\sqrt{2\sigma^{2}I(X;Y)}.$

Furthermore, if $f(x,Y)$ is $\sigma$-subgaussian for each $x$ and the
expectation below exists, then

$\mathbb{E}_{X,Y}\left[\left\lparen
f(X,Y)-\mathbb{E}_{\bar{Y}}[f(X,\bar{Y})]\right\rparen^{2}\right]\leq
4\sigma^{2}(I(X;Y)+\log 3),$

and for any $\epsilon>0$, we have

$\mathbb{P}\left\\{\left\lvert
f(X,Y)-\mathbb{E}_{\bar{Y}}[f(X,\bar{Y})]\right\rvert\geq\epsilon\right\\}\leq\frac{4\sigma^{2}(I(X;Y)+\log
3)}{\epsilon^{2}}.$

######  Lemma .7.

(Lemma 2 in [42]) Let $X$ be $\sigma$-subgaussian and $\mathbb{E}[X]=0$, then
for any $\lambda\in[0,1/4\sigma^{2})$:

$\mathbb{E}_{X}\left[e^{\lambda X^{2}}\right]\leq 1+8\lambda\sigma^{2}.$

######  Lemma .8.

(Donsker-Varadhan formula) Let $P$ and $Q$ be probability measures defined on
the same measurable space, where $P$ is absolutely continuous with respect to
$Q$. Then

$\mathrm{KL}(P\,\|\,Q)=\sup_{X}\left\\{\mathbb{E}_{P}[X]-\log\mathbb{E}_{Q}[e^{X}]\right\\},$

where $X$ is any random variable such that $e^{X}$ is $Q$-integrable and
$\mathbb{E}_{P}[X]$ exists.

######  Lemma .9.

Let $P$, and $Q$ be probability measures defined on the same measurable space.
Let $X\sim P$ and $X^{\prime}\sim Q$. If $f(X)$ is $\sigma$-subgaussian w.r.t
$X$ and the following expectations exists, then

$\displaystyle\left\lvert\mathbb{E}_{X^{\prime}}[f(X^{\prime})]-\mathbb{E}_{X}[f(X)]\right\rvert\leq\sqrt{2\sigma^{2}\mathrm{KL}(Q\,\|\,P)},$
$\displaystyle\mathbb{E}_{X^{\prime}}\left[\left\lparen
f(X^{\prime})-\mathbb{E}_{X}[f(X)]\right\rparen^{2}\right]\leq
4\sigma^{2}(\mathrm{KL}(Q\,\|\,P)+\log 3).$

Furthermore, by combining the results above and Markov’s inequality, we have
that for any $\epsilon>0$:

$\mathbb{P}\left\\{\lvert
f(X^{\prime})-\mathbb{E}_{X}[f(X)]\rvert\geq\epsilon\right\\}\leq\frac{4\sigma^{2}}{\epsilon^{2}}\left\lparen\mathrm{KL}(Q\,\|\,P)+\log
3\right\rparen.$

###### Proof.

Let $\lambda\in\mathbb{R}$ be any non-zero constant, then by the subgaussian
property of $f(X)$:

$\displaystyle\log\mathbb{E}_{X}\left[e^{\lambda(f(X)-\mathbb{E}_{X}[f(X)])}\right]$
$\displaystyle\leq\frac{\lambda^{2}\sigma^{2}}{2},$
$\displaystyle\log\mathbb{E}_{X}\left[e^{\lambda
f(X)}\right]-\lambda\mathbb{E}_{X}[f(X)]$
$\displaystyle\leq\frac{\lambda^{2}\sigma^{2}}{2}.$

By applying Lemma .8 with $X=\lambda f(X)$ we have

$\displaystyle\mathrm{KL}(Q\,\|\,P)$
$\displaystyle\geq\sup_{\lambda}\left\\{\mathbb{E}_{X^{\prime}}[\lambda
f(X^{\prime})]-\log\mathbb{E}_{X}\left[e^{\lambda f(X)}\right]\right\\}$
$\displaystyle\geq\sup_{\lambda}\left\\{\mathbb{E}_{X^{\prime}}[\lambda
f(X^{\prime})]-\lambda\mathbb{E}_{X}[f(X)]-\frac{\lambda^{2}\sigma^{2}}{2}\right\\}$
$\displaystyle=\frac{1}{2\sigma^{2}}\left\lparen\mathbb{E}_{X^{\prime}}[f(X^{\prime})]-\mathbb{E}_{X}[f(X)]\right\rparen^{2},$

where the supremum is taken by setting
$\lambda=\frac{1}{\sigma^{2}}(\mathbb{E}_{X^{\prime}}[f(X^{\prime})]-\mathbb{E}_{X}[f(X)])$.
This completes the proof of the first inequality.

To prove the second inequality, let $g(x)=(f(x)-\mathbb{E}_{X}[f(X)])^{2}$ and
$\lambda\in[0,1/4\sigma^{2})$. Apply Lemma .8 again with $X=\lambda g(X)$, we
have

$\displaystyle\mathrm{KL}(Q\,\|\,P)$
$\displaystyle\geq\sup_{\lambda}\left\\{\mathbb{E}_{X^{\prime}}[\lambda
g(X^{\prime})]-\log\mathbb{E}_{X}\left[e^{\lambda g(X)}\right]\right\\}$
$\displaystyle=\sup_{\lambda}\left\\{\mathbb{E}_{X^{\prime}}\left[\lambda\left\lparen
f(X^{\prime})-\mathbb{E}_{X}[f(X)]\right\rparen^{2}\right]-\log\mathbb{E}_{X}\left[e^{\lambda\left\lparen
f(X)-\mathbb{E}_{X}[f(X)]\right\rparen^{2}}\right]\right\\}$
$\displaystyle\geq\sup_{\lambda}\left\\{\mathbb{E}_{X^{\prime}}\left[\lambda\left\lparen
f(X^{\prime})-\mathbb{E}_{X}[f(X)]\right\rparen^{2}\right]-\log(1+8\lambda\sigma^{2})\right\\}$
$\displaystyle\geq\frac{1}{4\sigma^{2}}\mathbb{E}_{X^{\prime}}\left[\left\lparen
f(X^{\prime})-\mathbb{E}_{X}[f(X)]\right\rparen^{2}\right]-\log 3,$

where the second inequality follows by applying Lemma .7 and the last
inequality follows by taking $\lambda\rightarrow\frac{1}{4\sigma^{2}}$. This
finishes the proof of the second inequality.

Furthermore, by applying Markov’s inequality, we can get:

$\displaystyle\mathbb{P}\left\\{\lvert
f(X^{\prime})-\mathbb{E}_{X}[f(X)]\rvert\geq\epsilon\right\\}$
$\displaystyle=\mathbb{P}\left\\{\left\lparen
f(X^{\prime})-\mathbb{E}_{X}[f(X)]\right\rparen^{2}\geq\epsilon^{2}\right\\}$
$\displaystyle\leq\frac{1}{\epsilon^{2}}\mathbb{E}_{X^{\prime}}\left[\left\lparen
f(X^{\prime})-\mathbb{E}_{X}[f(X)]\right\rparen^{2}\right]$
$\displaystyle\leq\frac{4\sigma^{2}}{\epsilon^{2}}\left\lparen\mathrm{KL}(Q\,\|\,P)+\log
3\right\rparen,$

which completes the proof. ∎

######  Lemma .10.

(Kantorovich-Rubinstein Duality) Let $P$ and $Q$ be probability measures
defined on the same measurable space $\mathcal{X}$, then

$\mathbb{W}(P,Q)=\sup_{f\in
Lip_{1}}\left\\{\int_{\mathcal{X}}f\mathop{}\\!\mathrm{d}P-\int_{\mathcal{X}}f\mathop{}\\!\mathrm{d}Q\right\\},$

where $Lip_{1}$ denotes the set of $1$-Lipschitz functions in the metric $c$,
i.e. $\lvert f(x)-f(x^{\prime})\rvert\leq c(x,x^{\prime})$ for any $f\in
Lip_{1}$ and $x,x^{\prime}\in\mathcal{X}$.

######  Lemma .11.

(Pinsker’s Inequality) Let $P$ and $Q$ be probability measures defined on the
same space, then
$\mathrm{TV}(P,Q)\leq\sqrt{\frac{1}{2}\mathrm{KL}(Q\,\|\,P)}$.

######  Proposition .12.

For any constant $\lambda\in(0,1)$, we have

$\displaystyle\mathbb{P}\\{\lvert L_{s}(W)-L_{t}(W)\rvert\geq\epsilon\\}$
$\displaystyle\leq\mathbb{P}\\{\lvert
L_{s}(W)-L(W)\rvert\geq\lambda\epsilon\\}$
$\displaystyle\quad+\mathbb{P}\\{\lvert
L_{t}(W)-L(W)\rvert\geq(1-\lambda)\epsilon\\}.$

###### Proof.

Notice that $\lvert L_{s}(W)-L(W)\rvert\leq\lambda\epsilon$ and $\lvert
L_{t}(W)-L(W)\rvert\leq(1-\lambda)\epsilon$ together implies $\lvert
L_{s}(W)-L_{t}(W)\rvert\leq\epsilon$, we then have

$\mathbb{P}\\{\lvert
L_{s}(W)-L_{t}(W)\rvert\leq\epsilon\\}\geq\mathbb{P}\\{\lvert
L_{s}(W)-L(W)\rvert\leq\lambda\epsilon\bigcap\lvert
L_{t}(W)-L(W)\rvert\leq(1-\lambda)\epsilon\\}.$

This implies that

$\mathbb{P}\\{\lvert
L_{s}(W)-L_{t}(W)\rvert\geq\epsilon\\}\leq\mathbb{P}\\{\lvert
L_{s}(W)-L(W)\rvert\geq\lambda\epsilon\bigcup\lvert
L_{t}(W)-L(W)\rvert\geq(1-\lambda)\epsilon\\}.$

By applying Boole’s inequality, we then have

$\displaystyle\mathbb{P}\\{\lvert L_{s}(W)-L_{t}(W)\rvert\geq\epsilon\\}$
$\displaystyle\leq\mathbb{P}\\{\lvert
L_{s}(W)-L(W)\rvert\geq\lambda\epsilon\\}$
$\displaystyle\quad+\mathbb{P}\\{\lvert
L_{t}(W)-L(W)\rvert\geq(1-\lambda)\epsilon\\}.$

∎

## Omitted Proofs in Section III

###  -A Proof of Proposition III.1

######  Proposition III.1 (Restate).

Let $Q_{R|X}$ and $Q_{Y|R}$ denote the encoder and classifier characterized by
$f_{\phi}$ and $f_{\psi}$ respectively, and let $Q_{Y|X}$ denote the whole
model by $Q_{Y|X}=\int_{\mathcal{R}}Q_{Y|R}\mathop{}\\!\mathrm{d}Q_{R|X}$.
Then for any domain $d\in\mathcal{D}$:

$\displaystyle\mathrm{KL}(P_{Y|X,D}\,\|\,Q_{Y|X})\geq I(Y;D|X),$
$\displaystyle\mathrm{KL}(P_{Y|X,D=d}\,\|\,Q_{Y|X})\leq
I(X;Y|R,D=d)+\mathrm{KL}(P_{Y|R,D=d}\,\|\,Q_{Y|R}).$

###### Proof.

For better clarity, we abbreviate $P_{X}(X)$ as $P_{X}$ in the following
expectations for any random variable $X$.

$\displaystyle\mathrm{KL}(P_{Y|X,D}\,\|\,Q_{Y|X})$
$\displaystyle=\mathbb{E}_{D,X,Y}\left[\log\frac{P_{Y|X,D}}{Q_{Y|X}}\right]$
$\displaystyle=\mathbb{E}_{D,X,Y}\left[\log\frac{P_{Y|X,D}}{P_{Y|X}}\cdot\frac{P_{Y|X}}{Q_{Y|X}}\right]$
$\displaystyle=\mathbb{E}_{D,X,Y}\left[\log\frac{P_{Y,D|X}}{P_{Y|X}P_{D|X}}\right]+\mathbb{E}_{X,Y}\left[\log\frac{P_{Y|X}}{Q_{Y|X}}\right]$
$\displaystyle=I(Y;D|X)+\mathrm{KL}(P_{Y|X}\,\|\,Q_{Y|X})\geq I(Y;D|X).$

The last inequality is by the positiveness of the KL divergence. It holds with
equality if and only if $Q_{Y|X}=P_{Y|X}$.

To prove the second inequality, we apply Jensen’s inequality on the concave
logarithmic function:

$\displaystyle\mathrm{KL}(P_{Y|X,D=d}\,\|\,Q_{Y|X})$
$\displaystyle=\mathbb{E}_{X,Y|D=d}\left[\log\frac{P_{Y|X,D=d}}{Q_{Y|X}}\right]$
$\displaystyle=\mathbb{E}_{X,Y|D=d}\left[\log\frac{P_{Y|X,D=d}}{\mathbb{E}_{R|X=x}[Q_{Y|R}]}\right]$
$\displaystyle\leq\mathbb{E}_{X,Y|D=d}\mathbb{E}_{R|X=x}\left[\log\frac{P_{Y|X,D=d}}{Q_{Y|R}}\right]$
$\displaystyle=\mathbb{E}_{X,Y,R|D=d}\left[\log\frac{P_{Y|X,D=d}}{P_{Y|R,D=d}}\cdot\frac{P_{Y|R,D=d}}{Q_{Y|R}}\right]$
$\displaystyle=\mathbb{E}_{X,Y,R|D=d}\left[\log\frac{P_{Y,X|R,D=d}}{P_{Y|R,D=d}P_{X|R,D=d}}\right]+\mathbb{E}_{Y,R|D=d}\left[\log\frac{P_{Y|R,D=d}}{Q_{Y|R}}\right]$
$\displaystyle=I(X;Y|R,D=d)+\mathrm{KL}(P_{Y|R,D=d}\,\|\,Q_{Y|R}).$

The only inequality holds with equality when $\mathrm{Var}[Q_{R|X=x}]=0$ for
any $x\in\mathcal{X}$, i.e. $f_{\phi}$ is deterministic. This completes the
proof. ∎

###  -B Proof of Theorem III.2

######  Theorem III.2 (Restate).

If Assumption II.2 holds, then

$\displaystyle\left\lvert\mathbb{E}_{W,D_{s}}[L_{s}(W)]-\mathbb{E}_{W}[L(W)]\right\rvert\leq\frac{1}{m}\sum_{i=1}^{m}\sqrt{\frac{M^{2}}{2}I(W,D_{i})},$
$\displaystyle\mathbb{P}\left\\{\left\lvert
L_{s}(W)-L(W)\right\rvert\geq\epsilon\right\\}\leq\frac{M}{m\epsilon\sqrt{2}}\sum_{i=1}^{m}\sqrt{I(W,D_{i})}+\frac{1}{\epsilon}\mathbb{E}_{W,D}\lvert
L_{D}(W)-L(W)\rvert,$

where $D\sim\nu$ is independent of $W$.

###### Proof.

For any $D\in D_{s}$, let $\bar{D}$ be an independent copy of the marginal
distribution of $D$. Since $\ell(\cdot,\cdot)\in[0,M]$, we know that
$L_{\bar{D}}(W)$ is $\frac{M}{2}$-subgaussian. Then by setting $X=W$, $Y=D$
and $f(W,D)=L_{D}(W)$ in Lemma .6, we have

$\displaystyle\lvert\mathbb{E}_{W,D}[L_{D}(W)]-\mathbb{E}_{W}[L(W)]\rvert$
$\displaystyle=\lvert\mathbb{E}_{W,D}[L_{D}(W)]-\mathbb{E}_{W,\bar{D}}[L_{\bar{D}}(W)]\rvert\leq\sqrt{\frac{M^{2}}{2}I(W,D)}.$

By summing up the inequality above over each training domain, we can get

$\displaystyle\lvert\mathbb{E}_{W,D_{s}}[L_{s}(W)]-\mathbb{E}_{W}[L(W)]\rvert$
$\displaystyle=\left\lvert\frac{1}{m}\mathbb{E}_{W,D_{s}}\left[\sum_{i=1}^{m}L_{D_{i}}(W)\right]-\mathbb{E}_{W}[L(W)]\right\rvert$
$\displaystyle\leq\frac{1}{m}\sum_{i=1}^{m}\left\lvert\mathbb{E}_{W,D_{i}}[L_{D_{i}}(W)]-\mathbb{E}_{W}[L(W)]\right\rvert$
$\displaystyle\leq\frac{1}{m}\sum_{i=1}^{m}\sqrt{\frac{M^{2}}{2}I(W,D_{i})}.$

Similarly, for any $D\in D_{s}$, one can verify that $\lvert
L_{\bar{D}}(W)-L(W)\rvert\in[0,M]$ and is thus $\frac{M}{2}$-subgaussian. By
applying Lemma .6 with $X=W$, $Y=D$ and $f(W,D)=\lvert L_{D}(W)-L(W)\rvert$,
we obtain

$\displaystyle\mathbb{E}_{W,D}\lvert
L_{D}(W)-L(W)\rvert-\mathbb{E}_{W,\bar{D}}\lvert
L_{\bar{D}}(W)-L(W)\rvert\leq\sqrt{\frac{M^{2}}{2}I(W,D)}.$

Summing up this inequality over each training domain, we then have

$\displaystyle\mathbb{E}_{W,D_{s}}\lvert L_{s}(W)-L(W)\rvert$
$\displaystyle=\mathbb{E}_{W,D_{s}}\left\lvert\frac{1}{m}\sum_{i=1}^{m}L_{D_{i}}(W)-L(W)\right\rvert$
$\displaystyle\leq\frac{1}{m}\sum_{i=1}^{m}\mathbb{E}_{W,D_{i}}\lvert
L_{D_{i}}(W)-L(W)\rvert$
$\displaystyle\leq\frac{1}{m}\sum_{i=1}^{m}\sqrt{\frac{M^{2}}{2}I(W,D_{i})}+\mathbb{E}_{W,\bar{D}}\lvert
L_{\bar{D}}(W)-L(W)\rvert.$

By applying Markov’s inequality, we finally have

$\displaystyle\mathbb{P}\left\\{\lvert
L_{s}(W)-L(W)\rvert\geq\epsilon\right\\}\leq\frac{M}{m\epsilon\sqrt{2}}\sum_{i=1}^{m}\sqrt{I(W,D_{i})}+\frac{1}{\epsilon}\mathbb{E}_{W,\bar{D}}\lvert
L_{\bar{D}}(W)-L(W)\rvert.$

Additionally, by assuming that the training domains are independent, we have

$\displaystyle I(W;D_{s})$
$\displaystyle=I(W;\\{D_{i}\\}_{i=1}^{m})=I(W;D_{1})+I(W;\\{D_{i}\\}_{i=2}^{m}|D_{1})$
$\displaystyle=I(W;D_{1})+I(W;\\{D_{i}\\}_{i=2}^{m})-I(\\{D_{i}\\}_{i=2}^{m};D_{1})+I(\\{D_{i}\\}_{i=2}^{m};D_{1}|W)$
$\displaystyle=I(W;D_{1})+I(W;\\{D_{i}\\}_{i=2}^{m})+I(\\{D_{i}\\}_{i=2}^{m};D_{1}|W)$
$\displaystyle\geq I(W;D_{1})+I(W;\\{D_{i}\\}_{i=2}^{m})$
$\displaystyle\geq\cdots$ $\displaystyle\geq\sum_{i=1}^{m}I(W,D_{i}).$

∎

###  -C Proof of Theorem III.4

######  Theorem III.4 (Restate).

Let $W$ be the output of learning algorithm $\mathcal{A}$ under training
domains $D_{s}$. If $\ell(f_{w}(X),Y)$ is $\beta^{\prime}$-Lipschitz w.r.t
$w$, i.e.
$\lvert\ell(f_{w_{1}}(X),Y)-\ell(f_{w_{2}}(X),Y)\rvert\leq\beta^{\prime}c(w_{1},w_{2})$
for any $w_{1},w_{2}\in\mathcal{W}$, then

$\lvert\mathbb{E}_{W,D_{s}}[L_{s}(W)]-\mathbb{E}_{W}[L(W)]\rvert\leq\frac{\beta^{\prime}}{m}\sum_{i=1}^{m}\mathbb{E}_{D_{i}}[\mathbb{W}(P_{W|D_{i}},P_{W})].$

###### Proof.

For any $D_{i}\in D_{s}$, let $P=P_{W|D_{i}=d}$, $Q=P_{W}$ and
$f(w)=L_{D_{i}}(w)$ in Lemma .10, then

$\displaystyle\lvert\mathbb{E}_{W,D_{s}}[L_{s}(W)]-\mathbb{E}_{W}[L(W)]\rvert$
$\displaystyle\leq\frac{1}{m}\mathbb{E}_{D_{s}}\left[\sum_{i=1}^{m}\left\lvert\mathbb{E}_{W|D_{i}}\left[L_{D_{i}}(W)\right]-\mathbb{E}_{W}[L(W)]\right\rvert\right]$
$\displaystyle=\frac{1}{m}\mathbb{E}_{D_{s}}\left[\sum_{i=1}^{m}\lvert\mathbb{E}_{W|D_{i}}[L_{D_{i}}(W)]-\mathbb{E}_{W}[L_{D_{i}}(W)]\rvert\right]$
$\displaystyle\leq\frac{1}{m}\mathbb{E}_{D_{s}}\left[\sum_{i=1}^{m}\beta^{\prime}\mathbb{W}(P_{W|D_{i}},P_{W})\right]$
$\displaystyle=\frac{\beta^{\prime}}{m}\sum_{i=1}^{m}\mathbb{E}_{D_{i}}[\mathbb{W}(P_{W|D_{i}},P_{W})].$

When the metric $d$ is discrete, the Wasserstein distance is equal to the
total variation. Combining with Lemma .11, we have the following reductions:

$\displaystyle\mathbb{E}_{D_{i}}[\mathbb{W}(P_{W|D_{i}},P_{W})]$
$\displaystyle=\mathbb{E}_{D_{i}}[\mathrm{TV}(P_{W|D_{i}},P_{W})]$
$\displaystyle\leq\mathbb{E}_{D_{i}}\left[\sqrt{\frac{1}{2}\mathrm{KL}(P_{W|D_{i}}\,\|\,P_{W})}\right]$
$\displaystyle\leq\sqrt{\frac{1}{2}I(W;D_{i})},$

where the last inequality follows by applying Jensen’s inequality on the
concave square root function. ∎

###  -D Proof of Theorem III.5

######  Theorem III.5 (Restate).

For any $w\in\mathcal{W}$, $\mathbb{E}_{D_{t}}[L_{t}(w)]=L(w)$. Additionally
if Assumption II.3 holds, then

$\mathbb{P}\left\\{\lvert
L_{t}(w)-L(w)\rvert\geq\epsilon\right\\}\leq\frac{\sigma}{\epsilon}\sqrt{2I(Z;D)}.$

###### Proof.

By the identical marginal distribution of the test domains
$\mathcal{D}_{t}=\\{D_{k}\\}_{k=1}^{m^{\prime}}$, we have

$\displaystyle\mathbb{E}_{D_{t}}[L_{t}(w)]$
$\displaystyle=\frac{1}{m^{\prime}}\sum_{k=1}^{m^{\prime}}\mathbb{E}_{D_{k}}[L_{D_{k}}(w)]=\frac{1}{m^{\prime}}\sum_{k=1}^{m^{\prime}}\mathbb{E}_{D}[L_{D}(w)]$
$\displaystyle=\mathbb{E}_{D}[L_{D}(w)]=L(w).$

For any $d\in\mathcal{D}$, by applying Lemma .9 with $P=P_{Z}$, $Q=P_{Z|D=d}$
and $f(Z)=\ell(f_{w}(X),Y)$, we can get

$\displaystyle\lvert L_{d}(w)-L(w)\rvert$
$\displaystyle=\left\lvert\mathbb{E}_{Z|D=d}[\ell(f_{w}(X),Y)]-\mathbb{E}_{Z}[\ell(f_{w}(X),Y]\right\rvert$
$\displaystyle\leq\sqrt{2\sigma^{2}\mathrm{KL}(P_{Z|D=d}\,\|\,P_{Z})}.$

Taking the expectation over $D\sim\nu$, we can get

$\displaystyle\mathbb{E}_{D}\lvert L_{d}(w)-L(w)\rvert$
$\displaystyle\leq\mathbb{E}_{D}\sqrt{2\sigma^{2}\mathrm{KL}(P_{Z|D}\,\|\,P_{Z})}$
$\displaystyle\leq\sqrt{2\sigma^{2}\mathbb{E}_{D}[\mathrm{KL}(P_{Z|D}\,\|\,P_{Z})]}$
$\displaystyle=\sigma\sqrt{2I(Z;D)}.$

By summing up the inequality above over each test domain, we can get

$\displaystyle\mathbb{E}_{D_{t}}\lvert L_{t}(w)-L(w)\rvert$
$\displaystyle=\mathbb{E}_{D_{t}}\left\lvert\frac{1}{m^{\prime}}\sum_{k=1}^{m^{\prime}}L_{D_{k}}(w)-L(W)\right\rvert$
$\displaystyle\leq\frac{1}{m^{\prime}}\sum_{k=1}^{m^{\prime}}\mathbb{E}_{D_{k}}\left\lvert
L_{D_{k}}(w)-L(W)\right\rvert$
$\displaystyle\leq\frac{1}{m^{\prime}}\sum_{k=1}^{m^{\prime}}\sigma\sqrt{2I(Z;D_{k})}$
$\displaystyle=\sigma\sqrt{2I(Z;D)}.$

By applying Markov’s inequality, we finally have

$\displaystyle\mathbb{P}\left\\{\lvert
L_{t}(w)-L(w)\rvert\geq\epsilon\right\\}\leq\frac{\sigma}{\epsilon}\sqrt{2I(Z;D)}.$

∎

###  -E Proof of Theorem III.6

######  Theorem III.6 (Restate).

If Assumption II.2 and II.4 hold, then for any $w=(\phi,\psi)\in\mathcal{W}$,

$\mathbb{P}\left\\{L_{t}(\psi)-L(\psi)\geq\epsilon\right\\}\leq\frac{\sigma}{\epsilon}\sqrt{2I(R;D)}+\frac{2}{\epsilon}L^{*}.$

where $L^{*}=\min_{f^{*}:\mathcal{R}\mapsto\mathcal{Y}}L(f^{*})$.

###### Proof.

For any environment $d\in\mathcal{D}$, classifier $\psi$ and
$f^{*}:\mathcal{R}\mapsto\mathcal{Y}$, denote

$\displaystyle
L_{d}(\psi,f^{*})=\mathbb{E}_{R|D=d}[\ell(f_{\psi}(R),f^{*}(R))],$
$\displaystyle L(\psi,f^{*})=\mathbb{E}_{R}[\ell(f_{\psi}(R),f^{*}(R))].$

By setting $P=P_{R}$, $Q=P_{R|D=d}$ and $f(R)=\ell(f_{\psi}(R),f^{*}(R))$ and
applying Lemma .9, we have

$\left\lvert
L_{d}(\psi,f^{*})-L(\psi,f^{*})\right\rvert\leq\sqrt{2\sigma^{2}\mathrm{KL}(P_{R|D=d}\,\|\,P_{R})}.$

By taking the expectation over $D\sim\nu$, we get

$\displaystyle\mathbb{E}_{D}\left\lvert
L_{D}(\psi,f^{*})-L(\psi,f^{*})\right\rvert$
$\displaystyle\leq\mathbb{E}_{D}\sqrt{2\sigma^{2}\mathrm{KL}(P_{R|D}\,\|\,P_{R})}$
$\displaystyle\leq\sqrt{2\sigma^{2}I(R;D)}.$

If Assumption II.4 holds, then by the symmetry and triangle inequality of
$\ell(\cdot,\cdot)$, we have

$\displaystyle L_{d}(\psi,f^{*})$
$\displaystyle=\mathbb{E}_{R|D=d}[\ell(f_{\psi}(R),f^{*}(R))]$
$\displaystyle\leq\mathbb{E}_{R,Y|D=d}[\ell(f_{\psi}(R),Y)+\ell(f^{*}(R),Y)]$
$\displaystyle=L_{d}(\psi)+L_{d}(f^{*}).$ $\displaystyle L(\psi,f^{*})$
$\displaystyle\leq L(\psi)+L(f^{*}).$

Similarly, we can prove that

$\displaystyle L_{d}(\psi,f^{*})$
$\displaystyle=\mathbb{E}_{R|D=d}[\ell(f_{\psi}(R),f^{*}(R))]$
$\displaystyle\geq\mathbb{E}_{R,Y|D=d}[\ell(f_{\psi}(R),Y)-\ell(f^{*}(R),Y)]$
$\displaystyle=L_{d}(\psi)-L_{d}(f^{*}).$ $\displaystyle L(\psi,f^{*})$
$\displaystyle\geq L(\psi)-L(f^{*}).$

Combining the results above, we have

$\displaystyle L_{d}(\psi)-L(\psi)$ $\displaystyle\leq
L_{d}(\psi,f^{*})+L_{d}(f^{*})-L(\psi,f^{*})+L(f^{*}),$ $\displaystyle
L(\psi)-L_{d}(\psi)$ $\displaystyle\leq
L(\psi,f^{*})+L(f^{*})-L_{d}(\psi,f^{*})+L_{d}(f^{*}).$

Combining the two inequalities above, we then get

$\left\lvert L_{d}(\psi)-L(\psi)\right\rvert\leq\left\lvert
L_{d}(\psi,f^{*})-L(\psi,f^{*})\right\rvert+L_{d}(f^{*})+L(f^{*}).$

By taking the expectation over $D\sim\nu$, we obtain

$\displaystyle\mathbb{E}_{D}\left\lvert L_{D}(\psi)-L(\psi)\right\rvert$
$\displaystyle\leq\mathbb{E}_{D}\left\lvert
L_{D}(\psi,f^{*})-L(\psi,f^{*})\right\rvert+\mathbb{E}_{D}[L_{D}(f^{*})+L(f^{*})]$
$\displaystyle\leq\sqrt{2\sigma^{2}I(R;D)}+2L(f^{*}).$

Summing up the inequality above over each test domain, we can get

$\displaystyle\mathbb{E}_{D_{t}}\left\lvert L_{t}(\psi)-L(\psi)\right\rvert$
$\displaystyle\leq\mathbb{E}_{D_{t}}\left\lvert\frac{1}{m^{\prime}}\sum_{k=1}^{m^{\prime}}L_{D_{k}}(\psi)-L(\psi)\right\rvert$
$\displaystyle\leq\frac{1}{m^{\prime}}\sum_{k=1}^{m^{\prime}}\mathbb{E}_{D_{k}}\left\lvert
L_{D_{k}}(\psi)-L(\psi)\right\rvert$
$\displaystyle\leq\frac{1}{m^{\prime}}\sum_{k=1}^{m^{\prime}}\left\lparen\sqrt{2\sigma^{2}I(R;D)}+2L(f^{*})\right\rparen$
$\displaystyle=\sqrt{2\sigma^{2}I(R;D)}+2L(f^{*}).$

By applying Markov’s inequality, we finally have

$\displaystyle\mathbb{P}\left\\{L_{t}(\psi)-L(\psi)\geq\epsilon\right\\}\leq\frac{\sigma}{\epsilon}\sqrt{2I(R;D)}+\frac{2}{\epsilon}L(f^{*}).$

The proof is complete by taking the minimizer of $\min_{f^{*}}[L(f^{*})]$. ∎

###  -F Proof o Proposition III.7

######  Proposition III.7 (Restate).

Assume that $P_{R,D}\ll P_{R_{i},D_{i}}$ and $P_{R_{i},D_{i}}\ll P_{R,D}$,
then

$\mathrm{SKL}(P_{R,D}\,\|\,P_{R_{i},D_{i}})=O(\sqrt{I(W;D_{i})}).$

###### Proof.

The condition $P_{R,D}\ll P_{R_{i},D_{i}}$ and $P_{R_{i},D_{i}}\ll P_{R,D}$
implies that there exists $B>1$, such that for any $r\in\mathcal{R}$,
$d\in\mathcal{D}$, we have
$\frac{P_{R,D}(r,d)}{P_{R_{i},D_{i}}(r,d)}\in[\frac{1}{B},B]$. Therefore,
$\log\frac{P_{R,D}}{P_{R_{i},D_{i}}}\in[-\log(B),\log(B)]$ and is
$\log(B)$-subgaussian.

By applying Lemma .6 with $X=W$, $Y=D_{i}$ and
$f(W,D)=\mathbb{E}_{X|D}[\log\frac{P_{R,D}}{P_{R_{i},D_{i}}}]$, we have

$\displaystyle\left\lvert\mathbb{E}_{W,D_{i},R_{i}}\left[\log\frac{P_{R,D}(R_{i},D_{i})}{P_{R_{i},D_{i}}(R_{i},D_{i})}\right]-\mathbb{E}_{W,D,R}\left[\log\frac{P_{R,D}(R,D)}{P_{R_{i},D_{i}}(R,D)}\right]\right\rvert$
$\displaystyle\leq\sqrt{2\log^{2}(B)I(W;D_{i})}$
$\displaystyle\left\lvert\mathrm{KL}(P_{R,D}\,\|\,P_{R_{i},D_{i}})+\mathrm{KL}(P_{R_{i},D_{i}}\,\|\,P_{R,D})\right\rvert$
$\displaystyle\leq\sqrt{2\log^{2}(B)I(W;D_{i})}$
$\displaystyle\mathrm{SKL}(P_{R,D}\,\|\,P_{R_{i},D_{i}})$
$\displaystyle\leq\log(B)\sqrt{2I(W;D_{i})}.$

∎

## Omitted Proofs in Section IV

###  -G Proof of Theorem III.3

######  Theorem III.3 (Restate).

Let $G_{t}^{i}=-\eta_{t}g(W_{t-1},B_{t}^{i})$ and
$G_{t}=\sum_{i=1}^{m}G_{t}^{i}$, then

$I(W_{T};D_{i})\leq\sum_{t=1}^{T}I(G_{t};D_{i}|W_{t-1}).$

Additionally, if the training domains are independent, then

$I(W_{T};D_{i})\leq\sum_{t=1}^{T}I(G_{t}^{i};D_{i}|W_{t-1}).$

###### Proof.

Noticing the Markov chain relationship
$D_{i}\rightarrow(W_{T-1},G_{T})\rightarrow W_{T-1}+G_{T}$, then by the data
processing inequality

$\displaystyle I(W_{T};D_{i})$ $\displaystyle=I\left\lparen
W_{T-1}+G_{T};D_{i}\right\rparen$ $\displaystyle\leq I\left\lparen
W_{T-1},G_{T};D_{i}\right\rparen$
$\displaystyle=I(W_{T-1};D_{i})+I\left\lparen
G_{T};D_{i}|W_{T-1}\right\rparen.$

where the last equality is by the chain rule of conditional mutual
information. By applying the reduction steps above recursively, we can get

$\displaystyle I(W_{T};D_{i})$ $\displaystyle\leq
I(W_{T-1};D_{i})+I\left\lparen G_{T};D_{i}|W_{T-1}\right\rparen$
$\displaystyle\leq I(W_{T-2};D_{i})+I\left\lparen
G_{T-1};D_{i}|W_{T-2}\right\rparen$ $\displaystyle\qquad+I\left\lparen
G_{T};D_{i}|W_{T-1}\right\rparen$ $\displaystyle\leq\cdots$
$\displaystyle\leq\sum_{t=1}^{T}I\left\lparen
G_{t};D_{i}|W_{t-1}\right\rparen.$

When the training domains are independent, we additionally have

$\displaystyle I(W_{T};D_{i})$ $\displaystyle\leq
I(W_{T-1};D_{i})+I\left\lparen G_{T};D_{i}|W_{T-1}\right\rparen$
$\displaystyle\leq
I(W_{T-1};D_{i})+I\left\lparen\\{G_{T}^{k}\\}_{k=1}^{m};D_{i}|W_{T-1}\right\rparen$
$\displaystyle=I(W_{T-1};D_{i})+I\left\lparen
G_{T}^{i};D_{i}|W_{T-1}\right\rparen+I\left\lparen\\{G_{T}^{k}\\}_{k=1}^{m}\setminus
G_{T}^{i};D_{i}|W_{T-1},G_{T}^{i}\right\rparen$
$\displaystyle=I(W_{T-1};D_{i})+I\left\lparen
G_{T}^{i};D_{i}|W_{T-1}\right\rparen.$

Then by the same scheme of recursive reduction, we can prove that

$I(W_{T};D_{i})\leq\sum_{t=1}^{T}I\left\lparen
G_{t}^{i};D_{i}|W_{t-1}\right\rparen.$

∎

###  -H Proof of Theorem IV.1

######  Theorem IV.1 (Restate).

[Formal] Let $n$ be the dimensionality of the distribution and $b$ be the
number of data points, then

  * •

If $n>b+1$, then for any sampled data points $s=\\{x_{i}\\}_{i=1}^{b}$, there
exists infinite environments $d_{1}$, $d_{2}$, $\cdots$ such that
$p(S=s|D=d_{1})=p(S=s|D=d_{2})=\cdots$.

  * •

If $n>2b+1$, then for any two batch of sampled data points
$s_{1}=\\{x_{i}^{1}\\}_{i=1}^{b}$ and $s_{2}=\\{x_{i}^{2}\\}_{i=1}^{b}$, there
exists infinite environments $d_{1}$, $d_{2}$, $\cdots$ such that for each
$j\in[1,\infty)$, $p(S=s_{1}|D=d_{j})=p(S=s_{2}|D=d_{j})$.

###### Proof.

Without loss of generality, we assume that the data-generating distributions
$p(X|D)$ are Gaussian with zero means for simplicity, i.e.

$p(x|D=d)=\frac{1}{\sqrt{(2\pi)^{n}\lvert\Sigma_{d}\rvert}}\exp\left\lparen-\frac{1}{2}x^{\top}\Sigma_{d}x\right\rparen,$

where $\Sigma_{d}$ is the corresponding covariance matrix of environment $d$.
Let $X\in\mathbb{R}^{n\times b}$ be the data matrix of $S$ such that the
$i$-th column of $X$ equals $x_{i}$, we then have

$p(S=s|D=d)=\frac{1}{\sqrt{(2\pi)^{bn}\lvert\Sigma_{d}\rvert^{b}}}\exp\left\lparen-\frac{1}{2}\mathrm{tr}(X^{\top}\Sigma_{d}X)\right\rparen.$

Since the rank of $X$ is at most $b$, one can decompose
$\Sigma_{d}=\Sigma_{d}^{1}+\Sigma_{d}^{2}$ with
$\mathrm{rank}(\Sigma_{d}^{1})=b$ and $\mathrm{rank}(\Sigma_{d}^{2})=n-b\geq
2$ through eigenvalue decomposition, and let the eigenvector space of
$\Sigma_{d}^{1}$ cover the column space of $X$. Then we have

$\mathrm{tr}(X^{\top}\Sigma_{d}^{1}X)=\mathrm{tr}(X^{\top}\Sigma_{d}X),\quad\textrm{and}\quad\mathrm{tr}(X^{\top}\Sigma_{d}^{2}X)=0.$

Therefore, one can arbitrarily modify the eigenvector space of
$\Sigma_{d}^{2}$ as long as keeping it orthogonal to that of $\Sigma_{d}^{1}$,
without changing the value of $\mathrm{tr}(X^{\top}\Sigma_{d}X)$. This
finishes the proof of the first part.

To prove the second part, similarly we decompose $\Sigma_{d}$ by
$\Sigma_{d}^{1}+\Sigma_{d}^{2}$ such that $\mathrm{rank}(\Sigma_{d}^{1})=2b+1$
and $\mathrm{rank}(\Sigma_{d}^{2})=n-2b-1\geq 1$, and make the eigenvector
space of $\Sigma_{d}^{1}$ cover the column space of both $X_{1}$ and $X_{2}$,
where $X_{1}$ and $X_{2}$ are the data matrix of $S_{1}$ and $S_{2}$
respectively. We then have

$\displaystyle\mathrm{tr}(X_{1}^{\top}\Sigma_{d}^{1}X_{1})$
$\displaystyle=\mathrm{tr}(X_{1}^{\top}\Sigma_{d}X_{1}),$
$\displaystyle\mathrm{tr}(X_{2}^{\top}\Sigma_{d}^{1}X_{2})$
$\displaystyle=\mathrm{tr}(X_{2}^{\top}\Sigma_{d}X_{2}),$
$\displaystyle\textrm{and}\quad\mathrm{tr}(X_{1}^{\top}\Sigma_{d}^{2}X_{1})$
$\displaystyle=\mathrm{tr}(X_{2}^{\top}\Sigma_{d}^{2}X_{2})=0.$

Let $\Sigma_{d}^{1}=U_{d}^{\top}\Lambda_{d}U_{d}$ be the eigenvalue
decomposition of $\Sigma_{d}^{1}$, where $U_{d}\in\mathbb{R}^{(2b+1)\times n}$
and $\Lambda_{d}=\mathrm{diag}(\lambda_{1}^{d},\cdots,\lambda_{2b+1}^{d})$.
Notice that for any $x\in\mathbb{R}^{n}$, we have
$x^{\top}\Sigma_{d}x=(U_{d}x)^{\top}\Lambda_{d}(U_{d}x)=\sum_{i=1}^{2b+1}(U_{d}x)_{i}^{2}\lambda_{i}$.
By assuming that $p(S=s_{1}|D=d)=p(S=s_{2}|D=d)$, we have the following
homogeneous linear equations:

$\displaystyle
a_{1}^{1}\lambda_{1}+a_{1}^{2}\lambda_{2}+\cdots+a_{1}^{2b+1}\lambda_{2b+1}$
$\displaystyle=0,$ $\displaystyle
a_{2}^{1}\lambda_{1}+a_{2}^{2}\lambda_{2}+\cdots+a_{2}^{2b+1}\lambda_{2b+1}$
$\displaystyle=0,$ $\displaystyle\cdots$ $\displaystyle
a_{b}^{1}\lambda_{1}+a_{b}^{2}\lambda_{2}+\cdots+a_{b}^{2b+1}\lambda_{2b+1}$
$\displaystyle=0,$

where $a_{i}^{j}=(U_{d}x_{i}^{1})_{j}^{2}-(U_{d}x_{i}^{2})_{j}^{2}$. Since
$2b+1>b$, the linear system above has infinite non-zero solutions, which
finishes the proof of the second part. ∎

###  -I Proof of Theorem IV.2

######  Theorem IV.2 (Restate).

Let $f$ be a bijection mapping from $[1,b]$ to $[1,b]$, then
$\mathrm{KL}(\bar{P}\,\|\,\bar{Q})\leq\frac{1}{b}\sum_{i=1}^{b}\mathrm{KL}(P_{i}\,\|\,Q_{f(i)})$,
and
$\mathbb{W}(\bar{P},\bar{Q})\leq\frac{1}{b}\sum_{i=1}^{b}\mathbb{W}(P_{i},Q_{f(i)})$,
where $P_{i}$ is the probability measure defined by $p_{i}$ (respectively for
$Q_{i}$).

###### Proof.

Recall that $\bar{p}(x)=\frac{1}{b}\sum_{i=1}^{b}p_{i}(x)$ and
$\bar{q}(x)=\frac{1}{b}\sum_{i=1}^{b}q_{i}(x)$, we then have

$\displaystyle\mathrm{KL}(\bar{P}\,\|\,\bar{Q})$
$\displaystyle=\int_{\mathcal{X}}\bar{p}(x)\log\left\lparen\frac{\bar{p}(x)}{\bar{q}(x)}\right\rparen\mathop{}\\!\mathrm{d}x$
$\displaystyle=-\int_{\mathcal{X}}\bar{p}(x)\log\left\lparen\frac{1}{b}\sum_{i=1}^{b}\frac{p_{i}(x)}{\bar{p}(x)}\cdot\frac{q_{f(i)}(x)}{p_{i}(x)}\right\rparen\mathop{}\\!\mathrm{d}x$
$\displaystyle\leq-\int_{\mathcal{X}}\bar{p}(x)\frac{1}{b}\sum_{i=1}^{b}\frac{p_{i}(x)}{\bar{p}(x)}\log\left\lparen\frac{q_{f(i)}(x)}{p_{i}(x)}\right\rparen\mathop{}\\!\mathrm{d}x$
$\displaystyle=-\frac{1}{b}\sum_{i=1}^{b}\int_{\mathcal{X}}p_{i}(x)\log\left\lparen\frac{q_{f(i)}(x)}{p_{i}(x)}\right\rparen\mathop{}\\!\mathrm{d}x$
$\displaystyle=\frac{1}{b}\sum_{i=1}^{b}\mathrm{KL}(P_{i}\,\|\,Q_{f(i)}),$

where the only inequality follows by applying Jensen’s inequality on the
concave logarithmic function. This finishes the proof of the upper bound for
KL divergence.

To prove the counterpart for Wasserstein distance, we apply Lemma .10 on
$\bar{P}$ and $\bar{Q}$:

$\displaystyle\mathbb{W}(\bar{P},\bar{Q})$ $\displaystyle=\sup_{f\in
Lip_{1}}\left\\{\int_{\mathcal{X}}f\mathop{}\\!\mathrm{d}\bar{P}-\int_{\mathcal{X}}f\mathop{}\\!\mathrm{d}\bar{Q}\right\\}$
$\displaystyle=\sup_{f\in
Lip_{1}}\left\\{\int_{\mathcal{X}}f\mathop{}\\!\mathrm{d}\left\lparen\frac{1}{b}\sum_{i=1}^{b}P_{i}\right\rparen-\int_{\mathcal{X}}f\mathop{}\\!\mathrm{d}\left\lparen\frac{1}{b}\sum_{i=1}^{b}Q_{f(i)}\right\rparen\right\\}$
$\displaystyle\leq\frac{1}{b}\sum_{i=1}^{b}\sup_{f\in
Lip_{1}}\left\\{\int_{\mathcal{X}}f\mathop{}\\!\mathrm{d}P_{i}-\int_{\mathcal{X}}f\mathop{}\\!\mathrm{d}Q_{f(i)}\right\\}$
$\displaystyle=\frac{1}{b}\sum_{i=1}^{b}\mathbb{W}(P_{i},Q_{f(i)}).$

The proof is complete. ∎

###  -J Proof of Theorem IV.3

######  Theorem IV.3 (Restate).

Suppose that $\\{x_{i}^{1}\\}_{i=1}^{b}$, $\\{x_{i}^{2}\\}_{i=1}^{b}$ are
sorted in the same order, then $f(j)=j$ is the minimizer of
$\sum_{i=1}^{b}\mathrm{KL}(P_{i}\,\|\,Q_{f(i)})$ and
$\sum_{i=1}^{b}\mathbb{W}(P_{i},Q_{f(i)})$.

###### Proof.

For simplicity, we assume that all data points of $\\{x_{i}^{1}\\}_{i=1}^{b}$
and $\\{x_{i}^{2}\\}_{i=1}^{b}$ are different from each other. Since $P_{i}$
and $Q_{i}$ are Gaussian distributions with the same variance, the KL
divergence and Wasserstein distance between them could be analytically
acquired:

$\mathrm{KL}(P_{i}\,\|\,Q_{j})=\frac{(x_{i}^{1}-x_{j}^{2})^{2}}{2\sigma^{2}},\quad\textrm{and}\quad\mathbb{W}(P_{i},Q_{j})=\lvert
x_{i}^{1}-x_{j}^{2}\rvert.$

Suppose there exists $i\in[1,b]$ such that $f(i)\neq i$. Without loss of
generality, we assume that $f(i)>i$. Then by the pigeonhole principle, there
exists $j\in(i,b]$ that satisfies $f(j)<f(i)$. Suppose that
$\\{x_{i}^{1}\\}_{i=1}^{b}$, $\\{x_{i}^{2}\\}_{i=1}^{b}$ are both sorted in
ascending order, we have $x_{i}^{1}<x_{j}^{1}$ and
$x_{f(i)}^{2}>x_{f(j)}^{2}$. For any $p\in\\{1,2\\}$, the following $3$ cases
cover all possible equivalent combinations of the order of $x_{i}^{1}$,
$x_{j}^{1}$, $x_{f(j)}^{2}$ and $x_{f(i)}^{2}$:

  * •

When $x_{i}^{1}<x_{j}^{1}<x_{f(j)}^{2}<x_{f(i)}^{2}$ and $p=2$, we have

$\displaystyle(x_{i}^{1}-x_{f(i)}^{2})^{2}+(x_{j}^{1}-x_{f(j)}^{2})^{2}-(x_{i}^{1}-x_{f(j)}^{2})^{2}-(x_{j}^{1}-x_{f(i)}^{2})^{2}$
$\displaystyle=(2x_{i}^{1}-x_{f(i)}^{2}-x_{f(j)}^{2})(x_{f(j)}^{2}-x_{f(i)}^{2})-(2x_{j}^{1}-x_{f(j)}^{2}-x_{f(i)}^{2})(x_{f(j)}^{2}-x_{f(i)}^{2})$
$\displaystyle=(x_{f(j)}^{2}-x_{f(i)}^{2})(2x_{i}^{1}-2x_{j}^{1})>0.$

Elsewise when $p=1$, we have

$\lvert x_{i}^{1}-x_{f(i)}^{2}\rvert+\lvert
x_{j}^{1}-x_{f(j)}^{2}\rvert=\lvert x_{i}^{1}-x_{f(j)}^{2}\rvert+\lvert
x_{j}^{1}-x_{f(i)}^{2}\rvert.$

  * •

When $x_{i}^{1}<x_{f(j)}^{2}<x_{j}^{1}<x_{f(i)}^{2}$, we have

$\lvert x_{i}^{1}-x_{f(i)}^{2}\rvert^{p}>\lvert
x_{i}^{1}-x_{f(j)}^{2}\rvert^{p}+\lvert x_{j}^{1}-x_{f(i)}^{2}\rvert^{p}.$

  * •

When $x_{i}^{1}<x_{f(j)}^{2}<x_{f(i)}^{2}<x_{j}^{1}$, we have

$\displaystyle\lvert x_{i}^{1}-x_{f(i)}^{2}\rvert^{p}+\lvert
x_{j}^{1}-x_{f(j)}^{2}\rvert^{p}$ $\displaystyle\geq\lvert
x_{i}^{1}-x_{f(j)}^{2}\rvert^{p}+\lvert x_{f(i)}^{2}-x_{f(j)}^{2}\rvert^{p}$
$\displaystyle\quad+\lvert x_{j}^{1}-x_{f(i)}^{2}\rvert^{p}+\lvert
x_{f(i)}^{2}-x_{f(j)}^{2}\rvert^{p}$ $\displaystyle>\lvert
x_{i}^{1}-x_{f(j)}^{2}\rvert^{p}+\lvert x_{j}^{1}-x_{f(i)}^{2}\rvert^{p}.$

In conclusion, under all possible circumstances, we have $\lvert
x_{i}^{1}-x_{f(i)}^{2}\rvert^{p}+\lvert
x_{j}^{1}-x_{f(j)}^{2}\rvert^{p}\geq\lvert
x_{i}^{1}-x_{f(j)}^{2}\rvert^{p}+\lvert x_{j}^{1}-x_{f(i)}^{2}\rvert^{p}$,
which implies that by setting $f^{\prime}(i)=f(j)$, $f^{\prime}(j)=f(i)$ and
$f^{\prime}(k)=f(k)$ for $k\notin\\{i,j\\}$, $f^{\prime}$ will be a better
choice over $f$ to minimize $\mathrm{KL}(\bar{P}\,\|\,\bar{Q})$ or
$\mathbb{W}(\bar{P},\bar{Q})$. The proof is complete since the existence of a
minimizer is obvious. ∎

## Further Discussions

###  -K Information Bottleneck for Target-domain Generalization

Alternatively, one can also decompose the representation space distribution
shift from an anti-causal perspective:

$I(R,Y;D)\textrm{ (distribution shift)}=I(Y;D)\textrm{ (label
shift)}+I(R;D|Y)\textrm{ (concept shift)}.$

While label shift $I(Y;D)$ is an intrinsic property of the data distribution
and cannot be optimized by learning algorithms, the anti-causal concept shift
$I(R;D|Y)$ is closely connected to the IB principle: Notice the Markov chain
$(D,Y)-X-R$, we then have that by applying the data-processing inequality,

$I(R;D|Y)=I(R;D,Y)-I(R;Y)\leq I(R;X)-I(R;Y).$

Recall that the spirit of IB is to minimize $I(X;R)$ while maximizing
$I(R;Y)$, this target can be achieved by solely penalizing $I(R;X|Y)$ [47, 48,
49]. Therefore, when there is no label imbalance issues ($I(Y;D)=0$), the
anti-causal concept shift $I(R;D|Y)$ can be minimized by the IB principle:
Information bottleneck promotes test-domain generalization if
$I(Y;D)\rightarrow 0$. Notably, our analysis facilitates previous works [35]
utilizing IB to enhance the performance of OOD generalization. While the
analysis of [35] is primarily restricted to linear models, our results apply
to any encoder-classifier type network. Similar to Proposition III.7, we
provide the following result on the feasibility of utilizing the empirical IB
$I(R_{i};X_{i}|Y_{i})$ as a proxy to optimize $I(R;X|Y)$.

######  Proposition .13.

Under mild conditions, we have

$\mathrm{SKL}(P_{R,X,Y}\,\|\,P_{R_{i},X_{i},Y_{i}})=O\left\lparen\sqrt{I(W;D_{i})}\right\rparen.$

###### Proof.

The proof follows the same development as Proposition III.7. ∎

###  -L Leveraging the Independence Assumption

In the main text, we only assume that the test domains are independent of the
training domains, while the training domains are not necessarily independent
of each other (same for test domains). This assumption is much weaker than the
i.i.d domains assumption adopted in [16] by allowing correlations between
training domains, e.g. sampling from a finite set without replacement. While
this weaker assumption is preferable, we highlight that the test-domain
generalization bounds in Theorem III.5 and III.6 can be further tightened by a
factor of $1/m^{\prime}$ when the i.i.d condition of test domains is
incorporated. Therefore, one can now guarantee better generalization by
increasing the number of domains, which is consistent with real-world
observations.

######  Theorem .14.

If Assumption II.3 holds and the test domains $D_{t}$ are independent of each
other, then

$\mathbb{P}\left\\{\lvert
L_{t}(w)-L(w)\rvert\geq\epsilon\right\\}\leq\frac{2\sigma^{2}}{m^{\prime}\epsilon^{2}}I(Z;D).$

###### Proof.

If Assumption II.3 holds, then for any $d\in\mathcal{D}$, by setting
$P=P_{Z}$, $Q=P_{Z|D=d}$ and $f(z)=\ell(f_{w}(x),y)$ in Lemma .9, we have

$\displaystyle(L_{d}(w)-L(w))^{2}$
$\displaystyle=\left\lparen\mathbb{E}_{Z|D=d}[\ell(f_{w}(X),Y)]-\mathbb{E}_{Z}[\ell(f_{w}(X),Y]\right\rparen^{2}$
$\displaystyle\leq\left\lparen\sqrt{2\sigma^{2}\mathrm{KL}(P_{Z|D=d}\,\|\,P_{Z})}\right\rparen^{2}$
$\displaystyle\leq 2\sigma^{2}\mathrm{KL}(P_{Z|D=d}\,\|\,P_{Z}).$

Taking the expectation over $D\sim\nu$, we can get

$\displaystyle\mathbb{E}_{D}[(L_{D}(w)-L(w))^{2}]$
$\displaystyle\leq\mathbb{E}_{D}\left[2\sigma^{2}\mathrm{KL}(P_{Z|D=d}\,\|\,P_{Z})\right]$
$\displaystyle=2\sigma^{2}\mathrm{KL}(P_{Z|D}\,\|\,P_{Z})$
$\displaystyle=2\sigma^{2}I(Z;D).$

When the test domains $\\{D_{k}\\}_{i=k}^{m^{\prime}}$ are independent of each
other, they can be regarded as i.i.d copies of $D$, i.e.

$\displaystyle\mathrm{Var}_{D_{t}}[L_{t}(w)]$
$\displaystyle=\frac{1}{m^{\prime
2}}\sum_{k=1}^{m^{\prime}}\mathrm{Var}_{D_{k}}[L_{D_{k}}(w)]$
$\displaystyle=\frac{1}{m^{\prime
2}}\sum_{k=1}^{m^{\prime}}\mathbb{E}_{D_{k}}[(L_{D_{k}}(w)-L(w))^{2}]$
$\displaystyle\leq\frac{1}{m^{\prime
2}}\sum_{k=1}^{m^{\prime}}2\sigma^{2}I(Z;D)$
$\displaystyle=\frac{1}{m^{\prime}}2\sigma^{2}I(Z;D).$

Finally, by applying Chebyshev’s inequality, we can prove that

$\mathbb{P}\left\\{\lvert
L_{t}(w)-L(w)\rvert\geq\epsilon\right\\}\leq\frac{2\sigma^{2}}{m^{\prime}\epsilon^{2}}I(Z;D).$

∎

######  Theorem .15.

If Assumption II.2 and II.4 hold and the test domains $D_{t}$ are independent
of each other, then for any $w=(\phi,\psi)\in\mathcal{W}$,

$\mathbb{P}\left\\{L_{t}(\psi)-L(\psi)\geq\epsilon+L^{*}\right\\}\leq\frac{2\sigma^{2}}{m^{\prime}\epsilon^{2}}I(X;D),$

where $L^{*}=\min_{f^{*}}L_{t}(f^{*})+L(f^{*})$.

###### Proof.

For any environment $d\in\mathcal{D}$, classifier $\psi$ and
$f^{*}:\mathcal{R}\mapsto\mathcal{Y}$, denote

$\displaystyle
L_{d}(\psi,f^{*})=\mathbb{E}_{R|D=d}[\ell(f_{\psi}(R),f^{*}(R))],$
$\displaystyle L(\psi,f^{*})=\mathbb{E}_{R}[\ell(f_{\psi}(R),f^{*}(R))].$

By setting $P=P_{R}$, $Q=P_{R|D=d}$ and $f(R)=\ell(f_{\psi}(R),f^{*}(R))$ and
applying Lemma .9, we have

$\left\lparen L_{d}(\psi,f^{*})-L(\psi,f^{*})\right\rparen^{2}\leq
2\sigma^{2}\mathrm{KL}(P_{X|D=d}\,\|\,P_{X}).$

By taking the expectation over $D\sim\nu$, we get

$\displaystyle\mathbb{E}_{D}\left[\left\lparen
L_{d}(\psi,f^{*})-L(\psi,f^{*})\right\rparen^{2}\right]$ $\displaystyle\leq
2\sigma^{2}\mathrm{KL}(P_{X|D}\,\|\,P_{X})$ $\displaystyle=2\sigma^{2}I(X;D).$

Through a similar procedure of proving Theorem III.5, we have

$\mathbb{P}\left\\{\lvert
L_{d}(\psi,f^{*})-L(\psi,f^{*})\rvert\geq\epsilon\right\\}\leq\frac{2\sigma^{2}}{m^{\prime}\epsilon^{2}}I(X;D).$

Recall that in Theorem III.6 we proved

$\displaystyle L_{d}(\psi,f^{*})\leq L_{d}(\psi)+L_{d}(f^{*}).$ $\displaystyle
L(\psi,f^{*})\leq L(\psi)+L(f^{*}).$ $\displaystyle L_{d}(\psi,f^{*})\geq
L_{d}(\psi)-L_{d}(f^{*}).$ $\displaystyle L(\psi,f^{*})\geq L(\psi)-L(f^{*}).$

Combining the results above, we have that with probability at least
$1-\frac{2\sigma^{2}}{m^{\prime}\epsilon^{2}}I(X;D)$,

$\displaystyle L_{t}(\psi)$ $\displaystyle\leq L_{t}(\psi,f^{*})+L_{t}(f^{*})$
$\displaystyle\leq L(\psi,f^{*})+\epsilon+L_{t}(f^{*})$ $\displaystyle\leq
L(\psi)+L(f^{*})+\epsilon+L_{t}(f^{*}).$

By minimizing $L_{t}(f^{*})+L(f^{*})$, it follows that

$\mathbb{P}\left\\{L_{t}(\psi)-L(\psi)\geq\epsilon+\min_{f^{*}}(L_{t}(f^{*})+L(f^{*}))\right\\}\leq\frac{2\sigma^{2}}{m^{\prime}\epsilon^{2}}I(X;D),$

which finishes the proof. ∎

Furthermore, when the training domains satisfy the i.i.d condition, we prove
in Theorem .18 that $\sum_{i=1}^{m}I(W;D_{i})\leq I(W;D_{s})$. Otherwise, we
can only guarantee that for any $i\in[1,m]$, $I(W;D_{i})\leq I(W;D_{s})$. This
indicates that while the model achieves source-domain generalization by
letting $I(W;D_{i})\rightarrow 0$, it still learns from the training domains
$D_{s}$. Notably, having $I(W;D_{i})=0$ for each $i\in[1,m]$ does not
necessarily lead to $I(W;D_{s})=0$. To see this, we take $D_{i}$ as i.i.d
random binary variables such that
$\mathbb{P}(D_{i}=0)=\mathbb{P}(D_{i}=1)=\frac{1}{2}$, and let
$W=D_{1}\oplus\cdots\oplus D_{m}$, where $\oplus$ is the XOR operator. Then it
is easy to verify that $W$ is independent of each $D_{i}$ since
$P_{W|D_{i}}=P_{W}$, implying $I(W;D_{i})=0$. However, $I(W;D_{s})=H(W)$ is
strictly positive.

###  -M High-probability Problem Formulation

Another high-probability formulation of the DG problem is presented by [16],
namely Quantile Risk Minimization (QRM). Under our notations, the QRM
objective can be expressed as:

$\min_{w}\epsilon\quad
s.t.\quad\mathbb{P}\\{L_{t}(w)\geq\epsilon\\}\leq\delta.$

The main difference between our formulation in Problem II.6 and QRM is that we
not only consider the randomness of $D_{t}$, but also those of $D_{s}$ and
$W$. This randomized nature of training domains and the hypothesis serve as
the foundation of our information-theoretic generalization analysis. When the
training-domain risks have been observed, i.e. $L_{s}(W)$ is fixed, our
formulation reduces to QRM. The main advantage of our formulation is that it
could be directly optimized by learning algorithms without further assumptions
or simplifications. On the contrary, [16] needs to further adopt kernel
density estimation to approximate the quantile of the risks and transform the
QRM problem to the empirical one (EQRM). Further advantages of our formulation
include:

  * •

It is questionable to use the training risk distribution as a substitute for
test risk distribution to minimize the quantile objective. Since the
hypothesis and the training domains are dependent, there exists a gap between
these two distributions which is not addressed by [16].

  * •

Kernel density estimation required by QRM is challenging when the number of
training domains is not sufficiently large. For comparison, IDM could be
easily applied as long as there are at least $2$ training domains.

  * •

Our formulation aims to find the optimal learning algorithm instead of the
optimal hypothesis. This would be essential to analyze the correlations
between the hypothesis $W$ and training domains $D_{s}$, and also is more
suitable in robust learning settings when measuring the error bar.

  * •

Our formulation directly characterizes the trade-off between optimization and
generalization, which is the main challenge to achieve invariance across
different domains [9].

###  -N Tighter Bounds for Target-Domain Population Risk

In a similar vein, we provide the following two upper bounds for test-domain
generalization error in terms of Wasserstein distances. For the following
analysis, we assume the independence between training and test domains.

######  Theorem .16.

If Assumption II.5 holds, then for any $w\in\mathcal{W}$,

$\mathbb{P}\left\\{\lvert
L_{t}(w)-L(w)\rvert\geq\epsilon\right\\}\leq\frac{\beta^{2}}{m^{\prime}\epsilon^{2}}\mathbb{E}_{D}[\mathbb{W}^{2}(P_{Z|D=d},P_{Z})].$

Furthermore, if the metric $d$ is discrete, then

$\mathbb{P}\left\\{\lvert
L_{t}(w)-L(w)\rvert\geq\epsilon\right\\}\leq\frac{\beta^{2}}{2m^{\prime}\epsilon^{2}}I(Z;D).$

###### Proof.

For any $d\in\mathcal{D}$, by setting $P=P_{Z|D=d}$, $Q=P_{Z}$ and
$f(z)=\frac{1}{\beta}\ell(f_{w}(x),y)$ in Lemma .10, we have

$(L_{d}(w)-L(w))^{2}\leq\beta^{2}\mathbb{W}^{2}(P_{Z|D=d},P_{Z}).$

Following a similar procedure with the proof of Theorem III.5, we have

$\mathbb{P}\left\\{\lvert
L_{t}(w)-L(w)\rvert\geq\epsilon\right\\}\leq\frac{\beta^{2}}{m^{\prime}\epsilon^{2}}\mathbb{E}_{D}[\mathbb{W}^{2}(P_{Z|D=d},P_{Z})].$

When the metric $d$ is discrete, Wasserstein distance is equivalent to the
total variation. Therefore

$\displaystyle\mathbb{P}\left\\{\lvert
L_{t}(w)-L(w)\rvert\geq\epsilon\right\\}$
$\displaystyle\leq\frac{\beta^{2}}{m^{\prime}\epsilon^{2}}\mathbb{E}_{D}[\mathrm{TV}^{2}(P_{Z|D=d},P_{Z})]$
$\displaystyle\leq\frac{\beta^{2}}{m^{\prime}\epsilon^{2}}\mathbb{E}_{D}\left[\frac{1}{2}\mathrm{KL}(P_{Z|D=d}\,\|\,P_{Z})\right]$
$\displaystyle=\frac{\beta^{2}}{2m^{\prime}\epsilon^{2}}I(Z;D),$

where the second inequality is by applying Lemma .11. The proof is complete. ∎

######  Theorem .17.

If Assumption II.4 holds, and $\ell(f_{w}(X),f_{w^{\prime}}(X))$ is
$\beta$-Lipschitz for any $w,w^{\prime}\in\mathcal{W}$, then for any
$w\in\mathcal{W}$

$\mathbb{P}\left\\{L_{t}(w)-L(w)\geq\epsilon+L^{*}\right\\}\leq\frac{\beta^{2}}{m^{\prime}\epsilon^{2}}\mathbb{E}_{D}[\mathbb{W}^{2}(P_{X|D=d},P_{X})],$

where $L^{*}=\min_{w^{*}\in\mathcal{W}}(L_{t}(w^{*})+L(w^{*}))$. Furthermore,
if the metric $d$ is discrete, then

$\mathbb{P}\left\\{L_{t}(w)-L(w)\geq\epsilon+L^{*}\right\\}\leq\frac{\beta^{2}}{2m^{\prime}\epsilon^{2}}I(X;D).$

###### Proof.

Following the proof sketch of Theorem III.6, by setting $P=P_{X}$,
$Q=P_{X|D=d}$ and $f(x)=\ell(f_{w}(x),f_{w^{\prime}}(x))$ in Lemma .10, we
have

$\left\lparen
L_{d}(w,w^{\prime})-L(w,w^{\prime})\right\rparen^{2}\leq\beta^{2}\mathbb{W}^{2}(P_{X|D=d},P_{X}),$

for any $d\in\mathcal{D}$ and $w,w^{\prime}\in\mathcal{W}$. Similarly, by
applying the independence of $\\{D_{k}\\}_{k=1}^{m^{\prime}}$ and Chebyshev’s
inequality, we have

$\mathbb{P}\left\\{\lvert
L_{t}(w,w^{\prime})-L(w,w^{\prime})\rvert\geq\epsilon\right\\}\leq\frac{\beta^{2}}{m^{\prime}\epsilon^{2}}\mathbb{E}_{D}[\mathbb{W}^{2}(P_{X|D=d},P_{X})].$

Through a similar procedure of proving Theorem III.6 and Theorem .16, we can
get

$\displaystyle\mathbb{P}\left\\{L_{t}(w)-L(w)\geq\epsilon+\min_{w^{*}\in\mathcal{W}}(L_{t}(w^{*})+L(w^{*}))\right\\}$
$\displaystyle\leq\frac{\beta^{2}}{m^{\prime}\epsilon^{2}}\mathbb{E}_{D}[\mathbb{W}^{2}(P_{X|D=d},P_{X})]$
$\displaystyle\leq\frac{\beta^{2}}{2m^{\prime}\epsilon^{2}}I(X;D),$

which finishes the proof. ∎

The expected Wasserstein distances can be regarded as analogs to the mutual
information terms $I(Z;D)$ and $I(X;D)$ respectively, through a similar
reduction procedure as depicted in (2). This also provides alternative
perspectives to the distribution shift and the covariate shift in (1).

Moreover, these bounds can be further tightened by considering the shift in
the risks. Given a hypothesis $w\in\mathcal{W}$ and domain $d\in\mathcal{D}$,
let $L=\ell(f_{w}(X),Y)$ be the risk of predicting some random sample $Z\sim
P_{Z|D=d}$. Then by applying Lemma .9 with $P=P_{L}$, $Q=P_{L|D=d}$, $f(x)=x$
and Assumption II.2, we have

$\displaystyle\mathbb{E}_{D}[(L_{D}(w)-L(w))^{2}]$
$\displaystyle=\mathbb{E}_{D}[(\mathbb{E}_{L|D=d}[L]-\mathbb{E}_{L}[L])^{2}]$
$\displaystyle\leq\mathbb{E}_{D}\left[\left\lparen\sqrt{\frac{M^{2}}{2}\mathrm{KL}(P_{L|D=d}\,\|\,P_{L})}\right\rparen^{2}\right]$
$\displaystyle=\frac{M^{2}}{2}I(L;D).$

Through the similar sketch of proving Theorem III.5, we can prove that

$\mathbb{P}\\{\lvert
L_{t}(w)-L(w)\rvert\geq\epsilon\\}\leq\frac{M^{2}}{2\epsilon^{2}}I(L;D).$

By the Markov chain relationship
$D\rightarrow(X,Y)\rightarrow(f_{\phi}(X),Y)\rightarrow(f_{w}(X),Y)\rightarrow
L$, this upper bound is strictly tighter than Theorem III.5 which uses sample
space $I(Z;D)$ or representation space $I(R,Y;D)$ distribution shifts. Also,
notice that the mutual information $I(L;D)$ could be rewritten as
$\mathrm{KL}(P_{L|D}\,\|\,P_{L})$, this suggests that matching the inter-
domain distributions of the risks helps to generalize on test domains. This
observation facilitates the work of [15], which proposes to align the
empirical risks of distinct training domains. Still, the gap between the risk
shift $I(L;D)$ and its empirical counterpart $I(L_{i};D_{i})$ is not addressed
by [15], which further requires minimizing $I(W;D_{i})$. Considering that $L$
is a scalar while $R$ is a vector, aligning the distributions of the risks
avoids high-dimensional distribution matching, and thus enables more efficient
implementation. We will leave this for future research.

###  -O Generalization Bounds for Source-Domain Empirical Risk

The information-theoretic technique adopted in this paper to derive
generalization bounds is closely related to the recent advancements of
information-theoretic generalization analysis. Specifically, Theorem III.2 can
be viewed as a multi-domain version of the standard generalization error bound
in supervised learning [22, 41]. In this section, we further consider the
effect of finite training samples, which further raises a gap between domain-
level population and empirical risks. In addition to the generalization bounds
for population risks established in Section III, upper bounds for the
empirical risk can also be derived by incorporating Assumption II.2:

######  Theorem .18.

Let $W$ be the output of learning algorithm $\mathcal{A}$ with input $S$. If
Assumption II.2 holds, then

$\displaystyle\lvert\mathbb{E}_{W,D_{s},S}[L^{\prime}(W)]-\mathbb{E}_{W}[L(W)]\rvert$
$\displaystyle\leq\frac{1}{m}\sum_{i=1}^{m}\sqrt{\frac{M^{2}}{2}I(W;D_{i})}$
$\displaystyle\quad+\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\sqrt{\frac{M^{2}}{2}I(W;Z_{j}^{i}|D_{i})},$
$\displaystyle\mathbb{P}\left\\{\lvert\mathbb{E}_{W,D_{s},S}[L^{\prime}(W)]-\mathbb{E}_{W}[L(W)]\rvert\geq\epsilon\right\\}$
$\displaystyle\leq\frac{M^{2}}{mn\epsilon^{2}}\left\lparen I(W;S)+\log
3\right\rparen.$

###### Proof.

Recall that any random variables bounded by $[0,M]$ are
$\frac{M}{2}$-subgaussian. From assumption II.2, we know that
$\ell(f_{W}(X),Y)$ is $\frac{M}{2}$-subgaussian w.r.t $P_{W}\circ P_{Z}$. Then
by applying Lemma .9, we have

$\displaystyle\lvert\mathbb{E}_{W,D_{s},S}[L^{\prime}(W)]-\mathbb{E}_{W}[L(W)]\rvert$
$\displaystyle=\left\lvert\frac{1}{m}\sum_{i=1}^{m}\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}_{W,D_{i},Z_{j}^{i}}[\ell(f_{W}(X_{j}^{i}),Y_{j}^{i})]-\mathbb{E}_{W,D,Z}[\ell(f_{W}(X),Y)]\right\rvert$
$\displaystyle\leq\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\left\lvert\mathbb{E}_{W,D_{i},Z_{j}^{i}}[\ell(f_{W}(X_{j}^{i}),Y_{j}^{i})]-\mathbb{E}_{W,D,Z}[\ell(f_{W}(X),Y)]\right\rvert$
$\displaystyle\leq\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\sqrt{\frac{M^{2}}{2}\mathrm{KL}\left(P_{W,D_{i},Z_{j}^{i}}\,\middle\|\,P_{W}P_{D_{i},Z_{j}^{i}}\right)}.$

Notice that for any $D\in D_{s}$ and $Z\in S_{D}$,

$\displaystyle\mathrm{KL}\left(P_{W,D,Z}\,\middle\|\,P_{W}P_{D,Z}\right)$
$\displaystyle=\mathbb{E}_{W,D,Z}\left[\log\frac{P_{W,D,Z}}{P_{W}P_{D,Z}}\right]$
$\displaystyle=I(W;D,Z)$ $\displaystyle=I(W;D)+I(W;Z|D).$

Combining our results above, we then get

$\displaystyle\lvert\mathbb{E}_{W,D_{s},S}[L^{\prime}(W)]-\mathbb{E}_{W}[L(W)]\rvert$
$\displaystyle\leq\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\sqrt{\frac{M^{2}}{2}(I(W;D_{i})+I(W;Z_{j}^{i}|D_{i}))}$
$\displaystyle\leq\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\left\lparen\sqrt{\frac{M^{2}}{2}I(W;D_{i})}+\sqrt{\frac{M^{2}}{2}I(W;Z_{j}^{i}|D_{i})}\right\rparen$
$\displaystyle=\frac{1}{m}\sum_{i=1}^{m}\sqrt{\frac{M^{2}}{2}I(W;D_{i})}+\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\sqrt{\frac{M^{2}}{2}I(W;Z_{j}^{i}|D_{i})}.$

Similarly, we have

$\displaystyle\mathbb{E}_{W,S}\left[(L^{\prime}(W)-\mathbb{E}_{W}[L(W)])^{2}\right]$
$\displaystyle=\mathbb{E}_{W,S}\left[\left\lparen\frac{1}{m}\sum_{i=1}^{m}\frac{1}{n}\sum_{j=1}^{n}\ell(f_{W}(X_{j}^{i}),Y_{j}^{i})-\mathbb{E}_{W}[L(W)]\right\rparen^{2}\right]$
$\displaystyle=\frac{M^{2}}{mn}\left\lparen I(W;S)+\log 3\right\rparen.$

This further implies by Lemma .6 that

$\mathbb{P}\left\\{\lvert\mathbb{E}_{W,D_{s},S}[L^{\prime}(W)]-\mathbb{E}_{W}[L(W)]\rvert\geq\epsilon\right\\}\leq\frac{M^{2}}{mn\epsilon^{2}}\left\lparen
I(W;S)+\log 3\right\rparen,$

which completes the proof. ∎

######  Theorem .19.

Let $W$ be the output of learning algorithm $\mathcal{A}$ under training
domains $D_{s}$. If $\ell(f_{w}(X),Y)$ is $\beta^{\prime}$-Lipschitz w.r.t
$w$, then

$\displaystyle\lvert\mathbb{E}_{W,D_{s},S}[L^{\prime}(W)]-\mathbb{E}_{W}[L(W)]\rvert$
$\displaystyle\leq\frac{\beta^{\prime}}{m}\sum_{i=1}^{m}\mathbb{E}_{D_{i}}[\mathbb{W}(P_{W|D_{i}},P_{W})]$
$\displaystyle\quad+\frac{\beta^{\prime}}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\mathbb{E}_{D_{i},Z_{j}^{i}}[\mathbb{W}(P_{W|D_{i},Z_{j}^{i}},P_{W|D_{i}})].$

###### Proof.

Recall the proof of Theorem .18, we have

$\displaystyle\lvert\mathbb{E}_{W,D_{s},S}[L^{\prime}(W)]-\mathbb{E}_{W}[L(W)]\rvert$
$\displaystyle\leq\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\left\lvert\mathbb{E}_{W,D_{i},Z_{j}^{i}}[\ell(f_{W}(X_{j}^{i}),Y_{j}^{i})]-\mathbb{E}_{W,D,Z}[\ell(f_{W}(X),Y)]\right\rvert$
$\displaystyle\leq\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\left\lvert\mathbb{E}_{W,D_{i},Z_{j}^{i}}[\ell(f_{W}(X_{j}^{i}),Y_{j}^{i})]-\mathbb{E}_{W,D_{i},Z}[\ell(f_{W}(X),Y)]\right\rvert$
$\displaystyle\qquad+\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\left\lvert\mathbb{E}_{W,D_{i},Z}[\ell(f_{W}(X),Y)]-\mathbb{E}_{W,D,Z}[\ell(f_{W}(X),Y)]\right\rvert$
$\displaystyle\leq\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\mathbb{E}_{D_{i},Z_{j}^{i}}\left\lvert\mathbb{E}_{W|D_{i},Z_{j}^{i}}[\ell(f_{W}(X_{j}^{i}),Y_{j}^{i})]-\mathbb{E}_{W|D_{i}}[\ell(f_{W}(X_{j}^{i}),Y_{j}^{i})]\right\rvert$
$\displaystyle\qquad+\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\mathbb{E}_{D_{i}}\left\lvert\mathbb{E}_{W|D_{i}}[L_{D_{i}}(W)]-\mathbb{E}_{W}[L_{D_{i}}(W)]\right\rvert$
$\displaystyle\leq\frac{\beta^{\prime}}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\left\lparen\mathbb{E}_{D_{i},Z_{j}^{i}}[\mathbb{W}(P_{W|D_{i},Z_{j}^{i}},P_{W|D_{i}})]+\mathbb{E}_{D_{i}}[\mathbb{W}(P_{W|D_{i}},P_{W})]\right\rparen$
$\displaystyle=\frac{\beta^{\prime}}{m}\sum_{i=1}^{m}\mathbb{E}_{D_{i}}[\mathbb{W}(P_{W|D_{i}},P_{W})]+\frac{\beta^{\prime}}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}\mathbb{E}_{D_{i},Z_{j}^{i}}[\mathbb{W}(P_{W|D_{i},Z_{j}^{i}},P_{W|D_{i}})],$

where the last inequality is by applying Lemma .10. The proof is complete. ∎

The theorems above provide upper bounds for the empirical generalization risk
by exploiting the mutual information between the hypothesis and the samples
(or the Wasserstein distance counterparts). Compared to Theorems III.2 and
III.4, these results additionally consider the randomness of the sampled data
points $Z_{j}^{i}$, and indicate that traditional techniques that improve the
generalization of deep learning algorithms by minimizing $I(W;S)$, such as
gradient clipping [25, 34] and stochastic gradient perturbation [50, 51]
methods, also enhance the capability of learning algorithm $\mathcal{A}$ to
generalize on target domains under our high-probability problem setting by
preventing overfitting to training samples. This observation is also verified
in [34]. Relevant analysis may also motivate information-theoretic
generalization analysis for meta-learning tasks [52, 53, 54, 55]. We do not
consider these approaches in this paper, as they are beyond the scope of
solving Problem II.6.

## Experiment Details

In this paper, deep learning models are trained with an Intel Xeon CPU
(2.10GHz, 48 cores), 256GB memory, and 4 Nvidia Tesla V100 GPUs (32GB).

###  -P Implementation of IDM

We provide the pseudo-code for PDM in Algorithm 1, where the moving averages
$X_{ma}^{i}$ are initialized with $0$. The input data points for distribution
alignment are represented as matrices $X^{i}\in\mathbb{R}^{b\times d}$, where
$b$ denotes the batch size and $d$ represents the dimensionality. Each row of
$X$ then corresponds to an individual data point. We also present the pseudo-
code for IDM in Algorithm 2 for completeness.

Algorithm 1 PDM for distribution matching.

1: Input: Data matrices $\\{X^{i}\\}_{i=1}^{m}$, moving average $\gamma$.

2: Output: Penalty of distribution matching.

3: for $i$ from $1$ to $m$ do

4: Sort the elements of $X^{i}$ in each column in ascending order.

5: Calculate moving average $X_{ma}^{i}=\gamma X_{ma}^{i}+(1-\gamma)X^{i}$.

6: end for

7: Calculate the mean of data points across domains:
$X_{ma}=\frac{1}{m}\sum_{i=1}^{m}X_{ma}^{i}$.

8: Output: $\mathcal{L}_{\mathrm{PDM}}=\frac{1}{mdb}\sum_{i=1}^{m}\lVert
X_{ma}-X_{ma}^{i}\rVert_{F}^{2}$.

Algorithm 2 IDM for high-probability DG.

1: Input: Model $W$, training dataset $S$, hyper-parameters $\lambda_{1}$,
$\lambda_{2}$, $t_{1}$, $t_{2}$, $\gamma_{1}$, $\gamma_{2}$.

2: for $t$ from $1$ to #steps do

3: for $i$ from $1$ to $m$ do

4: Randomly sample a batch $B_{t}^{i}=(X_{t}^{i},Y_{t}^{i})$ from $S_{D_{i}}$
of size $b$.

5: Compute individual representations:
$(R_{t}^{i})_{j}=f_{\Phi}\left\lparen(X_{t}^{i})_{j}\right\rparen$, for
$j\in[1,b]$.

6: Compute individual risks: $(L_{t}^{i})_{j}=\ell\left\lparen
f_{\Psi}\left\lparen(R_{t}^{i})_{j}\right\rparen,(Y_{t}^{i})_{j}\right\rparen$,
for $j\in[1,b]$.

7: Compute individual gradients:
$(G_{t}^{i})_{j}=\nabla_{\Psi}(L_{t}^{i})_{j}$, for $j\in[1,b]$.

8: end for

9: Compute total empirical risk:
$\mathcal{L}_{\mathrm{IDM}}=\frac{1}{mn}\sum_{i=1}^{m}\sum_{j=1}^{n}(L_{t}^{i})_{j}$.

10: if $t\geq t_{1}$ then

11: Compute gradient alignment risk:
$\mathcal{L}_{\mathrm{G}}=\mathrm{PDM}(\\{G_{t}^{i}\\}_{i=1}^{m},\gamma_{1})$.

12:
$\mathcal{L}_{\mathrm{IDM}}=\mathcal{L}_{\mathrm{IDM}}+\lambda_{1}\mathcal{L}_{\mathrm{G}}$.

13: end if

14: if $t\geq t_{2}$ then

15: Compute representation alignment risk:
$\mathcal{L}_{\mathrm{R}}=\mathrm{PDM}(\\{R_{t}^{i}\\}_{i=1}^{m},\gamma_{2})$.

16:
$\mathcal{L}_{\mathrm{IDM}}=\mathcal{L}_{\mathrm{IDM}}+\lambda_{2}\mathcal{L}_{\mathrm{R}}$.

17: end if

18: Back-propagate gradients $\nabla_{W}\mathcal{L}_{\mathrm{IDM}}$ and update
the model $W$.

19: end for

We follow the experiment settings of [13] and utilize a moving average to
increase the equivalent number of data points for more accurate probability
density estimation in distribution alignments. This does not invalidate our
analysis in Theorem IV.1, as the maximum equivalent batch size
($b/(1-\gamma)\approx 32/(1-0.95)=640$) remains significantly smaller than the
dimensionality of the representation ($2048$ for ResNet-50 in DomainBed) or
the gradient ($2048\times c$, the number of classes) and satisfies $d>2b+1$.
Therefore, it is still impossible to distinguish different inter-domain
distributions as indicated by Theorem IV.1. However, this moving average
technique indeed helps to improve the empirical performance, as shown by our
ablation studies.

###  -Q Colored MNIST

The Colored MNIST dataset is a binary classification task introduced by IRM
[9]. The main difference between Colored MNIST and the original MNIST dataset
is the manually introduced strong correlation between the label and image
colors. Colored MNIST is generated according to the following procedure:

  * •

Give each sample an initial label by whether the digit is greater than $4$
(i.e. label $0$ for $0$-$4$ digits and label $1$ for $5$-$9$ digits.

  * •

Randomly flip the label with probability $0.25$, so an oracle predictor that
fully relies on the shape of the digits would achieve a $75\%$ accuracy.

  * •

Each environment is assigned a probability $P_{e}$, which characterizes the
correlation between the label and the color: samples with label $0$ have
$P_{e}$ chance to be red, and $1-P_{e}$ chance to be green, while samples with
label $1$ have $P_{e}$ chance to be green, and $1-P_{e}$ chance to be red.

The original environment setting of [9] includes two training domains
$D_{s}=\\{P_{1}=90\%,P_{2}=80\%\\}$, such that the predictive power of the
color superiors that of the actual digits. This correlation is reversed in the
test domain $D_{t}=\\{P_{3}=10\%\\}$, thus fooling algorithms without
causality inference abilities to overfit the color features and generalize
poorly on test environments.

The original
implementation222https://github.com/facebookresearch/InvariantRiskMinimization
uses a $3$-layer MLP network with ReLU activation. The model is trained for
$501$ epochs in a full gradient descent scheme, such that the batch size
equals the number of training samples $25,000$. We follow the hyper-parameter
selection strategy of [9] through a random search over $50$ independent
trials, as reported in Table III along with the parameters selected for IDM.
Considering that the covariate shift is not prominent according to the dataset
construction procedure, we only apply gradient alignment without feature
alignment in this experiment.

TABLE III: The hyper-parameters of Colored MNIST.

Parameter Random Distribution Selected Value dimension of hidden layer
$2^{\mathrm{Uniform}(6,9)}$ $433$ weight decay $10^{\mathrm{Uniform}(-2,-5)}$
$0.00034$ learning rate $10^{\mathrm{Uniform}(-2.5,-3.5)}$ $0.000449$ warmup
iterations $\mathrm{Uniform}(50,250)$ $154$ regularization strength
$10^{\mathrm{Uniform}(4,8)}$ $2888595.180638$

###  -R DomainBed Benchmark

DomainBed [21] is an extensive benchmark for both DA and DG algorithms, which
involves various synthetic and real-world datasets mainly focusing on image
classification:

  * •

Colored MNIST [9] is a variant of the MNIST dataset. As discussed previously,
Colored MNIST includes $3$ domains $\\{90\%,80\%,10\%\\}$, $70,000$ samples of
dimension $(2,28,28)$ and $2$ classes.

  * •

Rotated MNIST [56] is a variant of the MNIST dataset with $7$ domains
$\\{0,15,30,45,60,75\\}$ representing the rotation degrees, $70,000$ samples
of dimension $(28,28)$ and $10$ classes.

  * •

VLCS [57] includes $4$ domains
$\\{\mathrm{Caltech101},\mathrm{LabelMe},\mathrm{SUN09},\mathrm{VOC2007}\\}$,
$10,729$ samples of dimension $(3,224,224)$ and $5$ classes.

  * •

PACS [58] includes $4$ domains
$\\{\mathrm{art},\mathrm{cartoons},\mathrm{photos},\mathrm{sketches}\\}$,
$9,991$ samples of dimension $(3,224,224)$ and $7$ classes.

  * •

OfficeHome [59] includes $4$ domains
$\\{\mathrm{art},\mathrm{clipart},\mathrm{product},\mathrm{real}\\}$, $15,588$
samples of dimension $(3,224,224)$ and $65$ classes.

  * •

TerraIncognita [60] includes $4$ domains
$\\{\mathrm{L100},\mathrm{L38},\mathrm{L43},\mathrm{46}\\}$ representing
locations of photographs, $24,788$ samples of dimension $(3,224,224)$ and $10$
classes.

  * •

DomainNet [61] includes $6$ domains
$\\{\mathrm{clipart},\mathrm{infograph},\mathrm{painting},\mathrm{quickdraw},\mathrm{real},\mathrm{sketch}\\}$,
$586,575$ samples of dimension $(3,224,224)$ and $345$ classes.

We list all competitive DG approaches below. Note that some recent progress is
omitted [62, 16, 63, 64, 65, 66], which either contributes complementary
approaches, does not report full DomainBed results, or does not report the
test-domain validation scores. Due to the limitation of computational
resources, we are not able to reproduce the full results of these works on
DomainBed.

  * •

ERM: Empirical Risk Minimization.

  * •

IRM: Invariant Risk Minimization [9].

  * •

GroupDRO: Group Distributionally Robust Optimization [14].

  * •

Mixup: Interdomain Mixup [67].

  * •

MLDG: Meta Learning Domain Generalization [68].

  * •

CORAL: Deep CORAL [5].

  * •

MMD: Maximum Mean Discrepancy [6].

  * •

DANN: Domain Adversarial Neural Network [7].

  * •

CDANN: Conditional Domain Adversarial Neural Network [8].

  * •

MTL: Marginal Transfer Learning [17].

  * •

SagNet: Style Agnostic Networks [69].

  * •

ARM: Adaptive Risk Minimization [18].

  * •

V-REx: Variance Risk Extrapolation [15].

  * •

RSC: Representation Self-Challenging [70].

  * •

AND-mask: Learning Explanations that are Hard to Vary [20].

  * •

SAND-mask: Smoothed-AND mask [44].

  * •

Fish: Gradient Matching for Domain Generalization [12].

  * •

Fishr: Invariant Gradient Variances for Out-of-distribution Generalization
[13].

  * •

SelfReg: Self-supervised Contrastive Regularization [71].

  * •

CausIRL: Invariant Causal Mechanisms through Distribution Matching [10].

The same fine-tuning procedure is applied to all approaches: The network is a
multi-layer CNN for synthetic MNIST datasets and is a pre-trained ResNet-50
for other real-world datasets. The hyper-parameters are selected by a random
search over $20$ independent trials for each target domain, and each
evaluation score is reported after $3$ runs with different initialization
seeds333https://github.com/facebookresearch/DomainBed. The hyper-parameter
selection criteria are shown in Table IV. Note that warmup iterations and
moving average techniques are not adopted for representation alignment.

TABLE IV: The hyper-parameters of DomainBed.

Condition Parameter Default Value Random Distribution MNIST Datasets learning
rate $0.001$ $10^{\mathrm{Uniform}(-4.5,-3.5)}$ batch size $64$
$2^{\mathrm{Uniform}(3,9)}$ Real-world Datasets learning rate $0.00005$
$10^{\mathrm{Uniform}(-5,-3.5)}$ batch size $32$ $2^{\mathrm{Uniform}(3,5)}$
(DomainNet) / $2^{\mathrm{Uniform}(3,5.5)}$ (others) weight decay $0$
$10^{\mathrm{Uniform}(-6,-2)}$ dropout $0$ $\mathrm{Uniform}(\\{0,0.1,0.5\\})$
- steps $5000$ $5000$ IDM gradient penalty $1000$ $10^{\mathrm{Uniform}(1,5)}$
gradient warmup $1500$ $\mathrm{Uniform}(0,5000)$ representation penalty $1$
$10^{\mathrm{Uniform}(-1,1)}$ moving average $0.95$
$\mathrm{Uniform}(0.9,0.99)$

Note that although the same Colored MNIST dataset is adopted by DomainBed, the
experimental settings are completely different from the previous one [9]. The
main difference is the batch size ($25000$ for IRM, less than $512$ for
DomainBed), making it much harder to learn invariance for causality inference
and distribution matching methods since fewer samples are available for
probability density estimation. This explains the huge performance drop
between these two experiments using the same DG algorithms.

## Additional Experimental Results

###  -S Component Analysis

In this section, we conduct ablation studies to demonstrate the effect of each
component of the proposed IDM algorithm. Specifically, we analyze the effect
of gradient alignment (GA), representation alignment (RA), warmup iterations
(WU), moving average (MA), and the proposed PDM method for distribution
matching.

TABLE V: Component Analysis on ColoredMNIST of DomainBed.

Algorithm GA RA WU MA 90% 80% 10% Average ERM - 71.8 $\pm$ 0.4 72.9 $\pm$ 0.1
28.7 $\pm$ 0.5 57.8 IDM ✗ ✓ ✗ ✗ 71.9 $\pm$ 0.4 72.5 $\pm$ 0.0 28.8 $\pm$ 0.7
57.7 ✓ ✗ ✓ ✓ 73.1 $\pm$ 0.2 72.7 $\pm$ 0.3 67.4 $\pm$ 1.6 71.1 ✓ ✓ ✗ ✓ 72.9
$\pm$ 0.2 72.7 $\pm$ 0.1 60.8 $\pm$ 2.1 68.8 ✓ ✓ ✓ ✗ 72.0 $\pm$ 0.1 71.5 $\pm$
0.3 48.7 $\pm$ 7.1 64.0 ✓ ✓ ✓ ✓ 74.2 $\pm$ 0.6 73.5 $\pm$ 0.2 68.3 $\pm$ 2.5
72.0

TABLE VI: Component Analysis on OfficeHome of DomainBed.

Algorithm GA RA WU MA A C P R Average ERM - 61.7 $\pm$ 0.7 53.4 $\pm$ 0.3 74.1
$\pm$ 0.4 76.2 $\pm$ 0.6 66.4 IDM ✗ ✓ ✗ ✗ 64.7 $\pm$ 0.5 54.6 $\pm$ 0.3 76.2
$\pm$ 0.4 78.1 $\pm$ 0.5 68.4 ✓ ✗ ✓ ✓ 61.9 $\pm$ 0.4 53.0 $\pm$ 0.3 75.5 $\pm$
0.2 77.9 $\pm$ 0.2 67.1 ✓ ✓ ✗ ✓ 62.5 $\pm$ 0.1 53.0 $\pm$ 0.7 75.0 $\pm$ 0.4
77.2 $\pm$ 0.7 66.9 ✓ ✓ ✓ ✗ 64.2 $\pm$ 0.3 53.5 $\pm$ 0.6 76.1 $\pm$ 0.4 78.1
$\pm$ 0.4 68.0 ✓ ✓ ✓ ✓ 64.4 $\pm$ 0.3 54.4 $\pm$ 0.6 76.5 $\pm$ 0.3 78.0 $\pm$
0.4 68.3

####  -S1 Gradient Alignment

According to our theoretical analysis, gradient alignment promotes training-
domain generalization, especially when concept shift is prominent. As can be
seen in Table V, IDM without gradient alignment (57.7%) performs similarly to
ERM (57.8%), which is unable to learn invariance across training domains.
Gradient alignment also significantly boosts the performance on VLCS (77.4% to
78.1%) and PACS (86.8% to 87.6%), as seen in Table VII and VIII. However, for
datasets where concept shift is not prominent e.g. OfficeHome, gradient
alignment cannot help to improve performance as shown in Table VI. It is worth
noting that gradient alignment also penalizes a lower bound for the
representation space distribution shift: In the $t$-th step of gradient
descent, the Markov chain relationship $D_{i}\rightarrow
B_{t}^{i}\rightarrow(R_{t}^{i},Y_{t}^{i})\rightarrow G_{t}^{i}$ holds
conditioned on the current predictor $W_{t-1}$, which implies the lower bound
$I(G_{t}^{i};D_{i}|W_{t-1})\leq I(R_{t}^{i},Y_{t}^{i};D_{i}|W_{t-1})$ by the
data processing inequality. This indicates that gradient alignment also helps
to address the covariate shift, which explains the promising performance of
gradient-based DG algorithms e.g. Fish and Fishr. However, since this is a
lower bound rather than an upper bound, gradient manipulation is insufficient
to fully address representation space covariate shifts, as seen in the
following analysis for representation alignment.

TABLE VII: Effect of gradient alignment (GA) on VLCS of DomainBed.

Algorithm GA A C P S Average ERM - 97.6 $\pm$ 0.3 67.9 $\pm$ 0.7 70.9 $\pm$
0.2 74.0 $\pm$ 0.6 77.6 IDM ✗ 97.1 $\pm$ 0.7 67.2 $\pm$ 0.4 69.9 $\pm$ 0.4
75.6 $\pm$ 0.8 77.4 IDM ✓ 97.6 $\pm$ 0.3 66.9 $\pm$ 0.3 71.8 $\pm$ 0.5 76.0
$\pm$ 1.3 78.1

TABLE VIII: Effect of gradient alignment (GA) on PACS of DomainBed.

Algorithm GA A C P S Average ERM - 86.5 $\pm$ 1.0 81.3 $\pm$ 0.6 96.2 $\pm$
0.3 82.7 $\pm$ 1.1 86.7 IDM ✗ 87.8 $\pm$ 0.6 81.6 $\pm$ 0.3 97.4 $\pm$ 0.2
80.6 $\pm$ 1.3 86.8 IDM ✓ 88.0 $\pm$ 0.3 82.6 $\pm$ 0.6 97.6 $\pm$ 0.4 82.3
$\pm$ 0.6 87.6

####  -S2 Representation Alignment

Representation alignment promotes test-domain generalization by minimizing the
representation level covariate shift. As shown in Table V \- IX,
representation alignment is effective in OfficeHome (67.1% to 68.3%) and
RotatedMNIST (97.8% to 98.0%), and still enhances the performance even though
covariate shift is not prominent in ColoredMNIST (71.1% to 72.0%). This
verifies our claim that representation alignment complements gradient
alignment in solving Problem II.6, and is necessary for achieving high-
probability DG.

TABLE IX: Effect of representation alignment (RA) on RotatedMNIST of
DomainBed.

Algorithm RA 0 15 30 45 60 75 Average ERM - 95.3 $\pm$ 0.2 98.7 $\pm$ 0.1 98.9
$\pm$ 0.1 98.7 $\pm$ 0.2 98.9 $\pm$ 0.0 96.2 $\pm$ 0.2 97.8 IDM ✗ 95.6 $\pm$
0.1 98.4 $\pm$ 0.1 98.7 $\pm$ 0.2 99.1 $\pm$ 0.0 98.7 $\pm$ 0.1 96.6 $\pm$ 0.4
97.8 IDM ✓ 96.1 $\pm$ 0.3 98.7 $\pm$ 0.1 99.1 $\pm$ 0.1 98.9 $\pm$ 0.1 98.9
$\pm$ 0.1 96.6 $\pm$ 0.1 98.0

####  -S3 Warmup Iterations

Following the experimental settings of [9, 13], we do not apply the penalties
of gradient or representation alignment until the number of epochs reaches a
certain value. This is inspired by the observation that forcing invariance in
early steps may hinder the models from extracting useful correlations. By
incorporating these warmup iterations, predictors are allowed to extract all
possible correlations between the inputs and the labels at the beginning, and
then discard spurious ones in later updates. As can be seen in Table V and VI,
this strategy helps to enhance the final performances on ColoredMNIST (68.8%
to 72.0%) and OfficeHome (66.9% to 68.3%).

####  -S4 Moving Average

Following [13, 72], we use an exponential moving average when computing the
gradients or the representations. This strategy helps when the batch size is
not sufficiently large to sketch the probability distributions. In the IRM
experiment setup where the batch size is $25000$, Fishr (70.2%) and IDM
(70.5%) both easily achieve near-optimal accuracy compared to Oracle (71.0%).
In the DomainBed setup, the batch size $2^{\mathrm{Uniform}(3,9)}$ is
significantly diminished, resulting in worse test-domain accuracy of Fishr
(68.8%). As shown in Table V and VI, this moving average strategy greatly
enhances the performance of IDM on ColoredMNIST (64.0% to 72.0%) and
OfficeHome (68.0% to 68.3%).

####  -S5 PDM for Distribution Matching

We then demonstrate the superiority of our PDM method over moment-based
distribution alignment techniques. Specifically, we compare IGA [11] which
matches the empirical expectation of the gradients, Fishr [13] which proposes
to align the gradient variance, the combination of IGA + Fishr (i.e. aligning
the expectation and variance simultaneously), and our approach IDM (without
representation space alignment). The performance gain of IDM on the Colored
MNIST task in [9] is not significant, since it is relatively easier to learn
invariance with a large batch size ($25000$). In the DomainBed setting, the
batch size is significantly reduced ($8$-$512$), making this learning task
much harder. The results are reported in Table X.

TABLE X: Superiority of PDM on Colored MNIST of DomainBed.

Algorithm 90% 80% 10% Average ERM 71.8 $\pm$ 0.4 72.9 $\pm$ 0.1 28.7 $\pm$ 0.5
57.8 IGA 72.6 $\pm$ 0.3 72.9 $\pm$ 0.2 50.0 $\pm$ 1.2 65.2 Fishr 74.1 $\pm$
0.6 73.3 $\pm$ 0.1 58.9 $\pm$ 3.7 68.8 IGA + Fishr 73.3 $\pm$ 0.0 72.6 $\pm$
0.5 66.3 $\pm$ 2.9 70.7 IDM 74.2 $\pm$ 0.6 73.5 $\pm$ 0.2 68.3 $\pm$ 2.5 72.0

As can be seen, IDM achieves significantly higher performance on Colored MNIST
(72.0%) even compared to the combination of IGA + Fishr (70.7%). This verifies
our conclusion that matching the expectation and the variance is not
sufficient for complex probability distributions, and demonstrates the
superiority of the proposed PDM method for distribution alignment.

###  -T Running Time Comparison

Since IDM only stores historical gradients and representations for a single
batch from each training domain, the storage and computation overhead is
marginal compared to training the entire network. As shown in Table XI, the
training time is only 5% longer compared to ERM on the largest DomainNet
dataset.

TABLE XI: Computational overhead of IDM using default batch size. Dataset | Training Time (h) | Memory Requirement (GB)
---|---|---
ERM | IDM | Overhead | ERM | IDM | Overhead
ColoredMNIST | 0.076 | 0.088 | 14.6% | 0.138 | 0.139 | 0.2%
RotatedMNIST | 0.101 | 0.110 | 9.3% | 0.338 | 0.342 | 1.0%
VLCS | 0.730 | 0.744 | 2.0% | 8.189 | 8.199 | 0.1%
PACS | 0.584 | 0.593 | 1.5% | 8.189 | 8.201 | 0.1%
OfficeHome | 0.690 | 0.710 | 2.9% | 8.191 | 8.506 | 3.8%
TerraIncognita | 0.829 | 0.840 | 1.3% | 8.189 | 8.208 | 0.2%
DomainNet | 2.805 | 2.947 | 5.0% | 13.406 | 16.497 | 23.1%

###  -U Full DomainBed Results

In this paper, we focus on the test-domain model selection criterion, where
the validation set follows the same distribution as the test domains. Our
choice is well-motivated for the following reasons:

  * •

Test-domain validation is provided by the DomainBed benchmark as one of the
default model-selection methods, and is also widely adopted in the literature
in many significant works like IRM [9], V-Rex [15], and Fishr [13].

  * •

As suggested by Proposition III.1, any algorithm that fits well on training
domains will suffer from strictly positive risks in test domains once concept
shift is induced. Therefore, training-domain validation would result in sub-
optimal selection results.

  * •

Training-domain validation may render efforts to address concept shift
useless, as spurious features are often more predictive than invariant ones.
This is particularly unfair for algorithms that aim to tackle the concept
shift. As shown in Table 9 in [13], no algorithm can significantly outperform
ERM on Colored MNIST using training-domain validation (an exception is ARM
which uses test-time adaptation, and thus cannot be directly compared), even
though ERM is shown to perform much worse than random guessing (10% v.s. 50%
accuracy) for the last domain (see Table 1 in [9] and Appendix D.4.1 in [13]).
As a result, models selected by training-domain validation may not generalize
well when concept shift is substantial.

  * •

As mentioned by [73], training-domain validation suffers from
underspecification, where predictors with equivalently strong performances in
training domains may behave very differently during testing. It is also
emphasized by [74] that OOD performance cannot, by definition, be performed
with a validation set from the same distribution as the training data. This
further raises concerns about the validity of using training-domain accuracies
for validation purposes.

  * •

Moreover, test-domain validation is also applicable in practice, as it is
feasible to label a few test-domain samples for validation purposes. It is
also unrealistic to deploy models in target environments without any form of
verification, making such efforts necessary in practice.

Due to the reasons listed above, we believe that the test-domain validation
results are sufficient to demonstrate the effectiveness of our approach in
real-world learning scenarios. We report detailed results of IDM for each
domain in each dataset of the DomainBed benchmark under test-domain model
selection for a complete evaluation in Table XII \- XVIII. Note that detailed
scores of certain algorithms (Fish, CausIRL) are not available.

TABLE XII: Detailed results on Colored MNIST in DomainBed.

Algorithm 90% 80% 10% Average ERM 71.8 $\pm$ 0.4 72.9 $\pm$ 0.1 28.7 $\pm$ 0.5
57.8 IRM 72.0 $\pm$ 0.1 72.5 $\pm$ 0.3 58.5 $\pm$ 3.3 67.7 GroupDRO 73.5 $\pm$
0.3 73.0 $\pm$ 0.3 36.8 $\pm$ 2.8 61.1 Mixup 72.5 $\pm$ 0.2 73.9 $\pm$ 0.4
28.6 $\pm$ 0.2 58.4 MLDG 71.9 $\pm$ 0.3 73.5 $\pm$ 0.2 29.1 $\pm$ 0.9 58.2
CORAL 71.1 $\pm$ 0.2 73.4 $\pm$ 0.2 31.1 $\pm$ 1.6 58.6 MMD 69.0 $\pm$ 2.3
70.4 $\pm$ 1.6 50.6 $\pm$ 0.2 63.3 DANN 72.4 $\pm$ 0.5 73.9 $\pm$ 0.5 24.9
$\pm$ 2.7 57.0 CDANN 71.8 $\pm$ 0.5 72.9 $\pm$ 0.1 33.8 $\pm$ 6.4 59.5 MTL
71.2 $\pm$ 0.2 73.5 $\pm$ 0.2 28.0 $\pm$ 0.6 57.6 SagNet 72.1 $\pm$ 0.3 73.2
$\pm$ 0.3 29.4 $\pm$ 0.5 58.2 ARM 84.9 $\pm$ 0.9 76.8 $\pm$ 0.6 27.9 $\pm$ 2.1
63.2 V-REx 72.8 $\pm$ 0.3 73.0 $\pm$ 0.3 55.2 $\pm$ 4.0 67.0 RSC 72.0 $\pm$
0.1 73.2 $\pm$ 0.1 30.2 $\pm$ 1.6 58.5 AND-mask 71.9 $\pm$ 0.6 73.6 $\pm$ 0.5
30.2 $\pm$ 1.4 58.6 SAND-mask 79.9 $\pm$ 3.8 75.9 $\pm$ 1.6 31.6 $\pm$ 1.1
62.3 Fishr 74.1 $\pm$ 0.6 73.3 $\pm$ 0.1 58.9 $\pm$ 3.7 68.8 SelfReg 71.3
$\pm$ 0.4 73.4 $\pm$ 0.2 29.3 $\pm$ 2.1 58.0 IDM 74.2 $\pm$ 0.6 73.5 $\pm$ 0.2
68.3 $\pm$ 2.5 72.0

TABLE XIII: Detailed results on Rotated MNIST in DomainBed.

Algorithm 0 15 30 45 60 75 Average ERM 95.3 $\pm$ 0.2 98.7 $\pm$ 0.1 98.9
$\pm$ 0.1 98.7 $\pm$ 0.2 98.9 $\pm$ 0.0 96.2 $\pm$ 0.2 97.8 IRM 94.9 $\pm$ 0.6
98.7 $\pm$ 0.2 98.6 $\pm$ 0.1 98.6 $\pm$ 0.2 98.7 $\pm$ 0.1 95.2 $\pm$ 0.3
97.5 GroupDRO 95.9 $\pm$ 0.1 99.0 $\pm$ 0.1 98.9 $\pm$ 0.1 98.8 $\pm$ 0.1 98.6
$\pm$ 0.1 96.3 $\pm$ 0.4 97.9 Mixup 95.8 $\pm$ 0.3 98.7 $\pm$ 0.0 99.0 $\pm$
0.1 98.8 $\pm$ 0.1 98.8 $\pm$ 0.1 96.6 $\pm$ 0.2 98.0 MLDG 95.7 $\pm$ 0.2 98.9
$\pm$ 0.1 98.8 $\pm$ 0.1 98.9 $\pm$ 0.1 98.6 $\pm$ 0.1 95.8 $\pm$ 0.4 97.8
CORAL 96.2 $\pm$ 0.2 98.8 $\pm$ 0.1 98.8 $\pm$ 0.1 98.8 $\pm$ 0.1 98.9 $\pm$
0.1 96.4 $\pm$ 0.2 98.0 MMD 96.1 $\pm$ 0.2 98.9 $\pm$ 0.0 99.0 $\pm$ 0.0 98.8
$\pm$ 0.0 98.9 $\pm$ 0.0 96.4 $\pm$ 0.2 98.0 DANN 95.9 $\pm$ 0.1 98.9 $\pm$
0.1 98.6 $\pm$ 0.2 98.7 $\pm$ 0.1 98.9 $\pm$ 0.0 96.3 $\pm$ 0.3 97.9 CDANN
95.9 $\pm$ 0.2 98.8 $\pm$ 0.0 98.7 $\pm$ 0.1 98.9 $\pm$ 0.1 98.8 $\pm$ 0.1
96.1 $\pm$ 0.3 97.9 MTL 96.1 $\pm$ 0.2 98.9 $\pm$ 0.0 99.0 $\pm$ 0.0 98.7
$\pm$ 0.1 99.0 $\pm$ 0.0 95.8 $\pm$ 0.3 97.9 SagNet 95.9 $\pm$ 0.1 99.0 $\pm$
0.1 98.9 $\pm$ 0.1 98.6 $\pm$ 0.1 98.8 $\pm$ 0.1 96.3 $\pm$ 0.1 97.9 ARM 95.9
$\pm$ 0.4 99.0 $\pm$ 0.1 98.8 $\pm$ 0.1 98.9 $\pm$ 0.1 99.1 $\pm$ 0.1 96.7
$\pm$ 0.2 98.1 V-REx 95.5 $\pm$ 0.2 99.0 $\pm$ 0.0 98.7 $\pm$ 0.2 98.8 $\pm$
0.1 98.8 $\pm$ 0.0 96.4 $\pm$ 0.0 97.9 RSC 95.4 $\pm$ 0.1 98.6 $\pm$ 0.1 98.6
$\pm$ 0.1 98.9 $\pm$ 0.0 98.8 $\pm$ 0.1 95.4 $\pm$ 0.3 97.6 AND-mask 94.9
$\pm$ 0.1 98.8 $\pm$ 0.1 98.8 $\pm$ 0.1 98.7 $\pm$ 0.2 98.6 $\pm$ 0.2 95.5
$\pm$ 0.2 97.5 SAND-mask 94.7 $\pm$ 0.2 98.5 $\pm$ 0.2 98.6 $\pm$ 0.1 98.6
$\pm$ 0.1 98.5 $\pm$ 0.1 95.2 $\pm$ 0.1 97.4 Fishr 95.8 $\pm$ 0.1 98.3 $\pm$
0.1 98.8 $\pm$ 0.1 98.6 $\pm$ 0.3 98.7 $\pm$ 0.1 96.5 $\pm$ 0.1 97.8 SelfReg
96.0 $\pm$ 0.3 98.9 $\pm$ 0.1 98.9 $\pm$ 0.1 98.9 $\pm$ 0.1 98.9 $\pm$ 0.1
96.8 $\pm$ 0.1 98.1 IDM 96.1 $\pm$ 0.3 98.7 $\pm$ 0.1 99.1 $\pm$ 0.1 98.9
$\pm$ 0.1 98.9 $\pm$ 0.1 96.6 $\pm$ 0.1 98.0

TABLE XIV: Detailed results on VLCS in DomainBed.

Algorithm C L S V Average ERM 97.6 $\pm$ 0.3 67.9 $\pm$ 0.7 70.9 $\pm$ 0.2
74.0 $\pm$ 0.6 77.6 IRM 97.3 $\pm$ 0.2 66.7 $\pm$ 0.1 71.0 $\pm$ 2.3 72.8
$\pm$ 0.4 76.9 GroupDRO 97.7 $\pm$ 0.2 65.9 $\pm$ 0.2 72.8 $\pm$ 0.8 73.4
$\pm$ 1.3 77.4 Mixup 97.8 $\pm$ 0.4 67.2 $\pm$ 0.4 71.5 $\pm$ 0.2 75.7 $\pm$
0.6 78.1 MLDG 97.1 $\pm$ 0.5 66.6 $\pm$ 0.5 71.5 $\pm$ 0.1 75.0 $\pm$ 0.9 77.5
CORAL 97.3 $\pm$ 0.2 67.5 $\pm$ 0.6 71.6 $\pm$ 0.6 74.5 $\pm$ 0.0 77.7 MMD
98.8 $\pm$ 0.0 66.4 $\pm$ 0.4 70.8 $\pm$ 0.5 75.6 $\pm$ 0.4 77.9 DANN 99.0
$\pm$ 0.2 66.3 $\pm$ 1.2 73.4 $\pm$ 1.4 80.1 $\pm$ 0.5 79.7 CDANN 98.2 $\pm$
0.1 68.8 $\pm$ 0.5 74.3 $\pm$ 0.6 78.1 $\pm$ 0.5 79.9 MTL 97.9 $\pm$ 0.7 66.1
$\pm$ 0.7 72.0 $\pm$ 0.4 74.9 $\pm$ 1.1 77.7 SagNet 97.4 $\pm$ 0.3 66.4 $\pm$
0.4 71.6 $\pm$ 0.1 75.0 $\pm$ 0.8 77.6 ARM 97.6 $\pm$ 0.6 66.5 $\pm$ 0.3 72.7
$\pm$ 0.6 74.4 $\pm$ 0.7 77.8 V-REx 98.4 $\pm$ 0.2 66.4 $\pm$ 0.7 72.8 $\pm$
0.1 75.0 $\pm$ 1.4 78.1 RSC 98.0 $\pm$ 0.4 67.2 $\pm$ 0.3 70.3 $\pm$ 1.3 75.6
$\pm$ 0.4 77.8 AND-mask 98.3 $\pm$ 0.3 64.5 $\pm$ 0.2 69.3 $\pm$ 1.3 73.4
$\pm$ 1.3 76.4 SAND-mask 97.6 $\pm$ 0.3 64.5 $\pm$ 0.6 69.7 $\pm$ 0.6 73.0
$\pm$ 1.2 76.2 Fishr 97.6 $\pm$ 0.7 67.3 $\pm$ 0.5 72.2 $\pm$ 0.9 75.7 $\pm$
0.3 78.2 SelfReg 97.9 $\pm$ 0.4 66.7 $\pm$ 0.1 73.5 $\pm$ 0.7 74.7 $\pm$ 0.7
78.2 IDM 97.6 $\pm$ 0.3 66.9 $\pm$ 0.3 71.8 $\pm$ 0.5 76.0 $\pm$ 1.3 78.1

TABLE XV: Detailed results on PACS in DomainBed.

Algorithm A C P S Average ERM 86.5 $\pm$ 1.0 81.3 $\pm$ 0.6 96.2 $\pm$ 0.3
82.7 $\pm$ 1.1 86.7 IRM 84.2 $\pm$ 0.9 79.7 $\pm$ 1.5 95.9 $\pm$ 0.4 78.3
$\pm$ 2.1 84.5 GroupDRO 87.5 $\pm$ 0.5 82.9 $\pm$ 0.6 97.1 $\pm$ 0.3 81.1
$\pm$ 1.2 87.1 Mixup 87.5 $\pm$ 0.4 81.6 $\pm$ 0.7 97.4 $\pm$ 0.2 80.8 $\pm$
0.9 86.8 MLDG 87.0 $\pm$ 1.2 82.5 $\pm$ 0.9 96.7 $\pm$ 0.3 81.2 $\pm$ 0.6 86.8
CORAL 86.6 $\pm$ 0.8 81.8 $\pm$ 0.9 97.1 $\pm$ 0.5 82.7 $\pm$ 0.6 87.1 MMD
88.1 $\pm$ 0.8 82.6 $\pm$ 0.7 97.1 $\pm$ 0.5 81.2 $\pm$ 1.2 87.2 DANN 87.0
$\pm$ 0.4 80.3 $\pm$ 0.6 96.8 $\pm$ 0.3 76.9 $\pm$ 1.1 85.2 CDANN 87.7 $\pm$
0.6 80.7 $\pm$ 1.2 97.3 $\pm$ 0.4 77.6 $\pm$ 1.5 85.8 MTL 87.0 $\pm$ 0.2 82.7
$\pm$ 0.8 96.5 $\pm$ 0.7 80.5 $\pm$ 0.8 86.7 SagNet 87.4 $\pm$ 0.5 81.2 $\pm$
1.2 96.3 $\pm$ 0.8 80.7 $\pm$ 1.1 86.4 ARM 85.0 $\pm$ 1.2 81.4 $\pm$ 0.2 95.9
$\pm$ 0.3 80.9 $\pm$ 0.5 85.8 V-REx 87.8 $\pm$ 1.2 81.8 $\pm$ 0.7 97.4 $\pm$
0.2 82.1 $\pm$ 0.7 87.2 RSC 86.0 $\pm$ 0.7 81.8 $\pm$ 0.9 96.8 $\pm$ 0.7 80.4
$\pm$ 0.5 86.2 AND-mask 86.4 $\pm$ 1.1 80.8 $\pm$ 0.9 97.1 $\pm$ 0.2 81.3
$\pm$ 1.1 86.4 SAND-mask 86.1 $\pm$ 0.6 80.3 $\pm$ 1.0 97.1 $\pm$ 0.3 80.0
$\pm$ 1.3 85.9 Fishr 87.9 $\pm$ 0.6 80.8 $\pm$ 0.5 97.9 $\pm$ 0.4 81.1 $\pm$
0.8 86.9 SelfReg 87.5 $\pm$ 0.1 83.0 $\pm$ 0.1 97.6 $\pm$ 0.1 82.8 $\pm$ 0.2
87.7 IDM 88.0 $\pm$ 0.3 82.6 $\pm$ 0.6 97.6 $\pm$ 0.4 82.3 $\pm$ 0.6 87.6

TABLE XVI: Detailed results on OfficeHome in DomainBed.

Algorithm A C P R Average ERM 61.7 $\pm$ 0.7 53.4 $\pm$ 0.3 74.1 $\pm$ 0.4
76.2 $\pm$ 0.6 66.4 IRM 56.4 $\pm$ 3.2 51.2 $\pm$ 2.3 71.7 $\pm$ 2.7 72.7
$\pm$ 2.7 63.0 GroupDRO 60.5 $\pm$ 1.6 53.1 $\pm$ 0.3 75.5 $\pm$ 0.3 75.9
$\pm$ 0.7 66.2 Mixup 63.5 $\pm$ 0.2 54.6 $\pm$ 0.4 76.0 $\pm$ 0.3 78.0 $\pm$
0.7 68.0 MLDG 60.5 $\pm$ 0.7 54.2 $\pm$ 0.5 75.0 $\pm$ 0.2 76.7 $\pm$ 0.5 66.6
CORAL 64.8 $\pm$ 0.8 54.1 $\pm$ 0.9 76.5 $\pm$ 0.4 78.2 $\pm$ 0.4 68.4 MMD