The Matrix-Variate Normal (MVN) Distribution
Matrix-Variate Gaussian/Normal distributions (MVN distributions) have been first mentioned in [1] and are a generalization of Multivariate Gaussian Distributions. They formalize the idea of seeing an entire matrix as a random variable.
In contrast to describing each single entry as a Gaussian-distributed scalar that is independent from the other entries, the parametrisation of the MVN distribution allows for incorporating structural ties between the entries of a matrix.
A (real-valued) MVN Distribution for random variables is characterised by the matrix height , matrix width , mean matrix and covariance matrix components (determines row covariance) and (determines column covariance).
It is equivalent to the multivariate normal distribution
(1)
where vec() denotes the concatenation of columns of , and denotes the Kronecker product of and (the distribution’s covariance matrix):
(2)
The KL Divergence between MVN Distributions
The Kullback-Leibler (KL) divergence (also called “relative entropy”) is an asymmetric distance that, to put it bluntly, describes the “difference” between two probability distributions. When using the a probability distribution as a model where the actual distribution is (both over the domain ), the KL divergence quantifies the expected amount of excess “surprise” [2]. For continuous distributions it can be written as follows, referring to the distributions’ density functions and :
(3)
The probability density function for an MVN distribution on is given by
(4)
Where denotes the matrix trace and denotes the exponential function. Correspondingly, the KL divergence between two MVN distributions and on random variable square matrices can be derived as follows:
(5)
From [3], we obtain that the expectation terms can be rewritten as follows:
(6)
Where is an arbitrary matrix. Consequently, for the KL divergence, we obtain the following expectation-free term by further reshaping equation 5 :
(7)
This is really good, because we can now actually calculate the KL divergence between two arbitrary MVN distributions! But what is that good for?
Time-Series Modeling with Matrix-Shaped Latents
In [4], Hafner and his collaborators trained an agent in a Reinforcement Learning setting that consults a so-called Recurrent State Space Model (RSSM) to predict the outcome of possible actions and thus choose the optimal action to take. The RSSM combines recurrency (implemented with Gated Recurrent Units [5]) and stochastic modeling (by means of a Variational Autoencoder (VAE) [6]) to handle time-series of environment observations and predict possible future outcomes, and the proposed Planning Network (PlaNet) uses that RSSM as its model to achieve remarkable results on a variety of challenging tasks.
In my Master’s thesis, I had a closer look at the RSSM: The environment observations come in the form of images, which contain spatial information by nature. However, the RSSM’s VAE training procedure assumes that the latent variables follow a diagonal multivariate normal distribution, meaning that spatial information between the (encoded) pixels of the input observations could be lost in the process of latent variable encoding.
Therefore, part of my thesis dealt with the re-modeling of the latent space as a vector of independent random square matrices instead of random scalars. In order to train the resulting “spatial” VAE by means of Variational Inference [7], we need to calculate the KL divergence between the model’s posterior and prior distribution. And that’s where the formula for the KL divergence between two MVN distributions came in handy (albeit under a few factorization assumptions, since matrix inversion is expensive).
Code/Link to my thesis might appear here soon, but no guarantees…
References
- A. K. Gupta and D. K. Nagar: Matrix Variate Distributions. Chapman and Hall/CRC, 1999. ISBN 978-0-203-74928-9. doi: 10.1201/9780203749289
- https://en.wikipedia.org/wiki/Kullback-Leibler_divergence (accessed: Jan 08 2022) (yes, I actually cite Wikipedia here…)
- J. Li, H. Yan, J. Gao, D. Kong, L. Wang, S. Wang, and B. Yin. Matrix-variate variational auto-encoder with applications to image process. Journal of Visual Communication and Image Representation, 67:102750, Feb. 2020. ISSN 1047-3203. doi: 10.1016/j.jvcir.2019.102750.
- D. Hafner, T. P. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 2555–2565. PMLR, 2019.
- K. Cho, B. van Merrienboer, Ç. Gülçehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing, EMNLP 2014, October 25-29, 2014, Doha, Qatar
- D. P. Kingma and M. Welling. Auto-encoding variational bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada
- https://ermongroup.github.io/cs228-notes/inference/variational/ (accessed: Jan 08 2022)
To the best of my knowledge, I’ve cited all relevant non-trivial scientific contributions used in this post. If you spot any errors or false/missing citations please, let me know!