# Trying to implement the Jensen-Shannon Divergence for Multivariate Gaussians

Given two multivariate Gaussian distributions $P \equiv \mathcal{N}(\mu_p, \Sigma_p)$ and $Q \equiv \mathcal{N}(\mu_q, \Sigma_q)$, I am trying to calculate the Jensen-Shannon divergence between them.

I am following JSD-discussion for multivariate Gaussian in this discussion in this question. There it is suggested that one can approximate the midpoint measure $M$ using Monte Carlo sampling.

Specifically, it is pointed out that the JSD for continuous RVs (in my case Gaussian), is given by

$$\mathrm{JSD} = \frac{1}{2} (D_{KL}(P\,\|M)+ D_{KL}(Q\|M)) = h(M) - \frac{1}{2} (h(P) + h(Q)) \>,$$

where $h(P)$ and $h(Q)$ are just the differential entropies for the MVN. These properties are well known and we can calculate them easily, e.g.

$$h(P) = \frac{1}{2} \log_2\big((2\pi e)^n |\Sigma_p|\big)$$

What is causing me trouble is $M$. I believe I have misunderstood/not-implemented-correctly the Monte Carlo estimate for it.

User FrankD says that for the JSD approximation: $$JSD(P\|Q) = \frac{1}{2} (D_{KL}(P\|M)+ D_{KL}(Q\|M))$$ we can use Monte Carlo estimates for the individual components. The Kullback-Leibler divergence is defined as: $$D_{KL}(P|M) = \int P(x) log\big(\frac{P(x)}{M(x)}\big) dx$$ The Monte Carlo approximation of this is: $$D_{KL}^{approx}(P|M) = \frac{1}{n} \sum^n_i log\big(\frac{P(x_i)}{M(x_i)}\big)$$

where the $x_i$ have been sampled from $P(x)$, which is easy as it is a Gaussian in our case. As $n \to \infty, D_{KL}^{approx}(P|M) \to KLD(P|M)$. $M(x_i)$ can be calculated as

$$M(x_i) = \frac{1}{2}P(x_i) + \frac{1}{2}Q(x_i)$$.

Here is my attempt:

import numpy as np
from scipy.stats import multivariate_normal as MVN

def jsd(mu_1: np.array, sigma_1: np.ndarray, mu_2: np.array, sigma_2: np.ndarray):
"""
Monte carlo approximation to jensen shannon divergence for multivariate Gaussians.
"""
assert mu_1.shape == mu_2.shape, "Shape mismatch."
assert sigma_1.shape == sigma_2.shape, "Shape mismatch."

# Monte Carlo samples
MC_samples = 1000

# Take MC samples
P_samples = MVN.rvs(mean=mu_1, cov=sigma_1, size=MC_samples)
Q_samples = MVN.rvs(mean=mu_2, cov=sigma_2, size=MC_samples)

P = lambda x: MVN.pdf(x, mean=mu_1, cov=sigma_1)
Q = lambda x: MVN.pdf(x, mean=mu_2, cov=sigma_2)
M = lambda x: 0.5 * P(x) + 0.5 * Q(x)

P_div_M = lambda x: P(x) / M(x)
Q_div_M = lambda x: Q(x) / M(x)

D_KL_approx_PM = lambda x: (1 / MC_samples) * sum(np.log2(P_div_M(x)))
D_KL_approx_QM = lambda x: (1 / MC_samples) * sum(np.log2(Q_div_M(x)))

return 0.5 * D_KL_approx_PM(P_samples) + 0.5 * D_KL_approx_QM(Q_samples)


Suffice to say, this does not quite produce what it should.

• The code seems ok, although the definition of the KL-divergence uses natural log, not base 2 (but this should only mean that you're off by a factor of log(2)). You say it does not produce what it should -- how do you know? – FrankD Jun 8 '18 at 14:38
• I think problem might be that the PDF of M is not equal to 1/2(P+Q) – Shikun Huang Jun 18 at 8:31