3. Bayesian Optimization

Learn how to write a Bayesian optimization script to intelligently adjust RGB power levels to match a target color and visualize optimization results.

../../_images/grid-vs-random-vs-bayesian-ax-logo.png

An optimization efficiency comparison between grid search, random search, and Bayesian optimization using Meta’s Ax Platform, where Bayesian optimization is the most efficient at matching the target color.

A Gentle Introduction to Bayesian Optimization

Choosing the right optimization method depends on the specific needs and constraints of a project.

✅ Watch the Gentle Introduction to Bayesian Optimization below. This one is very important to understand the basics of Bayesian optimization and how it compares to traditional design of experiments.

✅ Watch Bayes rule by 3Blue1Brown. This will help you understand the most basic statistical concept that powers Bayesian optimization algorithms.

✅ Read the Why use quasi-random search? section from a quasi-random search guide

✅ Review the scipy documentation for two types of quasi-random search: Sobol and Latin Hypercube (LHE)

✅ Spend 20 minutes exploring the Ax Platform Docs

✅ Complete Ax’s Service API Tutorial [Colab link]

✅ Run the script below in a Python interpreter of your choice with the ax-platform dependency installed (i.e., pip install ax-platform). For convenience, you may use the Open in Colab badge in this companion notebook.

import math
from ax.service.ax_client import AxClient, ObjectiveProperties

obj1_name = "branin"

def branin(x1, x2):
    y = float(
        (x2 - 5.1 / (4 * math.pi**2) * x1**2 + 5.0 / math.pi * x1 - 6.0) ** 2
        + 10 * (1 - 1.0 / (8 * math.pi)) * math.cos(x1)
        + 10
    )

    return y

ax_client = AxClient()
ax_client.create_experiment(
    parameters=[
        {"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
        {"name": "x2", "type": "range", "bounds": [0.0, 15.0]},
    ],
    objectives={
        obj1_name: ObjectiveProperties(minimize=True),
    },
)

for _ in range(10):
    parameters, trial_index = ax_client.get_next_trial()
    results = branin(
        parameters["x1"],
        parameters["x2"],
    )
    ax_client.complete_trial(trial_index=trial_index, raw_data=results)

best_parameters, metrics = ax_client.get_best_parameters()

[WARNING 01-05 18:54:05] ax.service.utils.with_db_settings_base: Ax currently requires a sqlalchemy version below 2.0. This will be addressed in a future release. Disabling SQL storage in Ax for now, if you would like to use SQL storage please install Ax with mysql extras via `pip install ax-platform[mysql]`.
[INFO 01-05 18:54:05] ax.service.ax_client: Starting optimization with verbose logging. To disable logging, set the `verbose_logging` argument to `False`. Note that float values in the logs are rounded to 6 decimal points.
[INFO 01-05 18:54:05] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x1. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 01-05 18:54:05] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x2. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 01-05 18:54:05] ax.service.utils.instantiation: Created search space: SearchSpace(parameters=[RangeParameter(name='x1', parameter_type=FLOAT, range=[-5.0, 10.0]), RangeParameter(name='x2', parameter_type=FLOAT, range=[0.0, 15.0])], parameter_constraints=[]).
[INFO 01-05 18:54:05] ax.modelbridge.dispatch_utils: Using Models.BOTORCH_MODULAR since there are more ordered parameters than there are categories for the unordered categorical parameters.
[INFO 01-05 18:54:05] ax.modelbridge.dispatch_utils: Calculating the number of remaining initialization trials based on num_initialization_trials=None max_initialization_trials=None num_tunable_parameters=2 num_trials=None use_batch_trials=False
[INFO 01-05 18:54:05] ax.modelbridge.dispatch_utils: calculated num_initialization_trials=5
[INFO 01-05 18:54:05] ax.modelbridge.dispatch_utils: num_completed_initialization_trials=0 num_remaining_initialization_trials=5
[INFO 01-05 18:54:05] ax.modelbridge.dispatch_utils: `verbose`, `disable_progbar`, and `jit_compile` are not yet supported when using `choose_generation_strategy` with ModularBoTorchModel, dropping these arguments.
[INFO 01-05 18:54:05] ax.modelbridge.dispatch_utils: Using Bayesian Optimization generation strategy: GenerationStrategy(name='Sobol+BoTorch', steps=[Sobol for 5 trials, BoTorch for subsequent trials]). Iterations after 5 will take longer to generate due to model-fitting.
[INFO 01-05 18:54:05] ax.service.ax_client: Generated new trial 0 with parameters {'x1': 1.866556, 'x2': 12.412659}.
[INFO 01-05 18:54:05] ax.service.ax_client: Completed trial 0 with data: {'branin': (87.005011, None)}.
[INFO 01-05 18:54:05] ax.service.ax_client: Generated new trial 1 with parameters {'x1': 2.569913, 'x2': 9.499428}.
[INFO 01-05 18:54:05] ax.service.ax_client: Completed trial 1 with data: {'branin': (47.303482, None)}.
[INFO 01-05 18:54:05] ax.service.ax_client: Generated new trial 2 with parameters {'x1': 0.665525, 'x2': 5.085759}.
[INFO 01-05 18:54:05] ax.service.ax_client: Completed trial 2 with data: {'branin': (17.560653, None)}.
[INFO 01-05 18:54:05] ax.service.ax_client: Generated new trial 3 with parameters {'x1': 2.19972, 'x2': 1.36233}.
[INFO 01-05 18:54:05] ax.service.ax_client: Completed trial 3 with data: {'branin': (7.455261, None)}.
[INFO 01-05 18:54:05] ax.service.ax_client: Generated new trial 4 with parameters {'x1': 7.092661, 'x2': 7.112501}.
[INFO 01-05 18:54:05] ax.service.ax_client: Completed trial 4 with data: {'branin': (51.458909, None)}.
[INFO 01-05 18:54:06] ax.service.ax_client: Generated new trial 5 with parameters {'x1': -2.027617, 'x2': 0.0}.
[INFO 01-05 18:54:06] ax.service.ax_client: Completed trial 5 with data: {'branin': (100.986227, None)}.
[INFO 01-05 18:54:07] ax.service.ax_client: Generated new trial 6 with parameters {'x1': 3.122451, 'x2': 3.793926}.
[INFO 01-05 18:54:08] ax.service.ax_client: Completed trial 6 with data: {'branin': (2.661517, None)}.
[INFO 01-05 18:54:09] ax.service.ax_client: Generated new trial 7 with parameters {'x1': 4.977416, 'x2': 0.0}.
[INFO 01-05 18:54:09] ax.service.ax_client: Completed trial 7 with data: {'branin': (14.15021, None)}.
[INFO 01-05 18:54:11] ax.service.ax_client: Generated new trial 8 with parameters {'x1': 3.383846, 'x2': 2.136473}.
[INFO 01-05 18:54:11] ax.service.ax_client: Completed trial 8 with data: {'branin': (0.680103, None)}.
[INFO 01-05 18:54:13] ax.service.ax_client: Generated new trial 9 with parameters {'x1': 10.0, 'x2': 0.0}.
[INFO 01-05 18:54:13] ax.service.ax_client: Completed trial 9 with data: {'branin': (10.960889, None)}.
[WARNING 01-05 18:54:13] ax.modelbridge.cross_validation: Metric branin was unable to be reliably fit.
[WARNING 01-05 18:54:13] ax.service.utils.best_point: Model fit is poor; falling back on raw data for best point.
[WARNING 01-05 18:54:13] ax.service.utils.best_point: Model fit is poor and data on objective metric branin is noisy; interpret best points results carefully.

Additional Resources

For the mathematically inclined or those who would like to understand the theory behind Bayesian optimization, you are encouraged to work through A Visual Exploration of Gaussian Processes, which will give you a better sense of what Gaussian processes are and how Bayesian inference is used to perform Gaussian process regression. Note that Gaussian processes are the most widely used surrogate model in Bayesian optimization applications. Next, you can work through Exploring Bayesian Optimization. These two resources will help you to develop intuition about Bayesian optimization by understanding the underlying theory.

You may also be interested in the following resources:

That’s it! You’ve completed the tutorial 🎉. Return to the course website to do a knowledge check.