Make it possible to infer noise (#1345)

Summary:
Pull Request resolved: #1345

We currently set the noise to 0 for all benchmark problems even though we often infer the noise in practice. This diff adds an `infer_noise` flag that we can use to control whether different problems should set the noise to a fixed value or `None`, where the latter will trigger inferring the noise in the model. I'm setting `infer_noise=True` by default for all problems.

Reviewed By: saitcakmak, Balandat

Differential Revision: D39413742

fbshipit-source-id: 70075940a43379bc4cf2b35366bb3e00bdd24947

ax/benchmark/benchmark_problem.py

@@ -43,6 +43,7 @@ class BenchmarkProblem(Base):
     optimization_config: OptimizationConfig
     runner: Runner
     num_trials: int
+    infer_noise: bool
     tracking_metrics: List[Metric] = field(default_factory=list)
 
     @classmethod
@@ -51,6 +52,7 @@ def from_botorch(
         test_problem_class: Type[BaseTestProblem],
         test_problem_kwargs: Dict[str, Any],
         num_trials: int,
+        infer_noise: bool = True,
     ) -> "BenchmarkProblem":
         """Create a BenchmarkProblem from a BoTorch BaseTestProblem using specialized
         Metrics and Runners. The test problem's result will be computed on the Runner
@@ -76,7 +78,7 @@ def from_botorch(
             objective=Objective(
                 metric=BotorchTestProblemMetric(
                     name=f"{test_problem.__class__.__name__}",
-                    noise_sd=(test_problem.noise_std or 0),
+                    noise_sd=None if infer_noise else (test_problem.noise_std or 0),
                 ),
                 minimize=True,
             )
@@ -91,6 +93,7 @@ def from_botorch(
                 test_problem_kwargs=test_problem_kwargs,
             ),
             num_trials=num_trials,
+            infer_noise=infer_noise,
         )
 
 
@@ -102,10 +105,8 @@ class SingleObjectiveBenchmarkProblem(BenchmarkProblem):
 
     optimal_value: float = field()
 
-    # pyre-fixme[2]: Parameter must be annotated.
-    def __init__(self, optimal_value: float, **kwargs) -> None:
+    def __init__(self, optimal_value: float, **kwargs: Any) -> None:
         super().__init__(**kwargs)
-
         object.__setattr__(self, "optimal_value", optimal_value)
 
     @classmethod
@@ -114,6 +115,7 @@ def from_botorch_synthetic(
         test_problem_class: Type[SyntheticTestFunction],
         test_problem_kwargs: Dict[str, Any],
         num_trials: int,
+        infer_noise: bool = True,
     ) -> "SingleObjectiveBenchmarkProblem":
         """Create a BenchmarkProblem from a BoTorch BaseTestProblem using specialized
         Metrics and Runners. The test problem's result will be computed on the Runner
@@ -127,6 +129,7 @@ def from_botorch_synthetic(
             test_problem_class=test_problem_class,
             test_problem_kwargs=test_problem_kwargs,
             num_trials=num_trials,
+            infer_noise=infer_noise,
         )
 
         return cls(
@@ -135,6 +138,7 @@ def from_botorch_synthetic(
             optimization_config=problem.optimization_config,
             runner=problem.runner,
             num_trials=num_trials,
+            infer_noise=infer_noise,
             optimal_value=test_problem.optimal_value,
         )
 
@@ -153,11 +157,9 @@ def __init__(
         self,
         maximum_hypervolume: float,
         reference_point: List[float],
-        # pyre-fixme[2]: Parameter must be annotated.
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
-
         object.__setattr__(self, "maximum_hypervolume", maximum_hypervolume)
         object.__setattr__(self, "reference_point", reference_point)
 
@@ -167,6 +169,7 @@ def from_botorch_multi_objective(
         test_problem_class: Type[MultiObjectiveTestProblem],
         test_problem_kwargs: Dict[str, Any],
         num_trials: int,
+        infer_noise: bool = True,
     ) -> "MultiObjectiveBenchmarkProblem":
         """Create a BenchmarkProblem from a BoTorch BaseTestProblem using specialized
         Metrics and Runners. The test problem's result will be computed on the Runner
@@ -180,12 +183,13 @@ def from_botorch_multi_objective(
             test_problem_class=test_problem_class,
             test_problem_kwargs=test_problem_kwargs,
             num_trials=num_trials,
+            infer_noise=infer_noise,
         )
 
         metrics = [
             BotorchTestProblemMetric(
                 name=f"{test_problem.__class__.__name__}_{i}",
-                noise_sd=(test_problem.noise_std or 0),
+                noise_sd=None if infer_noise else (test_problem.noise_std or 0),
                 index=i,
             )
             for i in range(test_problem.num_objectives)
@@ -218,6 +222,7 @@ def from_botorch_multi_objective(
             optimization_config=optimization_config,
             runner=problem.runner,
             num_trials=num_trials,
+            infer_noise=infer_noise,
             maximum_hypervolume=test_problem.max_hv,
             reference_point=test_problem._ref_point,
         )

ax/benchmark/problems/hpo/pytorch_cnn.py

@@ -36,7 +36,12 @@ def __eq__(self, other: Base) -> bool:
 
     @classmethod
     def from_datasets(
-        cls, name: str, num_trials: int, train_set: Dataset, test_set: Dataset
+        cls,
+        name: str,
+        num_trials: int,
+        train_set: Dataset,
+        test_set: Dataset,
+        infer_noise: bool = True,
     ) -> "PyTorchCNNBenchmarkProblem":
         optimal_value = 1
 
@@ -73,7 +78,7 @@ def from_datasets(
         )
         optimization_config = OptimizationConfig(
             objective=Objective(
-                metric=PyTorchCNNMetric(),
+                metric=PyTorchCNNMetric(infer_noise=infer_noise),
                 minimize=False,
             )
         )
@@ -87,12 +92,14 @@ def from_datasets(
             optimization_config=optimization_config,
             runner=runner,
             num_trials=num_trials,
+            infer_noise=infer_noise,
         )
 
 
 class PyTorchCNNMetric(Metric):
-    def __init__(self) -> None:
+    def __init__(self, infer_noise: bool = True) -> None:
         super().__init__(name="accuracy")
+        self.infer_noise = infer_noise
 
     def fetch_trial_data(self, trial: BaseTrial, **kwargs: Any) -> MetricFetchResult:
         try:
@@ -105,7 +112,7 @@ def fetch_trial_data(self, trial: BaseTrial, **kwargs: Any) -> MetricFetchResult
                     "arm_name": [name for name, _ in trial.arms_by_name.items()],
                     "metric_name": self.name,
                     "mean": accuracy,
-                    "sem": 0,
+                    "sem": None if self.infer_noise else 0,
                     "trial_index": trial.index,
                 }
             )

ax/benchmark/problems/hpo/torchvision.py

@@ -25,7 +25,10 @@
 class PyTorchCNNTorchvisionBenchmarkProblem(PyTorchCNNBenchmarkProblem):
     @classmethod
     def from_dataset_name(
-        cls, name: str, num_trials: int
+        cls,
+        name: str,
+        num_trials: int,
+        infer_noise: bool = True,
     ) -> "PyTorchCNNTorchvisionBenchmarkProblem":
         if name not in _REGISTRY:
             raise UserInputError(
@@ -49,7 +52,11 @@ def from_dataset_name(
         )
 
         problem = cls.from_datasets(
-            name=name, num_trials=num_trials, train_set=train_set, test_set=test_set
+            name=name,
+            num_trials=num_trials,
+            train_set=train_set,
+            test_set=test_set,
+            infer_noise=infer_noise,
         )
         runner = PyTorchCNNTorchvisionRunner(
             name=name, train_set=train_set, test_set=test_set
@@ -61,6 +68,7 @@ def from_dataset_name(
             optimization_config=problem.optimization_config,
             runner=runner,
             num_trials=num_trials,
+            infer_noise=infer_noise,
             optimal_value=problem.optimal_value,
         )
 

ax/benchmark/problems/registry.py

@@ -31,6 +31,7 @@ class BenchmarkProblemRegistryEntry:
             "test_problem_class": Ackley,
             "test_problem_kwargs": {},
             "num_trials": 50,
+            "infer_noise": True,
         },
     ),
     "branin": BenchmarkProblemRegistryEntry(
@@ -39,6 +40,7 @@ class BenchmarkProblemRegistryEntry:
             "test_problem_class": Branin,
             "test_problem_kwargs": {},
             "num_trials": 30,
+            "infer_noise": True,
         },
     ),
     "branin_currin": BenchmarkProblemRegistryEntry(
@@ -47,12 +49,16 @@ class BenchmarkProblemRegistryEntry:
             "test_problem_class": BraninCurrin,
             "test_problem_kwargs": {},
             "num_trials": 30,
+            "infer_noise": True,
         },
     ),
     "branin_currin30": BenchmarkProblemRegistryEntry(
         factory_fn=lambda n: embed_higher_dimension(
             problem=MultiObjectiveBenchmarkProblem.from_botorch_multi_objective(
-                test_problem_class=BraninCurrin, test_problem_kwargs={}, num_trials=100
+                test_problem_class=BraninCurrin,
+                test_problem_kwargs={},
+                num_trials=100,
+                infer_noise=True,
             ),
             total_dimensionality=n,
         ),
@@ -64,6 +70,7 @@ class BenchmarkProblemRegistryEntry:
             "test_problem_class": Hartmann,
             "test_problem_kwargs": {"dim": 6},
             "num_trials": 50,
+            "infer_noise": True,
         },
     ),
     "hartmann30": BenchmarkProblemRegistryEntry(
@@ -72,31 +79,89 @@ class BenchmarkProblemRegistryEntry:
                 test_problem_class=Hartmann,
                 test_problem_kwargs={"dim": 6},
                 num_trials=100,
+                infer_noise=True,
             ),
             total_dimensionality=n,
         ),
         factory_kwargs={"n": 30},
     ),
     "hpo_pytorch_cnn_MNIST": BenchmarkProblemRegistryEntry(
         factory_fn=PyTorchCNNTorchvisionBenchmarkProblem.from_dataset_name,
-        factory_kwargs={"name": "MNIST", "num_trials": 50},
+        factory_kwargs={"name": "MNIST", "num_trials": 50, "infer_noise": True},
     ),
     "hpo_pytorch_cnn_FashionMNIST": BenchmarkProblemRegistryEntry(
         factory_fn=PyTorchCNNTorchvisionBenchmarkProblem.from_dataset_name,
-        factory_kwargs={"name": "FashionMNIST", "num_trials": 50},
+        factory_kwargs={"name": "FashionMNIST", "num_trials": 50, "infer_noise": True},
     ),
     "jenatton": BenchmarkProblemRegistryEntry(
         factory_fn=get_jenatton_benchmark_problem,
-        factory_kwargs={"num_trials": 50},
+        factory_kwargs={"num_trials": 50, "infer_noise": True},
     ),
     "powell": BenchmarkProblemRegistryEntry(
         factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic,
         factory_kwargs={
             "test_problem_class": Powell,
             "test_problem_kwargs": {},
             "num_trials": 50,
+            "infer_noise": True,
+        },
+    ),
+    # Problems without inferred noise
+    "branin_fixed_noise": BenchmarkProblemRegistryEntry(
+        factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic,
+        factory_kwargs={
+            "test_problem_class": Branin,
+            "test_problem_kwargs": {},
+            "num_trials": 30,
+            "infer_noise": False,
+        },
+    ),
+    "branin_currin_fixed_noise": BenchmarkProblemRegistryEntry(
+        factory_fn=MultiObjectiveBenchmarkProblem.from_botorch_multi_objective,
+        factory_kwargs={
+            "test_problem_class": BraninCurrin,
+            "test_problem_kwargs": {},
+            "num_trials": 30,
+            "infer_noise": False,
+        },
+    ),
+    "branin_currin30_fixed_noise": BenchmarkProblemRegistryEntry(
+        factory_fn=lambda n: embed_higher_dimension(
+            problem=MultiObjectiveBenchmarkProblem.from_botorch_multi_objective(
+                test_problem_class=BraninCurrin,
+                test_problem_kwargs={},
+                num_trials=100,
+                infer_noise=False,
+            ),
+            total_dimensionality=n,
+        ),
+        factory_kwargs={"n": 30},
+    ),
+    "hartmann6_fixed_noise": BenchmarkProblemRegistryEntry(
+        factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic,
+        factory_kwargs={
+            "test_problem_class": Hartmann,
+            "test_problem_kwargs": {"dim": 6},
+            "num_trials": 50,
+            "infer_noise": False,
         },
     ),
+    "hartmann30_fixed_noise": BenchmarkProblemRegistryEntry(
+        factory_fn=lambda n: embed_higher_dimension(
+            problem=SingleObjectiveBenchmarkProblem.from_botorch_synthetic(
+                test_problem_class=Hartmann,
+                test_problem_kwargs={"dim": 6},
+                num_trials=100,
+                infer_noise=False,
+            ),
+            total_dimensionality=n,
+        ),
+        factory_kwargs={"n": 30},
+    ),
+    "jenatton_fixed_noise": BenchmarkProblemRegistryEntry(
+        factory_fn=get_jenatton_benchmark_problem,
+        factory_kwargs={"num_trials": 50, "infer_noise": False},
+    ),
 }
 
 

ax/benchmark/problems/surrogate.py

@@ -5,8 +5,6 @@
 
 from typing import Any, Dict, Iterable, List, Set
 
-import numpy as np
-
 import pandas as pd
 import torch
 from ax.benchmark.benchmark_problem import SingleObjectiveBenchmarkProblem
@@ -45,6 +43,7 @@ def from_surrogate(
         minimize: bool,
         optimal_value: float,
         num_trials: int,
+        infer_noise: bool = True,
     ) -> "SurrogateBenchmarkProblem":
         return SurrogateBenchmarkProblem(
             name=name,
@@ -63,12 +62,14 @@ def from_surrogate(
             ),
             optimal_value=optimal_value,
             num_trials=num_trials,
+            infer_noise=infer_noise,
         )
 
 
 class SurrogateMetric(Metric):
-    def __init__(self) -> None:
+    def __init__(self, infer_noise: bool = True) -> None:
         super().__init__(name="prediction")
+        self.infer_noise = infer_noise
 
     # pyre-fixme[2]: Parameter must be annotated.
     def fetch_trial_data(self, trial: BaseTrial, **kwargs) -> MetricFetchResult:
@@ -82,7 +83,7 @@ def fetch_trial_data(self, trial: BaseTrial, **kwargs) -> MetricFetchResult:
                     "arm_name": [name for name, _ in trial.arms_by_name.items()],
                     "metric_name": self.name,
                     "mean": prediction,
-                    "sem": np.nan,
+                    "sem": None if self.infer_noise else 0,
                     "trial_index": trial.index,
                 }
             )

ax/benchmark/problems/synthetic/hss/jenatton.py

@@ -14,6 +14,7 @@
 
 def get_jenatton_benchmark_problem(
     num_trials: int = 50,
+    infer_noise: bool = True,
 ) -> SingleObjectiveBenchmarkProblem:
     search_space = HierarchicalSearchSpace(
         parameters=[
@@ -54,7 +55,9 @@ def get_jenatton_benchmark_problem(
     )
 
     optimization_config = OptimizationConfig(
-        objective=Objective(metric=JenattonMetric(), minimize=True)
+        objective=Objective(
+            metric=JenattonMetric(infer_noise=infer_noise), minimize=True
+        )
     )
 
     return SingleObjectiveBenchmarkProblem(
@@ -63,5 +66,6 @@ def get_jenatton_benchmark_problem(
         optimization_config=optimization_config,
         runner=SyntheticRunner(),
         num_trials=num_trials,
+        infer_noise=infer_noise,
         optimal_value=0.1,
     )