Generic FullyBayesianMultiTaskGP (meta-pytorch#3180)

Summary:

Creates generic FullyBayesianMultiTaskGP implementation and refactors the existing SAAS model as a special case

Differential Revision: D93037430

Author: Carl Hvarfner

botorch/fit.py

@@ -21,7 +21,7 @@
 from botorch.models import SingleTaskGP
 from botorch.models.approximate_gp import ApproximateGPyTorchModel
 from botorch.models.fully_bayesian import AbstractFullyBayesianSingleTaskGP
-from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
+from botorch.models.fully_bayesian_multitask import FullyBayesianMultiTaskGP
 from botorch.models.map_saas import get_map_saas_model
 from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.models.transforms.input import InputTransform
@@ -334,7 +334,7 @@ def _fit_fallback_approximate(
 
 
 def fit_fully_bayesian_model_nuts(
-    model: AbstractFullyBayesianSingleTaskGP | SaasFullyBayesianMultiTaskGP,
+    model: AbstractFullyBayesianSingleTaskGP | FullyBayesianMultiTaskGP,
     max_tree_depth: int = 6,
     warmup_steps: int = 512,
     num_samples: int = 256,

botorch/models/__init__.py

@@ -15,7 +15,10 @@
     PosteriorMeanModel,
 )
 from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
-from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
+from botorch.models.fully_bayesian_multitask import (
+    FullyBayesianMultiTaskGP,
+    SaasFullyBayesianMultiTaskGP,
+)
 from botorch.models.gp_regression import SingleTaskGP
 from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
 from botorch.models.gp_regression_mixed import MixedSingleTaskGP
@@ -37,6 +40,7 @@
     "AffineFidelityCostModel",
     "ApproximateGPyTorchModel",
     "EnsembleMapSaasSingleTaskGP",
+    "FullyBayesianMultiTaskGP",
     "GenericDeterministicModel",
     "HigherOrderGP",
     "KroneckerMultiTaskGP",

botorch/models/fully_bayesian.py

@@ -169,8 +169,8 @@ def set_inputs(
         train_X: Tensor,
         train_Y: Tensor,
         train_Yvar: Tensor | None = None,
-        task_feature: int | None = None,
-        task_rank: int | None = None,
+        task_feature: int | None = None,  # noqa: ARG002
+        task_rank: int | None = None,  # noqa: ARG002
     ) -> None:
         """Set the training data.
 

botorch/models/fully_bayesian_multitask.py

@@ -16,6 +16,7 @@
     matern52_kernel,
     MCMC_DIM,
     MIN_INFERRED_NOISE_LEVEL,
+    PyroModel,
     reshape_and_detach,
     SaasPyroModel,
 )
@@ -39,8 +40,8 @@
 from typing_extensions import Self
 
 # Can replace with Self type once 3.11 is the minimum version
-TSaasFullyBayesianMultiTaskGP = TypeVar(
-    "TSaasFullyBayesianMultiTaskGP", bound="SaasFullyBayesianMultiTaskGP"
+TFullyBayesianMultiTaskGP = TypeVar(
+    "TFullyBayesianMultiTaskGP", bound="FullyBayesianMultiTaskGP"
 )
 
 
@@ -260,19 +261,20 @@ def get_dummy_mcmc_samples(
 
 class MultitaskSaasPyroModel(LatentFeatureMultiTaskPyroMixin, SaasPyroModel):
     r"""
-    Multi-task SAAS model. Backward-compatible subclass that composes
-    ``LatentFeatureMultiTaskPyroMixin`` with ``SaasPyroModel``.
+    Multi-task SAAS model using latent task features. Backward-compatible
+    subclass that composes ``LatentFeatureMultiTaskPyroMixin`` with
+    ``SaasPyroModel``.
     """
 
     pass
 
 
-class SaasFullyBayesianMultiTaskGP(MultiTaskGP):
-    r"""A fully Bayesian multi-task GP model with the SAAS prior.
+class FullyBayesianMultiTaskGP(MultiTaskGP):
+    r"""A fully Bayesian multi-task GP model.
+
     This model assumes that the inputs have been normalized to [0, 1]^d and that the
     output has been stratified standardized to have zero mean and unit variance for
-    each task. The SAAS model [Eriksson2021saasbo]_ with a Matern-5/2 is used as data
-    kernel by default.
+    each task.
 
     You are expected to use ``fit_fully_bayesian_model_nuts`` to fit this model as it
     isn't compatible with ``fit_gpytorch_mll``.
@@ -285,11 +287,12 @@ class SaasFullyBayesianMultiTaskGP(MultiTaskGP):
         >>> ])
         >>> train_Y = torch.cat(f1(X1), f2(X2)).unsqueeze(-1)
         >>> train_Yvar = 0.01 * torch.ones_like(train_Y)
-        >>> mtsaas_gp = SaasFullyBayesianMultiTaskGP(
-        >>>     train_X, train_Y, train_Yvar, task_feature=-1,
+        >>> mt_gp = FullyBayesianMultiTaskGP(
+        >>>     train_X, train_Y, task_feature=-1,
+        >>>     pyro_model=MultitaskSaasPyroModel(),
         >>> )
-        >>> fit_fully_bayesian_model_nuts(mtsaas_gp)
-        >>> posterior = mtsaas_gp.posterior(test_X)
+        >>> fit_fully_bayesian_model_nuts(mt_gp)
+        >>> posterior = mt_gp.posterior(test_X)
     """
 
     _is_fully_bayesian = True
@@ -306,7 +309,7 @@ def __init__(
         all_tasks: list[int] | None = None,
         outcome_transform: OutcomeTransform | None = None,
         input_transform: InputTransform | None = None,
-        pyro_model: MultitaskSaasPyroModel | None = None,
+        pyro_model: PyroModel | None = None,
         validate_task_values: bool = True,
     ) -> None:
         r"""Initialize the fully Bayesian multi-task GP model.
@@ -333,8 +336,7 @@ def __init__(
                 instantiation of the model.
             input_transform: An input transform that is applied to the inputs ``X``
                 in the model's forward pass.
-            pyro_model: Optional ``PyroModel`` that has the same signature as
-                ``MultitaskSaasPyroModel``. Defaults to ``MultitaskSaasPyroModel``.
+            pyro_model: A ``PyroModel`` that inherits from ``MultiTaskPyroMixin``.
             validate_task_values: If True, validate that the task values supplied in the
                 input are expected tasks values. If false, unexpected task values
                 will be mapped to the first output_task if supplied.
@@ -384,7 +386,8 @@ def __init__(
         self.likelihood = None
         if pyro_model is None:
             pyro_model = MultitaskSaasPyroModel()
-        # apply task_mapper
+        if not isinstance(pyro_model, MultiTaskPyroMixin):
+            raise ValueError("pyro_model must be a multi-task model.")
         x_before, task_idcs, x_after = self._split_inputs(transformed_X)
         pyro_model.set_inputs(
             train_X=torch.cat([x_before, task_idcs, x_after], dim=-1),
@@ -394,15 +397,13 @@ def __init__(
             task_rank=self._rank,
             all_tasks=all_tasks,
         )
-        self.pyro_model: MultitaskSaasPyroModel = pyro_model
+        self.pyro_model: PyroModel = pyro_model
         if outcome_transform is not None:
             self.outcome_transform = outcome_transform
         if input_transform is not None:
             self.input_transform = input_transform
 
-    def train(
-        self, mode: bool = True, reset: bool = True
-    ) -> TSaasFullyBayesianMultiTaskGP:
+    def train(self, mode: bool = True, reset: bool = True) -> TFullyBayesianMultiTaskGP:
         r"""Puts the model in ``train`` mode.
 
         Args:
@@ -436,7 +437,7 @@ def num_mcmc_samples(self) -> int:
     @property
     def batch_shape(self) -> torch.Size:
         r"""Batch shape of the model, equal to the number of MCMC samples.
-        Note that ``SaasFullyBayesianMultiTaskGP`` does not support batching
+        Note that ``FullyBayesianMultiTaskGP`` does not support batching
         over input data at this point.
         """
         self._check_if_fitted()
@@ -513,22 +514,14 @@ def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
         r"""Custom logic for loading the state dict.
 
         The standard approach of calling ``load_state_dict`` currently doesn't
-        play well with the ``SaasFullyBayesianMultiTaskGP`` since the
+        play well with the ``FullyBayesianMultiTaskGP`` since the
         ``mean_module``, ``covar_module`` and ``likelihood`` aren't initialized
         until the model has been fitted. The reason for this is that we don't
         know the number of MCMC samples until NUTS is called. Given the state
         dict, we can initialize a new model with some dummy samples and then
-        load the state dict into this model. This currently only works for a
-        ``MultitaskSaasPyroModel`` and supporting more Pyro models likely
-        requires moving the model construction logic into the Pyro model itself.
-
-        TODO: If this were to inherit from ``SaasFullyBayesianSingleTaskGP``, we could
-        simplify this method and eliminate some others.
+        load the state dict into this model. The dummy samples are obtained
+        from ``pyro_model.get_dummy_mcmc_samples()``.
         """
-        if not isinstance(self.pyro_model, MultitaskSaasPyroModel):
-            raise NotImplementedError(  # pragma: no cover
-                "load_state_dict only works for MultitaskSaasPyroModel"
-            )
         raw_mean = state_dict["mean_module.base_means.0.raw_constant"]
         num_mcmc_samples = len(raw_mean)
         tkwargs = {"device": raw_mean.device, "dtype": raw_mean.dtype}
@@ -572,3 +565,7 @@ def condition_on_observations(
             X = X.repeat(*(Y.shape[:-2] + (1, 1)))
 
         return super().condition_on_observations(X, Y, **kwargs)
+
+
+class SaasFullyBayesianMultiTaskGP(FullyBayesianMultiTaskGP):
+    r"""A fully Bayesian multi-task GP model with the SAAS prior by default."""

test/models/test_fully_bayesian_multitask.py

@@ -39,6 +39,7 @@
     SaasPyroModel,
 )
 from botorch.models.fully_bayesian_multitask import (
+    FullyBayesianMultiTaskGP,
     LatentFeatureMultiTaskPyroMixin,
     MultiTaskPyroMixin,
     MultitaskSaasPyroModel,
@@ -915,6 +916,77 @@ def test_construct_inputs(self):
             self.assertEqual(model._task_feature, d)
             self.assertEqual(model.pyro_model.task_feature, d)
 
+    def test_constructor_validation_and_defaults(self):
+        """Test FullyBayesianMultiTaskGP constructor validation and defaults."""
+        tkwargs = {"device": self.device, "dtype": torch.double}
+        train_X, train_Y, train_Yvar = self._get_base_data(**tkwargs)
+
+        with self.subTest("rejects_single_task_pyro_model"):
+            with self.assertRaisesRegex(
+                ValueError, "pyro_model must be a multi-task model"
+            ):
+                FullyBayesianMultiTaskGP(
+                    train_X=train_X,
+                    train_Y=train_Y,
+                    train_Yvar=train_Yvar,
+                    task_feature=4,
+                    pyro_model=MaternPyroModel(),
+                )
+
+        with self.subTest("accepts_explicit_multitask_pyro_model"):
+            pyro_model = MultitaskSaasPyroModel()
+            model = FullyBayesianMultiTaskGP(
+                train_X=train_X,
+                train_Y=train_Y,
+                train_Yvar=train_Yvar,
+                task_feature=4,
+                pyro_model=pyro_model,
+            )
+            self.assertIs(model.pyro_model, pyro_model)
+            self.assertIsInstance(model.pyro_model, MultiTaskPyroMixin)
+
+    def test_non_saas_mt_model_load_state_dict(self):
+        """Test round-trip load_state_dict with a non-SAAS multi-task PyroModel."""
+        tkwargs = {"device": self.device, "dtype": torch.double}
+
+        class MultitaskMaternPyroModel(
+            LatentFeatureMultiTaskPyroMixin, MaternPyroModel
+        ):
+            pass
+
+        train_X, train_Y, train_Yvar = self._get_base_data(**tkwargs)
+
+        pyro_model = MultitaskMaternPyroModel()
+        model = FullyBayesianMultiTaskGP(
+            train_X=train_X,
+            train_Y=train_Y,
+            train_Yvar=train_Yvar,
+            task_feature=4,
+            pyro_model=pyro_model,
+        )
+        self.assertIsInstance(model.pyro_model, MultiTaskPyroMixin)
+
+        fit_fully_bayesian_model_nuts(
+            model, warmup_steps=8, num_samples=5, thinning=2, disable_progbar=True
+        )
+        state_dict = model.state_dict()
+        test_X = torch.rand(3, 4, **tkwargs)
+        preds1 = model.posterior(test_X)
+
+        pyro_model2 = MultitaskMaternPyroModel()
+        m_new = FullyBayesianMultiTaskGP(
+            train_X=train_X,
+            train_Y=train_Y,
+            train_Yvar=train_Yvar,
+            task_feature=4,
+            pyro_model=pyro_model2,
+        )
+        m_new.load_state_dict(state_dict)
+
+        preds2 = m_new.posterior(test_X)
+        self.assertTrue(torch.equal(preds1.mean, preds2.mean))
+        self.assertTrue(torch.equal(preds1.variance, preds2.variance))
+
 
 class TestPyroModelMultitaskMixin(BotorchTestCase):
     """Tests for the MultiTaskPyroMixin and LatentFeatureMultiTaskPyroMixin classes."""