D OMAIN -A DJUSTED R EGRESSION OR : ERM M AY A LREADY L EARN F EATURES S UFFICIENT FOR O UT - OF -D ISTRIBUTION G ENERALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
A common explanation for the failure of deep networks to generalize out-ofdistribution is that they fail to recover the “correct” features.
We challenge this notion with a simple experiment which suggests that ERM already learns sufﬁcient features and that the current bottleneck is not feature learning, but robust regression .
Our ﬁndings also imply that given a small amount of data from the target distribution, retraining only the last linear layer will give excellent performance.
We therefore argue that devising simpler methods for learning predictors on existing features is a promising direction for future research.
Towards this end, we introduce Domain-Adjusted Regression (DARE), a convex objective for learning a linear predictor that is provably robust under a new model of distribution shift.
Rather than learning one function, DARE performs a domain-speciﬁc adjustment to unify the domains in a canonical latent space and learns to predict in this space.
Under a natural model, we prove that the DARE solution is the minimax-optimal predictor for a constrained set of test distributions.
Further, we provide the ﬁrst ﬁnite-environment convergence guarantee to the minimax risk, improving over existing analyses which only yield minimax predictors after an environment threshold.
Evaluated on ﬁnetuned features, we ﬁnd that DARE compares favorably to prior methods, consistently achieving equal or better performance.
I NTRODUCTION
The historical motivation for deep learning focuses on the ability of deep neural networks to automatically learn rich, hierarchical features of complex data (LeCun et al., 2015; Goodfellow et al., 2016).
Simple Empirical Risk Minimization (ERM), with appropriate regularization, results in high-quality representations which surpass carefully hand-selected features on a wide variety of downstream tasks.
Despite these successes, or perhaps because of them, the dominant focus of late is on the shortcomings of this approach: recent work points to the failure of networks trained with ERM to generalize under even moderate distribution shift (Recht et al., 2019; Miller et al., 2020).
A common explanation for this phenomenon is reliance on “spurious correlations” or “shortcuts”, where a network makes predictions based on structure in the data which generalizes on average in the training set but may not persist in future test distributions (Poliak et al., 2018;
Geirhos et al., 2019; Xiao et al., 2021).
Many proposed solutions implicitly assume that this problem is due to the entire neural network: they suggest an alternate objective to be minimized over a deep network in an end-to-end fashion (Sun & Saenko, 2016; Ganin et al., 2016; Arjovsky et al., 2019).
These objectives are complex, poorly understood, and difﬁcult to optimize.
Indeed, the efﬁcacy of many such objectives was recently called into serious question (Zhao et al., 2019; Rosenfeld et al., 2021; Gulrajani & Lopez-Paz, 2021).
Though a neural network is often viewed as a deep feature embedder with a ﬁnal linear predictor applied to the features, it is still unclear—and to our knowledge has not been directly asked or tested —whether these issues are primarily because of (i) learning the wrong features or (ii) learning good features but failing to ﬁnd the best-generalizing linear predictor on top of them.
We begin with a simple experiment (Figure 1) to try to distinguish between these two possibilities: we train a deep network with ERM on several domain generalization benchmarks, where the task is to learn a predictor using a collection of distinct training domains and then perform well on a new,
A
C
P
R 0 .0 .
T e s t
A cc u r a c y
OﬃceHome
C
L
S
V
VLCS
A
C
P
S 0 .0 .0 .1 .
T e s t
A cc u r a c y
PACS c
i p
q r s 0 .0 .0 .
DomainNet
ERM
DARE
Linear † End-to-End †
Pretrained †
Cheating accuracy on this is due to ﬁnetuning on the overparameterization unseen domain. After training, we freeze the features and separately learn a linear classiﬁer on top of them. Crucially, when training this classiﬁer (i.e., retraining just the last linear layer), we give it an unreasonable advantage by optimizing on both the train and test domains—henceforth we refer to this as “cheating”.
Since we use just a linear classiﬁer, this process establishes a lower bound on what performance we could plausibly achieve using standard ERM features.
We then separately cheat while training the full network end-to-end, simulating the idealized setting with no distribution shift.
Note that in neither case do we train on the test points ; our cheating entails training on (different) samples from the test domain , which are assumed unavailable in domain generalization.
Notably, we ﬁnd that simple (cheating) logistic regression on frozen deep features learned via
ERM results in enormous improvements over current state of the art, on the order of 10-15%.
In fact, it usually performs comparably to the full cheating method—which learns both features and classiﬁer end-to-end with test domain access—sometimes even outperforming it.
Put another way, cheating while training the entire network rarely does signiﬁcantly better than cheating while training just the last linear layer.
One possible explanation for this is that the pretrained model is so overparametrized as to effectively be a kernel with universal approximation power; in this case, the outstanding performance of a cheating linear classiﬁer on top of these features would be unsurprising.
However, we ﬁnd that this cheating method does not ensure good performance on pretrained features, which implies that we are not yet in such a regime and that the effect we observe is indeed due to ﬁnetuning via ERM.
Collectively, these results suggest that training modern deep architectures with ERM and established in-distribution training and regularization practices may be “good enough” for out-of-distribution generalization and that the current bottleneck lies primarily in learning a simple, robust predictor.
Motivated by these ﬁndings, we propose a new objective, which we call Domain-Adjusted Regression (DARE).
The DARE objective is convex and it learns a linear predictor on frozen features.
Unlike invariant prediction (Peters et al., 2016), which projects out feature variation such that a single predictor performs acceptably on very different domains, DARE performs a domain-speciﬁc adjustment to unify the environmental features in a canonical latent space.
Based on the presumption that standard ERM features are good enough (made formal in Section 4), DARE enjoys strong theoretical guarantees: under a new model of distribution shift which captures ideas from invariant/non-invariant latent variable models, we precisely characterize the adversarial risk of the DARE solution against a natural perturbation set, and we prove that this risk is minimax .
We further provide the ﬁrst ﬁnite-environment convergence guarantee to the minimax risk, improving over existing results which merely demonstrate a threshold in the number of observed environments at which the solution is discovered (Rosenfeld et al., 2021; Chen et al., 2021; Wang et al., 2022).
Finally, we show how our objective can be modiﬁed to leverage access to unlabeled samples at test-time.
We use this to derive a method for provably effective “just-in-time” unsupervised domain adaptation, for which we provide a ﬁnite-sample excess risk bound.
Evaluated on ﬁnetuned features, we ﬁnd that DARE compares favorably to existing methods, consistently achieving equal or better performance.
We also ﬁnd that methods which previously underperformed on these benchmarks do much better in this frozen feature setting, often besting
ERM.
This suggests that these approaches are beneﬁcial for linear prediction (the setting in which they are understood and justiﬁed) but using them to train deep networks may result in worse features.
ERM L EARNS S URPRISINGLY U SEFUL F EATURES
Our experiments are motivated by a simple question: is the observed failure of deep neural networks to generalize out-of-distribution more appropriately attributed to inadequate feature learning or inadequate robust prediction ?
Both could undoubtedly be improved, but we are concerned with which currently serves as the primary bottleneck.
It’s typical to train the entire network end-to-end and then evaluate on a test distribution; in reality, this measures the quality of the interaction of the features and the classiﬁer, not of either one individually.
Using datasets and methodology from D OMAIN B ED (Gulrajani & Lopez-Paz, 2021), we ﬁnetune a
ResNet-50 (He et al., 2016) with ERM on the training domains to extract features.
Next, we cheat while learning a linear classiﬁer on top of these frozen features by optimizing on both the train and test domains.
We compare this cheating linear classiﬁer to a full cheating network trained on all domains end-to-end.
If it is the case that ERM learns features which do not generalize, we should expect that the cheating linear classiﬁer will not substantially improve over the current state of the art and will perform signiﬁcantly worse than cheating end-to-end, since the latter method can adapt the features to better suit the test domain.
Instead, we ﬁnd that simply giving the linear predictor access to all domains while training makes up the vast majority of the gap between current state of the art (ERM with heavy regularization) and the ideal setting where we train a network on all domains.
In other words, ERM produces features which are informative enough that a linear classiﬁer on top of these frozen features is—in principle—capable of generalizing almost as well as if we had access to the test domain when training the entire network.
1 Figure 2 in the Appendix depicts the evaluation methodology described above, along with a more detailed explanation.
We conjecture that this phenomenon occurs more broadly: our ﬁndings suggest that for many applications, existing features learned via ERM may be “good enough” for out-of-distribution generalization in deep learning.
Based on this idea, we posit that future work would beneﬁt from modularity, working to improve representation learning and robust classiﬁcation/regression separately.
There are several distinct advantages to this approach:
• More robust comparisons and better reproducibility: Current methods have myriad degrees of freedom which makes informative comparisons difﬁcult; evaluating feature learning and robust regression separately eliminates many of these sources of experimental variation.
• Less compute overhead: Training large networks to learn both features and a classiﬁer is expensive.
Benchmarks could include ﬁles with the weights for ready-to-use deep feature embedders trained with various objectives—these models can be much larger and trained on much more data than would be feasible for many.
Using these features, academic researchers could thus make faster progress on better classiﬁers with less compute.
• Modular theoretical guarantees: Conditioning on the frozen features, we can use more classical analyses to provide guarantees for the simpler parametric classiﬁers learned on top of these features.
For example, Bansal et al.
2021) derive a generalization bound for simpler classiﬁers which is agnostic to the complexity of the features on which they are trained.
We conclude by emphasizing that while the predictor in our experiments is linear, future methods need not be.
Rather, we are highlighting that there may be no need for complex, expensive, highly variable regularization of a deep network when much simpler approaches sufﬁce.
T HE D OMAIN -A DJUSTED R EGRESSION O BJECTIVE
The goal of many prior methods in deep domain adaptation or generalization is to learn a single network which does well on all environments simultaneously, often by throwing away non-invariant components (Peng et al., 2019; Arjovsky et al., 2019).
While invariance is a powerful framework, there are some clear drawbacks , such as the need to throw away possibly informative features.
In settings where we expect the distribution shift to be less than worst-case, this would be unnecessarily conservative.
Instead, we reframe the problem by thinking of each training domain as a distinct transformation from a shared canonical representation space.
In this framing, we can “adjust” each domain in order to undo these transformations, aligning the representations to learn a single robust predictor in this uniﬁed space.
Speciﬁcally, we propose to whiten each domain’s features to have zero mean and identity covariance, which can be thought of as a “domain-speciﬁc batchnorm” but using the full feature covariance.
This idea is not totally new: some prior works align the moments between domains to improve generalization.
The difference is that DARE does not learn a single featurizer which aligns the moments, but rather aligns the moments of already learned features ; the latter approach maintains useful variation between domains which would be eliminated by the former.
DARE is closer to methods which learn separate batchnorm parameters per domain over a deep network, possibly adjusting at test-time (Seo et al., 2019; Chang et al., 2019; Segù et al., 2020)—these methods perform well, but they are entirely heuristic based, difﬁcult to optimize, and come with no formal guarantees.
Our theoretical analysis thus serves as preliminary justiﬁcation for the observed beneﬁts of such methods, which have so far lacked serious grounding.
To begin, deﬁne E as the set of observed training environments, each of which is deﬁned by a distribution p e ( x, y ) over features x ∈ R d and class labels y ∈
[ k ] .
For each e ∈ E , denote the mean and covariance of the features as µ e , Σ e .
Our ﬁrst step is to adjust the features via the whitening transformation
Σ − 1 / 2 e ( x − µ e ) .
Unfortunately, we cannot undo the test transformation because we have no way of knowing what the test mean will be.
Interestingly, this is not a problem provided the predictor satisﬁes a simple constraint.
Suppose that the predictor’s output on the mean representation of each environment is a multiple of the all-ones vector.
As softmax is invariant to a constant offset, this enforces that the environment mean has no effect on the ﬁnal probability distribution, and thus there is no need to adjust for the test-time mean.
We therefore enforce this constraint during training, with the hope that the same invariance will approximately hold at test-time.
Formally, the DARE objective ﬁnds a matrix β ∈ R d × k which solves
[Equation]
subject to softmax [Equation] where (cid:96) is the multinomial logistic loss and we omit the bias for brevity.
For binary classiﬁcation, β is a vector and the softmax is replaced with the logistic function—the constraint is then equivalent to requiring the mean of the logits to be 0.
Thus, the DARE objective explicitly regresses on the adjusted features, while the constraint enforces that each environment mean has no effect on the output distribution to encourage predictions to also be invariant to test-time transformations.
The astute reader will point out that we also do not know the correct whitening matrix for the test data—instead, we adjust using our best guess for the test covariance: denoting this estimate as ¯Σ , our prediction on a new sample x is f ( x ; β ) = softmax (cid:0) β T ¯Σ − 1 / 2 x (cid:1) .
We prove that this prediction is minimax so long as our guess is “sufﬁciently close”, and in practice we ﬁnd that simply averaging the training domain adjustments performs well.
Table 2 in the Appendix shows this average is actually quite close to the sample covariance of the (unseen) test domain, explaining the good performance.
Unlike many prior methods, DARE does not enforce invariance of the features themselves.
Rather, it aligns the representations such that different domains share similar optimal predictors.
To support this claim, Table 3 displays the cosine similarity between optimal linear classiﬁers for individual domains—we observe a large increase in average similarity as a result of the feature adjustment.
Further, in Section 5 we demonstrate that the DARE objective is in fact minimizing the worst-case risk under a constrained set of possible distribution shifts, making it less conservative than methods which require complete feature invariance.
Though the objective (1) assumes a particular form of invariance, this speciﬁc choice is by no means a requirement for the approach we take.
We justify this by noting that the minimax predictor can always be written as the solution to an objective of this form, depending only on what we choose to assume holds constant:
Proposition 3.1 (informal) .
Assume the minimax-optimal predictor f ∗ lies in our hypothesis class F .
Let S : F × P ( X × Y )
→ R d be any “sufﬁcient statistic” which encodes all in- variances, i.e. S ( f ∗ , p e ) =
c ∀ e , for some constant c , and ∀ f ∈ F with lower risk than f ∗ on at least one environment,
[Equation]
Then as |E| → ∞ , f
[Equation]
=
[Equation]
This generic objective recovers IRM by deﬁning S ( f, p e ) =
E p e [ ∇ (cid:96) ( f ( x ) , y ))] , c [-= 
 (the-]=(the “sufﬁcient statistic” is the environment loss gradient), or CORAL with S ( f, p e )
= (cid:2)
E p e [ x ] , E p e [( x − µ )( x − µ ) T ] (cid:3) (the feature mean and covariance) and letting c be arbitrary.
DARE instead allows for domain-speciﬁc transformations T e , and we deﬁne S ( f, p e ) = softmax (cid:0) E p e [ f ( T − 1 e ( x ))] (cid:1) —here T − 1 e is the whitening operation detailed in the next section.
Implementation in practice.
Due to its convexity, the DARE objective is extremely simple to optimize.
In practice, we ﬁnetune a deep network over the training data with ERM and then extract the features.
Next, treating the frozen features of the training data as direct observations x , we minimize the empirical Lagrangian form:
[Equation]
where ˆ µ e , ˆΣ e are the usual sample estimates and n e is the number of samples in environment e .
We ﬁnd that the solution is incredibly robust to the choice of λ , but it is natural to wonder whether each of the components above is necessary for the performance gains we observe.
We ablate both the whitening operation and the use of the constraint (Appendix E) and see performance drops in both cases.
We also consider estimating the test-time mean rather than enforcing invariance, but this results in substantially worse accuracy—the environment feature means vary quite a bit, so the estimate is usually inaccurate.
For regression, we consider the same setup but minimize mean squared error on targets y ∈ R .
Here, the DARE solution is constrained such that the mean output has no effect on the prediction, meaning each domain’s mean prediction should be zero.
We discuss this in more detail in
Appendix A.1, along with an interesting connection to anchor regression (Rothenhäusler et al., 2021).
Another beneﬁt to our approach is that the adjustments for each domain do not depend on labels, so given unlabeled samples from the test domain we can often do even better.
Unlike methods which use those samples while training (or even for test-time training), this adjustment can be done just-in-time , without updating any parameters!
We name this task Just-in-Time Unsupervised Domain Adaptation (JIT-UDA), and we provide ﬁnite-sample risk bounds for the DARE objective in this setting.
A N EW M ODEL OF D ISTRIBUTION S HIFT
The DARE objective is based on the intuition that all domains jointly share a representation space and that they arise as unique transformations from this space.
To capture this notion mathematically, we model the joint distribution p e ( (cid:15), y ) over latents (cid:15) ∈ R d and label y ∈ { 0 , 1 } , along with an environment-speciﬁc transformation to observations x ∈ R d for each domain.
In the fully general case, this transformation can take an arbitrary form and can be written as x = T e ( (cid:15), y ) .
Our primary assumption is that p e ( y | (cid:15) ) is constant for all domains.
Our goal is thus to invert each transformation T e such that we are learning an invariant conditional p ( y | T − 1 e ( x ))
(throughout we assume T e is invertible).
One can also view this model as a generalization of covariate shift: where the usual assumption is constancy of p ( y | x ) , the inverse transformation gives a richer model which can more realistically capture real-world variation across domains.
It is important to note that this model generalizes (and has strictly weaker requirements than) both IRM and domain-invariant representation learning, which can be recovered by assuming T e is the same for all environments.
For a typical deep learning analysis we would model T e as a non-linear generative map from latents to high-dimensional observations.
However, our ﬁnding that ERM features are good enough suggests that modeling the learned features as a simple function of the “true latents” (cid:15) is not unreasonable.
In other words, we now consider x to represent the frozen features of the trained network, and we expect this network to have already “undone” the majority of the complexity, resulting in observations which are a simple function of the ground truth.
Accordingly, we consider the following model:
[Equation]
Here, (cid:15) 0 ∼ p e ( (cid:15) 0 ) , which we allow to be any domain-speciﬁc distribution ; we assume only that its mean is zero (such that E [ (cid:15) ]
= b e ) and that its covariance exists.
We ﬁx β ∗ ∈ R d for all domains and model η as logistic noise.
Finally, A e ∈ R d × d , b e ∈ R
d are domain-speciﬁc.
For regression, we model the same generative process for x, (cid:15) , while for the response, we have y ∈ R and η is zero-mean independent noise: y = β ∗ T (cid:15) + η .
We remark that the reason for separate deﬁnitions of p e ( (cid:15) 0 ) and b e is that our robustness guarantees are agnostic to the distribution of latent residuals p e ( (cid:15) 0 ) —the only aspect of the latent distribution p ( (cid:15) ) which affects our bounds is the environment mean b e .
Connection to Invariant Prediction.
Statistical models of varying and invariant (latent) features have recently become a popular tool for analyzing the behavior of deep representation learning algorithms (Rosenfeld et al., 2021;
Chen et al., 2021; Wald et al., 2021).
We see that Equation (2) can model such a setting by assuming, e.g., that a subspace of the columnspan of A e is constant for all e , while the remaining subspace can vary.
In such a case, the features x are expressed as the sum of a varying and an invariant component, and any minimax representation must remove the varying component, throwing away potentially useful information.
Instead, DARE realigns these components so that we can use them at test-time.
We illustrate this with the following running example:
Example 1 (Invariant Latent Subspace Recovery) .
Consider the model (2) with (cid:15) ∼ N (0 , I d 1 + d 2 ) .
Deﬁne Π := blockdiag ( I d 1 , 0 d 2 ) and assume A e =
blockdiag (Σ 1 / 2 , Σ 1 / 2 e ) for all e , where Σ ∈ R d 1 × d 1 is constant for all domains but Σ e ∈ R d 2 × d 2 varies.
Here we have a simple latent variable model: the features have the decomposition x
= A e (cid:15) =
Σ 1 / 2 Π (cid:15) + Σ 1 / 2 e ( I − Π) (cid:15) , where the ﬁrst component is constant across environments and the second varies.
It will be instructive at this point to analyze what an invariant prediction algorithm such as IRM would do in this setting.
Here, the IRM constraint would enforce learning a featurizer Φ such that Φ( x ) has an invariant relationship with the target.
Under the above model, the solution is Φ( x ) =
Π
A − 1 e x = Σ 1 / 2 Π (cid:15) , retainining only the component lying in span (Π) .
Crucially, removing these features actually results in worse performance on the training environments—IRM only does so because using these features could harm test performance in the worst case.
However, projecting out the varying component in this manner is unnecessarily conservative, as we may expect that for a future distribution, Σ e (cid:48) will not be too different from what we have seen before.
Instead of projecting out this component, DARE performs a more nuanced alignment of environment subspaces, and it is thus able to take advantage of this additional informaton.
We will see shortly the resulting beneﬁts.
T HEORETICAL A NALYSIS
Before we can analyze the DARE objective, we observe that there is a possible degeneracy in
Equation (2), since two identical observational distributions p ( x, y ) can have different regression vectors β ∗ .
We therefore begin with a simple non-degeneracy assumption:
Assumption 5.1.
Write the SVD of A e as U e S e V Te .
We assume V e = V ∀ e .
One special case where this holds is when A e is constant for all domains; this is very similar to the “additive intervention” setting of Rothenhäusler et al.
(2021), since only p e , b e can vary.
We assume WLOG that V
=
I , as any other value can be subsumed by the other parameters.
We further let E [ (cid:15) 0 (cid:15) T 0 ]
= I WLOG by the same reasoning (see Appendix A.2 for a discussion on these conditions).
With Assumption 5.1, we can uniquely recover A e =
U e S e via the eigendecomposition of the covariance Σ e =
A e A Te =
U e S 2 e U Te .
We therefore use the notation Σ 1 / 2 e to refer to A e recovered in this way.
We allow covariances to have zero eigenvalues, in which case we write the matrix inverse to implicitly refer to the pseudoinverse.
As is standard in domain generalization analysis, unless stated otherwise we assume full distribution access to the training domains, though standard concentration inequalities could easily be applied.
A remark on our assumptions.
Assumption 5.1 is not trivial.
Domain generalization is an exceptionally difﬁcult problem, and showing anything meaningful requires some assumption of consistency between train and test.
Our assumptions are only as strong as necessary to prove our results, but future work could relax them, resulting in weaker performance guarantees.
Experiments in Appendix E demonstrate that our covariance estimation is indeed accurate, and the fact that our method exceeds state of the art even with this strong constraint (and does worse without the constraint, see Figure 4) is further evidence that our assumptions are reasonable.
We begin by deriving the solution to Equations (1) and (3) .
Recall that the DARE constraint requires that the mean representation of each domain has no effect on our prediction.
To enforce this, the
DARE solution must project out the subspace in which the means vary.
Given a set of E training environments, deﬁne B as the d × E matrix whose columns are the environmental mean parameters b e .
Throughout this section, we make use of the matrix ˆΠ , which is deﬁned as the orthogonal projection onto the nullspace of B T : ˆΠ :=
I − BB † =
U ˆΠ S
ˆΠ U T
ˆΠ ∈ R d × d .
This matrix projects onto the DARE constraint set, and it turns out to be all that is necessary to state the solution:
Theorem 5.2.
(Closed-form solution to the DARE population objective).
Under model (2), the solution to the DARE population objective (3) for linear regression is ˆΠ β ∗ .
If (cid:15) is Gaussian, then the solution for logistic regression (1) is α ˆΠ β ∗ for some α ∈ (0 , 1] .
In Example 1, we saw how invariant prediction will discard a large subspace of the representation and why this is undesirable.
Instead, DARE undoes the environment transformation and regresses directly on (cid:15) =
T − 1 e ( x ) .
Because we are performing a separate transformation for each environment, we are aligning the varying subspaces rather than throwing them away, allowing us to make use of additional information.
Though the DARE solution also recovers β ∗ up to a projection, it is using the adjusted features; DARE therefore only removes what cannot be aligned.
In particular, whereas Π has rank d 1 in Example 1, ˆΠ =
I would have full rank—this retains strictly more information.
Indeed, in the ideal setting where we have a good estimate of Σ e (cid:48) (e.g., under mild distribution shift or when solving JIT-UDA), we can make the Bayes-optimal prediction as E [ y | x ] =
β ∗ T A − 1 e (cid:48) x .
Thus we see a clear advantage that DARE enjoys over invariant prediction.
5. T HE A DVERSARIAL R ISK OF DARE
Moving forward, we denote the DARE solution β ∗ ˆΠ := ˆΠ β ∗ , with β ∗ I − ˆΠ deﬁned analogously.
We next study the behavior of the DARE solution under worst-case distribution shift.
We consider the setting where an adversary directly observes our choices of ¯Σ , ˆ β and chooses new environmental parameters
A e (cid:48) , b e (cid:48)
so as to cause the greatest possible loss.
Speciﬁcally, we study the square loss, deﬁning the excess test risk of a predictor as R e (cid:48)
(ˆ β ) :
=
E p e (cid:48)
[(ˆ β T ¯Σ − 1 / 2 x − β ∗ T (cid:15) ) 2 ] (we leave the dependence on ¯Σ implicit).
For logistic regression we therefore analyze the squared error with respect to the log-odds.
With some abuse of notation, we also reference excess risk when only using a particular subspace, i.e. R Π e (cid:48) (ˆ β ) : = E p e (cid:48)
[(ˆ
[Equation]
T Π ¯Σ −
[Equation]
/ 2 x
[Equation]
β ∗ T Π (cid:15) ) 2 ] .
Because we guess ¯Σ before observing any data from this new distribution, ensuring success is impossible in the general case.
Instead, we consider a set of restrictions on the adversary which will make the problem tractable.
Deﬁne the error in our test domain adjustment as ∆ := Σ 1 / 2 e (cid:48)
¯Σ − 1 / 2 − I ; observe that if ¯Σ =
Σ e (cid:48) , then ∆ = 0 .
Our ﬁrst assumption 2 says that the effect of our adjustment error with respect to the interaction between subspaces ˆΠ and ( I − ˆΠ) is bounded:
Assumption 5.3.
For a ﬁxed constant B ≥ 0 , (cid:107) ( [Equation] ≤ B (cid:107) ˆΠ [Equation]
Remark 5.4 .
To see speciﬁc cases when this would hold, consider the decomposition of ∆ according to its components in the subspaces ˆΠ
and I − ˆΠ : U T ˆΠ [Equation] U ˆΠ =
(cid:20)
[Equation]
∆
[Equation]
(cid:21) .
A few settings automatically satisfy Assumption 5.3 with B = 0 , since U T ˆ Π ∆ U ˆΠ will be block-diagonal.
In particular, this will be the case if all domains share an invariant subspace—e.g., if Π A e Π is constant as in Example 1.
Below, we show that exact recovery of this subspace occurs once we observe rank( I − Π) environments—this matches (actually, it is one less than) the linear environment complexity of most invariant predictors (Rosenfeld et al., 2021), and therefore Assumption 5.3 with B = 0 is no stronger than assuming a linear number of environments.
Assumption 5.5.
Using only covariates in the non-varying subspace, the risk of the ground truth regressor β ∗ is less than that of the trivial zero predictor: R ˆΠ p e [Equation]
< [Equation]
The need for this restriction should be immediate—if our adjustment error were so large that this did not hold, even the oracle regression vector would do worse than simply always predicting ˆ y = 0 .
Assumption 5.5 is satisﬁed for example if (cid:107) ∆ˆΠ (cid:107) < 1 , which again is guaranteed if there is an invariant subspace.
Note that we make no restriction on the risk in the subspace I − ˆΠ —the adversary is allowed any amount of variation in directions where we have already seen variation in the mean terms b e , but introduction of new variation is assumed bounded.
This is a no-free-lunch necessity: if we have never seen a particular type of variation, we cannot possibly know how to use it at test-time.
With these restrictions on the adversary, our main result derives the supremum of the excess test risk of the DARE solution under adversarial distribution shift.
Furthermore, we prove that this risk is minimax : making no more restrictions on the adversary other than a global bound on the mean, the
DARE solution achieves the best performance we could possibly hope for at test-time:
Theorem 5.6 (DARE risk and minimaxity) .
Denote the set of possible test environments A ρ which contains all parameters ( A e (cid:48) , b e (cid:48) ) subject to Assumptions 5.3 and 5.5 and a bound on the mean: (cid:107) b e (cid:48) (cid:107) ≤ ρ .
For logistic or linear regression, let ˆ β
be the minimizer of the corresponding DARE objective as in Theorem 5.2.
Then, sup ( A e (cid:48) ,b e (cid:48) ) ∈A
[Equation]
R e (cid:48)
(ˆ β ) =
{+(1 + ρ 2 )( (cid:107) β ∗ (cid:107) 
 + 2 B (cid:107) β ∗ ˆΠ (cid:107)(cid:107) β ∗ I − ˆΠ (cid:107) ) .+}
Furthermore, the DARE solution is minimax :
[Equation]
A special case when our assumptions hold is when all domains share an invariant subspace and we only predict using that subspace, but this is often too conservative.
There are settings where allowing for (limited) new variation can improve our predictions, and Theorem 5.6 shows that DARE should outperform invariant prediction in such settings.
5. T HE E NVIRONMENT C OMPLEXITY OF DARE
An important new measure of domain generalization algorithms is their environment complexity , which describes how the test risk behaves as a function of the number of (possibly random) domains we observe.
In contrast to Example 1, for this analysis we assume an invariant subspace Π outside of which both A e and b e can vary arbitrarily—we formalize this as a prior over the b e whose covariance has the same span as I − Π .
Our next result demonstrates that DARE achieves the same threshold as prior methods, but we also prove the ﬁrst ﬁnite-environment convergence guarantee , quantifying how quickly the risk of the DARE predictor approaches that of the minimax-optimal predictor.
Theorem 5.7 (Environment complexity of DARE) .
Deﬁne the smallest gap between consecutive eigenvalues: ξ (Σ) :
= min i ∈ [ d − 1] λ i − λ
i +1 .
Fix test parameters A e (cid:48) , b e (cid:48) and guess ¯Σ .
Suppose we minimize the DARE regression objective (3) on environments whose means b e are Gaussian vectors with covariance Σ b , with span (Σ b )
= span ( I − Π) .
After seeing E training domains:
If E ≥ rank(Σ b ) then DARE recovers the minimax-optimal predictor almost surely: ˆ β = β ∗ Π . 2.
Otherwise, if E ≥ r (Σ b ) then with probability ≥ 1 − δ ,
[Equation]
where O ( · ) hides dependence on (cid:107) ∆ (cid:107) and r (Σ) :
= Tr(Σ) / (cid:107) Σ (cid:107) is the effective rank.
For coherence we present the ﬁrst item as a probabilistic statement, but it holds deterministically so long as there are rank(Σ b ) linearly independent observations of b e .
Remark 5.8 .
Prior analyses of invariant prediction methods only show a discontinuous threshold where the minimax predictor is discovered after seeing a ﬁxed number of environments—usually linear in the non-invariant latent dimension But one should expect that if the variation is not too large then we can do better, and indeed Theorem 5.7 shows that if the effective rank of Σ b is sufﬁciently small, the risk of the DARE predictor will approach that of the minimax predictor as O ( E − 1 / 2 ) .
A PPLYING DARE TO JIT-UDA
So far, we have only considered a setting with no knowledge of the test domain.
As discussed in
Section 3, we’d expect that estimating the adjustment via unlabeled samples will improve performance.
Prior works have extensively explored how to leverage access to unlabeled test samples for improved generalization—but while some suggest ways of using unlabeled samples at test-time, they are not truly “just-in time”, nor have the advantages been formally quantiﬁed.
Our ﬁnal theorem investigates the provable beneﬁts of using the empirical moments instead of enforcing invariance, giving a ﬁnite-sample convergence guarantee for the unconstrained DARE objective in the JIT-UDA setting:
Theorem 5.9 (JIT-UDA, shortened) .
Let Σ S , Σ T be the covariances of the source and target distributions, respectively.
Deﬁne m (Σ) :
= λ max [Equation] Σ) λ 3min (Σ) .
Assume we observe n s = Ω( m (Σ S ) d 2 ) source samples and analogous n T target samples.
Then after solving the DARE objective, with probability at least 1 − 3 d − 1 , the excess squared risk of our predictor on the new environment is bounded as
[Equation]
Experimentally, we found that DARE does not outperform methods speciﬁcally intended for UDA, possibly because n (cid:28) d 2 —but we believe this is a promising direction for future research, since it doesn’t require unlabeled samples at train-time and it can incorporate new data on-the-ﬂy.
E XPERIMENTS
[Table]
Most algorithms implemented in the D OMAIN B ED benchmark are only applicable to deep networks; many apply complex regularization to either the learned features or the network itself.
We instead compare to three popular algorithms which work for linear classiﬁers: ERM (Vapnik, 1999), IRM (Arjovsky et al., 2019) and GroupDRO (Sagawa et al., 2020).
We evaluate all approaches on four datasets: Ofﬁce-Home (Venkateswara et al., 2017), PACS (Li et al., 2017), VLCS (Fang et al., 2013), and DomainNet (Peng et al., 2019).
We ﬁnd that DARE consistently matches or surpasses prior methods. A detailed description of the evaluation and comparison methodology can be found in
Appendix D and evaluations of other end-to-end methods in Appendix E.
Prior methods now consistently outperform ERM.
Interestingly, the previously observed gap between ERM and alternatives disappears in this setting with a linear predictor.
Gulrajani & LopezPaz (2021) report IRM and GroupDRO performing much worse on DomainNet (5-10% lower accuracy on some domains) but they surpass ERM when using frozen features.
This suggests they are more difﬁcult to optimize over a deep network and when they do work, it is most likely due to learning a better linear classiﬁer.
This further motivates work on methods for learning simpler robust predictors.
R ELATED W ORK
A popular approach to domain generalization matches the domains in feature space, either by aligning moments (Sun & Saenko, 2016) or with an adversarial loss (Ganin et al., 2016), though these methods are known to be inadequate in general (Zhao et al., 2019).
DARE differs from these approaches in that the constraint requires only that the feature mean projection onto the vector β be invariant.
Domain-invariant projections (Baktashmotlagh et al., 2013) were recently analyzed by Chen &
Bühlmann (2021), though notably under fully observed features and only for domain adaptation.
There has been intense recent focus on invariant prediction, based on ideas from causality (Peters et al., 2016) and catalyzed by IRM (Arjovsky et al., 2019).
Though the goal of such methods is minimax-optimality under major distribution shift, later work identiﬁes critical failure modes of this approach (Rosenfeld et al., 2021; Kamath et al., 2021).
These methods eliminate features whose information is not invariant, which is often overly conservative.
DARE instead allows for limited new variation by aligning the non-invariant subspaces, enabling stronger theoretical guarantees.
Some prior works “normalize” each domain by learning separate batchnorm parameters but sharing the rest of the network.
This was initially suggested for UDA (Li et al., 2016; Bousmalis et al., 2016;
Chang et al., 2019), which is not directly comparable to DARE since it requires unlabeled test data.
This idea has also been applied to domain generalization (Seo et al., 2019; Segù et al., 2020) but in an ad-hoc manner. Because of the difﬁculty in training the network end-to-end, there is no consistent method for optimizing or validating the objective—in particular, all deep domain generalization methods were recently called into question when Gulrajani & Lopez-Paz (2021) gave convincing evidence that nothing beats ERM when evaluated fairly.
Nevertheless, our analysis provides an initial justiﬁcation for these methods, suggesting that this idea is worth exploring further.
L IMITATIONS
Our experiments demonstrate that traditional training already learns excellent features for generalizing to new distributions.
This suggests that we should focus on methods learning to use these features, but it also means that these features are likely to contain just as much spurious or sensitive information as the original inputs.
Thus, even if a method which uses frozen features performs better under shift (such as on minority subpopulations), care must still be taken to account for possible biases in the resulting classiﬁer.
R EPRODUCIBILITY S TATEMENT
A major advantage to evaluating only a linear predictor is that we can release the exact model weights which were used to extract the frozen features on which all algorithms were trained.
Code for our experiments will be released along with these models.
A D ISCUSSION A. C ONNECTION TO A NCHOR R EGRESSION
The DARE objective for linear regression is written
[Equation]
The idea of adjusting for domain projections has similarities to Anchor Regression (Rothenhäusler et al., 2021), an objective which linearly regresses separately on the projection and rejection of the data onto the span of a set of anchor variables .
These variables represent some (known) measure of variability across the data, and the resulting solution enjoys robustness to pointwise-additive shifts in the underlying SCM.
If we deﬁne the anchor variable to be a one-hot vector indicating a sample’s environment, the Anchor Regression objective minimizes
[Equation]
where (cid:96) is the squared loss and µ y,e
= E p e [ y ] .
Here we see that Anchor Regression is “adjusting” in a sense, by regressing on the residuals, though the objective also regresses the mean prediction onto the target mean.
Unfortunately, this requires access to the target mean, which is unavailable in logistic regression due to the lack of a good estimator for E (cid:104) log p ( y =1 | x ) p ( y = − 1 | x ) (cid:105) .
Nevertheless, if we (i) assume the feature covariance is ﬁxed for all environments and (ii) assume the target mean is zero for all environments, we observe that the above objective becomes equivalent to the Lagrangian form of DARE for linear regression (3).
Alternatively, we could imagine combining the two by keeping Anchor Regression’s use of separate target means while adding DARE’s feature covariance whitening which would give
[Equation]
However, this still requires us to estimate the target mean, so it is unclear if or how this objective could be applied to the task of classiﬁcation.
A. O N THE C ONDITIONS A SSUMED W ITHOUT L OSS OF G ENERALITY
It may not be immediately clear why it is reasonable to assume both V
=
I and Σ (cid:15)
=
I WLOG simultaneously, so we clarify this point here.
It is crucial to observe here that β ∗ and A e do not need to be directly identiﬁable, because we only care about the predictive distribution β ∗ T (cid:15) .
We only need
A e to be identiﬁable from x up to equivalence of this distribution.
So if for example we recover some invertible transformation ˆ A e =
A e M , this is not at all a problem because we also learn the corresponding ˆ β
= M T β ∗ such that ˆ β T
ˆ
A − 1 e x = β ∗ T MM − 1 A − 1 e x = β ∗ T (cid:15) .
In particular, suppose instead V (cid:54)
=
I and E [ (cid:15) 0 (cid:15) T 0 ] = Σ 0 .
Then we can simply reparameterize as (cid:15)
[Equation]
[Equation]
(cid:15)
[Equation]
, b
[Equation]
→
[Equation]
/ 2b e , β ∗ → Σ 1 / 2 0 β ∗ , A e → A e V − 1 Σ 1 / 2 0 .
It is easy to see this results in the same observed distribution over ( x, y ) , and further that learning β ∗ to predict on (cid:15) 0 is the same as learning Σ 1 / 2 0 β ∗ to predict on Σ − 1 / 2(cid:15) 0 .
So now we’ve reduced this to a setting where
{+E [ (cid:15) 
+}
(cid:15) T 0 ]
=
I but perhaps V (cid:54) =
I .
However, when V (cid:54) =
I it represents precisely the unidentiﬁable transformation M above, which does not pose a problem for prediction because it will not change in future environments.
B P ROOFS OF M AIN R ESULTS B. N OTATION
We use capital letters to denote matrices and lowercase to denote vectors or scalars, where the latter should be clear from context.
(cid:107)·(cid:107) refers to the usual vector norm, or spectral norm for matrices.
For a matrix M , we use λ max ( M ) to mean its maximum eigenvalue—the minimum is deﬁned analogously.
We write the pseudo-inverse as M † .
For a collection of samples { x i } ni =1 , we frequently make use of the sample mean, ˆ µ := 1 n (cid:80) x
i , and the sample covariance, ˆΣ := 1 n (cid:80) ( x i − ¯ µ )( x i − ¯ µ ) T .
The notation (cid:46) means less than or equal to up to constant factors.
B. S TATEMENT AND PROOF OF L EMMA B.1, AND DISCUSSION OF RELATED RESULTS
Lemma B.1.
Assume our data follows a logistic regression model with regression vector β ∗ and covariates z ∼ N (0 , I ) : log p ( y =1 | z ) p ( y = − 1 | z ) =
β ∗ T z .
Then the solution to the dimension-constrained logistic regression problem
[Equation]
where S ⊆
[ d ] indexes a subset of the dimensions, is equal to αβ ∗ S for some α ∈ (0 , 1] .
Proof.
The logistic regression model can be rewritten:
[Equation]
where (cid:15) is drawn from a standard logistic distribution.
If we are restricted to not use z S c , we can see that these can be modeled simply as an additional noise term.
Thus, our new model is
[Equation]
where τ := β ∗ TS c
z S c ∼ p is symmetric zero-mean noise, independent of z S .
Because we are now modeling the other dimensions as noise, moving forward we will drop the S subscript, writing simply β ∗ T z .
Deﬁne F, f as the CDF and PDF of the distribution of (cid:15) (cid:48) :
= (cid:15) + τ .
Then the MLE population objective can be written
[Equation]
) [ 1 { β ∗ T z + (cid:15) (cid:48) > 0 } log σ ( β T z )
+ 1 {− ( β
[Equation]
z + (cid:15) (cid:48) ) > 0 } log σ ( − β T z )]] . For a ﬁxed z , note that E (cid:15) (cid:48) [ 1 { β ∗ T z + (cid:15) (cid:48) > 0 } ]
= P ( (cid:15) (cid:48) ≥ − β ∗ T z ) =
F ( β ∗ T z ) (since f is symmetric), and therefore taking the derivative of this objective we get
[Equation]
Because f is symmetric, we have F ( z ) = 1 − F ( − z ) , giving
[Equation]
Consider the directional derivative of the loss in the direction β ∗ , at the point β
= αβ ∗ :
[Equation]
Because F is the CDF of a logistic distribution convolved with p , by Fubini’s theorem we have
[Equation]
Further, because p is symmetric [Equation] this is equal to 12 ( E τ ∼ p
σ ( z − τ )]
+ E τ ∼ p [ σ ( z + τ )]) .
Thus, we have
[Equation]
We ﬁrst consider the case where α = 1 .
When β ∗ T z > 0 , the term inside the expectation is positive for all τ (cid:54) = 0 , and vice-versa when β ∗ T z < 0 (this can be veriﬁed by writing the difference as a function of β ∗ T z, τ , and observing that all the terms are non-negative except for a factor of e β ∗ T z − 1 ).
It follows that at the point β = β ∗ , − β ∗ is a descent direction.
Furthermore, since the objective is continuous in α , we can follow this direction by reducing α (that is, moving in the direction − β ∗ ) until the directional derivative vanishes.
Next, consider α = 0 .
Then the directional derivative is
[Equation]
Here, when β ∗ T z > 0 the inner term is negative, and vice-versa for β ∗ T z < 0 , implying that the directional derivative is now negative.
Because the objective is convex, it follows that the optimal choice for α lies in (0 , 1] , being equal to
when τ = 0 almost surely.
It remains to show that the optimal vector has no other component orthogonal to β ∗ —in other words, that the solution is precisely αβ ∗ .
For isotropic Gaussian z , we have for any δ perpendicular to β ∗ that E [ δ T z | β ∗ T z ] = 0 .
Therefore, the gradient in the direction δ is
[Equation]
Since β ∗ and all orthogonal directions form a complete basis, it follows that ∇L ( αβ ∗ ) = 0 and therefore that αβ ∗ is the optimal solution.
Though we prove this lemma only for Gaussian z , we found empirically that the result approximately holds whenever z is dimension-wise independent and symmetric about the origin.
We believe this is a consequence of the Central Limit Theorem: our proof relies on the conditional expectation of inner products with z which converge to Gaussian in distribution as the dimensionality of z grows.
We observe that Li & Duan (1989) prove a much simpler result under a general “linear conditional expectation” condition which is similar to the property we exploit regarding zero-mean conditional expectation of orthogonal inner products with isotropic Gaussians.
Their result is more general, but it allows for any value of α , including negative.
In this case, we would actually be recovering the opposite effects of the ground truth, which is clearly insufﬁcient for test-time prediction.
Heagerty & Zeger (2000) give an analytical closed-form for the solution under a probit model with Gaussian noise; since this is not a logistic model, the Gaussian noise represents signiﬁcantly less model misspeciﬁcation, which explains why the exact closed-form is recoverable.
B. P ROOF OF T HEOREM 5.
Theorem 5.2.
(Closed-form solution to the DARE population objective).
Under model (2), the solution to the DARE population objective (3) for linear regression is ˆΠ β ∗ .
If (cid:15) is Gaussian, then the solution for logistic regression (1) is α ˆ Π β ∗ for some α ∈ (0 , 1] .
Proof.
Observe that under the constraint, regressing on the centered observations is equivalent to regressing on the non-centered observations (since the mean must have no effect on the output), so the solutions to these two objectives must be the same and have the same minimizers.
We therefore consider the solution to the DARE objective but on non-centered observations.
It is immediate that the unconstrained solution to the DARE population objective on non-centered data is β ∗ for both linear and logistic regression.
For linear regression, we observe that because the adjusted covariates in each environment have identity covariance, the excess training risk of a predictor β is exactly (cid:107) β − β ∗ (cid:107) 2 .
Therefore, the solution can be rewritten
[Equation]
Recalling that Σ − 1 / 2 e µ e = b e , the constraint can be written in matrix form as B T β = 0 , and thus we see that the solution is the (cid:96) 2 -norm projection of β ∗ onto the nullspace of B T (i.e., the intersection of the nullspaces of the b e ).
By deﬁnition, this is given by ( I − BB † ) β ∗ =
ˆ Π β ∗ .
To derive the closed-form for logistic regression, write the spectral decomposition B = UDV T , and consider regressing on U T (cid:15) instead of (cid:15) .
As the predictor only affects the objective through its linear projection, the solution to this objective will be U T times the solution to the original objective (that is, for all vectors v , ( U T β ) T U T v = β T v ).
We will denote parameters for the rotated objective with a tilde, e.g. ˜ v :=
U T v .
The constraint in Equation (1) is equivalent to requiring that the mean projection is a constant vector c 1 and, with the inclusion of a bias term, we can WLOG consider c = 0 .
Thus, the constraint can be written ˜ B T ˜ β
= V D ˜ β
=
⇐⇒ D ˜ β
= 0 .
We can therefore see that this constraint is the same as requiring that the dimensions of ˜ β corresponding to the non-zero dimensions of D are 0.
Noting that U T (cid:15) ∼ N (0 , I ) , we now apply Lemma B.1 to see that the solution will be ˜ β
= α ( I −
DD † )˜
β ∗ for some α ∈ (0 , 1] .
Finally, as argued above we can recover the solution to the original objective by rotating back, giving the solution β
= U ˜ β
= αU ( I − DD † ) U T β ∗ =
α ˆΠ β ∗ .
B. P ROOF OF T HEOREM 5.
Theorem B.2 (Theorem 5.6, restated) .
For any ρ ≥ 0 , denote the set of possible test environments
A ρ which contains all parameters ( A e (cid:48) , b e (cid:48) ) subject to Assumptions 5.3 and 5.5 and a bound on the mean: (cid:107) b e (cid:48) (cid:107) ≤
ρ .
For logistic or linear regression, let ˆ β
be the minimizer of the corresponding
DARE objective (1) or (3).
Then,
[Equation]
Furthermore, the DARE solution is minimax :
[Equation]
Proof.
Recall that in an environment e
[Equation]
E e [ y | x ] = β ∗ T Σ − 1 / 2 e x .
So, for an environment e (cid:48) and predictor ˆ β [Equation] we have the following excess risk decomposition:
[Equation]
Observe that term T 1 does not depend on the mean b e (cid:48) .
Term T 2 simpliﬁes to
[Equation]
and so we can write a supremum over T 2 as
[Equation]
Next, observe that T 1 simpliﬁes to
[Equation]
So, returning to the full loss and recalling that T 1 is independent of b e (cid:48) , we have
[Equation]
Of course, the ideal would be for a given environment e (cid:48) to set ˆ β :
=
¯Σ 1 / 2 Σ − 1 / 2 e (cid:48) β ∗ =
⇒ v = 0 , but we have to choose a single ˆ β for all possible environments e (cid:48) parameterized by ( A e (cid:48) , b e (cid:48) )
∈
A ρ .
We will show that the choice of ˆ β := α ˆΠ β ∗ is minimax-optimal under this set for any α ∈ (0 , 1] .
Leaving the supremum over adversary choices implicit, we can rewrite the squared norm of v as
[Equation]
By Lemma C.4, we can deﬁne an environment by deﬁning the block terms of U ˆΠ ∆ U T ˆΠ directly.
Consider the choice of ∆ 1 = λ
[Equation]
2 , ∆
[Equation]
= ∆ 12 = ∆ 21 = 0 .
This choice is in A ρ for any λ ∈ R since it is block-diagonal and (cid:107) ∆ 2 (cid:107) = 0 .
Now we can write
[Equation]
via Cauchy-Schwarz and the triangle inequality.
So, taking λ → ∞ means that the minimax risk is unbounded unless ˆ β I − ˆΠ =
⇐⇒ ˆ β = ˆ β ˆΠ .
For the remainder of the proof we consider only this case.
With this restriction
[Equation]
we have
[Equation]
Assumption 5.3 implies that
[Equation]
Consider the DARE solution ˆ β = αβ ∗ ˆΠ for α ∈ (0 , 1] .
Then using the equivalent supremized set from Assumption 5.5 and Lemma C. [Equation] the worst-case risk of this choice is
[Equation]
{+β ∗ I − ˆΠ (cid:107) 
+}
+ 2 B (cid:107) β ∗ ˆΠ (cid:107)(cid:107) β ∗
I − ˆΠ (cid:107) ) .
It remains to show that any other choice of ˆ β results in greater risk.
Observe that the second two terms of (4) are constant with respect to ∆ , so we focus on the ﬁrst term.
That is, we show that any other choice results in sup (cid:107)
∆ β ∗ ˆΠ (cid:107) 2 < (cid:107) β ∗ ˆΠ (cid:107) 2 (cid:107) (∆ + I )ˆ β ˆΠ − β ∗ ˆΠ (cid:107) 2 > (cid:107) β ∗ ˆΠ (cid:107) 2 (except for ˆ β = 0 , which we address at the end).
For any choice of ˆ β ˆΠ , decompose the vector into its projection and rejection onto β ∗ ˆΠ as ˆ β ˆΠ = αβ ∗ ˆΠ + δ , with δ T β ∗ ˆΠ = 0 .
The adversary can choose ∆ 2 = λδδ T , which lies in A ρ for any λ .
Then taking λ → ∞ causes risk to grow without bound—it follows that we must have δ = 0 .
Next, consider any choice α (cid:54)∈ (0 , 1] .
If α > 2 or α < 0 , choosing ∆ 2 = 0 makes this term ( α − 1) 2 (cid:107) β ∗ ˆΠ (cid:107) 2 > (cid:107) β ∗
ˆΠ (cid:107) 2 .
If 2 ≥ α > 1 , choosing ∆ 2 = 1 α I makes this term α 2 (cid:107) β ∗ ˆΠ (cid:107) 2 > (cid:107) β ∗
ˆΠ (cid:107) 2 .
Finally, if α = 0 then this term is equal to (cid:107) β ∗ ˆΠ (cid:107) 2 —however, this value is the supremum of the adversarial risk of the DARE solution and it cannot actually be attained.
B. P ROOF OF T HEOREM 5.
Theorem B.3 (Theorem 5.7, restated) .
Fix test environment parameters A e (cid:48) , b e (cid:48) and our guess ¯Σ .
Suppose we minimize the DARE regression objective (3) on environments whose means b e are
Gaussian vectors with covariance Σ b , with span (Σ b ) =
span ( I − Π) .
After seeing E training domains:
If E ≥ rank(Σ b ) then DARE recovers the minimax-optimal predictor almost surely: ˆ β = β ∗ Π . 2.
Otherwise, if E ≥ r (Σ b ) then with probability ≥ 1 − δ ,
[Equation]
where O ( · ) hides dependence on (cid:107) ∆ (cid:107) .
Proof.
Deﬁne ˆΠ E as the projection onto the nullspace of B T after having seen E environments.
Item 1 is immediate, since as soon as we observe E = rank(Σ b ) linearly independent b e
we have that span ( B ) =
span (Σ b ) and therefore ˆΠ E
=
Π (this occurs almost surely for any absolutely continuous distribution).
To prove item 2, we will write the solution learned after seeing E environments as ˆ β E :=
ˆΠ E β ∗ .
We can write the excess risk of the ground-truth minimax predictor β ∗ Π as
[Equation]
and likewise we have
[Equation]
Taking the difference,
[Equation]
Now we note that
[Equation]
and
[Equation]
These matrices are bounded by (cid:107) ∆ +
I (cid:107) 2 , (cid:107) ∆ + I (cid:107) respectively and are constant with respect to the training environments we sample.
It follows that we can bound the risk difference as
[Equation]
and all that remains is to bound the term (cid:107) ˆΠ E − Π (cid:107) .
Combining Theorems 4 and 5 from Koltchinskii & Lounici (2017) with the triangle inequality, we have that when r (Σ b ) ≤ E , with probability ≥ 1 − δ ,
[Equation]
Since r (Σ b ) ≤ E , the ﬁrst term of the second max dominates.
Further, Corrolary 3 of Yu et al.
(2014), a variant of the Davis-Kahan theorem, gives us
[Equation]
Combining these facts gives the result.
B. P ROOF OF T HEOREM 5.
Theorem B.4 (Theorem 5.9, fully stated) .
Assume the data follows model (2) with (cid:15) ∼ N (0 , I ) .
Observing n S samples from a source distribution S with covariance Σ S , we use half to estimate ˆΣ S and the other half to learn parameters β which minimize the unconstrained ( λ = 0 ), uncentered
DARE regression objective.
At test-time, given n T samples { x i } n T
i
=
1 from the target distribution T with mean and covariance µ T , Σ T , we predict f ( x ; β ) =
β T ˆΣ − 1 / 2 T x .
Deﬁne m (Σ) :=
λ max (Σ) λ 3min (Σ) , and assume n S = Ω( m (Σ S ) d 2 ) , n T = Ω( m (Σ T ) d 2 ) .
Then with probability at least 1 − 3 d − 1 , the excess squared error of our predictor on the new environment is bounded as
[Equation]
Proof.
We begin by bounding the error of our solution (cid:107) β − β ∗ (cid:107) .
Observe that with our estimate ˆΣ S of the source environment moments, we are solving linear regression with targets β ∗ T (cid:15)
i + η i and covariates
ˆ x i =
ˆΣ − 1 / 2 S x i =
ˆΣ − 1 / 2 S Σ S 1 / 2 (cid:15) i .
Thus, if we had access to the true gradient of the modiﬁed least-squares objective (which is not the same as assuming n S → ∞ , because in that case we would have ˆΣ S → Σ S ), the solution would be E
[Equation]
Moving forward we denote this solution to the modiﬁed objective as ˆ β := ˆΣ 1 / 2 S Σ S − 1 / 2 β ∗ , and further deﬁne ∆ S :
= Σ S 1 / 2 ˆΣ − 1 / 2 S , with ∆ T deﬁned analogously.
A classical result tells us that the
OLS solution is distributed as N (cid:16) ˆ β, σ 2 y n S M − 1 (cid:17) , where M is the modiﬁed covariance ∆ TS ∆ S and σ 2 y :=
E [ η 2 ] .
To show a rate of convergence to the OLS solution, we need to solve for the minimum eigenvalue λ min ( M ) —this will sufﬁce since the above fact implies ﬁnite-sample convergence of the
OLS estimator to the population solution.
The well-known bound for sub-Gaussian random vectors tells us that with probability ≥ 1 − δ 1 ,
[Equation]
and moving forward we condition on this event.
Now let γ S :
= (cid:112) d/n S , with γ T deﬁned analogously.
Since M (cid:23) 0 , it follows that
[Equation]
and further,
[Equation]
By Lemma C.1, we have that with probability ≥
[Equation]
This implies that our solution’s error can be bounded as
[Equation]
We assume (cid:107) β ∗ (cid:107) =
1 so we can avoid carrying it throughout the rest of the proof.
Bounding the second term with Lemma C.2 gives
[Equation]
On the target distribution [Equation] our excess risk with a predictor β is
[Equation]
Now [Equation] we have
[Equation]
Once again invoking Lemma C.2, with probability ≥ 1 − d − 1 ,
[Equation]
and using this plus the previous result, the triangle inequality, and ( a + b ) 2 ≤ 2( a 2 + b 2 ) gives
[Equation]
where the lower bound on n T enforces γ T (cid:112)
d m (Σ T ) ≤ 1 .
It follows that the excess risk can be bounded as
[Equation]
Letting δ 1 = 1 /d , combining all of the above via union bound, and eliminating lower-order terms, we get
[Equation]
with probability ≥ 1 − 3 d − 1 .
C T ECHNICAL L EMMAS
Lemma C.1.
Suppose we observe n ∈ Ω( d + log 1 /δ )
samples X ∼ N ( µ, Σ) with Σ (cid:23) 0 and estimate the inverse covariance matrix Σ − 1 with the inverse of the sample covariance matrix ¯Σ − 1 .
Then with probability ≥ 1 − δ , it holds that
[Equation]
n
[Equation]
Proof.
As ¯Σ − 1 / 2 Σ¯Σ − 1 / 2 and Σ 1 / 2 ¯Σ − 1 Σ 1 / 2 have the same spectrum, it sufﬁces to bound the latter.
Observe that
[Equation]
Now applying Theorem 10 of Kereta & Klock (2021) with A = B = Σ 1 / 2 yields the result.
Lemma C.2.
Assume the conditions of Lemma C.1.
Then under the same event as Lemma C.1, it holds that
[Equation]
where γ = (cid:113) d +log1 /δ n .
In particular, if δ = d − 1 , then
[Equation]
Proof.
We begin by deriving a bound for (cid:107) ¯Σ − 1 / 2 − Σ − 1 / 2 (cid:107) .
Deﬁne E = ¯Σ − 1 − Σ − 1 .
Observe that
[Equation]
, and now apply Lemma C.1 (since ¯Σ − 1 / 2 Σ¯Σ − 1 / 2 and Σ 1 / 2 ¯Σ − 1 Σ 1 / 2 have the same spectrum), giving
[Equation]
Let UDU T be the eigendecomposition of Σ , and deﬁne the matrix [ √· , α ] i,j = 1 √ D ii + √ D jj as in
Carlsson (2018).
The Daleckii-Krein theorem (Daleckii & Krein, 1965) tells us that
[Equation]
Note that max i, [Equation] [Equation] [Equation]
[Equation]
⇒
[Equation] [ √· , α ] [Equation] ≤ (cid:112) d (cid:107) Σ − 1 (cid:107) , and therefore by sub-multiplicativity of spectral norm under Hadamard product,
[Equation]
Finally, noting that
[Equation]
completes the main proof.
To see the second claim, note that n ≥ 2 (cid:107) Σ − 1 (cid:107) implies (cid:107) Σ − 1 (cid:107) 3 / 2 √ d ≥ (cid:107)
Σ − 1 (cid:107) 2 γ , meaning the ﬁrst of the two terms dominates.
Lemma C.3.
Let ¯ Σ be ﬁxed.
Then Assumption 5.5 is satisﬁed if and only if (cid:107) ∆ β ∗ ˆ Π (cid:107) 2 < (cid:107) β ∗ ˆ Π (cid:107) 2 .
Proof.
The claim follows by rewriting the risk terms.
Recall that ∆ =
Σ 1 / 2 e (cid:48)
¯Σ − 1 / 2 − I .
Writing out the excess risk of the ground truth β ∗ in the subspace ˆΠ ,
[Equation]
Next, the excess risk of the null vector ˆ β = 0 in the same subspace is
[Equation]
Lemma C.4.
For a ﬁxed ¯Σ , choosing an environmental covariance Σ e (cid:48) is equivalent to directly choosing the error terms ∆ 1 , ∆ 2 , ∆ 12 , ∆ 21 .
Proof.
For a ﬁxed ¯Σ , due to the unique deﬁnition of square root as a result of Assumption 5.1, the map Σ e (cid:48)
→ Σ 1 / 2 e (cid:48) ¯Σ − 1 / 2 −
I is one-to-one.
Recall that we can write
[Equation]
which deﬁnes a one-to-one map from ∆ to the error terms.
Since the composition of bijective functions is bijective, the claim follows.
D I MPLEMENTATION D ETAILS D. E VALUATION PIPELINE
[Figure]
Figure 2 depicts the overall pipeline for evaluation of the different approaches we describe in Section 2.
Our baselines are three distinct pipelines, each slightly different:
The ﬁrst pipeline (dark blue in Figure
is the standard method
of
evaluating ERM on a
new
domain. We train the
entire network (
the
combination
of a feature
embedder
comprising all layers except the last linear layer, and the last layer linear classiﬁer
end-to-end on the training domains, simultaneously learning the features and the linear classiﬁer on the training domains via backpropagation When evaluated on a new domain, this achieves state-of-the-art accuracy ( Gulrajani & Lopez
Paz,
, but it still performs quite poorly. The second pipeline (red) is very similar to the ﬁrst, but we include the test domain among the domains on which the network is trained end-to-end.
Because this means that the test distribution is one of the training domains, this simulates an “ideal” setting where no distribution shift occurs from train-time to test-time.
As such, it is unsurprising that this approach performs substantially better (though it still leaves a bit to be desired—this raises a separate question about failures of in-distribution learning with multiple domains).
As this approach requires cheating , it is completely unrealistic to expect current methods to even begin to approach this baseline but it gives a good sense of what would be the best performance to hope for. 3.
The ﬁnal pipeline (grey) is our main experimental contribution.
Here, the features (all but the last layer) are
learned without cheating, as in the ﬁrst pipeline.
Next, we freeze the features and cheat only while learning a linear classiﬁer on top of these features.
Not only does this method do signiﬁcantly better than the ﬁrst baseline, it actually performs almost as well as—and sometimes better than—the second pipeline.
Thus, we ﬁnd that standard features learned via ERM without cheating are already good enough for generalization and that the main bottleneck to reaching the accuracy of the second baseline is in learning a good linear classiﬁer only .
This has several important advantages to current methods which attempt to change the entire end-to-end process, which we explicate in Section 2 in the main body.
D. N UMERICAL C OMPARISONS
The features used for each trial are the default hyperparameter sweep of D OMAIN B ED .
For computational reasons, we used fewer random hyperparameter choices per trial, meaning the reported accuracies are not directly comparable.
Nevertheless, our results are signiﬁcant because they are consistently evaluated according to the same methodology.
All algorithms use the same set of features for each trial, so their performances are highly dependent.
To account for this, we perform a one-sided paired t-test between algorithms to determine the best performers.
The fact that DARE consistently bests ERM is somewhat surprising: it is intended to guard against a worst-case distribution shift and depends on the quality of our guess ¯Σ , so we would in general expect worse performance on some datasets.
D. C ODE DETAILS
All features were learned by ﬁnetuning a ResNet-50 using the default settings and hyperparameter sweeps of the D OMAIN B ED benchmark (Gulrajani & Lopez-Paz, 2021).
We extracted features from 3 trials, with 5 random hyperparameter choices per trial, picking the one with the best training domain validation accuracy.
We used the default random splits of 80% train / 20% test for each domain.
Using frozen features, the cheating linear classiﬁer was trained by minimizing the multinomial logistic loss with full-batch SGD with momentum for 3000 steps.
We did not use a validation set.
For the main experiments, all algorithms were trained using full-batch L-BFGS (Liu & Nocedal, 1989).
We used the exact same optimization hyperparameters for all methods; the default learning rate of 1 was unstable when optimizing IRM, frequently diverging, so we lowered the learning rate until IRM consistently converged (which occurred at a learning rate of 0.2).
Since the IRM penalty only makes sense with both a feature embedder and a classiﬁcation vector, we used an additional linear layer for IRM, making the objective non-convex.
Presumably due to their convexity for linear classiﬁers, all other methods were unaffected by this change.
For all methods, we halted optimization once 20 epochs occurred with no increase in training domain validation accuracy; the maximum validation accuracy typically occurred within the ﬁrst 5 epochs.
For stability when whitening (and because the number of samples per domain was often less than the feature dimension), in estimating ˆΣ e for each environment we shrank the sample covariance towards the identity.
Speciﬁcally, we deﬁne ˆΣ e = (1 − ρ ) 1 n (cid:80) ni
=1 x i x Ti + ρI , with ρ = 0 . 1 , and ρ = 0 .
01 for
DomainNet due to its much larger size.
We found that increasing λ beyond ∼ 1 had little-to-no effect on the accuracy, loss, or penalty of the DARE solution (see Appendix E for ablations).
However, we did observe that choosing a very large value for λ (e.g., 10 5 or higher) could result in poor conditioning of the objective, with the result being that the L-BFGS optimizer took several epochs for the loss to begin going down.
E A DDITIONAL E XPERIMENTS
Here we present additional experimental results.
All reported accuracies represent the mean of three trials, and all error bars (where present) display 90% conﬁdence intervals calculated as ± 1 .
645 ˆ σ √ n .
Note though that the results are not independent across methods, so simply checking if the error bars overlap is overly conservative in identifying methods which perform better.
E. A CCURACY OF THE T EST C OVARIANCE W HITENER
As we do not have access to the true test-time covariance, we estimate the adjustment with the average of the train-time adjustments.
Table 2 reports the normalized squared Frobenius norm of our estimate’s error to the sample covariance adjustment, deﬁned as (cid:107) ˆ W − W (cid:107) 2 F (cid:107) W (cid:107) 2 F , where W :
= ˆΣ − 1 / 2 e (cid:48) is the sample covariance adjustment and ˆ W := 1 |E| (cid:80) e ∈E Σ − 1 / 2 e is our averaging estimate.
We ﬁnd that this normalized error is consistently small, meaning our estimated adjustment is reasonably close to the “true” adjustment, relative to its spectrum (we put “true” in quotation marks because the best we can do is estimate it via the sample covariance).
This helps explain why our averaging adjustment performs so well in practice, though we expect future methods could improve on this estimate (particularly on the last domain of PACS).
E. A LIGNMENT OF D OMAIN -S PECIFIC O PTIMAL C LASSIFIERS
As discussed in Section 3, DARE does not project out varying subspaces but rather aligns them such that the adjusted domains have similar optimal classiﬁers simultaneously.
To verify that this is actually happening, we learn the optimal linear classiﬁer for each domain individually and evaluate the interdomain cosine similarity for these vectors for each class.
We see that without adjustment, different domains’ optimal linear decision boundaries have normals with small alignment on average, which explains why trying to learn a single linear classiﬁer which does well on all domains simultaneously performs poorly in most cases.
After alignment, the individually optimal classiﬁers are more aligned, which allows a single classiﬁer to perform better on all domains.
Furthermore, this is done without [Table] throwing away the varying component (as would be done by invariant prediction (Peters et al., 2016)), allowing DARE to use more information and resulting in higher test accuracy.
[Table]
E. R ESULTS ON A DDITIONAL E ND - TO -E ND M ETHODS
Though ERM represents approximate state-of-the-art accuracy, for completeness we include a few additional end-to-end methods to conﬁrm that this fact still holds under our experimental conditions.
We also include additional results on ColoredMNIST:
[Table] [Table]
E. A BLATIONS