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: #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_UCB function factory and passing it as kwargs to GenerationStep. And, I also passed the necessary input values of beta and maximize keywords (0.2, False) as kwargs and added them as argument to the make_and_optimize_acqf function. Code has run with no error, however I could not get any better results when repeating the study using get_UCB function 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 GPMES are 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 qNegIntegratedPosteriorVariance acquisition function. I implemented it by creating get_qNIPV factory function as in the UCB one . I am getting satisfactory results. I just want to ask what to set minimize keyword 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.

[samueljamesbell]

Sorry to revive this thread, but it seems like the most appropriate place to ask.

I have something quite similar to @ertandemiral's set-up above, though I can't work out what dimension mc_points should be.

My assumption was that it should be N x D - N sampled points from across the D-dimensional search space, but I think I might be missing a batch dim.

Can anyone advise what qNegIntegratedPosteriorVariance expects?

Thanks so much!

[ertandemiral]

Hi @samueljamesbell ,

I have updated recently my setup after the addition of BOTORCH_MODULAR feature to ax. I have added the new setup below and recommend using it:

from ax.modelbridge import get_sobol
import torch
from botorch.acquisition.active_learning import qNegIntegratedPosteriorVariance
from ax import ParameterType, RangeParameter, SearchSpace
from botorch.models.gp_regression import SingleTaskGP
from ax.models.torch.botorch_modular.surrogate import Surrogate
from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Models
from ax.service.ax_client import AxClient

def objective_function(x):
    f = x["x1"]**2 + x["x2"]**2 + x["x3"]**2
    return  {"f": (f, None)}

search_space = SearchSpace(parameters = [RangeParameter(name ="x1", lower = 0.0, upper = 1.0, parameter_type = ParameterType.FLOAT),
                                         RangeParameter(name ="x2", lower = 0.0, upper = 1.0, parameter_type = ParameterType.FLOAT),
                                         RangeParameter(name ="x3", lower = 0.0, upper = 1.0, parameter_type = ParameterType.FLOAT),
                                         ])
sobol = get_sobol(search_space)
mc_points = sobol.gen(1024).param_df.values
mcp = torch.tensor(mc_points)

model_kwargs_val = {"surrogate" : Surrogate(SingleTaskGP),
              "botorch_acqf_class" : qNegIntegratedPosteriorVariance,
              "acquisition_options" : {"mc_points" : mcp}}

gs = GenerationStrategy(steps = [GenerationStep(model = Models.SOBOL,           num_trials = 5),
                                 GenerationStep(model = Models.BOTORCH_MODULAR, num_trials = 15, model_kwargs = model_kwargs_val)])

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"}, 
    ],
    objective_name = "f",
    minimize = True)

for _ in range(20):
    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"])

To make run above code, input constructor for acquisition class qNegIntegratedPosteriorVariance should be registered in \botorch\acquisiton\input_constructer.py file of botorch library. So, the below code should also be appended in the corresponding file:

@acqf_input_constructor(qNegIntegratedPosteriorVariance)
def construct_inputs_qNIPV(
    model: Model,
    mc_points: Tensor,
    training_data: TrainingData,
    objective: Optional[ScalarizedObjective] = None,
    X_pending: Optional[Tensor] = None,
    sampler: Optional[MCSampler] = None,
    **kwargs: Any,
) -> Dict[str, Any]:
    
    if model.num_outputs == 1 : objective = None
    
    base_inputs = construct_inputs_mc_base(
        model=model,
        training_data=training_data,
        sampler=sampler,
        X_pending=X_pending,
        objective = objective,
    )

    return {**base_inputs, "mc_points": mc_points}

Dimension problem may occur due to the objective variable being considered as multi-output in NegIntegratedPosteriorVariance class if it is not equal to None. In above registration code, it is set to None for single output case. And, mc_points can be given in N x D format for this setup. I hope it helps and developers may correct me if anything is wrong.

[samueljamesbell]

Hey @ertandemiral, thanks so much for this!

Turns out this line:

    if model.num_outputs == 1 : objective = None

was exactly what I was missing. Thanks for the help!

[Kh-im]

Hello,

@ertandemiral solution is working great but did you see the slicing results ? Let's say I use exactly the same code but with objective function : f = x["x1"]**2.0,

The slicing results are :

Do you understand what's happening here ?

Thanks for the help !

[Kh-im]

#930

[ertandemiral]

Hello @Kh-im ,

I obtained properly fitting slicing plot after running experiment with your objective function: (installed ax version is 0.2.9)

[sgbaird]

xref: pytorch/botorch#422