Return TorchPosterior for non-MVN distributions in ApproximateGPyTorchModel (#3242)

Summary:
When using SingleTaskVariationalGP with non-Gaussian likelihoods (e.g., BetaLikelihood), the posterior distribution is not MultivariateNormal. Now returns TorchPosterior instead of GPyTorchPosterior in that case, since GPyTorchPosterior's methods assume MVN.

Closes #3066

Pull Request resolved: #3242

Reviewed By: esantorella

Differential Revision: D97317429

Pulled By: saitcakmak

fbshipit-source-id: 71792043c3bea6051b7d7dca6329bc15faeeaedb

Author: saitcakmak

botorch/models/approximate_gp.py

@@ -49,6 +49,7 @@
     InducingPointAllocator,
 )
 from botorch.posteriors.gpytorch import GPyTorchPosterior
+from botorch.posteriors.torch import TorchPosterior
 from gpytorch.distributions import MultivariateNormal
 from gpytorch.kernels import Kernel
 from gpytorch.likelihoods import (
@@ -150,7 +151,7 @@ def posterior(
         output_indices: list[int] | None = None,
         observation_noise: bool = False,
         posterior_transform: PosteriorTransform | None = None,
-    ) -> GPyTorchPosterior:
+    ) -> TorchPosterior:
         if output_indices is not None:
             raise NotImplementedError(  # pragma: no cover
                 f"{self.__class__.__name__}.posterior does not support output indices."
@@ -170,7 +171,10 @@ def posterior(
         if observation_noise:
             dist = self.likelihood(dist)
 
-        posterior = GPyTorchPosterior(distribution=dist)
+        if isinstance(dist, MultivariateNormal):
+            posterior = GPyTorchPosterior(distribution=dist)
+        else:
+            posterior = TorchPosterior(distribution=dist)
         if hasattr(self, "outcome_transform"):
             posterior = self.outcome_transform.untransform_posterior(posterior, X=X)
         if posterior_transform is not None:

test/models/test_approximate_gp.py

@@ -23,8 +23,13 @@
     GreedyVarianceReduction,
 )
 from botorch.posteriors import GPyTorchPosterior, TransformedPosterior
+from botorch.posteriors.torch import TorchPosterior
 from botorch.utils.testing import BotorchTestCase
-from gpytorch.likelihoods import GaussianLikelihood, MultitaskGaussianLikelihood
+from gpytorch.likelihoods import (
+    BetaLikelihood,
+    GaussianLikelihood,
+    MultitaskGaussianLikelihood,
+)
 from gpytorch.mlls import VariationalELBO
 from gpytorch.variational import (
     IndependentMultitaskVariationalStrategy,
@@ -342,6 +347,32 @@ def test_custom_inducing_point_init(self):
         self.assertAllClose(model_1_inducing, model_2_inducing)
         self.assertFalse(model_1_inducing[0, 0] == model_3_inducing[0, 0])
 
+    def test_non_gaussian_likelihood_posterior(self) -> None:
+        """Test that non-Gaussian likelihoods return TorchPosterior."""
+        train_X = torch.rand(10, 1, device=self.device)
+        model = SingleTaskVariationalGP(
+            train_X=train_X,
+            likelihood=BetaLikelihood(),
+        ).to(self.device)
+        test_X = torch.rand(5, 1, device=self.device)
+
+        # Without observation noise, the distribution is MVN (from the GP),
+        # so it should return GPyTorchPosterior.
+        posterior = model.posterior(test_X, observation_noise=False)
+        self.assertIsInstance(posterior, GPyTorchPosterior)
+
+        # With observation noise, the likelihood transforms the MVN into a
+        # Beta distribution, so it should return TorchPosterior.
+        posterior = model.posterior(test_X, observation_noise=True)
+        self.assertIsInstance(posterior, TorchPosterior)
+        self.assertNotIsInstance(posterior, GPyTorchPosterior)
+
+        # Verify that sampling works with the TorchPosterior and that the
+        # shape doesn't have a spurious trailing dim from GPyTorchPosterior.
+        # GPyTorchPosterior.rsample would unsqueeze(-1), adding an extra dim.
+        samples = posterior.rsample(sample_shape=torch.Size([2]))
+        self.assertEqual(samples.shape, torch.Size([2, 10, 5]))
+
     def test_input_transform(self) -> None:
         train_X = torch.linspace(1, 3, 10, dtype=torch.double)[:, None]
         y = -3 * train_X + 5