Wishart / InverseWishart / LKJ priors

[rachtsingh]

Would it be possible to implement Wishart / InverseWishart / LKJ priors?

gpytorch has them already, but when I tried mixing in TorchDistributionMixin to get something useable in Pyro, I realized that they don't have a .sample method.

I don't think it's easy to get efficient samplers for the InverseWishart (I think trying to build it via a TransformedDistribution might be too slow), but curious to see other approaches. There's a Tensorflow Probability tutorial on this here, which explains some of the underlying ideas, as well as a note here.

Great work by the way! Pyro is so easy to use it's incredible.

[eb8680]

You actually don't need a .sample method for a distribution unless you want to sample from it - see e.g. pyro.distributions.VonMises. Does your use case require sampling from the prior? Were you getting an error when you tried to use TorchDistributionMixin with the GPyTorch versions?

[rachtsingh]

I was trying to use HMC (NUTS), which I think requires sampling (I do get an NotImplementedError). On second thought, though this might belong in torch.distributions instead of here, so feel free to close.

[eb8680]

I was trying to use HMC (NUTS), which I think requires sampling

That's just for setting the initial value - you can hack around that for now by defining a dummy sample method for your prior that just returns an arbitrary appropriately shaped value that's in the prior's support.

[rachtsingh]

Got it. I just tried it with the patched method, and ran into two issues:

1. We can't merge GPyTorch Prior classes and TorchDistributionMixin since the former subclasses nn.Module but doesn't implement forward. So I just recreated the distribution by copying their code (see below).
2. This doesn't work either because:

NotImplementedError: Cannot transform _PositiveDefinite constraints

which makes sense.

Actually modeling-wise the LKJ prior would be more useful, but I tried that and it also runs into the same error.

Here's my attempt:

class LKJCorr(Distribution, TorchDistributionMixin):
    arg_constraints = {"n": constraints.positive_integer, "eta": constraints.positive}
    support = constraints.positive_definite
    _validate_args = True

    def __init__(self, n, eta, validate_args=False):
        if not isinstance(n, int) or n < 1:
            raise ValueError("n must be a positive integer")
        if isinstance(eta, Number):
            eta = torch.tensor(float(eta))
        self.n = torch.tensor(n, dtype=torch.long, device=eta.device)
        batch_shape = eta.shape
        event_shape = torch.Size([n, n])
        i = torch.arange(n, dtype=eta.dtype, device=eta.device)
        C = (((2 * eta.view(-1, 1) - 2 + i) * i).sum(1) * math.log(2)).view_as(eta)
        C += n * torch.sum(2 * torch.lgamma(i / 2 + 1) - torch.lgamma(i + 2))
        self.eta = eta
        self.C = C
        super(LKJCorr, self).__init__(batch_shape, event_shape, validate_args=validate_args)

    def log_prob(self, X):
        if any(s != self.n for s in X.shape[-2:]):
            raise ValueError("Correlation matrix is not of size n={}".format(self.n.item()))
        if not _is_valid_correlation_matrix(X):
            raise ValueError("Input is not a valid correlation matrix")
        log_diag_sum = torch.stack([p.cholesky(upper=True).diag().log().sum() for p in X.view(-1, *X.shape[-2:])])
        return self.C + (self.eta - 1) * 2 * log_diag_sum
    
    def sample(self, sample_shape=torch.Size()):
        shape = self._extended_shape(sample_shape)
        return torch.eye(self.n.item()).expand(shape)

class LKJCov(Distribution, TorchDistributionMixin):
    arg_constraints = {"n": constraints.positive_integer, "eta": constraints.positive}
    support = constraints.positive_definite
    _validate_args = True
    
    def __init__(self, n, eta, sd_prior, validate_args=False):
        correlation_prior = LKJCorr(n=n, eta=eta, validate_args=validate_args)
        self.correlation_prior = correlation_prior
        self.sd_prior = sd_prior
        super(LKJCov, self).__init__(self.correlation_prior._batch_shape, 
                                     self.correlation_prior._event_shape, 
                                     self.correlation_prior._validate_args)

    def log_prob(self, X):
        marginal_var = torch.diagonal(X, dim1=-2, dim2=-1)
        if not torch.all(marginal_var >= 0):
            raise ValueError("Variance(s) cannot be negative")
        marginal_sd = marginal_var.sqrt()
        sd_diag_mat = _batch_form_diag(1 / marginal_sd)
        correlations = torch.matmul(torch.matmul(sd_diag_mat, X), sd_diag_mat)
        log_prob_corr = self.correlation_prior.log_prob(correlations)
        log_prob_sd = self.sd_prior.log_prob(marginal_sd)
        return log_prob_corr + log_prob_sd

    def sample(self, sample_shape=torch.Size()):
        shape = self._extended_shape(sample_shape)
        return torch.eye(self.correlation_prior.n.item()).expand(shape)

[fehiepsi]

How about using LKJ prior for Cholesky? From Stan reference LKJ prior for cov will make "the code to run slower and consume more memory with more risk of numerical errors". I think that LowerCholeskyTransform has inverse method. There still needs an implementation for .log_abs_det_jacobian() method. Then it is ready for using in LKJ Cholesky distribution.

[rachtsingh]

Huh, that sounded right, but I think I'm missing something. Shouldn't the domain of that transform be constraints.positive_definite? And I don't think it should call torch.tril, but instead torch.cholesky (which now supports batching).

When that's figured out, the implementation looks like:

class LKJCholesky(LKJCorr):
    support = constraints.lower_cholesky

    def log_prob(self, L):
        log_diag_sum = torch.diagonal(L, dim1=-2, dim2=-1).log().sum(-1)
        return self.C + (self.eta - 1) * 2 * log_diag_sum

class LKJCholeskyCov(Distribution, TorchDistributionMixin):
    arg_constraints = {"n": constraints.positive_integer, "eta": constraints.positive}
    support = constraints.lower_cholesky
    
    def __init__(self, n, eta, sd_prior, validate_args=False):
        self.correlation_prior = LKJCholesky(n=n, eta=eta, validate_args=validate_args)
        self.sd_prior = sd_prior
        super(LKJCholeskyCov, self).__init__(self.correlation_prior._batch_shape, 
                                     self.correlation_prior._event_shape, 
                                     self.correlation_prior._validate_args)
    
    def log_prob(self, L):
        # we essentially have the (LD^{1/2}) part of the LDL decomposition
        marginal_std = torch.diagonal(X, dim1=-2, dim2=-1)
        sd_diag_mat = _batch_form_diag(1 / marginal_std)
        correlation_L = torch.matmul(L, sd_diag_mat) # or is that backwards??
        log_prob_corr = self.correlation_prior.log_prob(correlation_L)
        log_prob_sd = self.sd_prior.log_prob(marginal_std)
        return log_prob_corr + log_prob_sd
    
    def sample(self, sample_shape=torch.Size()):
        shape = self._extended_shape(sample_shape)
        return torch.eye(self.correlation_prior.n.item()).expand(shape)

[fritzo]

I'm missing something. Shouldn't the domain ...

https://github.com/pytorch/pytorch/blob/master/torch/distributions/transforms.py#L514 transforms a matrix with garbage in the upper triangle to a matrix with zeros in the upper triangle, hence the use of torch.tril rather than torch.cholesky. You should never need to directly access the unconstrained domain (except indirectly via a Pyro optimizer or HMC). To get a positive definite matrix, you should be able to define x=my_transform(u) as the constrained parameter, and set cov=x.mm(x.t()).

[fehiepsi]

At one day, we should find a clean way to define unconstrained domain (maybe vector instead of matrix) for LowerCholesky Transform. It will reduce the number of parameters to optimize for these priors.

[elbamos]

Just tossing in a +1 on this... I'm experimenting with pyro by converting a Stan model, and the lack of priors for covariance matrices is kinda an impediment.

[fehiepsi]

I guess the most complicated work for LKJ prior (which is more numerically stable than Wishart/InverseWishart) is to define a transform from unconstrained space to the space of correlation matrices. Stan reference gives a nice derivation for such transform, which based on the paper: https://www.sciencedirect.com/science/article/pii/S0047259X09000876. The tricky part (which requires loops) is to transform the canonical partial correlation to the Cholesky of correlation (transform from z to x in Stan reference). It is a simple math but I don't know how to code it efficiently to support batches. Maybe it is good to start with a "loop" version if someone is interested in.

[elbamos]

Does it need to be coded to efficiently support batches? If someone has a template (i.e., that shows what functions need to be filled in), I can help.

[fehiepsi]

@elbamos Here is a template which I come up with:

Step 1: define constraint as in this script

from torch.distributions.constraints import Constraint

class _CorrelationMatrix(Constraint):
    def check(self, value):
        # check if `value` is positive definite, diagonal is equal to 1, and each entries lie in [-1, 1]

# Public interface.
correlation_matrix = _CorrelationMatrix()

Step 2: define transform as in this script

from torch.distributions.transforms import Transform

class CorrelationMatrixTransform():
    domain = constraints.real
    codomain = constraints.correlation_matrix
    bijective = True
    sign = +1

    def __eq__(self, other):
        return isinstance(other, CorrelationMatrixTransform)

    def _call(self, x):
        # compute y from the vector x as in Stan reference:
        # https://mc-stan.org/docs/2_18/reference-manual/correlation-matrix-transform-section.html
        return y

    def _inverse(self, y):
        # see the above reference
        return x

    def log_abs_det_jacobian(self, x, y):
        # see the above reference
        return log_abs_det

Step 3: register bijective as in this script

Step 4: define distribution (as in PyTorch/Pyro distributions)

• Sample method: See the algorithm in the section 2.4. of https://www.sciencedirect.com/science/article/pii/S0047259X09000876. The first step is to use Beta distribution to generate partial correlation, which is a vector in the interval [-1, 1]. Then use the correlation matrix transform to transform it to a correlation matrix.
• Log prob: See Stan reference for unnormalized probability and the section 3.3 of the above paper on how to compute normalization constant.

Hope that help!

[elbamos]

Maybe I'm missing something here, but it seems to me we don't need to define a constraint on correlation matrices or a transformation from z to x.   Don't we only need the transformation from an unconstrained vector to the lower Cholesky factor of the matrix? This can follow the algorithm in the Stan manual (and we also need its inverse and log abs det jacobian). The constraint on the domain should be `lower_cholesky`, which is already defined in pytorch.  If someone wants to turn it into a correlation matrix (I guess to look at it after the fact), they can use x'x.  No? Also, I think the relevant Stan page is this one: https://mc-stan.org/docs/2_18/reference-manual/cholesky-factors-of-correlation-matrices-1.html (Also, if you do try to implement the algorithm on https://mc-stan.org/docs/2_18/reference-manual/cholesky-factors-of-correlation-matrices-1.html, be aware that the equation for the reverse transformation is missing a factor.)

[fritzo]

Don't we only need the transformation from an unconstrained vector to the lower Cholesky

@elbamos I think what you're describing are arbitrary covariance matrices. Correlation matrices have the additional constraint that the diagonal is all-ones and off diagonal entries are all in [-1,1].

EDIT While I don't think we can use PyTorch's lower_cholesky constraint directly, I think we could translate the stan parametrization to PyTorch. @elbamos thanks for the pointers!