Model-Based Domain Generalization

Introduction

Despite well-documented success in numerous applications , the complex prediction rules learned by modern machine learning methods can fail catastrophically when presented with out-of-distribution (OOD) data . Indeed, rapidly growing bodies of work conclusively show that state-of-the-art methods are vulnerable to distributional shifts arising from spurious correlations , adversarial attacks , sub-populations , and naturally-occurring variation . This failure mode is particularly pernicious in safety-critical applications, wherein the shifts that arise in fields such as medical imaging , autonomous driving , and robotics 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 .

In the last decade, the domain generalization community has emerged in an effort to improve the OOD performance of machine learning methods . 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 , it was recently shown that that no existing domain generalization algorithm can significantly outperform empirical risk minimization (ERM) over the training domains when ERM is properly tuned and equipped with state-of-the-art architectures and data augmentation techniques . 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 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.

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.

Related work

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 . Among such works, one prominent thrust has been to learn predictors which have internal feature representations that are consistent across domains . This approach is also popular in the field of unsupervised domain adaptation , wherein it is assumed that unlabeled data from the test domain is available during training . Also related are works that seek to learn a kernel-based embedding of each domain in an underlying feature space , and those that employ Model-Agnostic Meta Learning to adapt to unseen domains . Recently, another prominent direction has been to design weight-sharing and instance re-weighting schemes . Unlike any of these approaches, we explicitly enforce hard invariance-based constraints on the underlying statistical domain generalization problem.

Another approach toward improving OOD performance is to modify the available training data. Among such methods, perhaps the most common is to leverage various forms of data augmentation . Recently, several approaches have been proposed which use style-transfer techniques and image-to-image translation networks to augment the training domains with artificially-generated data . Alternatively, rather than generating new data, all seek to 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-based constraints on data generated by unsupervised generative models.

Domain generalization

To formalize this notion of environmental variation, we assume that data is drawn from a set of environments or domains $\mathcal{E}_{\text{all}}$ (see Figure 1(b)). Concretely, each domain $e\in\mathcal{E}_{\text{all}}$ can be identified with a pair of random variables $(X^{e},Y^{e})$, which together denote the observation of the random variable pair $(X,Y)$ in environment $e$. Given samples from a finite subset $\mathcal{E}_{\text{train}}\subsetneq\mathcal{E}_{\text{all}}$ 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)\approx Y$. This can be summarized as follows:

The growing domain generalization literature encompasses a great deal of past work, wherein both of these shifts have been studied in various contexts , resulting in numerous algorithms designed to solve Problem 3.1. Indeed, as this body of work has grown, new benchmark datasets have been developed which span the gamut between covariate and concept shift (see e.g. Figure 3 in and the discussion therein). 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 . As ERM is known to fail in the presence natural distribution shifts , this result highlights the critical need for new algorithms that can go beyond ERM toward solving Problem 3.1.

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.

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\in\mathcal{E}_{\text{all}}$ is generated from the underlying random variable pair $(X,Y)$ via an unknown function $G$.

Let $\delta_{e}$ denote a Dirac distribution for $e\in\mathcal{E}_{\text{all}}$. We assume that there existsCrucially, 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 7 and in Appendix E. a measurable function $G:\mathcal{X}\times\mathcal{E}_{\text{all}}\to\mathcal{X}$, which we refer to as a domain transformation model, that parameterizes the inter-domain covariate shift via

where $\#$ denotes the push-forward measure and $=^{d}$ denotes equality in distribution.

On the other hand, in our running medical imaging example, the label $y\sim 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 , the instance-conditional distributions can vary across domains, in this paper we assume that inter-domain variation is solely characterized by the domain shift parameterized by $G$.

2 A causal interpretation

3 Pulling back Problem 3.1

Under Assumptions 4.1 and 4.2, Problem 3.1 is equivalent to

4 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 the authors seek to learn an equipredictive representation $\Phi:\mathcal{X}\to\mathcal{Z}$ , i.e. an intermediate representation that satisfies

Despite compelling theoretical motivation for this approach, it has been shown that current algorithms which seek equipredictive representations do not significantly improve over ERM . 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$.

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 $\mathcal{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.

5 Formulating the MBDG optimization problem

The $G$-invariance property described in the previous section is the key toward reformulating the min-max problem in (3). Indeed, the following proposition follows from Assumptions 4.1 and 4.2 and from the definition of $G$-invariance.

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 (3), (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 B.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.

As in Problem 3.1, let $\mathcal{E}_{\text{train}}\subsetneq\mathcal{E}_{\text{all}}$ be a finite subset of training domains and assume that we have access to datasets $\mathcal{D}^{e}$ $\forall e\in\mathcal{E}_{\text{train}}$. 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).

6 Challenges in solving Problem 4.6

Problem 4.6 offers a new, theoretically-principled perspective on Problem 3.1 when data varies from domain to domain with respect to an underlying domain transformation model $G$. However, just as in general solving the min-max problem of Problem 3.1 is known to be difficult, the optimization problem in (MBDG) is also challenging to solve for several reasons:

Constrained optimization. Problem 4.6 is a constrained problem over an infinite dimensional functional space $\mathcal{F}$. While it is common to replace $\mathcal{F}$ with a parameterized function class, this approach creates further complications. Firstly, enforcing constraints on most modern, non-convex function classes such as the class of deep neural networks is known to be a challenging problem . Further, while a variety of heuristics exist for enforcing constraints on such classes (e.g. regularization), these approaches cannot guarantee constraint satisfaction for constrained problems .

Unknown domain transformation model. In general, we do not have access to the underlying domain transformation model $G$. While an analytic expression for $G$ may be known for simpler problems (e.g. rotations of the MNIST digits), analytic expressions for $G$ are most often difficult or impossible to obtain. For instance, obtaining a simple equation that describes the variation in color and contrast in Figure 1(c) would be challenging.

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 6.

Data-dependent duality gap for MBDG

In this section, we offer a principled analysis of Problem 4.6. In particular, we first address (C1) by introducing a relaxation of the $G$-invariance constraint that is compatible with modern notions of constrained PAC learnability . Next, to resolve the fundamental difficulty involved in solving constrained statistical problems highlighted in (C2), we follow by formulating 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).

Throughout this section, we have chosen to present our results somewhat informally by deferring preliminary results and regularity assumptions to the appendices. Proofs of each of the results in this section are provided in Appendix B.

Among the challenges inherent to solving Problem 4.6, one of the most fundamental is the difficulty of enforcing the $G$-invariance equality constraint. Indeed, it is not clear a priori how to enforce a hard invariance constraint on the class $\mathcal{F}$ of predictors. To alleviate some of this difficulty, we introduce the following relaxation of Problem 4.6:

While at first glance this problem may appear to be a significant relaxation of the MBDG optimization problem in (MBDG), when $\gamma=0$ and under mild conditions on $d$, the two problems are equivalent in the sense that $P^{\star}(0)=P^{\star}$ (see Proposition A.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 $\gamma$ is strictly larger than zero, under the assumption that the perturbation function $P^{\star}(\gamma)$ is $L$-Lipschitz continuous, we show in Remark A.2 that $|P^{\star}-P^{\star}(\gamma)|\leq L\gamma$, meaning that the gap between the problems is relatively small when $\gamma$ is chosen to be small. In particular, when strong duality holds for (MBDG), this Lipschitz constant $L$ is equal to the $L^{1}$ norm of the optimal dual variable $\left\|\nu^{\star}\right\|_{L^{1}}$ for (MBDG) (see Remark A.4).

2 Addressing (C2) by formulating the parameterized dual problem

As written, the relaxation in (5) is an infinite-dimensional constrained optimization problem over a functional space $\mathcal{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 $\cal F$, such as the class of deep neural networks . The approximation power of such a parameterization can be captured in the following definition:

The benefit of using such a parameterization is that optimization is generally more tractable in the parameterized space $\mathcal{A}_{\epsilon}:=\{\varphi(\theta,\cdot):\theta\in\mathcal{H}\}\subseteq\mathcal{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 , these approaches cannot provide guarantees on the underlying statistical problem of interest . Thus, to provide guarantees on the underlying statistical problem in Problem 4.6, given an $\epsilon$-parameterization $\varphi$ of $\mathcal{F}$, we consider the following saddle-point problem:

where $\mathcal{P}(\mathcal{E}_{\text{all}})$ is the space of normalized probability distributions over $\mathcal{E}_{\text{all}}$ and $\lambda\in\mathcal{P}(\mathcal{E}_{\text{all}})$ is the (semi-infinite) dual variable. Here we have slightly abused notation to write $R(\theta)=R(\varphi(\theta,\cdot))$ and $\mathcal{L}^{e}(\theta)=\mathcal{L}^{e}(\varphi(\theta,\cdot))$. One can think of (6) as the dual problem to (5) solved over the parametric space $\mathcal{A}_{\epsilon}$. Notice that unlike Problem 4.6, the problem in (6) is unconstrained, making it much more amenable for optimization over the class of deep neural networks. Moreover, under mild conditions, the optimality gap between (5) and (6) can be explicitly bounded as follows:

where $\lambda^{\star}_{\text{pert}}$ is the optimal dual variable for a perturbed version of (5) in which the constraints are tightened to hold with margin $\gamma-k\epsilon$.

In this way, solving the parameterized dual problem in (6) provides a solution that can be used to recover a close approximation of the primal problem in (5). To see this, observe that Prop. 5.2 implies that $|D_{\epsilon}^{\star}(\gamma)-P^{\star}(\gamma)|\leq\epsilon k(1+||\lambda_{\text{pert}}^{\star}||_{L^{1}})$. This tells us that the gap between $P^{\star}(\gamma)$ and $D^{\star}_{\epsilon}(\gamma)$ is small when we use a tight $\epsilon$-parameterization of $\mathcal{F}$.

3 Addressing (C3) by bounding the empirical duality gap

where $\hat{R}(\theta)$ and $\hat{}\mathcal{L}^{e}(\theta)$ are the empirical counterparts of $R(\theta)$ and $\mathcal{L}^{e}(\theta)$, i.e.

and $\hat{\Lambda}(\theta,\lambda)$ is the empirical Lagrangian. Notably, the duality gap between the solution to (8) and the original model-based problem in (MBDG) can be explicitly bounded as follows.

where $L$ is the Lipschitz constant of $P^{\star}(\gamma)$ and $k$ and $\lambda^{\star}_{\text{pert}}$ are as defined in Proposition 5.2.

Learning domain transformation models from data

Regarding challenge (C4), critical to our approach is having access to the underlying domain transformation model $G(x,e)$. 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 medical imaging dataset shown in Figure 1 would be challenging.

To address this challenge, we argue that a realistic approximation of the underlying domain transformation model can be learned from the instances drawn from the training datasets $\mathcal{D}^{e}$ for $e\in\mathcal{E}_{\text{train}}$. In this paper, to learn domain transformation models, we train multimodal image-to-image translation networks (MIITNs) on the instances drawn from the training domains. MIITNs are designed to transform samples from one dataset so that they resemble a diverse collection of images from another dataset. That is, the constraints used to train these models enforce that a diverse array of samples is outputted for each input image. This feature precludes the possibility of learning trivial maps between domains, such as the identity transformation.

As illustrated in Figure 3, these architectures generally consist of two components: a disentangled representation and a generative model. The role of the disentangled representation is to recover a sample $x$ generated according to $X$ from a instance $x^{e}$ observed in a particular domain $e\in\mathcal{E}_{\text{all}}$. In other words, for a fixed instance $x^{e}=G(x,e)$, the disentangled representation is designed to disentangle $x$ from $e$ via $(x,e)=H(x^{e})$. On the other hand, the role of the generative is to map each instance $x\sim X$ to a realization in a new environment $e^{\prime}$. Thus, given $x$ and $e$ at the output of the disentangled representation, we generate an instance from a new domain by replacing the environmental code $e$ with a different environmental parameter $e^{\prime}\in\mathcal{E}_{\text{all}}$ to produce the instance $x^{e^{\prime}}=G(x,e^{\prime})$. In this way, MIITNs are a natural framework for learning domain transformation models, as they facilitate 1) recovering samples from $X$ via the disentangled representation, and 2) generating instances from new domains in a multimodal fashion.

In each of the experiments in Section 8, we use the MUNIT architecture introduced in to parameterize MIITNs. As shown in Table 1 and in Appendix E, models trained using the MUNIT architecture learn accurate and diverse transformations of the training data, which often generalize to generate images from new domains. Notice that in this table, while the generated samples still retain the characteristic features of the input image (e.g. in the top row, the cell patterns are the same across the generated samples), there is clear variation between the generated samples. Although these learned models cannot be expected to capture the full range of inter-domain generalization in the unseen test domains $\mathcal{E}_{\text{all}}\backslash\mathcal{E}_{\text{train}}$, in our experiments, we show that these learned models are sufficient to significantly advance the state-of-the-art on several domain generalization benchmarks.

A principled algorithm for Model-Based Domain Generalization

Motivated by the theoretical results in Section 5 and the approach for learning domain transformation models in Section 6, we now introduce a new domain generalization algorithm designed to solve the empirical, parameterized dual problem in (8). We emphasize that while our theory relies on the assumption that inter-domain variation is solely characterized by covariate shift, our algorithm is broadly applicable to problems with or without covariate shift (see the experimental results in Section 8). In particular, assuming access to an appropriate learned domain transformation model $G$, we leverage $G$ toward solving the unconstrained dual optimization problem in (8) via a primal-dual iteration.

Given a learned approximation $G(x,e)$ of the underlying domain transformation model, the next step in our approach is to use a primal-dual iteration toward solving (8) using the training datasets $\mathcal{D}^{e}$. As we will show, the primal-dual iteration is a natural algorithmic choice for solving the empirical, parameterized dual problem in (8). Indeed, because the outer maximization in (8) is a linear program in $\lambda$, the primal-dual iteration can be characterized by alternating between the following steps:

Here $[\cdot]_{+}=\max\{0,\cdot\}$, $\eta>0$ is the dual step size, and $\rho\mbox{-}\operatorname*{argmin}$ denotes a solution that is $\rho$-close to being a minimizer, i.e. it holds that

For clarity, we refer to (11) as the primal step, and we call (12) the dual step.

The utility of running this primal-dual scheme is as follows. 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. In particular, this result is captured in the following theoremFor clarity, we state this theorem informally in the main text; a full statement of the theorem and proof are provided in Appendix B.6.:

Here $\kappa=\kappa(\epsilon)$ is a constant that captures the regularity of the parametric space $\mathcal{H}$ and $K(\rho,\kappa,\gamma)$ is a small constant depending linearly on $\rho$, $\kappa$, and $\gamma$.

This theorem means that by solving the empirical, parameterized dual problem in 8 for sufficiently many steps with small enough step size, we can reach a solution that is close to solving the Model-Based Domain Generalization problem in Problem 4.6. In essence, the proof of this fact is a corollary of Theorem 5.3 in conjunction with the recent literature concerning constrained PAC learning (see Appendix A.3).

2 Implementation of MBDG

In Algorithm 1, we outline two main procedures. In lines 12-15, we describe the GenerateImage($x^{e}$) procedure, which takes an image $x^{e}$ as input and returns an image that has been passed through a learned domain transformation model. The MUNIT architecture uses a normally distributed latent code to vary the environment of a given image. Thus, whenever GenerateImage is called, an environmental latent code $e^{\prime}\sim\mathcal{N}(0,I)$ is sampled and then passed through $G$ along with the disentangled input image.

In lines 4-8 of Algorithm 1, we show the main training loop for MBDG. In particular, after generating new images using the GenerateImage procedure, we calculate the loss term $\text{loss}(\theta)$ and the regularization term $\text{distReg}(\theta)$, both of which are defined in the empirical, parameterized dual problem in (8). Note that we choose to enforce the constraints between $x^{e}=G(x,e)$ and $x^{e^{\prime}}=G(x,e^{\prime})$, so that $\text{distReg}(\theta)=(1/m)\sum_{j=1}^{m}d(\varphi(\theta,x^{e}),\varphi(\theta,x^{e^{\prime}})$. We emphasize that this is completely equivalent to enforcing the constraints between $x^{e}=G(x,e)$ and $x$, in which the regulizer would be $\text{distReg}(\theta)=(1/m)\sum_{j=1}^{m}d(\varphi(\theta,x^{e}),\varphi(\theta,x)$. Next, in line 7, we perform the primal SGD step on $\theta$, and then in line 8, we perform the dual step on $\lambda$. Throughout, we use the KL-divergence for the distance function $d$ in the $G$-invariance term $\text{distReg}(\theta)$.

Experiments

We now evaluate the performance of MBDG on a range of standard domain generalization benchmarks. In the main text, we present results for ColoredMNIST, Camelyon17-WILDS, FMoW-WILDS, and PACS; we defer results for VLCS to the supplemental. For ColoredMNIST, PACS, and VLCS, we used the DomainBedhttps://github.com/facebookresearch/DomainBed package , 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 used the repository provided with the WILDS dataset suitehttps://github.com/p-lambda/wilds, and we performed 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 D.

We first consider the Camelyon17-WILDS and FMoW-WILDS datasets from the WILDS family of domain generalization benchmarks . Camelyon17 contains roughly 400k $96\times 96$ images of potentially cancerous cells taken at different hospitals, whereas FMoW-WILDS contains roughly 500k $224\times 224$ images of aerial scenes characterized by different forms of land use. Thus, both of these datasets are significantly larger than ColoredMNIST in both the number of images and the dimensionality of each image. In Table 2, we report classification accuracies for MBDG and a range of baselines on both Camelyon17-WILDS and FMOW-WILDS. Of particular interest is the fact that MBDG improves by more than 20 percentage points over the state-of-the-art baselines on Camelyon17-WILDS. On FMoW-WILDS, we report a relatively modest improvement of around one percentage point.

In essence, the significant improvement we achieve on Camelyon17-WILDS is due to the ability of the learned model to vary the coloration and brightness in the images. In the second row of Table 1, observe that the input image is transformed so that it resembles images from the other domains shown in Figure 1. Thus, the ability of MBDG to enforce invariance to the changes captured by the learned domain transformation model is the key toward achieving strong domain generalization on this benchmark. To further study the benefits of enforcing the $G$-invariance constraint, we consider two ablation studies on Camelyon17-WILDS.

In Section 4, we restricted our attention predictors satisfying the $G$-invariance condition. To test whether our algorithm successfully enforces $G$-invariance when a domain transformation model $G$ is learned from data, we measure the distribution of distReg$(\theta)$ over all of the instances from the training domains of Camelyon17-WILDS for ERM, IRM, and MBDG. In Figure 5, 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.

As shown in Table 1 and in Appendix E, accurate approximations of an underlying domain transformation model can often be learned from data drawn from the training domains. However, 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 (RA, RA-Geom, and RA-Color). More details regarding these data augmentation schemes are given in Appendix D. The bars in Figure 4 show that although these schemes offer strong performance in our MBDG framework, the learned model trained using MUNIT offers the best OOD accuracy.

2 ColoredMNIST

We next consider the ColoredMNIST dataset , which is a standard domain generalization benchmark created by colorizing subsets of the MNIST dataset . This dataset contains three domains, each of which is characterized by a different level of correlation between the label and digit color. The domains are constructed so that the colors are more strongly correlated with the labels than with the digits. Thus, as was argued in , stronger domain generalization on ColoredMNIST can be obtained by eliminating color as a predictive feature.

As shown in Table 3, despite the fact that the data generating procedure used to construct this dataset does not fulfill Assumptions 4.1 and 4.2 (see Figure 2(b)), the MBDG algorithm still improves over each baseline by nearly thirty percentage points. Indeed, due to way the ColoredMNIST dataset is constructed, the best possible result is an accuracy of 75%. Thus, the fact that MBDG achieves 68.5% accuracy when averaged over the domains means that it is close to achieving perfect domain generalization.

To understand the reasons behind this improvement, consider the first row of Table 1. Notice that whereas the input image shows a red ‘5’, samples from the learned domain transformation model show the same ‘5’ colored green. Thus, the $G$-invariance constraint calculated in line 5 of Algorithm 1 forces the classifier $f$ to predict the same label for both the red ‘5’ and the green ‘5’. Therefore, in essence the $G$-invariance constraint explicitly eliminates color as a predictive feature, resulting in the strong performance shown in Table 3. To further evaluate the MBDG algorithm and its performance on ColoredMNIST, we consider three ablation studies.

For the three MBDG classifiers selected by cross-validation at the bottom of Table 3, we plot the constraint term distReg$(\theta)$ and the corresponding dual variable at each training step in Figure 7(a). 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 show in the next subsection, without the dual update step, the constraints may never be satisfied (see Figure 7(b)). This underscores the message of Theorem 7.1, which is that the primal dual method can be used to enforce constraint satisfaction for Problem 4.6, resulting in stronger invariance across domains.

A common trick for encouraging constraint satisfaction in deep learning is to introduce soft constraints by adding a regularizer multiplied by a fixed penalty weight to the objective. While this approach yields a related problem to (8) where the dual variables are fixed (see Appendix A.4), there are few formal guarantees for this approach and tuning the penalty weight can require expert or domain-specific knowledge.

In Table 3, we show the performance of a regularized version of MBDG (MBDG-Reg in Table 3) where the dual variable is fixed during training (see Appendix C.2 for pseudocode). Note that while the performance of MBDG-Reg improves significantly over the baselines, it lags more than ten percentage points behind MBDG. Furthermore, consider that relative to Figure 7(a), the value of distReg($\theta$) shown in 7(b) is much larger than the margin of $\gamma=0.025$ used in Figure 7(a), meaning that the constraint is not being satisfied when running MBDG-Reg. Therefore, while regularization offers a heuristic alternative to MBDG, the primal-dual approach offers both stronger guarantees as well as superior performance.

To study the efficacy of the primal-dual approach taken by the MBDG algorithm toward improving the OOD accuracy on the test domain, 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 (MBDG-DA). We provide psuedocode and further discussion of both of these methods in Appendix C.1. As shown at the bottom of Table 3, 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 across domains.

3 PACS

In this subsection, we provide results for the standard PACS benchmark. This dataset contains four domains of $224\times 224$ images; the domains are “art/paining” (A), “cartoon” (C), “photo” (P), and “sketch” (S). In the fourth row of Table 1, we show several samples for one of the domain transformation models used for the PACS dataset. Further, Table 4 shows that MBDG achieves 85.6% classification accuracy (averaged across the domains), which is the best known result for PACS. In particular, this result is nearly two percentage points higher than any of the baselines, which represents a significant advancement in the state-of-the-art for this benchmark. In large part, this result is due to significant improvements on the “Sketch” (S) subset, wherein MBDG improves by nearly seven percentage points over all other baselines.

Conclusion

In this paper, we introduced a new framework for domain generalization called Model-Based Domain Generalization. In this framework, we showed that under a natural model of data generation and a concomitant notion of invariance, the classical domain generalization problem is equivalent to a semi-infinite constrained statistical learning problem. We then provide a theoretical, duality based perspective on problem, which results in a novel primal-dual style algorithm that improves by up to 30 percentage points over state-of-the-art baselines.

References

Appendix A Further theoretical results and discussion

In Section 5 of the main text, we claimed that the relaxation introduced in (5) was tight under certain conditions. In this section of the appendix, we formally enumerate the conditions under which the relaxation is tight. Further, we show that the tightness of the relaxation can be characterized by the margin parameter $\gamma$.

First, observe that because $d(\cdot,\cdot)$ is a metric, it is non-negative-valued. Then the following statement is trivial

Next, we claim that under the assumptions given in the statement of the proposition, $\mathcal{L}^{e}(f)=0$ is equivalent to the $G$-invariance condition. To verify this claim, for simplicity we start by defining the random variable

Altogether we have shown that $\mathcal{L}^{e}(f)\leq 0$ if and only if $f$ is $G$-invariant. Thus, when $\gamma=0$, the optimization problems in (MBDG) and (5) are equivalent, which implies that $P^{\star}(0)=P^{\star}$. ∎

A.2 The case when γ>0𝛾0\gamma>0

When $\gamma>0$, the relaxation is no longer tight. However, if the perturbation function $P^{\star}(\gamma)$ is assumed to be Lipschitz continuous, we can directly characterize the tightness of the bound.

Let us assume that the perturbation function $P^{\star}(\gamma)$ is $L$-Lipschitz continuous in $\gamma$. Then given Proposition A.1, it follows that $|P^{\star}-P^{\star}(\gamma)|\leq L\gamma$.

Observe that by Proposition A.1, we have that $P^{\star}=P^{\star}(0)$. It follows that

where the inequality in (21) follows by the definition of Lipschitz continuity. ∎

We note that in general the perturbation function $P^{\star}(\gamma)$ cannot be guaranteed to be Lipschitz. However, as we will show in Remark A.4, when strong duality holds for (MBDG), $P^{\star}(\gamma)$ turns out to be Lipschitz continuous with a Lipschitz constant equal to the $L^{1}$ norm of optimal dual variable for the dual problem to (MBDG). Before proving this result, we state a preliminary lemma from .

Assume that strong duality holds for this problem, and let $\lambda^{\star}$ denote an optimal dual variable. Define the perturbation function as follows:

This useful result, which follows from a simple one-line proof in §5.6.2 of , shows that the perturbation function $p^{\star}(u)$ can be related to the optimal value of the unperturbed problem via the optimal dual variable. We can readily use a semi-infinite version of this lemma to prove the following remark:

where $\mathcal{B}(\cdot)$ denotes the cone of non-regular, non-negative Borel measures supported on its argument . Assume that strong duality holds, and let $\nu^{\star}$ denote an optimal dual variable for this problem. Then it holds that

The idea here is to apply Lemma A.3 for the constant function defined by $u=u(e)=\gamma$ $\forall e\in\mathcal{E}_{\text{all}}$. To begin, let $\langle\cdot,\cdot\rangle$ denote the standard inner product on $L^{2}$; i.e. $\langle f,g\rangle=\int_{\mathcal{E}_{\text{all}}}f(e)g(e)\text{d}e$ for $f,g\in L^{2}(\mathcal{E}_{\text{all}})$. In this way, we find that

where the second inequality holds because for $\gamma$ strictly larger than zero, the relaxation in (5) corresponds to an expansion of the feasible set of relative to (MBDG). In this case, since $u$ is constant, a simple calculation shows that

where in the last step we have used the fact that the optimal dual variable $\nu^{\star}\succeq 0$. Now if we apply this result to (27), we find that

which directly implies the desired result. ∎

A.3 Relationship to constrained PAC learning

Recently, the authors of introduced the Probably Approximately Correct Constrained (PACC) framework, which extends the classical PAC framework to constrained problems. In particular, recall the following definition of agnostic PAC learnability:

The authors of extended this definition toward studying the learning theoretic properties of constrained optimization problems of the form

One of the main results in is that a hypothesis class $\mathcal{H}$ is PAC learnable if and only if it is PACC learnable.

Now if we consider the optimization problem in (31), we see that the admissible constraints are both inequality constraints. In contrast, the optimization problem in Problem 4.6 contains a family of equality constraints. Thus, in addition to easing the burden of enforcing hard $G$-invariance, the relaxation in (5) serves to manipulate the Model-Based Domain Generalization problem into a form compatible with (31). This is one of the key steps that sets the stage for deriving the learning theoretic guarantees for Model-Based Domain Generalization (e.g. Theorems 5.3 and 7.1).

A.4 Regularization vs. dual ascent

A common trick for encouraging constraint satisfaction is to introduce soft constraints by adding a regularizer multiplied by a fixed penalty weight to the objective. As noted in Section 8, this approach yields a similar optimization problem to (8). In particular, the regularized version of (8) is the following:

where $w(e)\geq 0$ $e\in\mathcal{E}_{\text{train}}$ are weights that are chosen as hyperparameters. From an optimization perspective, the benefit of such an objective is that gradient-based algorithms are known to converge to local minima given small enough step sizes (MBDG). However, classical results in learning theory can only provide generalization guarantees on the aggregated objective, rather than on each term individually. Furthermore, the choice of the penalty weights $w(e)$ is non-trivial and often requires significant domain knowledge, limiting the applicability of this approach.

In contrast, in primal-dual style algorithms, the weights $\lambda(e)$ are not fixed beforehand. Rather, the $\lambda(e)$ are updated iteratively via the dual ascent step described in line 8 of Algorithm 1. Furthermore, as we showed in the main text, the optimal value of the primal problem $P^{\star}$ can be directly related to the solution of the empirical dual problem in (8) via Theorem 5.3. Such guarantees are not possible in the regularization case, which underscores the benefits of the primal-dual iteration over the more standard regularization approach.

Appendix B Omitted proofs

In this appendix, we provide the proofs that were omitted in the main text. For ease of exposition, we restate each result before proving it so that the reader can avoid scrolling back and forth between the main text and the appendices.

Proposition 4.3. Under Assumptions 4.1 and 4.2, Problem 3.1 is equivalent to

where the final step again follows from the tower property of expectation. Therefore, by combining (39) and (43), we conclude that

which directly implies the statement of the proposition. ∎

B.2 Proof of Proposition 4.5

Proposition 4.5. Under Assumptions 4.1 and 4.2, if we restrict the feasible set to the set of $G$-invariant predictors, then Problem 3.1 is equivalent to the following semi-infinite constrained problem:

The main idea in this proof is simply to leverage the definition of $G$-invariance and the result of Prop. 4.3. Starting from Prop. 4.3, we see that by restricting the feasible set to the set of $G$ invariant predictors, the optimization problem in (3) can be written as

Now observe that due to the constraint, we can replace the $f(G(X,e))$ term in the objective with $f(X)$. Thus, the above problem is equivalent to

Now observe that the objective in (48) is free of the optimization variable $e\in\mathcal{E}_{\text{all}}$. Therefore, we can eliminate the inner maximization step in (48), which verifies the claim of the proposition. ∎

B.3 Proof of Proposition 5.2

The distance metric $d$ is non-negative, convex, and satisfies the following uniform Lipschitz-like inequality for some constant $L_{d}>0$:

Proposition 5.2. Let $\gamma>0$ be given. Then under Assumption B.1, it holds that

In this proof, we extend the results of to optimization problems with an infinite number of constraints. The key insight toward deriving the lower bound is to use the fact that maximizing over the $\epsilon$-parameterization of $\mathcal{F}$ yields a sub-optimal result vis-a-vis maximizing over $\mathcal{F}$. On the other hand, the upper bound, which requires slightly more machinery, leverages Jensen’s and Hölder’s inequalities along with the definition of the $\epsilon$-parameterization to over-approximate the parameter space via a Lipschitz $\epsilon$-ball covering argument.

Step 1. In the first step, we prove the lower bound in (52). To begin, we define the dual problem to the relaxed Model-Based Domain Generalization problem in (5) in the following way:

where with a slight abuse of notation, we redefine the Lagrangian $\Lambda$ from (6) in its first argument. Now recall that by assumption, there exists a predictor $f\in\mathcal{F}$ such that $\mathcal{L}(f)<\gamma$ $\forall e\in\mathcal{E}_{\text{all}}$. Thus, Slater’s condition holds , and therefore so too does strong duality. Now let $f^{\star}$ be optimal for the primal problem (5), and let $\lambda^{\star}\in\mathcal{B}(\mathcal{E}_{\text{all}})$ be dual optimal for the dual problem (54); that is,

At this early stage, it will be useful to state the following saddle-point relation, which is a direct result of strong duality:

which holds for all $f^{\prime}\in\mathcal{F}$ and for all $\lambda^{\prime}\in\mathcal{B}(\mathcal{E}_{\text{all}})$. Now consider that by the definition of the optimization problem in (6), we have that

Therefore, by choosing $\lambda^{\prime}=\lambda^{\star}$ in the above expression, and since $\mathcal{A}_{\epsilon}=\{\varphi(\theta,\cdot):\theta\in\mathcal{H}\}\subseteq\mathcal{F}$ by the definition of an $\epsilon$-parametric approximation, we have that

This concludes the proof of the lower bound: $P^{\star}(\gamma)\leq D^{\star}_{\epsilon}(\gamma)$.

Step 2. Next, we show that $D_{\epsilon}^{\star}(\gamma)$ is upper bounded by the optimal value of a perturbed version of the empirical dual problem. To begin, we add and subtract $\min_{f\in\mathcal{F}}\>\Lambda(f,\lambda)$ from the parameterized dual problem in (6).

Now let $\mu(e)$ denote any probability measure with support over $\mathcal{E}_{\text{all}}$. Consider the latter two terms in the above problem, and observe that we can write

where $(*)$ and $(**)$ follows from separate applications of Hölder’s ineqaulity , and $(\square)$ follows from an application of Minkowski’s inequality and from the fact that $\mu$ is a (normalized) probability distribution. Let us now consider the second term in the above product:

Now let $D^{\star}_{\text{pert}}(\gamma)$ denote the optimal value of the above problem; that is,

Step 3. In the final step, we prove the theorem. We begin with the perhaps unintuitive fact that the perturbed problem defined above is the dual problem to a perturbed version of the optimization problem in (5). More specifically, the perturbed problem in (79) is the dual of

Note that as this primal perturbed optimization problem is convex since (5) is convex, and by assumption strong duality also holds for this perturbed problem. Let $(f_{\text{pert}}^{\star},\lambda_{\text{pert}}^{\star})$ be primal-dual optimal for the perturbed problems we have defined above. The following saddle-point relation is evident from the fact that strong duality holds:

where the inequalities hold for all $f^{\prime}\in\mathcal{F}$ and for all $\lambda^{\prime}\in\mathcal{B}(\mathcal{E}_{\text{all}})$. Using this result for the choice of $f^{\prime}=f^{\star}$, where we recall that $f^{\star}$ is defined in (55) as the primal optimal solution to (5), it follows from (77) that

Now, recalling the original saddle-point relation in (77), it holds that $\Lambda(f^{\star},\lambda_{\text{pert}}^{\star})\leq\Lambda(f^{\star},\lambda^{\star})$. Using this fact along with (83) yields the following result:

B.4 Characterizing the empirical gap (used in Theorem 5.3)

In this proof, we use a similar approach as in [138, Prop. 2] to derive the generalization bound. Notably, we extend the ideas given in this proof to accommodate two problems with different constraints, wherein the constraints of one problem are a strict subset of the other problem.

To begin, let $(\theta_{\epsilon}^{\star},\lambda^{\star}_{\epsilon})$ and $(\theta_{\epsilon,N,\mathcal{E}_{\text{train}}}^{\star},\lambda_{\epsilon,N,\mathcal{E}_{\text{train}}}^{\star})$ be primal-dual optimal pairs for (6) and (8) that achieve $D^{\star}_{\epsilon}(\gamma)$ and $D^{\star}_{\epsilon,N,\mathcal{E}_{\text{train}}}(\gamma)$ respectively; that is,

are satisfied. Due to the optimality of these primal-dual pairs, both primal-dual pairs satisfy the KKT conditions . In particular, the complementary slackness condition implies that

Thus, as (88) indicates that the second term in the objective of (86) is zero, we can recharacterize the optimal value $D_{\epsilon}^{\star}(\gamma)$ via

and similarly from (89), can recharacterize the optimal value $D^{\star}_{\epsilon,N,\mathcal{E}_{\text{train}}}(\gamma)$ as

Ultimately, our goal is to bound the gap between $|D_{\epsilon}^{\star}(\gamma)-D^{\star}_{\epsilon,N,\mathcal{E}_{\text{train}}}(\gamma)|$. Combining (90) and (91), we see that this gap can be characterized in the following way

Now due to the optimality of the primal-optimal variables $\theta_{\epsilon}^{\star}$ and $\theta_{\epsilon,N,\mathcal{E}_{\text{train}}}^{\star}$, observe that

which, when combined with (92), implies that

To wrap up the proof, we simply leverage the classical VC-dimension bounds for both of the terms in (97). That is, following , it holds for all $\theta$ that with probability $1-\delta$,

As the bound in (98) holds $\forall\theta\in\mathcal{H}$, in particular it holds for $\theta_{\epsilon}^{\star}$ and $\theta_{\epsilon,N,\mathcal{E}_{\text{train}}}^{\star}$. This directly implies the bound in (85). ∎

B.5 Proof of Theorem 5.3

The proof of this theorem is a simple consequence of the triangle inequality. Indeed, by combining Remark A.2, Proposition 5.2, and Proposition B.2, we find that

B.6 Proof of Theorem 7.1

Finally, assume that there exists a parameter $\theta\in\mathcal{H}$ such that $\varphi(\theta,\cdot)$ is strictly feasible for (5), i.e. that

where $\nu$ is the constant from (104). Then it follows that the primal-dual pair $(\theta^{(T)},\lambda^{(T)})$ obtained after running the alternating primal-dual iteration in (11) and (12) for $T$ steps with step size $\eta$, where

where $\left\|\lambda^{\star}\right\|$ is the optimal dual variable for (6).

Observe that by the triangle inequality, we have

where the last step follows from Remark A.2. Then, from [149, Theorem 2], it directly follows that

Combining this with (110) completes the proof. ∎

Appendix C Algorithmic variants for MBDG

In Section 8, we considered several algorithmic variants of MBDG. Each variant offers a natural point of comparison to the MBDG algorithm, and for completeness, in this section we fully characterize these variants.

In Section 8, we did an ablation study concerning various data-augmentation alternatives to MBDG. In particular, in the experiments performed on ColoredMNIST, we compared results obtained with MBDG to two algorithms we called MBDA and MBDG-DA. For clarity, in what follows we describe each of them in more detail.

In the MDBA variant, we train using ERM with data augmentation through the learned domain transformation model $G(x,e)$. This procedure is summarized in Algorithm 2. Notice that in this algorithm, we do not consider the constraints engendered by the assumption of $G$-invariance. Rather, we simply seek to use follow the recent empirical evidence that suggests that ERM with proper tuning and data augmentation yields state-of-the-art performance in domain generalization . Note that in Table 3, the MBDA algorithm improves significantly over the baselines, but that it lags more than 20 percentage points behind results obtained using MBDG. This highlights the utility of enforcing constraints rather than performing data augmentation on the training objective.

In the MBDG-DA variant, we follow a similar procedure to the MBDG algorithm. The only modification is that we perform data augmentation through the learned model $G(x,e)$ on the training objective in addition to enforcing the $G$-invariance constraints. This procedure is summarized in Algorithm 3. As shown in Table 3, this procedure performs rather well on ColoredMNIST, beating all baselines by nearly 20 percentage points. However, this algorithm still does not reach the performance level of MBDG when the -90% domain is taken to be the test domain.

C.2 Regularization

In Section 8, we also compared the performance of MBDG to a regularized version of MBDG. In this regularized version, we sought to solve (37) using the algorithm described in Algorithm 4. In particular, in this algorithm we fix the weight $w>0$ as a hyperparameter, and we perform SGD on the regularized loss function $\text{loss}(\theta)+w\cdot\text{distReg}(\theta)$. Note that while this method performs well in practice (see Table 3), it is generally not possible to provide generalization guarantees for the regularized version of the problem.

Appendix D Additional experiments and experimental details

In this appendix, we record further experimental details beyond the results presented in Section 8. The experiments performed on ColoredMNIST, PACS, and VLCS were all performed using the DomainBed package. All of the default hyperparameters (e.g. learning rate, weight decay, etc.) were left unchanged from the standard DomainBed implementation. In Table 5, we record the additional hyperparameters used for MBDG and its variants as well as the random criteria by which hyperparameters were generated. For each of these DomainBed datasets, model-selection was performed via hold-one-out cross-validation, and the baseline accuracies were taken from commit 7df6f06 of the DomainBed repository. The experiments on the WILDS datasets used the hyperparameters recorded by the authors of ; these hyperparameters are recorded in Sections D.1 and D.2. Throughout the experiments, we use the KL-divergence as the distance metric $d$.

For the Camelyon17-WILDS dataset, we used the out-of-distribution validation set provided in the Camelyon17-WILDS dataset to tune the hyperparameters for each classifier. This validation set contains images from a hospital that is not represented in any of the training domains or the test domain. Following , we used the DenseNet-121 architecture and the Adam optimizer with a batch size of 200. We also used the same hyperparameter sweep as was described in Appendix B.4 of . In particular, when training using our algorithm, we used the the following grid for the (primal) learning rate: $\eta_{p}\in\{0.01,0.001,0.0001\}$. Because we use the same hyperparameter sweep, architecture, and optimizer, we report the classification accuracies recorded in Table 9 of to provide a fair comparison to past work. After selecting the hyperparameters based on the accuracy on the validation set, we trained classifiers using MBDG for 10 independent runs and reported the average accuracy and standard deviation across these trials in Table 2.

In Section 8, we performed an ablation study on Camelyon17-WILDS wherein the model $G$ was replaced by standard data augmentation transforms. For completeness, we describe each of the methods used in this plot below. For each method, invariance was enforced between a clean images drawn from the training domains and corresponding data that was varied according to a particular fixed transformation.

The PIL color transformationhttps://pillow.readthedocs.io/en/stable/reference/ImageEnhance.html#PIL.ImageEnhance.Color. See Figure 8 for samples.

PIL Brightnesshttps://pillow.readthedocs.io/en/stable/reference/ImageEnhance.html#PIL.ImageEnhance.Brightness and Contrasthttps://pillow.readthedocs.io/en/stable/reference/ImageEnhance.html#PIL.ImageEnhance.Contrast transformations. See Figure 9 for samples.

We use the data augmentation technique RandAugment , which randomly samples random transformations to be applied at training time. In particular, the following transformations are randomly sampled: AutoContrast, Equalize, Invert, Rotate, Posterize, Solarize, SolarizeAdd, Color, Constrast, Brightness, Sharpness, ShearX, ShearY, CutoutAbs, TranslateXabs, and TranslateYabs. We used an open-source implementation of RandAugment for this experimenthttps://github.com/ildoonet/pytorch-randaugment. See Figure 10 for samples.

We use the RandAugment scheme with a subset of the transformations mentioned in the previous paragraph. In particular, we use the following geometric transformations: Rotate, ShearX, ShearY, CutoutAbs, TranslateXabs, and TranslateYabs. See Figure 11 for samples.

We use the RandAugment scheme with a subset of transformations mentioned in the RandAugment paragraph. In particular, we use the following color-based transformations: AutoContrast, Equalize, Invert, Posterize, Solarize, SolarizeAdd, Color, Constrast, Brightness, Sharpness. See Figure 12 for samples.

We use an MUNIT model trained on the images from the training datasets; this is the procedure advocated for in the main text, i.e. in the GenerateImage(x) procedure. See Figure 13 for samples.

D.2 FMoW-WILDS

As with the Camelyon17-WILDS dataset, to facilitate a fair comparison, we again use the out-of-distribution validation set provided in . While the authors report the architecture, optimizer, and final hyperparameter choices used for the FMoW-WILDS dataset, they not report the grid used for hyperparameter search. For this reason, we rerun all baselines along with our algorithm over a grid of hyperparameters using the same architecture and optimizer as in . In particular, we follow by training a DenseNet-121 with the Adam optimizer with a batch size of 64. We selected the (primal) learning rate from $\eta_{p}\in\{0.05,0.01,0.005,0.001\}$. We selected the trade-off parameter $\lambda_{\text{IRM}}$ for IRM from the grid $\lambda_{\text{IRM}}\in\{0.1,0.5,1.0,10.0\}$. As before, the results in Table 2 list the average accuracy and standard deviation over ten independent runs attained by our algorithm as well as ERM, IRM, and ARM.

D.3 VLCS

In Table 6, we provide a full set of results for the VLCS dataset. As shown in this Table, MBDG offers competitive performance to other state-of-the-art method. Indeed, MBDG achieves the best results on the “LabelMe” (L) subset by nearly two percentage points.

Appendix E Further discussion of domain transformation models

In some applications, domain transformation models in the spirit of Assumption 4.1 are known a priori. To illustrate this, consider the classic domain generalization task in which the domains correspond to different fixed rotations of the data . In this setting, the underlying generative model is given by

where $R(e)$ is a one-dimensional rotation matrix parameterized by an angle $e$. In this way, each angle $e$ is identified with a different domain in $\mathcal{E}_{\text{all}}$. However, unlike in this simple example, for the vast majority of settings encountered in practice, the underlying domain transformation model is not known a priori and cannot be represented by concise mathematical expressions. For example, obtaining a closed-form expression for a generative model that captures the variation in coloration, brightness, and contrast in the Camelyon17-WILDS cancer cell dataset shown in Figure 1(a) would be very challenging.

In this appendix, we provide an extensive discussion concerning the means by which we used unlabeled data to learn domain transformation models using instances drawn from the training domains $\mathcal{E}_{\text{train}}$. In particular, we argue that it is not necessary to have access to the true underlying domain transformation model $G$ to achieve state-of-the-art results in domain generalization. We then give further details concerning how we used the MUNIT architecture to train domain transformation models for ColoredMNIST, Camelyon17-WILDS, FMoW-WILDS, PACS, and VLCS. Finally, we show further samples from these learned domain transformation models to demonstrate that high-quality samples can be obtained on this diverse array of datasets.

We emphasize that while our theoretical results rely on having access to the underlying domain transformation model, our algorithm and empirical results do not rely on having access to the true $G$. Indeed, although we did not have access to the true model in any of the experiments in Section 8, our empirical results show that we were able to achieve state-of-the-art results on several datasets.

E.2 Learning domain transformation models with MUNIT

In practice, to learn a domain transformation model, a number of methods from the deep generative modeling literature have been recently been proposed . In particular, throughout the remainder of this paper we will use the MUNIT architecture introduced in to parameterize learned domain transformation models. This architecture comprises two GANs and two autoencoding networks. In particular, the MUNIT architecture – along with many related works in the image-to-image translation literature – was designed to map images between two datasets $A$ and $B$. In this paper, rather than separating data we simply use $\mathcal{D}_{X}$ for both $A$ and $B$, meaning that we train MUNIT to map the training data back to itself. In this way, since $\mathcal{D}_{X}$ contains data from different domains $e\in\mathcal{E}_{\text{train}}$, the architecture is exposed to different environments during training, and thus seeks to map data between domains.

E.3 On the utility of multi-modal image-to-image translation networks.

In this paper, we chose the MUNIT framework because it is designed to learn a multimodal transformation that maps an image $x$ to a family of images with different levels of variation. Unlike methods that seek deterministic mappings, e.g. CycleGAN and its variants , this method will learn to generate diverse images, which allows us to more effectively enforce invariance over a wider class of images. In Figures 14, 15, and 16, we plot samples generated by sampling different style codes $e\sim\mathcal{N}(0,I)$ for MUNIT. Note that while the results for Camelyon17-WILDS and FMoW-WILDS are sampled using the model $G(x,e)$, the samples from PACS are all sampled from different models.