Support SAAS ensemble models in RFFs (#1530)

Summary:
Pull Request resolved: #1530

RFFs currently support batched models but not batched evaluation (without the model batch dimension unsqueezed) of RFFs constructed from batched models. This prevents their use in acquisition functions. The changes here make it possible to construct an RFF from a SAAS ensemble model and use it in acquisition functions.

Reviewed By: Balandat

Differential Revision: D41569212

fbshipit-source-id: 6d371a0c4b1ab3be24f7e790b3527cc5f29e6808

Author: saitcakmak

botorch/utils/gp_sampling.py

@@ -17,6 +17,7 @@
 from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.models.multitask import MultiTaskGP
 from botorch.utils.sampling import manual_seed
+from botorch.utils.transforms import is_fully_bayesian
 from gpytorch.kernels import Kernel, MaternKernel, RBFKernel, ScaleKernel
 from linear_operator.utils.cholesky import psd_safe_cholesky
 from torch import Tensor
@@ -228,7 +229,20 @@ def forward(self, X: Tensor) -> Tensor:
             a `batch_shape`, the output `batch_shape` will be
             `(sample_shape) x (kernel_batch_shape)`.
         """
-        self._check_forward_X_shape_compatibility(X)
+        try:
+            self._check_forward_X_shape_compatibility(X)
+        except ValueError as e:
+            # A workaround to support batched SAAS models.
+            # TODO: Support batch evaluation of multi-sample RFFs as well.
+            # Multi-sample RFFs have input batch as the 0-th dimension,
+            # which is different than other posteriors which would have
+            # the sample shape as the 0-th dimension.
+            if len(self.kernel_batch_shape) == 1:
+                X = X.unsqueeze(-3)
+                self._check_forward_X_shape_compatibility(X)
+            else:
+                raise e
+
         # X is of shape (additional_batch_shape) x (sample_shape)
         # x (kernel_batch_shape) x n x d.
         # Weights is of shape (sample_shape) x (kernel_batch_shape) x d x num_rff.
@@ -489,6 +503,7 @@ def get_gp_samples(
                 models[m].outcome_transform = _octf
             if _intf is not None:
                 base_gp_samples.models[m].input_transform = _intf
+        base_gp_samples.is_fully_bayesian = is_fully_bayesian(model=model)
         return base_gp_samples
     elif n_samples > 1:
         base_gp_samples = get_deterministic_model_multi_samples(
@@ -507,4 +522,5 @@ def get_gp_samples(
     if octf is not None:
         base_gp_samples.outcome_transform = octf
         model.outcome_transform = octf
+    base_gp_samples.is_fully_bayesian = is_fully_bayesian(model=model)
     return base_gp_samples

botorch/utils/transforms.py

@@ -192,11 +192,17 @@ def is_fully_bayesian(model: Model) -> bool:
         SaasFullyBayesianMultiTaskGP,
     ]
 
-    if any(isinstance(model, m_cls) for m_cls in full_bayesian_model_cls):
+    if any(
+        isinstance(model, m_cls) or getattr(model, "is_fully_bayesian", False)
+        for m_cls in full_bayesian_model_cls
+    ):
         return True
     elif isinstance(model, ModelList):
         for m in model.models:
-            if any(isinstance(m, m_cls) for m_cls in full_bayesian_model_cls):
+            if any(
+                isinstance(m, m_cls) or getattr(model, "is_fully_bayesian", False)
+                for m_cls in full_bayesian_model_cls
+            ):
                 return True
             elif isinstance(m, ModelListGP) and any(
                 isinstance(m_sub, m_cls)

test/utils/test_gp_sampling.py

@@ -11,6 +11,7 @@
 import torch
 from botorch.models.converter import batched_to_model_list
 from botorch.models.deterministic import DeterministicModel
+from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
 from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP
 from botorch.models.model import ModelList
 from botorch.models.multitask import MultiTaskGP
@@ -26,6 +27,7 @@
     RandomFourierFeatures,
 )
 from botorch.utils.testing import BotorchTestCase
+from botorch.utils.transforms import is_fully_bayesian
 from gpytorch.kernels import MaternKernel, PeriodicKernel, RBFKernel, ScaleKernel
 from torch.distributions import MultivariateNormal
 
@@ -661,3 +663,31 @@ def test_with_fixed_noise(self):
                 torch.Size([2, 1]) if n_samples == 1 else torch.Size([n_samples, 2, 1])
             )
             self.assertEqual(samples.shape, expected_shape)
+
+    def test_with_saas_models(self):
+        # Construct a SAAS model.
+        tkwargs = {"dtype": torch.double, "device": self.device}
+        num_samples = 4
+        model = SaasFullyBayesianSingleTaskGP(
+            train_X=torch.rand(10, 4, **tkwargs), train_Y=torch.randn(10, 1, **tkwargs)
+        )
+        mcmc_samples = {
+            "lengthscale": torch.rand(num_samples, 1, 4, **tkwargs),
+            "outputscale": torch.rand(num_samples, **tkwargs),
+            "mean": torch.randn(num_samples, **tkwargs),
+            "noise": torch.rand(num_samples, 1, **tkwargs),
+        }
+        model.load_mcmc_samples(mcmc_samples)
+        # Test proper setup & sampling support.
+        gp_samples = get_gp_samples(
+            model=model,
+            num_outputs=1,
+            n_samples=1,
+        )
+        self.assertTrue(is_fully_bayesian(gp_samples))
+        # Non-batch evaluation.
+        samples = gp_samples(torch.rand(2, 4, **tkwargs))
+        self.assertEqual(samples.shape, torch.Size([4, 2, 1]))
+        # Batch evaluation.
+        samples = gp_samples(torch.rand(5, 2, 4, **tkwargs))
+        self.assertEqual(samples.shape, torch.Size([5, 4, 2, 1]))