Positive definite issue in MvNormal during Cholesky decomposition

[junpenglao]

Similar issue has been report also in #2015.
While I love the new cholesky decomposed covariance matrix in MvNormal, setting it as a default is numerical quite unstable:

        if cov is not None:
            self.chol_cov = tt.slinalg.cholesky(tt.as_tensor_variable(cov))
        elif tau is not None:
            self.chol_tau = tt.slinalg.cholesky(tt.as_tensor_variable(tau))
        else:
            self.chol_cov = tt.as_tensor_variable(chol)

I have a few simulation parameterized on cov used to work fine, they now returns error during the calling of tt.slinalg.cholesky quite often.

Maybe we should use QR decomposition and reconstruct the cholesky factor? It should be more stable.

[junpenglao]

CC @jonathanhfriedman and @aseyboldt

[aseyboldt]

This is more complicated than I hoped.
I think the best solution would be to switch to a pivoted cholesky decomposition. There is one in lapack (dpstrf), but this is not exposed in theano or even scipy. Stan uses a custom implementation of a pivoted L^TDL decomposition. This would be quite a bit of work however. Using a qr decomposition (with pivoting?) would be much easier. Pivoted qr is in scipy.
Why do you want to reconstruct the cholesky factor from the qr decomposition? Couldn't we just us the qr directly? We need to solve a linear system and compute the det. But since det(Q) \in {-1, 1}, both should be easy. Getting the sign of the det would be nice (in order to return a logp of -inf if it is not positive). Unfortunately the scipy interface doesn't report that in any way, although it should be known. The sign depends on the number of householder reflections.
There is going to be some trouble when computing the gradient, too. The qr decomposition op does not define its gradient.
Do you know a good reference for the numerical properties of the various versions of (pivoted) cholesky/qr/svd?

[junpenglao]

Did some digging and it is indeed a pain: most of the operation in theano.nlinalg dont have grad.
Actually, I remembered wrong, QR and Cholesky is related directly only for solving A = X'X. But Cholesky can still be computed from SVD output. Unfortunately, there is no SVD in theano which gives the grad.
I cooked up some ugly code for myself right now to do cholesky to avoid positive definite error (mostly borrowing code from theano):

import numpy
try:
    import numpy.linalg
    imported_numpy = True
except ImportError:
    # some ops (e.g. Cholesky, Solve, A_Xinv_b) won't work
    imported_numpy = False
import theano
from theano import tensor
from theano.gof import Op, Apply
from theano.tensor import as_tensor_variable
from theano.tensor.slinalg import Solve

solve_upper_triangular = Solve(A_structure='upper_triangular', lower=False)
class QR_Chol(Op):
    """
    (mostly) Copy from theano:
    
    Incomplete QR Decomposition.
    Computes the QR decomposition of a matrix.
    Factor the matrix a as qr and return a single matrix.
    """
    _numop0 = staticmethod(numpy.linalg.svd)
    _numop = staticmethod(numpy.linalg.qr)
    __props__ = ('mode',)

    def __init__(self, mode):
        self.mode = mode

    def make_node(self, x):
        x = as_tensor_variable(x)
        assert x.ndim == 2, "The input should be a matrix."
        z = theano.tensor.matrix(dtype=x.dtype)
        return Apply(self, [x], [z])

    def perform(self, node, inputs, outputs):
        (x,) = inputs
        (z,) = outputs
        assert x.ndim == 2, "The input should be a matrix."
        w, u, v = self._numop0(x,'full_matrices', 'compute_uv')
        tmp = np.matmul(np.diag(np.sqrt(u)), w.T).astype(x.dtype)
        r = self._numop(tmp, self.mode)
        z[0] = r.T
        
    def grad(self, inputs, gradients):
        """
        Cholesky decomposition reverse-mode gradient update.
        Symbolic expression for reverse-mode Cholesky gradient taken from [0]_
        References
        ----------
        .. [0] I. Murray, "Differentiation of the Cholesky decomposition",
           http://arxiv.org/abs/1602.07527
        """

        x = inputs[0]
        dz = gradients[0]
        chol_x = self(x)

        def tril_and_halve_diagonal(mtx):
            """Extracts lower triangle of square matrix and halves diagonal."""
            return tensor.tril(mtx) - tensor.diag(tensor.diagonal(mtx) / 2.)

        def conjugate_solve_triangular(outer, inner):
            """Computes L^{-T} P L^{-1} for lower-triangular L."""
            return solve_upper_triangular(
                outer.T, solve_upper_triangular(outer.T, inner.T).T)

        s = conjugate_solve_triangular(
            chol_x, tril_and_halve_diagonal(chol_x.T.dot(dz)))

        return [tensor.tril(s + s.T) - tensor.diag(tensor.diagonal(s))]
        
def qr_chol(a, mode="r"):
    return QR_Chol(mode)(a)

But in a short team maybe we should not force cholesky in MvNormal.

[junpenglao]

And in the long run the pymc3.gp.GP probability will work better with cholesky

[twiecki]

Yes, let's implement a short-term fix that removes the chol default.

[aseyboldt]

Look what I found scipy/scipy#6831. This seems to be exactly what we need, right (+theano wrapper of course)?

So how about this: We go back to using the general purpose matrix solve in MvNormal for now – for those cases where we don't have a cholesky factor already. This will be slower and less precise, as we do some work twice (log det and solving both do a factorization) and we don't use the fact that our matrix is spd. But it should not have issues with uncatchable linalg errors. And when the scipy PR gets merged, we try again using that.

[aseyboldt]

@twiecki we wrote that at the same time I guess. I'm on it.

[twiecki]

Awesome!

[denadai2]

found this error as well :(

[aseyboldt]

@denadai2 Yeah, sorry about that. 😇
I'll probably finish #2028 the next couple of days, but in the meantime you could use the current state in there, it should be usable as it is. Or better yet, use LKJCholeskyCov, if your use case allows it.