allow different inferred noise levels for each task in MultitaskGP (#2997)

Summary:

X-link: facebook/Ax#4233

see title

Differential Revision: D81373665

Author: sdaulton

botorch/models/multitask.py

@@ -42,7 +42,6 @@
 from botorch.models.utils.assorted import get_task_value_remapping
 from botorch.models.utils.gpytorch_modules import (
     get_covar_module_with_dim_scaled_prior,
-    get_gaussian_likelihood_with_lognormal_prior,
     MIN_INFERRED_NOISE_LEVEL,
 )
 from botorch.posteriors.multitask import MultitaskGPPosterior
@@ -56,6 +55,7 @@
 from gpytorch.kernels.index_kernel import IndexKernel
 from gpytorch.kernels.multitask_kernel import MultitaskKernel
 from gpytorch.likelihoods.gaussian_likelihood import FixedNoiseGaussianLikelihood
+from gpytorch.likelihoods.hadamard_gaussian_likelihood import HadamardGaussianLikelihood
 from gpytorch.likelihoods.likelihood import Likelihood
 from gpytorch.likelihoods.multitask_gaussian_likelihood import (
     MultitaskGaussianLikelihood,
@@ -212,10 +212,20 @@ def __init__(
         self._output_tasks = output_tasks
         self._num_outputs = len(output_tasks)
 
-        # TODO (T41270962): Support task-specific noise levels in likelihood
         if likelihood is None:
             if train_Yvar is None:
-                likelihood = get_gaussian_likelihood_with_lognormal_prior()
+                noise_prior = LogNormalPrior(loc=-4.0, scale=1.0)
+                likelihood = HadamardGaussianLikelihood(
+                    num_tasks=self.num_tasks,
+                    batch_shape=torch.Size(),
+                    noise_prior=noise_prior,
+                    noise_constraint=GreaterThan(
+                        MIN_INFERRED_NOISE_LEVEL,
+                        transform=None,
+                        initial_value=noise_prior.mode,
+                    ),
+                    task_feature_index=task_feature,
+                )
             else:
                 likelihood = FixedNoiseGaussianLikelihood(noise=train_Yvar.squeeze(-1))
 

test/models/test_multitask.py

@@ -30,6 +30,7 @@
 from gpytorch.likelihoods import (
     FixedNoiseGaussianLikelihood,
     GaussianLikelihood,
+    HadamardGaussianLikelihood,
     MultitaskGaussianLikelihood,
 )
 from gpytorch.means import ConstantMean, MultitaskMean
@@ -154,7 +155,9 @@ def test_MultiTaskGP(self) -> None:
             if fixed_noise:
                 self.assertIsInstance(model.likelihood, FixedNoiseGaussianLikelihood)
             else:
-                self.assertIsInstance(model.likelihood, GaussianLikelihood)
+                self.assertIsInstance(model.likelihood, HadamardGaussianLikelihood)
+                self.assertEqual(model.likelihood.noise.shape, torch.Size([2]))
+                self.assertEqual(model.likelihood.task_feature_index, 0)
             data_covar_module, task_covar_module = model.covar_module.kernels
             self.assertIsInstance(model.mean_module, ConstantMean)
             self.assertIsInstance(data_covar_module, RBFKernel)
@@ -195,8 +198,8 @@ def test_MultiTaskGP(self) -> None:
                     torch.tensor([0.05, 0.1], **tkwargs).repeat_interleave(2)
                 ).expand(3, 4, 4)
             else:
-                noise_covar = model.likelihood.noise_covar.noise * torch.eye(
-                    4, **tkwargs
+                noise_covar = torch.diag(
+                    model.likelihood.noise_covar.noise.repeat_interleave(2)
                 ).expand(3, 4, 4)
             expected_y_covar = posterior_f.covariance_matrix + noise_covar
             self.assertTrue(
@@ -337,7 +340,7 @@ def test_MultiTaskGP_single_output(self) -> None:
             data_covar_module, task_covar_module = model.covar_module.kernels
             self.assertIsInstance(model, MultiTaskGP)
             self.assertEqual(model.num_outputs, 1)
-            self.assertIsInstance(model.likelihood, GaussianLikelihood)
+            self.assertIsInstance(model.likelihood, HadamardGaussianLikelihood)
             self.assertIsInstance(model.mean_module, ConstantMean)
             self.assertIsInstance(data_covar_module, RBFKernel)
             self.assertIsInstance(data_covar_module.lengthscale_prior, LogNormalPrior)