Gated Information Bottleneck for Generalization in Sequential Environments

TL;DR

This paper introduces Gated Information Bottleneck (GIB), a neural network method that dynamically filters out irrelevant features to improve generalization across different environments, especially in sequential settings.

Contribution

GIB is a novel neural network-based IB approach with a simple objective that adaptively selects task-relevant features without variational assumptions.

Findings

GIB outperforms other IB methods in adversarial robustness.

GIB improves out-of-distribution detection.

GIB performs well in sequential environment scenarios where IRM fails.

Abstract

Deep neural networks suffer from poor generalization to unseen environments when the underlying data distribution is different from that in the training set. By learning minimum sufficient representations from training data, the information bottleneck (IB) approach has demonstrated its effectiveness to improve generalization in different AI applications. In this work, we propose a new neural network-based IB approach, termed gated information bottleneck (GIB), that dynamically drops spurious correlations and progressively selects the most task-relevant features across different environments by a trainable soft mask (on raw features). GIB enjoys a simple and tractable objective, without any variational approximation or distributional assumption. We empirically demonstrate the superiority of GIB over other popular neural network-based IB approaches in adversarial robustness and out-of-distribution (OOD) detection. Meanwhile, we also establish the connection between IB theory and invariant causal representation learning, and observed that GIB demonstrates appealing performance when different environments arrive sequentially, a more practical scenario where invariant risk minimization (IRM) fails. Code of GIB is available at https://github.com/falesiani/GIB