Multi-fidelity optimization with KG and service API

Hi, First off, thanks for the great work! I’ve been trying to run multi-fidelity optimization using the one-shot-KG method from botorch with the service API. In principle, everything seems to run fine. However, I noticed that the optimizer keeps running in the lowest fidelity and does not explore in this parameter. To see I’ve this is an issue with my problem I tried to reproduce the MFKG example from the botorch documentation where the fidelity is clearly changed during the optimization, but I observe the same behaviour. This is the code I ran to reproduce the botorch example.

from ax.service.ax_client import AxClient
from botorch.test_functions.multi_fidelity import AugmentedHartmann
from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Models
import torch

problem = AugmentedHartmann(negate=True)
def objective(parameters):
    # x7 is the fidelity
    x = torch.tensor([parameters.get(f"x{i+1}") for i in range(7)])
    return {'f':(problem(x),0.0)}

gs = GenerationStrategy(
steps=[GenerationStep(model=Models.SOBOL,num_trials = 16),
GenerationStep(model=Models.GPKG,num_trials=-1),
])


ax_client = AxClient(generation_strategy=gs)
ax_client.create_experiment(
    name="hartmann_mf_experiment",
    parameters=[
        {
            "name": "x1",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "x2",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "x3",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "x4",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "x5",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "x6",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "x7",
            "type": "range",
            "bounds": [0.0, 1.0],
            "is_fidelity": True,
            "target_value": 1.
        },
    ],
    objective_name="f",
)
# Initial sobol samples
for i in range(16):
    parameters, trial_index = ax_client.get_next_trial()
    ax_client.complete_trial(trial_index=trial_index, raw_data=objective(parameters))

# KGBO
for i in range(6):
    q_p, q_t = [], []
    # Simulate batches
    for q in range(4):
        parameters, trial_index = ax_client.get_next_trial()
        q_p.append(parameters)
        q_t.append(trial_index)
    for q in range(4):
        pi = q_p[q]
        ti = q_t[q]
        ax_client.complete_trial(trial_index=ti, raw_data=objective(pi))

After the initial samples the fidelity stays at 0 except for two trial where it is very close to zero.

Are there any parameters that need to be specified for the acquisition function or is MFKG not yet supported with the service API?

Thanks!

[sgbaird]

@lena-kashtelyan Any thoughts on a multi-fidelity Ax tutorial? (e.g. based on what’s in this thread). Not sure if it’s worth it compared to other priorities, but figured I’d float the idea. Not a huge deal, just a drive-by comment.

[soerenjalas]

Hi, to follow up on this: I think I managed to get it working more like I would expect. The behaviour seems to be related to the cost_intercept value of the KG object:

GenerationStep(model=Models.GPKG,num_trials=-1,
               model_kwargs={'cost_intercept':5}, model_gen_kwargs={"num_fantasies":128})

After setting the parameters to match the named botorch example, the fidelity is better explored.