Sparsity-Inducing Categorical Prior Improves Robustness of the Information Bottleneck

TL;DR

This paper introduces a sparsity-inducing categorical prior within the information bottleneck framework, enabling flexible, robust representations that adapt to individual data points and improve out-of-distribution performance.

Contribution

It proposes a novel spike-slab categorical prior that allows data-dependent latent dimensions and models joint uncertainty, enhancing robustness and flexibility over fixed priors.

Findings

Improves accuracy on in-distribution data

Enhances robustness to out-of-distribution samples

Outperforms traditional fixed-dimensional priors

Abstract

The information bottleneck framework provides a systematic approach to learning representations that compress nuisance information in the input and extract semantically meaningful information about predictions. However, the choice of a prior distribution that fixes the dimensionality across all the data can restrict the flexibility of this approach for learning robust representations. We present a novel sparsity-inducing spike-slab categorical prior that uses sparsity as a mechanism to provide the flexibility that allows each data point to learn its own dimension distribution. In addition, it provides a mechanism for learning a joint distribution of the latent variable and the sparsity and hence can account for the complete uncertainty in the latent space. Through a series of experiments using in-distribution and out-of-distribution learning scenarios on the MNIST, CIFAR-10, and ImageNet data, we show that the proposed approach improves accuracy and robustness compared to traditional fixed-dimensional priors, as well as other sparsity induction mechanisms for latent variable models proposed in the literature.