Cache covariance matrix decomposition in frozen multivariate_normal

Currently, it seems that a frozen multivariate_normal distribution unnecessarily re-computes the root decomposition (and other properties such as evals for the logpdf) of the covariance matrix for each operation. For instance, when sampling scipy just calls the sampling of the underlying distribution with the full covariance matrix: https://github.com/scipy/scipy/blob/adc4f4f7bab120ccfab9383aba272954a0a12fb0/scipy/stats/_multivariate.py#L762

This is super wasteful, as most of the computation is in fact in computing this decomposition.

What should happen instead is that the frozen object should also have a factor or L attribute s.t. L @ L.T = cov, and then it would compute that only once (upon the first iteratio), and in later steps pass that into the sampling method instead of cov, avoiding a bunch of unnecessary compute.

Torch’s mvn does this (though it doesn’t allow singular matrices at this point): https://github.com/pytorch/pytorch/blob/master/torch/distributions/multivariate_normal.py#L146-L15

[rlucas7]

Might be nice to have A benchmark given that (AFAICS) this is a (very useful) speed up. Might be a good place to start.

Here is a good read if you’re not familiar with benchmarks in scipy

http://scipy.github.io/devdocs/dev/contributor/benchmarking.html

Sincerely,

-Lucas Roberts

On May 20, 2020, at 3:27 PM, siddhantwahal [e-mail] wrote:

The rvs method indeed requires recomputing the square root of the covariance.

One approach to fix this could be adding an extra attribute to _PSD, say L as: self.L = np.multiply(u, np.sqrt(s)), where s, u is the eigendecomposition of the covariance matrix.

Then, the square root of the covariance is available in multivariate_normal_frozen as self.cov_info.L, and can be employed in rvs.

I think only rvs needs modifying. As far as I can tell, logpdf uses self.U, which is a square root of the precision matrix.

Happy to work on this if no one else is.

— You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub, or unsubscribe.

[Balandat]

Thanks! Let me use the opportunity to make a shameless plug for a related (PyTorch) project for structured linear algebra that can also encode various structured covariance matrices: https://github.com/cornellius-gp/linear_operator