Cache covariance matrix decomposition in frozen multivariate_normal

[Balandat]

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:

scipy/scipy/stats/_multivariate.py

Line 762 in adc4f4f

return self._dist.rvs(self.mean, self.cov, size, random_state)

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

[siddhantwahal]

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.

[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 ***@***.***> 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.

[mdhaber]

I think this has been addressed by the Covariance class. Do you agree @tirthasheshpatel @tupui? Shall we close this one?

[tirthasheshpatel]

@mdhaber Yup, Covariance should address the issue. Feel free to close this!

@Balandat FYI: Now, we can specify the covariance using the factory methods provided by the stats.Covariance class! Here's an example:

import numpy as np
from scipy import stats

cov = np.eye(3)
mu = np.zeros(3)

# only compute the eigen-decomposition once
w, v = np.linalg.eigh(cov)
cov_object = stats.Covariance.from_eigendecomposition((w, v))

# use the cov_object instead
dist = stats.multivariate_normal(mu, cov_object)
dist.pdf([0, 0, 0])

# also, notice that the covariance object is diagonal
# we can use the special property to make computation
# of pdf, ... more efficient
cov_object = stats.Covariance.from_diagonal(np.diag(cov))
stats.multivariate_normal.pdf([0, 0, 0], mu, cov_object)  # uses the special property of the
                                                          # covariance matrix to compute pdf
                                                          # more efficiently.

Feel free to refer to the devdocs here for more examples.

[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