AIRAformers: Agent-Discovered Transformer Models

• AIRAformers are a family of Transformer-only models discovered by AI agents using an agentic architecture search, differing in the depthwise ordering of attention and MLP blocks.
• They are extrapolated from 16-layer prototypes to multi-billion parameter scales and consistently outperform baselines like Llama 3.2 under fixed token budgets.
• Their design leverages variable scheduling of attention and MLP primitives to optimize compute performance, offering both attention-heavy and balanced configurations.

Searching arXiv for the specified papers and closely related references.

AIRAformers are a family of Transformer-based LLM architectures discovered by AI agents within the AIRA-Compose framework. They are composed only of Transformer primitives—multi-head self-attention and MLP blocks—but depart from the standard decoder-only Transformer by varying the ratio and ordering of those primitives across depth rather than enforcing a rigid $1{:}1$ attention–MLP alternation. In the reported experiments, the family is extrapolated from 16-layer proxy architectures to 350M, 1B, and 3B parameter scales, and at 1B scale it consistently outperforms Llama 3.2 and Composer-found Transformer baselines under a fixed token budget; some variants also exhibit more favorable compute–performance scaling frontiers (Pepe et al., 15 May 2026).

1. Definition, scope, and nomenclature

AIRAformers are one of two architecture families produced by the agentic search in AIRA-Compose, the other being AIRAhybrids. The distinction is architectural: AIRAformers use only Transformer primitives, whereas AIRAhybrids combine Transformer primitives with Mamba2 State Space Model blocks. Within that taxonomy, AIRAformers are the Transformer-only outcomes of an agentic architecture search over layerwise arrangements of attention and MLP blocks (Pepe et al., 15 May 2026).

The defining departure from a standard decoder-only Transformer is global rather than intra-block. A conventional model uses a strict pattern of $[A-M]$ repeated at every layer. AIRAformers instead use learned 16-layer base patterns whose local motifs can include long contiguous runs of attention or attention-heavy segments, such as

$(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$

This suggests that the main architectural novelty lies in the depthwise scheduling of familiar primitives rather than in replacing attention or feed-forward mechanisms with new layer types.

Four named variants are reported: AIRAformer-A, AIRAformer-B, AIRAformer-C, and AIRAformer-D. Each is first identified as a 16-layer base architecture at small scale and then extrapolated to larger models. The paper associates AIRAformer-A and AIRAformer-B with balanced or slightly MLP-heavy configurations, and AIRAformer-C and AIRAformer-D with attention-heavy configurations.

2. Primitive structure and architectural variants

At the primitive level, the search space for AIRAformers contains only two blocks. The MLP primitive is a Transformer feed-forward block; at small scale it is a standard 2-layer MLP with ReLU and hidden dimension $h_{\text{mlp}}=258$, while at larger scales it becomes a SwiGLU-style feed-forward network with training FLOPs per token

$F_{\text{MLP}} = 18dh .$

The attention primitive is causal multi-head self-attention using grouped-query attention at scale, with training FLOPs per token

$$F_{\text{Attn}} = 6(2d^2 + 2d \cdot d_{\text{kv}}) + 12sd .$$

All AIRAformers keep the same intra-block design as the baselines, including typical normalization and positional encodings inherited from the training setup; the novelty is the global arrangement of $A$ and $M$ blocks (Pepe et al., 15 May 2026).

The four reported variants are summarized below.

Variant	16-layer base pattern	Characterization
AIRAformer-A	$(A + M) + (2A + 2M) + 4 \times (A + M) + 2M$	$7A:9M$, balanced or slightly MLP-heavy
AIRAformer-B	$[A-M]$0	$[A-M]$1, balanced or slightly MLP-heavy
AIRAformer-C	$[A-M]$2	$[A-M]$3, attention-heavy
AIRAformer-D	$[A-M]$4	$[A-M]$5, attention-heavy

Scaling from these 16-layer bases uses two strategies. In the Stacked strategy, the base pattern is repeated to reach the target depth. In the Stretched strategy, contiguous groups of $[A-M]$6 and $[A-M]$7 are proportionally expanded while preserving the approximate attention-to-MLP ratio. For example, AIRAformer-C at 1B has 48 layers under both strategies, but the distribution of contiguous attention and MLP segments differs between the stacked and stretched realizations.

A recurring empirical pattern is that AIRAformer-C and AIRAformer-D contain long attention-only tails or extended attention clusters, whereas AIRAformer-A and AIRAformer-B remain closer to a balanced composition. A plausible implication is that the search is exploiting compute-regime-specific trade-offs between the representational benefits of additional attention layers and the lower cost of MLP-heavy layouts.

3. Agentic discovery in AIRA-Compose

AIRAformers are discovered by AIRA-Compose, which recasts architecture search as a set of LLM-agent tasks. In the 2-primitive search relevant to AIRAformers, the search space contains 16 layers, each assigned either $[A-M]$8 or $[A-M]$9, for a total of

$(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$0

possible patterns. AIRA-Compose uses 11 agents to explore fundamental computational primitives under a 24-hour budget; more specifically, the setup includes 10 main agents derived from 5 base LLMs crossed with one-shot and greedy modes, plus additional greedy GPT-5 runs on MAD (Pepe et al., 15 May 2026).

The agent scaffold is organized around four stages: draft, debug, improve, and analyze. The agents emit architecture proposals as 16-token strings such as mh-attention mlp mlp ..., together with evaluation code, and those proposals are then trained and evaluated by an external fixed training script derived from Composer. Each greedy run lasts up to 24 hours on MAD or 60 hours on BabiStories and DCLM, on a single H200 GPU.

Three proxy tasks are used for small-scale evaluation:

Proxy task	Description	Metric
MAD	6 synthetic mechanistic tasks	average accuracy
BabiStories	synthetic children’s stories	cross-entropy loss
DCLM subset	10k train / 9,275 test samples	cross-entropy loss

Across all runs, the agents explored 2,307 unique 16-layer architectures, about 3.17% of the 2-primitive search space. The top architectures were then aggregated using clustering and the paper’s “N0/N1/N2” aggregation schemes to produce robust base patterns. This suggests that the reported AIRAformer variants are not single-run anomalies but clustered summaries of high-performing regions in the proxy search space.

4. Extrapolation, training regimes, and compute analysis

After base patterns are identified at small scale, they are extrapolated to 350M, 1B, and 3B parameter models and pretrained on DCLM. The paper evaluates the resulting models under two complementary regimes: an isoToken regime at 1B parameters and an isoFLOP regime spanning multiple parameter scales (Pepe et al., 15 May 2026).

In the 1B isoToken setting, all models are trained with batch size 4, sequence length 8,192, and DP=16, corresponding to approximately 37.5B tokens, or about 38 tokens per parameter, over 71,565 steps. Results are reported over 3 random seeds per architecture as mean $(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$1 standard deviation. This setup is used to compare AIRAformers directly to Llama 3.2 and Composer-based Transformer baselines under a fixed token budget.

The isoFLOP analysis uses five FLOP budgets: $(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$2 For each architecture and FLOP budget, validation loss is fit as a quadratic function of model size; the minimum of each parabola defines a compute-optimal frontier. The total step cost is determined by the number of attention and MLP layers in the architecture: $(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$3 and the number of training steps permitted under a FLOP budget $(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$4 is

$(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$5

This methodology makes the A/M ratio an explicit scaling variable. Because attention layers are more expensive than MLPs, attention-heavy AIRAformers can be disadvantaged at low compute even if they dominate under fixed-token or higher-compute conditions. The paper’s central analytical claim is therefore not merely that one pattern is universally superior, but that different patterns define different compute-optimal frontiers.

5. Reported performance and scaling behavior

At 1B scale under the isoToken regime, AIRAformer-D (Stretched) is the best-performing Transformer-only architecture reported in the paper. Its validation loss is $(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$6, compared with $(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$7 for Llama 3.2, $(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$8 for Composite (Stretched), and $(2A + M) + 3 \times (A + M) + (2A + M) + 4A .$9 for AIRAformer-D (Stacked). Other AIRAformer variants also outperform Llama 3.2, with AIRAformer-A (Stretched), AIRAformer-B (Stacked), AIRAformer-C (Stacked), and AIRAformer-C (Stretched) all reported in the narrow range from $h_{\text{mlp}}=258$0 to $h_{\text{mlp}}=258$1 validation loss (Pepe et al., 15 May 2026).

The same ranking appears in downstream task aggregates. On the six 0-shot tasks—ARC-C, ARC-E, HellaSwag, PIQA, SciQ, and WinoGrande—Llama 3.2 attains 57.5% raw and 58.4% normalized average accuracy, while AIRAformer-D (Stretched) reaches 59.7% raw and 60.8% normalized. The per-task values for AIRAformer-D (Stretched) are also reported: ARC-C 29.4% raw and 31.5 normalized, ARC-E 63.7% and 59.0 normalized, HellaSwag 43.7% and 56.4 normalized, PIQA 72.9% and 72.4 normalized, SciQ 89.4% and 84.5 normalized, and WinoGrande 58.9%.

On the DCLM Core score, which averages 14 few-shot benchmarks, Llama 3.2 scores $h_{\text{mlp}}=258$2, the best Composer Transformer scores $h_{\text{mlp}}=258$3, AIRAformer-C (Stacked) scores $h_{\text{mlp}}=258$4, and AIRAformer-D (Stretched) scores $h_{\text{mlp}}=258$5. The abstract summarizes this family-level result by stating that, on downstream tasks, AIRAformer-D improves accuracy by 2.4% over Llama 3.2.

The isoFLOP results reveal a more differentiated pattern. Balanced AIRAformers such as AIRAformer-A and AIRAformer-B often perform better in low-compute isoFLOP regimes. Attention-heavy AIRAformers such as AIRAformer-C and AIRAformer-D have slightly worse isoFLOP loss at small compute budgets, but their frontiers have steeper slopes than Llama 3.2 and the Composite baselines. The paper summarizes this as AIRAformer-C scaling 54% faster than Llama 3.2 and 71% faster than Composer’s best Transformer, while the AIRAformer-D family also scales better than those baselines.

6. Interpretation, limitations, and related terminology

The paper does not present detailed causal ablations for each AIRAformer motif, but it does identify several empirical regularities. First, attention-heavy patterns such as AIRAformer-C and AIRAformer-D consistently win in isoToken settings, where tokens are fixed and additional attention layers appear to improve dependency modeling per unit depth. Second, balanced patterns such as AIRAformer-A and AIRAformer-B are more FLOP-efficient at small compute because attention is more expensive than MLPs. Third, the steep scaling slopes of attention-heavy variants imply that architectural optimality is compute-regime dependent rather than absolute (Pepe et al., 15 May 2026).

Several limitations are explicit. AIRA-Compose still relies on small-scale proxy searches, so a proxy–target gap remains. The search space is restricted to primitive arrangements; it does not search over normalization types, positional encodings, activation variants, or broader hyperparameter choices. AIRA-Design, which tasks agents with writing new attention mechanisms and training scripts for Long Range Arena and Autoresearch, is complementary rather than integrated: its low-level discoveries are not fed back into the AIRAformer family in this work.

A common misconception is terminological. AIRAformers should not be conflated with Aformer, the Conformer-based end-to-end ASR model for low-resource accented speech recognition introduced in “Multi-pass Training and Cross-information Fusion for Low-resource End-to-end Accented Speech Recognition” (Wang et al., 2023). That work proposes a dual-encoder architecture with a general encoder, an accent encoder, multi-pass training, and fusion mechanisms for accented ASR. It is conceptually relevant as an example of non-standard Transformer/Conformer organization, but it does not mention “AIRAformer(s)” explicitly and belongs to a separate research line.

In that sense, AIRAformers occupy a specific place in the contemporary architecture-search literature: they are not new attention kernels, not hybrid SSM–Transformer models, and not domain-adaptation Conformer systems. They are agent-discovered Transformer-only LLMs whose main innovation is the learned depthwise composition of standard attention and MLP primitives, together with empirical evidence that such compositions can surpass hand-designed baselines and alter compute-optimal scaling behavior.