Represent Gaussian using square root of precision

[fehiepsi]

This is an effort to resolve the Cholesky issue as reported in #2017.

There is some work during the transition to this representation but I hope that it will be smooth.

TODO:

• Convert marginalize method to use prec_sqrt
• Convert matrix_and_mvn_to_gaussian, mvn_to_gaussian to this new representation. It is best to have a MVN distribution parameterized by prec_sqrt instead of using torch.inverse(scale_tril), but we can defer this to future improvements.
• Add tests for this new representation.
• Implement unbroadcasted gaussian_tensordot (instead of using pad + add + marginalize). Edit: this is not necessary, pad + add + marginalize would have the same performance.
• [ ] Update to match recent changes in Gaussian.
• Add logic to deal with shape blow up which is pointed out by @fritzo
• [ ] Add tests to see if this works for high dimensional latent hmm.
• Merge with Gaussian class, dispatching as necessary

[fehiepsi]

@fritzo The square root representation works smoothly and I believe that it will be more stable than the current representation. To make it easier to compare and test, I separated the implementation into a new file. In addition, I added the S notation to class and functions in this file to distinguish this implementation and the current implementation. Do you have other suggestion to name this new class and the corresponding functions.

Here are several properties of this representation:

• event_permute will permute the row vectors of prec_sqrt matrix
• add two gaussians will concatenate the column vectors of their prec_sqrt matrices
• event_pad will add more zero rows to prec_sqrt matrix

I believe this version will be a bit faster (faster at mvn_to_gaussianS function) and more stable. It would be faster and more stable if use pyro.param('prec_sqrt', constraint=lower_cholesky) and use prec_sqrt instead of scale_tril in MVN implementation.

[fehiepsi]

Btw, just realize that this is PR #2019 ! :D

[fritzo]

Hi @fehiepsi, I took a first pass through the code and it looks plausible. I have one major concern.

IIUC the sqrt form replaces cholesky solves with concatenation, which is numerically stable but which is not bounded in space. For example, consider the basic operation of the GaussianHMM over T time steps: we recursively call Gaussian.__add__() and Gaussian.marginalize(). Each op is called O(log2(T)) times. Consider the following identities

# Rank identities for GaussianS:
(g1 + g2).rank() == g1.rank() + g2.rank()  # rank is additive.
g.marginalize(...).rank() == g.rank()      # .marginalize() preserves rank.

Now in the released Gaussian implementation of the GaussianHMM with n-dimensional state, rank always n so total memory space is O(T * n^2). However in the proposed GaussianS implementation of GaussianHMM rank at recursive stage s is n * 2^s, leading to total memory space O(T * log2(T) * n^2), i.e. the proposed square root form would cost an extra factor of O(log2(T)) more space. Does this analysis appear correct to you?

[fritzo]

Why not simply

return self.prec_sqrt.new_zeros(self.batch_shape + (self.dim(), self.dim()))

[fehiepsi]

Here we need to return empty tensor, not zeros tensor. I think it is used to pass tests for constant Gaussian.

[fritzo]

.precision should be the zeros tensor, not the empty tensor.

[fehiepsi]

You are right. I thought we are multiplying two empty tensor, which will return an empty tensor. :D

[fehiepsi]

@fritzo The operator A.matmul(A.T) will remove all the expanded rank at each step. Assume that A has shape n x 4n (which is the case when using tensordot), then A.matmul(A.T) will have shape n x n. I will add a test to guarantee that the square root's shape does not expanded in tensordot op.

[fehiepsi]

Similarly, QR of A.T where A.shape = n x 4n has shape n x n.

[fritzo]

But you're not currently triangularizing Psqrt_a in .marginalize(). I think to achieve bounded rank, you'll need to manually triangularize at the end of .marginalize():

if prec_sqrt.size(-1) > prec_sqrt.size(-2):
    prec_sqrt = triangularize(prec_sqrt)

[fehiepsi]

You are right! I missed this point. Let me see if there is an easy way to achieve this. :(

[fehiepsi]

@fritzo Applying at the end of .marginalize() is correct but it seems that we have to use qr (which is more expensive than Cholesky) there because there is no requirement that the returned Gaussian has full rank. I want to consider other possibilities so I'll revisit this issue after numpyro 0.2 and GaussianGamma PR. Thanks for pointing out that rank issue!

[fehiepsi]

@fritzo It seems that to marginalize out x in an (x,y) Gaussian pair, we can just simply triangularize the prec_sqrt part and keep the bottom right block. This way, no need for using cholesky during marginalization and the rank is also bounded. More precisely, given a square root matrix,

[[A 0]
 [B C]]

the marginalized precision is (B@Bt+C@Ct) - B@At@(A@At)^-1@A@Bt, which is the same as C@Ct if we assume that A is invertible (which is required to be true if we want to marginalize x out).

This is the main trick used in square root information filters, there only QR is needed.

I can complete this work if this is needed by someone. Maybe we can also incorporate this fact into funsor.

[fehiepsi]

TODO: switch to whiten_vec as in funsor. Probably renaming the attributes to sri_matrix, sri_factor (to match SRIF literature) and log_factor