Original: O PTIMAL R EPRESENTATIONS FOR C OVARIATE S HIFTS
Revised: O PTIMAL R EPRESENTATIONS FOR C OVARIATE S HIFTS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Machine learning systems often experience a distribution shift between training and testing.
Revised: Machine learning systems often experience a distribution shift between training and testing.

Original: In this paper, we introduce a simple variational objective whose optima are exactly the set of  representations on which risk minimizers are guaranteed to be robust to any distribution shift that preserves the Bayes predictor, e.g., covariate shifts.
Revised: 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.

Original: Our objective has two components.
Revised: Our objective has two components.

Original: First, a representation must remain discriminative for the task, i.e., one predictor can simultaneously minimize the source and target risk.
Revised: First, a representation must remain discriminative for the task, i.e., some predictor must be able to simultaneously minimize the source and target risk.

Original: Second, the representation’s marginal support needs to be the same across source and target.
Revised: Second, the representation’s marginal support needs to be the same across source and target.

Original: We make thisobservation practical by designing self-supervised learning methods that  use unlabelled data and augmentations to train robust representations.
Revised: We make this practical by designing self-supervised learning methods that only use unlabelled data and augmentations to train robust representations.

Original: Our objectives achieve state-of-the-art results on
Revised: Our objectives achieve state-of-the-art results on

Original: DomainBed, and give insights into the robustness of recent methods, such as CLIP.
Revised: DomainBed, and give insights into the robustness of recent methods, such as CLIP.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: It is hard to build machine learning (ML) systems that are robust to distribution shifts between training (source) and testing (target).
Revised: It is hard to build machine learning (ML) systems that are robust to distribution shifts between a source (train) and target (test) domain.

Revised: One promising approach to domain generalization (DG) is learning robust representations from which predictors trained on source must perform well on target.

Original: No method for domain generalization (DG) uniformly outperforms empirical source-risk minimizers (ERM)in practice (Gulrajani & Lopez-Paz, 2021).
Revised: In practice, however, no current DG methods for learning representation uniformly outperform empirical source-risk minimizers (ERM) (Gulrajani & Lopez-Paz, 2021).

Original: Representation learning seems like a promising approach. For example, domain-invariant representation learning (e.g., DANN, Ganin et al., 2016) seeks representations that are invariant to distribution shifts, while retaining information about the source task.
Revised: Furthermore, our theoretical understanding of DG is still lacking.

Original: Unfortunately, this intuitive approach is not sound; there are examples of invariant representations that perform well on the source, but poorly on the target (Zhao et al., 2019; Johansson et al., 2019). While properties that would  imply robust representationsare known (Ben-David et al., 2010a), the minimal set of  requirements  is not  known.
Revised: Speciﬁcally, 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.

Original: We introduce the ﬁrst, simple, variational objective whose optima are exactly the set of  representations on which source risk minimizers are guaranteed to generalize across distribution shifts that preserve the Bayes predictor.
Revised: We introduce the ﬁrst, 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.

Original: We work in an idealized DG (IDG) setting; we assume that a learner has access to the source population risk.
Revised: We work in an idealized DG (IDG) setting; we assume that a learner has access to the source population risk.

Original: Our variational characterization implies that it is both sufﬁcient 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.
Revised: Our variational characterization implies that it is both sufﬁcient 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.

Original: This means that any optimal representation learning method must seek discriminative information about the target.
Revised: This means that any optimal representation learning method must seek discriminative information about the target.

Original: Even worse, we prove that without access to some knowledge about the target, any representation learning algorithm cannot uniformly (over all discriminative tasks) outperform a constant representation, which may explain why DG methods struggle to outperform ERM.
Revised: 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.

Original: We show how  these challenges can be overcome with access to target inputs and particulartypes of data augmentations that carry discriminative informationabout downstream tasks, but minimal domain-speciﬁc information.
Revised: 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-speciﬁc information.

Original: Text descriptions of images are examples of such augmentations, because they are informative for many downstream classiﬁcation tasks, but they remove a lot of domain-speciﬁc information.
Revised: Text descriptions of images are examples of such augmentations, as they are informative for many downstream classiﬁcation tasks, but they remove a lot of domain-speciﬁc information.

Original: With such augmentations, we design practical self-supervised learning (SSL) objectives for learning robust representations.
Revised: With such augmentations, we design practical self-supervised learning (SSL) objectives for learning robust representations.

Original: 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  results on DomainBed (Gulrajani & Lopez-Paz, 2021).
Revised: 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).

Original: To summarize, we:
Revised: To summarize, we:

Original: • provide all objectives whose optima achieve optimal domain generalization under covariate shift; • prove that it is impossible to learn useful representations without accessing target information; • provide practical objectives to learn  robust representations using speciﬁcdata augmentations; • get state-of-the-art results on typical domain generalization benchmarks.
Revised: • provide minimal sufﬁcient 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 speciﬁc augmentations; • get state-of-the-art results on typical domain generalization benchmarks.

Original:  B ACKGROUND : D OMAIN G ENERALIZATION AND REPRESENTATIONS
Revised:  B ACKGROUND : DOMAIN GENERALIZATION AND REPRESENTATIONS

Original: We are interested in predictions that are robust across distribution shifts.
Revised: We are interested in predictions that are robust across distribution shifts.

Original: We formalize this using domain generalization (DG) language.
Revised: We formalize this using domain generalization (DG) language.

Original: Given a distribution p X,Y | d s over inputs x ∈ X and labels y ∈ Y from the source domain d s ∈ D , we select a predictor f : X → Γ .
Revised: Given a distribution p X,Y | d s over inputs x ∈ X and labels y ∈ Y from the source domain d s ∈ D , we select a predictor f : X → Γ .

Original: The predictions γ ∈ Γ could for example be labels or distributions over labels.
Revised: The predictions γ ∈ Γ could for example be labels or distributions over labels.

Original: Despite being selected on the source domain, we would like f to achieve a small expected risk with respect to a loss function (cid:96) :
Revised: Despite being selected on the source domain, we would like f to achieve a small expected risk with respect to a loss function (cid:96) :

Original: Y × Γ → R ≥ 0 ,
Revised: Y × Γ → R ≥ 0 ,

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: on a distribution p X,Y | d from a target domain d = d t ∈ D , which is somehow related to d s .
Revised: on a distribution p X,Y | d from a target domain d = d t ∈ D , which is somehow related to d s .

Original: A common strategy for DG is to learn robust representations.
Revised: A common strategy for DG is to learn robust representations, which splits the problem into two.

Original: Speciﬁcally, the idea is to split the problem of learning robust predictors into two: ﬁrst, learn an encoder p Z | X , whichstochastically maps inputs X to representations Z .
Revised: First, learn an encoder p Z | X , which maps inputs X to representations Z .

Original: Then, learn a predictor h :
Revised: Then, learn a predictor h :

Original: Z → Γ from representations Z to labels Y using standard risk minimization.
Revised: Z → Γ from representations Z to labels Y using standard risk minimization.

Original: The goal is to design a robust representation, so that predictors h trained to minimize the source risk R d sh [ Y | Z ] also achieve low target risk R d th [ Y | Z ] .
Revised: The goal is to design a robust representation Z , so that predictors h trained to minimize the source risk R d sh [ Y | Z ] also achieve low target risk R d th [ Y | Z ] .

Original: Many methods have been proposed to try to learn suchrepresentations Z  by enforcing  invariance of the marginal p Z | d (e.g., Ganin et al., 2016) to the domain d .
Revised: Many methods have been proposed to try to learn such Z , e.g., by enforcing domain invariance of the marginal p Z | d (e.g., Ganin et al., 2016).

Original: Still, many of these proposals are not sound (Zhao et al., 2019; Johansson et al., 2019).
Revised: Still, many of these proposals are not sound (Zhao et al., 2019; Johansson et al., 2019).

Original: Furthermore, these methods rarely outperform source empirical risk minimization (ERM) in practice (Gulrajani & Lopez-Paz, 2021).
Revised: Furthermore, they rarely outperform source empirical risk minimization (ERM) in practice (Gulrajani & Lopez-Paz, 2021).

Original:  O PTIMAL REPRESENTATIONS FOR DOMAIN GENERALIZATION
Revised:  O PTIMAL REPRESENTATIONS FOR DOMAIN GENERALIZATION

Original: To separate domain generalization from ﬁnite sample generalization
Revised: To separate domain generalization from ﬁnite sample generalization

Original: , we consider an idealized DG (IDG) setting, where thelearner selects her predictor h according to the source population risk rather than empirical risk.
Revised: , we consider an idealized DG (IDG), where the predictor h is selected on the source population risk rather than empirical risk.

Original: In our conclusions, we discuss what is and is not lost in this idealization.

Original: We assume sample spaces X , Z , Y , D are discrete; formal statements and proofs are in Appxs.
Revised: We assume sample spaces X , Z , Y , D are discrete; formal statements and proofs are in Appxs.

Original: A and B.
Revised: A and B.

Original: 3. D EFINING O PTIMAL R EPRESENTATIONS FOR I DEALIZED D OMAIN G ENERALIZATION
Revised: 3. D EFINING O PTIMAL R EPRESENTATIONS FOR I DEALIZED D OMAIN G ENERALIZATION

Original: We are interested in evaluating the quality of a representation Z(i.e., encoder p Z | X ) of X for IDG.
Revised: We want to evaluate the quality of a representation Z of X .

Original: In ourmodel of IDG, the learner is given a random source D s ; she selects any source risk minimizer; andshe is scored according to her risk on a random target domain D t .
Revised: In our IDG, the learner is given a random source D s ; she selects any source risk minimizer; and is scored according to her risk on a random target domain D t .

Original: To give uniform guarantees while reﬂecting the uncertainty over which source-target pair she will be given, we measure the quality of Z as the expected value of the learner’s worst-case choice.
Revised: To give uniform guarantees while reﬂecting the uncertainty over the source-target pair ( D s , D t ) , we measure the quality of Z as the expected risk of the learner’s worst-case choice.

Original: Deﬁnition.
Revised: Deﬁnition.

Original: The idealized domain generalization risk (IDG risk) of an encoder p Z | X is the expected (over domains) worst-case (over source risk minimizers) target risk, i.e.,
Revised: The idealized domain generalization risk (IDG risk) of an encoder p Z | X is the expected (over domains) worst-case (over source risk minimizers) target risk, i.e.,

Original: [Equation]
Revised: [Equation]

Original: where H ∗ D s :=
Revised: where H ∗ D s :=

Original: arg min
Revised: arg min

Original: h R D s h
Revised: h R D s h

Original: [ Y | Z ] are the source risk minimizers.
Revised: [ Y | Z ] are the source risk minimizers.

Original: We call a representation Z ∗ (or its encoder) optimal for IDG if it minimizes the IDG risk: p Z ∗ | X ∈ arg min p Z | X R IDG
Revised: We call a representation Z ∗ (or its encoder) optimal for IDG if it minimizes the IDG risk: p Z ∗ | X ∈ arg min p Z | X R IDG

Original: [ Y | Z ] .
Revised: [ Y | Z ] .

Original: 3. C HARACTERIZING OPTIMAL REPRESENTATIONS FOR IDG UNDER COVARIATE SHIFT
Revised: 3. C HARACTERIZING OPTIMAL REPRESENTATIONS FOR IDG UNDER COVARIATE SHIFT

Original: 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.
Revised: 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.

Original: Under mild assumptions, we provide a simpliﬁed, equivalent objective, which is easier to optimize.
Revised: Under mild assumptions, we provide a simpliﬁed, equivalent objective, which is easier to optimize.

Original: For convenience, we assume that there is a unique Bayes predictor f ∗ , which minimizes the expected risk over domains, i.e., f ∗ =
Revised: For convenience, we assume that there is a unique Bayes predictor f ∗ , which minimizes the expected risk over domains, i.e., f ∗ =

Original: arg min f E p
Revised: arg min f E p

Original: Dt [R D t f
Revised: Dt [R D t f

Original: [ Y | X ]] .
Revised: [ Y | X ]] .

Original: This is satisﬁed by standard ML tasks p Y,X and losses (cid:96) .
Revised: This is satisﬁed by standard ML tasks p Y,X and losses (cid:96) .

Original: More importantly, we assume the following domain structure, which ensures the existence of optimal encoders and allows our simpliﬁcation.
Revised: More importantly, we assume the following domain structure, which ensures the existence of optimal encoders and allows our simpliﬁcation.

Original: Assumptions.
Revised: Assumptions.

Original: All domains d ∈ D we consider are related by the following assumptions:
Revised: All domains d ∈ D we consider are related by the following assumptions:

Original: Generalized covariate shift.
Revised: Generalized covariate shift.

Original: All domain-speciﬁc risk minimizers f ∈ arg min f
Revised: All domain-speciﬁc risk minimizers f ∈ arg min f

Original: [R d f [ Y | X ]] are equal to the Bayes predictor f ∗ on their support,
Revised: [R d f [ Y | X ]] are equal to the Bayes predictor f ∗ on their support,

Original: i
Revised: i

Original: e., f ( x ) =
Revised: e., f ( x ) =

Original: f ∗ ( x ) for all x ∈ supp( p X | d ) .[Table]
Revised: f ∗ ( x ) for all x ∈ supp( p X | d ) .

Original: +KqnqRZFQbz/t05uYXFpeWcyv51bX1jc3C1nZNJ6nCpIoTlqhGhDRhVJCqoYaRhlQE8YiRetS7GPr1O6I0TcS16UsScNQRNKYYGSvdyDC7vW+HehAWil7JG8GdJf6EFMu73teL+JaVsPDRaic45UQYzJDWTd+TJsiQMhQzMsi3Uk0kwj3UIU1LBeJEB9no4IF7YJW2GyfKljDuSP09kSGudZ9HtpMj09XT3lD8z2umJj4PMipkaojA40VxylyTuMPv3TZVBBvWtwRhRe2tLu4ihbCxGeVtCP70y7OkdlzyT0qnVzaNIxgjB3uwD4fgwxmU4RIqUAUMHB7gCZ4d5Tw6r87buHXOmczswB847z9TH5Rq</latexit>

Original: [Figure]

Original: Invariance of Bayes predictions.
Revised: Invariance of Bayes predictions.

Original: The set of Bayes predictions is the same for all domains, i.e., (cid:8) 
Revised: The set of Bayes predictions is the same for all domains, i e., (cid:8) f ∗ ( x ) | x ∈ supp( p X | d ) (cid:9) = { f ∗ ( x ) | x ∈ X} .

Revised: [Table]

Revised: +KqnqRZFQbz/t05uYXFpeWcyv51bX1jc3C1nZNJ6nCpIoTlqhGhDRhVJCqoYaRhlQE8YiRetS7GPr1O6I0TcS16UsScNQRNKYYGSvdyDC7vW+HehAWil7JG8GdJf6EFMu73teL+JaVsPDRaic45UQYzJDWTd+TJsiQMhQzMsi3Uk0kwj3UIU1LBeJEB9no4IF7YJW2GyfKljDuSP09kSGudZ9HtpMj09XT3lD8z2umJj4PMipkaojA40VxylyTuMPv3TZVBBvWtwRhRe2tLu4ihbCxGeVtCP70y7OkdlzyT0qnVzaNIxgjB3uwD4fgwxmU4RIqUAUMHB7gCZ4d5Tw6r87buHXOmczswB847z9TH5Rq</latexit>

Original: [Equation]
Revised: [Figure]

Original: Generalized covariate shift (GCS) ensures that f ∗ is simultaneously optimal on all domains.
Revised: Generalized covariate shift (GCS) ensures that f ∗ is simultaneously optimal on all domains.

Original: For log-loss (cid:96) it recovers standard covariate shift, i.e., p Y | x,d = p Y | x .
Revised: For log-loss (cid:96) it recovers standard covariate shift, i.e., p Y | x,d = p Y | x .

Original: For other losses, Asm. 1 is weaker, e.g., it only requires invariance of most likely labels for 0-1 loss, and of conditional expectations for
Revised: 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

Original: MSE.
Revised: MSE.

Original: Invariance of Bayes predictors is necessary to learn useful predictors using a single domain.
Revised: Invariance of Bayes predictors is necessary to learn useful predictors using a single domain.

Original: For example, for 0-1 loss it ensures that each label is seen at least once in each domain.
Revised: For example, for 0-1 loss it ensures that each label is seen at least once in each domain.

Original: The intuition behind our objective is that under GCS any source risk minimizer will make optimal predictions onall target samples x that are also in the source domain’s support.
Revised: 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.

Original: Thus,  optimal representationsfor IDG are exactly those that (a)ensure that all domains have the same support in Z , and (b) retain GCS from Z without sacriﬁcing the ability to predict Y optimally, which can be ensured by minimizing theachievable riskof from Z .
Revised: 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 sacriﬁcing the ability to predict Y , which can be ensured by minimizing the risk from Z .

Original: See Fig.
Revised: See Fig.

Original: Theorem 1.
Revised: Theorem 1.

Original: Under our assumptions, an encoder p Z ∗ | X is optimal for IDG if and only if it minimizes the risk R [ Y | Z ] := inf 
Revised: Under our assumptions, an encoder p Z ∗ | X is optimal for IDG if and only if it minimizes the risk R [ Y | Z ] := inf [Equation]

Original: h E p Dt (cid:2) R D t h
Revised: E p Dt (cid:2) R [Equation] t h

Original: [ Y | Z ] (cid:3) while matching the support of Z across domains, i.e.,
Revised: [ Y | Z ] (cid:3) while matching the support of Z across domains, i.e.,

Original: [Equation]
Revised: [Equation]

Original: Moreover, such encoders exist and their IDG risk is the Bayes risk
Revised: Moreover, such encoders exist and their IDG risk is the Bayes risk

Original: R IDG [ Y | Z ∗ ]
Revised: R IDG [ Y | Z ∗ ]

Original: =
Revised: =

Original: R [ Y | X ] .
Revised: R [ Y | X ] .

Original: Theorem 2 provides an objective to learn representations on which performing risk minimization using a single domain  is as good as performing risk minimization on all domains simultaneously, i.e., R IDG [ Y | Z ∗ ] = R [ Y | X] .
Revised: 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 .

Original: Other sufﬁcient objectives had previously beenproven or hinted towards (Ben-David et al., 2010a), e.g.,minimizing the risk while matching the representation’s marginal.
Revised: Other sufﬁcient conditions have previously been hinted towards, e.g., matching the marginal p Z | 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).

Original:  To our knowledge, Thm. 2 isnevertheless the ﬁrst that identiﬁes the necessary and sufﬁcient conditions for  representations Z ∗ that are optimal for IDG.
Revised: Previous conditions are nevertheless generally neither necessary nor achievable. To our knowledge, Thm. 1 is the ﬁrst characterization of necessary and sufﬁcient conditions for IDG optimal representations Z ∗ .

Original: This has the advantage of giving better insights into IDG and provides a framework for deriving all objectives that describe optimal IDG.
Revised: Theorem 1 thus gives better insights into IDG and provides a framework for deriving the least stringent objectives for optimal IDG.

Original: The risk minimization (Eq.
Revised: The risk minimization (Eq.

Original: (3)) shows that one must have some knowledge about the target domains to learn optimal representations for IDG.
Revised: (3)) shows that one must have some knowledge about the target domains to learn optimal representations for IDG.

Original: Access to targets might seem unrealistic, but without such knowledge or additional assumptions it is provably impossible to beat even constant representations.
Revised: Access to targets might seem unrealistic, but without such knowledge or additional assumptions it is provably impossible to beat even constant representations.

Original: Proposition 1 (No free lunch for IDG) .
Revised: Proposition 1 (No free lunch for IDG) .

Original: Let d s be any source domain, Z d s be any representation chosen on source d s , and
Revised: Let d s be any source domain, Z d s be any representation chosen on source d s , and

Original: C ∈ Z be a constant representation.
Revised: C ∈ Z be a constant representation.

Original: Under minor assumptions, for every “good” target domain outside the source’s support on which Z d s outperforms C for IDG, there are many “bad” target domains on which Z d s is strictly worse than C .
Revised: Under minor assumptions, for every “good” target domain outside the source’s support on which Z d s outperforms C for IDG, there are many “bad” target domains on which Z d s is strictly worse than C .

Original: Formal statement in Appx.
Revised: Formal statement in Appx.

Original: B.3.
Revised: B.3.

Original: Proposition 1 shows that target knowledge is necessary for learning useful representations in IDG.
Revised: Proposition 1 shows that target knowledge is necessary for learning useful representations in IDG.

Original: This may explain why previous  methods have been unable toto outperform ERM in standardDG benchmarks (Gulrajani & Lopez-Paz, 2021): the knowledge they have access to is insufﬁcient.
Revised: 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 insufﬁcient to generalize.

Original: Taken together, Prop. 1 and Thm.
Revised: Taken together, Prop. 1 and Thm.

Original: 2 say  either you allow yourself toconsider target domains d t , in which case you can achieve an IDG risk that matches supervised learning, or you do not access d t , in which caseand any representation learning algorithm can achieve worst IDG risk than a constant.[Figure]
Revised: 1 say that either you have access to target domains d t , in which case you can achieve an IDG risk that matches supervised learning, or you do not access d t , in which case any representation learning algorithm can achieve worse IDG risk than a constant.

Original:  L EARNING REPRESENTATIONS UNDER COVARIATE SHIFT 4. S ELF - SUPERVISED LEARNING USING DOMAIN - COVERING AUGMENTATIONS
Revised:  L EARNING REPRESENTATIONS UNDER COVARIATE SHIFT 4. S ELF - SUPERVISED LEARNING USING DOMAIN - AGNOSTIC AUGMENTATIONS

Original: Our characterization of optimal representations for IDG (Thm. 2) requires labeled data from all domains, which is impractical.
Revised: Our characterization of optimal representations for IDG (Thm. 1) requires labeled data from all domains, which is impractical.

Original: 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.
Revised: We show how this can be overcome with self-supervised learning [Figure] (SSL), which is a technique for training representations without direct access to labels, and a particular class of data augmentations.

Original: 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).
Revised: 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).

Original: 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).
Revised: 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).

Original: More precisely, an augmentation A is a random variable sampled conditionally from the input X .
Revised: More precisely, an augmentation A is a random variable sampled conditionally from the input X .

Original: The key requirement  that allows us to retain task information is, that, if any samples x, x (cid:48) ∈ X have the same augmentation conditional p
Revised: The key requirement is that augmentations retain task information. Speciﬁcally, if any samples x, x (cid:48) ∈ X have the same augmentation conditional p

Original: A | x = p
Revised: A | x = p

Original: A | x (cid:48) , then their Bayes predictions must be the same f ∗ ( x ) =
Revised: A | x (cid:48) , then their Bayes predictions must be the same f ∗ ( x ) =

Original: f ∗ ( x (cid:48) ) .
Revised: f ∗ ( x (cid:48) ) .

Original: With augmentations that preserve label information in this way, one can learn an encoder that minimizes the risk R [ Y | Z ] by instead maximizing mutual information
Revised: With such A , one can learn an encoder that minimizes the risk R [ Y | Z ] by instead maximizing mutual information

Original: I[ A ; Z ] .
Revised: I[ A ; Z ] .

Original: Intuitively, if Z has allinformation about the augmentation A , then it must have information about the conditional p
Revised: Intuitively, if Z has all augmentation information, then it must have information about the conditional p

Original: A | X , and thus the Bayes prediction f ∗ ( X ) .
Revised: A | X , and thus the Bayes prediction f ∗ ( X ) .

Original: This suggeststhe possibility of learning optimal representations for IDG by replacing Eq.
Revised: This suggests learning optimal representations for IDG by replacing Eq.

Original: (3) with a maximization of I[ A ; Z ] .
Revised: (3) with a maximization of I[ A ; Z ] .

Original: Unfortunately, fully optimizing I
Revised: Unfortunately, fully optimizing I

Original: [ A ; Z ] w.r.t.
Revised: [ A ; Z ] w.r.t.

Original: p Z | X is not generally possible under the support constraint Eq.
Revised: p Z | X is not generally possible under the support constraint Eq.

Original: This can be overcome under a domain-covering assumption, which requires that for every domain there is an example that gets mapped to every possible augmentation distribution, i.e., (cid:8) p
Revised: 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., (cid:8) p

Original: A | x | x ∈ supp( p X | d ) (cid:9) = (cid:8) p A | x | x ∈ X (cid:9) .
Revised: [Equation]

Original: Proposition 2.
Revised: Proposition 2.

Original: Let p A |
Revised: Let p A |

Original: X be a domain-covering augmenter.
Revised: X be a domain-agnostic augmenter.

Original: Then any optimal solution p Z ∗ | X of the following objective is optimal for IDG:
Revised: Then any optimal solution p Z ∗ | X of the following objective is optimal for IDG:

Original: [Equation]
Revised: [Equation]

Original: Proposition 2 shows that we can still learn  optimal representationsfor IDG without labels as long as we have access to the right augmentations.
Revised: Proposition 2 shows that we can still learn IDG optimal representations without labels if we have access to the right augmentations.

Original: But how realistic are those augmentations?
Revised: How realistic are those augmentations?

Original: For 0loss (cid:96) , the most likely label should be preserved, which is satisﬁed by standard image augmentations likemild rotations and color jittering.
Revised: For 0-1 loss (cid:96) , the most likely label should be preserved, which is satisﬁed by standard image augmentations like rotations and color jittering.

Original: Those augmentations are nevertheless not domain-covering for typical domains (e.g. sketches and photos), since outputs A arehighly correlated with the  domain D of the original input X , as seen in Fig.
Revised: 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.

Original: 2a.
Revised: 2a.

Original: A practical choice of augmentation that is nearly domain-covering, is a mapping from images to text descriptions, as with CLIP (Radford et al., 2021).
Revised: 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.

Original: Image-text augmentations have many advantages.
Revised: Image-text augmentations have many advantages.

Original: First, text augmentations preserve label information for many downstreamclassiﬁcation tasks.
Revised: First, text augmentations preserve label information for many downstream tasks.

Original: Second, they are close to being domain-covering, since images from different domains (e.g., sketches and photos) but similar semantics are often mapped to similar descriptions 1 (Fig. 2c).
Revised: 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.(Fig. 2c).

Original: This may explain the incredible robustness of CLIP’s representations when compared to other SSL methods (Chen et al., 2020; He et al., 2020; Grill et al., 2020).
Revised: This gives insights into the open question (Radford et al., 2021) about why CLIP’s representations are so robust compared to other SSL methods.

Original: Finally, image-text pairs are easy to access in practice given their abundance on the internet.
Revised: Finally, image-text pairs are easy to access in practice given their abundance on the internet.

Original: There is still the question of enforcing the support constraint for augmented data ( X, A ) (Eq.
Revised: 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.

Original: (5)), which may not come with domain information.
Revised: In practice, even the domain information D is usually unknown.

Original: In this case, one can  replace the support constraint with a stronger one, e.g., minimizing I[ Z ; X ] (see Sec. 4.2.2), that does not rely on the domain information.
Revised: One can nevertheless still optimize (Eq. (4)) by replace the support constraint with a stronger one that does not rely on D e.g., minimizing I[ Z ; X ] (see Sec. 4.2.2), .

Original: This highlights the potential of Prop. 2: if one can ﬁnd a large source of inputs X and domain-covering 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.
Revised: This highlights the potential of Prop. 2: if one can ﬁnd 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.

Original: 4. P RACTICAL OBJECTIVES
Revised: 4. P RACTICAL OBJECTIVES

Original: We now design practical objectives for learning optimal representations without labels.
Revised: We now design practical objectives for learning optimal representations without labels.

Original: Propositiondoes provide an objective but it is impractical as it involves constrained optimization.
Revised: Propositiondoes provide an objective but it is impractical as it involves constrained optimization.

Original: 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,
Revised: 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,

Original: [Equation]
Revised: [Equation]

Original: where H[ A | Z ] replaces I
Revised: where H[ A | Z ] replaces I

Original: [ A ; Z ]
Revised: [ A ; Z ]

Original: = H
Revised: = H

Original: [ A ] − H
Revised: [ A ] − H

Original: [ A | Z ] as H
Revised: [ A | Z ] as H

Original: [ A ] is a constant with respect to the encoder.
Revised: [ A ] is a constant with respect to the encoder.

Original: Eq.
Revised: Eq.

Original: (6) is a valid reformulation of Prop. 2 as long as minimizing B[ Z, D ] while maximizing I
Revised: (5) is a valid reformulation of Prop. 2 as long as minimizing B[ Z, D ] while maximizing I

Original: [ A ; Z ] enforces the support constraint  Eq.
Revised: [ A ; Z ] enforces the support constraint in Eq.

Original: Below, we provide different choices of such B[ Z, D ] each of which results in a different SSL objective.
Revised: Below, we provide different choices of such B[ Z, D ] each of which results in a different SSL objective.

Original: In practice, however, terms in Eq.
Revised: In practice, however, terms in Eq.

Original: (6) are hard to estimate from ﬁnite samples.
Revised: (5) are hard to estimate from ﬁnite samples.

Original: We now discuss two variational bounds that can be efﬁciently estimated and optimized with stochastic gradient descent (Bottou, 2010).
Revised: We now discuss two variational bounds that can be efﬁciently estimated and optimized with stochastic gradient descent (Bottou, 2010).

Original: For simplicity, we use a deterministic encoder e ϕ : X → Z for the rest of the paper.
Revised: For simplicity, we use a deterministic encoder e ϕ : X → Z for the rest of the paper.

Original: Detailed derivations are in Appx.
Revised: Detailed derivations are in Appx.

Original: C.
Revised: C.

Original: For both practical objectives we use an upper bound on H
Revised: For both practical objectives we use an upper bound on H

Original: [ A | Z ] ≤ E p A,Z [ − log q
Revised: [ A | Z ] ≤ E p A,Z [ − log q

Original: ( A | Z )] , where q is the contrastive variational distribution as in InfoNCE (Oord et al., 2018), which is standard in SSL.
Revised: ( A | Z )] , where q is the contrastive variational distribution as in InfoNCE (Oord et al., 2018), which is standard in SSL.

Original: Speciﬁcally, for a sample X , we construct X :=
Revised: Speciﬁcally, for a sample X , we construct X :=

Original: { X, X − 1 , . . .
Revised: { X, X − 1 , . . .

Original: , X − n } where each X − i are i.i.d.
Revised: , X − n } where each X − i are i.i.d.

Original: sampled from p X .
Revised: sampled from p X .

Original: Then we obtain a collection A :
Revised: Then we obtain a collection A :

Original: = { A + , A − 1 , . . .
Revised: = { A + , A − 1 , . . .

Original: , A − n } of one positive augmentation A + and n negatives A − i by independently sampling an augmentation from p A | X (cid:48) for each example X (cid:48)
Revised: , A − n } of one positive augmentation A + and n negatives A − i by independently sampling an augmentation from p A | X (cid:48) for each example X [Equation]

Original: ∈ X .
Revised: ∈ X .

Original: InfoNCE then uses a critic s ψ to score how likely each A (cid:48)
Revised: InfoNCE then uses a critic s ψ to score how likely each A (cid:48)

Original: ∈
Revised: ∈

Original: A is to be positive, i.e.,
Revised: A is to be positive, i.e.,

Original: H
Revised: H

Original: [ A | Z ] ≤ E p X, A ,Z [ − log q ψ, A ( A + | Z )] where q 
Revised: [ A | Z ] ≤ E p X, A ,Z [ − log q ψ, A ( A + | Z )] where q [Equation] | Z

Original: [Equation]
Revised: [Equation]

Original: 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 ) :=
Revised: 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 ) :=

Original: e ϕ ( A ) T Z .
Revised: e ϕ ( A ) T Z .

Original: Many previous DG regularizers (e.g., Ganin et al., 2016; Li et al., 2018b;a) could be validthe domain bottlenecks.
Revised: Many previous DG regularizers (e.g., Ganin et al., 2016; Li et al., 2018b;a) could be valid domain bottlenecks.

Original: In the following, we discuss two possible B[ Z, D ] , the ﬁrst of which is novel.
Revised: In the following, we discuss two possible B[ Z, D ] , the ﬁrst of which is novel.

Original: 4.2. C ONTRASTIVE A DVERSARIAL D OMAIN B OTTLENECK (CAD)
Revised: 4.2. C ONTRASTIVE A DVERSARIAL D OMAIN B OTTLENECK (CAD)

Original: Require: e
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: , D,
Revised: [Table] [Equation] [Table]

Original: [Equation]
Revised: [Equation]

Original: , n 1:
Revised: [Table] [Equation] [Table]

Original: [Equation]
Revised: [Equation]

Original: ( X ) 2:
Revised: [Table] [Equation] [Table] [Equation]

Original: A
Revised: [Table] [Equation] [Table]

Original: [Equation]
Revised: [Equation]

Original: ← sample ( p A | X )
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: sample ( p D,X,A ) 4: X , A ←{ X } ∪ (cid:8) X − i (cid:9) n i =1 , { A + } ∪ (cid:8) A − i (cid:9) n i = 5: X D ←{ X } ∪ (cid:8) X − i | D − i
Revised: [Table]

Original: = D, i ∈
Revised: [Equation]

Original: [ n ] (cid:9) 6: L aug ← − log exp s ψ
Revised: [Table]

Original: ( A + ,Z ) (cid:80)
Revised: [Equation]

Original: A (cid:48)∈ A exp s ψ ( A (cid:48) ,Z ) 7: L supp ← log (cid:80) X (cid:48)∈ X D exp e ϕ ( X (cid:48) )
Revised: [Table] [Equation] [Table] : L supp ← log (cid:80) X (cid:48)∈ X D exp e ϕ ( X (cid:48) )

Original: T Z (cid:80)
Revised: T Z (cid:80)

Original: X (cid:48)(cid:48)∈ X exp e ϕ
Revised: X (cid:48)(cid:48)∈ X exp e ϕ

Original: ( X (cid:48)(cid:48) )
Revised: ( X (cid:48)(cid:48) )

Revised: T Z (cid:46) I

Original: T Z  8: return L CAD =
Revised: [ Z ; D ] 8: return L CAD =

Original: L aug + λ L supp
Revised: L aug + λ L supp

Original: Our ﬁrst domain bottleneck minimizes  I[ Z ; D ] , which enforces support match using a KL divergence.
Revised: Our ﬁrst domain bottleneck minimizes B[ Z, D ] = I[ Z ; D ] , which enforces support match using a KL divergence.

Original: Dropping constants w.r.t.
Revised: Dropping constants w.r.t.

Original: Z we thus aim to maximize H
Revised: Z we thus aim to maximize H

Original: [ D | Z ] .
Revised: [ D | Z ] .

Original: Domain-adversarial neural network (DANN, Ganin et al., 2016) does so by ensuring that a domain classiﬁer q φ cannot predict domains from representations, i.e., it maximizes E p D,Z [ − log q φ ( D | Z )] ≥ H
Revised: Domain-adversarial neural network (DANN, Ganin et al., 2016) does so by ensuring that a domain classiﬁer q φ cannot predict domains from representations, i.e., it maximizes E p D,Z [ − log q φ ( D | Z )] ≥ H

Original: [ D | Z ] w.r.t.
Revised: [ D | Z ] w.r.t.

Original: encoder parameter ϕ
Revised: encoder parameter ϕ

Original: but minimizes it w.r.t.
Revised: but minimizes it w.r.t.

Original: φ .
Revised: φ .

Original: 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.
Revised: 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

Original: To overcome these issues, we construct q ( D | Z ) without introducing additional parameters and with a bound that is tight with enough samples.
Revised: To overcome these issues, we construct q ( D | Z ) without introducing additional parameters and with a bound that is tight with enough samples.

Original: In short, using the equality p D | Z = E p X | Z (cid:2) p D | X (cid:3) , we set our variational distribution to q ϕ, X , D
Revised: In short, using the equality p D | Z = E p X | Z (cid:2) p D | X (cid:3) , we set our variational distribution to q ϕ, X , D

Original: ( D | Z ) =
Revised: ( D | Z ) =

Original: E q ϕ, X [ˆ p ( D | X )] , where ˆ p is a count estimate of p D | X , and q ϕ, X ( X | Z ) is a contrastive family similar to Eq.
Revised: E q ϕ, X [ˆ p ( D | X )] , where ˆ p is a count estimate of p D | X , and q ϕ, X ( X | Z ) is a contrastive family similar to Eq.

Original: (7) but with critic e ϕ ( X ) T Z , Detailed derivations and explanations are in Appx.
Revised: (6) but with critic e ϕ ( X ) T Z , Detailed derivations and explanations are in Appx.

Original: C.3.
Revised: C.3.

Original: The resulting contrastive adversarial domain (CAD) objective is detailed in Algorithm 1.
Revised: The resulting contrastive adversarial domain (CAD) objective is detailed in Algorithm 1.

Original: First, sample domains D :=
Revised: First, sample domains D :=

Original: { D, D − 1 , . . .
Revised: { D, D − 1 , . . .

Original:  D − n } for each X (cid:48) ∈ X .
Revised: , D − n } for each X (cid:48) [Equation]

Original:  Then collect inputs associated  the current domain D , i.e., X D := { X } ∪{ X − i | D − i
Revised: X . Then collect inputs associated with the current domain D , i.e., X D := { X }∪{ X − i | D − i

Original: =
Revised: =

Original: D, i ∈
Revised: D, i ∈

Original: [ n ] } .
Revised: [ n ] } .

Original: Finally, compute q
Revised: Finally, compute q

Original: [Equation]
Revised: [Equation]

Revised: X D

Revised: [Equation]

Revised: ϕ, X

Revised: [Equation]

Revised: X

Revised: [Equation]

Original:  Z ) .
Revised: | Z ) .

Original: The resulting criterion is
Revised: The resulting criterion is

Original: [Equation]
Revised: [Equation]

Original: In Appx.
Revised: In Appx.

Original: 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.
Revised: 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.

Original: 4.2. E NTROPY B OTTLENECK (E NT )
Revised: 4.2. E NTROPY B OTTLENECK (E NT )

Original: Our second domain bottleneck is the entropy bottleneck (Ent) that minimizes H[ Z ] =
Revised: Our second domain bottleneck is the entropy bottleneck (Ent) that minimizes H[ Z ] =

Original: I[ Z ; X ] ≥
Revised: I[ Z ; X ] ≥

Original: I[ Z ; D ] , where the ﬁrst equality uses the determinism of the encoder.
Revised: I[ Z ; D ] , where the ﬁrst equality uses the encoder’s determinism.

Original: Ent enforces support match by removing all information that is not needed to maximize I[ Z ; A ] .
Revised: Ent enforces support match by removing all information that is not needed to maximize I[ Z ; A ] .

Original: In particular, minimizing I[ Z ; X ] is more stringent than I[ Z ; D ] , as it not only matches the representationsacross domains but also inside a domain.
Revised: In particular, minimizing I[ Z ; X ] is more stringent than I[ Z ; D ] , as it also matches the representations inside a domain.

Original: Although Entremoves more information than what is necessary compared to CAD, it has the advantage of not requiring access to domain samples D , which are rarely accessible in SSL.
Revised: The advantage of Ent is that it does not require domain samples D , which are rarely accessible in SSL.

Original: We consider the standard variational bound used in neural compression (Ballé et al., 2016; Theis et al., 2017), H[ Z ] ≤ E p Z [ − log q θ ( Z )] , where an entropy model q θ ( Z ) is used.
Revised: We consider the standard variational bound used in neural compression (Ballé et al., 2016; Theis et al., 2017), H[ Z ] ≤ E p Z [ − log q θ ( Z )] , where an entropy model q θ ( Z ) is used.

Original: This leads to,
Revised: This leads to,

Original: [Equation]
Revised: [Equation]

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Provably robust representationsZ under covariate shift.
Revised: Provably robust representations under covariate shift.

Revised: Previous work mostly focuses on domain generalization bounds for robust representations.

Original: Ben-David et al.
Revised: Ben-David et al.

Original: (2010a) bounds the target risk using  source risk and a divergence between source and target distributions. They do not consider representation learning, but in our setting, this implies that matching the marginal of Z while minimizing R [ Y | Z ] is sufﬁcient for optimality.
Revised: (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.

Original: Ben-David et al.
Revised: Mansour et al.

Original: (2010b) suggests that R [ Y | Z ] is not sufﬁcient, and
Revised: (2009) generalizes these results from 0-1 loss to more general losses.

Original: Zhao et al.
Revised: Johansson et al.

Original: (2019) also prove that one should minimize the joint R [ Y | Z ] instead of the source R d s [ Y | Z ] risk.
Revised: (2019) takes this further by deriving a support-based bound. In our setting, these bounds only hint towards a sufﬁcient condition for optimality, i.e., matching the marginal p Z | d or its support while minimizing R [ Y | Z ] . However, these bounds can often be loose and the implied sufﬁcient conditions are neither necessary nor generally achievable.

Original: Similarly, des Combes et al.
Revised: Ben-David et al.

Original: (2020) shows that matching the conditional p Y | Z,d = p Y | Z  is sufﬁcient. Johansson et al.
Revised: (2010b) suggests that separately minimizing R [ Y | Z ] or matching the marginal is not sufﬁcient, while Zhao et al.

Original: (2019) take this further by proving that matching only the support of Z  is also sufﬁcient.
Revised: (2019) also proves minimizing only the source risk R d s [ Y | Z ] is not sufﬁcient; but none of them proves the desired necessary condition.

Original: Our work distinguishesitself from those and other related work on three key aspects: (i) We are the ﬁrst to  provide the set of necessary and sufﬁcient conditions for robust representations; (ii) We prove that one can  learn optimal Z ∗ with SSL using only large samples of inputs X and domain-covering augmentations A . (iii) We consider a  general DG setting which deals with  a less stringent generalized covariate shift and works for all standard losses and Y in ML.
Revised: Our work distinguishes from previous work on three key aspects: (i) we are the ﬁrst to study and formalize optimally robust representations, and provide the achievable sufﬁcient 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.

Original: Still, our work is more speciﬁc than others, as we consider idealized DG and unrestricted predictors H .
Revised: Still, our work is more speciﬁc than others, as we consider idealized DG and unrestricted predictors H .

Original: Our theory could be combined with Dubois et al.’s
Revised: Our theory could be combined with Dubois et al.’s

Original: (2020), who provide conditions for optimal generalization from ﬁnite samples and constrained H in supervised learning.
Revised: (2020), who provide conditions for optimal generalization from ﬁnite samples and constrained H in supervised learning.

Original: Practical objectives for DG.
Revised: Practical objectives for DG.

Original: The most popular DG methods aim to learn domain-invariant representation by minimizing various divergernces between the conditionals p Z | d andmarginals p Z
Revised: The most popular DG methods aim to learn domain-invariant representation by minimizing various divergernces between the marginal distributions p Z | d and

Original:  (Long et al., 2015; Ganin et al., 2016; Sun & Saenko, 2016;
Revised: p Z (Long et al., 2015; Ganin et al., 2016; Sun & Saenko, 2016;

Original: Long et al., 2017; Li et al., 2018a;
Revised: Long et al., 2017; Li et al., 2018a;

Original: Shen et al., 2018; Nguyen et al., 2021).
Revised: Shen et al., 2018; Nguyen et al., 2021).

Original: Others propose matching the conditional p Z | y,d across domains instead (Gong et al., 2016;
Revised: Others propose matching the conditional p Z | y,d across domains instead (Gong et al., 2016;

Original: Li et al., 2018b; Tachet des Combes et al., 2020).
Revised: Li et al., 2018b; Tachet des Combes et al., 2020).

Original: These regularizers would all be valid domain bottlenecks B
Revised: These regularizers would all be valid domain bottlenecks B

Original: [ Z, D ] .
Revised: [ Z, D ] .

Original: Another line of work aims at learning Z with invariant predictors p Y | z,d across domains (e.g., Arjovsky et al., 2019; Krueger et al., 2021;
Revised: Another line of work aims at learning Z with invariant predictors p Y | z,d across domains (e.g., Arjovsky et al., 2019; Krueger et al., 2021;

Original: Li et al., 2021).
Revised: Li et al., 2021).

Original: However, none of these methods outperform ERM with fair model selections (Gulrajani & Lopez-Paz, 2021).
Revised: However, none of these methods outperform ERM with fair model selections (Gulrajani & Lopez-Paz, 2021).

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: 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.
Revised: 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.

Original: Unless stated otherwise, we consider a two-stage training setup.
Revised: Unless stated otherwise, we consider a two-stage training setup.

Original: First, the representation learner (“the representor”) trains an encoder p Z | X using a speciﬁed objective and freezes it.Then, the person performing predictions (“the learner”) Base Ent
Revised: First, the representation learner (“the representor”) trains an encoder p Z | X using a speciﬁed objective and freezes it.

Original: CAD

Original: R [ Y | Z ]

Original: H

Original: [ A | Z ] -5.1 ± 0.3 -0.4 ± 0.1 -0.7 ± 0.

Original: -3.5 ± 0.1 -0.0 ± 0.0 -0.8 ± 0.

Original: [-(a) Effect of different objectives 10101010 
-]
Revised: [-(a) Effect of different objectives 10101010 
-]

Original: L o g li k e li h oo d source target (b) Effect of λ

Original: DC

Original: Approx.

Original: Non-DC Augmentation type - - - - -

Original: L o g li k e li h

Original: oo d

Original: Supervised SingleDom

Original: ApproxDC IntraDom Standard (c) Effect of augmentations but approx. trains her predictor h from Z by minimizing the risk on source data.
Revised: Then, the person performing predictions (“the learner”) trains her predictor h from Z by minimizing the risk on source data.

Original: Finally, the representation Z and predictor h are evaluated on target data.
Revised: Finally, the representation Z and predictor h are evaluated on target data.

Original: In all experiments, the learner uses a linear classiﬁer for h .
Revised: In all experiments, the learner uses a linear classiﬁer for h .

Original: For the Ent bottleneck, we used Ballé et al.’s
Revised: For the Ent bottleneck, we used Ballé et al.’s

Original: (2018) entropy model.
Revised: (2018) entropy model.

Original: For the CAD bottleneck we used its conditional version whenever labels were available.
Revised: For the CAD bottleneck we used its conditional version whenever labels were available.

Original: When a model contains no domain bottleneck, we label it as “Base”.
Revised: When a model contains no domain bottleneck, we label it as “Base”.

Original: For experimental details and additional results see Appxs.
Revised: For experimental details and additional results see Appxs.

Original: D and E.
Revised: D and E.

Original: 6. S CIENTIFIC SETTING :  OPTIMAL REPRESENTATIONS FOR WORST - CASE DG
Revised: 6. S CIENTIFIC SETTING : EXPLORING OPTIMAL REPRESENTATIONS FOR WORST - CASE DG

Original: To validate our theory, we studied optimal representations in a scientiﬁc setup that is as close to our
Revised: To validate our theory, we studied optimal representations in a scientiﬁc setup that is as close to our

Original: IDG framework as possible with log-loss (cid:96) .
Revised: IDG framework as possible with log-loss (cid:96) .

Original: 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.
Revised: In particular, we used the PACS dataset (Li et al., 2017) Base Ent CAD

Revised: R [ Y | Z ]

Revised: H

Original: To approximate the worst-case source predictor, we followed Dubois et al.
Revised: [ A | Z ] -5.1 ± 0.3 -0.4 ± 0.1 -0.7 ± 0. [Figure] -3.5 ± 0. [Figure]

Revised: [Figure]

Revised: ±

Revised: [Figure]

Revised: [Figure]

Revised: ± 0.

Revised: (a) Effect of different objectives

Revised: [Figure]

Revised: ApproxDA

Revised: [Figure]

Revised: Standard (c) Effect of augmentations 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.

Revised: To approximate the worst-case source predictor, we followed Dubois et al.

Original: (2020) by incorporating the wrongly labeled target data to the source domain.
Revised: (2020) by incorporating the wrongly labeled target data to the source domain.

Original: 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 classiﬁer h on every possible pair of (source, target); (iii) the negative target risk (log likelihood) for each h is evaluated.
Revised: 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 classiﬁer h on every possible pair of (source, target); (iii) the negative target risk (log likelihood) for each h is evaluated.

Original: We reported the log likelihood averaged over 5 seeds.
Revised: We reported the log likelihood averaged over 5 seeds.

Original: For more realistic scenarios (i.e. non-idealized average-case DG) see Appx.
Revised: For more realistic scenarios (i.e. non-idealized average-case DG) see Appx.

Original: E.2 which replicates the following results.
Revised: E.2 which replicates the following results.

Original: Do our domain bottlenecks improve worst-case DG?
Revised: Do our domain bottlenecks improve worst-case DG?

Original: In Fig.
Revised: In Fig.

Original: 3a, we compare IDG performance of representations trained with (Ent, CAD) and without (Base) domain bottlenecks.
Revised: 3a, we compare IDG performance of representations trained with (Ent, CAD) and without (Base) domain bottlenecks.

Original: We see that both bottlenecks signiﬁcantly improve the worst-case DG, and nearly achieve the source-domain performance (0 log likelihood).
Revised: We see that both bottlenecks signiﬁcantly improve the worst-case DG, and nearly achieve the source-domain performance (0 log likelihood).

Original: This shows the importance of support match (Thm. 2) and the effectiveness of our bottlenecks to enforce it.
Revised: This shows the importance of support match (Thm. 2) and the effectiveness of our bottlenecks to enforce it.

Original: In Appx.
Revised: In Appx.

Original: 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% ).
Revised: 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% ).

Original: What is the effect of λ ?
Revised: What is the effect of λ ?

Original: Fig.
Revised: Fig.

Original: 3b shows the effect of the bottleneck weight λ on the worst-case target and source performance.
Revised: 3b shows the effect of the bottleneck weight λ on the worst-case target and source performance.

Original: We see that increasing λ will decrease the DG gap.
Revised: We see that increasing λ will decrease the DG gap.

Original: As a result the target performance improves until λ ≈ 10 2 , where source performance starts to decrease.
Revised: As a result the target performance improves until λ ≈ 10 2 , where source performance starts to decrease.

Original: What if the representor has access to domain-covering augmentations instead of labels?
Revised: What if the representor has access to domain-agnostic augmentations instead of labels?

Original: In
Revised: In

Original: Sec.
Revised: Sec.

Original: 4.2, we provide a contrastive objective for using augmentations.
Revised: 4.2, we provide a contrastive objective for using augmentations.

Original: To show the effectiveness of the objective, we compared minimizing H
Revised: To show the effectiveness of the objective, we compared minimizing H

Original: [ A | Z ] using Eq.
Revised: [ A | Z ] using Eq.

Original: (7) to standard supervised risk minimization
Revised: (6) to standard supervised risk minimization

Original: R [ Y | Z ] . We ensured domain coverage by using supervised augmentations (Fig. 2b).
Revised: R [ Y | Z ] and used the domain-agnostic supervised augmentations (Fig. 2b).

Original: The 1 st and 2 nd row of Fig.
Revised: The 1 st and 2 nd row of Fig.

Original: 3a show that our objective performs similarly to direct label prediction.
Revised: 3a show that our objective performs similarly to direct label prediction.

Original: How important is the choice of augmentations?
Revised: How important is the choice of augmentations?

Original: Prop. 2 shows that domain-covering (DC) augmentations are sufﬁcient for achieving IDG, but it does not give necessary conditions.
Revised: Prop. 2 shows that domain-agnostic (DA) augmentations are sufﬁcient for achieving IDG, but it does not give necessary conditions.

Original: Here we investigate the effect of using our loss with different choices of augmentations.
Revised: Here we investigate the effect of using our loss with different choices of augmentations.

Original: Speciﬁcally, we used
Revised: Speciﬁcally, we used

Original: L CAD with ﬁve augmentations.
Revised: L CAD with ﬁve augmentations.

Original: The ﬁrst two are DC.
Revised: The ﬁrst two are DA.

Original: ‘Supervised’: augment inputs inside the label class across all domains as in Fig.
Revised: ‘Supervised’: augment inputs inside the label class across all domains as in Fig.

Original: 2b; ‘SingleDom’: augment inputs to same label samples from a ﬁxed domain.
Revised: 2b; ‘SingleDom’: augment inputs to same label samples from a ﬁxed domain.

Original: The second two are not DC.
Revised: The second two are not DA.

Original: ‘Standard’: standard SSL augmentations (Chen et al., 2020) as in Fig.
Revised: ‘Standard’: standard SSL augmentations (Chen et al., 2020) as in Fig.

Original: 2a; ‘IntraDom’: augment inputs to same label and same domain samples.
Revised: 2a; ‘IntraDom’: augment inputs to same label and same domain samples.

Original: Finally, we consider ‘ApproxDC’, which is approximately DC by augmenting 10% of the time with ‘Supervised‘ and 90% of the time with ‘IntraDom‘.
Revised: Finally, we consider ‘ApproxDA’, which is approximately DA by augmenting 10% of the time with ‘Supervised‘ and 90% of the time with ‘IntraDom‘.

Original: Fig.
Revised: Fig.

Original: 3c shows that the non-DC augmentations give terrible results compared to DC.
Revised: 3c shows that the non-DA augmentations give terrible results compared to DA.

Original: Interestingly, ‘ApproxDC’ also performs very well, which suggests that approximately DC augmentations might be sufﬁcient to learn optimal representations in practice.
Revised: Interestingly, ‘ApproxDA’ also performs very well, which suggests that approximately DA augmentations might be sufﬁcient to learn optimal representations in practice.

Original: What if the representor does not have access to target domains?
Revised: What if the representor does not have access to target domains?

Original: Prop. 1 shows that DG without access to target domains is generally impossible.
Revised: Prop. 1 shows that DG without access to target domains is generally impossible.

Original: We empirically veriﬁed this by excluding a predeﬁned target d t domain from the representor’s training set, i.e., L CAD is optimized on 3 of thedomains.
Revised: We empirically veriﬁed this by excluding a predeﬁned target d t domain from the representor’s training set, i.e., L CAD is optimized on 3 of thedomains.

Original: The learner then trains a predictor h on each source.
Revised: The learner then trains a predictor h on each source.

Original: We ﬁnally evaluate each h on the target domain d t , and average over choices of d t .
Revised: We ﬁnally evaluate each h on the target domain d t , and average over choices of d t .

Original: The resulting worst-case log likelihood was − 4 .
Revised: The resulting worst-case log likelihood was − 4 .

Original: 2 ± 0 .
Revised: 2 ± 0 .

Original: 2 , which is signiﬁcantly worse than when the representor had access to all domains ( − 0 .
Revised: 2 , which is signiﬁcantly worse than when the representor had access to all domains ( − 0 .

Original: 8 ± 0 . 2 ).
Revised: 8 ± 0 . 2 ).

Original: 6. A PPROXIMATING OPTIMAL REPRESENTATIONS BY EXPLOITING PRETRAINED SSL
Revised: 6. A PPROXIMATING OPTIMAL REPRESENTATIONS BY EXPLOITING PRETRAINED SSL

Original: 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 covering augmentations A .
Revised: 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 .

Original: This is nearly how CLIP was pretrained (SSL with 400M image-text pairs) except it did not include a domain bottleneck.
Revised: This is nearly how CLIP was pretrained (SSL with 400M image-text pairs) except it did not include a domain bottleneck.

Original: In this section, we investigate how to take advantage of CLIP to approximate optimal representations for IDG.
Revised: In this section, we investigate how to take advantage of CLIP to approximate optimal representations for IDG.

Original: We did so in two simple steps.
Revised: We did so in two simple steps.

Original: First, we froze the pretrained CLIP and added a multi-layer perceptron (MLP) that could effectively ﬁnetune CLIP’s representations.
Revised: First, we froze the pretrained CLIP and added a multi-layer perceptron (MLP) that could effectively ﬁnetune CLIP’s representations.

Original: Then, we trained the MLP by minimizing our CAD bottleneck and R [ Y | Z ] on the available data.
Revised: Then, we trained the MLP by minimizing our CAD bottleneck and R [ Y | Z ] on the available data.

Original: In all experiments, we used the standard DomainBed benchmark (with non-MNIST datasets) and protocol (Gulrajani & Lopez-Paz, 2021).
Revised: In all experiments, we used the standard DomainBed benchmark (with non-MNIST datasets) and protocol (Gulrajani & Lopez-Paz, 2021).

Original: In particular, we left out a target domain for evaluation and used the union of other domains for training both the encoder and the classiﬁer.
Revised: In particular, we left out a target domain for evaluation and used the union of other domains for training both the encoder and the classiﬁer.

Original: Contrary to our scientiﬁc setting, the representor does not get access to the target domain.
Revised: Contrary to our scientiﬁc setting, the representor does not get access to the target domain.

Original: All our representations were evaluated by ﬁtting a linear classiﬁer on source domains with source validation selection.
Revised: All our representations were evaluated by ﬁtting a linear classiﬁer on source domains with source validation selection.

Original: 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.
Revised: 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.

Original: Due to space limit, we only included as baselines ‘ERM’ and ‘DomainBed SOTA’ which for each dataset is the best result over all baselines.
Revised: Due to space limit, we only included as baselines ‘ERM’ and ‘DomainBed SOTA’ which for each dataset is the best result over all baselines.

Original: The extended results and baselines are in Table 4.
Revised: The extended results and baselines are in Table 4.

Original: Details in Appx.
Revised: Details in Appx.

Original: D.3.
Revised: D.3.

Original: We investigated two pretrained CLIP models with different number of parameters.
Revised: We investigated two pretrained CLIP models with different number of parameters.

Original: The larger ViT-B/ denoted ‘CLIP L’ and the smaller ResNet-50 denoted ‘CLIP S’.
Revised: The larger ViT-B/32 denoted ‘CLIP L’ and the smaller ResNet-50 denoted ‘CLIP S’.

Original: [Table]
Revised: [Table]

Original: Can we approximate optimal representations by exploiting pretrained CLIP?
Revised: Can we approximate optimal representations by exploiting pretrained CLIP?

Original: The last row in
Revised: The last row in

Original: Table 1 shows that ﬁnetuning a large pretrained CLIP model with our CAD achieves SOTA on nearly all DomainBed benchmarks by a very large margin (see 2 nd row).
Revised: Table 1 shows that ﬁnetuning a large pretrained CLIP model with our CAD achieves SOTA on nearly all DomainBed benchmarks by a very large margin (see 2 nd row).

Original: Note that the poor performance on
Revised: Note that the poor performance on

Original: TerraIncognita is likely because CLIP’s dataset did not cover such images (camera traps monitoring animals).
Revised: TerraIncognita is likely because CLIP’s dataset did not cover such images (camera traps monitoring animals).

Original: In Appx.
Revised: In Appx.

Original: E.2, we estimated the non-idealized DG performance of optimal representations on PACS (with access to all-domain labeled data) to be 96 .
Revised: E.2, we estimated the non-idealized DG performance of optimal representations on PACS (with access to all-domain labeled data) to be 96 .

Original: 7% , which is only 2% higher than CLIP L + CAD.
Revised: 7% , which is only 2% higher than CLIP L + CAD.

Original: This suggests that our simple SSL encoder might already be close to optimal.
Revised: This suggests that our simple SSL encoder might already be close to optimal.

Original: Are gains due to the architectural differences?
Revised: Are gains due to the architectural differences?

Original: DomainBed’s baselines ﬁnetuned an ImageNet(Deng et al., 2009) pretrained ResNet-50.
Revised: DomainBed’s baselines ﬁnetuned an ImageNet pretrained ResNet-50.

Original: In contrast, CLIP L pretrained a larger ViT. To decouple gains due to our objective from architectural gains, we evaluated ResNet-50 pretrained.
Revised: 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

Original: Table 1 shows that CLIP S  still  outperforms DomainBed baselines.
Revised: CLIP S + CAD still signiﬁcantly outperforms DomainBed baselines.

Original: Our theory does not constrain the encoder and so we expect larger encoders to be better. Table 1 shows that CLIP L indeed outperforms CLIP S.
Revised: Note that our theory does not constrain the encoder and so we expect larger encoders to be better as seen in Table 1.

Original: What is the effect of domain bottlenecks?
Revised: What is the effect of domain bottlenecks?

Original: In the last ﬁve rows of Table 1, we investigated the effect of ﬁnetuning with our CAD bottleneck.
Revised: In the last ﬁve rows of Table 1, we investigated the effect of ﬁnetuning with our CAD bottleneck.

Original: We see that for both CLIP L and CLIP S, it improves results by around 1 ∼ 2% .
Revised: We see that for both CLIP L and CLIP S, it improves results by around 1 ∼ 2% .

Original: These gains are due to the bottleneck, rather than due to the additional MLP trained on source data as seen by ‘CLIP S + Base’.
Revised: These gains are due to the bottleneck, rather than ﬁnetuning on source data as seen by ‘CLIP S + Base’.

Revised: We believe the gains could potentially be much larger if CLIP was trained end-to-end with our bottleneck.

Original: Note thatthe raw CLIP S already signiﬁcantly outperforms baselines.
Revised: Note that raw CLIP S already signiﬁcantly outperforms baselines.

Original: We hypothesize that this could be because SGD training of neural networks favors support match, e.g., by minimizing I[ X ; Z ] as suggested by Shwartz-Ziv & Tishby (2017).
Revised: We hypothesize that this is because SGD acts as an information bottleneck that naturally favors support match (Shwartz-Ziv & Tishby, 2017).

Original: Which pretrained SSL model to use?
Revised: Which pretrained SSL model to use?

Original: Our theory suggests that we can exploit pretrained SSL models as long as their augmentations are domain-covering and their training set covers desired domains.
Revised: 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.

Original: We investigated the effect ofadapting SSL models that do not satisfy those properties by ﬁnetuning DINO (Caron et al., 2021), the current SOTA on SSL ImageNet.
Revised: We investigated adaption of SSL models that do not satisfy those properties by ﬁnetuning DINO (Caron et al., 2021), the current SOTA on SSL ImageNet.

Original: DINO only pretraiendon ImageNet using standard augmentations.
Revised: DINO is pretraiend using standard augmentations.

Original: As a result, Table 1 shows that the ﬁnetuned DINO+CAD signiﬁcantly underperforms compared to CLIP S and DomainBed baselines.
Revised: As a result, Table 1 shows that the ﬁnetuned DINO+CAD signiﬁcantly underperforms compared to CLIP S and DomainBed baselines.

Revised: This supports our hypothesis that CLIP is much more robust than other SSL methods due to its domain-agnostic augmentations.

Original: 6. T OWARDS GENERIC ROBUST REPRESENTATIONS WITH SSL
Revised: 6. T OWARDS GENERIC ROBUST REPRESENTATIONS WITH SSL

Original: In the previous section, we ﬁnetuned CLIP in a task speciﬁc fashion by optimizing R [ Y | Z ] and our
Revised: In the previous section, we ﬁnetuned CLIP in a task speciﬁc fashion by optimizing R [ Y | Z ] and our

Original: CAD bottleneck.
Revised: CAD bottleneck.

Original: To get generic (task agnostic) robust representations, one should instead directly use our objectives on a sufﬁciently large dataset with image-text augmentations.
Revised: To get generic (task agnostic) robust representations, one should instead directly use our objectives on a sufﬁciently large dataset with image-text augmentations.

Original: Unfortunately, we cannot fully train CLIP with our bottlenecks as we do not have access to CLIP’s original dataset and sufﬁcient compute.
Revised: Unfortunately, we cannot fully train CLIP with our bottlenecks as we do not have access to CLIP’s original dataset and sufﬁcient compute.

Original: In this section, we aim to emulate such training of generic robust representations.
Revised: In this section, we aim to emulate such training of generic robust representations.

Original: To do so we used LAION-400M (LAION, 2021) which is a public dataset that contains 400M web-crawled image-text pairs.
Revised: To do so we used LAION-400M (Schuhmann et al., 2021) that is a public dataset that contains 400M web-crawled image-text pairs.

Original: Due to our computational budget, we again froze the pretrained CLIP
Revised: Due to our computational budget, we again froze the pretrained CLIP

Original: L and only ﬁnetuned an additional MLP with our L Ent .
Revised: L and only ﬁnetuned an additional MLP with our L Ent .

Original: We used L Ent as it only requires access to paired image X and text A but no prior information about domain D .
Revised: We used L Ent as it only requires access to paired image X and text A but no prior information about domain D .

Original: As in CLIP’s paper, we evaluated the learned representation Z in Taori et al.’s
Revised: As in CLIP’s paper, we evaluated the learned representation Z in Taori et al.’s

Original: (2020) realistic setting, where a linear classiﬁer h from Z is trained on ImageNet and tested on 7 natural distribution shift datasets.
Revised: (2020) realistic setting, where a linear classiﬁer h from Z is trained on ImageNet and tested on 7 natural distribution shift datasets.

Original: Details in Appx.
Revised: Details in Appx.

Original: D.4.
Revised: D.4.

Original: Would training CLIP with a bottleneck have improved its robustness?
Revised: Would training CLIP with a bottleneck have improved its robustness?

Original: As shown in the lastrows of Table 2, ﬁnetuning CLIP L on LAION with L Ent (LAION + Ent) outperforms ﬁnetuning without bottleneck (LAION + Base) on all 7 distribution shift datasets.
Revised: As shown in the lastrows of Table 2, ﬁnetuning CLIP L on LAION with L Ent (LAION + Ent) outperforms ﬁnetuning without bottleneck (LAION + Base) on all 7 distribution shift datasets.

Original: This suggests that directly training CLIP with our Ent bottleneck would improve the robustness of learned representations.
Revised: This suggests that directly training CLIP with our Ent bottleneck would improve the robustness of learned representations.

Original: We hypothesize that the gains could be larger if SSL models trained L Ent end-to-end.
Revised: We hypothesize that the gains could be larger if SSL models trained L Ent end-to-end.

Original: In Appx.
Revised: In Appx.

Original: E.4, we show similar results on DomainBed.
Revised: E.4, we show similar results on DomainBed.

Original: 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.
Revised: 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.

Original: [Table]
Revised: [Table]

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: 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.
Revised: 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.

Original: Similar to previous work, our theory strongly implies the need for target information when learning representations for domain generalization.
Revised: Similar to previous work, our theory strongly implies the need for target information when learning representations for domain generalization.

Original: Nevertheless, we identiﬁed a domain-covering property of data augmentations that make it possible to learn optimal representations from unlabelled data.
Revised: Nevertheless, we identiﬁed a domain-agnostic property of data augmentations that make it possible to learn optimal representations from unlabelled data.

Original: Thus, we showed that it is possible to learn robust representations using only large sources of inputs X and augmentations A .
Revised: Thus, we showed that it is possible to learn robust representations using only large sources of inputs X and augmentations A .

Original: There are caveats that need to be addressed in future work.
Revised: There are caveats that need to be addressed in future work.

Original: First, we studied an idealized DG, which assumes access to the population distributions.
Revised: First, we studied an idealized DG, which assumes access to the population distributions.

Original: This gives insights into the challenges that are speciﬁc to DG, rather than ﬁnite sample challenges faced throughout ML.
Revised: This gives insights into the challenges that are speciﬁc to DG, rather than ﬁnite sample challenges faced throughout ML.

Original: Second, we considered risk minimizers from an unconstrained hypothesis class.
Revised: Second, we considered risk minimizers from an unconstrained hypothesis class.

Original: The support constraint can likely be weakened, if the hypothesis class is constrained.
Revised: The support constraint can likely be weakened, if the hypothesis class is constrained.

Original: Finally, we focus only on optimal representations, but it would be interesting to characterize approximately optimal representations.
Revised: Finally, we focus only on optimal representations, but it would be interesting to characterize approximately optimal representations.

Original: 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.
Revised: 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.

Original: Reproducibility Statement For our theoretical results, we include formal assumptions, statements, and proofs in Appxs.
Revised: Reproducibility Statement For our theoretical results, we include formal assumptions, statements, and proofs in Appxs.

Original: A and B. We include the detailed derivations of our algorithms in Appx.
Revised: A and B. We include the detailed derivations of our algorithms in Appx.

Original: C. For our experiments, we include experimental details for reproducing our results in Appx.
Revised: C. For our experiments, we include experimental details for reproducing our results in Appx.

Original: D and will release our code upon acceptance.
Revised: D and will release our code upon acceptance.

Original: A P RELIMINARIES A. N OTATION
Revised: A P RELIMINARIES A. N OTATION

Original: For the most part, we will assume that all spaces are discrete probability spaces.
Revised: For the most part, we will assume that all spaces are discrete probability spaces.

Original: A full list of assumptions is found at Appx.
Revised: A full list of assumptions is found at Appx.

Original: A.3.
Revised: A.3.

Original: General The image of a set A ⊆ X under a function f : X → Y is denoted f → ( A ) =
Revised: General The image of a set A ⊆ X under a function f : X → Y is denoted f → ( A ) =

Original: { f ( x ) | x ∈ A } .
Revised: { f ( x ) | x ∈ A } .

Original: The pre-image is denoted f ← ( B ) = { x ∈ X | f ( x ) ∈ B } for B ⊆ Y .
Revised: The pre-image is denoted f ← ( B ) = { x ∈ X | f ( x ) ∈ B } for B ⊆ Y .

Original: 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.
Revised: 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.

Original: The probability mass function (pmf) of a random variable X is denoted as p X .
Revised: The probability mass function (pmf) of a random variable X is denoted as p X .

Original: We use capital P instead of p to denote the measure under p .
Revised: We use capital P instead of p to denote the measure under p .

Original: The support supp( p X ) of a discrete distribution is the set of all points x ∈ X with positive probability, i.e., supp( p X ) =
Revised: The support supp( p X ) of a discrete distribution is the set of all points x ∈ X with positive probability, i.e., supp( p X ) =

Original: { x ∈ X | p X ( x ) > 0 } .
Revised: { x ∈ X | p X ( x ) > 0 } .

Original: The space of all probability distributions on X is denoted
Revised: The space of all probability distributions on X is denoted

Original: P ( X ) = (cid:8) p X | p X ( x ) ≥ 0 and (cid:80) x ∈X p X ( x ) = 1 (cid:9) .
Revised: P ( X ) = (cid:8) p X | p X ( x ) ≥ 0 and (cid:80) x ∈X p X ( x ) = 1 (cid:9) .

Original: When it is necessary to be explicit, we will denote ‘ X is distributed as p X ’ using the notation
Revised: When it is necessary to be explicit, we will denote ‘ X is distributed as p X ’ using the notation

Original: X d ∼ p X .
Revised: X d ∼ p X .

Original: Expectations are written as: E p X
Revised: Expectations are written as: E p X

Original: [ f ( X )] , independence of two r.v. as ·⊥⊥· , conditional independence as ·⊥⊥ · | · .
Revised: [ f ( X )] , independence of two r.v. as ·⊥⊥· , conditional independence as ·⊥⊥ · | · .

Original: For jointly distributed random variables ( X, Y ) taking value in (t.v.i.)
Revised: For jointly distributed random variables ( X, Y ) taking value in (t.v.i.)

Original: X × Y , the conditional distribution is denoted as p Y | X : Y × X
Revised: X × Y , the conditional distribution is denoted as p Y | X : Y × X

Original: → [0 , 1] .
Revised: → [0 , 1] .

Original: For convenience, let p Y | x = p Y | X ( · | x ) be the conditional distribution of Y given x .
Revised: For convenience, let p Y | x = p Y | X ( · | x ) be the conditional distribution of Y given x .

Original: All random variables are independently distributed, unless an explicit joint distribution or coupling is given.
Revised: All random variables are independently distributed, unless an explicit joint distribution or coupling is given.

Original: A. D EFINITIONS
Revised: A. D EFINITIONS

Original: We are interested in prediction problems with domain shift.
Revised: We are interested in prediction problems with domain shift.

Original: There are three random variables: the target domain D t , the input X , the label Y .
Revised: There are three random variables: the target domain D t , the input X , the label Y .

Original: They have the following joint distribution:
Revised: They have the following joint distribution:

Original: [Equation]
Revised: [Equation]

Original: where we drop the arguments of the probability densities for clarity.
Revised: where we drop the arguments of the probability densities for clarity.

Original: We make a variety of convenience assumptions on these random variables (Assumption 6).
Revised: We make a variety of convenience assumptions on these random variables (Assumption 6).

Original: Crucially, we will be making the Bayes invariance assumption on p D t ,X,Y that can be thought of as a generalized covariate shift assumption (Assumption 4).
Revised: Crucially, we will be making the Bayes invariance assumption on p D t ,X,Y that can be thought of as a generalized covariate shift assumption (Assumption 4).

Original: We will be studying the effect of changing the representation of the data.
Revised: We will be studying the effect of changing the representation of the data.

Original: This is done by “encoding”
Revised: This is done by “encoding”

Original: X into a representation Z using a conditional distribution p Z | X .
Revised: X into a representation Z using a conditional distribution p Z | X .

Original: Deﬁnition 1 (Encoder) .
Revised: Deﬁnition 1 (Encoder) .

Original: An encoder is a conditional distribution p Z | X : Z × X → [0 , 1] from the input space
Revised: An encoder is a conditional distribution p Z | X : Z × X → [0 , 1] from the input space

Original: X to the representation space Z .
Revised: X to the representation space Z .

Original: The data together with the representation has the following joint:
Revised: The data together with the representation has the following joint:

Original: [Equation]
Revised: [Equation]

Original: The key thing to notice here is that Z is conditionally independent of Y, D t given X .
Revised: The key thing to notice here is that Z is conditionally independent of Y, D t given X .

Original: In particular, the same encoder is used across all domains.
Revised: In particular, the same encoder is used across all domains.

Original: A.2. R ISK MINIMIZATION
Revised: A.2. R ISK MINIMIZATION

Original: Our ultimate goal is to predict Y from the representation Z of X in a manner that is robust to changes in the domain.
Revised: Our ultimate goal is to predict Y from the representation Z of X in a manner that is robust to changes in the domain.

Original: We formalize this in the standard way by making predictions γ ∈ Γ in a space of predictions or actions.
Revised: We formalize this in the standard way by making predictions γ ∈ Γ in a space of predictions or actions.

Original: For example the prediction space may be the set of all possible labels Γ =
Revised: For example the prediction space may be the set of all possible labels Γ =

Original: Y , in which case we would be predicting deterministic labels.
Revised: Y , in which case we would be predicting deterministic labels.

Original: 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 ) .
Revised: 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 ) .

Original: A predictor is a function mapping inputs to predictions, i.e., f : X → Γ , or representations to predictions, i.e., h : Z → Γ .
Revised: A predictor is a function mapping inputs to predictions, i.e., f : X → Γ , or representations to predictions, i.e., h : Z → Γ .

Original: For example, f may be a neural network that takes as input a sample x and outputs a vector of logits that parameterize a softmax distribution over ﬁnitely many labels.
Revised: For example, f may be a neural network that takes as input a sample x and outputs a vector of logits that parameterize a softmax distribution over ﬁnitely many labels.

Original: We select predictors according to the risk deﬁned via a loss function (cid:96) :
Revised: We select predictors according to the risk deﬁned via a loss function (cid:96) :

Original: Y × Γ → R ≥ 0 ∪ {∞} :
Revised: Y × Γ → R ≥ 0 ∪ {∞} :

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: In particular, we are interested in the Bayes (minimum) risk over all predictors:
Revised: In particular, we are interested in the Bayes (minimum) risk over all predictors:

Original: [Equation]
Revised: [Equation]

Original: We denote the set of all optimal predictors from X as
Revised: We denote the set of all optimal predictors from X as

Original: [Equation]
Revised: [Equation]

Original: Similarly, we deﬁne the risk R h
Revised: Similarly, we deﬁne the risk R h

Original: [ Y | Z ] , the Bayes risk R [ Y | Z ] , and the set of optimal predictors
Revised: [ Y | Z ] , the Bayes risk R [ Y | Z ] , and the set of optimal predictors

Original: [Equation]
Revised: [Equation]

Original: from Z , all of which vary as a function of the encoder p Z | X .
Revised: from Z , all of which vary as a function of the encoder p Z | X .

Original: Note, in the main body of the paper, we omitted the subscript Z from H ∗ Z for clarity, but we will keep it in the Appendices.
Revised: Note, in the main body of the paper, we omitted the subscript Z from H ∗ Z for clarity, but we will keep it in the Appendices.

Original: We assume that together our loss and prediction space always admit optima (Item 2 of Assumption 2), and thus
Revised: We assume that together our loss and prediction space always admit optima (Item 2 of Assumption 2), and thus

Original: F ∗ , H ∗ Z are always non-empty.
Revised: F ∗ , H ∗ Z are always non-empty.

Original: We will be assuming that the risk admits unique optimal prediction when predicting from X (Itemof Assumption 2).
Revised: We will be assuming that the risk admits unique optimal prediction when predicting from X (Itemof Assumption 2).

Original: Thus it makes sense to deﬁne the following:
Revised: Thus it makes sense to deﬁne the following:

Original: Deﬁnition 2 (The Bayes predictor) .
Revised: Deﬁnition 2 (The Bayes predictor) .

Original: The Bayes predictor f ∗ : X → Γ is the unique predictor that is optimal for all x ∈ X :
Revised: The Bayes predictor f ∗ : X → Γ is the unique predictor that is optimal for all x ∈ X :

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Deﬁnition 3 (The Bayes image) .
Revised: Deﬁnition 3 (The Bayes image) .

Original: The image of all the inputs under the Bayes predictor will be denoted as Γ ∗ = f ∗ → ( X ) and called the Bayes image .
Revised: The image of all the inputs under the Bayes predictor will be denoted as Γ ∗ = f ∗ → ( X ) and called the Bayes image .

Original: Note that F ∗ becomes a singleton { f ∗ } , but it is not necessarily the case for H ∗ Z since we will not be making any uniqueness assumption on optimal prediction from Z .
Revised: Note that F ∗ becomes a singleton { f ∗ } , but it is not necessarily the case for H ∗ Z since we will not be making any uniqueness assumption on optimal prediction from Z .

Original: A.2. D OMAIN GENERALIZATION
Revised: A.2. D OMAIN GENERALIZATION

Original: We are interested in controlling the risk in a domain generalization setting, and so we deﬁne the domain-conditional risk ,
Revised: We are interested in controlling the risk in a domain generalization setting, and so we deﬁne the domain-conditional risk ,

Original: [Equation]
Revised: [Equation]

Original: R d [ Y | X ] , F ∗ d are deﬁned as Eqs.
Revised: R d [ Y | X ] , F ∗ d are deﬁned as Eqs.

Original: (13) and (14), respectively, but with respect to R df .
Revised: (12) and (13), respectively, but with respect to R df .

Original: Similarly, deﬁne the Bayes image for domain d as
Revised: Similarly, deﬁne the Bayes image for domain d as

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: We also deﬁne domain-conditional quantities for prediction from a representation Z .
Revised: We also deﬁne domain-conditional quantities for prediction from a representation Z .

Original: The most important term which we will be investigating is an idealization of the domain generalization worstcase risk.
Revised: The most important term which we will be investigating is an idealization of the domain generalization worstcase risk.

Original: Deﬁnition 4 (IDG risk) .
Revised: Deﬁnition 4 (IDG risk) .

Original: Given an encoder p Z | X and a distribution p D t ,D s over a target domain D t and source domain D s , the idealized domain generalization worst-case risk, IDG risk for short, is the expected worst-case target risk taken over source minimizers, i.e.,
Revised: Given an encoder p Z | X and a distribution p D t ,D s over a target domain D t and source domain D s , the idealized domain generalization worst-case risk, IDG risk for short, is the expected worst-case target risk taken over source minimizers, i.e.,

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Note that the IDG risk is well-deﬁned because H ∗ Z,D s is non-empty by Assumption 2.
Revised: Note that the IDG risk is well-deﬁned because H ∗ Z,D s is non-empty by Assumption 2.

Original: The desired optimal representations, are then those that minimize the IDG risk.
Revised: The desired optimal representations, are then those that minimize the IDG risk.

Original: Deﬁnition 5 (Optimal representations for IDG) .
Revised: Deﬁnition 5 (Optimal representations for IDG) .

Original: An encoder p Z ∗ | X is optimal for idealized domain generalization if and only if it minimizes the IDG risk, i.e.,
Revised: An encoder p Z ∗ | X is optimal for idealized domain generalization if and only if it minimizes the IDG risk, i.e.,

Original: [Equation]
Revised: [Equation]

Original: A. A SSUMPTIONS
Revised: A. A SSUMPTIONS

Original: We make a the following assumptions throughout the paper.
Revised: We make a the following assumptions throughout the paper.

Original: All these assumptions should hold for practical settings.
Revised: All these assumptions should hold for practical settings.

Original: Assumption 1 (Convenience: discrete probability spaces) .
Revised: Assumption 1 (Convenience: discrete probability spaces) .

Original: All data spaces ( D , X , Y , Z , A ) are discrete spaces.
Revised: All data spaces ( D , X , Y , Z , A ) are discrete spaces.

Original: Because the distributions of X, Y, D are ﬁxed, we assume for convenience that supp( p X ) =
Revised: Because the distributions of X, Y, D are ﬁxed, we assume for convenience that supp( p X ) =

Original: X , supp( p Y )
Revised: X , supp( p Y )

Original: =
Revised: =

Original: Y , and supp( p D t ) =
Revised: Y , and supp( p D t ) =

Original: D .
Revised: D .

Original: Assumption 1 is a convenience assumption to avoid measure theory for the sake of clarity.
Revised: Assumption 1 is a convenience assumption to avoid measure theory for the sake of clarity.

Original: It always holds in practice due to ﬁniteness of computers, i.e., all spaces will be ﬁnite but arbitrarily large.
Revised: It always holds in practice due to ﬁniteness of computers, i.e., all spaces will be ﬁnite but arbitrarily large.

Original: We believe that our claims can nevertheless be generalized to typical continuous spaces with some minor technical assumptions.
Revised: We believe that our claims can nevertheless be generalized to typical continuous spaces with some minor technical assumptions.

Original: Assumption 2 (Properness of our losses) .
Revised: Assumption 2 (Properness of our losses) .

Original: We assume that our risk always admits optimal predictions:
Revised: We assume that our risk always admits optimal predictions:

Original: |
Revised: |

Original: [Equation]

Original: |

Original: [Equation]
Revised: [Equation]

Original: For all p Υ ∈ P ( Y ) , there exists γ ∗ ∈ Γ , such that
Revised: For all p Υ ∈ P ( Y ) , there exists γ ∗ ∈ Γ , such that

Original: [Equation]
Revised: [Equation]

Original: For all x ∈ X , there exist γ ∗ ∈ Γ , such that
Revised: For all x ∈ X , there exist γ ∗ ∈ Γ , such that

Original: [Equation]
Revised: [Equation]

Original: Note that for log-loss (cid:96) ( y, γ ) =
Revised: Note that for log-loss (cid:96) ( y, γ ) =

Original: − log γ ( y ) and ﬁnite Y , these assumptions are satisﬁed if Γ = P ( Y ) where the optimal prediction for Item 3 is γ ∗ =
Revised: − log γ ( y ) and ﬁnite Y , these assumptions are satisﬁed if Γ = P ( Y ) where the optimal prediction for Item 3 is γ ∗ =

Original: p Y | x by strict properness (Gneiting & Raftery, 2007).
Revised: p Y | x by strict properness (Gneiting & Raftery, 2007).

Original: If we consider the 0-1 loss (reverse accuracy) (cid:96) ( y, γ ) = 1 − 1 [ y = γ ] with Γ =
Revised: If we consider the 0-1 loss (reverse accuracy) (cid:96) ( y, γ ) = 1 − 1 [ y = γ ] with Γ =

Original: Y and a ﬁnite label space where the optimal prediction for Item 3 is γ ∗ = arg max y ∈Y p Y | x ( y ) , this assumption is mostly satisﬁed, except we assume that p
Revised: Y and a ﬁnite label space where the optimal prediction for Item 3 is γ ∗ = arg max y ∈Y p Y | x ( y ) , this assumption is mostly satisﬁed, except we assume that p

Original: [Equation]
Revised: [Equation]

Original: Assumption 2 serves two purposes: Item 2 ensures that for any representation the optimal predictors from Z exists such that the IDG risk is well-deﬁned as in Def. 5; Item 3 ensures a unique Bayes predictor from X , which simpliﬁes the analysis and is satisﬁed by common losses as described above.
Revised: Assumption 2 serves two purposes: Item 2 ensures that for any representation the optimal predictors from Z exists such that the IDG risk is well-deﬁned as in Def. 5; Item 3 ensures a unique Bayes predictor from X , which simpliﬁes the analysis and is satisﬁed by common losses as described above.

Original: Assumption 3 (Cardinalities) .
Revised: Assumption 3 (Cardinalities) .

Original: We assume that
Revised: We assume that

Original: [Equation]
Revised: [Equation]

Original: Assumption 3 is very weak and ensures that optimal representations always exists (Prop. 3).
Revised: Assumption 3 is very weak and ensures that optimal representations always exists (Prop. 3).

Original: Assumption 4 (Generalized covariate shift) .
Revised: Assumption 4 (Generalized covariate shift) .

Original: The Bayes predictor is optimal for all domains.
Revised: The Bayes predictor is optimal for all domains.

Original: I.e., for all ( x, d ) ∈ supp( p X,D t ) , γ ∈ Γ such that γ (cid:54) =
Revised: I.e., for all ( x, d ) ∈ supp( p X,D t ) , γ ∈ Γ such that γ (cid:54) =

Original: f ∗ ( x ) , we have
Revised: f ∗ ( x ) , we have

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: For example, in the case of strictly proper scoring rules, e.g. log loss, covariate shift p Y | X,D = p Y
Revised: For example, in the case of strictly proper scoring rules, e.g. log loss, covariate shift p Y | X,D = p Y

Original: | X is equivalent to the invariance of the Bayes predictor.
Revised: | X is equivalent to the invariance of the Bayes predictor.

Original: For the 0-1 loss, this is guaranteed by invariance of the most likely label.
Revised: For the 0-1 loss, this is guaranteed by invariance of the most likely label.

Original: For MSE it is guaranteed by the invariance of the expected label.
Revised: For MSE it is guaranteed by the invariance of the expected label.

Original: In the latter two cases, Assumption 4 is less stringent than the typical covariate shift assumption.
Revised: In the latter two cases, Assumption 4 is less stringent than the typical covariate shift assumption.

Original: Assumption 4 is the core assumption for our theoretical results.
Revised: Assumption 4 is the core assumption for our theoretical results.

Original: It ensures that source and target domains are related in a useful way that can be utilized by the representation.
Revised: It ensures that source and target domains are related in a useful way that can be utilized by the representation.

Original: Assumption 5 (Constant Bayes image) .
Revised: Assumption 5 (Constant Bayes image) .

Original: The Bayes image is invariant across domains, i.e., for all d ∈ D , Γ
Revised: The Bayes image is invariant across domains, i.e., for all d ∈ D , Γ

Original: [Equation]
Revised: [Equation]

Original: For the case of 0-1 loss, this simply means that the label set for all domains is the same, which is trivial.
Revised: For the case of 0-1 loss, this simply means that the label set for all domains is the same, which is trivial.

Original: For log-loss, this means that the set of possible conditional distributions Γ ∗ d = { p Y | x | x ∈ supp( p X | d ) } is the same across domains.
Revised: For log-loss, this means that the set of possible conditional distributions Γ ∗ d = { p Y | x | x ∈ supp( p X | d ) } is the same across domains.

Original: Assumption 5 is crucial to be able to learn.
Revised: Assumption 5 is crucial to be able to learn.

Original: 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.
Revised: 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.

Original: Assumption 5 is also necessary to guarantee the existence of optimal representations as in Prop. 3.
Revised: Assumption 5 is also necessary to guarantee the existence of optimal representations as in Prop. 3.

Original: Assumption 6 (Domain joint) .
Revised: Assumption 6 (Domain joint) .

Original: p D t ,D s is any distribution such that supp( p D t ,D s ) =
Revised: p D t ,D s is any distribution such that supp( p D t ,D s ) =

Original: D × D .
Revised: D × D .

Original: In a simpliﬁed scenario, one could deﬁne the source D s and target D t as i.i.d.
Revised: In a simpliﬁed scenario, one could deﬁne the source D s and target D t as i.i.d.

Original: r.v.
Revised: r.v.

Original: from p D t , where p D t ,D s =
Revised: from p D t , where p D t ,D s =

Original: p D t · p D s = p D t · p D t and Assumption 6 is trivially satisﬁed.
Revised: p D t · p D s = p D t · p D t and Assumption 6 is trivially satisﬁed.

Original: B P ROOFS B. L EMMAS FOR GENERAL LOSSES
Revised: B P ROOFS B. L EMMAS FOR GENERAL LOSSES

Original: An important result that we will be using is the generalized data processing inequality of Bayes risk (Xu & Raginsky, 2020; Dubois et al., 2021).
Revised: An important result that we will be using is the generalized data processing inequality of Bayes risk (Xu & Raginsky, 2020; Dubois et al., 2021).

Original: We include it here for completeness.
Revised: We include it here for completeness.

Original: Lemma 1 (Generalized DPI (Xu & Raginsky, 2020; Dubois et al., 2021)) .
Revised: Lemma 1 (Generalized DPI (Xu & Raginsky, 2020; Dubois et al., 2021)) .

Original: Let Z − X − Y be a
Revised: Let Z − X − Y be a

Original: Markov chain of random variables.
Revised: Markov chain of random variables.

Original: For any loss function (cid:96) ,
Revised: For any loss function (cid:96) ,

Original: [Equation]
Revised: [Equation]

Original: For the case of strictly proper losses (Assumption 2) we can go one step further.
Revised: For the case of strictly proper losses (Assumption 2) we can go one step further.

Original: Lemma 2.
Revised: Lemma 2.

Original: Let Z − X − Y be a Markov chain of random variables.
Revised: Let Z − X − Y be a Markov chain of random variables.

Original: Then, under Assumptionsand 2 we have that
Revised: Then, under Assumptionsand 2 we have that

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Proof.
Revised: Proof.

Original: Suppose that for all h ∗ ∈ H ∗ Z we have h ∗ ( z ) =
Revised: Suppose that for all h ∗ ∈ H ∗ Z we have h ∗ ( z ) =

Original: f ∗ ( x ) on the support of p X,Z .
Revised: f ∗ ( x ) on the support of p X,Z .

Original: Then,
Revised: Then,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original: Now suppose there exists a h ∗ ∈ H ∗ Z and a pair ( x (cid:48) , z (cid:48) )
Revised: Now suppose there exists a h ∗ ∈ H ∗ Z and a pair ( x (cid:48) , z (cid:48) )

Original: ∈ supp( p X,Z ) such that h ∗ ( z (cid:48) ) (cid:54) =
Revised: ∈ supp( p X,Z ) such that h ∗ ( z (cid:48) ) (cid:54) =

Original: f ∗ ( x (cid:48) ) .
Revised: f ∗ ( x (cid:48) ) .

Original: Then
Revised: Then

Original: R [ Y | Z ] (33) =
Revised: R [ Y | Z ] (32) =

Original: E p X,Z p Y | X
Revised: E p X,Z p Y | X

Original: [ (cid:96) ( Y, h ∗ ( Z ))] (34) =
Revised: [ (cid:96) ( Y, h ∗ ( Z ))] (33) =

Original: p X,Z ( x (cid:48) , z (cid:48) )
Revised: p X,Z ( x (cid:48) , z (cid:48) )

Original: E p Y | x (cid:48)
Revised: E p Y | x (cid:48)

Original: [ (cid:96) ( Y, h ∗ ( z (cid:48) ))]
Revised: [ (cid:96) ( Y, h ∗ ( z (cid:48) ))]

Original: + (cid:88) ( x
Revised: + (cid:88) ( x

Original: [Equation]
Revised: [Equation]

Original: z )(cid:54)
Revised: z )

Original: [Equation]
Revised: [Equation]

Original: ( x (cid:48)
Revised: ( x (cid:48)

Original: [Equation]
Revised: [Equation]

Original: z (cid:48) )
Revised: z (cid:48) )

Original: p X,Z ( x, z ) E p Y | x [ (cid:96) ( Y, h ∗ ( z ))] (35) ≥ p X,Z ( x (cid:48) , z (cid:48) )
Revised: p X,Z ( x, z ) E p Y | x [ (cid:96) ( Y, h ∗ ( z ))] (34) ≥ p X,Z ( x (cid:48) , z (cid:48) )

Original: E p Y | x (cid:48)
Revised: E p Y | x (cid:48)

Original: [ (cid:96) ( Y, h ∗ ( z (cid:48) ))]
Revised: [ (cid:96) ( Y, h ∗ ( z (cid:48) ))]

Original: + (cid:88) ( x,z ) (cid:54) =(
Revised: + (cid:88) ( x,z ) (cid:54) =(

Original: x (cid:48) ,z (cid:48) )
Revised: x (cid:48) ,z (cid:48) )

Original: p X,Z ( x, z ) E p Y | x [ (cid:96) ( Y, f ∗ ( x ))] (36) > p X,Z ( x (cid:48) , z (cid:48) )
Revised: p X,Z ( x, z ) E p Y | x [ (cid:96) ( Y, f ∗ ( x ))] (35) > p X,Z ( x (cid:48) , z (cid:48) )

Original: E p Y | x (cid:48)
Revised: E p Y | x (cid:48)

Original: [ (cid:96) ( Y, f ∗ ( x (cid:48) ))]
Revised: [ (cid:96) ( Y, f ∗ ( x (cid:48) ))]

Original: + (cid:88)( x
Revised: + (cid:88)

Original: [Equation]
Revised: [Equation]

Original: z
Revised: x

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: (cid:48)

Original: [Equation] [Equation][Equation]
Revised: [Equation] [Equation]

Original: Eq.
Revised: Eq.

Original: (36) follows by Item 3 of Assumption 2 along with the deﬁnition of f ∗ .
Revised: (35) follows by Item 3 of Assumption 2 along with the deﬁnition of f ∗ .

Original: Eq.
Revised: Eq.

Original: (37) follows by Itemof Assumption 2 and the fact that h ∗ ( z (cid:48) ) (cid:54)
Revised: (36) follows by Itemof Assumption 2 and the fact that h ∗ ( z (cid:48) ) (cid:54)

Original: =
Revised: =

Original: f ∗ ( x (cid:48) ) .
Revised: f ∗ ( x (cid:48) ) .

Original: This completes the proof, because Lemmaprevents R [ Y | Z ] < R [ Y | X ] .
Revised: This completes the proof, because Lemmaprevents R [ Y | Z ] < R [ Y | X ] .

Original: B. P ROOF OF T HEOREM
Revised: B. P ROOF OF T HEOREM

Original: 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.
Revised: 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.

Original: This is a direct extension of the example from Dubois et al.’s
Revised: This is a direct extension of the example from Dubois et al.’s

Original: Proposition 6, to the case of proper losses.
Revised: Proposition 6, to the case of proper losses.

Original: Proposition 3 (Existence of optimal representations) .
Revised: Proposition 3 (Existence of optimal representations) .

Original: Under Assumptions 1 to 5, there exists an encoder p Z ∗ | X that is optimal for Eqs.
Revised: Under Assumptions 1 to 5, there exists an encoder p Z ∗ | X that is optimal for Eq.

Original: (3) and (4) , i.e.,
Revised: (3) and ?? , i.e.,

Original: [Equation]
Revised: [Equation]

Original: Moreover, we have that
Revised: Moreover, we have that

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Because we assume arbitrary encoders p Z | X , the essence of this construction is simple: we embed the Bayes image into Z .
Revised: Because we assume arbitrary encoders p Z | X , the essence of this construction is simple: we embed the Bayes image into Z .

Original: 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).
Revised: 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).

Original: Then let Z ∗ =
Revised: Then let Z ∗ =

Original: φ ( f ∗ ( X )) .
Revised: φ ( f ∗ ( X )) .

Original: We now verify the properties of p Z ∗ | X .
Revised: We now verify the properties of p Z ∗ | X .

Original: Z ∗ satisﬁes R [ Y | X ]
Revised: Z ∗ satisﬁes R [ Y | X ]

Original: =
Revised: =

Original: R [ Y | Z ∗ ] .
Revised: R [ Y | Z ∗ ] .

Original: Indeed,
Revised: Indeed,

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (43) is by our construction of Z ∗ and Eq.
Revised: (42) is by our construction of Z ∗ and Eq.

Original: (44) is by the deﬁnition of the Bayes risk.
Revised: (43) is by the deﬁnition of the Bayes risk.

Original: Due to the data processing inequality of Bayes risk (Lemma 1) wealso have R [ Y | X ] ≤
Revised: Due to the data processing inequality of Bayes risk (Lemma 1) we

Original:  R [ Y | Z ∗ ] , from which we concludethat R [ Y | X ] =
Revised: also have R [ Y | X ] ≤ R [ Y | Z ∗ ] , from which we conclude

Revised: that

Original: R [ Y |  Z ∗ ] and that Eq.
Revised: R [ Y | X ] = R [ Y | Z ∗ ] and that Eq.

Original: (40) holds.
Revised: (39) holds.

Original: Recall that Γ ∗ = f ∗ → ( 
Revised: Recall that Γ ∗ = f ∗ → ( X ) and Γ ∗ d = f ∗

Original: X

Original: and Γ ∗ d = f ∗ →
Revised: →

Original: (cid:0)
Revised: (cid:0)

Original: supp( p X |
Revised: supp( p X |

Original: d
Revised: d

Original: (cid:1) .
Revised: (cid:1) .

Original: Now let us compute the
Revised: Now let us compute the

Original: desired 
Revised: desired support

Original: support for all d ∈ D
Revised: for all d ∈ D

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (46) is by Assumption 5.
Revised: (45) is by Assumption 5.

Original: Because R [ Y | X ] is the minimum achievable risk by any encoder regardless of constraint (this is by
Revised: Because R [ Y | X ] is the minimum achievable risk by any encoder regardless of constraint (this is by

Original: Lemma 1), this implies that p Z ∗ | X is an optimal encoder for Eqs.
Revised: Lemma 1), this implies that p Z ∗ | X is an optimal encoder for Eq.

Original: (3) and (4).
Revised: (3) and ?? .

Original: 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.
Revised: 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.

Original: Lemma 3.
Revised: Lemma 3.

Original: Let p Z | X be an encoder
Revised: Let p Z | X be an encoder

Original: such that R
Revised: such that R

Original: [ Y | Z ] =
Revised: [ Y | Z ] =

Original: R [ Y | X ] .
Revised: R [ Y | X ] .

Original: Under Assumptions 1 and 2, we have that for all z ∈ supp( p Z ) , there exists γ ∗ ∈ Γ such that
Revised: Under Assumptions 1 and 2, we have that for all z ∈ supp( p Z ) , there exists γ ∗ ∈ Γ such that

Original: [Equation]
Revised: [Equation]

Original: In other words, the restriction of any h ∗ ∈ H ∗ Z to supp( p Z ) is unique.
Revised: In other words, the restriction of any h ∗ ∈ H ∗ Z to supp( p Z ) is unique.

Original: If, in addition, Assumption 4 holds, then for all ( z, d ) ∈ supp( p Z,D t ) , γ ∈ Γ such that γ (cid:54) =
Revised: If, in addition, Assumption 4 holds, then for all ( z, d ) ∈ supp( p Z,D t ) , γ ∈ Γ such that γ (cid:54) =

Original: h ∗ ( z ) ,
Revised: h ∗ ( z ) ,

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: In other words the restrictionof unique and equal
Revised: In other words, the restriction

Original: Proof.
Revised: Proof.

Original: For the ﬁrst result, let z ∈ supp( p Z ) and consider x ∈ supp( p X | z ) .
Revised: For the ﬁrst result, let z ∈ supp( p Z ) and consider x ∈ supp( p X | z ) .

Original: By Lemma 2, it must be the case that f ∗ is constant on supp( p X | z ) .
Revised: By Lemma 2, it must be the case that f ∗ is constant on supp( p X | z ) .

Original: Thus, we can pick γ ∗ =
Revised: Thus, we can pick γ ∗ =

Original: f ∗ ( x ) .
Revised: f ∗ ( x ) .

Original: Now, let γ (cid:54) =
Revised: Now, let γ (cid:54) =

Original: γ ∗ .
Revised: γ ∗ .

Original: We have that,
Revised: We have that,

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (50) is due to the conditional independence of Y and Z given X .
Revised: (49) is due to the conditional independence of Y and Z given X .

Original: Eq.
Revised: Eq.

Original: (52) is due to Assumptionand the deﬁnition of the Bayes predictor.
Revised: (51) is due to Assumptionand the deﬁnition of the Bayes predictor.

Original: Let h ∗ : supp( p Z ) → Γ be the unique Bayes predictor from Z .
Revised: Let h ∗ : supp( p Z ) → Γ be the unique Bayes predictor from Z .

Original: Now, for the second result, note that
Revised: Now, for the second result, note that

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: where Eq.
Revised: where Eq.

Original: (59) is due to the deﬁnition of (domain-conditional) Bayes risk.
Revised: (58) is due to the deﬁnition of (domain-conditional) Bayes risk.

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: Lemma
Revised: Lemma

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Thus, any encoder that achieves R [ Y | Z ] =
Revised: Thus, any encoder that achieves R [ Y | Z ] =

Original: R [ Y | X ] also satisﬁes R d
Revised: R [ Y | X ] also satisﬁes R d

Original: [ Y | Z ] =
Revised: [ Y | Z ] =

Original: R d [ Y | X ] for all d ∈ D since we assume that supp( p D t )
Revised: R d [ Y | X ] for all d ∈ D since we assume that supp( p D t )

Original: =
Revised: =

Original: D in Assumption 1.
Revised: D in Assumption 1.

Original: Now, let d ∈ D .
Revised: Now, let d ∈ D .

Original: An argument analogous to Lemma 2 gives us,
Revised: An argument analogous to Lemma 2 gives us,

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (62) is derived from R d
Revised: (61) is derived from R d

Original: [ Y | Z ] =
Revised: [ Y | Z ] =

Original: R d [ Y | X ] using Assumption 4 in place of Item 3 of Assumption 2 for a speciﬁc domain d .
Revised: R d [ Y | X ] using Assumption 4 in place of Item 3 of Assumption 2 for a speciﬁc domain d .

Original: Let z ∈ supp( p Z | d ) and γ ∈ Γ such that γ (cid:54)
Revised: Let z ∈ supp( p Z | d ) and γ ∈ Γ such that γ (cid:54)

Original: =
Revised: =

Original: h ∗ ( z ) .
Revised: h ∗ ( z ) .

Original: Since supp( p X | z,d ) ⊆ supp( p X | z ) , f ∗ is a constant on supp( p X | z,d ) and equal to h ∗ .
Revised: Since supp( p X | z,d ) ⊆ supp( p X | z ) , f ∗ is a constant on supp( p X | z,d ) and equal to h ∗ .

Original: Now, as above, we have that
Revised: Now, as above, we have that

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (65) is due to Assumption 4.
Revised: (64) is due to Assumption 4.

Original: Corollary 1.
Revised: Corollary 1.

Original: Let p Z | X be an encoder such that R [ Y | Z ]
Revised: Let p Z | X be an encoder such that R [ Y | Z ]

Original: =
Revised: =

Original: R [ Y | X ] .
Revised: R [ Y | X ] .

Original: Under Assumptions 1, 2 and 4 we have that
Revised: Under Assumptions 1, 2 and 4 we have that

Original: [Equation]
Revised: [Equation]

Original: all d
Revised: all d

Original: [Equation]
Revised: [Equation]

Original: D and that for all ds
Revised: D and that for all d

Original: [Equation]
Revised: [Equation]

Original: d t ∈ D
Revised: d t ∈ D

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Proof.
Revised: Proof.

Original: H ∗ Z ⊆
Revised: H ∗ Z ⊆

Original: H ∗ Z,d is immediate from Lemma 3.
Revised: H ∗ Z,d is immediate from Lemma 3.

Original: Now, we have that R d th [ Y | Z ] ≥ R d t [ Y | Z ] .
Revised: Now, we have that R d th [ Y | Z ] ≥ R d t [ Y | Z ] .

Original: So, the result follows by taking any h ∈ H ∗ Z ⊆
Revised: So, the result follows by taking any h ∈ H ∗ Z ⊆

Original: H ∗ Z,ds in the inf of Eq.
Revised: H ∗ Z,ds in the inf of Eq.

Original: Theorem 2 (Characterizing optimal representations for IDG) .
Revised: Theorem 2 (Characterizing optimal representations for IDG, equiv. Theorem 1) .

Original: Under Assumptions 1 to 6, an encoder p Z | X is optimal for idealized domain generalization if and only if it minimizes the Bayes risk while matching the support of p Z | d and p Z for all d ∈ D , i.e.,
Revised: Under Assumptions 1 to 6, an encoder p Z | X is optimal for idealized domain generalization if and only if it minimizes the Bayes risk while matching the support of p Z | d and p Z for all d ∈ D , i.e.,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: The IDG risk is lower bounded by R [ Y | X ] :
Revised: The IDG risk is lower bounded by R [ Y | X ] :

Original: [Equation]
Revised: [Equation]

Original: Lemma

Original: [Equation]

Original: Assumption

Original: [Equation]

Original: We will now show that this lower bound is achieved by an encoder if and only if it satisﬁes Eqs.
Revised: We will now show that this lower bound is achieved by an encoder if and only if it satisﬁes Eqs.

Original: (68) and (69), which exist by Prop. 3.
Revised: (67) and (68), which exist by Prop. 3.

Original: Sufﬁciency ( ⇐ = ):
Revised: Sufﬁciency ( ⇐ = ):

Original: Let p Z | X be an encoder that satisﬁes Eqs.
Revised: Let p Z | X be an encoder that satisﬁes Eqs.

Original: (68) and (69).
Revised: (67) and (68).

Original: Note that R [ Y | Z ] =
Revised: Note that R [ Y | Z ] =

Original: R [ Y | X ] by Prop. 3.
Revised: R [ Y | X ] by Prop. 3.

Original: Let h ∗ ∈
Revised: Let h ∗ ∈

Original: H ∗ Z , then we have the following IDG risk
Revised: H ∗ Z , then we have the following IDG risk

Original: [Equation]
Revised: [Equation]

Original: Lemma 3 under matching support
Revised: Lemma 3 under matching support

Original: [Equation]
Revised: [Equation]

Original: Necessity ( = ⇒ ):
Revised: Necessity ( = ⇒ ):

Original: If the IDG risk is R [ Y | X ] , then it must be the case that
Revised: If the IDG risk is R [ Y | X ] , then it must be the case that

Original: [Equation]
Revised: [Equation]

Original: We will prove by contrapositive that Eq.
Revised: We will prove by contrapositive that Eq.

Original: (80) implies support match (Eq.
Revised: (79) implies support match (Eq.

Original: Suppose that the support match does not hold.
Revised: Suppose that the support match does not hold.

Original: Since supp( p Z ) = ∪ d ∈D supp( p Z | d ) and supp( p D s ,D t ) =
Revised: Since supp( p Z ) = ∪ d ∈D supp( p Z | d ) and supp( p D s ,D t ) =

Original: D × D
Revised: D × D

Original: (Assumption 6), there must exist ( d s , d t ) ∈ supp( p D s ,D t ) such that supp( p Z | d s ) (cid:54) =
Revised: (Assumption 6), there must exist ( d s , d t ) ∈ supp( p D s ,D t ) such that supp( p Z | d s ) (cid:54) =

Original: supp( p Z | d s ) .
Revised: supp( p Z | d s ) .

Original: Deﬁne the set S = supp
Revised: Deﬁne the set S = supp

Original: ( p Z | d s ) ∩ supp( p Z | d t ) and ¯ S = supp( p Z | d t ) \ supp( p Z | d s ) , let ρ =
Revised: ( p Z | d s ) ∩ supp( p Z | d t ) and ¯ S = supp( p Z | d t ) \ supp( p Z | d s ) , let ρ =

Original: P Z | d t ( S ) [Equation] and let h ∗ ∈ H ∗ Z .
Revised: P Z | d t ( S ) , and let h ∗ ∈ H ∗ Z .

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: Lem.
Revised: Lem.

Original: [Equation] [Equation][Equation]
Revised: [Equation] [Equation]

Original: >

Original: [Equation]
Revised: [Equation]

Original: Lem.
Revised: Lem.

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Eq.
Revised: Eq.

Original: (86) uses the following reasoning.
Revised: (85) uses the following reasoning.

Original: 1 − ρ > 0 due to support mismatch.
Revised: 1 − ρ > 0 due to support mismatch.

Original: For any h ∈ H ∗ Z,ds such that h (cid:54) =
Revised: For any h ∈ H ∗ Z,ds such that h (cid:54) =

Original: h ∗ on ¯ S (such an h exists by Item 1 of Assumption 2), we have that
Revised: h ∗ on ¯ S (such an h exists by Item 1 of Assumption 2), we have that

Original: [Equation]
Revised: [Equation]

Original: by Lemma 3.
Revised: by Lemma 3.

Original: As a corollary from the proof strategy we directly have that the optimal DG risk is simply R [ Y | X ] .
Revised: As a corollary from the proof strategy we directly have that the optimal DG risk is simply R [ Y | X ] .

Original: 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.
Revised: 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.

Original: Corollary 2 (Optimal IDG Risk) .
Revised: Corollary 2 (Optimal IDG Risk) .

Original: Under Assumptions 1 to 6, inf p Z |
Revised: Under Assumptions 1 to 6, inf p Z |

Original: X R IDG
Revised: X R IDG

Original: [ Y | Z ] =
Revised: [ Y | Z ] =

Original: R [ Y | X ] .
Revised: R [ Y | X ] .

Original: B. I MPOSSIBILITY RESULTS
Revised: B. I MPOSSIBILITY RESULTS

Original: As a direct corollary of Thm.
Revised: As a direct corollary of Thm.

Original: 2 we know that it is impossible to learn an optimal representation without knowledge or assumptions on the target domain.
Revised: 2 we know that it is impossible to learn an optimal representation without knowledge or assumptions on the target domain.

Original: We can actually prove the following much stronger negative result, which essentially states that it is impossible to ﬁnd a useful representation without having some information about the target domain.
Revised: We can actually prove the following much stronger negative result, which essentially states that it is impossible to ﬁnd a useful representation without having some information about the target domain.

Original: Speciﬁcally, we prove that if there exists a non-trivial target domain on which the representation is advantageous then there exists an inﬁnite amount of target domains on which it is disadvantageous compared to predicting from a constant.
Revised: Speciﬁcally, we prove that if there exists a non-trivial target domain on which the representation is advantageous then there exists an inﬁnite amount of target domains on which it is disadvantageous compared to predicting from a constant.

Original: 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.
Revised: 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.

Original: The key is that outside of the source domain, the label distribution is unconstrained because generalized covariate shift has no effect.
Revised: The key is that outside of the source domain, the label distribution is unconstrained because generalized covariate shift has no effect.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Proposition 4 (No free lunch  learning representations for DG) .
Revised: Proposition 4 (No free lunch for learning representations for IDG, equiv. Proposition 1) .

Original: Let (cid:96) be the 0-1 loss with prediction space Γ =
Revised: Let (cid:96) be the 0-1 loss with prediction space Γ =

Original: Y .
Revised: Y .

Original: Let Rep : P ( X , Y ) → P ( Z|X ) be any algorithm for choosing an encoder p Z | X from the data distribution p X,Y , C be any constant r.v.
Revised: Let Rep : P ( X , Y ) → P ( Z|X ) be any algorithm for choosing an encoder p Z | X from the data distribution p X,Y , C be any constant r.v.

Original: that t.v.i.
Revised: that t.v.i.

Original: Z , and p X,Y | d s be any desired source distribution such that
Revised: Z , and p X,Y | d s be any desired source distribution such that

Original: • there is a unique constant prediction γ C = arg min
Revised: • there is a unique constant prediction γ C = arg min

Original: y ∈Y E p Y | ds [ (cid:96) ( Y, y )] , • and |X \ supp( p X | d s ) | > 1 .
Revised: y ∈Y E p Y | ds [ (cid:96) ( Y, y )] , • and |X \ supp( p X | d s ) | > 1 .

Original: Let p Z ds | X :=
Revised: Let p Z ds | X :=

Original: Rep( p X,Y | d s ) be the chosen source encoder.
Revised: Rep( p X,Y | d s ) be the chosen source encoder.

Original: If there exists a target domain p X,Y | d gt such that
Revised: If there exists a target domain p X,Y | d gt such that

Original: • (Non-trivial support) ∅ (cid:54) =
Revised: • (Non-trivial support) ∅ (cid:54) =

Original: supp( p X | d gt ) ⊆
Revised: supp( p X | d gt ) ⊆

Original: X \ supp( p X | d s ) ; • (Satisﬁes Bayes image invariance)
Revised: X \ supp( p X | d s ) ; • (Satisﬁes Bayes image invariance)

Original: Γ ∗ d gt = Y , i e., there is at least one example for every possible label; • (Source encoder is useful)
Revised: Γ ∗ d gt = Y , i.e., there is at least one example for every possible label; • (Source encoder is useful)

Original: p Z ds | X performs better than a constant representation,
Revised: p Z ds | X performs better than a constant representation,

Original: [Equation]
Revised: [Equation]

Original: Then there exist multiple target domains d bt such that p Z ds |
Revised: Then there exist multiple target domains d bt such that p Z ds |

Original: X underperforms a constant encoder,
Revised: X underperforms a constant encoder,

Original: [Equation]
Revised: [Equation]

Revised: h ∈H ∗ C,d s R d bth

Original: R d bth [ Y | C ] .
Revised: [ Y [Equation] C ] .

Original: Proof.
Revised: Proof.

Original: Let h ∗ ∈
Revised: Let h ∗ ∈

Original: H ∗ Zds,ds be any source Bayes predictor corresponding to our encoder.
Revised: H ∗ Zds,ds be any source Bayes predictor corresponding to our encoder.

Original: Partition Z according to whether h ∗ predicts like the constant or not:
Revised: Partition Z according to whether h ∗ predicts like the constant or not:

Original: [Equation]
Revised: [Equation]

Original: We know by assumption that d gt is s.t.
Revised: We know by assumption that d gt is s.t.

Original: [Equation]
Revised: [Equation]

Original: which is clearly only possible if
Revised: which is clearly only possible if

Original: [Equation]
Revised: [Equation]

Original: In other words, there exists some input x (cid:54) =
Revised: In other words, there exists some input x (cid:54) =

Original: C ∈ X \ supp( p X | d s ) that will get represented outside of the constant region, i.e.,
Revised: C ∈ X \ supp( p X | d s ) that will get represented outside of the constant region, i.e.,

Original: [Equation]
Revised: [Equation]

Original: We will now construct the desired bad domain d bt by giving nearly all mass to this x (cid:54) =
Revised: We will now construct the desired bad domain d bt by giving nearly all mass to this x (cid:54) =

Original: C , speciﬁcally, let p X | d bt ( x (cid:54) =
Revised: C , speciﬁcally, let p X | d bt ( x (cid:54) =

Original: C )
Revised: C )

Original: = 1 − δ for some 0 < δ < 1 .
Revised: = 1 − δ for some 0 < δ < 1 .

Original: We assign this example to the constant label, i.e., p Y | x (cid:54) =
Revised: We assign this example to the constant label, i.e., p Y | x (cid:54) =

Original: C ,d bt ( γ C ) = 1 .
Revised: C ,d bt ( γ C ) = 1 .

Original: The rest of the target domain mass δ is distributed as with the source domain, i.e., p X,Y | d bt ( x, y )
Revised: The rest of the target domain mass δ is distributed as with the source domain, i.e., p X,Y | d bt ( x, y )

Original: =
Revised: =

Original: δ · p X,Y |
Revised: δ · p X,Y |

Original: d s ( x, y ) for all x, y ∈ supp( p X,Y | d s ) .
Revised: d s ( x, y ) for all x, y ∈ supp( p X,Y | d s ) .

Original: Importantly, the constructed domain d bt is valid.
Revised: Importantly, the constructed domain d bt is valid.

Original: Indeed, the Bayes image is the same as the source’s
Revised: Indeed, the Bayes image is the same as the source’s

Original: (Assumption 5), because we removed no prediction γ from the source’s Bayes image ( δ > 0 ).
Revised: (Assumption 5), because we removed no prediction γ from the source’s Bayes image ( δ > 0 ).

Original: We added no new prediction γ , because f ∗ ( x (cid:54) = C ) =
Revised: We added no new prediction γ , because f ∗ ( x (cid:54) = C ) =

Original: γ C ∈ Y which must already have been in Γ ∗ due to the validity of d gt .
Revised: γ C ∈ Y which must already have been in Γ ∗ due to the validity of d gt .

Original: Now let us compute the desired risk for that “bad” domain and show that the desired encoder performs worse than a constant encoder.
Revised: Now let us compute the desired risk for that “bad” domain and show that the desired encoder performs worse than a constant encoder.

Original: [Equation]
Revised: [Equation]

Original: In contrast, it is easy to show that sup h ∈H ∗ C,ds R d bth [ Y | C ] ≤ δ because the constant predictor would be perfect for x (cid:54)
Revised: In contrast, it is easy to show that sup h ∈H ∗ C,ds R d bth [ Y | C ] ≤ δ because the constant predictor would be perfect for x (cid:54)

Original: =
Revised: =

Original: C .
Revised: C .

Original: So any choice of 0 < δ <
Revised: So any choice of 0 < δ <

Original: [Equation]
Revised: [Equation]

Original: , would satisfy Eq.
Revised: , would satisfy Eq.

Original: [Equation]
Revised: [Equation]

Original: We conclude the proof by noting that there are inﬁnitely many such choices of δ , and any choice of those would result in a different valid bad domain d bt .
Revised: We conclude the proof by noting that there are inﬁnitely many such choices of δ , and any choice of those would result in a different valid bad domain d bt .

Original: Note that representations can often be much worse than using a constant r.v.
Revised: Note that representations can often be much worse than using a constant r.v.

Original: Speciﬁcally, if an encoder p Z | X maps an x outside of the source support then there exists an inﬁnite number of target domains where that representation is the worst possible representation.
Revised: Speciﬁcally, if an encoder p Z | X maps an x outside of the source support then there exists an inﬁnite number of target domains where that representation is the worst possible representation.

Original: Proposition 5 (Worst representation) .
Revised: Proposition 5 (Worst representation) .

Original: Let Rep , p Y,X | d s , p Z ds | X , (cid:96) be as in Prop. 1, and (cid:15) > 0 .
Revised: Let Rep , p Y,X | d s , p Z ds | X , (cid:96) be as in Prop. 4, and (cid:15) > 0 .

Original: If there exists an example x b ∈ X \ supp( p X | d s ) that is mapped outside of the source support, i.e., supp( p
Revised: If there exists an example x b ∈ X \ supp( p X | d s ) that is mapped outside of the source support, i.e., supp( p

Original: [Equation]
Revised: [Equation]

Original: ds
Revised: ds

Original: [Equation]
Revised: [Equation]

Original: b ) ∩ supp( p Z | d s ) = ∅ , then there exist many target domains p X,Y | d t s.t.
Revised: b ) ∩ supp( p Z | d s ) = ∅ , then there exist many target domains p X,Y | d t s.t.

Original: p Z ds |
Revised: p Z ds |

Original: X is (cid:15) close to the worst possible loss, i.e.,
Revised: X is (cid:15) close to the worst possible loss, i.e.,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: By assumption there exists an x b whose support is outside the source support.
Revised: By assumption there exists an x b whose support is outside the source support.

Original: Then similarly to Prop. 1 we construct a bad target domain d t by giving nearly all mass to that example p X | d t ( x b ) = 1 − δ where δ > 0 and assign with probability 1 to some label that is in the source Bayes image, i.e., p Y | x b ,d t ( γ b ) = 1 for some γ b ∈ Γ ∗ d s .
Revised: Then similarly to Prop. 4 we construct a bad target domain d t by giving nearly all mass to that example p X | d t ( x b ) = 1 − δ where δ > 0 and assign with probability 1 to some label that is in the source Bayes image, i.e., p Y | x b ,d t ( γ b ) = 1 for some γ b ∈ Γ ∗ d s .

Original: The rest of the target domain mass δ is distributed as in
Revised: The rest of the target domain mass δ is distributed as in

Original: Prop. 1 to the source inputs.
Revised: Prop. 4 to the source inputs.

Original: As in Prop. 1, such a target domain d t satisﬁes our assumptions.
Revised: As in Prop. 4, such a target domain d t satisﬁes our assumptions.

Original: Now let us compute the risk for that d t and show that the desired encoder performs arbitrarily bad.
Revised: Now let us compute the risk for that d t and show that the desired encoder performs arbitrarily bad.

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (102) uses the fact that H ∗ Zds,ds is unconstrained outside of the source support and that by assumption supp( p Z ds | x b ) ∩ supp( p Z ds | d s ) = ∅ .
Revised: (101) uses the fact that H ∗ Zds,ds is unconstrained outside of the source support and that by assumption supp( p Z ds | x b ) ∩ supp( p Z ds | d s ) = ∅ .

Original: To achieve the sup 1 − δ it then sufﬁces to predict an γ (cid:54) =
Revised: To achieve the sup 1 − δ it then sufﬁces to predict an γ (cid:54) =

Original: γ b ∈ Γ .
Revised: γ b ∈ Γ .

Original: We thus see that Eq.
Revised: We thus see that Eq.

Original: (98) holds for d t as long as 0 < δ < (cid:15) .
Revised: (97) holds for d t as long as 0 < δ < (cid:15) .

Original: We conclude the proof by noting that there is an inﬁnite possible choices of δ each of which give rise to a bad target domain.
Revised: We conclude the proof by noting that there is an inﬁnite possible choices of δ each of which give rise to a bad target domain.

Original: B. A UGMENTATIONS
Revised: B. A UGMENTATIONS

Original: Prop. 2 shows that the optimal representations for IDG can be learned with augmentations in a self-supervised fashion.
Revised: Proposition 2 shows that the optimal representations for IDG can be learned with augmentations in a self-supervised fashion.

Original: Here, we provide formal deﬁnitions, assumptions, and proofs.
Revised: Here, we provide formal deﬁnitions, assumptions, and proofs.

Original: Deﬁnition 6 (Augmenter) .
Revised: Deﬁnition 6 (Augmenter) .

Original: An augmenter is a conditional distribution p
Revised: An augmenter is a conditional distribution p

Original: A | X : A × X → [0 , 1] from the input space
Revised: A | X : A × X → [0 , 1] from the input space

Original: X to an augmentation space A .
Revised: X to an augmentation space A .

Original: For example, in CLIP X is the space of images and A is the space of text.
Revised: For example, in CLIP X is the space of images and A is the space of text.

Original: In standard SSL,
Revised: In standard SSL,

Original: A is typically the same as X (e.g., both X
Revised: A is typically the same as X (e.g., both X

Original: and A are the space of images).
Revised: and A are the space of images).

Original: Deﬁnition 7 (Augmentation conditional set) .
Revised: Deﬁnition 7 (Augmentation conditional set) .

Original: Given an augmenter p
Revised: Given an augmenter p

Original: A | X , deﬁne the augmentation conditional set as the set of conditionals of A given X :
Revised: A | X , deﬁne the augmentation conditional set as the set of conditionals of A given X :

Original: [Equation]
Revised: [Equation]

Original: Similarly, we can deﬁne the augmentation conditional set for domain d :
Revised: Similarly, we can deﬁne the augmentation conditional set for domain d :

Original: [Equation]
Revised: [Equation]

Original: These sets are clearly countable.
Revised: These sets are clearly countable.

Original: 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).
Revised: 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).

Original: Assumption 7 (Finite augmentation entropy) .
Revised: Assumption 7 (Finite augmentation entropy) .

Original: We consider the augmenter p
Revised: We consider the augmenter p

Original: A | X such that the entropy of the augmentation A is ﬁnite, i.e., H[ A ] < ∞ .
Revised: A | X such that the entropy of the augmentation A is ﬁnite, i.e., H[ A ] < ∞ .

Original: Assumption 8 (Cardinalities) .
Revised: Assumption 8 (Cardinalities) .

Original: We assume that
Revised: We assume that

Original: [Equation]
Revised: [Equation]

Original: This is a similar assumption as Assumption 3, which ensures the existence of optimal representations.
Revised: This is a similar assumption as Assumption 3, which ensures the existence of optimal representations.

Original: Assumption 9 (Domain-covering augmentation) .
Revised: Assumption 9 (Domain-agnostic augmentation) .

Original: We assume that the augmentation A is domaincovering , i.e., the augmentation conditional set is invariant across domains,
Revised: We assume that the augmentation A is domainagnostic , i.e., the augmentation conditional set is invariant across domains,

Original: [Equation]
Revised: [Equation]

Original: This assumption is generalized from the constant Bayes image assumption (Assumption 5), which guarantees the existence of optimal representations.
Revised: This assumption is generalized from the constant Bayes image assumption (Assumption 5), which guarantees the existence of optimal representations.

Original: Domain-covering augmentations essentially ensures that each augmentation conditional p
Revised: Domain-agnostic augmentations essentially ensures that each augmentation conditional p

Original: A | x ∈
Revised: A | x ∈

Original: P ∗ ( A | X ) is seen at least once in all domains.
Revised: P ∗ ( A | X ) is seen at least once in all domains.

Original: If we introduce an equivalence relation ∼ as x ∼ x (cid:48) iff p
Revised: If we introduce an equivalence relation ∼ as x ∼ x (cid:48) iff p

Original: A | x = p
Revised: A | x = p

Original: A | x (cid:48) and the equivalence class [ x ] : = { x (cid:48) ∈ X | x (cid:48) ∼ x } .
Revised: A | x (cid:48) and the equivalence class [ x ] : = { x (cid:48) ∈ X | x (cid:48) ∼ x } .

Original: Under this relation, it is easy to see that the above assumption is satisﬁed if and only if, for all possible equivalence classes [ x ] ∈{ [ x (cid:48) ] | x (cid:48) ∈ X
Revised: Under this relation, it is easy to see that the above assumption is satisﬁed if and only if, for all possible equivalence classes [ x ] ∈{ [ x (cid:48) ] | x (cid:48) ∈ X

Original: } , we have that [ x ] has intersections with all domains:
Revised: } , we have that [ x ] has intersections with all domains:

Original: [Equation]
Revised: [Equation]

Original: Not all augmentations are domain-covering.
Revised: Not all augmentations are domain-agnostic.

Original: In particular, the standard image augmentations used by typical SSL models like SimCLR are not domain-covering, but the text-image augmentations of CLIP nearly are, as discussed in the main body (Sec. 4).
Revised: 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).

Original: Assumption 10 (Bayes-preserving augmentation) .
Revised: Assumption 10 (Bayes-preserving augmentation) .

Original: We assume that the augmentation A is Bayespreserving , i.e.,
Revised: We assume that the augmentation A is Bayespreserving , i.

Original: [Equation]
Revised: [Equation]

Original:  Under the notion of equivalence relation in Assumption 9, this means that for each equivalence class [ x ] , all x (cid:48) ∈
Revised: [Equation] Under the notion of equivalence relation in Assumption 9, this means that for each equivalence class [ x ] , all x (cid:48) ∈

Original: [ x ] have the same Bayes prediction.
Revised: [ x ] have the same Bayes prediction.

Original: Note that most augmentations used in practice like standard image augmentations are Bayes-preserving.
Revised: Note that most augmentations used in practice like standard image augmentations are Bayes-preserving.

Original: Next, we show that under the above assumptions, we can learn optimal representations by maximizing the mutual information
Revised: Next, we show that under the above assumptions, we can learn optimal representations by maximizing the mutual information

Original: I[ A ; Z ] (in the case of log-loss (cid:96) ) under the support match constraint.
Revised: I[ A ; Z ] (in the case of log-loss (cid:96) ) under the support match constraint.

Original: We use log-loss simply because it is typically the loss used for training in practice.
Revised: We use log-loss simply because it is typically the loss used for training in practice.

Original: Note that the learned representations are optimal for any strict proper losses.
Revised: Note that the learned representations are optimal for any strict proper losses.

Original: Proposition 6 (Optimal encoders without labels) .
Revised: Proposition 6 (Learning optimal representations without labels, equiv. Proposition 2) .

Original: Let p A |
Revised: Let p A |

Original: X be an augmenter.
Revised: X be an augmenter.

Original: Under Assumptionsto 10, any encoder p Z | X such that
Revised: Under Assumptions 1 to 10, any encoder p Z | X such that

Original: [Equation]
Revised: [Equation]

Original: is optimal for idealized domain generalization.
Revised: is optimal for idealized domain generalization.

Original: Proof.
Revised: Proof.

Original: The support match constraint Eq.
Revised: The support match constraint Eq.

Original: (110) is equivalent to the support match constraint Eq.
Revised: (109) is equivalent to the support match constraint Eq.

Original: Thus, Prop. 3 and Thm. 2 state that we only need to prove that maximizing the mutual information of
Revised: Thus, Prop. 3 and Thm. 2 state that we only need to prove that maximizing the mutual information of

Original: A and Z under the support constraint implies that
Revised: A and Z under the support constraint implies that

Original: [Equation]
Revised: [Equation]

Original: We will prove this by constructing an optimal predictor h ∗ .
Revised: We will prove this by constructing an optimal predictor h ∗ .

Original: Since H[ A ] < ∞ (Assumption 7) we have that
Revised: Since H[ A ] < ∞ (Assumption 7) we have that

Original: [Equation]
Revised: [Equation]

Original: Note the fact that the conditional entropy is the Bayes risk under the log-loss (Gneiting & Raftery, 2007), i.e., H
Revised: Note the fact that the conditional entropy is the Bayes risk under the log-loss (Gneiting & Raftery, 2007), i.e., H

Original: [ A | Z ]
Revised: [ A | Z ]

Original: = R
Revised: = R

Original: [ A | Z ] .
Revised: [ A | Z ] .

Original: By construction, A satisﬁes covariate shift w.r.t.
Revised: By construction, A satisﬁes covariate shift w.r.t.

Original: X (thus Bayes invariant) since A − X − D forms a Markov chain.
Revised: X (thus Bayes invariant) since A − X − D forms a Markov chain.

Original: Together with Assumptions 1 and 7 to 9, it means that the optimization problem in Eqs.
Revised: Together with Assumptions 1 and 7 to 9, it means that the optimization problem in Eqs.

Original: (109) and (110) satisﬁes the assumptions of Prop. 3, with A in place of Y .
Revised: (108) and (109) satisﬁes the assumptions of Prop. 3, with A in place of Y .

Original: Thus, an optimal encoder satisﬁes R
Revised: Thus, an optimal encoder satisﬁes R

Original: [ A | Z ]
Revised: [ A | Z ]

Original: = R
Revised: = R

Original: [ A | X ] , which leads to
Revised: [ A | X ] , which leads to

Original: [Equation]
Revised: [Equation]

Original: By Assumption 7, we can invoke Lemma 2 with the fact that A − X − Z forms a Markove chain to show that for all ( x, z ) ∈ supp( p X,Z )[Equation]
Revised: By Assumption 7, we can invoke Lemma 2 with the fact that A − X − Z forms a Markove chain to show that for all ( x, z ) ∈ supp( p X,Z )

Original: [Equation]
Revised: [Equation]

Original:  as the conditional distributions are the Bayes optimal predictors due to strict properness of log-loss.
Revised: 113) as the conditional distributions are the Bayes optimal predictors due to strict properness of log-loss.

Original: Now, deﬁne the following equivalence relation on X ,
Revised: Now, deﬁne the following equivalence relation on X ,

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Because the number of equivalence classes under ∼ is countable, there exists a maximal invariant
Revised: Because the number of equivalence classes under ∼ is countable, there exists a maximal invariant

Original: M : X → N from X to the natural numbers (for our deﬁnition of a maximal invariant see Deﬁnition 2, Dubois et al., 2021).
Revised: M : X → N from X to the natural numbers (for our deﬁnition of a maximal invariant see Deﬁnition 2, Dubois et al., 2021).

Original: By Assumption 10, f ∗ is invariant on the equivalence classes [ x ] : =
Revised: By Assumption 10, f ∗ is invariant on the equivalence classes [ x ] : =

Original: { x (cid:48) ∈
Revised: { x (cid:48) ∈

Original: X | x (cid:48) ∼ x } for all x ∈ X .
Revised: X | x (cid:48) ∼ x } for all x ∈ X .

Original: Thus, there exists a function g : N → A such that f ∗ =
Revised: Thus, there exists a function g : N → A such that f ∗ =

Original: g ◦ M (Lemma 5, Dubois et al., 2021).
Revised: g ◦ M (Lemma 5, Dubois et al., 2021).

Original: Given z ∈ supp( p Z ) , we construct h ∗ in the following way.
Revised: Given z ∈ supp( p Z ) , we construct h ∗ in the following way.

Original: Let x z ∈ supp( p X | z ) be any input point that could have led to this representation z and deﬁne h ∗ ( z ) =
Revised: Let x z ∈ supp( p X | z ) be any input point that could have led to this representation z and deﬁne h ∗ ( z ) =

Original: g
Revised: g

Original: [Equation]
Revised: [Equation]

Original: By Eq.
Revised: By Eq.

Original: (114) we are guaranteed that all x ∈ supp( p X | z ) share the same value for f ∗ since they are in the same equivalence class.
Revised: (113) we are guaranteed that all x ∈ supp( p X | z ) share the same value for f ∗ since they are in the same equivalence class.

Original: Thus, by the deﬁnition of M we have that
Revised: Thus, by the deﬁnition of M we have that

Original: [Equation]
Revised: [Equation]

Original: for
Revised: for

Original: [Equation]
Revised: [Equation]

Original: Therefore,
Revised: Therefore,

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: C P RACTICAL OBJECTIVES
Revised: C P RACTICAL OBJECTIVES

Original: Proposition 6 provides an objective to obtain the desired optimal representations, compared to Thm.it is more practical in that it does not require access to the labels but right augmentations.
Revised: Proposition 6 provides an objective to obtain the desired optimal representations, compared to Thm.it is more practical in that it does not require access to the labels but right augmentations.

Original: There are nevertheless multiple remaining issues for deriving objectives that can be trained with in practice.
Revised: There are nevertheless multiple remaining issues for deriving objectives that can be trained with in practice.

Original: Speciﬁcally, (i) the support constraint is hard to satisfy in practice; (ii) mutual information
Revised: Speciﬁcally, (i) the support constraint is hard to satisfy in practice; (ii) mutual information

Original: I
Revised: I

Original: [ A ; Z ] is hard to estimate from samples (Poole et al., 2019); (iii) the objective is constrained which is harder to optimize.
Revised: [ A ; Z ] is hard to estimate from samples (Poole et al., 2019); (iii) the objective is constrained which is harder to optimize.

Original: 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.
Revised: 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.

Original: In contrast to the proofs of main theoretical results (previous section), here the derivations will be less formal.
Revised: In contrast to the proofs of main theoretical results (previous section), here the derivations will be less formal.

Original: As we have seen in Proposition 6, the optimal representation achieves I
Revised: As we have seen in Proposition 6, the optimal representation achieves I

Original: [ A ; Z ]
Revised: [ A ; Z ]

Original: = I
Revised: = I

Original: [ A ; X ] , and thus
Revised: [ A ; X ] , and thus

Original: H
Revised: H

Original: [ A | Z ]
Revised: [ A | Z ]

Original: = H
Revised: = H

Original: [ A | X ] (Eq.
Revised: [ A | X ] (Eq.

Original: In the following, we will rewrite the objective as the constrained optimization:
Revised: In the following, we will rewrite the objective as the constrained optimization:

Original: [Equation]
Revised: [Equation]

Original: 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).
Revised: 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).

Original: The requirement on the domain bottleneck objective is that minimizing Eq.
Revised: The requirement on the domain bottleneck objective is that minimizing Eq.

Original: (123) under Eq.
Revised: (122) under Eq.

Original: (124) implies that the support match constraint holds (and can be achieved by some encoder), which leads to optimal representations for idealized domain generalization.
Revised: (123) implies that the support match constraint holds (and can be achieved by some encoder), which leads to optimal representations for idealized domain generalization.

Original: Different domain bottlenecks will be derived later this section.
Revised: Different domain bottlenecks will be derived later this section.

Original: We can then use Lagrangian relaxation to get the following unconstrained objectives.
Revised: We can then use Lagrangian relaxation to get the following unconstrained objectives.

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst term can be easily optimized using variational objectives.
Revised: The ﬁrst term can be easily optimized using variational objectives.

Original: Throughout the paper, we will use a contrastive variational upper bound which is based on InfoNCE (Oord et al., 2018).
Revised: Throughout the paper, we will use a contrastive variational upper bound which is based on InfoNCE (Oord et al., 2018).

Original: Namely, let X := (cid:8) X, X − 1 , . . .
Revised: Namely, let X := (cid:8) X, X − 1 , . . .

Original: , X − n (cid:9) be a sequence of samples where each negative X − i is independently sampled from the marginal p X .
Revised: , X − n (cid:9) be a sequence of samples where each negative X − i is independently sampled from the marginal p X .

Original: Now let A :
Revised: Now let A :

Original: =
Revised: =

Original: (cid:8)
Revised: (cid:8)

Original: A + , A − 1 , . .
Revised: A + , A − 1 , . .

Original: A − n (cid:9) be the sequence of augmented examples that come from independently augmenting each X ∈ X using the augmenter p A | X .
Revised: A − n (cid:9) be the sequence of augmented examples that come from independently augmenting each X ∈ X using the augmenter p A | X .

Original: In particular, the augmented positive A + has the marginal p
Revised: In particular, the augmented positive A + has the marginal p

Original: A | X while the negatives A − i follow the marginal p A .
Revised: A | X while the negatives A − i follow the marginal p A .

Original: Let Z be the representation of X by passing it through the encoder p ϕ :
Revised: Let Z be the representation of X by passing it through the encoder p ϕ :

Original: =
Revised: =

Original: p Z | X parameterized by ϕ and s ψ the critic function parametrized by ψ used to score which A (cid:48) ∈
Revised: p Z | X parameterized by ϕ and s ψ the critic function parametrized by ψ used to score which A (cid:48) ∈

Original: A is the positive augmentation.
Revised: A is the positive augmentation.

Original: Then we have a variational distribution of p A | Z :
Revised: Then we have a variational distribution of p A | Z :

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: which leads to the following variational bound:
Revised: which leads to the following variational bound:

Original: [Equation]
Revised: [Equation]

Original: with equality if the variational families s ψ , p ϕ are unconstrained and and we use inﬁnite samples ( n → ∞ ).
Revised: with equality if the variational families s ψ , p ϕ are unconstrained and and we use inﬁnite samples ( n → ∞ ).

Original: Typically the critic is separable, i.e., s ψ ( A, Z ) :
Revised: Typically the critic is separable, i.e., s ψ ( A, Z ) :

Original: = g ψ ( A ) T h ψ ( Z ) .
Revised: = g ψ ( A ) T h ψ ( Z ) .

Original: As discussed in the main body, it can be tied with p ϕ when A = X .
Revised: As discussed in the main body, it can be tied with p ϕ when A = X .

Original: In the following we focus on the second term B[ Z, X, Y, D ] and discuss several choices.
Revised: In the following we focus on the second term B[ Z, X, Y, D ] and discuss several choices.

Original: C. M UTUAL INFORMATION BOTTLENECK B[ Z, X, Y, D ] = I[ Z ; X ]
Revised: C. M UTUAL INFORMATION BOTTLENECK B[ Z, X, Y, D ] = I[ Z ; X ]

Original: The ﬁrst bottleneck we consider is so called mutual information (MI) bottleneck B[ Z, X, Y, D ] =
Revised: The ﬁrst bottleneck we consider is so called mutual information (MI) bottleneck B[ Z, X, Y, D ] =

Original: I[ Z ; X ] , which was introduced by Tishby et al. (2000) to achieve a tradeoff between the predictive power and the complexity of representations.
Revised: I[ Z ; X ] , which was introduced by Tishby et al. (2000) to achieve a tradeoff between the predictive power and the complexity of representations.

Original: Intuitively, it tries to remove all information of Z that is not needed for maximizing I
Revised: Intuitively, it tries to remove all information of Z that is not needed for maximizing I

Original: [ Z ; A ] .
Revised: [ Z ; A ] .

Original: In particular, using the fact that Z − X − D forms a Markov chain and the chain rule of MI, we have I[ Z ; X ] =
Revised: In particular, using the fact that Z − X − D forms a Markov chain and the chain rule of MI, we have I[ Z ; X ] =

Original: I
Revised: I

Original: [ Z ; X, D ] =
Revised: [ Z ; X, D ] =

Original: I
Revised: I

Original: [ Z ; D ] + I
Revised: [ Z ; D ] + I

Original: [ Z ; X | D ] .
Revised: [ Z ; X | D ] .

Original: Thus, it not only minimizes I[ Z ; D ] , i.e., matches the representations’ distribution across domains, but also minimizes I
Revised: Thus, it not only minimizes I[ Z ; D ] , i.e., matches the representations’ distribution across domains, but also minimizes I

Original: [ Z ; X | D ] , i.e., matches the representations’ distribution inside domains.
Revised: [ Z ; X | D ] , i.e., matches the representations’ distribution inside domains.

Original: Why The key to show is that minimizing Eq.
Revised: Why The key to show is that minimizing Eq.

Original: (123), i.e., arg min p Z |
Revised: (122), i.e., arg min p Z |

Original: X I[ Z ; X ] under I[ A ; Z ] =
Revised: X I[ Z ; X ] under I[ A ; Z ] =

Original: I[ A ; X ] , implies the support match constraint.
Revised: I[ A ; X ] , implies the support match constraint.

Original: This can be seen as a speciﬁc subcase of Dubois et al.’s (2021)
Revised: This can be seen as a speciﬁc subcase of Dubois et al.’s (2021)

Original: Corollary 15 with A in place of Y and M ( X ) induced by p A |
Revised: Corollary 15 with A in place of Y and M ( X ) induced by p A |

Original: X
Revised: X

Original: as in the proof of
Revised: as in the proof of

Original: Prop. 6.
Revised: Prop. 6.

Original: From the corollary, we know that min p Z |
Revised: From the corollary, we know that min p Z |

Original: X I
Revised: X I

Original: [ Z ; X ] =
Revised: [ Z ; X ] =

Original: H[ M ( X )] which can be achieved by any Z s.t.
Revised: H[ M ( X )] which can be achieved by any Z s.t.

Original: p Z | x = p Z | x (cid:48)
Revised: p Z | x = p Z | x (cid:48)

Original: ⇐⇒ M ( x ) =
Revised: ⇐⇒ M ( x ) =

Original: M ( x (cid:48) ) .
Revised: M ( x (cid:48) ) .

Original: With the assumption of domain-covering augmentations (Assumption 9), we have that the set of maximal invariant { M ( x ) | x ∈ supp( p X | d ) } is invariant across domains.
Revised: With the assumption of domain-agnostic augmentations (Assumption 9), we have that the set of maximal invariant { M ( x ) | x ∈ supp( p X | d ) } is invariant across domains.

Original: Then we directly have supp( p Z | d ) = ∪ x ∈ supp( p X | d ) supp( p Z | x ) = ∪ x ∈ supp( p X ) supp( p Z | x ) =
Revised: Then we directly have supp( p Z | d ) = ∪ x ∈ supp( p X | d ) supp( p Z | x ) = ∪ x ∈ supp( p X ) supp( p Z | x ) =

Original: supp( p Z ) , where we use the fact that x within the same equivalence class has the the same p
Revised: supp( p Z ) , where we use the fact that x within the same equivalence class has the the same p

Original: Z | x .
Revised: Z | x .

Original: How Essentially, we can use any variational upper bound of mutual information.
Revised: How Essentially, we can use any variational upper bound of mutual information.

Original: We consider the one used by Variational Information Bottelenck (Alemi et al., 2016), i.e.,
Revised: We consider the one used by Variational Information Bottelenck (Alemi et al., 2016), i.e.,

Original: [Equation]
Revised: [Equation]

Original: where a variational distribution q θ is used to approximate p Z and is jointly optimized with p ϕ to minimize the bound.
Revised: where a variational distribution q θ is used to approximate p Z and is jointly optimized with p ϕ to minimize the bound.

Original: The approximation gap of the bound is D KL [ p Z ( Z ) (cid:107) q θ ( Z )] .
Revised: The approximation gap of the bound is D KL [ p Z ( Z ) (cid:107) q θ ( Z )] .

Original: Then the ﬁnal loss becomes
Revised: Then the ﬁnal loss becomes

Original: [Equation]
Revised: [Equation]

Original: which recovers the optimal encoder in the case of unconstrained variational families for p ϕ , q θ , s ψ , inﬁnite samples | A | , and any λ > 1 (Dubois et al., 2021).
Revised: which recovers the optimal encoder in the case of unconstrained variational families for p ϕ , q θ , s ψ , inﬁnite samples | A | , and any λ > 1 (Dubois et al., 2021).

Original: C. E NTROPY BOTTLENECK B[ Z, X, Y, D ] = H[ Z ]
Revised: C. E NTROPY BOTTLENECK B[ Z, X, Y, D ] = H[ Z ]

Original: 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 ϕ
Revised: 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 ϕ

Original: ( Z | x ) is a dirac delta function for all x ∈ X
Revised: ( Z | x ) is a dirac delta function for all x ∈ X

Original: and we denote by e ϕ ( x ) the deterministic encoder s.t.
Revised: and we denote by e ϕ ( x ) the deterministic encoder s.t.

Original: p ϕ ( e ϕ ( x ) |
Revised: p ϕ ( e ϕ ( x ) |

Original: x ) = 1 .
Revised: x ) = 1 .

Original: Why In the deterministic case, the MI bottleneck becomes the entropy bottleneck because I[ X ; Z ] =
Revised: Why In the deterministic case, the MI bottleneck becomes the entropy bottleneck because I[ X ; Z ] =

Original: H[ Z ] − H[ Z | X ] =
Revised: H[ Z ] − H[ Z | X ] =

Original: H[ Z ] , where we use the fact that H[ Z | X ] = 0 .
Revised: H[ Z ] , where we use the fact that H[ Z | X ] = 0 .

Original: Importantly, considering only deterministic encoders does not constrain our ability to learning optimal encoders.
Revised: Importantly, considering only deterministic encoders does not constrain our ability to learning optimal encoders.

Original: Indeed, just as with the MI bottleneck optimizing the objective with the entropy bottleneck under I[ A ; Z ] = I
Revised: Indeed, just as with the MI bottleneck optimizing the objective with the entropy bottleneck under I[ A ; Z ] = I

Original: [ A ; X ] will recover encoders s.t.
Revised: [ A ; X ] will recover encoders s.t.

Original: e ϕ ( x (cid:48) )
Revised: e ϕ ( x (cid:48) )

Original: =
Revised: =

Original: e ϕ ( x ) ⇐⇒ M ( x ) =
Revised: e ϕ ( x ) ⇐⇒ M ( x ) =

Original: M ( x (cid:48) ) , which also satisﬁes the support match constraint as discussed before.
Revised: M ( x (cid:48) ) , which also satisﬁes the support match constraint as discussed before.

Original: How Using the same derivation as the MI bottleneck, we can derive the variational upper bound on entropy
Revised: How Using the same derivation as the MI bottleneck, we can derive the variational upper bound on entropy

Original: [Equation]
Revised: [Equation]

Original: which is the standard variational bound used in neural compression (Ballé et al., 2016; Theis et al., 2017).
Revised: which is the standard variational bound used in neural compression (Ballé et al., 2016; Theis et al., 2017).

Original: Putting all together, we have
Revised: Putting all together, we have

Original: [Equation]
Revised: [Equation]

Original: which also recovers the optimal encoder with unconstrained variational families, inﬁnite samples, and λ > 1 as with the MI bottleneck.
Revised: which also recovers the optimal encoder with unconstrained variational families, inﬁnite samples, and λ > 1 as with the MI bottleneck.

Original: The detialed algorithm is provided in Algorithm 2.
Revised: The detialed algorithm is provided in Algorithm 2.

Original: Note that the discreteness of Z could lead to difﬁculty of gradient-based optimization, and we follow Ballé et al.
Revised: Note that the discreteness of Z could lead to difﬁculty of gradient-based optimization, and we follow Ballé et al.

Original: (2016) to add uniform noise to Z as a differentiable substitute for rounding during training.
Revised: (2016) to add uniform noise to Z as a differentiable substitute for rounding during training.

Original: In our experiments, we will mostly use the Ent bottleneck instead of the MI bottleneck to avoid introducing stochastic encoders.
Revised: In our experiments, we will mostly use the Ent bottleneck instead of the MI bottleneck to avoid introducing stochastic encoders.

Original: Algorithm 2 Ent objective
Revised: Algorithm 2 Ent objective

Revised: Require: e ϕ , s ψ , q θ

Original: [Equation]
Revised: [Equation]

Revised: X

Original: 7: return
Revised: [Equation] n 1: Z ← e ϕ [Equation] X )

Revised: [Equation]

Revised: : A

Revised: [Equation]

Revised: ← sample

Revised: [Equation]

Revised: p

Revised: [Equation]

Revised: 3: { ( X

Revised: [Equation]

Revised: sample (

Revised: [Equation]

Revised: 4: A

Revised: [Equation]

Revised: 5: L

Revised: [Equation]

Revised: ← −

Original: [Equation]
Revised: [Equation]

Original: C. C ONTRASTIVE ADVERSARIAL DOMAIN BOTTLENECK
Revised: 6: L [Equation] ← − log q

Original: [Equation]
Revised: [Equation]

Revised: 7: return L Ent =

Original:  ; D ]
Revised: L aug + λ L supp C. C ONTRASTIVE ADVERSARIAL DOMAIN BOTTLENECK B[ Z, X, Y, D ] = I[ Z ; D ]

Original: 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.
Revised: 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.

Original: Now we introduce a bottleneck B[ Z, X, Y, D ]
Revised: Now we introduce a bottleneck B[ Z, X, Y, D ]

Original: = I
Revised: = I

Original: [ Z ; D ] which we only seek to remove the information of Z about the domain D .
Revised: [ Z ; D ] which we only seek to remove the information of Z about the domain D .

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: For simplicity we consider the deterministic encoder e ϕ ( x ) as with the main body.
Revised: For simplicity we consider the deterministic encoder e ϕ ( x ) as with the main body.

Original: Why Similar to the previous analysis, we aim to show that arg min p Z
Revised: Why Similar to the previous analysis, we aim to show that arg min p Z

Original: | X I[ Z ; D ] under I
Revised: | X I[ Z ; D ] under I

Original: [ A ; Z ] =
Revised: [ A ; Z ] =

Original: I[ A ; X ] leads to the support match constraint.
Revised: I[ A ; X ] leads to the support match constraint.

Original: Using Eq.
Revised: Using Eq.

Original: (117) we have I[ Z ; D ] =
Revised: (116) we have I[ Z ; D ] =

Original: I[ Z, M ( X Z ); D ] =
Revised: I[ Z, M ( X Z ); D ] =

Original: I[ Z, M ( X ); D ] = I[ M ( X ); D ] + I
Revised: I[ Z, M ( X ); D ] = I[ M ( X ); D ] + I

Original: [ Z ; D | M ( X )] where the last equality uses the chain rule of mutual information.
Revised: [ Z ; D | M ( X )] where the last equality uses the chain rule of mutual information.

Original: Due to the non-negativity of (conditional) mutual information, we have that the minimum of I[ Z ; D ] under I[ A ; Z ] =
Revised: Due to the non-negativity of (conditional) mutual information, we have that the minimum of I[ Z ; D ] under I[ A ; Z ] =

Original: I
Revised: I

Original: [ A ; X ] is I
Revised: [ A ; X ] is I

Original: [ M ( X ); D ] .
Revised: [ M ( X ); D ] .

Original: Then we show the minimum is achievable by constructing the same optimal encoder e ϕ ( X ) as the Ent bottleneck which clearly satisﬁes I[ Z ; D | M ( X )] = 0 .
Revised: Then we show the minimum is achievable by constructing the same optimal encoder e ϕ ( X ) as the Ent bottleneck which clearly satisﬁes I[ Z ; D | M ( X )] = 0 .

Original: It is then easy to show that the support match constraint has to hold
Revised: It is then easy to show that the support match constraint has to hold

Original: when I[ Z ; D | M ( X )] = 0 by contrapositive.
Revised: when I[ Z ; D | M ( X )] = 0 by contrapositive.

Original: 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.
Revised: 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.

Original: How The typical way of minimizing I[ Z ; D ] is to derive the variational bound as 
Revised: How The typical way of minimizing I[ Z ; D ] is to derive the variational bound as [Equation] [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original: where a variational distribution (or domain classiﬁer) q φ is used to approximate p D | Z and jointly trained to maximize the bound.
Revised: where a variational distribution (or domain classiﬁer) q φ is used to approximate p D | Z and jointly trained to maximize the bound.

Original: This recovers the domain-adversarial training method as introduced in Ganin et al.
Revised: This recovers the domain-adversarial training method as introduced in Ganin et al.

Original: 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).
Revised: 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).

Original: We propose the contrastive adversarial domain (CAD) bottleneck, which is based on the above explicit version but uses a variational distribution q φ ( D | Z ) that is tied with other parts of the model , thus no need to learn a domain classiﬁer.
Revised: We propose the contrastive adversarial domain (CAD) bottleneck, which is based on the above explicit version but uses a variational distribution q φ ( D | Z ) that is tied with other parts of the model , thus no need to learn a domain classiﬁer.

Original: Speciﬁcally, as with Eq.
Revised: Speciﬁcally, as with Eq.

Original: (126), we ﬁrst introduce a contrastive variational distribution q ϕ, X
Revised: (125), we ﬁrst introduce a contrastive variational distribution q ϕ, X

Original: ( X | Z ) of p X | Z as:
Revised: ( X | Z ) of p X | Z as:

Original: [Equation]
Revised: [Equation]

Original: 138) where s ϕ  X, Z ) :=
Revised: (137) where s ϕ ( X, Z ) :

Original:  e ϕ ( X ) T Z is tied with the encoder e ϕ .
Revised: = e ϕ ( X ) T Z is tied with the encoder e ϕ .

Original: Since p D | Z can be rewritten as
Revised: Since p D | Z can be rewritten as

Original: E p X | Z (cid:2) p D | X (cid:3) using the fact that D − X − Z forms a Markov chain, we obtain can the following variational distribution:
Revised: E p X | Z (cid:2) p D | X (cid:3) using the fact that D − X − Z forms a Markov chain, we obtain can the following variational distribution:

Original: [Equation]
Revised: [Equation]

Original: (139) which recovers p D | Z when q ϕ, X =
Revised: (138) which recovers p D | Z when q ϕ, X =

Original: p X
Revised: p X

Original: | Z .
Revised: | Z .

Original: This is the case where s
Revised: This is the case where s

Original: [Equation]
Revised: [Equation]

Original: X, Z ) ∝ log p X,Z ( X, Z ) and inﬁnite samples n → ∞ .
Revised: X, Z ) ∝ log p X,Z ( X, Z ) and inﬁnite samples n → ∞ .

Original: Note that p D | X is still not available, we can use a count estimate ˆ p D , X in practice.
Revised: Note that p D | X is still not available, we can use a count estimate ˆ p D , X in practice.

Original: In particular, we obtain a collection D by taking each X (cid:48)
Revised: In particular, we obtain a collection D by taking each X (cid:48)

Original: ∈ X and independently sampling D (cid:48) from p D | X (cid:48) to get D :=
Revised: ∈ X and independently sampling D (cid:48) from p D | X (cid:48) to get D :=

Original: { D, D − 1 , . . .
Revised: { D, D − 1 , . . .

Original: , D − n } .
Revised: , D − n } .

Original: In other words, ( D, X ) and ( D − i , X − i ) for i ∈ [ n ] := { 1 , . . . , n } are all i.i.d.
Revised: In other words, ( D, X ) and ( D − i , X − i ) for i ∈ [ n ] := { 1 , . . . , n } are all i.i.d.

Original: sampled from p D,X .
Revised: sampled from p D,X .

Original: Then we use a count estimate
Revised: Then we use a count estimate

Original: [Equation]
Revised: [Equation]

Original: which is an accurate estimate with inﬁnite samples.
Revised: which is an accurate estimate with inﬁnite samples.

Original: This leads to our ﬁnal variational distribution:
Revised: This leads to our ﬁnal variational distribution:

Original: [Equation]
Revised: [Equation]

Original: Putting all together we get that the ﬁnal loss:
Revised: Putting all together we get that the ﬁnal loss:

Original: [Equation]
Revised: [Equation]

Original: In practice, ˆ p D , X
Revised: In practice, ˆ p D , X

Original: ( D | X ) is typically a dirac delta function since it is rare to have the same samples in a batch.
Revised: ( D | X ) is typically a dirac delta function since it is rare to have the same samples in a batch.

Original: Thus, in Eq.
Revised: Thus, in Eq.

Original: (141) we only need to sum q ϕ, X ( X (cid:48) | Z ) over those associated with the same domain label D as X , which reduces Eq.
Revised: (140) we only need to sum q ϕ, X ( X (cid:48) | Z ) over those associated with the same domain label D as X , which reduces Eq.

Original: (142) to Eq.
Revised: (141) to Eq.

Original: (8) described in the main body with a detailed algorithm in Algorithm 1.
Revised: (7) described in the main body with a detailed algorithm in Algorithm 1.

Original: Note that it is easy to generalize Algorithm 1 to parallel computation within a batch of samples.
Revised: Note that it is easy to generalize Algorithm 1 to parallel computation within a batch of samples.

Original: Indeed, for each sample in the batch, we can view all other samples in the batch as negatives and compute the loss efﬁciently in parallel.
Revised: Indeed, for each sample in the batch, we can view all other samples in the batch as negatives and compute the loss efﬁciently in parallel.

Original: C. 
Revised: C.

Original: [Equation]
Revised: [Equation]

Original: 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 ) .
Revised: 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 ) .

Original: However, since M ( X ) is typically not available in practice, we consider the special case where M ( X ) =
Revised: However, since M ( X ) is typically not available in practice, we consider the special case where M ( X ) =

Original: Y .
Revised: Y .

Original: In particular, this is the case where the labels are available and the supervised augmentations are used (see Fig. 2b).
Revised: In particular, this is the case where the labels are available and the supervised augmentations are used (see Fig. 2b).

Original: This reduces the bottleneck to B[ Z, X, Y, D ] =
Revised: This reduces the bottleneck to B[ Z, X, Y, D ] =

Original: I
Revised: I

Original: [ Z ; D | Y ] which is related to the conditional version of the domain-adversarial neural network (Li et al., 2018b).
Revised: [ Z ; D | Y ] which is related to the conditional version of the domain-adversarial neural network (Li et al., 2018b).

Original: 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
Revised: 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

Original: Y .
Revised: Y .

Original: In the following, we derive the conditional CAD (C 2 AD) bottleneck using a similar idea as CAD.
Revised: In the following, we derive the conditional CAD (C 2 AD) bottleneck using a similar idea as CAD.

Original: How In this case, we want to minimize
Revised: How In this case, we want to minimize

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: where q ( D | Z, Y ) is a variational distribution of p D | Z,Y .
Revised: where q ( D | Z, Y ) is a variational distribution of p D | Z,Y .

Original: Similar to the unconditional case, we also aim to use a non-parametric approximation that is tied with other parts of the model, and we obtain it using the fact p D | Z,Y =
Revised: Similar to the unconditional case, we also aim to use a non-parametric approximation that is tied with other parts of the model, and we obtain it using the fact p D | Z,Y =

Original: E p X | Z,Y (cid:2)
Revised: E p X | Z,Y (cid:2)

Original: p D | X (cid:3) .
Revised: p D | X (cid:3) .

Original: Speciﬁcally, let Y :=
Revised: Speciﬁcally, let Y :=

Original: { Y, Y − 1 , . . .
Revised: { Y, Y − 1 , . . .

Original: , Y − n } be the collection of labels obtained by independently sampling the label from p Y | X (cid:48) for each X (cid:48)
Revised: , Y − n } be the collection of labels obtained by independently sampling the label from p Y | X (cid:48) for each X (cid:48)

Original: ∈ X .
Revised: ∈ X .

Original: We collect samples associated with the label Y , i.e., X Y :=
Revised: We collect samples associated with the label Y , i.e., X Y :=

Original: { X } ∪{ X − i | Y − i
Revised: { X } ∪{ X − i | Y − i

Original: =
Revised: =

Original: Y, i ∈ [ n ] } and obtain a variational distribution of p
Revised: Y, i ∈ [ n ] } and obtain a variational distribution of p

Original: [Equation]
Revised: [Equation]

Original: where we use the same critic s ϕ ( X, Z ) :
Revised: where we use the same critic s ϕ ( X, Z ) :

Original: = e ϕ ( X ) T Z that is tied with the encoder e ϕ as before, but only take softmax over those samples with the same label Y .
Revised: = e ϕ ( X ) T Z that is tied with the encoder e ϕ as before, but only take softmax over those samples with the same label Y .

Original: For the term p D | X , we use the same count estimate ˆ p D , X in Eq.
Revised: For the term p D | X , we use the same count estimate ˆ p D , X in Eq.

Original: Then we obtain the variational distribution of p D | Z,Y :
Revised: Then we obtain the variational distribution of p D | Z,Y :

Original: [Equation]
Revised: [Equation]

Original: Putting all together we get that the ﬁnal loss:
Revised: Putting all together we get that the ﬁnal loss:

Original: [Equation]
Revised: [Equation]

Original: Again, since in practice ˆ p D , X ( D | X ) is typically a dirac delta function, the summation in Eq.
Revised: Again, since in practice ˆ p D , X ( D | X ) is typically a dirac delta function, the summation in Eq.

Original: (147) can be done only over those associated with the same label Y and the same domain label D as X .
Revised: (146) can be done only over those associated with the same label Y and the same domain label D as X .

Original: In particular, we collect samples X Y,D := { X } ∪{ X − i | Y − i
Revised: In particular, we collect samples X Y,D := { X } ∪{ X − i | Y − i

Original: =
Revised: =

Original: Y, D − i
Revised: Y, D − i

Original: =
Revised: =

Original: D, i ∈
Revised: D, i ∈

Original: [ n ] } and obtained the simpliﬁed loss:
Revised: [ n ] } and obtained the simpliﬁed loss:

Original: [Equation]
Revised: [Equation]

Original: A detailed algorithm is in Algorithm 
Revised: A detailed algorithm is in Algorithm 3. Algorithm

Revised: Require: e ϕ , s ψ , D, X, Y, n 1: Z ← e

Original: [Equation]
Revised: [Equation]

Original: Require: e
Revised: X ) 2: A + ← sample

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: , D

Original: [Equation] [Equation][Equation]
Revised: [Equation] [Equation]

Original: ← sample

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation][Table] [Equation] [Table]
Revised: [Equation]

Original: [Equation]

Original: [Table][Equation] [Table] [Equation]
Revised: [Table]

Original: [Table]

Original: [Equation][Table] [Equation] [Table]
Revised: [Equation]

Original: [Equation] [Table] [Equation] [Table]
Revised: [Table]

Original: [Equation][Table] [Equation] [Table]
Revised: [Equation]

Original: [Equation]

Original: [Table]
Revised: [Table]

Original: C D E XPERIMENTAL D ETAILS D. S CIENTIFIC
Revised: D E XPERIMENTAL D ETAILS D. S CIENTIFIC

Original: In both the scientiﬁc setting and the following bridge setting, we consider rather unrealistic setups for verifying our theory where we have access to labels from all domains.
Revised: In both the scientiﬁc setting and the following bridge setting, we consider rather unrealistic setups for verifying our theory where we have access to labels from all domains.

Original: We can choose to directly minimize the risk R [ Y | Z ] with the cross-entropy loss (denoted as CE henceafter), or minimize
Revised: We can choose to directly minimize the risk R [ Y | Z ] with the cross-entropy loss (denoted as CE henceafter), or minimize

Original: H
Revised: H

Original: [ A | Z ] (i.e., maximize I
Revised: [ A | Z ] (i.e., maximize I

Original: [ A ; Z ] ) with supervised augmentations as in Fig.
Revised: [ A ; Z ] ) with supervised augmentations as in Fig.

Original: 2b detailed below.
Revised: 2b detailed below.

Original: Implementation of supervised augmentations When using supervised augmentations, for each sample we obtain its augmentations from within its label class across all domains.
Revised: Implementation of supervised augmentations When using supervised augmentations, for each sample we obtain its augmentations from within its label class across all domains.

Original: A constrastive loss with such augmentations will essentially reduce to the supervised contrast loss (SupCon, Khosla
Revised: A constrastive loss with such augmentations will essentially reduce to the supervised contrast loss (SupCon, Khosla

Original: [Figure]
Revised: [Figure]

Original: et al., 2020).
Revised: et al., 2020).

Original: 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.
Revised: 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.

Original: In Khosla et al.
Revised: In Khosla et al.

Original: , 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.
Revised: , 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.

Original: (3) in Khosla et al.
Revised: (3) in Khosla et al.

Original: (2020)) or outside (SupCon-Out, Eq.
Revised: (2020)) or outside (SupCon-Out, Eq.

Original: (2) in Khosla et al.
Revised: (2) in Khosla et al.

Original: the log.
Revised: the log.

Original: Though Khosla et al.
Revised: Though Khosla et al.

Original: (2020) chose SupCon-Out because it worked better than SupCon-
Revised: (2020) chose SupCon-Out because it worked better than SupCon-

Original: In, we hypothesized that this is because SupCon-Out has an implicit bottleneck effect.
Revised: In, we hypothesized that this is because SupCon-Out has an implicit bottleneck effect.

Original: 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.
Revised: 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.

Original: Since this might confound with the effect of our bottlenecks, we chose to use SupCon-
Revised: Since this might confound with the effect of our bottlenecks, we chose to use SupCon-

Original: In though it performed slightly worse in out initial experiments.
Revised: In though it performed slightly worse in out initial experiments.

Original: For the implementation of SupCon, we followed Khosla et al.
Revised: For the implementation of SupCon, we followed Khosla et al.

Original: (2020) except that no projection was used.
Revised: (2020) except that no projection was used.

Original: Speciﬁcally, the temperature was set to 0.1, and normalization was applied when computing the logits.
Revised: Speciﬁcally, the temperature was set to 0.1, and normalization was applied when computing the logits.

Original: In the scientiﬁc setting, we tried to simulate our theory to the greatest extent.
Revised: In the scientiﬁc setting, we tried to simulate our theory to the greatest extent.

Original: In particular, we had two special considerations as detailed below:
Revised: In particular, we had two special considerations as detailed below:

Original: 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.
Revised: 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.

Original: Speciﬁcally, we treated the dataset as the population distribution and used the same dataset for training the encoder and training/evaluating the predictor.
Revised: Speciﬁcally, we treated the dataset as the population distribution and used the same dataset for training the encoder and training/evaluating the predictor.

Original: The ResNetencoder was trained to 300 epochs without any regularization, using the Adam optimizer (Kingma &
Revised: The ResNetencoder was trained to 300 epochs without any regularization, using the Adam optimizer (Kingma &

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: As the sample weight increases, the target log likelihood (neg. risk) ﬁrst decreases and then increases.
Revised: As the sample weight increases, the target log likelihood (neg. risk) ﬁrst decreases and then increases.

Original: We hypothesized that the increasing trend was due to that the source performance was already not optimal (though not visible from the ﬁgure), thus we selected the weight close to the turning point and 10 −seemed to be reasonable for both CE-Base and SupCon-Base.
Revised: We hypothesized that the increasing trend was due to that the source performance was already not optimal (though not visible from the ﬁgure), thus we selected the weight close to the turning point and 10 − 5 seemed to be reasonable for both CE-Base and SupCon-Base.

Original: Although we did not adaptively select the sample weight for each setup due to the computational cost, the pre-speciﬁed sample turned out to be reasonable for all other losses and different lambda combinations.
Revised: Although we did not adaptively select the sample weight for each setup due to the computational cost, the pre-speciﬁed sample turned out to be reasonable for all other losses and different λ combinations.

Original: Furthermore, we also removed regularization when training the linear classiﬁer and initialized the linear weight i.i.d.
Revised: Furthermore, we also removed regularization when training the linear classiﬁer and initialized the linear weight i.i.d.

Original: from N (0 , 1) .
Revised: from N (0 , 1) .

Original: Next, we provide other experimental details for reproducibility:
Revised: Next, we provide other experimental details for reproducibility:

Original: Implementation of standard augmentations We followed SimCLR (Chen et al., 2020) for implementing standard image augmentations.
Revised: Implementation of standard augmentations We followed SimCLR (Chen et al., 2020) for implementing standard image augmentations.

Original: For a fair comparison between the cases when using standard augmentations (SimCLR) and supervised augmentations (SupCon), we kept the total batch size the same and also used the same conﬁgurations for computing the SupCon loss, i.e., temperature set to 0.1, no projection, and normalization applied.
Revised: For a fair comparison between the cases when using standard augmentations (SimCLR) and supervised augmentations (SupCon), we kept the total batch size the same and also used the same conﬁgurations for computing the SupCon loss, i.e., temperature set to 0.1, no projection, and normalization applied.

Original: Details of Fig.
Revised: Details of Fig.

Original: 3c In Fig.
Revised: 3c In Fig.

Original: 3c, we considered different choices of augmentations.
Revised: 3c, we considered different choices of augmentations.

Original: The ‘Standard‘ augmentation implementation is described above (Appx.
Revised: The ‘Standard‘ augmentation implementation is described above (Appx.

Original: D.1).
Revised: D.1).

Original: The ‘Supervised’ augmentation was essentially implemented using the SupCon loss as described in Appx.
Revised: The ‘Supervised’ augmentation was essentially implemented using the SupCon loss as described in Appx.

Original: D.1.
Revised: D.1.

Original: For other augmentations considered, we implemented them by dropout inter-domain supervised augmentations in SupCon.
Revised: For other augmentations considered, we implemented them by dropout inter-domain supervised augmentations in SupCon.

Original: Speciﬁcally, for each sample in the batch, we randomly masked the samples from different domains (i.e., both inter-domain positives and negatives) i.i.d.
Revised: Speciﬁcally, for each sample in the batch, we randomly masked the samples from different domains (i.e., both inter-domain positives and negatives) i.i.d.

Original: with the speciﬁed dropout probability, while samples within the same domain were always kept. ‘IntraDom’ and ‘ApproxDC’ correspond to dropout probability 1 and 0.9, respectively.
Revised: with the speciﬁed dropout probability, while samples within the same domain were always kept. ‘IntraDom’ and ‘ApproxDA’ correspond to dropout probability 1 and 0.9, respectively.

Original: ‘SingleDom’ were implemented by dropout all interdomain samples with probability 1 except for a ﬁxed domain (the ‘A’ domain of PACS in our case).
Revised: ‘SingleDom’ were implemented by dropout all interdomain samples with probability 1 except for a ﬁxed domain (the ‘A’ domain of PACS in our case).

Original: D. B RIDGE
Revised: D. B RIDGE

Original: In the bridge setting (see Appx.
Revised: In the bridge setting (see Appx.

Original: E.2), we aimed to bridge the gap between our theoretical setup to the practical setup.
Revised: E.2), we aimed to bridge the gap between our theoretical setup to the practical setup.

Original: The main differences from the scientiﬁc setups are that the empirical generalization gap is considered and the average-case source predictor is used, as detailed below:
Revised: The main differences from the scientiﬁc setups are that the empirical generalization gap is considered and the average-case source predictor is used, as detailed below:

Original: Incorporating empirical generalization In practice, empirical-generalization gap should also be considered besides the source-target generalization gap.
Revised: Incorporating empirical generalization In practice, empirical-generalization gap should also be considered besides the source-target generalization gap.

Original: Thus, we randomly split the PACS dataset to 80% training and 20% validation splits for each domain.
Revised: Thus, we randomly split the PACS dataset to 80% training and 20% validation splits for each domain.

Original: 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.
Revised: 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.

Original: We used the ResNet-50 model as the encoder and initialized it from ImageNet pretrained model.
Revised: We used the ResNet-50 model as the encoder and initialized it from ImageNet pretrained model.

Original: 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.
Revised: 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.

Original: Using average-case source predictor Instead of approximating the worst-case source predictor in the scientiﬁc setting, we considered the average-case 2 source predictor which is closer to the common practice.
Revised: Using average-case source predictor Instead of approximating the worst-case source predictor in the scientiﬁc setting, we considered the average-case 2 source predictor which is closer to the common practice.

Original: Speciﬁcally, we freezed the encoder and trained a SVM classiﬁer with L2 regularization on the source training split.
Revised: Speciﬁcally, we freezed the encoder and trained a SVM classiﬁer with L2 regularization on the source training split.