Use Standardize/Normalize by default for SingleTaskGP (#2458)

Summary:
X-link: facebook/Ax#2630

Pull Request resolved: #2458

D60080819 recently updated the default `SingleTaskGP` BoTorch priors. One significant change was to remove the use of an outputscale, which may not work well if the outputs aren't standardized. This diff changes the `SingleTaskGP` to use `Standardize` and `Normalize` by default if no input/outcome transforms are specified (this allows users to explicitly pass in `None` if they don't want to use any transforms).

Differential Revision: D60492937

Author: David Eriksson

botorch/acquisition/analytic.py

@@ -1091,15 +1091,17 @@ def _get_noiseless_fantasy_model(
     # are used across all batches (by default, a GP with batched training data
     # uses independent hyperparameters for each batch).
 
-    # Don't apply `outcome_transform` and `input_transform` here,
-    # since the data being passed has already been transformed.
-    # So we will instead set them afterwards.
+    # We don't want to use the true `outcome_transform` and `input_transform` here
+    # since the data being passed has already been transformed. We thus pass `None`
+    # and will instead set them afterwards.
     fantasy_model = SingleTaskGP(
         train_X=model.train_inputs[0],
         train_Y=model.train_targets.unsqueeze(-1),
         train_Yvar=model.likelihood.noise_covar.noise.unsqueeze(-1),
         covar_module=deepcopy(model.covar_module),
         mean_module=deepcopy(model.mean_module),
+        outcome_transform=None,
+        input_transform=None,
     )
 
     Yvar = torch.full_like(Y_fantasized, 1e-7)

botorch/acquisition/multi_objective/max_value_entropy_search.py

@@ -64,7 +64,7 @@ class qMultiObjectiveMaxValueEntropy(
         _default_sample_shape: The `sample_shape` for the default sampler.
 
     Example:
-        >>> model = SingleTaskGP(train_X, train_Y)
+        >>> model = SingleTaskGP(train_X, train_Y, outcome_transform=None)
         >>> MESMO = qMultiObjectiveMaxValueEntropy(model, sample_pfs)
         >>> mesmo = MESMO(test_X)
     """

botorch/models/contextual.py

@@ -102,7 +102,13 @@ def __init__(
                 dimension is set to 1 for each categorical variable.
             context_weight_dict: Known population weights of each context.
         """
-        super().__init__(train_X=train_X, train_Y=train_Y, train_Yvar=train_Yvar)
+        super().__init__(
+            train_X=train_X,
+            train_Y=train_Y,
+            train_Yvar=train_Yvar,
+            input_transform=None,
+            outcome_transform=None,
+        )
         self.covar_module = LCEAKernel(
             decomposition=decomposition,
             batch_shape=self._aug_batch_shape,

botorch/models/converter.py

@@ -17,6 +17,7 @@
 import torch
 from botorch.exceptions import UnsupportedError
 from botorch.exceptions.warnings import BotorchWarning
+from botorch.models import SingleTaskGP
 from botorch.models.gp_regression import HeteroskedasticSingleTaskGP
 from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
 from botorch.models.gp_regression_mixed import MixedSingleTaskGP
@@ -179,6 +180,8 @@ def model_list_to_batched(model_list: ModelListGP) -> BatchedMultiOutputGPyTorch
         batch_length = len(models)
         covar_module = _batched_kernel(models[0].covar_module, batch_length)
         kwargs["covar_module"] = covar_module
+    if isinstance(models[0], SingleTaskGP):
+        kwargs["outcome_transform"] = None
 
     # construct the batched GP model
     input_transform = getattr(models[0], "input_transform", None)
@@ -331,6 +334,8 @@ def batched_to_model_list(batch_model: BatchedMultiOutputGPyTorchModel) -> Model
             )
         if isinstance(batch_model, SingleTaskMultiFidelityGP):
             kwargs.update(batch_model._init_args)
+        if isinstance(batch_model, SingleTaskGP):
+            kwargs["outcome_transform"] = None
         # NOTE: Adding outcome transform to kwargs to avoid the multiple
         # values for same kwarg issue with SingleTaskMultiFidelityGP.
         if outcome_transform is not None:
@@ -418,6 +423,8 @@ def batched_multi_output_to_single_output(
         kwargs["train_Yvar"] = noise_covar.noise.clone().unsqueeze(-1)
     if isinstance(batch_mo_model, SingleTaskMultiFidelityGP):
         kwargs.update(batch_mo_model._init_args)
+    if isinstance(batch_mo_model, SingleTaskGP):
+        kwargs["outcome_transform"] = None
     single_outcome_model = batch_mo_model.__class__(
         input_transform=input_transform, **kwargs
     )

botorch/models/gp_regression.py

@@ -36,8 +36,8 @@
 import torch
 from botorch.models.gpytorch import BatchedMultiOutputGPyTorchModel
 from botorch.models.model import FantasizeMixin
-from botorch.models.transforms.input import InputTransform
-from botorch.models.transforms.outcome import Log, OutcomeTransform
+from botorch.models.transforms.input import InputTransform, Normalize
+from botorch.models.transforms.outcome import Log, OutcomeTransform, Standardize
 from botorch.models.utils import validate_input_scaling
 from botorch.models.utils.gpytorch_modules import (
     get_covar_module_with_dim_scaled_prior,
@@ -46,6 +46,7 @@
 )
 from botorch.utils.containers import BotorchContainer
 from botorch.utils.datasets import SupervisedDataset
+from botorch.utils.types import _DefaultType, DEFAULT
 from gpytorch.constraints.constraints import GreaterThan
 from gpytorch.distributions.multivariate_normal import MultivariateNormal
 from gpytorch.likelihoods.gaussian_likelihood import (
@@ -134,8 +135,8 @@ def __init__(
         likelihood: Optional[Likelihood] = None,
         covar_module: Optional[Module] = None,
         mean_module: Optional[Mean] = None,
-        outcome_transform: Optional[OutcomeTransform] = None,
-        input_transform: Optional[InputTransform] = None,
+        outcome_transform: Optional[Union[OutcomeTransform, _DefaultType]] = DEFAULT,
+        input_transform: Optional[Union[InputTransform, _DefaultType]] = DEFAULT,
     ) -> None:
         r"""
         Args:
@@ -154,16 +155,27 @@ def __init__(
             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 in the model's
-                forward pass.
+                `.posterior` on the model will be on the original scale). We use a
+                `Standardize` transform if no `outcome_transform` is specified.
+                Pass down `None` to use no outcome transform.
+            input_transform: An input transform that is applied in the model's forward
+                pass. We use a `Normalize` transform if no `input_transform` is
+                specified. Pass down `None` to use no input transform.
         """
+        self._validate_tensor_args(X=train_X, Y=train_Y, Yvar=train_Yvar)
+        if outcome_transform == DEFAULT:
+            outcome_transform = Standardize(
+                m=train_Y.shape[-1], batch_shape=train_X.shape[:-2]
+            )
+        if input_transform == DEFAULT:
+            input_transform = Normalize(d=train_X.shape[-1], transform_on_train=True)
         with torch.no_grad():
             transformed_X = self.transform_inputs(
                 X=train_X, input_transform=input_transform
             )
         if outcome_transform is not None:
             train_Y, train_Yvar = outcome_transform(train_Y, train_Yvar)
+        # Validate again after applying the transforms
         self._validate_tensor_args(X=transformed_X, Y=train_Y, Yvar=train_Yvar)
         ignore_X_dims = getattr(self, "_ignore_X_dims_scaling_check", None)
         validate_input_scaling(
@@ -352,6 +364,7 @@ def __init__(
             train_X=train_X,
             train_Y=train_Y,
             likelihood=likelihood,
+            outcome_transform=None,
             input_transform=input_transform,
         )
         self.register_added_loss_term("noise_added_loss")

test/acquisition/multi_objective/test_max_value_entropy_search.py

@@ -71,14 +71,14 @@ def test_multi_objective_max_value_entropy(self):
             # test batched model
             train_X = torch.rand(1, 1, 2, dtype=dtype, device=self.device)
             train_Y = torch.rand(1, 1, m, dtype=dtype, device=self.device)
-            model = SingleTaskGP(train_X, train_Y)
+            model = SingleTaskGP(train_X, train_Y, outcome_transform=None)
             with self.assertRaises(NotImplementedError):
                 qMultiObjectiveMaxValueEntropy(model, dummy_sample_pareto_frontiers)
             # test initialization
             train_X = torch.rand(4, 2, dtype=dtype, device=self.device)
             train_Y = torch.rand(4, m, dtype=dtype, device=self.device)
             # test batched MO model
-            model = SingleTaskGP(train_X, train_Y)
+            model = SingleTaskGP(train_X, train_Y, outcome_transform=None)
             mesmo = qMultiObjectiveMaxValueEntropy(model, dummy_sample_pareto_frontiers)
             self.assertEqual(mesmo.num_fantasies, 16)
             # Initialize the sampler.
@@ -98,7 +98,10 @@ def test_multi_objective_max_value_entropy(self):
             )
             # test ModelListGP
             model = ModelListGP(
-                *[SingleTaskGP(train_X, train_Y[:, i : i + 1]) for i in range(m)]
+                *[
+                    SingleTaskGP(train_X, train_Y[:, i : i + 1], outcome_transform=None)
+                    for i in range(m)
+                ]
             )
             mock_sample_pfs = mock.Mock()
             mock_sample_pfs.return_value = dummy_sample_pareto_frontiers(model=model)
@@ -156,7 +159,7 @@ def test_multi_objective_max_value_entropy(self):
                 ],
                 dim=1,
             )
-            fantasy_model = SingleTaskGP(fant_X, fant_Y)
+            fantasy_model = SingleTaskGP(fant_X, fant_Y, outcome_transform=None)
 
             # test with X_pending is not None
             with mock.patch.object(

test/acquisition/test_proximal.py

@@ -245,7 +245,9 @@ def test_proximal_model_list(self):
             train_X = torch.rand(5, 3, device=self.device, dtype=dtype)
             train_Y = train_X.norm(dim=-1, keepdim=True)
 
-            gp = SingleTaskGP(train_X, train_Y).to(device=self.device)
+            gp = SingleTaskGP(train_X, train_Y, input_transform=None).to(
+                device=self.device
+            )
             model = ModelListGP(gp, gp)
 
             scalarized_posterior_transform = ScalarizedPosteriorTransform(