Model-Based Domain Generalization

1 Introduction . Despite well-documented success in numerous applications [ 1–4 ] , the complex prediction rules learned by modern machine learning methods can fail catastrophically when presented with outof-distribution ( OOD ) data [ 5–9 ] . Indeed , rapidly growing bodies of work conclusively show that state-of-the-art methods are vulnerable to distributional shifts arising from spurious correlations [ 10–12 ] , adversarial attacks [ 13–17 ] , sub-populations [ 18–21 ] , and naturally-occurring variation [ 22–24 ] . This failure mode is particularly pernicious in safety-critical applications , wherein the shifts that arise in fields such as medical imaging [ 25–28 ] , autonomous driving [ 29–31 ] , and robotics [ 32–34 ] are known to lead to unsafe behavior . And while some progress has been made toward addressing these vulnerabilities , the inability of modern machine learning methods to generalize to OOD data is one of the most significant barriers to deployment in safety-critical applications [ 35 , 36 ] . In the last decade , the domain generalization community has emerged in an effort to improve the OOD performance of machine learning methods [ 37–40 ] . In this field , predictors are trained on data drawn from a family of related training domains and then evaluated on a distinct and unseen test domain . Although a variety of approaches have been proposed for this setting [ 41 , 42 ] , it was recently shown that that no existing domain generalization algorithm can significantly outperform empirical risk minimization ( ERM ) [ 43 ] over the training domains when ERM is properly tuned and equipped with state-of-the-art architectures [ 44 , 45 ] and data augmentation techniques [ 46 ] . Therefore , due to the prevalence of OOD data in safety critical applications , it is of the utmost importance that new algorithms be proposed which can improve the OOD performance of machine learning methods . In this paper , we introduce a new framework for domain generalization which we call Model-Based Domain Generalization ( MBDG ) . The key idea in our framework is to first learn transformations 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . that map data between domains and then to subsequently enforce invariance to these transformations . Under a general model of covariate shift and a novel notion of invariance to learned transformations , we use this framework to rigorously re-formulate the domain generalization problem as a semi-infinite constrained optimization problem . We then use this re-formulation to prove that a tight approximation of the domain generalization problem can be obtained by solving the empirical , parameterized dual for this semi-infinite problem . Finally , motivated by these theoretical insights , we propose a new algorithm for domain generalization ; extensive experimental evidence shows that our algorithm advances the state-of-the-art on a range of benchmarks by up to thirty percentage points . Contributions . Our contributions can be summarized as follows : • We propose a new framework for domain generalization in which invariance is enforced to underlying transformations of data which capture inter-domain variation . • Under a general model of covariate shift , we rigorously prove the equivalence of the domain generalization problem to a novel semi-infinite constrained statistical learning problem . • We derive data-dependent duality gap bounds for the empirical parameterized dual of this semi-infinite problem , proving that tight approximations of the domain generalization problem can be obtained by solving this dual problem under the covariate shift assumption . • We introduce a primal-dual style algorithm for domain generalization in which invariance is enforced over unsupervised generative models trained on data from the training domains . • We empirically show that our algorithm significantly outperforms state-of-the-art baselines on several standard benchmarks , including ColoredMNIST , Camelyon17-WILDS , and PACS . 2 Related work . Domain generalization . The rapid acceleration of domain generalization research has led to an abundance of principled algorithms , many of which distill knowledge from an array of disparate fields toward resolving OOD failure modes [ 47–50 ] . Among such works , one prominent thrust has been to learn predictors which have internal feature representations that are consistent across domains [ 51– 62 ] . This approach is also popular in the field of unsupervised domain adaptation [ 63–67 ] , wherein it is assumed that unlabeled data from the test domain is available during training [ 68–70 ] . Also related are works that seek to learn a kernel-based embedding of each domain in an underlying feature space [ 71 , 72 ] , and those that employ Model-Agnostic Meta Learning [ 73 ] to adapt to unseen domains [ 42 , 74–81 ] . Recently , another prominent direction has been to design weight-sharing [ 82–85 ] and instance re-weighting schemes [ 86–88 ] . Unlike any of these approaches , we explicitly enforce hard invariance-based constraints on the underlying statistical domain generalization problem . Data augmentation . Another approach to improve OOD performance is to augment the available training data . Among such methods , perhaps the most common is to leverage various forms of data augmentation [ 89–96 ] . Recently , several approaches have used style-transfer techniques and imageto-image translation networks [ 97–104 ] to augment the training domains with artificially-generated data [ 105–112 ] . Alternatively , rather than generating new data , [ 113–115 ] all remove textural features in the data to encourage domain invariance . Unlike the majority of these works , we do not perform data augmentation directly on the training objective ; rather , we derive a principled primal-dual style algorithm which enforces invariance constraints on data generated by unsupervised generative models . 3 Domain generalization . The domain generalization setting is characterized by a pair of random variables ( X , Y ) over instances x 2 X ✓ Rd and corresponding labels y 2 Y , where ( X , Y ) is jointly distributed according to an unknown probability distribution P ( X , Y ) . Ultimately , the objective in this setting is to learn a predictor f such that f ( X ) ⇡ Y , meaning that f should be able to predict the labels y of corresponding instances x for each ( x , y ) ⇠ P ( X , Y ) . However , unlike in standard supervised learning tasks , the domain generalization problem is complicated by the assumption that one can not sample directly from P ( X , Y ) . Rather , it is assumed that we can only measure ( X , Y ) under different environmental conditions , each of which corrupts or varies the data in a different way . For example , in medical imaging tasks , these environmental conditions might correspond to the imaging techniques and stain patterns used at different hospitals ; this is illustrated in Figure 1a . To formalize this notion of environmental variation , we assume that data is drawn from a set of environments or domains Eall ( see Figure 1b ) . Concretely , each domain e 2 Eall can be identified with a pair of random variables ( Xe , Y e ) , which together denote the observation of the random variable pair ( X , Y ) in environment e. Given samples from a finite subset Etrain ( Eall of domains , the goal of the domain generalization problem is to learn a predictor f that generalizes across all possible environments , implying that f ( X ) ⇡ Y . This can be summarized as follows : Problem 3.1 ( Domain generalization ) . Let Etrain ( Eall be a finite subset of training domains , and assume that for each e 2 Etrain , we have access to a dataset De : = { ( xej , yej ) } ne j=1 sampled i.i.d . from P ( Xe , Y e ) . Given a function class F and a loss function ` : Y ⇥ Y ! R 0 , our goal is to learn a predictor f 2 F using the data from the datasets De that minimizes the worst-case risk over the entire family of domains Eall . That is , we want to solve the following optimization problem : minimize f2F max e2Eall EP ( Xe , Y e ) ` ( f ( Xe ) , Y e ) . ( DG ) In essence , in Problem 3.1 we seek a predictor f 2 F that generalizes from the finite set of training domains Etrain to perform well on the set of all domains Eall . However , note that while the inner maximization in ( DG ) is over the set of all training domains Eall , by assumption we do not have access to data from any of the domains e 2 Eall\Etrain , making this problem challenging to solve . Indeed , as generalizing to arbitrary test domains is impossible [ 116 ] , further structure is often assumed on the topology of Eall and on the corresponding distributions P ( Xe , Y e ) . Disentangling the sources of variation across environments . The difficulty of a particular domain generalization task can be characterized by the extent to which the distribution of data in the unseen test domains Eall\Etrain resembles the distribution of data in the training domains Etrain . For instance , if the domains are assumed to be convex combinations of the training domains , as is often the case in multi-source domain generalization [ 117–119 ] , Problem 3.1 can be seen as an instance of distributionally robust optimization [ 120 ] . More generally , in a similar spirit to [ 116 ] , we identify two forms of variation across domains : covariate shift and concept shift . These shifts characterize the extent to which the marginal distributions over instances P ( Xe ) and the instance-conditional distributions P ( Y e|Xe ) differ between domains . We capture these shifts in the following definition : Definition 3.2 ( Covariate shift & concept shift ) . Problem 3.1 is said to experience covariate shift if environmental variation is due to differences between the set of marginal distributions over instances { P ( Xe ) } e2Eall . On the other hand , Problem 3.1 is said to experience concept shift if environmental variation is due to changes amongst the instance-conditional distributions { P ( Y e|Xe ) } e2Eall . The growing domain generalization literature encompasses a great deal of past work , wherein both of these shifts have been studied in various contexts [ 121–125 ] . Indeed , as this literature has grown , new benchmarks have been developed which span the gamut between covariate and concept shift [ 126 ] . However , a large-scale empirical study recently showed that no existing algorithm can significantly outperform ERM across these standard domain generalization benchmarks when ERM is carefully implemented [ 46 ] . As ERM is known to fail in the presence natural distribution shifts [ 127 ] , this result highlights the critical need for new algorithms that can go beyond ERM toward solving Problem 3.1 . 4 Model-based domain generalization . In what follows , we introduce a new framework for domain generalization that we call Model-Based Domain Generalization ( MBDG ) . In particular , we prove that when Problem 3.1 is characterized solely by covariate shift , then under a natural invariance-based condition , Problem 3.1 is equivalent to an infinite-dimensional constrained statistical learning problem , which forms the basis of MBDG . Formal assumptions for MBDG . In general , domain generalization tasks can be characterized by both covariate and concept shift . However , in this paper , we restrict the scope of our theoretical analysis to focus on problems in which inter-domain variation is due solely to covariate shift through an underlying model of data generation . Formally , we assume that the data in each domain e 2 Eall is generated from the underlying random variable pair ( X , Y ) via an unknown function G. Assumption 4.1 . Let e denote a Dirac distribution for e 2 Eall . We assume that there exists1 a measurable function G : X ⇥ Eall ! X , which we refer to as a domain transformation model , that parameterizes the inter-domain covariate shift via P ( Xe ) =d G # ( P ( X ) ⇥ e ) 8e 2 Eall , where # denotes the push-forward measure and =d denotes equality in distribution . Informally , this assumption specifies that there should exist a function G that relates the random variables X and Xe via X 7 ! G ( X , e ) = Xe . In past work , this particular setting in which the instances Xe measured in an environment e are related to the underlying random variable X has been referred to as domain shift [ 128 , §1.8 ] . In our medical imaging example , the domain shift captured by a domain transformation model would characterize the mapping from the underlying distribution P ( X ) over different cells to the distribution P ( Xe ) of images of these cells observed at a particular hospital ; this is illustrated in Figure 1c , wherein inter-domain variation is due to varying colors and stain patterns encountered at different hospitals . On the other hand , in this example example , the label y ⇠ Y describing whether a given cell contains a cancerous tumor should not depend on the lighting and stain patterns used at different hospitals . In this sense , while in other applications , e.g . the datasets introduced in [ 10 ] , the instance-conditional distributions can vary across domains , in this paper we assume that inter-domain variation is solely characterized by the domain shift due to G. Assumption 4.2 ( Domain shift ) . We assume that inter-domain variation is solely characterized by domain shift in the marginal distributions P ( Xe ) , as described in Assumption 4.1 . As a consequence , we assume that the instance-conditional distributions P ( Y e|Xe ) are stable across domains , meaning that Y e and Y are equivalent in distribution and that for each x 2 X and y 2 Y , it holds that P ( Y = y|X = x ) = P ( Y e = y|Xe = G ( x , e ) ) 8e 2 Eall . ( 1 ) Pulling back Problem 3.1 . The structure imposed on Problem 3.1 by Assumptions 4.1 and 4.2 provides a concrete way of parameterizing large families of distributional shifts in domain generalization problems . Indeed , the utility of these assumptions is that when taken together , they provide the basis for pulling-back Problem 3.1 onto the underlying distribution P ( X , Y ) via the domain transformation model G. This insight is captured in the following proposition : Proposition 4.3 . Under Assumptions 4.1 and 4.2 , Problem 3.1 is equivalent to minimize f2F max e2Eall EP ( X , Y ) ` ( f ( G ( X , e ) ) , Y ) . ( 2 ) The proof of this fact is a straightforward consequence of the decomposition P ( Xe , Y e ) = P ( Y e|Xe ) · P ( Xe ) in conjunction with Assumptions 4.1 and 4.2 ( see Appendix C.2 ) . Note that this result allows us to implicitly absorb each of the domain distributions P ( Xe , Y e ) into the domain transformation model . Thus , the outer expectation in ( 2 ) is defined over the underlying distribution P ( X , Y ) . On the other hand , just as in ( DG ) , this problem is still a challenging statistical min-max problem . To this end , we next introduce a new notion of invariance with respect to domain transformation models , which allows us to reformulate the problem in ( 2 ) as a semi-infinite constrained optimization problem . 1Crucially , although we assume the existence of a domain transformation model G , we emphasize that for many problems , it may be impossible to obtain or derive a simple analytic expression for G. This topic will be discussed at length in Section 6 and in Appendix G. A new notion of model-based invariance . Common to much of the domain generalization literature is the idea that predictors should be invariant to inter-domain changes . For instance , in [ 10 ] the authors seek to learn an equipredictive representation : X ! Z [ 129 ] , i.e . an intermediate representation that satisfies P ( Y e1 | ( Xe1 ) ) = P ( Y e2 | ( Xe2 ) ) 8e1 , e2 2 Eall . Despite compelling theoretical motivation for this approach , it has been shown that current algorithms which seek equipredictive representations do not significantly improve over ERM [ 130–133 ] . With this in mind and given the additional structure introduced in Assumptions 4.1 and 4.2 , we introduce a new definition of invariance with respect to the variation captured by the underlying domain transformation model G. Definition 4.4 ( G-invariance ) . Given a domain transformation model G , we say a classifier f is G-invariant if it holds for all e 2 Eall that f ( x ) = f ( G ( x , e ) ) almost surely when x ⇠ P ( X ) . Concretely , this definition says that a predictor f is G-invariant if environmental changes under G ( x , e ) can not change the prediction returned by f . Intuitively , this notion of invariance couples with the definition of domain shift , in the sense that we expect that a prediction should return the same prediction for any realization of data under G. Thus , whereas equipredictive representations are designed to enforce invariance of in an intermediate representation space Z , Definition 4.4 is designed to enforce invariance directly on the predictions made by f . In this way , in the setting of Figure 1 , G-invariance would imply that the predictor f would return the same label for a given cluster of cells regardless of the hospital at which these cells were imaged . Formulating the MBDG optimization problem . The G-invariance property described in the previous section is the key toward reformulating the min-max problem in ( 2 ) . Indeed , the following proposition follows from Assumptions 4.1 and 4.2 and from the definition of G-invariance . Proposition 4.5 . Under Assumptions 4.1 and 4.2 , if we restrict the domain F of Problem 3.1 to the set of G-invariant predictors , then Problem 3.1 is equivalent to the following constrained problem : P ? , minimize f2F R ( f ) , EP ( X , Y ) ` ( f ( X ) , Y ) ( MBDG ) subject to f ( x ) = f ( G ( x , e ) ) a.e . x ⇠ P ( X ) 8e 2 Eall . Here a.e . stands for “ almost everywhere ” and R ( f ) is the statistical risk of a predictor f with respect to the underlying random variable pair ( X , Y ) . Note that unlike ( 2 ) , ( MBDG ) is not a composite optimization problem , meaning that the inner maximization has been eliminated . In essence , the proof of Proposition 4.6 relies on the fact that G-invariance implies that predictions should not change across domains ( see Appendix C.2 ) . The optimization problem in ( MBDG ) forms the basis of our Model-Based Domain Generalization framework . To explicitly contrast this problem to Problem 3.1 , we introduce the following concrete problem formulation for Model-Based Domain Generalization . Problem 4.6 ( Model-Based Domain Generalization ) . As in Problem 3.1 , let Etrain ( Eall be a finite subset of training domains and assume that we have access to datasets De 8e 2 Etrain . Then under Assumptions 4.1 and 4.2 , the goal of Model-Based Domain Generalization is to use the data from the training datasets to solve the semi-infinite constrained optimization problem in ( MBDG ) . Problem 4.6 offers a principled perspective on Problem 3.1 when data varies WRT an underlying domain transformation model . However , just as solving the min-max problem of Problem 3.1 is known to be difficult , the problem in ( MBDG ) is also challenging to solve for several reasons : ( C1 ) The G-invariance constraint in ( MBDG ) is strict and thus challenging to enforce . ( C2 ) Problem 4.6 is a constrained problem over an infinite-dimensional functional space F . ( C3 ) We do have access to the set of all domains Eall or to the underlying distribution P ( X , Y ) . ( C4 ) We also generally do not have access to the underlying domain transformation model G. In the ensuing sections , we explicitly address each of these challenges toward developing a tractable method for approximately solving Problem 4.6 with guarantees on optimality . In particular , we discuss challenges ( C1 ) , ( C2 ) , and ( C3 ) in Section 5 . We then discuss ( C4 ) in Section F .

Summary: This paper proposes a framework for domain generalization which involves learning a transformation between domains and optimizing the worst case loss over these transformations. They formulate this as a min max optimization problem and eventually solve the dual of a relaxed version of the problem. Empirically, they find that their algorithm improves a lot for ColoredMNIST and Camelyon-WILDS especially, which are datasets where the domain transformations are somewhat simple (color shift). They get a modest gain on a satellite dataset FMoW-WILDS.