Implementing MC-Based Acquisition function

[qres]

Implementing MC-Based Acquisition function

I try to understand how MC-Based Acquisition functions are implemented. However, I'm unsure how the dimensions in qExpectedImprovement work.

My general setup looks like this:

import torch

from botorch.models import SingleTaskGP
from botorch.optim import optimize_acqf
from botorch.acquisition import qExpectedImprovement
from torch.quasirandom import SobolEngine

def func(x):
    return 1 + torch.sin(x[..., :1] )

DIM = 1
N_INIT = 2
BATCH_SIZE = 3

X = SobolEngine(dimension=DIM, scramble=True, seed=0).draw(N_INIT)
Y = func(X)
model = SingleTaskGP(X, Y)
acq = qExpectedImprovement(model=model, best_f=Y.min())

candidates_, acq_values = optimize_acqf(
    acq,
    bounds=torch.cat((-torch.ones(1, DIM), torch.ones(1, DIM))),
    q=BATCH_SIZE,
    raw_samples=512,
    num_restarts=20,
)

In the expected improvement acquisition function, I get the shape [512, 512, 3, 1] for samples once and then [512, 20, 3, 1] afterwards.

Judging from the implemented formula, which takes the mean over the first dimension, I assume that the first dimension is "number of samples used to approximate the acquisition functions value". And its value defaults to raw_samples.

Is this correct? And can I choose a different value for the number of starting points and for the number of samples?

Thank you!

[saitcakmak]

Hi @qres. You can parse the shape of the posterior samples as sample shape x t-batch shape x q-batch shape x number of outcomes.

For single objective acquisition functions, the default sampler has sample_shape = torch.Size([512]).
In optimize_acqf, raw_samples candidates (total shape of X raw_samples x q-batch x dim, with raw_samples corresponding to the t-batch shape) get randomly generated, the acquisition function is evaluated with these, and num_restarts candidates are selected from these with Boltzmann sampling. So, that's the 512 / 20 dimension you see.
The remaining dimensions are simply the q-batch shape = BATCH_SIZE and the number of outcomes of the model, which is 1.

If you're trying to understand the internals of an acquisition function, I'd recommend manually evaluating the various parts rather than using a utility like optimize_acqf -- since this hides a bunch of the intermediate steps. To see what the samples look like, you can do

X = torch.rand(32, 4, 1)
posterior = model.posterior(X)
samples = acq.get_posterior_samples(posterior)  # This should be 512 x 32 x 4 x 1 if using the default sampler.

[qres]

Thank you! Having the same number for the sample shape and number of raw samples by chance put me on the wrong track.
I now set the number of samples by explicitly setting the sampler in for the acquisition function.