Ax 1.0: Open Platform for Efficient Adaptive Experimentation
Meta has released Ax 1.0, an open-source platform that uses machine learning to automate complex, resource-intensive experimentation. The platform is designed for situations where evaluating a single configuration is costly—covering domains from AI model development to physical hardware design. Ax employs Bayesian optimization for adaptive experimentation, actively proposing new configurations based on insights from previous evaluations rather than relying on static search strategies.
The platform has been applied across Meta's engineering disciplines for tasks including hyperparameter optimization, architecture search, production infrastructure tuning, and optimizing design parameters for AR/VR devices. Recent work has also used Ax to address GenAI challenges, such as discovering optimal data mixtures for training models. A companion paper, "Ax: A Platform for Adaptive Experimentation," details the platform's architecture and methodology.
Installation is straightforward: pip install ax-platform. Documentation, tutorials, and deep dives on the underlying methods are available on the Ax website.
Built for Real-World Constraints
Running adaptive experiments presents challenges beyond just applying sophisticated ML methods. Ax provides specialized infrastructure for managing experiment state, automating orchestration, and offering analysis and diagnostics. The platform also handles the nuance of real-world experiment goals, which often involve balancing multiple objective metrics subject to constraints and guardrails, rather than simply improving a single number.
Optimization and Understanding
Ax serves a dual purpose: finding efficient configurations and explaining the system being optimized. The platform provides analysis tools such as plots and tables to track optimization progress over time, visualize trade-offs between metrics via a Pareto frontier, and display the effects of specific parameters across the input space. Sensitivity analysis quantifies how much each input parameter contributes to results, helping users walk away with deep system knowledge along with an optimal configuration.
By default, Ax relies on Bayesian optimization via BoTorch. This method balances exploration—learning how new configurations perform—with exploitation—refining configurations known to be good. The iterative process works without assumptions about the form of the underlying function being optimized.
The Bayesian Optimization Loop
Under typical settings, Ax uses a Gaussian process (GP) as the surrogate model. GPs can make predictions while quantifying uncertainty, making them effective even with few data points. Ax then applies an acquisition function from the expected improvement (EI) family to determine which configurations to evaluate next, capturing the expected value of a new point relative to the best one observed so far.
This loop scales from one-dimensional examples to problems with hundreds of tunable parameters and outcomes. In high-dimensional settings, where covering the search space exponentially increases cost, the surrogate-based approach offers significant advantages over other methods. A visual overview of the process is available on the Ax website's Introduction to Bayesian Optimization page.
Production Usage at Scale
Thousands of developers at Meta use Ax for tuning online recommender and ranking models, running simulation optimizations for hardware design, and improving infrastructure. The platform supports nuanced objectives, such as multi-objective optimization to improve model accuracy while minimizing resource usage, and constrained optimization for recommender systems that tunes key metrics without regressing others.
Notable applications include optimizing natural language models to balance size and performance for Ray-Ban Stories devices, and designing faster-curing, low-carbon concrete mixes deployed at Meta data center construction sites—work that supports Meta's goal of net zero emissions by 2030.
Tackling Complexity with Data-Driven Design
Ax addresses problems across domains where the ultimate quality of a system depends on parameters with interactions too complex to reason about without experimentation, and where running those experiments carries real cost. By adapting the experimental approach as data unfolds, the platform efficiently solves these challenges.
Open Source and Ongoing Development
Ax is open source under the MIT license, and development continues on features for experiment design, new optimization methods, and integrations with external platforms. Contributions are invited from both practitioners and researchers, including new surrogate models, acquisition functions, research extensions, or core bug fixes. The team can be reached via GitHub Issues.
