[Paper Review] The Representation Jensen-Shannon Divergence

Jhoan K. Hoyos-Osorio, Luis G. Sanchez-Giraldo|arXiv (Cornell University)|May 25, 2023

Statistical Mechanics and Entropy4 citations

TL;DR

This paper proposes the Representation Jensen-Shannon Divergence (RJSD), a novel kernel-based divergence that avoids density estimation by embedding data into a reproducing kernel Hilbert space (RKHS) using uncentered covariance operators. The method enables scalable, differentiable, and minibatch-compatible estimation via Fourier features or kernel matrices, achieving state-of-the-art performance in two-sample testing and mitigating mode collapse in generative modeling with theoretical guarantees as a lower bound on the classical JSD.

ABSTRACT

Quantifying the difference between probability distributions is crucial in machine learning. However, estimating statistical divergences from empirical samples is challenging due to unknown underlying distributions. This work proposes the representation Jensen-Shannon divergence (RJSD), a novel measure inspired by the traditional Jensen-Shannon divergence. Our approach embeds data into a reproducing kernel Hilbert space (RKHS), representing distributions through uncentered covariance operators. We then compute the Jensen-Shannon divergence between these operators, thereby establishing a proper divergence measure between probability distributions in the input space. We provide estimators based on kernel matrices and empirical covariance matrices using Fourier features. Theoretical analysis reveals that RJSD is a lower bound on the Jensen-Shannon divergence, enabling variational estimation. Additionally, we show that RJSD is a higher-order extension of the maximum mean discrepancy (MMD), providing a more sensitive measure of distributional differences. Our experimental results demonstrate RJSD's superiority in two-sample testing, distribution shift detection, and unsupervised domain adaptation, outperforming state-of-the-art techniques. RJSD's versatility and effectiveness make it a promising tool for machine learning research and applications.

Motivation & Objective

• To address the challenge of estimating statistical divergences from empirical samples without density estimation.
• To develop a divergence that leverages second-order moments in RKHS for improved robustness and theoretical grounding.
• To provide a scalable, differentiable, and minibatch-compatible estimator suitable for deep learning optimization.
• To establish RJSD as a lower bound on the classical Jensen-Shannon divergence, enabling variational estimation with statistical guarantees.
• To demonstrate superior performance in two-sample testing and mode collapse mitigation in generative modeling.

Proposed method

• Proposes the Representation Jensen-Shannon Divergence (RJSD) by embedding probability distributions into an RKHS using uncentered covariance operators.
• Employs explicit Fourier feature mapping to construct an empirical estimator from sample covariance matrices, enabling differentiability and scalability.
• Provides an alternative kernel matrix-based estimator without explicit feature mapping, preserving theoretical consistency.
• Derives consistency and sample complexity bounds for the proposed estimators under mild regularity conditions.
• Leverages the connection between kernel-based entropy and Shannon entropy to relate RJSD to classical information-theoretic quantities.
• Optimizes RJSD using Adam with joint learning of Fourier features, kernel bandwidth, and network parameters in end-to-end training.

Figure 1: Jensen-Shannon Divergence estimation for two set of samples following Cauchy distributions (N = 512). We compare the following estimators: NWJ (Nguyen et al., 2010 ) , infoNCE (Oord et al., 2018 ) , CLUB (Cheng et al., 2020 ) , MINE (Belghazi et al., 2018 ) , KNIFE (Pichler et al., 2022 )

Experimental results

Research questions

• RQ1 Can a divergence be constructed in RKHS using second-order moments that avoids direct density estimation?
• RQ2 Is the proposed RJSD estimator consistent and sample-efficient under empirical sampling?
• RQ3 Can RJSD serve as a lower bound on the classical Jensen-Shannon divergence, enabling variational estimation with guarantees?
• RQ4 Does RJSD outperform state-of-the-art methods in two-sample testing across diverse data distributions?
• RQ5 Can RJSD effectively prevent mode collapse and improve sample diversity in generative adversarial networks?

Key findings

• RJSD achieves state-of-the-art performance in two-sample testing across multiple datasets, including MNIST, Higgs, and HDGM, with higher statistical power than MMD, C2ST, and other baselines.
• The RJSD-based GAN training effectively mitigates mode collapse and generates more diverse, high-quality samples compared to baseline GANs.
• The RJSD estimator is a lower bound on the classical Jensen-Shannon divergence, enabling a variational estimator with theoretical consistency guarantees.
• The Fourier feature-based estimator demonstrates strong scalability and compatibility with minibatch optimization, enabling efficient training in deep generative models.
• Empirical results show that using D << N Fourier features often yields better performance than high-dimensional feature spaces, contrary to typical kernel theory expectations.
• The method remains robust under varying data scales, though rank-inconsistency in kernel matrices necessitates careful hyperparameter tuning or use of explicit feature maps.

Figure 3: Generated samples using rep JSD.