AIRA-Compose: Agentic Discovery of Neural Architectures
Abstract
We introduce AIRA-Compose — a framework that enables AI research agents to autonomously explore and discover novel neural architectures. AIRA-Compose leverages the previously-proposed Composer abstraction to define a combinatorial design space over fundamental model primitives (Attention, MLP, and Mamba), then uses a guided search policy to efficiently navigate this space for novel neural architecture design. Our approach operates in two stages: (1) agents iteratively designs and evaluates candidate architectures at the million-parameter scale, then (2) top-performing candidates are scaled to billion parameters for pre-training. Overall, AIRA-Compose discovers 14 novel architectures spanning two families — AIRAformers (Transformer-based) and AIRAhybrids (Transformer-Mamba Hybrid-based). When pre-trained at 1B scale under a fixed token budget, agent-discovered top-performing architectures consistently outperform both Llama 3.2 and Composer-found alternatives. On downstream tasks, AIRAformer-D and AIRAhybrid-D improve accuracy by 2.4% and 3.8% over Llama 3.2, respectively. AIRA-Compose also finds novel model architectures that achieve steeper, more efficient compute-optimal scaling frontiers. AIRAformer-C scales 54% and 71% faster than Llama 3.2 and the best Composer-found Transformer, while AIRAhybrid-C scales 23% and 37% faster than the modified Nemotron-2 and the best Composer-found hybrid, respectively. For the first time, AIRA-Compose demonstrates that AI research agents can discover hybrid architectures that surpass hand-designed baselines on the accuracy–compute frontier. It establishes a flexible paradigm for LLMs to discover next generation foundation models, marking a step towards recursive self-improvement.