Acquisition Functions

Chapter 11 built the Bayesian optimization loop and ran it with one rule for choosing the next evaluation: the posterior mean plus a multiple of the posterior standard deviation. The rule worked, but its weight was a knob, and Figure 11.3 showed that turning the knob too far either way costs a great deal. A reader may reasonably ask whether there is a principled way to decide how much uncertainty is worth.

This chapter gives several answers. Each acquisition function starts from a statement about what an evaluation is for: beating the best value seen so far, improving the final recommendation, or learning where the maximum is. From that statement and the Gaussian posterior, the formula follows, and the balance between exploring and exploiting comes out of the derivation instead of being set by hand. We derive the classic acquisition functions one at a time, look at all of them on the same posterior in one and then two dimensions, and end with the problem every one of them leaves behind: finding the maximum of the acquisition function itself, which in many dimensions is a hard optimization problem of its own.

12.1 What an acquisition function is #

Recall the setting. After $n$ evaluations the data are $\D_n = \{(\vx_i, y_i)\}_{i=1}^n$, and the Gaussian process posterior gives every input $\vx$ a Gaussian belief about $f(\vx)$, with mean $\mu_n(\vx)$ and standard deviation $\sigma_n(\vx)$ (Chapter 8). An acquisition function $a_n(\vx)$ scores each input, and the loop evaluates the objective where the score is highest (Definition 11.1).

The cleanest way to build such a score is to say what we would be happy to have at the end and then ask how much one more evaluation is expected to add. Write $u(\D)$ for the utility of a data set: a number that says how good our situation is if we stop with data $\D$. If we evaluate at $\vx$ and observe $y$, the utility changes from $u(\D_n)$ to $u(\D_n \cup \{(\vx, y)\})$. We do not know $y$ before evaluating, but the posterior says what it is likely to be, so we can average over it.

Different utilities give different acquisition functions. Table 12.1 previews the ones in this chapter. Two of them, the upper confidence bound and Thompson sampling, do not come from a utility at all; they come from the bandit problems of Section 13.2, where they have guarantees of a different kind (Chapter 13).

12.1.1 Why look only one step ahead #

Equation (12.1) looks one evaluation ahead, as if the next evaluation were the last. That is a simplification. With $N - n$ evaluations left, the optimal choice now depends on what we will be able to do afterward, and the next evaluation is worth more if it sets up good later ones. Writing that out gives a dynamic program in which every future outcome branches into every future choice. It is the optimal policy, and it is also, in Kushner's words, "virtually impossible to compute" (Kushner, 1964; quoted in Garnett, 2023, sec. 12.3).

So almost all acquisition functions in use are myopic: they are optimal if the next evaluation is the last one, and only approximately optimal otherwise (Frazier, 2018). The approximation may be better than it sounds. In a few special problems where the optimal multi-step policy can be computed, the myopic rules come close; in one such study, cited by Frazier (2018), the knowledge gradient came within 98% of optimal. Myopia also has a cost that is easy to see, though: a rule that ignores the future undervalues exploration, because the payoff of exploring arrives later. Several of the rules below add exploration back in one way or another.

12.2 Probability of improvement #

The oldest rule asks the simplest question: how likely is it that $\vx$ beats the best value seen so far (Kushner, 1964)? Write $f^*_n = \max_{i \le n} y_i$ for that value, the incumbent. With exact evaluations, it is the value of the best input we could recommend right now.

The posterior says that $f(\vx)$ is Gaussian with mean $\mu_n(\vx)$ and standard deviation $\sigma_n(\vx)$. The probability that it exceeds the incumbent by at least a margin $\xi \ge 0$ follows in three steps.

In terms of Definition 12.1, PI is the expected gain of a utility that is 1 if the best value improves by at least $\xi$ and 0 otherwise.

PI has a well-known flaw: it counts improvements but ignores their size. With $\xi = 0$, an input whose mean sits just above the incumbent and whose standard deviation is tiny has PI close to 1, and the rule prefers it to an uncertain input that might improve on the incumbent by a lot. The result is a search that creeps along in tiny steps near the best point. The margin $\xi$ is the fix. Raising it asks for improvements that are worth having, which pushes the mean term in Equation (12.2) below zero for most inputs; then the only way to have a sizable probability is a large $\sigma_n(\vx)$, and the rule explores.

How large should $\xi$ be? Kushner suggested starting high and lowering it as the search proceeds (Brochu et al., 2010), and his papers discussed adjusting it by hand during the search (Garnett, 2023, sec. 12.3). Jones found the method "extremely sensitive to the choice of the target": too small and the search stays local, too large and it never refines a promising solution (Jones, 2001; quoted in Brochu et al., 2010). The next section's rule makes the margin far less important, because it accounts for the size of the improvement directly.

12.3 Expected improvement #

Instead of asking whether the incumbent improves, ask by how much, on average. Evaluating at $\vx$ raises the best value seen from $f^*_n$ to $\max(f^*_n, f(\vx))$. The gain is the improvement $\max(f(\vx) - f^*_n, 0)$, which is zero when $f(\vx)$ falls short, and its expectation under the posterior is the expected improvement:

This is Equation (12.1) with the utility $u(\D) = \max_i y_i$, the best value observed. EI is usually credited to Močkus and his colleagues (Močkus, 1975; Frazier, 2018), with an earlier explicit formula traced to Šaltenis in 1971 (Garnett, 2023, sec. 12.3), and it became standard through the EGO algorithm of Jones et al. (1998).

The difference $f(\vx) - f^*_n$ is a Gaussian variable, since $f^*_n$ is a known number. So everything reduces to one fact about Gaussians, which we state for any Gaussian $D$ because Section 19.4 uses it for a different one.

When $s = 0$, $D$ is the constant $\delta$ and the expectation is $\max(\delta, 0)$, which is also the limit of Equation (12.4) as $s \to 0$.

For expected improvement, take $D = f(\vx) - f^*_n - \xi$, with the same optional margin $\xi$ as in PI. Its mean is $\delta = \mu_n(\vx) - f^*_n - \xi$ and its standard deviation is $s = \sigma_n(\vx)$:

This is the closed form that Jones et al. (1998) made standard, with the margin $\xi$ as a later addition. Brochu et al. (2010) report experiments by Lizotte suggesting that $\xi = 0.01$, scaled by the signal variance if necessary, works well in almost all cases.

12.3.1 Reading the formula #

The two terms of Equation (12.5) are exploitation and exploration in one expression. The first is large where the mean is above the incumbent. The second is large where the standard deviation is large. Unlike the upper confidence bound, nobody chose the exchange rate between them: it comes out of the integral.

Three properties make this precise. Write $\EI(\delta, s)$ for the right-hand side of Equation (12.4).

• EI is never below the improvement of the mean. Since $\max(\cdot, 0)$ is convex, Jensen's inequality gives $\E[\max(D, 0)] \ge \max(\E[D], 0) = \max(\delta, 0)$. Uncertainty can only add value.
• At a mean equal to the incumbent, EI is about $0.4\,s$. With $\delta = 0$, Equation (12.4) gives $\EI = s\,\phi(0) = s / \sqrt{2\pi} \approx 0.399\,s$. This is the example of Section 2.6, where the improvement of the mean is zero but the expected improvement is not.
• EI increases with both the mean and the uncertainty. The partial derivatives are $\partial \EI / \partial \delta = \Phi(\delta/s)$ and $\partial \EI / \partial s = \phi(\delta/s)$ (Exercise 12.1), and both are positive. More uncertainty always makes EI larger, wherever the mean is.

The last property is where EI and PI part ways. PI's derivative with respect to $s$ is $-(\delta/s^2)\,\phi(\delta/s)$, which is negative whenever $\delta > 0$: once the mean is above the incumbent, PI prefers certainty. Figure 12.1 shows both quantities for a single input.

Three settings show the difference.

A mean below the incumbent. The figure opens with $\delta = -0.3$. Both PI and EI grow as $s$ grows: the only way to beat the incumbent is for the function to be higher than the model expects, and more uncertainty makes that more likely.

A mean above the incumbent. Set $\delta = 0.5$. Now PI falls as $s$ grows, because a wider belief puts more weight on falling short. EI still rises: the extra weight on large improvements outweighs the extra weight on falling short, which costs nothing, since improvement is never negative.

A tiny, certain improvement. Set $\delta = 0.05$ and $s = 0.05$. PI is about 0.84, while EI is only about 0.05. An input with $\delta = -0.3$ and $s = 1$ has a PI of only 0.38, but its EI is about 0.27, about five times larger. EI ranks the second input higher; PI ranks the first.

Two practical notes complete the picture. First, Equation (12.5) assumes exact evaluations, so that $f^*_n$ is a known value. With noise the best observed value is itself uncertain and biased upward, and several noisy variants replace it; Section 14.2 compares them. Second, a real application shows the formula at work. Snoek et al. (2012) tuned machine learning algorithms with expected improvement, averaging Equation (12.5) over samples of the Gaussian process hyperparameters instead of fixing them. They also modeled the training time with a second Gaussian process and maximized expected improvement per second, which prefers inputs that are both promising and quick to evaluate. Chapter 22 replays a problem of this kind.

12.4 Upper confidence bounds #

The upper confidence bound is the rule Chapter 11 already used:

Its principle comes from the bandit literature (Section 13.2) and is called optimism in the face of uncertainty: act as if the world were as good as it plausibly could be, then let the evaluation correct you. If the optimism was justified, the evaluation finds a good input; if not, the evaluation shrinks $\sigma_n(\vx)$ there and the inflated estimate deflates, so the rule moves on.

The weight has a probabilistic reading. Under the posterior, $\Prob(f(\vx) \le \mu_n(\vx) + \beta^{1/2} \sigma_n(\vx)) = \Phi(\beta^{1/2})$, so Equation (12.6) is a quantile of the belief about $f(\vx)$. With $\beta^{1/2} = 2$ it is the 97.7% quantile, and maximizing it picks the input whose plausible best case is highest.

What should $\beta$ be? Srinivas et al. (2010) answered with a schedule that grows slowly with the number of evaluations $t$. For a finite domain $\X$ of candidate inputs and a confidence parameter $\delta \in (0, 1)$, their GP-UCB rule uses

and they proved that with this schedule the regret grows sublinearly, so the average gap to the maximum goes to zero (Section 13.4 gives the theorem and its proof). The schedule is just large enough that, with probability at least $1 - \delta$, the bands $\mu \pm \beta_t^{1/2} \sigma$ contain the true function at every candidate and every step at once. For continuous domains, the schedule gains a term proportional to $d \log t$ (Srinivas et al., 2010).

The numbers show how cautious the theory is. With $|\X| = 1000$ candidates and $\delta = 0.1$, Equation (12.7) gives $\beta_t^{1/2} \approx 5.4$ at $t = 10$ and $\approx 6.2$ at $t = 100$: five to six standard deviations of optimism. Figure 11.3 found weights of 1 to 2 best on the running example, and weights near 6 noticeably worse. The authors themselves found that their algorithm improved when $\beta_t$ was scaled down by a factor of 5 (Srinivas et al., 2010). In practice, libraries take a constant $\beta$ from the user; BoTorch's UpperConfidenceBound computes the mean plus $\sqrt{\beta}$ times the standard deviation. The schedule matters for the proof, which needs exploration that never switches off. A constant that works well over a fixed budget is a different, practical question.

12.5 Thompson sampling #

The oldest idea in this chapter is from 1933. Thompson (1933) considered two medical treatments of unknown effectiveness and proposed assigning each new patient a treatment with the probability that it is the better one, given the evidence so far. For a Gaussian process the rule is:

1. Draw one function $g$ from the posterior given $\D_n$ (Section 8.5).
2. Evaluate the objective where that sample is largest: $\vx_{n+1} \in \argmax_{\vx} g(\vx)$.

The first step is the only random one, and it is what makes the rule work. Given the data, the sample $g$ and the unknown $f$ have the same distribution, so the maximizer of $g$ has the same distribution as the maximizer of $f$:

This is called probability matching: Thompson sampling evaluates each region exactly as often as the model believes the maximum lies there. A region the model is sure is poor is almost never chosen. A region that could hide the maximum is chosen in proportion to how likely that is, however uncertain the rest of the model is.

Thompson sampling has no weight or margin to tune, and it parallelizes naturally: to choose ten evaluations at once, draw ten samples and take each one's maximizer. Its guarantees are close to those of UCB. Russo and Van Roy (2014) established a connection between posterior sampling and upper confidence bound algorithms that converts regret bounds proved for UCB algorithms into Bayesian regret bounds for posterior sampling, including one for Gaussian process models, and Chowdhury and Gopalan (2017) proved a regret bound for a Gaussian process version when the unknown function is fixed rather than drawn from the prior.

The cost lies in step 1. A sample on a grid of $m$ inputs needs a Cholesky factorization of the $m \times m$ posterior covariance, $O(m^3)$, which limits exact sampling to a few thousand inputs. In more dimensions, implementations draw the sample on a candidate set, a few thousand points chosen to cover the promising regions, or draw approximate sample functions that can be evaluated anywhere, built from a finite set of random basis functions (random features) or from a prior sample corrected by the data (pathwise updates), and maximize them with gradients (Rahimi and Recht, 2007; Wilson et al., 2020). Section 12.9 explains why the choice of candidate set becomes the hard part in many dimensions.

12.6 Knowledge gradient #

Expected improvement makes a quiet assumption: at the end, we recommend one of the inputs we evaluated, the one with the best observed value (Frazier, 2018). Often we would happily recommend an input we never evaluated, if the model is confident it is good. And with noisy evaluations, no observed value can be trusted at face value anyway, so the natural recommendation is the input with the highest posterior mean (Section 11.2.3).

The knowledge gradient takes that recommendation seriously. Its utility is the value of the final recommendation: if we stopped now, we would recommend the maximizer of the posterior mean, and its expected value under the posterior is

One more evaluation at $\vx$ changes the posterior mean everywhere, not only at $\vx$, so it can raise $\mu^*$ even if $f(\vx)$ itself turns out to be mediocre. Plugging this utility into Equation (12.1) gives the knowledge gradient:

The value of a query in this sense is the expected value of the final recommendation after one more answer. A query that maximizes it is one-step Bayes optimal: if the session ended after this one evaluation, no other choice would leave a better recommendation in expectation, and the knowledge gradient is the acquisition function that picks it. Section 19.4.1 relies on exactly this property when queries are pairs of options.

To compute Equation (12.9) we need to know how the posterior mean moves when one observation arrives.

Every point of the new posterior mean moves by a multiple of the same standard normal $Z$, because a single number, the outcome $y$, moves them all. Inputs strongly correlated with $\vx$ move a lot; inputs far from $\vx$ hardly move. So the knowledge gradient is

Three consequences follow from this form.

• KG is never negative. The maximum of functions is convex, and $Z$ has mean zero, so by Jensen's inequality the expected maximum is at least the maximum at $Z = 0$, which is $\mu^*_n$. Information never hurts in expectation.
• An exact repeat is worth nothing. If $\vx$ was already evaluated without noise, $\sigma_n(\vx) = 0$ and $k_n(\vx', \vx) = 0$ for every $\vx'$, so nothing moves and $\mathrm{KG}_n(\vx) = 0$. With noise ($\sigma_\varepsilon > 0$), repeating an evaluation can still be worth something, which is why KG handles noisy problems gracefully (Frazier, 2018).
• With two inputs in play, KG is Clark's formula. If the maximum in Equation (12.11) runs over only two inputs, it is the maximum of two jointly Gaussian values, $a_1 + b_1 Z$ and $a_2 + b_2 Z$, and its expectation is the formula of Clark (1961) (Section B.5). The same formula returns in preferential Bayesian optimization (PBO), where a query is a pair of options and the acquisition function is the expected utility of the better one (EUBO, Equation (19.3)); with noise-free answers it has the knowledge gradient's one-step optimality (Section 19.4.1).

The knowledge gradient also clarifies expected improvement. If the recommendation must be an evaluated input and evaluations are exact, the utility in Equation (12.1) becomes the best observed value, and the one-step gain is exactly EI (Frazier, 2018). EI is the knowledge gradient of a decision maker who will only recommend what has been measured.

Figure 12.2 makes the lookahead visible.

Look where the violet curves fan out. At the default candidate, in the unexplored gap on the right, the fantasies spread widely: a high outcome would put a new maximum of the mean there, a low one would leave the old maximum in place. Since a low outcome cannot lower $\mu^*$ (the old maximum is still available) while a high one raises it, the average rises.

Move the candidate onto an evaluated point. The violet curves collapse onto the blue one and the readout shows KG equal to zero, the second consequence above. Now move it a little to either side: KG jumps back up. Two exact evaluations very close together reveal the slope of $f$ between them, and near the top of a bump a slope can move the maximum of the mean. The narrow notches in the strip are this effect.

Compare with the strip. The knowledge gradient is largest around the current best region, not in the middle of the widest gap. Near the best point, an evaluation directly moves the maximum of the mean; in the far gap, an outcome must be large to matter at all. KG explores less than UCB with $\beta^{1/2} = 2$, a pattern Figure 12.3 shows again.

Add noise. Set the noise to 0.2. The fantasies spread less, because a noisy outcome moves the mean less, and the notches widen into valleys. At the evaluated inputs on the broad bump, KG is now slightly positive: a second noisy measurement there could still move the maximum of the mean. At the evaluated input in the dip near $x = 0.45$ it stays at zero, since no outcome there could make that region the best.

The knowledge gradient was introduced for choosing among a finite set of alternatives with independent beliefs (Frazier et al., 2008) and extended to correlated beliefs (Frazier et al., 2009), the setting of a Gaussian process on a grid. On a finite set, Equation (12.11) can be computed exactly: the maximum of the lines $a_j + b_j z$ is a piecewise linear function of $z$, and integrating each piece against the normal density gives a closed form (Frazier et al., 2009); Figure 12.2 does exactly this. On a continuous domain the inner maximum has no closed form, and implementations estimate KG by simulating outcomes, re-solving the inner maximization for each (Frazier, 2018), or, as in BoTorch, by optimizing the candidate together with one maximizer per simulated outcome in a single "one-shot" problem (Balandat et al., 2020). (Section 11.5 tells how the knowledge gradient and expected improvement were entangled at their origin.)

12.7 Entropy search #

The rules so far value an evaluation by what it does for a final answer. A different family values it by what it teaches about the maximum. The posterior over $f$ induces a distribution over the location of the maximum, $p(\vx^\star \mid \D_n)$: draw a function from the posterior, find where it is largest, and repeat. Where this distribution is spread out, we do not know where the maximum is. An evaluation is valuable if it is expected to concentrate the distribution.

Entropy, from Section 6.1, measures the spread of a distribution in a way that depends only on probabilities, not on distances. Entropy search picks the evaluation that most reduces the entropy of $p(\vx^\star \mid \D)$ in expectation:

The expected reduction in entropy is the mutual information between the outcome $y$ and $\vx^\star$ (Section 6.3), and Equation (12.12) is Equation (12.1) with the negative entropy of $\vx^\star$ as the utility. The idea was proposed by Villemonteix et al. (2009) and independently by Hennig and Schuler (2012), who coined the name (Garnett, 2023, sec. 12.4).

Computing Equation (12.12) is hard: the distribution of $\vx^\star$ has no closed form, its entropy must be approximated, and that approximation must be repeated for every hypothetical outcome $y$. Two reformulations made the idea practical.

Predictive entropy search uses the symmetry of mutual information. The information that $y$ carries about $\vx^\star$ equals the information that $\vx^\star$ carries about $y$, so

The first term is the entropy of a Gaussian, available in closed form; the second averages over sampled maximizers, each requiring an approximation of how knowing $\vx^\star$ would change the prediction at $\vx$ (Hernández-Lobato et al., 2014). Exact ES and PES are the same function; their approximations differ (Frazier, 2018).

Max-value entropy search (MES) changes the target: it asks for information about the maximum value $f^\star = f(\vx^\star)$, a single number, instead of its location (Wang and Jegelka, 2017). Knowing $f^\star$ tells us one simple thing about $f(\vx)$: it cannot exceed $f^\star$. So the belief about $f(\vx)$ given $f^\star$ is the Gaussian posterior cut off above $f^\star$, a truncated Gaussian whose entropy has a closed form. Averaging over $K$ sampled maximum values $f^\star_1, \dots, f^\star_K$ gives

which is equation (6) of Wang and Jegelka (2017). The samples $f^\star_k$ can be taken as the maxima of posterior sample functions, or drawn from a cheaper approximation of their distribution. Each term is large when $\gamma_k$ is small, that is, when the sampled maximum is not far above the mean at $\vx$, measured in standard deviations: an evaluation there could reveal whether the maximum is about that high. The authors report that MES matches or improves on ES and PES at a fraction of the cost, and that it is much less sensitive to the number of samples (Wang and Jegelka, 2017).

12.8 Comparing them #

Each rule has now been derived on its own. To see how differently they behave, Figure 12.3 puts all six on one posterior of the running objective. You play the optimizer: click the plot to evaluate the objective anywhere, and watch where each rule would go next.

The default posterior has four evaluations: two on the broad bump, one in the dip after it, and one at the far right. The tall bump near $x = 0.73$ sits in an unexplored gap. Some guided experiments:

The rules disagree. PI, EI, MES, and KG choose near the broad bump, at $x$ between about 0.19 and 0.27, where the mean is already close to the incumbent and a modest improvement is likely. UCB with $\beta^{1/2} = 2$ goes into the gap, to $x \approx 0.68$, where the upper edge of the band is highest. Thompson sampling depends on its sample: press New samples a few times and its choice jumps among the broad bump, the gap, and the left edge, roughly in proportion to the violet bars.

Raise the margin. Move $\xi$ up from 0.01. At $\xi = 0.1$, EI already switches to the gap: asking for an improvement of at least 0.1 makes the small, likely gains near the broad bump worthless. PI holds on to the broad bump until $\xi$ reaches 0.39. Both rules change their choice abruptly at some margin, which is the sensitivity that Jones warned about.

Run the loop. Pick a rule and press Evaluate its maximum repeatedly, then press Reset and try another. Every rule reaches the tall bump eventually, by different routes: UCB goes there first, EI, KG, and MES after one more evaluation on the broad bump, PI after two, and Thompson sampling later still. On this posterior PI happens to come within 0.05 of the maximum fastest, and KG, which is trying to improve the maximum of the mean rather than the best observed value, is slowest by the best-observed yardstick. One run on one problem ranks nothing.

Watch the violet bars. Once the tall bump has been evaluated, the estimated distribution of $x^\star$ collapses onto it. Thompson sampling and MES then concentrate their evaluations there; UCB keeps visiting other regions while their bands are wide.

No rule is best on every problem. Comparisons on benchmark functions favor different rules on different problems, and the regret bounds of Section 13.5.2 are proved in different settings for different rules, so they do not rank them either (inference). Table 12.2 lists the practical differences that do hold.

12.8.1 In two dimensions #

One-dimensional pictures hide an important fact: the number of places to look grows exponentially with dimension, while the posterior is informative only near the data. Figure 12.4 repeats the comparison on a two-dimensional problem, the Branin function, a standard test function with three global maxima of equal height (rescaled here to the unit square and negated so that larger is better).