Invariant Information Bottleneck for Domain Generalization

Invariant risk minimization (IRM) has recently emerged as a promising alternative for domain generalization. Nevertheless, the loss function is difficult to optimize for nonlinear classifiers and the original optimization objective could fail when pseudo-invariant features and geometric skews exist. Inspired by IRM, in this paper we propose a novel formulation for domain generalization, dubbed invariant information bottleneck (IIB). IIB aims at minimizing invariant risks for nonlinear classifiers and simultaneously mitigating the impact of pseudo-invariant features and geometric skews. Specifically, we first present a novel formulation for invariant causal prediction via mutual information. Then we adopt the variational formulation of the mutual information to develop a tractable loss function for nonlinear classifiers. To overcome the failure modes of IRM, we propose to minimize the mutual information between the inputs and the corresponding representations. IIB significantly outperforms IRM on synthetic datasets, where the pseudo-invariant features and geometric skews occur, showing the effectiveness of proposed formulation in overcoming failure modes of IRM. Furthermore, experiments on DomainBed show that IIB outperforms 13 baselines by 0.9% on average across 7 real datasets. Introduction In most statistical machine learning algorithms, a fundamental assumption is that the training data and test data are independently and identically distributed (i.i.d.). However, the data we have in many real-world applications are not i.i.d. Distributional shifts are ubiquitous. Under such circumstances, classic statistical learning paradigms with strong generalization guarantees, e.g., Empirical Risk Minimization (ERM) (Vapnik 1999), often fail to generalize due to the violation of the i.i.d. assumption. It has been widely observed that the performance of a model often deteriorates dramatically when it is faced with samples from a different domain, even under a mild distributional shift (Arjovsky et al. 2019). On the other hand, collecting training samples from all possible future scenarios is essentially infeasible. Hence, understanding and improving the generalization of models on out-of-distribution data is crucial. Domain generalization (DG), which aims to learn a model from several different domains so that it generalizes to unCopyright © 2022, Association for the Advancement of Artificial Intelligence (www.aaai.org). All rights reserved. seen related domains, has recently received much attention. From the perspective of representation learning, there are several paradigms towards this goal, including invariant representation learning (Muandet, Balduzzi, and Schölkopf 2013; Zhao et al. 2018; Tachet des Combes et al. 2020), invariant causality prediction (Arjovsky et al. 2019; Krueger et al. 2020b), meta-learning (Balaji, Sankaranarayanan, and Chellappa 2018; Du et al. 2020), and feature disentanglement (Du et al. 2020; Peng et al. 2019). Of particular interest is the invariant learning methods. Some early works, e.g., DANN (Ganin et al. 2017), CDANN (Long et al. 2018), aim at finding representations that are invariant across domains. Nevertheless, learning invariant representations fails for domain adaptation or generalization when the marginal label distributions change between source and target domains (Zhao et al. 2019a). Recently, Invariant Causal Prediction (ICP), and its follow-up Invariant Risk Minimization (IRM), have attracted much interest. ICP assumes that the data are generated according to a structural causal model (SCM) (Pearl 2010). The causal mechanism for the data generating process is the same across domains, while the interventions can vary among different domains. Under such data generative assumptions, IRM (Arjovsky et al. 2019) attempts to learn an optimal classifier that is invariant across domains. ICP then argues that under the SCM assumption, such a classifier can generalize across domains. Despite the intuitive motivations, IRM falls short in several aspects. First, the proposed loss function in (Arjovsky et al. 2019) is difficult to optimize when the classifier is nonlinear. Furthermore, it has been shown that IRM fails when the pseudo-invariant features (Rosenfeld, Ravikumar, and Risteski 2020) or geometric skews exist (Nagarajan, Andreassen, and Neyshabur 2021). Under such circumstances, the classifier will utilize both the causal and spurious features, leading to a violation of invariant causal prediction. To address the first issue, we propose an information-theoretical formulation of invariant causal prediction and adopt a variational approximation to ease the optimization procedure. To tackle the second issue, we emphasize that the use of pseudo-invariant features or geometric skews will inevitably increase the mutual information between the inputs and the representations. Thus, to mitigate the impact of pseudoinvariant features and geometric skews, we propose to constrain this mutual information, which naturally leads to a ar X iv :2 10 6. 06 33 3v 5 [ cs .L G ] 1 0 D ec 2 02 1 formulation of information bottleneck. Our empirical results show that the proposed approach can effectively improve the accuracy when the pseudo-invariant features and geometric skews exist. Contributions: We propose a novel information-theoretic formulation for domain generalization, termed as invariant information bottleneck (IIB). IIB aims at minimizing invariant risks while at the same time mitigating the impact of pseudo-invariant features and geometric skews. Specifically, our contributions can be summarized as follows: (1) We propose a novel formulation for invariant causal prediction via mutual information. We further adopt variational approximation to develop tractable loss functions for nonlinear classifiers. (2) To mitigate the impact of pseudo-invariant features and geometric skews, inspired by the information bottleneck principle, we propose to constrain the mutual information between the inputs and the representations. The effectiveness is verified by the synthetic experiments of failure modes (Ahuja et al. 2021; Nagarajan, Andreassen, and Neyshabur 2021), where IIB significantly improves the performance of IRM. (3) Empirically, we analyze IIB’s performance with extensive experiments on both synthetic and large-scale benchmarks. We show that IIB is able to eliminate the spurious information better than other existing DG methods, and achieves consistent improvements on 7 datasets by 0.7% on DomainBed (Gulrajani and Lopez-Paz 2020). Related Work Domain Generalization Existing methods of DG can be divided into three categories: (1) Data Manipulation: Machine learning models typically rely on diverse training data to enhance the generalization ability. Data manipulation/augmentation methods (Nazari and Kovashka 2020; Riemer et al. 2019) aim to increase the diversity of existing training data with operations including flipping, rotation, etc. Domain randomization (Borrego et al. 2018; Yue et al. 2019; Zakharov, Kehl, and Ilic 2019) provides more complex operations for image data, such as altering the location/texture of objects, replicating and resizing objects. In addition, there are some methods (Riemer et al. 2019; Qiao, Zhao, and Peng 2020; Liu et al. 2018; Truong et al. 2019; Zhao et al. 2019b) that exploits generated data samples to enhance the model generalization ability. (2) Ensemble Learning methods (Mancini et al. 2018; Segù, Tonioni, and Tombari 2020) assume that any sample in the test domain can be regarded as an integrated sample of the multiple-source domains, so the overall prediction should be inferred by a combination of the models trained on different domains. (3) Meta-Learning aims at learning a general model from multiple domains. In terms of domain generalization, MLDG (Li et al. 2018a) divides data from the multiple domains into meta-train and meta-test to simulate the domain shift situation to learn the general representations. In particular, Meta-Reg (Balaji, Sankaranarayanan, and Chellappa 2018) learns a meta-regularizer for the classifier, and Meta-VIB (Du et al. 2020) learns to generate the weights in the meta-learning paradigm by regularizing the KL divergence between marginal distributions of representations of the same category but from different domains. Mutual Information-based Domain Adaptation Domain Adaptation is an important topic in the direction of transfer learning (Long et al. 2015; Ganin et al. 2016; Tzeng et al. 2017; Long et al. 2018; Zhao et al. 2021, 2020c,b; Li et al. 2020a). The mutual information-based approaches have been widely applied in this area. The key idea is to learn a domain-invariant representation that are informative to the label, which can be formulated as (Zhao et al. 2020a; Li et al. 2020b) max Z I(Z,Y ) − λI(Z,A) (1) where A is the identity of domains, Z denotes the representation, and Y denotes the labels. Commonly adopted implementations of (1) are DANN (Ganin et al. 2017) and CDANN (Long et al. 2018). These implementations are also often adopted in domain generalization as baselines (Gulrajani and Lopez-Paz 2020). Invariant Risk Minimization The above approaches enforces the invariance of the learned representations. On the other hand, Invariant Risk Minimization (IRM) suggests the invariance of featureconditioned label distribution. Specifically, IRM seeks for an invariant causal prediction such that E[Y ∣Φ(X)] = E[Y e ′ ∣Φ(X ′ )], for all e, e′ ∈ E . The objective of IRM is given by min w,Φ ∑ e∈Etrain R(w ○Φ),