In this tutorial, we explore adaptive experimentation using Meta’s Ax with the modern Client API. We work through a complete workflow where we tune a RandomForest model on a synthetic classification dataset while balancing predictive accuracy against model footprint. We begin by defining a mixed search space with integer, float, log-scaled, and categorical parameters, then use Ax’s ask-tell optimization loop to run constrained Bayesian optimization, multi-objective optimization, and parameter-constrained experimentation. Along the way, we visualize convergence, inspect the Pareto frontier, use Ax’s built-in analysis tools, and persist the experiment for future reuse.
import importlib, subprocess, sys
def _ensure(module, pip_name=None):
try:
importlib.import_module(module)
except ImportError:
print(f"Installing {pip_name or module} ...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", pip_name or module])
_ensure("ax", "ax-platform")
_ensure("sklearn", "scikit-learn")
import logging, warnings, time
import numpy as np
import matplotlib.pyplot as plt
warnings.filterwarnings("ignore")
logging.getLogger("ax").setLevel(logging.WARNING)
from ax.api.client import Client
from ax.api.configs import RangeParameterConfig, ChoiceParameterConfig
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import StratifiedKFold, cross_val_score
np.random.seed(0)
We begin by preparing the Colab environment and installing the required packages for Ax and scikit-learn. We import the core libraries for optimization, machine learning, plotting, logging, and reproducibility. We also configure warnings and Ax logging to keep the notebook output clean and focused on the experimental results.
X, y = make_classification(
n_samples=1400, n_features=20, n_informative=8, n_redundant=4,
n_classes=3, random_state=0,
)
CV = StratifiedKFold(n_splits=3, shuffle=True, random_state=0)
def evaluate(p):
n_est, depth = int(p["n_estimators"]), int(p["max_depth"])
clf = RandomForestClassifier(
n_estimators=n_est,
max_depth=depth,
max_features=float(p["max_features"]),
min_samples_leaf=int(p["min_samples_leaf"]),
criterion=p["criterion"],
ccp_alpha=float(p["ccp_alpha"]),
n_jobs=-1,
random_state=0,
)
accuracy = cross_val_score(clf, X, y, cv=CV, scoring="accuracy").mean()
model_size = n_est * depth
return {"accuracy": float(accuracy), "model_size": float(model_size)}
SEARCH_SPACE = [
RangeParameterConfig(name="n_estimators", bounds=(50, 300), parameter_type="int"),
RangeParameterConfig(name="max_depth", bounds=(3, 24), parameter_type="int"),
RangeParameterConfig(name="max_features", bounds=(0.2, 1.0), parameter_type="float"),
RangeParameterConfig(name="min_samples_leaf",bounds=(1, 12), parameter_type="int"),
RangeParameterConfig(name="ccp_alpha", bounds=(1e-5, 1e-1), parameter_type="float", scaling="log"),
ChoiceParameterConfig(name="criterion", values=["gini", "entropy", "log_loss"],
parameter_type="str", is_ordered=False),
]
def run_study(client, total_trials, metric_keys, batch=4):
records = []
while len(records) < total_trials:
trials = client.get_next_trials(max_trials=min(batch, total_trials - len(records)))
if not trials:
break
for idx, params in trials.items():
full = evaluate(params)
raw = {k: full[k] for k in metric_keys}
client.complete_trial(trial_index=idx, raw_data=raw)
records.append({"trial": idx, "params": params, **full})
return records
We create a synthetic multi-class classification dataset and define a cross-validation strategy to evaluate Random Forest models. We build an evaluation function that returns both accuracy and model size, allowing us to measure performance and cost together. We then define a mixed search space with integer, float, log-scaled, and categorical parameters, along with a reusable ask-tell study runner.
print("\n=== Study 1: constrained single-objective Bayesian optimization ===")
c1 = Client()
c1.configure_experiment(parameters=SEARCH_SPACE, name="rf_constrained")
c1.configure_optimization(objective="accuracy",
outcome_constraints=["model_size <= 2500"])
rec1 = run_study(c1, total_trials=24, metric_keys=["accuracy", "model_size"])
best_params, prediction, best_idx, best_arm = c1.get_best_parameterization()
print("\nBest feasible configuration found:")
for k, v in best_params.items():
print(f" {k:>16}: {v}")
print(" predicted:", prediction)
feasible = [(r["trial"], r["accuracy"]) for r in rec1 if r["model_size"] <= 2500]
best_so_far, cur = [], -np.inf
for _, acc in feasible:
cur = max(cur, acc); best_so_far.append(cur)
plt.figure(figsize=(7, 4))
plt.plot(range(1, len(best_so_far) + 1), best_so_far, "o-")
plt.xlabel("feasible trial #"); plt.ylabel("best accuracy so far")
plt.title("Study 1 — convergence (subject to model_size <= 2500)")
plt.grid(alpha=0.3); plt.tight_layout(); plt.show()
We run a constrained single-objective Bayesian optimization study where we maximize accuracy while keeping model size below a fixed threshold. We use Ax to suggest hyperparameter configurations, evaluate them, and report both accuracy and model size back to the optimizer. We then extract the best feasible configuration and plot the best accuracy achieved over feasible trials.
print("\n=== Study 2: multi-objective (accuracy vs. model_size) ===")
c2 = Client()
c2.configure_experiment(parameters=SEARCH_SPACE, name="rf_multiobjective")
c2.configure_optimization(objective="accuracy, -model_size")
rec2 = run_study(c2, total_trials=28, metric_keys=["accuracy", "model_size"])
try:
frontier = c2.get_pareto_frontier()
print(f"Ax identified {len(frontier)} Pareto-optimal configurations.")
except Exception as e:
frontier = None
print("get_pareto_frontier unavailable in this version:", e)
acc = np.array([r["accuracy"] for r in rec2])
size = np.array([r["model_size"] for r in rec2])
order = np.argsort(size)
pareto_idx, best_acc = [], -np.inf
for i in order:
if acc[i] > best_acc:
best_acc = acc[i]; pareto_idx.append(i)
plt.figure(figsize=(7, 5))
plt.scatter(size, acc, c="lightgray", label="all trials")
plt.scatter(size[pareto_idx], acc[pareto_idx], c="crimson", zorder=3, label="Pareto front")
plt.plot(size[pareto_idx], acc[pareto_idx], "--", c="crimson", alpha=0.6)
plt.xlabel("model_size (lower = cheaper)"); plt.ylabel("accuracy (higher = better)")
plt.title("Study 2 — accuracy vs. model size trade-off")
plt.legend(); plt.grid(alpha=0.3); plt.tight_layout(); plt.show()
We move from single-objective optimization to multi-objective optimization by jointly maximizing accuracy and minimizing model size. We use Ax to search for configurations that represent strong trade-offs between predictive performance and computational footprint. We then calculate and visualize the empirical Pareto frontier to understand how accuracy varies with model size.
print("\n=== Study 3: parameter constraints on a synthetic surface ===")
c3 = Client()
c3.configure_experiment(
parameters=[
RangeParameterConfig(name="x1", bounds=(0.0, 1.0), parameter_type="float"),
RangeParameterConfig(name="x2", bounds=(0.0, 1.0), parameter_type="float"),
],
parameter_constraints=["x1 + x2 <= 1.5"],
name="constrained_surface",
)
c3.configure_optimization(objective="-dist")
for _ in range(14):
for idx, p in c3.get_next_trials(max_trials=1).items():
dist = (p["x1"] - 0.9) ** 2 + (p["x2"] - 0.9) ** 2
c3.complete_trial(trial_index=idx, raw_data={"dist": float(dist)})
bp, _, _, _ = c3.get_best_parameterization()
print(f"Best point: x1={bp['x1']:.3f}, x2={bp['x2']:.3f}, "
f"sum={bp['x1'] + bp['x2']:.3f} (constraint: <= 1.5)")
print("Unconstrained optimum would be (0.9, 0.9); Ax respects the boundary.")
We demonstrate parameter constraints using a simple two-dimensional synthetic optimization problem. We ask Ax to minimize the distance to a target point while enforcing the input constraint that the sum of the two variables remains below a boundary. We observe that the optimizer respects the constraint and finds the best feasible point near the constrained optimum.
print("\n=== Ax built-in analyses for Study 1 ===")
try:
import plotly.io as pio
if "google.colab" in sys.modules:
pio.renderers.default = "colab"
cards = c1.compute_analyses(display=True)
print(f"Computed {len(cards)} analysis cards.")
except Exception as e:
print("Interactive analyses didn't render in this environment:", e)
print("(The matplotlib plots above already capture the key results.)")
print("\n=== Saving / loading the experiment ===")
try:
c1.save_to_json_file("ax_study1.json")
reloaded = Client.load_from_json_file("ax_study1.json")
print("Saved to ax_study1.json and reloaded successfully.")
rp, _, _, _ = reloaded.get_best_parameterization()
print("Best params from reloaded client match:", rp == best_params)
except Exception as e:
print("JSON persistence API differs in this version:", e)
print("See: https://ax.dev/docs/recipes/experiment-to-json")
print("\nDone. You optimized a mixed-type search space with constraints, "
"traced a Pareto frontier, and persisted in the experiment.")
We use Ax’s built-in analysis tools to generate diagnostic cards, such as sensitivity, cross-validation, and other experiment insights, when the environment supports them. We then save the completed experiment to a JSON file and reload it to verify that the optimization state is preserved. We finish by confirming that the tutorial covers constrained optimization, multi-objective trade-offs, analysis, and experiment persistence.
In conclusion, we developed a practical understanding of how Ax helps us run efficient and structured hyperparameter optimization experiments. We optimized a mixed-type search space, enforced both outcome and parameter constraints, compared accuracy against model size through multi-objective optimization, and identified trade-offs using an empirical Pareto frontier. We also used Ax’s analysis and persistence features to make the experimentation workflow more interpretable and reproducible.
Check out the Full Codes here. Also, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us
Sana Hassan, a consulting intern at Marktechpost and dual-degree student at IIT Madras, is passionate about applying technology and AI to address real-world challenges. With a keen interest in solving practical problems, he brings a fresh perspective to the intersection of AI and real-life solutions.
Credit: Source link
