Year: 2021

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 learning methods that only use unlabelled data and augmentations to train robust representations. Our objectives achieve state-of-the-art results on DomainBed, and give insights into the robustness of recent methods, such as CLIP.

text: 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 learning methods that only use unlabelled data and augmentations to train robust representations. Our objectives achieve state-of-the-art results on DomainBed, and give insights into the robustness of recent methods, such as CLIP.

heading: 1 INTRODUCTION
text: It is hard to build machine learning (ML) systems that are robust to distribution shifts between a source (train) and target (test) domain. One promising approach to domain generalization (DG) is 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) (Gulrajani & 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-David et al., 2007; 2010a; Zhao et al., 2019; Johansson et 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 representations on which source risk minimizers are guaranteed to generalize across distribution shifts that preserve the Bayes predictor. We work in an idealized DG (IDG) setting; we assume that a learner has access to the source population risk. Our variational characterization implies that it 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 all 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-specific information. Text descriptions of images are examples of such augmentations, as they are informative for many downstream classification tasks, but they remove a lot of domain-specific information. With such augmentations, we design practical self-supervised learning (SSL) objectives for learning robust representations. Our objectives give insights into the robustness of CLIP (Radford et al., 2021), and lead to improved CLIP-based representations that achieve state-of-the-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 is impossible to learn useful representations without accessing target information; • provide practical objectives to learn optimally robust representations using specific augmentations; • get state-of-the-art results on typical domain generalization benchmarks.

heading: 2 BACKGROUND: DOMAIN GENERALIZATION AND REPRESENTATIONS
text: We are interested in predictions that are robust across distribution shifts. We formalize this using domain generalization (DG) language. Given a distribution pX,Y | ds over inputs x ∈ X and labels y ∈ Y from the source domain ds ∈ D, we select a predictor f : X → Γ. The predictions γ ∈ Γ could for example be labels or distributions over labels. Despite being selected on the source domain, we would like f to achieve a small expected risk with respect to a loss function ` : Y × Γ→ R≥0,
Rdf [Y |X] := EpX,Y | d [`(Y, f(X))] , (1) on a distribution pX,Y | d from a target domain d = dt ∈ D, which is somehow related to ds. 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 : Z → Γ from representations Z to labels Y using standard risk minimization. The goal is to design a robust representation Z, so that predictors h trained to minimize the source risk Rdsh [Y |Z] also achieve low target risk Rdth [Y |Z]. Many methods have been proposed to try to learn such Z, e.g., by enforcing domain invariance of the marginal pZ | d (e.g., Ganin et al., 2016). Still, many of these proposals are not sound (Zhao et al., 2019; Johansson et al., 2019). Furthermore, they rarely outperform source empirical risk minimization (ERM) in practice (Gulrajani & Lopez-Paz, 2021).

heading: 3 OPTIMAL REPRESENTATIONS FOR DOMAIN GENERALIZATION
text: 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 ,Z,Y,D are discrete; formal statements and proofs are in Appxs. A and B.

heading: 3.1 DEFINING OPTIMAL REPRESENTATIONS FOR IDEALIZED DOMAIN GENERALIZATION
text: We want to evaluate the quality of a representation Z of X . In our IDG, the learner is given a random source Ds; she selects any source risk minimizer; and is scored according to her risk on a random target domain Dt. To give uniform guarantees while reflecting the uncertainty over the source-target pair (Ds, Dt), we measure the quality of Z as the expected risk of the learner’s worst-case choice.
Definition. 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.,
RIDG [Y |Z] := EpDs,Dt
[ sup
h∈H∗Ds RDth [Y |Z]
] (2)
whereH∗Ds := arg minh R Ds h [Y |Z] are the source risk minimizers. We call a representation Z∗ (or its encoder) optimal for IDG if it minimizes the IDG risk: pZ∗ |X ∈ arg minpZ |X RIDG [Y |Z].

heading: 3.2 CHARACTERIZING OPTIMAL REPRESENTATIONS FOR IDG UNDER COVARIATE SHIFT
text: The IDG risk is useful to evaluate representations, but gives few insights into IDG and is impractical to optimize due to 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 predictor f∗, which minimizes the expected risk over domains, i.e., f∗ = arg minf EpDt [R Dt f [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 ∈ D we consider are related by the following assumptions: 1. Generalized covariate shift. All domain-specific risk minimizers f ∈ arg minf [Rdf [Y |X]] are
equal to the Bayes predictor f∗ on their support, i.e., f(x) = f∗(x) for all x ∈ supp(pX | d). 2. Invariance of Bayes predictions. The set of Bayes predictions is the same for all domains, i.e.,{
f∗(x) |x ∈ supp(pX | d) } ={f∗(x) |x ∈ X}.
Under review as a conference paper at ICLR 2022
<latexit sha1_base64="nv5kgG6Q4MLitfs23A69VllAD+g=">AAAB6nicdVC7SgNBFJ2NrxhfUUtBBoOQapldspoUYsDGMkHzkGQJs5NJMmT2wcysEJaUljYWithaW+c77PwGf8JJoqCiBy4czrmXe+71Is6kQujNSC0sLi2vpFcza+sbm1vZ7Z26DGNBaI2EPBRND0vKWUBriilOm5Gg2Pc4bXjDs6nfuKZCsjC4VKOIuj7uB6zHCFZaurg6QZ1sDpnILjoFGyLTdlDJKmniIKt0VICWiWbInb5Mqu83+5NKJ/va7oYk9mmgCMdStiwUKTfBQjHC6TjTjiWNMBniPm1pGmCfSjeZRR3DQ610YS8UugIFZ+r3iQT7Uo58T3f6WA3kb28q/uW1YtUrugkLoljRgMwX9WIOVQind8MuE5QoPtIEE8F0VkgGWGCi9Hcy+glfl8L/Sd02rYLpVFGunAdzpMEeOAB5YIFjUAbnoAJqgIA+uAX34MHgxp3xaDzNW1PG58wu+AHj+QO7hpHM</latexit>
<latexit sha1_base64="8sdK5eqgt6vtuNYDok13axAjw7g=">AAAB6nicdVC7SgNBFJ31GeMrainIYBBSLbNxF00hBmwsEzQPSZYwO5lNhsw+mJkVwpLS0sZCEVtr63yHnd/gTzhJFFT0wIXDOfdyz71ezJlUCL0Zc/MLi0vLmZXs6tr6xmZua7suo0QQWiMRj0TTw5JyFtKaYorTZiwoDjxOG97gbOI3rqmQLAov1TCmboB7IfMZwUpLF1cnVieXR6aDrJJTgsgsOqhk25ogC6FDG1ommiJ/+jKuvt/sjSud3Gu7G5EkoKEiHEvZslCs3BQLxQino2w7kTTGZIB7tKVpiAMq3XQadQQPtNKFfiR0hQpO1e8TKQ6kHAae7gyw6svf3kT8y2slyj92UxbGiaIhmS3yEw5VBCd3wy4TlCg+1AQTwXRWSPpYYKL0d7L6CV+Xwv9JvWhatulUUb5cADNkwC7YBwVggSNQBuegAmqAgB64BffgweDGnfFoPM1a54zPmR3wA8bzB632kcI=</latexit>
<latexit sha1_base64="NzJ/0OJqTfPqJLiM8d9EWMJJFFU=">AAAB8HicbVDJSgNBEK1xjXGLy83LYBA8SJgRRY8BLx4jmEWTYejp9CRNunua7h4hjPkKLx5c8OrnePMX9CfsLAdNfFDweK+KqnqRZFQbz/t05uYXFpeWcyv51bX1jc3C1nZNJ6nCpIoTlqhGhDRhVJCqoYaRhlQE8YiRetS7GPr1O6I0TcS16UsScNQRNKYYGSvdyDC7vW+HZhAWil7JG8GdJf6EFMu73teL+JaVsPDRaic45UQYzJDWTd+TJsiQMhQzMsi3Uk0kwj3UIU1LBeJEB9no4IF7YJW2GyfKljDuSP09kSGudZ9HtpMj09XT3lD8z2umJj4PMipkaojA40VxylyTuMPv3TZVBBvWtwRhRe2tLu4ihbCxGeVtCP70y7OkdlzyT0qnVzaNIxgjB3uwD4fgwxmU4RIqUAUMHB7gCZ4d5Tw6r87buHXOmczswB847z9UpJRr</latexit>
<latexit sha1_base64="un3pkM+QCvsHMLw7+AWvfpJS2xM=">AAAB8HicbVDJSgNBEK1xjXGLy83LYBA8SJgRRY8BLx4jmEWTYejp9CRNunua7h4hjPkKLx5c8OrnePMX9CfsLAdNfFDweK+KqnqRZFQbz/t05uYXFpeWcyv51bX1jc3C1nZNJ6nCpIoTlqhGhDRhVJCqoYaRhlQE8YiRetS7GPr1O6I0TcS16UsScNQRNKYYGSvdyDC7vW+HehAWil7JG8GdJf6EFMu73teL+JaVsPDRaic45UQYzJDWTd+TJsiQMhQzMsi3Uk0kwj3UIU1LBeJEB9no4IF7YJW2GyfKljDuSP09kSGudZ9HtpMj09XT3lD8z2umJj4PMipkaojA40VxylyTuMPv3TZVBBvWtwRhRe2tLu4ihbCxGeVtCP70y7OkdlzyT0qnVzaNIxgjB3uwD4fgwxmU4RIqUAUMHB7gCZ4d5Tw6r87buHXOmczswB847z9TH5Rq</latexit>
Generalized covariate shift (GCS) ensures that f∗ is simultaneously optimal on all domains. For log-loss ` it recovers standard covariate shift, i.e., pY | x,d = pY | x. For other losses, GCS is weaker, e.g., it only requires invariance of most likely labels for 0-1 loss, and of conditional expectations for MSE. Invariance of Bayes predictors is 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 on target samples x that are also in the source. Thus, IDG optimal representations are exactly those that (a) have the same support in Z for all domain, and (b) retain GCS from Z without sacrificing the ability to predict Y , which can be ensured by minimizing the risk from Z. See Fig. 1.
Theorem 1. Under our assumptions, an encoder pZ∗ |X is optimal for IDG if and only if it minimizes the risk R [Y |Z] := infh EpDt [ RDth [Y |Z] ] while matching the support of Z across domains, i.e.,
pZ∗ |X ∈ arg max pZ |X R [Y |Z] s.t. ∀ d ∈ D, supp(pZ | d) = supp(pZ) (3)
Moreover, such encoders exist and their IDG risk is the Bayes risk RIDG [Y |Z∗] = R [Y |X].
Theorem 1 provides an 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 | d instead of its support (e.g., Ben-David et al., 2010a) which is the focus of most DG methods (e.g., Ganin et al., 2016). Previous conditions are nevertheless generally neither necessary nor achievable. To our knowledge, Thm. 1 is the first characterization of necessary and sufficient conditions for IDG optimal representations Z∗. Theorem 1 thus gives better insights into IDG and provides a framework for deriving the least stringent objectives for optimal IDG.
The risk minimization (Eq. (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 or additional assumptions it is provably impossible to beat even constant representations.
Proposition 1 (No free lunch for IDG). Let ds be any source domain, Zds be any representation chosen on source ds, and C ∈ Z be a constant representation. Under minor assumptions, for every “good” target domain outside the source’s support on which Zds outperforms C for IDG, there are many “bad” target domains on which Zds is strictly worse than C. Formal statement in Appx. B.3.
Proposition 1 shows that target knowledge is necessary for learning useful representations in IDG. This may explain why previous DG methods have been unable to outperform ERM in standard benchmarks (Gulrajani & Lopez-Paz, 2021): the knowledge they have access to is insufficient to generalize. Taken together, Prop. 1 and Thm. 1 say that either you have access to target domains dt, in which case you can achieve an IDG risk that matches supervised learning, or you do not access dt, in which case any representation learning algorithm can achieve worse IDG risk than a constant.

heading: 4 LEARNING REPRESENTATIONS UNDER COVARIATE SHIFT

heading: 4.1 SELF-SUPERVISED LEARNING USING DOMAIN-AGNOSTIC AUGMENTATIONS
text: Our characterization of optimal representations for IDG (Thm. 1) requires labeled data from all domains, which is impractical. We show how this can be overcome with self-supervised learning
(SSL), which is a technique for training representations without direct access to 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., 2021).
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, if any samples x, x′ ∈ X have the same augmentation conditional pA | x = pA | x′ , then their Bayes predictions must be the same f∗(x) = f∗(x′). With such A, one can learn an encoder that minimizes the risk R [Y |Z] by instead maximizing mutual information I[A;Z]. Intuitively, if Z has all augmentation information, then it must have information about the conditional pA |X , and thus the Bayes prediction f∗(X).
This suggests learning optimal representations for IDG by replacing Eq. (3) with a maximization of I[A;Z]. Unfortunately, fully optimizing I[A;Z] w.r.t. pZ |X is not generally possible under 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 ∈ supp(pX | d) } = { pA | x |x ∈ X } .
Proposition 2. Let pA |X be a domain-agnostic augmenter. Then any optimal solution pZ∗ |X of the following objective is optimal for IDG:
pZ∗ |X ∈ arg max pZ |X I[A;Z] s.t. ∀ d ∈ D, supp(pZ | d) = supp(pZ) (4)
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 `, 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 domain D. See Fig. 2a.
A practical choice of augmentation that is nearly domain-agnostic, is a mapping from images to text descriptions, as with CLIP (Radford et al., 2021) 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 domain-agnostic, since images from different domains (e.g., sketches and photos) but similar semantics are often mapped to similar descriptions.1 (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. We note that many other multi-modal augmentations, e.g., audio-video (Wang et al., 2021), would also likely be domain-agonostic, which one can explore in practice.
In practice, even the domain information D is usually unknown. One can nevertheless still optimize (Eq. (4)) by replace the support constraint with a stronger one that does not rely onD e.g., minimizing I[Z;X] (see Sec. 4.2.2), . This highlights the potential of Prop. 2: if one can find a large source of inputs X and domain-agnostic augmentations A (e.g., the 400M image-text pairs of CLIP) then one can, in principle, learn optimal representations for IDG on any downstream task Y that A preserves.
1Although 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.

heading: 4.2 PRACTICAL OBJECTIVES
text: We now design practical objectives for learning optimal representations without labels. Proposition 2 does provide an objective but it is impractical as it involves constrained optimization. We can nevertheless convert it to the following unconstrained objective by using a Lagrangian relaxation and introducing a domain bottleneck B[Z,D] that enforces support match,
arg min pZ |X
H[A |Z] + λB[Z,D] , (5)
where H[A |Z] replaces I[A;Z] = H[A]−H[A |Z] as H[A] is a constant with respect to the encoder. Eq. (5) is a valid reformulation of Prop. 2 as long as minimizing B[Z,D] while maximizing I[A;Z] enforces the support constraint in Eq. (4). Below, we provide different choices of such B[Z,D] each of which results in a different SSL objective. In practice, however, terms in Eq. (5) 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 a deterministic encoder eϕ : X → Z for the rest of the paper. Detailed derivations are in Appx. C. For both practical objectives we use an upper bound on H[A |Z] ≤ EpA,Z [− log q(A |Z)], where q is the contrastive variational distribution as in InfoNCE (Oord et al., 2018), which is standard in SSL. Specifically, for a sample X , we construct X :={X,X−1 , . . . , X−n } where each X−i are i.i.d. sampled from pX . Then we obtain a collection A :={A+, A−1 , . . . , A−n} of one positive augmentation A+ and n negatives A−i by independently sampling an augmentation from pA |X′ for each example X
′ ∈ X. InfoNCE then uses a critic sψ to score how likely each A′ ∈ A is to be positive, i.e.,
H[A |Z] ≤ EpX,A,Z [− log qψ,A(A+ |Z)] where qψ,A(A+ |Z) := exp sψ(A +, Z)∑ A′∈A exp sψ(A ′, Z) . (6) When A = X , one can tie the critics’ and encoders’ parameters by passing augmentations through the encoder and taking an inner product, i.e., sψ(A,Z) := eϕ(A)TZ.
Many previous DG regularizers (e.g., Ganin et al., 2016; Li et al., 2018b;a) could be valid domain bottlenecks. In the following, we discuss two possible B[Z,D] , the first of which is novel.

heading: 4.2.1 CONTRASTIVE ADVERSARIAL DOMAIN BOTTLENECK (CAD)
text: Algorithm 1 CAD objective Require: eϕ, sψ, D,X, n 1: Z ← eϕ(X) 2: A+ ← sample(pA |X) 3: { (D−i , X − i , A − i ) }n i=1
i.i.d.←−− sample(pD,X,A) 4: X,A←{X} ∪ { X−i }n i=1 ,{A+} ∪ { A−i }n i=1
5: XD ←{X} ∪ { X−i |D − i = D, i ∈ [n] } 6: Laug ← − log exp sψ(A +,Z)∑
A′∈A exp sψ(A ′,Z) . H[A |Z]
7: Lsupp ← log ∑ X′∈XD exp eϕ(X ′)TZ∑
X′′∈X exp eϕ(X ′′)TZ
. I[Z;D]
8: return LCAD = Laug + λLsupp Our first domain bottleneck minimizes B[Z,D] = I[Z;D], which enforces support match using a KL divergence. Dropping constants w.r.t. Z we thus aim to maximize H[D |Z]. Domain-adversarial neural network (DANN, Ganin et al., 2016) does so by ensuring that a domain classifier qφ cannot predict domains from representations, i.e., it maximizes EpD,Z [− log qφ(D |Z)] ≥ 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 challenging in practice.
To overcome these issues, we construct q(D |Z) without introducing additional parameters and with a bound that is tight with enough samples. In short, using the equality pD |Z = EpX |Z [ pD |X ] , we set our variational distribution to qϕ,X,D(D |Z) = Eqϕ,X [p̂(D |X)], where p̂ is a count estimate of pD |X , and qϕ,X(X |Z) is a contrastive family similar to Eq. (6) but with critic eϕ(X)TZ, Detailed derivations and explanations are in Appx. C.3. The resulting contrastive adversarial domain (CAD) objective is detailed in Algorithm 1. First, sample domains D :={D,D−1 , . . . , D−n} for eachX ′ ∈ X. Then collect inputs associated with the current domainD, i.e., XD :={X}∪{X−i |D−i = D, i ∈ [n]}. Finally, compute qϕ,X,D(D |Z) = ∑ X′∈XD qϕ,X(X ′ |Z). The resulting criterion is
LCAD(ϕ,ψ) := EpD,X,A,Z [ − log qψ,A(A+ |Z) + λ log ( ∑ X′∈XD qϕ,X(X ′ |Z) )] . (7)
In Appx. C.4, we also derive a conditional variation of CAD that minimizes I[Z;D |Y ], which can be used when labels are available and supervised augmentations are used.

heading: 4.2.2 ENTROPY BOTTLENECK (ENT)
text: Our second domain bottleneck is the entropy bottleneck (Ent) that minimizes H[Z] = I[Z;X] ≥ I[Z;D], where the first equality uses the encoder’s determinism. Ent enforces support match by removing all information that is not needed to maximize I[Z;A]. In particular, minimizing I[Z;X] is more stringent than I[Z;D], as it also matches the representations inside a domain. The advantage of Ent is that it 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., 2016; Theis et al., 2017), H[Z] ≤ EpZ [− log qθ(Z)], where an entropy model qθ(Z) is used. This leads to,
LEnt(ψ,ϕ, θ) := EpX,A,Z [− log qψ,A(A+ |Z)− λ log qθ(Z)] . (8)

heading: 5 RELATED WORK
text: Provably robust representations under covariate shift. Previous work mostly focuses on domain generalization bounds for robust representations. Ben-David et al. (2007; 2010a) bound the target risk using the source risk, a 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 a sufficient condition for optimality, i.e., matching the marginal pZ | d or its support while minimizing R [Y |Z]. However, these bounds can often be loose and the implied sufficient conditions are neither necessary nor generally achievable. Ben-David 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 Rds [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. Our theory could be combined with Dubois et al.’s (2020), who provide conditions for optimal generalization from finite samples and constrainedH in supervised learning. Practical objectives for DG. The most popular DG methods aim to learn domain-invariant representation by minimizing various divergernces between the marginal distributions pZ | d and pZ (Long et al., 2015; Ganin et al., 2016; Sun & Saenko, 2016; Long et al., 2017; Li et al., 2018a; Shen et al., 2018; Nguyen et al., 2021). Others propose matching the conditional pZ | y,d across domains instead (Gong et al., 2016; Li et al., 2018b; Tachet des Combes et al., 2020). These regularizers would all be valid domain bottlenecks B[Z,D] . Another line of work aims at learning Z with invariant predictors pY | z,d across domains (e.g., Arjovsky et al., 2019; Krueger et al., 2021; Li et al., 2021). However, none of these methods outperform ERM with fair model selections (Gulrajani & Lopez-Paz, 2021).

heading: 6 EXPERIMENTS
text: In our experiments, we aimed to: (i) verify our theoretical results in practice; (ii) investigate our proposed representation learning objectives in practical DG; (iii) take advantage of pretrained SSL models (in particular, CLIP) to achieve powerful models for DG. Unless stated otherwise, we consider a two-stage training setup. First, the representation learner (“the representor”) trains an encoder pZ |X using a specified objective and freezes it. Then, the person performing predictions (“the learner”) trains her predictor h from Z by minimizing the risk on source data. Finally, the representation Z and predictor h are evaluated on target data. In all experiments, the learner uses a linear classifier for h. For the Ent bottleneck, we used Ballé et al.’s (2018) entropy model. For the CAD bottleneck we used its conditional version whenever labels were available. When a model contains no domain bottleneck, we label it as “Base”. For experimental details and additional results see Appxs. D and E.

heading: 6.1 SCIENTIFIC SETTING: EXPLORING OPTIMAL REPRESENTATIONS FOR WORST-CASE DG
text: 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 `. In particular, we used the PACS dataset (Li et al., 2017)
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 a ResNet-18 (He et al., 2016) to minimize the objective on labeled data from all domains; (ii) the learner trains a 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. E.2 which replicates the following results.
Do our domain bottlenecks improve worst-case DG? In Fig. 3a, we 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. E.2, we show that bottlenecks also helps in practical scenarios, i.e., non-idealized average-case DG evaluated with accuracy (95.9%→ 96.7%).
What is the effect of λ? Fig. 3b shows the effect of the bottleneck weight λ on the worst-case target and source performance. We see that increasing λ will decrease the DG gap. As a result the target performance improves until λ ≈ 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 a 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-agnostic supervised augmentations (Fig. 2b). The 1st and 2nd row of Fig. 3a show that our objective performs similarly to direct label prediction.
How important is the choice of augmentations? Prop. 2 shows that domain-agnostic (DA) augmentations 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 LCAD with five augmentations. The first two are DA. ‘Supervised’: augment inputs inside the label class across all domains as in Fig. 2b; ‘SingleDom’: augment inputs to same label samples from a fixed domain. The second two are 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‘. Fig. 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.
What if the representor does not have access to target domains? Prop. 1 shows that DG without access to target domains is generally impossible. We empirically verified this by excluding a predefined target dt domain from the representor’s training set, i.e., LCAD is optimized on 3 of the 4 domains. The learner then trains a predictor h on each source. We finally evaluate each h on the target domain dt, and average over choices of dt. The resulting worst-case log likelihood was −4.2± 0.2, which is significantly worse than when the representor had access to all domains (−0.8± 0.2).

heading: 6.2 APPROXIMATING OPTIMAL REPRESENTATIONS BY EXPLOITING PRETRAINED SSL
text: As discussed in Sec. 4.1, one can learn optimal representations for IDG by performing SSL with a domain bottleneck on a 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 to 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 a 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 a 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. 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 Appx. D.3. We investigated two pretrained CLIP models with different number of parameters. The larger ViT-B/32 denoted ‘CLIP L’ and the smaller ResNet-50 denoted ‘CLIP S’.
Can we approximate optimal representations by exploiting pretrained CLIP? The last row in Table 1 shows 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 TerraIncognita is likely because CLIP’s dataset did not cover such images (camera traps monitoring animals). In Appx. E.2, we estimated the non-idealized DG performance of optimal representations on PACS (with access to all-domain labeled data) to be 96.7%, which is only 2% higher than CLIP L + CAD. This suggests that our simple SSL encoder might already be close to optimal.
Are gains due to the architectural differences? DomainBed’s baselines 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 1 shows that CLIP S + CAD still significantly outperforms DomainBed baselines. Note that our theory does not constrain the encoder and so we expect larger encoders to be better as seen in Table 1.
What is the effect of domain bottlenecks? In the last five rows of Table 1, we investigated the effect of finetuning with our CAD bottleneck. We see that for both CLIP L and CLIP S, it improves results by around 1 ∼ 2%. These gains are due to the bottleneck, rather than 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 raw CLIP S already significantly outperforms baselines. We hypothesize that this is because SGD acts as an information bottleneck that naturally favors support match (Shwartz-Ziv & Tishby, 2017).
Which pretrained SSL model to use? Our theory suggests that we can exploit pretrained SSL 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 a result, Table 1 shows that the finetuned DINO + CAD significantly underperforms 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-agnostic augmentations.

heading: 6.3 TOWARDS GENERIC ROBUST REPRESENTATIONS WITH SSL
text: In the previous section, we finetuned CLIP in a 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 a 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., 2021) that is a public dataset that contains 400M web-crawled image-text pairs. Due to our computational budget, we again froze the pretrained CLIP L and only finetuned an additional MLP with our LEnt. We used LEnt as it only requires access to paired image X and text A but no prior information about domain D. As in CLIP’s paper, we evaluated the learned representation Z in Taori et al.’s (2020) realistic setting, where a linear classifier h from Z is trained on ImageNet and tested on 7 natural distribution shift datasets. Details in Appx. D.4.
Would training CLIP with a bottleneck have improved its robustness? As shown in the last 2 rows of Table 2, finetuning CLIP L on LAION with LEnt (LAION + Ent) outperforms finetuning without bottleneck (LAION + Base) on all 7 distribution shift datasets. This suggests that directly training CLIP with our Ent bottleneck would improve the robustness of learned representations. We hypothesize that the gains could be larger if SSL models trained LEnt end-to-end. In Appx. E.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.

heading: 7 CONCLUSION
text: We gave a 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-agnostic property of data augmentations that make it possible to learn optimal representations from unlabelled data. Thus, we showed that it is possible to learn robust representations using only large sources of inputs X and augmentations A.
There are caveats that need to be addressed in 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 an unconstrained hypothesis class. The support constraint can likely be weakened, if 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 idealized setting, our characterization is a springboard from which all future objectives can be derived, and, in general, it brings us closer to the goal of robust machine learning systems.
Reproducibility Statement For our theoretical results, we include formal assumptions, statements, and proofs in Appxs. A and B. We include the detailed derivations of our algorithms in Appx. C. For our experiments, we include experimental details for reproducing our results in Appx. D and will release our code upon acceptance.

heading: A PRELIMINARIES

heading: A.1 NOTATION
text: For the most part, we will assume that all spaces are discrete probability spaces. A full list of assumptions is found at Appx. A.3.
General The image of a set A ⊆ X under a function f : X → Y is denoted f→(A) = {f(x) |x ∈ A}. The pre-image is denoted f←(B) ={x ∈ X | f(x) ∈ B} for B ⊆ Y . Probability Random variables (r.v.) are denoted by uppercase letters (e.g., X), and their sample space and realizations are denoted by the corresponding calligraphic (e.g., X ) and lowercase letters (e.g., x) respectively. The probability mass function (pmf) of a random variable X is denoted as pX . We use capital P instead of p to denote the measure under p. The support supp(pX) of a discrete distribution is the set of all points x ∈ X with positive probability, i.e., supp(pX) = {x ∈ X | pX(x) > 0}. The space of all probability distributions on X is denoted P(X ) = { pX | pX(x) ≥ 0 and ∑ x∈X pX(x) = 1 } .
When it is necessary to be explicit, we will denote ‘X is distributed as pX ’ using the notation X d∼ pX . Expectations are written as: EpX [f(X)], independence of two r.v. as ·⊥⊥·, conditional independence as ·⊥⊥ · | ·. For jointly distributed random variables (X,Y ) taking value in (t.v.i.) X × Y , the conditional distribution is denoted as pY |X : Y × X → [0, 1]. For convenience, let pY | x = pY |X( · |x) be the conditional distribution of Y given x. All random variables are independently distributed, unless an explicit joint distribution or coupling is given.

heading: A.2 DEFINITIONS
text: We are interested in prediction problems with domain shift. There are three random variables: the target domain Dt, the input X , the label Y . They have the following joint distribution:
(Dt, X, Y ) d∼ pDt · pX,Y |Dt (9)
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 pDt,X,Y that can be thought of as a generalized covariate shift assumption (Assumption 4).
We will be studying the effect of changing the representation of the data. This is done by “encoding” X into a representation Z using a conditional distribution pZ |X .
Definition 1 (Encoder). An encoder is a 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: (Dt, X, Y, Z)
d∼ pDt · pX,Y |Dt · pZ |X (10) The key thing to notice here is that Z is conditionally independent of Y,Dt given X . In particular, the same encoder is used across all domains.

heading: A.2.1 RISK MINIMIZATION
text: Our ultimate goal is to predict Y from the representation Z of X in a manner that is robust to changes in the domain.
We formalize this in the standard way by making predictions γ ∈ Γ in a space of predictions or actions. For example the prediction space may be the set of all possible labels Γ = Y , in which case we would be predicting deterministic labels. Or we may predict a distribution over labels, in which case the prediction space would be the set of all probability distributions on Y , i.e. Γ = P(Y). A predictor is a function mapping inputs to predictions, i.e., f : X → Γ, or representations to predictions, i.e., h : Z → Γ. For example, f may be a neural network that takes as input a sample x and outputs a vector of logits that parameterize a softmax distribution over finitely many labels.
We select predictors according to the risk defined via a loss function ` : Y × Γ→ R≥0 ∪ {∞}: Rf [Y |X] := EpX,Y [`(Y, f(X))] . (11)
In particular, we are interested in the Bayes (minimum) risk over all predictors:
R [Y |X] := inf f Rf [Y |X] , (12)
We denote the set of all optimal predictors from X as
F∗ :={f | Rf [Y |X] = R [Y |X]} (13) Similarly, we define the risk Rh [Y |Z], the Bayes risk R [Y |Z], and the set of optimal predictors H∗Z :={h | Rh [Y |Z] = R [Y |Z]} (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 Z fromH∗Z for clarity, but we will keep it in the Appendices. We assume that together our loss and prediction space always admit optima (Item 2 of Assumption 2), and thus F∗,H∗Z are always non-empty. We will be assuming that the risk admits unique optimal prediction when predicting from X (Item 3 of Assumption 2). Thus it makes sense to define the following:
Definition 2 (The Bayes predictor). The Bayes predictor f∗ : X → Γ is the unique predictor that is optimal for all x ∈ X :
f∗(x) = arg min γ∈Γ EpY | x [`(Y, γ)] (15)
Definition 3 (The Bayes image). The image of all the inputs under the Bayes predictor will be denoted as Γ∗ = f∗→(X ) and called the Bayes image.
Note that F∗ becomes a singleton {f∗}, but it is not necessarily the case forH∗Z since we will not be making any uniqueness assumption on optimal prediction from Z.

heading: A.2.2 DOMAIN GENERALIZATION
text: We are interested in controlling the risk in a domain generalization setting, and so we define the domain-conditional risk,
Rdf [Y |X] := EpX,Y | d [`(Y, f(X))] . (16)
Rd [Y |X] ,F∗d are defined as Eqs. (12) and (13), respectively, but with respect to Rdf . Similarly, define the Bayes image for domain d as
Γ∗d := f ∗→(supp(pX | d)) . (17)
We also define domain-conditional quantities for prediction from a representation Z. The most important term which we will be investigating is an idealization of the domain generalization worstcase risk.
Definition 4 (IDG risk). Given an encoder pZ |X and a distribution pDt,Ds over a target domain Dt and source domain Ds, the idealized domain generalization worst-case risk, IDG risk for short, is the expected worst-case target risk taken over source minimizers, i.e.,
RIDG [Y |Z] := EpDt,Ds
[ sup
h∈H∗Z,Ds RDth [Y |Z]
] (18)
Note that the IDG risk is well-defined becauseH∗Z,Ds is non-empty by Assumption 2. The desired optimal representations, are then those that minimize the IDG risk.
Definition 5 (Optimal representations for IDG). An encoder pZ∗ |X is optimal for idealized domain generalization if and only if it minimizes the IDG risk, i.e.,
RIDG [Y |Z∗] = inf pZ |X RIDG [Y |Z] (19)

heading: A.3 ASSUMPTIONS
text: We make a the following assumptions throughout the paper. All these assumptions should hold for practical settings.
Assumption 1 (Convenience: discrete probability spaces). All data spaces (D,X ,Y,Z,A) are discrete spaces. Because the distributions of X,Y,D are fixed, we assume for convenience that supp(pX) = X , supp(pY ) = Y , and supp(pDt) = D.
Assumption 1 is a convenience assumption to 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 (Properness of our losses). We assume that our risk always admits optimal predictions:
1. |Γ| > 1. 2. For all pΥ ∈ P(Y), there exists γ∗ ∈ Γ, such that
EpΥ [`(Υ, γ∗)] ≤ EpΥ [`(Υ, γ)] ∀ γ ∈ Γ. (20)
3. For all x ∈ X , there exist γ∗ ∈ Γ, such that EpY | x [`(Y, γ ∗)] < EpY | x [`(Y, γ)] ∀ γ 6= γ∗. (21)
Note that for log-loss `(y, γ) = − log γ(y) and finite Y , these assumptions are satisfied if Γ = P(Y) where the optimal prediction for Item 3 is γ∗ = pY | x by strict properness (Gneiting & Raftery, 2007). If we consider the 0-1 loss (reverse accuracy) `(y, γ) = 1− 1[y = γ] with Γ = Y and a finite label space where the optimal prediction for Item 3 is γ∗ = arg maxy∈Y pY | x(y), this assumption is mostly satisfied, except we assume that pY | x has a unique mode.
Assumption 2 serves two purposes: Item 2 ensures that for any representation the optimal predictors from Z exists such that the IDG risk is well-defined as in Def. 5; Item 3 ensures a unique Bayes predictor from X , which simplifies the analysis and is satisfied by common losses as described above.
Assumption 3 (Cardinalities). We assume that |Z| ≥ |Γ∗| ≥ 2 (22)
Assumption 3 is very weak and ensures that optimal representations always exists (Prop. 3).
Assumption 4 (Generalized covariate shift). The Bayes predictor is optimal for all domains. I.e., for all (x, d) ∈ supp(pX,Dt), γ ∈ Γ such that γ 6= f∗(x), we have
EpY | x,d [`(Y, f ∗(x))] < EpY | x,d [`(Y, γ)] . (23)
For example, in the case of strictly proper scoring rules, e.g. log loss, covariate shift pY |X,D = pY |X is equivalent to the invariance of the Bayes predictor. For the 0-1 loss, this is guaranteed by invariance of the most likely label. For MSE it is guaranteed by the invariance of the expected label. In the latter two cases, Assumption 4 is less stringent than the typical covariate shift assumption.
Assumption 4 is the core assumption for our theoretical results. It ensures that source and target domains are related in a useful way that can be utilized by the representation.
Assumption 5 (Constant Bayes image). The Bayes image is invariant across domains, i.e., for all d ∈ D,
Γ∗d = Γ ∗. (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 Γ∗d = {pY | x |x ∈ supp(pX | d)} is the same across domains.
Assumption 5 is crucial to be able to learn. Without it, in the extreme case, one could set each domain to be all examples associated with a single element from the label set (or the Bayes image set) in which case it is impossible to generalize across different domains. Assumption 5 is also necessary to guarantee the existence of optimal representations as in Prop. 3.
Assumption 6 (Domain joint). pDt,Ds is any distribution such that supp(pDt,Ds) = D ×D.
In a simplified scenario, one could define the source Ds and target Dt as i.i.d. r.v. from pDt , where pDt,Ds = pDt · pDs = pDt · pDt and Assumption 6 is trivially satisfied.

heading: B PROOFS

heading: B.1 LEMMAS FOR GENERAL LOSSES
text: An important result that we will be using is the generalized data processing inequality of Bayes risk (Xu & Raginsky, 2020; Dubois et al., 2021). We include it here for completeness.
Lemma 1 (Generalized DPI (Xu & Raginsky, 2020; Dubois et al., 2021)). Let Z − X − Y be a Markov chain of random variables. For any loss function `,
R [Y |X] ≤ R [Y |Z] . (25)
For the case of strictly proper losses (Assumption 2) we can go one step further.
Lemma 2. Let Z −X − Y be a Markov chain of random variables. Then, under Assumptions 1 and 2 we have that
R [Y |Z] = R [Y |X] ⇐⇒ ∀h∗ ∈ H∗Z ,∀(x, z) ∈ supp(pX,Z), h∗(z) = f∗(x). (26)
Proof. Suppose that for all h∗ ∈ H∗Z we have h∗(z) = f∗(x) on the support of pX,Z . Then,
R [Y |X] = EpX,Y [`(Y, f∗(X))] (27) = EpX,Y pZ |X [`(Y, f ∗(X))] (28)
= EpX,Y pZ |X [`(Y, h ∗(Z))] (29) = EpZ,Y [`(Y, h∗(Z))] (30) = R [Y |Z] . (31)
Now suppose there exists a h∗ ∈ H∗Z and a pair (x′, z′) ∈ supp(pX,Z) such that h∗(z′) 6= f∗(x′). Then
R [Y |Z] (32) = EpX,ZpY |X [`(Y, h ∗(Z))] (33)
= pX,Z(x ′, z′)EpY | x′ [`(Y, h
∗(z′))] + ∑
(x,z)6=(x′,z′) pX,Z(x, z)EpY | x [`(Y, h
∗(z))] (34)
≥ pX,Z(x′, z′)EpY | x′ [`(Y, h∗(z′))] + ∑
(x,z)6=(x′,z′) pX,Z(x, z)EpY | x [`(Y, f
∗(x))] (35)
> pX,Z(x ′, z′)EpY | x′ [`(Y, f
∗(x′))] + ∑
(x,z)6=(x′,z′) pX,Z(x, z)EpY | x [`(Y, f
∗(x))] (36)
= R [Y |X] (37)
Eq. (35) follows by Item 3 of Assumption 2 along with the definition of f∗. Eq. (36) follows by Item 3 of Assumption 2 and the fact that h∗(z′) 6= f∗(x′). This completes the proof, because Lemma 1 prevents R [Y |Z] < R [Y |X].

heading: B.2 PROOF OF THEOREM 1
text: First we will show that the desired representation exists by taking all inputs for which the Bayes predictor predicts similarly and “bucketing” them to the same representation. This is a direct extension of the example from Dubois et al.’s (2020) Proposition 6, to the case of proper losses.
Proposition 3 (Existence of optimal representations). Under Assumptions 1 to 5, there exists an encoder pZ∗ |X that is optimal for Eq. (3) and ??, i.e.,
pZ∗ |X ∈ arg min pZ |X R [Y |Z] s.t. ∀ d ∈ D, supp(pZ | d) = supp(pZ). (38)
Moreover, we have that R [Y |X] = R [Y |Z∗] . (39)
Proof. Because we assume arbitrary encoders pZ |X , the essence of this construction is simple: we embed the Bayes image into Z . Indeed, let φ : Γ∗ → Z be any one-to-one function, which exists due to Assumption 3 (here we use deterministic one-to-one function for simplicity, the construction can be easily extended to stochastic case). Then let Z∗ = φ(f∗(X)). We now verify the properties of pZ∗ |X .
1. Z∗ satisfies R [Y |X] = R [Y |Z∗]. Indeed, R [Y |X] = EpX,Y [`(Y, f∗(X))] (40)
= EpX,Y pZ∗ |X [`(Y, f ∗(X))] (41) = EpZ∗,Y [ `(Y, φ−1(Z∗)) ] (42)
≥ R [Y |Z∗] . (43) Eq. (42) is by our construction of Z∗ and Eq. (43) is by the definition of the Bayes risk. Due to the data processing inequality of Bayes risk (Lemma 1) we also have R [Y |X] ≤ R [Y |Z∗], from which we conclude that R [Y |X] = R [Y |Z∗] and that Eq. (39) holds.
2. Recall that Γ∗ = f∗→(X ) and Γ∗d = f∗→ ( supp(pX | d) ) . Now let us compute the desired
support for all d ∈ D: supp(pZ∗ | d) = φ →(Γ∗d) (44)
= φ→(Γ∗) (45) = supp(pZ∗). (46)
Eq. (45) is by Assumption 5.
Because R [Y |X] is the minimum achievable risk by any encoder regardless of constraint (this is by Lemma 1), this implies that pZ∗ |X is an optimal encoder for Eq. (3) and ??.
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 |Z] = R [Y |X]. Under Assumptions 1 and 2, we have that for all z ∈ supp(pZ), there exists γ∗ ∈ Γ such that
EpY | z [`(Y, γ ∗)] < EpY | z [`(Y, γ)] ∀ γ 6= γ∗. (47)
In other words, the restriction of any h∗ ∈ H∗Z to supp(pZ) is unique. If, in addition, Assumption 4 holds, then for all (z, d) ∈ supp(pZ,Dt), γ ∈ Γ such that γ 6= h∗(z),
EpY | z,d [`(Y, h ∗(z)] < EpY | z,d [`(Y, γ)] . (48)
In other words, the restriction of any h ∈ H∗Z,d to supp(pZ | d) is unique and equal to h∗.
Proof. For the first result, let z ∈ supp(pZ) and consider x ∈ supp(pX | z). By Lemma 2, it must be the case that f∗ is constant on supp(pX | z). Thus, we can pick γ∗ = f∗(x). Now, let γ 6= γ∗. We have that,
EpY | z [`(Y, γ ∗)] = EpX | zpY |X [`(Y, γ ∗)] (49)
= EpX | zpY |X [`(Y, f ∗(X))] (50)
< EpX | zpY |X [`(Y, γ)] (51) = EpY | z [`(Y, γ)] . (52)
Eq. (49) is due to the conditional independence of Y and Z given X . Eq. (51) is due to Assumption 2 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 [Y |X] = Rf∗ [Y |X] (53) = ∑ d∈D pDt(d) R d f∗ [Y |X] (54)
= ∑ d∈D pDt(d) R d [Y |X] , Assumption 4 (55)
and
R [Y |Z] = Rh∗ [Y |Z] (56) = ∑ d∈D pDt(d) R d h∗ [Y |Z] (57)
≥ ∑ d∈D pDt(d) R d [Y |Z] , (58)
where Eq. (58) is due to the definition of (domain-conditional) Bayes risk. Then R [Y |Z]− R [Y |X] ≥ ∑ d∈D pDt(d) ( Rd [Y |Z]− Rd [Y |X] )
(59)
≥ 0. Lemma 1 conditioned on d (60)
Thus, any encoder that achieves R [Y |Z] = R [Y |X] also satisfies Rd [Y |Z] = Rd [Y |X] for all d ∈ D since we assume that supp(pDt) = D in Assumption 1. Now, let d ∈ D. An argument analogous to Lemma 2 gives us,
∀h ∈ H∗Z,d,∀(x, z) ∈ supp(pX,Z | d), h(z) = f∗(x) = h∗(z). (61)
Eq. (61) is derived from Rd [Y |Z] = Rd [Y |X] using Assumption 4 in place of Item 3 of Assumption 2 for a specific domain d. Let z ∈ supp(pZ | d) and γ ∈ Γ such that γ 6= h∗(z). Since supp(pX | z,d) ⊆ supp(pX | z), f∗ is a constant on supp(pX | z,d) and equal to h∗. Now, as above, we have that
EpY | z,d [`(Y, h ∗(z))] = EpX | z,dpY |X,d [`(Y, h ∗(z))] (62)
= EpX | z,dpY |X,d [`(Y, f ∗(X))] (63)
< EpX | z,dpY |X,d [`(Y, γ)] (64) = EpY | z,d [`(Y, γ)] . (65)
Eq. (64) is due to Assumption 4.
Corollary 1. Let pZ |X be an encoder such that R [Y |Z] = R [Y |X]. Under Assumptions 1, 2 and 4 we have thatH∗Z ⊆ H∗Z,d for all d ∈ D and that for all ds, dt ∈ D
inf h∈H∗Z,ds
Rdth [Y |Z] = Rdt [Y |Z] (66)
Proof. H∗Z ⊆ H∗Z,d is immediate from Lemma 3. Now, we have that Rdth [Y |Z] ≥ Rdt [Y |Z]. So, the result follows by taking any h ∈ H∗Z ⊆ H∗Z,ds in the inf of Eq. (66).
Theorem 2 (Characterizing optimal representations for IDG, equiv. Theorem 1). Under Assumptions 1 to 6, an encoder pZ |X is optimal for idealized domain generalization if and only if it minimizes the Bayes risk while matching the support of pZ | d and pZ for all d ∈ D, i.e.,
pZ |X ∈ arg min pZ |X R [Y |Z] (67)
s.t. ∀ d ∈ D, supp(pZ | d) = supp(pZ) (68)
Proof. The IDG risk is lower bounded by R [Y |X]:
RIDG [Y |Z] ≥ EpDs,Dt
[ inf
h∈H∗Z,Ds RDth [Y |Z]
] (69)
≥ EpDs,Dt [ RDt [Y |Z] ] (70)
≥ EpDs,Dt [ RDt [Y |X] ] Lemma 1 (71)
= R [Y |X] Assumption 4 (72) We will now show that this lower bound is achieved by an encoder if and only if it satisfies Eqs. (67) and (68), which exist by Prop. 3.
Sufficiency (⇐= ): Let pZ |X be an encoder that satisfies Eqs. (67) and (68). Note that R [Y |Z] = R [Y |X] by Prop. 3. Let h∗ ∈ H∗Z , then we have the following IDG risk
RIDG [Y |Z] (73)
= EpDs,Dt
[ sup
h∈H∗Z,Ds EpZ,Y |Dt [`(Y, h(Z))]
] (74)
= EpDs,Dt
[ sup
h∈H∗Z,Ds EpZ,Y |Dt [`(Y, h
∗(Z))] ] Lemma 3 under matching support (75)
= EpDt [ EpZ,Y |Dt [`(Y, h ∗(Z))] ]
constant w.r.t Ds (76)
= R [Y |Z] = R [Y |X] (77)
Necessity ( =⇒ ): If the IDG risk is R [Y |X], then it must be the case that R [Y |Z] = R [Y |X] (78)
sup h∈H∗Z,ds
Rdth [Y |Z] = Rdt [Y |Z] ∀(ds, dt) ∈ supp(pDs,Dt) (79)
We will prove by contrapositive that Eq. (79) implies support match (Eq. (68)). Suppose that the support match does not hold. Since supp(pZ) = ∪d∈Dsupp(pZ | d) and supp(pDs,Dt) = D × D (Assumption 6), there must exist (ds, dt) ∈ supp(pDs,Dt) such that supp(pZ | ds) 6= supp(pZ | ds). Define the set S = supp(pZ | ds) ∩ supp(pZ | dt) and S̄ = supp(pZ | dt) \ supp(pZ | ds), let ρ = PZ | dt(S), and let h ∗ ∈ H∗Z . Then,
sup h∈H∗Z,ds
Rdth [Y |Z] (80)
= sup h∈H∗Z,ds
ρEpY,Z |S,dt [`(Y, h(Z))] + (1− ρ)EpY,Z | S̄,dt [`(Y, h(Z))] (81)
= sup h∈H∗Z,ds
ρEpY,Z |S,dt [`(Y, h ∗(Z))] + (1− ρ)EpY,Z | S̄,dt [`(Y, h(Z))] Lem. 3 (82)
= ρEpY,Z |S,dt [`(Y, h ∗(Z))] + (1− ρ) sup
h∈H∗Z,ds EpY,Z | S̄,dt [`(Y, h(Z))] (83)
= Rdt [Y |Z] + (1− ρ) sup h∈H∗Z,ds EpY,Z | S̄,dt [`(Y, h(Z))− `(Y, h ∗(Z))] (84)
> Rdt [Y |Z] Lem. 3 (85) Eq. (85) uses the following reasoning. 1− ρ > 0 due to support mismatch. For any h ∈ H∗Z,ds such that h 6= h∗ on S̄ (such an h exists by Item 1 of Assumption 2), we have that
EpY,Z | S̄,dt [`(Y, h(Z))− `(Y, h ∗(Z))] > 0 (86)
by Lemma 3.
As a corollary from the proof strategy we directly have that the optimal DG risk is simply R [Y |X]. This means that using the optimal encoder one can actually perform just as well by training on the source as if you were to directly train on the target using the raw data.
Corollary 2 (Optimal IDG Risk). Under Assumptions 1 to 6, infpZ |X RIDG [Y |Z] = R [Y |X].
B.3 IMPOSSIBILITY RESULTS
As a direct corollary of Thm. 2 we know that it is impossible to learn an optimal representation without knowledge or assumptions on the target domain. We can actually prove the following much stronger negative result, which essentially states that it is impossible to find a useful representation without having some information about the target domain. Specifically, we prove that if there exists a non-trivial target domain on which the representation is advantageous then there exists an infinite amount of target domains on which it is disadvantageous compared to predicting from a constant.
For clarity, we will focus on the proof for the standard accuracy (0-1 loss) which is much shorter and simpler to understand, but note that we can generalize the proof to all losses with the right assumptions.
The key is that outside of the source domain, the label distribution is unconstrained because generalized covariate shift has no effect. In other words, for any domain which gives some probability mass on an example that has not been seen during training, then all possible labels for that example gives a valid domain. Furthermore, if there exists one domain on which the representation is good, then one can construct a domain on which the representation is bad simply by labelling this point as the constant prediction.
Proposition 4 (No free lunch for learning representations for IDG, equiv. Proposition 1). Let ` be the 0-1 loss with prediction space Γ = Y . Let Rep : P(X ,Y) → P(Z|X ) be any algorithm for choosing an encoder pZ |X from the data distribution pX,Y , C be any constant r.v. that t.v.i. Z , and pX,Y | ds be any desired source distribution such that
• there is a unique constant prediction γC = arg miny∈Y EpY | ds [`(Y, y)],
• and |X \ supp(pX | ds)| > 1.
Let pZds |X := Rep(pX,Y | ds) be the chosen source encoder. If there exists a target domain pX,Y | dgt such that
• (Non-trivial support) ∅ 6= supp(pX | dgt ) ⊆ X \ supp(pX | ds);
• (Satisfies Bayes image invariance) Γ∗ dgt = Y , i.e., there is at least one example for every possible label;
• (Source encoder is useful) pZds |X performs better than a constant representation,
sup h∈H∗Zds ,ds
R dgt h [Y |Zds ] < sup h∈H∗C,ds R dgt h [Y |C] , (87)
Then there exist multiple target domains dbt such that pZds |X underperforms a constant encoder,
sup h∈H∗Zds ,ds
R dbt h [Y |Zds ] > sup h∈H∗C,ds R dbt h [Y |C] . (88)
Proof. Let h∗ ∈ H∗Zds ,ds be any source Bayes predictor corresponding to our encoder. Partition Z according to whether h∗ predicts like the constant or not:
ZC :={z ∈ Z |h∗(z) = γC} Z 6=C := Z \ ZC . (89) We know by assumption that dgt is s.t.
sup h∈H∗Zds ,ds
R dgt h [Y |Zds ] < sup h∈H∗C,ds R dgt h [Y |C] , (90)
which is clearly only possible if PZds | dgt (Z 6=C) > 0. (91) In other words, there exists some input x6=C ∈ X \ supp(pX | ds) that will get represented outside of the constant region, i.e.,
PZds | x 6=C (Z 6=C) > 0. (92)
We will now construct the desired bad domain dbt by giving nearly all mass to this x 6=C , specifically, let pX | dbt (x6=C) = 1 − δ for some 0 < δ < 1. We assign this example to the constant label, i.e., pY | x 6=C ,dbt (γC) = 1. The rest of the target domain mass δ is distributed as with the source domain, i.e., pX,Y | dbt (x, y) = δ · pX,Y | ds(x, y) for all x, y ∈ supp(pX,Y | ds). Importantly, the constructed domain dbt is valid. Indeed, the Bayes image is the same as the source’s (Assumption 5), because we removed no prediction γ from the source’s Bayes image (δ > 0). We added no new prediction γ, because f∗(x 6=C) = γC ∈ Y which must already have been in Γ∗ due to the validity of dgt . Now let us compute the desired risk for that “bad” domain and show that the desired encoder performs worse than a constant encoder.
sup h∈H∗Zds ,ds
R dbt h [Y |Zds ] (93)
= sup h∈H∗Zds ,ds
(1− δ)EpZds | x 6=C [1− 1[γC = h(Zds)]] + δR ds h [Y |Zds ] (94)
≥ (1− δ)(1− PZds | x6=C (ZC)) (95) = (1− δ)PZds | x 6=C (Z 6=C) (96)
In contrast, it is easy to show that suph∈H∗C,ds R dbt h [Y |C] ≤ δ because the constant predictor would be perfect for x 6=C . So any choice of 0 < δ < PZds | x6=C
(Z 6=C) 1+PZds | x 6=C (Z 6=C) , 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 dbt .
Note that representations can often be much worse than using a constant r.v. Specifically, if an encoder pZ |X 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 | ds , pZds |X , ` be as in Prop. 4, and > 0. If there exists an example xb ∈ X \ supp(pX | ds) that is mapped outside of the source support, i.e., supp(pZds | xb) ∩ supp(pZ | ds) = ∅, then there exist many target domains pX,Y | dt s.t. pZds |X is close to the worst possible loss, i.e.,
sup h∈H∗Zds ,ds
Rdth [Y |Zds ] ≥ 1− . (97)
Proof. By assumption there exists an xb whose support is outside the source support. Then similarly to Prop. 4 we construct a bad target domain dt by giving nearly all mass to that example pX | dt(xb) = 1 − δ where δ > 0 and assign with probability 1 to some label that is in the source Bayes image, i.e., pY | xb,dt(γb) = 1 for some γb ∈ Γ∗ds . The rest of the target domain mass δ is distributed as in
Prop. 4 to the source inputs. As in Prop. 4, such a target domain dt satisfies our assumptions. Now let us compute the risk for that dt and show that the desired encoder performs arbitrarily bad.
sup h∈H∗Zds ,ds
Rdth [Y |Zds ] (98)
= sup h∈H∗Zds ,ds
(1− δ)EpZds | xb [1− 1[γb = h(Zds)]] + δR ds h [Y |Zds ] Eq. (94) (99)
≥ sup h∈H∗Zds ,ds (1− δ)EpZds | xb [1− 1[γb = h(Zds)]] (100)
= 1− δ (101) Eq. (101) uses the fact that H∗Zds ,ds is unconstrained outside of the source support and that by assumption supp(pZds | xb)∩ supp(pZds | ds) = ∅. To achieve the sup 1− δ it then suffices to predict an γ 6= γb ∈ Γ. We thus see that Eq. (97) holds for dt as long as 0 < δ < . We conclude the proof by noting that there is an infinite possible choices of δ each of which give rise to a bad target domain.

heading: B.4 AUGMENTATIONS
text: 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 pA |X : A×X → [0, 1] from the input space X to an augmentation spaceA. For example, in CLIP X is the space of images andA 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 |X , define the augmentation conditional set as the set of conditionals of A given X:
P∗(A |X) := { pA | x |x ∈ X } (102)
Similarly, we can define the augmentation conditional set for domain d: P∗d (A |X) := { pA | x |x ∈ supp(pX | d) } (103)
These sets are clearly countable. Note that the augmentation conditional set can be seen as a special case of the Bayes image (Def. 3) if we view the augmentation A as the label and consider the log-loss where the conditional distribution is the Bayes optimal predictor due to its strict properness (Gneiting & Raftery, 2007).
Assumption 7 (Finite augmentation entropy). We consider the augmenter pA |X such that the entropy of the augmentation A is finite, i.e., H[A] <∞.
Assumption 8 (Cardinalities). We assume that
|Z| ≥ |P∗(A |X)| (104) This is a similar assumption as Assumption 3, which ensures the existence of optimal representations.
Assumption 9 (Domain-agnostic augmentation). We assume that the augmentation A is domainagnostic, i.e., the augmentation conditional set is invariant across domains,
P∗d (A |X) = P∗(A |X), ∀d ∈ D (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 pA | x ∈ P∗(A |X) is seen at least once in all domains. If we introduce an equivalence relation ∼ as x ∼ x′ iff pA | x = pA | x′ and the equivalence class [x] := {x′ ∈ X |x′ ∼ x}. Under this relation, it is
easy to see that the above assumption is satisfied if and only if, for all possible equivalence classes [x] ∈{[x′] |x′ ∈ X}, we have that [x] has intersections with all domains:
[x] ∩ supp(pX | d) 6= ∅, ∀d ∈ D (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 of CLIP nearly are, as discussed in the main body (Sec. 4).
Assumption 10 (Bayes-preserving augmentation). We assume that the augmentation A is Bayespreserving, i.e., ∀x, x′ ∈ X ,
pA | x = pA | x′ =⇒ f∗(x) = f∗(x′). (107) Under the notion of equivalence relation in Assumption 9, this means that for each equivalence class [x], all x′ ∈ [x] have the same Bayes prediction. Note that most augmentations used in practice like standard image augmentations are Bayes-preserving.
Next, we show that under the above assumptions, we can learn optimal representations by maximizing the mutual information I[A;Z] (in the case of log-loss `) under the support match constraint. We use log-loss simply because it is typically the loss used for training in practice. Note that the learned representations are optimal for any strict proper losses.
Proposition 6 (Learning optimal representations without labels, equiv. Proposition 2). Let pA |X be an augmenter. Under Assumptions 1 to 10, any encoder pZ |X such that
pZ |X ∈ arg max pZ |X I[A;Z] (108)
s.t. ∀ d ∈ D, supp(pZ | d) = supp(pZ) (109) is optimal for idealized domain generalization.
Proof. The support match constraint Eq. (109) is equivalent to the support match constraint Eq. (68). Thus, Prop. 3 and Thm. 2 state that we only need to prove that maximizing the mutual information of A and Z under the support constraint implies that
R [Y |Z] = R [Y |X] . (110) We will prove this by constructing an optimal predictor h∗.
Since H[A] <∞ (Assumption 7) we have that arg max
pZ |X I[A;Z] = arg min pZ |X H[A |Z] . (111)
Note the fact that the conditional entropy is the Bayes risk under the log-loss (Gneiting & Raftery, 2007), i.e., H[A |Z] = R [A |Z]. By construction, A satisfies covariate shift w.r.t. X (thus Bayes invariant) since A−X −D forms a Markov chain. Together with Assumptions 1 and 7 to 9, it means that the optimization problem in Eqs. (108) and (109) satisfies the assumptions of Prop. 3, with A in place of Y . Thus, an optimal encoder satisfies R [A |Z] = R [A |X], which leads to
H[A |Z] = H[A |X] . (112) By Assumption 7, we can invoke Lemma 2 with the fact that A−X − Z forms a Markove chain to show that for all (x, z) ∈ supp(pX,Z)
pA | z = pA | x, (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′ ⇐⇒ pA | x = pA | x′ . (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 a maximal invariant see Definition 2, Dubois et al., 2021). By Assumption 10, f∗ is invariant on the equivalence classes [x] := {x′ ∈
X |x′ ∼ x} for all x ∈ X . Thus, there exists a function g : N → A such that f∗ = g ◦ M (Lemma 5, Dubois et al., 2021). Given z ∈ supp(pZ), we construct h∗ in the following way. Let xz ∈ supp(pX | z) be any input point that could have led to this representation z and define
h∗(z) = g(M(xz)). (115)
By Eq. (113) we are guaranteed that all x ∈ supp(pX | z) share the same value for f∗ since they are in the same equivalence class. Thus, by the definition of M we have that
M(xZ) = M(X) for (X,Z) ∼ pX,Z . (116) Therefore,
Rh∗ [Y |Z] = EpY,Z [`(Y, h∗(Z)] (117) = EpY |XpX,Z [`(Y, h ∗(Z)] (118)
= EpY |XpX,Z [`(Y, g(M(xZ)))] Eq. (115) (119) = EpY |XpX,Z [`(Y, g(M(X)))] Eq. (116) (120) = EpY |XpX,Z [`(Y, f ∗(X))] = R [Y |X] . (121)

heading: C PRACTICAL OBJECTIVES
text: Proposition 6 provides an objective to obtain the desired optimal representations, compared to Thm. 2 it is more practical in that it does not require access to the labels but right augmentations. There are nevertheless multiple remaining issues for deriving objectives that can be trained with in practice. Specifically, (i) the support constraint is hard to satisfy in practice; (ii) mutual information I[A;Z] is hard to estimate from samples (Poole et al., 2019); (iii) the objective is constrained which is harder to optimize. We will now show different objectives and variational bounds of them that do not suffer from these issues, and could still recover the desired encoders in their optima. In contrast to the proofs of main theoretical results (previous section), here the derivations will be less formal.
As we have seen in Proposition 6, the optimal representation achieves I[A;Z] = I[A;X], and thus H[A |Z] = H[A |X] (Eq. (112)). In the following, we will rewrite the objective as the constrained optimization:
pZ |X ∈ arg min pZ |X B[Z,X, Y,D] (122)
s.t. H[A |Z] = H[A |X] (123) where we introduce the domain bottleneck B[Z,X, Y,D] as the objective for enforcing support match (which we denote as B[Z,D] in the main body for simplicity). The requirement on the domain bottleneck objective is that minimizing Eq. (122) under Eq. (123) implies that the support match constraint holds (and can be achieved by some encoder), which leads to optimal representations for idealized domain generalization. Different domain bottlenecks will be derived later this section. We can then use Lagrangian relaxation to get the following unconstrained objectives.
arg min pZ |X
H[A |Z] + λB[Z,X, Y,D] (124)
The first term can be easily optimized using variational objectives. Throughout the paper, we will use a contrastive variational upper bound which is based on InfoNCE (Oord et al., 2018). Namely, let X := { X,X−1 , . . . , X − n } be a sequence of samples where each negative X−i is independently
sampled from the marginal pX . Now let A := { A+, A−1 , . . . , A − n } be the sequence of augmented examples that come from independently augmenting each X ∈ X using the augmenter pA |X . In particular, the augmented positive A+ has the marginal pA |X while the negatives A − i follow the marginal pA. Let Z be the representation of X by passing it through the encoder pϕ := pZ |X parameterized by ϕ and sψ the critic function parametrized by ψ used to score which A′ ∈ A is the positive augmentation. Then we have a variational distribution of pA |Z :
qψ,A(A |Z) := exp sψ(A,Z)∑
A′∈A exp sψ(A ′, Z)
(125)
which leads to the following variational bound: H[A |Z] = EpX,Z [ − log pA |Z(A |Z) ] ≤ EpA,X,Z [− log qψ,A(A+ |Z)] (126) with equality if the variational families sψ, pϕ are unconstrained and and we use infinite samples (n → ∞). Typically the critic is separable, i.e., sψ(A,Z) := gψ(A)Thψ(Z). As discussed in the main body, it can be tied with pϕ when A = X . In the following we focus on the second term B[Z,X, Y,D] and discuss several choices.
C.1 MUTUAL INFORMATION BOTTLENECK B[Z,X, Y,D] = I[Z;X]
The first bottleneck we consider is so called mutual information (MI) bottleneck B[Z,X, Y,D] = I[Z;X], which was introduced by Tishby et al. (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 I[Z;A]. In particular, using the fact that Z −X −D forms a Markov chain and the chain rule of MI, we have I[Z;X] = I[Z;X,D] = I[Z;D] + I[Z;X |D]. Thus, it not only minimizes I[Z;D], i.e., matches the representations’ distribution across domains, but also minimizes I[Z;X |D], i.e., matches the representations’ distribution inside domains. Why The key to show is that minimizing Eq. (122), i.e., arg minpZ |X I[Z;X] under I[A;Z] = I[A;X], implies the support match constraint. This can be seen as a specific subcase of Dubois et al.’s (2021) Corollary 15 with A in place of Y and M(X) induced by pA |X as in the proof of Prop. 6. From the corollary, we know that minpZ |X I[Z;X] = H[M(X)] which can be achieved by any Z s.t. pZ | x = pZ | x′ ⇐⇒ M(x) = M(x′). With the assumption of domain-agnostic augmentations (Assumption 9), we have that the set of maximal invariant {M(x) |x ∈ supp(pX | d)} is invariant across domains. Then we directly have supp(pZ | d) = ∪x∈supp(pX | d)supp(pZ | x) = ∪x∈supp(pX)supp(pZ | x) = supp(pZ), where we use the fact that x within the same equivalence class has the the same pZ | x.
How Essentially, we can use any variational upper bound of mutual information. We consider the one used by Variational Information Bottelenck (Alemi et al., 2016), i.e.,
I[Z;X] = EpX,Z [ log
pϕ(X |Z) pZ(Z)
] (127)
= EpX,Z [ log
pϕ(X |Z) qθ(Z)
] −DKL[pZ(Z)‖qθ(Z)] (128)
≤ EpX,Z [ log
pϕ(X |Z) qθ(Z)
] (129)
= EpX [DKL[pϕ(Z |X)‖qθ(Z)]] (130) where a variational distribution qθ is used to approximate pZ and is jointly optimized with pϕ to minimize the bound. The approximation gap of the bound is DKL[pZ(Z)‖qθ(Z)]. Then the final loss becomes
LMI(ψ,ϕ, θ) := EpX,A,Z [ − log exp sψ(A
+, Z)∑ A′∈A exp sψ(A ′, Z) + λDKL[pϕ(Z |X)‖qθ(Z)]
] (131)
which recovers the optimal encoder in the case of unconstrained variational families for pϕ, qθ, sψ, infinite samples |A|, and any λ > 1 (Dubois et al., 2021).
C.2 ENTROPY BOTTLENECK B[Z,X, Y,D] = H[Z]
The entropy (Ent) bottleneck introduced in the main body is a special case of the MI bottleneck, where the encoder is a deterministic mapping, i.e., pϕ(Z |x) is a dirac delta function for all x ∈ X and we denote by eϕ(x) the deterministic encoder s.t. pϕ(eϕ(x) |x) = 1. Why In the deterministic case, the MI bottleneck becomes the entropy bottleneck because I[X;Z] = H[Z]− H[Z |X] = H[Z], where we use the fact that H[Z |X] = 0. Importantly, considering only deterministic encoders does not constrain our ability to learning optimal encoders. Indeed, just as with the MI bottleneck optimizing the objective with the entropy bottleneck under I[A;Z] = I[A;X] will recover encoders s.t. eϕ(x′) = eϕ(x) ⇐⇒ M(x) = M(x′), 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[Z] ≤ EpZ [− log qθ(Z)] (132) which is the standard variational bound used in neural compression (Ballé et al., 2016; Theis et al., 2017). Putting all together, we have
LEnt(ψ, θ, ϕ) := EpX,A,Z [ − log exp sψ(A
+, Z)∑ A′∈A exp sψ(A ′, Z) − λ log qθ(Z)
] (133)
which also recovers the optimal encoder with unconstrained variational families, infinite samples, and λ > 1 as with the MI bottleneck. The detialed algorithm is provided in Algorithm 2. Note that the discreteness of Z could lead to difficulty of gradient-based optimization, and we follow Ballé et al. (2016) to add uniform noise to Z as a differentiable substitute for rounding during training. In our experiments, we will mostly use the Ent bottleneck instead of the MI bottleneck to avoid introducing stochastic encoders.
Algorithm 2 Ent objective Require: eϕ, sψ, qθ, X, n
1: Z ← eϕ(X) 2: A+ ← sample(pA |X) 3: {(X−i , A−i )} n i=1
i.i.d.←−− sample(pX,A) 4: A←{A+} ∪{A−i } n i=1 5: Laug ← − log exp sψ(A +,Z)∑
A′∈A exp sψ(A ′,Z) . H[A |Z]
6: Lsupp ← − log qθ(Z) . H[Z] 7: return LEnt = Laug + λLsupp
C.3 CONTRASTIVE ADVERSARIAL DOMAIN BOTTLENECK B[Z,X, Y,D] = I[Z;D]
The previous two bottlenecks require to remove the information of Z (about X) as much as possible, which seems to be unnecessary since our ultimate goal is to match the support of Z across domains. Now we introduce a bottleneck B[Z,X, Y,D] = I[Z;D] which we only seek to remove the information of Z about the domain D. This is very related to the work on invariant representation learning for domain generalization/adaptation (e.g., Ganin et al., 2016; Li et al., 2018a). We derive a new variational bound called the contrastive adversarial domain (CAD) bottleneck that is more stable to train and leads to better empirical performance. For simplicity we consider the deterministic encoder eϕ(x) as with the main body.
Why Similar to the previous analysis, we aim to show that arg minpZ |X I[Z;D] under I[A;Z] = I[A;X] leads to the support match constraint. Using Eq. (116) we have I[Z;D] = I[Z,M(XZ);D] = I[Z,M(X);D] = I[M(X);D] + I[Z;D |M(X)] where the last equality uses the chain rule of mutual information. Due to the non-negativity of (conditional) mutual information, we have that the minimum of I[Z;D] under I[A;Z] = I[A;X] is I[M(X);D]. Then we show the minimum is achievable by constructing the same optimal encoder eϕ(X) as the Ent bottleneck which clearly satisfies I[Z;D |M(X)] = 0. It is then easy to show that the support match constraint has to hold when I[Z;D |M(X)] = 0 by contrapositive. Indeed, suppose that the support constraint does not hold then it must be true that I[Z;D |M(X)] > 0 and so the encoder cannot be optimal. How The typical way of minimizing I[Z;D] is to derive the variational bound as
I[Z;D] = H[D]−H[D |Z] (134) = (const)− EpD,Z [ − log pD |Z(D |Z) ] (135)
≥ (const)− EpD,Z [− log qφ(D |Z)] (136) where a variational distribution (or domain classifier) qφ is used to approximate pD |Z and jointly trained to maximize the bound. This recovers the domain-adversarial training method as introduced in Ganin et al. (2016). However, this has two potential issues: 1) it gives a lower bound instead of the
desired upper bound on I[Z;D]; 2) it requires adversarial training which is not stable (Goodfellow, 2016; Kodali et al., 2017).
We propose the contrastive adversarial domain (CAD) bottleneck, which is based on the above explicit version but uses a variational distribution qφ(D |Z) that is tied with other parts of the model, thus no need to learn a domain classifier. Specifically, as with Eq. (125), we first introduce a contrastive variational distribution qϕ,X(X |Z) of pX |Z as:
qϕ,X(X |Z) := exp sϕ(X,Z)∑
X′∈X exp sϕ(X ′, Z)
(137)
where sϕ(X,Z) := eϕ(X)TZ is tied with the encoder eϕ. Since pD |Z can be rewritten as EpX |Z [ pD |X ] using the fact that D −X − Z forms a Markov chain, we obtain can the following variational distribution: qϕ,X(D |Z) = Eqϕ,X [ pD |X(D |X) ] (138) which recovers pD |Z when qϕ,X = pX |Z . This is the case where sϕ(X,Z) ∝ log pX,Z(X,Z) and infinite samples n → ∞. Note that pD |X is still not available, we can use a count estimate p̂D,X in practice. In particular, we obtain a collection D by taking each X ′ ∈ X and independently sampling D′ from pD |X′ to get D :={D,D−1 , . . . , D−n}. In other words, (D,X) and (D−i , X−i ) for i ∈ [n] :={1, . . . , n} are all i.i.d. sampled from pD,X . Then we use a count estimate
p̂D,X(d |x) = I (X = x,D = d) +
∑n i=1 I (X − i = x,D − i = d)
I (X = x) + ∑n i=1 I (X − i = x)
(139)
which is an accurate estimate with infinite samples. This leads to our final variational distribution: qϕ,X,D(D |Z) = ∑ X′∈X qϕ,X(X ′ |Z)p̂D,X(D |X ′) (140)
Putting all together we get that the final loss:
LCAD(ϕ,ψ) := EpD,X,A.Z [ − log qψ,A(A+ |Z) + λ log ( ∑ X′∈X qϕ,X(X ′ |Z)p̂D,X(D |X ′) )] .
(141) In practice, p̂D,X(D |X) is typically a dirac delta function since it is rare to have the same samples in a batch. Thus, in Eq. (140) we only need to sum qϕ,X(X ′ |Z) over those associated with the same domain label D as X , which reduces Eq. (141) to Eq. (7) described in the main body with a detailed algorithm in Algorithm 1. Note that it is easy to generalize Algorithm 1 to parallel computation within a batch of samples. Indeed, for each sample in the batch, we can view all other samples in the batch as negatives and compute the loss efficiently in parallel.
C.4 CONDITIONAL CAD B[Z,X, Y,D] = I[Z;D |Y ]
The analysis of the CAD bottleneck also implies that we can minimize the conditional mutual information I[Z;D |M(X)] if we have access to M(X). However, since M(X) is typically not available in practice, we consider the special case where M(X) = Y . In particular, this is the case where the labels are available and the supervised augmentations are used (see Fig. 2b). This reduces the bottleneck to B[Z,X, Y,D] = I[Z;D |Y ] which is related to the conditional version of the domain-adversarial neural network (Li et al., 2018b). In practice, minimizing I[Z;D |Y ] could be easier for optimization than I[Z;D], as it does not require to remove the information that D has about Y . In the following, we derive the conditional CAD (C2AD) bottleneck using a similar idea as CAD.
How In this case, we want to minimize
I[Z;D |Y ] = H[D |Y ]−H[D |Z, Y ] (142) = (const)−H[D |Z, Y ] (143) ≥ (const)− EpD,Z,Y [− log q(D |Z, Y )] (144)
where q(D |Z, Y ) is a variational distribution of pD |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,Y [ pD |X ] . Specifically, let Y :={Y, Y −1 , . . . , Y −n } be the collection of labels obtained by independently sampling the label from pY |X′ for each X ′ ∈ X. We collect samples associated with the label Y , i.e., XY :={X} ∪{X−i |Y −i = Y, i ∈ [n]} and obtain a variational distribution of pX |Z,Y :
qϕ,X,Y(X |Z, Y ) := exp sϕ(X,Z)∑
X′∈XY exp sϕ(X ′, Z)
(145)
where we use the same critic sϕ(X,Z) := eϕ(X)TZ that is tied with the encoder eϕ as before, but only take softmax over those samples with the same label Y . For the term pD |X , we use the same count estimate p̂D,X in Eq. (139). Then we obtain the variational distribution of pD |Z,Y :
qϕ,X,D,Y(D |Z, Y ) = ∑
X′∈XY qϕ,X,Y(X
′ |Z, Y )p̂D,X(D |X ′) (146)
Putting all together we get that the final loss:
LC2AD(ϕ,ψ) := EpD,X,A,Y,Z [ − log qψ,A(A+ |Z) + λ log ( ∑ X′∈XY qϕ,X,Y(X ′ |Z, Y )p̂D,X(D |X ′) )] . (147) Again, since in practice p̂D,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 . In particular, we collect samples XY,D :={X} ∪{X−i |Y −i = Y,D−i = D, i ∈ [n]} and obtained the simplified loss:
LC2AD(ϕ,ψ) := EpD,X,A,Y,Z − log qψ,A(A+ |Z) + λ log  ∑ X′∈XY,D qϕ,X,Y(X ′ |Z, Y )  . (148) A detailed algorithm is in Algorithm 3.
Algorithm 3 conditional CAD (C2AD) objective Require: eϕ, sψ, D,X, Y, n 1: Z ← eϕ(X) 2: A+ ← sample(pA |X) 3: { (D−i , X − i , A − i , Y − i ) }n i=1
i.i.d.←−− sample(pD,X,A,Y ) 4: X,A←{X} ∪ { X−i }n i=1 ,{A+} ∪ { A−i }n i=1
5: XY ←{X} ∪ { X−i |Y − i = Y, i ∈ [n] } 6: XY,D ←{X} ∪ { X−i |Y − i = Y,D − i = D, i ∈ [n]
} 7: Laug ← − log exp sψ(A +,Z)∑
A′∈A exp sψ(A ′,Z) . H[A |Z]
8: Lsupp ← log ∑ X′∈XY,D exp eϕ(X ′)TZ∑
X′′∈XY exp eϕ(X′′)TZ . I[Z;D |Y ] 9: return LC2AD = Laug + λLsupp

heading: D EXPERIMENTAL DETAILS

heading: D.1 SCIENTIFIC
text: In both the scientific setting and the following bridge setting, we consider rather unrealistic setups for verifying our theory where we have access to labels from all domains. We can choose to directly minimize the risk R [Y |Z] with the cross-entropy loss (denoted as CE henceafter), or minimize H[A |Z] (i.e., maximize I[A;Z]) with supervised augmentations as in Fig. 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., 2020). In particular, for a single sample in a batch, all samples in the batch with the same labels can be used as the positives (could come from the same domain or different domains) and others as the negatives. In Khosla et al. (2020), two variants of SupCon loss were introduced for solving the issue of multi-positives depending on whether the summation over multi-positives was located inside (SupCon-In, Eq. (3) in Khosla et al. (2020)) or outside (SupCon-Out, Eq. (2) in Khosla et 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 an implicit bottleneck effect. Intuitively, SupCon-Out upper bounds SupCon-In and achieves its optima only if the logits with positive samples are all the same by Jensen’s inequality, which may encourage positive samples from different domains to get clustered. Since this might confound with the effect of our bottlenecks, we chose to use SupCon-In though it performed slightly worse in out initial experiments. For the implementation of SupCon, we followed Khosla et al. (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 as detailed below:
Eliminating empirical generalization As our theory focuses on the idealized domain generalization that assumes access to population distribution, we considered the setup where the empirical generalization was eliminated. 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, 2014) with a learning rate of 5e-5, a batch size of 192 (48 for each domain), and a cosine learning rate decay schedule.
Worst-case approximation To approximate the worst-case source predictor, we included the target data with randomly assigned wrong labels to the training set for training the source predictor. The target data samples were down-weighted with a sample weight that maximizes the target risk while keeping the source risk close to optima (which is 0). We selected the sample weight by sweeping over [10−10, 1] with a logarithmic scale using CE-Base and SupCon-Base, as shown in Fig. 4. As 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 we selected the weight close to the turning point and 10−5 seemed to be reasonable for both CE-Base and SupCon-Base. Although we did not adaptively select the sample weight for each setup due to the computational cost, the pre-specified sample turned out to 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 N (0, 1). Next, we provide other experimental details for reproducibility:
Implementation of standard augmentations We followed SimCLR (Chen et al., 2020) for implementing standard image augmentations. For a fair comparison between the cases when using standard
augmentations (SimCLR) and supervised augmentations (SupCon), we kept the total batch size the same and also used the same 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. D.1). The ‘Supervised’ augmentation was essentially implemented using the SupCon loss as described in Appx. D.1. For other augmentations considered, we implemented them by dropout inter-domain supervised augmentations in SupCon. Specifically, for each sample in the batch, we randomly masked the samples from different domains (i.e., both inter-domain positives and negatives) i.i.d. with the specified dropout probability, while samples within the same domain were always kept. ‘IntraDom’ and ‘ApproxDA’ correspond to dropout probability 1 and 0.9, respectively. ‘SingleDom’ were implemented by dropout all interdomain samples with probability 1 except for a fixed domain (the ‘A’ domain of PACS in our case).

heading: D.2 BRIDGE
text: In the bridge setting (see Appx. E.2), we aimed to bridge the gap between our theoretical setup to the practical setup. The main differences from the scientific setups are that the empirical generalization gap is considered and the average-case source predictor is used, as detailed below:
Incorporating empirical generalization In practice, empirical-generalization gap should also be considered besides the source-target generalization gap. Thus, we randomly split the PACS dataset to 80% training and 20% validation splits for each domain. The training splits were used to train both the encoder and the source predictor, and the validation splits were used for encoder and source predictor selection as well as evaluation on target domains. We used the ResNet-50 model as the encoder and initialized it from ImageNet pretrained model. The encoder was trained to a maximum of 50 epochs with a 1e-5 weight decay, using the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 5e-5, a batch size of 112 (28 for each domain), and a cosine learning rate decay schedule.
Using average-case source predictor Instead of approximating the worst-case source predictor in the scientific setting, we considered the average-case2 source predictor which is closer to the common practice. Specifically, we freezed the encoder and trained a SVM classifier with L2 regularization on the source training split. The regularization parameter was tuned over {1e-4, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2, 1e3} with the source validation accuracy.
Next, we provide other experimental details for reproducibility:
Selection of λ For all different setups considered in bridge settings, the CAD bottleneck was used and the λ was tuned over {1e-3, 1e-2, 1e-1, 1, 1e1} independently for each.

heading: D.3 DOMAINBED
text: Datasets We used non-MNIST datasets on DomainBed that were non-synthetic, including VLCS (Fang et al., 2013), PACS (Li et al., 2017), OfficeHome (Venkateswara et al., 2017), TerraIncognita (Beery et al., 2018), and DomainNet (Peng et al., 2019). For each dataset, we split it to 80%/20% training/validation set according to DomainBed.
SSL-based models & Training For all models based on pretrained SSL models (either CLIP-based or DINO-based) with finetuning in this experiment, we freezed the pretrained SSL model and added on top a 1-layer MLP with hidden size 1024, and residual connection. We used CLIP ResNet-50 (CLIP S) to obtain the best possible fair comparison with baselines from DomainBed, and CLIP ViT-B/32 (CLIP L) to achieve the best results. Note that the ResNet-50 model of CLIP S was modified as described in Radford et al. (2021) and contained 38M parameters (more than 23M of the original CLIP). The model was trained to 300 epochs for DomainNet and 50 epochs on other datasets (an epoch is defined as a single pass over the smallest domain according to DomainBed). No data augmentation was used and the temperature for scaling the logits in CAD was fixed to 0.05. We
2Here we have a 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 ‘best-case’ since the max-margin classifier (SVM) was used.
used the Adam optimizer with a 1e-5 weight dacay, and a cosine learning rate decay schedule. The hyperparameter search space is:
• Learning rate: discrete set {1e-4, 3e-4, 1e-3, 3e-3}
• Batch size: discrete set {128, 256, 512} for DomainNet and OfficeHome, and {64, 128, 256} for other datasets
• MLP dropout: discrete set {0., 0.1, 0.5}
• Learning rate warmup: discrete set {True, False}
End-to-end models & Training In Table 1, we also included an end-to-end trained model without any pretrained SSL models. We used exactly the same model architecture (the original ResNet50, initialized from ImageNet pretrained model), training procedure and evaluation protocal as baselines on DomainBed. Importantly, the linear 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 [1e-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, 1e-2]
• Temperature: discrete set {0.05, 0.1}
Linear Probe Evaluation In all the experiments except for the end-to-end training setup, we always followed the procedure of two-stage training, where we 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, 1e-2, 1e-1, 1, 1e1, 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 a batch size 512, the Adam optimizer with a learning rate 5e-4 and early stopping. Note that an alternative was to just use the linear head 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 do that since we wanted to stick to the representation learning protocol with two-stage training. We did that in our end-to-end training setup since we wanted it to be compeletely comparable to baselines on DomainBed (which did not do refitting).
Selection of λ In our experiments, we treated λ as a special hyperparamter. For each model, we selected λ on the PACS dataset, and then used the same λ value for all other datasets, because our bottleneck was fairly robust to the choice of λ. The λ values chosen for SSL-based models (i.e., CLIP S, CLIP L, DINO) and end-to-end ResNet-50 were 1e-1 and 1e-5, respectively.

heading: D.4 LAION
text: Model We used the CLIP L model (i.e., CLIP ViT-B/32) with an additional network on top for finetuning. The additional network were two blocks of 2-layer MLP, each with hidden size 2048, pre-activation batch normalization, residual connection, and dropout probability 0.1. Note that the original CLIP L model was frozen and only the additional network was trained.
Dataset We used the LAION-400M dataset which contained 400 million image-text pairs for training. Though the dataset might not be as clean as the original CLIP training data (as evidenced by our experimental results), it was the largest publicly available image-text-pair dataset that we could get access to. As we froze the CLIP L model and only did finetuning, we used the 1TB preprocessed embeddings provided by LAION-400M3. No further preprocessing was applied.
3See https://laion.ai/laion-400-open-dataset/ for details.
Training We used the image-text contrastive loss as introduced in Radford et al. (2021) for training model. The temperature was learnable which was initialized as 0.07 and clipped with a minimum 0.01. The model was trained for 1 epoch using the Adam optimizer with a batch size of 16384 and a cosine learning rate decay schedule. The learning rate was tuned over the set {3e-5, 1e-4, 3e-4, 1e-3, 3e-3, 1e-2} and the λ value for the Ent bottleneck was tuned over {1e-3, 1e-2, 1e-1, 1, 1e1}.
Evaluation For the evaluation on the ImageNet-related datasets, we followed a similar procedure in Radford et al. (2021), where a linear classifier was fitted on ImageNet using the model representations and evaluated on 7 natural distribution shift datasets. In particular, we fitted a logistic regression classifier with 1e-5 L2 regularization on ImageNet training set which was trained with a batch size 512, the Adam optimizer with a learning rate 3e-4 and early stopping. Note that this was different from Radford et al. (2021), where a 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 testbed4. The evaluation datasets and their abbreviations used in Table 2 were: ImageNetV2 (IN-V2, Recht et al., 2019), ImageNet-Sketch (IN-S, Wang et al., 2019), Youtube-BB (YT-BB, Shankar et al., 2019), ImageNet-Vid (IN-Vid, Shankar et al., 2019), ObjectNet (Barbu et al., 2019), ImageNet Adversarial (IN-A, Hendrycks et al., 2021), and ImageNet Rendition (IN-R Hendrycks et al., 2020).

heading: E ADDITIONAL EXPERIMENTAL RESULTS

heading: E.1 SCIENTIFIC
text: What’s the effect of λ for different objectives on the worst-case DG performance? In Fig. 5, the worst-case target log likelihood versus λ values for different objectives is shown. We found that Ent is much more sensitive to the choice of λ than CAD, which was part of the reason why we used the latter in most of our experiments. Note that for SupCon-Ent with small λ values, it was worse than SupCon-Base because of the discretization introduced by the Ent bottleneck, which we verified by observing that setting λ = 0 lead to similar results.

heading: E.2 BRIDGE
text: The scientific setup is closer to our theory than what we do in practice in that worst-case predictor was considered and empirical generalization gap was ignored. Here we bridged these gaps with a more practical setup. In particular, we split the PACS dataset to training and validation splits for each domain and considered the setting: the representor trains the encoder on all-domain training splits with a validation loss selection; the learner trains the SVM predictor (average-case) on the source training split which is selected over the source validation split, and evaluates on the validation splits of other target domains. The target validation accuracy averaged over all (source, target) setups was reported. For simplicity, we will use CE to denote the objective with the cross-entropy loss that uses
4https://github.com/modestyachts/imagenet-testbed
labels to minimize R [Y |Z], and SupCon for the contrastive loss that uses supervised augmentations to minimize H[A |Z]. We will use CAD in following experiments unless otherwise specified (chosen with initial experiments). Details in Appx. D.2.
Does domain bottleneck improve the average-case DG performance? Though our theory focuses on the worst-case DG, we empirically showed that adding bottlenecks to enforce support match can also improve the average-case DG performance by comparing CE-Base and CE-CAD in Table 3.
What if the representor only has access to source domains? Similar to what we did in the 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 Table 3, which is much worse then CE-CAD. This shows the necessity of getting access to target domain information for DG.
What if the representor only has access to domain-agnostic augmentations? In Table 3, we also compared SupCon-CAD which used supervised augmentations through the labels with CECAD 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 our objectives. We empirically showed this by using augmentations of SimCLR (see Appx. D.1 for details) with our objectives (SimCLR-CAD). In Table 3, we indeed observed that using standard augmentations performed much worse than using desired augmentations (SupCon-CAD).
How do augmentations matter? Besides investigating the ‘Supervised’ augmentations (SupConCAD) and ‘Standard’ augmentations (SimCLR-CAD) above, we also compared other three augmentations as in the scientific section. Specifically, we 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 with ‘Supervised’, ‘SingleDom’, and ‘IntraDom’ augmentations gave 84.7 ± 0.4, 83.2 ± 0.3, and 77.5 ± 2.3, respectively. This shows the importance of using domain-agnostic augmentations in practice.
Do standard augmentations affect source performance? Previously, we showed that using standard augmentations hurt the DG performance measured by the average target accuracy. It is natural to ask whether using standard augmentations also hurt the source performance since we should also be interested in the ‘effective robustness’ (Taori et 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 if we consider the source-target gap which was 0.2 for SupCon-CAD and 28.4 for SimCLR-CAD, which still verified that the non-domain-agnostic standard augmentations were harder to force support match. To be even more convincing, we did the same analysis on VLCS, and the average source accuracy of SupCon-CAD and SimCLR-CAD were 86.6 ± 0.1 and 84.6 ± 0.5 which were fairly close, but the average target accuracy were 84.7 ± 0.4 and 57.5 ± 1.7, respectively.

heading: E.3 DOMAINBED
text: What is the impact of CLIP pretraining? To ensure that our gains are not only due to a novel CAD bottleneck, but the synergy between enforcing support constraint and using desired SSL models, we investigated CAD using the standard DomainBed protocol denoted as CAD in the table. It shows that CAD on its own performs similarly with DomainBed baselines (see Table 4 for a full comparison).
Full result of Table 1 We included the full result of Table 1 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 et 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 S with linear probe. But we observed that it was outperformed by our CLIP S + CAD.
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 in DomainBed (in the sense that we needed target domain information to achieve IDG). Here, we also provided results with ‘source validation’ selection in Table 5. Source validation selection relies on the assumption that source and target data follow similar distributions (Gulrajani & Lopez-Paz, 2021) thus source and target accuracy are highly correlated, which is not really true in practice. We found some issues with source validation selection results:
• The selected model with the highest source validation accuracy tends to 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 to diminish the effect of bottlenecks. This can be seen by the fact that the gap between CLIP + Base and CLIP + CAD of source validation selection is much smaller than that of oracle selection;

heading: E.4 LAION
text: Evaluation results on DomainBed We included the evaluation results of trained models on DomainBed in Table 6, where we followed exactly the same linear evaluation protocal discussed in Appx. D.3. We observed similar results as Table 2: the CLIP L model trained with the Ent bottleneck on LAION (LAION + Ent) outperformed the one without ((LAION + Base) on all DomainBed datasets, but slightly underperformed the original CLIP L model (which might be due to quality of the LAION-400M dataset).

Review summaries:
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.

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. 

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.  To 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.  They 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.