Acquisition functions

gpax.acquisition.UCB(rng_key, model, X, beta=0.25, maximize=False, n=1, noiseless=False, penalty=None, recent_points=None, grid_indices=None, penalty_factor=1.0, **kwargs)

Upper confidence bound

Given a probabilistic model \(m\) that models the objective function \(f\), the Upper Confidence Bound at an input point \(x\) is defined as:

\[UCB(x) = \mu(x) + \kappa \sigma(x)\]

where \(\mu(x)\) is the predictive mean, \(\sigma(x)\) is the predictive standard deviation, and \(\kappa\) is the exploration-exploitation trade-off parameter.

In the case of HMC, the function leverages multiple predictive posteriors, each associated with a different HMC sample of the GP model parameters, to capture both prediction uncertainty and hyperparameter uncertainty. In this setup, the uncertainty in parameters of probabilistic mean function (if any) also contributes to the acquisition function values.

Parameters:

• rng_key (Array) – JAX random number generator key

• model (Type[ExactGP]) – trained model

• X (Array) – new inputs

• beta (float) – coefficient balancing exploration-exploitation trade-off

• maximize (bool) – If True, assumes that BO is solving maximization problem

• n (int) – number of samples drawn from each MVN distribution (number of distributions is equal to the number of HMC samples)

• noiseless (bool) – Noise-free prediction. It is set to False by default as new/unseen data is assumed to follow the same distribution as the training data. Hence, since we introduce a model noise for the training data, we also want to include that noise in our prediction.

• penalty (Optional[str]) –

  Penalty applied to the acquisition function to discourage re-evaluation at or near points that were recently evaluated. Options are:

  • ’delta’:

  The infinite penalty is applied to the recently visited points.

  • ’inverse_distance’:

  Modifies the acquisition function by penalizing points near the recent points.

  For the ‘inverse_distance’, the acqusition function is penalized as:

  \[\alpha - \lambda \cdot \pi(X, r)\]

  where \(\pi(X, r)\) computes a penalty for points in \(X\) based on their distance to recent points \(r\), \(\alpha\) represents the acquisition function, and \(\lambda\) represents the penalty factor.

• recent_points (Array) – An array of recently visited points [oldest, …, newest] provided by user

• grid_indices (Array) – Grid indices of data points in X array for the penalty term calculation. For example, if each data point is an image patch, the indices could correspond to the (i, j) pixel coordinates of their centers in the original image.

• penalty_factor (float) – Penalty factor \(\lambda\) in \(\alpha - \lambda \cdot \pi(X, r)\)

• **jitter – Small positive term added to the diagonal part of a covariance matrix for numerical stability (Default: 1e-6)

Return type:

Array

gpax.acquisition.EI(rng_key, model, X, best_f=None, maximize=False, n=1, noiseless=False, penalty=None, recent_points=None, grid_indices=None, penalty_factor=1.0, **kwargs)

Expected Improvement

Given a probabilistic model \(m\) that models the objective function \(f\), the Expected Improvement at an input point \(x\) is defined as:

\[\begin{split}EI(x) = \begin{cases} (\mu(x) - f^+) \Phi(Z) + \sigma(x) \phi(Z) & \text{if } \sigma(x) > 0 \\ 0 & \text{if } \sigma(x) = 0 \end{cases}\end{split}\]

where \(\mu(x)\) is the predictive mean, \(\sigma(x)\) is the predictive standard deviation, \(f^+\) is the value of the best observed sample. \(Z\) is defined as:

\[Z = \frac{\mu(x) - f^+}{\sigma(x)}\]

provided \(\sigma(x) > 0\).

In the case of HMC, the function leverages multiple predictive posteriors, each associated with a different HMC sample of the GP model parameters, to capture both prediction uncertainty and hyperparameter uncertainty. In this setup, the uncertainty in parameters of probabilistic mean function (if any) also contributes to the acquisition function values.

Parameters:

• rng_key (Array) – JAX random number generator key

• model (Type[ExactGP]) – trained model

• X (Array) – new inputs

• best_f (float) – Best function value observed so far. Derived from the predictive mean when not provided by a user.

• maximize (bool) – If True, assumes that BO is solving maximization problem

• n (int) – number of samples drawn from each MVN distribution (number of distributions is equal to the number of HMC samples)

• noiseless (bool) – Noise-free prediction. It is set to False by default as new/unseen data is assumed to follow the same distribution as the training data. Hence, since we introduce a model noise for the training data, we also want to include that noise in our prediction.

• penalty (Optional[str]) –

  Penalty applied to the acquisition function to discourage re-evaluation at or near points that were recently evaluated. Options are:

  • ’delta’:

  The infinite penalty is applied to the recently visited points.

  • ’inverse_distance’:

  Modifies the acquisition function by penalizing points near the recent points.

  For the ‘inverse_distance’, the acqusition function is penalized as:

  \[\alpha - \lambda \cdot \pi(X, r)\]

  where \(\pi(X, r)\) computes a penalty for points in \(X\) based on their distance to recent points \(r\), \(\alpha\) represents the acquisition function, and \(\lambda\) represents the penalty factor.

• recent_points (Array) – An array of recently visited points [oldest, …, newest] provided by user

• grid_indices (Array) – Grid indices of data points in X array for the penalty term calculation. For example, if each data point is an image patch, the indices could correspond to the (i, j) pixel coordinates of their centers in the original image.

• penalty_factor (float) – Penalty factor \(\lambda\) in \(\alpha - \lambda \cdot \pi(X, r)\)

• **jitter – Small positive term added to the diagonal part of a covariance matrix for numerical stability (Default: 1e-6)

Return type:

Array

gpax.acquisition.POI(rng_key, model, X, best_f=None, xi=0.01, maximize=False, n=1, noiseless=False, penalty=None, recent_points=None, grid_indices=None, penalty_factor=1.0, **kwargs)

Probability of Improvement

Given a probabilistic model \(m\) that models the objective function \(f\), the Probability of Improvement at an input point \(x\) is defined as:

\[PI(x) = \Phi\left(\frac{\mu(x) - f^+ - \xi}{\sigma(x)}\right)\]

where \(\mu(x)\) is the predictive mean, \(\sigma(x)\) is the predictive standard deviation, \(f^+\) is the value of the best observed sample, \(\xi\) is a small positive “jitter” term to encourage more exploration, and \(\Phi\) is the cumulative distribution function (CDF) of the standard normal distribution.

In the case of HMC, the function leverages multiple predictive posteriors, each associated with a different HMC sample of the GP model parameters, to capture both prediction uncertainty and hyperparameter uncertainty. In this setup, the uncertainty in parameters of probabilistic mean function (if any) also contributes to the acquisition function values.

Parameters:

• rng_key (Array) – JAX random number generator key

• model (Type[ExactGP]) – trained model

• X (Array) – new inputs

• best_f (float) – Best function value observed so far. Derived from the predictive mean when not provided by a user.

• xi (float) – coefficient affecting exploration-exploitation trade-off

• maximize (bool) – If True, assumes that BO is solving maximization problem

• n (int) – number of samples drawn from each MVN distribution (number of distributions is equal to the number of HMC samples)

• noiseless (bool) – Noise-free prediction. It is set to False by default as new/unseen data is assumed to follow the same distribution as the training data. Hence, since we introduce a model noise for the training data, we also want to include that noise in our prediction.

• penalty (Optional[str]) –

  Penalty applied to the acquisition function to discourage re-evaluation at or near points that were recently evaluated. Options are:

  • ’delta’:

  The infinite penalty is applied to the recently visited points.

  • ’inverse_distance’:

  Modifies the acquisition function by penalizing points near the recent points.

  For the ‘inverse_distance’, the acqusition function is penalized as:

  \[\alpha - \lambda \cdot \pi(X, r)\]

  where \(\pi(X, r)\) computes a penalty for points in \(X\) based on their distance to recent points \(r\), \(\alpha\) represents the acquisition function, and \(\lambda\) represents the penalty factor.

• recent_points (Array) – An array of recently visited points [oldest, …, newest] provided by user

• grid_indices (Array) – Grid indices of data points in X array for the penalty term calculation. For example, if each data point is an image patch, the indices could correspond to the (i, j) pixel coordinates of their centers in the original image.

• penalty_factor (float) – Penalty factor \(\lambda\) in \(\alpha - \lambda \cdot \pi(X, r)\)

• **jitter – Small positive term added to the diagonal part of a covariance matrix for numerical stability (Default: 1e-6)

Return type:

Array

gpax.acquisition.UE(rng_key, model, X, n=1, noiseless=False, penalty=None, recent_points=None, grid_indices=None, penalty_factor=1.0, **kwargs)

Uncertainty-based exploration

Given a probabilistic model \(m\) that models the objective function \(f\), the Uncertainty-based Exploration (UE) at an input point \(x\) targets regions where the model’s predictions are most uncertain. It quantifies this uncertainty as:

\[UE(x) = \sigma^2(x)\]

where \(\sigma^2(x)\) is the predictive variance of the model at the input point \(x\).

In the case of HMC, the function leverages multiple predictive posteriors, each associated with a different HMC sample of the GP model parameters, to capture both prediction uncertainty and hyperparameter uncertainty. In this setup, the uncertainty in parameters of probabilistic mean function (if any) also contributes to the acquisition function values.

Parameters:

• rng_key (Array) – JAX random number generator key

• model (Type[ExactGP]) – trained model

• X (Array) – new inputs

• n (int) – number of samples drawn from each MVN distribution (number of distributions is equal to the number of HMC samples)

• noiseless (bool) – Noise-free prediction. It is set to False by default as new/unseen data is assumed to follow the same distribution as the training data. Hence, since we introduce a model noise for the training data, we also want to include that noise in our prediction.

• penalty (Optional[str]) –

  Penalty applied to the acquisition function to discourage re-evaluation at or near points that were recently evaluated. Options are:

  • ’delta’:

  The infinite penalty is applied to the recently visited points.

  • ’inverse_distance’:

  Modifies the acquisition function by penalizing points near the recent points.

  For the ‘inverse_distance’, the acqusition function is penalized as:

  \[\alpha - \lambda \cdot \pi(X, r)\]

  where \(\pi(X, r)\) computes a penalty for points in \(X\) based on their distance to recent points \(r\), \(\alpha\) represents the acquisition function, and \(\lambda\) represents the penalty factor.

• recent_points (Array) – An array of recently visited points [oldest, …, newest] provided by user

• grid_indices (Array) – Grid indices of data points in X array for the penalty term calculation. For example, if each data point is an image patch, the indices could correspond to the (i, j) pixel coordinates of their centers in the original image.

• penalty_factor (float) – Penalty factor \(\lambda\) in \(\alpha - \lambda \cdot \pi(X, r)\)

• **jitter – Small positive term added to the diagonal part of a covariance matrix for numerical stability (Default: 1e-6)

Return type:

Array

gpax.acquisition.Thompson(rng_key, model, X, n=1, noiseless=False, **kwargs)

Thompson sampling.

For MAP approximation, it draws a single sample of a function from the posterior predictive distribution. In the case of HMC, it draws a single posterior sample from the HMC samples of GP model parameters and then samples a function from it.

Parameters:

• rng_key (Array) – JAX random number generator key

• model (Type[ExactGP]) – trained model

• X (Array) – new inputs

• n (int) – number of samples drawn from the randomly selected MVN distribution

• noiseless (bool) – Noise-free prediction. It is set to False by default as new/unseen data is assumed to follow the same distribution as the training data. Hence, since we introduce a model noise for the training data, we also want to include that noise in our prediction.

• **jitter – Small positive term added to the diagonal part of a covariance matrix for numerical stability (Default: 1e-6)

Return type:

Array

gpax.acquisition.KG(rng_key, model, X, n=1, maximize=False, noiseless=False, penalty=None, recent_points=None, grid_indices=None, penalty_factor=1.0, **kwargs)

Knowledge gradient

Given a probabilistic model \(m\) that models the objective function \(f\), the Knowledge Gradient (KG) at an input point \(x\) quantifies the expected improvement in the optimal decision after observing the function value at \(x\).

The KG value is defined as:

\[KG(x) = \mathbb{E}[V_{n+1}^* - V_n^* | x]\]

where \(V_{n+1}^*\) is the optimal expected value of the objective function after (n+1) observations and \(V_n^*\) is the optimal expected value of the objective function based on the current (n) observations.

Parameters:

• rng_key (Array) – JAX random number generator key for sampling simulated observations

• model (Type[ExactGP]) – Trained model

• X (Array) – New inputs

• n (int) – Number of simulated samples for each point in X

• maximize (bool) – If True, assumes that BO is solving maximization problem

• noiseless (bool) – Noise-free prediction. It is set to False by default as new/unseen data is assumed to follow the same distribution as the training data. Hence, since we introduce a model noise for the training data, we also want to include that noise in our prediction.

• penalty (Optional[str]) –

  Penalty applied to the acquisition function to discourage re-evaluation at or near points that were recently evaluated. Options are:

  • ’delta’:

  The infinite penalty is applied to the recently visited points.

  • ’inverse_distance’:

  Modifies the acquisition function by penalizing points near the recent points.

  For the ‘inverse_distance’, the acqusition function is penalized as:

  \[\alpha - \lambda \cdot \pi(X, r)\]

  where \(\pi(X, r)\) computes a penalty for points in \(X\) based on their distance to recent points \(r\), \(\alpha\) represents the acquisition function, and \(\lambda\) represents the penalty factor.

• recent_points (Array) – An array of recently visited points [oldest, …, newest] provided by user

• grid_indices (Array) – Grid indices of data points in X array for the penalty term calculation. For example, if each data point is an image patch, the indices could correspond to the (i, j) pixel coordinates of their centers in the original image.

• penalty_factor (float) – Penalty factor \(\lambda\) in \(\alpha - \lambda \cdot \pi(X, r)\)

• **jitter – Small positive term added to the diagonal part of a covariance matrix for numerical stability (Default: 1e-6)

Return type:

Array

gpax.acquisition.qUCB(rng_key, model, X, beta=0.25, maximize=False, noiseless=False, maximize_distance=False, subsample_size=1, n_evals=10, indices=None, **kwargs)

Batch-mode Upper Confidence Bound

qUCB computes the Upper Confidence Bound values for given input points X using multiple randomly drawn samples from the HMC-inferred model’s posterior. If maximize_distance is enabled, qUCB considers diversity among the posterior samples by maximizing the mean distance between samples that give the highest acquisition values across multiple evaluations.

Parameters:

• rng_key (Array) – random number generator key

• model (Type[ExactGP]) – trained model

• X (Array) – new inputs

• best_f – Best function value observed so far. Derived from the predictive mean when not provided by a user.

• maximize (bool) – If True, assumes that BO is solving maximization problem

• noiseless (bool) – Noise-free prediction. It is set to False by default as new/unseen data is assumed to follow the same distribution as the training data. Hence, since we introduce a model noise for the training data, we also want to include that noise in our prediction.

• maximize_distance (bool) – If set to True, it means we want our batch to contain points that are as far apart as possible in the acquisition function space. This encourages diversity in the batch.

• subsample_size (int) – Size of the subsample from the GP model’s MCMC samples.

• n_evals (int) – Number of evaluations (how many times a ramdom subsample is drawn) when maximizing distance between maxima of different EIs in a batch.

• indices (Optional[Array]) – Indices of the input points.

Return type:

Array

Returns:

The computed batch Expected Improvement values at the provided input points X.

gpax.acquisition.qEI(rng_key, model, X, best_f=None, maximize=False, noiseless=False, maximize_distance=False, subsample_size=1, n_evals=10, indices=None, **kwargs)

Batch-mode Expected Improvement

qEI computes the Expected Improvement values for given input points X using multiple randomly drawn samples from the HMC-inferred model’s posterior. If maximize_distance is enabled, qEI considers diversity among the posterior samples by maximizing the mean distance between samples that give the highest acquisition values across multiple evaluations.

Parameters:

• rng_key (Array) – random number generator key

• model (Type[ExactGP]) – trained model

• X (Array) – new inputs

• best_f (float) – Best function value observed so far. Derived from the predictive mean when not provided by a user.

• maximize (bool) – If True, assumes that BO is solving maximization problem

• noiseless (bool) – Noise-free prediction. It is set to False by default as new/unseen data is assumed to follow the same distribution as the training data. Hence, since we introduce a model noise for the training data, we also want to include that noise in our prediction.

• maximize_distance (bool) – If set to True, it means we want our batch to contain points that are as far apart as possible in the acquisition function space. This encourages diversity in the batch.

• subsample_size (int) – Size of the subsample from the GP model’s MCMC samples.

• n_evals (int) – Number of evaluations (how many times a ramdom subsample is drawn) when maximizing distance between maxima of different EIs in a batch.

• indices (Optional[Array]) – Indices of the input points.

Return type:

Array

Returns:

The computed batch Expected Improvement values at the provided input points X.

gpax.acquisition.qPOI(rng_key, model, X, best_f=None, maximize=False, noiseless=False, maximize_distance=False, subsample_size=1, n_evals=10, indices=None, **kwargs)

Batch-mode Probability of Improvement

qPOI computes the Probability of Improvement values for given input points X using multiple randomly drawn samples from the HMC-inferred model’s posterior. If maximize_distance is enabled, qPOI considers diversity among the posterior samples by maximizing the mean distance between samples that give the highest acquisition values across multiple evaluations.

Parameters:

• rng_key (Array) – random number generator key

• model (Type[ExactGP]) – trained model

• X (Array) – new inputs

• best_f (float) – Best function value observed so far. Derived from the predictive mean when not provided by a user.

• maximize (bool) – If True, assumes that BO is solving maximization problem

• noiseless (bool) – Noise-free prediction. It is set to False by default as new/unseen data is assumed to follow the same distribution as the training data. Hence, since we introduce a model noise for the training data, we also want to include that noise in our prediction.

• maximize_distance (bool) – If set to True, it means we want our batch to contain points that are as far apart as possible in the acquisition function space. This encourages diversity in the batch.

• subsample_size (int) – Size of the subsample from the GP model’s MCMC samples.

• n_evals (int) – Number of evaluations (how many times a ramdom subsample is drawn) when maximizing distance between maxima of different EIs in a batch.

• indices (Optional[Array]) – Indices of the input points.

Return type:

Array

Returns:

The computed batch Expected Improvement values at the provided input points X.

gpax.acquisition.qKG(rng_key, model, X, n=10, maximize=False, noiseless=False, maximize_distance=False, subsample_size=1, n_evals=10, indices=None, **kwargs)

Batch-mode Knowledge Gradient

qKG computes the Knowledge Gradient values for given input points X using multiple randomly drawn samples from the HMC-inferred model’s posterior. If maximize_distance is enabled, qKG considers diversity among the posterior samples by maximizing the mean distance between samples that give the highest acquisition values across multiple evaluations.

Parameters:

• rng_key (Array) – random number generator key

• model (Type[ExactGP]) – trained model

• X (Array) – new inputs

• n (int) – number of simulated samples for each point in X

• maximize (bool) – If True, assumes that BO is solving maximization problem

• noiseless (bool) – Noise-free prediction. It is set to False by default as new/unseen data is assumed to follow the same distribution as the training data. Hence, since we introduce a model noise for the training data, we also want to include that noise in our prediction.

• maximize_distance (bool) – If set to True, it means we want our batch to contain points that are as far apart as possible in the acquisition function space. This encourages diversity in the batch.

• subsample_size (int) – Size of the subsample from the GP model’s MCMC samples.

• n_evals (int) – Number of evaluations (how many times a ramdom subsample is drawn) when maximizing distance between maxima of different EIs in a batch.

• indices (Optional[Array]) – Indices of the input points.

Return type:

Array

Returns:

The computed batch Knowledge Gradient values at the provided input points X.

gpax.acquisition.optimize_acq(rng_key, model, acq_fn, num_initial_guesses, lower_bound, upper_bound, **kwargs)

Optimizes an acquisition function for a given Gaussian Process model using the JAXopt library.

This function finds the point that maximizes the acquisition function within the specified bounds. It uses L-BFGS-B algorithm through ScipyBoundedMinimize from JAXopt.

Parameters:

• rng_key (Array) – A JAX random key for stochastic processes.

• model (Type[ExactGP]) – The Gaussian Process model to be used.

• acq_fn (Callable) – The acquisition function to be maximized.

• num_initial_guesses (int) – Number of random initial guesses for the optimization.

• lower_bound (Union[List, Tuple, float, ndarray, Array]) – Lower bounds for the optimization.

• upper_bound (Union[List, Tuple, float, ndarray, Array]) – Upper bounds for the optimization.

• **kwargs – Additional keyword arguments to be passed to the acquisition function.

Return type:

Array

Returns:

Parameter(s) that maximize the acquisition function within the specified bounds.

Note

Ensure JAXopt is installed to use this function (pip install jaxopt). The acquisition function is minimized using its negative value to find the maximum.

Examples

Optimize EI given a trained GP model for 1D problem

>>> acq_fn = gpax.acquisition.EI
>>> num_initial_guesses = 10
>>> lower_bound = -2.0
>>> upper_bound = 2.0
>>> x_next = gpax.acquisition.optimize_acq(
>>>    rng_key, gp_model, acq_fn,
>>>    num_initial_guesses, lower_bound, upper_bound,
>>>    maximize=False, noiseless=True)