Implements Multi-Fidelity GIBBON (Lower Bound MES) acquisition. (#1185)

Summary:
<!--
Thank you for sending the PR! We appreciate you spending the time to make BoTorch better.

Help us understand your motivation by explaining why you decided to make this change.

You can learn more about contributing to BoTorch here: https://github.com/pytorch/botorch/blob/main/CONTRIBUTING.md
-->

## Motivation

Since `qLowerBoundMaxValueEntropy` provides a cheap approximation to `qMaxValueEntropy`, this PR implements the multi-fidelity version of `qMultiFidelityMaxValueEntropy` to be able to use the approximation in a multi-fidelity setting as well.

### Have you read the [Contributing Guidelines on pull requests](https://github.com/pytorch/botorch/blob/main/CONTRIBUTING.md#pull-requests)?

Yes.

Pull Request resolved: #1185

Test Plan: Copied the unit test for `qMultiFidelityMaxValueEntropy` without changes (besides the class name). The `qMultiFidelityLowerBoundMaxValueEntropy` class is the same as `qMultiFidelityMaxValueEntropy` with a different compute_information_gain function, so doesn't require a different set of tests.

Reviewed By: dme65

Differential Revision: D35746956

Pulled By: Balandat

fbshipit-source-id: f937399cc2c83d28cd7ab3b6f44533ae6013a061

Author: npielawski

botorch/acquisition/__init__.py

@@ -34,6 +34,7 @@
 from botorch.acquisition.max_value_entropy_search import (
     MaxValueBase,
     qLowerBoundMaxValueEntropy,
+    qMultiFidelityLowerBoundMaxValueEntropy,
     qMaxValueEntropy,
     qMultiFidelityMaxValueEntropy,
 )
@@ -81,6 +82,7 @@
     "MaxValueBase",
     "qMultiFidelityKnowledgeGradient",
     "qMaxValueEntropy",
+    "qMultiFidelityLowerBoundMaxValueEntropy",
     "qLowerBoundMaxValueEntropy",
     "qMultiFidelityMaxValueEntropy",
     "qMultiStepLookahead",

botorch/acquisition/max_value_entropy_search.py

@@ -795,6 +795,56 @@ def forward(self, X: Tensor) -> Tensor:
         return ig.mean(dim=0)  # average over the fantasies
 
 
+class qMultiFidelityLowerBoundMaxValueEntropy(qMultiFidelityMaxValueEntropy):
+    r"""Multi-fidelity acquisition function for General-purpose Information-Based
+    Bayesian optimization (GIBBON).
+
+    The acquisition function for multi-fidelity max-value entropy search
+    with support for trace observations. See [Takeno2020mfmves]_
+    for a detailed discussion of the basic ideas on multi-fidelity MES
+    (note that this implementation is somewhat different). This acquisition function
+    is similar to `qMultiFidelityMaxValueEntropy` but computes the information gain
+    from the lower bound described in [Moss2021gibbon].
+
+    The model must be single-outcome, unless using a PosteriorTransform.
+    The batch case `q > 1` is supported through cyclic optimization and fantasies.
+
+    Example:
+        >>> model = SingleTaskGP(train_X, train_Y)
+        >>> candidate_set = torch.rand(1000, bounds.size(1))
+        >>> candidate_set = bounds[0] + (bounds[1] - bounds[0]) * candidate_set
+        >>> MF_qGIBBON = qMultiFidelityLowerBoundMaxValueEntropy(model, candidate_set)
+        >>> mf_gibbon = MF_qGIBBON(test_X)
+    """
+
+    def _compute_information_gain(
+        self, X: Tensor, mean_M: Tensor, variance_M: Tensor, covar_mM: Tensor
+    ) -> Tensor:
+        r"""Compute GIBBON's approximation of information gain at the design points `X`.
+
+        When using GIBBON for batch optimization (i.e `q > 1`), we calculate the
+        additional information provided by adding a new candidate point to the current
+        batch of design points (`X_pending`), rather than calculating the information
+        provided by the whole batch. This allows a modest computational saving.
+
+        Args:
+            X: A `batch_shape x 1 x d`-dim Tensor of `batch_shape` t-batches
+                with `1` `d`-dim design point each.
+            mean_M: A `batch_shape x 1`-dim Tensor of means.
+            variance_M: A `batch_shape x 1`-dim Tensor of variances
+                consisting of `batch_shape` t-batches with `num_fantasies` fantasies.
+            covar_mM: A `batch_shape x num_fantasies x (1 + num_trace_observations)`
+                -dim Tensor of covariances.
+
+        Returns:
+            A `num_fantasies x batch_shape`-dim Tensor of information gains at the
+            given design points `X`.
+        """
+        return qLowerBoundMaxValueEntropy._compute_information_gain(
+            self, X=X, mean_M=mean_M, variance_M=variance_M, covar_mM=covar_mM
+        )
+
+
 def _sample_max_value_Thompson(
     model: Model,
     candidate_set: Tensor,

test/acquisition/test_max_value_entropy_search.py

@@ -15,6 +15,7 @@
     qLowerBoundMaxValueEntropy,
     qMaxValueEntropy,
     qMultiFidelityMaxValueEntropy,
+    qMultiFidelityLowerBoundMaxValueEntropy,
 )
 from botorch.acquisition.objective import (
     PosteriorTransform,
@@ -241,14 +242,16 @@ def test_q_lower_bound_max_value_entropy(self):
             with self.assertRaisesRegex(UnsupportedError, "X_pending is not None"):
                 qGIBBON(X)
 
-    def test_q_multi_fidelity_max_value_entropy(self):
+    def test_q_multi_fidelity_max_value_entropy(
+        self, acqf_class=qMultiFidelityMaxValueEntropy
+    ):
         for dtype in (torch.float, torch.double):
             torch.manual_seed(7)
             mm = MESMockModel()
             train_inputs = torch.rand(10, 2, device=self.device, dtype=dtype)
             mm.train_inputs = (train_inputs,)
             candidate_set = torch.rand(10, 2, device=self.device, dtype=dtype)
-            qMF_MVE = qMultiFidelityMaxValueEntropy(
+            qMF_MVE = acqf_class(
                 model=mm, candidate_set=candidate_set, num_mv_samples=10
             )
 
@@ -277,7 +280,7 @@ def test_q_multi_fidelity_max_value_entropy(self):
             pt = ScalarizedPosteriorTransform(
                 weights=torch.ones(2, device=self.device, dtype=dtype)
             )
-            qMF_MVE = qMultiFidelityMaxValueEntropy(
+            qMF_MVE = acqf_class(
                 model=mm,
                 candidate_set=candidate_set,
                 num_mv_samples=10,
@@ -286,6 +289,13 @@ def test_q_multi_fidelity_max_value_entropy(self):
             X = torch.rand(1, 2, device=self.device, dtype=dtype)
             self.assertEqual(qMF_MVE(X).shape, torch.Size([1]))
 
+    def test_q_multi_fidelity_lower_bound_max_value_entropy(self):
+        # Same test as for MF-MES since GIBBON only changes in the way it computes the
+        # information gain.
+        self.test_q_multi_fidelity_max_value_entropy(
+            acqf_class=qMultiFidelityLowerBoundMaxValueEntropy
+        )
+
     def test_sample_max_value_Gumbel(self):
         for dtype in (torch.float, torch.double):
             torch.manual_seed(7)