On Variational Bounds of Mutual Information (1905.06922v1)

Abstract: 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.

• The paper unifies existing MI estimators into a single variational framework to address challenges in high-dimensional data.
• It introduces a continuum of multi-sample lower bounds that flexibly trade off bias and variance in estimation.
• Empirical evaluations on the dSprites dataset demonstrate improved decoder-free disentangled representation learning.

On Variational Bounds of Mutual Information

Introduction

The paper "On Variational Bounds of Mutual Information" addresses the critical challenges associated with estimating and optimizing Mutual Information (MI) in high-dimensional settings. Mutual information, denoted $I(X; Y)$, quantifies the dependency between variables $X$ and $Y$, and is pivotal in diverse fields such as computational neuroscience, Bayesian experiment design, and representation learning. The authors explore variational bounds on MI parametrized by neural networks, highlighting the limitations of existing approaches that suffer from high bias or high variance, particularly when MI is large. They propose a novel continuum of lower bounds capable of flexibly trading off bias and variance, thereby unifying various existing bounds under a single framework.

Methodological Contributions

The authors present several significant contributions:

1. Review and Unification: They review existing MI estimators and establish relationships among them, providing a unified framework that incorporates these estimators.
2. New Continuum of Lower Bounds: They introduce a novel continuum of multi-sample lower bounds that generalize previous bounds and enable a flexible trade-off between bias and variance.
3. Leveraging Conditional Structures: The paper demonstrates leveraging known conditional structures to derive simple lower and upper bounds that effectively sandwich MI in representation learning contexts where $p_\theta(y|x)$ is tractable.
4. Systematic Empirical Evaluation: They empirically characterize the bias and variance of MI estimators and their gradients on controlled high-dimensional problems.
5. Applications in Representation Learning: The effectiveness of these new bounds is demonstrated in decoder-free disentangled representation learning on the dSprites dataset.

Detailed Examination of Variational Bounds

Normalized Upper and Lower Bounds

The authors begin by presenting the classic upper and lower bounds on MI proposed by Barber and Agakov. They discuss the variational upper bound, which is derived by introducing a variational approximation $q(y)$ to the intractable marginal $p(y)$:

$I(X; Y) \le E_{p(x)}[KL(p(y|x) \| q(y))]$

The corresponding lower bound is derived by replacing the intractable conditional distribution $p(x|y)$ with a variational distribution $q(x|y)$, yielding:

$I(X; Y) \ge E_{p(x,y)}[\log q(x|y)] + h(X)$

Unnormalized Lower Bounds

To avoid the intractable entropy term in the lower bound, the authors turn to unnormalized variational families for $q(x|y)$ parameterized by a critic $f(x,y)$, leading to bounds such as:

$E_{p(x,y)}[f(x,y)] - E_{p(y)}[\log Z(y)]$

This formulation allows the use of Jensen's inequality to derive various practical lower bounds, including the TUBA bound, which introduces variational parameters $a(y)$ for tighter bounds.

Multi-sample Bounds

To further reduce variance, the authors extend these unnormalized bounds to a multi-sample setting. By introducing additional samples, they develop the InfoNCE bound:

$$I(X; Y) \ge \mathbb{E}\left[\frac{1}{K}\sum_{i=1}^K \log \frac{e^{f(x_i, y_i)}}{\frac{1}{K} \sum_{j=1}^K e^{f(x_i, y_j)}}\right]$$

This bound is shown to be upper-bounded by $\log K$, thus effective only when $I(X; Y) \le \log K$.

Interpolated Bounds

By interpolating between the batch mixture and learned marginal, the authors present a new bound that can interpolate between the low-variance InfoNCE and high-variance unnormalized bounds, allowing a flexible trade-off that can be tuned based on the batch size and MI level.

Empirical Evaluations and Applications

Bias-Variance Trade-off Studies

The empirical evaluations delve into the bias and variance of MI estimators, using critics optimized for different bounds. Their findings align with theoretical expectations:

• Single-sample bounds (e.g., TUBA) show higher variance but less bias.
• Multi-sample bounds (e.g., InfoNCE) show lower variance but increased bias when MI is high.

Representation Learning

In disentangled representation learning on dSprites, the proposed bounds show utility beyond existing methods by avoiding adversarial or moment-matching techniques and leveraging MI estimates directly. The regularized InfoNCE objective used demonstrates the ability to recover meaningful features such as position and scale without a decoder.

Implications and Future Directions

The paper's findings significantly impact both theoretical and practical aspects of MI estimation. The continuum of bounds introduced not only provides a better understanding of the trade-offs involved but also equips practitioners with more robust tools for MI-based optimization tasks. Future work should explore more efficient estimators that maintain low bias and variance across different regimes and task types. Additionally, validating these bounds on larger-scale datasets and more complex representation learning tasks remains an open and promising avenue.

In conclusion, this paper makes substantial contributions to the field of MI estimation, demonstrating the efficacy and flexibility of variational bounds in tackling high-dimensional machine learning problems.