Ax: how to force models to make more exploration

[ertandemiral]

Hi,

I would like to thank everyone who contributed to this great library. It enables easy use of Bayesian optimization to solve problems with the state of the art algorithms.

I have implemented Ax for my single objective design optimization study. Here is the code snippet:
`
from ax.service.ax_client import AxClient
from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Models

def objective_function(x):
# region of f calculation
# gives 'ErrorDesign' in case of error, otherwise float.
return {"f": (f, 0.0)}

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

ax_client = AxClient(generation_strategy=gs)
ax_client.create_experiment(
name="single_objective_design",
parameters=[
{"name": "x1", "type": "range","bounds": [0.2, 1.0],"value_type": "float"},
{"name": "x2", "type": "range","bounds": [2.0, 6.0],"value_type": "float"},
{"name": "x3", "type": "range","bounds": [0.2, 1.0],"value_type": "float"},
{"name": "x4", "type": "range","bounds": [1.7, 8.7],"value_type": "float"},
{"name": "x5", "type": "range","bounds": [ 0, 25],"value_type": "int"},
{"name": "x6"," type": "range","bounds": [4.0,12.0],"value_type": "float"},
{"name": "x7", "type": "range","bounds": [2.0, 5.0],"value_type": "float"},
{"name": "x8", "type": "range","bounds": [0.2, 1.0],"value_type": "float"},
{"name": "x9", "type": "range","bounds": [80., 95.],"value_type": "float"},
{"name": "x10","type": "range","bounds": [ 0, 25],"value_type": "int"},
{"name": "x11","type": "choice","values":["4","8","12","16"],"value_type": "str"},
{"name": "x12","type": "choice","values":["4","8","12","16"],"value_type": "str"},
],
objective_name="f",
minimize=True)

for _ in range(200):
trial_params, trial_index = ax_client.get_next_trial()
data = objective_function(trial_params)
if data["f"][0] == 'ErrorDesign':
ax_client.log_trial_failure(trial_index=trial_index)
else:
ax_client.complete_trial(trial_index=trial_index, raw_data=data["f"])
`

I have 12 design parameters (10 ranges, 2 choices) to be optimized and benefit service API with generation strategies ([sobol + gpmes, sobol + gpei, sobol + botorch, sobol + gpkg]) as seen in the code snippet. I am using python3.8 and the latest versions of botorch, gpytorch, and torch libraries.

Below is the history plot showing objective values with respect to iteration number for different models after running code respectively. I have also added the history of design parameters for the GPEI Model.

My question is about the non-explorative search behavior of the models after 20 sobol iterations. As you see from the objective history figure, successive designs have close objective values. Indeed, I would expect the code to do more exploration since the search space is quite large, but each model quickly converge some local minimum and continue to search around that minimum. By the way, the global minimum of the objective function is around -3.6.

I have tried the followings, but code behavior is not much affected:

• Repeated runs with different sobol initializations
• Increasing sobol trial number
• Increasing num_fantasies, num_mv_samples, num_y_samples, candidate_size

Any help to force these generation strategies into making more exploration would be appreciated.
Thanks in advance.

[Balandat]

So I guess it really depends on the behavior of f how much additional exploration you'd want. By default we use NoisyExpectedImprovement as the acquisition function, which generally works well but there may be cases where something like an upper confidence bound method works better. If it were me, I would probably try to see how UCB does here.

Orthogonal to that, another challenge may be the choice parameters. There is a know issue with our optimization strategy that can result in excessive repeated evaluations of the same choice values. We have a fix for this in the works, but in the meantime maybe you can help identify whether that's the issue for you here. If you know the optimal choice parameter values, could you remove them from the search space (hard-code them to the optimal values in your evaluation function), and see whether you still observe the same behavior? If not, then this issue is likely the culprit here.

[ertandemiral]

Thank you for the quick response,

Models.BOTORCH (red line circle) seems to use noisy expected improvement as acquisition function by default as you said. It is usually giving better results (lower objective values) compared to other models, however, there is no significant difference between objective value behaviors. I would like to try UCB for my study, but I could not find related model under ax.modelbridge.registry.Models. Is there any easy way to use UCB in the generation strategy?

I will remove choice parameters and hardcode them as you recommend. Evaluating the objective function is a little costly, so it may take a few days to finish. I will post the results as soon as possible.

Thank you.

[lena-kashtelyan]

@ertandemiral, we are currently wrapping up work on a new BoTorchModel setup in Ax that will allow you to plug in any BoTorch AcquisitionFunction into Ax more easily. Discussion of that is here: https://github.com/facebook/Ax/issues/363, and I'll post a new and improved tutorial for the new setup in that issue shortly.

[ertandemiral]

Thank you @lena-kashtelyan ,

I managed to plug in UCB acquisition function into Ax by creating get_UCBfunction factory and passing it as kwargs to GenerationStep. And, I also passed the necessary input values of betaand maximizekeywords (0.2, False) as kwargs and added them as argument to the make_and_optimize_acqffunction. Code has run with no error, however I could not get any better results when repeating the study using get_UCBfunction factory in the presence of choice parameters. Here is the code snippet:

def get_UCB(
    model: Model,
    beta: Union[float, Tensor],
    maximize: bool,
    **kwargs: Any,
) -> AcquisitionFunction:
    return UpperConfidenceBound(model=model, beta=beta, maximize=maximize)

beta_value=0.2
maximize_value=False

gs = GenerationStrategy(
 steps=[
    GenerationStep(model=Models.SOBOL,num_trials = 20),
    GenerationStep(model=Models.BOTORCH,num_trials = -1,model_kwargs={"acqf_constructor": get_UCB,"kwargs":[beta_value,maximize_value]}),
      ]
 )

On the other hand, when the choice parameters are excluded from the search space (as @Balandat suggested), searching behavior of the models are affected positively such that they always converge the global minimum point with increased exploration behavior (below figure). Almost all models especially GPMESare giving superior results in that case. I will rerun the study when the fix regarding choice parameters is completed.

I have another question about the issue of easily plugging in the BoTorch acquisition functions. I also want to conduct active learning experiment in Ax utilizing BoTorch qNegIntegratedPosteriorVarianceacquisition function. I implemented it by creating get_qNIPVfactory function as in the UCB one . I am getting satisfactory results. I just want to ask what to set minimizekeyword of the create_experiment. I don't think keyword value will affect the experiment, but I would appreciate it if you could comment on this.
Here is the code snippet:

def objective_function(x):
    # skipped --> region of f calculation
    return  {"f": (f, 0.0)}

from botorch.acquisition.active_learning import qNegIntegratedPosteriorVariance
def get_qNIPV(
    model: Model,
    mc_points: Tensor,
    **kwargs: Any,
) -> AcquisitionFunction:
    return qNegIntegratedPosteriorVariance(model=model,mc_points=mc_points)

# skipped --> 1000 mc_points are generated using sobol sequence considering the search space  

gs = GenerationStrategy(
    steps=[GenerationStep(model=Models.SOBOL,num_trials = 5),
           GenerationStep(model=Models.BOTORCH,num_trials = -1, model_kwargs={"acqf_constructor": get_qNIPV, "kwargs":[mc_points,"sil"]}),
          ])

ax_client = AxClient(generation_strategy=gs)
ax_client.create_experiment(
    name="active_learning_experiment",
    parameters=[      
        {"name": "x1", "type": "range","bounds": [0.0, 1.0],"value_type": "float"},
        {"name": "x2", "type": "range","bounds": [0.0, 1.0],"value_type": "float"},
        {"name": "x3", "type": "range","bounds": [0.0, 1.0],"value_type": "float"},
        {"name": "x4", "type": "range","bounds": [0.0, 1.0],"value_type": "float"},
        {"name": "x5", "type": "range","bounds": [0.0, 1.0],"value_type": "float"},
        {"name": "x6"," type": "range","bounds": [0.0, 1.0],"value_type": "float"},    
    ],
    objective_name="f",
    minimize=True)

for _ in range(500):
    trial_params, trial_index = ax_client.get_next_trial()
    data = objective_function(trial_params)
    ax_client.complete_trial(trial_index=trial_index, raw_data=data["f"])

Thank you.

[ldworkin]

@ertandemiral I'm glad to hear that removing the choice parameters seems to work for you!

re: the minimize keyword, that's a good question -- given your custom setup, it seems likely to me that this parameter has no effect, but maybe @lena-kashtelyan has a better idea.

[lena-kashtelyan]

@ertandemiral, I believe @ldworkin is correct, minimize will take no effect on the optimization itself. That keyword results in setting the direction of the Objective object in Ax, which in turn sets the objective direction in BoTorch, but given that your objective is constant, I don't think that setting will have a meaning in this case.

cc @Balandat in case he has thoughts on the setup

[Balandat]

Yeah qNIPV is agnostic to the direction, the goal is to minimize a global measure of uncertainty of the model, so there is no better or worse w.r.t. the function values.

[ertandemiral]

Thank you for your comments.