[WIP] Heteroskedastic Likelihood

This is a proof of concept of how heteroskedastic likelihoods may work.

gpytorch/functions/__init__.py

@@ -64,14 +64,15 @@ def dsmm(sparse_mat, dense_mat):
     return DSMM(sparse_mat)(dense_mat)
 
 
-def exact_predictive_mean(full_covar, full_mean, train_labels, num_train, likelihood, precomputed_cache=None):
+def exact_predictive_mean(full_covar, full_mean, train_inputs, train_labels, num_train, likelihood, precomputed_cache=None):
     """
     Computes the posterior predictive mean of a GP
 
     Args:
     - full_covar ( (n+t) x (n+t) ) - the block prior covariance matrix of training and testing points
         - [ K_XX, K_XX*; K_X*X, K_X*X* ]
     - full_mean (n + t) - the training and test prior means, stacked on top of each other
+    - train_inputs TODO
     - train_labels (n) - the training labels minus the training prior mean
     - noise (1) - the observed noise (from the likelihood)
     - precomputed_cache - speeds up subsequent computations (default: None)
@@ -86,10 +87,10 @@ def exact_predictive_mean(full_covar, full_mean, train_labels, num_train, likeli
         from ..lazy.non_lazy_tensor import NonLazyTensor
 
         full_covar = NonLazyTensor(full_covar)
-    return full_covar.exact_predictive_mean(full_mean, train_labels, num_train, likelihood, precomputed_cache)
+    return full_covar.exact_predictive_mean(full_mean, train_inputs, train_labels, num_train, likelihood, precomputed_cache)
 
 
-def exact_predictive_covar(full_covar, num_train, likelihood, precomputed_cache=None):
+def exact_predictive_covar(full_covar, train_inputs, num_train, likelihood, precomputed_cache=None):
     """
     Computes the posterior predictive covariance of a GP
 
@@ -110,7 +111,7 @@ def exact_predictive_covar(full_covar, num_train, likelihood, precomputed_cache=
         from ..lazy.non_lazy_tensor import NonLazyTensor
 
         full_covar = NonLazyTensor(full_covar)
-    return full_covar.exact_predictive_covar(num_train, likelihood, precomputed_cache)
+    return full_covar.exact_predictive_covar(train_inputs, num_train, likelihood, precomputed_cache)
 
 
 def log_normal_cdf(x):

gpytorch/lazy/interpolated_lazy_tensor.py

@@ -365,7 +365,7 @@ def diag(self):
                 res = res.view(batch_size, n_data, -1).sum(-1)
             return res
 
-    def exact_predictive_mean(self, full_mean, train_labels, num_train, likelihood, precomputed_cache=None):
+    def exact_predictive_mean(self, full_mean, train_inputs, train_labels, num_train, likelihood, precomputed_cache=None):
         from ..distributions import MultivariateNormal
 
         if precomputed_cache is None:
@@ -383,7 +383,7 @@ def exact_predictive_mean(self, full_mean, train_labels, num_train, likelihood,
 
             train_mean = full_mean.narrow(-1, 0, train_train_covar.size(-1))
 
-            mvn = likelihood(MultivariateNormal(train_mean, train_train_covar))
+            mvn = likelihood(MultivariateNormal(train_mean, train_train_covar), train_inputs)
             train_mean, train_train_covar = mvn.mean, mvn.lazy_covariance_matrix
 
             train_train_covar_inv_labels = train_train_covar.inv_matmul((train_labels - train_mean).unsqueeze(-1))
@@ -423,11 +423,11 @@ def _exact_predictive_covar_inv_quad_form_root(self, precomputed_cache, test_tra
         res = left_interp(test_interp_indices, test_interp_values, precomputed_cache)
         return res
 
-    def exact_predictive_covar(self, num_train, likelihood, precomputed_cache=None):
+    def exact_predictive_covar(self, train_inputs, num_train, likelihood, precomputed_cache=None):
         from ..distributions import MultivariateNormal
 
         if not beta_features.fast_pred_var.on() and not beta_features.fast_pred_samples.on():
-            return super(InterpolatedLazyTensor, self).exact_predictive_covar(num_train, likelihood, precomputed_cache)
+            return super(InterpolatedLazyTensor, self).exact_predictive_covar(train_inputs, num_train, likelihood, precomputed_cache)
 
         n_test = self.size(-2) - num_train
         train_interp_indices = self.left_interp_indices.narrow(-2, 0, num_train)
@@ -453,7 +453,7 @@ def exact_predictive_covar(self, num_train, likelihood, precomputed_cache=None):
             )
 
             grv = MultivariateNormal(torch.zeros(1), train_train_covar)
-            train_train_covar = likelihood(grv).lazy_covariance_matrix
+            train_train_covar = likelihood(grv, train_inputs).lazy_covariance_matrix
 
             # Get probe vectors for inverse root
             num_probe_vectors = beta_features.fast_pred_var.num_probe_vectors()

gpytorch/lazy/lazy_evaluated_kernel_tensor.py

@@ -168,22 +168,22 @@ def representation_tree(self):
     def evaluate(self):
         return self.evaluate_kernel().evaluate()
 
-    def exact_predictive_mean(self, full_mean, train_labels, num_train, likelihood, precomputed_cache=None):
+    def exact_predictive_mean(self, full_mean, train_inputs, train_labels, num_train, likelihood, precomputed_cache=None):
         if self.kernel.has_custom_exact_predictions:
             return self.evaluate_kernel().exact_predictive_mean(
-                full_mean, train_labels, num_train, likelihood, precomputed_cache
+                full_mean, train_inputs, train_labels, num_train, likelihood, precomputed_cache
             )
         else:
             return super(LazyEvaluatedKernelTensor, self).exact_predictive_mean(
-                full_mean, train_labels, num_train, likelihood, precomputed_cache
+                full_mean, train_inputs, train_labels, num_train, likelihood, precomputed_cache
             )
 
-    def exact_predictive_covar(self, num_train, likelihood, precomputed_cache=None):
+    def exact_predictive_covar(self, train_inputs, num_train, likelihood, precomputed_cache=None):
         if self.kernel.has_custom_exact_predictions:
-            return self.evaluate_kernel().exact_predictive_covar(num_train, likelihood, precomputed_cache)
+            return self.evaluate_kernel().exact_predictive_covar(train_inputs, num_train, likelihood, precomputed_cache)
         else:
             return super(LazyEvaluatedKernelTensor, self).exact_predictive_covar(
-                num_train, likelihood, precomputed_cache
+                train_inputs, num_train, likelihood, precomputed_cache
             )
 
     def repeat(self, *sizes):

gpytorch/lazy/lazy_tensor.py

@@ -428,14 +428,15 @@ def evaluate_kernel(self):
         """
         return self.representation_tree()(*self.representation())
 
-    def exact_predictive_mean(self, full_mean, train_labels, num_train, likelihood, precomputed_cache=None):
+    def exact_predictive_mean(self, full_mean, train_inputs, train_labels, num_train, likelihood, precomputed_cache=None):
         """
         Computes the posterior predictive covariance of a GP
         Assumes that self is the block prior covariance matrix of training and testing points
         [ K_XX, K_XX*; K_X*X, K_X*X* ]
 
         Args:
             full_mean (:obj:`torch.tensor`): the training and test prior means, stacked on top of each other
+            train_inputs TODO
             train_labels (:obj:`torch.tensor`): the training labels minus the training prior mean
             noise (:obj:`torch.tensor`): the observed noise (from the likelihood)
             precomputed_cache (optional): speeds up subsequent computations (default: None)
@@ -453,7 +454,7 @@ def exact_predictive_mean(self, full_mean, train_labels, num_train, likelihood,
                 train_train_covar = self[:num_train, :num_train]
 
             train_mean = full_mean.narrow(-1, 0, train_train_covar.size(-1))
-            mvn = likelihood(MultivariateNormal(train_mean, train_train_covar))
+            mvn = likelihood(MultivariateNormal(train_mean, train_train_covar), train_inputs)
             train_mean, train_train_covar = mvn.mean, mvn.lazy_covariance_matrix
 
             train_labels_offset = train_labels - train_mean
@@ -472,13 +473,14 @@ def exact_predictive_mean(self, full_mean, train_labels, num_train, likelihood,
         res = res + test_mean
         return res, precomputed_cache.detach()
 
-    def exact_predictive_covar(self, num_train, likelihood, precomputed_cache=None):
+    def exact_predictive_covar(self, train_inputs, num_train, likelihood, precomputed_cache=None):
         """
         Computes the posterior predictive covariance of a GP
         Assumes that self is the block prior covariance matrix of training and testing points
         [ K_XX, K_XX*; K_X*X, K_X*X* ]
 
         Args:
+            train_inputs TODO - CAN GET RID OF num_train arg here as well!
             num_train (int): The number of training points in the full covariance matrix
             noise (scalar): The observed noise (from the likelihood)
             precomputed_cache (optional): speeds up subsequent computations (default: None)
@@ -498,7 +500,7 @@ def exact_predictive_covar(self, num_train, likelihood, precomputed_cache=None):
             test_train_covar = self[num_train:, :num_train]
             test_test_covar = self[num_train:, num_train:]
 
-        train_train_covar = likelihood(MultivariateNormal(torch.zeros(1), train_train_covar)).lazy_covariance_matrix
+        train_train_covar = likelihood(MultivariateNormal(torch.zeros(1), train_train_covar), train_inputs).lazy_covariance_matrix
         if not beta_features.fast_pred_var.on():
             from .matmul_lazy_tensor import MatmulLazyTensor
 

gpytorch/likelihoods/__init__.py

@@ -1,18 +1,18 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from __future__ import unicode_literals
+from __future__ import absolute_import, division, print_function, unicode_literals
 
+from .bernoulli_likelihood import BernoulliLikelihood
+from .gaussian_likelihood import GaussianLikelihood, HeteroskedasticGaussianLikelihood, HomoskedasticGaussianLikelihood
 from .likelihood import Likelihood
-from .gaussian_likelihood import GaussianLikelihood
 from .multitask_gaussian_likelihood import MultitaskGaussianLikelihood
-from .bernoulli_likelihood import BernoulliLikelihood
 from .softmax_likelihood import SoftmaxLikelihood
 
+
 __all__ = [
-    "Likelihood",
+    "BernoulliLikelihood",
     "GaussianLikelihood",
+    "HeteroskedasticGaussianLikelihood",
+    "HomoskedasticGaussianLikelihood",
+    "Likelihood",
     "MultitaskGaussianLikelihood",
-    "BernoulliLikelihood",
     "SoftmaxLikelihood",
 ]

gpytorch/likelihoods/gaussian_likelihood.py

@@ -1,42 +1,34 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import math
-import torch
-from ..distributions import MultivariateNormal
-from ..functions import add_diag
-from ..likelihoods import Likelihood
+
 from .. import settings
+from ..distributions import MultivariateNormal
+from ..lazy import DiagLazyTensor
+from .homoskedastic_noise import HomoskedasticNoise
+from .likelihood import Likelihood
 
 
 class GaussianLikelihood(Likelihood):
-    r"""
-    """
+    def __init__(self, noise_covar):
+        Likelihood.__init__(self)
+        self.noise_covar = noise_covar
 
-    def __init__(self, log_noise_prior=None, batch_size=1):
-        super(GaussianLikelihood, self).__init__()
-        self.register_parameter(
-            name="log_noise", parameter=torch.nn.Parameter(torch.zeros(batch_size, 1)), prior=log_noise_prior
-        )
-
-    @property
-    def noise(self):
-        return self.log_noise.exp()
-
-    def forward(self, input):
+    def forward(self, input, *params):
         if not isinstance(input, MultivariateNormal):
             raise ValueError("GaussianLikelihood requires a MultivariateNormal input")
         mean, covar = input.mean, input.lazy_covariance_matrix
-        noise = self.noise
-        if covar.ndimension() == 2:
-            if settings.debug.on() and noise.size(0) > 1:
-                raise RuntimeError("With batch_size > 1, expected a batched MultivariateNormal distribution.")
-            noise = noise.squeeze(0)
+        return input.__class__(mean, covar + self.noise_covar(*params))
 
-        return input.__class__(mean, add_diag(covar, noise))
+
+class HomoskedasticGaussianLikelihood(GaussianLikelihood):
+    def __init__(self, log_noise_prior=None, batch_size=1):
+        noise_covar = HomoskedasticNoise(log_noise_prior=log_noise_prior, batch_size=1)
+        super(HomoskedasticGaussianLikelihood, self).__init__(noise_covar)
 
     def variational_log_probability(self, input, target):
         mean, variance = input.mean, input.variance
-        log_noise = self.log_noise
+        log_noise = self.noise_covar.log_noise
         if variance.ndimension() == 1:
             if settings.debug.on() and log_noise.size(0) > 1:
                 raise RuntimeError("With batch_size > 1, expected a batched MultivariateNormal distribution.")
@@ -45,3 +37,18 @@ def variational_log_probability(self, input, target):
         res = -0.5 * ((target - mean) ** 2 + variance) / log_noise.exp()
         res += -0.5 * log_noise - 0.5 * math.log(2 * math.pi)
         return res.sum(0)
+
+
+class HeteroskedasticGaussianLikelihood(GaussianLikelihood):
+    def __init__(self, log_noise_model):
+        Likelihood.__init__(self)
+        self.log_noise_model = log_noise_model
+
+    def forward(self, input, *params):
+        if not isinstance(input, MultivariateNormal):
+            raise ValueError("HeteroskedasticGaussianLikelihood requires a MultivariateNormal input")
+        mean, covar = input.mean, input.lazy_covariance_matrix
+        # TODO: This is inefficient, fix it! Allow for non-diagonal outputs
+        log_noise_covar = self.log_noise_model(*params).diag()
+        noise_covar = DiagLazyTensor(log_noise_covar.exp())
+        return input.__class__(mean, covar + noise_covar)

gpytorch/likelihoods/homoskedastic_noise.py

@@ -0,0 +1,26 @@
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import torch
+from torch.nn import Parameter
+
+from ..lazy import DiagLazyTensor
+from ..module import Module
+
+
+class HomoskedasticNoise(Module):
+    def __init__(self, log_noise_prior=None, batch_size=1):
+        super(HomoskedasticNoise, self).__init__()
+        self.register_parameter(
+            name="log_noise", parameter=Parameter(torch.zeros(batch_size, 1)), prior=log_noise_prior
+        )
+
+    def forward(self, params):
+        noise = self.log_noise.exp()
+        if isinstance(params, list):
+            variance_shape = params[0].shape[:-2] + params[0].shape[-1:]
+        else:
+            variance_shape = params.shape[:-2] + params.shape[-1:]
+        if len(variance_shape) == 1:
+            noise = noise.squeeze(0)
+        variances = noise * torch.ones(*variance_shape, dtype=noise.dtype, device=noise.device)
+        return DiagLazyTensor(variances)

gpytorch/mlls/exact_marginal_log_likelihood.py

@@ -23,21 +23,22 @@ def __init__(self, likelihood, model):
             raise RuntimeError("Likelihood must be Gaussian for exact inference")
         super(ExactMarginalLogLikelihood, self).__init__(likelihood, model)
 
-    def forward(self, output, target):
+    def forward(self, output, target, *params):
         if not isinstance(output, MultivariateNormal):
             raise RuntimeError("ExactMarginalLogLikelihood can only operate on Gaussian random variables")
 
         # Get the log prob of the marginal distribution
-        output = self.likelihood(output)
+        output = self.likelihood(output, *params)
         res = output.log_prob(target)
 
         # Add terms for SGPR / when inducing points are learned
         trace_diff = torch.zeros_like(res)
         for variational_strategy in self.model.variational_strategies():
             if isinstance(variational_strategy, MVNVariationalStrategy):
                 trace_diff = trace_diff.add(variational_strategy.trace_diff())
-        trace_diff = trace_diff / self.likelihood.log_noise.exp()
-        res = res.add(0.5, trace_diff)
+        if hasattr(self.likelihood, "log_noise"):
+            trace_diff = trace_diff / self.likelihood.log_noise.exp()
+            res = res.add(0.5, trace_diff)
 
         # Add log probs of priors on the parameters
         for _, param, prior in self.named_parameter_priors():

gpytorch/models/__init__.py

@@ -1,12 +1,16 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from __future__ import unicode_literals
+from __future__ import absolute_import, division, print_function, unicode_literals
 
-from .gp import GP
+from .additive_grid_inducing_variational_gp import AdditiveGridInducingVariationalGP
 from .exact_gp import ExactGP
-from .variational_gp import VariationalGP
+from .gp import GP
 from .grid_inducing_variational_gp import GridInducingVariationalGP
-from .additive_grid_inducing_variational_gp import AdditiveGridInducingVariationalGP
+from .variational_gp import VariationalGP
+
 
-__all__ = ["GP", "ExactGP", "VariationalGP", "GridInducingVariationalGP", "AdditiveGridInducingVariationalGP"]
+__all__ = [
+    "AdditiveGridInducingVariationalGP",
+    "ExactGP",
+    "GP",
+    "VariationalGP",
+    "GridInducingVariationalGP",
+]

gpytorch/models/exact_gp.py

@@ -148,13 +148,15 @@ def __call__(self, *args, **kwargs):
             predictive_mean, mean_cache = exact_predictive_mean(
                 full_covar=full_covar,
                 full_mean=full_mean,
+                train_inputs=train_inputs,
                 train_labels=train_targets,
                 num_train=num_train,
                 likelihood=self.likelihood,
                 precomputed_cache=self.mean_cache,
             )
             predictive_covar, covar_cache = exact_predictive_covar(
                 full_covar=full_covar,
+                train_inputs=train_inputs,
                 num_train=num_train,
                 likelihood=self.likelihood,
                 precomputed_cache=self.covar_cache,