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.

  • 1
    $\begingroup$ 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? $\endgroup$
    – FrankD
    Jun 8 '18 at 14:38
  • $\begingroup$ I think problem might be that the PDF of M is not equal to 1/2(P+Q) $\endgroup$ Jun 18 '19 at 8:31

Actually, using the answer in https://stackoverflow.com/questions/26079881/kl-divergence-of-two-gmms (and the fact, that the author factored out the 1/2 from the logarithm, made the montecarlo approximation sample from both distributions to average the result), I would say, that the symmetrized numerical code for jensen shannon divergence using monte carlo integration, even for general scikit.stats distributions (_p and _q), should look like this:

def distributions_js(distribution_p, distribution_q, n_samples=10 ** 5):
    # jensen shannon divergence. (Jensen shannon distance is the square root of the divergence)
    # all the logarithms are defined as log2 (because of information entrophy)
    X = distribution_p.rvs(n_samples)
    p_X = distribution_p.pdf(X)
    q_X = distribution_q.pdf(X)
    log_mix_X = np.log2(p_X + q_X)

    Y = distribution_q.rvs(n_samples)
    p_Y = distribution_p.pdf(Y)
    q_Y = distribution_q.pdf(Y)
    log_mix_Y = np.log2(p_Y + q_Y)

    return (np.log2(p_X).mean() - (log_mix_X.mean() - np.log2(2))
            + np.log2(q_Y).mean() - (log_mix_Y.mean() - np.log2(2))) / 2

print("should be different:")
print(distributions_js(st.norm(loc=10000), st.norm(loc=0)))
print("should be same:")
print(distributions_js(st.norm(loc=0), st.norm(loc=0)))

For noncontinuous, change .pdf to probabilities of samples.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.