Efficiently computing pairwise KL divergence between multiple diagonal-covariance Gaussian distributions

Question

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 $$

Answer 1

I figured it out, if this is useful to anyone stumbling across this in the future. For the case of a diagonal covariance Gaussian, note that the Mahalanbois-looking term simplifies to:

$$ \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 - 2\mathbf{x}_i^T \Sigma_j^{-1} \mathbf{x}_j + \mathbf{x}_j^T \Sigma_j^{-1} \mathbf{x}_j $$

It's straightforward to calculate the last two terms on the right hand side of the above equation, using the same logic as calculating the Gramian in this question, calculating the first term is simpler than I thought, and is clear from using the form:

$$ S_{ij} = \mathbf{x}_i^T \Sigma_j^{-1} \mathbf{x}_i = \sum_{k=1}^D \left(\sigma_j^{(k)}\right)^{-2} \left(x_i^{(k)}\right)^2 $$

To construct the matrix $S_{ij}$ in a vectorized manner, if we have a matrix holding observations as rows and a matrix where each row holds the diagonal of the covariance matrix, then we can do the following (in torch, but should be simple to generalize):

B, D = 128, 8
x, inv_var_diag = torch.randn([B,D]), torch.randn([B,D])
S_ij = x**2 @ inv_var_diag.T 

Trying to compute this for a (2048,2048) matrix results in a runtime of over 10 minutes when naively iterating over each element, compared to 300 ms when computing it in vectorized form!