Skip to content

jaxglobal.gp

Gaussian-process predictions with relaxations that keep the kernel structure.

gp_mean is a primitive: the posterior mean mu(x) = k(x)^T alpha. Relaxed term by term, the sum of N kernels with signed weights loses their cancellation (||alpha||_1 can be many times ||mu||_H): the interval stays several times the true range on every box. For kernels differentiable at 0 (matern52, se) the rule intersects that relaxation with a first-order Taylor model whose remainder is bounded in the RKHS of the kernel. mu is in it with ||mu||_H^2 = alpha^T K_f alpha (K_f the noise-free training covariance), so for c the box center, delta = x - c and s = ||delta / l||^2,

|mu(x) - mu(c) - grad mu(c).delta| <= ||mu||_H ||k_x - k_c - D k_c delta||_H = ||mu||_H e(s),
e(s)^2 = 2 sf2 (1 - phi(s) + 2 s phi'(s) - s phi'(0)),

phi the unit-variance kernel as a function of s. e is nondecreasing in s (its derivative is phi'(s) - phi'(0) + 2 s phi''(s) >= 0, phi' being nondecreasing and ``phi''

= 0for these kernels), so its box maximum is at the farthest corner, ande = O(s) = O(w^2): the model converges to the true range as the box shrinks, and does not involve the conditioning ofK. The zeroth-order bound|mu(x) - mu(c)| <= ||mu||_H sqrt(2 sf2 (1 - phi(s)))holds for every kernel but is not used: on large boxes it makes the relaxation flat, its minimizer then says nothing about wheremuis low, and the local solves started from it found the incumbent later (more nodes inbenchmarks/gp_mean_ab.py``, docs/PLAN-acquisition-relaxations.md).

gp_std is a primitive too: the posterior standard deviation

sigma(x) = sqrt(max(v(x), eps)),   v(x) = sf2 - k(x)^T A k(x) = sf2 - ||W^T k(x)||^2,

with k_i(x) = sf2 * kernel(sum_j ((x_j - Z_ij) / l_j)**2), A the inverse of the (noisy) training covariance and A = W W^T its Cholesky factor. Relaxed as k^T A k, the quadratic form is a sum of bilinear products, and any bound through A inherits its conditioning (~1/noise). The rule combines two relaxations, keeping the max of the convex and the min of the concave sides:

  • sqrt(max(sf2 - ||W^T k||^2, eps)), squares of linear forms in k: good from above;
  • a Lipschitz bound from the GP itself. The posterior is a Gaussian process, so for any x and the box center c, |sigma(x) - sigma(c)| <= sd_post(f(x) - f(c)) <= sd_prior(f(x) - f(c)) = sqrt(2 sf2 (1 - kernel(d(x, c)))) (the triangle inequality for standard deviations; conditioning never increases the variance of a linear functional). Over the box this gives sigma(c) -+ h with h from the farthest corner: first order in the box size, independent of A;
  • for matern52 and se, a first-order model (_std_taylor): the same argument with f(c) + grad f(c).delta in place of f(c) gives |sigma(x) - sqrt(q(delta))| <= e(s), e the prior remainder of the mean's model above (without ||mu||_H) and q the posterior variance of f(c) + grad f(c).delta, a PSD quadratic in closed form, q = sf2 + sum_j D_j delta_j^2 - ||w + J_w delta||^2 (w = W^T k(c), J_w its Jacobian, D_j = -2 sf2 phi'(0) / l_j^2). sqrt(q) is convex: its convex side is a max of tangent planes, its concave side sqrt of a concave overestimator of q. Second order: the interval converges to the true range as the box shrinks.

v >= 0 in exact arithmetic (A is positive semidefinite); eps keeps the square root differentiable.

gp_mean 

gp_mean(x, Z, lengthscale, alpha, sf2, kernel='matern52')

Posterior mean k(x)^T alpha (alpha = A y) at one point x.

Z (N, n) training inputs, lengthscale (n,), sf2 the signal variance, alpha the weights (any vector: the RKHS norm the relaxation uses is computed from it).

gp_std 

gp_std(x, Z, lengthscale, A, sf2, kernel='matern52', eps=1e-12)

Posterior standard deviation sqrt(max(sf2 - k(x)^T A k(x), eps)) at one point x.

Z (N, n) training inputs, lengthscale (n,), A the inverse of the training covariance including noise (symmetric positive definite), sf2 the signal variance.

expected_improvement 

expected_improvement(mu, sigma, fmin)

Expected improvement below fmin of a Gaussian with mean mu and std sigma (elementwise).

Same value as (fmin - mu) Phi(z) + sigma phi(z); relaxed as the jointly convex, monotone function it is, instead of through Phi, phi and the division one by one.

probability_of_improvement 

probability_of_improvement(mu, sigma, target)

Probability that a Gaussian with mean mu and std sigma falls below target (elementwise).

Same value as Phi((target - mu) / sigma); relaxed from its monotonicity in mu and sigma instead of through Phi and the division separately.