Iterative Feature Matching: Toward Provable Domain Generalization with Logarithmic Environments
Anonymous Author(s)
AffiliationAddress email
Abstract
Domain generalization aims at performing well on unseen test environments withdata from a limited number of training environments.
Despite a proliferation ofproposed algorithms for this task, assessing their performance both theoreticallyand empirically is still very challenging.
Distributional matching algorithms suchas (Conditional) Domain Adversarial Networks [12, 28] are popular and enjoyempirical success, but they lack formal guarantees.
Other approaches such as
Invariant Risk Minimization (IRM) require a prohibitively large number of trainingenvironments—linear in the dimension of the spurious feature space d s —evenon simple data models like the one proposed by Rosenfeld et al.
Under avariant of this model, we show that ERM and IRM can fail to find the optimalinvariant predictor with o ( d s ) environments.
We then present an iterative featurematching algorithm that is guaranteed with high probability to find the optimalinvariant predictor after seeing only O (log d s ) environments.
Our results providethe first theoretical justification for distribution-matching algorithms widely usedin practice under a concrete nontrivial data model.
Introduction
Domain generalization aims at performing well on unseen environments using labeled data from alimited number of training environments [8].
In contrast to transfer learning or domain adaptation,domain generalization assumes that neither labeled nor unlabeled data from the test environmentsis available at training time.
In the empirical literature, invariance of a “signal” feature distributionconditioned on the label, i e. P (Φ( x ) | y ) , is the underlying assumption in widely adopted algorithmssuch as Correlation Alignment (CORAL) [43, 42], Maximum Mean Discrepancy (MMD) [13, 25] and(Conditional)
Domain Adversarial Networks [12, 28].
Distribution-matching algorithms regularizethe distribution of intermediate representations of examples from different environments to be closerin some metric.
They are among the top-performing methods for domain generalization benchmarkslike PACS [24], VLCS [11] and DomainNet [34], but they lack formal guarantees.
Previous empiricalworks usually characterize the performance of these algorithms using generalization bounds basedon divergence between domain distributions [7, 31], but those bounds are vacuous for the typicalbenchmarks for domain generalization and thus cannot explain their success.
For example, whendifferent environments have disjoint supports (as is the case in the above-mentioned benchmarks)the H -divergence [7] is 1 even for the linear hypothesis class.
To obtain useful guarantees, it isnecessary to study data models that encode structure reflective of settings of interest.
Prior worksattempting to theoretically characterize the performance of feature matching algorithms emphasizelower bounds [49, 44].
In this work, we seek to give the first positive theoretical justification forfeature matching algorithms.
Another common assumption in the literature is invariance of the label distribution conditioned onthe signal features.
Invariant Risk Minimization (IRM) [3] assumes E [ y | Φ( x )] is invariant, andfollow-up works assume invariance of higher moments [46, 18, 30, 6].
However, empirical resultsfor these algorithms are mixed: Gulrajani and Lopez-Paz [15], Aubin et al.
[4] present experimentalevidence that these methods do not consistently outperform ERM for either realistic or simple lineardata models, when fairly evaluated.
Recent theoretical works [37, 19] also question the theoretical foundations of IRM and its variants,shedding light on their failure conditions.
These works study specific data generative models;a common assumption is that, conditioned on the label, some invariant features have identicaldistribution for all environments, and other spurious features have varying distributions acrossenvironments.
The goal of these methods is then to obtain an invariant predictor , i.e. a classifierwhich uses only the invariant features.
These works also often assume each training environmentcontains infinite samples.
Thus, the central measure of domain generalization is the numberof environments needed to recover an invariant predictor—we refer to this measure as theenvironment complexity of a learning algorithm.
Rosenfeld et al.
[37] show even for a simplegenerative model and linear classifiers, the environment complexity of IRM—and other objectivesbased on the same principle of invariance—is at least as large as the dimension of the spuriouslatent features, d s .
Further results by Kamath et al.
[19], Ahuja et al.
[2] also point to a linearenvironment complexity.
Although the models in these works are simple, they help elucidate whyexisting algorithms fail and can help inform better algorithmic design.
IRM’s linear environment complexity is prohibitive for realistic applications.
Domain generalizationbenchmarks have fewer training environments than spurious features dimensions—PACS, VLCS,
DomainNet have 4-6 total environments but the stylistic variations between domains are likely highdimensional.
In this paper, we show that a variant of distributional matching algorithm (Algorithm 1)finds the optimal invariant predictor with sublinear (even logarithmic) environment complexity.
1.1 Our contributions
Conceptually, we propose a “smoothed covariance” extension of the data model in Rosenfeld et al.[37] , which is rich enough to establish an environment complexity separation between IRM/ERMand feature matching algorithms.
More precisely, we show that ERM and IRM can fail to find theoptimal invariant predictor with fewer than d s + 1 training environments (Theorems 4.2, 4.3), whilea relatively simple and natural algorithm based on iterative feature matching (IFM) finds the optimalinvariant predictor with O (log d s ) environments.
Moreover, our analysis highlights the value of iterative matching, much similar to matching atmultiple layers in a deep neural network.
IFM iteratively projects the features to a lower dimensionwhile matching the label-conditioned feature distributions on a small, disjoint subset of the trainingenvironments.
A projection inducing invariance in the non-invariant features across one subset ofenvironments is unlikely to do so for a different subset, so each projection only removes (a constantfraction of) spurious feature dimensions with high probability.
Intuitively, the iterative schemeprevents the different environments from “colluding” to create a solution which depends on spuriousfeatures.
As a result, IFM recovers the optimal invariant predictor after O (log d s ) rounds and uses
O (1) environments per round, thus requiring O (log d s ) environments.
Our techniques for proving the upper and lower bounds may be of independent interest.
The upperbound (Theorem 5.1) is proved by showing that any projection matrix that uses many spuriousdimensions cannot match feature covariances in a large set of training environments.
This is done viaintricate 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 anon-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 ofcomplementary dimensions k, d − k is even when certain tranversality conditions hold, implying thatthe origin cannot be the only solution.
Finally, to corroborate the advantages of the proposed algorithm, we perform experiments on a
Gaussian dataset and a semi-synthetic Noised MNIST dataset [23], where the background noisespuriously correlates with the label.
Our results in Section 6 suggest that practitioners may benefit from feature matching algorithms when the distinguishing property of the signal feature is indeedconditional distributional invariance, and may get additional advantage via matching at multiplelayers with diminishing dimensions, echoing existing empirical observations [27, 29].
Preliminaries2.1 Domain generalization
In domain generalization, we are given a set of E training environments E tr indexed by e ∈ [ E ] , 1 anda set of test environments E ts .
For environment e we have n examples { ( X ei , Y ei ) } ni
=1 drawn fromthe distribution P e .
In this work we study the infinite sample limit n → ∞ so as to separate the effectof limited training environments from that of limited samples per environment, as is done in previoustheoretical works [37, 19].
Let X , P , Y denote the space of inputs, intermediate features, and labels.
For a featurizer Φ : X → P and classifier w : P → Y , their risk on environment e is denoted by
R e Φ ,w = E ( X,Y ) ∼ P e
[ l ( w ◦ Φ( X ) , Y )] for any common loss function l .
In this paper we focus on Y ={± 1 } , linear featurizers Φ( X ) =
UX for U ∈ R k × d , and unit-norm predictors (cid:98) Y = sgn ( w ⊤ Ux )where w ∈ R k and ∥ w ⊤ U ∥ 2 = 1 for some feature dimension k ≤ d chosen by the algorithm.
Apredictor’s 0-1 risk on environment e is denoted by R e U,w = Pr ( X,Y ) ∼ P e [ sgn ( w ⊤ UX )
̸
= Y ] .
Wefocus on unit-norm predictors because we evaluate on the 0-1 risk on test environments, which areinvariant to 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, ERMand IRM.
ERM learns a classifier that minimizes the average loss over all training environments,where l is any common training loss such as the logistic loss:
[Equation]
IRM learns a featurizer Φ( X
) ∈ R k
such that the optimal classifier on top of the featurizer is invariantacross training environments.
As we focus on linear classifier, it is equivalent to learning a lineartransformation
U ∈ R k × d such that it induces a classifier w that is optimal for all e ∈ E tr :
[Equation]
This is objective is not the same as feature distribution matching; IRM only tries to match the firstmoment.
Observe that this constrained objective is intended to solve a minimax domain generalizationproblem, as opposed to ERM which is typically viewed as minimizing the risk in expectation.
Problem setup
We first recall the data model from Rosenfeld et al.
We assume without loss of generality thatthe label
Y is uniformly randomly drawn from {± 1 } (extension of our theorems to Y = 1 withprobability η ̸ = 0 .
5 is straightforward).
Latent variable Z consists of invariant features Z 1 ∈ R rand spurious features Z 2 ∈ R d s where d s = d − r .
The number of spurious features can be muchlarger than the number of invariant features, i.e. d s ≫ r .
The input X ∈ R
d is generated via a lineartransformation of latent variable Z , i.e. X = SZ for a matrix S ∈ R d × d such that its left r columnshave rank r (so that there are r invariant dimensions).
For each training environment indexed by e ∈ [ E ] , the invariant features conditioned on Y are drawnfrom a Gaussian distribution with mean Y · µ 1 ∈ R r and nonsingular covariance Σ 1 ∈ R r × r .
Thespurious features conditioned on Y have mean Y · µ e 2 ∈ R d e and covariance Σ e 2 ∈ R d e × d e whereµ e ’s and Σ e ’s vary across e ∈
[ E ] .
The assumption of symmetric class center with respect to theorigin can also be relaxed.
Define µ e =
[ µ 1 , µ e 2 ] and Σ e =
[Σ 1 , 0 r × d s ; 0 d s × r , Σ e 2 ] .
The overall datamodel for training environments is summarized below:
[Equation]
Since the goal of invariant feature learning is to learn a predictor that only uses the invariant features,one reasonable measure for domain generalization is a predictor’s performance on test environmentswhere the spurious features Z 2 are drawn from a different distribution—in particular, they are usuallychosen adversarially.
A classifier that predicts using the spurious features will perform badly on suchtest environments.
When modeling the test environments, we consider the difficult scenario wherethere is one corresponding test environment for each training environment, whose parameters arethe same except that the spurious means are flipped.
Formally, for each environment e ∈ E tr weconstruct a corresponding test environment
e ′ ∈
E test where
[Equation]
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.
[Equation]
d s , where σ e is a scalar for an environment indexed by e .
We consider a generalized model where the covariances of spurious features for each environment isa generic random PSD matrix, instead of only random in scaling:
Assumption 3.2 (Data model for spurious covariances in this work) .Σ
e 2 ∼ Σ
e 2 + G e G ⊤ e , where Σ e 2 ∈ R d s × d s is arbitrary (and can be adversarial), and
[ G e ] i,j iid ∼ N (0 , 1) for all i, j ∈ [ d s ] .
Furthermore, max e ∥ Σ e 2 ∥ 22 ≤ D .
Assumption 3.2 allows the covariance matrix of the spurious features to be almost worst-case : itis an arbitrary matrix, plus a Gaussian perturbation.
This is a common assumption in algorithmiccomplexity called smoothed analysis [41], as we allow the parameters to be a random smoothing ofarbitrary parameters.
While our data model is simple, it has already been used as a sandbox to understand algorithms like
IRM [37, 19] because it captures some important aspects of real-life data like latent variables andcorrelations between the labels and the spurious features.
In the next section, we show that baseline algorithms ERM and IRM still have Ω( d s ) environmentcomplexity under assumption 3.2, whereas our iterative feature matching algorithm (Algorithm 1)requires only O (log d s ) training environments.
Note that the environment complexity of our algorithmonly 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 algorithmbased on iterative feature matching.
We then provide formal guarantees for its environment complexityin comparison to ERM and IRM.
Require: Invariant feature dimension r , total feature dimension d , number of training environments
E = |E tr | , infinite samples { ( X ei , Y ei ) } ∞ i =1 ∼ P e from each environment e ∈ E tr , an array of integers indicating the number of training environments to use at each iteration E 1 , E 2 , . . . . 1: r 0 ← d , t ← 0 . 2: while r t > r do 3: t ← t + 1 . Uniformly randomly sample E t training environments without replacement.
Binary search between 1 and r t − 1 to find the maximum dimension r t such that there exists orthonormal U t ∈ R r t × r t − 1 and C t ∈ R r t × r t , where for all e ∈ E t ,
[Equation]
(4.1) 6: Return a classifier on projected features that minimizes the average risk
[Equation]
4.1 Iterative feature matching algorithm
We hope to recover the invariant features by imposing constraints which are satisfied by onlythose features (i.e., they are not satisfied by the spurious features).
A natural idea is to match thefeature means and covariances across E tr .
Since µ 1 , Σ 1 are constant, any orthonormal featurizer
U ∈ R
r × d such that US has only non-zero entries in the first r rows yields invariant means USµ e andcovariances US Σ e S ⊤ U ⊤ .
Thus we need E large enough such that any U
′ ∈
R r × d using spuriousdimensions cannot match the means and covariances.
Informally, for each e ∈ [ E ] we get r × requations from matching covariances U Σ
e U ⊤ =
C , and we have r × d parameters to estimate in U .
Rough parameter counting suggests that if we match covariances of all E environments jointly, weneed at least E > d/r environments to find a unique solution.
Our key observation is that, due to theindependence of randomness in Σ e 2 , we can split E environments into disjoint groups E 1 , . . .
, E T , anduse E t to train an orthonormal featurizer that shrinks the feature dimensions from r t − 1 to r t .
Thus, ineach round we shrink the dimension by a constant factor using a constant number of environments.
The main theoretical challenge that remains is to show that in each iteration, with high probability,only spurious features are projected out.
This brings us to IFM (Algorithm 1), which proceeds in T
=
O (log d s ) rounds.
Starting with aninput dimension r 0 =
d , each round we learn an orthonormal matrix U t projecting features from r t − 1to r t dimensions so that the feature covariances after projection match across a fresh set of trainingenvironments.
In practice, for each choice of r t , starting from random initialization, we perform SGDon objective min U t ∈ R rt × rt − 1 (cid:80) e,e ′ ∈E t
∥ U t (Σ e − Σ e ′ ) U Tt ∥ 2 F + λ ∥ U ⊤ t U t − I ∥ 2 F until the objective isϵ -close to 0. To ensure that all invariant dimensions are preserved, we always find the projectionwith the maximum possible dimension r t that still matches the covariances (i.e. admits an ϵ -optimalsolution U t ), until we are left with only r dimensions.
Among algorithms used in practice, IFM is most similar to CORAL [43].
We study IFM since it ismore amenable to theoretical analysis.
The differences between IFM and CORAL are: first, CORALdoes not enforce that the featurizer is orthonormal; second, IFM learns to extract features in anunsupervised manner, whereas CORAL jointly minimizes the supervised loss and feature distributiondiscrepancy; third, IFM matches the feature distributions at multiple layers and uses a disjoint setof environments for each layer—this iterative process is necessary for the theoretical guarantees weprovide—but CORAL matches only at the last layer.
Despite these differences, our theoretical resultsserve as a justification for using feature matching algorithms in general, when the distinguishingattribute of signal vs. spurious features is that the former have invariant distributions across allenvironments.
In section 6 we empirically show that adding the core features of IFM to CORAL canimprove test accuracy.
The following theorem states that the environment complexity of IFM is logarithmic in the spuriousfeature 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 |E t | = ˜Ω(1) 199training environments, with probability 1 − exp ( − Ω( d s )) , IFM terminates in O (log d s ) roundsand outputs (cid:98) w = w ∗ .
As a remark, we note that there is in fact a simple algorithm that achieves O (1) environment complexity under Assumption 3.2 (see Appendix D).
However, this algorithm is extremely brittle and relianton very specific aspects of the data model (such as Gaussianity) and it cannot be extended to othersettings, whereas feature matching algorithms are regularly applied to general model architecturesand real-world datasets.
The goal of this paper is to provide theoretical justification for distributionmatching algorithms and to investigate why they may outperform ERM and IRM, rather than to solvea 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 additionalstructure 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 complexity under Assumption 3.2.
The first theorem says there are hard instances where the ERM solutionhas worse-than-random performance on the test environments.
Theorem 4.2 (ERM lower bound) .
Suppose E ≤ d s , parameters µ 1 ∈ R r , Σ
=
σ 21 I r , µ e 2 ∈ R d s ,Σ e 2 =
σ 22 I d s (recall Σ e 2 ∼ Σ e 2 + G e G ⊤ e ).
Then any unit-norm linear classifier which achievesaccuracy ≥ Φ (cid:16) 2 ∥ µ 1 ∥ min( σ 1 ,σ 2 ) (cid:17) on all training environments will suffer 0-1 error at least 12 on every testenvironment with flipped spurious mean, where Φ is the standard Normal CDF.
A complete proof of Theorem 4.2 is in Appendix B.3.
Note that it is quite reasonable to assume thatthe ERM solution satisfies the accuracy condition.
In particular, it is common to model the spuriousfeatures as having much greater magnitude than the invariant features, since they have much greaterdimensionality.
For example, with a unit-norm mean we would expect ∥ µ 1 ∥ 2 ≈ r/d, ∥ µ e 2 ∥
≈ d s /d .
Then for r ≪ d and σ 1 , σ 2 = ω (1 / √ d ) one can verify that 2 ∥ µ 1 ∥ / min( σ 1 , σ 2 )
=
o ( (cid:112) r/d ) is veryclose to 0, meaning the lower bound Φ (cid:16)
[Equation]
is only slightly larger than 12 .
IRM fails to learn invariant features Our next theorem proves that even under Assumption 3.2, IRMis still not guaranteed to find an invariant predictor.
We can show this by proving that when E ≤ d s ,we can find a featurizer that only uses spurious dimensions, i.e., u s ∈ R d s , such that u ⊤ s Σ
e 2 u s = u ⊤ s µ e 2for all e ∈ E tr .
If so, the optimal predictor on top of features u ⊤ s Z 2 , (cid:99) w e
=
( u ⊤ s Σ e 2 u s ) − 1 u ⊤ s µ e isinvariant across all e , and can therefore be the preferred solution IRM when the spurious featureshave greater magnitude than the invariant features on the training environments.
Theorem 4.3 (IRM lower bound) .
Suppose E ≤ d s .
If µ 12 , . .
. , µ E 2 ∈ R d
s are linearly independent,then there exists u s ∈ R d s , ∥ u s ∥ 2 > 0 , such that u ⊤ s Σ e 2 u s =
u ⊤ s µ e 2 for all e ∈
[ E ] .
Proof Sketch.
Observe that each environment provides an ellipsoidal constraint E e
= { u s ∈ R d s :u ⊤ s Σ e 2 u s − u ⊤ 2 µ e 2 = 0 } .
The origin is a trivial intersection.
We prove the existence of a non-trivialintersection using tools from differential topology.
The key lemma is that the total number ofintersection points between two manifolds of complementary dimensions k, d − k is even whencertain tranversality conditions hold.
Using these techniques, we show that | (cid:84) e E e | ≥ 2 for almostall matrices Σ 12 , . . .
, Σ E 2 , as long as the means are linearly independent.
A complete proof of Theorem 4.3 is in Appendix B.4.
Proof sketch for the main upper bound Theorem 4.1
[-the main upper bound Theorem 4.
-]
To argue that IFM outputs a featurizer U 1 . . .
U T that does not use the spurious features, we need toshow that the right d s columns of matrix U T . . .
U 1 S are all-zero.
The main lemma below says thatthis happens with high probability if we match ˜Ω(1) environments at every iteration,
Lemma 5.1.
If for all 1 ≤ t ≤ T , |E t | =
E t =
Ω ( c (log ( D/d s ) + log d s )) , and U 1 , . .
, U T arethe orthonormal matrices returned by IFM, then with probability 1 − exp ( − Ω( d s )) , r t − r <( r t − 1 − r + 1) /c
for all t , and if we write U T . . .
U 1 S =
[ A, B ] , where B ∈ R
r × d s , then B =
r × d s .
Theorem 4.1 follows from Lemma 5.1 as follows: We take c = 2 and the algorithm terminates in
T =
O (log d s ) rounds.
Therefore with an environment complexity of O (log d s ) , we learn a featureextractor U
= U T . . .
U 1 that does not use any spurious dimensions.
Since U is orthonormal, it mustcontain all signal dimensions.
The predictor on top of this representation uses all and only signaldimensions, so with high probability, IFM outputs (cid:98) w = w ∗ .
The first step towards proving Lemma 5.1 is to show that with high probability, any one-layerfeaturizer Q 1 ∈ R k 1 × d s that uses only spurious dimensions cannot match feature covariancesfrom ˜Ω( d s /k 1 ) environments.
If a featurizer U 1 ∈ R r 1 × d uses k 1 spurious dimensions, there is acorresponding rank- k 1 featurizer Q 1 ∈ R k
1 ×
d s that uses only spurious dimensions.
So Lemma 5.2implies that any U 1 that matches covariances in E 1 must use at most d s /E
1 spurious dimensions.
Wewill then apply this argument recursively until we have 0 spurious dimensions.
Lemma 5.2 (Informal version of Lemma B.3) .
For any integer 2 ≤ k 1 ≤ d s / 2 , when |E 1 | =
E 1 = Ω (cid:16)
[Equation]
max
[Equation]
1 , log
[Equation]
(cid:17) , log (cid:16)
[Equation]
s k 1 − 1 (cid:17)(cid:111)(cid:17) , with probability 1 − O (exp ( − d s )) ,
noorthonormal Q ∈ R k 1 ×
d s satisfies that for some constant C 1 ∈ R k 1 × k 1 .
[Equation]
Section B.1 gives a proof sketch.
The formal statement Lemma B.3 and its proof can be found in
Appendix B.2.
The next claim says that Lemma 5.2 can be applied iteratively, i.e. fixing a featurizer from previousiterations that uses k t − 1 spurious dimensions, with high probability, any U t that matches featuresfrom Ω ( k t − 1 /k t ) new environments uses at most k t spurious dimensions.
Corollary 5.3 (Informal version of Corollary B.7) .
Suppose 2 ≤ k t ≤ k t − 1 / 2 ≤ d s / 2 .
When |E t | =
E t = Ω (cid:16)
[Equation]
t [Equation] [Equation] 1 max (cid:110) 1 , log (cid:16)
D [Equation]
k t − 1) d s (cid:17) , log (cid:16)
d s k t − 1 (cid:17)(cid:111)(cid:17) , for fixed orthonormal
P ∈ R k t − 1 ×
d s , with probability 1 − O (exp ( − d s )) , no orthonormal Q ∈ R k t × k t − 1 satisfies
∀ e ∈
[ E t ] ,
QP Σ e 2 P ⊤ Q ⊤ =
C t for some constant C t ∈ R r t × r t .
The formal statement Corollary B.7 and its proof can be found in Appendix B.2.
Lemma 5.1 followsfrom iterative application of Corollary 5.3, as shown in Appendix B.2.
[-Designing a testing statistic To show that for fixed Q , Q ∆ e Q ⊤ = 
-]
[-( q ⊤ i ∆ e 2 q j ) 
-]
Experiments
In light of the differences between IFM and CORAL discussed in section 4.1, we test several questionsinspired by our theory: (Q1)
Do feature matching algorithms (IFM and CORAL) have much smallerenvironment complexity compared to ERM and IRM, with finite samples drawn from data modelssimilar to our assumptions?
(Q2) Can decoupling feature matching and supervised training of theclassifier (IFM) improve over joint training (CORAL)?
(Q3)
For neural network featurizers, canmatching 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 matchingall environments at all layers?
(Q5) Is it important to shrink feature dimensions?
We use two tasks toinvestigate those questions empirically.
Appendix C contains additional details.
Gaussian dataset is a binary classification task that closely reflects our assumptions in section 3.
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 trainingenvironments.
In test environments, the noise is uncorrelated with the label.
Algorithms and architectures.
For Gaussian dataset we use linear predictors.
IRM follows theimplementation in Arjovsky et al.
[3]; CORAL jointly minimizes average supervised loss on trainingenvironments and L coral , which is the average of squared distances in conditonal feature means (in l 2norm) and covariances (in Frobenius norm) between adjacent training environments; CORAL+ONadds orthonormal penalty loss L on ( U ) =
∥ UU ⊤ − I ∥ 2 F where U is the featurizer; IFM is our
Algorithm 1, where for each layer U t , the training objective is L t ( U t ) =
λ 1 L coral + λ 2 L on .
Wetest IFM with 1 vs. 3-layer featurizers, either matching all ( match-all ) or a disjoint set of trainingenvironments ( match-disjoint ) at each layer.
For Noised MNIST we use ReLU networks with 1-6 layers.
Although our theory for IFM is onlyfor the linear setting, in practice we can extend the idea to deep network by training the classifierhead together with feature matching at different layers.
CORAL (match-all) and CORAL (matchdisjoint) in Figure 2 are natural extensions of IFM to nonlinear models.
They match at all layerspost-activation, using either all ( match-all ) or a disjoint subset of training environments ( matchdisjoint ) per layer.
Figure 2 shows that they outperform naive CORAL , which matches features onlyat the last layer.
Thus, the core ideas of IFM may be used to improve performance of existing featurematching algorithms with nonlinear models and realistic datasets (e.g. DANN, MMD).
Results
(Q1) Figures 1 and 2 show that IFM and CORAL have much smaller environmentcomplexity 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 matchall, CORAL match-disjoint) improves over matching at only the last layer (CORAL).
(Q4)
In bothdatasets, matching disjoint sets of environments at each layer (IFM match-disjoint, CORAL matchdisjoint) is almost as good as matching all environments at all layers (IFM match-all, CORAL matchall) while saving computation.
(Q5)
For Noised MNIST (Table 2 in Appendix C), shrinking featuredimensions is crucial for the advantage of feature matching at multiple layers, e.g. matching featuresat 3 layers with widths [24 , 24 , 24] does not significantly improve over matching features at the lastlayer (CORAL).
Overall, our results suggest that practitioners may benefit from feature matchingalgorithms when the data is similar to our assumed model, and may get additional advantage viamatching at multiple layers with diminishing dimensions, echoing existing empirical works [27, 29].
Conclusion
This work presents the first domain generalization algorithm which provably recovers an invariantpredictor 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”is possible, and based on our theory we believe the use of an iterative approach is a key insight whichcould lead to additional positive results for out-of-distribution generalization.
Notably, this work alsorepresents the first theoretical justification for the empirical success of existing algorithms which usefeature distribution matching.
However, there remains much room for improvement: our results are for a linear data model with fairly special assumptions.
It would be interesting to analyze data modelsin which the observables are a non-linear function of the latent variables.
Checklist
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 1.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] SeeSection E.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
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 4,Appendix B.
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 attachedin supplemental file.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Section 6, Appendix C.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See figures 1, 2.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix C.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] We cited LeCun andCortes [23].
(b) Did you mention the license of the assets?
[Yes] See Appendix C.
(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]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]