[GENERAL SUPPORT]: fixed_features

[omrisch]

Question

Issue:
When using the fixed_features keyword argument with get_next_trials under GPEI, the initial suggested parameterizations consistently fall on the extremes or corners of the defined search space. Only after providing results for these initial trials does the client begin generating more meaningful or high-quality parameterizations.

Expected behavior:
Early suggestions should incorporate fixed features while still producing diverse, informative parameterizations rather than defaulting to boundary values.

Steps to reproduce:

Define a search space with at least 3 range parameters

generate sufficient Explore parameterizations and provide the values of these parameterizations using complete_trial

Call get_next_trials using GPEI, provide a fixed_features kwarg for one of the range parameters

Observe that the first few suggestions cluster at the search space corners. Provide the values of these parameterizations using complete_trial

Please provide any relevant code snippet if applicable.

Code of Conduct

• I agree to follow this Ax's Code of Conduct

[mpolson64]

Hi @omrisch I was not able to reproduce this behavior on my machine, but hopefully I can offer some advice. This is the code I ran, based largely off the Quickstart tutorial on ax.dev:

from ax.api.client import Client
from ax.api.configs import RangeParameterConfig

# 1. Initialize the Client.
client = Client()

# 2. Configure where Ax will search.
client.configure_experiment(
    name="booth_function",
    parameters=[
        RangeParameterConfig(
            name=f"x{i + 1}",
            bounds=(-10.0, 10.0),
            parameter_type="float",
        ) for i in range(3)
    ],
)

client.configure_optimization(objective="-1 * booth")

for _ in range(6):
    # Use higher value of `max_trials` to run trials in parallel.
    for trial_index, parameters in client.get_next_trials(max_trials=1).items():
        client.complete_trial(
            trial_index=trial_index,
            raw_data={
                "booth": (parameters["x1"] + 2 * parameters["x2"] - 7) ** 2
                + (2 * parameters["x1"] + parameters["x2"] - 5) ** 2 + parameters["x3"]
            },
        )

for _ in range(5):
    # Use higher value of `max_trials` to run trials in parallel.
    for trial_index, parameters in client.get_next_trials(max_trials=1, fixed_parameters={"x3": 0}).items():
        client.complete_trial(
            trial_index=trial_index,
            raw_data={
                "booth": (parameters["x1"] + 2 * parameters["x2"] - 7) ** 2
                + (2 * parameters["x1"] + parameters["x2"] - 5) ** 2 + parameters["x3"]
            },
        )

client.summarize()

	trial_index	arm_name	trial_status	generation_node	booth	x1	x2	x3
0	0	0_0	COMPLETED	CenterOfSearchSpace	74	0	0	0
1	1	1_0	COMPLETED	Sobol	403.583	3.86475	-8.09843	1.02716
2	2	2_0	COMPLETED	Sobol	22.8164	-3.16963	6.24109	-8.52288
3	3	3_0	COMPLETED	Sobol	1490.85	-9.92876	-4.19687	5.46414
4	4	4_0	COMPLETED	Sobol	210.229	8.0459	2.3396	-2.9501
5	5	5_0	COMPLETED	MBM	48.6569	0.988403	6.28112	-4.86803
6	6	6_0	COMPLETED	MBM	2.89228	-0.0404347	4.26681	0
7	7	7_0	COMPLETED	MBM	139.922	-1.38374	10	0
8	8	8_0	COMPLETED	MBM	234	-10	10	0
9	9	9_0	COMPLETED	MBM	19.2575	-1.69583	4.04527	0
10	10	10_0	COMPLETED	MBM	3.93108	-0.304396	3.62675	0

Regardless, it can be common for Bayesian optimization techniques to sample the boundaries of the search space -- there is often more uncertainty in the GP close to the bounds and EI will tend to sample in places of high uncertainty if there is not a particularly promising point which beats the incumbent. This is especially true for a parameter which does not have a big effect on the metric. Considering this, it doesn't surprise me that fixing some parameters (i.e. forcing Ax to sample from a hyperplane within the search space with comparatively high uncertainty) results in Ax sampling the boundaries of the free parameters.

...suggestions should incorporate fixed features while still producing diverse, informative parameterizations

While I would argue that these boundary values are useful to help the GP understand the space, if you would like to avoid them one solution could be to reset the GenerationStrategy to Sobol (pseudorandom) before you being generating candidate trials with fixed parameterizations by calling client.configure_generation_strategy(initialization_budget=9999, initialize_with_center=False) -- this will result in more diverse parameterizations.

Let me know if you found this helpful or if you have any follow up questions.

[mpolson64]

Code sample with the adjustment:

from ax.api.client import Client
from ax.api.configs import RangeParameterConfig

# 1. Initialize the Client.
client = Client()

# 2. Configure where Ax will search.
client.configure_experiment(
    name="booth_function",
    parameters=[
        RangeParameterConfig(
            name=f"x{i + 1}",
            bounds=(-10.0, 10.0),
            parameter_type="float",
        ) for i in range(3)
    ],
)

client.configure_optimization(objective="-1 * booth")

for _ in range(5):
    # Use higher value of `max_trials` to run trials in parallel.
    for trial_index, parameters in client.get_next_trials(max_trials=1).items():
        client.complete_trial(
            trial_index=trial_index,
            raw_data={
                "booth": (parameters["x1"] + 2 * parameters["x2"] - 7) ** 2
                + (2 * parameters["x1"] + parameters["x2"] - 5) ** 2 + parameters["x3"]
            },
        )

# Add this line here to return to Sobol sampling
client.configure_generation_strategy(initialization_budget=9999, initialize_with_center=False)

for _ in range(5):
    # Use higher value of `max_trials` to run trials in parallel.
    for trial_index, parameters in client.get_next_trials(max_trials=1, fixed_parameters={"x3": 0}).items():
        client.complete_trial(
            trial_index=trial_index,
            raw_data={
                "booth": (parameters["x1"] + 2 * parameters["x2"] - 7) ** 2
                + (2 * parameters["x1"] + parameters["x2"] - 5) ** 2 + parameters["x3"]
            },
        )

client.summarize()

[omrisch]

Thanks, that matches our thinking.
Is it possible to replace EI with a more exploitative acquisition function, and would that help direct suggestions toward higher-certainty points away from the edges of the hyperplane?

[mpolson64]

You can follow this tutorial for how to use different acquisition functions from Botorch within Ax https://ax.dev/docs/tutorials/modular_botorch/