Agentic Discovery of Neural Architectures: AIRA-Compose 与 AIRA-Design
本研究探索利用 LLM 智能体 自主设计超越标准 Transformer 的基础模型,迈向递归自我改进。我们提出双框架方法:AIRA-Compose 负责高层架构搜索,AIRA-Design 负责低层机制实现。 AIRA-Compose 使用 11 个智能体在 24 小时预算内探索基本计算原语。智能体评估百万参数候选,并将最优设计外推至 350M、1B 和 3B 规模,产生两个系列的 14 种架构:AIRAformer(基于 Transformer)和 AIRAhybrid(Transformer-Mamba 混合)。在 1B 规模预训练后,这些模型一致超越 Llama 3.2 和 Composer 基线。在下游任务上,AIRAformer-D 和 AIRAhybrid-D 相比 Llama 3.2 准确率分别提升 2.4% 和 3.8%。此外,AIRA-Compose 发现的模型具有高效扩展前沿:AIRAformer-C 缩放速度比 Llama 3.2 和 Composer 最佳 Transformer 分别快 54% 和 71%,而 AIRAhybrid-C 比 Nemotron-2 快 23% 且比 Composer 最佳混合模型快 37%。 AIRA-Design 则指派 20 个智能体编写新颖的注意力机制以处理长程依赖,并生成高性能训练脚本。在 Long Range Arena 基准上,智能体设计的架构在文档匹配和文本分类任务上达到人类最优性能的 97.7% 和 97.4%。在 Autoresearch 基准上,Greedy Opus 4.5 在固定时间预算内取得 0.968 验证位·每字节,超越已发表的最小值。 综上所述,这些框架表明 AI 智能体能够自主发现架构和算法优化,匹配乃至超越手工设计的基线,为发现下一代基础模型建立了强大范式,标志着向递归自我改进迈出了明确一步。
论文精读
TL;DR LLM 智能体通过双框架自主设计新型基础模型架构,预训练后性能超越 Llama 3.2 及人工设计基线,迈向递归自我改进。
问题
问题背景
当前基础模型架构仍以 Transformer 为主,但其计算效率(尤其处理长序列时)和 scaling 行为渐显瓶颈。业界正积极寻求超越标准 Transformer 的新型架构,以降低训练与推理成本,并探索自动化的设计范式。
现有方法局限
- 手工设计依赖专家经验:如 Mamba、混合专家(MoE) 等新架构的提出,往往凝练多年领域知识,搜索空间受限,迭代缓慢。
- 神经架构搜索(NAS)泛化性不足:传统 NAS 多面向特定任务和小规模模型,搜索到的架构难以直接扩展至十亿参数级预训练;其评估方式(如基于代理任务或权重共享)常引入排名误差,在小尺度上的性能优劣不一定保持到 target scale。
- 搜索效率低下:组合空间(注意力、MLP、状态空间模型等原语)数量巨大,强化学习或演化算法需大量试错,计算成本高昂,且不具备跨尺度外推的稳定规律。
为什么这个问题难/重要
- 组合爆炸与评估瓶颈:基础运算原语的排列组合数随着原语种类和层数呈指数增长;每个候选需在百万参数规模完整训练验证,再外推到 1B+ 参数,过程耗时且排序不一致(rank inconsistency)风险高。
- 追求计算最优 scaling law:除了下游精度,业界更关注模型的 scaling frontier,即如何以更少 FLOPs 达到给定性能,这对控制大模型训练开销意义重大。例如,本文发现的架构可实现比 Llama 3.2 快 54% 的 scaling 速度。
- 迈向递归自我改进:若能利用 AI 代理自主设计下一代基础模型,则意味着模型进化可脱离人类带宽限制,形成自动加速的研发闭环,是迈向 AGI 的关键一步。
行业类比
如同 AutoML 在计算机视觉领域催生了 EfficientNet 等高效架构,将自动架构搜索引入大规模语言模型设计,有望从底层重塑基础模型效率,降低全行业对海量算力的依赖。
核心洞察
- 不同于传统**神经架构搜索 (NAS)** 依赖离散算子或梯度优化,**AIRA-Compose** 让一组 LLM 智能体在受限 24 小时算力下自主探索基元组合空间。 智能体利用先验知识生成候选并评估,避免了纯随机或进化搜索的低效,能在有限预算内找到可外推至 **1B 参数的强架构**。这证明 LLM 驱动的架构设计可以匹配甚至超越人工设计,且搜索出的架构家族(**AIRAformers** 和 **AIRAhybrids**)在下游任务上 **超过 Llama 3.2 基线 2.4–3.8% 准确率**,为自动化发现下一代基础模型提供了可扩展的范式。
- **AIRA-Compose** 与 **AIRA-Design** 的双框架设计分别从高层架构组合和低层机制实现两个维度推动模型创新,而多数现有工作只聚焦单一层面。 - 高层:AIRA-Compose 搜索出的 **AIRAformer-C** 和 **AIRAhybrid-C** 不仅在下游任务表现更好,其**计算最优扩展前沿** 比 **Llama 3.2** 快 54–71%,比 **Composer** 最佳模型快 37–71%,这对大模型的高效训练至关重要。 - 低层:AIRA-Design 通过 20 个智能体直接编写新注意力机制,在 **Long Range Arena** 基准上接近人类 SOTA(差距 2.3–2.6%),并在 **Autoresearch** 上以 `0.968 bpb` 超越已发布最佳结果。 这种联合优化展示了 AI 可以同时发明新算子并组装高效宏架构,向递归自我改进迈出了明确一步。
方法
输入与搜索空间定义
AIRA 系统基于预定义的计算原语(Attention、MLP、Mamba)及其组合方式构建高层架构搜索空间。对于低层机制设计,输入为任务描述和性能基准(如 Long Range Arena),要求 Agent 直接生成注意力机制代码和训练脚本。
关键模块:Agent 驱动的双框架设计
AIRA-Compose:高层架构搜索
- 多 Agent 协同:11 个 LLM Agent 并行提出候选架构,每个 Agent 独立生成模型配置(原语类型、层数、连接顺序等),并在小规模(百万参数)快速评估困惑度等指标,迭代优化设计。
- 两阶段外推:第一阶段在有限计算预算(24 小时)内筛选出最具潜力的架构,随后将 top 设计线性外推至 350M、1B、3B 参数规模,进行完整预训练验证,确保可扩展性。
- 效率-性能权衡:通过 IsoFLOP 预算方法绘制缩放曲线,识别计算最优前沿架构(AIRAformer-C 和 AIRAhybrid-C),其性能增速显著高于 Llama 3.2 和人工设计的混合模型。
AIRA-Design:低层机制创新
- 代码生成式设计:最多 20 个 Agent 直接编写新型注意力算子,针对长程依赖任务(如文档匹配、文本分类)进行定制化开发。Agent 通过执行训练脚本、分析验证指标(准确率、bits-per-byte)来自动调优。
- 基准验证:在 LRA 基准上,Agent 设计的注意力机制达到与人类 SOTA 接近的精度;在 Autoresearch 上,基于 Greedy Opus 4.5 优化的训练方案超越了已公布的最佳 bits-per-byte 记录。
输出:可部署的模型架构与算法
搜索过程最终产出 14 个新型架构家族(AIRAformers 和 AIRAhybrids),其中 AIRAformer-D 和 AIRAhybrid-D 在下游任务上平均准确率超越 Llama 3.2 达 2.4% 和 3.8%。AIRA-Compose 还提供了计算高效的缩放曲线,为后续训练资源分配提供参考。
与以往 NAS 方法不同,AIRA 不依赖固定的搜索算法(如强化学习或进化),而是利用 LLM Agent 的生成与推理能力,同时在两个抽象层次(宏观架构与微观算子)进行自主发现,实现了递归自我改进的初步雏形。
实验
实验设计
AIRA-Compose 使用 11 个 LLM agent 组成搜索团队,在 24 小时固定计算预算 下探索由 Attention、MLP、Mamba 三种计算原语组合而成的架构空间。搜索分两阶段:先评估百万参数级的候选模型,再对表现最优的设计外推到 350M、1B、3B 参数规模,并在等 token 预算下预训练。
AIRA-Design 则聚焦低层机制创新,部署 最多 20 个 agent 直接编写面向长程依赖的新型注意力机制及训练脚本,在 Long Range Arena (LRA) 和 Autoresearch 基准上验证。
关键发现
- 搜索得到 14 种新架构,分为 AIRAformer (纯 Transformer) 和 AIRAhybrid (Transformer-Mamba 混合) 两类。
- 在 1B 参数预训练后,AIRAformer-D 和 AIRAhybrid-D 下游准确率分别超越 Llama 3.2 基线 2.4% 和 3.8%。
- 缩放效率显著提升:AIRAformer-C 的损失-计算曲线斜率比 Llama 3.2 陡峭 54%,比 Composer 最优 Transformer 陡峭 71%;AIRAhybrid-C 则比 Nemotron-2 快 23%,比 Composer 最优混合模型快 37%。
- 在 LRA 的文档匹配和文本分类任务上,agent 设计的架构与人类 SOTA 的差距仅 2.3% 和 2.6%。
- Autoresearch 中,Greedy Opus 4.5 在固定时间内将验证 bits-per-byte 压至 0.968,超越此前发表的最优值。
基线对比解读
Agent 设计方法的优势源于 不受人类先验约束的全自动搜索。传统人工设计(如 Llama、Nemotron)依赖工程师对注意力机制或混合策略的直觉,而 AIRA 框架 通过组合爆炸的进化式探索,更易发现非直观的高效模式。
值得注意的是,混合架构的缩放增益 (AIRAhybrid-C) 显著超越同类人工基线 (Nemotron-2),说明 Mamba 与 Attention 的最优协同配比 很难手工调参获得。这提示工程实践上,针对大规模训练的前沿搜索,用 agent 替代部分人工试错能更高效地逼近算力最优边界。同时,AIRA-Design 在 LRA 上逼近 SOTA 表明 agent 已具备 直接生产可用算子代码 的能力,为递归自我改进打下基础。
行业影响
落地场景
AIRA-Compose 和 AIRA-Design 使 AI 代理能够自主设计基础模型架构与底层注意力机制,该能力可直接嵌入大模型研发管线、自动机器学习(AutoML)平台以及云原生模型服务中。
- 模型工厂:云厂商或 AI 实验室可将该流程作为“架构探索引擎”,针对不同业务场景(如长文本理解、多模态对齐)自动产出定制化骨干网络,代替人工试错。
- 垂直行业模型定制:在内容平台(视频理解、长文档摘要)和电商搜索/推荐等场景,代理可根据任务特性重设计算基元组合(Attention、MLP、Mamba),用固定计算预算快速收敛到高效架构。
- 边缘/端侧部署:通过代理搜索计算最优扩展前沿(如 AIRAformer-C 较 Llama 3.2 扩展速度提升 54%),可在给定延迟/内存约束下生成 Pareto 最优模型。
商业价值
核心理路是缩短从任务定义到高性能模型上线的时间成本。
- 降本:传统架构搜索依赖大规模人工实验,耗时数周至数月。AIRA-Compose 仅用 24 小时计算预算和 11 个代理即完成百万参数级搜索,并可靠外推至 1B–3B 规模,大幅降低人力与算力开销。
- 增收:在 1B 参数预训练后,AIRAformer-D 和 AIRAhybrid-D 下游任务准确率分别较 Llama 3.2 提升 2.4% 和 3.8%;更陡峭的扩展前沿意味着同等计算量下模型表现更强,能为 SaaS 或 API 服务提供性能溢价。
- 体验提升:AIRA-Design 在 Long Range Arena 上实现文档匹配准确率与人类最优差距仅 2.3%,直接优化端侧长文本理解体验;在 Autoresearch 任务中,Greedy Opus 4.5 优化训练至 0.968 bits-per-byte,超越己知基线,提升归一化模型质量。
与现有产品/工作流的接口
AIRA 框架并非替代现有训练栈,而是作为自动化架构生成层介入。
- 集成方式:可封装为 Kubernetes 原生 Job 或 ML Pipeline 组件(如 Kubeflow、Argo),接收任务描述与计算预算,输出候选架构配置文件(
model_config.yaml)及权重初始化方案。 - 现有工具链对接:代理搜索空间覆盖 PyTorch / JAX 生态常见算子,生成结果可直接导入 Hugging Face Transformers 或 vLLM 等推理框架;训练脚本产出为可直接运行的代码,与现有 CI/CD 流程兼容。
- 实际场景:
- 电商推荐系统:多模态推荐模型常需平衡序列建模与交叉特征交互,AIRA 代理可在 24 小时内搜索 Transformer-Mamba 混合架构,无缝替换现有 Base Model,再通过 1–2 轮微调上线,快速验证业务指标。
- 企业知识库问答:长文档问答依赖高效注意力机制,AIRA-Design 代理可自动生成针对 Long Range Arena 优化的新型注意力算子,作为插件插入现有 RAG 流水线,提升检索增强生成质量。
整体上,AIRA 展示的“代理自主发现优于人工设计的架构”范式,标志着基础模型研发正从手工匠时代迈向自动化工程时代,可被任何有模型定制需求的 AI 产品快速吸收。
局限
- 搜索空间由预定义的计算原语(**Attention**、**MLP**、**Mamba**)构成,智能体只能对其进行组合与调参,无法产生**超出原语集合的全新算子**。这使最终架构仍属 Transformer 或 Mamba 变体,限制了在需要非常规结构(如图数据处理或长程记忆)时的创新空间。
- 实验集中于 1B 参数规模,外推至 3B 的**缩放规律**仅基于少量数据点拟合,缺乏大规模验证。下游评估任务数量有限,未涉及**长上下文、多模态或推理密集型场景**,故方法在更大规模或更复杂应用中的泛化表现尚不明确,实际部署风险较高。
- 搜索过程在 **24 小时固定计算预算**下采用单一智能体编队,未探索更灵活的资源调度或与**进化算法、强化学习**等传统方法的融合,可能存在搜索效率瓶颈。同时,LLM 智能体自身的**推理偏差**(如对热门结构的偏好)未被量化分析,可能导致搜索陷入局部最优,影响最终发现的全局优度。