Domain-Adjusted Regression or: ERM May Already Learn Features Sufficient for Out-of-Distribution Generalization

TL;DR

The paper investigates whether out-of-distribution generalization hinges on learning different features or simply on learning a robust predictor on top of existing features. It shows that using frozen ERM features, a robust linear predictor can achieve near-ideal performance, motivating Domain-Adjusted Regression (DARE), a convex objective that per-domain whitens features and learns a unified predictor in a canonical space with guarantees. Theoretical contributions include closed-form population solutions, a minimax risk characterization under a constrained shift model, and finite-environment convergence guarantees, along with a JIT-UDA extension. Empirically, DARE matches or outperforms prior domain-generalization methods on several benchmarks when operating on fixed features, supporting a modular view: robust prediction on pre-learned features can substantially close the gap to ideal generalization while reducing computational cost. Overall, the work advocates simpler, well-grounded approaches to robust prediction that leverage existing representations rather than relying on end-to-end complex invariance training.

Abstract

A common explanation for the failure of deep networks to generalize out-of-distribution is that they fail to recover the "correct" features. We challenge this notion with a simple experiment which suggests that ERM already learns sufficient features and that the current bottleneck is not feature learning, but robust regression. Our findings 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-specific 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 first finite-environment convergence guarantee to the minimax risk, improving over existing analyses which only yield minimax predictors after an environment threshold. Evaluated on finetuned features, we find that DARE compares favorably to prior methods, consistently achieving equal or better performance.

Figures (5)

• Figure 1: Accuracy via "cheating": (mean over 3 trials) dagger ($\dagger$) denotes access to test domain at train-time. Each letter is a domain. Dark blue is approximate SOTA, orange is our proposed DARE objective, light grey represents cheating while retraining the linear classifier only. All three methods use the same features, attained without cheating. Dark grey is "ideal" accuracy, cheating while training the entire deep network. Surprisingly, cheating only for the linear classifier rivals cheating for the whole network. Cheating accuracy on pretrained features (light blue) makes clear that this effect is due to finetuning on the train domains, and not simply overparameterization (i.e., a very large number of features).
• Figure 2: Depiction of evaluation pipeline. Standard training on train domains leads to poor performance. Cheating while training the full network (features and classifier) does substantially better. However, cheating on just the linear classifier does almost as well, implying that the features learned without cheating are already good enough for massive improvements over SOTA.
• Figure 3: Demonstration of the effect of whitening. NoSigmaDARE is the exact same algorithm as DARE but with no covariance whitening. In almost all cases, covariance whitening + guessing at test-time results in better performance. We expect under much larger distribution shift that this pattern may reverse.
• Figure 4: Effect of penalty term $\lambda$ on the two algorithms which use it. $\lambda=0$ corresponds to no constraint, and the lower performance demonstrates that this invariance requirement is essential to the quality of the learned classifier. For $\lambda\neq0$, DARE accuracy is extremely robust, effectively constant for all $\lambda \geq 1$; in practice we also found the penalty term itself to always be $\sim$0. In contrast, IRM accuracy appears to decrease with increasing $\lambda$, implying that the observed benefit of IRM primarily comes from the domain reweighting as in our Reweighted ERM method.
• Figure 5: Effect of final feature bottleneck dimensionality on cheating accuracy. Reducing the dimensionality reduces accuracy of all methods to varying degrees, though in some cases it actually increases test accuracy. We observe that the main pattern persists, though the gap between cheating on finetuned features and traditional ERM shrinks as the dimensionality is reduced substantially. To reduce dimensionality of the pretrained features we use a random projection; unsurprisingly, the quality dramatically falls as the number of dimensions is reduced.

Theorems & Definitions (29)

• Example 4.1: Invariant Latent Subspace Recovery
• Theorem 5.2
• Remark 5.4
• Theorem 5.6: DARE risk and minimaxity
• proof : Proof sketch
• Definition 5.7
• Definition 5.8
• Theorem 5.9: Environment complexity of DARE
• proof : Proof sketch
• Remark 5.10