Ax for AutoML: From Research to Production

AutoML requires solving resource-intensive black-box optimization (BBO) problems which appear at all stages of the machine learning pipeline. For example: in the early stages of the pipeline feature selection, architecture search, and hyperparameter optimization (HPO) often require time- and resource-intensive training workflows, and once a candidate ML model is ready for deployment, practitioners often invest additional resources in optimizing it for inference and serving efficiency.

Enabling AutoML in real-world settings depends on solving these pervasive BBO problems in a sample efficient manner. This is challenging and often necessitates research to identify novel surrogate models, acquisition functions, pruning strategies, etc. Furthermore, validating the results via experimentation is non-trivial, requiring complex state management, data preprocessing, integrations with external software systems, storage solutions, and more. As ML models grow ever more computationally costly, and the issue of scale increases, these considerations become even more critical.

The demands for a flexible research system and a simple, robust experimentation platform can often be in tension with each other. At Meta, we have built Ax as a solution to this challenge. Ax is designed to be capable in both settings, with features necessary for productionizing state-of-the-art research at industry-scale.

Ax provides an expressive API which handles complex search spaces, multiple objectives, constraints on both parameters and outcomes, and noisy observations, with or without observed noise levels. It supports suggesting multiple designs to evaluate in parallel (both synchronously and asynchronously) and can stop evaluations early to save resources. The platform allows researchers deep control over optimization algorithms, empowering both researchers interested in testing their custom algorithms on real-world data without worrying about designing bespoke solutions for orchestration, data preprocessing, storage, data fetching and deployment, error handling, and other essential aspects of experiment management. It also provides sensible defaults, facilitating access to advanced techniques that are typically reserved for experts.

This tutorial will provide an overview of Ax and its applications at various stages of the AutoML pipeline, discuss the relationship between Ax and Botorch as well as Ax’s motivating features, and walk though a number of example applications of Ax with increasing complexity.

Part 1: Introduction and Discussion

• Discussion of BBO as part of the full AutoML pipeline, beyond hyperparameter optimization
• Introduction to Ax
• Overview of the Ax API
• Dissection the relationship between Ax and BoTorch

Part 2: Interactive Demonstration

• Using Ax for an AutoML task
• Configuring Ax to solve a typical HPO problem
• Optimising for complex goals (ex. Constraints on outcomes, multiple objectives, etc.)
• Using a custom method for candidate generation
• Wrap a surrogate and/or acquisition function implemented in BoTorch in Ax’s GenerationNode using Modular BoTorch
• Using ax.analysis to examine differences in performance
• Using early stopping to save resources and examine its performance
• Q&A

We invite the audience to follow along with the live notebook demonstration and to ask questions both during the walkthrough and during the subsequent Q&A.

Miles Olson is a software engineer at Meta in the Adaptive Experimentation group, working on Ax, Meta’s open source platform for understanding, managing, deploying, and automating adaptive experiments. He is primarily focused on applying state-of-the-art methods and new modes of experimentation at scale, enabling practitioners to utilize cutting-edge Bayesian optimization techniques in their problems, and accelerating the research to production pipeline.

Mia Garrard is a Research Engineer at Meta in the Adaptive Experimentation group, which develops two open source platforms: Ax and BoTorch. Her research interests include Bayesian optimization, causal inference, and ML-driven experimentation. She is passionate about leveraging basic research in applied systems for real-world impact across domains. Mia has co-authored works such as “Looper: An end-to-end ML platform for product decisions” (KDD), “Interpretable personalized experimentation” (KDD), and “Rapid intraoperative histology of unprocessed surgical specimens via fibre-laser-based stimulated Raman scattering microscopy” (Nature Biomedical Engineering). She believes in the democratization of new methods and contributes through the development of open source software, and is excited to collaborate with the broader community.