Skip to main content
Version: Next

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 tt which allows the function to produce timeseries-like data where the value returned is closer and closer to Hartmann6's true value as tt increases. At t=100t = 100 the function will simply return Hartmann6's unaltered value.

f(x,t)=hartmann6(x)log2(t/100) 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

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.api.client import Client
from ax.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 xi(0,1)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),
)
Output:
(np.float64(-0.4878737485613134), np.float64(-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
Output:
[INFO 04-18 05:09:09] ax.early_stopping.strategies.percentile: Early stoppinging trial 13: Trial objective value 3.2740222171309075 is worse than 50.0-th percentile (3.0713281027689607) across comparable trials..
[INFO 04-18 05:09:11] ax.early_stopping.strategies.percentile: Early stoppinging trial 14: Trial objective value 3.2729364184272933 is worse than 50.0-th percentile (3.075379092733468) across comparable trials..
[INFO 04-18 05:09:13] ax.early_stopping.strategies.percentile: Early stoppinging trial 15: Trial objective value 3.273689301279302 is worse than 50.0-th percentile (3.1560662454248485) across comparable trials..
[INFO 04-18 05:09:15] ax.early_stopping.strategies.percentile: Early stoppinging trial 16: Trial objective value 3.273991560512199 is worse than 50.0-th percentile (3.236753398116229) across comparable trials..
[INFO 04-18 05:09:17] ax.early_stopping.strategies.percentile: Early stoppinging trial 17: Trial objective value 3.2741568725993986 is worse than 50.0-th percentile (3.2441586256113206) across comparable trials..
[INFO 04-18 05:09:19] ax.early_stopping.strategies.percentile: Early stoppinging trial 18: Trial objective value 3.273987129788082 is worse than 50.0-th percentile (3.2515638531064117) across comparable trials..
[INFO 04-18 05:09:21] ax.early_stopping.strategies.percentile: Early stoppinging trial 19: Trial objective value 3.272320604302441 is worse than 50.0-th percentile (3.2619422287044264) across comparable trials..
[INFO 04-18 05:09:22] ax.early_stopping.strategies.percentile: Early stoppinging trial 20: Trial objective value 3.275720660640375 is worse than 50.0-th percentile (3.272320604302441) across comparable trials..
[INFO 04-18 05:09:24] ax.early_stopping.strategies.percentile: Early stoppinging trial 21: Trial objective value 3.2733981138934745 is worse than 50.0-th percentile (3.272628511364867) across comparable trials..
[INFO 04-18 05:09:26] ax.early_stopping.strategies.percentile: Early stoppinging trial 22: Trial objective value 3.27451683932821 is worse than 50.0-th percentile (3.2729364184272933) across comparable trials..
[INFO 04-18 05:09:28] ax.early_stopping.strategies.percentile: Early stoppinging trial 23: Trial objective value 3.273664241514862 is worse than 50.0-th percentile (3.273167266160384) across comparable trials..
[INFO 04-18 05:09:30] ax.early_stopping.strategies.percentile: Early stoppinging trial 24: Trial objective value 3.2738250297449016 is worse than 50.0-th percentile (3.2733981138934745) across comparable trials..
[INFO 04-18 05:09:32] ax.early_stopping.strategies.percentile: Early stoppinging trial 25: Trial objective value 3.135663191531393 is worse than 50.0-th percentile (2.933824579019026) across comparable trials..
[INFO 04-18 05:09:34] ax.early_stopping.strategies.percentile: Early stoppinging trial 26: Trial objective value 3.135781554432402 is worse than 50.0-th percentile (2.937875568983533) across comparable trials..
[INFO 04-18 05:09:36] ax.early_stopping.strategies.percentile: Early stoppinging trial 27: Trial objective value 3.2750822456959616 is worse than 50.0-th percentile (3.2733415960379055) across comparable trials..
[INFO 04-18 05:09:38] ax.early_stopping.strategies.percentile: Early stoppinging trial 28: Trial objective value 3.2735210012046814 is worse than 50.0-th percentile (3.2733981138934745) across comparable trials..
[INFO 04-18 05:09:40] ax.early_stopping.strategies.percentile: Early stoppinging trial 29: Trial objective value 3.1357210259869013 is worse than 50.0-th percentile (3.018562721674914) 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)
Output:
Best Parameters: {'x1': 0.5618872077822922, 'x2': 0.36631463069954034, 'x3': 0.20237361440475782, 'x4': 0.32275977695315977, 'x5': 0.24499881345218716, 'x6': 0.5484114810071069}
Prediction (mean, variance): {'hartmann6': (np.float64(-1.477791941765223), np.float64(0.0005824669947379483))}

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
# display=True instructs Ax to sort then render the resulting analyses
cards = client.compute_analyses(display=True)

Parallel Coordinates for hartmann6

The parallel coordinates plot displays multi-dimensional data by representing each parameter as a parallel axis. This plot helps in assessing how thoroughly the search space has been explored and in identifying patterns or clusterings associated with high-performing (good) or low-performing (bad) arms. By tracing lines across the axes, one can observe correlations and interactions between parameters, gaining insights into the relationships that contribute to the success or failure of different configurations within the experiment.

loading...

Summary for hartmann6_experiment

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

trial_indexarm_nametrial_statusgeneration_nodehartmann6x1x2x3x4x5x6
000_0COMPLETEDSobol-0.2465490.0796450.2983530.4771480.1606590.6714250.189673
111_0COMPLETEDSobol-0.0426280.9688050.9472140.7747270.7273930.3389990.960393
222_0COMPLETEDSobol-0.6133570.7309840.1450920.0626030.4685330.1611020.727277
333_0COMPLETEDSobol-0.0467820.3408480.6063350.6736580.8941550.8284890.498001
444_0COMPLETEDSobol-0.0050530.4049070.0300090.880990.8358610.9254420.029374
555_0COMPLETEDMBM-0.070364000001
666_0COMPLETEDMBM-0.0851750.951087000.4497180.2050040.260562
777_0COMPLETEDMBM-0.2546510.7981590.06554200.518840.1183280.806149
888_0COMPLETEDMBM-0.7986480.7045090.1765480.0871740.4489310.177970.69652
999_0COMPLETEDMBM-1.118640.6604980.228850.1275620.4161870.2056220.645939
101010_0COMPLETEDMBM-1.333290.6223240.2900340.1654070.3860440.2245030.609328
111111_0COMPLETEDMBM-1.48250.5618870.3663150.2023740.322760.2449990.548411
121212_0COMPLETEDMBM-0.6256830.5546540.5001730.413510.4831390.1724530
131313_0EARLY_STOPPEDMBM3.274020.8241360.563250.223160.6688720.2129841
141414_0EARLY_STOPPEDMBM3.272940.8229430.5636310.223170.6663530.2133071
151515_0EARLY_STOPPEDMBM3.273690.8238470.5646420.2231450.6676230.2132751
161616_0EARLY_STOPPEDMBM3.273990.8240960.5626160.223170.6689870.2128091
171717_0EARLY_STOPPEDMBM3.274160.8244410.5645890.2231390.6687190.2133431
181818_0EARLY_STOPPEDMBM3.273990.8242510.5627270.2231680.6688970.2129081
191919_0EARLY_STOPPEDMBM3.272320.822320.5631090.2231820.6652040.213451
202020_0EARLY_STOPPEDMBM3.275720.8260990.5645780.2231210.672250.2128271
212121_0EARLY_STOPPEDMBM3.27340.8235050.5635230.2231660.6673810.2131941
222222_0EARLY_STOPPEDMBM3.274520.8246760.5643270.2231380.6696660.2131571
232323_0EARLY_STOPPEDMBM3.273660.8237630.563930.2231530.6678720.2132611
242424_0EARLY_STOPPEDMBM3.273820.8238530.5651760.2231360.667870.2134781
252525_0EARLY_STOPPEDMBM3.135660.8231850.5631990.2231710.6669970.2131591
262626_0EARLY_STOPPEDMBM3.135780.8233240.5642520.2231560.666970.2134821
272727_0EARLY_STOPPEDMBM3.275080.8253780.5653280.2231210.6705890.2132191
282828_0EARLY_STOPPEDMBM3.273520.8237180.5635780.2231620.6676180.2132291
292929_0EARLY_STOPPEDMBM3.135720.8231090.5640960.2231580.6669080.2133451

hartmann6 by progression

The progression plot tracks the evolution of each metric over the course of the experiment. This visualization is typically used to monitor the improvement of metrics over Trial iterations, but can also be useful in informing decisions about early stopping for Trials.

loading...

Sensitivity Analysis for hartmann6

Understand how each parameter affects hartmann6 according to a second-order sensitivity analysis.

loading...

x3 vs. hartmann6

The slice plot provides a one-dimensional view of predicted outcomes for hartmann6 as a function of a single parameter, while keeping all other parameters fixed at their status_quo value (or mean value if status_quo is unavailable). This visualization helps in understanding the sensitivity and impact of changes in the selected parameter on the predicted metric outcomes.

loading...

x3, x4 vs. hartmann6

The contour plot visualizes the predicted outcomes for hartmann6 across a two-dimensional parameter space, with other parameters held fixed at their status_quo value (or mean value if status_quo is unavailable). This plot helps in identifying regions of optimal performance and understanding how changes in the selected parameters influence the predicted outcomes. Contour lines represent levels of constant predicted values, providing insights into the gradient and potential optima within the parameter space.

loading...

x3, x6 vs. hartmann6

The contour plot visualizes the predicted outcomes for hartmann6 across a two-dimensional parameter space, with other parameters held fixed at their status_quo value (or mean value if status_quo is unavailable). This plot helps in identifying regions of optimal performance and understanding how changes in the selected parameters influence the predicted outcomes. Contour lines represent levels of constant predicted values, providing insights into the gradient and potential optima within the parameter space.

loading...

x2, x3 vs. hartmann6

The contour plot visualizes the predicted outcomes for hartmann6 across a two-dimensional parameter space, with other parameters held fixed at their status_quo value (or mean value if status_quo is unavailable). This plot helps in identifying regions of optimal performance and understanding how changes in the selected parameters influence the predicted outcomes. Contour lines represent levels of constant predicted values, providing insights into the gradient and potential optima within the parameter space.

loading...

x3, x5 vs. hartmann6

The contour plot visualizes the predicted outcomes for hartmann6 across a two-dimensional parameter space, with other parameters held fixed at their status_quo value (or mean value if status_quo is unavailable). This plot helps in identifying regions of optimal performance and understanding how changes in the selected parameters influence the predicted outcomes. Contour lines represent levels of constant predicted values, providing insights into the gradient and potential optima within the parameter space.

loading...

Cross Validation for hartmann6

The cross-validation plot displays the model fit for each metric in the experiment. It employs a leave-one-out approach, where the model is trained on all data except one sample, which is used for validation. The plot shows the predicted outcome for the validation set on the y-axis against its actual value on the x-axis. Points that align closely with the dotted diagonal line indicate a strong model fit, signifying accurate predictions. Additionally, the plot includes 95% confidence intervals that provide insight into the noise in observations and the uncertainty in model predictions. A horizontal, flat line of predictions indicates that the model has not picked up on sufficient signal in the data, and instead is just predicting the mean.

loading...

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.