Open in GitHubRun in Google Colab

Ask-tell experimentation with trial-level early stopping

Trial-level early stopping aims to monitor the results of expensive evaluations with timeseries-like data and terminate those that are unlikely to produce promising results prior to completing that evaluation. This reduces computational waste, and enables the same amount of resources to explore more configurations. Early stopping is useful for expensive to evaluate problems where stepwise information is available on the way to the final measurement.

Like the ask-tell tutorial we'll be minimizing the Hartmann6 function, but this time we've modified it to incorporate a new parameter $t$ which allows the function to produce timeseries-like data where the value returned is closer and closer to Hartmann6's true value as $t$ increases. At $t = 100$ the function will simply return Hartmann6's unaltered value.

$$f(x, t) = hartmann6(x) - log_2(t/100)$$

While the function is synthetic, the workflow captures the intended principles for this tutorial and is similar to the process of training typical machine learning models.

Learning Objectives

• Understand when time-series-like data can be used in an optimization experiment
• Run a simple optimization experiment with early stopping
• Configure details of an early stopping strategy
• Analyze the results of the optimization

Prerequisites

• Familiarity with Python and basic programming concepts
• Understanding of adaptive experimentation and Bayesian optimization
• Ask-tell Optimization of Python Functions

Step 1: Import Necessary Modules

First, ensure you have all the necessary imports:

import numpy as np
import plotly.express as px
import plotly.graph_objects as go
from ax.early_stopping.strategies import PercentileEarlyStoppingStrategy
from ax.preview.api.client import Client
from ax.preview.api.configs import ExperimentConfig, ParameterType, RangeParameterConfig

Step 2: Initialize the Client

Create an instance of the Client to manage the state of your experiment.

client = Client()

Step 3: Configure the Experiment

The Client instance can be configured with a series of Configs that define how the experiment will be run.

The Hartmann6 problem is usually evaluated on the hypercube $x_i \in (0, 1)$, so we will define six identical RangeParameterConfigs with these bounds and add these to an ExperimentConfig along with other metadata about the experiment.

You may specify additional features like parameter constraints to further refine the search space and parameter scaling to help navigate parameters with nonuniform effects. For more on configuring experiments, see this recipe.

# Define six float parameters for the Hartmann6 function
parameters = [
    RangeParameterConfig(
        name=f"x{i + 1}", parameter_type=ParameterType.FLOAT, bounds=(0, 1)
    )
    for i in range(6)
]

# Create an experiment configuration
experiment_config = ExperimentConfig(
    name="hartmann6_experiment",
    parameters=parameters,
    # The following arguments are optional
    description="Optimization of the Hartmann6 function",
    owner="developer",
)

# Apply the experiment configuration to the client
client.configure_experiment(experiment_config=experiment_config)

Step 4: Configure Optimization

Now, we must configure the objective for this optimization, which we do using Client.configure_optimization. This method expects a string objective, an expression containing either a single metric to maximize, a linear combination of metrics to maximize, or a tuple of multiple metrics to jointly maximize. These expressions are parsed using SymPy. For example:

• "score" would direct Ax to maximize a metric named score
• "-loss" would direct Ax to Ax to minimize a metric named loss
• "task_0 + 0.5 * task_1" would direct Ax to maximize the sum of two task scores, downweighting task_1 by a factor of 0.5
• "score, -flops" would direct Ax to simultaneously maximize score while minimizing flops

For more information on configuring objectives and outcome constraints, see this recipe.

client.configure_optimization(objective="-hartmann6")

Step 5: Run Trials with early stopping

Here, we will configure the ask-tell loop.

We begin by defining our Hartmann6 function as written above. Remember, this is just an example problem and any Python function can be substituted here.

Then we will iteratively do the following:

• Call client.get_next_trials to "ask" Ax for a parameterization to evaluate
• Evaluate hartmann6_curve using those parameters in an inner loop to simulate the generation of timeseries data
• "Tell" Ax the partial result using client.attach_data
• Query whether the trial should be stopped via client.should_stop_trial_early
• Stop the underperforming trial and report back to Ax that is has been stopped

This loop will run multiple trials to optimize the function.

Ax will configure an EarlyStoppingStrategy when should_stop_trial_early is called for the first time. By default Ax uses a Percentile early stopping strategy which will terminate a trial early if its performance falls below a percentile threshold when compared to other trials at the same step. Early stopping can only occur after a minimum number of progressions to prevent premature early stopping. This validates that both enough data is gathered to make a decision and there is a minimum number of completed trials with curve data; these completed trials establish a baseline.

# Hartmann6 function
def hartmann6(x1, x2, x3, x4, x5, x6):
    alpha = np.array([1.0, 1.2, 3.0, 3.2])
    A = np.array(
        [
            [10, 3, 17, 3.5, 1.7, 8],
            [0.05, 10, 17, 0.1, 8, 14],
            [3, 3.5, 1.7, 10, 17, 8],
            [17, 8, 0.05, 10, 0.1, 14],
        ]
    )
    P = 10**-4 * np.array(
        [
            [1312, 1696, 5569, 124, 8283, 5886],
            [2329, 4135, 8307, 3736, 1004, 9991],
            [2348, 1451, 3522, 2883, 3047, 6650],
            [4047, 8828, 8732, 5743, 1091, 381],
        ]
    )

    outer = 0.0
    for i in range(4):
        inner = 0.0
        for j, x in enumerate([x1, x2, x3, x4, x5, x6]):
            inner += A[i, j] * (x - P[i, j]) ** 2
        outer += alpha[i] * np.exp(-inner)
    return -outer


# Hartmann6 function with additional t term such that
# hartmann6(X) == hartmann6_curve(X, t=100)
def hartmann6_curve(x1, x2, x3, x4, x5, x6, t):
    return hartmann6(x1, x2, x3, x4, x5, x6) - np.log2(t / 100)


(
    hartmann6(0.1, 0.45, 0.8, 0.25, 0.552, 1.0),
    hartmann6_curve(0.1, 0.45, 0.8, 0.25, 0.552, 1.0, 100),
)

Out:

(-0.4878737485613134, -0.4878737485613134)

maximum_progressions = 100  # Observe hartmann6_curve over 100 progressions

for _ in range(30):  # Run 30 trials
    trials = client.get_next_trials(maximum_trials=1)
    for trial_index, parameters in trials.items():
        for t in range(1, maximum_progressions + 1):
            raw_data = {"hartmann6": hartmann6_curve(t=t, **parameters)}

            # On the final reading call complete_trial and break, else call attach_data
            if t == maximum_progressions:
                client.complete_trial(
                    trial_index=trial_index, raw_data=raw_data, progression=t
                )
                break

            client.attach_data(
                trial_index=trial_index, raw_data=raw_data, progression=t
            )

            # If the trial is underperforming, stop it
            if client.should_stop_trial_early(trial_index=trial_index):
                client.mark_trial_early_stopped(trial_index=trial_index)
                break

Out:

[INFO 03-14 05:09:31] ax.early_stopping.strategies.percentile: Early stoppinging trial 18: Trial objective value 2.3523120805459214 is worse than 50.0-th percentile (2.2222933918958923) across comparable trials..

Out:

[INFO 03-14 05:09:33] ax.early_stopping.strategies.percentile: Early stoppinging trial 19: Trial objective value 2.348788737582197 is worse than 50.0-th percentile (2.2855410647390446) across comparable trials..

Out:

[INFO 03-14 05:09:34] ax.early_stopping.strategies.percentile: Early stoppinging trial 20: Trial objective value 2.183045056882433 is worse than 50.0-th percentile (2.0847898681459576) across comparable trials..

Out:

[INFO 03-14 05:09:36] ax.early_stopping.strategies.percentile: Early stoppinging trial 21: Trial objective value 2.3445778143344795 is worse than 50.0-th percentile (2.332563197483424) across comparable trials..

Out:

[INFO 03-14 05:09:38] ax.early_stopping.strategies.percentile: Early stoppinging trial 22: Trial objective value 2.34746345690617 is worse than 50.0-th percentile (2.3445778143344795) across comparable trials..

Out:

[INFO 03-14 05:09:39] ax.early_stopping.strategies.percentile: Early stoppinging trial 23: Trial objective value 2.1656337631409786 is worse than 50.0-th percentile (2.125211815643468) across comparable trials..

Out:

[INFO 03-14 05:09:41] ax.early_stopping.strategies.percentile: Early stoppinging trial 24: Trial objective value 2.366761630772938 is worse than 50.0-th percentile (2.3445778143344795) across comparable trials..

Out:

[INFO 03-14 05:09:42] ax.early_stopping.strategies.percentile: Early stoppinging trial 25: Trial objective value 2.0343078723589256 is worse than 50.0-th percentile (1.9592589860620986) across comparable trials..

Out:

[INFO 03-14 05:09:44] ax.early_stopping.strategies.percentile: Early stoppinging trial 26: Trial objective value 2.19863836545895 is worse than 50.0-th percentile (2.1627362587918815) across comparable trials..

Out:

[INFO 03-14 05:09:46] ax.early_stopping.strategies.percentile: Early stoppinging trial 27: Trial objective value 2.3568249094213334 is worse than 50.0-th percentile (2.340359851771682) across comparable trials..

Out:

[INFO 03-14 05:09:47] ax.early_stopping.strategies.percentile: Early stoppinging trial 28: Trial objective value 2.1681458509119547 is worse than 50.0-th percentile (2.1656337631409786) across comparable trials..

Out:

[INFO 03-14 05:09:48] ax.early_stopping.strategies.percentile: Early stoppinging trial 29: Trial objective value 2.360109465443389 is worse than 50.0-th percentile (2.340359851771682) across comparable trials..

Step 6: Analyze Results

After running trials, you can analyze the results. Most commonly this means extracting the parameterization from the best performing trial you conducted.

best_parameters, prediction, index, name = client.get_best_parameterization()
print("Best Parameters:", best_parameters)
print("Prediction (mean, variance):", prediction)

Out:

Best Parameters: {'x1': 0.44143969602748834, 'x2': 0.2132962520317021, 'x3': 0.3326155761025044, 'x4': 0.2712898299798851, 'x5': 0.2751418364985259, 'x6': 0.6328874129133936}
Prediction (mean, variance): {'hartmann6': (-2.5748502808245615, 0.0032479046786129126)}

Step 7: Compute Analyses

Ax can also produce a number of analyses to help interpret the results of the experiment via client.compute_analyses. Users can manually select which analyses to run, or can allow Ax to select which would be most relevant. In this case Ax selects the following:

• Parrellel Coordinates Plot shows which parameterizations were evaluated and what metric values were observed -- this is useful for getting a high level overview of how thoroughly the search space was explored and which regions tend to produce which outcomes
• Interaction Analysis Plot shows which parameters have the largest affect on the function and plots the most important parameters as 1 or 2 dimensional surfaces
• Summary lists all trials generated along with their parameterizations, observations, and miscellaneous metadata

client.compute_analyses(display=True) # By default Ax will display the AnalysisCards produced by compute_analyses

Parallel Coordinates for hartmann6

View arm parameterizations with their respective metric values

loading...

hartmann6 by progression

Observe how the metric changes as each trial progresses

loading...

Interaction Analysis for hartmann6

Understand an Experiment's data as one- or two-dimensional additive components with sparsity. Important components are visualized through slice or contour plots

loading...

Summary for hartmann6_experiment

High-level summary of the Trial-s in this Experiment

	trial_index	arm_name	trial_status	generation_method	generation_node	hartmann6	x1	x2	x3	x4	x5	x6
0	0	0_0	COMPLETED	Sobol	Sobol	-0.106925	0.177323	0.7554	0.742666	0.440512	0.705899	0.80327
1	1	1_0	COMPLETED	Sobol	Sobol	-0.132358	0.969821	0.049505	0.141375	0.533487	0.279822	0.352548
2	2	2_0	COMPLETED	Sobol	Sobol	-0.028089	0.735593	0.500031	0.98977	0.100572	0.160783	0.057986
3	3	3_0	COMPLETED	Sobol	Sobol	-0.021849	0.411818	0.30395	0.403563	0.882461	0.852518	0.598567
4	4	4_0	COMPLETED	Sobol	Sobol	-0.000315	0.316204	0.739374	0.055304	0.786278	0.948584	0.885846
5	5	5_0	COMPLETED	BoTorch	MBM	-1.24719	0.65442	0.14301	0.435167	0.466066	0.350256	0.645672
6	6	6_0	COMPLETED	BoTorch	MBM	-1.3499	0.632291	0.063884	0.43263	0.421793	0.284825	0.789439
7	7	7_0	COMPLETED	BoTorch	MBM	-0.565801	0.643751	0.035596	0.154409	0.376237	0.417006	0.974714
8	8	8_0	COMPLETED	BoTorch	MBM	-1.43145	0.640835	0.026037	0.480317	0.229587	0.31064	0.805952
9	9	9_0	COMPLETED	BoTorch	MBM	-0.427434	0.585275	0.27551	0.497422	0	0.198016	1
10	10	10_0	COMPLETED	BoTorch	MBM	-0.572103	0.701303	0	0.642484	0	0.275841	0.725476
11	11	11_0	COMPLETED	BoTorch	MBM	-1.09963	0.662206	0	0.509523	0.327381	0.448667	0.656137
12	12	12_0	COMPLETED	BoTorch	MBM	-2.19419	0.494814	0	0.493684	0.325783	0.318969	0.726258
13	13	13_0	COMPLETED	BoTorch	MBM	-1.18138	0.457048	0	0.507847	0.454661	0.356154	0.890672
14	14	14_0	COMPLETED	BoTorch	MBM	-2.42008	0.459385	0.002023	0.468699	0.286651	0.270283	0.653674
15	15	15_0	COMPLETED	BoTorch	MBM	-1.41663	0.425329	0	0.904664	0.298363	0.250027	0.608034
16	16	16_0	COMPLETED	BoTorch	MBM	-2.60479	0.44144	0.213296	0.332616	0.27129	0.275142	0.632887
17	17	17_0	COMPLETED	BoTorch	MBM	-2.2387	0.429377	0	0.149733	0.276908	0.262049	0.633281
18	18	18_0	EARLY_STOPPED	BoTorch	MBM	2.35231	0.431879	0.687592	0.383725	0.279572	0.275458	0.600176
19	19	19_0	EARLY_STOPPED	BoTorch	MBM	2.34879	0.431663	0.687103	0.382516	0.278803	0.276868	0.600813
20	20	20_0	EARLY_STOPPED	BoTorch	MBM	2.18304	0.431801	0.67869	0.383534	0.279341	0.27536	0.599902
21	21	21_0	EARLY_STOPPED	BoTorch	MBM	2.34458	0.432071	0.685054	0.383612	0.2805	0.274564	0.599833
22	22	22_0	EARLY_STOPPED	BoTorch	MBM	2.34746	0.431716	0.686642	0.382629	0.279018	0.276622	0.600718
23	23	23_0	EARLY_STOPPED	BoTorch	MBM	2.16563	0.432042	0.673404	0.384113	0.280238	0.274141	0.599344
24	24	24_0	EARLY_STOPPED	BoTorch	MBM	2.36676	0.432329	0.691018	0.383611	0.281585	0.273867	0.599837
25	25	25_0	EARLY_STOPPED	BoTorch	MBM	2.03431	0.43152	0.672934	0.382822	0.27799	0.276972	0.600532
26	26	26_0	EARLY_STOPPED	BoTorch	MBM	2.19864	0.431871	0.682976	0.383442	0.279683	0.275262	0.599991
27	27	27_0	EARLY_STOPPED	BoTorch	MBM	2.35683	0.431864	0.689052	0.38281	0.279662	0.27604	0.600567
28	28	28_0	EARLY_STOPPED	BoTorch	MBM	2.16815	0.431823	0.674669	0.383397	0.279309	0.27558	0.600015
29	29	29_0	EARLY_STOPPED	BoTorch	MBM	2.36011	0.431921	0.689622	0.383413	0.280011	0.27507	0.60003

Cross Validation for hartmann6

Out-of-sample predictions using leave-one-out CV

loading...

Out:

[<ax.analysis.plotly.plotly_analysis.PlotlyAnalysisCard at 0x7fd1cc793110>,
<ax.analysis.plotly.plotly_analysis.PlotlyAnalysisCard at 0x7fd1c3f0cc20>,
<ax.analysis.plotly.plotly_analysis.PlotlyAnalysisCard at 0x7fd1c3707710>,
<ax.analysis.plotly.plotly_analysis.PlotlyAnalysisCard at 0x7fd1c37738f0>,
<ax.analysis.analysis.AnalysisCard at 0x7fd1c3ce1010>]

Conclusion

This tutorial demonstates Ax's early stopping capabilities, which utilize timeseries-like data to monitor the results of expensive evaluations and terminate those that are unlikely to produce promising results, freeing up resources to explore more configurations. This can be used in a number of applications, and is especially useful in machine learning contexts.