Enable unobserved task support in MultiTaskGP (meta-pytorch#3145)

Summary:
Pull Request resolved: meta-pytorch#3145

Permits an MTGP to predict on an unobserved task, addressing these issues:
meta-pytorch#2360
meta-pytorch#3085
To do this, we assume that the unobserved task is maximally correlated with the target tasks (equally with each, by averaging the elements). Exact heuristic on correlation is definitely up for discussion, but this seems like a decent default assumption.
Will come in handy for TL initialization.

Differential Revision: D90769576

Author: Carl Hvarfner

botorch/models/fully_bayesian_multitask.py

@@ -7,7 +7,7 @@
 r"""Multi-task Gaussian Process Regression models with fully Bayesian inference."""
 
 from collections.abc import Mapping
-from typing import Any, NoReturn, TypeVar
+from typing import Any, NoReturn, Self, TypeVar
 
 import pyro
 import torch
@@ -19,7 +19,10 @@
     reshape_and_detach,
     SaasPyroModel,
 )
-from botorch.models.gpytorch import BatchedMultiOutputGPyTorchModel
+from botorch.models.gpytorch import (
+    BatchedMultiOutputGPyTorchModel,
+    MultiTaskGPyTorchModel,
+)
 from botorch.models.multitask import MultiTaskGP
 from botorch.models.transforms.input import InputTransform
 from botorch.models.transforms.outcome import OutcomeTransform
@@ -55,6 +58,7 @@ def set_inputs(
         train_Yvar: Tensor | None,
         task_feature: int,
         task_rank: int | None = None,
+        all_tasks: list[int] | None = None,
     ) -> None:
         """Set the training data.
 
@@ -73,7 +77,11 @@ def set_inputs(
         task_feature = task_feature % train_X.shape[-1]
         super().set_inputs(train_X, train_Y, train_Yvar)
         # obtain a list of task indicies
-        all_tasks = train_X[:, task_feature].unique().to(dtype=torch.long).tolist()
+        all_tasks = (
+            train_X[:, task_feature].unique().to(dtype=torch.long).tolist()
+            if all_tasks is None
+            else all_tasks
+        )
         self.task_feature = task_feature
         self.num_tasks = len(all_tasks)
         self.task_rank = task_rank or self.num_tasks
@@ -242,7 +250,10 @@ 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!
+            all_tasks: A list of all task indices. If omitted, all tasks will be
+                inferred from the task feature column of the training data. Used to
+                inform the model about the total number of tasks, including any
+                unobserved tasks.
             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
@@ -310,6 +321,7 @@ def __init__(
             train_Yvar=train_Yvar,
             task_feature=task_feature,
             task_rank=self._rank,
+            all_tasks=all_tasks,
         )
         self.pyro_model: MultitaskSaasPyroModel = pyro_model
         if outcome_transform is not None:
@@ -383,6 +395,20 @@ def load_mcmc_samples(self, mcmc_samples: dict[str, Tensor]) -> None:
             _,
         ) = self.pyro_model.load_mcmc_samples(mcmc_samples=mcmc_samples)
 
+    def eval(self) -> Self:
+        r"""Puts the model in eval mode.
+
+        Circumvents the need to call MultiTaskGP.eval(), which computes the
+        task_covar_matrix for non-observed tasks. This is not needed for fully
+        Bayesian models, since the non-observed tasks' covar factors are instead
+        sampled.
+
+        Returns:
+            The model itself.
+        """
+        self._check_if_fitted()
+        return MultiTaskGPyTorchModel.eval(self)
+
     def posterior(
         self,
         X: Tensor,

botorch/models/multitask.py

@@ -30,7 +30,7 @@
 from __future__ import annotations
 
 import math
-from typing import Any
+from typing import Any, Self
 
 import torch
 from botorch.acquisition.objective import PosteriorTransform
@@ -238,7 +238,24 @@ def __init__(
                 "This is not allowed as it will lead to errors during model training."
             )
         all_tasks = all_tasks or all_tasks_inferred
-        self.num_tasks = len(all_tasks_inferred)
+        # Compute observed and unobserved task indices when all_tasks includes
+        # unobserved tasks
+        sorted_all_tasks = sorted(all_tasks)
+        if set(all_tasks) != set(all_tasks_inferred):
+            observed_set = set(all_tasks_inferred)
+            self._observed_task_indices = torch.tensor(
+                [i for i, t in enumerate(sorted_all_tasks) if t in observed_set],
+                dtype=torch.long,
+            )
+            self._unobserved_task_indices = torch.tensor(
+                [i for i, t in enumerate(sorted_all_tasks) if t not in observed_set],
+                dtype=torch.long,
+            )
+        else:
+            # All tasks are observed - set observed indices to all tasks
+            self._observed_task_indices = torch.arange(len(all_tasks), dtype=torch.long)
+            self._unobserved_task_indices = torch.tensor([], dtype=torch.long)
+        self.num_tasks = len(all_tasks)
         if outcome_transform == DEFAULT:
             outcome_transform = Standardize(m=1, batch_shape=train_X.shape[:-2])
         if outcome_transform is not None:
@@ -321,7 +338,7 @@ def __init__(
             default_task_value=None if output_tasks is None else output_tasks[0],
         )
         self.register_buffer("_task_mapper", task_mapper)
-        self._expected_task_values = set(all_tasks_inferred)
+        self._expected_task_values = set(all_tasks)
         if input_transform is not None:
             self.input_transform = input_transform
         if outcome_transform is not None:
@@ -407,6 +424,28 @@ def forward(self, x: Tensor) -> MultivariateNormal:
         covar_x = self.covar_module(x_covar)
         return MultivariateNormal(mean_x, covar_x)
 
+    def eval(self) -> Self:
+        r"""Puts the model in ``eval`` mode.
+
+        When unobserved tasks are present (i.e., ``all_tasks`` includes tasks not in
+        the training data), this method sets the covariance factor for unobserved tasks
+        to the mean of the observed tasks' covariance factors. This provides a
+        reasonable initialization for prediction on unobserved tasks.
+        """
+        if len(self._unobserved_task_indices) > 0:
+            task_covar_module = self.covar_module.kernels[1]
+            # Get the current covar_factor (transformed from raw_covar_factor)
+            covar_factor = task_covar_module.covar_factor
+            # Compute mean of observed tasks' covar_factor rows
+            observed_covar_factor = covar_factor[self._observed_task_indices]
+            mean_covar_factor = observed_covar_factor.mean(dim=0)
+            # Create new covar_factor with unobserved tasks set to mean
+            new_covar_factor = covar_factor.clone()
+            new_covar_factor[self._unobserved_task_indices] = mean_covar_factor
+            # Set the new covar_factor (this applies inverse_transform internally)
+            task_covar_module._set_covar_factor(new_covar_factor)
+        return super().eval()
+
     @classmethod
     def get_all_tasks(
         cls,

test/models/test_fully_bayesian_multitask.py

@@ -529,6 +529,8 @@ def test_fit_model_with_task_mapper(self) -> None:
         self.assertTrue(
             torch.equal(model._task_mapper, torch.tensor([0, 1, 1], **tkwargs))
         )
+        # Verify the pyro_model has the correct number of tasks (3, not 2)
+        self.assertEqual(model.pyro_model.num_tasks, 3)
         self.test_fit_model(
             use_outcome_transform=True,
             all_tasks=all_tasks,

test/models/test_multitask.py

@@ -445,9 +445,12 @@ def test_all_tasks_input(self) -> None:
         model = MultiTaskGP(
             train_X=train_X, train_Y=train_Y, task_feature=0, all_tasks=[0, 1, 2, 3]
         )
-        self.assertEqual(model.num_tasks, 2)
+        self.assertEqual(model.num_tasks, 4)
         # Check that PositiveIndexKernel knows of all tasks.
-        self.assertEqual(model.covar_module.kernels[1].raw_covar_factor.shape[0], 2)
+        self.assertEqual(model.covar_module.kernels[1].raw_covar_factor.shape[0], 4)
+        # Check that observed and unobserved task indices are computed correctly.
+        self.assertEqual(model._observed_task_indices.tolist(), [0, 1])
+        self.assertEqual(model._unobserved_task_indices.tolist(), [2, 3])
 
     def test_MultiTaskGP_construct_inputs(self) -> None:
         for dtype, fixed_noise, skip_task_features_in_datasets in zip(
@@ -540,13 +543,98 @@ def test_validatation_of_task_values(self) -> None:
             validate_task_values=True,
         )
 
+        # Task 2 is in all_tasks, so it should be valid even with validation enabled
+        self.assertTrue(
+            torch.equal(
+                torch.tensor([1], **tkwargs),
+                model._map_tasks(task_values=torch.tensor([2], **tkwargs)),
+            )
+        )
+
+        # Task 3 is NOT in all_tasks, so it should raise an error
         with self.assertRaisesRegex(
             ValueError,
             "Received invalid raw task values. Expected raw value to be in"
-            r" \{0, 1\}, but got unexpected task"
-            r" values: \{2\}.",
+            r" \{0, 1, 2\}, but got unexpected task"
+            r" values: \{3\}.",
         ):
-            model._map_tasks(task_values=torch.tensor([2], **tkwargs))
+            model._map_tasks(task_values=torch.tensor([3], **tkwargs))
+
+    def test_multitask_gp_unobserved_tasks(self) -> None:
+        """Test MultiTaskGP with unobserved tasks.
+
+        This test verifies that:
+        1. Creating a model with all_tasks including unobserved tasks works
+        2. In train mode, unobserved task covar_factor is at random initialization
+        3. In eval mode, unobserved task covar_factor is set to mean of observed
+        4. Predictions work for the unobserved task
+        """
+        tkwargs = {"device": self.device, "dtype": torch.double}
+
+        # Create data for tasks 0 and 2 only (task 1 is unobserved)
+        _, (train_X, train_Y, _) = gen_multi_task_dataset(task_values=[0, 2], **tkwargs)
+
+        # Create model with all_tasks=[0, 1, 2] including unobserved task 1
+        model = MultiTaskGP(
+            train_X=train_X,
+            train_Y=train_Y,
+            task_feature=0,
+            all_tasks=[0, 1, 2],
+        )
+        model.to(**tkwargs)
+
+        # Verify model.num_tasks == 3
+        self.assertEqual(model.num_tasks, 3)
+
+        # Verify observed and unobserved task indices are correctly set
+        self.assertEqual(model._observed_task_indices.tolist(), [0, 2])
+        self.assertEqual(model._unobserved_task_indices.tolist(), [1])
+
+        # Get the task covariance module
+        task_covar_module = model.covar_module.kernels[1]
+
+        # In train mode, get the covar_factor for unobserved task (index 1)
+        model.train()
+        train_covar_factor = task_covar_module.covar_factor.clone()
+        unobserved_train_covar = train_covar_factor[1]
+        observed_train_covar = train_covar_factor[[0, 2]]
+        mean_observed_train = observed_train_covar.mean(dim=0)
+
+        # Unobserved task covar_factor should be at random init in train mode
+        # (very unlikely to be exactly equal to mean of observed)
+        self.assertFalse(
+            torch.allclose(unobserved_train_covar, mean_observed_train, atol=1e-6)
+        )
+
+        # Switch to eval mode
+        model.eval()
+
+        # In eval mode, get the covar_factor for unobserved task
+        eval_covar_factor = task_covar_module.covar_factor.clone()
+        unobserved_eval_covar = eval_covar_factor[1]
+        observed_eval_covar = eval_covar_factor[[0, 2]]
+        mean_observed_eval = observed_eval_covar.mean(dim=0)
+
+        # Unobserved task covar_factor should equal mean of observed in eval mode
+        self.assertTrue(
+            torch.allclose(unobserved_eval_covar, mean_observed_eval, atol=1e-6)
+        )
+
+        # Verify predictions work for the unobserved task
+        # Create test input for unobserved task (task 1)
+        test_X = torch.rand(3, 2, **tkwargs)
+        test_X[:, 0] = 1.0  # Set task feature to 1 (unobserved task)
+
+        with torch.no_grad():
+            posterior = model.posterior(X=test_X)
+
+        # Verify posterior has expected shape
+        self.assertEqual(posterior.mean.shape, torch.Size([3, 1]))
+        self.assertEqual(posterior.variance.shape, torch.Size([3, 1]))
+
+        # Verify we can sample from the posterior
+        samples = posterior.rsample(sample_shape=torch.Size([2]))
+        self.assertEqual(samples.shape, torch.Size([2, 3, 1]))
 
 
 class TestKroneckerMultiTaskGP(BotorchTestCase):