Source-linked AI summary

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

Alberto Pepe, Chien-Yu Lin, Despoina Magka, Bilge Acun, Yannan Nellie Wu, Anton Protopopov, Carole-Jean Wu, Yoram Bachrach

arXiv:2605.15871v1cs.AI

TL;DR

The paper addresses whether AI agents can systematically discover neural architectures and training optimizations beyond manually designed Transformers. It introduces AIRA-Compose and AIRA-Design, whose agents produce architectures and optimization strategies that match or surpass hand-designed baselines across several evaluations.

  • Problem

    As neural-model design spaces become vast and combinatorial, manual exploration may overlook high-performing configurations and new computational primitives.

    Source: §1 · p. 2
  • Method

    The paper introduces AIRA-Compose for arranging predefined computational primitives and AIRA-Design for open-ended implementation of novel models, organized as AIRS-Bench agentic tasks.

    Source: §1 · p. 3
  • Results

    Across architecture discovery, long-range modeling, and training-loop optimization, agents achieve favorable performance, including a 0.968 validation BPB from Greedy Opus 4.5.

    Source: §1 · p. 3, §7 · p. 18
  • Takeaways & Limitations

    The results support agent-based discovery as a viable approach for engineering high-performing architectures and optimization strategies toward recursive self-improvement.

    Source: §1 · p. 3, §7 · p. 18

    AIRA-Compose relies on small-scale proxies, non-agentic aggregation and extrapolation, and a single dataset per task, limiting reliable large-scale prediction and robustness.

    Source: Limitations and Future Work · p. 18

Abstract

from arXiv · show

Toward recursive self-improvement, we investigate LLM agents autonomously designing foundation models beyond standard Transformers. We introduce a dual-framework approach: AIRA-Compose for high-level architecture search, and AIRA-Design for low-level mechanistic implementation. AIRA-Compose uses 11 agents to explore fundamental computational primitives under a 24-hour budget. Agents evaluate million-parameter candidates, extrapolating top designs to 350M, 1B, and 3B scales. This yields 14 architectures across two families: AIRAformers (Transformer-based) and AIRAhybrids (Transformer-Mamba). Pre-trained at 1B scale, these consistently outperform Llama 3.2 and Composer-found baselines. On downstream tasks, AIRAformer-D and AIRAhybrid-D improve accuracy by 2.4% and 3.8% over Llama 3.2. Furthermore, AIRA-Compose finds models with highly efficient scaling frontiers: AIRAformer-C scales 54% and 71% faster than Llama 3.2 and Composer's best Transformer, while AIRAhybrid-C outscales Nemotron-2 by 23% and Composer's best hybrid by 37%. AIRA-Design tasks 20 agents with writing novel attention mechanisms for long-range dependencies and high-performing training scripts. On the Long Range Arena benchmark, agent-designed architectures reach within 2.3% and 2.6% of human state-of-the-art on document matching and text classification. On the Autoresearch benchmark, Greedy Opus 4.5 achieves 0.968 validation bits-per-byte under a fixed time budget, surpassing the published minimum. Together, these frameworks show AI agents can autonomously discover architectures and algorithmic optimizations matching or surpassing hand-designed baselines. This establishes a powerful paradigm for discovering next-generation foundation models, marking a clear step toward recursive self-improvement.

1 Introduction

The paper presents AIRA-Compose and AIRA-Design as complementary agentic frameworks for recursive self-improvement through neural architecture search and low-level model design. Agents discover performant hybrid architectures, attention mechanisms, and training optimizations that outperform established baselines or approach human state of the art.

  • The frameworks address Transformer limitations including quadratic attention complexity and inference-time key-value caching, as hybrid model design creates a vast combinatorial search space.The paper argues that agent-guided NAS can systematically explore non-obvious configurations and new computational primitives beyond manual, domain-expertise-driven design.Source: §1 · p. 2
  • AIRA-Compose searches arrangements of predefined computational primitives at small scale and scales only top-performing architectures, while AIRA-Design implements novel primitives and trains them from scratch.Both approaches are implemented as AIRS-Bench tasks supporting autonomous research loops, with AIRA-Compose restricting agent freedom to primitive arrangement and AIRA-Design enabling open-ended code generation.Source: §1 · pp. 2–3
  • The authors introduce 12 AIRS-Bench agentic tasks through flexible frameworks that can use different agent harnesses and reasoning models.AIRA-Compose discovers novel architectures from predefined primitives, whereas AIRA-Design targets Long Range Arena and Autoresearch through low-level mechanistic implementation.Source: §1 · pp. 2–3
  • 2.3pp and 2.6pp are the gaps between agent-designed models and human state of the art on document matching and text classification, respectively, in Long Range Arena.The reported human-state-of-the-art accuracies are 82% for document matching and 91% for text classification; four agents exceed 0.3 average normalized score across three tasks.Source: §1 · p. 3
  • 0.968 validation BPB is achieved by the best agent on Autoresearch, surpassing the published reference minimum; larger-than-reported baseline deltas are reproduced in 20% of experiments.Agents iteratively optimize a GPT training script under a 5-minute wall-clock budget, and curated literature and code repositories shift optimization strategies toward the lowest BPB across 100 runs.Source: §1 · p. 3
  • Agent-discovered AIRAformer and AIRAhybrid models consistently outperform Llama 3.2, approximated Nemotron-2, and Composer-found alternatives across validation loss, normalized accuracy, and scaling evaluations.These families use original interleavings of Transformer, MLP, and, for AIRAhybrids, Mamba2 SSM blocks, with models extrapolated to 350M, 1B, and 3B parameters.Source: §1 · p. 3

2 Methodology

The methodology frames architecture search and mechanistic design as standardized AIRS-Bench tasks executed through AIRA-dojo, an agent harness for tree-based code exploration. RSI tasks adapt submission and evaluation artifacts to assess complete architectures and training pipelines under consistent conditions.

Agent harness

  • AIRA-dojo evolves candidate Python solutions through structured tree search using Draft, Debug, Improve, and Analyze operators.The harness supports policies such as greedy search and Monte Carlo Tree Search.Source: §2 · p. 4

Search scaffolds

  • One-shot scaffolds generate one candidate, whereas greedy scaffolds explore multiple solutions, improving the highest-fitness candidate and debugging solutions identified as buggy.Greedy runs begin with five drafted solutions and can span tens to hundreds of search steps, each with validation and independently obtained test fitness.Source: §2 · p. 4

Task formulation

  • Architecture search and mechanistic design are formulated as AIRS-Bench tasks defined by a problem, dataset, and fitness metric.This triplet specifies the challenge, data, and performance measure for each task.Source: §2 · p. 4

RSI task suite

  • The 12 RSI tasks comprise AIRA-Compose architecture submissions and AIRA-Design LRA or Autoresearch artifacts requiring submission.csv, model.py, or train.py.AIRA-Compose covers four tasks across three datasets and primitive pools, while AIRA-Design targets architectural implementations and training scripts.Source: Table 1 · p. 5

Evaluation design

  • RSI evaluation replaces prediction files and isolated metric calculations with complete architectures or training loops evaluated by full pipelines ported from Composer, LRA, and Autoresearch.This structure standardizes evaluation across agent-generated architectures and baselines while removing the need for agents to implement training loops.Source: §2 · p. 5

3 The AIRA-Compose Pipeline

AIRA-Compose transforms Composer’s architecture-search stages into agentic AIRS-Bench tasks, enabling agents to propose, evaluate, and iteratively refine 16-layer designs. Aggregated discoveries are then scaled to 350M, 1B, and 3B parameters for large-scale evaluation across proxy-predictive datasets.

Aggregation and scaling

  • Composer’s aggregation uses layer-wise clustering to select frequent primitives and reduce proxy-training noise before stretching or stacking architectures to target parameter counts.Extrapolation supports 350M, 1B, and 3B models by proportionally expanding contiguous blocks or repeating the discovered architecture.Source: §3 · p. 6

Agentic pipeline

  • AIRA-Compose recasts Composer’s search and evaluation stages as AIRS-Bench tasks, while retaining aggregation and extrapolation for large-scale evaluation.Agents freely formulate structural hypotheses, propose primitive arrangements, evaluate them, and refine designs; discovered architectures are subsequently aggregated and scaled.Source: §3 · pp. 5–6
  • Greedy tree search selects the highest-scoring node for further improvement, allowing agents to use reasoning and prior scores rather than fixed mutation operators.Each node produces a candidate architecture and evaluation script, while improve operations propose new designs informed by the parent’s reasoning and score.Source: Figure 3 · p. 6, §3 · p. 6

Primitive search spaces

  • The search spaces contain 65,536 two-primitive and approximately 43M three-primitive 16-layer arrangements built from MLP, multi-head Attention, and Mamba SSM blocks.Configurations are defined at small scale and extrapolated to 350M, 1B, and 3B parameter scales.Source: §3 · p. 6

Proxy datasets

  • Proxy evaluation uses MAD accuracy, BabiStories cross-entropy, and DCLM cross-entropy, with agents validating hypotheses before isolated test-set evaluation.Training sizes are 800, 927,158, and 10,000 samples respectively; submitted architectures are trained from scratch on full training sets and tested on withheld data.Source: §3 · pp. 6–7

4 AIRA-Design: low-level mechanistic design

AIRA-Design recasts LRA and Autoresearch as agent-driven low-level design tasks, requiring agents to implement neural architectures and optimize training loops from executable code. The framework evaluates long-range sequence modeling and language-model training efficiency under configurable task settings and fixed resource constraints.

4 AIRA-Design: low-level mechanistic design

  • AIRA-Design uses up to 20 agents and 12 LLMs to write novel computational primitives or complete training loops into model.py or train.py files.Tasks run with one-shot and greedy AIRA-dojo scaffolds across equivalent LRA and Autoresearch benchmarks.Source: Figure 4 · p. 7
  • Agents implement memory-efficient, sub-quadratic attention mechanisms for sequences exceeding 2000–4000 tokens without materializing full O(n2) attention matrices.Each submission defines a compatible CustomEncoder or CustomDualEncoder in an executable model.py file.Source: §4 · p. 7
  • Configurable LRA tasks let agents tune global variables, whereas Non-Configurable tasks fix default parameters to isolate architectural design from hyperparameter tuning.Autoresearch also includes a With Literature variant providing 41 papers and 14 reference repositories as structured resources.Source: §4 · pp. 7–8
  • Autoresearch asks agents to minimize validation bits per byte within a fixed 5-minute, single-GPU budget by modifying architecture, optimization, and training configuration while leaving data and evaluation fixed.The baseline includes multi-head causal attention, rotary embeddings, flash attention, and a hybrid Muon–AdamW optimizer; agents may alter model and training-loop components but not the data pipeline, tokenizer, or harness.Source: §4 · p. 7

4.1 Datasets

  • LRA evaluation covers IMDB sentiment, ListOps, and AAN retrieval with accuracy, using validation during search and a held-out test split for final evaluation.Autoresearch uses ClimbMix web text, an approximately 8192-token BPE vocabulary, and validation bits per byte, where lower is better.Source: §4.1 · p. 8

5 Experiments

Experiments follow AIRS-Bench (Lupidi et al., 2026) using bounded multi-seed agent runs, H200 resources, and task-specific evaluation metrics. AIRA-Compose emphasizes scaled-model performance, while AIRA-Design additionally measures submission validity and normalized progress.

Experimental setup

  • Experiments use 20 seeds for one-shot agents and 10 for greedy agents, with runs capped at 24 hours or 500 steps and one H200 GPU per agent.BabiStories and DCLM searches receive 60 hours per run; large-scale AIRA-Compose training uses 8 H200 GPUs at 350M parameters and 16 at 1B or 3B.Source: §5 · p. 8

AIRA-Compose metrics

  • AIRA-Compose ranks small models by accuracy or cross-entropy loss and evaluates scaled models using validation loss, six-task zero-shot accuracy, and the 14-task DCLM Core score.The six downstream tasks are HellaSwag, WinoGrande, ARC-Easy, ARC-Challenge, PIQA, and SciQ; task metrics include accuracy, normalized accuracy, and F1 where specified.Source: §5 · p. 8

AIRA-Design metrics

  • AIRA-Design reports raw score, valid submission rate, and normalized score, with normalized score mapping worst performance to 0 and literature SOTA to 1.Normalized score exceeds 1 when an agent surpasses SOTA, while failed or invalid submissions receive a normalized score of 0.Source: §5 · p. 9

Task-specific reporting

  • Because LRA tasks require novel functional code, all three AIRA-Design metrics are reported; AIRA-Compose reports raw and scaled-up metrics, while Autoresearch uses validation bits-per-byte.AIRA-Compose submissions map to predefined blocks, making valid submission rate virtually 100%.Source: §5 · p. 9

6 Results

Section 6 evaluates agentic architecture search, mechanistic design, and optimization strategies across competitive model discovery and autoresearch. The results examine scaling and latency frontiers, solution novelty, iterative validation, task difficulty, agent choice, literature access, and optimization progress.

6.1.1 2 primitives (M, mA)

  • 2.734 validation loss, 59.7% average 0-shot accuracy, and 48.9% DCLM Core score make AIRAformer-D Stretched the strongest reported two-primitive model.It outperforms Llama 3.2 and Composer-found models across validation loss and downstream performance, with low variance across three seeds.Source: Table 2 · p. 10, §6.1.1 · p. 10

6.1.2 3 primitives (M, mA, Mb)

  • 2.719 validation loss and 60.5% average 0-shot accuracy make AIRAhybrid-D Stretched the strongest reported hybrid, generally outperforming Mamba and approximated Nemotron baselines.At 1B scale and approximately 38 TPP, AIRAhybrids often outperform established baselines across linguistic and reasoning benchmarks.Source: Table 3 · p. 11, §6.1.2 · p. 11

6.1.1 2 primitives (M, mA) / 6.1.2 3 primitives (M, mA, Mb)

  • AIRAformer-C/D and AIRAhybrid-C show steeper scaling, but balanced architectures are more compute-efficient under isoFLOP constraints.AIRAformer-C/D are less efficient when compute is binding, while AIRAhybrid-C Stretched reaches ∆q = −0.14 despite a higher intercept; balanced hybrids achieve lower validation loss across sizes and budgets.Source: §6.1.1 · pp. 10–11, Figure 5 · p. 10, §6.1.2 · p. 12

6.1.2 3 primitives (M, mA, Mb)

  • AIRAhybrid-D Stacked and Stretched advance Composer’s 1B latency–validation-loss Pareto frontier, achieving lower validation loss at strictly lower latency.The comparison uses the isoToken protocol and H100-profiled block latencies.Source: Figure 7 · p. 12, §6.1.2 · p. 12

6.2.1 Long Range Arena

  • Greedy Gemini 3 Pro achieves the highest mean best accuracy across nearly all Long Range Arena task-configuration combinations, but agents mainly recombine established mechanisms rather than inventing novel attention methods.Solutions use linear attention, hierarchical pooling, recurrent-memory local attention, depthwise convolutions, and gated linear units; one-shot agents produce no valid submissions, whereas iterative refinement is essential.Source: Figure 8 · p. 13, §6.2.1 · p. 13, Figure 9 · p. 14
  • 0.257 ± 0.054 normalized score makes ListOps the most tractable task, while all LRA tasks fall in AIRS-Bench’s hard category and configurable search generally lowers performance.Low scores reflect difficulties with JAX/Flax, complete low-level attention implementations, fixed compute budgets, out-of-memory errors, and numerical instability; submitted accuracy usually remains close to best observed accuracy.Source: §6.2.1 · pp. 13, 15, Table 4 · p. 15

6.2.2 Autoresearch

  • 0.990 median BPB for Opus 4.6 +Lit, 0.991 for Opus 4.6, and 0.992 for Opus 4.5 substantially improve on the 1.0121 RADv1 baseline.Opus families converge more reliably than Gemini agents, but literature access does not uniformly help; the greedy scaffold regenerates full scripts, causing crashes and larger, less-isolated changes.Source: Figure 12 · p. 16, §6.2.2 · p. 16
  • Δ = 0.036 was the largest single improvement, achieved by Opus 4.5 +Literature after introducing focal loss around step 15.This result is reported among the best-seed improvement steps for the Autoresearch variants.Source: §6.2.2 · p. 17, Table 5 · p. 17
  • Opus 4.6 +Literature recorded the most improvement steps, with moderate gains from tuning depth, batch size, learning rates, and MLP width.Its trajectory contrasted with the larger, more targeted improvement produced by Opus 4.5 +Literature.Source: §6.2.2 · p. 17, Table 5 · p. 17
  • The improvement trajectories show that literature access can steer agents toward qualitatively different optimization strategies across Autoresearch runs.The complete 100-run progression is provided in Figure 28, while the lowest-validation-BPB solution is implemented in Appendix H.Source: §6.2.2 · p. 17

7 Conclusion · Appendix · A Related Work

AIRA-Compose and AIRA-Design demonstrate agentic discovery of architectures and training methods, while revealing limitations in scalability, innovation, and interpretability. Related work situates these efforts within hybrid language models, neural architecture search, and recursive self-improvement.

7 Conclusion

  • Greedy Opus 4.5 surpassed the published Autoresearch reference baseline with final validation BPB of 0.968, while Greedy Gemini 3 Pro and Greedy Opus 4.6 reached within 3% of human SOTA on LRA.AIRA-Design evaluated 20 agents on LRA and 7 agents on Autoresearch; literature and repository augmentation improved optimization strategies and convergence.Source: §7 · p. 18

Limitations and Future Work

  • AIRA-Compose remains limited by small-scale proxy fidelity, non-agentic aggregation and extrapolation, single-dataset evaluation, and a search space focused mainly on primitive arrangements.Future extensions include multi-GPU harnesses, autonomous aggregation and scaling, validation across datamixes, and tunable hyperparameters or normalization choices.Source: Limitations and Future Work · p. 18
  • AIRA-Design agents mainly recombined prior mechanisms rather than producing paradigm-shifting theory, while one-shot generation failed and weaker agents degraded under expanded configurability.PyTorch bias constrained JAX/Flax performance; full-file regeneration hindered causal attribution, and literature access helped agents inconsistently.Source: Limitations and Future Work · p. 19

A Related Work

  • Hybrid language models interleave attention and state-space primitives to address self-attention’s quadratic cost while balancing efficiency and expressiveness.The related work identifies Mamba-2, Gated DeltaNet, Samba, and Hymba as linear-complexity alternatives, alongside newer hybrid models including Qwen3-Next and Nemotron variants.Source: A Related Work · p. 26
  • Neural architecture search addresses hybrid models’ combinatorial design space, but traditional LLM-scale methods remain computationally prohibitive and therefore rely on approximations.Prior approaches include supernets, zero-cost proxies, PostNAS, analytical layer allocation, and Composer’s small-scale proxy search.Source: A Related Work · p. 26
  • Agentic research now spans autonomous scientific discovery, engineering benchmarks, paper reproduction, and iterative optimization toward recursive self-improvement.Related efforts include AIRS-Bench, MLGym, MLE-bench, PaperBench, Automated LLM Speedrunning, Autoresearch, and the nanochat environment.Source: A Related Work · p. 26

B Small Scale Ranking

The section ranks agent-discovered architectures from greedy search on MAD using two- and three-primitive configurations, with each architecture shown alongside its test accuracy and discovering agent.

MAD 2-Primitives

  • Top 120 two-primitive architectures are ranked by test performance, displaying primitive sequences, accuracies, and the discovering agents.The configurations use M and mA primitives.Source: Figure 14 · p. 27

MAD 3-Primitives

  • Top 120 three-primitive architectures are ranked by test performance, displaying primitive sequences, accuracies, and the discovering agents.The configurations use M, mA, and Mb primitives.Source: Figure 15 · p. 28

C Additional Details · D Aggregation and Scaling Patterns

The additional analyses provide per-task evaluation details, characterize validation–test generalization gaps, and show that balanced compute allocation generally improves isoFLOP validation loss across architecture families and scales.

C Additional Details

  • Table 6 documents environment requirements for AIRA-Compose tasks, while the expanded DCLM breakdown reports performance across all 14 evaluation tasks.The per-task table uses unweighted averages for DCLM Core Score and identifies best and second-best results within primitive groups.Source: Table 6 · p. 29, C Additional Details · p. 29, Table 7 · p. 31
  • Running-maximum trajectories compare submitted and best solutions across agents, with 95% confidence intervals and explicit validation–test generalization gaps.The 3-primitive search is constrained to M, mA, and Mb, and the best-performing agent is highlighted.Source: Figure 16 · p. 29, Figure 17 · p. 29, C Additional Details · p. 29
  • Validation–test gaps persist throughout greedy architecture search, highlighting generalization challenges when validation fitness selects agent-designed architectures.The 2-primitive trajectories average 10 seeds, while the 3-primitive trajectories use the same validation/test trajectory framing.Source: Figure 20 · p. 30, Figure 21 · p. 31
  • The DCLM Core results report an overall average of 49.3% at 1B scale; AIRAhybrid-B leads WinoGrande, XWinograd, HellaSwag, and PIQA.AIRAhybrid-E leads OpenBookQA and ties for best on AGIEval, with scores of 37.4 and 23.0, respectively.Source: Table 7 · p. 31, C Additional Details · p. 32
  • AIRAformer variants outperform Composite and Llama baselines on CoQA, SQuADv2, and HellaSwag, with AIRAformer-D Stretched reaching a 48.9% average.BoolQ remains near chance at approximately 37.9% across architectures under 10-shot evaluation.Source: C Additional Details · p. 32
  • 2.7415 is AIRAformer-A Stretched’s validation loss at 1B scale and 4 × 10^20 FLOPs, versus 2.7647–2.7732 for attention-heavy variants.Across scales and budgets, balanced 2-primitive architectures outperform attention-heavy AIRAformer-C and AIRAformer-D variants.Source: C Additional Details · p. 32, Figure 23 · p. 33
  • Balanced 3-primitive designs such as AIRAhybrid-C Stretched and Composer (2Mb-M-3A) consistently achieve the lowest validation losses, whereas Mamba-only and attention-dominated hybrids underperform.The comparison spans 350M, 1B, and 3B models across five FLOP budgets.Source: C Additional Details · p. 32, Figure 24 · p. 34

D Aggregation and Scaling Patterns

  • The scaling analysis compares validation-loss frontiers for 350M, 1B, and 3B architectures across five FLOP budgets, including approximated Nemotron-2 and Nemotron-H references.The full 3-primitive comparison includes 11 architectures and an optimal-frontier analysis.Source: C Additional Details · p. 32, Figure 23 · p. 33, Figure 24 · p. 34

D.1 Aggregation methods

D.1 defines eight aggregated architectures across two- and three-primitive spaces, using rank-weighted dataset aggregation, N0/N1/N2 aggregation, and clustering-based selection.

Two-primitive space

  • AIRAformer-A aggregates top-20 architectures from each of three datasets with rank-based exponential weighting and equal dataset weights.Its base pattern is (A + M) + (2A + 2M) + 4 × (A + M) + 2M.Source: §D.1 · p. 36
  • AIRAformer-B, AIRAformer-C, and AIRAformer-D use N1 or N2 aggregation over high-performing MAD architectures, with C and D selected from three k-means clusters.Their base patterns are respectively 2A + 5 × (A + M) + 4M; (2A + M) + 3 × (A + M) + (2A + M) + 4A; and 5 × (2A + M) + A.Source: §D.1 · p. 36

Three-primitive space

  • AIRAhybrid-A, AIRAhybrid-B, and AIRAhybrid-C apply N0, N1, or N2 aggregation to the top cluster from three-cluster k-means over greedy GPT-5 architectures.Their base patterns are 2Mb + M + 11Mb + 2M; 3 × (2Mb + M + A) + 2Mb + 2M; and 2Mb + A + 2 × (Mb + A) + M + 2 × (A + Mb + A + M).Source: §D.1 · p. 36
  • AIRAhybrid-D and AIRAhybrid-E aggregate the top cluster from three-cluster k-means over architectures found by all six agents, using N1 and N2 respectively.Their base patterns are 2Mb + M + 2 × (Mb + M) + A + M + Mb + M + A + M + Mb + M + A + M and 5 × (Mb + M + A) + M.Source: §D.1 · p. 36

D.2 Scaling patterns

The section specifies how the architecture primitives are implemented at 16-layer proxy scale and at 350M, 1B, and 3B scales. It distinguishes MLP, attention, and Mamba configurations while referring to Table 8 for complete layer patterns and parameter summaries.

Scaling setup

  • Implementations use a 16-layer proxy with d = 128, while large-scale configurations set d, attention heads, head dimensions, and hidden dimensions per scale.The large-scale settings follow the IsoFLOP methodology.Source: §D.2 · p. 36

MLP

  • MLP uses a two-layer ReLU network with hmlp = 258 at small scale and a gated SwiGLU variant with scale-dependent expansion at large scale.The large-scale variant uses gate and up projections d → hmlp and a down projection hmlp → d.Source: §D.2 · p. 36

Multi-head attention

  • Attention uses 16 query and 16 key-value heads with dh = 8 at small scale, then switches to GQA with nkv = 8 at large scale.Large-scale grouped-query attention shares key-value heads across more query heads, with dkv = nkv × dh.Source: §D.2 · p. 36

Mamba

  • Mamba uses a Mamba-2 selective SSM with linear sequence complexity, data-dependent gating, and selective scan; its dimensions and state settings vary by scale.At small scale, essm = 1.25 gives dssm = 160, with ns = 4, k = 4, and ng = 1; large scale uses a rounded dssm formula and dh = 64.Source: §D.2 · pp. 36–37

Architecture configurations

  • Table 8 summarizes the complete two-primitive layer patterns, depths, parameter counts with and without embeddings, and approximate Attention-to-MLP ratios across scales.The supplied passage identifies these reported quantities but does not provide their numerical table entries.Source: Table 8 · p. 37

D.3 IsoFLOP Budget Methodology

The isoFLOP methodology converts primitive-level training costs into comparable training-step budgets using standard FLOP accounting, fixed model configurations, and a constant token batch size. It then allocates steps by architecture-specific per-step costs for attention, MLP, and Mamba components.

FLOP accounting

  • 6 · d_in · d_out FLOPs per projection per token are obtained by charging 2MKN for the forward pass and multiplying by 3 for training.The factor accounts for one forward and two backward passes.Source: §D.3 · p. 37

Model configurations

  • All three scales use sequence length s = 8,192, n_kv = 8 key-value heads, and a constant batch size of 524,288 tokens per step.The batch size is derived as DP × local_bs × acc_steps × s.Source: §D.3 · p. 37

Per-primitive FLOPs

  • Per-token training costs separately account for grouped-query attention, SwiGLU MLP projections, and the four Mamba2 SSM components.Attention includes four linear projections and quadratic attention computation; Mamba2 includes input projection, depthwise convolution, selective scan, and output projection.Source: §D.3 · p. 40

Training steps from FLOP budget

  • Training-step cost is C_step = (L_Attn · F_Attn + L_MLP · F_MLP + L_SSM · F_SSM) × B, and steps for budget F are T = floor(F/C_step).This formulation accounts for each architecture’s layer composition while holding the token batch size fixed.Source: §D.3 · p. 40

Worked examples

  • For a budget of 2 × 10^19 FLOPs, the two reported configurations require 4,552 and 5,308 steps, respectively, with both values matching cited tables.The calculations use per-step costs of 4.39 × 10^15 and 3.77 × 10^15 FLOPs, respectively.Source: §D.3 · p. 40

E AIRA-Design: LRA - Task Details

AIRA-Design agents extensively tuned three configurable LRA tasks, producing 7,537 valid evaluation steps, while results were tracked across six configurable and non-configurable task setups.

Task Details

  • 7,537 valid evaluation steps were generated across ListOps, Text, and Retrieval during hyperparameter tuning by the greedy agents.The totals were 3,189 for ListOps, 2,642 for Text, and 1,706 for Retrieval.Source: E AIRA-Design: LRA - Task Details · p. 43
  • Hyperparameter exploration covered value ranges and unique evaluated settings for each configuration parameter across the three configurable LRA tasks.These exploration statistics are reported for the agents’ configurable-task searches.Source: Table 13 · p. 43, E AIRA-Design: LRA - Task Details · p. 43
  • The best generated solutions are broken down for six LRA AIRA-Design tasks in configurable and non-configurable setups, including generating agents, scores, and architecture summaries.For non-configurable tasks, the breakdown identifies configurations that differ from the defaults.Source: Table 14 · p. 44

F AIRA-Design: Autoresearch - Task Details

The Autoresearch task augments agents with structured literature and code context to support low-cost, high-impact optimization within a fixed 5-minute training budget. Its review spans architecture improvements, training strategies, and optimizers, with additional implementation resources for selected papers.

Literature context

  • 41 curated papers provide structured summaries covering architecture improvements, training strategies, and optimizers, with 14 also offering working GitHub repositories.The summaries include task specifications, pipelines, datasets, architectures, hyperparameters, and experiment metadata extracted from papers’ .tex files.Source: F AIRA-Design: Autoresearch - Task… · p. 45, Table 16 · p. 45, Table 17 · p. 46

Task setup

  • Within the fixed 5-minute training budget, agents are instructed to consult the summaries, inspect reference code, and prioritize low-cost changes with high convergence impact.The literature context is supplied through an additional pwc/ folder containing task_info.jsonl and, for selected papers, code under pwc/code/.Source: F AIRA-Design: Autoresearch - Task… · p. 45

Literature context

  • The literature spans papers published from 2021 through 2026, including 2 from 2021, 7 from 2023, 21 from 2024, 8 from 2025, and 3 from 2026.The papers are organized into Architecture Improvements, Training Strategies, and Optimizers.Source: F AIRA-Design: Autoresearch - Task… · p. 45

G AIRA-Design: Autoresearch - Additional Results · H Best Autoresearch Solution (Greedy Opus 4.5 + Literature access)

The additional results analyze how agents improve validation bits-per-byte through compound architectural and hyperparameter changes, while the best solution combines value embeddings, sliding-window attention, focal loss, and fused AdamW/Muon optimization. Across 156 improvements, architectural and attention-pattern changes are especially influential, but full-file regeneration limits causal attribution.

G AIRA-Design: Autoresearch - Additional Results

  • Across 156 improvement steps, depth and width changes occurred in 46% of steps, while attention-pattern changes had the highest median gain of 0.009 across 37% of steps.Depth and width changes had a median Δ of 0.007.Source: G AIRA-Design: Autoresearch -… · p. 48
  • Learning-rate changes co-occurred with 49% of improvements, showing that agents frequently modify optimization alongside architecture.The analysis extracts 27 structured features from regenerated train.py files and compares each improvement against the previous best.Source: G AIRA-Design: Autoresearch -… · p. 48
  • ReLU2 → SwiGLU transitions improved BPB more than reversions, with median Δ values of 0.007 versus 0.003, suggesting many SwiGLU trials were exploratory.The most frequent transitions were ReLU2 ↔ SwiGLU, occurring 14 and 13 times.Source: G AIRA-Design: Autoresearch -… · p. 50
  • Value embeddings and EMA additions produced the strongest boolean-feature gains, with median Δ values of 0.014 and 0.015, respectively.Removing z-loss also had a median Δ of 0.013, while removing gradient clipping had a median Δ of 0.004.Source: G AIRA-Design: Autoresearch -… · p. 50
  • Only 7 of 27 features were ever the sole change in an improvement, because agents regenerate the entire file and typically make compound modifications.This regeneration paradigm prevents clean single-variable ablations and complicates causal attribution.Source: G AIRA-Design: Autoresearch -… · pp. 48, 50

H Best Autoresearch Solution (Greedy Opus 4.5 + Literature access)

  • The best solution uses alternating short and long attention windows, value embeddings gated into attention values, ReLU2 MLPs, and rotary-normalized causal attention.Its implementation also uses RMS normalization, residual and initial-state scaling parameters, and a softcapped output head.Source: H Best Autoresearch Solution (Greedy… · pp. 51–52
  • Training uses separate AdamW groups for embeddings and scalars plus Muon updates for matrix parameters, with fused compiled optimizer steps and shape-dependent normalization.The training script specifies a 0.55 warmdown ratio, zero final learning-rate fraction, and a fixed time-budget stopping condition.Source: H Best Autoresearch Solution (Greedy… · pp. 53–55
Loading 2605.15871v1…