Source code for botorch.acquisition.proximal

#!/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.

A wrapper around AcquisitionFunctions to add proximal weighting of the
acquisition function.

from __future__ import annotations

from typing import Optional

import torch
from botorch.acquisition import AcquisitionFunction
from botorch.exceptions.errors import UnsupportedError
from botorch.models import ModelListGP
from botorch.models.model import Model
from botorch.models.transforms.input import InputTransform
from botorch.utils import t_batch_mode_transform
from torch import Tensor
from torch.nn import Module

[docs]class ProximalAcquisitionFunction(AcquisitionFunction): """A wrapper around AcquisitionFunctions to add proximal weighting of the acquisition function. Acquisition function is weighted via a squared exponential centered at the last training point, with varying lengthscales corresponding to `proximal_weights`. Can only be used with acquisition functions based on single batch models. Small values of `proximal_weights` corresponds to strong biasing towards recently observed points, which smoothes optimization with a small potential decrese in convergence rate. Example: >>> model = SingleTaskGP(train_X, train_Y) >>> EI = ExpectedImprovement(model, best_f=0.0) >>> proximal_weights = torch.ones(d) >>> EI_proximal = ProximalAcquisitionFunction(EI, proximal_weights) >>> eip = EI_proximal(test_X) """ def __init__( self, acq_function: AcquisitionFunction, proximal_weights: Tensor, ) -> None: r"""Derived Acquisition Function weighted by proximity to recently observed point. Args: acq_function: The base acquisition function, operating on input tensors of feature dimension `d`. proximal_weights: A `d` dim tensor used to bias locality along each axis. """ Module.__init__(self) self.acq_func = acq_function model = self.acq_func.model if hasattr(acq_function, "X_pending"): if acq_function.X_pending is not None: raise UnsupportedError( "Proximal acquisition function requires `X_pending` to be None." ) self.X_pending = acq_function.X_pending self.register_buffer("proximal_weights", proximal_weights) _validate_model(model, proximal_weights)
[docs] @t_batch_mode_transform(expected_q=1, assert_output_shape=False) def forward(self, X: Tensor) -> Tensor: r"""Evaluate base acquisition function with proximal weighting. Args: X: Input tensor of feature dimension `d` . Returns: Base acquisition function evaluated on tensor `X` multiplied by proximal weighting. """ model = self.acq_func.model train_inputs = model.train_inputs[0] if isinstance(model, ModelListGP): train_inputs = train_inputs[0] model = model.models[0] input_transform = _get_input_transform(model) last_X = train_inputs[-1].reshape(1, 1, -1) # un-transform last_X if input_transform is not None: last_X = input_transform.untransform(last_X) diff = X - last_X M = torch.linalg.norm(diff / self.proximal_weights, dim=-1) ** 2 proximal_acq_weight = torch.exp(-0.5 * M) return self.acq_func(X) * proximal_acq_weight.flatten()
def _validate_model(model: Model, proximal_weights: Tensor) -> None: r"""Validate model Perform vaidation checks on model used in base acquisition function to make sure it is compatible with proximal weighting. Args: model: Model associated with base acquisition function to be validated. proximal_weights: A `d` dim tensor used to bias locality along each axis. """ # check model for train_inputs and single batch if not hasattr(model, "train_inputs"): raise UnsupportedError("Acquisition function model must have `train_inputs`.") # get train inputs for each type of possible model if isinstance(model, ModelListGP): # ModelListGP models # check to make sure that the training inputs and input transformers for each # model match and are reversible train_inputs = model.train_inputs[0][0] input_transform = _get_input_transform(model.models[0]) for i in range(len(model.train_inputs)): if not torch.equal(train_inputs, model.train_inputs[i][0]): raise UnsupportedError( "Proximal acquisition function does not support unequal " "training inputs" ) if not input_transform == _get_input_transform(model.models[i]): raise UnsupportedError( "Proximal acquisition function does not support non-identical " "input transforms" ) else: # any non-ModelListGP model train_inputs = model.train_inputs[0] # check to make sure that the model is single t-batch (q-batches are allowed) if model.batch_shape != torch.Size([]) and train_inputs.shape[1] != 1: raise UnsupportedError( "Proximal acquisition function requires a single batch model" ) # check to make sure that weights match the training data shape if ( len(proximal_weights.shape) != 1 or proximal_weights.shape[0] != train_inputs.shape[-1] ): raise ValueError( "`proximal_weights` must be a one dimensional tensor with " "same feature dimension as model." ) def _get_input_transform(model: Model) -> Optional[InputTransform]: """get input transform if defined""" try: return model.input_transform except AttributeError: return None