Counterfactual Supervision-based Information Bottleneck for Out-of-Distribution Generalization

Abstract

Learning invariant (causal) features for out-of-distribution (OOD) generalization has attracted extensive attention recently, and among the proposals invariant risk minimization (IRM) is a notable solution. In spite of its theoretical promise for linear regression, the challenges of using IRM in linear classification problems remain. By introducing the information bottleneck (IB) principle into the learning of IRM, IB-IRM approach has demonstrated its power to solve these challenges. In this paper, we further improve IB-IRM from two aspects. First, we show that the key assumption of support overlap of invariant features used in IB-IRM is strong for the guarantee of OOD generalization and it is still possible to achieve the optimal solution without this assumption. Second, we illustrate two failure modes that IB-IRM (and IRM) could fail for learning the invariant features, and to address such failures, we propose a \textit{Counterfactual Supervision-based Information Bottleneck (CSIB)} learning algorithm that provably recovers the invariant features. By requiring counterfactual inference, CSIB works even when accessing data from a single environment. Empirical experiments on several datasets verify our theoretical results.

Theoretical Understanding of OOD Generalization

References in corpus (10)

• Unsupervised Domain Adaptation by Backpropagation
• Shortcut Learning in Deep Neural Networks
• Domain Generalization via Invariant Feature Representation
• Out-of-Distribution Generalization via Risk Extrapolation (REx)
• Learning Causal Semantic Representation for Out-of-Distribution Prediction
• Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization
• Gradient Starvation: A Learning Proclivity in Neural Networks
• Towards a Theoretical Framework of Out-of-Distribution Generalization
• Does Invariant Risk Minimization Capture Invariance?
• Provable Domain Generalization via Invariant-Feature Subspace Recovery