[Paper Review] How Does Information Bottleneck Help Deep Learning?

Kenji Kawaguchi, Zhun Deng+2 more|arXiv (Cornell University)|May 30, 2023

Stochastic Gradient Optimization Techniques15 citations

TL;DR

The paper provides the first rigorous generalization bounds linking information bottleneck regularization to generalization in deep learning, including scenarios where the encoder is learned with training data, and validates the theory with experiments across architectures.

ABSTRACT

Numerous deep learning algorithms have been inspired by and understood via the notion of information bottleneck, where unnecessary information is (often implicitly) minimized while task-relevant information is maximized. However, a rigorous argument for justifying why it is desirable to control information bottlenecks has been elusive. In this paper, we provide the first rigorous learning theory for justifying the benefit of information bottleneck in deep learning by mathematically relating information bottleneck to generalization errors. Our theory proves that controlling information bottleneck is one way to control generalization errors in deep learning, although it is not the only or necessary way. We investigate the merit of our new mathematical findings with experiments across a range of architectures and learning settings. In many cases, generalization errors are shown to correlate with the degree of information bottleneck: i.e., the amount of the unnecessary information at hidden layers. This paper provides a theoretical foundation for current and future methods through the lens of information bottleneck. Our new generalization bounds scale with the degree of information bottleneck, unlike the previous bounds that scale with the number of parameters, VC dimension, Rademacher complexity, stability or robustness. Our code is publicly available at: https://github.com/xu-ji/information-bottleneck

Motivation & Objective

• Provide a rigorous learning theory linking information bottleneck to generalization in deep learning.
• Show that controlling information bottleneck can bound generalization error in end-to-end learning of intermediate representations.
• Improve upon previous conjectures by deriving bounds that depend on conditional mutual information I(X;Z|Y) and encoder-data dependence I(φ(S);S).
• Demonstrate through experiments that generalization correlates with information bottleneck measures across architectures and settings.

Proposed method

• Model the neural network as a composition f^s = g_l^s o φ_l^s with φ_l^s the encoder and g_l^s the remainder of the network.
• Derive generalization bounds that involve I(X;Z_l^s|Y) and, in the learned-encoder case, I(φ_l^S;S) as measures of information preservation and overfitting.
• Provide two main results: Theorem 1 for fixed encoders independent of s, and Theorem 2 for encoders learned with s, linking generalization to information quantities.
• Show that replacing 2^{I(X;Z)} with I(X;Z|Y) yields tighter, linear-in-information bounds.
• Address infinite-domain issues and propose corollaries to handle binning-based estimations of mutual information.
• Support theoretical findings with experiments on CIFAR10 and other datasets, comparing representations and model compression metrics as predictors of generalization.

Experimental results

Research questions

• RQ1 How does information bottleneck regularization relate to generalization error in deep neural networks?
• RQ2 Can a rigorous generalization bound be established when the encoder is learned from the training data?
• RQ3 Is using conditional mutual information I(X;Z|Y) a better predictor of generalization than I(X;Z) or other complexity measures?
• RQ4 Do empirical estimates of information bottleneck quantities and encoder-data dependence predict generalization across architectures?

Key findings

• A novel generalization bound shows that simplicity in both the representation and the representation function supports generalization.
• Replacing exponential dependence on mutual information with linear dependence (I(X;Z|Y)) yields tighter bounds.
• When the encoder is learned with data, the bound involves I(X;Z|Y) plus I(φ(S);S), capturing both representation compression and encoder overfitting.
• Empirical results on CIFAR10 and MNIST indicate that bounds based on combined representation and model compression outperform those based on representation compression alone.
• Taking the minimum over layers of the combined term I(S;θ_l^S) + I(X;Z_l^s|Y) provides strong predictive power for generalization gaps.
• The approach resolves binning arbitrariness for mutual information and remains valid under common estimators and transfer learning settings.