Repository: cornellius-gp/gpytorch

Issue:
[Bug] psd_safe_cholesky less stable with small noise
# 🐛 Bug

I'm getting errors (see below) now from `psd_safe_cholesky` which weren't there a few commits back.

* Current commit has this error: 6a235fb6ec339e77a86a3cb20c715071c71ac9fc
* A recent commit without this error: 308a18094d31d5cd2fbb7d942212669141dd8cec. I'm not sure what the *most* recent commit without this error is, though.

## To reproduce

```python
import torch
import gpytorch

device = "cuda"

train_x = torch.tensor([
    [-55.5527, -42.9031],
    [ 63.1439, -36.1096],
    [-44.4039, -51.3083],
    [-37.9279,  44.9540],
    [-54.8894,   1.7219],
    [-26.8499, -40.2266],
    [ 17.5439,  26.0574],
    [-31.0171,  -2.4689],
    [ 37.8813, -39.7865],
    [-10.5718,  67.5785],
    [ 13.2326, -45.9986],
    [ 51.3409, -21.4855],
    [ 64.9224, -52.7161],
    [ 13.1864,  63.9087],
    [-67.8713, -59.4553],
    [  3.3758,  13.9108],
    [ 12.1702,  35.1921],
    [ 47.5776,  69.5157],
    [ -7.5170,  59.2905],
    [-34.0096, -51.0971],
    [-43.5709,  14.1764],
    [-17.3790,  57.2735],
    [ 46.1217,   4.4776],
    [ 17.8831,  68.3338],
    [ -0.1228,  61.4203],
    [ 46.7972,  27.2138],
    [-35.2764, -19.0872],
    [ 36.9366,  45.3593],
    [-20.7411,  40.9940],
    [ 16.2726, -29.3061]
], device=device)
train_y = torch.tensor([59.1276,  -53.7113,   -1.7330, -105.8712,   10.8977,  -16.1213,
         -43.1551,   55.5215,   32.3309,   87.8219,    4.2356,   28.3534,
         -60.4486,  115.4285,   68.5259,  -64.0631,  -12.9961,   34.0834,
          61.4407,  -28.7934,  -31.1267,   18.5751,  -70.5619,  121.7277,
          90.2278,   73.1477,   36.8967,   48.5388,  -68.2077,    1.3766], device=device)

class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y):
        likelihood = gpytorch.likelihoods.GaussianLikelihood()
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.MaternKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

def train(model, train_x, train_y, train_steps):
    model.set_train_data(train_x, train_y, strict=False)
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

    mll = gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model)

    for i in range(train_steps):
        optimizer.zero_grad()
        output = model(train_x)
        loss = -mll(output, train_y)
        loss.backward()
        optimizer.step()
    model.eval()

model = ExactGPModel(train_x, train_y)
model.to(device)
model.likelihood.initialize(noise=1e-5)
model.likelihood.noise_covar.raw_noise.requires_grad_(False)
train(model, train_x, train_y, train_steps=100)
```

## Stack traces

If `device = "cpu"`"

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-5-dcd721edef0b> in <module>
     73 model.likelihood.initialize(noise=1e-5)
     74 model.likelihood.noise_covar.raw_noise.requires_grad_(False)
---> 75 train(model, train_x, train_y, train_steps=100)

<ipython-input-5-dcd721edef0b> in train(model, train_x, train_y, train_steps)
     64         optimizer.zero_grad()
     65         output = model(train_x)
---> 66         loss = -mll(output, train_y)
     67         loss.backward()
     68         optimizer.step()

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)
     20 
     21     def __call__(self, *inputs, **kwargs):
---> 22         outputs = self.forward(*inputs, **kwargs)
     23         if isinstance(outputs, list):
     24             return [_validate_module_outputs(output) for output in outputs]

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/mlls/exact_marginal_log_likelihood.py in forward(self, output, target, *params)
     26         # Get the log prob of the marginal distribution
     27         output = self.likelihood(output, *params)
---> 28         res = output.log_prob(target)
     29 
     30         # Add terms for SGPR / when inducing points are learned

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/distributions/multivariate_normal.py in log_prob(self, value)
    127 
    128         # Get log determininat and first part of quadratic form
--> 129         inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
    130 
    131         res = -0.5 * sum([inv_quad, logdet, diff.size(-1) * math.log(2 * math.pi)])

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in inv_quad_logdet(self, inv_quad_rhs, logdet, reduce_inv_quad)
    990             from .chol_lazy_tensor import CholLazyTensor
    991 
--> 992             cholesky = CholLazyTensor(self.cholesky())
    993             return cholesky.inv_quad_logdet(inv_quad_rhs=inv_quad_rhs, logdet=logdet, reduce_inv_quad=reduce_inv_quad)
    994 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in cholesky(self, upper)
    716             (LazyTensor) Cholesky factor (lower triangular)
    717         """
--> 718         res = self._cholesky()
    719         if upper:
    720             res = res.transpose(-1, -2)

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
     32         cache_name = name if name is not None else method
     33         if not is_in_cache(self, cache_name):
---> 34             add_to_cache(self, cache_name, method(self, *args, **kwargs))
     35         return get_from_cache(self, cache_name)
     36 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in _cholesky(self)
    401             evaluated_mat.register_hook(_ensure_symmetric_grad)
    402 
--> 403         cholesky = psd_safe_cholesky(evaluated_mat.double()).to(self.dtype)
    404         return NonLazyTensor(cholesky)
    405 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/utils/cholesky.py in psd_safe_cholesky(A, upper, out, jitter)
     44                 continue
     45 
---> 46         raise e
     47 
     48 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/utils/cholesky.py in psd_safe_cholesky(A, upper, out, jitter)
     19     """
     20     try:
---> 21         L = torch.cholesky(A, upper=upper, out=out)
     22         # TODO: Remove once fixed in pytorch (#16780)
     23         if A.dim() > 2 and A.is_cuda:

RuntimeError: Lapack Error in potrf : the leading minor of order 1 is not positive definite at /pytorch/aten/src/TH/generic/THTensorLapack.cpp:658
```

If `device = "cuda"`

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-6-2b962e725013> in <module>
     73 model.likelihood.initialize(noise=1e-5)
     74 model.likelihood.noise_covar.raw_noise.requires_grad_(False)
---> 75 train(model, train_x, train_y, train_steps=100)

<ipython-input-6-2b962e725013> in train(model, train_x, train_y, train_steps)
     64         optimizer.zero_grad()
     65         output = model(train_x)
---> 66         loss = -mll(output, train_y)
     67         loss.backward()
     68         optimizer.step()

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)
     20 
     21     def __call__(self, *inputs, **kwargs):
---> 22         outputs = self.forward(*inputs, **kwargs)
     23         if isinstance(outputs, list):
     24             return [_validate_module_outputs(output) for output in outputs]

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/mlls/exact_marginal_log_likelihood.py in forward(self, output, target, *params)
     26         # Get the log prob of the marginal distribution
     27         output = self.likelihood(output, *params)
---> 28         res = output.log_prob(target)
     29 
     30         # Add terms for SGPR / when inducing points are learned

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/distributions/multivariate_normal.py in log_prob(self, value)
    127 
    128         # Get log determininat and first part of quadratic form
--> 129         inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
    130 
    131         res = -0.5 * sum([inv_quad, logdet, diff.size(-1) * math.log(2 * math.pi)])

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in inv_quad_logdet(self, inv_quad_rhs, logdet, reduce_inv_quad)
    990             from .chol_lazy_tensor import CholLazyTensor
    991 
--> 992             cholesky = CholLazyTensor(self.cholesky())
    993             return cholesky.inv_quad_logdet(inv_quad_rhs=inv_quad_rhs, logdet=logdet, reduce_inv_quad=reduce_inv_quad)
    994 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in cholesky(self, upper)
    716             (LazyTensor) Cholesky factor (lower triangular)
    717         """
--> 718         res = self._cholesky()
    719         if upper:
    720             res = res.transpose(-1, -2)

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
     32         cache_name = name if name is not None else method
     33         if not is_in_cache(self, cache_name):
---> 34             add_to_cache(self, cache_name, method(self, *args, **kwargs))
     35         return get_from_cache(self, cache_name)
     36 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in _cholesky(self)
    401             evaluated_mat.register_hook(_ensure_symmetric_grad)
    402 
--> 403         cholesky = psd_safe_cholesky(evaluated_mat.double()).to(self.dtype)
    404         return NonLazyTensor(cholesky)
    405 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/utils/cholesky.py in psd_safe_cholesky(A, upper, out, jitter)
     44                 continue
     45 
---> 46         raise e
     47 
     48 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/utils/cholesky.py in psd_safe_cholesky(A, upper, out, jitter)
     19     """
     20     try:
---> 21         L = torch.cholesky(A, upper=upper, out=out)
     22         # TODO: Remove once fixed in pytorch (#16780)
     23         if A.dim() > 2 and A.is_cuda:

RuntimeError: MAGMA potrf : A(1,1) is 0, A cannot be factorized at /pytorch/aten/src/THC/generic/THCTensorMathMagma.cu:597
```

## Expected Behavior

Training is no less stable than it was at, e.g., commit 308a18094d31d5cd2fbb7d942212669141dd8cec.

## System information

**Please complete the following information:**
- GPyTorch Version: commit 6a235fb6ec339e77a86a3cb20c715071c71ac9fc
- PyTorch Version: 1.0.1.post2
- OS: Ubuntu 16.04.5

## Additional context

Removing the `model.likelihood` lines makes this work, so it only happens with small noise.


Discussion:
I know exactly what is happening here actually -- the issue is that we've introduced a default lower bound on the noise, and currently trying to set the noise lower than that lower bound ends up setting it to nan rather than either erroring or automatically adjusting the lower bound.

I'm working on a fix right now.

Resolving patch:
diff --git a/gpytorch/constraints/constraints.py b/gpytorch/constraints/constraints.py
--- a/gpytorch/constraints/constraints.py
+++ b/gpytorch/constraints/constraints.py
@@ -44,7 +44,13 @@ def enforced(self):
         return self._transform is not None
 
     def check(self, tensor):
-        return torch.all(tensor <= self.upper_bound) and torch.all(tensor >= self.lower_bound)
+        return bool(torch.all(tensor <= self.upper_bound) and torch.all(tensor >= self.lower_bound))
+
+    def check_raw(self, tensor):
+        return bool(
+            torch.all((self.transform(tensor) <= self.upper_bound))
+            and torch.all(self.transform(tensor) >= self.lower_bound)
+        )
 
     def intersect(self, other):
         """
diff --git a/gpytorch/module.py b/gpytorch/module.py
--- a/gpytorch/module.py
+++ b/gpytorch/module.py
@@ -77,12 +77,28 @@ def initialize(self, **kwargs):
             elif name not in self._parameters:
                 setattr(self, name, val)
             elif torch.is_tensor(val):
+                constraint = self.constraint_for_parameter_name(name)
+                if constraint is not None and not constraint.check_raw(val):
+                    raise RuntimeError(
+                        "Attempting to manually set a parameter value that is out of bounds of "
+                        f"its current constraints, {constraint}. "
+                        "Most likely, you want to do the following:\n likelihood = GaussianLikelihood"
+                        "(noise_constraint=gpytorch.constraints.GreaterThan(better_lower_bound))"
+                    )
                 try:
                     self.__getattr__(name).data.copy_(val.expand_as(self.__getattr__(name)))
                 except RuntimeError:
                     self.__getattr__(name).data.copy_(val.view_as(self.__getattr__(name)))
 
             elif isinstance(val, float):
+                constraint = self.constraint_for_parameter_name(name)
+                if constraint is not None and not constraint.check_raw(val):
+                    raise RuntimeError(
+                        "Attempting to manually set a parameter value that is out of bounds of "
+                        f"its current constraints, {constraint}. "
+                        "Most likely, you want to do the following:\n likelihood = GaussianLikelihood"
+                        "(noise_constraint=gpytorch.constraints.GreaterThan(better_lower_bound))"
+                    )
                 self.__getattr__(name).data.fill_(val)
             else:
                 raise AttributeError("Type {t} not valid for initializing parameter {p}".format(t=type(val), p=name))


Test patch:
diff --git a/test/kernels/_base_kernel_test_case.py b/test/kernels/_base_kernel_test_case.py
--- a/test/kernels/_base_kernel_test_case.py
+++ b/test/kernels/_base_kernel_test_case.py
@@ -20,7 +20,7 @@ def test_active_dims_list(self):
         kernel_basic = self.create_kernel_no_ard()
         covar_mat_actual = kernel_basic(x[:, [0, 2, 4, 6]]).evaluate_kernel().evaluate()
 
-        self.assertLess(torch.norm(covar_mat - covar_mat_actual), 1e-4)
+        self.assertLess(torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(), 1e-4)
 
     def test_active_dims_range(self):
         active_dims = list(range(3, 9))
@@ -30,7 +30,7 @@ def test_active_dims_range(self):
         kernel_basic = self.create_kernel_no_ard()
         covar_mat_actual = kernel_basic(x[:, active_dims]).evaluate_kernel().evaluate()
 
-        self.assertLess(torch.norm(covar_mat - covar_mat_actual), 1e-4)
+        self.assertLess(torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(), 1e-4)
 
     def test_no_batch_kernel_single_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard()
@@ -41,7 +41,7 @@ def test_no_batch_kernel_single_batch_x_no_ard(self):
         actual_mat_2 = kernel(x[1]).evaluate_kernel().evaluate()
         actual_covar_mat = torch.cat([actual_mat_1.unsqueeze(0), actual_mat_2.unsqueeze(0)])
 
-        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
 
     def test_single_batch_kernel_single_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
@@ -52,7 +52,7 @@ def test_single_batch_kernel_single_batch_x_no_ard(self):
         actual_mat_2 = kernel(x[1]).evaluate_kernel().evaluate()
         actual_covar_mat = torch.cat([actual_mat_1.unsqueeze(0), actual_mat_2.unsqueeze(0)])
 
-        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
 
     def test_no_batch_kernel_double_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
@@ -68,7 +68,7 @@ def test_no_batch_kernel_double_batch_x_no_ard(self):
 
         actual_covar_mat = torch.cat([ac.unsqueeze(0) for ac in ij_actual_covars])
 
-        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
 
     def test_no_batch_kernel_double_batch_x_ard(self):
         try:
@@ -88,7 +88,7 @@ def test_no_batch_kernel_double_batch_x_ard(self):
 
         actual_covar_mat = torch.cat([ac.unsqueeze(0) for ac in ij_actual_covars])
 
-        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
 
     def test_smoke_double_batch_kernel_double_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([3, 2]))