Let's say I want to compute the pairwise KL divergence between a large number (O(100)) of multivariate Gaussian distributions with diagonal covariance. The mean parameters for each Gaussian are stored within a matrix, where the rows correspond to the mean vectors, and the same applies for the diagonal of the covariance matrix of each Gaussian.
I could use a nested for loop to achieve this, but this seems slightly wasteful. Is there a more efficient way (preferably using numpy
or something Pythonic), to return a matrix $D$ where the $(i,j)$-th entry $D_{ij} = D_{KL}\left(\mathcal{N}(\mu_i, \Sigma_i) \Vert \mathcal{N}(\mu_j, \Sigma_j)\right)$ corresponds to the KLD between distributions $i$ and $j$?
I know that it is possible to compute this efficiently for the case where we assume homoscedasticity $\Sigma_i = \sigma^2 \mathbb{1} \; \forall \; i$, as it is possible to calculate the form of the KL via a pairwise distance matrix, as in this question, but not sure how to generalize this to the Mahalanobis distance - looking term in the KLD term.
For reference, here is the KLD between two diagonal-covariance Gaussians of dimension $D$:
$$ D_{KL}\left(\mathcal{N}(\mu_i, \Sigma_i) \Vert \mathcal{N}(\mu_j, \Sigma_j)\right) = \frac{1}{2} \left[ Tr \log \Sigma_j - Tr\log \Sigma_i + (\mu_i - \mu_j)^T \Sigma_j^{-1} (\mu_i-\mu_j) + Tr\left(\Sigma_j^{-1} \Sigma_i\right) - D\right] $$
Edit: Because of the expansion of the Mahalanbois-looking term into a sum of bilinears:
$$ \left(\mathbf{x}_i - \mathbf{x}_j\right)^T \Sigma_j^{-1}(\mathbf{x}_I - \mathbf{x}_j) = \mathbf{x}_i^T \Sigma_j^{-1} \mathbf{x}_i - \mathbf{x}_i^T \Sigma_j^{-1} \mathbf{x}_j - \mathbf{x}_j^T \Sigma_j^{-1} \mathbf{x}_i + \mathbf{x}_j^T \Sigma_j^{-1} \mathbf{x}_j $$
I suppose the problem can be reduced to how to efficiently compute the Gramian matrix $G_{ij}$ in the presence of a scaling matrix $\Sigma$. e.g. for the standard Gramian, where $X$ is a matrix with observations as columns,
$$X = \left( \mathbf{x}_1 \vert \mathbf{x}_2 \vert \ldots \vert \mathbf{x}_N \right)$$
$$ G = X^T X $$