Multivariate normal with linear operator support - #245
Conversation
|
@csuter I think you mentioned you were working on this, right? |
|
Hi @Bonnevie, thank you so much for following up with this update to the previous PR! I have not done any actual work on it, so we're not overlapping efforts or anything (in case that was your concern in asking if I was working on it). I am traveling this week and may be time constrained but will take a good look at this ASAP. Anyone else from the team, is of course, welcome to comment. @brianwa84 and @langmore may wish to comment (as they were looking at the previous iteration). |
|
@csuter yeah, on rereading the old thread I was a bit worried that you'd had been putting time into it too. Hope it comes close to the design you were envisioning. There was a problem with sampling in the new class which I have just fixed - apparently the tests don't verify that sampling works? |
brianwa84
left a comment
brianwa84
left a comment
There was a problem hiding this comment.
can you walk through the upsides of this? (I'll go look over the PR again now, too)
| mvn = tfd.MultivariateNormalFullCovariance( | ||
| loc=mu, | ||
| covariance=covariance_matrix) | ||
| covariance=covariance) |
There was a problem hiding this comment.
the param can't be renamed
| tf.matrix_transpose(covariance_matrix), | ||
| message="Matrix was not symmetric") | ||
| ], covariance_matrix) | ||
| # LinearOperator applies cholesky which will fail if not PSD |
There was a problem hiding this comment.
Should probably say MVNLinearOperator, not LinearOperator
| # is called by the Bijector. | ||
| # However, cholesky() ignores the upper triangular part, so we do need | ||
| # to separately assert symmetric. | ||
| scale_tril = tf.cholesky(covariance_matrix) |
There was a problem hiding this comment.
I'm a bit concerned that in eager mode we will be recomputing this a lot.
| seed = seed_stream.SeedStream(seed, salt="multivariate normal") | ||
|
|
||
| loc = _broadcast_to_shape(self.loc, self._sample_shape()) | ||
| mvn = mvn_tril.MultivariateNormalTriL( |
There was a problem hiding this comment.
it seems like this won't be able to (efficiently) sample a low rank + diag because of the cholesky.
|
Hi @Bonnevie, I am having a problem with the original MultivariateNormalFullCovariance code (Cholesky decomposition was not successful) and I think it might be a bug. Do you think your code could solve this problem? Thanks! |
5 participants
TL;DR new version of
MultivariateNormalFullCovariancethat is backwards-compatible but also able to takeLinearOperator-type covariance matrices as input.I have previously touched on the need for a way to work with structured covariance matrices in both issue #161 and the closed PR #184 which demonstrated a prototype of this design. The concerns raised in the PR was that there were already too many variants of the
MultivariateNormal. As suggested in the PR, this new version instead revamps the standardMultivariateNormalFullCovariancewith a version that can take both tensors and linear operators as input when specifying the covariance. I have not touched the existing tests, so as it stands this new version should be completely compatible with existing code (assuming test coverage).Design decisions
I was unsure whether it would be worthwhile to inherit from any of the existing multivariate normal classes as
MultivariateNormalFullCovarianceused to do. I ultimately decided not to, basing it on the similarMultivariateStudentTinstead, which was already using aLinearOperator-based design. This was mostly to ensure that a cholesky-factor was not computed in cases where that would not be efficient or desirable.Incorporating
MultivariateNormalTrilThis could eventually subsume the need for
MultivariateNormalTril, but currentlyMultivariateNormalFullCovariancedepends on it for sampling which is most easily done with a Cholesky factor. In addition,MultivariateNormalTrilis slightly more efficient since it only has to solve the systeminv(L)*xonce and then take a norm, as opposed to a linear operator-based solve which will always solve twice even if it knows the cholesky decompositioninv(L*L')*x. SinceLinearOperatoris getting cholesky support, it might make sense to also have special methods forx'*A*xandx'*inv(A)*xas well as a PSD linear operator specified via a cholesky factor.