Domain-Adjusted Regression or: ERM May Already Learn Features Sufﬁcient for Out-of-Distribution Generalization
Abstract
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.
The full version of this paper is available on ar χ iv at this https URL.
Introduction
The historical motivation for deep learning focuses on the ability of deep neural networks to automatically learn rich, hierarchical features of complex data [21, 13].
Simple Empirical Risk Minimization (ERM), with appropriate regularization, results in high-quality representations.
Despite these successes, the dominant focus 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 [30, 26].
A common explanation for this phenomenon is reliance on “spurious correlations” or “shortcuts” [29, 12, 41].
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 [36, 11, 1].
These objectives are complex, poorly understood, and difﬁcult to optimize.
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 Workshop on Distribution Shifts, 36th Conference on Neural Information Processing Systems (NeurIPS 2022).
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 pretrained features (light
) makes clear that this effect is due to ﬁnetuning on the train domains, and not simply overparameterization (i.e., a very large number of features). to learn a predictor using a collection of distinct training domains and then perform well on a new, 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.
We also ﬁnd that this cheating method does not ensure good performance on pretrained features, which implies that we are not yet in the “universal approximation” regime and that the effect we observe is indeed due to ﬁnetuning via ERM.
These results suggest that features learned by ERM may be “good enough” for out-of-distribution generalization in deep learning.
For a more detailed description of the experimental setup, see Appendices B & H.
Motivated by these ﬁndings, we propose a new objective, which we call Domain-Adjusted Regression (DARE).
DARE enjoys strong theoretical guarantees: under a new model of distribution shift 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 [31, 7, 40].
Evaluated on ﬁnetuned features, we ﬁnd that DARE compares favorably to existing methods, consistently achieving equal or better performance.
We hope these results will encourage going “back to basics”, with future work focusing on simpler, easier to understand methods for robust prediction.
The Domain-Adjusted Regression Objective
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 [27, 1].
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.
DARE is similar to methods which learn separate batchnorm parameters per domain over a deep network, possibly adjusting at test-time [35, 6, 34].
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.
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.
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.
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.
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.
We also do not know precisely 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.
Note that unlike prior methods, DARE does not enforce invariance of the features themselves.
Rather, it aligns the representations such that different domains share optimal predictors.
The appendix contains discussion of implementation of the DARE pipeline in practice, as well as ablations of various design and constraint choices (Appendix I).
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.
Theorem D.9 provides ﬁnite-sample risk bounds for the
DARE objective in this setting.
A New Model of Distribution Shift
The DARE objective is based on the intuition that all domains arise as unique transformation from a joint representation space.
To capture this, 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.
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) + η .
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 [31, 7, 39].
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 an example in Appendix C.
Theoretical Analysis
Due to space constraints, we brieﬂy list our theoretical results in this section.
Results are stated fully in Appendix D, along with a more involved discussion .
First, we solve for the closed-form solution to the DARE objective under our model:
Theorem 4.1.
Under model (2), the solution to the DARE population objective (5) for linear regression is ˆΠ β ∗ .
If (cid:15) is Gaussian, the solution for (1) is α ˆΠ β ∗ for some α ∈ (0 , 1] .
Next, under some natural assumptions on the strength of the adversary choosing the test distribution, we solve for the exact risk of the DARE predictor. Further, we prove that this risk is minimax:
Theorem 4.2.
Denote the set of possible test environments A ρ which contains all parameters ( A e (cid:48) , b e (cid:48) ) subject to Assumptions D.3 and D.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 D.2.
Then,
[Equation] [Equation]
Furthermore, the DARE solution is minimax :
[Equation]
4.2 The Environment Complexity 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.
Prior methods exclusively give thresholds at which the optimal invariant predictor is recovered, usually linear in the ambient dimension of the varying subspace. However, their behavior when this threshold has not been reached is not well understood. Our next result demonstrates that DARE achieves this same threshold, but we also prove the ﬁrst ﬁnite-environment convergence guarantee , quantifying how quickly the risk of the DARE predictor approaches that of the minimaxoptimal predictor.
Theorem 4.3 (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 ¯Σ .
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]
O (cid:18) (cid:107) Σ b (cid:107) ξ (Σ b ) (cid:18)(cid:113) r
(Σ b ) E + max (cid:26)(cid:113) log1 /δ E , log1 /δ
E (cid:27)(cid:19)(cid:19) where O ( · ) hides dependence on (cid:107) ∆ (cid:107) and r (Σ) :
= Tr(Σ) (cid:107) Σ (cid:107) is the effective rank.
Remark 4.4 . Thus, DARE achieves the same threshold for minimax optimality. But we should expect that if the shift in the test distribution is not too large then we can do better, and indeed Theorem D. 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 ) 4.3 Applying DARE to JIT-UDA
Finally, we provide a ﬁnite-sample risk bound given access to unlabeled test data at test time , with no retraining necessary.
We’d expect that estimating the adjustment via unlabeled samples will improve the 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 testtime, the advantages have not 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 4.5 (JIT-UDA, shortened) . Let Σ S , Σ T be the covariances of the source and target distributions, respectively.
Deﬁne m (Σ) : = λ max (Σ) λ 3min (Σ) . Assume we observe n s = Ω( m (Σ S ) d 2 ) source samples and analogous n T target samples. Then after solving the DARE objective
[Table]
with probability at least 1 − 3 d −
[Table]
the excess squared risk of our predictor on the new environment is bounded as
[Equation]
Experiments
[Table]
mean according to one-sided paired t-tests at p = 0 . 1 signiﬁcance. If not the overall highest, underline indicates higher mean than ERM under the same test.
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 [37], IRM [1], and
GroupDRO [33].
We also discovered that simply rescaling ERM by downweighting each domain proportionally to its size serves as a strong baseline.
Concurrent work makes a similar observation, exploring this baseline when training the whole network rather than just the last layer [17].
We denote this method “Reweighted ERM”.
We evaluate all approaches on four datasets: Ofﬁce-Home [38], PACS [22], VLCS [10], and DomainNet [27].
The features 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.
Additional comparisons to other end-to-end methods can be found in Appendix I.
Results are reported in Table 1.
We ﬁnd that DARE consistently matches or surpasses prior methods.
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.
Prior methods now consistently outperform ERM.
It is interesting that the previously observed gap between ERM and alternatives disappears in this setting with a linear predictor.
For example,
Gulrajani & Lopez-Paz [14] 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 could be because they learn worse features, or perhaps they are just more difﬁcult to optimize over a deep network.
This also implies that when these methods 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.
Limitations
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.
A Related Work
A popular approach to domain generalization matches the domains in feature space, either by aligning moments [36] or with an adversarial loss [11], though these methods are known to be inadequate in general [43].
DARE differs from these approaches in that the constraint requires only that the feature mean projection onto the vector β be invariant.
Domain-invariant projections [2] were recently analyzed by Chen & Bühlmann [8], though notably under fully observed features and only for domain adaptation.
They analyze other invariances as well, and we expect combining these methods with domain adjustment can serve as a direction for future study.
There has been intense recent focus on invariant prediction, based on ideas from causality [28] and catalyzed by IRM [1].
Though the goal of such methods is minimax-optimality under major distribution shift, later work identiﬁes critical failure modes of this approach [31, 18].
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.
Initially suggested for UDA [24, 4, 6], this idea has also been applied to domain generalization [35, 34] but in a somewhat ad-hoc manner—this is problematic because deep domain generalization methods were recently called into question when Gulrajani & Lopez-Paz [14] gave convincing evidence that no method beats ERM when evaluated fairly.
Our analysis provides an initial justiﬁcation for these methods, suggesting that this idea is worth exploring further.
B ERM Learns Surprisingly Useful Features
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 [14], we ﬁnetune a ResNet-50 [15] 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.
[3] 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.
C Example Comparing Invariant Prediction to DARE
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.
D Detailed Theoretical Results
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 D.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.
[32], 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.
With Assumption D.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 D.1 and the constraint in Equation (1) are 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 but possibly still reasonable performance guarantees.
Experiments in Appendix I 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 (5) .
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 D.2.
(Closed-form solution to the DARE population objective).
Under model (2), the solution to the DARE population objective (5) 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.
D.1 The Adversarial Risk 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)
[(ˆ β T Π ¯Σ − 1 / 2 x − β ∗ 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 D.3.
For a ﬁxed constant B ≥ 0 , (cid:107) ( I − ˆΠ)∆ˆΠˆ β (cid:107) ≤ B (cid:107) ˆΠ β ∗ (cid:107) .
Remark D.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 ˆΠ ∆ U ˆΠ =
(cid:20) ∆∆ 12 ∆∆(cid:21) , where ∆ 1 ∈ R rank(ˆΠ)
× rank(ˆΠ) .
A few settings automatically satisfy Assumption D.3 with B = 0 , due to the fact that 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 in this setting, 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 [31], and therefore Assumption D.3 with B = 0 is no stronger than assuming a linear number of environments.
This will similarly occur if U e is shared across all environments (e.g., under ﬁxed A e as described above) and we use any sort of “averaging” guess of ¯Σ .
Assumption D.5 (Bound on adjustment error in non-varying subspace) .
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 (cid:48) ( β ∗ ) <
R ˆΠ p e (cid:48) ( 0 ) .
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 D.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 D.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 D.3 and D.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 D.2.
Then, sup
[Equation]
Furthermore, the DARE solution is minimax : ˆ β ∈ arg
[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 D.6 shows that DARE should outperform invariant prediction in such settings.
D.2 The Environment Complexity 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 D.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 (5) 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 D.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 D.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 ) .
D.3 Applying DARE to JIT-UDA
So far, we have only considered a setting with no knowledge of the test domain.
As discussed in
Section 2, 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 “just-in time” (i.e., only at test-time), the advantages have not 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 D.9 (JIT-UDA, shortened) .
Let Σ S , Σ T be the covariances of the source and target distributions, respectively.
Deﬁne m (Σ) :
= λ max (Σ) λ 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]
E Connection to Anchor Regression
The DARE objective for linear regression is written
[Equation]
The idea of adjusting for domain projections has similarities to Anchor Regression [32], 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 (5).
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.
F Proofs of Main Results F.1 Notation
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.
F.2 Statement and proof of Lemma F.1, and discussion of related results
Lemma F.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]
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 β ∗
[Equation] at the point β = αβ ∗ :
[Equation]
Because F is the CDF of a logistic distribution convolved with p [Equation] 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 [23] 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 [16] 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.
F.3 Proof of Theorem D.
Theorem D.2.
(Closed-form solution to the DARE population objective).
Under model (2), the solution to the DARE population objective (5) 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 F.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 β ∗ =
α ˆΠ β ∗ .
F.4 Proof of Theorem D.
Theorem F.2 (Theorem D.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 D.3 and D.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 (5).
Then,
[Equation]
Furthermore, the DARE solution is minimax :
[Equation]
Proof.
Recall that in an environment e , E e
[ y | x ] = β ∗ T Σ − 1 / 2 e x .
So, for an environment e (cid:48) and predictor ˆ β , 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 sup
[Equation]
Of course, the ideal would be for a given environment e (cid:48) to set ˆ β :=
¯Σ
[Equation] 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 G.4, we can deﬁne an environment by deﬁning the block terms of U ˆΠ ∆ U T ˆΠ directly.
Consider the choice of ∆ 1 = λ
ˆ β
I − ˆΠ ˆ β TI − ˆΠ (cid:107) ˆ β
[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, we have
[Equation]
Assumption D.3 implies that
[Equation]
Consider the DARE solution ˆ β = αβ ∗ ˆΠ for α ∈ (0 , 1] .
Then using the equivalent supremized set from Assumption D.5 and Lemma G.3, the worst-case risk of this choice is
[Equation]
It remains to show that any other choice of ˆ β results in greater risk.
Observe that the second two terms of (6) 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.
F.5 Proof of Theorem D.
Theorem F.3 (Theorem D.7, restated) .
Fix test environment parameters A e (cid:48) , b e (cid:48) and our guess ¯Σ .
Suppose we minimize the DARE regression objective (5) 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: ˆ β = β ∗ Π .
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 [20] with the triangle inequality, we have that when r (Σ b ) ≤ E , with probability ≥ 1 − δ , (cid:107)
[Equation]
max
[Equation]
Since r (Σ b ) ≤ E , the ﬁrst term of the second max dominates.
Further, Corrolary 3 of Yu et al.
[42], a variant of the Davis-Kahan theorem, gives us
[Equation]
Combining these facts gives the result.
F.6 Proof of Theorem D.
Theorem F.4 (Theorem D.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
[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 G.1, we have that with probability ≥ 1 − d −[-
-]
[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 G.2 gives
[Equation]
On the target distribution [Equation] our excess risk with a predictor β is
[Equation]
Now, we have
[Equation]
Once again invoking Lemma G.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 .
G Technical Lemmas
Lemma G.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]
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 [19] with A = B = Σ 1 / 2 yields the result.
Lemma G.2.
Assume the conditions of Lemma G.1.
Then under the same event as Lemma G.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 G.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 [5].
The Daleckii-Krein theorem [9] tells us that
[Equation]
Note that max i, [Equation] | [ [Equation]
[Equation] , 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 G.3.
Let ¯Σ be ﬁxed.
Then Assumption D.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 G.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 D.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.
H Implementation Details H.1 Evaluation pipeline
[Figure]
Figure 2 depicts the overall pipeline for evaluation of the different approaches we describe in
Appendix B. Our baselines are three distinct pipelines, each slightly different:
The ﬁrst pipeline (dark blue in Figure 2) 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 [14], 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 Appendix B in the main body.
H.2 Code details
All features were learned by ﬁnetuning a ResNet-50 using the default settings and hyperparameter sweeps of the D OMAIN B ED benchmark [14].
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 [25].
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 I 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.
I Additional Experiments
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.
I.1 Accuracy of the Test Covariance Whitener
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).
[Table]
I.2 Alignment of Domain-Speciﬁc Optimal Classiﬁers
As discussed in Section [Table] 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 throwing away the varying component (as would be done by invariant prediction [28]), allowing
DARE to use more information and resulting in higher test accuracy.
[Table]
I.3 Results on Additional End-to-End Methods
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]
I.4 Ablations