Learning Robust Representations via Nuisance-extended Information Bottleneck

Jongheon Jeong, Sihyun Yu, Hankook Lee, Jinwoo Shin

Keywords: out-of-distribution robustness, information bottleneck, representation learning, autoencoder

Abstract: The information bottleneck (IB) is a principled approach to obtain a succinct representation $\mathbf{x} \rightarrow \mathbf{z}$ for a given downstream task $\mathbf{x} \rightarrow \mathbf{y}$: namely, it finds $\mathbf{z}$ that (a) maximizes the (task-relevant) mutual information $I(\mathbf{z}; \mathbf{y})$, while (b) minimizing $I(\mathbf{x}; \mathbf{z})$ to constrain the capacity of $\mathbf{z}$ for better generalization. In practical scenarios where the training data is limited, however, many predictive-yet-compressible signals in the data can be rather from some biases in data acquisition (i.e., less generalizable), so that even the IB objective cannot prevent $\mathbf{z}$ from co-adapting on such (so-called) "shortcut" signals. To bypass such a failure mode, we consider an adversarial threat model of $\mathbf{x}$ under constraint on the mutual information $I(\mathbf{x}; \mathbf{y})$. This motivates us to extend IB to additionally model the nuisance information against $\mathbf{z}$, namely $\mathbf{z}_n$, so that $(\mathbf{z}, \mathbf{z}_n)$ can reconstruct $\mathbf{x}$. To enable the idea, we propose an auto-encoder based training upon the variational IB framework, as well as practical encoder designs to facilitate the proposed hybrid discriminative-generative training considering both convolutional- and Transformer-based architectures. Our experimental results show that the proposed scheme improves robustness of learned representations (remarkably without using any domain-specific knowledge), with respect to multiple challenging modern security measures including novelty detection, corruption (or natural) robustness and certified adversarial robustness.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Submission Guidelines: Yes

Please Choose The Closest Area That Your Submission Falls Into: Deep Learning and representational learning

TL;DR: We propose to model the nuisance of information bottleneck for out-of-distribution generalization.