Skip to content

Expected Improvement

decijax.acquisition_functions.expected_improvement

Expected improvement acquisition functions.

ExpectedImprovement

Bases: AbstractSinglePointAcquisitionFunctionBuilder

Standard Expected Improvement acquisition function.

As introduced by MoÄŤkus, 1974. The "best" incumbent value is defined as the highest posterior mean value evaluated at the previously observed points. This enables the acquisition function to be utilised with noisy observations.

build_acquisition_function 

build_acquisition_function(
    models: Mapping[str, ProbabilisticModel], key: KeyArray
) -> SinglePointAcquisitionFunction

Build the Expected Improvement acquisition function.

This computes the expected improvement over the "best" of the previously observed points, utilising the posterior distribution of the surrogate model. For posterior distribution $f(\cdot)$, and best incumbent value $\eta$, this is defined as:

\alpha_{\text{EI}}(\mathbf{x})
= \mathbb{E}\left[\max(0, f(\mathbf{x}) - \eta)\right]

For models carrying a leading sample axis (e.g. fully Bayesian GPs), the expected improvement is computed per sample and averaged, which is the correct marginalisation $\mathbb{E}_\theta[\alpha_{\text{EI},\theta}(\mathbf{x})]$.

Parameters:

  • models (Mapping[str, ProbabilisticModel]) –

    Dictionary of models used to form the acquisition function. One model must correspond to the OBJECTIVE key and support Gaussian prediction, as we use the objective posterior to form the acquisition function.

  • key (KeyArray) –

    JAX PRNG key used for random number generation. Since the expected improvement is computed deterministically, the key is not used.

Returns:

LogExpectedImprovement

Bases: AbstractSinglePointAcquisitionFunctionBuilder

Numerically stable Log Expected Improvement acquisition function [1].

Given a predictive posterior distribution of the objective function \(f\), the log expected improvement at a test point \(x\) is defined as:

\[\text{LogEI}(x) = \log \mathbb{E}\left[\max(0, f(x) - f(x^*))\right]\]

where \(x^*\) is the maximiser of the posterior mean at previously observed values (to handle noisy observations).

Being a strictly increasing transform of the expected improvement, this shares its maximiser exactly, but is far better behaved as an optimisation target: expected improvement vanishes to exactly zero, gradient included, once the scaled improvement falls below roughly \(-40\) in double precision, and those flat regions come to dominate the search space in higher dimensions.

References:

[1] Ament, S., Daulton, S., Eriksson, D., Balandat, M., & Bakshy, E. (2023). Unexpected improvement to expected improvement for Bayesian optimization. Advances in Neural Information Processing Systems, 36.

build_acquisition_function 

build_acquisition_function(
    models: Mapping[str, ProbabilisticModel], key: KeyArray
) -> SinglePointAcquisitionFunction

Build the Log Expected Improvement acquisition function.

The expected improvement factorises as $\sigma \cdot h(z)$, for scaled improvement $z = \frac{\mu - \eta}{\sigma}$ and $h(z) = \phi(z) + z\Phi(z)$, so that:

\alpha_{\text{LogEI}}(\mathbf{x}) = \log \sigma(\mathbf{x}) + \log h(z)

with the second term computed by _log_ei_helper. For models carrying a leading sample axis (e.g. fully Bayesian GPs), it is computed per sample and reduced with a log-mean-exp, the correct marginalisation $\log \mathbb{E}_\theta[\alpha_{\text{EI},\theta}(\mathbf{x})]$.

Parameters:

  • models (Mapping[str, ProbabilisticModel]) –

    Dictionary of models used to form the acquisition function. One model must correspond to the OBJECTIVE key and support Gaussian prediction, as we use the objective posterior to form the acquisition function.

  • key (KeyArray) –

    JAX PRNG key used for random number generation. Since the log expected improvement is computed deterministically, the key is not used.

Returns: