#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
r"""
Acquisition functions for max-value entropy search for multi-objective
Bayesian optimization (MESMO).
References
.. [Belakaria2019]
S. Belakaria, A. Deshwal, J. R. Doppa. Max-value Entropy Search
for Multi-Objective Bayesian Optimization. Advances in Neural
Information Processing Systems, 32. 2019.
"""
from __future__ import annotations
from math import pi
from typing import Callable, Optional, Tuple, Union
import torch
from botorch.acquisition.max_value_entropy_search import qMaxValueEntropy
from botorch.acquisition.multi_objective.joint_entropy_search import (
LowerBoundMultiObjectiveEntropySearch,
)
from botorch.acquisition.multi_objective.monte_carlo import (
MultiObjectiveMCAcquisitionFunction,
)
from botorch.models.converter import (
batched_multi_output_to_single_output,
model_list_to_batched,
)
from botorch.models.model import Model
from botorch.models.model_list_gp_regression import ModelListGP
from botorch.posteriors.gpytorch import GPyTorchPosterior
from botorch.sampling.base import MCSampler
from botorch.sampling.normal import SobolQMCNormalSampler
from botorch.utils.transforms import concatenate_pending_points, t_batch_mode_transform
from torch import Tensor
[docs]
class qMultiObjectiveMaxValueEntropy(
qMaxValueEntropy, MultiObjectiveMCAcquisitionFunction
):
r"""The acquisition function for MESMO.
This acquisition function computes the mutual information of
Pareto frontier and a candidate point. See [Belakaria2019]_ for
a detailed discussion.
q > 1 is supported through cyclic optimization and fantasies.
Noisy observations are support by computing information gain with
observation noise as in Appendix C in [Takeno2020mfmves]_.
Note: this only supports maximization.
Attributes:
_default_sample_shape: The `sample_shape` for the default sampler.
Example:
>>> model = SingleTaskGP(train_X, train_Y)
>>> MESMO = qMultiObjectiveMaxValueEntropy(model, sample_pfs)
>>> mesmo = MESMO(test_X)
"""
_default_sample_shape = torch.Size([128])
def __init__(
self,
model: Model,
sample_pareto_frontiers: Callable[[Model], Tensor],
num_fantasies: int = 16,
X_pending: Optional[Tensor] = None,
sampler: Optional[MCSampler] = None,
) -> None:
r"""Multi-objective max-value entropy search acquisition function.
Args:
model: A fitted multi-output model.
sample_pareto_frontiers: A callable that takes a model and returns a
`num_samples x n' x m`-dim tensor of outcomes to use for constructing
`num_samples` sampled Pareto frontiers.
num_fantasies: Number of fantasies to generate. The higher this
number the more accurate the model (at the expense of model
complexity, wall time and memory). Ignored if `X_pending` is `None`.
X_pending: A `m x d`-dim Tensor of `m` design points that have been
submitted for function evaluation, but have not yet been evaluated.
"""
MultiObjectiveMCAcquisitionFunction.__init__(self, model=model, sampler=sampler)
# Batch GP models (e.g. fantasized models) are not currently supported
if isinstance(model, ModelListGP):
train_X = model.models[0].train_inputs[0]
else:
train_X = model.train_inputs[0]
if train_X.ndim > 3:
raise NotImplementedError(
"Batch GP models (e.g. fantasized models) "
"are not yet supported by qMultiObjectiveMaxValueEntropy"
)
# convert to batched MO model
batched_mo_model = (
model_list_to_batched(model) if isinstance(model, ModelListGP) else model
)
self._init_model = batched_mo_model
self.fantasies_sampler = SobolQMCNormalSampler(
sample_shape=torch.Size([num_fantasies])
)
self.num_fantasies = num_fantasies
# weight is used in _compute_information_gain
self.maximize = True
self.weight = 1.0
self.sample_pareto_frontiers = sample_pareto_frontiers
# Set X_pending, register converted model and sample max values.
self.set_X_pending(X_pending)
# This avoids attribute errors in qMaxValueEntropy code.
self.posterior_transform = None
[docs]
def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
r"""Set pending points.
Informs the acquisition function about pending design points,
fantasizes the model on the pending points and draws max-value samples
from the fantasized model posterior.
Args:
X_pending: `m x d` Tensor with `m` `d`-dim design points that have
been submitted for evaluation but have not yet been evaluated.
"""
MultiObjectiveMCAcquisitionFunction.set_X_pending(self, X_pending=X_pending)
if X_pending is not None:
# fantasize the model
fantasy_model = self._init_model.fantasize(
X=X_pending,
sampler=self.fantasies_sampler,
)
self.mo_model = fantasy_model
# convert model to batched single outcome model.
self.model = batched_multi_output_to_single_output(
batch_mo_model=self.mo_model
)
self._sample_max_values()
else:
# This is mainly for setting the model to the original model
# after the sequential optimization at q > 1
self.mo_model = self._init_model
self.model = batched_multi_output_to_single_output(
batch_mo_model=self.mo_model
)
self._sample_max_values()
def _sample_max_values(self) -> None:
"""Sample max values for MC approximation of the expectation in MES.
Sets self.posterior_max_values."""
with torch.no_grad():
# num_samples x (num_fantasies) x n_pareto_points x m
sampled_pfs = self.sample_pareto_frontiers(self.mo_model)
if sampled_pfs.ndim == 3:
# add fantasy dim
sampled_pfs = sampled_pfs.unsqueeze(-3)
# take component-wise max value
self.posterior_max_values = sampled_pfs.max(dim=-2).values
[docs]
@t_batch_mode_transform(expected_q=1)
def forward(self, X: Tensor) -> Tensor:
r"""Compute max-value entropy at the design points `X`.
Args:
X: A `batch_shape x 1 x d`-dim Tensor of `batch_shape` t-batches
with `1` `d`-dim design points each.
Returns:
A `batch_shape`-dim Tensor of MVE values at the given design points `X`.
"""
# `m` dim tensor of information gains
# unsqueeze X to add a batch-dim for the batched model
igs = qMaxValueEntropy.forward(self, X=X.unsqueeze(-3))
# sum over objectives
return igs.sum(dim=-1)
[docs]
class qLowerBoundMultiObjectiveMaxValueEntropySearch(
LowerBoundMultiObjectiveEntropySearch
):
r"""The acquisition function for the multi-objective Max-value Entropy Search,
where the batches `q > 1` are supported through the lower bound formulation.
This acquisition function computes the mutual information between the observation
at a candidate point `X` and the Pareto optimal outputs.
See [Tu2022]_ for a discussion on the estimation procedure.
NOTES:
(i) The estimated acquisition value could be negative.
(ii) The lower bound batch acquisition function might not be monotone in the
sense that adding more elements to the batch does not necessarily increase the
acquisition value. Specifically, the acquisition value can become smaller when
more inputs are added.
"""
def __init__(
self,
model: Model,
hypercell_bounds: Tensor,
X_pending: Optional[Tensor] = None,
estimation_type: str = "LB",
num_samples: int = 64,
) -> None:
r"""Lower bound multi-objective max-value entropy search acquisition function.
Args:
model: A fitted batch model with 'M' number of outputs.
hypercell_bounds: A `num_pareto_samples x 2 x J x M`-dim Tensor
containing the hyper-rectangle bounds for integration, where `J` is
the number of hyper-rectangles. In the unconstrained case, this gives
the partition of the dominated space. In the constrained case, this
gives the partition of the feasible dominated space union the
infeasible space.
X_pending: A `m x d`-dim Tensor of `m` design points that have been
submitted for function evaluation, but have not yet been evaluated.
estimation_type: A string to determine which entropy estimate is
computed: "0", "LB", "LB2", or "MC".
num_samples: The number of Monte Carlo samples for the Monte Carlo
estimate.
"""
super().__init__(
model=model,
pareto_sets=None,
pareto_fronts=None,
hypercell_bounds=hypercell_bounds,
X_pending=X_pending,
estimation_type=estimation_type,
num_samples=num_samples,
)
def _compute_posterior_statistics(
self, X: Tensor
) -> dict[str, Union[GPyTorchPosterior, Tensor]]:
r"""Compute the posterior statistics.
Args:
X: A `batch_shape x q x d`-dim Tensor of inputs.
Returns:
A dictionary containing the posterior variables used to estimate the
entropy.
- "initial_entropy": A `batch_shape`-dim Tensor containing the entropy of
the Gaussian random variable `p(Y| X, D_n)`.
- "posterior_mean": A `batch_shape x num_pareto_samples x q x 1 x M`-dim
Tensor containing the posterior mean at the input `X`.
- "posterior_variance": A `batch_shape x num_pareto_samples x q x 1 x
M`-dim Tensor containing the posterior variance at the input `X`
excluding the observation noise.
- "observation_noise": A `batch_shape x num_pareto_samples x q x 1 x M`
-dim Tensor containing the observation noise at the input `X`.
- "posterior_with_noise": The posterior distribution at `X` which
includes the observation noise. This is used to compute the marginal
log-probabilities with respect to `p(y| x, D_n)` for `x` in `X`.
"""
tkwargs = {"dtype": X.dtype, "device": X.device}
CLAMP_LB = torch.finfo(tkwargs["dtype"]).eps
# Compute the initial entropy term depending on `X`.
# TODO: Below we compute posterior_plus_noise twice:
# (1) Firstly, we compute p(Y| X, D_n) when computing the initial entropy
# (2) Secondly, we compute p(y| x, D_n) for x in X in order to compute
# log(p(y|x, D_n)) for x in X in the Monte Carlo estimate..
# This could be simplified if we could evaluate log(p(y|x, D_n)) using the
# the posterior p(Y| X, D_n)
posterior_plus_noise = self.initial_model.posterior(X, observation_noise=True)
# Additional constant term.
add_term = (
0.5
* self.model.num_outputs
* (1 + torch.log(2 * pi * torch.ones(1, **tkwargs)))
)
# The variance initially has shape `batch_shape x (q*M) x (q*M)`
# prior_entropy has shape `batch_shape x num_fantasies`
initial_entropy = add_term + 0.5 * torch.logdet(
posterior_plus_noise.mvn.covariance_matrix
)
posterior_statistics = {"initial_entropy": initial_entropy}
# Compute the posterior entropy term.
posterior_plus_noise = self.model.posterior(
X.unsqueeze(-2), observation_noise=True
)
# `batch_shape x q x 1 x M`
mean = posterior_plus_noise.mean
var_plus_noise = posterior_plus_noise.variance.clamp_min(CLAMP_LB)
# Expand shapes to `batch_shape x num_pareto_samples x q x 1 x M`
new_shape = (
mean.shape[:-3] + torch.Size([self.num_pareto_samples]) + mean.shape[-3:]
)
mean = mean.unsqueeze(-4).expand(new_shape)
var_plus_noise = var_plus_noise.unsqueeze(-4).expand(new_shape)
# TODO: This computes the observation noise via a second evaluation of the
# posterior. This step could be done better.
posterior = self.model.posterior(X.unsqueeze(-2), observation_noise=False)
var = posterior.variance.clamp_min(CLAMP_LB)
var = var.unsqueeze(-4).expand(new_shape)
obs_noise = var_plus_noise - var
posterior_statistics["posterior_mean"] = mean
posterior_statistics["posterior_variance"] = var
posterior_statistics["observation_noise"] = obs_noise
posterior_statistics["posterior_with_noise"] = posterior_plus_noise
return posterior_statistics
def _compute_monte_carlo_variables(
self, posterior: GPyTorchPosterior
) -> Tuple[Tensor, Tensor]:
r"""Compute the samples and log-probability associated with a posterior
distribution.
Args:
posterior: The posterior distribution, which includes the observation
noise.
Returns:
A two-element tuple containing
- samples: A `num_mc_samples x batch_shape x num_pareto_samples x q x 1
x M`-dim Tensor containing the Monte Carlo samples.
- samples_log_prob: A `num_mc_samples x batch_shape x num_pareto_samples
x q`-dim Tensor containing the log-probabilities of the Monte Carlo
samples.
"""
# `num_mc_samples x batch_shape x q x 1 x M`
samples = self.get_posterior_samples(posterior)
# `num_mc_samples x batch_shape x q`
if self.model.num_outputs == 1:
samples_log_prob = posterior.mvn.log_prob(samples.squeeze(-1))
else:
samples_log_prob = posterior.mvn.log_prob(samples)
# Expand shape to `num_mc_samples x batch_shape x num_pareto_samples x
# q x 1 x M`
new_shape = (
samples.shape[:-3]
+ torch.Size([self.num_pareto_samples])
+ samples.shape[-3:]
)
samples = samples.unsqueeze(-4).expand(new_shape)
# Expand shape to `num_mc_samples x batch_shape x num_pareto_samples x q`
new_shape = (
samples_log_prob.shape[:-1]
+ torch.Size([self.num_pareto_samples])
+ samples_log_prob.shape[-1:]
)
samples_log_prob = samples_log_prob.unsqueeze(-2).expand(new_shape)
return samples, samples_log_prob
[docs]
@concatenate_pending_points
@t_batch_mode_transform()
def forward(self, X: Tensor) -> Tensor:
r"""Evaluates qLowerBoundMultiObjectiveMaxValueEntropySearch at the design
points `X`.
Args:
X: A `batch_shape x q x d`-dim Tensor of `batch_shape` t-batches with `q`
`d`-dim design points each.
Returns:
A `batch_shape`-dim Tensor of acquisition values at the given design
points `X`.
"""
return self._compute_lower_bound_information_gain(X)
