Combining multi-fidelity and risk-adverse BO

[samuelkim16]

I'm trying to combine multi-fidelity and risk-aware BO, taken from this tutorial and this tutorial, respectively.

The acquisition function for multi-fidelity is a two-parter:

def get_mfkg(model):

    # Fix fidelity to highest value in order to evaluate information gain
    curr_val_acqf = FixedFeatureAcquisitionFunction(
        acq_function=PosteriorMean(model),
        d=problem_dim+1,
        columns=[problem_dim],
        values=[1],
    )

    _, current_value = optimize_acqf(
        acq_function=curr_val_acqf,
        bounds=bounds[:, :-1],
        q=1,
        num_restarts=10 if not SMOKE_TEST else 2,
        raw_samples=1024 if not SMOKE_TEST else 4,
        options={"batch_limit": 10, "maxiter": 200},
    )

    return qMultiFidelityKnowledgeGradient(
        model=model,
        num_fantasies=128 if not SMOKE_TEST else 2,
        current_value=current_value,
        cost_aware_utility=cost_aware_utility,
        project=project,
    )

def optimize_mfkg_and_get_observation(mfkg_acqf):
    """Optimizes MFKG and returns a new candidate, observation, and cost."""

    # generate new candidates
    candidates, _ = optimize_acqf_mixed(
        acq_function=mfkg_acqf,
        bounds=bounds,
        fixed_features_list=[{problem_dim: 0.0}, {problem_dim: 1.0}],
        q=BATCH_SIZE,
        num_restarts=NUM_RESTARTS,
        raw_samples=RAW_SAMPLES,
        # batch_initial_conditions=X_init,
        options={"batch_limit": 5, "maxiter": 200},
    )

    # observe new values
    cost = cost_model(candidates).sum()
    new_x = candidates.detach()
    new_obj = fitness_fun(new_x)
    print(f"candidates:\n{new_x}\n")
    print(f"observations:\n{new_obj}\n\n")
    return new_x, new_obj, cost

and the acquisition function for risk-adverse BO looks like this:

def optimize_acqf_and_get_observation():
    r"""Optimizes the acquisition function, and returns a new candidate and observation."""
    acqf = qNoisyExpectedImprovement(
        model=model,
        X_baseline=train_X,
        sampler=SobolQMCNormalSampler(sample_shape=torch.Size([128])),
        objective=risk_measure,
        prune_baseline=True,
    )

    candidate, _ = optimize_acqf(
        acq_function=acqf,
        bounds=bounds,
        q=BATCH_SIZE,
        num_restarts=NUM_RESTARTS,
        raw_samples=RAW_SAMPLES,
    )

    new_observations = fitness_fun(candidate)
    return candidate, new_observations

My first thought was to replace PosteriorMean(model) with acqf = qNoisyExpectedImprovement(...) but it gives a shape error:

AssertionError: Expected the output shape to match either the t-batch shape of X, or the `model.batch_shape` in the case of acquisition functions using batch models; but got output with shape torch.Size([128, 5, 16]) for X with shape torch.Size([128, 5, 1, 3]).

What would be the correct way to combine these acquisition functions?

[saitcakmak]

Hi @samuelkim16. I haven't worked with risk-averse BO in a while and I never tried to combine it with multi-fidelity. This is a non-standard use case and will likely require some debugging to get working. Here are some thoughts I have:

• qNEI seems like a strange choice for the value function to me (it measures improvement rather than value). A more direct replacement to PosteriorMean would be qSimpleRegret, its MC counterpart. qSimpleRegret should work with the objective=risk_measure argument to compute the risk measure as the value function.
• The error message that you shared catches for a use case where the acquisition function utilizes a model that has a one-to-many input transform but it lacks an objective that reduces the samples back to the original q-batch shape (this would be the risk measure). Knowing where the error is raised would be helpful here. Looking at qMultiFidelityKnowledgeGradient, I see that it internally constructs the value function with the given valfunc_cls, objective and other arguments. I'd try constructing qMultiFidelityKnowledgeGradient with these arguments and go from there.

[samuelkim16]

Thank you! I have swapped out the value function for qSimpleRegret and added the objective=risk_measure argument to qMultiFidelityKnowledgeGradient, so that my acquisition function looks like this:

# Acquisition
def get_mfkg(model, train_X):

    acqf = qSimpleRegret(
        model=model,
        sampler=SobolQMCNormalSampler(sample_shape=torch.Size([128])),
        objective=risk_measure,
    )

    # Fix fidelity to highest value in order to evaluate information gain
    curr_val_acqf = FixedFeatureAcquisitionFunction(
        acq_function=acqf,
        d=problem_dim+1,
        columns=[problem_dim],
        values=[1],
    )

    _, current_value = optimize_acqf(
        acq_function=curr_val_acqf,
        bounds=bounds[:, :-1],
        q=1,
        num_restarts=10 if not SMOKE_TEST else 2,
        raw_samples=1024 if not SMOKE_TEST else 4,
        options={"batch_limit": 10, "maxiter": 200},
    )

    return qMultiFidelityKnowledgeGradient(
        model=model,
        num_fantasies=128 if not SMOKE_TEST else 2,
        objective=risk_measure,
        current_value=current_value,
        cost_aware_utility=cost_aware_utility,
        project=project,
    )

It runs now with no error! I don't quite understand what the valfunc_cls does. Is this supposed to be the risk_measure rather than objective?

It also takes a while to run in each iteration. I'm guessing it is because it has to evaluate the risk measure inside an inner loop of optimizing the knowledge gradient? Do you have recommendations for speeding it up?

[saitcakmak]

I don't quite understand what the valfunc_cls does.

This is where you'd pass in qSimpleRegret but I believe it is set as default when an objective is provided.

It also takes a while to run in each iteration. I'm guessing it is because it has to evaluate the risk measure inside an inner loop of optimizing the knowledge gradient?

This is expected. KG is expensive to begin with and adding risk measures will only make it more expensive, since you have to joinly sample from multiple points to compute the risk measure.

Do you have recommendations for speeding it up?

Reducing num_fantasies should help. You can also pass in a sampler that uses a smaller number of samples. I believe the default sample shape is torch.Size([512]). You can try something like sampler=SobolQMCNormalSampler(sample_shape=torch.Size([128]) and decrease it further if it is still too slow.