OK so why Gaussian Processes?

Here’s the problem you face: you have some noisy data generated by an unknown (continuous, real-valued, sexy) function. You can’t abide the mystery and need to map out the unobserved regions of the function. In particular, you want to find its maximum value. But it’s not straightforward, because for whatever reason the mystery function is expensive to run.

mystery_function <- function(x) {
  cost(dollars = 100)
  1.1 * sin(sqrt(x))
}

Imagine it’s API charges, token consumption, or tariffs or something. The point is you can’t just brute force your way through this one. You need help.

Here are the observations you have so far. There’s a little noise on them too, just to make life that little bit harder.

Code

domain <- c(0, 40)

observations <- tibble(x = c(0, 8, 15, 20, 35)) |>
  mutate(
    y = mystery_function(x) + rnorm(length(x), mean = 0, sd = 0.1),
  )

ggplot(observations) +
  aes(x = x, y = y) +
  scale_x_continuous(limits = domain) +
  geom_point(colour = "red", size = 3) +
  geom_function(fun = mystery_function, colour = "black", linetype = "dashed")

You could just fit a line, sure. If you have a good idea about the function that’s producing your data, that’s the simplest and most obvious thing to do.

Code

ggplot(observations) +
  aes(x = x, y = y) +
  scale_x_continuous(limits = domain) +
  geom_point(colour = "red", size = 3) +
  geom_function(fun = mystery_function, colour = "black", linetype = "dashed") +
  geom_smooth(method = "lm", formula = y ~ sin(x**0.5))

If we weren’t masters of this toy universe however, we might not have any idea what function to fit. We could fit anything else, like a polynomial. But whatever we fit, we’re making a lot of assumptions about the function and could end up quite far from the mark, especially if we don’t have many observations.

Code

ggplot(observations) +
  aes(x = x, y = y) +
  scale_x_continuous(limits = domain) +
  geom_point(colour = "red", size = 3) +
  geom_function(fun = mystery_function, colour = "black", linetype = "dashed") +
  geom_smooth(method = "lm", formula = "y ~ poly(x, 3)")

The Bayesian aproach

What we’ve done so far is the non-Bayesian approach to the problem. Why bother with the Bayesian approach (Gaussian Processes) at all?

• First, because we don’t necessarily have a good idea what our function is. It could be very nonlinear. A Gaussian Process is a non-parametric model, i.e. it doesn’t assume any specific form of the function.

• Second, as usual with Bayesian models, you can get a richer understanding of uncertainty that you can with linear models and confidence intervals. It’s also more suited to low-data scenarios.

Getting your head around the Bayesian approach can be a little challenging though. We have to forget the idea of optimising coefficients of some linear formula to fit our model. Instead we are estimating possible locations of the line. We don’t care about the formula any more.

Rather than a formula, we can use a Gaussian Process (GP). The GP is a distribution that generates “functions” (lines). How so? It’s a multivariate normal distribution, that can generate a random value for many dimensions.

If you imagine that each dimension is a point on the x-axis, a sample from the multivariate normal distribution describe a random line, or some unknown function, to look at it the other way.

Code

n.dims <- 16

# Sigma specifies the covariance matrix of the variables.
# This identify matrix means each dim is completely unrelated.
sigma <- matrix(0, nrow = n.dims, ncol = n.dims)
for (i in 1:n.dims) {
  sigma[i, i] = 1
}

random.draw <- MASS::mvrnorm(
  n = 1, # n.samples
  mu = rep(0, n.dims), # all centered on 0
  Sigma = sigma
)

ggplot(
  data.frame(x = 1:n.dims, y = random.draw)
) +
  aes(x = x, y = y) +
  geom_line()

The kernel (sigma in the code above) controls the shape of the functions generated. In the example above it’s an identify matrix, but a more conventional choice would be squared exponential covariance. This generates a nice variety of functions. The two hyperparameters alpha and lambda control something like the amplitude and frequency of the generated functions.

Code

# Kernel
sq_exp_cov <- function(x, lambda, alpha) {
  n <- length(x)
  K <- matrix(0, n, n)
  for (i in 1:n) {
    for (j in 1:n) {
      diff <- sqrt(sum((x[i] - x[j])^2))
      K[i, j] <- alpha^2 * exp(-diff^2 / (2 * lambda^2))
    }
  }
  K
}

plot_random_gp <- function(lambda, alpha) {
  n.samples <- 5
  x <- seq(0, 9, 0.1)
  samples <- MASS::mvrnorm(
    n = n.samples,
    mu = rep(0, length(x)),
    Sigma = sq_exp_cov(x, lambda = lambda, alpha = alpha)
  )

  t(samples) |>
    as_tibble() |>
    mutate(x = x) |>
    pivot_longer(-x, names_to = "sample", values_to = "y") |>
    ggplot() +
    aes(x = x, y = y, group = sample, colour = sample) +
    geom_line()
}
plot_random_gp(lambda = 2, alpha = 1)

So we have this mystery function generator! It can propose possible functions, and we can see how well they fit the data. We can also predict values for these functions at points that we don’t have observations for.

Time for Stan. As per the last blog, this Stan code was taken from an excellent course on Carpentries.

I maintain that the best way to read Stan is back-to-front, so here is some invalid Stan that I’ve inverted and annotated.

```{stan}
model {
  // f_x is sampled from the GP.
  // Remember it's a vector of f(x), not parameters describing f.
  f_x ~ multi_normal(rep_vector(0, n), K);

  // Likelihood is evaluated on f(x) plus normally distributed noise.
  y_obs ~ normal(f[1 : n_obs], sigma);

}
parameters {
  // Our outputs will be f(x) for all x (our observations, x_obs, and the predictions we want, x_pred)
  vector[n] f_x;
}
data {
  // Observed data
  int n_obs;
  array[n_obs] real x_obs;
  array[n_obs] real y_obs;
  
  // Observation error - pick a value
  real<lower=0> sigma;
  
  // x values for which we aim to predict f(x)
  int n_pred;
  array[n_pred] real x_pred;
  
  // Hyperparameters for the kernel
  real alpha;
  real lambda;
}
transformed data {
  // We join the x observations and desired x prediction points
  // because we want f(x) for both observed data and new predictions
  int n = n_obs + n_pred;
  
  array[n] real x;
  x[1 : n_obs] = x_obs;
  x[(n_obs + 1): n] = x_pred;
  
  // We calculate the Kernel values for all observed x
  matrix[n, n] K;
  K = gp_exp_quad_cov(x, alpha, lambda);
  
  // Add "nugget" on diagonal for numerical stability
  for (i in 1 : n) {
    K[i, i] = K[i, i] + 1e-6;
  }
}

```

The actual valid Stan file is here. An optimised version with Cholesky decomposition is here, which is the one we’ll run.

x.pred <- seq(domain[1], domain[2], by = 0.25)
x.vals <- c(observations$x, x.pred)

model <- cmdstan_model(stan_file = "gp-cholesky.stan", exe = "gp.stan.bin")
sample <- model$sample(
  seed = SEED,
  list(
    n_data = nrow(observations),
    x_data = as.array(observations$x),
    y_data = as.array(observations$y),
    # Sigma here is our guess about the noise level.
    sigma = 0.2,
    n_pred = length(x.pred),
    x_pred = x.pred,
    alpha = 1,
    lambda = 2
  ),
  parallel_chains = 4,
  max_treedepth = 20,
  show_messages = FALSE # disabled to avoid polluting the blog post, should be TRUE
)

sample

 variable   mean median   sd  mad     q5    q95 rhat ess_bulk ess_tail
   lp__   -83.63 -83.23 8.92 8.77 -99.08 -69.46 1.00     1503     2005
   eta[1]   0.13   0.13 0.19 0.18  -0.19   0.45 1.00     6394     2417
   eta[2]   0.27   0.27 0.19 0.20  -0.05   0.58 1.00     6378     2810
   eta[3]  -0.67  -0.68 0.20 0.20  -1.00  -0.35 1.00     6100     3056
   eta[4]  -0.94  -0.94 0.20 0.21  -1.27  -0.62 1.00     5973     2324
   eta[5]  -0.34  -0.34 0.20 0.19  -0.66  -0.02 1.00     6341     2760
   eta[6]  -0.01  -0.02 1.03 1.02  -1.69   1.69 1.00     6262     2636
   eta[7]  -0.01  -0.02 0.99 0.99  -1.65   1.65 1.00     6661     2761
   eta[8]   0.01   0.03 0.97 0.99  -1.59   1.62 1.00     5957     3165
   eta[9]   0.00   0.01 1.02 1.02  -1.69   1.66 1.00     5475     2417

 # showing 10 of 333 rows (change via 'max_rows' argument or 'cmdstanr_max_rows' option)

The Rhat and ESS are healthy (≤1.01 and >>100 respectively) so we can be happy. We wrangle the samples into (x, y) values and plot them alongside the original observations. It now looks a bit like a Christmas decoration.

Code

tidy_sample <- function(sample) {
  mat <- sample$draws(format = "draws_matrix")
  as_tibble(mat) |>
    # Every row is a draw, which we number
    mutate(draw = 1:n()) |>
    # Since each column is an observation f(x) for x indices, we pivot
    pivot_longer(starts_with("f"), names_to = "x", values_to = "y") |>
    # And map the x index back to an x value
    mutate(
      idx = as.numeric(str_extract(x, "[0-9]+")),
      x = x.vals[idx],
      y = as.numeric(y)
    )
}

plot_draws <- function(sample) {
  draws <- tidy_sample(sample)
  ggplot(observations) +
    aes(x = x, y = y) +
    geom_line(
      data = draws,
      mapping = aes(group = draw),
      alpha = 0.01,
      colour = "darkgreen"
    ) +
    geom_point(colour = "red", size = 3) +
    geom_function(fun = mystery_function, colour = "black", linetype = "dashed")
}

plot_draws(sample)

The Bayesian model has a much richer model of uncertainty than the linear fit. For example it’s much less uncertain around the positions where we have data.

Remember our original goal of finding the maximum value of the mystery function? We could narrow it down with more data. It would be most useful is if we had a way to determine the next best observation to acquire. Say we’re seeking the maximum value of . If we could add one more observation at the value of our choice, what should that be?

Enter the Acquisition Function

This is the acquisition function, which we maximise to find the best next to acquire. There are two objectives to balance: where is improvement likely, and where is there a potentially large improvement.

The most popular acquisition function is Expected Improvement, which does exactly that. For each it calculates the expected value of improvement.

Since we have a Bayesian model with many posterior samples, the maths is quite easy. We just look at each point on and take the mean improvement over the best (i.e. maximum if we’re aiming to maximise ). Whichever gave the biggest mean improvement is the lucky winner and should be acquired next in the experiment.

y.best <- max(observations$y)

expected_improvement <- function(sample, x.vals, y.best) {
  mat <- sample$draws(format = "draws_matrix")
  mat <- mat[, str_starts(colnames(mat), "f")]
  mat <- t(mat)

  n.points <- nrow(mat)
  ei <- numeric(length = n.points)

  for (i in 1:n.points) {
    samples <- mat[i, ]
    improvements <- pmax(0, samples - y.best)
    ei[i] <- mean(improvements)
  }

  x.vals[which.max(ei)]
}

x.next <- expected_improvement(sample, x.vals, y.best)
sprintf("Next x: %.2f", x.next)

[1] "Next x: 5.00"

Et voila, the model says we should measure at x=5. The maximum of the mystery function when solved analytically is in fact about 2.47, so this is intuitively a step in the right direction. Not bad, acquisition function!

Bayesian Optimisation

What if we do measure again, what then? We could rinse and repeat, until we converge on a maximum. In essence, that’s all there is to Bayesian Optimisation: fit the model, calculate your next acquisition with the expected improvement function, and fit again.

Of course the R community giveth; you can use tune_bayes rather than rolling your own optimiser.

So, there we have it. We fitted a Gaussian Process model on noisy observations from a mystery function, and used the expected improvement acquisition function to determine the next observation we should obtain if we want to find the function maximum, which is the first step in Bayesian Optimisation.

I look forward to seeing you again in six months when I have forgotten this all and need to relearn it once more!