Advice on Multi-Target Optimization (MultiFidelitySingleTaskGP + qMultiFidelityKnowledgeGradient)

[AlexanderMouton]

Hi,

I am currently working on an optimization problem involving a 13-dimensional vector. This vector comprises:

• 8 continuous dimensions representing problem variables.
• 3 binary dimensions indicating the active status of three distinct components (e.g., 1 1 0 means the first two components are active).
• 2 continuous fidelity dimensions.

My objective is to optimize across all 8 possible combinations of the binary dimensions (representing the active/inactive states of the components).

I'm utilizing the MultiFidelitySingleTaskGP model and the qMultiFidelityKnowledgeGradient acquisition function (AF), with a projection to the highest fidelity. From my understanding, the Knowledge Gradient AF aims to identify points that maximize information gain concerning the target fidelity. However, my goal is different: I want to optimize the information gain across all binary combinations simultaneously for the target fidelity, not just the target fidelity.

I realized the projection capability of the AF isn't limited to just fidelity dimensions - It is possible to project the AF onto a specific binary setting along with the highest fidelity (like 1 1 0 1.0 1.0).

Is it feasible to create an AF for each binary target and then combine them into a single AF that can be passed into optimize_acqf()?

Any insights or suggestions would be appreciated, thanks!
Alex

[AlexanderMouton]

Hi,

For the time being I've made a slightly altered version of the qMultiFidelityKnowledgeGradient class that seems to do the trick:

from __future__ import annotations

from copy import deepcopy
from typing import Any, Callable, Dict, Optional, Type, List

import torch
from botorch import settings
from botorch.acquisition.acquisition import (
    AcquisitionFunction,
)
from botorch.acquisition.cost_aware import CostAwareUtility
from botorch.acquisition.objective import MCAcquisitionObjective, PosteriorTransform
from botorch.exceptions.errors import UnsupportedError
from botorch.models.model import Model
from botorch.sampling.base import MCSampler
from botorch.utils.transforms import (
    match_batch_shape,
    t_batch_mode_transform,
)
from torch import Tensor

from botorch.acquisition.knowledge_gradient import qKnowledgeGradient, _split_fantasy_points, _get_value_function

class qMultiTargetMultiFidelityKnowledgeGradient(qKnowledgeGradient):
    r"""Batch Knowledge Gradient for multi-target multi-fidelity optimization.

    A version of `qKnowledgeGradient` that supports multi-target multi-fidelity
    optimization via a `CostAwareUtility` and the `project` and `expand` operators.
    If none of these are set, this acquisition function reduces to `qKnowledgeGradient`.
    Through `valfunc_cls` and `valfunc_argfac`, this can be changed into a custom
    multi-fidelity acquisition function (it is only KG if the terminal value is
    computed using a posterior mean).
    """

    def __init__(
        self,
        model: Model,
        num_fantasies: Optional[int] = 64,
        sampler: Optional[MCSampler] = None,
        objective: Optional[MCAcquisitionObjective] = None,
        posterior_transform: Optional[PosteriorTransform] = None,
        inner_sampler: Optional[MCSampler] = None,
        X_pending: Optional[Tensor] = None,
        current_value: Optional[Tensor] = None,
        cost_aware_utility: Optional[CostAwareUtility] = None,
        projections: List[Callable[[Tensor], Tensor]] = [lambda X: X],
        expand: Callable[[Tensor], Tensor] = lambda X: X,
        valfunc_cls: Optional[Type[AcquisitionFunction]] = None,
        valfunc_argfac: Optional[Callable[[Model], Dict[str, Any]]] = None,
    ) -> None:
        r"""Multi-Fidelity q-Knowledge Gradient (one-shot optimization).

        Args:
            model: A fitted model. Must support fantasizing.
            num_fantasies: The number of fantasy points to use. More fantasy
                points result in a better approximation, at the expense of
                memory and wall time. Unused if `sampler` is specified.
            sampler: The sampler used to sample fantasy observations. Optional
                if `num_fantasies` is specified.
            objective: The objective under which the samples are evaluated. If
                `None`, then the analytic posterior mean is used. Otherwise, the
                objective is MC-evaluated (using inner_sampler).
            posterior_transform: An optional PosteriorTransform. If given, this
                transforms the posterior before evaluation. If `objective is None`,
                then the analytic posterior mean of the transformed posterior is
                used. If `objective` is given, the `inner_sampler` is used to draw
                samples from the transformed posterior, which are then evaluated under
                the `objective`.
            inner_sampler: The sampler used for inner sampling. Ignored if the
                objective is `None`.
            X_pending: A `m x d`-dim Tensor of `m` design points that have
                points that have been submitted for function evaluation
                but have not yet been evaluated.
            current_value: The current value, i.e. the expected best objective
                given the observed points `D`. If omitted, forward will not
                return the actual KG value, but the expected best objective
                given the data set `D u X`.
            cost_aware_utility: A CostAwareUtility computing the cost-transformed
                utility from a candidate set and samples of increases in utility.
            projections: A List of callables mapping a `batch_shape x q x d` tensor
                of design points to a tensor with shape `batch_shape x q_term x d`
                projected to the desired target set (e.g. the target fidelities in
                case of multi-fidelity optimization). For the basic case,
                `q_term = q`.
            expand: A callable mapping a `batch_shape x q x d` input tensor to
                a `batch_shape x (q + q_e)' x d`-dim output tensor, where the
                `q_e` additional points in each q-batch correspond to
                additional ("trace") observations.
            valfunc_cls: An acquisition function class to be used as the terminal
                value function.
            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`).
        """
        if current_value is None and cost_aware_utility is not None:
            raise UnsupportedError(
                "Cost-aware KG requires current_value to be specified."
            )
        super().__init__(
            model=model,
            num_fantasies=num_fantasies,
            sampler=sampler,
            objective=objective,
            posterior_transform=posterior_transform,
            inner_sampler=inner_sampler,
            X_pending=X_pending,
            current_value=current_value,
        )
        self.cost_aware_utility = cost_aware_utility
        self.projections = projections
        self.expand = expand
        self._cost_sampler = None
        self.valfunc_cls = valfunc_cls
        self.valfunc_argfac = valfunc_argfac

    @property
    def cost_sampler(self):
        if self._cost_sampler is None:
            # Note: Using the deepcopy here is essential. Removing this poses a
            # problem if the base model and the cost model have a different number
            # of outputs or test points (this would be caused by expand), as this
            # would trigger re-sampling the base samples in the fantasy sampler.
            # By cloning the sampler here, the right thing will happen if the
            # the sizes are compatible, if they are not this will result in
            # samples being drawn using different base samples, but it will at
            # least avoid changing state of the fantasy sampler.
            self._cost_sampler = deepcopy(self.sampler)
        return self._cost_sampler

    @t_batch_mode_transform()
    def forward(self, X: Tensor) -> Tensor:
        r"""Evaluate qMultiFidelityKnowledgeGradient on the candidate set `X`.

        Args:
            X: A `b x (q + num_fantasies) x d` Tensor with `b` t-batches of
                `q + num_fantasies` design points each. We split this X tensor
                into two parts in the `q` dimension (`dim=-2`). The first `q`
                are the q-batch of design points and the last num_fantasies are
                the current solutions of the inner optimization problem.

                `X_fantasies = X[..., -num_fantasies:, :]`
                `X_fantasies.shape = b x num_fantasies x d`

                `X_actual = X[..., :-num_fantasies, :]`
                `X_actual.shape = b x q x d`

                In addition, `X` may be augmented with fidelity parameteres as
                part of thee `d`-dimension. Projecting fidelities to the target
                fidelity is handled by `project`.

        Returns:
            A Tensor of shape `b`. For t-batch b, the q-KG value of the design
                `X_actual[b]` is averaged across the fantasy models, where
                `X_fantasies[b, i]` is chosen as the final selection for the
                `i`-th fantasy model.
                NOTE: If `current_value` is not provided, then this is not the
                true KG value of `X_actual[b]`, and `X_fantasies[b, : ]` must be
                maximized at fixed `X_actual[b]`.
        """
        X_actual, X_fantasies = _split_fantasy_points(X=X, n_f=self.num_fantasies)

        # We only concatenate X_pending into the X part after splitting
        if self.X_pending is not None:
            X_eval = torch.cat(
                [X_actual, match_batch_shape(self.X_pending, X_actual)], dim=-2
            )
        else:
            X_eval = X_actual

        # construct the fantasy model of shape `num_fantasies x b`
        # expand X (to potentially add trace observations)
        fantasy_model = self.model.fantasize(
            X=self.expand(X_eval),
            sampler=self.sampler,
        )

        values_list = []
        for project in self.projections:
            # get the value function for the current projection
            value_function = _get_value_function(
                model=fantasy_model,
                objective=self.objective,
                posterior_transform=self.posterior_transform,
                sampler=self.inner_sampler,
                project=project,
                valfunc_cls=self.valfunc_cls,
                valfunc_argfac=self.valfunc_argfac,
            )

            # make sure to propagate gradients to the fantasy model train inputs
            # project the fantasy points
            with settings.propagate_grads(True):
                values = value_function(X=X_fantasies)  # num_fantasies x b

            if self.current_value is not None:
                values = values - self.current_value

            if self.cost_aware_utility is not None:
                values = self.cost_aware_utility(
                    X=X_actual, deltas=values, sampler=self.cost_sampler
                )
            
            values_list.append(values.mean(dim=0))

        # return average over the fantasy samples
        return torch.stack(values_list, dim=0).mean(dim=0)

The alteration just allows passing in multiple projection functions and averaging the AF value for all of them.

Would this change be worth making a PR for?
The same can be applied to the qMultiFidelityMaxValueEntropy and qMultiFidelityLowerBoundMaxValueEntropy classes.

[AlexanderMouton]

customAcquisition.zip

Multi-target versions of the qMFKG, qMFMVE and qMFLBMVE acquisition functions.

[sdaulton]

Cool problem! Could you share more details about your problem? I am curious how often this pops up. cc @Balandat

[AlexanderMouton]

Hi @sdaulton (cc @Balandat )

I am tuning a parallel MCTS algorithm on combinatorial games with various enhancements that can be turned on/off (e.g. MAST/RAVE/CMC). The fidelity dimensions relate to the time given per move and the number of games to play, i.e. the target would be to optimise for 1s/move and 50 games, but smaller time durations and/or fewer games are cheaper to evaluate.

A single enhancement may entirely outperform the other enhancements on a given combinatorial game, meaning that the BO algorithm will usually focus only on maximising for that enhancement; so projecting to each possible binary configuration (hopefully) solves this problem.

If no enhancements are turned on, only the first dimension influences the outcome; if only the first enhancement is turned on, the first 3 dimensions influence the outcome; if only the second enhancement is turned on, dimensions 0, 3 and 4 influence the outcome; etc.

[Balandat]

Hmm very interesting setting - I haven't seen this particular one before. The approach makes sense to me (although I haven't spent a ton of time thinking through any of the details).

The overall problem appears to be related to a contextual optimization problem (in your case the particular set of enhancements are the contexts) in which one can leverage the context observations to increase the model quality (i.e. borrow strength across contexts) but the goal is to a policy or a set of policies that works well for all/most contexts. @AlexanderMouton would that be a fair characterization?

[Balandat]

Would this change be worth making a PR for?

Hmm I'd probably be hesitant to put this in the core package since it is a bit of a non-standard use case. But we have a new directory for community contributions (#2167) for which this could be a good fit.

[Balandat]

cc @bletham, @qingfeng10 in case this or a similar problem has come up / will come up in the (pun not intended) context of contextual BO.

[AlexanderMouton]

Hi @Balandat, thanks for the feedback!

Yes, that would be a fair characterization.

I will get round to making a community contribution PR once I've submitted the next section of my thesis to my supervisors!

[AlexanderMouton]

@Balandat cc @sdaulton

One side effect of projecting to multiple contexts that comes to mind is that, when constructing the qMultiFidelityKnowledgeGradient AF, the current_value, or 'reference' value, should be calculated differently.

Usually, the reference value is found by fixing the fidelity dimensions to the highest fidelity (as done in this tutorial) and then the AF is optimized via optimize_acqf, but now the reference value will have to be calculated as the average acquisition value after running optimize_acqf for each of the contexts to keep things consistent. Does this seem sensible?

Thanks!