Model-Based Domain Generalization
Year: 2021

Abstract:
Despite remarkable success in a variety of applications, it is well-known that deep learning can fail catastrophically when presented with out-of-distribution data. Toward addressing this challenge, we consider the domain generalization problem, wherein predictors are trained using data drawn from a family of related training domains and then evaluated on a distinct and unseen test domain. We show that under a natural model of data generation and a concomitant invariance condition, the domain generalization problem is equivalent to an infinite-dimensional constrained statistical learning problem; this problem forms the basis of our approach, which we call Model-Based Domain Generalization. Due to the inherent challenges in solving constrained optimization problems in deep learning, we exploit nonconvex duality theory to develop unconstrained relaxations of this statistical problem with tight bounds on the duality gap. Based on this theoretical motivation, we propose a novel domain generalization algorithm with convergence guarantees. In our experiments, we report improvements of up to 30% over state-of-the-art domain generalization baselines on several benchmarks including ColoredMNIST, Camelyon17-WILDS, FMoW-WILDS, and PACS. Our code is publicly available at the following link: https://github.com/arobey1/mbdg.

heading: 1 Introduction
text: 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.

heading: 2 Related work
text: 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.

heading: 3 Domain generalization
text: 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 cannot 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.

heading: 4 Model-based domain generalization
text: 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) cannot 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.

heading: 5 Data-dependent duality gap for MBDG
text: In this section, we offer a principled analysis of Problem 4.6. In particular, we first address (C1) by introducing a tight relaxation of the G-invariance constraint. Next, to resolve the fundamental
difficulty involved in solving constrained statistical problems highlighted in (C2), we formulate the parameterized dual problem, which is unconstrained and thus more suitable for learning with deep neural networks. Finally, to address (C3), we introduce an empirical version of the parameterized dual problem and explicitly characterize the data-dependent duality gap between this problem and Problem 4.6. At a high level, this analysis results in an unconstrained optimization problem which is guaranteed to produce a solution that is close to the solution of Problem 3.1 (see Theorem 5.3). In this section, we have chosen to present our results somewhat informally by deferring preliminary results, regularity assumptions, and proofs to the appendices.
Addressing (C1) by relaxing the G-invariance constraint. One of the most fundamental challenges in solving Problem 4.6 is the difficulty of enforcing the G-invariance equality constraint. To alleviate some of this difficulty, we introduce the following relaxation of Problem 4.6:
P ?( ) , minimize f2F R(f) s.t. Le(f) , EP(X) d f(X), f(G(X, e))  8e 2 Eall (3)
where > 0 is a fixed margin the controls the extent to which we enforce G-invariance and d : P(Y)⇥P(Y)! R 0 is a distance metric over the space of probability distributions on Y . While at first glance this problem may appear to be a significant relaxation of the MBDG optimization problem in (MBDG), when = 0 and under mild conditions on d, the two problems are equivalent in the sense that P ?(0) = P ? (see Proposition B.1). Indeed, we note that the conditions we require on d are not restrictive, and include the KL-divergence and more generally the family of f -divergences. Moreover, when the margin is strictly larger than zero, under the assumption that the perturbation function P ?( ) is L-Lipschitz continuous, we show in Remark B.2 that |P ? P ?( )|  L , meaning that the gap between the problems is relatively small when is chosen to be small. In particular, when strong duality holds for (MBDG), this Lipschitz constant L is equal to the L1 norm of the optimal dual variable for (MBDG) (see Remark B.4).
Addressing (C2) by formulating the parameterized dual problem. As written, the relaxation in (3) is an infinite-dimensional constrained optimization problem over a functional space F (e.g. L 2 or the space of continuous functions). Optimization in this infinite-dimensional function space is not tractable, and thus we follow the standard convention by leveraging a finite-dimensional parameterization of F , such as the class of deep neural networks [134, 135]. The approximation power of such a parameterization can be captured in the following definition: Definition 5.1 (✏-parameterization). Let H ✓ Rp be a finite-dimensional parameter space. For ✏ > 0, a function ' : H⇥ X ! Y is said to be an ✏-parameterization of F if it holds that for each f 2 F , there exists a parameter ✓ 2 H such that EP(X) k'(✓, x) f(x)k1  ✏.
The benefit of using such a parameterization is that optimization is generally more tractable in the parameterized space A✏ := {'(✓, ·) : ✓ 2 H} ✓ F . However, typical parameterizations often lead to nonconvex problems, wherein methods such as SGD cannot guarantee constraint satisfaction. And while several heuristic algorithms have been designed to enforce constraints over common parametric classes [136–141], these approaches cannot provide guarantees on the underlying statistical problem of interest [142]. Thus, to provide guarantees on the underlying statistical problem in Problem 4.6, given an ✏-parameterization ' of F , we consider the following saddle-point problem:
D ? ✏ ( ) , maximize
2P(Eall) min ✓2H R(✓) +
Z
Eall
[Le(✓) ] d (e). (4)
where P(Eall) is the space of normalized probability distributions over Eall and 2 P(Eall) is the (semi-infinite) dual variable. Here we have slightly abused notation to write R(✓) = R('(✓, ·)) and Le(✓) = Le('(✓, ·)). One can think of (4) as the dual problem to (3) solved over the parametric space A✏. Notice that unlike Problem 4.6, the problem in (4) is unconstrained, making it much more amenable for optimization over the class of deep neural networks. Moreover, under mild conditions, the optimality gap between (3) and (4) can be explicitly bounded as follows: Proposition 5.2 (Parameterization gap). Let > 0 be given. Under mild regularity assumptions (see Assumption C.1 in Appendix C.3) on ` and d, there exists a small universal constant k such that
P ?( )  D?✏ ( )  P ?( ) + ✏k
⇣ 1 +
?pert L1
⌘ , (5)
where ?pert is the optimal dual variable for a perturbed version of (3) in which the constraints are tightened to hold with margin k✏.
Algorithm 1 Model-Based Domain Generalization (MBDG) 1: Hyperparameters: Primal step size ⌘p > 0, dual step size ⌘d 0, margin > 0 2: repeat 3: for minibatch {(xj , yj)}mj=1 in training dataset do 4: x̃j GENERATEIMAGE(xj) 8j 2 [m] . Generate model-based images 5: distReg(✓) (1/m) Pm j=1 d('(✓, xj),'(✓, x̃j)) . Calculate distance regularizer
6: loss(✓) (1/m) Pm
j=1 ` (xj , yj ;'(✓, ·)) . Calculate classification loss 7: ✓ ✓ ⌘pr✓[ loss(✓) + · distReg(✓) ] . Primal step for ✓ 8: [ + ⌘d (distReg(✓) )]+ . Dual step for 9: end for 10: until convergence 11: 12: procedure GENERATEIMAGE(x) 13: Sample e ⇠ N (0, I) . e is a latent code for MUNIT 14: return G(x, e) . Return image produced by MUNIT 15: end procedure
In this way, solving the parameterized dual problem in (4) provides a solution that can be used to recover a close approximation of the primal problem in (3). To see this, observe that Prop. 5.2 implies that |D?✏ ( ) P ?( )|  ✏k(1 + || ?pert||L1). This tells us that the gap between P ?( ) and D?✏ ( ) is small when we use a tight ✏-parameterization of F . Addressing (C3) by bounding the empirical duality gap. The parameterized dual problem in (4) gives us a principled way to address Problem 4.6 in the context of deep learning. However, complicating matters is the fact that we do not have access to the full distribution P(X,Y ) or to data from any of the domains in Eall\Etrain. In practice, it is ubiquitous to solve optimization problems such as (4) over a finite sample of N data points drawn from P(X,Y )2. More specifically, given {(xj , yj)}Nj=1 drawn i.i.d. according to (X,Y ), we consider the empirical counterpart of (4):
D ? ✏,N,Etrain( ) , maximize (e) 0, e2Etrain min ✓2H ⇤̂(✓, ) , R̂(✓) + 1|Etrain| X
e2Etrain
h L̂e(✓)
i (e) (6)
where R̂(✓) and L̂e(✓) are the empirical counterparts of R(f) and L(f). Notably, the duality gap between the solution to (6) and (MBDG) can be explicitly bounded as follows. Theorem 5.3 (Data-dependent duality gap). Let ✏ > 0 be given, and let ' be an ✏-parameterization of F . Under mild regularity assumptions on ` and d and assuming that A✏ has finite VC-dimension, with probability 1 over the N samples from P(X,Y ) we have that
|P ? D?✏,N,Etrain( )|  L + ✏k ⇣ 1 + ?pert L1 ⌘ +O ⇣p log(N)/N ⌘ (7)
where L is the Lipschitz constant of P ?( ) and k and ?pert are as defined in Proposition 5.2.
The key message to take away from Theorem 5.3 is that given samples from P(X,Y ), the duality gap incurred by solving the empirical problem in (6) is small when (a) the G-invariance margin is small, (b) the parametric space A✏ is a close approximation of F , and (c) we have access to sufficiently many samples. Thus, assuming that Assumptions 4.1 and 4.2 hold, the solution to Problem 3.1 is closely-approximated by the solution to the empirical, parameterized dual problem in (6).

heading: 6 MBDG: A principled algorithm for domain generalization
text: Motivated by these theoretical insights, we now introduce a new domain generalization algorithm which is widely applicable to problems with or without covariate shift. Our algorithm consists of two steps. First, we learn an approximation of the underlying domain transformation model G(x, e) using the data from the training domains Etrain. Next, we leverage G toward solving the unconstrained dual optimization problem in (6) via a primal-dual iteration.
2Indeed, in practice we do not have access to any samples from P(X,Y ). In Section 6, we argue that the N samples from P(X,Y ) can be replaced by the
P e2Etrain ne samples drawn from the training datasets D e.
Learning domain transformation models from data. Regarding challenge (C4), critical to our approach is having access to the domain transformation model G. For the vast majority of settings, the underlying function G(x, e) is not known a priori and cannot be represented by a simple expression. For example, obtaining a closed-form expression for a model that captures the variation in coloration, brightness, and contrast in the dataset shown in Figure 1 would be very challenging. While in general it is impossible to learn the true underlying domain transformation model when one only has access to data from the training domains, we argue that a realistic approximation of the underlying model can be learned from this data. To this end, to learn a domain transformation model, we train multimodal image-to-image translation networks on the training data. These networks are designed to transform samples from one dataset so that they resemble a diverse collection of images from another dataset. In particular, in each of the experiments in Section 7, we use the MUNIT architecture introduced in [102] to parameterize learned domain transformation models. As shown in Figure 5 and in Appendix G, models trained using the MUNIT architecture learn accurate and diverse transformations of the training data, which often generalize to generate images from new domains.
Primal-dual iteration. Given a learned approximation G(x, e) of the underlying domain transformation model, the next step in our procedure is to use a primal-dual iteration [143] toward solving (6) using the datasets De. We note that while our theory calls for data drawn from P(X,Y ), the G-invariance condition implies that when (6) is feasible, '(✓, x) ⇡ '(✓, xe) when x ⇠ P(X), x e ⇠ Pe(X), and xe = G(x, e). Therefore, the data from [e2EtrainDe is a useful proxy for data drawn from P(X,Y ). As the outer maximization in (6) is a linear program in , the primal-dual iteration can be characterized by alternating between the following steps:
✓ (t+1) 2 ⇢- argmin ✓2H ⇤̂(✓, (t)) (8) (t+1)(e)
h (t)(e) + ⌘ ⇣ L̂e(✓)
⌘i
+ (9)
Here [·]+ = max{0, ·}, ⌘ is the dual step size, and ⇢- argmin denotes a solution that is ⇢-close to being a minimizer, i.e. we should that have ⇤̂(✓(t+1), (t))  min✓2H ⇤̂(✓, (t)) + ⇢. We call (8) the primal step, and we call (9) the dual step. Furthermore, it can be shown that if this iteration is run for sufficiently many steps and with small enough step size, the iteration convergences with high probability to a solution which closely approximates the solution to Problem 4.6. Theorem 6.1 (Primal-dual convergence). Assuming that ` and d are [0, B]-bounded, H has finite VC-dimension, and under mild regularity conditions on (6), the primal-dual pair (✓(T ), (T )) obtained after running the alternating primal-dual iteration in (8) and (9) for T steps with step size ⌘, where
T =
⇠ 1
2⌘ ⇡ + 1 and ⌘  2|Etrain|B2
(10)
satisfies |P ? ⇤̂(✓(T ), µ(T ))|  K(⇢,, ✏) +O( p
log(N)/N). Here  is a constant that captures the regularity of (6) (see Appendix C.6) and K(⇢,, ✏) is a small constant depending on ⇢, , and ✏.
This means that by solving the empirical dual problem for sufficiently many steps, we can reach a solution that is close to solving the Model-Based Domain Generalization problem in Problem 4.6.
Implementation of MBDG. In practice, because (a) it may not be tractable to find a ⇢-minimizer over H at each iteration and (b) there may be a large number of domains in Etrain, we propose two modifications of the primal-dual iteration in which we replace (8) with a stochastic gradient step and we use only one dual variable for all of the domains; we call this algorithm MBDG (see Algorithm 1). We provide results in Appendix E where one dual variable is used per training domain.

heading: 7 Experiments
text: We now evaluate the performance of MBDG on a range of standard domain generalization benchmarks. In the main text, we present results on ColoredMNIST, Camelyon17-WILDS, FMoW-WILDS, and PACS; we defer results on VLCS to the supplemental. For ColoredMNIST, PACS, and VLCS, we use the DomainBed package [46], facilitating comparison to a range of baselines. Model selection for each of these datasets was performed using hold-one-out cross-validation. For Camelyon17-WILDS and FMoW-WILDS, we use the repository provided with the WILDS dataset suite, and we perform model-selection using the out-of-distribution validation set provided in the WILDS repository. Further details concerning hyperparameter tuning and model selection are deferred to Appendix E.
We first consider the ColoredMNIST dataset [10], which is a standard domain generalization benchmark created by colorizing subsets of the MNIST dataset [144]. This dataset contains three domains, each of which is characterized by a different level of correlation between the label and digit color. As shown in Table 1, the MBDG algorithm improves over each baseline by nearly 30%. To understand the reasons behind this improvement, we consider three ablation studies on ColoredMNIST.
Tracking the dual variables. For the three MBDG classifiers in Table 1, we plot the regularization term distReg(✓) and the corresponding dual variable at each training step in Figure 2. Observe that for the +90% and +80% domains, the dual variables decay to zero, as the constraint is satisfied early on in training.
On the other hand, the constraint for the -90% domain is not satisfied early on in training, and in response, the dual variable increases, gradually forcing constraint satisfaction. As we shown in the next subsection, without the dual update step, the constraints may never be satisfied (see Figure 3).
Regularization vs. dual ascent. A common trick for encouraging constraint satisfaction is to introduce soft constraints by adding a regularizer multiplied by a fixed multiplier to the objective. While this approach yields a related problem to (6) (see Appendix B.4) where the dual variables are fixed, there are few formal guarantees for this approach. Moreover, we show in Table 1 that when the dual variable is fixed during training (MBDG-Reg in Table 1), the performance drops significantly vis-a-vis MBDG. Notice that relative to Figure 2, the value of distReg(✓) is much larger than the margin, meaning that the constraint is not being satisfied.
Ablation on data augmentation. To study the efficacy of the MBDG algorithm, we consider two natural alternatives MBDG: (1) ERM with data augmentation through the learned model G(x, e) (MBDA); and (2) MBDG with data augmentation through G(x, e) on the training objective (MBDGDA). As shown at the bottom of Table 1, while these variants significantly outperform the baselines, they not perform nearly as well as MBDG. Thus, while data augmentation can in some cases improve performance, the primal-dual iteration is a much more effective tool for enforcing invariance.

heading: 7.2 Camelyon17-WILDS and FMoW-WILDS
text: We next consider the Camelyon17-WILDS and FMoW-WILDS datasets from the WILDS family of domain generalization baselines [20]. Table 2 shows that on Camelyon17-WILDS, MBDG improves by more than 20 percentage points over the state-of-the-art baselines.
Measuring G-invariance of trained classifiers. In Section 4, we restricted our attention predictors satisfying the G-invariance condition. To test whether our algorithm successfully enforces Ginvariance when G is learned from data, we measure the distribution of distReg(✓) over all of the instances from the training domains of Camelyon17-WILDS for ERM, IRM, and MBDG. In Figure 6, observe that whereas MBDG is quite robust to changes under G, ERM and IRM are not nearly as robust. This property is key to the ability of MBDG to learn invariant representations across domains.
Ablation on learning models vs. data augmentation. Rather than learning G from data, a heuristic alternative is to replace the GENERATEIMAGE procedure in Algorithm 1 with standard data augmentation transformations. In Figure 4, we investigate this approach with five different forms of data augmentation: B+C (brightness and contrast), CJ (color jitter), and three variants of RandAugment [145] (RA, RA-Geom, and RA-Color). More details concerning these data augmentation schemes are given in Appendix E. The bars in Figure 4 show that although these schemes offer strong performance in our MBDG framework, the learned model trained using MUNIT performs best.
7.3 PACS
In this subsection, we highlight selected results for the PACS dataset. Due to spatial limitations, we report the top-performing baselines in the main text, and defer the full set of results to Appendix E. Notably, MBDG beats the current SOTA by nearly 2% when averaged over the four domains. Of note in Table 3 is
the result on the “sketch” (S) subset, wherein MBDG improves by more than 10% over the baselines.

heading: 8 Acknowledgements and disclosure of funding
text: The authors would like to thank Samuel Sokota for his helpful feedback.
The research of Alexander Robey and Hamed Hassani is supported by NSF Grants 1837253, 1943064, 1934876, AFOSR Grant FA9550-20-1-0111, and DCIST-CRA. Alexander Robey and George J. Pappas are supported by the NSF-Simons Program on the Mathematics of Deep Learning, ARL CRA DCIST W911NF-17-2-0181, and NSF 2038873.
Checklist 1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] (c) Did you discuss any potential negative societal impacts of your work? [Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to
them? [Yes] 2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] (b) Did you include complete proofs of all theoretical results? [Yes]
3. If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)? [Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)? [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)? [Yes] 4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] (b) Did you mention the license of the assets? [N/A] (c) Did you include any new assets either in the supplemental material or as a URL? [Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating? [N/A] All of the external code used for this project is open-sourced and freely-available.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation? [N/A]

Review summaries:
This paper considers the domain generalization setting where data in all domains can be transformed into one another by a transformation G. Then, authors then propose to learn a predictor invariant to this transformation. Experimental results show some encouraging results. Overall approach presented is interesting, stimulating and the use of different optimization techniques to formulate the problem is pleasant. However, there is a ‘fatal’ problem in the framework, either by a mathematical mistake (i.e., moving from Eq 31 to 32 in the appendix), or authors have failed to declare a very strong assumption that that p(y|x) = p(y | G(x,e)) (see detailed comments below). If this point can be clarified, I would be willing to reconsider my rating.  ----- post-rebuttal I want to thank the reviewer for their thorough rebuttal. While I still want to see some solid analysis on the effect of the transform G and the assumption, the new discussions made sense and early empirical results on testing far domains are encouraging. The authors, please, make sure you factor them into the paper as they will not only enhance the understanding but also increase the validity of your work. 

This paper proposes a new model specifically designed for Out-of-Domain (OOD) detection. The authors start by setting up the domain generalization problem by explaining that (x,y) ~ p_e(x,y) is sampled from different domains ‘e’ in the training and test phase. Making certain assumptions such as (1) the inter-domain variation is only because of p(x), and (2) existing a transformative network g(x,e) that can map observation x to different domains e, the authors propose to slightly modify the problem definition which they call “Model-Based Domain Generalization”. This problem is slightly easier to deal with and under certain conditions, guarantees optimality. The authors finally do an extensive set of experiments to demonstrate the proposed algorithms significantly improves OOD detection results over previous models, as well as the role of regularizing theta with the distance regularizer. 

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.