Calculate KL Divergence Between Two Multivariate Gaussian Distributions

Problem

Compute the Kullback-Leibler divergence between two multivariate Gaussian distributions. Given the mean vectors and covariance matrices of two distributions P and Q, calculate KL(P || Q).

Solution

import numpy as np

def kl_divergence_gaussian(mu_p: np.ndarray, cov_p: np.ndarray, mu_q: np.ndarray, cov_q: np.ndarray) -> float:
    k = mu_p.shape[0]

    cov_q_inv = np.linalg.inv(cov_q)

    # KL(P || Q) = 0.5 * (tr(Sigma_q^{-1} Sigma_p) + (mu_q - mu_p)^T Sigma_q^{-1} (mu_q - mu_p) - k + ln(det(Sigma_q) / det(Sigma_p)))
    diff = mu_q - mu_p

    trace_term = np.trace(cov_q_inv @ cov_p)
    mahalanobis_term = diff.T @ cov_q_inv @ diff
    log_det_term = np.log(np.linalg.det(cov_q) / np.linalg.det(cov_p))

    kl = 0.5 * (trace_term + mahalanobis_term - k + log_det_term)
    return float(kl)

Explanation

1. The KL divergence between two multivariate Gaussians has a closed-form solution.
2. Trace term: tr(Sigma_q^{-1} @ Sigma_p) measures how well Q's covariance covers P's.
3. Mahalanobis term: Measures the squared distance between means weighted by Q's precision matrix.
4. Log determinant term: ln(det(Sigma_q) / det(Sigma_p)) accounts for the volume difference.
5. The -k term is a correction for the dimensionality.

Complexity

• Time: O(k^3) for matrix inversion and determinant computation
• Space: O(k^2) for the inverse covariance matrix