AIRAhybrids: Agent-Discovered Hybrid LMs
- AIRAhybrids are a family of decoder-style causal language models that combine attention, MLP, and Mamba blocks to overcome limitations of standard Transformers.
- They are discovered via the AIRA-Compose framework using agentic search over a 16-layer primitive space, enabling novel hybrid scheduling strategies.
- They demonstrate competitive scaling behavior and efficiency, achieving strong downstream task performance and improved compute-accuracy trade-offs.
Searching arXiv for the cited AIRAhybrids source paper and closely related context papers. AIRAhybrids are a family of agent-discovered hybrid LLM architectures introduced within the AIRA-Compose framework as part of a broader investigation into autonomous foundation-model design beyond standard Transformers. In the paper that defines them, AIRAhybrids are specifically Transformer–Mamba hybrids discovered in a 3-primitive search space comprising MLP blocks, multi-head attention blocks, and Mamba-2 state-space model blocks. They are presented as one of two discovered architecture families—alongside AIRAformers—and are evaluated after extrapolation from 16-layer proxy patterns to 350M, 1B, and 3B scales (Pepe et al., 15 May 2026).
1. Definition and family position
Within the terminology of the source paper, AIRAhybrids are decoder-style causal LLMs assembled from three standard primitive layer types: attention block (), MLP block (), and Mamba block (). Their defining property is not a new fused operator, but an agent-discovered serial ordering of known primitives. The paper is explicit that the hybridization occurs at the architecture-composition level: AIRAhybrids are “not a new primitive,” but “a new family of layer schedules” over attention, feed-forward, and Mamba layers (Pepe et al., 15 May 2026).
The family is motivated by a post-Transformer design hypothesis stated in the paper. Pure Transformers are framed as having limitations associated with the quadratic complexity of attention and KV-cache costs, whereas state-space models such as Mamba offer linear-time sequence modeling. AIRAhybrids therefore occupy the design space in which attention-heavy and SSM-heavy computation are deliberately interleaved rather than used in isolation. In the paper’s taxonomy, this distinguishes them from AIRAformers, which search only over attention and MLP blocks, whereas AIRAhybrids search over attention, MLP, and Mamba blocks.
The paper’s interpretation of the family is also regime-specific. It reports that more attention-heavy models tend to do better under fixed-token training, whereas more balanced models tend to be better under isoFLOP efficiency. For AIRAhybrids specifically, balanced Mamba-to-attention ratios are associated with lower validation loss at fixed FLOP budgets, while the most attention-heavy hybrid exhibits the steepest scaling slope.
2. Architectural composition
At small scale, the primitive implementations are simple. The MLP block is a 2-layer MLP with ReLU, the attention block is standard causal multi-head self-attention, and the Mamba block is a Mamba-2 selective state-space model. At large scale, the paper states that the MLP becomes a SwiGLU feed-forward block and the attention implementation uses grouped-query attention. The model family is therefore hybrid only in schedule, not in the sense of introducing new mechanistic sublayers (Pepe et al., 15 May 2026).
The five named base AIRAhybrid families are defined as aggregated 16-layer proxy patterns:
| Variant | Base 16-layer pattern | Character |
|---|---|---|
| AIRAhybrid-A | $2Mb + M + 11Mb + 2M$ | Pure-SSM / Mamba-heavy, no attention |
| AIRAhybrid-B | Balanced recurring hybrid | |
| AIRAhybrid-C | Attention-heavy hybrid | |
| AIRAhybrid-D | Balanced, MLP-heavy | |
| AIRAhybrid-E | Regular triplet-cycling hybrid |
These patterns are later scaled through two explicit extrapolation strategies. Stacking repeats the discovered base pattern, whereas stretching proportionally expands contiguous layer groups. The evaluated large-scale variants are AIRAhybrid-A (Stretched), AIRAhybrid-B (Stacked), AIRAhybrid-B (Stretched), AIRAhybrid-C (Stretched), AIRAhybrid-D (Stacked), AIRAhybrid-D (Stretched), AIRAhybrid-E (Stacked), and AIRAhybrid-E (Stretched).
At 1B scale, the paper reports a broad range of compositions. AIRAhybrid-A (Stretched) has , , 0, 1, with parameters 2B. AIRAhybrid-C (Stretched) has 3, 4, 5, 6, with parameters 7B. AIRAhybrid-D (Stretched) has 8, 9, 0, 1, with parameters 2B. The paper uses these counts to emphasize that the family spans SSM-dominated, balanced, and attention-dominated schedules rather than a single canonical hybrid ratio.
3. Discovery within AIRA-Compose
AIRAhybrids are discovered by AIRA-Compose, the paper’s high-level agentic architecture-search framework built on Composer. Composer originally has four stages—search, proxy evaluation, aggregation, and extrapolation—and AIRA-Compose replaces the first two with agents that propose architecture strings, train and evaluate them on proxy tasks, rank them, aggregate high-performing patterns, and scale them to larger models (Pepe et al., 15 May 2026).
For the 3-primitive AIRAhybrid search, the search space is all 16-layer sequences over 3, yielding
4
The paper reports that the agents explored 2,248 unique architectures, corresponding to only 0.0052\% of that search space. This sparsity is one of the paper’s principal claims about search efficiency.
The proxy dataset for the 3-primitive search is MAD, chosen because it had “better agreement between small-scale ranking and large-scale performance.” MAD is described as a synthetic suite of six token-manipulation tasks with 800 training samples, 1,280 test samples, and average accuracy as the metric. For the hybrid search specifically, the paper reports 170 agentic runs on MAD. Across the broader AIRA-Compose effort, the abstract states that the framework deploys an ensemble of 11 agents.
Candidate models are represented as submission.csv files containing strings of 16 primitives such as mh-attention, mlp, and mamba. In greedy mode, an agent drafts five initial candidates, iteratively improves the current best one, and debugs malformed or out-of-memory candidates when necessary. The paper emphasizes that AIRAhybrids were not hand-seeded alternating templates; rather, they emerged from agent-proposed combinations and subsequent aggregation.
Aggregation is also architecture-specific. For AIRAhybrids, the paper reports clustering with 3-means / k-means with 3 clusters, selection of the top cluster, and then application of one of the aggregation operators 5, 6, or 7. AIRAhybrid-A, B, and C are obtained from greedy GPT-5 architectures via 8, 9, and $2Mb + M + 11Mb + 2M$0, respectively, while AIRAhybrid-D and E are obtained from all 6 agents via $2Mb + M + 11Mb + 2M$1 and $2Mb + M + 11Mb + 2M$2. The exact formulas for $2Mb + M + 11Mb + 2M$3, $2Mb + M + 11Mb + 2M$4, and $2Mb + M + 11Mb + 2M$5 are not given in the provided text.
4. Training protocol and evaluation regime
The discovery pipeline begins with million-parameter proxy models. The paper states that AIRAhybrids originate from 16-layer small-scale architectures with model dimension $2Mb + M + 11Mb + 2M$6, evaluated on MAD. Selected architectures are then extrapolated and actually trained at 350M, 1B, and 3B scales (Pepe et al., 15 May 2026).
Hardware allocation is scale-dependent. The paper reports 8 H200 GPUs for 350M models and 16 H200 GPUs for 1B and 3B models. For the 1B fixed-token, or “isoToken,” comparison, the setup is: 71,565 steps, batch size 4, sequence length 8,192, data parallelism 16, and about 37.5B tokens, denoted “$2Mb + M + 11Mb + 2M$7 TPP.” The paper is explicit that, unlike the 2-primitive results, the 3-primitive isotoken results are reported from a single seed.
The isoFLOP study evaluates hybrids at 350M, 1B, and 3B under five FLOP budgets: $2Mb + M + 11Mb + 2M$8 The batch size in tokens per step is fixed at
$2Mb + M + 11Mb + 2M$9
and the shared sequence length is
0
The number of training steps is computed from the per-step cost according to
1
Pretraining and validation use a DCLM subset, with validation loss measured on 1000 samples of DCLM validation. Reported downstream zero-shot tasks are HellaSwag, WinoGrande, ARC-Easy, ARC-Challenge, PIQA, and SciQ. The paper also uses a DCLM Core score, defined as an average over 14 tasks with 0-shot, 3-shot, and 10-shot evaluations. Optimizer settings and hyperparameters are said to be kept consistent with prior Composer work, but the provided text does not specify the optimizer name, learning rate, warmup, or regularization parameters for AIRAhybrid pretraining.
5. Measured results, scaling behavior, and efficiency
At 1B scale under fixed-token pretraining, the paper evaluates 12 hybrid models: 8 AIRAhybrids, 3 baselines—Mamba, Nemotron-H approx., Nemotron-2 approx.—and 1 Composer-found hybrid. In the main comparison, AIRAhybrid-D (Stretched) achieves the best AIRAhybrid validation loss at 2.719 and the best AIRAhybrid average 0-shot accuracy at 60.5 raw and 62.2 normalized. AIRAhybrid-B (Stretched) achieves the best DCLM Core among AIRAhybrids with 49.1, while the highest DCLM Core overall in that hybrid table is 49.3 for Composer (2Mb-M-3A) (Pepe et al., 15 May 2026).
The paper also reports variant-specific strengths. AIRAhybrid-C (Stretched) is explicitly noted as showing localized strength on SciQ, where it reaches 90.2. Appendix-level results cited in the provided text attribute strong task-level values to B variants, including WinoGrande 59.9, XWinograd 82.6, HellaSwag 58.5, and PIQA 75.0. AIRAhybrid-E (Stretched) is reported to lead on OpenBookQA with 37.4 and to tie best on AGIEval with 23.0.
Relative to Llama 3.2, the paper’s headline summary states that AIRAhybrid-D improves accuracy by 3.8\% on downstream tasks. This is a central comparative claim in the abstract and summary discussion.
The isoFLOP analysis is the strongest distinct scaling claim for the family. The paper fits isoFLOP parabolas of validation loss versus model size and derives compute-optimal frontiers characterized by slope difference 2 and intercept offset 3, with lower values stated to be better. In this analysis, balanced Mamba/attention models such as Nemotron-2, AIRAhybrid-B (Stretched), and AIRAhybrid-D (Stretched) tend to achieve the lowest validation loss across sizes and FLOP budgets. However, AIRAhybrid-C (Stretched) has the steepest scaling slope, with
4
relative to Nemotron-2. The paper summarizes this by stating that AIRAhybrid-C scales 23\% faster than the modified Nemotron-2 and 37\% faster than Composer’s best hybrid. It also reports that AIRAhybrid-A (Stretched), the pure-SSM design, scales slightly better than Nemotron-2 with
5
Latency is treated separately from scaling-law fits. At 1B, the paper reports a latency–validation-loss Pareto analysis in which AIRAhybrid-D variants advance the Pareto frontier relative to Composer, achieving lower validation loss at strictly lower latencies than the Composer counterpart. Latency is estimated from block counts multiplied by GPU(H100)-profiled block latency and summed across all layers.
6. Interpretation, caveats, and significance
The paper presents AIRAhybrids as strong evidence that sparse agentic search can discover competitive hybrid schedules in a large combinatorial space. It repeatedly emphasizes that the agents examined only 0.0052\% of the 43-million-architecture search space, yet found hybrids that outperform Mamba-only, Nemotron-H approx., Nemotron-2 approx., and in some metrics Composer and Llama 3.2 (Pepe et al., 15 May 2026).
At the same time, the results are explicitly qualified. AIRAhybrids are not universally best on every metric. The paper notes that Composer’s hybrid baseline still has the top DCLM Core score in the main hybrid comparison, at 49.3, and that Nemotron-2 approx. remains very competitive. It also states that small-scale proxy evaluation does not always faithfully predict large-scale performance, identifying proxy-to-target mismatch as a limitation of AIRA-Compose.
Another caveat concerns the baselines. Nemotron-H and Nemotron-2 are described as approximated baselines, with smaller-scale versions and MoEs replaced by MLPs. Consequently, some comparisons are not to exact production Nemotron models. The paper also stresses that AIRAhybrids are novel at the architecture level, not at the primitive level; low-level mechanistic novelty belongs more to AIRA-Design than to AIRAhybrids themselves.
The paper’s own explanation for why these hybrids work remains high-level. It reports that attention-heavy architectures do better under fixed-token training, balanced architectures are more compute-efficient under isoFLOP constraints, and a mix of Mamba and attention can outperform either alone. A stronger mechanistic explanation is not provided. This suggests that the principal contribution of AIRAhybrids is architectural discovery under agentic search, rather than a new theoretical account of attention–SSM complementarity.
In practical terms, the paper identifies two especially important members of the family. AIRAhybrid-D (Stretched) is the strongest measured 1B fixed-token performer, while AIRAhybrid-C (Stretched) is the most notable asymptotic scaling result. Taken together, the family is presented as a concrete demonstration that agent-discovered Transformer–Mamba schedules can be competitive in downstream accuracy, validation loss, latency–loss trade-offs, and fitted scaling frontiers.