Archai: Microsoft's Modular Toolkit for Neural Architecture Search

The problem Archai targets: separating the search loop from the model code

Most teams that try neural architecture search end up writing three things at once: a parameterised model, an evaluation harness, and a loop that decides which candidate to try next. Archai's contribution is that it treats those as three separate objects with published interfaces. The README describes the library as accelerating NAS through fast, reproducible and modular research, and the quickstart bears that out: a search space object, a SearchObjectives container, and an algorithm object are constructed independently and then combined. The intended user is a researcher or an ML engineer who already has a PyTorch model family in mind and wants to explore variants of it under constraints, rather than someone looking for a pretrained model. The repository also ships two end-to-end tasks, text generation and face segmentation, which serve as worked examples of the full pipeline rather than as products.

How the discrete search API is wired together

The central abstraction in the quickstart is a search space, instantiated here as TransformerFlexSearchSpace("gpt2"). That object is passed both to the algorithm and to two of the evaluators, which means the search space is expected to know how to materialise a concrete architecture on demand. Objectives live in a SearchObjectives container. Each add_objective call takes a name, an evaluator instance, a direction flag, and optionally a constraint tuple and a compute_intensive flag. In the README example, non_embedding_params is marked higher_is_better=True with constraint=(1e6, 1e9), while onnx_latency and onnx_memory are both higher_is_better=False. The algorithm, EvolutionParetoSearch, is constructed with the space, the objectives, a third positional argument the README passes as None, an output directory string ("tmp"), num_iters=5, init_num_models=10 and seed=1234. Calling algo.search() then runs the loop and, per the README, produces a frontier of Pareto-optimal results. The seed argument is the concrete mechanism behind the reproducibility claim: the same seed and the same objective set should give the same search trajectory. The evaluators are also a data-flow decision point. TransformerFlexOnnxLatency and TransformerFlexOnnxMemory are named for ONNX, so the latency and memory numbers come from an exported ONNX graph rather than from a CUDA timing loop, which is what makes them cheap enough to mark compute_intensive=False.

Getting it installed and a search running

Installation is a single PyPI command, pip install archai, and the README recommends doing it inside a conda or pyenv virtual environment. The stated runtime requirements are Python 3.8 or newer and PyTorch 1.7.0 or newer. There is no separate extras step in the README for the NLP evaluators, so if the transformer_flex modules pull in ONNX tooling, that dependency resolution happens through the package itself or fails at import time. The quickstart is four code blocks that can be pasted in order: import TransformerFlexSearchSpace and build space from "gpt2"; build a SearchObjectives instance and add the three objectives; construct EvolutionParetoSearch with the arguments listed above; call algo.search(). The one argument worth pausing on is the third positional parameter, which the README fills with None and never names. In the snippet it sits between the objectives and the output directory, and since the example works with None, whatever it represents is optional for this configuration. Anyone adapting the snippet should keep that positional slot in mind rather than assuming the signature is (space, objectives, output_dir).

Where the documentation stops short

The quickstart is a single example, and it leaves several things unstated. It never says what the third positional argument to EvolutionParetoSearch is, nor what the algorithm does when an objective has no constraint versus when it has one like the (1e6, 1e9) tuple on non_embedding_params. It does not describe the return value of algo.search(), only that a Pareto frontier is produced. It does not give a runtime figure, a hardware requirement beyond PyTorch, or a sense of how long five iterations over ten initial models takes. The README points to the documentation site and a notebooks page for more, and to the tasks directory for the two end-to-end examples, so the honest position is that the quickstart is an entry point rather than a tutorial. Anyone evaluating Archai for a production decision will need to read the task code, not the README, to understand checkpointing, resumption and how results are persisted.

Release cadence is the main practical risk

The most recent tagged release is v1.0.0, dated 15 September 2023. Before that, v0.7.0 and v0.6.9 both landed on 18 January 2023. The repository is not archived and the last push timestamp is 31 August 2026, so commits continue, but the gap between the newest tag and the present is the number that matters if you depend on Archai through pip. Installing from PyPI gives you the tagged version, not the branch head. The README's stated floor of PyTorch 1.7.0 also tells you the library was written against an older torch API surface, and the NLP evaluators reach into transformer internals through a flex search space, which is exactly the kind of code that breaks when upstream model implementations change. None of this is disqualifying for research use, but it does mean the upgrade cost is not zero: you are adopting a snapshot, and moving to a newer commit means testing the search space and evaluator imports yourself.

Archai against single-shot NAS libraries

The clearest contrast is with supernet and one-shot approaches such as the DARTS-style methods listed among the repository topics. A one-shot method trains a single over-parameterised network and then extracts a sub-architecture by weight sharing, so the search cost is dominated by one training run and the search space is usually fixed by the method's authors. Archai's discrete search path does the opposite: EvolutionParetoSearch evaluates distinct candidate architectures across iterations, with explicit objectives and a population size set by init_num_models. That makes the cost scale with the number of candidates and the number of iterations, and it makes multi-objective trade-offs, like parameters against ONNX latency against ONNX memory, first-class rather than something bolted on. The trade-off is that you supply the search space and the evaluators. Archai gives you the loop, the objective container and the Pareto bookkeeping; it does not give you a pretrained supernet to fine-tune. If your goal is a single accuracy number on a standard benchmark, a one-shot method will get you there with less glue code. If your goal is a frontier over deployment-relevant metrics for a model family you control, the discrete API is the part of Archai worth the setup.

Licence, contributions and what that means for forks

Archai is released under the MIT License, and the README points to the LICENSE file in the repository for the full text. MIT is permissive, so incorporating the library into a larger system, including a commercial one, is consistent with the licence terms as written. Two things sit alongside the licence and are worth reading before you fork. First, contributions require a Contributor License Agreement, handled by a CLA bot on pull requests, so upstreaming a fix is not a plain pull request. Second, the README carries a trademark section stating that use of Microsoft trademarks and logos in modified versions must not imply Microsoft sponsorship. If you fork Archai and ship it under your own name, the code is yours to use but the branding is not. This is a description of what the repository states, not legal advice; read the LICENSE file and the trademark guidelines yourself if the distinction matters to your organisation.

Who should pick this up, and what to check first

Archai fits a specific shape of work: you have a PyTorch model family, you care about more than one metric, and you want the search loop and Pareto frontier handling written for you. The quickstart path through TransformerFlexSearchSpace, SearchObjectives and EvolutionParetoSearch is short enough to evaluate in an afternoon. It fits poorly if you need a maintained dependency with a recent release cadence, or if your architecture decisions are better expressed as hyperparameters over a fixed model, in which case a hyperparameter optimiser is the smaller tool. The first thing to verify is import health: run the four quickstart blocks against your installed torch and confirm that archai.discrete_search.search_spaces.nlp.transformer_flex.search_space and the two evaluator modules under archai.discrete_search.evaluators.nlp resolve. The second is the third positional argument to EvolutionParetoSearch, which the README passes as None and never documents; read the constructor signature before you adapt the snippet. The third is the output written to the directory passed as "tmp", since the README says a frontier is produced but not in what format.

Editorial conclusion

Adopt Archai if you already have PyTorch training and evaluation code and want to plug a search space and objective set into an evolutionary loop without writing the loop yourself, and if you are comfortable pinning to the v1.0.0 tag from September 2023. Do not adopt it if you need an actively released dependency, because the README states Python 3.8+ and PyTorch 1.7.0+ and the newest tagged release predates the current PyTorch line by several versions. Before committing, verify that the archai.discrete_search.search_spaces.nlp.transformer_flex modules still import against your installed torch, since the quickstart imports TransformerFlexSearchSpace, TransformerFlexOnnxLatency and TransformerFlexOnnxMemory directly from that namespace.

Official sources

1. License: MIT
2. microsoft/archai on GitHub
3. Project website
4. README
5. Releases