[Paper Review] Learning Representations that Support Robust Transfer of Predictors

Yilun Xu, Tommi Jaakkola|arXiv (Cornell University)|Oct 19, 2021

Domain Adaptation and Few-Shot Learning38 references4 citations

TL;DR

This paper proposes Transfer Risk Minimization (TRM), a novel robust estimation criterion that optimizes representation learning to improve transferability of predictors across unseen environments. By minimizing the risk of applying optimal per-environment predictors to other environments, TRM decomposes into a direct transfer term and a weighted gradient-matching term, outperforming IRM and baselines in out-of-distribution generalization on controlled and real-world datasets like PACS and Office-Home.

ABSTRACT

Ensuring generalization to unseen environments remains a challenge. Domain shift can lead to substantially degraded performance unless shifts are well-exercised within the available training environments. We introduce a simple robust estimation criterion -- transfer risk -- that is specifically geared towards optimizing transfer to new environments. Effectively, the criterion amounts to finding a representation that minimizes the risk of applying any optimal predictor trained on one environment to another. The transfer risk essentially decomposes into two terms, a direct transfer term and a weighted gradient-matching term arising from the optimality of per-environment predictors. Although inspired by IRM, we show that transfer risk serves as a better out-of-distribution generalization criterion, both theoretically and empirically. We further demonstrate the impact of optimizing such transfer risk on two controlled settings, each representing a different pattern of environment shift, as well as on two real-world datasets. Experimentally, the approach outperforms baselines across various out-of-distribution generalization tasks. Code is available at \url{https://github.com/Newbeeer/TRM}.

Motivation & Objective

• To address the challenge of out-of-distribution generalization under domain shift, where standard ERM fails due to spurious, environment-specific features.
• To develop a learning objective that directly optimizes for robust transfer of predictors to unseen environments.
• To improve upon existing invariance-based methods like IRM by formulating a criterion that better supports generalization in nonlinear settings.
• To evaluate the method on controlled shift patterns and real-world benchmarks, demonstrating superior robustness against spurious features.

Proposed method

• Proposes Transfer Risk Minimization (TRM), a criterion that measures the average risk of applying an optimal predictor trained on one environment to another adversarially chosen environment.
• Decomposes transfer risk into two components: a direct transfer term and a weighted gradient-matching term that encourages consistency of predictors across environments.
• Develops an alternating optimization algorithm to jointly learn the feature representation and per-environment predictors under the TRM objective.
• Uses a two-stage training process: first optimizing the predictor for each environment, then updating the representation to minimize transfer risk.
• Applies the method to both controlled datasets (10C-CMNIST, SceneCOCO) and real-world benchmarks (PACS, Office-Home) with standard backbones like ResNet18 and ResNet50.
• Adapts TRM for group distributional robustness by combining groups to preserve invariant feature learning, even with spurious correlations.

Experimental results

Research questions

• RQ1 Can a representation learning objective be designed to directly optimize for robust transfer of predictors across unseen environments?
• RQ2 How does TRM compare to IRM and other invariance-based methods in terms of out-of-distribution generalization performance?
• RQ3 Does TRM effectively mitigate the negative impact of spurious, environment-specific features in training data?
• RQ4 Can TRM generalize well in both controlled shift settings and real-world domain generalization benchmarks?

Key findings

• On the PACS dataset with ResNet50, TRM achieves a test accuracy of 82.9% on the average and 70.9% on the worst group, outperforming ERM, IRM, and GroupDRO.
• On the Office-Home dataset, TRM achieves 65.5% average accuracy, outperforming ERM (64.1%), IRM (64.3%), and REx (64.8%), with strong performance on the worst group.
• In the CelebA group distributional robustness experiment, TRM achieves a worse-group accuracy of 90.3%, surpassing GroupDRO (90.0%) and Reweight (89.3%), while maintaining high average accuracy.
• On the 10C-CMNIST and SceneCOCO controlled settings, TRM shows superior robustness when non-causal features are present, outperforming baselines in out-of-distribution generalization.
• TRM’s training time per epoch is 8.69s, only moderately higher than ERM (5.81s), indicating efficient optimization despite additional computation for gradient matching.
• The ablation study confirms that the weighted gradient-matching term is essential for robustness, especially under non-linear feature transformations.