Layer-wise Learning of Stochastic Neural Networks with Information Bottleneck

The Information Bottleneck (IB) method is an unsupervised non-parametric technique that extracts from one variable the relevant information about another variable. In this work, we introduce Parametric Information Bottleneck (PIB), a parameterized version of the IB principle to neural networks in a layer-wise manner. In particular, PIB solely relies on the IB principle to learn all layers of a neural network for supervised learning tasks. This framework attempts at maximizing the layer-wise informativeness about the target variable under a compression constraint. We develop an efficient variational approximation to the intractable mutual information and show that the resulting PIB principle, as compared to the standard Maximum Likelihood Estimate (MLE) principle, improves the generalization, interpretation, and adversarial robustness of neural networks.