Invariant Information Bottleneck for Domain Generalization

TL;DR

This paper introduces the invariant information bottleneck (IIB), a novel approach for domain generalization that improves upon invariant risk minimization (IRM) by better handling pseudo-invariant features and geometric skews, leading to superior performance.

Contribution

The paper proposes a new formulation called invariant information bottleneck (IIB) that enhances IRM by using mutual information to improve domain generalization for nonlinear classifiers.

Findings

IIB outperforms IRM on synthetic datasets with pseudo-invariant features.

IIB surpasses 13 baselines by 0.9% on average across 7 real datasets.

The method effectively mitigates failure modes of IRM.

Abstract

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.