Source code for botorch.acquisition.multi_objective.max_value_entropy_search
#!/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 typing import Any, Callable, Optional
import torch
from botorch.acquisition.max_value_entropy_search import qMaxValueEntropy
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.sampling.samplers import MCSampler, SobolQMCNormalSampler
from botorch.utils.transforms import 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.
Example:
>>> model = SingleTaskGP(train_X, train_Y)
>>> MESMO = qMultiObjectiveMaxValueEntropy(model, sample_pfs)
>>> mesmo = MESMO(test_X)
"""
def __init__(
self,
model: Model,
sample_pareto_frontiers: Callable[[Model], Tensor],
num_fantasies: int = 16,
X_pending: Optional[Tensor] = None,
sampler: Optional[MCSampler] = None,
**kwargs: Any,
) -> 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.mo_model = batched_mo_model
self.model = batched_multi_output_to_single_output(
batch_mo_model=batched_mo_model
)
self.fantasies_sampler = SobolQMCNormalSampler(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
# this avoids unnecessary model conversion if X_pending is None
if X_pending is None:
self._sample_max_values()
else:
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, observation_noise=True
)
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:
r"""Sample max values for MC approximation of the expectation in MES"""
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)
