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 [10, 21] 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 [7].
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) [32, 31], Maximum Mean Discrepancy (MMD)
and (Conditional) Domain Adversarial Networks [10, 21].
For years, these algorithms dominateddomain generalization benchmarks like PACS [18], VLCS [9] and DomainNet [26] but they lackformal guarantees.
Previous empirical works usually characterize the performance of these algorithmsusing generalization bounds based on divergence between domain distributions [6, 24], but thosebounds are vacuous for the typical benchmarks for domain generalization and thus cannot explaintheir success.
For example, when different environments have disjoint supports (as is the case inthe above-mentioned benchmarks) the H -divergence [6] is 1 even for the linear hypothesis class.
To obtain useful guarantees, it is necessary to study data models that encode structure reflective ofsettings of interest.
Prior works attempting to theoretically characterize the performance of featurematching algorithms emphasize lower bounds [36, 33].
In this work, we seek to give the first positivetheoretical justification for feature matching algorithms.
The other major alternative assumption in the literature is enforcing invariance of the label distributionconditioned on the signal features.
Invariant Risk Minimization (IRM) [2] assumes E [ y | Φ( x )]is invariant, and follow-up works assume invariance of higher moments [35, 14, 23, 5].
However, empirical results for these algorithms are mixed: Gulrajani and Lopez-Paz [13], Aubin et al.
[3]present experimental evidence that these methods do not consistently outperform ERM for eitherrealistic or simple linear data models, when fairly evaluated.
Recent theoretical works [27, 15] 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.
[27] 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.
[15], Ahuja et al.
[1] 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.[27], 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 [8].
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 [17], where the background noisespuriously correlates with the label.
Our results in Section 7 suggest that practitioners may benefitfrom 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 [20, 22].
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 [27, 15].
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 eU,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 the origin 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.
[27]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.
[27]) .Σ e 2 =
σ 2 e I 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) .
[Equation]
, where Σ e 2 ∈ R d s × d s is arbitrary (and can be adversarial), and[
[Equation] ] 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 [30], as we allow the parameters to be a random smoothing ofarbitrary parameters.
The real-world interpretation is that our analysis applies to all but a measure-zeroset of covariances.
While our data model is simple, it has already been used as a sandbox to understand algorithms like
IRM [27, 15] 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.
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 the
Algorithm 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 e ∈ E tr .
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]
[Equation]
feature 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.
To ensure that all invariant dimensions are preserved, we always find the projectionwith the maximum possible dimension r t that still matches the covariances, until we are left withonly r dimensions.
Among algorithms used in practice, IFM is most similar to CORAL [32].
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 7 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) 196training 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 C).
However, this algorithm is extremely brittle and reliant on 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)
[Equation] min( [Equation] 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 A.2.
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] [Table] [Equation]
[Table]
[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 A.3.
Proof sketch for
{+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 feature extractor
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 A.3) .
For any integer 2 ≤ k 1 ≤ d s / 2 , when |E 1 | =
E 1 = Ω (cid:16)
[Equation]
[Equation]
[Equation]
(cid:16)
[Equation]
(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 6 gives a proof sketch.
The formal statement Lemma A.3 and its proof can be found in
Appendix A.1.
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 A.7) .
Suppose 2 ≤ k t ≤ k t − 1 / 2 ≤ d s / 2 .
When |E t | =
E t = Ω (cid:16) k t − 1 − k t k t − 1 max (cid:110) 1 , log (cid:16)
D ( 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 A.7 and its proof can be found in Appendix A.1.
Lemma 5.1 followsfrom iterative application of Corollary 5.3, as shown in Appendix A.1.  Proof sketch of Lemma 5.2
This section gives a proof sketch of the main lemma 5.2.
Our goal is to show that when E = Ω ( d s /k ) ,with high probability, no orthonormal Q ∈ R k × d s satisfies Q ∆ e Q ⊤ = 0 for all e , where we definethe difference between the covariances of spurious features in two adjacent environments e, e + 1 as∆
[Equation] = [Equation] = (
[Equation]
) + ( G
[Equation]
G
[Equation]
G
[Equation]
G ⊤ e +1 ) .
Step 1: Discretization To show this, we discretize over the space of orthonormal matrices Q , andshow that for fixed Q , the probability that Q ∆ e 2 Q ⊤ = 0 for all e is small.
We then union bound overthe covering.
Step 2:
{+Designing a testing statistic To show that for fixed Q , Q ∆ e Q ⊤ = 
+}
k × k is unlikely, wefocus on showing that the sum of squares of all entries in Q ∆ e Q ⊤ is bounded away from 0.
Moreformally, let q i be the i -th row of Q .
Define Z ije =
{+( q ⊤ i ∆ e 2 q j ) 
+}
so it is (cid:2) Q ∆ e Q ⊤ (cid:3) 2 i,j .
Define the nonrandom part A ije =
q ⊤ i (Σ e 2 − Σ e + ) q j , and
A =
(cid:80)
A 2 ije , we write(cid:88)
[Equation]
We will eventually show that with high probability, (cid:80) Z ije =
Ω (cid:0) A + Ek 2 d s (cid:1) .
Step 3: Decoupling One difficulty is that (cid:80) Z ije is not a sum of independent random variables,because q ⊤ i G e G ⊤ e q j and q ⊤ i G e G ⊤ e q j ′ are dependent.
The key insight is to use a decoupling toolfrom de la Peña and Montgomery-Smith [8].
In more details, define V i,e = G e q i .
For fixed Q , V i,e ∼ N (0 , I d s ) and the ensemble { V i,e } i ∈ [ k ] ,e ∈
[ E ] 286 is independent.
Therefore
[Equation]
For further simplification, we define X i,e =
[ V i,e ; V i,e +1 ]
∈ R 2 d s , and I ∗ =
[ I d s , 0 ; 0 , − I d s ] , so
[Equation]
Note that X ⊤ i,e I ∗ X j,e and X ⊤ i,e I ∗ X j ′ ,e are dependent for j ′ ̸
= j, j ′ ̸
= i, i ̸
=
j .
To resolve this, Weapply Lemma
A.4 to decouple them.
Lemma A.4 says to show concentration of sum of Z ije we justneed to show concentration of the sum of some other random variables Z ′ ije , where
[Equation]
Here Y i,e and X i,e are identically distributed.
We then turn { Z ′ ije } into their identically distributed counterparts { Z ′′ ije } , where
[Equation]
Step 4: High probability norm bounds for { Y i,e } We first consider the randomness in { Y i,e } , andprove that with high probability { Y i,e } satisfies two norm bounds; we then show the concentration of 295 (cid:80)
Z ′′ ije conditioned on the event that { Y i,e } satisfies these bounds.
We pack { Y i,e } as rows of a matrix
Y e ∈ R k × 2 d .
We design event E 1 so that conditioned on E 1 , for all Q ∈ Q and e ,
∥ Y e ∥ 2 =
O ( √ d s ) .
Event E 2 denotes the event that for all cover elements in ˜ Q , all e , ∥ Y e
∥ 2 F = Θ( Ekd s ) .
Step 5: Conditioned on Y e , proving that (cid:80) Z ′′ ije concentrates Once we fix Y e , the randomvariables { P ei = (cid:80) j Z ′′ ije } are independent, so the concentration of their sum is immediate.
[Equation]
We show concentration of the second and the third terms of equation (6.1) using Hoeffding’s inequalityand Hanson-Wright inequality.
Conditioned on good events E 1 and E 2 , for fixed Q ,
Pr
[Equation]
In other words, with high probability all Q ∈ ˜ Q satisfies (cid:80) e,i P ei =
O (cid:0) A + Ek 2 d s (cid:1) .
We choose ϵ to make the discretization error smaller than O ( Ek 2 d s ) .
Therefore, with high probability,for all Q ∈ Q , (cid:80) eij
Z ije =
Ω (cid:0) A + Ek 2 d s (cid:1) , i.e. the testing statistic is bounded away from zero.
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 B 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 [17].
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.
[2]; CORAL jointly minimizes average supervised loss on training environments 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 B), 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 [20, 22].
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 arefor 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 8 and Appendix D.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Section D.
(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 A.
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 7, Appendix B.
(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., type of GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix B.
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 and Cortes [17].
(b) Did you mention the license of the assets?
[Yes] See Appendix B.
(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
[Table]
[Table]
and the total
[Table]
participant compensation?
[N [Table]