MARS:基于反思搜索的模块化智能体用于自动化AI研究 MARS: Modular Agent with Reflective Search for Automated AI Research
预算感知MCTS+模块化构建+比较反思记忆,自动化MLE研究
前置知识
蒙特卡洛树搜索 (MCTS)
MCTS是一种用于在大规模搜索空间中寻找最优解的启发式搜索算法,广泛应用于博弈论和规划问题。它通过四个迭代阶段运作:选择(Selection)使用UCT公式在已探索节点中选择最有潜力的子节点;扩展(Expansion)在选定节点添加新的子节点;模拟(Simulation)从新节点进行随机或启发式评估;回传(Backpropagation)将评估结果沿路径向上更新。MCTS的核心优势在于它能平衡探索(exploration)与利用(exploitation),通过上置信界(Upper Confidence Bound)公式 $UCT = \bar{X}_j + c\sqrt{\frac{\ln N}{n_j}}$ 在高收益节点和未充分探索节点间做出权衡。在代码生成场景中,MCTS可用于探索不同的代码实现策略。
MARS将MCTS作为核心搜索框架,并创新性地引入预算感知机制,理解标准MCTS的工作原理是理解本文技术贡献的前提
机器学习工程 (MLE)
MLE指构建完整机器学习流水线的工程实践,包括数据加载、特征工程、模型架构设计、训练循环、验证和推理等环节。与传统软件工程不同,MLE具有三个独特挑战:评估计算成本高昂(模型训练可能需要数小时甚至数天);性能归因不透明(很难确定哪个改动带来了性能提升);需要同时处理代码正确性和模型性能。MLE-Bench是评估智能体MLE能力的标准基准,包含75个Kaggle竞赛任务,覆盖NLP、CV和表格数据分析。
MLE是本文要解决的核心问题场景,论文的所有方法设计都是为了应对MLE的特殊挑战
模块化代码生成
模块化代码生成是指将复杂的代码任务分解为多个独立、可测试的功能模块,每个模块封装特定子任务(如数据预处理、模型定义、训练引擎),通过主编排脚本协调各模块工作。与生成单一大型脚本的单体化(monolithic)方法相比,模块化方法具有多项优势:绕过token输出长度限制;减少上下文噪音提高精度;支持模块级缓存和复用;便于定位调试。这与专业软件工程的工作流程一致,是工程化代码生成的重要范式。
MARS用模块化范式替代了现有智能体的单体脚本生成方式,这是其核心架构创新之一
信用分配问题 (Credit Assignment Problem)
信用分配问题是强化学习和迭代优化中的经典难题:当一个实验获得更好结果时,很难确定是哪个具体改动(而非其他因素)导致了性能提升。在MLE场景中,这一问题尤为突出,因为代码改动、数据处理变化、随机种子等因素都会影响最终指标。现有基于记忆的智能体缺乏有效机制来隔离因果因素,往往从嘈杂的日志中过度泛化,无法从过去的试验中有效学习。
MARS通过比较反思记忆直接解决信用分配问题,这是论文最核心的技术贡献
上置信界 (UCT/UCB)
UCT(Upper Confidence bound applied to Trees)是MCTS中用于节点选择的标准算法。它为每个节点维护一个价值估计,综合考虑该节点的历史平均收益(利用项)和访问次数相对于父节点的稀有程度(探索项)。具体公式为 $UCT(v) = \bar{R}(v) + c\sqrt{\frac{\ln N(parent)}{N(v)}}$,其中 $\bar{R}(v)$ 是节点 $v$ 的平均奖励,$c$ 是探索系数,$N$ 是访问次数。这确保了算法不会过早收敛到局部最优,同时也不会在明显劣质的分支上浪费过多计算资源。
MARS使用UCT作为节点选择策略的基础,并在此基础上加入了效率引导的奖励函数
研究动机
当前基于LLM的智能体在处理复杂机器学习工程(MLE)任务时面临三大核心瓶颈。首先是成本忽视问题:现有框架将问题求解视为纯代码挑战,忽略研究的经济现实。例如,一个智能体可能将模型准确率提升0.1%,但将训练时间从1小时增加到10小时,这在实践中毫无用处,但标准搜索算法会优先选择它。其次是单体化脆弱性:现有智能体生成的庞大单文件脚本缺乏模块性,难以应对研究代码库的复杂性——数据加载、模型架构和训练循环需要无缝交互,单体脚本无法有效管理这种复杂性。最后是信用分配失败:研究进展是迭代且不透明的,当新实验产生更好结果时,很难隔离因果因素。标准记忆机制(如Reflexion的自我纠正或简单的历史记录)无法解决这一问题,导致智能体无法从过去试验中有效学习。实验证据表明,完全移除记忆机制会导致性能急剧下降,说明现有记忆架构虽有必要但远非充分。
本文的目标是MARS旨在构建一个专门为自动化AI研究优化的框架,能够自主完成复杂的MLE任务。具体目标包括三个方面:第一,在固定计算预算内最大化任务性能,实现性能与效率的平衡;第二,生成模块化、可维护、可测试的代码库,而非脆弱的单体脚本;第三,通过有效的反思机制从过去的成功和失败中提取高质量洞察,解决信用分配问题。最终,MARS希望在MLE-Bench基准上达到开源框架的最优水平,并与全球排行榜上的顶级方法保持竞争力。
与已有工作不同的是,MARS的独特切入角度在于它将三个正交的创新点有机结合,形成了一个协同工作的框架。与现有方法的本质区别在于:(1)在搜索策略层面,现有方法要么使用贪心搜索(如AIDE),要么使用标准MCTS,都只优化任务性能而忽略计算成本。MARS引入预算感知MCTS,将效率引导的奖励函数直接集成到搜索树中,显式平衡性能与执行开销。(2)在代码生成范式层面,现有系统全部采用单体脚本范式,MARS强制执行仓库级模块化范式,通过设计-分解-实现流水线管理复杂性。(3)在记忆机制层面,现有方法要么简单记录所有历史(如AIDE),要么只进行经验分析(如AIRA的Scoped Memory)。MARS提出比较反思记忆,通过分析当前解与已知最优解之间的差异来隔离因果因素,这类似于人类工程师进行消融实验的思维方式。如论文Table 1所示,MARS是唯一同时具备这三个特性的框架。
核心方法
MARS将长期AI研究问题形式化为在约束预算 $B$ 下寻找最优软件仓库的搜索问题。给定问题实例 $P = (I, E, O)$,其中 $I$ 是自然语言指令,$E$ 是环境(如数据集),$O$ 是目标函数(如准确率),目标是找到解 $s^* = \arg\max_s O(s, E)$,满足 $Cost(s) \leq B$。MARS的整体框架包含三个协同模块:(A)资源感知规划模块使用预算感知MCTS在解空间中进行战略性导航,从起草新架构、调试运行时错误、改迚有效解三种动作中选择;(B)模块化分解模块用设计-分解-实现流水线将方案架构化为独立、可测试的模块,支持基于差异的精确修改;(C)反思记忆模块将原始执行日志提炼为结构化的调试课程和解决方案课程,主动防止错误重复并加速后续迭代的收敛。这三个模块在MARS主循环中交替执行,形成一个迭代进化的解决方案搜索过程。
MARS的核心创新是比较反思记忆(Comparative Reflective Memory)机制,它直接解决了信用分配问题。与现有方法的本质区别在于:标准记忆机制(如Reflexion)只从错误中学习(自我纠正),或者简单记录执行日志供未来参考。MARS的比较反思则通过分析当前解与先前最佳已知解之间的代码差异(delta),精确定位驱动性能变化的具体算法改动。这一过程分为两个阶段:首先,经验分析智能体(Empirical Analysis Agent)审查执行日志提取客观发现(如指标趋势);然后,课程蒸馏智能体(Lesson Distillation Agent)执行比较反思,分析当前解与最佳解的差异,产出包含(1)隔离的因果改动、(2)比较影响分析、(3)通用规则的结构化课程。实验数据表明,63%的被利用课程来自跨分支转移,说明智能体能有效跨搜索路径泛化洞察。人工评估确认90%的课程因果准确率,LLM审计显示88.34%的课程正确归因了性能变化。
方法步骤详情
MARS的执行流程可分为三个阶段。第一阶段是任务准备:多智能体系统提取任务元数据,形式化优化目标,准备训练、验证和测试数据集;探索性数据分析(EDA)智能体生成报告指导后续特征工程。第二阶段是MARS主循环,迭代执行以下步骤:(1)MCTS选择阶段从根节点开始,使用UCT算法递归选择子节点,直到找到未完全展开的候选节点;(2)扩展阶段根据节点状态执行三种操作之一——对新问题起草(Drafting)生成全新解,对有效节点执行改进(Improvement),对失败节点执行调试(Debugging);(3)模块化实现阶段,Idea Agent生成自然语言计划,Modular Agent将其分解为逻辑独立的功能模块,Coding Agent顺序实现每个模块并通过验证脚本调试,最后通过主编排脚本串联所有模块;(4)反思记忆阶段,对有效解执行比较反思提取解决方案课程,对失败执行调试分析提取调试课程,Review Agent评估新课程过滤冗余;(5)回传阶段,使用效率引导奖励函数 $R(v) = G(v) \cdot [t(v)/L(v)]^w$ 更新节点价值。第三阶段输出得分最高的仓库作为最终解。关键超参数包括:记忆课程上限 $K_m = 30$,单次失败最大调试次数 $N_d = 10$,有效节点分支因子 $N_i = 2$,惩罚权重 $w = -0.07$。
技术新颖性
MARS在技术层面有四项新颖贡献。第一是预算感知MCTS的效率引导奖励函数: $$R(v) := G(v) \cdot [t(v)/L(v)]^w$$ 其中 $G(v)$ 是全局归一化性能分数,$t(v)$ 是执行时间,$L(v)$ 是时间限制,$w$ 是惩罚权重。当 $w=0$ 时退化为标准MCTS,$w<0$ 时惩罚高耗时方案。实验表明最优 $w=-0.07$,移除惩罚($w=0$)导致性能下降,过强惩罚($w=-0.15$)则过度偏向简单快速节点。第二是模块化仓库范式的形式化定义:解 $s_n = \langle\{M_j\}^l_{j=1}, \pi_{main}\rangle$,由 $l$ 个独立模块和一个编排脚本组成。第三是比较反思的两阶段蒸馏流程,通过代码差异分析而非日志摘要来隔离因果因素。第四是差异化编辑(Diff-Based Editing)机制,代码修改以标准化差异格式指定目标文件、替换块和新代码,支持单次推理完成多文件原子更新,避免了全仓库重新生成的浪费。这些创新相互增强:模块化使得差异编辑成为可能,差异分析使得比较反思更精确,预算感知搜索则确保整个过程在资源约束下高效运行。
实验结果
MARS在MLE-Bench基准上的75个Kaggle竞赛任务中进行了全面评估,采用严格控制的实验环境(单块NVIDIA A100 40GB GPU、12 vCPU、220GB RAM、1TB SSD,24小时时间预算)。在受控评估中,MARS与Gemini-3-Pro-Preview配合使用时,在所有开源框架中取得最优成绩:提交率98.7%、高于中位数率65.8%、铜牌率9.3%、银牌率15.6%、金牌率31.1%、任何奖牌率56.0%。其中金牌率31.1%是所有报告智能体中的最高值。扩展版本MARS+(使用2倍H100 GPU和48 vCPU运行两棵并行搜索树)进一步提升了表现:高于中位数率74.2%、任何奖牌率62.7%、金牌率33.8%,超越了包括ML-Master 2.0在内的强力竞争对手。消融研究验证了每个组件的必要性:移除模块化分解导致整体成功率下降;完全移除记忆机制导致性能急剧下降;比较组件相比仅经验分析提供了一致的性能提升。树搜索策略比较显示预算感知MCTS随时间推移持续优于贪心搜索和标准MCTS($w=0$)。有效解率方面,预算感知MCTS达到19.5%±1.5%,高于标准MCTS的16.1%±1.3%。课程利用率达65.8%±1.1%,跨分支转移率达63.0%±1.8%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MLE-Bench (Gemini-2.5-Pro, 受控评估) | Any Medal Rate (%) | 43.1 ± 1.6 | AIDE 23.1 ± 0.4 / AIRA-dojo 24.4 ± 1.2 | 相对AIDE提升86.6%,相对AIRA-dojo提升76.6% |
| MLE-Bench (Gemini-3-Pro-Preview, 受控评估) | Gold Medal Rate (%) | 31.1 ± 0.4 | AIDE 16.4 ± 1.8 / AIRA-dojo 24.0 ± 1.5 | 相对AIDE提升89.6%,相对AIRA-dojo提升29.6% |
| MLE-Bench (MARS+, Gemini-3-Pro-Preview) | Above Median Rate (%) | 74.2 ± 0.9 | ML-Master 2.0 63.1 ± 1.2 | 相对ML-Master 2.0提升17.6% |
| MLE-Bench Lite (Gemini-3-Pro-Preview) | Any Medal Rate (%) | 74.2 ± 1.5 | AIRA-dojo 56.1 ± 1.5 / AIDE 53.0 ± 6.1 | 相对AIRA-dojo提升32.3%,相对AIDE提升40.0% |
| MLE-Bench Medium (Gemini-3-Pro-Preview) | Any Medal Rate (%) | 52.6 ± 3.0 | AIRA-dojo 29.8 ± 3.8 / AIDE 26.3 ± 3.0 | 相对AIRA-dojo提升76.5%,相对AIDE提升99.6% |
| MLE-Bench High (Gemini-3-Pro-Preview) | Any Medal Rate (%) | 37.8 ± 2.2 | AIRA-dojo 31.1 ± 4.4 / AIDE 17.8 ± 2.2 | 相对AIRA-dojo提升21.5%,相对AIDE提升112.4% |
| 成本效益分析 (4小时MARS) | Any Medal Rate (%) | 28.4 | AIRA-dojo 24小时版本 24.4 | 用1/6时间($9.6 vs $39.0)超越24小时基线 |
局限与改进
论文承认的局限性包括:(1)成本问题:24小时版MARS的成本为$60.5,高于AIRA-dojo的$39.0,虽然性能提升显著(任何奖牌率从24.4%提升到43.1%),但对于预算敏感的用户仍需考虑;(2)任务范围:评估仅限于MLE-Bench的Kaggle竞赛任务,尚未扩展到更广泛的科学发现领域;(3)课程质量依赖LLM能力:比较反思的准确性依赖底层LLM的推理能力,不同LLM可能产生不同质量的课程。作者自己的观察还包括:(4)搜索空间效率:虽然预算感知MCTS提高了有效解率(19.5% vs 16.1%),但仍有约80%的探索未能改进,说明搜索效率仍有提升空间;(5)模块化粒度固定:当前的模块分解由智能体自动决定,没有机制根据任务特性动态调整模块粒度;(6)并行化有限:MARS+通过运行两棵并行搜索树来扩展,但搜索树内部仍是串行的,更细粒度的并行化可能带来更大收益。
独立分析的弱点
MARS存在几个值得深入分析的弱点。首先是记忆容量限制:课程池大小固定为 $K_m = 30$,这可能在长期任务中丢失有价值的早期课程。改进方向可以是引入课程重要性评分机制,保留高价值课程而非简单的时间窗口。其次是模块化分解的黑箱性:Idea Agent和Modular Agent的分解决策缺乏可解释性和可控性,用户无法干预或指导分解策略。可引入人机协作机制,允许用户指定模块边界或约束。第三是比较反思的时间开销:两阶段蒸馏(经验分析+比较反思)增加了每轮迭代的延迟。可探索异步课程蒸馏或增量更新机制。第四是调试循环的固定上限:最多 $N_d = 10$ 次调试尝试可能不足以解决某些复杂错误,也可能对简单错误过度分配资源。动态调整调试预算(如基于错误类型的启发式)可能更高效。最后是奖励函数的线性假设:效率引导奖励假设性能和时间成本之间是简单的幂律关系,更复杂的非线性建模可能更好地捕捉实际权衡。
未来方向
论文作者提出未来工作将聚焦两个方向:(1)将MARS扩展到更广泛的科学发现领域,如药物发现、材料科学等,这需要适配不同的环境接口和评估指标;(2)通过高级上下文缓存和提前停止机制优化框架的经济可行性。基于现有成果可延伸的方向包括:(3)自适应课程管理:根据任务进展动态调整课程池大小和检索策略;(4)多智能体协作:将MARS的模块化范式扩展为真正的多智能体系统,不同专门化智能体并行处理不同模块;(5)元学习搜索策略:从历史任务中学习最优的搜索超参数(如 $w$、$N_i$、$N_d$),实现跨任务的策略迁移;(6)连续学习:在多个相关任务间积累和迁移课程知识,构建跨项目的通用MLE最佳实践库;(7)人机协作模式:允许人类专家在关键决策点介入,如审核模块分解方案或验证蒸馏课程。
复现评估
MARS在复现方面提供了较好的支持。代码、轨迹和提示词已在 https://github.com/jfc43/MARS 开源发布。实验环境详细说明:标准节点配置为单块NVIDIA A100 40GB GPU、12 vCPU、220GB RAM、1TB SSD,24小时时间预算。底层LLM使用Gemini-2.5-Pro和Gemini-3-Pro-Preview两种模型。评估遵循标准MLE-Bench协议,结果报告为三次独立运行的均值±标准误。关键超参数已明确:$K_m=30$、$N_d=10$、$N_i=2$、$w=-0.07$。复现的主要挑战包括:(1)LLM API成本:完整评估需要大量API调用,24小时版成本约$60.5/竞赛;(2)计算资源:需要A100级GPU和充足内存;(3)MLE-Bench数据集获取:需要Kaggle账号和数据集下载许可;(4)随机性:虽然报告了三次运行的统计量,但LLM的非确定性可能导致不同环境下的结果差异。总体而言,复现难度中等偏高,主要受限于计算资源和API成本。
论文图表
该图展示了MARS在iMet-2020-FGVC7任务上的验证性能演进轨迹。图中以节点ID标注了关键的策略课程触发点,显示智能体如何从轻量级残差网络逐步演进到模型集成技术。现有方法无法达到奖牌级性能,而MARS通过持续的策略细化最终获得银牌。图中标注了多个关键节点(ID 1-19),每个节点代表一次解决方案迭代,箭头指示了性能提升的触发路径。
这张图直观展示了MARS的核心价值主张——通过反思学习实现顿悟时刻,是理解论文动机和贡献的关键视觉锚点