Optimal Representations for Covariate Shifts
Abstract
Machine learning often experiences distribution shifts between training and testing.
We introduce a simple objective whose optima are exactly all representations on which risk minimizers are guaranteed to be robust to Bayes-preserving shifts, e.g., covariate shifts.
Our objective has two components.
First, a representation must remain discriminative, i.e., some predictor must be able to minimize the source and target risk.
Second, the representation’s support should be invariant across source and target.
We make this practical by designing self-supervised methods that only use unlabelled data and augmentations.
Our objectives achieve SOTA on
DomainBed, and give insights into the robustness of recent methods, e.g., CLIP.
Introduction
It is hard to build machine learning (ML) systems that are robust to distribution shifts between a source (train) and target (test) domain.
One promising approach to DG is to learn representations from which predictors trained on source must perform well on target.
In practice, however, no current
DG method uniformly outperforms empirical source-risk minimizers (ERM) [22].
Our theoretical understanding is also lacking: while previous work has studied properties that are or are not sufﬁcient for robust representations [7, 58, 27], the minimal set of requirements is not yet known.
We introduce the ﬁrst simple, objective whose optima are exactly the set of all representations on which source risk minimizers are guaranteed to generalize across distribution shifts that preserve the
Bayes predictor.
Our characterization implies that it is sufﬁcient and necessary that an optimal representation: (a) remains discriminative, i e., there must exist predictors that simultaneously minimize both source and target risk; and (b) keeps the support of its marginal distribution shift-invariant.
Optimal representations must thus seek discriminative information of targets.
Even worse, we prove that without target knowledge, no representation can uniformly outperform constant representations, which may explain why DG methods struggle to outperform ERM.
We show how to overcome these challenges using only a large set of unlabeled examples and particular data augmentations that retain discriminative information but minimal domain-speciﬁc information.
Text descriptions of images are such augmentations, as they are informative for many classiﬁcation tasks, but remove domain-speciﬁc information.
With those augmentations, we design practical self-supervised (SSL) objectives for learning robust representations.
Our objectives give insights into CLIP’s robustness [43], and lead to improved CLIP-based representations that achieve SOTA on DomainBed [22].
Problem statement: idealized domain generalization (IDG)
We want to learn representations Z of inputs X that are robust across distribution shifts.
Speciﬁcally, we want an encoder p Z | X that ensures that predictors h from representations Z to labels Y , which are trained on a source D s ∈ D distribution p X,Y | D s , will perform well on a target D t ∈ D
[Table]
+KqnqRZFQbz/t05uYXFpeWcyv51bX1jc3C1nZNJ6nCpIoTlqhGhDRhVJCqoYaRhlQE8YiRetS7GPr1O6I0TcS16UsScNQRNKYYGSvdyDC7vW+HehAWil7JG8GdJf6EFMu73teL+JaVsPDRaic45UQYzJDWTd+TJsiQMhQzMsi3Uk0kwj3UIU1LBeJEB9no4IF7YJW2GyfKljDuSP09kSGudZ9HtpMj09XT3lD8z2umJj4PMipkaojA40VxylyTuMPv3TZVBBvWtwRhRe2tLu4ihbCxGeVtCP70y7OkdlzyT0qnVzaNIxgjB3uwD4fgwxmU4RIqUAUMHB7gCZ4d5Tw6r87buHXOmczswB847z9TH5Rq</latexit> distributions p X,Y | D t .
We evaluate h ’s performance using expected target risk R D t h
[ Y | Z ] :=
E p X,Y |
Dt [ (cid:96) ( Y, f ( Z ))] with respect to a well behaved loss function (cid:96) , e.g., log-loss, MSE, or 0-1 loss.
To separate domain generalization from ﬁnite sample generalization
, we consider an idealized DG (IDG), where predictors h are selected by minimizing the source population rather than empirical risk.
To give uniform guarantees while reﬂecting uncertainty over source-target
( D s , D t ) pairs, we score Z using the expected risk of the worst source minimizer h ∈ H ∗ D s :=
arg min
h R D s h
[ Y | Z ] .
Deﬁnition.
The idealized domain generalization risk (IDG risk) of an encoder p Z | X is the expected (over domains) worst-case (over source risk minimizers) target risk, i.e.
[Equation]
A representation Z ∗ is optimal for IDG if it minimizes Eq.
(1): p Z ∗ | X ∈ arg min p Z |
X R IDG
[ Y | Z ] .
Interesting DG is clearly only possible when target and source domains are related.
We assume that domains are related by the following generalized covariate shift (GCS), which says that the inputs’
Bayes predictor f ∗ = arg min f E p
Dt [R D t f
[ Y | X ]] is uniquely optimal on all domains.
Assumption (Generalized covariate shift) .
Domain risk minimizers f d ∈ arg min f
[R d f [ Y | X ]] are equal to the Bayes predictor on their support: f d ( x ) =
f ∗ ( x ) for all d, x ∈ supp( p D,X ) .
For log-loss (cid:96) GCS recovers standard covariate shift, i.e., p Y | x,d = p Y | x .
For other losses, GCS is weaker, e.g., it only requires invariance of most likely labels for 0-1 loss, and of conditional expectations for MSE.
For minor assumptions, formal statements and proofs see Appcs.
A and B.
Characterizing optimal representations for IDG under covariate shift
IDG risk is useful to evaluate representations but gives few insights into IDG and is impractical to optimize due to the sup.
in Eq.
Assuming GCS we can provide a simpliﬁed, equivalent objective that is easier to optimize.
The intuition is that under GCS any source risk minimizer will also make optimal predictions on all target samples x that are in source domain’s support.
Thus, optimal representations for IDG are exactly those that (a) ensure that all domains have the same support in Z , and (b) retain GCS from Z without sacriﬁcing the ability to predict Y optimally.
See Fig.
Theorem 1.
Under our assumptions, an encoder p Z ∗ | X is optimal for IDG if and only if it minimizes the risk R [ Y | Z ] := inf
h E p Dt (cid:2) R D t h
[ Y | Z ] (cid:3) while matching the support of Z across domains, i.e.,
[Equation]
Moreover, such encoders exist and their IDG risk is the Bayes risk
R IDG [ Y | Z ∗ ]
=
R [ Y | X ] .
Theorem 1 provides an objective to learn representations on which performing risk minimization using a single domain is as good as performing risk minimization on all domains simultaneously.
Other sufﬁcient objectives have previously been proven or hinted towards [7, 58, 13, 27], e.g., minimizing the risk while matching the representation’s marginal.
To our knowledge, Thm. 1 is nevertheless the ﬁrst to identify the necessary and sufﬁcient conditions for optimal representations.
This gives better insights into IDG and provides a framework for deriving all objectives that describe optimal IDG.
Theorem 1 shows that one must know the target domains to learn optimal representations for IDG.
Access to target domain might seem unrealistic, but without it, it is provably impossible to learn useful representations for IDG.
Speciﬁcally, the following holds under minor assumptions (see Appx.
B.3).
[Figure]
Proposition 1 (No free lunch for IDG) .
Let Z d s be any representation chosen on source d s .
For every “good” target domain outside the source’s support on which Z d s outperforms a constant representation C ∈ Z , there are many “bad” target domains on which Z d s is strictly worse than C .
Proposition 1 may explain why previous DG methods fail to outperform ERM [22]: the knowledge they have access to is insufﬁcient.
So either you access target domain d t and achieve an IDG risk that matches supervised learning (Thm. 1), or you do not and cannot do better than a constant (Prop. 1).
Learning optimal representations with domain-covering augmentations
Thm. 1 requires labels from all domains, which is impractical.
We overcome this with self-supervised learning (SSL) and particular data augmentations p A | X from inputs X to augmentations A .
The key requirement for A is to retain discriminative information about labels Y , i.e.,
if samples x, x (cid:48) ∈ X have same augmentations p
A | x = p
A | x (cid:48) , then they have the same Bayes predictions f ∗ ( x ) =
f ∗ ( x (cid:48) ) .
With such Bayes-preserving A , one can minimize R [ Y | Z ] by instead maximizing mutual information
I[ A ; Z ] [15].
However, fully optimizing I[ A ; Z ] is not generally possible under the support constraint in Eq.
This can be addressed by a domain-covering assumption: if for each d ∈ D , there is an input that is mapped to every augmentation distribution, i.e., (cid:8) p
A | x | x ∈ supp( p X | d ) (cid:9) =
(cid:8) p
A | x | x ∈ X (cid:9) .
Proposition 2.
Let p A |
X be a domain-covering augmenter.
Then any optimal solution p Z ∗ | X of the following objective is optimal for IDG:
[Equation]
Proposition 2 shows that we can still learn optimal representations for IDG without labels if we use right augmentations.
But how realistic are those augmentations?
Standard image augmentations like cropping and color jittering are generally Bayes-preserving, but not domain-covering for typical domains (e.g. sketches and photos), since outputs A are highly correlated with the domain D of the original input X , as seen in Fig. 2b.
A practical choice of A that is nearly domain-covering, is a mapping from images to text descriptions, as with CLIP [43].
Image-text augmentations have many advantages: they (i) preserve label information for many downstream tasks; (ii) are close to being domain-covering, as images from different domains but similar semantics are often mapped to similar descriptions (Fig. 2a); (iii) are easy to access in practice given their abundance on the internet.
This may explain the incredible robustness of CLIP compared to other SSL methods [11, 24, 21].
We now design practical objectives for optimizing Eq.
(3) by using a Lagrangian relaxation and introducing a domain bottleneck B[ Z, D ] that enforces support match.
Speciﬁcally, we convert Eq.
(3) to a unconstrained objective arg min p Z
| X H
[ A | Z ] + λ B[ Z, D ] .
where H[ A | Z ] replaces I
[ A ; Z ] as
H[ A ] is a constant w.r.t.
p Z | X .
A valid B[ Z, D ] ensures that minimizing B[ Z, D ] while maximizing
I[ A ; Z ] enforces the support constraint which we will introduce later.
We discuss variational bounds for the objective that can be efﬁciently estimated from samples and optimized with SGD [9].
For simplicity, we use a deterministic encoder e ϕ : X → Z .
Detailed derivations are in Appx.
C.
For the ﬁrst term, we use an upper bound on H
[ A | Z ] ≤ E p A,Z [ − log q
( A | Z )] , where q is the contrastive variational distribution as InfoNCE [40] that is standard in SSL.
Speciﬁcally, for a sample
X , we obtain a collection A :=
{ A + , A − 1 , . . .
, A − n } of one positive augmentation A + sampled from p A | X and n negatives A − i sampled from p A .
InfoNCE then constructs q ( A | Z ) using a critic s ψ to score how likely each A (cid:48)
∈
A is to be positive, see Line 6 of Algorithm 1.
For the second term B[ Z, D ] , we introduce a novel contrastive adversarial domain (CAD) bottleneck and discuss more choices in Appx.
C. Our CAD bottleneck aims to minimize I[ Z ; D ] , which enforces
[Table] [Table]
support match divergence
Dropping constants w Z we thus aim to maximize H
[ D | Z ] .
Domain-adversarial neural network [DANN, 17] does so by ensuring that a domain classiﬁer q φ cannot predict D from Z , i e., it maximizes E p D,Z [ − log q φ ( D | Z )] ≥ H[ D | Z ] w.r.t.
encoder ϕ
but minimizes it w.r t. φ .
However, it maximizes an upper bound on the desired term and suffers from unstable adversarial training
[Table]
A ←{ X } ∪ (cid:8) X − i (cid:9) n i =1 , { A + } ∪ (cid:8)
A − i (cid:9) n i = 5: X D ←{ X } ∪ (cid:8) X −
i
| D − i
=
D, i ∈
[ n ] (cid:9) 6: L aug ← − log exp s ψ
( A + ,Z ) (cid:80)
A (cid:48)∈ A exp s ψ ( A (cid:48) ,Z )
(cid:46) H
[ A | Z ] 7: L supp ← log (cid:80) X (cid:48)∈ X D exp e ϕ ( X (cid:48) )
T Z (cid:80)
X (cid:48)(cid:48)∈ X exp e ϕ
( X (cid:48)(cid:48) )
T Z (cid:46) I
[ Z ; D ] 8: return L CAD =
L aug + λ L supp
To overcome these issues, we construct q ( D | Z ) without introducing additional parameters and with a bound that is tight with enough samples.
We include the full derivation in Appx.
C.3 and brieﬂy present the algorithm here.
Given a sample
X with domain D and n tuples { ( D − i , X − i ) } ni =1 i.i.d.
sampled from p D,X , we ﬁrst obtain a variational distribution q
( X | Z ) as InfoNCE using a non-parametric critic e ϕ ( X ) T Z tied with e ϕ .
Then we collect those inputs X D associated with the current domain D , and sum q ( X (cid:48) | Z ) over
X (cid:48) ∈
X D to compute q ( D | Z ) .
See Algorithm 1 for details.
In Appx.
C.4, we derive a conditional variation of CAD that minimizes I[ Z ; D | Y ] , which can be used when labels Y are available.
In
Appx.
C.2, we also introduce an entropy bottleneck that minimizes H
[ Z ] , which does not require domain labels D that are rarely accessible in SSL.
Experiments
We aim to empirically verify our theoretical results and investigate our proposed SSL objectives in practical DG.
See Appcs.
E and F for experimental details and additional results.
Optimal representations for worst-case IDG
We validated our theory in an idealized DG setup on PACS [34].
Speciﬁcally, we (i) trained encoders on all PACS data, i.e., all domains, labels, and both train/test data; (ii) set a source D s and target D t ; (iii) selected worst-case predictors h by minimizing source risk but maximizing target risk (see [14]); (iv) repeated the last 2 steps over all source-target pairs and averaged each h ’s negative target risk (log likelihood).
Results are in Table 1.
How important is support match?
Representations trained with CAD (2 nd row) signiﬁcantly outperform those trained without bottleneck (Base, 1 st row) thus supporting Thm. 1.
Can we learn optimal representations without labels?
For domain-covering augmentations, minimizing the SSL H
[ A | Z ] (3 rd row) performs similarly to the supervised R [ Y | Z ] thus supporting Prop. 2.
Are standard augmentations sufﬁcient?
Representations trained with CAD and standard augmentations perform poorly (4 th row).
This shows the importance of domain-covering augmentations.
How important is target knowledge?
Excluding the target domain from the encoder’s training signiﬁcantly decreases performance (5 th row compared to 3 rd row), which supports Prop. 1.
Approximating optimal representations with pretrained SSL In practice, we can learn IDG optimal representations by performing SSL using a large source of unlabelled inputs X and domaincovering augmentations A .
This is nearly how CLIP was pretrained (SSL with 400M image-text pairs) except it did not include a domain bottleneck.
Here, we exploit CLIP to learn robust representations, by: (i) freezing the pretrained CLIP and adding an MLP on top of it; (ii) training the MLP with our CAD bottleneck and R [ Y | Z ] on available data.
We used the DomainBed benchmark and its protocol [22].
Due to space limit, we only report average target accuracy over PACS [34], VLCS [16],
OfﬁceHome [54], and DomainNet [41] datasets and include as baselines ‘ERM’ and ‘DomainBed
SOTA’, which is the best baseline on each dataset.
Results are in Table 2.
Can we approx.
optimal representations using CLIP?
Fine-tuning a large (ViT-B/32) CLIP with our
CAD achieves SOTA on DomainBed (7 th vs 2 nd row).
In particular, CLIP L’s performance on PACS ( 94 . 7% ) is close to the 96 .
7% estimated performance for optimal representations (see Appx.
F.2).
Are gains due to architectural differences?
DomainBed’s baseline use smaller ResNet50 models.
Finetuning a smaller ResNet50 (CLIP S) still outperforms baselines (5 th vs 2 nd row).
Our theory does not constrain the encoder and so we expect larger encoders to be better (6 th vs 5 th row).
How important is our bottleneck?
Finetuning CLIP S with our CAD bottleneck outperforms ﬁnetuning without bottlenecks and CLIP S without ﬁnetuning (5 th vs 4 th and 3 rd row).
A Preliminaries A.1 Notation
For the most part, we will assume that all spaces are discrete probability spaces.
A full list of assumptions is found at Appx.
A.3.
General The image of a set A ⊆ X under a function f : X → Y is denoted f → ( A ) =
{ f ( x ) | x ∈ A } .
The pre-image is denoted f ← ( B ) = { x ∈ X | f ( x ) ∈ B } for B ⊆ Y .
Probability Random variables (r.v.) are denoted by uppercase letters (e.g., X ), and their sample space and realizations are denoted by the corresponding calligraphic (e.g., X ) and lowercase letters (e.g., x ) respectively.
The probability mass function (pmf) of a random variable X is denoted as p X .
We use capital P instead of p to denote the measure under p .
The support supp( p X ) of a discrete distribution is the set of all points x ∈ X with positive probability, i.e., supp( p X ) =
{ x ∈ X | p X ( x ) > 0 } .
The space of all probability distributions on X is denoted
P ( X ) = (cid:8) p X | p X ( x ) ≥ 0 and (cid:80) x ∈X p X ( x ) = 1 (cid:9) .
When it is necessary to be explicit, we will denote ‘ X is distributed as p X ’ using the notation
X d ∼ p X .
Expectations are written as: E p X
[ f ( X )] , independence of two r.v. as ·⊥⊥· , conditional independence as ·⊥⊥ · | · .
For jointly distributed random variables ( X, Y ) taking value in (t.v.i.)
X × Y , the conditional distribution is denoted as p Y | X : Y × X
→ [0 , 1] .
For convenience, let p Y | x = p Y | X ( · | x ) be the conditional distribution of Y given x .
All random variables are independently distributed, unless an explicit joint distribution or coupling is given.
A.2 Deﬁnitions
We are interested in prediction problems with domain shift.
There are three random variables: the target domain D t , the input X , the label Y .
They have the following joint distribution:
[Equation]
where we drop the arguments of the probability densities for clarity.
We make a variety of convenience assumptions on these random variables (Assumption 6).
Crucially, we will be making the Bayes invariance assumption on p D t ,X,Y that can be thought of as a generalized covariate shift assumption (Assumption 4).
We will be studying the effect of changing the representation of the data.
This is done by “encoding”
X into a representation Z using a conditional distribution p Z | X .
Deﬁnition 1 (Encoder) .
An encoder is a conditional distribution p Z | X : Z × X → [0 , 1] from the input space
X to the representation space Z .
The data together with the representation has the following joint:
[Equation]
The key thing to notice here is that Z is conditionally independent of Y, D t given X .
In particular, the same encoder is used across all domains.
A.2.1 Risk minimization
Our ultimate goal is to predict Y from the representation Z of X in a manner that is robust to changes in the domain.
We formalize this in the standard way by making predictions γ ∈ Γ in a space of predictions or actions.
For example the prediction space may be the set of all possible labels Γ =
Y , in which case we would be predicting deterministic labels.
Or we may predict a distribution over labels, in which case the prediction space would be the set of all probability distributions on Y , i.e. Γ = P ( Y ) .
A predictor is a function mapping inputs to predictions, i.e., f : X → Γ , or representations to predictions, i.e., h : Z → Γ .
For example, f may be a neural network that takes as input a sample x and outputs a vector of logits that parameterize a softmax distribution over ﬁnitely many labels.
We select predictors according to the risk deﬁned via a loss function (cid:96) :
Y × Γ → R ≥ 0 ∪ {
[Equation] [Equation]
In particular, we are interested in the Bayes (minimum) risk over all predictors:
[Equation]
We denote the set of all optimal predictors from X as
[Equation]
Similarly, we deﬁne the risk R h
[ Y | Z ] , the Bayes risk R [ Y | Z ] , and the set of optimal predictors
[Equation]
from Z , all of which vary as a function of the encoder p Z | X .
Note, in the main body of the paper, we omitted the subscript Z from H ∗ Z for clarity, but we will keep it in the Appendices.
We assume that together our loss and prediction space always admit optima (Item 2 of Assumption 2), and thus
F ∗ , H ∗ Z are always non-empty.
We will be assuming that the risk admits unique optimal prediction when predicting from X (Itemof Assumption 2).
Thus it makes sense to deﬁne the following:
Deﬁnition 2 (The Bayes predictor) .
The Bayes predictor f ∗ : X → Γ is the unique predictor that is optimal for all x ∈ X :
[Equation]
Deﬁnition 3 (The Bayes image) .
The image of all the inputs under the Bayes predictor will be denoted as Γ ∗ = f ∗ → ( X ) and called the Bayes image .
Note that F ∗ becomes a singleton { f ∗ } , but it is not necessarily the case for H ∗ Z since we will not be making any uniqueness assumption on optimal prediction from Z .
[Equation]
.2 Domain generalization
We are interested in controlling the risk in a domain generalization setting, and so we deﬁne the domain-conditional risk ,
[Equation]
[Equation]
R d [ Y | X ] , F ∗ d are deﬁned as Eqs.
(7) and (8), respectively, but with respect to R df .
Similarly, deﬁne the Bayes image for domain d as
[Equation]
We also deﬁne domain-conditional quantities for prediction from a representation Z .
The most important term which we will be investigating is an idealization of the domain generalization worstcase risk.
Deﬁnition 4 (IDG risk) .
Given an encoder p Z | X and a distribution p D t ,D s over a target domain D t and source domain D s , the idealized domain generalization worst-case risk, IDG risk for short, is the expected worst-case target risk taken over source minimizers, i.e.,
[Equation]
[Equation]
Note that the IDG risk is well-deﬁned because H ∗ Z,Ds is non-empty by Assumption 2.
The desired optimal representations, are then those that minimize the IDG risk.
Deﬁnition 5 (Optimal representations for IDG) .
An encoder p Z ∗ | X is optimal for idealized domain generalization if and only if it minimizes the IDG risk, i.e.,
[Equation]
A.3 Assumptions
We make a the following assumptions throughout the paper.
All these assumptions should hold for practical settings.
Assumption 1 (Convenience: discrete probability spaces) .
All data spaces ( D , X , Y , Z , A ) are discrete spaces.
Because the distributions of X, Y, D are ﬁxed, we assume for convenience that supp( p X ) =
X , supp( p Y )
=
Y , and supp( p D t ) =
D .
Assumption 1 is a convenience assumption to avoid measure theory for the sake of clarity.
It always holds in practice due to ﬁniteness of computers, i.e., all spaces will be ﬁnite but arbitrarily large.
We believe that our claims can nevertheless be generalized to typical continuous spaces with some minor technical assumptions.
Assumption 2 (Properness of our losses) .
We assume that our risk always admits optimal predictions:
| Γ | > 1 . 2.
For all p Υ ∈ P ( Y ) , there exists γ ∗ ∈ Γ , such that
[Equation]
For all x ∈ X , there exist γ ∗ ∈ Γ , such that
[Equation]
Note that for log-loss (cid:96) ( y, γ ) =
− log γ ( y ) and ﬁnite Y , these assumptions are satisﬁed if Γ = P ( Y ) where the optimal prediction for Item 3 is γ ∗ =
p Y | x by strict properness [18].
If we consider the 0loss (reverse accuracy) (cid:96) ( y, γ ) = 1 − 1 [ y = γ ] with Γ =
Y and a ﬁnite label space where the optimal prediction for Item 3 is γ ∗ = arg max y ∈Y p Y | x ( y ) , this assumption is mostly satisﬁed, except we assume that p
Y | x has a unique mode.
Assumption 2 serves two purposes: Item 2 ensures that for any representation the optimal predictors from Z exists such that the IDG risk is well-deﬁned as in Def. 5; Item 3 ensures a unique Bayes predictor from X , which simpliﬁes the analysis and is satisﬁed by common losses as described above.
Assumption 3 (Cardinalities) .
We assume that
[Equation] [Equation]
Assumption 3 is very weak and ensures that optimal representations always exists (Prop. 3).
Assumption 4 (Generalized covariate shift) .
The Bayes predictor is optimal for all domains.
I.e., for all ( x, d ) ∈ supp( p X,D t ) , γ ∈ Γ such that γ (cid:54) =
f ∗ ( x ) , we have
[Equation]
[Equation]
For example, in the case of strictly proper scoring rules, e.g. log loss, covariate shift p Y | X,D = p Y
| X is equivalent to the invariance of the Bayes predictor.
For the 0-1 loss, this is guaranteed by invariance of the most likely label.
For MSE it is guaranteed by the invariance of the expected label.
In the latter two cases, Assumption 4 is less stringent than the typical covariate shift assumption.
Assumption 4 is the core assumption for our theoretical results.
It ensures that source and target domains are related in a useful way that can be utilized by the representation.
Assumption 5 (Constant Bayes image) .
The Bayes image is invariant across domains, i.e., for all d ∈ D ,
[Equation]
For the case of 0-1 loss, this simply means that the label set for all domains is the same, which is trivial.
For log-loss, this means that the set of possible conditional distributions Γ ∗ d = { p Y | x | x ∈ supp( p X | d ) } is the same across domains.
Assumption 5 is crucial to be able to learn.
Without it, in the extreme case, one could set each domain to be all examples associated with a single element from the label set (or the Bayes image set) in which case it is impossible to generalize across different domains.
Assumption 5 is also necessary to guarantee the existence of optimal representations as in Prop. 3.
Assumption 6 (Domain joint) .
p D t ,D s is any distribution such that supp( p D t ,D s ) =
D × D .
In a simpliﬁed scenario, one could deﬁne the source D s and target D t as i.i.d.
r.v.
from p D t , where p D t ,D s =
p D t · p D s = p D t · p D t and Assumption 6 is trivially satisﬁed.
B Proofs B.1 Lemmas for general losses
An important result that we will be using is the generalized data processing inequality of Bayes risk [56, 15].
We include it here for completeness.
Lemma 1 (Generalized DPI [56, 15]) .
Let Z − X − Y be a Markov chain of random variables.
For any loss function (cid:96) ,
[Equation]
For the case of strictly proper losses (Assumption 2) we can go one step further.
Lemma 2.
Let Z − X − Y be a Markov chain of random variables.
Then, under Assumptionsand 2 we have that
[Equation]
Proof.
Suppose that for all h ∗ ∈ H ∗ Z we have h ∗ ( z ) =
f ∗ ( x ) on the support of p X,Z .
Then,
[Equation] [Equation]
[Equation] [Equation]
[Equation]
Now suppose there exists a h ∗ ∈ H ∗ Z and a pair ( x (cid:48) , z (cid:48) )
∈ supp( p X,Z ) such that h ∗ ( z (cid:48) ) (cid:54) =
f ∗ ( x (cid:48) ) .
Then
R [ Y | Z ] (27) =
E p
[Equation]
[ (cid:96)
[Equation]
)] (28) =
p X
[Equation]
( x
[Equation]
z
[Equation]
[ (cid:96) (
[Equation]
h ∗ ( z
[Equation]
] + (cid:88)
[Equation]
) p
[Equation]
| x [ (cid:96)
[Equation]
z ))] (29) ≥ p X
[Equation]
Z ( x
[Equation] [Equation]
))] + (cid:88) ( x
[Equation]
) p
[Equation]
( x
[Equation]
z ) E
[Equation]
[ (cid:96) (
[Equation]
, f
[Equation]
( x ))] (30) > p X,Z ( x (cid:48) , z (cid:48) )
E p Y | x (cid:48)
[ (cid:96) ( Y, f ∗ ( x (cid:48) ))]
+ (cid:88) ( x
[Equation]
z
[Equation]
(cid:54) =( x (cid:48)
[Equation]
z (cid:48) )
p X
[Equation]
Z ( x
[Equation]
Eq.
(30) follows by Item 3 of Assumption 2 along with the deﬁnition of f ∗ .
Eq.
(31) follows by Itemof Assumption 2 and the fact that h ∗ ( z (cid:48) ) (cid:54)
=
f ∗ ( x (cid:48) ) .
This completes the proof, because Lemmaprevents R [ Y | Z ] < R [ Y | X ] .
B.2 Proof of Theorem
First we will show that the desired representation exists by taking all inputs for which the Bayes predictor predicts similarly and “bucketing” them to the same representation.
This is a direct extension of the example from Dubois et al.’s
Proposition 6, to the case of proper losses.
Proposition 3 (Existence of optimal representations) .
Under Assumptions 1 to 5, there exists an encoder p Z ∗ | X that is optimal for IDG, i.e.,
[Equation]
Moreover, we have that
[Equation]
Proof.
Because we assume arbitrary encoders p Z | X , the essence of this construction is simple: we embed the Bayes image into Z .
Indeed, let φ : Γ ∗ → Z be any one-to-one function, which exists due to Assumption 3 (here we use deterministic one-to-one function for simplicity, the construction can be easily extended to stochastic case).
Then let Z ∗ =
φ ( f ∗ ( X )) .
We now verify the properties of p Z ∗ | X .
Z ∗ satisﬁes R [ Y | X ]
=
R [ Y | Z ∗ ] .
Indeed,
[Equation]
Eq.
(37) is by our construction of Z ∗ and Eq.
(38) is by the deﬁnition of the Bayes risk.
Due to the data processing inequality of Bayes risk (Lemma 1) we also have R [ Y | X ] ≤
R [ Y | Z ∗ ] , from which we conclude that R [ Y | X ] =
R [ Y | Z ∗ ] and that Eq.
(34) holds.
Recall that Γ ∗ = f ∗ → ( X ) and Γ ∗ d = f ∗ → (cid:0) supp( p X | d ) (cid:1) .
Now let us compute the desired support for all d ∈ D :
[Equation]
Eq.
(40) is by Assumption 5.
Because R [ Y | X ] is the minimum achievable risk by any encoder regardless of constraint (this is by
Lemma 1), this implies that p Z ∗ | X is an optimal encoder for IDG.
The following lemma essentially says that when R [ Y | Z ] is minimized, then the optimal predictors for each domain all agree on the intersection of their support.
Lemma 3.
Let p Z | X be an encoder
such that R
[ Y | Z ] =
R [ Y | X ] .
Under Assumptions 1 and 2, we have that for all z ∈ supp( p Z ) , there exists γ ∗ ∈ Γ such that
[Equation]
In other words, the restriction of any h ∗ ∈ H ∗ Z to supp( p Z ) is unique.
If, in addition, Assumption 4 holds, then for all ( z, d ) ∈ supp( p Z,D t ) , γ ∈ Γ such that γ (cid:54) =
h ∗ ( z ) ,
[Equation]
In other words, the restriction of any h ∈ H ∗ Z,d to supp( p Z | d ) is unique and equal to h ∗ .
Proof.
For the ﬁrst result, let z ∈ supp( p Z ) and consider x ∈ supp( p X | z ) .
By Lemma 2, it must be the case that f ∗ is constant on supp( p X | z ) .
Thus, we can pick γ ∗ =
f ∗ ( x ) .
Now, let γ (cid:54) =
γ ∗ .
We have that,
[Equation] [Equation]
[Equation] [Equation]
Eq.
(44) is due to the conditional independence of Y and Z given X .
Eq.
(46) is due to Assumptionand the deﬁnition of the Bayes predictor.
Let h ∗ : supp( p Z ) → Γ be the unique Bayes predictor from Z .
Now, for the second result, note that
[Equation]
Assumption
[Equation]
and
[Equation]
where Eq.
(53) is due to the deﬁnition of (domain-conditional) Bayes risk.
Then
[Equation]
Lemma 1 conditioned on d
[Equation] [Equation]
Thus, any encoder that achieves R [ Y | Z ] =
R [ Y | X ] also satisﬁes R d
[ Y | Z ] =
R d [ Y | X ] for all d ∈ D since we assume that supp( p D t )
=
D in Assumption 1.
Now, let d ∈ D .
An argument analogous to Lemma 2 gives us,
[Equation] [Equation]
Eq.
(56) is derived from R d
[ Y | Z ] =
R d [ Y | X ] using Assumption 4 in place of Item 3 of Assumption 2 for a speciﬁc domain d .
Let z ∈ supp( p Z | d ) and γ ∈ Γ such that γ (cid:54)
=
h ∗ ( z ) .
Since supp( p X | z,d ) ⊆ supp( p X | z ) , f ∗ is a constant on supp( p X | z,d ) and equal to h ∗ .
Now, as above, we have that
[Equation]
Eq.
(59) is due to Assumption 4.
Corollary 1.
Let p Z | X be an encoder such that R [ Y | Z ]
=
R [ Y | X ] .
Under Assumptions 1, 2 and 4 we have that
[Equation]
D and that for
[Equation]
Proof.
H ∗ Z ⊆
H ∗ Z,d is immediate from Lemma 3.
Now, we have that R d th [ Y | Z ] ≥ R d t [ Y | Z ] .
So, the result follows by taking any h ∈ H ∗ Z ⊆
H ∗ Z,ds in the inf of Eq.
Theorem 2 (Characterizing optimal representations for IDG) .
Under Assumptions 1 to 6, an encoder p Z | X is optimal for idealized domain generalization if and only if it minimizes the Bayes risk while matching the support of p Z | d and p Z for all d ∈ D , i.e.,
[Equation]
Proof.
The IDG risk is lower bounded by R [ Y | X ] :
[Equation]
We will now show that this lower bound is achieved by an encoder if and only if it satisﬁes Eqs.
(62) and (63), which exist by Prop. 3.
Sufﬁciency ( ⇐ = ):
Let p Z | X be an encoder that satisﬁes Eqs.
(62) and (63).
Note that R [ Y | Z ] =
R [ Y | X ] by Prop. 3.
Let h ∗ ∈
H ∗ Z , then we have the following IDG risk
[Equation]
Lemma 3 under matching support
[Equation]
Necessity ( = ⇒ ):
If the IDG risk is R [ Y | X ] , then it must be the case that
[Equation]
We will prove by contrapositive that Eq.
(74) implies support match (Eq.
Suppose that the support match does not hold.
Since supp( p Z ) = ∪ d ∈D supp( p Z | d ) and supp( p D s ,D t ) =
D × D
(Assumption 6), there must exist ( d s , d t ) ∈ supp( p D s ,D t ) such that supp( p Z | d s ) (cid:54) =
supp( p Z | d s ) .
Deﬁne the set S = supp
( p Z | d s ) ∩ supp( p Z | d t ) and ¯ S = supp( p Z | d t ) \ supp( p Z | d s ) , let ρ =
P Z | d t ( S )
[Equation]
and let h ∗ ∈ H ∗ Z .
Then,
[Equation]
Lem.
[Equation]
Eq.
(80) uses the following reasoning.
1 − ρ > 0 due to support mismatch.
For any h ∈ H ∗ Z,ds such that h (cid:54) =
h ∗ on ¯ S (such an h exists by Item 1 of Assumption 2), we have that
[Equation]
by Lemma 3.
As a corollary from the proof strategy we directly have that the optimal DG risk is simply R [ Y | X ] .
This means that using the optimal encoder one can actually perform just as well by training on the source as if you were to directly train on the target using the raw data.
Corollary 2 (Optimal IDG Risk) .
Under Assumptions 1 to 6, inf p Z |
X R IDG
[ Y | Z ] =
R [ Y | X ] .
B.3 Impossibility results
As a direct corollary of Thm.
2 we know that it is impossible to learn an optimal representation without knowledge or assumptions on the target domain.
We can actually prove the following much stronger negative result, which essentially states that it is impossible to ﬁnd a useful representation without having some information about the target domain.
Speciﬁcally, we prove that if there exists a non-trivial target domain on which the representation is advantageous then there exists an inﬁnite amount of target domains on which it is disadvantageous compared to predicting from a constant.
For clarity, we will focus on the proof for the standard accuracy (0-1 loss) which is much shorter and simpler to understand, but note that we can generalize the proof to all losses with the right assumptions.
The key is that outside of the source domain, the label distribution is unconstrained because generalized covariate shift has no effect.
In other words, for any domain which gives some probability mass on an example that has not been seen during training, then all possible labels for that example gives a valid domain.
Furthermore, if there exists one domain on which the representation is good, then one can construct a domain on which the representation is bad simply by labelling this point as the constant prediction.
Proposition 4 (No free lunch learning representations for DG) .
Let (cid:96) be the 0-1 loss with prediction space Γ =
Y .
Let Rep : P ( X , Y ) → P ( Z|X ) be any algorithm for choosing an encoder p Z | X from the data distribution p
X,Y , C be any constant r.v.
that t.v.i.
Z , and p X,Y | d s be any desired source distribution such that
• there is a unique constant prediction γ C = arg min
y ∈Y E p Y | ds [ (cid:96) ( Y, y )] , • and |X \ supp( p X | d s ) | > 1 .
Let p Z ds | X :=
Rep( p X,Y | d s ) be the chosen source encoder.
If there exists a target domain p X,Y | d gt such that
• (Non-trivial support) ∅ (cid:54) =
supp( p X | d gt ) ⊆
X \ supp( p X | d s ) ; • (Satisﬁes Bayes image invariance)
Γ ∗ d gt = Y , i.e., there is at least one example for every possible label; • (Source encoder is useful)
p Z ds | X performs better than a constant representation,
[Equation]
Then there exist multiple target domains d bt such that p Z ds |
X underperforms a constant encoder,
[Equation]
Proof.
Let h ∗ ∈
H ∗ Zds,ds be any source Bayes predictor corresponding to our encoder.
Partition Z according to whether h ∗ predicts like the constant or not:
[Equation]
We know by assumption that d gt is s.t.
[Equation]
which is clearly only possible if
[Equation] [Equation]
In other words, there exists some input x (cid:54) =
C ∈ X \ supp( p X | d s ) that will get represented outside of the constant region, i.e.,
[Equation]
We will now construct the desired bad domain d bt by giving nearly all mass to this x (cid:54) =
C , speciﬁcally, let p X | d bt ( x (cid:54) =
C )
= 1 − δ for some 0 < δ < 1 .
We assign this example to the constant label, i.e., p Y | x (cid:54) =
C ,d bt ( γ C ) = 1 .
The rest of the target domain mass δ is distributed as with the source domain, i.e., p X,Y | d bt ( x, y )
=
δ · p X,Y |
d s ( x, y ) for all x, y ∈ supp( p X,Y | d s ) .
Importantly, the constructed domain d bt is valid.
Indeed, the Bayes image is the same as the source’s
(Assumption 5), because we removed no prediction γ from the source’s Bayes image ( δ > 0 ).
We added no new prediction γ , because f ∗ ( x (cid:54) = C ) =
γ C ∈ Y which must already have been in Γ ∗ due to the validity of d gt .
Now let us compute the desired risk for that “bad” domain and show that the desired encoder performs worse than a constant encoder.
[Equation]
In contrast, it is easy to show that sup h ∈H ∗ C,ds R d bth [ Y | C ] ≤ δ because the constant predictor would be perfect for x (cid:54)
=
C .
So any choice of 0 < δ <
[Equation]
) , would satisfy Eq.
We conclude the proof by noting that there are inﬁnitely many such choices of δ , and any choice of those would result in a different valid bad domain d bt .
Note that representations can often be much worse than using a constant r.v.
Speciﬁcally, if an encoder p Z | X maps an x outside of the source support then there exists an inﬁnite number of target domains where that representation is the worst possible representation.
Proposition 5 (Worst representation) .
Let Rep , p Y,X | d s , p Z ds | X , (cid:96) be as in Prop. 1, and (cid:15) > 0 .
If there exists an example x b ∈ X \ supp( p X | d s ) that is mapped outside of the source support, i.e., supp( p Z ds | x b ) ∩ supp( p Z | d s ) = ∅ , then there exist many target domains p X,Y | d t s.t.
p Z ds |
X is (cid:15) close to the worst possible loss, i.e.,
[Equation]
Proof.
By assumption there exists an x b whose support is outside the source support.
Then similarly to Prop. 1 we construct a bad target domain d t by giving nearly all mass to that example p X | d t ( x b ) = 1 − δ where δ > 0 and assign with probability 1 to some label that is in the source Bayes image, i.e., p Y | x b ,d t ( γ b ) = 1 for some γ b ∈ Γ ∗ d s .
The rest of the target domain mass δ is distributed as in Prop. 1 to the source inputs.
As in Prop. 1, such a target domain d t satisﬁes our assumptions.
Now let us compute the risk for that d t and show that the desired encoder performs arbitrarily bad.
[Equation]
Eq.
(96) uses the fact that H ∗ Zds,ds is unconstrained outside of the source support and that by assumption supp( p Z ds | x b ) ∩ supp( p Z ds | d s ) = ∅ .
To achieve the sup 1 − δ it then sufﬁces to predict an γ (cid:54) =
γ b ∈ Γ .
We thus see that Eq.
(92) holds for d t as long as 0 < δ < (cid:15) .
We conclude the proof by noting that there is an inﬁnite possible choices of δ each of which give rise to a bad target domain.
B.4 Augmentations
Prop. 2 shows that the optimal representations for IDG can be learned with augmentations in a self-supervised fashion.
Here, we provide formal deﬁnitions, assumptions, and proofs.
Deﬁnition 6 (Augmenter) .
An augmenter is a conditional distribution p
A | X : A × X → [0 , 1] from the input space
X to an augmentation space A .
For example, in CLIP X is the space of images and A is the space of text.
In standard SSL,
A is typically the same as X (e.g., both X
and A are the space of images).
Deﬁnition 7 (Augmentation conditional set) .
Given an augmenter p
A | X , deﬁne the augmentation conditional set as the set of conditionals of A given X :
[Equation]
[Equation]
Similarly, we can deﬁne the augmentation conditional set for domain d :
[Equation] [Equation]
These sets are clearly countable.
Note that the augmentation conditional set can be seen as a special case of the Bayes image (Def. 3) if we view the augmentation A as the label and consider the log-loss where the conditional distribution is the Bayes optimal predictor due to its strict properness [18].
Assumption 7 (Finite augmentation entropy) .
We consider the augmenter p
A | X such that the entropy of the augmentation A is ﬁnite, i.e., H[ A ] < ∞ .
Assumption 8 (Cardinalities) .
We assume that
[Equation] [Equation]
This is a similar assumption as Assumption 3, which ensures the existence of optimal representations.
Assumption 9 (Domain-covering augmentation) .
We assume that the augmentation A is domaincovering , i.e., the augmentation conditional set is invariant across domains,
[Equation]
This assumption is generalized from the constant Bayes image assumption (Assumption 5), which guarantees the existence of optimal representations.
Domain-covering augmentations essentially ensures that each augmentation conditional p
A | x ∈
P ∗ ( A | X ) is seen at least once in all domains.
If we introduce an equivalence relation ∼ as x ∼ x (cid:48) iff p
A | x = p
A | x (cid:48) and the equivalence class [ x ] := { x (cid:48) ∈ X | x (cid:48) ∼ x } .
Under this relation, it is easy to see that the above assumption is satisﬁed if and only if, for all possible equivalence classes [ x ] ∈{ [ x (cid:48) ] | x (cid:48) ∈ X
} , we have that [ x ] has intersections with all domains:
[Equation]
Not all augmentations are domain-covering.
In particular, the standard image augmentations used by typical SSL models like SimCLR are not domain-covering, but the text-image augmentations of
CLIP nearly are, as discussed in the main body (Sec. 4).
Assumption 10 (Bayes-preserving augmentation) .
We assume that the augmentation A is Bayespreserving , i.e., ∀ x, x (cid:48)
∈ X ,
[Equation]
Under the notion of equivalence relation in Assumption 9, this means that for each equivalence class [ x ] , all x (cid:48) ∈
[ x ] have the same Bayes prediction.
Note that most augmentations used in practice like standard image augmentations are Bayes-preserving.
Next, we show that under the above assumptions, we can learn optimal representations by maximizing the mutual information
I[ A ; Z ] (in the case of log-loss (cid:96) ) under the support match constraint.
We use log-loss simply because it is typically the loss used for training in practice.
Note that the learned representations are optimal for any strict proper losses.
Proposition 6 (Optimal encoders without labels) .
Let p A |
X be an augmenter.
Under Assumptionsto 10, any encoder p Z | X such that
[Equation]
is optimal for idealized domain generalization.
Proof.
The support match constraint Eq.
(104) is equivalent to the support match constraint Eq.
Thus, Prop. 3 and Thm. 2 state that we only need to prove that maximizing the mutual information of
A and Z under the support constraint implies that
[Equation]
We will prove this by constructing an optimal predictor h ∗ .
Since H[ A ] < ∞ (Assumption 7) we have that
[Equation]
Note the fact that the conditional entropy is the Bayes risk under the log-loss [18], i.e., H
[ A | Z ] =
R [ A | Z ] .
By construction, A satisﬁes covariate shift w.r.t.
X (thus Bayes invariant) since A − X − D forms a Markov chain.
Together with Assumptions 1 and 7 to 9, it means that the optimization problem in Eqs.
(103) and (104) satisﬁes the assumptions of Prop. 3, with A in place of Y .
Thus, an optimal encoder satisﬁes R
[ A | Z ]
= R
[ A | X ] , which leads to
[Equation]
By Assumption 7, we can invoke Lemma 2 with the fact that A − X − Z forms a Markove chain to show that for all ( x, z ) ∈ supp( p X,Z )
[Equation]
as the conditional distributions are the Bayes optimal predictors due to strict properness of log-loss.
Now, deﬁne the following equivalence relation on X ,
[Equation]
[Equation]
Because the number of equivalence classes under ∼ is countable, there exists a maximal invariant
M : X → N from X to the natural numbers [for our deﬁnition of a maximal invariant see Deﬁnition 2, 15].
By Assumption 10, f ∗ is invariant on the equivalence classes [ x ] :=
{ x (cid:48) ∈ X | x (cid:48) ∼ x } for all x ∈ X .
Thus, there exists a function g : N → A such that f ∗ =
g ◦ M [Lemma 5, 15].
Given z ∈ supp( p Z ) , we construct h ∗ in the following way.
Let x z ∈ supp( p X | z ) be any input point that could have led to this representation z and deﬁne
[Equation]
By Eq.
(108) we are guaranteed that all x ∈ supp( p X | z ) share the same value for f ∗ since they are in the same equivalence class.
Thus, by the deﬁnition of M we have that
M
[Equation]
x Z )
=
M ( X ) for ( X, Z ) ∼ p X,Z .
Therefore,
R
[Equation] [Equation]
[Equation]
Eq
[Equation] [Equation]
Eq
[Equation] [Equation]
C Practical objectives
Proposition 6 provides an objective to obtain the desired optimal representations, compared to Thm.it is more practical in that it does not require access to the labels but right augmentations.
There are nevertheless multiple remaining issues for deriving objectives that can be trained with in practice.
Speciﬁcally, (i) the support constraint is hard to satisfy in practice; (ii) mutual information
I
[ A ; Z ] is hard to estimate from samples [42]; (iii) the objective is constrained which is harder to optimize.
We will now show different objectives and variational bounds of them that do not suffer from these issues, and could still recover the desired encoders in their optima.
In contrast to the proofs of main theoretical results (previous section), here the derivations will be less formal.
As we have seen in Proposition 6, the optimal representation achieves I
[ A ; Z ]
= I
[ A ; X ] , and thus
H
[ A | Z ]
= H
[ A | X ] (Eq. (107)).
In the following, we will rewrite the objective as the constrained optimization:
[Equation] [Equation]
[Equation]
where we introduce the domain bottleneck B[ Z, X, Y, D ] as the objective for enforcing support match (which we denote as B[ Z, D ] in the main body for simplicity).
The requirement on the domain bottleneck objective is that minimizing Eq.
(117) under Eq.
(118) implies that the support match constraint holds (and can be achieved by some encoder), which leads to optimal representations for idealized domain generalization.
Different domain bottlenecks will be derived later this section.
We can then use Lagrangian relaxation to get the following unconstrained objectives.
[Equation] [Equation]
The ﬁrst term can be easily optimized using variational objectives.
Throughout the paper, we will use a contrastive variational upper bound which is based on InfoNCE [40].
Namely, let X := (cid:8)
X, X − 1 , . . .
, X − n (cid:9) be a sequence of samples where each negative X − i is independently sampled from the marginal p X .
Now let A :
=
(cid:8)
A + , A − 1 , . .
A − n (cid:9) be the sequence of augmented examples that come from independently augmenting each X ∈ X using the augmenter p A | X .
In particular, the augmented positive A + has the marginal p
A | X while the negatives A − i follow the marginal p A .
Let
Z be the representation of X by passing it through the encoder p ϕ :
= p Z | X parameterized by ϕ and s ψ the critic function parametrized by ψ used to score which A (cid:
) ∈
A is the positive augmentation.
Then we have a variational distribution of p A | Z :
[Equation]
[Equation]
which leads to the following variational bound:
[Equation]
(121) with equality if the variational families s ψ , p ϕ are unconstrained and and we use inﬁnite samples ( n → ∞ ).
Typically the critic is separable, i.e., s ψ ( A, Z ) :
= g ψ ( A ) T h ψ ( Z ) .
As discussed in the main body, it can be tied with p ϕ when A = X .
In the following we focus on the second term B[ Z, X, Y, D ] and discuss several choices.
C.1 Mutual information bottleneck B[ Z, X, Y, D ] = I[ Z ; X ]
The ﬁrst bottleneck we consider is so called mutual information (MI) bottleneck B[ Z, X, Y, D ] =
I[ Z ; X ] , which was introduced by Tishby et al.
[53] to achieve a tradeoff between the predictive power and the complexity of representations.
Intuitively, it tries to remove all information of Z that is not needed for maximizing I
[ Z ; A ] .
In particular, using the fact that Z − X − D forms a Markov chain and the chain rule of MI, we have I[ Z ; X ] =
I
[ Z ; X, D ] =
I
[ Z ; D ] + I
[ Z ; X | D ] .
Thus, it not only minimizes I[ Z ; D ] , i.e., matches the representations’ distribution across domains, but also minimizes I
[ Z ; X | D ] , i.e., matches the representations’ distribution inside domains.
Why The key to show is that minimizing Eq.
(117), i.e., arg min p Z |
X I[ Z ; X ] under I[ A ; Z ] =
I[ A ; X ] , implies the support match constraint.
This can be seen as a speciﬁc subcase of Dubois et al.’s (2021)
Corollary 15 with A in place of Y and M ( X ) induced by p A |
X
as in the proof of
Prop. 6.
From the corollary, we know that min p Z |
X I
[ Z ; X ] =
H[ M ( X )] which can be achieved by any Z s.t.
p Z | x = p Z | x (cid:48)
⇐⇒ M ( x ) =
M ( x (cid:48) ) .
With the assumption of domain-covering augmentations (Assumption 9), we have that the set of maximal invariant { M ( x ) | x ∈ supp( p X | d ) } is invariant across domains.
Then we directly have supp( p Z | d ) = ∪ x ∈ supp( p X | d ) supp( p Z | x ) = ∪ x ∈ supp( p X ) supp( p Z | x ) =
supp( p Z ) , where we use the fact that x within the same equivalence class has the the same p
Z | x .
How Essentially, we can use any variational upper bound of mutual information.
We consider the one used by Variational Information Bottelenck [1], i.e.,
[Equation] [Equation]
[Equation]
where a variational distribution q θ is used to approximate p Z and is jointly optimized with p ϕ to minimize the bound.
The approximation gap of the bound is D KL [ p Z ( Z ) (cid:107) q θ ( Z )] .
Then the ﬁnal loss becomes
[Equation]
which recovers the optimal encoder in the case of unconstrained variational families for p ϕ , q θ , s ψ , inﬁnite samples | A | , and any λ > 1 [15].
C.2 Entropy bottleneck B[ Z, X, Y, D ] = H[ Z ]
The entropy (Ent) bottleneck introduced in the main body is a special case of the MI bottleneck, where the encoder is a deterministic mapping, i.e., p ϕ
( Z | x ) is a dirac delta function for all x ∈ X
and we denote by e ϕ ( x ) the deterministic encoder s.
[Equation]
p ϕ
[Equation]
= 1 .
Why In the deterministic case, the MI bottleneck becomes the entropy bottleneck because I[ X ; Z ] =
H[ Z ] − H[ Z | X ] =
H[ Z ] , where we use the fact that H[ Z | X ] = 0 .
Importantly, considering only deterministic encoders does not constrain our ability to learning optimal encoders.
Indeed, just as with the MI bottleneck optimizing the objective with the entropy bottleneck under I[ A ; Z ] = I
[ A ; X ] will recover encoders s.t.
e ϕ ( x (cid:48) )
=
e ϕ ( x ) ⇐⇒ M ( x ) =
M ( x (cid:48) ) , which also satisﬁes the support match constraint as discussed before.
How Using the same derivation as the MI bottleneck, we can derive the variational upper bound on entropy
[Equation]
which is the standard variational bound used in neural compression [3, 52].
Putting all together, we have
[Equation]
which also recovers the optimal encoder with unconstrained variational families, inﬁnite samples, and λ > 1 as with the MI bottleneck.
The detialed algorithm is provided in Algorithm 2.
Note that the discreteness of Z could lead to difﬁculty of gradient-based optimization, and we follow Ballé et al.
[3] to add uniform noise to Z as a differentiable substitute for rounding during training.
In our experiments, we will mostly use the Ent bottleneck instead of the MI bottleneck to avoid introducing stochastic encoders.
Algorithm
Require: e
[Equation]
[Equation] [Equation]
← sample
[Equation]
p
[Equation] [Equation]
sample
[Equation]
[Equation] [Equation]
←
[Equation]
exp s
[Equation] [Equation]
A (cid:48)∈ A exp
[Equation] [Equation]
(cid:48)
[Equation]
Z )
(cid:46) H
[ A | Z ] 6: L supp ← − log q θ ( Z ) (cid:46) H
[ Z ] 7: return L Ent =
L aug + λ L supp
C.3 Contrastive adversarial domain bottleneck B[ Z, X, Y, D ] = I[ Z ; D ]
The previous two bottlenecks require to remove the information of Z (about X ) as much as possible, which seems to be unnecessary since our ultimate goal is to match the support of Z across domains.
Now we introduce a bottleneck B[ Z, X, Y, D ]
= I
[ Z ; D ] which we only seek to remove the information of Z about the domain D .
This is very related to the work on invariant representation learning for domain generalization/adaptation [e.g., 17, 35].
We derive a new variational bound called the contrastive adversarial domain (CAD) bottleneck that is more stable to train and leads to better empirical performance.
For simplicity we consider the deterministic encoder e ϕ ( x ) as with the main body.
Why Similar to the previous analysis, we aim to show that arg min p Z
| X I[ Z ; D ] under I
[ A ; Z ] =
I[ A ; X ] leads to the support match constraint.
Using Eq.
(111) we have I[ Z ; D ] =
I[ Z, M ( X Z ); D ] =
I[ Z, M ( X ); D ] = I[ M ( X ); D ] + I
[ Z ; D | M ( X )] where the last equality uses the chain rule of mutual information.
Due to the non-negativity of (conditional) mutual information, we have that the minimum of I[ Z ; D ] under I[ A ; Z ] =
I
[ A ; X ] is I
[ M ( X ); D ] .
Then we show the minimum is achievable by constructing the same optimal encoder e ϕ ( X ) as the Ent bottleneck which clearly satisﬁes I[ Z ; D | M ( X )] = 0 .
It is then easy to show that the support match constraint has to hold
when I[ Z ; D | M ( X )] = 0 by contrapositive.
Indeed, suppose that the support constraint does not hold then it must be true that I[ Z ; D | M ( X )] > 0 and so the encoder cannot be optimal.
How The typical way of minimizing I[ Z ; D ] is to derive the variational bound as
[Equation]
where a variational distribution (or domain classiﬁer) q φ is used to approximate p D | Z and jointly trained to maximize the bound.
This recovers the domain-adversarial training method as introduced in Ganin et al.
However, this has two potential issues: 1) it gives a lower bound instead of the desired upper bound on I[ Z ; D ] ; 2) it requires adversarial training which is not stable [20, 30].
We propose the contrastive adversarial domain (CAD) bottleneck, which is based on the above explicit version but uses a variational distribution q φ ( D | Z ) that is tied with other parts of the model , thus no need to learn a domain classiﬁer.
Speciﬁcally, as with Eq.
(120), we ﬁrst introduce a contrastive variational distribution q ϕ, X
( X | Z ) of p X | Z as:
[Equation]
where s ϕ ( X, Z ) :
= e ϕ ( X ) T Z is tied with the encoder e ϕ .
Since p D | Z can be rewritten as
E p X | Z (cid:2) p D | X (cid:3) using the fact that D − X − Z forms a Markov chain, we obtain can the following variational distribution:
[Equation]
which recovers p D | Z when q ϕ, X
= p X | Z .
This is the case where s ϕ
( X, Z ) ∝ log p X,Z ( X, Z ) and inﬁnite samples n → ∞ .
Note that p D | X is still not available, we can use a count estimate
ˆ p D , X in practice.
In particular, we obtain a collection D by taking each X (cid:48)
∈ X and independently sampling D (cid:48) from p D | X (cid:48) to get D :=
{ D, D − 1 , . . .
, D − n } .
In other words, ( D, X ) and ( D − i , X − i ) for i ∈ [ n ] := { 1 , . . . , n } are all i.i.d.
sampled from p D,X .
Then we use a count estimate
[Equation] [Equation]
which is an accurate estimate with inﬁnite samples.
This leads to our ﬁnal variational distribution:
[Equation]
Putting all together we get that the ﬁnal loss:
[Equation]
In practice, ˆ p D , X
( D | X ) is typically a dirac delta function since it is rare to have the same samples in a batch.
Thus, in Eq.
(135) we only need to sum q ϕ, X ( X (cid:48) | Z ) over those associated with the same domain label D as X , which reduces Eq.
(136) to
[Equation]
with a detailed algorithm in Algorithm 1.
Note that it is easy to generalize Algorithm 1 to parallel computation within a batch of samples.
Indeed, for each sample in the batch, we can view all other samples in the batch as negatives and compute the loss efﬁciently in parallel.
C.4 Conditional CAD B[ Z, X, Y, D ] = I[ Z ; D | Y ]
The analysis of the CAD bottleneck also implies that we can minimize the conditional mutual information I[ Z ; D | M ( X )] if we have access to M ( X ) .
However, since M ( X ) is typically not available in practice, we consider the special case where M ( X ) =
Y .
In particular, this is the case where the labels are available and the supervised augmentations are used (see Fig. 2c).
This reduces the bottleneck to B[ Z, X, Y, D ] =
I
[ Z ; D | Y ] which is related to the conditional version of the domain-adversarial neural network [36].
In practice, minimizing I[ Z ; D | Y ] could be easier for optimization than I[ Z ; D ] , as it does not require to remove the information that D has about Y .
In the following, we derive the conditional CAD (C 2 AD) bottleneck using a similar idea as CAD.
How In this case, we want to minimize
[Equation] [Equation]
[Equation]
(140) where q ( D | Z, Y ) is a variational distribution of p D | Z,Y .
Similar to the unconditional case, we also aim to use a non-parametric approximation that is tied with other parts of the model, and we obtain it using the fact p D | Z,Y =
E p X | Z,Y (cid:2)
p D | X (cid:3) .
Speciﬁcally, let Y :=
{ Y, Y − 1 , . . .
, Y − n } be the collection of labels obtained by independently sampling the label from p Y | X (cid:48) for each X (cid:48)
∈ X .
We collect samples associated with the label Y , i.e., X Y :=
{ X } ∪{ X − i | Y − i
=
Y, i ∈ [ n ] } and obtain a variational distribution of p
[Equation]
where we use the same critic s ϕ ( X, Z ) :
= e ϕ ( X ) T Z that is tied with the encoder e ϕ as before, but only take softmax over those samples with the same label Y .
For the term p D | X , we use the same count estimate ˆ p D , X in Eq.
Then we obtain the variational distribution of p D | Z,Y :
[Equation]
Putting all together we get that the ﬁnal loss:
[Equation]
Again, since in practice ˆ p D , X ( D | X ) is typically a dirac delta function
[Equation]
the summation in Eq.
(142) can be done only over those associated with the same label Y and the same domain label D as X .
In particular
[Equation]
we collect samples X Y,D := { X } ∪{ X − i | Y − i
=
Y, D − i
=
D, i ∈
[ n ] } and obtained the simpliﬁed loss:
[Equation]
A detailed algorithm is in Algorithm
[Equation]
Require:
[Equation]
D
[Equation]
X
[Equation]
n
[Equation]
[Equation] [Equation]
A
[Equation]
← sample ( p
[Equation]
X
[Equation] [Equation]
(cid:8) ( D
[Equation]
i
[Equation]
X
[Equation]
i
[Equation]
A
[Equation]
i
[Equation]
Y
[Equation]
i ) (cid:9) n i =1 i.i.d.
←−− sample ( p D
[Equation]
X
[Equation]
A
[Equation]
Y ) 4:
X
[Equation]
A ←{ X } ∪ (cid:8) X − i (cid:9) n i =1 , { A + } ∪ (cid:8)
A − i (cid:9) n i = 5: X Y ←{ X } ∪ (cid:8) X − i
[Table]
D Related work
Provably robust representations Z under covariate shift.
Ben-David et al.
[7] bounds the target risk using source risk and a divergence between source and target distributions.
They do not consider representation learning, but in our setting, this implies that matching the marginal of Z while minimizing R [ Y | Z ] is sufﬁcient for optimality.
Ben-David et al.
[8] suggests that R [ Y | Z ] is not sufﬁcient, and
Zhao et al.
[58] also prove that one should minimize the joint R [ Y | Z ] instead of the source R d s [ Y | Z ] risk.
Similarly, des Combes et al.
[13] shows that matching the conditional p Y | Z,d = p Y | Z is sufﬁcient.
Johansson et al.
[27] take this further by proving that matching only the support of Z is also sufﬁcient.
Our work distinguishes itself from those and other related work on three key aspects: (i)
We are the ﬁrst to provide the set of necessary and sufﬁcient conditions for robust representations; (ii) We prove that one can learn optimal Z ∗ with SSL using only large samples of inputs X and domain-covering augmentations A .
(iii)
We consider a general DG setting which deals with a less stringent generalized covariate shift and works for all standard losses and Y in ML.
Still, our work is more speciﬁc than others, as we consider idealized DG and unrestricted predictors H .
Our theory could be combined with Dubois et al.’s
[2020], who provide conditions for optimal generalization from ﬁnite samples and constrained H in supervised learning.
Practical objectives for DG.
The most popular DG methods aim to learn domain-invariant representation by minimizing various divergernces between the conditionals p Z | d and marginals p Z [37, 17, 49, 38, 35, 47, 39].
Others propose matching the conditional p Z | y,d across domains instead [19, 36, 50].
These regularizers would all be valid domain bottlenecks B
[ Z, D ] .
Another line of work aims at learning Z with invariant predictors p Y | z,d across domains [e.g., 2, 31, 33].
However, none of these methods outperform ERM with fair model selections [22].
[Figure]
E Experimental Details E.1 Scientiﬁc: exploring optimal representations for worst-case DG
In both the scientiﬁc setting and the following bridge setting, we consider rather unrealistic setups for verifying our theory where we have access to labels from all domains.
We can choose to directly minimize the risk R [ Y | Z ] with the cross-entropy loss (denoted as CE henceafter), or minimize
H
[ A | Z ] (i.e., maximize I
[ A ; Z ] ) with supervised augmentations as in Fig.
2c detailed below.
Implementation of supervised augmentations When using supervised augmentations, for each sample we obtain its augmentations from within its label class across all domains.
A constrastive loss with such augmentations will essentially reduce to the supervised contrast loss [SupCon, 28].
In particular, for a single sample in a batch, all samples in the batch with the same labels can be used as the positives (could come from the same domain or different domains) and others as the negatives.
In
Khosla et al.
[28], two variants of SupCon loss were introduced for solving the issue of multi-positives depending on whether the summation over multi-positives was located inside (SupCon-In, Eq.
(3) in
Khosla et al.
[28]) or outside (SupCon-Out, Eq.
(2) in Khosla et al.
[28]) the log.
Though Khosla et al.
[28] chose SupCon-Out because it worked better than SupCon-
In, we hypothesized that this is because SupCon-Out has an implicit bottleneck effect.
Intuitively, SupCon-Out upper bounds
SupCon-In and achieves its optima only if the logits with positive samples are all the same by Jensen’s inequality, which may encourage positive samples from different domains to get clustered.
Since this might confound with the effect of our bottlenecks, we chose to use SupCon-
In though it performed slightly worse in out initial experiments.
For the implementation of SupCon, we followed Khosla et al.
[28] except that no projection was used.
Speciﬁcally, the temperature was set to 0.1, and normalization was applied when computing the logits.
In the scientiﬁc setting, we tried to simulate our theory to the greatest extent.
In particular, we had two special considerations as detailed below:
Eliminating empirical generalization As our theory focuses on the idealized domain generalization that assumes access to population distribution, we considered the setup where the empirical generalization was eliminated.
Speciﬁcally, we treated the dataset as the population distribution and used the same dataset for training the encoder and training/evaluating the predictor.
The ResNetencoder was trained to 300 epochs without any regularization, using the Adam optimizer [29] with a learning rate of 5e-5, a batch size of 192 (48 for each domain), and a cosine learning rate decay schedule.
Worst-case approximation To approximate the worst-case source predictor, we included the target data with randomly assigned wrong labels to the training set for training the source predictor.
The target data samples were down-weighted with a sample weight that maximizes the target risk while keeping the source risk close to optima (which is 0).
We selected the sample weight by sweeping over [10 − 10 , 1] with a logarithmic scale using CE-Base and SupCon-Base, as shown in Fig.
As the sample weight increases, the target log likelihood (neg. risk) ﬁrst decreases and then increases.
We hypothesized that the increasing trend was due to that the source performance was already not optimal (though not visible from the ﬁgure), thus we selected the weight close to the turning point and 10 −seemed to be reasonable for both CE-Base and SupCon-Base.
Although we did not adaptively select the sample weight for each setup due to the computational cost, the pre-speciﬁed sample turned out to be reasonable for all other losses and different lambda combinations.
Furthermore, we also removed regularization when training the linear classiﬁer and initialized the linear weight i.i.d.
from N (0 , 1) .
Next, we provide other experimental details for reproducibility:
Implementation of standard augmentations We followed SimCLR [11] for implementing standard image augmentations.
For a fair comparison between the cases when using standard augmentations (SimCLR) and supervised augmentations (SupCon), we kept the total batch size the same and also used the same conﬁgurations for computing the SupCon loss, i.e., temperature set to 0.1, no projection, and normalization applied.
Details of Fig.
4c In Fig.
4c, we considered different choices of augmentations.
The ‘Standard‘ augmentation implementation is described above (Appx.
E.1).
The ‘Supervised’ augmentation was essentially implemented using the SupCon loss as described in Appx.
E.1.
For other augmentations considered, we implemented them by dropout inter-domain supervised augmentations in SupCon.
Speciﬁcally, for each sample in the batch, we randomly masked the samples from different domains (i.e., both inter-domain positives and negatives) i.i.d.
with the speciﬁed dropout probability, while samples within the same domain were always kept. ‘IntraDom’ and ‘ApproxDC’ correspond to dropout probability 1 and 0.9, respectively.
‘SingleDom’ were implemented by dropout all interdomain samples with probability 1 except for a ﬁxed domain (the ‘A’ domain of PACS in our case).
E.2 Bridge: understanding how to learn optimal representations in practice
In the bridge setting (see Appx.
F.2), we aimed to bridge the gap between our theoretical setup to the practical setup.
The main differences from the scientiﬁc setups are that the empirical generalization gap is considered and the average-case source predictor is used, as detailed below:
Incorporating empirical generalization In practice, empirical-generalization gap should also be considered besides the source-target generalization gap.
Thus, we randomly split the PACS dataset to 80% training and 20% validation splits for each domain.
The training splits were used to train both the encoder and the source predictor, and the validation splits were used for encoder and source predictor selection as well as evaluation on target domains.
We used the ResNet-50 model as the encoder and initialized it from ImageNet pretrained model.
The encoder was trained to a maximum of 50 epochs with a 1e-5 weight decay, using the Adam optimizer [29] with a learning rate of 5e-5, a batch size of 112 (28 for each domain), and a cosine learning rate decay schedule.
Using average-case source predictor Instead of approximating the worst-case source predictor in the scientiﬁc setting, we considered the average-case 2 source predictor which is closer to the common practice.
Speciﬁcally, we freezed the encoder and trained a SVM classiﬁer with L2 regularization on the source training split.
The regularization parameter was tuned over {1e-4, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2, 1e3} with the source validation accuracy.
Next, we provide other experimental details for reproducibility:
Selection of λ
For all different setups considered in bridge settings, the CAD bottleneck was used and the λ was tuned over {1e-3, 1e-2, 1e-1, 1, 1e1} independently for each.
Staggered training setup We ﬁrst trained model without any bottlenecks (i.e., with CE-Base) in the same way as described before.
Then we freezed the trained model, added a 1-
E
Approximating optimal representations
with
pretrained SSL
Datasets We used non-MNIST datasets on DomainBed that were non-synthetic, including VLCS [16], PACS [34], OfﬁceHome [54], TerraIncognita [6], and DomainNet [41].
For each dataset, we split it to 80%/20% training/validation set according to DomainBed.
SSL-based models & Training For all models based on pretrained SSL models (either CLIP-based or DINO-based) with staggered training in this experiment
, we freezed the pretrained SSL model and added on top a 1-layer MLP with hidden size 1024, and residual connection.
We used CLIP
ResNet-50 (CLIP S) to obtain the best possible fair comparison with baselines from DomainBed, and CLIP ViT-B/32 (CLIP L) to achieve the best results.
Note that the ResNet-50 model of CLIP
S was modiﬁed as described in Radford et al.
[43] and contained 38M parameters (more than 23M of the original CLIP).
The model was trained to 300 epochs for DomainNet and 50 epochs on other datasets (an epoch is deﬁned as a single pass over the smallest domain according to DomainBed).
No data augmentation was used and the temperature for scaling the logits in CAD was ﬁxed to 0.05.
We used the Adam optimizer with a 1e-5 weight dacay, and a cosine learning rate decay schedule.
The hyperparameter search space is:
• Learning rate: discrete set {1e-4, 3e-4, 1e-3, 3e-3} • Batch size: discrete set {128, 256, 512} for DomainNet and OfﬁceHome, and {64, 128, 256} for other datasets • MLP dropout: discrete set {0., 0.1, 0.5} • Learning rate warmup: discrete set {True, False}
End-to-end models & Training In Table 4, we also included an end-to-end trained model without any pretrained SSL models.
We used exactly the same model architecture (the original ResNet50, initialized from ImageNet pretrained model), training procedure and evaluation protocal as baselines on DomainBed.
Importantly, the linear classiﬁer was jointly trained with the encoder, and no reﬁtting was applied.
The model was trained to a maximum of 5000 steps on each dataset, and data augmentations were applied.
The Adam optimizer was used without any particular learning rate schedule.
The hyperparameter search space is (same as DomainBed except we added the temperature):
• Learning rate: log-uniform over [1e-5, 1e-3.5] • Batch size: log-uniform over [8, 64] for DomainNet, and [8, 2 5 . 5 ] for other datasets • MLP dropout: discrete set {0., 0.1, 0.5} • Weight decay: log-uniform over [1e-6, 1e-2] • Temperature: discrete set {0.05, 0.1}
Linear Probe Evaluation
In all the experiments except for the end-to-end training setup, we always followed the procedure of two-stage training, where we ﬁrst trained the encoder with speciﬁed objectives, and then reﬁt the classiﬁer.
For datasets except DomainNet, we ﬁtted the SVM classiﬁer and tuned the regularization parameter over {1e-4, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2, 1e3} with source validation selection.
Since DomainNet was too large and SVM cannot ﬁt it efﬁciently, we used the logistic regression classiﬁer which was trained with a batch size 512, the Adam optimizer with a learning rate 5e-4 and early stopping.
Note that an alternative was to just use the linear head ﬁtted when training the representor (as we used CE loss with source labels), and we found this could work better than reﬁtting since the classiﬁer was less overﬁtted to the source domain.
However, we didn’t do that since we wanted to stick to the representation learning protocol with two-stage training.
We did that in our end-to-end training setup since we wanted it to be compeletely comparable to baselines on DomainBed (which did not do reﬁtting).
Selection of λ
In our experiments, we treated λ as a special hyperparamter.
For each model, we selected λ on the PACS dataset, and then used the same λ value for all other datasets, because our bottleneck was fairly robust to the choice of λ .
The λ values chosen for SSL-based models (i.e., CLIP
S, CLIP L, DINO) and end-to-end ResNet-50 were 1e-1 and 1e-5, respectively.
E.4 Towards generic robust representations with SSL
Model We used the CLIP L model (i.e., CLIP ViT-B/32) with an additional network on top for staggered training.
The additional network were two blocks of 2-layer MLP, each with hidden size 2048, pre-activation batch normalization, residual connection, and dropout probability 0.1.
Note that the original CLIP L model was frozen and only the additional network was trained.
Dataset We used the LAION-400M dataset which contained 400 million image-text pairs for training.
Though the dataset might not be as clean as the original CLIP training data (as evidenced by our experimental results), it was the largest publicly available image-text-pair dataset that we could get access to.
As we froze the CLIP L model and only did staggered training, we used the 1TB preprocessed embeddings provided by LAION-400M 3 .
No further preprocessing was applied.
Training We used the image-text contrastive loss as introduced in Radford et al.
[43] for training model.
The temperature was learnable which was initialized as 0.07 and clipped with a minimum 0.01.
The model was trained for 1 epoch using the Adam optimizer with a batch size of 16384 and a cosine learning rate decay schedule.
The learning rate was tuned over the set {3e-5, 1e-4, 3e-4, 1e-3, 3e-3, 1e-2} and the λ value for the Ent bottleneck was tuned over {1e-3, 1e-2, 1e-1, 1, 1e1}.
Evaluation For the evaluation on the ImageNet-related datasets, we followed a similar procedure in
Radford et al.
[43], where a linear classiﬁer was ﬁtted on ImageNet using the model representations and evaluated on 7 natural distribution shift datasets.
In particular, we ﬁtted a logistic regression classiﬁer with 1e-5 L2 regularization on ImageNet training set which was trained with a batch size 512, the Adam optimizer with a learning rate 3e-4 and early stopping.
Note that this was different from Radford et al.
[43], where a logistic regression classiﬁer was ﬁtted using full-batch data with decent hyperparameter tuning, due to our computational budget.
For evaluation on natural distribution shift datasets, we followed Taori et al.
[51] and used their released testbed 4 .
The evaluation datasets and their abbreviations used in Table 6 were: ImageNetV2 [IN-V2, 44], ImageNet-Sketch [IN-S, 55],
Youtube-BB [YT-BB, 46], ImageNet-Vid [IN-Vid, 46], ObjectNet [5], ImageNet Adversarial [IN-A, 26], and ImageNet Rendition [IN-R 25].
F Additional Experimental Results
In our experiments, we aimed to: (i) verify our theoretical results in practice; (ii) investigate our proposed representation learning objectives in practical DG; (iii) take advantage of pretrained SSL models (in particular, CLIP) to achieve powerful models for DG.
Unless stated otherwise, we consider a two-stage training setup.
First, the representation learner (“the representor”) trains an encoder p Z | X using a speciﬁed objective and freezes it.
Then, the person performing predictions (“the learner”) trains her predictor h from Z by minimizing the risk on source data.
Finally, the representation Z and predictor h are evaluated on target data.
In all experiments, the learner uses a linear classiﬁer for h .
For the Ent bottleneck, we used Ballé et al.’s
(2018) entropy model.
For the CAD bottleneck we used its conditional version whenever labels were available.
When a model contains no domain bottleneck, we label it as “Base”.
For experimental details, see Appx.
E.
F.1 Scientiﬁc: exploring optimal representations for worst-case DG
To validate our theory, we studied optimal representations in a scientiﬁc setup that is as close to our IDG framework as possible with log-loss (cid:96) .
In particular, we used the PACS dataset [34] and approximated the idealized DG by treating the dataset as the population distribution, i.e., we did not split datasets into train and test sets.
To approximate the worst-case source predictor, we followed
Dubois et al.
[14] by incorporating the wrongly labeled target data to the source domain.
The experimental setup goes as follows: (i) the representor trains a ResNet-18 [23] to minimize the objective on labeled data from all domains; (ii) the learner trains a worst-case source classiﬁer h on every possible pair of (source, target); (iii) the negative target risk (log likelihood) for each h is evaluated.
We reported the log likelihood averaged over 5 seeds.
For more realistic scenarios (i.e. non-idealized average-case DG) see Appx.
F.2 which replicates the following results.
Base
Ent
CAD
R [ Y | Z ]
H
[ A | Z ] -5.1 ± 0.3 -0.4 ± 0.1 -0.7 ± 0.
-3.5 ± 0.1 -0.0 ± 0.0 -0.8 ± 0.
(a) Effect of different objectives 10101010
L o g li k e li h oo d source target (b) Effect of λ
DC
Approx.
Non-DC Augmentation type - - - - -
L o g li k e li h
oo d
Supervised SingleDom
ApproxDC IntraDom
Standard
[Figure]
Do our domain bottlenecks improve worst-case DG?
In Fig.
4a, we compare IDG performance of representations trained with (Ent, CAD) and without (Base) domain bottlenecks.
We see that both bottlenecks signiﬁcantly improve the worst-case DG, and nearly achieve the source-domain performance (0 log likelihood).
This shows the importance of support match (Thm. 2) and the effectiveness of our bottlenecks to enforce it.
In Appx.
F.2, we show that bottlenecks also helps in practical scenarios, i.e., non-idealized average-case DG evaluated with accuracy ( 95 . 9% → 96 . 7% ).
What is the effect of λ ?
Fig.
4b shows the effect of the bottleneck weight λ on the worst-case target and source performance.
We see that increasing λ will decrease the DG gap.
As a result the target performance improves until λ ≈ 10 2 , where source performance starts to decrease.
What if the representor has access to domain-covering augmentations instead of labels?
In
Sec. 4, we provide a contrastive objective for using augmentations.
To show the effectiveness of the objective, we compared minimizing H
[ A | Z ] using Eq.
(121) to standard supervised risk minimization
R [ Y | Z ] .
We ensured domain coverage by using supervised augmentations (Fig. 2c).
The 1 st and 2 nd row of Fig.
4a show that our objective performs similarly to direct label prediction.
How important is the choice of augmentations?
Prop. 2 shows that domain-covering (DC) augmentations are sufﬁcient for achieving IDG, but it does not give necessary conditions.
Here we investigate the effect of using our loss with different choices of augmentations.
Speciﬁcally, we used L CAD with ﬁve augmentations.
The ﬁrst two are DC.
‘Supervised’: augment inputs inside the label class across all domains as in Fig.
2c; ‘SingleDom’: augment inputs to same label samples from a ﬁxed domain.
The second two are not DC.
‘Standard’: standard SSL augmentations [11] as in Fig. 2b; ‘IntraDom’: augment inputs to same label and same domain samples.
Finally, we consider ‘ApproxDC’, which is approximately DC by augmenting 10% of the time with ‘Supervised‘ and 90% of the time with ‘IntraDom‘.
Fig.
4c shows that the non-DC augmentations give terrible results compared to DC.
Interestingly, ‘ApproxDC’ also performs very well, which suggests that approximately DC augmentations might be sufﬁcient to learn optimal representations in practice.
What if the representor does not have access to target domains?
Prop. 1 shows that DG without access to target domains is generally impossible.
We empirically veriﬁed this by excluding a predeﬁned target d t domain from the representor’s training set, i.e., L CAD is optimized on 3 of thedomains.
The learner then trains a predictor h on each source.
We ﬁnally evaluate each h on the target domain d t , and average over choices of d t .
The resulting worst-case log likelihood was − 4 .
2 ± 0 .
2 , which is signiﬁcantly worse than when the representor had access to all domains ( − 0 .
8 ± 0 . 2 ).
What’s the effect of λ for different objectives on the worst-case DG performance?
In Fig.
5, the worst-case target log likelihood versus λ values for different objectives is shown.
We found that
Ent is much more sensitive to the choice of λ than CAD, which was part of the reason why we used the latter in most of our experiments.
Note that for SupCon-Ent with small λ values, it was worse than SupCon-Base because of the discretization introduced by the Ent bottleneck, which we veriﬁed by observing that setting λ = 0 lead to similar results.
[Figure]
The scientiﬁc setup is closer to our theory than what we do in practice in that worst-case predictor was considered and empirical generalization gap was ignored.
Here we bridged these gaps with a more practical setup.
In particular, we split the PACS dataset to training and validation splits for each domain and considered the setting: the representor trains the encoder on all-domain training splits with a validation loss selection; the learner trains the SVM predictor (average-case) on the source training split which is selected over the source validation split, and evaluates on the validation splits of other target domains.
The target validation accuracy averaged over all (source, target) setups was reported.
For simplicity, we will use CE to denote the objective with the cross-entropy loss that uses labels to minimize R [ Y | Z ] , and SupCon for the contrastive loss that uses supervised augmentations to minimize H
[ A | Z ] .
We will use CAD in following experiments unless otherwise speciﬁed (chosen with initial experiments).
Details in Appx.
E.2.
[Table]
Does domain bottleneck improve the average-case DG performance?
Though our theory focuses on the worst-case DG, we empirically showed that adding bottlenecks to enforce support match can also improve the average-case DG performance by comparing CE-Base and CE-CAD in Table 3.
What if the representor only has access to source domains?
Similar to what we did in the scientiﬁc setting, we considered the setup where one single domain is speciﬁed as the target domain and excluded from the training set of the representor and used for evaluation with source predictors trained on other domains.
This is denoted as CE+CAD (partial domains) in Table 3, which is much worse then CE-CAD.
This shows the necessity of getting access to target domain information for DG.
What if the representor only has access to domain-covering augmentations?
In Table 3, we also compared SupCon-CAD which used supervised augmentations through the labels with CE-
CAD and they achieved the same performance.
This shows that the representor can still learn good representations without labels but only domain-covering augmentations in practice.
Can we use standard augmentations?
In Fig.
2, we point out that standard augmentations are not domain-covering and thus not suitable for SSL with our objectives.
We empirically showed this by using augmentations of SimCLR (see Appx.
E.1 for details) with our objectives (SimCLR-CAD).
In
Table 3, we indeed observed that using standard augmentations performed much worse than using desired augmentations (SupCon-CAD).
How do augmentations matter?
Besides investigating the ‘Supervised’ augmentations (SupConCAD) and ‘Standard’ augmentations (SimCLR-CAD) above, we also compared other three augmentations as in the scientiﬁc section.
Speciﬁcally, we considered the ‘SingleDom’, ‘IntraDom’, and ‘ApproxDC’ augmentations.
As shown in Table 3, SupCon-CAD (SingleDom) and (ApproxDC) maintained the DG performance but SupCon-CAD (IntraDom) was slightly worse (0.5 accuracy drop).
We assumed the small gap was due to the speciﬁc dataset that we used (PACS).
We did the same analysis on VLCS, and SupCon-CAD with ‘Supervised’, ‘SingleDom’, and ‘IntraDom’ augmentations gave 84.7 ± 0.4, 83.2 ± 0.3, and 77.5 ± 2.3, respectively.
This shows the importance of using domain-covering augmentations in practice.
How is end-to-end training compared to staggered training?
Since getting access to target domain data is not practical in practice, we can utilize SSL models pretrained on a very large support (as we did in practical and realistic settings).
However, typical SSL models are trained without any bottlenecks and we can adopt the staggered training where we freeze the SSL models and train a small network on top with bottlenecks.
To get an idea of how it compares to end-to-end training with bottlenecks, we considered a staggered version of CE-CAD in our setup.
As shown in Table 3, it performed similarly to CE-CAD.
Do standard augmentations affect source performance?
Previously, we showed that using standard augmentations hurt the DG performance measured by the average target accuracy.
It is natural to ask whether using standard augmentations also hurt the source performance since we should also be interested in the ‘effective robustness’ [51].
Thus we also reported the average source accuracy of
SupCon-CAD and SimCLR-CAD which were 96.9 ± 0.2 and 90.1 ± 0.2, respectively.
The source performance using standard augmentations was indeed worse, but if we consider the source-target gap which was 0.2 for SupCon-CAD and 28.4 for SimCLR-CAD, which still veriﬁed that the nondomain-covering standard augmentations were harder to force support match.
To be even more convincing, we did the same analysis on VLCS, and the average source accuracy of SupCon-CAD and SimCLR-CAD were 86.6 ± 0.1 and 84.6 ± 0.5 which were fairly close, but the average target accuracy were 84.7 ± 0.4 and 57.5 ± 1.7, respectively.
F.3 Approximating optimal representations with pretrained SSL
In this experiment, we used the standard DomainBed benchmark (with non-MNIST datasets) and protocol [22].
In particular, we left out a target domain for evaluation and used the union of other domains for training both the encoder and the classiﬁer.
Contrary to our scientiﬁc setting, the representor does not get access to the target domain.
All our representations were evaluated by ﬁtting a linear classiﬁer on source domains with source validation selection.
As in DomainBed we selected the encoder based on ‘oracle selection’ over 10 hyperparameters, and reported the target accuracy averaged over all choices of targets and 5 random seeds.
Details in Appx.
E.3.
We included the full result of Table 2 with all baselines on DomainBed as in Table 4.
We considered most representative baselines from DomainBed, most of which considered learning invariant representations or optimal classiﬁers across domains.
Speciﬁcally, we included IRM [2], GroupDRO [45],
Mixup [57], CORAL [49], MMD [35], DANN [17], CDANN [36], and VREx [31].
We also included the result pretrained CLIP S model with a zero-shot classiﬁer using text representations (CLIP S Zero
Shot), which demonstrated better DG performance than CLIP S with linear probe.
But we observed that it was outperformed by our CLIP S + CAD.
Can we approximate optimal representations by exploiting pretrained CLIP?
The last row in
Table 4 shows that ﬁnetuning a large pretrained CLIP model with our CAD achieves SOTA on nearly
[Table]
all DomainBed benchmarks by a very large margin (see 2 nd row).
Note that the poor performance on
TerraIncognita is likely because CLIP’s dataset did not cover such images (camera traps monitoring animals).
In Appx.
F.2, we estimated the non-idealized DG performance of optimal representations on PACS (with access to all-domain labeled data) to be 96 .
7% , which is only 2% higher than CLIP L + CAD.
This suggests that our simple SSL encoder might already be close to optimal.
Are gains due to the architectural differences?
DomainBed’s baselines ﬁnetuned an ImageNet [12] pretrained ResNet-50.
In contrast, CLIP L pretrained a larger ViT. To decouple gains due to our objective from architectural gains, we evaluated ResNet-50 pretrained.
Table 4 shows that CLIP S still outperforms DomainBed baselines.
Our theory does not constrain the encoder and so we expect larger encoders to be better.
Table 4 shows that CLIP L indeed
outperforms CLIP S.
What is the effect of domain bottlenecks?
In the last six rows of Table 4, we investigated the effect of ﬁnetuning with our CAD bottleneck.
We see that for both CLIP L and CLIP S, it improves results by around 1 ∼ 2% .
These gains are due to the bottleneck, rather than due to the additional
MLP trained on source data as seen by ‘CLIP S + Base’.
Note that the raw CLIP S already signiﬁcantly outperforms baselines.
We hypothesize that this could be because SGD training of neural networks favors support match, e.g., by minimizing I[ X ; Z ] as suggested by Shwartz-Ziv & Tishby [48].
Which pretrained SSL model to use?
Our theory suggests that we can exploit pretrained SSL models as long as their augmentations are domain-covering and their training set covers desired domains.
We investigated the effect of adapting SSL models that do not satisfy those properties by ﬁnetuning DINO [10], the current SOTA on SSL ImageNet.
DINO only pretraiend on ImageNet using standard augmentations.
As a result, Table 4 shows that the ﬁnetuned DINO+CAD signiﬁcantly underperforms compared to CLIP S and DomainBed baselines.
What is the impact of CLIP pretraining?
To ensure that our gains are not only due to a novel
CAD bottleneck, but the synergy between enforcing support constraint and using desired SSL models, we investigated CAD using the standard DomainBed protocol denoted as CAD in the table.
It shows that CAD on its own performs similarly with DomainBed baselines (see Table 4 for a full comparison).
tends to overﬁt more to the source domain and diminish the effect of bottlenecks.
[Table]
Why ‘oracle’ selection?
In the main body
[Table]
we provided the results with ‘oracle selection’ which was the closest to our theory among the model selection methods in DomainBed (in the sense that we needed target domain information to achieve IDG).
Here
[Table]
we also provided results with ‘source validation’ selection in Table 5.
Source validation selection relies on the assumption that source and target data follow similar distributions [22] thus source and target accuracy are highly correlated
[Table]
which is not really true in practice.
We found some issues with source validation selection results:
• The selected model with the highest source validation accuracy tends to overﬁt the source domain, thus possibly leads to worse performance on the target domain.
This can be probed by the fact that the staggered trained CLIP models (CLIP + Base or CLIP + CAD) were generally worse than the original CLIP model; • Selecting model with source validation accuracy tends to diminish the effect of bottlenecks.
This can be seen by the fact that the gap between CLIP + Base and CLIP + CAD of source validation selection is much smaller than that of oracle selection; • The source accuracy is not a good indicator of target accuracy thus its result has a larger variance.
F.4 Towards generic robust representations with SSL
In the previous section, we ﬁnetuned CLIP in a task speciﬁc fashion by optimizing R [ Y | Z ] and our
CAD bottleneck.
To get generic (task agnostic) robust representations, one should instead directly use our objectives on a sufﬁciently large dataset with image-text augmentations.
Unfortunately, we cannot fully train CLIP with our bottlenecks as we do not have access to CLIP’s original dataset and sufﬁcient compute.
In this section, we aim to emulate such training of generic robust representations.
To do so we used LAION-400M [32] which is a public dataset that contains 400M web-crawled image-text pairs.
Due to our computational budget, we again froze the pretrained CLIP L and only ﬁnetuned an additional MLP with our L Ent .
We used L Ent as it only requires access to paired image X and text A but no prior information about domain D .
As in CLIP’s paper, we evaluated the learned representation Z in Taori et al.’s
(2020) realistic setting, where a linear classiﬁer h from Z is trained on ImageNet and tested on 7 natural distribution shift datasets.
Details in Appx.
E.4.
Would training CLIP with a bottleneck have improved its robustness?
As shown in the lastrows of Table 6, ﬁnetuning CLIP L on LAION with L Ent (LAION + Ent) outperforms ﬁnetuning without bottleneck (LAION + Base) on all 7 distribution shift datasets.
This suggests that directly training CLIP with our Ent bottleneck would improve the robustness of learned representations.
We hypothesize that the gains could be larger if SSL models trained L Ent end-to-end.
In Table 7, we show similar results on DomainBed, where we followed exactly the same linear evaluation protocal discussed in Appx.
E.3.
Note that both models underperform the original CLIP L, likely due to non-end-to-end training and LAION data with (possibly) lower quality than CLIP’s data.
[Table] [Table]