MARS: Modular Agent with Reflective Search for Automated AI Research¶

会议: ICML 2026
arXiv: 2602.02660
代码: https://github.com/jfc43/MARS (有)
领域: LLM Agent / 自动化机器学习 (AutoML) / MLE Agent
关键词: MLE-Bench, 模块化 Agent, Budget-Aware MCTS, Comparative Reflective Memory, Lesson Learning

一句话总结¶

MARS 把自动化 AI 研究重构成"在软件仓库空间中搜索最优解"的问题，用 预算感知 MCTS + 模块化"设计-分解-实现"流水线 + 比较式反思记忆 三根支柱，在 MLE-Bench 上拿到开源框架 SOTA，金牌率 31.1%（Gemini-3-Pro-Preview），并出现 63% 的跨分支课程迁移这种"Aha! moment"。

研究背景与动机¶

领域现状：LLM agent 在通用软件工程上已经很强（修 GitHub issue、写测试），最近一批工作（AIDE、AIRA、R&D-Agent、ML-Master、InternAgent 等）开始把它们用到自动化 AI 研究的核心瓶颈 —— Machine Learning Engineering (MLE) 任务上，并在 OpenAI 的 MLE-Bench（75 个 Kaggle 竞赛 × 24h × 1×A100）上比拼金/银/铜牌率。

现有痛点：作者总结现有 MLE agent 三个结构性缺陷：

忽略执行成本：当前搜索（greedy、vanilla MCTS、evolutionary）只优化任务性能，不管 wall-clock。一个把准确率提 0.1%、把训练从 1h 拉到 10h 的方案在 24h 预算里是灾难，但标准 UCT 还是会偏向它。
monolithic 脚本脆弱：现有 agent 多数生成一个大 Python 文件，token 上限挤压逻辑、改一处要重写全部、debug 定位困难，无法承载真实研究 repo 那种"数据-模型-训练循环"多模块耦合。
memory 解决不了 credit assignment：实验结果变好时，到底是哪一行代码改动起的作用？Reflexion/MemGPT 这类 verbal reflection 或轨迹缓存只能"记住做过什么"，无法定位因果。

核心矛盾：MLE ≠ 通用编程。前者是"昂贵评估 + 不透明归因 + 高架构复杂度"的概率性长程任务，需要的是带预算的策略式搜索而不是更聪明的单脚本生成器。

本文目标：设计一个 agent scaffolding，同时回答三个子问题 ——（1）怎么在搜索时显式 trade-off 性能与成本；（2）怎么让 agent 产出 repo-level 而非脚本的解；（3）怎么把"成功 vs 失败"的差异蒸馏成可迁移的因果经验。

切入角度：把 MLE 形式化为 \(s^* = \arg\max_s \mathcal{O}(s, \mathcal{E})\) s.t. \(\text{Cost}(s) \le B\)，将解空间从"所有可能 Python 程序"重新定义为"所有可能模块化仓库 \(s_n = \langle \{\mathcal{M}_j\}_{j=1}^{l}, \pi_{\text{main}}\rangle\)"，让搜索、记忆、奖励函数都围绕这个 repo-level 表示设计。

核心 idea：用 预算感知 MCTS 在 repo 空间内搜索，用 Design-Decompose-Implement 替代单脚本生成，用 比较式反思记忆（对比当前解和 best-known 解的 diff）解决 credit assignment —— 三者协同形成 long-horizon "Aha! moment"。

方法详解¶

整体框架¶

MARS 是一个迭代 loop。输入是 MLE 任务三元组 \(\mathcal{P} = (\mathcal{I}, \mathcal{E}, \mathcal{O})\)（指令 / 环境 / 目标）和预算 \(B\)；输出是一个在 \(B\) 内最大化 \(\mathcal{O}\) 的模块化仓库。

流程分两阶段： 1. Task Preparation：多 agent 系统提取任务元数据，做 Exploratory Data Analysis（EDA），生成报告指导后续特征工程，并准备 train/val/test 划分。 2. MARS Loop：在 MCTS 树上反复迭代三个协同模块 —— Module A（Resource-Aware Planning，决定下一步动作）→ Module B（Modular Decomposition，生成/修改子模块）→ Module C（Reflective Memory，从 trajectory 蒸馏 lesson 喂回 Module B）。最终输出树上 score 最高的 leaf 对应的仓库。

每个 MCTS 节点表示一份候选解 \(s_n = \langle \{\mathcal{M}_j\}_{j=1}^{l}, \pi_{\text{main}}\rangle\)，三种 expansion 动作：Drafting（root 处从零写新解）、Improvement（在 valid 节点上改模块）、Debugging（在 buggy 节点上修 runtime error，最多 \(N_d=10\) 次 debug 循环）。

关键设计¶

Modular Decomposition（Design-Decompose-Implement 三阶段流水线）:
- 功能：把"agent 写代码"从"吐一个大脚本"变成"先架构、再拆模块、再逐个写并验证"，输出真正的多文件仓库。
- 核心思路：三个专门 agent 串行 —— Idea Generation Agent 用自然语言写完整方案 plan；Modular Agent 把 plan 拆成若干独立功能模块 \(\{\mathcal{M}_j\}\)（如 dataset.py、model.py、engine.py、losses.py、utils.py、config.py 等）；Coding Agent 顺序实现每个 \(\mathcal{M}_j\)，每个模块写完跑独立 validation 脚本验证，再由 \(\pi_{\text{main}}\) 编排端到端 pipeline。修改时采用 Diff-Based Editing：用标准 diff 格式指定"目标文件 + 待替换 block + 新代码"，让一次 LLM 推理就能原子地多文件更新，避免每次重写整个仓库。
- 设计动机：分散到多文件直接绕开 token output 上限；focus 在小逻辑单元上降低 context 噪声，复杂逻辑准确率更高；validated 模块可以缓存复用；debug 定位到单文件而不是全脚本扫描。实测 Table 4 显示开了 modular 之后平均 LOC 从 474.8 涨到 1103.9，文件数从 1.0 → 6.7，说明 agent 真的能产出更复杂结构化的方案，而不是只把一个文件拆成多个。
Comparative Reflective Memory（Lesson Learning，定位因果而非记轨迹）:
- 功能：解决"哪个改动让指标涨/跌"的 credit assignment 问题，把执行 trajectory 蒸馏成结构化、可检索的 lesson pool。
- 核心思路：对成功的 valid 解走两步 —— Empirical Analysis Agent 先从执行日志抽客观发现（如 metric 变化趋势）；Lesson Distillation Agent 再做比较式反思，对比当前解和"已知最佳解"的 code-level delta，输出含三要素的 lesson：(1) 被隔离出的因果改动、(2) 比较影响分析、(3) 可用于未来迭代的泛化规则。对失败的 buggy 解，专门 agent 分析 buggy code + error log + 已应用的 fix，输出"fix 是否生效 + 故障逻辑 + 如何提前识别同类错误"的 debugging lesson。lesson pool 由 Review Agent 用 LLM 推理过滤冗余，agent 工作时只保留最近 \(K_m = 30\) 条 lesson 进 context，并被强制在使用时显式 cite 是哪一条 lesson。
- 设计动机：现有 memory（Reflexion、MemGPT、AIDE/AIRA 的轨迹缓存）只总结"发生了什么"，agent 容易从噪声 log 里 over-generalize；diff-based 比较把"算法改动"从噪声 log 里隔离出来，等价于 agent 自己做了一次 ablation 实验。实测：lesson-utilization rate 65.8%，lesson-transfer rate 63.0%（即被使用的 lesson 里 63% 来自不同 tree 分支），这是"Aha!" 现象的量化证据 —— agent 真的在跨分支迁移经验，而不是局部贪心。
Budget-Aware MCTS（效率感知的奖励函数）:
- 功能：在标准 MCTS 的 UCT 选择基础上，把"执行时间"显式写进 reward，让搜索系统性地偏向"快且好"的节点。
- 核心思路：节点 \(v\) 的性能 \(M(v)\) 先做全局归一化 \(G(v) = (M(v) - M_{\min}) / (M_{\max} - M_{\min})\)（探索过节点集合 \(\mathcal{V}\) 上的 min/max），效率感知 reward 定义为 \(R(v) = G(v) \cdot [t(v)/L(v)]^w\)，其中 \(t(v)\) 是实际执行时间，\(L(v)\) 是时间上限，\(w\) 是负惩罚权重（默认 \(w = -0.07\)）。直觉：两个解 accuracy 相同时，跑得快的 \(t/L\) 比值小，乘上负指数后 reward 更高。节点扩展规则也做了 MLE 定制：buggy 节点直接标 fully-expanded（不再扩 sibling，只走 debug loop 最多 \(N_d=10\) 步），valid 节点扩到 \(N_i = 2\) 个 improvement child 后封口，root 在"\(n_s\) 个 valid 节点都没改进 best 解"时重新激活允许新 draft —— 形成"探索 ↔ 重启"的自适应机制。
- 设计动机：MLE 是 24h wall-clock 强约束的任务，普通 MCTS 会把预算耗在小幅提升但极慢的方案上。\(w = -0.07\) 在消融里被验证是甜蜜点：\(w=0\)（纯 vanilla MCTS）性能明显下降，\(w=-0.15\) 又会让 search 偏向 trivial 快节点。Budget-aware 把 effective solution rate 从 16.1% 提到 19.5%，定量证明它是"剪枝低效轨迹"的有效启发式。

损失函数 / 训练策略¶

MARS 是无训练的 scaffolding（用预训练 LLM 做 backbone，主要试了 Gemini-2.5-Pro 与 Gemini-3-Pro-Preview），所有"学习"都发生在 inference 阶段的 MCTS + lesson pool 演化。关键超参：\(K_m=30\)（memory 最大 lesson 数）、\(N_d=10\)（debug 上限）、\(N_i=2\)（valid 节点 improvement 分支因子）、\(w=-0.07\)（reward 惩罚权重）；MLE-Bench 默认 24h × 1×A100（MARS+ 扩到 2 棵并行树 × 2×H100 × 48 vCPU）。

实验关键数据¶

主实验：MLE-Bench 75 任务、3 次独立运行均值±SEM（%）¶

Agent	Model	Above Median	Bronze	Silver	Gold	Any Medal
AIDE	Gemini-3-Pro-Prev	48.0	4.9	11.1	16.4	32.4
AIRA-dojo	Gemini-3-Pro-Prev	55.6	5.8	8.0	24.0	37.8
ML-Master 2.0 (leaderboard)	Deepseek-V3.2-Speciale	63.1	11.1	25.8	19.6	56.4
MARS	Gemini-3-Pro-Prev	65.8	9.3	15.6	31.1	56.0
MARS+ (2×H100)	Gemini-3-Pro-Prev	74.2	12.4	16.4	33.8	62.7

在受控（同 LLM、同环境）对比下，MARS 全面碾压 AIDE / AIRA-dojo（Any Medal 56.0 vs 37.8 / 32.4）；对官方 leaderboard，MARS 在更少资源下拿到所有方法中最高的 Gold 31.1%，MARS+ 在所有 metric 上同时刷新。按难度切分（Table 3），MARS Gemini-3-Pro-Prev 在 Lite / Medium / High 三个 split 上 Any Medal 分别 74.2 / 52.6 / 37.8，全面优于 AIRA-dojo 的 56.1 / 29.8 / 31.1。

消融实验（MLE-Bench Lite，22 个竞赛）¶

配置	关键发现	说明
Full MARS	baseline	三模块全开
w/o Modular Decomposition	显著下降	验证模块化对降低 context 噪声、提升 testability 的作用；LOC 从 1103.9 跌到 474.8，文件数从 6.7 跌到 1.0
w/o Memory（彻底去掉）	drastic drop	没有 lesson pool 时 agent 几乎学不到任何东西
Memory 但只用 Empirical Analysis（去掉 comparative delta）	比 full 差但优于无 memory	证明"对比 best-known 的 diff 分析"是 lesson 质量的核心
Greedy Search	显著差	贪心只盯当前最好 val metric，缺少探索
Vanilla MCTS（\(w=0\)）	中等	有 MCTS 探索但无预算意识
Budget-Aware MCTS（\(w=-0.07\)，默认）	最优	effective solution rate 19.5% vs vanilla 16.1%
\(w=-0.15\)（更狠的惩罚）	退化	过度偏向快但弱的节点

关键发现¶

Comparative memory 是核心贡献：去掉 comparative 只留 empirical analysis 就掉点，说明 lesson 的价值来自"code-level diff 隔离因果改动"，而不是简单总结日志。
Lesson 真的会跨分支迁移：lesson-transfer rate 63%，配合 Figure 1 在 iMet-2020-FGVC7 上从轻量 ResNet 一路演化到 model ensemble 拿银牌的轨迹，是"Aha! moment"的定量+定性证据。
Budget-awareness 是剪枝启发式：相同 accuracy 时 reward 偏向快节点，effective solution rate 提 3.4 个点，等价于在 24h 预算下多探索 ~20% 有效解。
Modular 让 agent 写真复杂代码：开了 Modular 之后 LOC 翻倍、文件数 ×6.7，说明这不是"把同样代码拆文件"，而是 agent 真正在架构层面思考。
scaling 还能继续：MARS+ 双树并行 + 更大算力把 Above Median 从 65.8 推到 74.2，说明 MCTS 框架天然可并行。

亮点与洞察¶

把 MLE 重新定义为 "search over modular repositories"：这是范式转换。之前所有 agent 都在"搜索更好的脚本"，MARS 直接把搜索单元抬到了仓库级，配套的 modular decomposition + diff-based edit 让这个抽象在工程上可行。这个抽象可以直接迁移到 SWE-Bench、RepoCoder 这类 repo-level 任务。
Comparative reflection 等价于"agent 自己做 ablation"：人类做 ML 研究本来就靠 ablation 隔离因果，MARS 把这个动作内化进 memory pipeline。这个套路可以直接搬到任何"长程 + 多 trial + 指标可比"的 agent 任务（科学实验、超参搜索、prompt 优化）。
Budget 进 reward 是被低估的简单 trick：\(R = G \cdot (t/L)^w\) 只多一个超参，但把 vanilla MCTS 变成预算感知，effective solution rate +3.4 个点。任何带 wall-clock 的 LLM agent 搜索（tool-use、web research）都该抄。
Lesson cite 强制可解释：要求 agent 在用 lesson 时显式 cite 是哪一条，让 agent 行为可审计、可调试，这种"强制溯源"是工程上很务实的设计。
跨分支迁移 63% 是真实 emergent 行为：不是 sequence-level reflexion，而是把 lesson 当全局知识库跨分支调用 —— 这个数字本身就是 MCTS + memory 协同的有力证据。

局限与展望¶

作者承认的局限：MLE-Bench 只覆盖 75 个 Kaggle-like 任务，benchmark 的多样性有限；真实科研还需要假设生成、文献阅读、跨论文复现等能力，本文没覆盖。
强依赖大 LLM：实验全在 Gemini-2.5-Pro / Gemini-3-Pro-Preview 上跑，开源小模型能不能撑住三个 agent 串行（Idea / Modular / Coding）+ 反思 + lesson distill，需要专门验证。Lesson distill 本身就是高 reasoning 负载。
MCTS 树宽固定：\(N_i = 2\)、\(N_d = 10\) 是手动调的，没做自适应；任务难度差异大时这种固定分支因子可能浪费预算或剪得太狠。
lesson pool \(K_m=30\) 是上限而非检索：单纯 LRU 截断，没用 embedding-based retrieval，跨任务 lesson 复用（meta-MLE agent）几乎没探索。
缺少 cost vs quality 的 Pareto 曲线：论文只报最终 metric，没系统画出"算 1h / 6h / 24h 各拿到什么金银铜分布"，这对工业部署其实更重要。
可改进方向：把 lesson pool 升级为 RAG-style 带 embedding 的长期记忆，跨任务共享 lesson；把 Modular Decomposition 的 module 接口标准化做成可复用 zoo；把 \(w\) 做成 task-aware 自适应；把 search 扩到 multi-agent 协作（多树共享 lesson pool）。

评分¶

新颖性: ⭐⭐⭐⭐ 三大支柱单独看都有前作雏形（MCTS、modular code、reflection），但"把三者围绕 cost-constrained repo-level search 一致地组合起来"且每个支柱都做了 MLE 专属定制，是有原创性的系统级贡献。
实验充分度: ⭐⭐⭐⭐⭐ 在 MLE-Bench 全 75 任务 × 3 次 run × 多 LLM × 多 baseline 对照，并提供受控环境 vs leaderboard 双视角，含模块化 / 记忆 / 搜索策略 / 惩罚权重四项消融，外加 lesson-transfer rate 等定量指标。
写作质量: ⭐⭐⭐⭐ Figure 2 的整体框架图清晰，公式简洁，三模块层次分明；Related Work 的 Table 1 直接对比每个 agent 在 modular / budget / memory 三轴上的取舍，定位非常明确。
价值: ⭐⭐⭐⭐⭐ MLE-Bench 开源 SOTA + 完整 prompt 与代码开源（github.com/jfc43/MARS），三大设计可单独抽取迁移到其他 long-horizon agent 任务，对工业自动化 AI 研究是 baseline-setter。