Repository: cornellius-gp/gpytorch
Pull request #593 (2019-03-25 21:57:12)

Fixed noise GaussianLikelihood
Draft of a super simple fixed noise Gaussian Likelihood that takes noise levels upon instantiation and does not infer any noise levels during fitting. 

Note that the typical gpytorch pattern of using `self.likelihood(model(test_x))` does not make sense in this case. One will have to define a separate likelihood for the output. This separate output likelihood could also be a fixed noise likelihood (if for some reason observation noise levels at the test data were known). It could be a GaussianLikelihood with noise level initialized to some reasonable values (e.g. to the mean of the observation noise of the test data assuming homoskedasticity).

Balandat: @jacobrgardner this is WIP, but do you have any comments on how this may be used in conjunction with `get_fantasy_strategy` (at least without having to do bunch of instance checks on `self.likelihood`)? Right now this would error out [here](https://github.com/cornellius-gp/gpytorch/blob/f690c25de7713419e95212877be3cde1979ef6c7/gpytorch/models/exact_prediction_strategies.py#L119)
jacobrgardner: @Balandat I would make three changes:

1. The likelihood takes an `observation_noise` kwarg that, if specified, gets added instead of the fixed noise saved in the likelihood. This gives an interface for making predictions `p(y*)` on test data that has known observation noise, and will also simply fantasy observations.
2. For fantasy observations, we make two changes:
    - Likelihood defines a `get_fantasy_likelihood` similar to `get_fantasy_strategy` and `get_fantasy_model`. The default behavior is a no-op, the behavior in this case is to add the new noise observations to the fixed noise.
     - The call you linked to simply passes along kwargs to that likelihood call, in this case the observation noise for the fantasy point(s).
Balandat: sounds good, will work on incorporating this.
Balandat: argh apparently I suck at git 
Balandat: Added the observation_noise kwarg. I was surprised all tests were passing. Then I checked test coverage for likelihoods and it looks pretty dire....

Anyway, lmk if this looks good so far so I can take a stab at the fantasy likelihood
Balandat: Ok so now only FixedNoiseGaussianLikelihood uses the `observation_noise` kwarg.

I did try to get the fantasies working with the likelihood but I'm running into singularity issues inside the construction of the fast fantasies. I haven't dug into this too deeply yet. @jacobrgardner, maybe you can take a look since  you know that code well?

Here's a nb with tests and fanatasy repro failure: [dev_fixed_noise_likelihood.ipynb.txt](https://github.com/cornellius-gp/gpytorch/files/3036769/dev_fixed_noise_likelihood.ipynb.txt)

Balandat: I reduced the scope of this, removing the fantasy parts. Also did the honors of adding one of the first likelihood-specific tests...

If this looks good let's get this in, we can do the fantasy part on top of that.
jacobrgardner: @Balandat lgtm other than my one comment about `lengthscale -> noise`

Patch:
diff --git a/gpytorch/likelihoods/__init__.py b/gpytorch/likelihoods/__init__.py
index f83092cd..d0e0e2aa 100644
--- a/gpytorch/likelihoods/__init__.py
+++ b/gpytorch/likelihoods/__init__.py
@@ -2,7 +2,7 @@
 
 from .likelihood import Likelihood
 from .bernoulli_likelihood import BernoulliLikelihood
-from .gaussian_likelihood import GaussianLikelihood, _GaussianLikelihoodBase
+from .gaussian_likelihood import FixedNoiseGaussianLikelihood, GaussianLikelihood, _GaussianLikelihoodBase
 from .likelihood_list import LikelihoodList
 from .multitask_gaussian_likelihood import (
     MultitaskGaussianLikelihood,
@@ -17,6 +17,7 @@ __all__ = [
     "_GaussianLikelihoodBase",
     "_MultitaskGaussianLikelihoodBase",
     "BernoulliLikelihood",
+    "FixedNoiseGaussianLikelihood",
     "GaussianLikelihood",
     "HeteroskedasticNoise",
     "Likelihood",
diff --git a/gpytorch/likelihoods/bernoulli_likelihood.py b/gpytorch/likelihoods/bernoulli_likelihood.py
index ee638a26..11bf4f39 100644
--- a/gpytorch/likelihoods/bernoulli_likelihood.py
+++ b/gpytorch/likelihoods/bernoulli_likelihood.py
@@ -19,6 +19,7 @@ class BernoulliLikelihood(Likelihood):
             p(Y=y|f)=\Phi(yf)
         \end{equation*}
     """
+
     def forward(self, function_samples, **kwargs):
         output_probs = base_distributions.Normal(0, 1).cdf(function_samples)
         return base_distributions.Bernoulli(probs=output_probs)
@@ -34,7 +35,8 @@ class BernoulliLikelihood(Likelihood):
         if torch.any(observations.eq(-1)):
             warnings.warn(
                 "BernoulliLikelihood.expected_log_prob expects observations with labels in {0, 1}. "
-                "Observations with labels in {-1, 1} are deprecated.", DeprecationWarning
+                "Observations with labels in {-1, 1} are deprecated.",
+                DeprecationWarning,
             )
         else:
             observations = observations.mul(2).sub(1)
diff --git a/gpytorch/likelihoods/gaussian_likelihood.py b/gpytorch/likelihoods/gaussian_likelihood.py
index b4ca09f6..ae4b6569 100644
--- a/gpytorch/likelihoods/gaussian_likelihood.py
+++ b/gpytorch/likelihoods/gaussian_likelihood.py
@@ -2,17 +2,22 @@
 
 import math
 import warnings
+from typing import Any
+
+import torch
+from torch import Tensor
 
 from .. import settings
+from ..distributions import MultivariateNormal, base_distributions
 from ..likelihoods import Likelihood
-from ..distributions import base_distributions
-from .noise_models import HomoskedasticNoise
+from .noise_models import FixedGaussianNoise, HomoskedasticNoise, Noise
 
 
 class _GaussianLikelihoodBase(Likelihood):
-    """Base class for Gaussian Likelihoods, supporting general heteroskedastic noise models. """
+    """Base class for Gaussian Likelihoods, supporting general heteroskedastic noise models."""
+
+    def __init__(self, noise_covar: Noise, **kwargs: Any) -> None:
 
-    def __init__(self, noise_covar, **kwargs):
         super().__init__()
         param_transform = kwargs.get("param_transform")
         if param_transform is not None:
@@ -21,23 +26,23 @@ class _GaussianLikelihoodBase(Likelihood):
 
         self.noise_covar = noise_covar
 
-    def _shaped_noise_covar(self, base_shape, *params):
+    def _shaped_noise_covar(self, base_shape: torch.Size, *params: Any, **kwargs: Any):
         if len(params) > 0:
             # we can infer the shape from the params
             shape = None
         else:
             # here shape[:-1] is the batch shape requested, and shape[-1] is `n`, the number of points
             shape = base_shape
-        return self.noise_covar(*params, shape=shape)
+        return self.noise_covar(*params, shape=shape, **kwargs)
 
-    def forward(self, function_samples, *params, **kwargs):
-        return base_distributions.Normal(
-            function_samples, self._shaped_noise_covar(function_samples.shape, *params).diag()
-        )
+    def forward(self, function_samples: Tensor, *params: Any, **kwargs: Any) -> base_distributions.Normal:
+        observation_noise = self._shaped_noise_covar(function_samples.shape, *params, **kwargs).diag()
+        return base_distributions.Normal(function_samples, observation_noise)
 
-    def marginal(self, function_dist, *params, **kwargs):
+    def marginal(self, function_dist: MultivariateNormal, *params: Any, **kwargs: Any) -> MultivariateNormal:
         mean, covar = function_dist.mean, function_dist.lazy_covariance_matrix
-        full_covar = covar + self._shaped_noise_covar(mean.shape, *params)
+        noise_covar = self._shaped_noise_covar(mean.shape, *params, **kwargs)
+        full_covar = covar + noise_covar
         return function_dist.__class__(mean, full_covar)
 
 
@@ -51,22 +56,22 @@ class GaussianLikelihood(_GaussianLikelihoodBase):
         super().__init__(noise_covar=noise_covar)
 
     @property
-    def noise(self):
+    def noise(self) -> Tensor:
         return self.noise_covar.noise
 
     @noise.setter
-    def noise(self, value):
+    def noise(self, value: Tensor) -> None:
         self.noise_covar.initialize(noise=value)
 
     @property
-    def raw_noise(self):
+    def raw_noise(self) -> Tensor:
         return self.noise_covar.raw_noise
 
     @raw_noise.setter
-    def raw_noise(self, value):
+    def raw_noise(self, value: Tensor) -> None:
         self.noise_covar.initialize(raw_noise=value)
 
-    def expected_log_prob(self, target, input, *params, **kwargs):
+    def expected_log_prob(self, target: Tensor, input: MultivariateNormal, *params: Any, **kwargs: Any) -> Tensor:
         mean, variance = input.mean, input.variance
         noise = self.noise_covar.noise
 
@@ -81,3 +86,16 @@ class GaussianLikelihood(_GaussianLikelihoodBase):
         res = -0.5 * ((target - mean) ** 2 + variance) / noise
         res += -0.5 * noise.log() - 0.5 * math.log(2 * math.pi)
         return res.sum(-1)
+
+
+class FixedNoiseGaussianLikelihood(_GaussianLikelihoodBase):
+    def __init__(self, noise: Tensor, **kwargs: Any) -> None:
+        super().__init__(noise_covar=FixedGaussianNoise(noise=noise))
+
+    @property
+    def noise(self) -> Tensor:
+        return self.noise_covar.noise
+
+    @noise.setter
+    def noise(self, value: Tensor) -> None:
+        self.noise_covar.initialize(noise=value)
diff --git a/gpytorch/likelihoods/noise_models.py b/gpytorch/likelihoods/noise_models.py
index df50302c..b3e400fc 100644
--- a/gpytorch/likelihoods/noise_models.py
+++ b/gpytorch/likelihoods/noise_models.py
@@ -1,6 +1,9 @@
 #!/usr/bin/env python3
 
+from typing import Any, Optional
+
 import torch
+from torch import Tensor
 from torch.nn import Parameter
 
 from ..constraints import GreaterThan
@@ -10,7 +13,11 @@ from ..module import Module
 from ..utils.broadcasting import _mul_broadcast_shape
 
 
-class _HomoskedasticNoiseBase(Module):
+class Noise(Module):
+    pass
+
+
+class _HomoskedasticNoiseBase(Noise):
     def __init__(self, noise_prior=None, noise_constraint=None, batch_size=1, num_tasks=1):
         super().__init__()
         if noise_constraint is None:
@@ -27,15 +34,15 @@ class _HomoskedasticNoiseBase(Module):
         return self.raw_noise_constraint.transform(self.raw_noise)
 
     @noise.setter
-    def noise(self, value):
+    def noise(self, value: Tensor) -> None:
         self._set_noise(value)
 
-    def _set_noise(self, value):
+    def _set_noise(self, value: Tensor) -> None:
         if not torch.is_tensor(value):
             value = torch.tensor(value)
         self.initialize(raw_noise=self.raw_noise_constraint.inverse_transform(value))
 
-    def forward(self, *params, shape=None):
+    def forward(self, *params: Any, shape: Optional[torch.Size] = None) -> DiagLazyTensor:
         """In the homoskedastic case, the parameters are only used to infer the required shape.
         Here are the possible scenarios:
         - non-batched noise, non-batched input, non-MT -> noise_diag shape is `n`
@@ -89,7 +96,7 @@ class MultitaskHomoskedasticNoise(_HomoskedasticNoiseBase):
         )
 
 
-class HeteroskedasticNoise(Module):
+class HeteroskedasticNoise(Noise):
     def __init__(self, noise_model, noise_indices=None, noise_constraint=None):
         if noise_constraint is None:
             noise_constraint = GreaterThan(1e-4)
@@ -98,7 +105,9 @@ class HeteroskedasticNoise(Module):
         self._noise_constraint = noise_constraint
         self._noise_indices = noise_indices
 
-    def forward(self, *params, batch_shape=None, shape=None):
+    def forward(
+        self, *params: Any, batch_shape: Optional[torch.Size] = None, shape: Optional[torch.Size] = None
+    ) -> DiagLazyTensor:
         if len(params) == 1 and not torch.is_tensor(params[0]):
             output = self.noise_model(*params[0])
         else:
@@ -109,3 +118,14 @@ class HeteroskedasticNoise(Module):
         # will return a batched DiagLazyTensors of size n x num_tasks x num_tasks
         noise_diag = output.mean if self._noise_indices is None else output.mean[..., self._noise_indices]
         return DiagLazyTensor(self._noise_constraint.transform(noise_diag))
+
+
+class FixedGaussianNoise(Module):
+    def __init__(self, noise: Tensor) -> None:
+        super().__init__()
+        self.register_buffer("noise", noise)
+
+    def forward(self, *params: Any, observation_noise: Optional[Tensor] = None, **kwargs: Any) -> DiagLazyTensor:
+        if observation_noise is not None:
+            return DiagLazyTensor(observation_noise)
+        return DiagLazyTensor(self.noise)


Test patch:
diff --git a/test/likelihoods/test_gaussian_likelihood.py b/test/likelihoods/test_gaussian_likelihood.py
new file mode 100644
index 00000000..b1d68800
--- /dev/null
+++ b/test/likelihoods/test_gaussian_likelihood.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+
+import unittest
+
+import torch
+from gpytorch.distributions import MultivariateNormal
+from gpytorch.lazy import DiagLazyTensor
+from gpytorch.likelihoods import FixedNoiseGaussianLikelihood
+from gpytorch.likelihoods.noise_models import FixedGaussianNoise
+
+
+class TestFixedNoiseGaussianLikelihood(unittest.TestCase):
+    def test_fixed_noise_gaussian_likelihood(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            noise = 0.1 + torch.rand(5, device=device, dtype=dtype)
+            lkhd = FixedNoiseGaussianLikelihood(noise=noise)
+            # test basics
+            self.assertIsInstance(lkhd.noise_covar, FixedGaussianNoise)
+            self.assertTrue(torch.equal(noise, lkhd.noise))
+            new_noise = 0.1 + torch.rand(4, device=device, dtype=dtype)
+            lkhd.noise = new_noise
+            self.assertTrue(torch.equal(lkhd.noise, new_noise))
+            # test __call__
+            mean = torch.zeros(4, device=device, dtype=dtype)
+            covar = DiagLazyTensor(torch.ones(4, device=device, dtype=dtype))
+            mvn = MultivariateNormal(mean, covar)
+            out = lkhd(mvn)
+            self.assertTrue(torch.allclose(out.variance, 1 + new_noise))
+            # things should break if dimensions mismatch
+            mean = torch.zeros(5, device=device, dtype=dtype)
+            covar = DiagLazyTensor(torch.ones(5, device=device, dtype=dtype))
+            mvn = MultivariateNormal(mean, covar)
+            with self.assertRaises(RuntimeError):
+                lkhd(mvn)
+            # test __call__ w/ observation noise
+            obs_noise = 0.1 + torch.rand(5, device=device, dtype=dtype)
+            out = lkhd(mvn, observation_noise=obs_noise)
+            self.assertTrue(torch.allclose(out.variance, 1 + obs_noise))
+
+
+if __name__ == "__main__":
+    unittest.main()