Optimal Representations for Covariate Shift
Authors: Yangjun Ruan
Yann Dubois
Chris J. Maddison
Venue: ICLR 2022 Poster (2022)
Keywords: distribution shift
domain generalization
representation learning
self-supervised learning
invariance
robustness

Abstract: Machine learning systems often experience a distribution shift between training and testing. In this paper, we introduce a simple variational objective whose optima are exactly the set of all representations on which risk minimizers are guaranteed to be robust to any distribution shift that preserves the Bayes predictor, e.g., covariate shifts. Our objective has two components. First, a representation must remain discriminative for the task, i.e., some predictor must be able to simultaneously minimize the source and target risk. Second, the representation's marginal support needs to be the same across source and target. We make this practical by designing self-supervised objectives that only use unlabelled data and augmentations to train robust representations. 
Our objectives give insights into the robustness of CLIP, and further improve CLIP's representations to achieve SOTA results on DomainBed.

Published as a conference paper at ICLR 2022
OPTIMAL REPRESENTATIONS FOR COVARIATE SHIFT
Yangjun Ruan" Yann D Chrris J. Maddison 12 University of Toronto & Vector Institute
(yjruan, yanndubois, emaddi s (css toronto..dddt
ABSTRACT
Machine learning systems often experience a distribution shift between training and testing. In this paper, introduce simple variational objective whose optima are exactly the set of all representations on which risk minimizers are guaranteed to be robust to any distribution shift that preserves the Bayes predictor, c.g.. covariate shifts. Our objective has two components. First, a representation must remain discriminative for the task, i.e., some predictor must be able to simultaneously minimize the source and target risk. Second, the representation's 's marginal support needs be the same across source and target. We make this practical by designing self-supervised objectives that only use unlabelled data and augmentations to train robust representations. Our objectives give insights into the robustness of CLIP, and further improve CLIP's representations achieve SOTA results on DomainBed.
INTRODUCTION
hard build machine learning (ML) systems that are robust to distribution shifts between a source (train) and target (test) domain. One promising approach to domain generalization (DG) learning robust representations from which predictors trained on source must perform well on target. In practice, however, no current DG methods for learning representation uniformly outperform empirical source-risk minimizers (ERM) (Gulrajami & Lopez- Paz 2021). Furthermore, our theoretical understanding of DG is still lacking. Specifically, while previous work have studied properties that would or would not imply robust representations (Ben- -aaid et al.) 2007] 2010a Zhao et al. 2019) Johansson al. 2019) the minimal set of achievable requirements for perfect DG is not yet known.
We introduce the first, simple, variational objective whose optima are exactly the set of all represen- tations on which source risk minimizers are guaranteed generalize across distribution shifts that preserve the Bayes predictor. We work in an idealized DG (IDG) setting; we assume that a learner has access the source population risk. Our variational characterization implies that is both sufficient and necessary for optimal IDG that a representation: (a) remains discriminative for the learning task, i.e., there must exist predictors from the representation to the labels that can simultaneously minimize both source and target risk; and (b) keeps the support of its marginal distribution invariant to shifts.
This means that any optimal representation learning method must seek discriminative information about the target. Even worse, we prove that without access to some knowledge about the target, any representation learning algorithm cannot uniformly (over target domains) outperform a constant representation, which may explain why DG methods struggle to outperform ERM.
We show, in theory and practice, how to overcome these challenges using only a large set of unlabeled examples and particular data augmentations that retain all discriminative information but minimal domain-speciffe information. Text descriptions of images are examples of such augmentations, as they informative for many downstream classification tasks, but they remove lot of domain-specific information. With such augmentations, we design practical sulf-ssue learning (SSL) objectives for learning robust representations. Our objectives give insights into the robustness of CLIP (Radford) et al. 202T) over other SSL methods, and lead to improved CLIP--based representations that achieve of-te-of-thh-art (SOTA) results on DomainBed (Gulrajani & Lopez-Paz 2021 To summarize, we:
provide minimal sufficient objectives whose optima achieve optimal DG under covariate shift; prove that it impossible learn useful representations without accessing target information; provide practical objectives to learn optimally robust representations using specific augmentations; obtain SOTA results on typical domain generalization benchmarks!
AAuthors contributed equally.
Our implementation is released at

---

Published as a conference paper at ICLR 2022
2 BACKGROUND: DOMAIN GENERALIZATION AND REPRESENTATIONS
We are interested in predictions that are robust across distribution shifts. We formalize this using domain generalization (DG) language. Given distribution Px,y Id. over inputs E X and labels y E ) from the source domain d, E 'D we select a predictor f : X + r. The predictions E could for example be labels distributions over labels. Despite being selected the source domain, we would like f to achieve small expected risk with respect to a loss function / y x + R20,
R Y X] == Exx.r ((Y, (X))),
on a distribution Px,y d from a target domain d = de E D, which somehow related to dg.
A common strategy for DG is to learn robust representations, which splits the problem into two. First, learn an encoder Pz x, which maps inputs X to representations Z. Then, learn a predictor h : 2 + I from representations Z to labels Y using standard risk minimization. The goal is to design robust representation Z, so that predictors h trained to minimize the source risk R [Y] Z also achieve low target risk R [Y] Z]. Many methods have been proposed to try to learn such Z. e.g., by enforcing domain invariance of the marginal P2la a (e.g.. Ganin et al. 2016). Still, many of these proposals are not sound (Zhao al. 2019 Johansson et al., 2019). Furthermore, they rarely outperform source empirical risk minimization (ERM) in practice
Gulrajani & Lopez- Paz 2021).
OPTIMAL REPRESENTATIONS FOR DOMAIN GENERALIZATION
To separate domain generalization from finite sample generalization, we consider an idealized DG (IDG), where the predictor h is selected on the source population risk rather than empirical risk. We assume sample spaces X' 2, D are discrete; formal statements and proofs are in Appxs. A and
3.1 DEFINING OPTIMAL REPRESENTATIONS FOR IDEALIZED DOMAIN GENERALIZATION
We want to evaluate the quality of a representation Z of X. In our IDG, the learner is given random source Da; she selects any source risk minimizer; and is scored according to her risk on a random target domain D.. To give uniform guarantees while reflecting the uncertainty over the souree-target pair (D., D.), we measure the quality of Z as the expected risk of the learner's worst-case choice.
The idealized domain generalization risk (IDG risk) of an encoder pz X is the expected (over domains) worst-case (over source risk minimizers) target risk, i.e.,
Definition.
REn Y [Z] := W. Epp,,.
sup KEHb,
R [Y z
where H6, arg min, R, [Y| Z] are the source risk minimizers, and PD..D, is any joint distribu- tion that has full support over 'D x D. We call a representation 2" (or its encoder) optimal for IDG if it minimizes the IDG risk.
3.2 CHARACTERIZING OPTIMAL REPRESENTATION FOR IDG UNDER COVARIATE SHIFT
The IDG risk is useful to evaluate representations, but gives few insights into IDG and is impractical to optimize due the supremum in Eq. (2) Under mild assumptions, we provide a simplified, equivalent objective, which is easier to optimize. For convenience, we assume that there is a unique Bayes pre- dictor f*, which minimizes the expected risk over domains, i.c., f = arg min, y Eps, [R. [Y X]]. This is satisfied by standard ML tasks Py.x and losses é. More importantly, we assume the following domain structure, which ensures the existence of optimal encoders and allows our simplification.
Assumptions. All domains d E DD we consider are related by the following assumptions:
1. Generalized covariate shift. All domain-specifie risk minimizers f E arg min, (R² [Y |X] equal the Bayes predictor f" on their support, i.e., f(x) (x) for all supp(px a)-
Invariance of Bayes predictions. The set of Bayes predictions is the same for all domains, i.c., (x) x E supp(px w)} (x) x E x).
2

---

Published as a conference paper at ICLR 2022
source P2 d,
target P2(4,
(a) discriminative & support match
(b) only support match
(c) only discriminative
Figure 1: (a) Optimal representations for IDG must have invariant supports while being simultaneously discriminative on all domains: (b) without the discriminative requirement, source-risk minimizer can mispredict the target, and (c) without support match, some risk minimizer will perform poorly.
Generalized covariate shift (GCS) ensures that I* is simultaneously optimal on all domains. For log-loss ft recovers standard covariate shift, i... Jr.d = | 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. Invariance of Bayes predictors necessary to learn useful predictors using a single domain. For example, for 0-1 loss it ensures that each label is seen at least once in each domain.
The intuition behind our objective is that under GCS any source risk minimizer will make optimal predictions target samples that are also in the source. Thus, IDG optimal representations exactly those that (a) have the same support in Z for all domain, and (b) retain GCS from Z without sacrificing the ability predict Y, which be ensured by minimizing the risk from Z. See Fig.
Theorem 1. Under our assumptions, an encoder pz- xx is optimal for IDG if and only if minimizes the risk R [Y| Z] = infn Ep,, P [Y]
Z]] while matching the support of Z across domains, i.e.,
pz- |X E arg min R Z] s.t. EdED, supp(pz a) = =supp(p2)
(3)
P2IX
Moreover, such encoders exist and their IDG risk the Bayes risk RypG YY - = R XX.
Theorem I provides objective to learn representations on which performing risk minimization using a single domain and Z is as good as performing risk minimization on the target domain from inputs X. Other sufficient conditions have previously been hinted towards, e.g., matching the marginal pz instead of its support -ev-DDii et al. 2010a) which is the focus of most DG methods (e.g.. Ganin al. 2016). Note that previous conditions are nevertheless generally not necessary and could be too stringent to be achievable. To our knowledge, [Thm. 1 is the first characterization of necessary and sufficient conditions, which gives better insights into the essential goal for optimal IDG and provides guide for deriving the least stringent objectives in practice.
The risk minimization (Bq. (3)) shows that one must have some knowledge about the target domains to learn optimal representations for IDG. Access to targets might seem unrealistic, but without such knowledge additional assumptions provably impossible beat even constant representations.
Proposition 1 (No free lunch for IDG). Let de be any source domain, Za, be any representation chosen on source de, and C. E Z be a constant representation. Under minor assumptions, for every "good" target domain outside the source's 's support on which Za, outperforms C. for IDG, there are many "bad" target domains on which Za, is strictly worse than C. Formal statement in Appx. B.3
Proposition I shows that target knowledge is necessary for learning useful representations in IDG. This may explain why previous DG methods have been unable outperform ERM in standard benchmarks (Gulrajani & Lopez-Paz, 2021): the knowledge they have access to is insufficient to generalize. Taken together, orgepher.IProp. and Thm. say that either you have access to target domains d, in which case you can achieve an IDG risk that matches supervised learning, or you do not access de which case any representation learning algorithm can achieve worse IDG risk than a constant.
LEARNING REPRESENTATION UNDER COVARIATE SHIFT
4.1 SELF-SUPERVISED LEARNING USING DOMAIN-AGNOSTIC AUGMENTATIONS
Our characterization of optimal representations for IDG TTm. I) requires labeled data from all domains, which is impractical. We show how this can be overcome with self-supervised learning
3

---

Published as a conference paper at ICLR 2022
"A dog with floppy ears."
AA pointy- dog."
(a) standard augmentations
(b) supervised augmentations
(c) image-text augmentations
Figure 2: Image-text augmentations are practical domain- agnostic augmentations. Arrows denote augmenters. Bubbles denote inputs that have the same representations, as induced by predicting the augmentations. (a) Standard augmentations are not domain domain-agnostin agnostic. (b) Supervised augmentations uniformly augment inputs inside their label class, irrespective of domains. (c) Image-text augmenta- tions are (nearly) domain-agnostic as they map images across domains to similar descriptions.
(SSL), which is technique for training representations without direct access labels, and a particular class of data augmentations. E.g, in CLIP, images are augmented with alt-text collected on the internet and invariance is enforced between the representations of the image and its text pair (Radford et al.) 2021). Representations learned like this preserve discriminative information about all downstream tasks Y whose label information is preserved by the augmentation (e.g.. Dubois et al. 202T).
More precisely, an augmentation A is a random variable sampled conditionally from the input X. The key requirement is that augmentations retain task information. Specifically, any samples x; a' E X have the same augmentation conditional PA æ PA| z', then their Bayes predictions must be the same f (z) (z'). With such A, one can learn an encoder that minimizes the risk R [Y Z] by instead maximizing mutual information [A; Z]. Intuitively, if Z has all augmentation information, then it must have information about the conditional PA X. and thus the Bayes prediction (X)
This suggests learning optimal representations for IDG by replacing Eq. (3) with a maximiza- tion of 1[A; ZJ. Unfortunately, fully optimizing I(A; Z] w.r.t. P is not generally possible un- der the support constraint Eq. (3) This can be overcome under a domain-agnostic assumption, which requires that the set of possible augmentation distributions is the same across domains, i.e.,
[PA] x x E supp(px a)} = [PA]= x E X].
Proposition 2. Let PAIX x be a domain-agnostic augmenter. Then any optimal solution pz- 1x of the following objective is optimal for IDG:
Pz. X E arg max I[A 2] S.L. dd D, supp(pz (2) = supp(pz))
(4)
121X
Proposition 2] shows that we can still learn IDG optimal representations without labels if we have access to the right augmentations. How realistic are those augmentations? For 0-1 loss E, the most likely label should be preserved, which is satisfied by standard image augmentations like rotations and color jittering. Those augmentations are nevertheless not domain-agnostic for typical domains (e.g. sketches and photos), since outputs A are correlated with the input's 's domain D. See Fig. 2a
practical choice of augmentation that is nearly domain-agnostic, is a mapping from images to text descriptions, with CLIP (Radford et al. (02T) which uses text-image pairs. Image-text augmentations have many advantages. First, text augmentations preserve label information for many downstream tasks. Second, they are close to being fomain-agnostic, since images from different domains (c.g.. sketches and photos) but similar semantics are often mapped to similar descriptions (Fig. 2c). This gives insights into the open question (Radford et al. 2021) about why CLIP's representations are so robust compared to other SSL methods. Finally, image-text pairs are easy to access in practice given their abundance on the internet. Many other multi-mode augmentations, e.g.. audio-video (Wang et al. (02T). are also likely domain-agnostic and can be explored in practice.
In practice, even the domain information D is usually unknown. One can nevertheless still optimize (Fq.. (4)p by replace the support constraint with stronger one that does not rely on D e.g., minimizing [[; (see Sec. 4.2.2).. This highlights the potential of Prop. 2] if one can find large source of inputs X and domain domain-gnostic augmentations A (c.g., the 400M image-text pairs of CLIP) then one can, in principle, Iearn optimal representations for IDG on any downstream task Y that A preserves.
2 A Athough text descriptions might contain domain information (e.g., referring to "sketch"), they are still much better than standard augmentations that rarely map together images from different domains.
4

---

Published as a conference paper at ICLR 2022
4.2 PRACTICAL OBJECTIVES
We now design practical objectives for learning optimal representations without labels. Proposition 2 does provide objective but it impractical as it involves constrained optimization. We can nevertheless convert it to the following unconstrained objective by using Lagrangian relaxation and introducing a domain bottleneck B Biz, D] that enforces support match,
arg min
[[A: 2] + AB|Z,
(5)
P2
Eq. (5) is valid reformulation of Prop. 2 as long minimizing B B D) while maximizing T[ A[A; Z
enforces the support constraint in Eq. (4) Below, we provide different choices of such B|Z, D]
each of which results in different SSL objective. In practice, however, terms in Eq. ] are hard to estimate from finite samples. We now discuss two variational bounds that can be efficiently estimated and optimized with stochastic gradient descent (Bottou, 2010). For simplicity, we use deterministic encoder eip X Z for the rest of the paper. Detailed derivations are in Appx. C
For both practical objectives we use a contrastive variational lower bound on T[A; 2] based on InfoNCE (Oord et al. 2018), which is standard in SSL.. Specifically, for a sample X, we first obtain the augmented 'positive' A by sampling from PAlX- We then obtain n augmented 'negatives' {4 |= i.i.d. from the marginal P. by first independently sampling X {= } Ji=e from px and then sampling A from PAI PA. We denote A A, ,.... A. InfoNCE then uses critic s, to score how likely each A' E A is to be positive, resulting in the following variational bound,
exp 8 (A,
A[ Z] 2 log(n
+ Exx.x.2 log EA'e exp exps (A'' Z)
When A = X. one can tie the parameters of the critic and the encoder by passing augmentations through the encoder and taking an inner product, i.c., (A, Z) == e, (A)² Z.
2016 Li al. 2018b a) could be valid domain Many previous DG regularizers (e.g.. Ganin et al. bottlenecks. In the following, we discuss two possible le B Z, D), the first of which is novel.
4.2.1 CONTRASTIVE ADVERSARIAL DOMAIN BOTTLENECK (CAD)
Our first domain bottleneck minimizes BiZ, D] I[ D]. which enforces support match using KI. divergence. Dropping constants w.r.t. Z we thus aim to maximize H(D Z).
Algorithm 1 CAD objective
Require: C,, D, X,
1: Z (X)
neural network (DANN, Ganin et al. 2016) does so by ensuring that domain classifier 9a cannot predict domains from representations, i.e., it max- imizes Epp.2 1-o22llogg q. (D] 2)] > H[D] Z] w.r.t. encoder parameter but minimizes it w.r.t. é. However, DANN suffers from two issues: (i) it maximizes an upper bound on the desired term; (ii) it requires adversarial training, which is chal- lenging in practice.
2: A sample(palx)
{(D,((())) A.) }) 1=1
aid sample(pn.x.a)
4: X, A < (x7 } 1- ( (AjU(AT k
5: X-D + (X7 D Di |m]
6: Luge +- log exp (A.2)
b --1[A
DA'EA exp (a 7: Luup -lg Ex'xx_ (x
D[Z
Σx"Ex ex (X")
8: return LEAD Luve XC
To overcome these issues, we construct qDD Z) without introducing additional parameters and with bound that is tight with enough samples. In short, using the equality PD z Ex [PD]x]
our variational distribution q(D) Z) Ex,.x ((D) X)], where q..x (X (Z) is contrastive variational distribution of Px constructed with samples and critic e,, (X)" tied with the encoder, is a count estimate of PD x- Detailed derivations and explanations are in Appx. C.3 The resulting contrastive adversarial domain (CAD) objective is in Algoorithm First, sample domains D = [D, |"= for each X' E X. Then collect inputs associated with a different domain from the current domain D, i.e., X D = [X; ; D, + D. E [n]]. Ignoring constants, the final loss is
Les 1b = E EPD.X...
log EASA exp 8, (A', 2)
exp s. (A, Z)
- log
Σ X'EX
lx Xx (x)22)
(7)
In Appx. C.4 we also derive a conditional variation of CAD that minimizes I 1 D Y]. which can be used when labels are available and supervised augmentations used.
5

---

Published as a conference paper at ICLR 2022
4.2.2 ENTROPY BOTTLENECK (ENT)
Our second domain bottleneck is the entropy bottleneck (Ent) that minimizes H|2] 1[ >
I[Z; DJ. where the first equality uses the encoder's ' determinism. Ent enforces support match by removing all information that not needed to maximize I A]. In particular, minimizing I[Z; is more stringent than I[Z; D), as it also matches the representations inside a domain. The advantage of Ent is that does not require domain samples D, which are rarely accessible in SSL. We consider the standard variational bound used in neural compression (Ballé et al. [016;] Theis et al. 2017).
[[Z] S Papz
log 9o (Z)), where an entropy model go (z(Z) is used. This leads to
exps, (A,
(e, (2 = E.x.A. logg ΣNEA (A',)
AAlog go (z)
(8)
5 RELATED WORK
Provably robust representations under covariate shift. Previous work mostly focuses on domain generalization bounds for robust representations. [Ben-David et al. (2007) (010a) bound the target risk using the source risk, divergence between source and target distributions, and the joint optimal risk over source and target domains. Mansour et al. (2009) generalizes these results from 0-1 loss to more general losses. Johansson et al. (2019) takes this further by deriving a support-based bound. In our setting, these bounds only hint towards sufficient condition for optimality, i.e., matching the marginal pz j or support while minimizing R. [Y] Z]. However, these bounds can often be loose and the implied sufficient conditions are neither necessary nor generally achievable. [Bnn-Dvvd et al. (2010b) suggests that separately minimizing R [Y Z] or matching the marginal is not sufficient, while Zhao et al. (2019) also proves minimizing only the source risk Rd [Y Z] is not sufficient; but none of them proves the desired necessary condition. Our work distinguishes from previous work on three key aspects:
(i) we are the first to study and formalize optimally robust representations, and provide the achievable sufficient and necessary conditions; (ii) we prove that one can practically learn optimal Z with SSL using domain-agnostic augmentations; (iii) we consider a more general framework with any standard losses and a less stringent generalized covariate shift assumption, Still, our work is more specific than others, as we consider idealized DG and unrestricted predictors H.
Practical objectives for DG. The most popular DG methods aim to learn domain- invariant represen- tation by minimizing various divergernces between the marginal distributions pz la and pz (Long) et al. 2015) Ganin et al. 2016)] [Sun Saenko 2016 Long al., 2017]] [i et al. 2018a) [Shen et al. 2018 Nguyen al. 202 l.. Others propose matching the conditional pz v.d across domains instead
Gong et al. 2016 Li et al. 2018b 8b Tachet des Combes et al. 2020). These regularizers would all be
valid domain bottlenecks BIZ, DJ Another line of work aims at learning Z with invariant predictors
Py| a.d across domains (c.g., Arjovsky et al. 2019: Krueger et al. 2021 [i et al.) 202I). However,
none of these methods outperform ERM with fair model selections (Gulrajam & Lopez-Pat Paz 2021.
6 EXPERIMENTS
In our experiments, we aimed to: (i) verify our theoretical results in practice; (ii) investigate 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 pz x using a specified objective and freezes it. Then, the person performing predictions ("the learner") trains her predictor / 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 classifier for h. For the Ent bottleneck, we used Ballé et al. (2018) entropy model. For the CAD bottleneck we used its conditional version whenever labels were available. When model contains no domain bottleneck, we label it as "Base". For experimental details and additional results see Appxs. E and
6.1 SCIENTIFIC SETTING: EXPLORING OPTIMAL REPRESENTATIONS FOR WORST-CASE DG
To validate our theory, we studied optimal representations in a scientific setup that is as close to our IDG framework as possible with log-loss f. In particular, used the PACS dataset (Li al. 2017

---

Published as a conference paper at ICLR 2022
Sugervised
-5.12003
-0.40.11
-0.7200.1
Log likelhhod
Log likellhood
SingleDoes
ApproaDA
-2
-4
IntraDom
Standard
-3.5++0.1
-0.0+0.0
+08
soutte target
Base
Eat
CAD
10
10"
10
10
-10
DA
Approx, Non: DA
Augmentation type
(a) Effect different objectives
(b) Effect of
(c) Effect of augmentations
Figure 3: (a) Adding bottlenecks significantly improves the worst-case DG performance and using domain-agnostic (DA) augmentations (H[A) ZT) performs as well as with labels (R [Y ZD) ) (b) In- creasing the domain bottleneck weight À will improve target performance until it decreases source performance. (c) DA augmentations are crucial but approx. DA aug. might be also be sufficient.
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. (2020) by incorporating the wrongly labeled target data to the source domain. The experimental setup goes as follows: (i) the representor trains ResNet- 18 (He al. 2016) to minimize the objective on labeled data from all domains: (ii) the learner trains worst-case source classifier 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 following results.
Do domain bottlenecks improve worst-case DG? In Fig. 3a compare IDG performance of representations trained with (Ent, CAD) and without (Base) domain bottlenecks. We see that both bottlenecks significantly 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 X? Fig. 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 result the target performance improves until X R 102, where source performance starts to decrease.
What if the representor has access to domain-agnostic augmentations instead of labels? In Sec. 4.2 we provide contrastive objective for using augmentations. To show the effectiveness of the objective, we compared minimizing H[A| Z] using Eq. (6) to standard supervised risk minimization
R Y Z] and used the domain-agnostir supervised augmentations (Fig. 2bp. The 1st and 2nd row of Fig. 3a show that our objective performs similarly to direct label prediction.
How important the choice of augmentations? Prop. 2 mentations are sufficient for achieving IDG, but it does not give necessary conditions. Here we investigate the effect of using our loss with different choices of augmentations. Specifically, we used Leso with five augmentations. The first two are DA. 'Supervised': augment inputs inside the label class across all domains as Fig. 26) 'SingleDom': augment inputs to same label samples from fixed domain. The second two not DA. 'Standard': standard SSL augmentations (Chen et al.) 2020) as in Fig. 2a} 'IntraDom': augment inputs to same label and same domain samples. Finally, we consider 'ApproxDA'', which is approximately DA by augmenting 10% of the time with 'Supervised and 90% of the time with 'IntraDom', FFg. 3c shows that the non-DA augmentations give terrible results compared to DA. Interestingly, "ApproxDA'' also performs very well, which suggests that approximately DA augmentations might be sufficient to learn optimal representations in practice.
shows that domain aggnsicc (DA) aug-
What if the representor does not have access target domains? Prop. Isshows that DG without access to target domains is generally impossible. We empirically verified this by excluding predefined target de domain from the representor's 's training set, i.e., Lex is optimized on 3 of the domains. The learner then trains predictor h on each source. We finally evaluate each h the target domain d, and average over choices of d, The resulting worst-case log likelihood was -.2 + 0.2, which is significantly worse than when the representor had access to all domains (-0.8 + 0.2).
7

---

Published as a conference paper at ICLR 2022
6.2 APPROXIMATING OPTIMAL REPRESENTATIONS BY EXPLOITING PRETRAINED SSL
As discussed in Sec. 4.1, one can learn optimal representations for IDG by performing SSL with domain bottleneck on large sample of inputs X and domain-agnostic augmentations A. This is nearly how CLIP was pretrained (SSL with 400M image-text pairs) except it did not include a domain bottleneck. In this section, we investigate how take advantage of CLIP to approximate optimal representations for IDG. We did so in two simple steps. First, we froze the pretrained CLIP and added multi-layer perceptron (MLP) that could effectively finetune CLIP's representations. Then, we trained the MLP by minimizing our CAD bottleneck and R. (Y] Z] on the available data.
In all experiments, we used the standard DomainBed benchmark (with non-MNIST datasets) and protocol (Gulrajani & Lopez-Paz 2021). In particular, we left out target domain for evaluation and used the union of other domains for training both the encoder and the classifier. Contrary to our scientific setting, the representor does not get access to the target domain. All our representations were evaluated by fitting a linear classifier 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 with standard errors. Note that using 'oracle selection' is more consistent with our theory since it gets access the necessary target information (for model selection), as discussed in AApxx. F.3 Due to space limit, we only included as baselines 'ERM' and 'DomainBed SOTA' which for each dataset is the best result over all baselines. The extended results and baselines are in Table 4) Details in Appxx E.3 We investigated two pretrained CLIP models with different number of parameters. The larger ViT-B/32 denoted 'CLIP I.' and the smaller ResNet-50 denoted 'CLIP S'.
Table I: CLIP significantly outperforms the previous SOTA result DomainBed, as supported by our theoretical analysis. Finetuning CLIP with our CAD bottleneck consistently improves the robustness of its representations and achieves SOTA performance.
Algorithm
VLCS
PACS
OfficeHome
TerraIncognita
DomainNet
ERM
77.6 = 0.3
86.7 = 0.3
66.4 = 0.5
53.0 + 0.3 54.4 + 0.3
41.3 0.1 41.8 + 0.1
DomainBed SOTA
79.9 0.2
87.2 + 0.1
68.4 + 0.2
DINO CAD
69.6 - 0.6
76.1 = 0.1
56.9 = 0.5
25.9 + 1.2
33.6 + 0.1
CLIP S
81.1 + 0.5
90.3 + 0.2
70.6 + 0.1
29.6 + 0.8
47.7 + 0.0
CLIP $ Base
81.3 0.5
91.2 + 0.3
70.6 + 0.1
36.4 + 0.7
46.8 0.2
CLIP S + CAD
82.3 0.3
92.0 + 0.2
71.9 + 0.2
36.2 + 0.8
48.8 + 0.1
CLIP
80.7 0.4
93.7 =
79.6 + 0.1
36.9 0.6
52.8 + 0.1
CLIP L + CAD
81.6 - 0.1
94.9 = 0.3
80.0 + 0.2
40.6 1 1.1
53.7 1 0.1
Approx. Optimal z"
86.8 0.6
97.2 + 0.6 86.3 + 1.6
76.5 + 4.1
66.7 + 0.2
The row 'CLIP L+ CAD' in Table I] [sshows that finetuning a large pretrained CLIP model with our CAD achieves SOTA on nearly all DomainBed benchmarks by a very large margin (see 2nd row). Note that the poor performance on Terralncognita is likely because CLIP's dataset does not cover such images (camera traps monitoring animals). The last row essentially shows an optimal representation, which we approximate by finetuning CLIP L with our CAD on all domains including the target. The gap between CLIP L + CAD and the upper-bound suggests that one can still learn better representations. We hypothesize that end-to-end training of our objective would greatly shrink this gap.
Can we approximate optimal representations by exploiting pretrained CLIP?
Are gains due to the architectural differences? DomainBed's 's baselines finetuned an ImageNet 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 CLIP S. Table I shsows that CLIP $ + CAD still significantly outperforms DomainBed baselines. Note that our theory does not constrain the encoder and so we expect larger encoders be better as seen in Table I
What is the effect of domain bottlenecks? In the "CLIP" rows of Table we investigated the effect of finetuning CLIP with our CAD bottleneck. We see that for both CLIP L and CLIP S, consistently improves results by around ~ 2%. These gains are due to the bottleneck, rather than finetuning on source data as seen by 'CLIP S Base'. We believe the gains could potentially be much larger if CLIP was trained end-to-end with our bottleneck. Note that CLIP S already significantly
8

---

Published as a conference paper at ICLR 2022
outperforms baselines. We hypothesize that this is because SGD acts information bottleneck that naturally favors support match (Shwartz-Ziy & Tishby 2017).
Which pretrained SSL model to use? Our theory suggests that models as long as their augmentations are domain-agnostic and their training set covers desired domains. We investigated adaption of SSL models that do not satisfy those properties by finetuning DINO (Caron et al. 2021). the current SOTA on SSL ImageNet. DINO is pretraiend using standard augmentations. As result, Table I shows that the finetuned DINO + CAD significantly underper- forms compared to CLIP S and DomainBed baselines. This supports our hypothesis that CLIP is much more robust than other SSL methods due to its domain agnnssic augmentations.
exploit pretrained SSL
6.3 TOWARDS GENERIC ROBUST REPRESENTATIONS WITH SSL
In the previous section, we finetuned CLIP in task specific 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 sufficiently 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 sufficient compute. In this section, we aim to emulate such training of generic robust representations.
To do so we used LAION-400M (Schuhmann et al. 202T) that a public dataset that contains 400M web-crawled image-text pairs. Due to computational budget, again froze the pretrained CLIP and only finetuned an additional MLP with our Le.. We used L as it only requires access to paired image X and text A but prior information about domain D. As in CLIPPs paper, we evaluated the learned representation Z in Taori et al.] |s (2020) realistic setting, where a linear classifier from Z is trained on ImageNet and tested on 7 natural distribution shift datasets. Details in Appxx E.4
Would training CLIP with a bottleneck have improved its robustness? As shown in the last 2 rows Table 2. finetuning CLIP L on LAION with L... (Tuned w/ Ent) outperforms finetuning without bottleneck (Tuned w/o Ent) 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 Le,. end-to-end, In [Appx.. F.4, we show similar results on DomainBed. 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 2: Finetuning CLIP on LAION with an entropy bottleneck improves its robustness compared to finetuning without on 7 distribution shift datasets. The pretrained CLIP L. is still better likely due to end-to-end training with higher quality data. IN denotes ImageNet.
IN
IN-V2
IN-S
YT-BB
IN-
ObjectNet
IN-A
IN-R
Avg.
CLIP
75.2
64.2
41.0
58.4
71.6
42.8
27.5
62.9
52.6
Tuned w/o Ent
73.8
62.1
37.0
56.9
68.8
41.3
26.0
58.1
50.0
Tuned Ent
74.2
62.7
38.9
58.1
70.1
42.1
26.2
60.8
51.3
7 CONCLUSION
We gave simple variational characterization of all representations on which source-risk minimizers are guaranteed to generalize to target domains that preserve the Bayes predictor. Similar to previous work, our theory strongly implies the need for target information when learning representations for domain generalization. Nevertheless, we identified a domain domain-agnostic property of data augmentations that make possible learn optimal representations from unlabelled data. Thus, we showed that is possible to learn robust representations using only large sources of inputs X and augmentations A.
There are caveats that need to be addressed future work. First, we studied an idealized DG, which assumes access to the population distributions. This gives insights into the challenges that are specific to DG, rather than finite sample challenges faced throughout ML. Second, we considered risk minimizers from unconstrained hypothesis class. The support constraint can likely be weakened, the hypothesis class is constrained. Finally, we focus only on optimal representations, but it would be interesting to characterize approximately optimal representations. Nevertheless, in this idcalized setting, our characterization is springboard from which all future objectives can be derived, and, in general, it brings closer the goal of robust machine learning systems.

---

Published as a conference paper at ICLR 2022
Acknowledgement We would like to thank Elliot Creager, Roger Grosse, Elan Rosenfeld, Guodong Zhang, Han Zhao, and anonymous reviewers for their helpful feedbacks and encouragements. Re- sources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute. We acknowl- edge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), RGPIN-2021-03445.
Reproducibility For our theoretical results, we include formal assumptions, statements, and proofs in Appxs. A] and [ We include the detailed derivations of our algorithms in Appx. C For our experiments, we include experimental details for reproducing our results in Appx. E and have released our code at pottps://gitaie cow.ryoun////pptooao.o
REFERENCES
Alexander A Alemi, Jan Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint /612.00611.000 2016.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint 1rXXv:1907.0289. 2019.
Johannes Ballé, Valero Laparra, and Eero P Simoncelli. End-to-end optimized image compression. arXiv preprini arXiv: 2016.
Johannes Ballé, David Minnen, Saurabh Singh, Sung Jin Hwang, and Nick Johnston. Variational image compression with scale hyperprior. arXiv preprint arXiv:1802.014553 2018.
Andrei Barbu, David Mayo, Julian Alverio, William Luo, Christopher Wang, Danny Gutfreund, Joshua Tenenbaum, and Boris Katz. Objectnet: A large-scale bias-controlled dataset for pushing the limits of object recognition models. 2019.
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European conference on computer vision (ECCV), pp. 456-473, 2018.
Shai Ben-David, John Blitzer, Koby Crammer, Fernando Pereira, al. Analysis of representations for domain adaptation. Advances in neural information processing systems, 19:137, 2007.
Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Percira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine Learning, 79(1):111-17) 2010a.
Shai Ben-David, Tyler Lu, Teresa Luu, and David Pal. Impossibility theorems for domain adaptation. In Yee Whye Teh and Mike Titterington (eds.), Proceedings of the Thirteenth International Conference Artificial Intelligence and Statisties, volume 9 of Proceedings of Machine Learning Research, pp. 129-136, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010b. PMLR. URL https: https::/prrceeet: ngs. ml www.srrss//eeeee v /tdp///vviid.... 0a htm
Léon Bottou. Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT pp. 177-186. Springer, 2010.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. arXiv preprint
arXiv:2104. 14294, 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1597-1607. PMLR, 2020.
Yann Dubois, Douwe Kiela, David Schwab, and Ramakrishna Vedantam. Learning optimal repre- sentations with the decodable information bottleneck. In H. Larochelle, M. Ranzato, R. Hadsell,
M. F. Balcan, and H. Lin (eds.). Advances in Neural Information Processing Systems, volume 33, pp. 18674-186888 18690. Curran Associates, Inc., 2020. URL htRps:::::: /ttp::///r/cceeeeeee wew.eeediiiid....m ing.inn..u...................... ips ps cc cc/paper//e
Paper.pdf
Yann Dubois, Benjamin Bloem- Bloem-Reddy, Karen Ullrich, and Chris Maddison. Lossy compression for lossless prediction. arXiv preprint arXiv:2 106. 10800, 2021.
10

---

Published as a conference paper at ICLR 2022
Chen Fang, Ye Xu, and Daniel N Rockmore. Unbiased metrie learning: On the utilization of multiple datasets and web images for softening bias. In Proceedings of the IEEE International Conference Computer Viston, pp. 1657-1664, 2013.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. The journal of machine learning research, 17(1):2096 2016.
Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American statistical Association, 102(477):359-378 2007.
Mingming Gong, Kun Zhang, Tongliang Liu, Dacheng Tao, Clark Glymour, and Bernhard Schölkopf.
Domain adaptation with conditional transferable components. In International conference machine learning, pp. 2839-2849 PMLR, 2016.
Ian Goodfellow. Nips 2016 tutorial:
Generative adversarial networks.
arXiv preprint
arXiv:1701.00160, 2016.
Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In International Conference Learning Representations, 2021. /tthtttss//opeeeeeeeee forum?
1d=I@dXeXDoRtI
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, et al. The many faces of robustness: A critical analysis of distr--disttiuution generalization. arXiv preprint arXiv:2006. /6244),66/221,, 2020.
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adver- sarial examples. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 15262-15271, 2021.
Fredrik D) Johansson, David Sontag, and Rajesh Ranganath. Support and invertibility in domain- invariant representations. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 527-536. PMLR, 2019.
Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. arXiv preprint arXiv:2004. 1120000.11366 2020.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Naveen Kodali, Jacob Abernethy, James Hays, and Zsolt Kira. On convergence and stability of gans. arXiv preprint arXiv: arXiv: 1705.07215 2017.
David Krueger, Ethan Caballero, Joemn Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distributioo generalization via risk extrapola- tion (rex). In International Conference on Machine Learning, pp. 5815-5826. PMLR, 2021.
Li, Yifei Shen, Yezhen Wang, Wenzhen Zhu, Colorado Reed, Jun Zhang, Dongsheng Li, Kurt Keutzer, and Han Zhao. Invariant information bottleneck for domain generalization. arXiv preprint arXiv:2106.06333. 2021.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. Proceedings of the IEEE international conference on computer vision, pp. 5542-5550, 2017.
Haoliang Li, Sinno Jialin Pan, Shiqi Wang. and Alex C Kot. Domain generalization with adver- sarial feature learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5400-5409, 201 8a.
11

---

Published as a conference paper at ICLR 2022
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 624 639, 2018b.
Mingsheng Long, Yue Cao, Jianmin Wang, and Michael Jordan. Learning transferable features with deep adaptation networks. In International conference on machine learning, pp. 97-105. PMLR, 2015.
Mingsheng Long, Han Zhu, Jianmin Wang, and Michael I Jordan. Deep transfer learning with joint adaptation networks. In International conference on machine learning, pp. 2208-2217. PMLR, 2017.
Yishay Mansour, Mehryar Mohri, and Afshin Rostamizadeh. Domain adaptation: Learning bounds and algorithms. arXiv preprint arXiv:0902.343M0 2009.
A Tuan Nguyen, Toan Tran, Yarin Gal, Philip HS Torr, and Atilmm Güney Baydin. Kl guided domain adaptation. arXiv preprint arXiv:2106.07780, 2021.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807.03748 2018.
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1406-1415, 2019.
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International Conference Machine Learning, pp. 5171-5180.
PMLR, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8748 8763. PMLR, 18 24 Jul 2021. URL https: ptttps://proceedi:: ingeeeddiiigg.... mlr www.sp..ss///.em......e..r vt////111/////// (tt.339/rrfordd22111 html
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International Conference on Machine Learning, pp. 5389-5400. PMIR, 2019.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv: 191 911111111 108731, 2019.
Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar. A Theoretical Analysis of Contrastive Unsupervised Representation Learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 5628-5628666 5638- PMLR, 2019. URL http:
//proceedings. ml press / v97 ssunshi 1 9a. html
Christoph Schuhmann, Richard Vencu, Romain Beaumont, Robert Kaczmarczyk, Clayton Mullis, Aarush Katta, Theo Coombes, Jenia Jitsev, and Aran Komatsuzaki. Laion-400m: Open dataset of clip-filtered 400 million image-text pairs. arXiv preprint arXiv:21 1..0:: 2021.
Ohad Shamir, Sivan Sabato, and Naftali Tishby. Leaninng and generalization with the information bottleneck. Theor Comput. Sci., 41 411(22-33):2666222222 2711 2010. doi: 10.10166f.ess2210.44.006 URL /tttps:///oi.oot
2010. 04. 006
Vaishaal Shankar, Achal Dave, Rebecca Roelofs, Deva Ramanan, Benjamin Recht, and Ludwig Schmidt. Do image classifiers generalize across time? arXiv preprint arXiv: trXivv/906.02//// 2019.
Jian Shen, Yanru Qu, Weinan Zhang, and Yong Yu. Wasserstein distance guided representation leamning for domain adaptation. In Shirly-Second AAAI Conference on Artificial Intelligence, 2018.
12

---

Published as a conference paper at ICLR 2022
Ravid Shwartz-Zii Ziv and Naftali Tishby. Opening the black box of deep neural networks via information. CoRR, abs/1 770000000,, 10, 2017. UtL.ttt::/axxxi..rrr http:/bb/17000 00810
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pp. 443-450. Springer, 2016.
Remi Tachet des Combes, Han Zhao, Yu-Xiang Wang, and Geoffrey Gordon, Domain adaptation
with conditional distribution matching and generalized label shift. Advances in Neural Information Processing Systems, 33, 2020.
Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht, and Ludwig Schmidt. Measuring robustness to natural distribution shifts in image classification. arXiv preprint arXiv:2007.00644 2020.
Lucas Theis, Wenzhe Shi, Andrew Cunningham, and Ferenc Huszár. Lossy image compression with compressive autoencoders. arXiv preprint arXiv:1703 703.00395, 2017.
Naftali Tishby, Fernando Pereira, and William Bialek. The information bottleneck method. arXiv preprinl physics/0004057, 2000.
Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 5018550110000 2027,2017
Haohan Wang, Songwei Ge, Eric Xing, and Zachary C Lipton, Learning robust global representa- tions by penalizing local predictive power. arXiv preprint arXiv: 13549, 2019.
Luyu Wang, Pauline Luc, Adria Recasens, Jean-Baptiste Alayrac, and Aaron van den Oord. Multi- modal self-supervised learning of general audio representations. arXiv preprint arXiv:2104.12807, 2021.
Aolin Xu and Maxim Raginsky. Minimum excess risk in bayesian learning. arXiv preprint arXiv:2012. 14868, 2020.
Shen Yan, Huan Song, Nanxiang Li. Lincan Zou, and Liu Ren. Improve unsupervised domain adaptation with mixup training. arXiv preprint arXiv:2001.0061 2020.
Han Zhao, Remi Tachet Des Combes, Kun Zhang, and Geoffrey Gordon. On learning invariant representations for domain adaptation. In International Conference on Machine Learning, pp. 7523-7532, PMLR, 2019.
13

---

Published as a conference paper at ICLR 2022
Appendix
Table of Contents
Preliminaries
15
A. Notation
15
A22 Definitions
15
A.3 Assumptions
17
B
Proofs
19
Lemmas general losses
Proof Theorem
B.3
Impossibility results
23
B.
Augmentations
25
C
Practical objectives
27
C.I
Mutual information bottleneck B Z, X.Y = 12:X
27
Entropy bottleneck B, X, Y.D
Z
28
Contrastive adversarial domain bottleneck B H,Y,YY -1;D
29
C.4
Conditional CAD B X, 12: D Y
30
D
Extended Related Work
32
E
Experimental Details
32
E.
Scientifie
32
E.2
Bridge
33
DomainBed
34
E.4
LAION
35
F
Additional Experimental Results
36
F1
Scientific
36
F2
Bridge
36
F3
Ded
37
LAION
38
14

---

Published as a conference paper at ICLR 2022
A PRELIMINARIES
A.1 NOTATION
For the most part, we will assume that all spaces are discrete probability spaces. full list of assumptions is found at Appx. A.3
General The image of set A S X' under a function f : X' + y is denoted (-) = {f(z) E A}. The pre-image denoted (B) {a e X f(z) E B} for B y.
Probability Random variables (r.v.) are denoted by uppercase letters (e.g., X), and their sam- ple space and realizations are denoted by the corresponding calligraphic (e.g., X) and lowercase letters (e.g., 2) respectively. The probability mass function (pmf) of random variable X is denoted as px. We use capital instead of p to denote the measure under p. The support supp(px) ) of discrete distribution is the set of all points E X' with positive probability, i... supp(px) = {or E px (x) > 0}. The space of all probability distributions on X' is denoted P(X) Px (x) and Exex Px( =1].
When it is necessary to be explicit, we will denote "X" is distributed as Px' using the notation X A px. Expectations are written as: Epx [X(]], independence two r.V. as I.L,, conditional independence as -Ll.
For jointly distributed random variables (X, Y) taking value in (t.vi.) X' x V, the conditional distribution is denoted as py x : x X - 1]. For convenience, let py 1 Py -x conditional distribution of Y given æ. All random variables are independently distributed, unless an explicit joint distribution coupling is given.
(- [x) be the
A.2 DEFINTTIONS
We are interested in prediction problems with domain shift. There are three random variables: the target domain D,, the input X, the label Y. They have the following joint distribution:
X,,) A pn. Px,Y 1D:
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 PD..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 using a conditional distribution pz X-
Definition (Encoder). An encoder is conditional distribution pz 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:
(De, X,Y, Ap, Px.Y D. PZIX The key thing to notice here is that Z is conditionally independent of Y, D; given X. In particular, the same encoder used across all domains.
(10)
A.2.1 RISK MINIMIZATION
Our ultimate goal is to predict Y from the representation Z of X in manner that is robust to changes in the domain.
We formalize this in the standard way by making predictions actions. For example the prediction space may be the set of all possible labels = Y, in which case would be predicting deterministic labels. Or may predict distribution case the prediction space would be the set all probability distributions on ), i.e. P())).
I' in space of predictions or
labels, in which
A predictor is a function mapping inputs to predictions, i.e.. f X + r, representations to predictions, i.c., h : + I'. For example, f may be a neural network that takes as input a sample x and outputs vector of logits that parameterize softmax distribution finitely many labels.
15

---

Published as a conference paper at ICLR 2022
We select predictors according to the risk defined via a loss function (: x R20 foo):
Ry YY |X] == Epx,y [e... f(X)).
(11)
In particular, we are interested the Bayes (minimum) risk over all predictors:
R R X = inf R, [Y [X],
(12)
/
We denote the set of all optimal predictors from X as FF := (f R, Y X] R X]}
(13)
Similarly, we define the risk R. | Z]. the Bayes risk [Y Z], and the set optimal predictors
H : (h) R. Y 2] = R 2]}
(14)
from Z, all of which vary as a function of the encoder Pz x- Note, in the main body of the paper, we omitted the subscript from Hi for clarity, but we will keep the Appendices. We assume loss and prediction space always admit optima [Item of Assumptionn 2). and thus FF, H22 are always non-enn non-empty.
that together
We will be assuming that the risk admits unique optimal prediction when predicting from X (Tem 3 of Assumptiio 2). Thus it makes sense to define the following:
Definition 2 (The Bayes predictor). The Bayes predictor f" X - I is the unique predictor that is optimal for all x X:
(x) arg min L yel
[[(
(15)
Definition 3 (The Bayes image). The image of all the inputs under the Bayes predictor will be denoted as I f='("( and called the Bayes image.
Note that J becomes singleton { but it is not necessarily the case for H since we will be making any uniqueness assumption on optimal prediction from Z.
A.2.2 DOMAIN GENERALIZATION
We are interested controlling the risk in a domain generalization setting, and so we define the domain-conditional risk,
R Y x]
Exxy [LY ((X)]
(16)
Rd Y F) are defined as Eqs. (12) and (13) respectively, but with respect to Rý. Similarly, define the Bayes image for domain d as
r = f" (supp(px]a))
(17)
We also define domain-conditional quantities for prediction from a representation Z. The most important term which we will be investigating an idealization of the domain generalization worst- case risk.
Definition 4 (IDG risk). Given an encoder pz x and distribution PD, target domain D. and source domain D., the idealized domain generalization worst-case risk, IDG risk for short, is the expected worst-case target risk taken over source minimizers, i.c..
RYpO [Y] Z == Epp,
sup R.D [Y | hEH?
(18)
Note that the IDG risk is well-defined because H: Hy is non-empty by Assumption 2] optimal representations, are then those that minimize the IDG risk.
The desired
Definition 5 (Optimal representations for IDG). An encoder pz* |x is optimal for idealized domain generalization and only if it minimizes the IDG risk, i.e.,
RID Y 7 = inf RIdo [Y] 2)
(19)
P2
16

---

Published as a conference paper at ICLR 2022
A.3 ASSUMPTIONS
We make the following assumptions throughout the paper. All these assumptions should hold for practical settings.
Assumption (Convenience: discrete probability spaces), All data spaces discrete spaces. Because the distributions of X, Y, D. D fixed, assume for convenience that supp(px)= = X' supp(py) y, and supp(pp.) = D.
are
Assumption is convenience assumption avoid measure theory for the sake of clarity. It always holds in practice due to finiteness of computers, i.e., all spaces will be finite but arbitrarily large. We believe that our claims can nevertheless be generalized to typical continuous spaces with some minor technical assumptions.
Assumption 2 (Losses admit optima). We assume that our risk always admits optimal predictions:
1. [T] >> 1.
2. For all pr E P(V), there exists 7" E r, such that Epp ((((,,")) x E [(())))]
AyET.
(20)
3. For all x E X, there exist y E T, such that Ep [[)(("")) <
Vo
(21)
Note that for log-loss E(v, where the optimal prediction for Item is 7- = Py by strict properness (Gneiting & Raftery 2007). If we consider the 0-1 loss (reverse accuracy) Ely, 7) = - 1ly = 7] with I 3 and finite label space where the optimal prediction for Item 3 is 7 = arg maxyey PY (v), this assumption is mostly satisfied, except we assume that py i has a unique mode.
log (1) and finite V, these assumptions satisfied if 1' P())
Assumption 2 serves two purposes: Item eensures that for any representation the optimal predictors from Z exists such that the IDG risk is well-defined as in Def. 5 [tem 3 ensures a unique Bayes predictor from X, which simplifies the analysis and is satisfied by common losses as described above.
Assumption (Cardinalities). We assume that
2
2
(22)
Assumption i 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 (I,d) E supp(px.D.), 7 e such that 7 * f (x), we have
Ie f*))) < Epr 1z.a [()
(23)
For example, in the case of strictly proper scoring rules, e.g. log loss, covariate shift py X.D Py xx 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 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 isis the core assumption for our theoretical results. It ensures that source and target domains are related in useful way that can be utilized by the representation.
Assumption (Constant Bayes image). The Bayes image is invariant across domains, i.e., for all E D.
r"
(24)
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 r {pr supp(px } is the same across domains.
17

---

Published as a conference paper at ICLR 2022
Assumption is crucial to be able to learn. Without it, in the extreme case, could set each domain to all examples associated with a single element from the label set (or the Bayes image set) in which case it is impossible generalize across different domains. Assumption 5 is also necessary to guarantee the existence of optimal representations as in Prop. 3
Assumption 6 (Domain joint). PD,.D, is any distribution such that supp(pD..D.) D D.
In simplified scenario, one could define the source D, and target De as i.i.d. r.v. from pp,, where D, PD, PD. PD, PD, and Assumption ils trivially satisfied.
18

---

Published as a conference paper at ICLR 2022
B PROOFS
B.1 LEMMAS GENERAL LOSSES
An important result that we will be using is the generalized data processing inequality of Bayes risk
(Xu Raginsky 2020 Dubois al. 202T). We include it here for completeness.
Lemma 1 (Generalized DPI (Xu & Raginsky 2020 Dubois et al. 2021 ). Let Z - - Y be Markov chain of random variables. For any loss function
R
R [Y 2].
(25)
For the case of strictly proper losses [Assumption
2 can go one step further.
Lemma 2. Let X Y a Markov chain of random variables. Then, under Assumptions
and we have that
R R 2 R
vh H2 V(r, E supp(px,z), (=)=)= (x).
(26)
Proof. Suppose that for all h" Hy have h² (z) (xx) the support of px,z- Then,
R Y X| = Epx, [
(27)
Epx.YPA x
(28)
-Epx (0(Y,
(29)
= =Epz, I h (Z))
(30)
= =
(31)
Now suppose there exists ht H and pair (æ', x') supp(px,z) such that h (=()') + f (x').
Then
R [
%]
(32)
Exx.z 1x (E(Y,1c (Z))
(33)
= Px,z(x' z') Bpy 10Y, +
Σ
(x, E,,,
(34)
2 -xx,2(x' 2') Epy (((,n((<))))
Σ
Px,z(I, 2) Epy ((X,
(35)
(x2))
> px.2(xx Epy f(x)))
Σ
px,2(1, Epy [YY, f()))
(36)
R x
(37)
Eg. (35) follows by Item of Assumption 2along with the definition of f* Eq. (36) follows by Item
of Assumption 2 and the fact that h2 h²(2") f (x'). This completes the proof, because Lemma
prevents R < 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 al. (200)) Proposition 6, to the case of proper losses.
Proposition 3 (Existence of optimal representations). Under Assumptions encoder pz- that optimal for Fgg. (3) i.e.
to S
there exists an
Pz. X E arg min R Y Z]
s.t. Ad D, supp(pz a) supp(pz).
(38)
PZIX
Moreover, we have that
R X] RY 27)
(39)
19

---

Published as a conference paper at ICLR 2022
Proof. Because assume arbitrary encoders pz x the essence of this construction is simple: embed the Bayes image into Z. Indeed, let Φ I - Z be any one-to-one function, which exists due to Assumption 3 (here use deterministic one-to-one function for simplicity, the construction can be easily extended to stochastic case). Then let 7" = (X). We now verify the properties of Pz. xx-
1. Z satisfies R
R LY |*). Indeed,
R [Y | = Exx, (X))]
(40)
= =EpxYPZ2 x (
(41)
Epxe, 1 (Z"))
(42)
2 R 1z)
(43)
Eq. (42) is by our construction of Z" and Eq. (43) is by the definition of the Bayes risk.
Due the data processing inequality of Bayes risk Lemma 1 we also have R [ X <
R Y 2*), from which we conclude that R. [Y X] R Z and that Eq. (39) holds.
2. Recall that 1* = --- (2) and I = S 1 (supp(pxxppp ())) Now let us compute the desired support for all d E D:
supp(pz- (d) x- - (T2)
(44)
& (T*)
(45)
= supp(pz-).
(46)
Eq. (45) is by Assumption
Because R. [Y| X is the minimum achievable risk by any encoder regardless constraint (this is by Lemma I), this implies that pz. x is an optimal encoder for for (3)
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 pz x be an encoder such that R (Y] / R X). Under Assumptions / |and] we have that for all E supp(pz). there exists x E I such that
2
1 (((Y,))
47
(47)
In other words, the restriction of any ht E Hi to supp(pz) is unique. If, in addition, Assumption holds, then for all (z,d) d) E supp(pz.D, ). E such that # (2),,,
Epp a.d [ (x)) < | [,(Y
(48)
In other words, the restriction of any E Hy. to supp(pz aa) is unique and equal to h".
Proof. For the first result, let a E supp(pz) and consider æ E supp(px: z). By Lemma 2 it must be
the case that f is constant on supp(px 2=). Thus, we can pick 7" = (x:). Now, let 7 ,² We have that,
[<((,''))
EpxisPr [[(Y,>")]
(49)
-Exx [((Y)]((X))
(50)
< Epx Epxi-PP 1uy x [((Y,))]
(51)
= 1
(52)
Eq. (49) is due to the conditional independence of Y and Z given X. Eq. (51) is due to to/Asssuptiinn
and the definition of the Bayes predictor. Let h²" supp(pz) - be the unique Bayes predictor from Z.
Now, for the second result, note that
R R X]
R, YY XX]
(53)
20

---

Published as a conference paper at ICLR 2022
E (d) [X]]
(54)
dED
Σ (d) ( [x|
Assumption
(55)
dED
and
R 2 = R | Z]
(56)
Σ
PD. (d) R [Y -
(57)
dED
(d) Rd Y Z].
(58)
dED
where Eq. (58)
due to the definition of (domain-conditional) Bayes risk. Then
R 2| -
R
X 2 Σ
(d)
R YY Z] - R [ x1)
(59)
DED
2 0.
Lemma conditioned on d
(60)
Thus, any encoder that achieves R Y 2 R[Y X also satisfies R [Y 2] = T for all DD since we assume that supp(pp,) in Assumption I Now, let D. An argument analogous to Lemma 2 gives us,
Yh E Hz. V(x, 2) E supp(px,2 a), M(2) = f())=
(61)
Eq. (61) is derived from R. [Y 2 = R Y X using Assumption 4 in place of Item 3 of As
sumption for specific domain d. Let 2 E supp(pz (a) and y I' such that y # (z). Since
supp(px [=,a) S supp(px a), f is constant supp(px a.d)) and equal to h² Now, as above,
we have that
Epy [((Y,>> Exx aP x.d [((Y)>
(62)
Exx PP xx. [f
(63)
Epx LPP (...)(((,))))
(64)
Epy (0(Y,))).
(65)
Eq. (64) is due to Assumption 4
Corollary 1. Let pz x be an encoder such that YY Z] R Y []] Under Assumptions and 4 we have that H C H², for all d E D and that for all d,, de E DD
inf KER
de Y Z
(66)
Proof. H% C H., is immediate from Lemma Now, have that R Y 2 > Rd Y the result follows by taking any h H H in the inf of Eq. (66)
Z So.
Theorem (Characterizing optimal representations for IDG, equiv. Theorem J Under Assump lions 1 an encoder pz is optimal for idealized domain generulization if and only if minimizes The Bayes risk while matching the support pz and pz for all d D, i.e.
Pz x arg min [Y Z]
(67)
PZIX s.t.
V E D, supp(pz a) = supp(pz)
(68)
Proof. The IDG risk is lower bounded by R [Y X]:
Ris [Y| [1] 2 E Epp..
inf hEH
R [Y 2)
(69)
21

---

Published as a conference paper at ICLR 2022
E Epn..
R [Y]
Z]]
(70)
x]] = R |
2 E Pp Dt
R [Y
Lemma
(71)
Assumption 4
(72)
We will now show that this lower bound achieved by an encoder if and only if it satisfies Eqs. (67)
and (68] which exist by Propp.
Sufficiency = ): Let pzj x be an encoder that satisfies Eqs. (67) and (68) Note that [Y Z]
|X] by Prop. 3
Let h² E H2 then we have the following IDG risk
RIDG Y Z
(73)
Epn...:
sup NEW2.D,
EpZYID [((Y,((Z))]
(74)
Pp
sup hEH',
E PZIIID, [e(Y, (Z))
Lemma under matching support
(75)
-EpD Epsy 1
[e(Y, h (2))]
constant W.E. Dr
(76)
= =Y 7] [
(77)
Necessity
( ): If the IDG risk is R [Y X], then it must be the case that
R Σ = R [Y [
(78)
sup
du | = R Y Z] v(d,,d)) supp(pD..D.)
(79)
NER
We will prove by contrapositive that Eq. (79) implies support match (Eq. (68) Suppose that the
support match does not hold. Since supp(pz) = Udepsupp(Pzia aa) and supp(PD..D. = D x D
(Assumption 6). there must exist (da, d,) E supp(PD..D.) such that supp(pz (,,) # supp(pz a
Define the set S = supp(pz (a)) n supp(p2 a.) and 5 supp(P2]d) \ supp(p2 (. let p Pzid, (S), and let h² Hyz Then,
sup kE
Rd 2
(80)
sup Oppr [[Y,,(Z))] - (1 p) Epy. s.4, [((Y, kEHtz.d.
(81)
sup Bpy |C(Y, ((Y,)(())]] + (1 Bpy. s.4 [((Y,
Lem. 3
(82)
heHly,,
-PBpy.zss. [e(Y, + P)) sup hEH
Epy,z [((Y,((Z))
(83)
R YY 1 p) sup hEM
Epy.219.4 -
(84)
Rd de Y z]
Lem.
3
(85)
1-p> due to support mismatch. For any h E H2., such that h + h" on (such an h exists by Item of Assumption 2). we have that
Eq. (85) uses the following reasoning.
E By 18,, 1(Y (Z)) ,
(86)
by Lemma 3
As corollary from the proof strategy we directly have that the optimal DG risk simply R [Y] X]. This means that using the optimal encoder one can actually perform just well by training on the source as if you were to directly train on the target using the raw data.
Corollary (Optimal IDG Risk). Under Assumptions 7 to
inf PZIX Ride [V] 2) R | x].
22

---

Published as a conference paper at ICLR 2022
B.3 IMPOSSIBILITY RESULTS
As a direct corollary of TTm. 2 we know that it is impossible to learn an optimal representation without knowledge or assumptions the target domain. We actually prove the following much stronger negative result, which essentially states that it is impossible to find a useful representation without having some information about the target domain. Specifically, we prove that if there exists non-trivial target domain on which the representation is advantageous then there exists an infinite amount of target domains on which it 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 understand, but note that we can generalize the proof all losses with the right assumptions.
The key that outside the source domain, the label distribution is unconstrained because general- ized 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 valid domain. Furthermore, if there exists one domain on which the representation is good, then one can construct a domain on which the representation bad simply by labelling this point as the constant prediction.
Proposition 4 (No free lunch for learning representations for IDG, equiv. Proposition I). Let E be the 0-1 loss with prediction space I' V. Let Rep P(X, + 2 be any algorithum for choosing an encoder pz from the data distribution px,,, be any constant r.v. that t.v.i. Z, and Px, ld, be any desired source distribution such that
* there is a unique constant prediction "cc arg minyey Epy 14, i4, 1YY v)).
and (X supp(px(a.) > 1.
Let pza, |x Rep(px,y (d,) be the chosen source encoder. If there exists a target domain Px,y such that
(Non-trivial support) 9 + supp(px S X supplpx aa):
(Satisfies Bayes image invariance) lo = y, i.e,, there is at least one example for every
possible label;
(Source encoder is useful) pza, |x performs better than constant representation,
sup R [Y Za.l < NE do
sup NEHi
R [Y
(87)
Then there exist multiple target domains d; such that Pza, |x underperforms constant encoder,
sup NEN 2,,
R Y
Za.
sup
R [Y].
(88)
hE
Proof. Let h² His be any source Bayes predictor corresponding our encoder. Partition Z according to whether h² predicts like the constant or not:
Zc : sa E Z h" (2) = Yc}
3 = 2 Zc.
(89)
We know by assumption that di is s.t.
sup R YY Zz. <
sup
R CY
(90)
NEN do
hEN
which is clearly only possible
Pz., 14 (2xc) >
(91)
In other words, there exists some input I&C X { supp(px (d,) that will get represented outside of the constant region, ie..
Pza, uxe (2xc) 0.
(92)
23

---

Published as a conference paper at ICLR 2022
We will now construct the desired bad domain di by giving nearly all mass to this r4c. specifically, let Px l (x+c) = 1 -1-6 for some 0 < 6 < 1. We assign this example to the constant label, i.e., Py Iz (1C) = 1. The rest of the target domain mass ō is distributed as with the source domain, i.e. Px. dt (I. :) 6 Px,Y [d, [d, (2., y) for all 2, E supp(px,x (d,). Importantly, the constructed
domain d, 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 (6 > 0). We added no new prediction 7, because (* = T0 E y which must already have been in r due to the validity of d'.
Now let us compute the desired risk for that "bad" domain and show that the desired encoder performs worse than a constant encoder.
sup
R
[Y] Za,]
(93)
hEH id
sup (1- 6) Epza, 1xye [1 1/e h((a,))] + +8Rg [Y] Za,]
(94)
he 1% As
2 (1 5)(1 - P2d. 1#0 (2c))
(95)
= (1 62, 1220 (24c)
(96)
In contrast, it is easy to show that suphe R [Y]C] S 5 because the constant predictor would
Pza. "vc (Zxo)
be perfect for ryc. So any choice of 0 8 1+Pza, 1= (2xc)' would satisfy Eq. (88)
We
conclude the proof by noting that there are infinitely many such choices of and any choice of those would result in a different valid bad domain d".
Note that representations can often be much worse than using constant r.v. Specifically, if an encoder pz maps an x outside of the source support then there exists an infinite number of target domains where that representation is the worst possible representation.
Proposition 5 (Worst representation). Let Rep, PY.X PZa xx, be as in Prop. 4 and 6 > O. If there exists an example Fp E X supp(px d.) that is mapped outside of the source support, i.e., supp(pz,, axs) supp((p d,) = (, then there exist many target domains p (d, S.t. pza, x is close to the worst possible loss, i.e..
sup hE 2
R Y
(97)
Proof. By assumption there exists an Tb whose support is outside the source support. Then similarly to Prop. 4 we construct a bad target domain d, by giving nearly all mass to that example px d, (2x)
1 where& 8 > and assign with probability I to some label that is in the source Bayes image, i.c., Py [x (76) 1 for some 76 1' The rest of the target domain mass 8 is distributed as in
Prop. 4 to the source inputs. As in Prop. 4) such a target domain dg satisfies our assumptions. Now let
us compute the risk for that d, and show that the desired encoder performs arbitrarily bad.
sup
R Y Za.l
(98)
hEH2
sup
(1- 8) Epza, 1ng [1 1/70 h(Za.)ll + +8R [Y Za,]
Eq. (94)
(99)
ke 2d,
sup he Hza de
(1- 8) Bpza, Ing [1 1/70 h(Za.)ll
(100)
(101)
Eq. (101) uses the fact that H2,, a, is unconstrained outside of the source support and that by
assumption supp(pz,, (,)) nsupp(pz. (4,) = 0. To achieve the sup - 5 it then suffices to predict
an 7 * 76 E T. We thus that Eq. (97) holds for d, long as 0 < 5 < We conclude the
proof by noting that there is an infinite possible choices of each of which give rise to a bad target domain.
24

---

Published as a conference paper at ICLR 2022
B.4 AUGMENTATIONS
Proposition 2 shows that the optimal representations for IDG can be learned with augmentations in a self-supervised fashion. Here, we provide formal definitions, assumptions, and proofs.
Definition 6 (Augmenter). An augmenter is a conditional distribution PAIX X : A x X' - [0, 1] from the input space X' to 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).
Definition 7 (Augmentation conditional set). Given an augmenter PA define the augmentation conditional set the set of conditionals of given X:
P" (A x) == (PAIs x E }}
(102)
Similarly, we define the augmentation conditional for domain d:
(A ( == (PA] x supp(px
(103)
These sets clearly countable. Note that the augmentation conditional set 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 striet properness (Gneiting & Raftery 2007).
Assumption 7 (Finite augmentation entropy). We consider the augmenter PA xx such that the entropy of the augmentation A is finite, i.e., H[A] < oc.
Assumption (Cardinalities). We assume that
(104)
This is similar assumption as Assumption
which ensures the existence of optimal representations.
Assumption aggoaii-a augmentation). We assume that the augmentation A is domain- agnostic, i.c. the augmentation conditional set is invariant across domains,
PA = (A) Ad E
(105)
This assumption is generalized from the constant Bayes image assumption (Assumption 5) which guarantees the existence of optimal representations.
Domain-agnostic augmentations essentially ensures that each augmentation conditional PAI= E P ( is seen at least once all domains. If we introduce an equivalence relation ~ as x x' iff PAIæ PAll=' and the equivalence class [xc] := (x' E X x' ~ æ). Under this relation, it is easy to see that the above assumption is satisfied and only if, for all possible equivalence classes [x] E [<''] x' E X). we have that [z]] has intersections with all domains:
(106) Not all augmentations are domain-agnostic. In particular, the standard image augmentations used by typical SSL models like SimCLR are not domain-agnostic, but the text-image augmentations CLIP nearly are, as discussed in the main body (Sec. 4).
[x]] n supp((x d) + 0, E
Assumption 10 (Bayes-preserving augmentation). We assume that the augmentation A is Bayes- preserving, i.c., Va, xf E X,
PA = PAi= fz)
(107)
Under the notion of equivalence relation in Assumption 9 this means that for each equivalence class [z], all x' e [z]] have the same Bayes prediction. Note that most augmentations used in practice like standard image augmentations are Bayes- Sayee-preserving.
Next, we show that under the above assumptions, we can learn optimal representations by maximizing the mutual information I[A; 2] (in the case of log-loss f) 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 optimal for any strict proper losses.
25

---

Published as a conference paper at ICLR 2022
Proposition 6 (Learning optimal representations without labels, wqui.. Proposition 2). an augmenter. Under Assuupttssstto 1 to 10 any encoder pz such that
Let PA x be
PZIX E arg max [[A; 2]
(108)
PZIX
s.t. VdED, D, supp(pzia) = supp(ppz)
(109)
optimal for idealized domain generalization.
Proof. The support match constraint Eq. (109) is equivalent to the support match constraint Eqq. (68)
Thus, Prop. annddThn 2 state that we only need to prove that maximizing the mutual information of A and Z under the support constraint implies that
R Y Z R X].
(110)
We will prove this by constructing an optimal predictor lt" we have that
Since H[A] oc (Assumption' 7)
arg max I[A; Z] arg min H[A) Z]. P
(111)
P2
Note the fact that the conditional entropy is the Bayes risk under the log-loss (Gneiting & Raftery 2007), i.e., H[4 Z] = R [ Z]. By construction, satisfies covariate shift w.r.t. (thus Bayes invariant) since A - X D forms Markov chain. Together fitthmAssuptions Iadd 7 to it means that the optimization problem in Eqs. (108) and [109)] satisfies the assumptions of Prop. with A in place of Y. Thus, optimal encoder satisfies R
[AA 2 = R X]. which leads to
H[A 2 H[A X].
(112)
By Assumption 7 we can invoke Lemma 2 with the fact that X forms a Markove chain to show that for all (I, 2) supp(px.z)
PAI = PAI=
(113)
as the conditional distributions are the Bayes optimal predictors due to strict properness of log-loss. Now, define the following equivalence relation on X,
x ~ x'
PAI= = 2'
(114)
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 definition of maximal invariant Definition
2. Dubois et al., 2021). By Assumption 10 f* is invariant on the equivalence classes [ := (x' E for all (Lemmaa 5, Dubois et al. 2021). Given E supp(pz), we construct h* in the following way. Let be any input point that could have led to this representation and define r, E supp(px [a)
X x
X. Thus, there exists function g
+ A such that f g o M
(=) = (x=)))
(115)
By Eq. (113) we are guaranteed that all supp(px (=) share the same value for f since they are in the same equivalence class. Thus, by the definition of M we have that
M(zz) = ((X) for (X, ~ px.z.
(116)
Therefore,
R₂- Y Z Epy.2 1e h² (Z))
(117)
Epy xpx.2 [e(Y, (2))
(118)
Epy xpx,2 ((YYg((M())))
Eq. (115)
(119)
Epp xPx.2 ((())()(X)))
Eq. (116)
(120)
Epy xpxxa ((Y = (XY,
(121)
26

---

Published as a conference paper at ICLR 2022
C PRACTICAL OBJECTIVES
Proposition 6 provides an objective to obtain the desired optimal representations, compared to Thm. 2 is more practical that it does not require direct access to the labels and that it can use augmentations under appropriate assumptions. There are nevertheless multiple remaining issues for deriving objectives that can be trained with in practice. Specifically, (i) the support constraint is hard to satisfy practice; (ii) mutual information [[A: Z] is hard to estimate from samples (Poole et al. 2019); (ii) 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 1[4; 2] = X]; In the following, will rewrite the objective as the constrained optimization:
pz X E arg min BIZ, YYY, PZIX
(122)
s.t. 1[A: 21 - [[A
(123)
where we introduce the domain bottleneck BIZ, 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. (122) under EEq. [123]implics that the support match constraint holds (and can be achieved by some encoder), which leads to optimal representations for IDG. Different domain bottlenecks will be derived later this section. We can then use Lagrangian relaxation to get the following unconstrained objectives.
arg min PZIX
I[4; + ABIZZ
(124)
The first term can be easily optimized using variational bounds on MI. Throughout the paper, we will use a contrastive variational lower bound which is based on InfoNCE (Oord et al. 2018). Namely, let X be the input sample and A be the 'positive' augmentation sampled from PA x- We then obtain n 'negative' augmentations {A, )i= by first independently sampling { {X; === from the marginal px and then sampling A, from PAIX, It is easy see that the negatives A, follow the marginal pa-
We construct A : (A, AT. ..... A }. Let Z be the representation of X by passing it through the encoder p. := pz | x parameterized by o and s. the critic function parametrized by eb used to score which A' E A the positive augmentation. Then we have the following variational lower bound (Poole et al. 2019):
exp 8xp& (A, (A, Z)
T[A: 2 2 log(n + 1) + Epx.x. log
(125)
EA'EA exp s. (A', Z)
In the case of unconstrained variational families se, P. and infinite samples (n + oc). the above variational bound recovers T[A; Z] up to a constant (see Oord et al. (2018): Dubois et al. (2021) )
Typically the critic separable, i.e., (A, Z) := I (A) h (Z). As discussed in the main body, can be tied with the encoder p. when A = X.
In the following we focus on the second term B[2, X,Y,D) and discuss several choices.
Throughout this section, the function M X N is the maximal invariant defined in Prop. via the equivalence relation defined in Eq. (114)
C.1 MUTUAL INFORMATION BOTTLENECK B X1,,X,Y I[Z
The first bottleneck we consider is so called mutual information (MI) bottleneck B Y, 1[; X], which was introduced by Tishby et al. (2000) 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 1[Z; A). In particular, using the fact that Z - X D forms a Markov chain and the chain rule of MI, have I[2; X] I X, D) 1[2 DJ I[Z; X D). Thus, it not only minimizes 1[2; D]. ie.. matches the representations' distribution across domains, but also minimizes 1[;; X D). i.e., matches the representations' distribution inside domains.
27

---

Published as a conference paper at ICLR 2022
Why The key to show is that minimizing Eq. (122) i.e., arg mrg [[ under I[A; Z|
I[A; X], implies the support match constraint. This can be seen as specific subcase of Dubois et al. is (2021) Corollary 15 with A in place of Y and M(X) induced by PAIX as the proof of
Prop. 6 From the corollary, we know that minp x 1[2; X = H[M( which can be achieved
by any Z s.t. Pz = az = ((x) (x'). With the assumption of domain aggais-a
augmentations (Assumption 9). we have that the set of maximal invariant (M(x) | E supp(px(a)
invariant across domains. Then we directly have supp(pz)a) a) Uxesupp(px (supp(pz
Dresupp(px)suppe supp((z |=) = supp(pz), where we use the fact that x within the same equivalence class has the the same pz ax
How Essentially, we can use any variational upper bound of mutual information. We consider the one used by Variational Information Bottelenck (Alemi al. 2016), i...
log (Z XX)
I[Z; X] = E Px.
(126)
pz (z(
log PZ X)
Exx.2
DRI [enz (2)(()
(127)
2o(Z)
(Z X)
S E Pxxx
log
(128)
go(Z)
= Epx [DxL (p.lpp22 | (x]l9.
(129)
where a variational distribution Tu used to approximate pz and is jointly optimized with P. to minimize the bound. The approximation gap of the bound is (x..p2(())ll Ignoring the constant, the final loss becomes
exp S. (A,
(xi p, 0) == Epx.A.2 log Σ A exp s, (A', 2)
+ (2|
(130)
which recovers the optimal encoder in the case of unconstrained variational families for p.: 9os S.t infinite samples - oo, and any À > -1 (Dubois et al. 2021).
C.2 ENTROPY BOTTLENECK BIZ, 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 dirac delta function for all x E X and we denote by (x) the deterministic encoder s.t. P. (e, (x) x)) 1.
Why the deterministic case, the MI bottleneck becomes the entropy bottleneck because I[X; 2] = H[Z] H[Z| X] H[Z), where we use the fact that HiZ 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,, (xr) = e,p (x)) è M(x) = (xx), which also satisfies 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
H < E- log g. (Z))
(131)
which is the standard variational bound used in neural compression (Ballé et al. 2016 Theis et al. 2017). Putting all together, we have
exp (A, (A, Z)
Lmn (w,, 0,50) = E Exx.A. -log log
ΣvEA (A', - -Alog -Alog
(132)
which also recovers the optimal encoder with unconstrained variational families, infinite samples, and A >>> as with the MI bottleneck. The detialed algorithm is provided in Algoorithm 2) Note that the discreteness of Z could lead to difficulty of gradient- bradeent-based optimization, and we follow Ballé et al. (2016) to add uniform noise to Z as 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.
28

---

Published as a conference paper at ICLR 2022
Algorithm 2 Ent objective
Require: eu. Su: Go. X,
I: Z f (X)
2: A <- sample(PAIx)
3: ((XF,AA))
Lid.
sample(px.a)
4: A + [A] U(AF)
Fi=1
5: Luve f - -log En'Ea exp (2
exp (A.
D --1[A; 2]
6: Lupp f- log go (2)
D H[Z
7: return Lee Lus + Xhuup
C.3 CONTRASTIVE ADVERSARIAL DOMAIN BOTTLENECK X, = D]
The previous two bottlenecks require removing the information of Z (ahout X) as much possible, which seems to be unnecessary since our ultimate goal is to match the support of across domains. Now we introduce a bottleneck B|Z, X, Y, D] 12; DJ which we only seek to remove the informa- tion of Z about the domain D. This is very related to the work on invariant representation learning for domain seneralization//aapttttio (e.g.. Ganin et al.) 2016 Li et al. 2018a). We derive 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 (1)) as with the main body.
Why Similar to the previous analysis, we aim to show that arg minpz| p 1[2 D] under |A: Z] = I[A; X] leads to the support match constraint. Using Eq. (116) we have 1[2; D] = I[2, M (X Xz); DJ I[2, MX);; D] (X); DJ + I[Z; D M(X) where the last equality uses the chain rule of mutual information. Due to the non-ne negaatiivy of (conditional) mutual information, we have that the minimum of 1[2; D] under I[A; 2] I[A: X is (X); D). Then we show the minimum is achievable by constructing the same optimal encoder (X) as the Ent bottleneck which clearly is then easy to show that the support match constraint has hold
satisfies I D M())
when I[Z; D M = 0 by contrapositive. Indeed, suppose that the support constraint does not hold then it must be true that I[2; D M (X)] > O and so the encoder cannot be optimal.
How The typical way of minimizing TZZ D] is derive the variational bound
1|2: D] = H[D] H 2]
(133)
= (const) Pp x [ log PD1 z [D] Z
(134)
2 (const) Ep.. lo (D) Z)]
(135)
where a variational distribution (or domain classifier) 9a is used to approximate PD z and jointly trained maximize the bound. This recovers the domain-adversee
training method as introduced in Ganin et al. (2016). However, this has two potential issues: I) it gives lower bound instead of the desired upper bound on I 12; D): 2) it requires adversarial training which is not stable in practice (Goodfellow 2016 Kodali et al. 2017).
We propose the contrastive adversarial domain (CAD) bottleneck, which is based on the above explicit version but uses variational distribution gu [D] Z) that is tied with other parts of the model, thus no need to learn a domain classifier. Suppose we have access to a set of inputs X, we first introduce a contrastive variational distribution 9x.x (X] Z) of px z
exp (xpss(X22 Z)
I. x (X) Z) Exex exp (X',
(136)
where s, (X, := (X)T is tied with the encoder e,, Note that x, has support over X, and equals Px| z when s_(X Z) x logpx.z 2) and X recovers X. In practice, we use variety of crude approximations. Our first crude approximation is that we use the minibatch of samples, i.c. the independently sampled XX, Ji=1
as X. Now, since PD z can be rewritten as E,x [PD] x using the fact that D - X - 2 forms a Markov chain, we obtain the following variational distribution:
L. x (D) 7) B.g. [ (D) x)]
(137)
29

---

Published as a conference paper at ICLR 2022
which recovers PD| z when 9..x = Px z- Note that PD x is still not available. For our second crude approximation, we use a count estimate PD. x. In particular, we obtain a collection D by taking each X' E X and independently sampling D' from PD x to get D == =D; ) In other words, ( X)))== are all i.i.d. sampled from PD.X- Then we use a count estimate
x x) = E [X, =x, = Σ IX, =
(138)
which is accurate estimate with infinite samples. This leads our final variational distribution:
Ge d (D 2) Σ
gxx (X' |Xx0xx (DIX'
(139)
X
Putting all together we get that the loss:
]
exp 8. (A,
le 10 = EpD.x.A. log ENEA (A',, + >llog
Σ xxx x x") XXXX
(140)
In practice, x X) is typically dirac delta function since it is rare to have the same samples in a minibatch. Thus, in Eq. (139) we only need to sum xxx' Z) over those associated with the same domain label D as L.e. X, == (X; Di- D, [m]] where [n] == (1.....7) n). This
leads to the simplified loss:
exp s. (A, Z) LeADlep, v6 == E E.A... llg ΣveA exps, (A,Z)) Alog
] (141)
ga. Xxxx(X (X' X'EXD
Σ
In practice, we find that the second term that minimizes the log probability leads to numerical instability. Intuitively, this could be seen by the exploding gradient of the function log(p) when p -- O. We thus replace it with log((1 - p) which has the same optima. L.c. in practice we maximize the log of the probablity summed over X-D := X Xp. This reduces Eg. (141) to Eq. (7) described in the main body with detailed algorithm in Algoriihhm I Note that it is easy generalize
Algorithm 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 efficiently in parallel.
C.4 CONDITIONAL CAD BIZ, X,Y, IZ;
The analysis of the CAD bottleneck also implies that we can minimize the conditional mutual information 1[2; D [[)] if we have access to (X). However, since M (X) is typically not available in practice, we consider the special case where M = Y. In particular, this is the
where the labels are available and the supervised augmentations are used (see Fig. 2b). This reduces the bottleneck B[Z, X, Y, D] 1 D Y which is related the conditional version of the lomain-adversarte neural network (Li et al. 2018b). In practice, minimizing T[Z: D YJ could be easier for optimization than 1[; D). as does not require to remove the information that D has about Y. In the following, we derive the conditional CAD (C² AD) bottleneck using a similar idea as CAD.
How In this case, we want to minimize
T[: D D|Y] = H(D| Y] H[D Y
(142)
(const) H(D| Y
(143)
> (const) (- log (DD Z,Y)]
(144)
where qDD Y, variational distribution of Pn 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 PD z,y = Epx z., [PDIX] Specifically, let Y be the label of input X sampled from py x and : (Y= ........ Y, } be the collection of labels obtained by independently sampling the label from | xx' for each X' X. We collect samples associated with the label Y, : (x |Y, Y,i E [n]] and obtain variational distribution of Px z,Y i.c., x=:
cxps, (X, Z)
9x (
2,Y) = Exeex, exps, (X's,
(145)
30

---

Published as a conference paper at ICLR 2022
where we use the same critic s. Z) (X) that is tied with the encoder e, as before, but only take softmax over those samples with the same label Y. For the term PDI we use the same count estimate PD.x in Eq. (138) Then obtain the variational distribution of pD! z,y
X.DY 2,
E Xx.xx( (xx))("12))))xxxx0 X'EXY
x,Y) (x)
(146)
Putting all together we get that the final loss:
xx 2, D.D.D. (D) x3) )
exps 8, (A,
fer v) == EPD.XAYY log
+ Aloo Σ X'yEXy
EN'A cxps,(A', Z)
(147)
Again, since in practice PD.x [D] X) is typically a dirac delta function, the summation in Eq. (146) can be done only over those associated with the same label Y and the same domain label D as X, i.e. Xy.D = {X; Y,, Y, Dj- D, E [n]]. Similarly, instead of minimizing the log of the probability summed over Xy.D. we maximize the log of the probability summed over Xy,-D = Xy XY,D (X; 1Y, = D, * D, [n]]. Finally obtaine the simplified loss:
exp (A,
(e) EPD.X.A..2
llg EA'EA (A,,
Alog X'EXY,
Σ
lx z2
(148)
detailed algorithm is in Algorithm 3
Algorithm 3 conditional CAD AC) objective
Require: Co. 8u, P.X.Y,n
I: Z (X)
sample(PAix)
3: {(D, .X, .4, } sample(pn.x.A.r)
4: X,A ( F.
5: Xy + (X, =Yi
XY, + (x, 11, = Y,D, #DiE [w]
7: Lus f log EA'EA exp (d
expoglA.2)
1-|[A
8: Lupp E- log Ex'eXy, (X
> +[Z:D|||
Ex"'Ex, (X")7 return CAPAD Lune XLun
31

---

Published as a conference paper at ICLR 2022
D EXTENDED RELATED WORK
Provably optimal representations. Many previous work have theoretically studied advantages of representations in various two-stage settings (representation learning followed by standard training of predictors) by bounding downstream performance (c.g.. /c.....en-Jl. BBen-David et al. 2007, Shamir et al. 2010 Saunshi et al. 2019). As learning theoretical bounds can be loose, it is hard to know whether they give the right insights into the problem. Our work instead proves the properties of optimal representations, which ensure best downstream performance. Those properties need to be approximated but give the right insights into what to aim for. This perspective and our proofs were inspired by Dubois et al. (2020) who gives sufficient conditions for optimal representations in supervised learning.
E EXPERIMENTAL DETAILS
E. SCIENTIFIC
In both the scientific setting and the following bridge setting, we consider rather unrealistic setups for verifying our theory where we have access labels from all domains. We can choose to directly minimize the risk R [Y] 2 with the cross-entropy loss (denoted as CE henceafter), or minimize H[A| 2 (i.e., maximize I[A; Z)) with supervised augmentations as in Fig. 2b 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, Khosla et al. (022)). In particular, for 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 al. (2020), two variants of SupCon loss were introduced for solving the issue of multi- mostii-pssitiee depending on whether the summation over multi-positive was located inside (SupCon-In In, Eq. (3) in Khosla et al. (2020) or outside (SupCon- Eq. (2) in Khosla let al.] (2020)) the log. Though Khosla et al. (2020) chose SupCon-Out because it worked better than SupCon-In, we hypothesized that this is because SupCon-Out has implicit bottleneck effect. Intuitively, SupCon-Out upper bounds SupCon-In 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.) (2020) except that no projection was used. Specifically, the temperature was set to 0.1, and normalization was applied when computing the logits.
In the scientific setting, we tried to simulate our theory to the greatest extent. In particular, we had two special considerations detailed below:
Eliminating empirical generalization As theory focuses the idealized domain generaliza- tion that assumes access to population distribution, we considered the setup where the empirical generalization was eliminated. Specifically, we treated the dataset as the population distribution and used the same dataset for training the encoder and training/evaluating the predictor. The ResNet- 18 encoder was trained to 300 epochs without any regularization, using the Adam optimizer (Kingma & Ba 20144) 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-wee 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 the sample weight increases, the target log likelihood (neg. risk) first 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 figure), thus selected the weight close to the turning point and 10 5 seemed to be reasonable for both CE- Base and SupCon- BBsse. Although we did not adaptively select the sample weight for each setup due to the computational cost, the pre-specified sample turned
as shown in Fig. 4 As
32

---

Published as a conference paper at ICLR 2022
Log likelihood
Log likelihood
-2
source Larget
-6
10 9
10
10
10 3
10-1
10 9
10
10
10-3
10-1
sample weight
sample weight
(a) BasB
(b) SupCon-Base
Figure 4: Sweeping the sample weight using CE-Base and SupCon-Base. We selected 10-5 which seemed to be reasonble for both cases.
out be reasonable for all other losses and different À combinations. Furthermore, we also removed regularization when training the linear classifier and initialized the linear weight i.i.d. from from.N (000, 1).
Next, we provide other experimental details for reproducibility:
Implementation of standard augmentations We followed SimCLR (Chen al. 2020 for imple- menting standard image augmentations. For fair comparison between the cases when using standard augmentations (SimCLR) and supervised augmentations (SupCon), kept the total batch size the same and also used the same configurations for computing the SupCon loss, i.e., temperature set to 0.1, no projection, and normalization applied.
Details of Fig. 3c In Fig. 3c we considered different choices of augmentations. The 'Standard' augmentation implementation is described above (Appx. E. The 'Supervised" augmentation was essentially implemented using the SupCon loss described in Anpxx. E.1] For other augmentations considered, we implemented them by dropout inter-domain supervised augmentations Specifically, for each sample in the batch, randomly masked the samples from different domains (i.e., both inter-domain positives and negatives) i.i.d. with the specified dropout probability, while samples within the same domain were always kept. 'IntraDom' and 'ApproxDA' correspond to dropout probability and 0.9, respectively. 'SingleDom'' were implemented by dropout all inter- domain samples with probability except for fixed domain (the 'A' domain of PACS in our case).
SupCon.
E.2 BRIDGE
we aimed to bridge the gap between our theoretical setup practical setup. The main differences from the scientific setups are that the empirical generalization gap is considered and the average-case source predictor used, as detailed below:
In the bridge setting (see Appx. F.2
Incorporating empirical generalization In practice, gnnnriicllg gap should also be considered besides the source-target generalization gap. Thus, we randomly split the PACS dataset 80% training and 20% validation splits for each domain. The training splits were used 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 maximum of 50 epochs with a le-5 weight decay, using the Adam optimizer (Kingma & Ba 2014) with learning rate of 5c-5, batch size of 112 (28 for each domain), and cosine learning rate decay schedule.
Using average-case source predictor Instead of approximating the worst-case source predictor in the scientific setting, we considered the average-case source predictor which closer the common practice. Specifically, we freezed the encoder and trained SVM classifier with L2 regularization on
Herr we have slight abuse use of the phrase 'average-case to distinguish from the 'worst-case' that we use in the scientific setting. In fact, the source predictor could be close to the "hest-case' since the max-margin classifier (SVM) was used.
33

---

Published as a conference paper at ICLR 2022
the source training split. The regularization parameter tuned over (le-4, 1e-3, 1e-2, le-1, 1, lel, 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 (le-3, le-2, le-1, 1, lel] independently for each.
E.3 DOMAINBED
Datasets We used non-MNIST datasets DomainBed that were non-synthetic, including VLCS (Fng et al. 2013), PACS (Li et al. 2017), OfficeHome (Venkateswara et al. (017), Terralncognita (Beery 2018). and DomainNet (Peng et 2019). For each dataset, split it 80%/20%
training/validation set according to DomainBed.
SSL-based models Training For all models based on pretrained SSL models (cither CLIP-based or DINO-based) with finetuning in this experiment, we freezed the pretrained SSL model and added on top I-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 modified described in [Radford et al. (202T) 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 defined data augmentation was used and the temperature for scaling the logits in CAD was fixed to 0.05. We used the Adam optimizer with le-5 weight dacay, and a cosine learning rate decay schedule. The hyperparameter search space is:
single pass over the smallest domain according DomainBed), No
Learning rate: discrete set (le-4, 3e-4, le-3, 3e-3)
Batch size: discrete set (128, 256, 512} for DomainNet and OfficeHome, and {64, 128, 256} for other datasets
MLP dropout: discrete set (O., 0.1, 0.5]
Learning rate warmup: discrete set (True, False)
End-to-end models & Training In Table ] we also included an end-to-end trained model without any pretrained SSL models. We used exactly the same model architecture (the original ResNet- 50, initialized from ImageNet pretrained model), training procedure and evaluation protocal as baselines DomainBed. Importantly, the linear classifier was jointly trained with the encoder, and no refitting 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 [le-5, 1e-3.5]
Batch size: log-uniform over [8, 64] for DomainNet, and [8, 25.5] for other datasets
MLP dropout: discrete set (0., 0.1, 0.5]
Weight decay: log-uniform over [1e-6, le-2]
Temperature: discrete set (0.05,0. 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 first trained the encoder with specified objectives, and then refit the classifier. For datasets except DomainNet, we fitted the SVM classifier and tuned the regularization parameter over (1e-4, 1e-3, le-2, 1e-1, 1, lel, 1e2, 1e3) with source validation selection. Since DomainNet was too large and SVM cannot fit it efficiently, we used the logistic regression classifier which was trained with batch size 512, the Adam optimizer with learning rate 5e-4 and early stopping. Note that an altermative was to just use the linear head fitted when training the representor (as we used CE loss with source labels), and we found this could work better than refitting since the classifier was less overfitted to the source domain. However, we didn't 't do that since wanted stick to the representation learning protocol with two-stage training. We
34

---

Published as a conference paper at ICLR 2022
did that in our end-to-end training setup since wanted it to be compeletely comparable to baselines on DomainBed (which did not do refitting).
Selection of X our experiments, we treated À as a special hyperparamter. For each model, we used the same À selected on PACS on all datasets except DomainNet, because our bottleneck is fairly robust the choice of X. For the large-scale DomainNet dataset, we selected its À individually. The À values chosen for each model were:
CLIP S: 1 on DomainNet and le-2 on other datasets
CLIP L: le-1 on DomainNet and le-2 on other datasets
DINO: le-1 on all datasets
End-to-end ResNet-50: le-5 on all datasets
E.4 LAION
Model We used the CLIP I. model (i.e., CLIP ViT-B/32) with an additional network on top for finetuning. The additional network were two blocks 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), was the largest publicly available image-text text-pait dataset that we could get access to. As we froze the CLIP L. model and only did finetuning, we used the 1 1TTB preprocessed embeddings provided by LAION-400M No further preprocessing was applied.
Training We used the image-text contrastive loss introduced in Radford al. (202)) 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 batch size of 16384 and cosine learning rate decay schedule. The learning rate was tuned over the set [3e-5, lc-4, 3c-4, le-3, 3e-3, 1e-2) and the À value for the Ent bottleneck was tuned over [le-3, 1e-2, 1e-1, 1, lel).
Evaluation For the evaluation the ImageNet- datasets, followed similar procedure in Radford et al. 2021 where linear classifier was fitted on ImageNet using the model representations and evaluated
fitted logistic regression classifier with le-5 L2 regularization on ImageNet training set which was trained with a batch size 512, the Adam optimizer with learning rate 3e-4 and early stopping. Note that this was different from Radford et al. (202T). where logistic regression classifier was fitted 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. (2020) and used their released testbed The evaluation datasets and their abbreviations used in Table 2 were: ImageNetV2 (IN- Recht et al.) 2019), ImageNet-Sketehl (IN-S, Wang et al., 2019), Yourube-BB (YT-BB, Shankar et al. 2019). ImageNet- Vid (IN- Vid, Shankar et al. 2019), ObjectNet (Barbu et al. 2019). ImageNet Adversarial
natural distribution shift datasets. In particular,
(IN-A, Hendrycks et al. (022)) and ImageNet Rendition Hendryckss et al. 2020
*See Sssee/ttssr/abooo//////do0oereer for details. htteee//ithoo./////rrrerrrrttttt
35

---

Published as a conference paper at ICLR 2022
F ADDITIONAL EXPERIMENTAL RESULTS
F.1 SCIENTIFIC
Base
Base
Log likelihood
CAD
Log likelihood
CAD
-2
Ent
Ent
2
-4
-3
-5
-6
10-2
10
10
10*
10-2
10"
102
104
À
À
(a) CE (Y Z)) objectives
(b) SupCon (H[A) Z)) objectives
Figure 5: The worst-case DG performance of Ent bottleneck is more sensitive than CAD
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-En with small À values, it was worse than SupCon- CasCeBeeeseee because the discretization introduced by the Ent bottleneck, which we verified by observing that setting À = 0 lead to similar results.
F.2 BRIDGE
The scientific setup closer our theory than what do practice in that worst-case predictor was considered and empirical generalization gap was ignored. Here we bridged these gaps with 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- caaee) 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, labels to minimize R [Y] 7]. and SupCon for the contrastive loss that uses supervised augmentations minimize H[A] 2]. We will use CAD following experiments unless otherwise specified (chosen with initial experiments). Details in Appx. E.2
will use CE to denote the objective with the cross-entropy loss that
Table 3:
We repeated most empirical analysis (in the scientific setting) in the more practical bridge setting and observed similar results.
Setup
Avg. target acc.
CE-Base
95.9 + 0.5
CE-CAD
96.7 + 0.2
CE-CAD (partial domains)
82.6 + 0.5
SupCon-CAD
96.7 + 0.4
SupCon-CAD (SingleDom)
96.7 + 0.3
SupCon-CAD (ApproxDA)
96.6 1.6 0.3
SupCon-CAD (IntraDom)
96.2 + 0.7
SimCLR-CAD
61.7 + 0.8
36

---

Published as a conference paper at ICLR 2022
Does domain bottleneck improve the average-case DG performance? Though theory fo- cuses 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
What if the representor only has access to source domains? Similar to what we did in the scientific setting, we considered the setup where one single domain is specified 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 TTble 3] which is much worse then CE-CAD. This shows the necessity getting access to target domain information for DG.
What if the representor only has access to domain -again- 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-agnostic augmentations in practice.
Can we use standard augmentations? In Fig. 2. we point out that standard augmentations are not domain-agnostic and thus not suitable for SSL. with objectives. We empirically showed this by using augmentations of SimCLR (see (Apxx. E. for details) with our objectives (SimCLR-CAD In Table 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 (SupCon- CAD) and 'Standard' augmentations (SimCLR-CAM above, we also compared other three aug- mentations as in the scientific section. Specifically, considered the 'SingleDom'' 'IntraDom' and 'ApproxDA' augmentations. As shown in Table 3] SupCon-CAD (SingleDom) and (ApproxDA) maintained the DG performance but SupCon-CAD (IntraDom) was slightly worse (0.5 accuracy drop). We assumed the small gap was due to the specific dataset that we used (PACS). We did the same analysis on VLCS, and SupCon-CAD 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 agnasicagnostic augmentations in practice.
Do standard augmentations affect source performance? Previously, we showed that using stan- dard 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' (Taori al. 2020). 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 we consider the source-targe gap which was 0.2 for SupCon-CAl CAD and 28.4 for SimCLR-CAD, which still verified that the wow-ddmmi--ggnssss
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-CA R-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 DOMAINBED
Full result Table We included the full result Table 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 classifiers across domains. Specifically, we included
IRM (Arjovsky al. 2019). GroupDRO (Sagawa et al. 2019), Mixup (Yan et al. 2020), CORAL Sun & Saenko 2016), MMD (Li et al. 2018a). DANN (Ganin et al.) 2016). CDANN (Li et al. 2018b), and VREx (Krueger et al. 2021), We also included the result pretrained CLIP S model with a zero-shot classifier using text representations (CLIP& S Zero Shot), which demonstrated better DG performance than CLIP with linear probe. But we observed that it was outperformed by CLIP S + CAD.
What is the impact of CLIP pretraining? To ensure that our gains are not only due to novel CAD bottleneck, but the synergy between enforcing support constraint and using desired SSL models, 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).
37

---

Published as a conference paper at ICLR 2022
Table 4: Full results on DomainBed with 'oracle selection' method.
Algorithm
VLCS
PACS
OfficeHome
Terralncognita
DomainNet
ERM
0.
0.3
66.4
53.0 + 0.3
41.3 = 0.1
IRM
0.6
1.1
2.70
0.7
28.0
GroupDRO
0.5
8.1
6..2
5.1
33,4 0.3
Mixup
0.3
0.3
68.0
0.3
39.6
CORAL
0.2
0.5
6..4
52.8
41.8 0.1
MMD
7.1
8.2
6.3
0.4
23.5 9.4
DANN
0.5
8.2
6.8
5..6
38.3
CDANN
79,9 0.2
0.8
6.5
0.6
38.5
VREx
78.1 0.2
87.2 0.6
65.7
51,4
30.1 = 3.7
CAD
78.0 + 0.1
87.3 0.2
67.0 0.5
53.5 + 0.9
41.5 0.1
+ CAD
69.6 0.6
76.1 0.1
56.9 0.5
25.9 1.2
33.6 0.1
CLIP $
81.1 0.5
90.3 0.2
70.6 0.1
29.6
47.7
CLIP S (Zero-Shot)
0.1
9.1
70.4 + 0.2
19.1
46.9 0.0
CLIP Š Base
9.3 0.5
0.1
0.74
46.8 0.2
CLIP S CAD
$2.3 0.3
92.0
71.9 0.2
36.2
48.8 0.1
CLIP L
80.7 0.4
93.7 0.8
79.6 + 0.1
36.9 + 0.6
52.8 0.1
CLIP L CAD
81.6 0.1
94.9 - 0.3
80.0 0.2
40.6 + 1.1
53.7 0.1
Why 'oracle" selection? In the main body, we provided the results with 'oracle selection' which was the closest to our theory among the model selection methods DomainBed (in the sense that we needed target domain information to achieve IDG). Here, we also provided results with 'source validation' selection Table 5 Source validation selection relies on the assumption that source and target data follow similar distributions (Gulrajani & Lopez-Paz 2021) thus source and target accuracy are highly correlated, which 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 overfit the source domain, thus possibly leads to worse performance on the target domain. This can be probed by the fact that the finetuned CLIP models (CLIP + Base or CLIP + CAD) were generally worse than the original CLIP model:
Selecting model with source validation accuracy tends 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 good indicator of target accuracy thus its result has larger variance.
F.4 LAION
Evaluation results on DomainBed We included the evaluation results of trained models on Do- mainBed Table 6) where we followed exactly the same linear evaluation protocal discussed in
Appx. E.3 We observed similar results as Table 2) the CLIP model trained with the Ent bottle- neck on LAION (Tuned w/ Ent) outperformed the one without (Tuned w/o Ent) on all DomainBed datasets, but slightly underperformed the original CLIP L model (which might be due to quality of the LAION-400M dataset).
38

---

Published as a conference paper at ICLR 2022
Table 5: Results on DomainBed with 'source validation' selection. Source validation selected model tends overfit more to the source domain and diminish the effect of bottlenecks.
Algorithm
VLCS
PACS
OfficeHome
TerraIncognita
DomainNet
ERM
77.5
0.4
85.5
66.5 0.3
46.1 1.8
40.9 0.1
IRM
78.5 + 0.5
83.5 0.8
64.3 + 2.2
47.6 + 0.8
33.9 - 2.8
GroupDRO
76.7 0.6
8.8
6.70
43.2
= 0.2
Mixup
0.6
0.6
6.3
47.9
39,2 = 0.1
CORAL
0.6
0.3
6.3
4.0
41.5 0.1
MMD
0.9
0.5
6.1
1.6
23.4 9.5
DANN
7..6
8..6
6.69
0.5
38.3 0.1
CDANN
77.5 + 0.1
0.9
1.3
1.6
38.3 0.3
VREx
78.3 + 0.2
84.9 +
6.6
46.4 0.6
33.6 -
CAD
78.0 + 0.5
85.2 + 0.9
67.4 0.2
47.3 2.2
41.0 0.1
DINO + CAD
68.9 0.9
75.4
56.4
23,6 + 1.2
31.0 =1.022 2.3
CLIP $
81.1 + 0.5
90.3 + 0.2
70.6 + 0.1
29.6 + 0.8
47.7 + 0.0
CLIP Š (Zero-Shot)
80.9 + 0.1
91.8
70.4 0.2
19.1 0.1
46.9 0.0
CLIP S Base
0.44
8.76
70.4
3.2
44.6 1.6
CLIP S CAD
81.2 + 0.6
9.6
70.5 0.3
30.3 0.9
45.5 2.1
CLIP L
80.6 + 0.7
9.
79.4
37.5 + 0.7
50.1 1.1
CLIP L+ CAD
80.8 + 0.7
9.7
79.7 0.2
37.4 1.2
51.7 1.4
Table Finetuning CLIP on LAION with an entropy bottleneck performs better on DomainBed than finetuning without.
Algorithm
VLCS
PACS
OfficeHome
TerraIncognita
DomainNet
CLIP L
80.7 + 0.4
93.7 + 0.8
79.9 + 0.1
36.9 + 0.6
52.8 + 0.1
Tuned w/o Ent
79.2 + 0.7
93.4 + 0.3
77.2 + 0.5
36.1 0.4
51.2 - 0.1
Tuned w/ Ent
80.7 + 0.4
94.3 + 0.8
78.2 + 0.2
36.8 + 0.4
52.2 + 0.1
39


Reviews:
[{"review_id": "vkiAsVCXDoX", "reviewer": "Reviewer_1s9j", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "*Strong point(s) of the paper*\n\n- The paper contributes some lucid new theoretical insights to the problem of encoder design that is conducive to trained models which are domain robust, and does an effective job of linking this basic theory up with practical training procedures.\n\n- The writing is solid. High-level theoretical points are quite clear in the main paper (details in the appendix), the relation to existing work is described, and the experimental results are distilled into new questions/insights which are easily parsed by the reader.\n\n*Weak point(s) of the paper*\n\n- The theoretical results are stimulating, but they are quite idealized and leave a substantial gap between what we can know and check in practice, and what we want to know formally. For example, while the notion of a domain-covering augmentation is very simple mathematically, but one expects that evaluating whether we have something close to such an augmentation in practice becomes quite challenging when tasks are difficult to describe in words or in simple visual terms.\n\n*Questions*\n\n- Most of the notation is immaculate, but I would like to confirm about  $\\\\mathrm{R}_{h}^{d}[Y \\\\,|\\\\, Z]$. Is the assumption that $Z \\\\sim p\\_{Z \\\\,|\\\\, X}$ and the expectation is now with respect to $Y$, $X$, and the randomness in sampling $Z$? (instead of just $Y$ and $X$)", "summary_of_the_paper": "At a high level, this work considers the problem of designing machine learning systems that generalize well even when the \"domain\" (i.e., the data-generating process) under which the system is tested (the \"target\") does not match that under which it was trained (the \"source\"). More specifically, the authors look to provide an answer to the question of what makes a \"good\" representation or encoding (a stochastic transformation of the original features), from the perspective of achieving a small risk (expected loss) on the target domain, assuming the learner only has data (and thus representations) from the source domain.\n\nTo this question, they provide one answer by showing that in a generalized covariance shift scenario, a natural notion of representation optimality (their \"IDG optimality\") is only possible if a representation minimizes the (best) risk on the average target distribution, while ensuring that the support of the representations is constant across domains. This is obviously a strong requirement, and they show that one cannot expect to satisfy it without domain knowledge going beyond the source.\n\nThey also show that there is some hope for special cases of domains in which we have augmentations that preserve label information and \"cover\" the support of the input distribution across all domains. In such a case, the authors show that their strong optimality requirement can be satisfied by designing the encoder to maximize the mutual information with the augmented data. Relaxing this objective leads to objectives that are more practical, and the authors complement their theoretical analysis with a rather in-depth set of empirical tests that evaluate the efficacy of their methodology for representation learning.", "summary_of_the_review": "The problem is fundamental to machine learning, the theory developed by the authors is stimulating and to the best of my knowledge new, the empirical analysis is informative, and the writing throughout is very strong. I vote to accept.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "VHwhogyKIuw", "reviewer": "Reviewer_hrZK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Novelty of the contribution: It seems to me that the main contribution of the paper is to provide informative formulations that characterize the optimal representation, which provide guidance on designing a representation learning scheme under covariate shifts. This is a potentially important result, though I am not sure how much this distinguishes from the literature. From the related work discussed in Section 5, it has been known that it is sufficient to match the support across domains. The authors emphasize that the condition they proposed is also necessary, but how non-trivial is the necessity? Intuitively, if the representation of the source and target could have different support, then it is hard to guarantee good performance in the target. Perhaps the main contribution here is to make this intuition concrete, or perhaps crafting the minimax optimality in IDG? The authors might need to clarify and provide a more convincing argument to support their novelty. \n\nRequiring access to lots of information: As discussed by the authors, one limitation of their framework is to assume access to an overly demanding amount of information that seems not practical. Would a more meaningful framework be to look at representation trained only on the source domain (as in the setting of Section 6.2 which is more practical) and some degree of information from other domains, and consider the characterization of optimality under this setting. In this latter setting, is there a way to generalize Theorem 1 and if so how much would the characterization change? (I understand a full answer may be difficult, but at least some indications to boost my confidence would be helpful; see also the next point that has a similar flavor)\n\nAssumption on the domain-covering augmenter: Proposition 2 relies on the existence of a domain-covering augmenter. It is argued that image-text augmentations are nearly domain-covering, but it seems hard to expect that they are strictly domain-covering. So there appears a gap in the theory: The theory requires a strict domain-covering augmenter, while in practice this can only be nearly satisfied. Nevertheless, I think the proposition gives insights how the augmenter could be useful..\n\nNumerical results: It seems that the benefit of domain bottlenecks is not very significant. In Table 1, sometimes the improvement of CLIP S/L + CAD over CLIP S/L is not significant enough, e.g., 94.7\\pm 0.4 of CLIP L+CAD vs 93.7\\pm 0.8 given in the PACS column. The only significant improvement is seen in the DomainNet column where the improvement is about 1%. \n\nA few minor issues: \nPage 1 the first summary point: it is said that the paper provides “all” objectives that have the correct optima. Why couldn’t there exist other objectives not given in this paper?\nPage 3, the paragraph after Thm 1: it seems that the cross references to Thm 2 should be corrected to refer to Thm 1 (and similarly for other instances).\nPage 3, some minor typos in the paragraph after Proposition 1: e.g., “to to”, “worst IDG”.\nTable 1: it would be better to explain the level of the confidence interval (CI) in each entry, e.g., are they 95% CIs?", "summary_of_the_paper": "The paper studies representation learning under covariate shift. Under the IDG setting and the proposed assumptions, the paper gives a so-called variational characterization of the optimal representation. This characterization shows that the optimal representation should remain discriminative while has the same support across domains. It is argued that without any target information, no representation can do uniformly well over constant representation, thus supporting the necessity of target knowledge. The paper provides practical objectives of the proposed variational characterization by self-supervised learning using domain-covering augmentations. The proposed representation learning scheme is tested on several datasets.", "summary_of_the_review": "The paper provides insights on the optimal representation under covariate shifts, which is a potentially important problem. I think the authors did concrete work to develop the theory and set up the numerical experiments, though there are also limitations discussed in the main review. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "0w92bT1frNN", "reviewer": "Reviewer_b5Jo", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength\n\n- Theoretical analysis on the conditions for idealized domain generalization\n\nWeaknesses\n\n- Writing is poor\n- Empirical improvements come from external datasets rather than the method itself\n\nThe paper focuses on learning representations for domain generalization. The main idea is to learn a representation that matches the support across domains while maintaining discriminative information (Thm.1). Starting from idealized domain generalization risk (page 2), the paper goes through a long path of derivations and assumptions to derive practical objectives. The theoretical understanding is interesting, but the exposition is poor and empirical results may not be as promising as it seems.\n\n1. Technicalities (and required clarifications for some parts)\n\n- It is surprising to see that the current version does not discuss hypothesis class in question. For example, after (2), what's the candidate hypothesis set for h (argmin_h) when defining the set of source risk minimizers? The hypothesis class usually plays a central role in domain adaptation (e.g., Cortes et al. (2011; 2019)). Additionally, a motivating example would be great for introducing IDG risk and why we should minimize it.\n- Page 3 mentioned different scenarios of Assumption 1 (GCS) for different losses like log-loss, 0-1 loss and MSE. Can you elaborate or provide some derivations?\n- After Thm.1, it is claimed that \"risk minimization using a single domain is as good as performing risk minimization on all domains\". However, both R_{IDG} and R involve expectations over domains or pairs of domains. None is an optimization over a single domain.\n- It is not clear how \"the key requirement\" can be satisfied by the given augmentation. Specifically, how to design such augmentation, other than CLIP?\n\nRef:\n\n- Cortes, C. and Mohri, M. Domain adaptation in regression. In International Conference on Algorithmic Learning Theory, 2011.\n- Cortes, C., Mohri, M., and Medina, A. M. Adaptation based on generalized discrepancy. The Journal of Machine Learning Research, 2019.\n\n2. The writing is poor. Some parts of the paper are challenging to follow\n\n- Notations are not consistent. Sometimes d_t (after Eq.(1)), sometimes D_t (Eq.(2), Theorem 1 etc)\n- This paper fails to properly introduce the necessary background on CLIP.\n- Exposition in Sec.4.1 is poor. It would be much better if the paper can provide a running example about A, X and Z.\n    - A is said to be a RV. If I understand it correctly, an augmentation A can take value in, for example, {image rotation, cropping, mirroring} as shown in Fig.2a. Then p_{A|x} is the distribution of the above three augmentations conditioned on x. It is not clear why p_{A|x}=p_{A|x'} implies they have the same Bayes predictions. If two images are destined to be mirrored with probability one, then they must have the same prediction? If A is supposed to be the actual mirror image of x, then it is more than an arbitrary RV and should be specified more precisely. It is unclear where this \"key requirement\" comes from.\n    - What is the task/classes in Fig.2? Is it classifying \"floppy ear\" versus \"pointed-eared\"? If so, the augmentations in Fig.2b look arbitrary and do not make sense.\n    - \"Enforcing the support constraint for augmented data (X, A)\", but the constraint of Eq.(5) does not involve the augmentation A.\n- Sec.4.2\n    - What is the definition of the domain bottleneck B[Z, D]? And why it is related to the support constraint in Eq.(5)?\n    - What is \\mathcal{A} after Eq.(7)?\n\n3. Experiments\n\n- What is the \"fixed domain\" for SingleDom? And why it is DC?\n- For approximate DC, what's the result if SingleDom is used instead of Supervised?\n- The improvement in Table 1 seems to be mainly due to the fact that CLIP is pre-trained using a much larger dataset. Thus the proposed methods are only marginally better than vanilla CLIP if I understand it correctly. The results on TerraIncognita can justify this. To add some discussion, how can we ensure domain-covering in practice when even CLIP may not have sufficient coverage?", "summary_of_the_paper": "This work focuses on learning representations for domain generalization. It proposes to minimize the idealized domain generalization risk (IDG risk; defined on page 2). Under several assumptions (importantly, domain covering augmentations), the IDG risk objective can be altered (Thm.1 -> Prop.2 -> Eq.(6)) to a more practical one. With additional variational techniques (Sec.4.2), the proposed algorithms approximately minimize conditional entropy and domain bottleneck in (6). Experiments on standard benchmarks show that the proposed method can out-perform SOTA alternatives. ", "summary_of_the_review": "In summary, the paper theoretically analyzes the conditions under which domain generalization would be possible. However, the writing is not clear at places and the empirical results are not very convincing.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper aims at characterizing conditions for optimal representations required for the domain generalization problem under covariate shift. Under the Idealized Domain Generalization (IDG), the paper provides a variational characterization of the optimal representation and shows a number of intriguing results: (i) optimal representation should remain discriminative  across domains, (ii)  the representation’s marginal support needs to be the same across source and target. (ii) It is also shown that without any target information, no representation can do uniformly well over constant representation, thus supporting the necessity of target knowledge. Finally, the paper provides practical objectives of the proposed variational characterization by self-supervised learning using data-augmentation with experimental results. 

The reviewers had raised a number of concerns, many of which were alleviated through the responses provided by the authors. All in all, in the discussion period, the reviewers indicated the novelty of the results and their importance in learning good representations for the domain generalization problem. The reviewers still had reservations about the following points, and I strongly recommend to the authors to address these points in the revised version (please see the reviews for more details): 

(i) For the experiments, it is clear that there is a noticeable gap between the derived theory and the experimental results. It was argued that alt-text augmentation of CLIP is one of the practical choices for (approximate) domain-agnostic augmentation, but it is difficult to verify this.  

(ii) The empirical gain for the proposed method over CLIP is not very significant. As shown in Table 1 (or the complete results in Table 4), the performance of the proposed method is tightly related to the original CLIP method. In the case where the pre-trained representations of CLIP fail to cover the target set (see the TerraIncognita column in the table), the proposed method can be much worse than other alternatives. Thus I'm worried about the usefulness of the proposed method in practice.

(iii) One of the reviewers had asked about the importance of the necessity condition and its implications (e.g. in practice). Please, make sure to address this in the final version. 

Also, there has been several recent works on learning disentangled representations for the domain generalization problem, using e.g. weakly-supervised approaches (Matsuura et al, '20), or model-based approaches (Robey et al, '21). It would be interesting to see how the results/ideas in this paper would connect to/improve those settings.
Decision: Accept (Poster)