Use softplus to constrain Cholesky diagonal to be positive
#!/usr/bin/env python3
import torch
-from ..lazy import CholLazyTensor, DiagLazyTensor, NonLazyTensor
+from ..functions import add_diag
+from ..lazy import CholLazyTensor, NonLazyTensor
from ..distributions import MultivariateNormal
from .variational_distribution import VariationalDistribution
@@ -47,16 +48,16 @@ class CholeskyVariationalDistribution(VariationalDistribution):
that the diagonal remains positive.
"""
chol_variational_covar = self.chol_variational_covar
+ diagonal = NonLazyTensor(chol_variational_covar).diag()
dtype = chol_variational_covar.dtype
device = chol_variational_covar.device
- mask = torch.ones(self.chol_variational_covar.shape[-2:], dtype=dtype, device=device).triu()
- chol_variational_covar = chol_variational_covar.mul(mask)
- # Make sure that the diagonal is positive
- diagonal = NonLazyTensor(chol_variational_covar).diag()
- diag_correction = torch.clamp(diagonal, 1e-6, 1e10)
- chol_variational_covar = chol_variational_covar + DiagLazyTensor(diag_correction - diagonal).evaluate()
+ # First make the cholesky factor is upper triangular
+ # And has a positive diagonal
+ strictly_lower_mask = torch.ones(self.chol_variational_covar.shape[-2:], dtype=dtype, device=device).triu(1)
+ diagonal = torch.nn.functional.softplus(diagonal)
+ chol_variational_covar = add_diag(chol_variational_covar.mul(strictly_lower_mask), diagonal)
- # NOw construct the actual matrix
+ # Now construct the actual matrix
variational_covar = CholLazyTensor(chol_variational_covar.transpose(-1, -2))
return MultivariateNormal(self.variational_mean, variational_covar)
