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

Original: Anonymous authors

Original: 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 all 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., some predictor must be able to 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 this practical by designing self-supervised learning methods that only use unlabelled data and augmentations to train robust representations.
Revised: We make this practical by designing self-supervised objectives that only use unlabelled data and augmentations to train robust representations.

Original: Our objectivesachieve state-of-the-art results on DomainBed, and give insights into the robustness of recent methods, such as CLIP.
Revised: Our objectives give insights into the robustness of CLIP and further improve CLIP’s representations to achieve SOTA results on DomainBed

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

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

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

Original: In practice, however, no current DG methods for learning representation uniformly outperform empirical source-risk minimizers (ERM) (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: Furthermore, our theoretical understanding of DG is still lacking.
Revised: Furthermore, our theoretical understanding of DG is still lacking.

Original: Speciﬁcally, while previous work have studied properties that would or would not imply robust representations (Ben-David et al., 2007; 2010a;
Revised: Speciﬁcally, while previous work have studied properties that would or would not imply robust representations (Ben-David et al., 2007; 2010a;

Original: Zhao et al., 2019;
Revised: Zhao et al., 2019;

Original: Johansson et al., 2019), the minimal set of achievable requirements for perfect DG is not yet known.
Revised: 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 all 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 target domains) 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, 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.
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, as 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 (SOTA) results on DomainBed (Gulrajani & Lopez-Paz, 2021).
Revised: Our objectives give insights into the robustness of CLIP (Radford et al., 2021) over other SSL methods, 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 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. 
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; • obtain SOTA results on typical domain generalization benchmarks. Our implementation is released at https://github.com/ryoungj/optdom .

Original:  B ACKGROUND : DOMAIN GENERALIZATION 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 ,

Original: [Equation]
Revised: [Equation]

Revised: X ))] ,

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, which splits the problem into two.
Revised: A common strategy for DG is to learn robust representations, which splits the problem into two.

Original: First, learn an encoder p Z | X , which 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 Z , so that predictors h trained to minimize the source risk R d sh
Revised: The goal is to design a robust representation Z , so that predictors h trained to minimize the source risk R d sh

Original: [ Y | Z ] also achieve low target risk R d th [ Y | Z ] .
Revised: [ Y | Z ] also achieve low target risk R d th [ Y | Z ] .

Original: 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).
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, they 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), where the predictor h is selected on 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: 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 want to evaluate the quality of a representation Z of X .
Revised: We want to evaluate the quality of a representation Z of X .

Original: 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 .
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 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.
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, and p D s ,D t is any joint distribu- tion that has full support over

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 [ Y | Z ] .
Revised: D × D . We call a representation Z ∗ (or its encoder) optimal for IDG if it minimizes the IDG risk.

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, i e., f ( x ) =
Revised: [R d f [ Y | X ]] are equal to the Bayes predictor f ∗ on their support, i.e., f ( x ) =

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

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) f ∗ ( x ) | x ∈ supp( p X | d ) (cid:9) =
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) =

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

Original: [Table]
Revised: [Table]

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

Original: [Figure]
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, GCS 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 on target samples x that are also in the source.
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, 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 .
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

Original: [Equation]
Revised: [Equation]

Original: E p
Revised: E p

Original: Dt (cid:2) R

Original: [Equation]
Revised: [Equation]

Original: th
Revised: t

Revised: [Equation]

Original: [ Y | Z ] (cid:3) while matching the support of Z across domains, i.e.,
Revised: [Equation] | Z ] [Equation] 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 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 .
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 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).
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: Previous conditions are nevertheless generally neither necessary nor achievable.
Revised: Note that previous conditions are nevertheless generally not necessary and could be too stringent to be achievable.

Original: To our knowledge, Thm. 1 is the ﬁrst characterization of necessary and sufﬁcient conditions for IDG optimal representations Z ∗ . Theorem 1 thus gives better insights into  IDG and provides a framework for deriving the least stringent objectives for optimal IDG.
Revised: To our knowledge, Thm. 1 is the ﬁrst characterization of necessary and sufﬁcient conditions, which gives better insights into the essential goal for optimal IDG and provides a guide for deriving the least stringent objectives in practice.

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 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.
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: 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.
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 - AGNOSTIC 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. 1) 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
Revised: We show how this can be overcome with self-supervised learning

Original: [Figure]
Revised: [Figure]

Original: (SSL), which is a technique for training representations without direct access to labels, and a particular class of data augmentations.
Revised: (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 is that augmentations retain task information.
Revised: The key requirement is that augmentations retain task information.

Original: Speciﬁcally, if any samples x, x (cid:48) ∈ X have the same augmentation conditional p
Revised: 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 such A , 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 all augmentation information, 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 suggests 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-agnostic assumption, which requires that the set of possible augmentation distributions is the same across domains, 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)

Original: [Equation]
Revised: [Equation]

Original: Proposition 2.
Revised: Proposition 2.

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

Original: X be a domain-agnostic 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 IDG optimal representations without labels if 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: How realistic are those augmentations?
Revised: How realistic are those augmentations?

Original: 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.
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-agnostic for typical domains (e.g. sketches and photos), since outputs A are correlated with the input’s domain D .
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 .

Original: See Fig.
Revised: See Fig.

Original: 2a.
Revised: 2a.

Original: 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.
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 downstream tasks.
Revised: First, text augmentations preserve label information for many downstream tasks.

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

Original: 2c).
Revised: 2c).

Original: This gives insights into the open question (Radford et al., 2021) about why CLIP’s representations are so robust compared to other SSL methods.
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: 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.
Revised: Many other multi-modal augmentations, e.g., audio-video (Wang et al., 2021), are also likely domain-agnostic and can be explored in practice.

Original: In practice, even the domain information D is usually unknown.
Revised: In practice, even the domain information D is usually unknown.

Original: One can nevertheless still optimize (Eq.
Revised: One can nevertheless still optimize (Eq.

Original: (4)) by replace the support constraint with a stronger one that does not rely on D e.g., minimizing
Revised: (4)) by replace the support constraint with a stronger one that does not rely on D e.g., minimizing

Original: I[ Z ; X ] (see Sec. 4.2.2), .
Revised: 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-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.
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

Original: [ A ; Z ]

Original: = H

Original: [ A ] − H

Original: [ A | Z ] as H

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

Original: Eq.
Revised: Eq.

Original: (5) 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 in 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: (5) 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 [ A | Z ] ≤ E p A,Z [ − log q
Revised: For both practical objectives we use a contrastive variational lower bound on I[ A ; Z ] based on

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

Original: Speciﬁcally, for a sample X , we construct X := { X, X − 1 , . . . , X − n } where each X − i are i.i.d. sampled from p  X .
Revised: Speciﬁcally, for a sample X , we ﬁrst obtain the augmented ‘positive’ A by sampling from p A | X .

Original: Then we obtain a collection A :
Revised: We then obtain n augmented ‘negatives’ (cid:8)

Original: = { A + , A − 1 , . . .
Revised: A − i (cid:9) n i =1 i.i.d.

Original: , A  −  n } of one positive augmentation A + and n negatives A − iby independently sampling an augmentation from p A | X (cid:48) for each example X
Revised: from the marginal p A by ﬁrst independently sampling X := (cid:8) X − i (cid:9) n i =1 from p X and then sampling A − i from p A | X − i . We denote A :=

Original: [Equation]
Revised: (cid:8) A, A − 1 , . . .

Original: ∈ X .
Revised: , A − n (cid:9) .

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., H [ A | Z ] ≤ E p X, A ,Z [ − log q ψ, A ( A + | Z )] where q [Equation]
Revised: A is to be positive, resulting in the following variational bound

Original: | Z

Original: [Equation]
Revised: [Equation]

Original: When A = X , one can tie thecritics’ 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 parameters of the critic and the encoder 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 valid 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: [Table]
Revised: [Table]

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]
Revised: [Equation]

Original: [Table]
Revised: [Table]

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]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

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

Original: [Table]
Revised: [Table]

Original: [Equation]

Original: [Table] [Equation] [Table] [Equation]
Revised: return

Original: [Table]
Revised: [Table]

Original: : L supp ← log (cid:80) X (cid:48)∈ X D exp e ϕ ( X (cid:48) )

Original: T Z (cid:80)

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

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

Original: T Z (cid:46) I

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

Original: L aug + λ L supp

Original: Our ﬁrst domain bottleneck minimizes B[ Z, D ] =
Revised: Our ﬁrst domain bottleneck minimizes B[ Z, D ] =

Original: I[ Z ; D ] , which enforces support match using a KL divergence.
Revised: 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;
Revised: However, DANN suffers from two issues: (i) it maximizes an upper bound on the desired term;

Original: (ii) it requires adversarial training, which is challenging in practice
Revised: (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 ( D | Z ) =

Original: ( D | Z ) = 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 q ϕ, X ( X | Z ) is a contrastive variational distribution of p X | Z constructed with samples X and a critic e ϕ ( X ) T Z tied with the encoder, ˆ p is a count estimate of p D | X .

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

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

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

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

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

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

Original: [ n ] } .

Original: Finally, compute q

Original: [Equation]
Revised: [Equation]

Original: X D [Equation] ϕ, X
Revised: n ] } . Ignoring constants, the ﬁnal loss is

Original: [Equation]
Revised: [Equation]

Original: X [Equation] | Z ) . The resulting criterion is
Revised: CAD

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 encoder’s determinism.
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 also matches the representations 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: The advantage of
Revised: The advantage of

Original: Ent is that it does not require domain samples D , which are rarely accessible in SSL.
Revised: 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),
Revised: We consider the standard variational bound used in neural compression (Ballé et al., 2016; Theis et al., 2017),

Original: H[ Z ] ≤ E p Z [ − log q θ ( Z )] , where an entropy model q θ ( Z ) is used.
Revised: 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 representations under covariate shift.
Revised: Provably robust representations under covariate shift.

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

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

Original: (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.
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: Mansour et al.
Revised: Mansour et al.

Original: (2009) generalizes these results from 0-1 loss to more general losses.
Revised: (2009) generalizes these results from 0-1 loss to more general losses.

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

Original: (2019) takes this further by deriving a support-based bound.
Revised: (2019) takes this further by deriving a support-based bound.

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

Original: However, these bounds can often be loose and the implied sufﬁcient conditions are neither necessary nor generally achievable.
Revised: However, these bounds can often be loose and the implied sufﬁcient conditions are neither necessary nor generally achievable.

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

Original: (2010b) suggests that separately minimizing R [ Y | Z ] or matching the marginal is not sufﬁcient, while Zhao 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) 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.
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 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. Still, our work is more speciﬁc than others, as we consider idealized DG and unrestricted predictors H .Our theory could be combined with Dubois et al.’s (2020), who provide conditions for optimal generalization from ﬁnite samples and constrained H in supervised learning.
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, Still, our work is more speciﬁc than others, as we consider idealized DG and unrestricted predictors H .

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 marginal distributions p Z | d and
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: p Z (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.
Revised: First, the representation learner (“the representor”) trains an encoder p Z | X using a speciﬁed objective and freezes it.

Original: Then, the person performing predictions (“the learner”) 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: E and F.

Original: 6. S CIENTIFIC SETTING : EXPLORING 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)
Revised: In particular, we used the PACS dataset (Li et al., 2017)

Original: Base
Revised: Base

Original: Ent
Revised: Ent

Original: CAD
Revised: CAD

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

Original: H
Revised: H

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

Original: [Figure]

Original: -3.5 ± 0. [Figure] [Figure] ± 
Revised: -3.5 ± 0.1 -0.0 ± 0.0 -0.8 ± 0.

Original: [Figure]{+(a) Effect of different objectives 10101010 
 L o g li k e li h oo d source target (b) Effect of λ+}
Revised: {+(a) Effect of different objectives 10101010 
 L o g li k e li h oo d source target (b) Effect of λ+}

Original: [Figure]
Revised: DA

Original: ± 0.
Revised: Approx.

Original: (a) Effect of different objectives
Revised: Non-DA Augmentation type - - - - -

Original: [Figure]
Revised: L o g li k e li h oo d Supervised SingleDom

Original: ApproxDA
Revised: ApproxDA

Original: [Figure]
Revised: IntraDom

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

Original: To approximate the worst-case source predictor, we followed Dubois et al.
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: F.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: F.2, we show that bottlenecks also helps in practical scenarios, i.e., non-idealized average-case DG evaluated with accuracy ( 95 . 9% → 96 . 7% ).

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-agnostic 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: (6) to standard supervised risk minimization
Revised: (6) to standard supervised risk minimization

Original: R [ Y | Z ] and used the domain-agnostic 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
Revised: The 1 st and 2 nd row of

Original: Fig.
Revised: 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-agnostic (DA) 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 DA.
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 DA.
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 ‘ApproxDA’, which is approximately DA 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-DA augmentations give terrible results compared to DA.
Revised: 3c shows that the non-DA augmentations give terrible results compared to DA.

Original: Interestingly, ‘ApproxDA’ also performs very well, which suggests that approximately DA 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-agnostic 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 with standard errors.
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.

Revised: Note that using ‘oracle selection’ is more consistent with our theory since it gets access to the necessary target information (for model selection), as discussed in Appx.

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: F.3. Due to space limit, we only included as baselines ‘ERM’ and ‘DomainBed SOTA’ which for each dataset is the best result over all baselines.

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: E.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/32 denoted ‘CLIP L’ and the smaller ResNet-50 denoted ‘CLIP S’.
Revised: The larger ViT-B/ denoted ‘CLIP L’ and the smaller ResNet-50 denoted ‘CLIP S’.

Revised: outperforms the previous SOTA result on DomainBed, as supported by our theoretical analysis.

Revised: Finetuning CLIP with our CAD bottleneck consistently improves the robustness of its representations and achieves SOTA performance

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: Thelast row in
Revised: The row ‘CLIP

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: L + CAD’ in 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 TerraIncognita is likely because CLIP’s dataset did not cover such images (camera traps monitoring animals).
Revised: Note that the poor performance on TerraIncognita is likely because CLIP’s dataset does not cover such images (camera traps monitoring animals).

Original: In Appx. E.2, we estimated the non-idealized DG performance of optimal representations on PACS (with access to all-domain labeled data) to be 96 .
Revised: The last row essentially shows an optimal representation, which we approximate by ﬁnetuning CLIP L with our CAD on all domains including the target.

Original: 7% , which is only 2% higher than CLIP L + CAD.
Revised: The gap between CLIP L + CAD and the upper-bound suggests that one can still learn better representations.

Original: This suggests that  our simple SSL encoder might already be close to optimal.
Revised: We hypothesize that end-to-end training of our objective would greatly shrink this gap.

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

Original: DomainBed’s baselines ﬁnetuned an ImageNet 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 CLIP S. Table 1 shows that
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: CLIP S + CAD still signiﬁcantly outperforms DomainBed baselines.
Revised: CLIP S + CAD still signiﬁcantly outperforms DomainBed baselines.

Original: Note that our theory does not constrain the encoder and so we expect larger encoders to be better as seen in Table 1.
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 “CLIP” rows of Table 1, we investigated the effect of ﬁnetuning CLIP 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 consistently improves results by around 1 ∼ 2% .

Original: These gains are due to the bottleneck, rather than ﬁnetuning 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’.

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

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

Original: We hypothesize that this is because SGD acts as an information bottleneck that naturally favors support match (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-agnostic 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 adaption of SSL models that do not satisfy those properties by ﬁnetuning
Revised: We investigated adaption of SSL models that do not satisfy those properties by ﬁnetuning

Original: DINO (Caron et al., 2021), the current SOTA on SSL ImageNet.
Revised: DINO (Caron et al., 2021), the current SOTA on SSL ImageNet.

Original: DINO is pretraiend 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.

Original: This supports our hypothesis that CLIP is much more robust than other SSL methods due to its domain-agnostic augmentations.
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 (Schuhmann et al., 2021) that 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: E.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 last 2 rows of Table 2, ﬁnetuning CLIP L on LAION with L Ent (Tuned w/ Ent) outperforms ﬁnetuning without bottleneck (Tuned w/o Ent) 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: F.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-agnostic 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.

Revised: Acknowledgement We would like to thank Elliot Creager, Roger Grosse, Elan Rosenfeld, Guodong

Revised: Zhang, Han Zhao, and anonymous reviewers for their helpful feedbacks and encouragements.

Revised: Resources used in preparing this research were provided, in part, by the Province of Ontario, the

Revised: Government of Canada through CIFAR, and companies sponsoring the Vector Institute.

Revised: We acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC),

Original:  ReproducibilityStatement For our theoretical results, we include formal assumptions, statements, and proofs in Appxs.
Revised: RGPIN-2021-03445. Reproducibility 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: E and have released our code at https://github.com/ryoungj/optdom . Appendix Table of Contents A Preliminaries A.1 Notation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.2 Deﬁnitions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: A P RELIMINARIES A. N OTATION
Revised: A.3 Assumptions . . . B Proofs B.1 Lemmas for general losses . . B.2 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Revised: B.3 Impossibility results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Revised: B.4 Augmentations . .

Revised: C

Revised: Practical objectives

Revised: C.1 Mutual information bottleneck B[ Z, X, Y, D

Revised: = I

Revised: [ Z ; X

Revised: C.2 Entropy bottleneck B[ Z, X, Y, D ] =

Revised: H

Revised: [ Z

Revised: C.3 Contrastive adversarial domain bottleneck B[ Z, X, Y, D ] = I

Revised: [ Z ; D

Revised: C.4 Conditional CAD B[ Z, X, Y, D ] = I

Revised: [ Z ; D | Y

Revised: D Extended Related

Revised: Work

Revised: E Experimental Details

Revised: E.1 Scientiﬁc . .

Revised: E.2 Bridge . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Revised: E.3 DomainBed . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Revised: E.4 LAION . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Revised: F Additional Experimental Results

Revised: F.1 Scientiﬁc . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Revised: F.2 Bridge . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Revised: F.3 DomainBed . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Revised: F.4 LAION . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

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 ≥ [Equation] ∪ {∞} :

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: (12) and (13), 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]

Revised: [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 (Losses admit optima) .

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

Original: | [Equation]
Revised: | Γ | > 1 .

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: Y | x has a unique mode.

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]
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 , Γ [Equation] = Γ

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: 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]
Revised: [Equation]

Original: [Equation]

Original: [Equation] [Equation] [Equation] 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 ] (32) =
Revised: R [ Y | Z ] (32) =

Original: E p X,Z p Y | X

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

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

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

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

Original: + (cid:88) ( x

Original: [Equation]

Original: z )

Original: [Equation]

Original: ( x (cid:48)

Original: [Equation]

Original: z (cid:48) )

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

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

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

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

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

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

Original: + (cid:88)

Original: [Equation]
Revised: [Equation]

Original: x

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: Eq.
Revised: Eq.

Original: (35) 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: (36) 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 Eq.
Revised: Under Assumptions 1 to 5, there exists an encoder p Z ∗ | X that is optimal for Eq.

Original: (3)and ?? , i.e.,
Revised: (3) , 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ﬁesR [ Y | X ]
Revised: Z ∗ satisﬁes

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

Original: Indeed,
Revised: Indeed,

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

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

Original: (43) 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) we 
Revised: Due to the data processing inequality of Bayes risk (Lemma 1) we also have R [ Y | X ] ≤

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

Original: that

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

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

Original: Recall that Γ ∗ = f ∗ → ( X ) and Γ ∗ d = f ∗ → (cid:0) supp( p 
Revised: Recall that Γ ∗ = f ∗ → ( X ) and Γ ∗ d = f ∗ → (cid:0) supp( p X | d ) (cid:1) .

Original: X

Original: |

Original: d

Original: (cid:1) .

Original: Now let us

Original: compute

Original: the

Original:  desired support for all d ∈ D 
Revised: Now let us compute the desired support for all d ∈ D :

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (45) 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 Eq.
Revised: Lemma 1), this implies that p Z ∗ | X is an optimal encoder for Eq.

Original: (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 ) ,

Original: [Equation]

Original: [Equation]
Revised: [Equation]

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

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]

Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (49) 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: (51) 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: [Equation] Assumption

Original: [Equation]
Revised: [Equation]

Revised: 55) and

Original: [Equation]
Revised: [Equation]

Original: where Eq.
Revised: where Eq.

Original: (58) 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 1 conditioned on d

Original: [Equation]

Original: [Equation]
Revised: [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: (61) 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: (64) 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]

Original: all d

Original: [Equation]
Revised: [Equation]

Original: D and that for alld
Revised: that for all

Original: [Equation]
Revised: [Equation]

Original: d t ∈ D

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, equiv.
Revised: Theorem 2 (Characterizing optimal representations for IDG, equiv.

Original: Theorem 1) .
Revised: 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., [Equation] [Equation]

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 ] :

Revised: [Equation]

Revised: [Equation]

Revised: Lemma

Original: [Equation]
Revised: [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: (67) and (68), 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: (67) and (68).
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]

Revised: constant w.r.t D

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: (79) 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 ) , and let h ∗ ∈ H ∗ Z .
Revised: P Z | d t ( S ) , and let h ∗ ∈ H ∗ Z .

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Original: Lem.
Revised: Lem.

Original: [Equation]
Revised: [Equation]

Revised: =

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

Revised: > R d t

Revised: [ Y | Z ]

Original: Lem.
Revised: Lem.

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Eq.
Revised: Eq.

Original: (85) 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 for learning representations for IDG, equiv.
Revised: Proposition 4 (No free lunch for learning representations for IDG, equiv.

Original: Proposition 1) .
Revised: 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
Revised: Let Rep : P ( X , Y ) → P ( Z|X ) be any algorithm for choosing an encoder p

Original: Z | X from the data distribution p X,Y , C be any constant r.v.
Revised: 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 b t such that p Z ds |

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

Original: [Equation]
Revised: [Equation]

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

Original: [ Y

Original: [Equation]

Original: 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]

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. 4, 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 [Equation] ds 
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 Z ds | x b ) ∩ supp( p Z | d s ) = ∅ , then there exist many target domains p X,Y | d t s.t.

Original: [Equation]

Original: 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. 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 .
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 Prop. 4 to the source inputs.
Revised: The rest of the target domain mass δ is distributed as in Prop. 4 to the source inputs.

Original: As in Prop. 4, 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: (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 ) = ∅ .
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: (97) 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: Proposition 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-agnostic augmentation) .
Revised: Assumption 9 (Domain-agnostic augmentation) .

Original: We assume that the augmentation A is domainagnostic , 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-agnostic 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-agnostic.
Revised: Not all augmentations are domain-agnostic.

Original: 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).
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.
Revised: We assume that the augmentation A is Bayespreserving , i.e., ∀ x, x (cid:48)

Original: [Equation]
Revised: [Equation]

Revised: X ,

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: 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 (Learning optimal representations without labels, equiv.
Revised: Proposition 6 (Learning optimal representations without labels, equiv.

Original: Proposition 2) .
Revised: Proposition 2) .

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

Original: X be an augmenter.
Revised: X be an augmenter.

Original: Under Assumptions 1 to 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: (109) 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: (108) and (109) 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

Revised: [Equation]

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 )
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: 113) 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]
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ﬁneh ∗ ( z ) = g
Revised: Let x z ∈ supp( p X | z ) be any input point that could have led to this representation z and deﬁne

Original: [Equation]
Revised: [Equation]

Original: By Eq.
Revised: By Eq.

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

Original: [Equation]
Revised: [Equation]

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Revised: ∗ ( X ))]

Revised: =

Revised: R [ Y | X ] .

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{+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 direct access to the labels and in that it can use augmentations under appropriate assumptions.

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

Original: H

Original: [ A | Z ]

Original: = H

Original: [ A | X ] (Eq.
Revised: [ A ; X ] .

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: (122) under Eq.
Revised: (122) under Eq.

Original: (123) 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 IDG.

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 bounds on MI.

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 lower bound which is based on InfoNCE (Oord et al., 2018).

Original: Namely, let X := (cid:8) X, X− 1 , . . . 
Revised: Namely, let X be the input sample and A be the ‘positive’ augmentation sampled from p A | X . We then obtain n ‘negative’ augmentations (cid:8)

Original: , X − n (cid:9) be a sequence of samples where each negative X − i is independently sampled from the marginal p X  .
Revised: A − i (cid:9) n i =1 by ﬁrst independently sampling (cid:8) X − i (cid:9) n i =1 from the marginal p X and then sampling A − i from p A | X − i . It is easy to see that the negatives A − i follow the marginal p A .

Original: Now let A : =
Revised: We construct A :=

Original: (cid:8) A + , A − 1 , . . 
Revised: (cid:8) 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 . In particular, the augmented positive A + has the marginal p A | X while the negatives A − i follow the marginal p A .
Revised: , A − n (cid:9) .

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: 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 havea variational distribution of p A | Z : [Equation] [Equation] which leads to the following variational bound:
Revised: Then we have the following variational lower bound (Poole et al., 2019):

Original: [Equation]
Revised: [Equation]

Original: with equality if the  variational families s ψ , p ϕare unconstrained and andwe use inﬁnite samples ( n → ∞ ).
Revised: In the case of unconstrained variational families s ψ , p ϕ and inﬁnite samples ( n → ∞ ), the above variational bound recovers I

Revised: [ A ; Z ] up to a constant (see Oord et al.

Revised: ; Dubois et al.

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 the encoder 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.

Revised: Throughout this section, the function M : X → N is the maximal invariant deﬁned in Prop. 6 via the equivalence relation deﬁned in Eq.

Original: C. M UTUAL INFORMATION BOTTLENECK B[ Z,X, Y, D ] = I[ Z ; X ]
Revised: C. M UTUAL INFORMATION BOTTLENECK B[ Z,

Revised: [Equation]

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: (122), 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-agnostic 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., I[ Z ; X [Equation] p X,

Original: [Equation]
Revised: [Equation]

Revised: (cid:20)

Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Revised: q

Revised: [Equation]

Revised: (cid:21) (128) =

Original:  where a variational distribution q θ is used to approximate p Z and is jointly optimized with p ϕ to minimize the bound.
Revised: [Equation] 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: Ignoring the constant, 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 n → ∞ , 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.

Revised: e ϕ ( x ) =

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

Original: =
Revised: [Equation] M

Original: e ϕ ( x ) ⇐⇒ M ( x ) = 
Revised: [Equation] x ) = M [Equation]

Original: M ( x (cid:48) ) , which also satisﬁes the support match constraint as discussed before.
Revised: 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: Algorithm2 Ent objective Require: e ϕ , s ψ , q θ [Equation] X
Revised: Algorithm

Original: [Equation]

Original: n 1: Z ← e ϕ

Original: [Equation]

Original: X )

Original: [Equation] : A [Equation] ← sample
Revised: objective

Original: [Equation] p [Equation] 3: { ( X 
Revised: Require: e ϕ , s ψ , q θ , X, n 1: Z ← e ϕ ( X ) 2: A ← sample ( p A | X ) 3: { ( X − i , A − i ) } ni =1 i.i.d.

Original: [Equation] sample ( [Equation] 4: A 
Revised: ←−− sample ( p X,A ) 4: A ←{ A } ∪{ A − i } ni = 5:

Original: [Equation] 5: L [Equation] ← − 
Revised: L aug ← − log exp s ψ ( A,Z ) (cid:80)

Original: [Equation] 6: L [Equation] ← − log q 
Revised: A (cid:48)∈ A exp s ψ ( A (cid:48) ,Z ) (cid:46) − I[ A ; Z ] 6: L supp ← − log q θ ( Z ) (cid:46) H

Original: [Equation] 7: return L Ent =
Revised: [ Z ] 7: return L Ent =

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

Original: C. C ONTRASTIVE ADVERSARIAL DOMAIN BOTTLENECK B[ Z, X, Y, D ] = I[ Z ; D ]
Revised: 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 removing 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: (116) 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

Original: [Equation]
Revised: [Equation]

Original: [Equation][Equation]
Revised: [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: (135) 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 in practice (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. (125), we ﬁrst introduce a contrastive variational distribution q ϕ, X
Revised: Suppose we have access to a set of inputs X , 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: (137) where s ϕ ( X, Z ) :
Revised: 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: Note that q ϕ, X has support over X , and equals p X | Z when s ϕ ( X, Z ) ∝ log p X,Z ( X, Z ) and X recovers X . In practice, we use a variety of crude approximations. Our ﬁrst crude approximation is that we use the minibatch of samples, i.e., the independently sampled (cid:8) X − i (cid:9) n i =1 as X .

Original:  E p X | Z (cid:2) p D | X (cid:3) using the fact that D − X − Z forms a Markov chain, we obtaincan the following variational distribution:
Revised: Now, since p D | Z can be rewritten as E p X | Z (cid:2) p D | X (cid:3) using the fact that D − X − Z forms a Markov chain, we obtain the following variational distribution:

Original: [Equation]
Revised: [Equation]

Original: (138) which recovers p D | Z when q ϕ, X=
Revised: which recovers p D | Z when q ϕ, X

Original:  p X | Z .
Revised: = p X | Z .

Original: This is the case where s

Original: [Equation] X, Z ) ∝ log p X,Z ( X, Z ) and inﬁnite samples n → ∞ . Note that p D | X is still not available, wecan use a count estimate ˆ p D , X in practice.
Revised: Note that p D | X is still not available. For our second crude approximation, we use a count estimate ˆ p D , X .

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 − i } ni =1 .

Original: , D − n } . 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 − i , X − i ) } ni =1 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 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 minibatch.

Original: Thus, in Eq.
Revised: Thus, in Eq.

Original: (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.
Revised: (139) we only need to sum q ϕ, X ( X (cid:48) | Z ) over those associated with the same domain label D as X , i.e., X D := { X − i | D − i = D, i ∈

Original: (141) to Eq.
Revised: [ n ] } where [ n ] := { 1 , . . . , n } . This leads to the simpliﬁed loss: [Equation]

Revised: In practice, we ﬁnd that the second term that minimizes the log probability leads to numerical instability.

Revised: Intuitively, this could be seen by the exploding gradient of the function log( p ) when p → 0 .

Revised: We thus replace it with − log(1 − p ) which has the same optima.

Revised: I.e. in practice we maximize the log of the probablity summed over X ¬ D :=

Revised: X \ X D .

Revised: This reduces Eq.

Revised: (141) to Eq.

Original: (7) 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.C.
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: [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: 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 be the label of input X sampled from p Y | X and Y :=

Original: {Y, Y − 1 , . . .
Revised: { 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 − 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 X | Z,Y : [Equation]

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: (146) 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 , i.e., X Y,D :=

Original: In particular, we collect samples X Y,D := { X} ∪{ X − i | Y − i
Revised: { 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 ] } . Similarly, instead of minimizing the log of the probability summed over X Y,D , we maximize the log of the probability summed over X Y, ¬ D : =

Original: [Equation]
Revised: X Y \ X Y,D = { X − i | Y − i =

Original: A detailed algorithm is in Algorithm 3.
Revised: Y, D − i (cid:54) = D, i ∈ [ n ] } .

Original: Algorithm Require: e ϕ , s ψ , D, X, Y, n 1: Z ← e
Revised: Finally we obtaine the simpliﬁed loss: [Equation] 

Original: [Equation]
Revised: [Equation]

Original: X ) 2: A + ← sample
Revised: 

Original: [Equation]
Revised:    .

Original: [Equation] [Equation] [Equation] [Equation]
Revised: A detailed algorithm is in Algorithm

Original: [Table]
Revised: [Table]

Original: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: Provably optimal representations. Many previous work have theoretically studied advantages of representations in various two-stage settings (representation learning followed by standard training of predictors) by bounding downstream performance (e.g., Ben-David et al., 2007; Shamir et al., 2010;

Original: [Table]
Revised: Saunshi et al., 2019).

Original: [Equation]
Revised: As learning theoretical bounds can be loose, it is hard to know whether they give the right insights into the problem. Our work instead proves the properties of optimal representations, which ensure best downstream performance.

Original: [Table] [Equation]
Revised: Those properties need to be approximated but give the right insights into what to aim for.

Original: [Table] [Equation]
Revised: This perspective and our proofs were inspired by Dubois et al.

Original: [Table]
Revised: (2020) who gives sufﬁcient conditions for optimal representations in supervised learning.

Original: D E XPERIMENTAL D ETAILS D. S CIENTIFIC
Revised: E E XPERIMENTAL D ETAILS E. 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[Figure] et al., 2020).
Revised: A constrastive loss with such augmentations will essentially reduce to the supervised contrast loss (SupCon, Khosla 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 − 5 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 λ 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 [Figure] 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: E.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: E.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 ‘ApproxDA’ 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).