Support custom all_tasks for MTGPs (#2271)

Summary:
Pull Request resolved: #2271

This allows creation of MTGPs that support inference from tasks that don't appear in the training data. See #2265 for some discussion on how the task covariance behaves in the absence of task specific data.

Reviewed By: esantorella

Differential Revision: D53029681

fbshipit-source-id: 3df8c910ff03c828ee0a317c29a5067d90e7f769

Author: saitcakmak

botorch/models/contextual_multioutput.py

@@ -45,6 +45,7 @@ def __init__(
         context_emb_feature: Optional[Tensor] = None,
         embs_dim_list: Optional[List[int]] = None,
         output_tasks: Optional[List[int]] = None,
+        all_tasks: Optional[List[int]] = None,
         input_transform: Optional[InputTransform] = None,
         outcome_transform: Optional[OutcomeTransform] = None,
     ) -> None:
@@ -67,22 +68,31 @@ def __init__(
                 for each categorical variable.
             output_tasks: A list of task indices for which to compute model
                 outputs for. If omitted, return outputs for all task indices.
-
+            all_tasks: By default, multi-task GPs infer the list of all tasks from
+                the task features in `train_X`. This is an experimental feature that
+                enables creation of multi-task GPs with tasks that don't appear in the
+                training data. Note that when a task is not observed, the corresponding
+                task covariance will heavily depend on random initialization and may
+                behave unexpectedly.
         """
         super().__init__(
             train_X=train_X,
             train_Y=train_Y,
             task_feature=task_feature,
             train_Yvar=train_Yvar,
             output_tasks=output_tasks,
+            all_tasks=all_tasks,
             input_transform=input_transform,
             outcome_transform=outcome_transform,
         )
         self.device = train_X.device
-        #  context indices
-        all_tasks = train_X[:, task_feature].unique()
-        self.all_tasks = all_tasks.to(dtype=torch.long).tolist()
-        self.all_tasks.sort()  # unique in python does automatic sort; add for safety
+        if all_tasks is None:
+            all_tasks = train_X[:, task_feature].unique()
+            self.all_tasks = all_tasks.to(dtype=torch.long).tolist()
+        else:
+            all_tasks = torch.tensor(all_tasks, dtype=torch.long)
+            self.all_tasks = all_tasks
+        self.all_tasks.sort()  # These are the context indices.
 
         if context_cat_feature is None:
             context_cat_feature = all_tasks.unsqueeze(-1).to(device=self.device)

botorch/models/fully_bayesian_multitask.py

@@ -8,23 +8,21 @@
 """
 
 
-from typing import Any, Dict, List, Mapping, NoReturn, Optional, Tuple, Union
+from typing import Any, Dict, List, Mapping, NoReturn, Optional, Tuple
 
 import pyro
 import torch
 from botorch.acquisition.objective import PosteriorTransform
 from botorch.models.fully_bayesian import (
     matern52_kernel,
     MIN_INFERRED_NOISE_LEVEL,
-    PyroModel,
     reshape_and_detach,
     SaasPyroModel,
 )
 from botorch.models.multitask import MultiTaskGP
 from botorch.models.transforms.input import InputTransform
 from botorch.models.transforms.outcome import OutcomeTransform
 from botorch.posteriors.fully_bayesian import GaussianMixturePosterior, MCMC_DIM
-from botorch.utils.datasets import MultiTaskDataset, SupervisedDataset
 from gpytorch.distributions.multivariate_normal import MultivariateNormal
 from gpytorch.kernels import MaternKernel
 from gpytorch.kernels.kernel import Kernel
@@ -200,9 +198,10 @@ def __init__(
         train_Yvar: Optional[Tensor] = None,
         output_tasks: Optional[List[int]] = None,
         rank: Optional[int] = None,
+        all_tasks: Optional[List[int]] = None,
         outcome_transform: Optional[OutcomeTransform] = None,
         input_transform: Optional[InputTransform] = None,
-        pyro_model: Optional[PyroModel] = None,
+        pyro_model: Optional[MultitaskSaasPyroModel] = None,
     ) -> None:
         r"""Initialize the fully Bayesian multi-task GP model.
 
@@ -216,13 +215,15 @@ def __init__(
                 outputs for. If omitted, return outputs for all task indices.
             rank: The num of learned task embeddings to be used in the task kernel.
                 If omitted, use a full rank (i.e. number of tasks) kernel.
+            all_tasks: NOT SUPPORTED!
             outcome_transform: An outcome transform that is applied to the
                 training data during instantiation and to the posterior during
                 inference (that is, the `Posterior` obtained by calling
                 `.posterior` on the model will be on the original scale).
             input_transform: An input transform that is applied to the inputs `X`
                 in the model's forward pass.
-            pyro_model: Optional `PyroModel`, defaults to `MultitaskSaasPyroModel`.
+            pyro_model: Optional `PyroModel` that has the same signature as
+                `MultitaskSaasPyroModel`. Defaults to `MultitaskSaasPyroModel`.
         """
         if not (
             train_X.ndim == train_Y.ndim == 2
@@ -253,6 +254,12 @@ def __init__(
             output_tasks=output_tasks,
             rank=rank,
         )
+        if all_tasks is not None and self._expected_task_values != set(all_tasks):
+            raise NotImplementedError(
+                "The `all_tasks` argument is not supported by SAAS MTGP. "
+                f"The training data includes tasks {self._expected_task_values}, "
+                f"got {all_tasks=}."
+            )
         self.to(train_X)
 
         self.mean_module = None
@@ -383,29 +390,6 @@ def forward(self, X: Tensor) -> MultivariateNormal:
         covar = covar_x.mul(covar_i)
         return MultivariateNormal(mean_x, covar)
 
-    @classmethod
-    def construct_inputs(
-        cls,
-        training_data: Union[SupervisedDataset, MultiTaskDataset],
-        task_feature: int,
-        rank: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Dict[str, Any]:
-        r"""Construct `Model` keyword arguments from a dataset and other args.
-
-        Args:
-            training_data: A `SupervisedDataset` or a `MultiTaskDataset`.
-            task_feature: Column index of embedded task indicator features.
-            rank: The rank of the cross-task covariance matrix.
-        """
-        inputs = super().construct_inputs(
-            training_data=training_data, task_feature=task_feature, rank=rank, **kwargs
-        )
-        inputs.pop("task_covar_prior")
-        if "train_Yvar" not in inputs:
-            inputs["train_Yvar"] = None
-        return inputs
-
     def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
         r"""Custom logic for loading the state dict.
 

botorch/models/multitask.py

@@ -149,6 +149,7 @@ def __init__(
         task_covar_prior: Optional[Prior] = None,
         output_tasks: Optional[List[int]] = None,
         rank: Optional[int] = None,
+        all_tasks: Optional[List[int]] = None,
         input_transform: Optional[InputTransform] = None,
         outcome_transform: Optional[OutcomeTransform] = None,
     ) -> None:
@@ -176,6 +177,12 @@ def __init__(
                 full rank (i.e. number of tasks) kernel.
             task_covar_prior : A Prior on the task covariance matrix. Must operate
                 on p.s.d. matrices. A common prior for this is the `LKJ` prior.
+            all_tasks: By default, multi-task GPs infer the list of all tasks from
+                the task features in `train_X`. This is an experimental feature that
+                enables creation of multi-task GPs with tasks that don't appear in the
+                training data. Note that when a task is not observed, the corresponding
+                task covariance will heavily depend on random initialization and may
+                behave unexpectedly.
             input_transform: An input transform that is applied in the model's
                 forward pass.
             outcome_transform: An outcome transform that is applied to the
@@ -197,9 +204,12 @@ def __init__(
                 X=train_X, input_transform=input_transform
             )
         self._validate_tensor_args(X=transformed_X, Y=train_Y, Yvar=train_Yvar)
-        all_tasks, task_feature, self.num_non_task_features = self.get_all_tasks(
-            transformed_X, task_feature, output_tasks
-        )
+        (
+            all_tasks_inferred,
+            task_feature,
+            self.num_non_task_features,
+        ) = self.get_all_tasks(transformed_X, task_feature, output_tasks)
+        all_tasks = all_tasks or all_tasks_inferred
         self.num_tasks = len(all_tasks)
         if outcome_transform is not None:
             train_Y, train_Yvar = outcome_transform(Y=train_Y, Yvar=train_Yvar)
@@ -360,13 +370,16 @@ def construct_inputs(
         base_inputs = super().construct_inputs(
             training_data=training_data, task_feature=task_feature, **kwargs
         )
-        return {
-            **base_inputs,
-            "task_feature": task_feature,
-            "output_tasks": output_tasks,
-            "task_covar_prior": task_covar_prior,
-            "rank": rank,
-        }
+        if isinstance(training_data, MultiTaskDataset):
+            all_tasks = list(range(len(training_data.datasets)))
+            base_inputs["all_tasks"] = all_tasks
+        if task_covar_prior is not None:
+            base_inputs["task_covar_prior"] = task_covar_prior
+        if rank is not None:
+            base_inputs["rank"] = rank
+        base_inputs["task_feature"] = task_feature
+        base_inputs["output_tasks"] = output_tasks
+        return base_inputs
 
 
 class FixedNoiseMultiTaskGP(MultiTaskGP):
@@ -428,6 +441,7 @@ def __init__(
             "When `train_Yvar` is specified, `MultiTaskGP` behaves the same "
             "as the `FixedNoiseMultiTaskGP`.",
             DeprecationWarning,
+            stacklevel=2,
         )
         super().__init__(
             train_X=train_X,

test/models/test_fully_bayesian_multitask.py

@@ -583,7 +583,10 @@ def test_construct_inputs(self):
             )
             self.assertTrue(torch.equal(data_dict["train_X"], train_X))
             self.assertTrue(torch.equal(data_dict["train_Y"], train_Y))
-            self.assertAllClose(data_dict["train_Yvar"], train_Yvar)
+            if train_Yvar is not None:
+                self.assertAllClose(data_dict["train_Yvar"], train_Yvar)
+            else:
+                self.assertNotIn("train_Yvar", data_dict)
             self.assertEqual(data_dict["task_feature"], task_feature)
             self.assertEqual(data_dict["rank"], 1)
             self.assertTrue("task_covar_prior" not in data_dict)