Use HadamardGaussianLikelihood in HeterogeneousMTGP for inferred noise (#3176)

Summary:
Backlog task: T201650561 (kind of, homogenizing the TL stack)


Remove the custom likelihood override in HeterogeneousMTGP so that the parent MultiTaskGP's default HadamardGaussianLikelihood is used when noise variance is not provided. This enables per-task noise inference.

Reviewed By: sdaulton

Differential Revision: D92836693

Author: Carl Hvarfner

botorch/models/heterogeneous_mtgp.py

@@ -27,9 +27,6 @@
 from botorch.models.multitask import MultiTaskGP
 from botorch.models.transforms.input import InputTransform
 from botorch.models.transforms.outcome import OutcomeTransform
-from botorch.models.utils.gpytorch_modules import (
-    get_gaussian_likelihood_with_gamma_prior,
-)
 from botorch.posteriors.gpytorch import GPyTorchPosterior
 from botorch.posteriors.transformed import TransformedPosterior
 from botorch.utils.datasets import MultiTaskDataset
@@ -127,11 +124,7 @@ def __init__(
         # The features that are forward passed through the kernel should include
         # the task dim
         covar_module.active_dims = torch.arange(full_feature_dim + 1)
-        likelihood = (
-            None  # Constructed in MultiTaskGP.
-            if full_Yvar is not None
-            else get_gaussian_likelihood_with_gamma_prior()
-        )
+        likelihood = None  # Constructed in MultiTaskGP.
         super().__init__(
             train_X=full_X,
             train_Y=full_Y,

test/models/test_heterogeneous_mtgp.py

@@ -13,6 +13,8 @@
 from botorch.utils.testing import BotorchTestCase
 from gpytorch.distributions.multivariate_normal import MultivariateNormal
 from gpytorch.kernels import IndexKernel, ProductKernel
+from gpytorch.likelihoods.gaussian_likelihood import FixedNoiseGaussianLikelihood
+from gpytorch.likelihoods.hadamard_gaussian_likelihood import HadamardGaussianLikelihood
 
 
 class TestHeterogeneousMTGP(BotorchTestCase):
@@ -96,7 +98,7 @@ def test_input_constructor(self) -> None:
         self.assertIsNone(model_inputs["rank"])
 
     def test_standard_heterogeneous_mtgp(self) -> None:
-        # Construct the model.
+        # Construct the model (inferred noise: train_Yvars is None).
         model_inputs = HeterogeneousMTGP.construct_inputs(training_data=self.mtds)
         model = HeterogeneousMTGP(**model_inputs)
         self.assertEqual(model.train_inputs[0].shape, torch.Size([10, 6]))
@@ -113,6 +115,13 @@ def test_standard_heterogeneous_mtgp(self) -> None:
             data_covar_module.binary_map, [[1, 1, 0], [1, 0, 0], [1, 0, 1]]
         )
 
+        with self.subTest("inferred_noise_uses_hadamard_likelihood"):
+            self.assertIsInstance(model.likelihood, HadamardGaussianLikelihood)
+            # HadamardGaussianLikelihood should have per-task noise parameters.
+            self.assertEqual(
+                model.likelihood.noise_covar.noise.shape[-1], model.num_tasks
+            )
+
         # Evaluate the posterior.
         with self.assertRaisesRegex(UnsupportedError, "output_indices"):
             model.posterior(self.ds1.X, output_indices=[0, 1])
@@ -147,6 +156,25 @@ def test_standard_heterogeneous_mtgp(self) -> None:
         self.assertIsInstance(posterior.distribution, MultivariateNormal)
         self.assertEqual(posterior.mean.shape, torch.Size([5, 1]))
 
+    def test_fixed_noise_likelihood(self) -> None:
+        datasets = [
+            SupervisedDataset(
+                X=ds.X,
+                Y=ds.Y,
+                Yvar=torch.full_like(ds.Y, 0.1 * (i + 1)),
+                feature_names=ds.feature_names,
+                outcome_names=ds.outcome_names,
+            )
+            for i, ds in enumerate([self.ds1, self.ds2, self.ds3])
+        ]
+        mtds = MultiTaskDataset(
+            datasets=datasets, target_outcome_name="task0", task_feature_index=-1
+        )
+        model = HeterogeneousMTGP(
+            **HeterogeneousMTGP.construct_inputs(training_data=mtds)
+        )
+        self.assertIsInstance(model.likelihood, FixedNoiseGaussianLikelihood)
+
     def test_identical_search_space(self) -> None:
         # Check that the model works fine with identical search spaces.
         mtds = MultiTaskDataset(