Optimal Representations for Covariate Shift

Machine learning systems often experience a distribution shift between training and testing . In this paper , we introduce a simple variational objective whose optima are exactly the set of all representations on which risk minimizers are guaranteed to be robust to any distribution shift that preserves the Bayes predictor , e.g. , covariate shifts . Our objective has two components . First , a representation must remain discriminative for the task , i.e. , some predictor must be able to simultaneously minimize the source and target risk . Second , the representation ’ s marginal support needs to be the same across source and target . We make this practical by designing self-supervised learning methods that only use unlabelled data and augmentations to train robust representations . Our objectives achieve state-of-the-art results on DomainBed , and give insights into the robustness of recent methods , such as CLIP . 1 INTRODUCTION . It is hard to build machine learning ( ML ) systems that are robust to distribution shifts between a source ( train ) and target ( test ) domain . One promising approach to domain generalization ( DG ) is learning robust representations from which predictors trained on source must perform well on target . In practice , however , no current DG methods for learning representation uniformly outperform empirical source-risk minimizers ( ERM ) ( Gulrajani & Lopez-Paz , 2021 ) . Furthermore , our theoretical understanding of DG is still lacking . Specifically , while previous work have studied properties that would or would not imply robust representations ( Ben-David et al. , 2007 ; 2010a ; Zhao et al. , 2019 ; Johansson et al. , 2019 ) , the minimal set of achievable requirements for perfect DG is not yet known . We introduce the first , simple , variational objective whose optima are exactly the set of all representations on which source risk minimizers are guaranteed to generalize across distribution shifts that preserve the Bayes predictor . We work in an idealized DG ( IDG ) setting ; we assume that a learner has access to the source population risk . Our variational characterization implies that it is both sufficient and necessary for optimal IDG that a representation : ( a ) remains discriminative for the learning task , i.e. , there must exist predictors from the representation to the labels that can simultaneously minimize both source and target risk ; and ( b ) keeps the support of its marginal distribution invariant to shifts . This means that any optimal representation learning method must seek discriminative information about the target . Even worse , we prove that without access to some knowledge about the target , any representation learning algorithm can not uniformly ( over all target domains ) outperform a constant representation , which may explain why DG methods struggle to outperform ERM . We show , in theory and practice , how to overcome these challenges using only a large set of unlabeled examples and particular data augmentations that retain all discriminative information but minimal domain-specific information . Text descriptions of images are examples of such augmentations , as they are informative for many downstream classification tasks , but they remove a lot of domain-specific information . With such augmentations , we design practical self-supervised learning ( SSL ) objectives for learning robust representations . Our objectives give insights into the robustness of CLIP ( Radford et al. , 2021 ) , and lead to improved CLIP-based representations that achieve state-of-the-art ( SOTA ) results on DomainBed ( Gulrajani & Lopez-Paz , 2021 ) . To summarize , we : • provide minimal sufficient objectives whose optima achieve optimal DG under covariate shift ; • prove that it is impossible to learn useful representations without accessing target information ; • provide practical objectives to learn optimally robust representations using specific augmentations ; • get state-of-the-art results on typical domain generalization benchmarks . 2 BACKGROUND : DOMAIN GENERALIZATION AND REPRESENTATIONS . We are interested in predictions that are robust across distribution shifts . We formalize this using domain generalization ( DG ) language . Given a distribution pX , Y | ds over inputs x ∈ X and labels y ∈ Y from the source domain ds ∈ D , we select a predictor f : X → Γ . The predictions γ ∈ Γ could for example be labels or distributions over labels . Despite being selected on the source domain , we would like f to achieve a small expected risk with respect to a loss function ` : Y × Γ→ R≥0 , Rdf [ Y |X ] : = EpX , Y | d [ ` ( Y , f ( X ) ) ] , ( 1 ) on a distribution pX , Y | d from a target domain d = dt ∈ D , which is somehow related to ds . A common strategy for DG is to learn robust representations , which splits the problem into two . First , learn an encoder pZ |X , which maps inputs X to representations Z . Then , learn a predictor h : Z → Γ from representations Z to labels Y using standard risk minimization . The goal is to design a robust representation Z , so that predictors h trained to minimize the source risk Rdsh [ Y |Z ] also achieve low target risk Rdth [ Y |Z ] . Many methods have been proposed to try to learn such Z , e.g. , by enforcing domain invariance of the marginal pZ | d ( e.g. , Ganin et al. , 2016 ) . Still , many of these proposals are not sound ( Zhao et al. , 2019 ; Johansson et al. , 2019 ) . Furthermore , they rarely outperform source empirical risk minimization ( ERM ) in practice ( Gulrajani & Lopez-Paz , 2021 ) . 3 OPTIMAL REPRESENTATIONS FOR DOMAIN GENERALIZATION . To separate domain generalization from finite sample generalization , we consider an idealized DG ( IDG ) , where the predictor h is selected on the source population risk rather than empirical risk . We assume sample spaces X , Z , Y , D are discrete ; formal statements and proofs are in Appxs . A and B . 3.1 DEFINING OPTIMAL REPRESENTATIONS FOR IDEALIZED DOMAIN GENERALIZATION . We want to evaluate the quality of a representation Z of X . In our IDG , the learner is given a random source Ds ; she selects any source risk minimizer ; and is scored according to her risk on a random target domain Dt . To give uniform guarantees while reflecting the uncertainty over the source-target pair ( Ds , Dt ) , we measure the quality of Z as the expected risk of the learner ’ s worst-case choice . Definition . The idealized domain generalization risk ( IDG risk ) of an encoder pZ |X is the expected ( over domains ) worst-case ( over source risk minimizers ) target risk , i.e. , RIDG [ Y |Z ] : = EpDs , Dt [ sup h∈H∗Ds RDth [ Y |Z ] ] ( 2 ) whereH∗Ds : = arg minh R Ds h [ Y |Z ] are the source risk minimizers . We call a representation Z∗ ( or its encoder ) optimal for IDG if it minimizes the IDG risk : pZ∗ |X ∈ arg minpZ |X RIDG [ Y |Z ] . 3.2 CHARACTERIZING OPTIMAL REPRESENTATIONS FOR IDG UNDER COVARIATE SHIFT . The IDG risk is useful to evaluate representations , but gives few insights into IDG and is impractical to optimize due to the supremum in Eq . ( 2 ) . Under mild assumptions , we provide a simplified , equivalent objective , which is easier to optimize . For convenience , we assume that there is a unique Bayes predictor f∗ , which minimizes the expected risk over domains , i.e. , f∗ = arg minf EpDt [ R Dt f [ Y |X ] ] . This is satisfied by standard ML tasks pY , X and losses ` . More importantly , we assume the following domain structure , which ensures the existence of optimal encoders and allows our simplification . Assumptions . All domains d ∈ D we consider are related by the following assumptions : 1 . Generalized covariate shift . All domain-specific risk minimizers f ∈ arg minf [ Rdf [ Y |X ] ] are equal to the Bayes predictor f∗ on their support , i.e. , f ( x ) = f∗ ( x ) for all x ∈ supp ( pX | d ) . 2 . Invariance of Bayes predictions . The set of Bayes predictions is the same for all domains , i.e. , { f∗ ( x ) |x ∈ supp ( pX | d ) } = { f∗ ( x ) |x ∈ X } . Under review as a conference paper at ICLR 2022 < latexit sha1_base64= '' nv5kgG6Q4MLitfs23A69VllAD+g= '' > AAAB6nicdVC7SgNBFJ2NrxhfUUtBBoOQapldspoUYsDGMkHzkGQJs5NJMmT2wcysEJaUljYWithaW+c77PwGf8JJoqCiBy4czrmXe+71Is6kQujNSC0sLi2vpFcza+sbm1vZ7Z26DGNBaI2EPBRND0vKWUBriilOm5Gg2Pc4bXjDs6nfuKZCsjC4VKOIuj7uB6zHCFZaurg6QZ1sDpnILjoFGyLTdlDJKmniIKt0VICWiWbInb5Mqu83+5NKJ/va7oYk9mmgCMdStiwUKTfBQjHC6TjTjiWNMBniPm1pGmCfSjeZRR3DQ610YS8UugIFZ+r3iQT7Uo58T3f6WA3kb28q/uW1YtUrugkLoljRgMwX9WIOVQind8MuE5QoPtIEE8F0VkgGWGCi9Hcy+glfl8L/Sd02rYLpVFGunAdzpMEeOAB5YIFjUAbnoAJqgIA+uAX34MHgxp3xaDzNW1PG58wu+AHj+QO7hpHM < /latexit > < latexit sha1_base64= '' 8sdK5eqgt6vtuNYDok13axAjw7g= '' > AAAB6nicdVC7SgNBFJ31GeMrainIYBBSLbNxF00hBmwsEzQPSZYwO5lNhsw+mJkVwpLS0sZCEVtr63yHnd/gTzhJFFT0wIXDOfdyz71ezJlUCL0Zc/MLi0vLmZXs6tr6xmZua7suo0QQWiMRj0TTw5JyFtKaYorTZiwoDjxOG97gbOI3rqmQLAov1TCmboB7IfMZwUpLF1cnVieXR6aDrJJTgsgsOqhk25ogC6FDG1ommiJ/+jKuvt/sjSud3Gu7G5EkoKEiHEvZslCs3BQLxQino2w7kTTGZIB7tKVpiAMq3XQadQQPtNKFfiR0hQpO1e8TKQ6kHAae7gyw6svf3kT8y2slyj92UxbGiaIhmS3yEw5VBCd3wy4TlCg+1AQTwXRWSPpYYKL0d7L6CV+Xwv9JvWhatulUUb5cADNkwC7YBwVggSNQBuegAmqAgB64BffgweDGnfFoPM1a54zPmR3wA8bzB632kcI= < /latexit > < latexit sha1_base64= '' NzJ/0OJqTfPqJLiM8d9EWMJJFFU= '' > AAAB8HicbVDJSgNBEK1xjXGLy83LYBA8SJgRRY8BLx4jmEWTYejp9CRNunua7h4hjPkKLx5c8OrnePMX9CfsLAdNfFDweK+KqnqRZFQbz/t05uYXFpeWcyv51bX1jc3C1nZNJ6nCpIoTlqhGhDRhVJCqoYaRhlQE8YiRetS7GPr1O6I0TcS16UsScNQRNKYYGSvdyDC7vW+HZhAWil7JG8GdJf6EFMu73teL+JaVsPDRaic45UQYzJDWTd+TJsiQMhQzMsi3Uk0kwj3UIU1LBeJEB9no4IF7YJW2GyfKljDuSP09kSGudZ9HtpMj09XT3lD8z2umJj4PMipkaojA40VxylyTuMPv3TZVBBvWtwRhRe2tLu4ihbCxGeVtCP70y7OkdlzyT0qnVzaNIxgjB3uwD4fgwxmU4RIqUAUMHB7gCZ4d5Tw6r87buHXOmczswB847z9UpJRr < /latexit > < latexit sha1_base64= '' un3pkM+QCvsHMLw7+AWvfpJS2xM= '' > AAAB8HicbVDJSgNBEK1xjXGLy83LYBA8SJgRRY8BLx4jmEWTYejp9CRNunua7h4hjPkKLx5c8OrnePMX9CfsLAdNfFDweK+KqnqRZFQbz/t05uYXFpeWcyv51bX1jc3C1nZNJ6nCpIoTlqhGhDRhVJCqoYaRhlQE8YiRetS7GPr1O6I0TcS16UsScNQRNKYYGSvdyDC7vW+HehAWil7JG8GdJf6EFMu73teL+JaVsPDRaic45UQYzJDWTd+TJsiQMhQzMsi3Uk0kwj3UIU1LBeJEB9no4IF7YJW2GyfKljDuSP09kSGudZ9HtpMj09XT3lD8z2umJj4PMipkaojA40VxylyTuMPv3TZVBBvWtwRhRe2tLu4ihbCxGeVtCP70y7OkdlzyT0qnVzaNIxgjB3uwD4fgwxmU4RIqUAUMHB7gCZ4d5Tw6r87buHXOmczswB847z9TH5Rq < /latexit > Generalized covariate shift ( GCS ) ensures that f∗ is simultaneously optimal on all domains . For log-loss ` it recovers standard covariate shift , i.e. , pY | x , d = pY | x . For other losses , GCS is weaker , e.g. , it only requires invariance of most likely labels for 0-1 loss , and of conditional expectations for MSE . Invariance of Bayes predictors is necessary to learn useful predictors using a single domain . For example , for 0-1 loss it ensures that each label is seen at least once in each domain . The intuition behind our objective is that under GCS any source risk minimizer will make optimal predictions on target samples x that are also in the source . Thus , IDG optimal representations are exactly those that ( a ) have the same support in Z for all domain , and ( b ) retain GCS from Z without sacrificing the ability to predict Y , which can be ensured by minimizing the risk from Z . See Fig . 1 . Theorem 1 . Under our assumptions , an encoder pZ∗ |X is optimal for IDG if and only if it minimizes the risk R [ Y |Z ] : = infh EpDt [ RDth [ Y |Z ] ] while matching the support of Z across domains , i.e. , pZ∗ |X ∈ arg max pZ |X R [ Y |Z ] s.t . ∀ d ∈ D , supp ( pZ | d ) = supp ( pZ ) ( 3 ) Moreover , such encoders exist and their IDG risk is the Bayes risk RIDG [ Y |Z∗ ] = R [ Y |X ] . Theorem 1 provides an objective to learn representations on which performing risk minimization using a single domain and Z∗ is as good as performing risk minimization on the target domain from inputs X . Other sufficient conditions have previously been hinted towards , e.g. , matching the marginal pZ | d instead of its support ( e.g. , Ben-David et al. , 2010a ) which is the focus of most DG methods ( e.g. , Ganin et al. , 2016 ) . Previous conditions are nevertheless generally neither necessary nor achievable . To our knowledge , Thm . 1 is the first characterization of necessary and sufficient conditions for IDG optimal representations Z∗ . Theorem 1 thus gives better insights into IDG and provides a framework for deriving the least stringent objectives for optimal IDG . The risk minimization ( Eq . ( 3 ) ) shows that one must have some knowledge about the target domains to learn optimal representations for IDG . Access to targets might seem unrealistic , but without such knowledge or additional assumptions it is provably impossible to beat even constant representations . Proposition 1 ( No free lunch for IDG ) . Let ds be any source domain , Zds be any representation chosen on source ds , and C ∈ Z be a constant representation . Under minor assumptions , for every “ good ” target domain outside the source ’ s support on which Zds outperforms C for IDG , there are many “ bad ” target domains on which Zds is strictly worse than C. Formal statement in Appx . B.3 . Proposition 1 shows that target knowledge is necessary for learning useful representations in IDG . This may explain why previous DG methods have been unable to outperform ERM in standard benchmarks ( Gulrajani & Lopez-Paz , 2021 ) : the knowledge they have access to is insufficient to generalize . Taken together , Prop . 1 and Thm . 1 say that either you have access to target domains dt , in which case you can achieve an IDG risk that matches supervised learning , or you do not access dt , in which case any representation learning algorithm can achieve worse IDG risk than a constant .

Summary: At a high level, this work considers the problem of designing machine learning systems that generalize well even when the "domain" (i.e., the data-generating process) under which the system is tested (the "target") does not match that under which it was trained (the "source"). More specifically, the authors look to provide an answer to the question of what makes a "good" representation or encoding (a stochastic transformation of the original features), from the perspective of achieving a small risk (expected loss) on the target domain, assuming the learner only has data (and thus representations) from the source domain.  To this question, they provide one answer by showing that in a generalized covariance shift scenario, a natural notion of representation optimality (their "IDG optimality") is only possible if a representation minimizes the (best) risk on the average target distribution, while ensuring that the support of the representations is constant across domains. This is obviously a strong requirement, and they show that one cannot expect to satisfy it without domain knowledge going beyond the source.  They also show that there is some hope for special cases of domains in which we have augmentations that preserve label information and "cover" the support of the input distribution across all domains. In such a case, the authors show that their strong optimality requirement can be satisfied by designing the encoder to maximize the mutual information with the augmented data. Relaxing this objective leads to objectives that are more practical, and the authors complement their theoretical analysis with a rather in-depth set of empirical tests that evaluate the efficacy of their methodology for representation learning.