Improving Mutual Information Estimation with Annealed and Energy-Based Bounds
Rob Brekelmans, Sicong Huang, Marzyeh Ghassemi, Greg Ver Steeg, Roger Grosse, Alireza Makhzani
TL;DR
This work tackles the core challenge of estimating mutual information $I(\mathbf X;\mathbf Z)$ when densities are intractable. It introduces a unifying importance-sampling framework that yields both lower and upper bounds via extended-state-space constructions, and develops novel bounds including Multi-Sample AIS, GIWAE, and MINE-AIS. GIWAE generalizes variational and contrastive bounds (InfoNCE, IWAE, Barber-Agakov), while MINE-AIS expands MI estimation to settings with unknown $p(\mathbf X|\mathbf Z)$ by training an energy-based posterior and evaluating with Multi-Sample AIS. The methods are demonstrated on VAEs and GANs trained on MNIST and CIFAR, showing significant gains in tightness and practicality for large MI regimes. Collectively, the paper provides practical tools and theoretical insight for MI estimation in deep generative modeling and related applications.
Abstract
Mutual information (MI) is a fundamental quantity in information theory and machine learning. However, direct estimation of MI is intractable, even if the true joint probability density for the variables of interest is known, as it involves estimating a potentially high-dimensional log partition function. In this work, we present a unifying view of existing MI bounds from the perspective of importance sampling, and propose three novel bounds based on this approach. Since accurate estimation of MI without density information requires a sample size exponential in the true MI, we assume either a single marginal or the full joint density information is known. In settings where the full joint density is available, we propose Multi-Sample Annealed Importance Sampling (AIS) bounds on MI, which we demonstrate can tightly estimate large values of MI in our experiments. In settings where only a single marginal distribution is known, we propose Generalized IWAE (GIWAE) and MINE-AIS bounds. Our GIWAE bound unifies variational and contrastive bounds in a single framework that generalizes InfoNCE, IWAE, and Barber-Agakov bounds. Our MINE-AIS method improves upon existing energy-based methods such as MINE-DV and MINE-F by directly optimizing a tighter lower bound on MI. MINE-AIS uses MCMC sampling to estimate gradients for training and Multi-Sample AIS for evaluating the bound. Our methods are particularly suitable for evaluating MI in deep generative models, since explicit forms of the marginal or joint densities are often available. We evaluate our bounds on estimating the MI of VAEs and GANs trained on the MNIST and CIFAR datasets, and showcase significant gains over existing bounds in these challenging settings with high ground truth MI.