Change list surrogate construction (#1247)

Summary:
Pull Request resolved: #1247

Change `submodel_outcome_transforms`, `submodel_input_transforms`, `submodel_covar_module_class`, `submodel_covar_module_options`, `submodel_likelihood_class`, `submodel_likelihood_options` to only accept one input for all models.

Reviewed By: lena-kashtelyan

Differential Revision: D40164298

fbshipit-source-id: 2799f300069600a7cfa7ebee3d6783c7ae03d346

Author: David Eriksson

ax/models/torch/botorch_modular/list_surrogate.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import inspect
+from copy import deepcopy
 
 from logging import Logger
 from typing import Any, Dict, List, Optional, Type
@@ -49,16 +50,22 @@ class ListSurrogate(Surrogate):
             ``submodel_outions_per_outcome[submodel_outcome]`` (individual).
         mll_class: ``MarginalLogLikelihood`` class to use for model-fitting.
         mll_options: Dictionary of options / kwargs for the MLL.
-        submodel_outcome_transforms: A dictionary mapping each outcome to a
-            BoTorch outcome transform. Gets passed down to the BoTorch ``Model``s.
+        submodel_outcome_transforms: An outcome transform that will be used
+            by all outcomes. Gets passed down to the BoTorch ``Model``s.
             To use multiple outcome transforms on a submodel, chain them
             together using ``ChainedOutcomeTransform``.
-        submodel_input_transforms: A dictionary mapping each outcome to a
-            BoTorch input transform. Gets passed down to the BoTorch ``Model``.
+        submodel_input_transforms: An input transform that will be used
+            by all outcomes. Gets passed down to the BoTorch ``Model``.
             If sharing a single ``InputTransform`` object across submodels is
             preferred, pass in a dictionary where each outcome key references the
             same ``InputTransform`` object. To use multiple input transfroms on
             a submodel, chain them together using ``ChainedInputTransform``.
+        submodel_covar_module_class: A covar module that will be used by all outcomes.
+        submodel_covar_module_options: Options for a BoTorch covar module or options
+            that will be used by all outcomes.
+        submodel_likelihood_class: A likelihood that will be used by all outcomes.
+        submodel_likelihood_options: Options for a BoTorch likelihood or options that
+            will be used by all outcomes.
     """
 
     botorch_submodel_class_per_outcome: Dict[str, Type[Model]]
@@ -67,12 +74,12 @@ class ListSurrogate(Surrogate):
     submodel_options: Dict[str, Any]
     mll_class: Type[MarginalLogLikelihood]
     mll_options: Dict[str, Any]
-    submodel_outcome_transforms: Dict[str, OutcomeTransform]
-    submodel_input_transforms: Dict[str, InputTransform]
-    submodel_covar_module_class: Dict[str, Type[Kernel]]
-    submodel_covar_module_options: Dict[str, Dict[str, Any]]
-    submodel_likelihood_class: Dict[str, Type[Likelihood]]
-    submodel_likelihood_options: Dict[str, Dict[str, Any]]
+    submodel_outcome_transforms: Optional[OutcomeTransform]
+    submodel_input_transforms: Optional[InputTransform]
+    submodel_covar_module_class: Optional[Type[Kernel]]
+    submodel_covar_module_options: Dict[str, Any]
+    submodel_likelihood_class: Optional[Type[Likelihood]]
+    submodel_likelihood_options: Dict[str, Any]
     _model: Optional[Model] = None
     # Special setting for surrogates instantiated via `Surrogate.from_botorch`,
     # to avoid re-constructing the underlying BoTorch model on `Surrogate.fit`
@@ -87,12 +94,12 @@ def __init__(
         submodel_options: Optional[Dict[str, Any]] = None,
         mll_class: Type[MarginalLogLikelihood] = ExactMarginalLogLikelihood,
         mll_options: Optional[Dict[str, Any]] = None,
-        submodel_outcome_transforms: Optional[Dict[str, OutcomeTransform]] = None,
-        submodel_input_transforms: Optional[Dict[str, InputTransform]] = None,
-        submodel_covar_module_class: Optional[Dict[str, Type[Kernel]]] = None,
-        submodel_covar_module_options: Optional[Dict[str, Dict[str, Any]]] = None,
-        submodel_likelihood_class: Optional[Dict[str, Type[Likelihood]]] = None,
-        submodel_likelihood_options: Optional[Dict[str, Dict[str, Any]]] = None,
+        submodel_outcome_transforms: Optional[OutcomeTransform] = None,
+        submodel_input_transforms: Optional[InputTransform] = None,
+        submodel_covar_module_class: Optional[Type[Kernel]] = None,
+        submodel_covar_module_options: Optional[Dict[str, Any]] = None,
+        submodel_likelihood_class: Optional[Type[Likelihood]] = None,
+        submodel_likelihood_options: Optional[Dict[str, Any]] = None,
     ) -> None:
         if not bool(botorch_submodel_class_per_outcome) ^ bool(botorch_submodel_class):
             raise ValueError(  # pragma: no cover
@@ -106,11 +113,11 @@ def __init__(
         self.botorch_submodel_class = botorch_submodel_class
         self.submodel_options_per_outcome = submodel_options_per_outcome or {}
         self.submodel_options = submodel_options or {}
-        self.submodel_outcome_transforms = submodel_outcome_transforms or {}
-        self.submodel_input_transforms = submodel_input_transforms or {}
-        self.submodel_covar_module_class = submodel_covar_module_class or {}
+        self.submodel_outcome_transforms = submodel_outcome_transforms
+        self.submodel_input_transforms = submodel_input_transforms
+        self.submodel_covar_module_class = submodel_covar_module_class
         self.submodel_covar_module_options = submodel_covar_module_options or {}
-        self.submodel_likelihood_class = submodel_likelihood_class or {}
+        self.submodel_likelihood_class = submodel_likelihood_class
         self.submodel_likelihood_options = submodel_likelihood_options or {}
         super().__init__(
             botorch_model_class=ModelListGP,
@@ -159,7 +166,6 @@ def construct(
         # Construct input perturbation if doing robust optimization.
         # NOTE: Doing this here rather than in `_set_formatted_inputs` to make sure
         # we use the same perturbations for each sub-model.
-        submodel_input_transforms = self.submodel_input_transforms.copy()
         robust_digest: Optional[Dict[str, Any]] = kwargs.get("robust_digest", None)
         if robust_digest is not None:
             if len(robust_digest["environmental_variables"]):
@@ -176,15 +182,15 @@ def construct(
                 perturbation_set=samples, multiplicative=robust_digest["multiplicative"]
             )
 
-            for m in metric_names:
-                if submodel_input_transforms.get(m) is not None:
-                    # TODO: Support mixing with user supplied transforms.
-                    raise NotImplementedError(
-                        "User supplied input transforms are not supported "
-                        "in robust optimization."
-                    )
-                else:
-                    submodel_input_transforms[m] = perturbation
+            if self.submodel_input_transforms is not None:
+                # TODO: Support mixing with user supplied transforms.
+                raise NotImplementedError(
+                    "User supplied input transforms are not supported "
+                    "in robust optimization."
+                )
+            submodel_input_transforms = perturbation
+        else:
+            submodel_input_transforms = self.submodel_input_transforms
 
         submodels = []
         for m, dataset in zip(metric_names, datasets):
@@ -218,20 +224,33 @@ def construct(
             # way to filter the arguments. See the comment in `Surrogate.construct`
             # regarding potential use of a `ModelFactory` in the future.
             model_cls_args = inspect.getfullargspec(model_cls).args
-            covar_module_class = self.submodel_covar_module_class.get(m)
-            covar_module_options = self.submodel_covar_module_options.get(m)
-            likelihood_class = self.submodel_likelihood_class.get(m)
-            likelihood_options = self.submodel_likelihood_options.get(m)
-            outcome_transform = self.submodel_outcome_transforms.get(m)
-            input_transform = submodel_input_transforms.get(m)
-
             self._set_formatted_inputs(
                 formatted_model_inputs=formatted_model_inputs,
                 inputs=[
-                    ["covar_module", covar_module_class, covar_module_options, None],
-                    ["likelihood", likelihood_class, likelihood_options, None],
-                    ["outcome_transform", None, None, outcome_transform],
-                    ["input_transform", None, None, input_transform],
+                    [
+                        "covar_module",
+                        self.submodel_covar_module_class,
+                        self.submodel_covar_module_options,
+                        None,
+                    ],
+                    [
+                        "likelihood",
+                        self.submodel_likelihood_class,
+                        self.submodel_likelihood_options,
+                        None,
+                    ],
+                    [
+                        "outcome_transform",
+                        None,
+                        None,
+                        deepcopy(self.submodel_outcome_transforms),
+                    ],
+                    [
+                        "input_transform",
+                        None,
+                        None,
+                        deepcopy(submodel_input_transforms),
+                    ],
                 ],
                 dataset=dataset,
                 botorch_model_class_args=model_cls_args,

ax/models/torch/tests/test_list_surrogate.py

@@ -4,10 +4,10 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import math
 from unittest.mock import Mock, patch
 
 import numpy as np
-
 import torch
 from ax.core.search_space import SearchSpaceDigest
 from ax.exceptions.core import UserInputError
@@ -37,7 +37,8 @@
     GaussianLikelihood,
     Likelihood,  # noqa: F401
 )
-from gpytorch.mlls import ExactMarginalLogLikelihood, LeaveOneOutPseudoLikelihood
+from gpytorch.mlls import ExactMarginalLogLikelihood
+
 
 SURROGATE_PATH = f"{Surrogate.__module__}"
 UTILS_PATH = f"{choose_model_class.__module__}"
@@ -58,6 +59,9 @@ def setUp(self) -> None:
         Xs1, Ys1, Yvars1, bounds, _, _, _ = get_torch_test_data(
             dtype=self.dtype, task_features=self.search_space_digest.task_features
         )
+        # Change the inputs/outputs a bit so the data isn't identical
+        Xs1[0] *= 2
+        Ys1[0] += 1
         Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(
             dtype=self.dtype, task_features=self.search_space_digest.task_features
         )
@@ -352,19 +356,12 @@ def test_fit(
             )
 
     def test_with_botorch_transforms(self) -> None:
-        input_transforms = {"outcome_1": Normalize(d=3), "outcome_2": Normalize(d=3)}
-        outcome_transforms = {
-            "outcome_1": Standardize(m=1),
-            "outcome_2": Standardize(m=1),
-        }
+        input_transforms = Normalize(d=3)
+        outcome_transforms = Standardize(m=1)
         surrogate = ListSurrogate(
             botorch_submodel_class=SingleTaskGPWithDifferentConstructor,
             mll_class=ExactMarginalLogLikelihood,
-            # pyre-fixme[6]: For 3rd param expected `Optional[Dict[str,
-            #  OutcomeTransform]]` but got `Dict[str, Standardize]`.
             submodel_outcome_transforms=outcome_transforms,
-            # pyre-fixme[6]: For 4th param expected `Optional[Dict[str,
-            #  InputTransform]]` but got `Dict[str, Normalize]`.
             submodel_input_transforms=input_transforms,
         )
         with self.assertRaisesRegex(UserInputError, "The BoTorch model class"):
@@ -375,23 +372,34 @@ def test_with_botorch_transforms(self) -> None:
         surrogate = ListSurrogate(
             botorch_submodel_class=SingleTaskGP,
             mll_class=ExactMarginalLogLikelihood,
-            # pyre-fixme[6]: For 3rd param expected `Optional[Dict[str,
-            #  OutcomeTransform]]` but got `Dict[str, Standardize]`.
             submodel_outcome_transforms=outcome_transforms,
-            # pyre-fixme[6]: For 4th param expected `Optional[Dict[str,
-            #  InputTransform]]` but got `Dict[str, Normalize]`.
             submodel_input_transforms=input_transforms,
         )
         surrogate.construct(
             datasets=self.supervised_training_data,
             metric_names=self.outcomes,
         )
-        models = surrogate.model.models
-        for i, outcome in enumerate(("outcome_1", "outcome_2")):
-            # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C._TensorBase...
-            self.assertIs(models[i].outcome_transform, outcome_transforms[outcome])
-            # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C._TensorBase...
-            self.assertIs(models[i].input_transform, input_transforms[outcome])
+        # pyre-ignore [9]
+        models: torch.nn.modules.container.ModuleList = surrogate.model.models
+        for i in range(2):
+            self.assertIsInstance(models[i].outcome_transform, Standardize)
+            self.assertIsInstance(models[i].input_transform, Normalize)
+        self.assertEqual(models[0].outcome_transform.means.item(), 4.5)
+        self.assertEqual(models[1].outcome_transform.means.item(), 3.5)
+        self.assertAlmostEqual(
+            models[0].outcome_transform.stdvs.item(), 1 / math.sqrt(2)
+        )
+        self.assertAlmostEqual(
+            models[1].outcome_transform.stdvs.item(), 1 / math.sqrt(2)
+        )
+        self.assertTrue(
+            torch.all(
+                torch.isclose(
+                    models[0].input_transform.bounds,
+                    2 * models[1].input_transform.bounds,  # pyre-ignore
+                )
+            )
+        )
 
     def test_serialize_attributes_as_kwargs(self) -> None:
         expected = self.surrogate.__dict__
@@ -411,48 +419,64 @@ def test_serialize_attributes_as_kwargs(self) -> None:
         self.assertEqual(self.surrogate._serialize_attributes_as_kwargs(), expected)
 
     def test_construct_custom_model(self) -> None:
-        noise_con1, noise_con2 = Interval(1e-6, 1e-1), GreaterThan(1e-4)
-        surrogate = ListSurrogate(
-            botorch_submodel_class=SingleTaskGP,
-            mll_class=LeaveOneOutPseudoLikelihood,
-            submodel_covar_module_class={
-                "outcome_1": RBFKernel,
-                "outcome_2": MaternKernel,
-            },
-            submodel_covar_module_options={
-                "outcome_1": {"ard_num_dims": 1},
-                "outcome_2": {"ard_num_dims": 3},
-            },
-            submodel_likelihood_class={
-                "outcome_1": GaussianLikelihood,
-                "outcome_2": GaussianLikelihood,
-            },
-            submodel_likelihood_options={
-                "outcome_1": {"noise_constraint": noise_con1},
-                "outcome_2": {"noise_constraint": noise_con2},
-            },
-        )
-        surrogate.construct(
-            datasets=self.supervised_training_data,
-            metric_names=self.outcomes,
-        )
-        # pyre-fixme[16]: Optional type has no attribute `models`.
-        self.assertEqual(len(surrogate._model.models), 2)
-        self.assertEqual(surrogate.mll_class, LeaveOneOutPseudoLikelihood)
-        for i, m in enumerate(surrogate._model.models):
-            self.assertEqual(type(m.likelihood), GaussianLikelihood)
-            if i == 0:
-                self.assertEqual(type(m.covar_module), RBFKernel)
-                self.assertEqual(m.covar_module.ard_num_dims, 1)
-                self.assertEqual(
-                    m.likelihood.noise_covar.raw_noise_constraint, noise_con1
-                )
-            else:
+        noise_constraint = Interval(1e-4, 10.0)
+        for submodel_covar_module_options, submodel_likelihood_options in [
+            [{"ard_num_dims": 3}, {"noise_constraint": noise_constraint}],
+            [{}, {}],
+        ]:
+            surrogate = ListSurrogate(
+                botorch_submodel_class=SingleTaskGP,
+                mll_class=ExactMarginalLogLikelihood,
+                submodel_covar_module_class=MaternKernel,
+                submodel_covar_module_options=submodel_covar_module_options,
+                submodel_likelihood_class=GaussianLikelihood,
+                submodel_likelihood_options=submodel_likelihood_options,
+                submodel_input_transforms=Normalize(d=3),
+                submodel_outcome_transforms=Standardize(m=1),
+            )
+            surrogate.construct(
+                datasets=self.supervised_training_data,
+                metric_names=self.outcomes,
+            )
+            # pyre-fixme[16]: Optional type has no attribute `models`.
+            self.assertEqual(len(surrogate._model.models), 2)
+            self.assertEqual(surrogate.mll_class, ExactMarginalLogLikelihood)
+            # Make sure we properly copied the transforms
+            self.assertNotEqual(
+                id(surrogate._model.models[0].input_transform),
+                id(surrogate._model.models[1].input_transform),
+            )
+            self.assertNotEqual(
+                id(surrogate._model.models[0].outcome_transform),
+                id(surrogate._model.models[1].outcome_transform),
+            )
+
+            for m in surrogate._model.models:
+                self.assertEqual(type(m.likelihood), GaussianLikelihood)
                 self.assertEqual(type(m.covar_module), MaternKernel)
-                self.assertEqual(m.covar_module.ard_num_dims, 3)
-                self.assertEqual(
-                    m.likelihood.noise_covar.raw_noise_constraint, noise_con2
-                )
+                if submodel_covar_module_options:
+                    self.assertEqual(m.covar_module.ard_num_dims, 3)
+                else:
+                    self.assertEqual(m.covar_module.ard_num_dims, None)
+                if submodel_likelihood_options:
+                    self.assertEqual(
+                        type(m.likelihood.noise_covar.raw_noise_constraint), Interval
+                    )
+                    self.assertEqual(
+                        m.likelihood.noise_covar.raw_noise_constraint.lower_bound,
+                        noise_constraint.lower_bound,
+                    )
+                    self.assertEqual(
+                        m.likelihood.noise_covar.raw_noise_constraint.upper_bound,
+                        noise_constraint.upper_bound,
+                    )
+                else:
+                    self.assertEqual(
+                        type(m.likelihood.noise_covar.raw_noise_constraint), GreaterThan
+                    )
+                    self.assertEqual(
+                        m.likelihood.noise_covar.raw_noise_constraint.lower_bound, 1e-4
+                    )
 
     def test_w_robust_digest(self) -> None:
         surrogate = ListSurrogate(
@@ -470,7 +494,7 @@ def test_w_robust_digest(self) -> None:
             "environmental_variables": [],
             "multiplicative": False,
         }
-        surrogate.submodel_input_transforms = {self.outcomes[0]: Normalize(d=1)}
+        surrogate.submodel_input_transforms = Normalize(d=1)
         with self.assertRaisesRegex(NotImplementedError, "input transforms"):
             surrogate.construct(
                 datasets=self.supervised_training_data,