Ax: Reducing Domain / Bounds Each Run

[kjanoudi]

Hello,

I'm coming from a simpler package, and they have a concept of "Domain Reduction", which when used with my black box function has been very successful in rapidly finding near optimal ranges for 26 parameters, some Choice and some Int/Float Range. I've linked it here:

https://github.com/fmfn/BayesianOptimization/blob/master/bayes_opt/domain_reduction.py
https://github.com/fmfn/BayesianOptimization/blob/master/examples/domain_reduction.ipynb

I've tried running Ax on the same data, configured in a similar fashion without the "Domain Reduction" extra above, and gotten significantly worse results. My attempts were run with: Sobol 30 steps, 150 GPEI trials, as well as 250 trials. It never came close to optimizing the target.

I've been struggling to understand two things:

1) Is this somewhat already implemented (bounds tightening, domain reduction, in any sense) in this project?
2) If not, what would be my best path to adding something like this within this project?

Thanks

[2timesjay]

@kjanoudi

We don't have anything like Domain reduction directly built in, but using the Developer API you can achieve this by directly editing the search space.

If you take that tutorial and replace the code under https://ax.dev/tutorials/gpei_hartmann_developer.html#4.-Perform-Optimization with the following:

from ax.service.utils.best_point import get_best_raw_objective_point
from ax.core.experiment import Experiment

def trim_range(best: float, lower: float, upper: float):
    shrink_factor = 0.8
    trimmed_lower = best - shrink_factor * (best - lower)
    trimmed_upper = best + shrink_factor * (upper - best)
    return {"lower": trimmed_lower, "upper": trimmed_upper}

def trim_search_space(original_search_space: SearchSpace, experiment: Experiment) -> SearchSpace:
    """Given a search space an an experiment, return a trimmed search space centered on the best arm so far."""
    best_params = get_best_raw_objective_point(experiment)[0]
    trimmed_search_space = SearchSpace(
        parameters=[
            RangeParameter(
                name=param.name, parameter_type=param.parameter_type, 
                **trim_range(best=best_params[param.name], lower=param.lower, upper=param.upper)
            )
            if isinstance(param, RangeParameter) else param
            for param in original_search_space.parameters.values()
        ]
    )
    return trimmed_search_space

print(f"Running Sobol initialization trials...")
sobol = Models.SOBOL(exp.search_space)
for i in range(5):
    exp.new_trial(generator_run=sobol.gen(1))
search_space = hartmann_search_space

for domain_reduction_iters in range(5):
    search_space_volume = np.prod([param.upper-param.lower for param in search_space.parameters.values()])
    print(f"Searching space with volume {search_space_volume}")
    for i in range(5):
        print(f"Running GP+EI optimization trial {i+1}/5 on domain reduction #{domain_reduction_iters+1}/5...")
        # Reinitialize GP+EI model at each step with updated data.
        gpei = Models.BOTORCH(experiment=exp, data=exp.eval(), search_space=search_space, fit_out_of_design=True)
        batch = exp.new_trial(generator_run=gpei.gen(1))


    search_space = trim_search_space(original_search_space=search_space, experiment=exp)    

print("Done!")

You'll see a very crude version of Domain reduction applied to your problem.

We also have some not-yet-released features that may more closely match what you need - I'll follow up with a roadmap on those. (@sdaulton, @Balandat)

[kjanoudi]

Thank you greatly for the detailed help, looking forward to the roadmap

[kjanoudi]

Is there a way to mix the Developer API with the service API? I am taking advantage of the experiment, trial, and metric data being saved to SQL using the AxClient, and would like to keep that process in place using the code you provided above. Would the best way be to fork the repo, or is there a better option?

I was considering:

ax_client = AxClient(db_settings=db_settings, generation_strategy=gs)
ax_client.create_experiment(...params)
exp = ax_client.experiment
exp.search_space = trim_search_space(original_search_space=search_space, experiment=exp)

[lena-kashtelyan]

@kjanoudi, you should definitely be able to do what you suggest above. The experiment you get from ax_client.experiment is exactly the same as the Dev API Experiment, so anything you can do to the Dev API experiment, you can do through the Service API.

[Balandat]

cc @dme65 for TuRBO

[dme65]

@kjanoudi As 26 dimensions is quite high for vanilla BO you can try to use TuRBO which was developed for the high-dimensional setting: code, paper. TuRBO isn't available in BoTorch, but we are looking into a few options to support it.

[lena-kashtelyan]

Since formalized domain reduction is not currently on our roadmap and a viable workaround has been provided by @2timesjay, closing this issue. Feel free to reopen if you have further questions, @kjanoudi!