Agentic Discovery of Neural Architectures: AIRA-Compose and AIRA-Design

Key points

• Problem: Most LLMs today are Transformers that interleave attention and MLP layers 1:1, and moving to hybrid designs that mix in other primitives such as Mamba makes the design space vast and combinatorial. The authors argue that relying solely on manual exploration may overlook strong non-obvious arrangements or entirely new primitives. They therefore hand this search to research agents and frame it as a step toward recursive self-improvement, where agents discover and optimize the architectures that power them.
• Input and output: This is a framework paper that builds 12 tasks in the AIRS-Bench format, a task standard for evaluating an agent's ability to conduct independent research, and runs experiments on them. Given a task description and data, the agent submits a 16-primitive architecture string (submission.csv) in AIRA-Compose, or a model.py for LRA or a train.py for Autoresearch in AIRA-Design. An isolated evaluation script trains the submission from scratch and scores it on a held-out test split or a pinned validation shard.
• Method: An agent is an LLM plus a scaffold (search policy), run inside AIRA-dojo, a harness that evolves code solutions through tree-based exploration, with four operators: Draft, Debug, Improve, and Analyze. The greedy scaffold drafts 5 initial solutions and keeps applying Improve to the one with the best validation score. AIRA-Compose turns the search and evaluation steps of Composer, an existing framework that searches hybrid architectures with Bayesian optimization, into agent tasks, then uses Composer's aggregation and extrapolation steps to scale the found 16-layer patterns to 350M, 1B, and 3B parameters by stretching (proportionally expanding contiguous blocks) or stacking (repeating the whole pattern).
• Implementation: Each agent gets one H200 GPU, and a run lasts at most 24 hours or 500 steps, with 60 hours per run for the larger BabiStories and DCLM searches. AIRA-Compose uses 6 LLMs, while in AIRA-Design the LRA experiments build 20 agents from 12 LLMs including Opus 4.6 and Gemini 3 Pro and the Autoresearch experiments build greedy agents from 7 of those LLMs. Scaled-up training ran 350M models on 8 H200 GPUs and 1B and 3B models on 16 H200 GPUs.
• Results: Among architectures built from the two primitives MLP and attention and trained at the 1B scale on 37.5B tokens, AIRAformer-D (Stretched) reached a validation loss of 2.734 and 59.7% average 0-shot accuracy, ahead of Llama 3.2 and the Composer-found models, and in the 3-primitive comparison that adds Mamba, AIRAhybrid-D (Stretched) reached a validation loss of 2.719 from a single seed. On LRA, the best accuracy of agent-designed models was 2.3pp (percentage points) below the human-designed SOTA (state of the art) of 82% on document matching and 2.6pp below the SOTA of 91% on text classification. On Autoresearch, Greedy Opus 4.5 with literature access reached a validation BPB (bits per byte, lower is better) of 0.968 in its best seed, below the original Autoresearch best of 0.9775.
• Limits: The authors state that AIRA-Compose relies on small-scale proxy evaluation that does not always predict large-scale performance, and that the aggregation and extrapolation steps remain non-agentic fixed procedures. On LRA, the agent-designed architectures largely recombine prior ideas such as Performer, Longformer, and Conformer. The Autoresearch runs used H200 GPUs in a sandbox without internet access, unlike the original H100 setup, and the authors note this may explain why the same train.py started at a BPB of 1.012 instead of the original 0.998.

Paper links

External research summaries. These are not HDATF publications or measured product results.