Iterative Feature Matching: Toward Provable Domain Generalization with Logarithmic Environments
Authors: Yining Chen
Elan Rosenfeld
Mark Sellke
Tengyu Ma
Andrej Risteski
Venue: NeurIPS 2022 Accept (2022)
Keywords: Domain generalization
out-of-distribution generalization
invariant risk minimization
IRM
deep learning theory
domain generalization theory

Abstract: Domain generalization aims at performing well on unseen test environments with data from a limited number of training environments. Despite a proliferation of proposed algorithms for this task, assessing their performance both theoretically and empirically is still very challenging. Distributional matching algorithms such as (Conditional) Domain Adversarial Networks [Ganin et al., 2016, Long et al., 2018] are popular and enjoy empirical success, but they lack formal guarantees. Other approaches such as Invariant Risk Minimization (IRM) require a prohibitively large number of training environments---linear in the dimension of the spurious feature space $d_s$---even on simple data models like the one proposed by [Rosenfeld et al., 2021]. Under a variant of this model, we show that ERM and IRM can fail to find the optimal invariant predictor with $o(d_s)$ environments. We then present an iterative feature matching algorithm that is guaranteed with high probability to find the optimal invariant predictor after seeing only $O(\log d_s)$ environments. Our results provide the first theoretical justification for distribution-matching algorithms widely used in practice under a concrete nontrivial data model.

Iterative Feature Matching: Toward Provable Domain Generalization with Logarithmic Environments
Yining Chen Stanford University cynnj justtns edu
Elan Rosenfeld
Mark Sellke Stanford University
Carnegie Mellon University elan@cmu. edu
msellke@stanford.
Tengyu Ma
Andrej Risteski Carnegie Mellon University aristeskGandres. cmu. edu
Stanford University tengyuma@staafordd.co edu
Abstract
Domain generalization aims at performing well on unseen test environments with data from limited number of training environments. Despite proliferation of proposed algorithms for this task, assessing their performance both theoretically and empirically is still very challenging. Distributional matching algorithms such as (Conditional) Domain Adversarial Networks [12] (8) are popular and enjoy empirical success, but they lack formal guarantees. Other approaches such Invariant Risk Minimization (IRM) require a prohibitively large number of training environments linear in the dimension of the spurious feature space d,- even simple data models like the one proposed by Rosenfeld et al. 377. Under a variant of this model, we show that ERM and IRM can fail find the optimal invariant predictor with o(d,) environments. We then present an iterative feature matching algorithm that is guaranteed with high probability find the optimal invariant predictor after seeing only Ollog d,) environments. Our results provide the first theoretical justification for distribution -atttitttooiiinnggggoe algorithms widely used in practice under concrete nontrivial data model.
Introduction
Domain generalization aims at performing well on unseen environments using labeled data from limited number of training environments [8] In contrast transfer learning domain adaptation, domain generalization assumes that neither labeled nor unlabeled data from the test environments available at training time. In the empirical literature, invariance of "signal" feature distribution conditioned on the label, i.. (()(x) (x) v). is the underlying assumption in widely adopted algorithms such as Correlation Aligmment (CORAL) 42 Maximum Mcan Discrepancy (MMD) K 25 and
algorithms regularize the distribution of intermediate representations of examples from different environments to closer in some metric. They are among the top-pent nop-performing methods for domain generalization benchmarks like PACS (24) VLCS [1. and DomainNet (34) but they lack formal guarantees. Previous empirical works usually characterize the performance of these algorithms using generalization bounds based on divergence between domain distributions 7 3T, but those bounds are vacuous for the typical benchmarks for domain generalization and thus cannot explain their success. For example, when different environments have disjoint supports (as the case in the above-mentioned benchmarks) the H-divergence [7] is even for the linear hypothesis class. To obtain useful guarantees, it is necessary to study data models that encode structure reflective of settings of interest. Prior works attempting theoretically characterize the performance of feature matching algorithms emphasize
(Conditional) Domain Adversarial Networks 12 28 Distribution-mattee
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

lower bounds 449 44) In this work, we seek to give the first positive theoretical justification for feature matching algorithms.
Another common assumption in the literature is invariance of the label distribution conditioned on
the signal features. Invariant Risk Minimization (IRM) 3 assumes Ely (xx))] is invariant, and
follow-up works assume invariance of higher moments [46 18 30 6]. However, empirical results
for these algorithms are mixed: Guirajani and Lopez-Paz [15, Aubin et al. [4] present experimental
evidence that these methods do not consistently outperform ERM for either realistic or simple linear
data models, when fairly evaluated.
Recent theoretical works [37 [9] also question the theoretical foundations of IRM and its variants, shedding light their failure conditions. These works study specific data generative models; a common assumption is that, conditioned on the label, some invariant features have identical distribution for all environments, and other spurious features have varying distributions across environments. The goal of these methods is then to obtain an invariant predictor, i.e. a classifier which uses only the invariant features. These works also often assume each training environment contains infinite samples. Thus, the central measure of domain generalization is the number the
of environments needed to recover an invariant predictor-we refer to this measure environment complexity of a learning algorithm. Rosenfeld et al. 37 show even for a simple generative model and linear classifiers, the environment complexity of IRM- -add other objectives based on the same principle of invariance- is at least as large as the dimension of the spurious latent features, dz. Further results by Kamath et al. [9], Ahuja et al. [2] also point to a linear environment complexity. Although the models in these works are simple, they help elucidate why existing algorithms fail and can help inform better algorithmic design.
IRM's linear environment complexity is prohibitive for realistic applications. Domain generalization benchmarks have fewer training environments than spurious features dimensions-PAACS VLCS, DomainNet have 4-6 total environments but the stylistic variations between domains are likely high- dimensional. In this paper, we show that a variant of distributional matching algorithm (Algorithm IT finds the optimal invariant predictor with sublinear (even logarithmic) environment complexity.
1.1 Our contributions
Conceptually, we propose "smoothed covariance" extension of the data model in Rosenfeld et al. 337, which is rich enough to establish an environment complexity separation between IRM/ERM and feature matching algorithms. More precisely, we show that ERM and IRM can fail find the optimal invariant predictor with fewer than d, + 1 training environments (Theorems relatively simple and natural algorithm based iterative feature matching (IFM) finds the optimal invariant predictor with Ollog d.) environments.
( 4.2 4.3). while
Moreover, analysis highlights the value of iterative matching, much similar to matching at multiple layers in a deep neural network. IFM iteratively projects the features to a lower dimension while matching the aabel-conditioned feature distributions environments. A projection inducing invariance in the non-invariant features across one subset of environments is unlikely to do for a different subset, each projection only removes (a constant fraction of) spurious feature dimensions with high probability. Intuitively, the iterative scheme prevents the different environments from "colluding" to create solution which depends on spurious features. As a result, IFM recovers the optimal invariant predictor after O(log d.) rounds and uses O(1) environments per round, thus requiring O(log d.) environments.
small, disjoint subset of the training
Our techniques for proving the upper and lower bounds may be of independent interest. The upper bound (Theorem 5.1 is proved by showing that any projection matrix that uses many spurious dimensions cannot match feature covariances in large set of training environments. This is done via intricate matrix concentration bounds and decoupling inequalities [10].
We derive the exact lower bound for IRM (Theorem 4.3) using tools from differential topology. Each environment provides an ellipsoidal constraint on the solution, and we prove that there exists a non-trivial intersection of these constraints (besides the origin, which corresponds to the "intended" solution). Our key lemma shows that the total number of intersections between two manifolds of complementary dimensions k d - k even when certain tranversality conditions hold, implying that the origin cannot be the only solution.

---

Finally, to corroborate the advantages of the proposed algorithm, we perform experiments on Gaussian dataset and a semi-synthetic Noised MNIST dataset [23, where the background noise spuriously correlates with the label. Our results in Section from feature matching algorithms when the distinguishing property of the signal feature is indeed conditional distributional invariance, and may get additional advantage via matching at multiple layers with diminishing dimensions, echoing existing empirical observations 27,29)
6 suggest that practitioners may benefit
Preliminaries
2.1 Domain generalization
In domain generalization, we are given a set of training environments E,y bndexed E [ and set of test environments Ets- For environment we have examples {(XF, Y drawn from the distribution Pe. In this work we study the infinite sample limit n + oc so as to separate the effect of limited training environments from that of limited samples per environment, is done in previous theoretical works [37 [9]. Let X. P. y denote the space of inputs, intermediate features, and labels. For featurizer + X P and classifier w
V, their risk environment is denoted by Re = Ex [ (w o ((X), Y)] for any common loss function 1. In this paper we focus on y = {+1]. linear featurizers (XX) UX for U Rexdd and unit-norm predictors Ý = sgn (w UU)) where E R and || U 2= for some feature dimension k S d chosen by the algorithm. A predictor's 0-1 risk environment is denoted by Ri Prçx, Pr [sgn (w UX) * Yj. We focus on unit-norm predictors because we evaluate on the 0-1 risk on test environments, which are invariant the scaling of w U under our data model.
2.2 Baseline algorithms
We analyze the performance of our proposed method and compare it to two baseline algorithms, ERM and IRM. ERM learns classifier that minimizes the average loss over all training environments, where is any common training loss such as the logistic loss:
Σ E Exxyy> [ww X,Y).. RE(
min WE
IRM learns a featurizer (XX) E R such that the optimal classifier on top of the featurizer invariant across training environments. As we focus on linear classifier, it is equivalent to learning linear transformation k2 such that it induces a classifier w that optimal for all Ezy:
Σ eE(E
min UERExd WERK T U1
s.t. w E min B ((wd" UX),Y)) E E w'ERt
This is objective is not the same feature distribution matching; IRM only tries to match the first moment. Observe that this constrained objective is intended solve minimax domain generalization problem, as opposed to ERM which is typically viewed as minimizing the risk in expectation.
3 Problem setup
We first recall the data model from Rosenfeld et al. [37] We assume without loss of generality that the label Y is uniformly randomly drawn from {+1} (extension of our theorems to Y 1 with probability 7 * 0.5 is straightforward), Latent variable Z consists of invariant features Z, E R" and spurious features Zz E Rdd. where d, d The number of spurious features can be much larger than the number invariant features, i.e. d, >> r. The input X Rd is generated via a linear
We define [n]] = {L...., nn:: Ow E T denotes an all-zero matrix; s is the unit sphere in ued+1; sgn (c) (#1, 0} the sign of scalar E R. i denotes the
pseudo-inverse.

---

transformation of latent variable Z, i.c. X = SZ for matrix S E Rdxd such that its left r columns have rank (so that there are invariant dimensions).
For each training environment indexed by e E [F]. the invariant features conditioned on Y are drawn from Gaussian distribution with mean Y 'A1 RR" and nonsingular covariance Σn Rexr. The spurious features conditioned on Y have mean Y μu E I and covariance 2 E Igdexde where He's 's and 2 vary across E E [E]. The assumption of symmetric class center with respect the origin can also be relaxed. Define μ = [1,, and X [,, xrxd, Oa, 22 The overall data model for training environments is summarized below:
Y unif{=1}
/ N(Y 111, E1) E I"
/ ~ N(Y .15. En) E ERd.
Z = [Z Z2] E Rd
SZ.
Since the goal of invariant feature learning is to learn predictor that only uses the invariant features, one reasonable measure for domain generalization a predictor's performance on test environments where the spurious features Zo are drawn from a different distribution particular, they are usually chosen adversarially. A classifier that predicts using the spurious features will perform badly on such test environments. When modeling the test environments, we consider the difficult scenario where there is one corresponding test environment for each training environment, whose parameters are the same except that the spurious means are flipped. Formally, for each environment Etr construct a corresponding test environment e' E Etest where
Ze N(-Y 12 E5) e Rd.
In this setting where the observations X are a linear function of the latents Z, Rosenfeld et al. 37 assume that the covariances of spurious features are isotropic and vary only in magnitude: Assumption 3.1 (Data model for spurious covariances in Rosenfeld et al. 37).
E$ = I where de scalar for an environment indexed by
We consider generalized model where the covariances of spurious features for each environment is generic random PSD matrix, instead of only random in scaling:
A Assumption 3.2 (Data model for spurious covariances in this work).
Σg Σg + G.G., where Σn E Rd.xd, d, is arbitrary (and can be udversarial), and [Gali sid N for all i.j E [d,]. Furthermore, E> D.
Assumption 3.2 allows the covariance matrix the spurious features to be almost worst-case: an arbitrary matrix, plus a Gaussian perturbation. This is common assumption in algorithmic complexity called smoothed analysis [11 as we allow the parameters to be a random smoothing of arbitrary parameters.
While our data model is simple, it has already been used as a sandbox to understand algorithms like IRM 37 [9 because it captures some important aspects of real-life data like latent variables and correlations between the labels and the spurious features.
In the next section, we show that baseline algorithms ERM and IRM still have ((d,) environment complexity under assumption 3.22, whereas our iterative feature matching algorithm (Algorithm I requires only O(log d,) training environments. Note that the environment complexity of our algorithm only depends logarithmically on the norm bound D.
Main results
Armed with assumption 3.2 we now present our main results. We begin by presenting our algorithm based on iterative feature matching. We then provide formal guarantees for its environment complexity in comparison to ERM and IRM.

---

Algorithm 1 Iterative Feature Matching (IFM) algorithm
Require:
Invariant feature dimension r, total feature dimension d, number of training environments
E [1,g]. infinite samples {(X, ) Pe from each environment E Etr, an array of integers indicating the number of training environments to use at each iteration E1, Eg.....
ro f- d. 6
2: while ri > r do
3:
++
4:
Uniformly randomly sample E, training environments without replacement.
5: Binary search between I and r to find the maximum dimension r such that there exists orthonormal U, E Rr and C₁ Rr:xr; where for all E &s,
Y C..
(4.1)
6: Return a classifier on projected features that minimizes the average risk
min WES
Σ
Exx.y> U... Ui Y)
F.
RE(A
4.1 Iterative feature matching algorithm
We hope to recover the invariant features by imposing constraints which are satisfied by only those features (i.e., they are not satisfied by the spurious features). natural idea is match the feature means and covariances across &tr. Since J(, Ei are constant, any orthonormal featurizer U E Rr such that U. S has only non-zero entries in the first r rows yields invariant means US and covariances U S UT Thus we need E large enough such that any U' E Reed using spurious dimensions cannot match the means and covariances. Informally, for each e E [E] we get r x r equations from matching covariances U C, and we have Rough parameter counting suggests that if we match covariances of all E environments jointly, we need least
parameters to estimate U.
d/r environments find a unique solution. Our key observation that, due the independence of randomness in Σs, we can split E environments into disjoint groups E Er, and use & to train an orthonormal featurizer that shrinks the feature dimensions from ft-1 to rt. Thus, in each round we shrink the dimension by a constant factor using a constant number of environments. The main theoretical challenge that remains to show that in each iteration, with high probability, only spurious features are projected out.
This brings IFM (Algorithm D). which proceeds in T O(log de) rounds. Starting with an input dimension ro = d. each round we learn an orthonormal matrix U2 projecting features from 7t-1 to r, dimensions so that the feature covariances after projection match across fresh set of training environments. In practice, for each choice of rt, starting from random initialization, we perform SGD on objective min mi En EE, | - Ee YU? I|? +AIU, U, - TI until the objective is e-close to O. To ensure that all invariant dimensions are preserved, we always find the projection with the maximum possible dimension 7, that still matches the covariances (i.e. admits an e-optimal solution U,). until we are left with only dimensions.
Among algorithms used in practice, IFM is most similar CORAL 43
We study IFM since it is more amenable to theoretical analysis. The differences between IFM and CORAL are: first, CORAL does not enforce that the featurizer is orthonormal; second, IFM learns to extract features in an unsupervised manner, whereas CORAL jointly minimizes the supervised loss and feature distribution discrepancy; third, IFM matches the feature distributions multiple layers and uses disjoint of environments for each layer-this iterative process is necessary for the theoretical guarantees we provide- but CORAL matches only at the last layer. Despite these differences. our theoretical results serve as justification for using feature matching algorithms in general, when the distinguishing attribute of signal vs. spurious features that the former have invariant distributions across all environments. In section 6 we empirically show that adding the core features of IFM to CORAL can improve test accuracy.
The following theorem states that the environment complexity of IFM is logarithmic in the spurious feature dimension. A proof sketch is given in Section
5

---

Theorem 4.1 (IFM upper bound). Under assumption 3.2 suppose at each round IFM uses - (1(1
training environments, with probability 1 exp (d,)), IFM terminates (log d,) rounds and outputs @ = w².
As remark, we note that there is in fact simple algorithm that achieves ((1) environment complex- ity under Assumption [.2 (see Appendix D). However, this algorithm is extremely brittle and reliant on very specific aspects of the data model (such as Gaussianity) and it cannot be extended other settings, whereas feature matching algorithms are regularly applied to general model architectures and real-world datasets. The goal of this paper provide theoretical justification for distribution matching algorithms and to investigate why they may outperform ERM and IRM, rather than to solve specific data model.
4.2 ERM and IRM still have linear environment complexity
In the previous section, we showed that IFM has low environment complexity thanks to the additional structure we imposed in our model. However, it is possible that this additional structure also allows ERM and IRM to succeed. These next two results demonstrate that this is not the case.
ERM has low test accuracy In contrast to IFM, ERM still suffers from linear environment complex- ity under Assumption 3.2 The first theorem says there are hard instances where the ERM solution has worse-than- thannrandoo performance on the test environments.
Theorem 4.2 (ERM lower bound). Suppose E S de parameters P1
E, = μu E R Eg + G.G... Then any unit-norm linear classifier which achieves on all training environments will suffer 0-1 error at least on every test environment with flipped spurious mean, where & is the standard Normal CDF.
Σ -oI (recall 2
accuracy 2 +
21pp: (a)
A complete proof of Theorem 4.2 is in Appendix B.3 Note that it is quite reasonable to assume that the ERM solution satisfies the accuracy condition. In particular, it is common to model the spurious features having much greater magnitude than the invariant features, since they have much greater dimensionalityy. For example, with a unit-norm mean we would expect || r/d, || |= ds/ d./d.
Then for r < d and 01, 02 = w11vvd) one can verify that 2/ /I l/m/iiiinn..... σ2) ) = ol /r/d) is very
close to 0, meaning the lower bound +b
2////111 is only slightly larger than 2- min((r).77))
IRM fails to learn invariant features Our next theorem proves that even under Assumption 3.2 IRM still not guaranteed to find an invariant predictor. We can show this by proving that when E S d,. we can find a featurizer that only uses spurious dimensions, i.e., U, E Rd,, such that u 1S, = U, μ
for all E Ex.. If so, the optimal predictor on top of features u. Zz, wº (u, - -l,,"""" μr is invariant across all e, and can therefore be the preferred solution IRM when the spurious features have greater magnitude than the invariant features on the training environments.
Theorem 4.3 (IRM lower bound). Suppose E. S dg. If -/.2..., ur E Rdd are linearly independent, then there exists U E Rdd, ||u,||2> 0, such that u 2u, u, for all e [E].
Proof Sketch. Observe that each environment provides an ellipsoidal constraint E, = {u, E Rd u Ezu, u² μu3 = 0). The origin is trivial intersection. We prove the existence of a non-trivial intersection using tools from differential topology. The key lemma is that the total number of intersection points between two manifolds of complementary dimensions k d - k is even when certain tranversality conditions hold. Using these techniques, show that In, Eel > for almost all matrices Σ) ... ΣE, long the means are linearly independent.
A complete proof of Theorem 4.3 is in Appendix B.4
20x()) hides logarithmic factors in D, r, de-

---

Proof sketch for the main upper bound Theorem 4.1
To argue that IFM outputs featurizer U1 ... UT that does not the spurious features, need to show that the right d, columns of matrix Ur U are all-zero. The main lemma below says that this happens with high probability if we match St(1) environments at every iteration,
Lemma 5.1. for all 1 St S T, | E. Z (c(log (D/d,) + log d=)), and U.,..., UT are the orthonormal matrices returned by IFM, then with probability 1 exp (-S((d))), Pt < (ft-1-r c for all I, and if we write Ur...Ui U,S = [A, B), where B E Ikrxd, then B = Orxd,
Theorem 4.1 follows from Lemma 5. as follows: We take 2 and the algorithm terminates in O d.) rounds. Therefore with an environment complexity of O(log de). we learn a feature extractor U UT ... U1 that does not use any spurious dimensions. Since U is orthonormal, it must contain all signal dimensions. The predictor on top of this representation uses all and only signal dimensions, so with high probability, IFM outputs i w".
The first step towards proving Lemma 5.1 to show that with high probability, any one-layer pro featurizer Q1 E R that uses only spurious dimensions cannot match feature covariances from ((dd//k kk) environments. If a featurizer U1 E Rrixd uses kn spurious dimensions, there is corresponding rank-kj featurizer Q1 E & that uses only spurious dimensions. So Lemma 5.2 implics that any U1 that matches covariances in & must use most d./ E2 spurious dimensions. We will then apply this argument recursively until we have O spurious dimensions.
Lemma 5.2 (Informal version of Lemma B.3 For any integer 2 S kn S d,/2, when |&11=
Ej Ω + max 1o, (15)
log ))) with probability 1 (exx)--
orthonormal Q E xd. satisfies that for some constant C1 E Rkiik,
Ve E [E1],
QEςQT= C1.
(5.1)
Section B. g appendix|B.2
gives a proof sketch. The formal statement Lemma B.3 and its proof can be found in
The next claim says that Lemma 5.2 can be applied iteratively, i.e. fixing a featurizer from previous iterations that uses kz-1 spurious dimensions, with high probability, any U, that matches features from Ω (k/-1/kg) new environments uses at most ky spurious dimensions.
Corollary 5.3 (Informal version of Corollary B.7). Suppose 2 I le S 14-1/2 S d./2. kk-1-ke 10*K** ( k
When | E. 2
d. ( for fixed orthonormal max 1, log (Ak- log k.-1)
D
PE R 1 with probability 1 (exp (-d,)), ), no orthonormal Q R satisfies Ve E [Eg], QPE:PTQ C. for some constant C2 Rr
The formal statement Corollary B.7 and proof can be found in Appendix B.2 Lemma 5.1 follows from iterative application of Corollary 5.3
shown in Appendix B.2
6 Experiments
In light of the differences between IFM and CORAL discussed section 4.1 we test several questions inspired by our theory: (Q1) Do feature matching algorithms (IFM and CORAL) have much smaller environment complexity compared to ERM and IRM, with finite samples drawn from data models similar to our assumptions?
(Q2) Can decoupling feature matching and supervised training of the classifier (IFM) improve over joint training (CORAL)? (Q3) For neural network featurizers, can matching feature distributions at multiple layers improve over matching at only the last layer (naive CORAL)? (Q4) Can matching disjoint sets of environments at each layer perform as well as matching all environments all layers? (Q5) Is it important to shrink feature dimensions? We use two tasks to investigate those questions empirically. Appendix C contains additional details.
Gaussian dataset is a binary classification task that closely reflects our assumptions in section
Noised MNIST is a 10-way semi-synthetic classification task modified from LeCun and Cortes 23 Conditioned on the class, we add noise with identical mean but changing covariances across training environments. In test environments, the noise is uncorrelated with the label.
7

---

Avg Test Accuracy for Gaussian Dataset
Avg Test Accuracy for Noised MNIST
80
100
80
ERM
75
BRM
RM
RM
CORAL
70
CCRAL (anly
60
4ORAL
match last
IFM (1-layer,
65
layerl
motch all
CGRAL
40
Imatch-disjainti
(3-ayer.
60
CORAL
20
IFM (3-layer,
(match-at al)
match-disciint
55
0-
50
3
6
9
12
15
1
3
6
Number of Training Envs
Nurber of Layers in MLP
Figure 1: For Gaussian dataset, our algorithm IFM achieves highest test accuracy with the same number of training environments.
Figure For Noised MNIST, matching feature distributions from multiple layers improves over naive CORAL across different architectures.
Algorithms and architectures. For Gaussian dataset we use linear predictors. IRM follows the implementation in Arjovsky et al. [3: CORAL jointly minimizes average supervised loss on training environments and Leoral, which is the average of squared distances in conditonal feature means (in 12 norm) and covariances (in Frobenius norm) between adjacent training environments; CORAL+ON adds orthonormal penalty loss Lon(U)
||U' 11 where U is the featurizer; IFM is our Algorithm 1. where for each layer U, the training objective is LU,U,) = Lc + Ag We test IFM with I vs. 3-layer featurizers, either matching all (match-all) or a disjoint set of training environments (match-disjoint) at each layer.
For Noised MNIST we use ReLU networks with 1-6 layers. Although our theory for IFM is only for the linear setting, in practice we can extend the idea to deep network by training the classifier head together with feature matching at different disjoint) in Figure 2 are natural extensions of IFM to nonlinear models. They match at all layers post-activation, using either all (match-all) or a disjoint subset of training environments (match- disjoint) per layer. Figure
CORAL (match-all) and CORAL (match-
that they outperform naive CORAL, which matches features only the last layer. Thus, the core ideas of IFM may be used to improve performance of existing feature matching algorithms with nonlinear models and realistic datasets (e.g. DANN, MMD).
Results. (Q1) Figures I and 2 show that IFM and CORAL have much smaller environment complexity compared to ERM and IRM in both datasets. (Q2) In Gaussian dataset, IFM improves over CORAL. (Q3) For Noised MNIST, matching feature distributions at multiple layers (CORAL match- all, CORAL match-disjoint) improves over matching at only the last layer (CORAL). (Q4) In both datasets, matching disjoint sets of environments each layer (IFM match-disjoint, CORAL match- disjoint) is almost as good as matching all environments at all layers (IFM match-all, CORAL match- all) while saving computation. (Q5) For Noised MNIST (Table 2 in Appendix C). shrinking feature dimensions is crucial for the advantage of feature matching at multiple layers, c.g. matching features layers with widths [24, 24, 24] does not significantly improve over matching features at the last layer (CORAL). Overall, our results suggest that practitioners may benefit from feature matching algorithms when the data is similar to our assumed model, and may get additional advantage via matching at multiple layers with diminishing dimensions, echoing existing empirical works [27 29
7 Conclusion
This work presents the first domain generalization algorithm which provably recovers an invariant predictor with a number of environments that scales sub-linearly with the spurious feature dimension. Our results demonstrate that generalization which does not suffer from the "curse of dimensionality" possible, and based on our theory we believe the use an iterative approach is key insight which could lead to additional positive results for out-of-distribution generalization. Notably, this work also represents the first theoretical justification for the empirical success of existing algorithms which use feature distribution matching. However, there remains much room for improvement: our results
8

---

for linear data model with fairly special assumptions, It would be interesting to analyze data models in which the observables are non-linear function of the latent variables.
Acknowledgements
We would like to thank Colin Wei, Sang Michael Xie, Ananya Kumar, and Ciprian Manolescu for helpful discussions. YC supported by Stanford Graduate Fellowship. MS is supported by NSF and Stanford Graduate Fellowships. TM acknowledges support of Google Faculty Award and NSF IIS 2045685, AR and ER acknowledge support from the NSF via IIS-221 1907 and the CMU/PwC DT&I Center. AR acknowledges support of an Amazon Research Award.
References
[1] K. Ahuja, E. Caballero, D. Zhang, J.-C. Bognon-Auder Y. Bengio, I. Mitliagkas, and I. Rish. In- variance principle meets information bottleneck for -fitt--distiitttion generalization. Advances in Neural Information Processing Systems, 34:3438 3450, 2021.
[2] K. Abuja, Wang, A. Dhurandhar, Shanmugam, and K. R. Varshney, Empirical or invariant risk minimization? sample complexity perspective. In International Conference Learning Representations, 2021. URL htttp::/ http:://ppparrreeeee..eerrrrrdrd jx-joAAAAcy
[3] M. Arjovsky, L.. Botton, I. Guirajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv: 1907.02893 2019.
[4] B. Aubin, A. Stowik, M. Arjovsky, L. Bottou, and D. Lopez-Paz. Linear unit-tests for invariance discovery. arXiv preprint arXiv:2102.108677 2021.
[5] S. Beery, G. Van Horn, and Perona. Recognition in terra incognita. Proceedings of the European Conference on Computer Vision (ECCV), September 2018.
[6] A. Bellot and M. van der Schaar. Generalization and invariances the presence of unobserved confounding, arXiv preprint arXiv:2007.10653, 2020.
[7] S. Ben-David, Blitzer, K. Crammer, A. Kulesza, Pereira, and J. W. Vaughan. A theory of learning from different domains. Machine learning, (9(1 11:11)):5111155 2010.
[8] G. Blanchard, G. Lee, and C. Scott. Generalizing from several related classification tasks to a new unlabeled sample. In Advances in Neural Information Processing Systems, volume 24. Curran Associates, Inc., 2011. URL http:::/
www..//p..ccppppee//000.0 pdf
file/ bbtb:/eeeceeaaa0222222eeeeaa 16897 1633
[9] M. Chevalley, C. Bunne, A. Krause, and S. Bauer. Invariant causal mechanisms through distribution matching. arXiv preprint arXiv:2206. //66, 2022.
[10] V. H. de la Peña and S. J. Montgomery- Decoupling inequalities for the tail probabilities of multivariate u-statistics. The Annals of Probability, pages 806-816, 1995.
[11] C. Fang, Y. Xu, and D. N. Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. In Proceedings the IEEE International Conference on Computer Vision, pages 1657-1664, 2013.
[12] Y. Ganin, E. Ustinova, H. Ajakan, P. Germain, H. Larochelle, F. Laviolette, M. Marchand, and V. Lempitsky. Domain-adversarial training of neural networks. The journal of machine learning research, 17(1):209622)) 2030, 2016.
[13] A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Schölkopf. and A. Smola. A kernel two-sample test. J. Mach. Learn. Res., 13(m71)):233-733, Mar. 2012. ISSN 1532-4435.
[14] V. Guillemin and Pollack. Differential topology, volume 370. American Mathematical Soc., 2010.

---

[15] I. Gulrajani and D. Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2021. URL htppp//////ooeeeeiiee wtw.eenrriiiiin.eefrrr.rr net/forum?
1d= Id=1QaXeXDoWEI
[16] A. Hyvarinen and H. Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ica. Advances in neural information processing systems, 29, 2016.
[17] A. Hyvarinen, H. Sasaki, and R. Turner. Nonlinear ica using auxiliary variables and gener- alized contrastive learning, In K. Chaudhuri and M. Sugiyama, editors, Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pages 859-868. PMLR, 16-18 Apr 2019.
[18] W. Jin, R. Barzilay, and T. Jaakkola. Domain extrapolation via regret minimization. arXiv preprint arXiv:2006.03906 2020.
[19] P. Kamath, A. Tangella, D. Sutherland, and N. Srebro. Does invariant risk minimization capture invariance? In A. Banerjee and K. Fukumizu, editors, Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pages 4069-4077. PMLR, 13-15 Apr 2021. URL http: /proceedings..o mlr wttpp////////aaaaatit111111 /
html
[20] I. Khemakhem, D. Kingma, R. Monti, and A. Hyvarinen. Variational autoencoders and nonlinear iea: A unifying framework, International Conference on Artificial Intelligence and Statistics, pages 2207-2217. PMLR, 2020.
[21] D. Krueger, E. Caballero, J.-H. Jacobsen, A. Zhang, J. Binas, D. Zhang, R. Le Priol, and generalization via risk extrapolation (rex). In International
A. Courville. Out-of-distrrbure Conference on Machine Learning, pages 5815-5826. PMLR, 2021.
[22] B. Laurent and P. Massart, Adaptive estimation of a quadratic functional by model selection. Ann. Statist., 2865) 10 2000. doi: 10.121 ///11/////////////// URL http:///x
doi.org/ 10. 120.1121//00/11555573999
[23] LeCun and C. Cortes. MNIST handwritten digit database. 2010. URL http://yann. lecun com/ homm/exd//aie ///e///mmmmmm//eeeeeeeeeeeeeee/e mnist/
[24] D. Li, Y. Yang, Y.-Z. Song, and T. M. Hospedales. Deeper, broader and artier domain gen- cralization. In Proceedings of the IEEE international conference on computer vision, pages 5542-5550, 2017.
[25] H. Li, S. J. Pan, S. Wang, and A. C. Kot. Domain generalization with adversarial feature learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 5400-5409, 2018.
[26] Y. Li and Y. Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. In Advances in Neural Information Processing Systems, pages 8157-8166. 2018.
[27] M. Long, Cao, Wang, and M. Jordan. Learning transferable features with deep adaptation networks. In F. Bach and D. Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 97 105, Lille, France, 07-09 Jul 2015. PMLR. URL htt.:tttt///prrceding.s mlr www..prsss/v3.7 long15. html
[28] M. Long, Z. CAO, J. Wang, and M. I. Jordan. Conditional adversarial domain adap- tation. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Cur- ran Associates, Inc., 2018 htt. pdf
[29] Z. Luo, Y. Zou, J. Hoffman, and L. F. Fei-Fei. Label efficient learning of transferable rep- resentations acrosss domains and tasks. In I. Guyon, U. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URI http::/ttttt::://teeetsee cc/paper/2017/Fin http:///0000oo.oooooooooeoaeeeaa Papple-aaappppeee pdf
neurips neurips
10

---

[30] D. Mahajan, S. Tople, and A. Sharma. Domain generalization using causal matching. arXiv preprint arXiv:2006.07500 07500, 2020.
[31] Y. Mansour, M. Mohri, and A. Rostamizadeh. Multiple source adaptation and the rényi divergence. arXiv preprint Ntt5:2200.2222 2012.
[32] K. Muandet, D. Balduzzi, and B. Schölkopf. Domain generalization via invariant feature representation. S. Dasgupta and D. McAllester, editors, Proceedings of the 30th International Conference Machine Learning, volume 28 of Proceedings of Machine Learning Research, pages 10-18, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL URL htLLittt://prcceeding mlr. ppess/v28. 21128//mannee133 .hmml
[33] A. Pajor. Metric entropy of the grassmann manifold. Convex Geometric Analysis, 34:181-188, 1998.
[34] X. Peng. Q. Bai, X. Xia. Z. Huang, K. Saenko, and B. Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE International Conference on Computer Vision, pages 1406-1415, 2019.
[35] A. Rame, C. Dancette, and M. Cord. Fishr: Invariant gradient variances for generalization. International Conference on Machine Learning, pages 18347-18377. PMLR, 2022.
[36] A. Robey, G. J. Pappas, and H. Hassani. Model-based domain generalization. Advances in Neural Information Prucessing Systems, 34:20210-20229, 2021.
[37] E. Rosenfeld, P. Ravikumar, and A. Risteski. The risks invariant risk minimization. In International Conference on Learning Representations, 2021. URL /ttpsttss///oonrrrview net/forum?dd -tt.://dd--IIIPPJ--42
[38] M. Rudelson and R. Vershynin. Hanson- Hanson-Wright Wright inequality and sub- gaussian concentration. Electronic Communications in Probability, 18(none): 9, 2013. doi: 10.1214/ECPv18-2200
URL htttts://too...... org/ 10. 1214/ECP v18-2865
[39] S. Sagawa, P. W. Koh, T. B. Hashimoto, and P. Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint
2019.
[40] D. Serre. Matrices: Theory and Applications. Graduate Texts in Mathematics. Springer, 2010. ISBN 9781441930101. URL.
[41] D. A. Spielman and S.-H. Teng. Smoothed analysis of algorithms: Why the simplex algorithm usually takes polynomial time. ACM, 51(3):385 May 2004. ISSN 0004-5411. doi: 10.1 145/990308..0 990310. URL htt.:ttts///oi.orggggg./ 10. 1145/990308 990310
[42] B. Sun and K. Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pages 143-450. Springer, 2016.
[43] B. Sun, J. Feng, and K. Saenko. Return of frustratingly easy domain adaptation. In Proceedings of the Thirtieth AAAI Conference on Artificial Intelligence, AAAI'16, 16, page 2058-2065. AAAI Press, 2016.
[44] R. Tachet des Combes, H. Zhao, Y.-.. Wang, and G. J. Gordon. Domain adaptation with condi- tional distribution matching and generalized label shift. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, vol- ume 33, pages 19276-19289, Curran Associates, Inc., 2020. URL http:ttts////oceeeiiiii heurips. htwp://...ee.eeooooo./erdaaaaaed
pdf
[45] R. Vershynin. Introduction to the non-asymptotic analysis of random matrices. Y. Eldar and G. Kutyniok, editors, Compressed Sensing. Theory and Applications, chapter 5, pages 210-268. Cambridge University Press, 2012.
[46] C. Xie, Chen, Y. Liu, and Li. Risk variance penalization: From distributional robustness to causality, arXiv preprint arXiv:2006.07544, 2020.

---

[47] G. Zhang, H. Zhao, Y. Yu, and P. Poupart. Quantifying and improving transferabiliity in domain generalization. Advances in Neural Information Processing Systems, 34: 10957-10997, 2021.
[48] M. Zhang, H. Marklund, N. Dhawan, A. Gupta, S. Levine, and C. Finn. Adaptive risk minimiza- tion: Learning to adapt domain shift. Advances in Neural Information Processing Systems, 34:23664-2344
2021.
[49] H. Zhao, R. T. D. Combes, K. Zhang, and G. Gordon. On learning invariant representations for domain adaptation. In Proceedings of the 36th International Conference on Machine Learning. pages 7523-7532. PMLR, 09-15 Jun 2019. URL zhao 19a0119.... html
mlr http:g/...........mee.eee press/v97)
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] See Section T.1
(b) Did you describe the limitations of your work? [Yes] See Section 7 and Appendix E
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See
Section B
(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] See Section 3
(b) Did you include complete proofs of all theoretical results? [Yes] See Section Appendix
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] Code attached
in supplemental file.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section 0
Appendix C
(c) Did you report error bars (c.g.. with respect the random seed after running experi-
ments multiple times)? [Yes] See figures 2
(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] See Appendix C
4. If you using existing assets (e.g., code, data, models) curating/releasing
assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cited LeCun and Cortes 233
(b) Did you mention the license of the assets? [Yes] See Appendix
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(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 participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
12


Reviews:
[{"review_id": "T3tqXWhWbrJ", "reviewer": "Reviewer_Fzzp", "summary": "This paper analyzes the environmental complexity in domain generalization: how many environments are needed in order to recover an invariant predictor? While most existing papers give pessimistic results on domain generalization algorithms, this work shows positive theoretical results on the sublinear environment complexity. However, the data model may not be practical for real usage.\n\n--post rebuttal--\nThank you for the detailed response and the effort for improving the draft and adding discussions. From the new experiments, it seems that the IFM + CORAL does not improve the performance of CORAL (or very slightly). Therefore, it remains questionable whether the proposed method can be applied in practice, or it could be that the data model studied is too simple to represent more realistic tasks such as PACS/VLCS/OfficeHome. Although the theoretical analysis seems interesting, my main concerns are not addressed. Finally, understanding the non-linear feature embedding is more interesting (though challenging) and I would expect some theoretical results on this case in the revised draft.\n\n--post post rebuttal--\nAfter checking the response of Reviewer D1nV I decide to update my score due to the proof of the IRM lower bound which uses tools from differential topology. I haven't checked the proof in too much detail though.", "questions": "See the questions above.", "limitations": "As the authors acknowledge in Appendix D, the main limitation is the applicability of IFM in realistic datasets. ", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. This work is novel and interesting to the domain generalization field, and analyzes the theoretical model proposed in Rosenfeld et al. The data generation is anti-casual, we first generate the (binary) uniform label distribution and then generate the invariant and spurious features. Finally, the input variable is generated. The novelty part in this work is the addition of a Gram matrix from Gaussian.\n2. Based on the data model, the authors propose Iterative Feature Matching, or IFM, for finding invariant featurizer and classifier in $O(\\log d_s)$ rounds with $\\tilde{O}(1)$ training environments per round. Therefore the environmental complexity is smaller than the $O(d_s)$ complexity of ERM/IRM.\n3. The presentation is clear with solid theoretical results, it involves interesting techniques from probability theory, statistics and differential topology, which might be useful for later theoretical results in this field.\n\nWeaknesses:\n1. The setting borrows from Rosenfeld et al, which shows the reason why IRM may fail in some simple cases. However, the goal of the current paper is to show an DG algorithm that works in practice. As much as I appreciate the theoretical contributions about the positive results of IFM, in practice (e.g. on MNIST), the algorithm does not work. Moreover, as mentioned in Appendix C, there is even a $O(1)$ environmental complexity algorithm based on Assumption 3.2, which further adds doubts to the applicability of this assumption. If IFM does provide intuitions into algorithmic design, the authors need to at least show some more practical examples.\n2. The lower bounds of ERM/IRM do not seem novel compared to Rosenfeld et al. The only difference is the assumption 3.2 vs. assumption 3.1. In Line 149-150, the authors claim that the analysis applies to all but a measure-zero set of covariances, but I'm not quite convinced since the Gram matrix is built from standard normal and the $\\bar{\\Sigma}$ is bounded.\n3. A practical algorithm for finding the maximum dimension $r_t$ in Alg 1 is needed. This is very important for practitioners to apply this algorithm.\n4. In the experimental section, the main contribution part is 1) addition of orthonomality of the featurizer in the Gaussian model; 2) matching the sufficient statistics of each layer vs. only matching the last layer. This experiment may be too far away from the linear featurizer and the data model presented in the theoretical analysis. If matching the representation of each layer works for CORAL/DANN/MMD in practice, could the authors try some experiments or prove some theoretical results? \n5. In line 203-205, the goal is to provide a theoretical justification of the feature matching algorithm, rather than to solve a specific data model, but the main theoretical results are based on this model.\n6. The IFM algorithm resembles Independent Component Analysis (ICA). Discussions and comparisons of the similarities and differences are needed.\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "F5GHhG8QdQZ", "reviewer": "Reviewer_D1nV", "summary": "This paper studies the domain generalization (DG) setting.  The authors construct a simple yet fundamental data generating procedure for DG, which shares similarities with that of _The Risks of Invariant Risk Minimization_ (henceforth RoIRM).  The authors propose an iterative algorithm which requires only logarithmically many training domains in the spurious feature dimension $d_2$ to generalize to particular test environments given infinitely many samples per domain.  This result stands in contrast to results from other works (including RoIRM)l, which show that IRM needs $O(d_s)$ environments to generalize in the linear setting.  The authors also provide several toy experiments to illustrate their results.", "questions": "**Test environment generation.** I won't list this as a weakness per se, but I do have questions regarding how the test environments are generated.  One natural question is whether (or to what extent) this model of test environment generation is necessary for achieving OOD generalization in log(d_s) training environments.  Let's say that I was to select completely new means $\\mu^e_2$ and covariances $\\Sigma_2^e$; then let \n\n$$Z_2 \\sim\\mathcal{N}(Y \\cdot \\mu_2^e , \\Sigma^e_2).$$\n\nWould there be any hope of generalizing in this setting?  I think that this question is important because in practice, we would expect to receive data from *new domains* with new underlying parameters.  This is not to say that the setting analyzed here is less relevant; on the contrary, it seems evident that there needs to be some sort of connection between training and test domains to achieve OOD generalization.\n\n**Dimension of style.** The authors say that\n\n> \"stylistic variations between domains are likely high-dimensional.\"\n\nIs there any way to justify this?  In the case of the standard ColoredMNIST benchmark, this seems to be somewhat untrue given that the background color correlation could be controlled by a very low-dimensional parameter.  Furthermore, given the rise of large generative models which can capture semantic changes in low-dimensional feature spaces, I would tend to push back on this statement ever so slightly.\n\n**Tranversality.**  It would be helpful if the authors could explain a bit more what they mean by this _tranversality_ condition.  I anticipate that the definition is technical, but even some intuition for this condition would be helpful.\n\n**I.I.D.-ness.** I assume that the labels are being drawn i.i.d. from a Bernoulli, as are the latents.  This should perhaps be stated more explicitly for completeness.\n\n**Why does $r_t$ shrink to $r$?**  It would be great if the authors could help me to understand why $r_t$ will eventually shrink to $r$.  What is stopping it from failing to shrink?\n\n**Uniqueness.** One small question I have WRT Thm. 1 is whether $w^\\star$ is necessarily unique.\n\n**\"Simple\" algorithm.**  What is the basis for including this so-called \"simple\" alternative algorithm in the main text?  If it is mentioned, I would expect the authors to compare with it in the experiments section.  But the inclusion of this algorithm seems to be a bit orthogonal to the contribution of the paper.\n\n**Class imbalance.**  If $\\eta\\neq 1/2$, how would Thm. 4.2 change?  Would the condition that the test risk under the 0-1 loss be larger than $1/2$ become $\\eta$ (or perhaps $(1-\\eta)$?\n\n**ERM results.** In Thm. 4.2, I thought that we were operating under the assumption that $\\bar{\\Sigma^e_2}$ could be _adversarially_ chosen, and that it would not be diagonal.  It seems that one would want the assumptions of the three theorems to align; so unless I am misunderstanding something, this discrepancy calls into question whether all three results can hold under a common assumption.\n", "limitations": "The discussion in Appendix D is sufficient in my eyes.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n---\n\n**Problem setting.**  I liked quite a bit about this paper.  I felt that first and foremost, it attempts to theoretically study what makes domain generalization hard, and in what settings we expect to generalize OOD.  This is essential, as it is not well understood in the literature when domain generalization is even possible.  The data generating procedure discussed in Section 3 seems quite natural, and it certainly extends the framework of RoIRM to the setting of (more) general covariance matrices.  The other innovation over RoIRM here seems to be the generation of the test environments.  \n\n**Definition of environmental complexity.**  This seems like quite a natural definition of complexity in DG; certainly removing the caveat that in practice we get a finite sample from each environment helps to illuminate the structure of the problem.  It would be interesting to see future works study this definition for other algorithms/data generating models.\n\n**Theoretical work.** Thm. 1 seems to be a strong result here, because it proves the existence of an algorithm which generalizes with $\\log(d_s)$ in a particular linear DG setting.  To my knowledge, such a result was not known in the literature before this work.  I believe that the proof techniques may be of independent interest here, which adds to the contribution of this paper.\n\n### Weaknesses\n\n---\n\n**Over-simplification of the field of DG in the intro.**  One minor point: The introduction reads as if the *standard* assumption to make in DG is that the distributions of labels given features should be environment dependent (lines 21-22).  Furthermore, there is a claim that algorithms that rely on these assumptions \"dominate\" DG benchmarks (lines 23-24).  I think that both of these claims are in doubt.  There are many other works that make different assumptions on the set of stable or invariant features (see e.g. [Zhang et al., 2021] and [Robey et al., 2021], which assume domain/measurement shift).  I think this points to the fact that this paper could benefit from a related work section, which is notably absent in this submission.  WRT the point about feature-matching algorithms \"dominating\" the leaderboards, this seems untrue based on recent surveys: [Gulrajani & Lopez-Paz, 2020] and [Koh & Sagawa et al., 2021].  Both of these works show that ERM is a rather strong baseline relative to all other algorithms, especially relative to feature-matching algorithms.\n\n**The term \"feature matching.\"**  One point of confusion throughout the work was referring to \"matching.\"  I think it would help if the authors could formally define what it means to \"match\" features.  The reason I bring this up is that it gets rather confusing in the non-linear case how features are being \"matched.\"  The paper seems to assume that the reader is a priori going to know what this means, but I imagine that many readers may become confused at this point.\n\n**Algorithm description.**  The description of the algorithm that is actually run in practice is confusing, and crucial details are missing.  One question is: How are the subsets selected?  Is $T$ a hyperparameter?  If so, how is it tuned?  Furthermore, which Alg. 1 shows us the basic steps of the algorithms, it's not clear how one would actually go about implementing this method.  In particular, the step in line 5 seems challenging.  In the end, the authors appear to tweak the implementation of CORAL by removing the ERM term.  I think it would be reasonable to expect a more robust discussion of the implementation of these methods -- which also seems to be missing from the appendices.  I am listing this as a weakness because if a contribution of this paper is to propose this algorithm as a new method, which seems to be the case based on this:\n\n> \"This work presents the first domain generalization algorithm which provably recovers an invariant predictor with a number of environments that scales sub-linearly with the spurious feature dimension.\"\n\nthen the paper should be responsible for connecting the pseudocode in Alg. 1 to an implementable optimization problem.\n\n**Why doesn't CORAL work?** Related to the previous point, after reading the paper I am still unsure of why CORAL fails in the experiments whereas IFM succeeds.  Is it the fact that IFM optimizes for each $U_t$ separately in rounds, or is it the dropping of the ERM term that helps?  Perhaps this is already answered somewhere, but it was not obvious to me after reading.\n\n**Experiments.** The experiments feel like a weak point of this paper, with the caveat that this paper seems to be more of a theoretical paper than an experimental one.  Even still, I think it wouldn't be unreasonable to expect a stronger set of experiments here.  Some comments:\n\n* The Noisy MNIST dataset doesn't seem to satisfy the assumption concerning $Z_2$ below line 136.\n* The authors claim that\n\n> \"Our results in Section 7 suggest that practitioners may benefit from feature matching algorithms when the distinguishing property of the signal feature is indeed conditional distributional invariance\"\n\nHowever, the authors do not run their proposed algorithm on any standard benchmarks, making it difficult to discern whether or not IFM would be useful for practitioners.  Indeed, it looks like IFM isn't even run on the Noised MNIST dataset, which raises concerns about whether IFM on its own works (without the ERM loss term).\n* The authors refer to $U_t$ as a \"layer\" (line 325), which is confusing to me.  Is this supposed to be related to the layers of a DNN?\n* It's unclear how the algorithm works in the nonlinear case, i.e. when you need to do feature matching for different layers.  As I mentioned above, a better understanding of how the algorithms are implemented would lead to stronger insights from the experiments.\n\n### Overall evaluation\n\n---\n\nOverall, I liked this paper.  I think that it makes a strong theoretical contribution, and that it studies an important and understudied aspect of DG.  I believe that the paper is generally well written, and I appreciate that although this is more of a theoretical paper, the authors ran some experiments.  This being said, there are a number of drawbacks that I would like to see addressed.  Firstly, the algorithmic description and experimental evaluation are weak compared to the theory, but I think both of these points could be addressed (see above).  Another weakness is the lack of a related work section, which ideally would more broadly discuss recent developments and algorithms proposed for the DG setting.  All of this being said, I'm leaning toward accepting this paper.\n\n### Citations\n\n---\n\n[Zhang et al., 2021] Zhang, Marvin, et al. \"Adaptive risk minimization: Learning to adapt to domain shift.\" Advances in Neural Information Processing Systems 34 (2021): 23664-23678.\n\n[Robey et al., 2021] Robey, Alexander, George J. Pappas, and Hamed Hassani. \"Model-based domain generalization.\" Advances in Neural Information Processing Systems 34 (2021): 20210-20229.\n\n[Gulrajani & Lopez-Paz, 2020] Gulrajani, Ishaan, and David Lopez-Paz. \"In search of lost domain generalization.\" arXiv preprint arXiv:2007.01434 (2020).\n\n[Koh & Sagawa et al., 2021] Koh, Pang Wei, et al. \"Wilds: A benchmark of in-the-wild distribution shifts.\" International Conference on Machine Learning. PMLR, 2021.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EVCagPwO9lL", "reviewer": "Reviewer_ebXq", "summary": "This paper proposes an iterative feature matching algorithm to solve domain generalization problem with logarithmic environment complexity, compared to linear complexity of ERM and IRM. The authors also provide theoretical justification and empirical evaluation.", "questions": "Refer to Weakness", "limitations": "Refer to Weakness", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength: \n(a)\tThe investigated problem is of great importance in invariance-based domain generalization and out-of-distribution generalization.\n(b)\tThe proposed method has sound theoretical justification.\nWeakness:\n(a)\tThe experimental part is relatively weak, several more baselines (refer to Domainbed) and benchmarks (Colored-MNIST or other real-world datasets) should be included for better insight and significance.\n(b)\tThe organization of the manuscript could be improved by moving proofs into appendix rather than piling them up in main body.\n(c)\tThe authors should elaborate more on the gap between linear setting in theoretical analysis and real world applications with more complex models.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept