log values for HVKG using LogEHVI (#2974)

Summary:
Pull Request resolved: #2974

HKVG to output log values when _log=True using the LogEHVI value function.

Also adds a _log flag to InverseCostWeightedUtility to output utilities in log-space.

if _log=True, it assumes that
* current_value is in log-space
* cost_aware_utility outputs in log-space (raises an error if cost_aware_utility does not have a _log flag or if its _log=False).

Note that InverseCostWeightedUtility does a logarithmic transform for the inputted costs; assumes that inputted costs are in the original space. This is so that one does not need to make any direct modification to the cost fn, and because InverseCostWeightedUtility already does some pre-processing to the costs (e.g. clipping).

tldr: HKVG assumes all of its inputs are logged, but InverseCostWeightedUtility does not.

Reviewed By: SebastianAment

Differential Revision: D80263869

fbshipit-source-id: 95c29aac348a5541bf79de3a27208896936e424b

Author: Ethan Che

botorch/acquisition/cost_aware.py

@@ -11,7 +11,6 @@
 
 from __future__ import annotations
 
-import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 
@@ -21,7 +20,6 @@
     IdentityMCObjective,
     MCAcquisitionObjective,
 )
-from botorch.exceptions.warnings import CostAwareWarning
 from botorch.models.deterministic import DeterministicModel
 from botorch.models.gpytorch import GPyTorchModel
 from botorch.sampling.base import MCSampler
@@ -112,7 +110,7 @@ def __init__(
         cost_model: DeterministicModel | GPyTorchModel,
         use_mean: bool = True,
         cost_objective: MCAcquisitionObjective | None = None,
-        min_cost: float = 1e-2,
+        log: bool = False,
     ) -> None:
         r"""Cost-aware utility that weights increase in utility by inverse cost.
         For negative increases in utility, the utility is instead scaled by the
@@ -130,7 +128,8 @@ def __init__(
                 un-transform predictions/samples of a cost model fit on the
                 log-transformed cost (often done to ensure non-negativity). If the
                 cost model is multi-output, then by default this will sum the cost
-                across outputs.
+                across outputs. NOTE: Keep in mind that cost_objective must output
+                strictly positive values; forward will raise a ValueError otherwise.
             min_cost: A value used to clamp the cost samples so that they are not
                 too close to zero, which may cause numerical issues.
         Returns:
@@ -147,7 +146,7 @@ def __init__(
         self.cost_model = cost_model
         self.cost_objective: MCAcquisitionObjective = cost_objective
         self._use_mean = use_mean
-        self._min_cost = min_cost
+        self._log = log
 
     def forward(
         self,
@@ -202,18 +201,21 @@ def forward(
             cost = none_throws(sampler)(cost_posterior)
         cost = self.cost_objective(cost)
 
-        # Ensure non-negativity of the cost
-        if torch.any(cost < -1e-7):
-            warnings.warn(
-                "Encountered negative cost values in InverseCostWeightedUtility",
-                CostAwareWarning,
-                stacklevel=2,
+        # Ensure that costs are positive
+        if not torch.all(cost > 0.0):
+            raise ValueError(
+                "Costs must be strictly positive. Consider clamping cost_objective."
             )
-        # clamp (away from zero) and sum cost across elements of the q-batch -
-        # this will be of shape `num_fantasies x batch_shape` or `batch_shape`
-        cost = cost.clamp_min(self._min_cost).sum(dim=-1)
+
+        # sum costs along q-batch
+        cost = cost.sum(dim=-1)
 
         # compute and return the ratio on the sample level - If `use_mean=True`
         # this operation involves broadcasting the cost across fantasies.
-        # We multiply by the cost if the deltas are <= 0, see discussion #2914
-        return torch.where(deltas > 0, deltas / cost, deltas * cost)
+        if self._log:
+            # if _log is True then input deltas are in log space
+            # so original deltas cannot be <= 0
+            return deltas - torch.log(cost)
+        else:
+            # We multiply by the cost if the deltas are <= 0, see discussion #2914
+            return torch.where(deltas > 0, deltas / cost, deltas * cost)

botorch/acquisition/multi_objective/hypervolume_knowledge_gradient.py

@@ -32,11 +32,12 @@
 from botorch.acquisition.decoupled import DecoupledAcquisitionFunction
 from botorch.acquisition.knowledge_gradient import ProjectedAcquisitionFunction
 from botorch.acquisition.multi_objective.base import MultiObjectiveMCAcquisitionFunction
+from botorch.acquisition.multi_objective.logei import qLogExpectedHypervolumeImprovement
 from botorch.acquisition.multi_objective.monte_carlo import (
     qExpectedHypervolumeImprovement,
 )
 from botorch.acquisition.multi_objective.objective import MCMultiOutputObjective
-from botorch.exceptions.errors import UnsupportedError
+from botorch.exceptions.errors import BotorchError, UnsupportedError
 from botorch.exceptions.warnings import NumericsWarning
 from botorch.models.deterministic import PosteriorMeanModel
 from botorch.models.model import Model
@@ -47,6 +48,7 @@
 from botorch.utils.multi_objective.box_decompositions.non_dominated import (
     FastNondominatedPartitioning,
 )
+from botorch.utils.safe_math import logdiffexp, logmeanexp
 from botorch.utils.transforms import (
     average_over_ensemble_models,
     match_batch_shape,
@@ -91,6 +93,7 @@ def __init__(
         current_value: Tensor | None = None,
         use_posterior_mean: bool = True,
         cost_aware_utility: CostAwareUtility | None = None,
+        log: bool = False,
     ) -> None:
         r"""q-Hypervolume Knowledge Gradient.
 
@@ -133,6 +136,9 @@ def __init__(
                 [Daulton2023hvkg]_ for details.
             cost_aware_utility: A CostAwareUtility specifying the cost function for
                 evaluating the `X` on the objectives indicated by `evaluation_mask`.
+            log: If True, then returns the log of the HVKG value. If True, then it
+                expects current_value to be in log-space and cost_aware_utility to
+                output log utilities.
         """
         if sampler is None:
             # base samples should be fixed for joint optimization over X, X_fantasies
@@ -170,6 +176,8 @@ def __init__(
         self.cost_aware_utility = cost_aware_utility
         self._cost_sampler = None
 
+        self._log = log
+
     @property
     def cost_sampler(self):
         if self._cost_sampler is None:
@@ -242,6 +250,7 @@ def forward(self, X: Tensor) -> Tensor:
             objective=self.objective,
             sampler=self.inner_sampler,
             use_posterior_mean=self.use_posterior_mean,
+            log=self._log,
         )
 
         # make sure to propagate gradients to the fantasy model train inputs
@@ -259,9 +268,23 @@ def forward(self, X: Tensor) -> Tensor:
             values = value_function(X=X_fantasies.reshape(shape))  # num_fantasies x b
 
         if self.current_value is not None:
-            values = values - self.current_value
+            if self._log:
+                values = logdiffexp(self.current_value, values)
+            else:
+                values = values - self.current_value
 
         if self.cost_aware_utility is not None:
+            if self._log:
+                # check whether cost_aware_utility has a _log flag
+                # raises an error if it does not or if _log is False
+                if (
+                    not hasattr(self.cost_aware_utility, "_log")
+                    or not self.cost_aware_utility._log
+                ):
+                    raise BotorchError(
+                        "Cost-aware HVKG has _log=True and requires cost_aware_utility"
+                        "to output log utilities."
+                    )
             values = self.cost_aware_utility(
                 # exclude pending points
                 X=X_actual[..., :q, :],
@@ -271,7 +294,10 @@ def forward(self, X: Tensor) -> Tensor:
             )
 
         # return average over the fantasy samples
-        return values.mean(dim=0)
+        if self._log:
+            return logmeanexp(values, dim=0)
+        else:
+            return values.mean(dim=0)
 
     def get_augmented_q_batch_size(self, q: int) -> int:
         r"""Get augmented q batch size for one-shot optimization.
@@ -329,6 +355,7 @@ def __init__(
         valfunc_cls: type[AcquisitionFunction] | None = None,
         valfunc_argfac: Callable[[Model], dict[str, Any]] | None = None,
         use_posterior_mean: bool = True,
+        log: bool = False,
         **kwargs: Any,
     ) -> None:
         r"""Multi-Fidelity q-Knowledge Gradient (one-shot optimization).
@@ -376,6 +403,9 @@ def __init__(
             valfunc_argfac: An argument factory, i.e. callable that maps a `Model`
                 to a dictionary of kwargs for the terminal value function (e.g.
                 `best_f` for `ExpectedImprovement`).
+            log: If True, then returns the log of the HVKG value. If True, then it
+                expects current_value to be in log-space and cost_aware_utility to
+                output log utilities.
         """
 
         super().__init__(
@@ -392,6 +422,7 @@ def __init__(
             current_value=current_value,
             use_posterior_mean=use_posterior_mean,
             cost_aware_utility=cost_aware_utility,
+            log=log,
         )
         self.project = project
         if kwargs.get("expand") is not None:
@@ -465,6 +496,7 @@ def forward(self, X: Tensor) -> Tensor:
             valfunc_cls=self.valfunc_cls,
             valfunc_argfac=self.valfunc_argfac,
             use_posterior_mean=self.use_posterior_mean,
+            log=self._log,
         )
 
         # make sure to propagate gradients to the fantasy model train inputs
@@ -481,9 +513,24 @@ def forward(self, X: Tensor) -> Tensor:
             )
             values = value_function(X=X_fantasies.reshape(shape))  # num_fantasies x b
         if self.current_value is not None:
-            values = values - self.current_value
+            if self._log:
+                # Assumes current value is in log-space
+                values = logdiffexp(self.current_value, values)
+            else:
+                values = values - self.current_value
 
         if self.cost_aware_utility is not None:
+            if self._log:
+                # check whether cost_aware_utility has a _log flag
+                # raises an error if it does not or if _log is False
+                if (
+                    not hasattr(self.cost_aware_utility, "_log")
+                    or not self.cost_aware_utility._log
+                ):
+                    raise BotorchError(
+                        "Cost-aware HVKG has _log=True and requires cost_aware_utility"
+                        "to output log utilities."
+                    )
             values = self.cost_aware_utility(
                 # exclude pending points
                 X=X_actual[..., :q, :],
@@ -493,7 +540,7 @@ def forward(self, X: Tensor) -> Tensor:
             )
 
         # return average over the fantasy samples
-        return values.mean(dim=0)
+        return logmeanexp(values, dim=0) if self._log else values.mean(dim=0)
 
 
 def _get_hv_value_function(
@@ -505,6 +552,7 @@ def _get_hv_value_function(
     valfunc_cls: type[AcquisitionFunction] | None = None,
     valfunc_argfac: Callable[[Model], dict[str, Any]] | None = None,
     use_posterior_mean: bool = False,
+    log: bool = False,
 ) -> AcquisitionFunction:
     r"""Construct value function (i.e. inner acquisition function).
     This is a method for computing hypervolume.
@@ -518,7 +566,13 @@ def _get_hv_value_function(
             action="ignore",
             category=NumericsWarning,
         )
-        base_value_function = qExpectedHypervolumeImprovement(
+
+        base_value_function_class = (
+            qLogExpectedHypervolumeImprovement
+            if log
+            else qExpectedHypervolumeImprovement
+        )
+        base_value_function = base_value_function_class(
             model=model,
             ref_point=ref_point,
             partitioning=FastNondominatedPartitioning(