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

AI summary

Overview

Research area: Machine learning theory and representation learning, specifically the theory of mutual information (MI) estimation and the relationship between divergences used as training objectives.

Technical level: Advanced. The paper relies on f-divergences, the joint range of divergences, variational representations, and the Harremoës–Vajda theorem. The empirical sections and the high-level argument are accessible to an intermediate reader.

Scope (one sentence): The paper derives a new, optimal lower bound on the Kullback–Leibler divergence (KLD) in terms of the Jensen–Shannon divergence (JSD), uses it to prove that maximizing a JSD-based dependence measure increases a guaranteed lower bound on mutual information, and validates the bound empirically as an MI estimator and inside the Information Bottleneck framework.

What This Paper Is About

Many representation learning methods maximize the Jensen–Shannon divergence between the joint distribution and the product of marginals, often by training a binary classifier (discriminator) to tell true pairs from scrambled pairs. This is done without a clear theoretical guarantee that it increases the standard mutual information, which is defined as a Kullback–Leibler divergence. The paper closes that gap by proving an explicit, tight, monotone relationship between JSD and KLD, and by showing that minimizing a discriminator's cross-entropy loss raises a provable lower bound on MI.

Key Contributions

1. A new optimal lower bound on KLD in terms of JSD. For any two distributions p and q, the authors prove Ξ(D_JS[p||q]) ≤ D_KL[p||q], where Ξ is a strictly increasing function defined implicitly by its inverse, Ξ⁻¹: y ↦ D_JS[B(1)||B(exp(−y))], built from Bernoulli distributions. The bound is described as optimal, meaning equality is achieved for a family of distributions.

2. A provable link between JSD-based dependence and mutual information. Specializing the bound to joint and product-of-marginals distributions shows that maximizing I_JS[U;V] = D_JS[p_UV || p_U ⊗ p_V] increases a guaranteed lower bound on the true MI.

3. A cross-entropy interpretation of the JSD objective. The authors show that the f-divergence variational lower bound for JSD is equivalent to log 2 − L_CE ≤ D_JS[p||q] (Eq. 4), where L_CE is the expected binary cross-entropy loss of a classifier trained to separate joint from marginal pairs — recovering a connection originally noted in the GAN paper and making the suboptimal-discriminator gap explicit.

4. A combined variational lower bound and estimator. Chaining the above gives I_CE(θ) ≐ Ξ(log 2 − L_CE(θ)) ≤ Ξ(I_JS[U;V]) ≤ I[U;V] (Eq. 27). The paper also derives a two-step MI estimator, I[U;V] = E_{p_UV}[𝕃(p̃(z=1|u,v))] (Eq. 28), where 𝕃 is the Logit transform, which the authors state is equivalent to GAN-DIME.

Main Findings

• The bound is tight in the discrete setting: Across a parameterized family of joint distributions P_UV^(α) = (1−α)P_U ⊗ P_V + α diag(P_U) with uniform categorical marginals, varying α ∈ [0,1] and number of categories k ∈ {2, 3, …, 500}, the authors observe that for any given JSD value there exist discrete distributions whose true MI nearly coincides with the bound Ξ(x). An appendix result shows an infinite family of discrete joint–marginal pairs lies exactly on the bound.

• A usable closed-form approximation exists: Although Ξ has no closed-form expression, the paper gives the smooth approximation Ξ(x) ≈ 1.15 * 𝕃(.5(x/log 2 + 1.0)). Because Ξ is strictly increasing, maximizing Ξ(log 2 − L_CE) is equivalent to minimizing L_CE, so the approximation is not needed during optimization.

• Competitive MI estimation on Gaussian and cubic benchmarks: Using U ~ N(0, I_d), N ~ N(0, I_d), V = ρU + sqrt(1−ρ²)N (with MI = −(d/2)log(1−ρ²)), and a cubic variant v ↦ v³ that preserves MI while making marginals non-Gaussian, the proposed JSD-LB gives a tight, low-variance lower bound. MINE and NWJ show high variance, and MINE can overestimate its own lower bound and sometimes exceed the true MI; CPC remains limited by the contrastive bound log(b). The derived two-step estimator is reported as the most accurate MI estimate on these benchmarks.

• Results extend to harder continuous settings: The authors repeat the more challenging benchmark of Letizia et al. with complex Gaussian transformations (half-cube, asinh at d = 5) and non-Gaussian distributions (Student at d = 5, uniform), reporting that their bound and derived MI estimates remain competitive with lower variance.

• State-of-the-art Information Bottleneck results on MNIST: Replacing the MI estimator with JSD-LB in the IB framework gives the best reported values among the compared methods.

  • Generalization accuracy: JSD-LB 98.8%, versus VIB 97.6, NIB 97.2, squared-VIB 96.2, squared-NIB 93.3, DisenIB 98.2.
  • Adversarial robustness under a one-step gradient attack, training-time perturbations ε ∈ {0.1, 0.2, 0.3}: JSD-LB 99.5 / 96.0 / 91.4, versus DisenIB 94.3 / 81.5 / 68.4. Test-time ε = 0.1 / 0.2 / 0.3: JSD-LB 94.6 / 89.6 / 86.1, versus DisenIB 90.2 / 80.0 / 67.8.
  • Out-of-distribution detection on synthetic Gaussian noise: JSD-LB 100.0 AUROC, 99.9 AUPR In, 99.9 AUPR Out, 0.1 detection error, and TPR95 of 0.00 (lower is better), versus DisenIB 99.4 AUROC, 99.6 AUPR In, 98.9 AUPR Out, 1.7 detection error.

• The theory unifies existing practice: Minimizing the cross-entropy of a joint-vs-marginal discriminator raises a guaranteed lower bound on MI, providing a theoretical justification for JSD-based discriminative representation learning methods. The authors note this bound is tight in the non-parametric limit of infinite data and model capacity.

• Reported limitations of prior bounds: The paper states that McAllester and Stratos showed VLB estimators of the MINE/NWJ type are fundamentally upper-bounded by log(n), where n is the number of data samples, and that InfoNCE is upper-bounded by log(b), where b is the batch size.

Methodology in Plain English

The authors start from a piece of classical information theory: for two convex functions f and g, the set of all achievable pairs of f-divergences (the "joint range") is convex, and there is an optimal lower bound relating the two. By the Harremoës–Vajda result, that whole set can be characterized using only Bernoulli (binary) distributions, which drastically simplifies the problem.

They then look at the specific pair (JSD, KLD). Mapping pairs of Bernoulli distributions through both divergences and examining the triangle of parameters, they show that the relevant extremal edge is the case where one Bernoulli parameter equals 1. As the other parameter goes to 0, the KLD goes to +∞ while the JSD rises to at most log 2. Inverting that relationship gives the strictly increasing function Ξ, whose inverse has a closed form related to the JSD expression in Eq. 17. Since Ξ is defined by its inverse, they supply a Logit-based smooth approximation for practical use.

For the discriminator part, they set up a mixture where a binary label Z selects either a true joint pair (Z = 1) or an independent pair (Z = 0), each with probability 1/2. Plugging a reparameterized variational function t(u,v) = log(2 q_θ(z=1|u,v)) into the standard variational representation of JSD yields an expression whose maximization is exactly minimization of binary cross-entropy. Chaining the two results gives the MI lower bound.

Empirically, they check the bound in three ways: (1) a discrete setup where true MI and JSD are computed exactly, so tightness can be judged without neural approximation error; (2) neural estimation on benchmarks with analytically known MI, using a fully connected discriminator with input dimension 2d, two hidden layers of 256 ReLU units, one scalar output, Adam for 4000 steps, and batch size N = 64, matching prior work's architecture and hyperparameters; and (3) an Information Bottleneck experiment on MNIST replicating the setup of Pan et al., evaluating generalization, adversarial robustness, and OOD detection.

Why This Matters

The paper supplies the missing theoretical justification for a widespread practical recipe: train a classifier to distinguish joint from independent samples, and use that as your representation learning objective. It also offers a theoretically grounded alternative to variational lower bounds like MINE and NWJ, which the paper characterizes as unstable and high-variance, and to contrastive methods such as InfoNCE, which are capped at log(b) by batch size. The code repository is given as https://github.com/ReubenDo/JSDlowerbound.

Real-world applications (framed by the settings the paper evaluates or cites):

• Representation learning for classification. The MNIST Information Bottleneck experiments show improved generalization accuracy when the MI term is optimized via the JSD-based lower bound.
• Adversarially robust models. The reported large gains under one-step gradient perturbations (for example 91.4% versus 68.4% at training ε = 0.3) point to security-sensitive deployments.
• Out-of-distribution and anomaly detection. The near-perfect AUROC/AUPR and low detection error on synthetic noise suggest use in monitoring and safety pipelines.
• Domains where contrastive methods dominate. The related work notes strong contrastive predictive coding results in audio, vision, and NLP; a bound that does not require large batches is relevant where batch size is computationally constrained.

Industry relevance: any pipeline that maximizes mutual information between representations — self-supervised pretraining, contrastive learning, information bottleneck regularizers — can adopt this objective with a single discriminator and end-to-end training, and the paper notes this contrasts with two-step estimators that must be retrained whenever the joint distribution changes, which is restrictive when latent variables evolve during optimization.

Future Directions

• Scaling to large, high-dimensional data. The experiments cover discrete distributions with up to k = 500 categories and neural estimation with input dimension 2d, two 256-unit hidden layers, and batch size 64; behavior at modern scale is not reported.
• Non-variational and end-to-end estimation during training. The authors note that two-step estimators cannot be reused as plug-in MI estimators during representation learning, implying a need for further work on stable single-discriminator estimation as latent variables change.
• Extending the joint-range analysis to other divergence pairs. The framework is general over pairs of f-divergences; the paper only instantiates it for (JSD, KLD).
• Characterizing the suboptimal-discriminator gap. The paper states that its derivation "makes explicit the gap due to suboptimal discriminators" and includes an appendix proof quantifying the gap between the JSD lower bound and the cross-entropy one; how this gap behaves in practice under finite model capacity and finite data is left as an open empirical question.

Target Audience

Machine learning researchers working on information-theoretic representation learning, self-supervised and contrastive learning, and mutual information estimation; theorists interested in f-divergence inequalities and the joint range of divergences; and practitioners in generative or adversarial modeling who want a justification and a concrete implementation for discriminator-based objectives. Readers should be comfortable with probability distributions, KL and Jensen–Shannon divergences, and basic variational bound arguments.

Authors’ 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.