log values for HVKG using LogEHVI

Ethan Che · facebook-github-bot · commit bf7a89be9ab3 · 2025-08-14T08:05:17.000-07:00
Summary:
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.

Rollback Plan:

Differential Revision: D80263869
diff --git a/botorch/acquisition/cost_aware.py b/botorch/acquisition/cost_aware.py
@@ -25,6 +25,7 @@
 from botorch.models.deterministic import DeterministicModel
 from botorch.models.gpytorch import GPyTorchModel
 from botorch.sampling.base import MCSampler
+from botorch.utils.safe_math import log as safe_log
 from pyre_extensions import none_throws
 from torch import Tensor
 from torch.nn import Module
@@ -113,6 +114,7 @@ def __init__(
         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
@@ -148,6 +150,7 @@ def __init__(
         self.cost_objective: MCAcquisitionObjective = cost_objective
         self._use_mean = use_mean
         self._min_cost = min_cost
+        self._log = log
 
     def forward(
         self,
@@ -215,5 +218,10 @@ def forward(
 
         # 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 - safe_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)
diff --git a/botorch/acquisition/multi_objective/hypervolume_knowledge_gradient.py b/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,7 @@ 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.
         """
 
         super().__init__(
@@ -392,6 +420,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 +494,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 +511,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 +538,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_hv_value_function(
@@ -505,6 +553,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,20 +567,27 @@ def _get_hv_value_function(
             action="ignore",
             category=NumericsWarning,
         )
-        base_value_function = qExpectedHypervolumeImprovement(
-            model=model,
-            ref_point=ref_point,
-            partitioning=FastNondominatedPartitioning(
+
+        value_fn_kwargs = {
+            "model": model,
+            "ref_point": ref_point,
+            "partitioning": FastNondominatedPartitioning(
                 ref_point=ref_point,
                 Y=torch.empty(
                     (0, ref_point.shape[0]),
                     dtype=ref_point.dtype,
                     device=ref_point.device,
                 ),
             ),  # create empty partitioning
-            sampler=sampler,
-            objective=objective,
-        )
+            "sampler": sampler,
+            "objective": objective,
+        }
+
+        if log:
+            base_value_function = qLogExpectedHypervolumeImprovement(**value_fn_kwargs)
+        else:
+            base_value_function = qExpectedHypervolumeImprovement(**value_fn_kwargs)
+
     # ProjectedAcquisitionFunction requires this
     base_value_function.posterior_transform = None
 
diff --git a/test/acquisition/multi_objective/test_hypervolume_knowledge_gradient.py b/test/acquisition/multi_objective/test_hypervolume_knowledge_gradient.py
