Source-linked AI summary
On Variational Bounds of Mutual Information
Ben Poole, Sherjil Ozair, Aaron van den Oord, Alexander A. Alemi, George Tucker
TL;DR
High-dimensional mutual-information estimation and optimization remain difficult, while the relationships and tradeoffs among neural variational bounds are unclear. The paper unifies these bounds, introduces a continuum of lower bounds, and evaluates their bias, variance, and gradients. The new bounds trade off bias and variance and support estimation and representation learning, but low-variance, low-bias estimation remains unavailable for large MI with small batches.
Problem
High-dimensional MI estimation and optimization are challenging, and existing methods scale poorly or have brittle behavior while variational estimators remain insufficiently understood.
Method
The paper unifies variational MI bounds, derives a continuum of multi-sample lower bounds, and leverages tractable conditional structure for additional bounds.
Results
The new interpolated bounds trade off bias for variance, while controlled high-dimensional experiments characterize estimator and gradient behavior and applications demonstrate utility in representation learning.
Takeaways & Limitations
Variational MI bounds can be organized and selected through explicit tractability and bias–variance tradeoffs for estimation, optimization, and representation learning.
Takeaways & Limitations
No considered approach provides low-variance, low-bias MI estimates when MI is large and batch size is small.
Abstract
from arXiv · showhide
Estimating and optimizing Mutual Information (MI) is core to many problems in machine learning; however, bounding MI in high dimensions is challenging. To establish tractable and scalable objectives, recent work has turned to variational bounds parameterized by neural networks, but the relationships and tradeoffs between these bounds remains unclear. In this work, we unify these recent developments in a single framework. We find that the existing variational lower bounds degrade when the MI is large, exhibiting either high bias or high variance. To address this problem, we introduce a continuum of lower bounds that encompasses previous bounds and flexibly trades off bias and variance. On high-dimensional, controlled problems, we empirically characterize the bias and variance of the bounds and their gradients and demonstrate the effectiveness of our new bounds for estimation and representation learning.
1. Introduction
Mutual information is useful across machine learning, but existing estimators remain difficult to scale and poorly understood. This paper unifies variational bounds, analyzes their tradeoffs, and introduces new bounds for estimation and representation learning.
- Mutual information measures dependency and supports applications including neuroscience, experimental design, neural-network analysis, and representation learning.
- Existing sample-based MI estimators are brittle because estimator hyperparameters can affect scientific conclusions.
- Few existing parametric and non-parametric MI techniques scale to modern machine-learning dataset sizes and dimensionalities.
- Recent work combines neural-network parameterizations with variational bounds to make MI estimation differentiable and tractable.
- The paper reviews relationships and tradeoffs among variational estimators and introduces results that clarify their properties.
- The new multi-sample lower bounds generalize prior bounds and flexibly trade off bias and variance.
- Known conditional structure yields tractable MI bounds, while controlled high-dimensional experiments evaluate estimator and gradient bias and variance.
- The bounds are applied to decoder-free disentangled representation learning on dSprites.
2. Variational bounds of MI
The paper unifies variational MI bounds, characterizes their bias–variance tradeoffs, and introduces multi-sample bounds that interpolate between existing estimators. It also derives structured upper and lower bounds when conditional distributions are tractable.
- Unified framework: Existing variational MI bounds are organized in a unified framework spanning normalized, unnormalized, and multi-sample constructions.The framework derives earlier bounds and motivates interpolated bounds from the multi-sample extension.
- Normalized upper bounds: A tractable variational upper bound uses q(y) to approximate the intractable marginal p(y), is tight when q(y) = p(y), and requires tractable log q(y).In representation learning, this upper bound can regularize the information capacity of stochastic representations.
- Unnormalized lower bounds: Unnormalized bounds are tractable and tight with an optimal critic, but can have high variance from estimating upper bounds on log partition functions.Monte Carlo evaluation can also fail to preserve a finite-sample MI bound.
- Multi-sample lower bounds: Multi-sample extensions reduce variance and recover INCE, which is a valid MI lower bound but becomes loose when I(X; Y) > log K.INCE’s finite batch-size limitation is especially relevant when mutual information is high.
- Interpolated lower bounds: The interpolated lower bounds bridge low-bias, high-variance INWJ and high-bias, low-variance INCE through a tunable parameter α.This continuum is designed to trade off bias and variance rather than selecting only one endpoint.
- Structured bounds: When p(y|x) is tractable, known conditional structure yields variational upper bounds and can sandwich MI without learned variational distributions.Combining structured upper and lower bounds also provides tractable upper bounds on total correlation in suitable representation-learning settings.
3. Experiments
Experiments characterize estimator and gradient tradeoffs across controlled MI problems, batch sizes, critic architectures, and representation learning. The Iα bounds improve the bias–variance balance and support decoder-free disentangled representation learning.
- Controlled MI estimation: The experiments evaluate MI bounds on tractable Gaussian and nonlinear transformed problems, then analyze estimator and gradient bias–variance tradeoffs.The nonlinear problem applies a full-rank random linear transformation followed by a cubic nonlinearity.
- Critic architectures: Separable critics increase estimator variance and generally underperform joint critics for INWJ and IJS, although joint critics scale poorly with batch size.Separable critics require 2N forward passes for batch size N, whereas joint critics require N^2 forward passes.
- Bias–variance tradeoff: INWJ is unbiased but high-variance at large MI and small batch sizes, whereas INCE has low variance but bias that grows linearly with MI.INCE is upper bounded by log(batch size), causing saturation when MI continues increasing.
- Bias–variance tradeoff: The Iα bounds can achieve lower MSE than INWJ or INCE by combining less bias than INCE with lower variance than INWJ.The leave-one-out upper bound also has large bias and variance when the batch size is too small.
- Gradient accuracy: For encoder gradients, α near 1 is preferred at smaller MI and batch sizes, while α near 0 is preferred at larger MI and batch sizes.The reduced gradient MSE indicates that Iα bounds can serve as objectives for training encoders in the InfoMax setting.
- Decoder-free representation learning: On dSprites, the variational regularized InfoMax objective learns position and scale but not rotation in decoder-free representation learning.The method uses an IJS lower bound with a structured critic and leave-one-out upper bounds for total correlation.
4. Discussion
The paper’s interpolated bounds trade bias for variance, but no considered approach achieves low-bias, low-variance MI estimates when MI is large and batches are small. The experiments also show regularized InfoMax learns position and scale, but not rotation, while key finite-data and objective comparisons remain open.
- 4. Discussion: Interpolated bounds trade off bias and variance, but none provides low-bias, low-variance estimates when MI is large and batch size is small.The paper identifies whether such estimators are impossible or require distributional assumptions or neural-network inductive biases as future work.
- 4. Discussion: Regularized InfoMax learns disentangled features for x/y position and scale, but not rotation on dSprites.Figure 5 visualizes position tuning by latent activation and plots scale or rotation against latent values.
- 4. Discussion: Estimating MI gradients may be easier than estimating MI itself when the objective contains an unknown constant data entropy.The paper notes that maximizing IBA is feasible despite lacking access to that constant.
- 4. Discussion: The analysis and experiments assume an infinite dataset with no overfitting, leaving finite-sample and train-versus-heldout biases for future work.This scope excludes differences in MI between training and heldout data.
- 4. Discussion: Whether MI maximization is more useful than other unsupervised or self-supervised objectives for representation learning remains an open question.The paper notes that alternative objectives may provide other mechanisms for learning downstream-useful features.
A. Summary of mutual information lower bounds
This section summarizes how MI lower bounds and estimators differ in objective and gradient tractability, batch-size dependence, variance, and critic normalization. It also points to the parameterizations and objectives collected in the accompanying tables.
- A. Summary of mutual information lower bounds: Table 2 lists the parameters and objectives used for the mutual information estimators.The paper presents these parameterizations alongside the lower-bound characteristics summarized in Table 1.
- A. Summary of mutual information lower bounds: Table 1 characterizes MI lower bounds by objective and gradient tractability, batch-size dependence, variance, and critic normalization.The table distinguishes tractable versus intractable objectives and gradients, dependent versus independent batch-size behavior, high versus low variance, and normalized versus unnormalized critics.
B. Experimental details
The controlled experiments use 20-dimensional correlated Gaussian data, with a cubic nonlinear variant, and compare separable and joint neural critics. Table 3 reports hyperparameter-optimized results for the toy problems.
- B. Experimental details: The toy dataset samples 20-dimensional x and y from correlated Gaussians with mean 0 and correlation ρ.The true mutual information is computed from the correlation and dimensionality, and ρ is increased over time in Figure 2.
- B. Experimental details: The experiments compare separable critics that embed x and y independently with joint critics that process concatenated pairs.Separable critics require only 2N forward passes through the neural network and are described as more efficient in practice.
- B. Experimental details: Table 3 reports hyperparameter-optimized results on the toy Gaussian and Cubic problems of Figure 2.The table is identified as the summary of these optimized toy-problem results.
C.1. Exhaustive hyperparameter sweep.
An extensive sweep over architectures, nonlinearities, learning rates, and batch sizes finds that Iα outperforms the other approaches across all tested MI levels on both toy problems.
- C.1. Exhaustive hyperparameter sweep.: Iα outperforms all approaches at every tested mutual-information level on both the Gaussian and Cubic problems.The ordering remains qualitatively unchanged after hyperparameter optimization, despite improved absolute estimates.
- C.1. Exhaustive hyperparameter sweep.: The sweep optimizes learning rate, critic architecture, depth, hidden units, nonlinearity, and batch size for each bound.Architectures vary between separable and joint critics, with 1–3 hidden layers, 256–2048 hidden units, ReLU or Tanh, and batch sizes 64–512.
- C.1. Exhaustive hyperparameter sweep.: A bound using known p(y|x) is highly accurate because it only needs to learn the marginal q(y).This result comes from the additional experiments leveraging known conditional distributions.
C.2. Effective bias-variance tradeoffs with Iα
Iα provides a controllable bias–variance tradeoff, interpolating between INCE and INWJ. Compared with alternative interpolation schemes, it reduces bias or variance more effectively across mutual-information levels.
- Iα effectively interpolates between high-bias, low-variance INCE and low-bias, high-variance INWJ.The comparison spans three mutual-information levels and varied architecture settings on a toy 20-dimensional Gaussian problem.
- IJS is competitive at high mutual-information rates but has higher bias and variance than Iα at lower rates.
- Neither product nor linear interpolation reduces bias or variance as well as Iα.The alternatives were evaluated with varied hyperparameters at three mutual-information levels in the same setting as Figure 6.
D. IJS derivation
The IJS derivation uses a Jensen–Shannon-trained density-ratio estimate to construct a KL lower-bound critic. For mutual information, this yields a lower bound involving expectations under the joint distribution and the product of marginals.
- Optimizing the critic with a Jensen–Shannon lower bound can avoid the challenge posed by INWJ’s high variance.The resulting density-ratio estimate is used to construct a critic for the KL lower bound.
- The f-GAN softplus critic exposes a density-ratio estimate through its real-valued logits V(x).
- Using TKL(x) = 1 + V(x) in INWJ preserves a KL lower bound, unlike directly plugging the estimate into a Monte Carlo f-divergence approximation.
- When V(x) exactly equals the log density ratio, the exponential expectation under q is one and the first term equals KL(p∥q).
- For mutual information, p is the joint distribution p(x,y) and q is the product of marginals p(x)p(y).
- I(X; Y) ≥ 1 + E_p(x,y)[V(x,y)] − E_p(x)p(y)[exp(V(x,y))] ≜ IJS.
E. Alternative derivation of ITNCE
This appendix presents a direct derivation of ITNCE from iid minibatch samples, a stochastic encoder, and a minibatch estimate of the intractable marginal. The derivation rewrites mutual information as a sum of equal-expectation terms and uses the minibatch mixture in the logarithm.
- ITNCE is derived directly using K iid samples from the data distribution and a stochastic encoder p(y|x).
- The intractable marginal p(y) is replaced in the derivation by a minibatch estimate formed from the conditional distributions p(y|x_i).
- Mutual information is written as a sum of K terms, each having an expectation equal to the mutual information.
- Multiplying and dividing by the minibatch mixture m(y; x1:K) separates the resulting logarithmic expression into mixture-based terms.