Bayesian Optimization, in depth: Gaussian process surrogates, expected improvement and UCB, a NumPy implementation, kernel fitting and failure modes

The problem Bayesian optimization solves

Formally, we want x* = argmax f(x) over a bounded box X in d dimensions, where one call to f costs minutes to days. The budget is small, typically 20 to 200 evaluations, so the number of evaluations is the cost that matters and the overhead of the optimizer is nearly free by comparison.

BO is the right tool when d is small (roughly up to 10 to 20 continuous dimensions), f is reasonably smooth, and evaluations are expensive. It is the wrong tool when f is cheap (use an evolutionary method or plain random search with many samples), when gradients are available (use them), or when early partial results are informative, in which case multi-fidelity schedulers such as Hyperband usually win, and can be combined with BO.

The loop

The loop is short. Evaluate a handful of initial points chosen by a space-filling design. Fit the surrogate to all observations. Find the point that maximizes the acquisition function. Evaluate f there, append the result, and repeat until the budget is gone. Everything interesting lives in the surrogate and the acquisition.

The Gaussian process surrogate

A Gaussian process is a distribution over functions defined by a mean function (usually a constant after standardizing y) and a kernel k(x, x'), which says how strongly the values at two points are correlated. Nearby points are highly correlated, distant points nearly independent, and the lengthscale sets what counts as nearby. The Matern 5/2 kernel is the usual default for BO because it assumes twice-differentiable functions, which is less optimistic than the infinitely smooth squared-exponential kernel; Snoek, Larochelle and Adams recommended it in their 2012 hyperparameter-tuning paper.

Given n observations with kernel matrix K, observation noise variance s2 and the vector k* of kernel values between a new point x and the data, the posterior at x is Gaussian with mean mu(x) = k*^T (K + s2 I)^-1 y and variance sigma2(x) = k(x, x) - k*^T (K + s2 I)^-1 k*. The mean interpolates the data; the variance is the prior variance minus what the data explains, so it shrinks near observations and returns to the prior far from them.

Never form the inverse. Compute the Cholesky factor L of K + s2 I once per fit, which costs O(n^3), then solve triangular systems: the mean costs O(n) per point and the variance O(n^2). With a few hundred observations this is milliseconds. Adding a small jitter such as 1e-6 to the diagonal keeps the factorization stable when two points nearly coincide.

Acquisition functions

Let f+ be the best value observed so far. Three acquisition functions cover most practice. Probability of improvement, PI(x) = Phi(z), asks how likely x is to beat f+. Expected improvement, EI, asks how much it is expected to beat it by. For maximization with z = (mu - f+ - xi) / sigma, EI(x) = (mu - f+ - xi) Phi(z) + sigma phi(z), where Phi and phi are the standard normal CDF and density, and EI is defined as 0 where sigma is 0. Upper confidence bound, UCB(x) = mu + sqrt(beta) sigma, adds an explicit optimism bonus.

Each formula balances exploitation, a high mean, against exploration, a high standard deviation. EI does it naturally: the first term rewards points predicted to beat f+, the second rewards uncertainty. The margin xi shifts the balance toward exploration. UCB makes the balance explicit through beta and has regret guarantees when beta grows slowly with the iteration count (Srinivas and colleagues, 2010).

A complete implementation

The implementation below is complete and dependency-light: NumPy for the linear algebra and SciPy for the normal distribution and triangular solves. It works in the unit cube, standardizes y, and maximizes EI over random candidates. Production libraries use multi-start gradient ascent on the acquisition instead, but random candidates are honest and easy to debug in low dimensions.

import numpy as np
from scipy.linalg import cho_solve, solve_triangular
from scipy.stats import norm

def matern52(A, B, ls, var):
    r = np.sqrt((((A[:, None, :] - B[None, :, :]) / ls) ** 2).sum(-1))
    s5r = np.sqrt(5.0) * r
    return var * (1.0 + s5r + 5.0 * r ** 2 / 3.0) * np.exp(-s5r)

class GP:
    def __init__(self, ls=0.2, var=1.0, noise=1e-6):
        self.ls, self.var, self.noise = ls, var, noise

    def fit(self, X, y):
        self.X, self.m, self.s = X, y.mean(), y.std() + 1e-12
        yn = (y - self.m) / self.s
        K = matern52(X, X, self.ls, self.var) + self.noise * np.eye(len(X))
        self.L = np.linalg.cholesky(K)
        self.alpha = cho_solve((self.L, True), yn)
        return self

    def predict(self, Xs):
        Ks = matern52(Xs, self.X, self.ls, self.var)
        mu = Ks @ self.alpha
        v = solve_triangular(self.L, Ks.T, lower=True)
        var = np.clip(self.var - (v ** 2).sum(0), 1e-12, None)
        return mu * self.s + self.m, np.sqrt(var) * self.s

def expected_improvement(mu, sd, best, xi=0.01):
    imp = mu - best - xi
    z = imp / sd
    return imp * norm.cdf(z) + sd * norm.pdf(z)

def bayes_opt(f, bounds, n_init=6, n_iter=30, seed=0):
    rng = np.random.default_rng(seed)
    lo, hi = bounds[:, 0], bounds[:, 1]
    U = rng.uniform(size=(n_init, len(lo)))          # work in [0, 1]^d
    y = np.array([f(lo + u * (hi - lo)) for u in U])
    for _ in range(n_iter):
        gp = GP().fit(U, y)
        cand = rng.uniform(size=(4096, len(lo)))
        mu, sd = gp.predict(cand)
        u_next = cand[np.argmax(expected_improvement(mu, sd, y.max()))]
        U = np.vstack([U, u_next])
        y = np.append(y, f(lo + u_next * (hi - lo)))
    i = np.argmax(y)
    return lo + U[i] * (hi - lo), y[i]

Fitting the kernel

The code above fixes the lengthscale at 0.2 of the box. That is the single most important number in the whole method, and a fixed guess is the most common reason a home-grown BO loop performs no better than random search. Too long a lengthscale and the GP believes the function is flat, so it over-exploits; too short and every point looks independent, so EI degenerates into a space-filling search. Fit it, together with the signal variance and noise, by maximizing the log marginal likelihood after every new observation.

from scipy.optimize import minimize

def neg_log_ml(log_theta, X, yn):
    ls, var, noise = np.exp(log_theta)
    K = matern52(X, X, ls, var) + (noise + 1e-8) * np.eye(len(X))
    try:
        L = np.linalg.cholesky(K)
    except np.linalg.LinAlgError:
        return 1e10
    alpha = cho_solve((L, True), yn)
    return 0.5 * yn @ alpha + np.log(np.diag(L)).sum() + 0.5 * len(X) * np.log(2 * np.pi)

def fit_hypers(X, yn, restarts=5, rng=np.random.default_rng(0)):
    box = [(np.log(0.01), np.log(2.0)), (np.log(0.05), np.log(20.0)), (np.log(1e-6), np.log(1.0))]
    best = None
    for _ in range(restarts):
        x0 = [rng.uniform(a, b) for a, b in box]
        r = minimize(neg_log_ml, x0, args=(X, yn), method="L-BFGS-B", bounds=box)
        if best is None or r.fun < best.fun:
            best = r
    return np.exp(best.x)       # ls, var, noise

Bound the search. A lengthscale prior or box keeps the optimizer away from degenerate solutions such as a tiny lengthscale with zero noise that interpolates everything. With more than one input, use one lengthscale per dimension (automatic relevance determination); the fitted lengthscales then double as a sensitivity report, since a dimension with a huge lengthscale barely affects f. Transform inputs before fitting: learning rates, weight decay and other scale parameters belong on a log axis, otherwise the GP spends its budget on the top decade of the range.

Worked example: one acquisition decision

Suppose we are maximizing validation accuracy and the best result so far is f+ = 0.80. The GP offers two candidates. Candidate A sits in an unexplored region: mu = 0.79, sigma = 0.05. Candidate B sits next to the incumbent: mu = 0.82, sigma = 0.005.

With xi = 0, for A the improvement term is -0.01 and z = -0.2, so EI = -0.01 x 0.4207 + 0.05 x 0.3910 = 0.0153. For B the improvement is 0.02 and z = 4, so EI is about 0.0200. EI picks B, a safe local step. Now set xi = 0.01. A has z = -0.4 and EI = -0.02 x 0.3446 + 0.05 x 0.3683 = 0.0115, while B has z = 2 and EI = 0.01 x 0.9772 + 0.005 x 0.0540 = 0.0100. A small margin flips the decision to exploration. UCB with beta = 4 gives A 0.89 against B 0.83, an even stronger preference for A, while PI gives B about 1.0 against A 0.42, the greedy extreme.

The lesson is practical: the acquisition settings are a dial between refining what you have and searching for something better, and the same posterior can send you to opposite ends of the space.

Noise, batches, categories and dimension

Noise. Real validation scores vary with the seed. With noisy y, the best observed value is biased upward by luck, so EI chases it. Model the noise explicitly, use the best posterior mean as f+, or use noisy EI as implemented in BoTorch.

Parallel workers. Plain EI proposes one point. To fill eight GPUs, pick a point, pretend its outcome equals the posterior mean (the kriging believer heuristic) or a fixed constant (constant liar), refit and pick again; or optimize a joint batch criterion such as q-EI; or draw Thompson samples, which parallelize naturally.

Categorical and integer inputs. A GP over one-hot encodings works poorly. Use a tree-structured Parzen estimator, which is the default sampler in Optuna, use a random forest surrogate as SMAC does, or use kernels built for mixed spaces. Round integers inside f, not in the GP, and accept some wasted proposals.

High dimensions. A global GP needs data that grows exponentially with d to pin down a function, so vanilla BO typically degrades past about 20 dimensions. TuRBO keeps local trust regions, SAASBO uses sparsity priors on lengthscales, and random embedding methods assume only a few directions matter. Often the better fix is to tune fewer parameters.

Failure modes

Trade-offs and tools

Compared with random search, BO wins clearly with small budgets and smooth, low-dimensional objectives, and the gap closes as the budget grows or the dimension rises. Its overhead is O(n^3) per fit, irrelevant for hundreds of points. It is inherently sequential: on a large cluster, forty random trials in parallel may finish before five careful BO rounds.

For production work, use a maintained library rather than this sketch. BoTorch and Ax provide GP models, batch and noisy acquisitions and constraint handling; Optuna offers TPE by default and also ships a GP-based sampler. The value of building the loop yourself once is that you will recognize the failure modes above in a dashboard. For the surrounding workflow of tuning on GPUs, read hyperparameter search for GPU training and GPU hyperparameter tuning; for the broader landscape see hyperparameter optimization. The exploration dilemma is the same one studied in multi-armed bandits, and simulated annealing is a useful contrast for cheap objectives.

What to do next

1. Write down the objective, its noise level, the budget in evaluations and how many can run in parallel; if f is cheap, stop and use random search.
2. Reduce the search space to the parameters that matter, put scale parameters on a log axis and map everything to the unit cube.
3. Start with 2 x d to 10 x d space-filling initial points so the GP has something to fit.
4. Fit kernel hyperparameters by marginal likelihood each iteration and log the lengthscales.
5. Use EI with a small margin for serial runs; use a batch method or Thompson sampling for parallel workers.
6. Track best-so-far against a random-search baseline with the same budget; if BO is not ahead by a third of the budget, inspect the lengthscales and input transforms.
7. Re-evaluate the winning configuration with fresh seeds before you ship it.