Connecting Jensen-Shannon and Kullback-Leibler Divergences: A New Bound for Representation Learning

TL;DR

This paper establishes a principled link between Jensen–Shannon divergence and mutual information by deriving a tight, tractable lower bound Xi(D_JS) ≤ D_KL. It shows that maximizing the JSD-based dependence I_JS provably increases a lower bound on MI and connects discriminative training losses to MI via a cross-entropy bound. The authors develop a differentiable approximation for Xi and demonstrate, through extensive discrete and continuous experiments, that the bound is tight and yields low-variance MI estimates, outperforming several variational bounds in practice. They further apply the bound within the Information Bottleneck framework, achieving state-of-the-art performance on MNIST tasks, thereby providing both theoretical justification and practical tools for discriminative representation learning. Overall, the work justifies and enables end-to-end MI optimization using discriminative JSD objectives with strong theoretical guarantees and practical impact.

Abstract

Mutual Information (MI) is a fundamental measure of statistical dependence widely used in representation learning. While direct optimization of MI via its definition as a Kullback-Leibler divergence (KLD) is often intractable, many recent methods have instead maximized alternative dependence measures, most notably, the Jensen-Shannon divergence (JSD) between joint and product of marginal distributions via discriminative losses. However, the connection between these surrogate objectives and MI remains poorly understood. In this work, we bridge this gap by deriving a new, tight, and tractable lower bound on KLD as a function of JSD in the general case. By specializing this bound to joint and marginal distributions, we demonstrate that maximizing the JSD-based information increases a guaranteed lower bound on mutual information. Furthermore, we revisit the practical implementation of JSD-based objectives and observe that minimizing the cross-entropy loss of a binary classifier trained to distinguish joint from marginal pairs recovers a known variational lower bound on the JSD. Extensive experiments demonstrate that our lower bound is tight when applied to MI estimation. We compared our lower bound to state-of-the-art neural estimators of variational lower bound across a range of established reference scenarios. Our lower bound estimator consistently provides a stable, low-variance estimate of a tight lower bound on MI. We also demonstrate its practical usefulness in the context of the Information Bottleneck framework. Taken together, our results provide new theoretical justifications and strong empirical evidence for using discriminative learning in MI-based representation learning.

Figures (6)

• Figure 1: Left: Parameter space of two Bernoulli distributions. Right: Image (in gray) of the mapping $\phi$, showing the joint range of the Jensen–Shannon and the Kullback–Leibler divergences. The red curve is our new lower bound.
• Figure 2: Mutual information and its JSD-based lower bound $\Xi({\rm I_{\rm JS}})$ for a parameterized family of discrete joint distributions with known MI and JSD, varying in dependence strength ($\alpha$) and number of categories ($k$). The presence of MIs near the lower bound across settings empirically demonstrates the tightness of our JSD-based estimate.
• Figure 3: Staircase MI estimation comparison for $d = 5$ and batch size $b=64$. The Gaussian case is reported in the top row, while the cubic case is shown in the bottom row.
• Figure 5: Visualization of the negative of the determinant of the Jacobian. We can observe that it is constantly positive, meaning that the determinant of the Jacobian is always negative.
• Figure 6: Comparison between the numerical evaluation of $\Xi$ using Algorithm \ref{['lst:xi_function']} and the proposed differentiable approximation from Eq. \ref{['eq:approximation_appendix']}.

Theorems & Definitions (18)

• Theorem 3.1: Harremoës-Vajda, 2011
• Theorem 4.1
• Lemma B.1
• proof
• proof
• Lemma B.3
• proof
• Lemma B.4
• proof
• Conjecture B.5