智能体驱动的神经架构发现：AIRA-Compose 与 AIRA-Design

用LLM智能体自主设计混合LLM架构，在1B规模上以原创的AIRAformer和AIRAhybrid结构超越Llama 3.2与Nemotron。

前置知识

Transformer架构与多头注意力(MHA)

Transformer是当前LLM的事实标准，由multi-head attention层与MLP层以1:1比例交替堆叠而成。其核心的多头注意力通过Q、K、V矩阵计算token间的成对相似度，复杂度为序列长度的二次方$O(n^2)$，在长序列推理时KV缓存成本高昂。理解这一点是看懂为什么需要混合架构的前提。

本文的AIRAformer-D Stretched等架构本质上是Transformer变体，Llama 3.2基线也是标准Transformer，必须理解Transformer才能识别出AIRAformer在Attention-MLP比例和排列方式上的创新。

状态空间模型(SSM/Mamba2)

Mamba2是基于状态空间理论的序列建模架构，具备线性复杂度$O(n)$和常数大小的状态缓存，擅长捕获长距离依赖且推理时内存占用低。它通过选择性扫描机制自适应地控制信息流动，常作为Attention的低复杂度替代品嵌入混合模型中。

AIRAhybrid家族的第三种基本原语Mb就是Mamba2 SSM。三种原语(M, mA, Mb)的排列组合是3.16亿种可能(3^16)，理解Mamba的特性才能明白为什么把它与Attention混合能取得性能提升。

神经架构搜索(NAS)与计算最优(Compute-Optimal)缩放

NAS通过自动化的方式在庞大的架构空间中寻找高性能模型，传统方法依赖强化学习、进化算法或贝叶斯优化。Compute-Optimal缩放指在给定FLOPs预算下找到验证损失最低的模型规模与数据组合，由Chinchilla等论文定义。IsoFLOP分析即固定FLOPs预算，扫描不同模型规模，找到损失最低点。

本文的核心问题就是NAS如何在混合LLM的组合爆炸空间中有效搜索，并且用IsoFLOP缩放定律来量化AIRAformer的扩展效率提升(54%和71% faster scaling)。

LLM智能体与树搜索Scaffold

LLM智能体是LLM加上一套算法脚手架(scaffold)的组合，能够自主调用工具、修改代码并迭代优化解。树搜索scaffold(如greedy或MCTS)在每一步生成多个候选解并保留最优者继续分支，常用算子包括Draft(生成)、Improve(优化)、Debug(纠错)、Analyze(分析)。

AIRA-Compose采用greedy GPT-5智能体在搜索树中生成上百个候选架构，理解scaffold如何运作才能看懂图3中每个节点的设计思路、提交物和评估分数。

研究动机

当前LLM几乎全部依赖Transformer，其自注意力的二次方复杂度和KV缓存的内存成本限制了长序列处理能力。为此，社区正快速转向hybrid LLM——将Attention、MLP、SSM等不同计算原语按特定模式交错排列。然而手工调优在组合爆炸的搜索空间面前已经力不从心：以3种原语、16层为例，可能的排列数高达$3^{16} \approx 4.3 \times 10^7$种，传统NAS方法如贝叶斯优化、强化学习或进化算法在LLM规模上要么开销过大，要么依赖预定义模板与手工变异算子，无法系统探索非显而易见的配置或发现全新的计算原语。与此同时，LLM驱动的智能体已能端到端完成数学发现、ML内核设计等复杂研究任务，但尚未被用于'设计驱动智能体本身的下一代架构'这一更具颠覆性的任务。

本文的目标是本文旨在系统评估LLM智能体自主设计下一代基础模型架构的能力，朝着递归自我改进(Recursive Self-Improvement, RSI)迈出具体一步。具体目标分为两条互补路径：(1)AIRA-Compose负责高层架构搜索，让智能体在Attention、MLP、Mamba三种预定义原语中自由组合并自动评估，找到比Llama 3.2和Composer-found基线更优的混合模式；(2)AIRA-Design负责底层机制设计，让智能体从头编写新型注意力机制和高效训练脚本，在Long Range Arena (LRA)和Autoresearch基准上逼近人类SOTA。最终输出14个新架构(AIRAformer和AIRAhybrid家族)以及12个AIRS-Bench标准化任务模板，可被任何LLM/scaffold组合复用。

与已有工作不同的是，本文的独特切入角度是把NAS从'优化算法驱动的搜索'重构为'具备领域知识的智能体驱动的假设检验'。传统Composer框架使用固定贝叶斯优化加增量层搜索，结构模板受限；而本文的智能体能够用自然语言为每个候选架构写出设计理由(如'mamba corridor followed by dual-attention cells')，结合对长程依赖、计算效率等的先验知识进行有意义的探索，并且能在搜索过程中自我调试、自我改进。此外，本文首次把LRA和Autoresearch这两个原本需要人工调优的复杂代码任务封装成AIRS-Bench任务，使智能体能直接生成可执行的model.py和train.py，触及'机制创新'这一比单纯组合原语更深层的自动化设计问题。

核心方法

本文的核心思路是用LLM智能体替代传统NAS的优化器，把'在巨大搜索空间中寻找最优架构'转化为'智能体根据提示词自主生成、评估、调试候选方案'。整体方法基于AIRS-Bench任务标准，将每个设计问题抽象为{问题、数据集、指标}三元组，让智能体提交一个产物(submission.csv或model.py)，由独立的evaluate.py脚本进行统一评估。在执行层面，所有智能体运行在AIRA-dojo这一基于树搜索的harness上，使用one-shot(单次提交)或greedy(从5个种子出发，逐层improve)两种scaffold，配以Draft/Improve/Debug/Analyze四个算子。整个工作流分为两条平行流水线：AIRA-Compose仅允许智能体从{MLP, Multi-head Attention, Mamba2}中挑选16层排列(2原语空间$2^{16}=65536$种，3原语空间$3^{16}\approx 4.3 \times 10^7$种)，在小规模代理数据集上评估后由后处理步骤聚合并外推到350M/1B/3B；AIRA-Design则允许智能体直接生成model.py或train.py代码，在LRA(长程依赖)和Autoresearch(GPT训练循环优化)基准上展开开放式探索。

核心创新点是把搜索过程中每个节点的'决策依据'显式化、可语言化，让搜索从黑盒优化走向可解释的假设检验。传统NAS的贝叶斯优化只关心fitness数值，本文智能体在每个节点都会写出设计理由(如'周期性Attention锚点+Mamba-Duo段+MLP门控的沙堡微块')，并在improve操作中把这些理由连同父节点的fitness一起作为上下文，生成新架构。这意味着智能体实际是把NAS作为多轮对话式的科研实验在执行——能基于历史教训(例如'上次OOM了，这次减小隐藏维度')进行因果推理，而不仅仅依赖变异算子。在AIRA-Design侧，关键洞察是把'底层机制设计'从手工架构创新降级为可被LLM代码生成能力直接攻克的代码补全任务：让智能体直接写JAX/Flax的CustomEncoder或nanochat的train.py，并由独立的Python脚本执行验证。两条流水线共享同一AIRS-Bench任务外壳，从而可以横向对比智能体在不同粒度上的能力。

方法步骤详情

整体方法分为六个步骤：(1)任务形式化：每个AIRA任务被建模为{problem, dataset, metric}三元组，外加project_description.md提示词、prepare.py数据准备脚本、evaluate.py独立评分脚本和metadata.yaml配置。(2)智能体配置：在AIRA-dojo harness中为每个任务选定scaffold(one-shot或greedy)和LLM(GPT-5、o3、Opus 4.6、Gemini 3 Pro等)，greedy scaffold首先生成5个种子候选，每步挑选fitness最高的节点继续improve。(3)AIRA-Compose提交：智能体生成一个16元素的字符串如'mA-M-Mb-M-mA-Mb-M-Mb-mA-M-Mb-Mb-M-mA-Mb-M'映射到预定义PyTorch块；每个候选在MAD/BabiStories/DCLM等小规模代理数据上训练并评估，10-60小时/seed的预算内通常覆盖100-200个架构。(4)聚合与外推：搜索结束后，使用Composer的layer-wise clustering把同一fitness阈值下出现频率最高的原语作为该位置代表，形成稳定的16层基础模式，然后通过stacking(整段重复)或stretching(按比例扩展连续块)放大到350M/1B/3B规模，在固定37.5B token预算下预训练并对比。(5)AIRA-Design提交：智能体直接写model.py(JAX/Flax的CustomEncoder)或train.py(修改nanochat的GPT脚本)；LRA任务在24 GPU小时内于3个文本任务(ListOps 0.51、Text 0.88、Retrieval 0.79)上训练评估，Autoresearch任务在5分钟wall-clock、单H200 GPU预算内跑完训练并以验证BPB为指标。(6)结果评估：AIRA-Compose用验证损失、6任务0-shot准确率和14任务DCLM Core分数三个指标衡量；AIRA-Design的LRA任务额外引入Valid Submission Rate(VSR)和归一化得分NS(对SOTA用march of 9s变换后归一到[0,1])。

技术新颖性

技术新颖性体现在四个层面：第一，把NAS的组合搜索问题完全重构为AIRS-Bench风格的智能体任务，使评估脚本、提交格式、训练管道解耦，这是与传统Composer框架最关键的范式差异；第二，提出共享的搜索外壳支持任意粒度，从高层原语排列到完整的Python模型/训练文件，覆盖整个设计栈；第三，用单一harness支持20种LLM+2种scaffold的可插拔组合，并通过greedy GPT-5等智能体在3原语搜索空间0.0052%的探查量下获得14个高性能架构(平均优于Llama 3.2)，证明语义驱动的探索可以大幅压缩组合爆炸；第四，引入march of 9s变换$\phi_t(s) = -\log_{10}(|s - s^{opt}_t|)$来归一化不同LRA任务的得分，使跨任务对比直观，NS>1代表超过人类SOTA。

**Figure 2**: The AIRA-Compose pipeline.

**Figure 3**: Greedy GPT-5 search (partial) on the 3-primitive NAS task within AIRA-Compose.

**Figure 4**: AIRA-Design: agent-driven, low-level mechanistic design.

实验结果

AIRA-Compose在1B规模、37.5B token固定预算下取得了清晰的性能优势。在2原语(M, mA)空间中，AIRAformer-D Stretched以2.734验证损失刷新全场最低，相对Llama 3.2(2.815)降低0.081；6任务平均归一化准确率达到60.8%，比Llama 3.2的58.4%高出2.4个百分点；14任务DCLM Core得分48.9%相比Llama 3.2的46.9%提升2.0分，AIRAformer-D Stretched在所有9个被对比模型中三项指标同时最佳。在3原语(M, mA, Mb)空间，AIRAhybrid-D Stretched验证损失低至2.719，6任务平均准确率60.5%(归一化62.2%)，相对近似Nemotron-2的61.5%高出0.7pp，相对Llama 3.2高出3.8pp。IsoFLOP缩放分析显示注意力占比高的AIRAformer-C比Llama 3.2快54%、比最佳Composer-found Transformer快71%；AIRAhybrid-C比近似Nemotron-2快23%、比最佳Composer-found hybrid快37%。Latency-Pareto前沿分析进一步证实AIRAhybrid-D Stretched在更低延迟下达到更低验证损失，扩展了Composer既有前沿。AIRA-Design方面，LRA基准上Greedy Gemini 3 Pro和Greedy Opus 4.6分别在ListOps(0.51)、Retrieval(0.79)和Text(0.88)上接近人类SOTA(0.64/0.82/0.91)，差距控制在2.3-2.6pp以内；Autoresearch上Greedy Opus 4.5 +Literature以BPB=0.9683超过原Autoresearch参考基线(0.9980)，验证集BPB绝对降低0.03，所有100次run中Opus 4.6 +Lit中位数为0.990，显著优于基线1.0121。综合来看，14个AIRAformer/AIRAhybrid架构全部优于或持平Llama 3.2基线，且缩放斜率更陡。

**Table 1**: The 12 RSI tasks introduced in this manuscript.

**Table 2**: Pretraining results of 2-primitive architectures at 1B scale with fixed token budget (37.5B tokens).

**Table 3**: Pretraining results of 3-primitive hybrid architectures at 1B scale with fixed token budget (37.5B tokens).

**Table 4**: Average Normalized Score per task.

**Table 5**: Improvement steps for the best seed of Greedy Opus 4.5, Opus 4.6, and the two +Literature variants on the Autoresearch task.

**Figure 1**: AIRA-Compose and AIRA-Design: agentic frameworks for neural architecture search and model design.

**Figure 5**: IsoFLOP scaling curves and optimal frontier comparison (M, mA).

**Figure 6**: IsoFLOP scaling curves and optimal frontier comparison (M, mA, Mb).

**Figure 7**: Latency–validation loss Pareto frontiers at the 1B-parameter scale.

**Figure 8**: Mean best test accuracy achieved by the 8 greedy agents across the 6 low-level mechanistic design tasks.

**Figure 12**: Autoresearch: distribution of best validation BPB per seed.

**Figure 13**: Autoresearch: best-seed BPB progression for Greedy Opus 4.5 and Opus 4.6, with and without literature.

查看结构化数据

任务	指标	本文	基线	提升
1B模型预训练-验证损失(2原语空间)	Validation Loss(越低越好)	AIRAformer-D Stretched = 2.734	Llama 3.2 = 2.815, Composite Stretched = 2.759	比Llama 3.2降低0.081，相对降幅2.9%；比Composite Stretched降低0.025
1B模型-6任务0-shot平均归一化准确率(2原语空间)	Normalized Accuracy(%)	AIRAformer-D Stretched = 60.8%	Llama 3.2 = 58.4%, Composite Stretched = 59.8%	比Llama 3.2提升2.4pp，比Composite Stretched提升1.0pp
1B模型-14任务DCLM Core分数(2原语空间)	DCLM Core Score(%)	AIRAformer-D Stretched = 48.9%	Llama 3.2 = 46.9%, Composite Stretched = 47.3%	比Llama 3.2提升2.0pp，比Composite Stretched提升1.6pp
1B模型预训练-验证损失(3原语空间)	Validation Loss(越低越好)	AIRAhybrid-D Stretched = 2.719	近似Nemotron-2 = 2.732, Mamba = 2.771	比Nemotron-2降低0.013，比Mamba降低0.052
1B模型-6任务0-shot平均归一化准确率(3原语空间)	Normalized Accuracy(%)	AIRAhybrid-D Stretched = 62.2%	近似Nemotron-2 = 61.5%, Mamba = 60.9%	比Nemotron-2提升0.7pp，比Mamba提升1.3pp；AIRAhybrid-A归一化准确率62.1%同样优于Nemotron-2
IsoFLOP缩放前沿(2原语空间)	缩放斜率差Δq	AIRAformer-C Stretched Δq = -0.34	Llama 3.2 q = -0.63(参考线)	相对Llama 3.2 scaling law快54%；相对最佳Composer-found Transformer快71%
IsoFLOP缩放前沿(3原语空间)	缩放斜率差Δq	AIRAhybrid-C Stretched Δq = -0.14	近似Nemotron-2 q = -0.59(参考线)	比Nemotron-2快23%；比最佳Composer-found hybrid快37%
LRA-ListOps(Non-Config.)	Best Test Accuracy(越高越好)	Greedy Gemini 3 Pro = 0.51	人类SOTA = 0.64	距人类SOTA 13pp差距(0.51 vs 0.64)，是智能体最优成绩
LRA-Text IMDB(Non-Config.)	Best Test Accuracy(越高越好)	Greedy Gemini 3 Pro = 0.88	人类SOTA = 0.91	距人类SOTA仅2.6pp
LRA-Retrieval AAN(Non-Config.)	Best Test Accuracy(越高越好)	Greedy Opus 4.6 = 0.79	人类SOTA = 0.82	距人类SOTA仅2.3pp
Autoresearch-验证BPB	Bits per Byte(越低越好)	Greedy Opus 4.5 +Literature = 0.9683	Autoresearch published reference = 0.9980; RADv1 baseline = 1.0121	绝对降低0.03，相对降幅约3.0%，超过原Autoresearch发表基线

局限与改进

作者明确承认几个关键局限：第一，AIRA-Compose完全依赖小规模代理评估(MAD/BabiStories/DCLM子集)，小规模排名并不能完全预测大规模表现，proxy-target gap仍是未知数，且16层小规模到1B规模的stacking/stretching仍需手工调优；第二，AIRA-Design的智能体在LRA上的'设计'几乎都是对Performer、Longformer、Conformer等已有高效注意力机制的重组和适配，没有出现真正新的理论洞见——NS分数普遍低于1.0(最高也只到0.257)说明与人类SOTA仍有显著差距；第三，one-shot智能体在LRA任务上VSR为0%，完全无法产出可用提交，说明单轮代码生成不足以完成底层机制设计，必须依赖多轮improve+debug；第四，可配置(Configurable)LRA变体对弱智能体(如CWM、o3-mini、Devstral)反而是灾难(回归11-32pp)，说明当前智能体不会有效利用更大的超参数搜索空间；第五，Autoresearch中Literature访问对GPT-5和Opus 4.5反而是负优化，只有Opus 4.6能稳定从文献中获益，说明智能体整合先验知识仍不成熟；第六，greedy scaffold的整文件重写范式让每次improve引入多重修改，因果归因困难。个人额外观察：搜索预算24小时/H200 GPU/seed相对Composer原论文的固定贝叶斯搜索并不便宜，要复现本文结果需相当算力投入(1B规模16卡H200 x 37.5B token预训练本身就要数千卡时)。

独立分析的弱点

独立分析三个关键弱点：(1)小规模代理评估的可靠性仍未被严格验证——本文报告AIRAformer-D在MAD 2-primitive ranking中并非第一(0.733 accuracy落后若干架构)，但放大到1B后却表现最好，说明小规模排序与大规模性能之间存在不完美对应；改进方向应是在多数据集(DCLM+一个代码语料+一个多语言语料)上做交叉验证，并显式建模proxy-target之间的标定函数。(2)Greedy智能体的improve算子接受整段历史作为上下文，导致每次修改叠加多个变量(Autoresearch的step 1常常同时改depth、aspect_ratio、head_dim、window_pattern)，后续分析无法归因；改进方向是引入差异化的候选生成(如固定其他变量、只让一个变量变化)或专门的ablation harness。(3)AIRA-Design的可配置变体反而让弱智能体崩溃，说明当前LLM在面对增大搜索空间时缺乏分而治之策略；改进方向是在scaffold中加入显式的hyperparameter budget控制(例如每轮最多调1-2个变量)，或基于VSR自适应收窄搜索空间。这些弱点并非不可克服，但需要在harness层面而非单纯更换LLM来改进。

未来方向

作者明确指出四条未来路线：(1)将AIRA-Compose部署到支持多GPU的AIRA2 harness，从而把proxy规模从百万参数推到数十亿参数，直接缩小proxy-target gap；(2)把聚合与外推步骤也变为agentic任务，让智能体自主选择stacking vs stretching并决策每层宽度，而非用固定规则；(3)在LRA任务中引入tool-use和持久状态能力以支持交互式调试，并扩展到更长程基准(支持2M token的现代SOTA模型)；(4)开发支持增量、可解释代码修改的scaffold，替代greedy的全文件重写。基于成果可延伸的方向还包括：把AIRA-Compose从原语排列扩展到包含Normalization选择、激活函数、位置编码等更丰富设计维度；用AIRA-Design去发现混合LLM中Mamba与Attention的最优比例分配；在Autoresearch中引入更多样化的智能体群(异质化LLM)以加速收敛；以及最关键的——把这些能力组合起来，让智能体同时设计自己的架构和自己的训练循环，朝递归自我改进迈出下一步。

复现评估

复现评估中等偏上。代码方面，AIRA-Compose/Design任务都基于AIRS-Bench和AIRA-dojo这两个Meta的开源框架搭建，任务文件结构(project_description.md、prepare.py、evaluate.py)标准化，理论上可以替换不同LLM/scaffold复现。论文在附录D-F提供了详细的聚合策略、模型配置表(Table 7-12)和文献资源清单(Table 16-18)，附录H甚至给出了最低BPB模型train.py的完整实现。数据方面使用了DCLM、BabiStories、MAD、IMDB、ListOps、AAN、ClimbMix等公开数据集，tokenizer也是Karpathy nanochat的BPE(8192 vocab)。算力是主要门槛：AIRA-Compose大规模实验需要16卡H200 x 37.5B token(约数千GPU小时)来对比9-12个模型，加上340次24小时search run和300次60小时search run，总算力消耗极大；AIRA-Design侧相对便宜(LRA每智能体每任务在H200上24 GPU小时，Autoresearch每seed 5分钟)。复现难度主要不在算法正确性而在算力可承受性——若仅复现小规模search结果，相对容易；若要复现1B规模扩展，则需要专业集群支持。

智能体驱动的神经架构发现：AIRA-Compose 与 AIRA-Design Agentic Discovery of Neural Architectures: AIRA-Compose and AIRA-Design