Implementing Expected Further Improvement Acquisition Function

[samuelkim16]

I am interested in implementing the Expected Further Improvement (EFI) acquisition function that was proposed in this paper for multi-fidelity Bayesian optimization: A multi-fidelity Bayesian optimization approach based on the expected further improvement. The acquisition function reduces to normal EI for points at the highest fidelity. For lower fidelities, it is the expected value of the reduction of EI at the highest fidelity when conditioned on adding the lower fidelity data point.

My attempt at implementing this acquisition function in BoTorch. I have two questions on this:

1. It is hacked together and feels messy. I know the documentation for the forward() function in the SampleReducingMCAcquisitionFunction class said not to overwrite it, although I did not see a way around it. Any recommendations on how to clean this up or if there is a more proper way to do it?
2. I get several warnings when trying to use this acquisition function in an optimization loop (shown below the implementation code). I'm guessing this is because the gradients are not propagating properly or I implemented something incorrectly. Any tips on how to troubleshoot this?

class qMultiFidelityExpectedFurtherImprovement(qExpectedImprovement):
    r"""Batch Expected Further Improvement for multi-fidelity optimization, as described in Shu et al. (2021)

    A version of `qExpectedImprovement` that supports multi-fidelity optimization
    via a `CostAwareUtility` and the `project` and `expand` operators. If none
    of these are set, this acquisition function reduces to `qExpectedImprovement`.
    
    Citation:
        Shu, Leshi, Ping Jiang, and Yan Wang. 
        "A multi-fidelity Bayesian optimization approach based on the expected further improvement." 
        Structural and Multidisciplinary Optimization 63.4 (2021): 1709-1719.
    """

    def __init__(
        self,
        model: Model,
        best_f: Union[float, Tensor],
        num_fantasies: Optional[int] = 64,
        sampler: Optional[MCSampler] = None,
        objective: Optional[MCAcquisitionObjective] = None,
        posterior_transform: Optional[PosteriorTransform] = None,
        X_pending: Optional[Tensor] = None,
        cost_aware_utility: Optional[CostAwareUtility] = None,
        constraints: Optional[list[Callable[[Tensor], Tensor]]] = None,
        eta: Union[Tensor, float] = 1e-3,
        project: Callable[[Tensor], Tensor] = lambda X: X,
    ) -> None:
        r"""Multi-Fidelity Expected Further Improvement.

        Args:
            model: A fitted model. Must support fantasizing.
            best_f: The best objective value observed so far (assumed noiseless). Can be
                a scalar, or a `batch_shape`-dim tensor. In case of a batched model, the
                tensor can specify different values for each element of the batch.
            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.
                Defaults to `IdentityMCObjective()`.
            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`.
            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.
            cost_aware_utility: A CostAwareUtility computing the cost-transformed
                utility from a candidate set and samples of increases in utility.
            project: A callable 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`.
        """
        if best_f is None and cost_aware_utility is not None:
            raise UnsupportedError(
                "Cost-aware EI requires best_f to be specified."
            )
        super().__init__(
            model=model,
            best_f=best_f,
            sampler=sampler,
            objective=objective,
            posterior_transform=posterior_transform,
            X_pending=X_pending,
            constraints=constraints,
            eta=eta,
        )
        
        if sampler is None:
            if num_fantasies is None:
                raise ValueError(
                    "Must specify `num_fantasies` if no `sampler` is provided."
                )
            # base samples should be fixed for joint optimization over X, X_fantasies
            sampler = SobolQMCNormalSampler(sample_shape=torch.Size([num_fantasies]))
        elif num_fantasies is not None:
            if sampler.sample_shape != torch.Size([num_fantasies]):
                raise ValueError(
                    f"The sampler shape must match num_fantasies={num_fantasies}."
                )
        else:
            num_fantasies = sampler.sample_shape[0]
        MCSamplerMixin.__init__(self, sampler=sampler)
            
        self.cost_aware_utility = cost_aware_utility
        self.project = project
        self._cost_sampler = None

    @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

    def _get_samples_and_objectives(self, X: Tensor, model: Model=None) -> tuple[Tensor, Tensor]:
        """Computes posterior samples and objective values at input X.

        Args:
            X: A `batch_shape x q x d`-dim Tensor of model inputs.

        Returns:
            A two-tuple `(samples, obj)`, where `samples` is a tensor of posterior
            samples with shape `sample_shape x batch_shape x q x m`, and `obj` is a
            tensor of MC objective values with shape `sample_shape x batch_shape x q`.
        """
        if model is None:
            model = self.model
        
        posterior = model.posterior(
            X=X, posterior_transform=self.posterior_transform
        )
        samples = self.get_posterior_samples(posterior)
        obj = self.objective(samples=samples, X=X)
        return samples, obj

    def _non_reduced_forward(self, X: Tensor, model: Model=None) -> Tensor:
        """Compute the constrained acquisition values at the MC-sample, q level.

        Args:
            X: A `batch_shape x q x d` Tensor of t-batches with `q` `d`-dim
                design points each.

        Returns:
            A Tensor with shape `sample_sample x batch_shape x q`.
        """
        if model is None:
            model = self.model
        
        samples, obj = self._get_samples_and_objectives(X, model=model)
        samples = repeat_to_match_aug_dim(target_tensor=samples, reference_tensor=obj)
        acqval = self._sample_forward(obj)  # `sample_sample x batch_shape x q`
        return self._apply_constraints(acqval=acqval, samples=samples)

    @concatenate_pending_points
    @t_batch_mode_transform()
    def forward(self, X: Tensor) -> Tensor:
        r"""Evaluate qMultiFidelityExpectedImprovement on the candidate set `obj`.
        """
        
        non_reduced_acqval_current = self._non_reduced_forward(X)

        # construct the fantasy model of shape `num_fantasies x b`
        fantasy_model = self.model.fantasize(
            X=X,
            sampler=self.sampler,
        )
        
        X_proj = self.project(X)
        
        non_reduced_acqval_conditioned = self._non_reduced_forward(X=X_proj, model=fantasy_model)
        non_reduced_acqval_conditioned = torch.mean(non_reduced_acqval_conditioned, dim=1)
        
        non_reduced_acqval = non_reduced_acqval_current - non_reduced_acqval_conditioned
        
        values = self._sample_reduction(self._q_reduction(non_reduced_acqval))

        if self.cost_aware_utility is not None:
            values = self.cost_aware_utility(
                X=X, deltas=values, sampler=self.cost_sampler
            )

        return values

c:\Users\username\Documents\project_dir\.venv\Lib\site-packages\botorch\optim\initializers.py:433: BadInitialCandidatesWarning: Unable to find non-zero acquisition function values - initial conditions are being selected randomly.
  warnings.warn(
c:\Users\username\Documents\project_dir\.venv\Lib\site-packages\botorch\optim\optimize.py:568: RuntimeWarning: Optimization failed in `gen_candidates_scipy` with the following warning(s):
[OptimizationWarning('Optimization failed within `scipy.optimize.minimize` with status 2 and message ABNORMAL_TERMINATION_IN_LNSRCH.')]
Trying again with a new set of initial conditions.
  return _optimize_acqf_batch(opt_inputs=opt_inputs)
c:\Users\username\Documents\project_dir\.venv\Lib\site-packages\botorch\optim\optimize.py:568: RuntimeWarning: Optimization failed on the second try, after generating a new set of initial conditions.
  return _optimize_acqf_batch(opt_inputs=opt_inputs)

[Balandat]

I know the documentation for the forward() function in the SampleReducingMCAcquisitionFunction class said not to overwrite it, although I did not see a way around it. Any recommendations on how to clean this up or if there is a more proper way to do it?

You're right that with the current setup that does not work. The proper way to handle this is to support cost-weighting of the acquisition function in the base SampleReducingMCAcquisitionFunction class, essentially applying this block

        if self.cost_aware_utility is not None:
            values = self.cost_aware_utility(
                X=X, deltas=values, sampler=self.cost_sampler
            )

in SampleReducingMCAcquisitionFunction.forward(). @SebastianAment, @ltiao does that sound right? We will need to support this anyway (either in this or some other form).

I get several warnings when trying to use this acquisition function in an optimization loop (shown below the implementation code). I'm guessing this is because the gradients are not propagating properly or I implemented something incorrectly. Any tips on how to troubleshoot this?

The "Unable to find non-zero acquisition function values" suggests that your acquisition function is numerically zero most everywhere. This is well known and fundamental problem of EI, see our recent paper for a detailed discussion. BoTorch has qLogExpectedImprovement and the noisy variant qLogNoisyExpectedImprovement for this. Cost-weighting either of those is relatively straightforward (we'd do logEI(x) - log(cost(x)) rather than log(EI(x)/cost(x))) - the equivalent here would be that we'd like to have something like logEFI that expresses the log-further-expected-improvement in a less numerically dicey fashion.

[saitcakmak]

The proper way to handle this is to support cost-weighting of the acquisition function in the base SampleReducingMCAcquisitionFunction class, essentially applying this block in SampleReducingMCAcquisitionFunction.forward().

This is what I have in my latest prototype as well, though it doesn't extend to qLogNEHVI, since it overwrites forward.

I have been thinking of some alternatives for applying these acquisition utilities more broadly. I don't really like having to update each acqf definition just to apply a utility to the acquisition value.