MARS: Modular Agent with Reflective Search for Automated AI Research

AI summary

Overview

Research area: Autonomous AI research / machine learning engineering (MLE) agents built on large language models, with a focus on tree search, modular code generation, and agent memory.

Technical level: Advanced. The paper assumes familiarity with LLM agent scaffolding, Monte Carlo Tree Search, reward shaping, and the MLE-Bench evaluation protocol.

Scope: This paper introduces MARS (Modular Agent with Reflective Search), a framework that reformulates long-horizon AI research tasks as a budget-constrained search for an optimal software repository, and evaluates it on the MLE-Bench benchmark.

What This Paper Is About

Automating AI research is bottlenecked by machine learning engineering (MLE) tasks, which are expensive to evaluate (they require model training) and hard to attribute causally (it is unclear which code change caused a performance shift). Existing LLM agents attack these tasks by writing single monolithic scripts, ignoring execution cost and producing fragile code. MARS's goal is to build an agent that plans under an explicit compute budget, constructs solutions as modular repositories rather than one file, and learns distilled "lessons" by comparing its current solution against its best-known solution.

Key Contributions

1. MARS framework. A combination of three mechanisms for autonomous AI research: Budget-Aware Monte Carlo Tree Search (MCTS), a modular implementation pipeline, and Comparative Reflective Memory.
2. Budget-Aware MCTS. A cost-constrained search that optimizes an efficiency-guided reward explicitly balancing task performance against execution time (a latency penalty weight w, set to −0.07).
3. Modular "Design-Decompose-Implement" pipeline. An Idea Generation Agent, a Modular Agent, and a Coding Agent that produce solutions as sets of independent modules plus an orchestration script, updated via Diff-Based Editing rather than full regeneration.
4. Comparative Reflective Memory. A "Lesson Learning" mechanism that compares the current solution against the previous best-known solution to isolate causal code changes, plus Debugging Lessons for failures, reviewed and filtered by a Review Agent.
5. Open release. Prompts are provided in Appendix H, and MARS-generated code and trajectories are released at https://github.com/jfc43/MARS.

Main Findings

• State-of-the-art among open-source frameworks under comparable settings. On MLE-Bench (75 Kaggle competitions spanning NLP, computer vision, and tabular data), MARS with Gemini-3-Pro-Preview reached 98.7% ± 0.0 Valid Submission, 65.8% ± 1.6 Above Median, 9.3% ± 0.0 Bronze, 15.6% ± 1.2 Silver, 31.1% ± 0.4 Gold, and 56.0% ± 1.5 Any Medal across three runs.
• Scaled variant MARS+ leads reported metrics. MARS+, configured with two concurrent search trees and increased compute (2× H100 GPUs and 48 vCPUs), reached 100.0% ± 0.0 Valid Submission, 74.2% ± 0.9 Above Median, 12.4% ± 1.9 Bronze, 16.4% ± 1.2 Silver, 33.8% ± 0.4 Gold, and 62.7% ± 0.8 Any Medal.
• Best reported Gold Medal rate. The paper reports that the standard MARS achieves the highest Gold Medal rate (31.1%) among all reported agents in its comparison; MARS+ reports 33.8%.
• Consistent advantage across task complexity splits. On MLE-Bench Lite/Medium/High, MARS with Gemini-3-Pro-Preview reached Any Medal rates of 74.2% ± 1.5, 52.6% ± 3.0, and 37.8% ± 2.2 respectively, versus AIDE (53.0% ± 6.1, 26.3% ± 3.0, 17.8% ± 2.2) and AIRA-dojo (56.1% ± 1.5, 29.8% ± 3.8, 31.1% ± 4.4) with the same model.
• Modular decomposition produces larger, more structured repositories. Best solutions from MARS averaged 1,103.9 ± 35.9 lines of code across 6.7 ± 0.1 files, versus 474.8 ± 13.5 lines across 1.0 ± 0.0 files without modular decomposition on MLE-Bench Lite.
• Budget-Aware MCTS improves effective solution discovery. Its effective solution rate was 19.5% ± 1.5%, compared to 16.1% ± 1.3% for Vanilla MCTS.
• Ablations confirm each component. Removing modular decomposition degrades performance; removing memory mechanisms causes a drastic drop; the comparative component provides a consistent boost over empirical analysis alone; Budget-Aware MCTS outperforms Greedy Search and Vanilla MCTS (w = 0); and the default penalty weight w = −0.07 is optimal, with w = −0.15 excessively biasing the search toward latency.
• Lessons transfer across search branches. MARS achieved a lesson-utilization rate of 65.8% ± 1.1% and a lesson-transfer rate of 63.0% ± 1.8% on MLE-Bench.
• Distilled lessons are largely causally correct. An audit with Claude 4.6 Sonnet of all 3,611 solution lessons generated by Gemini-3-Pro-Preview found 88.34% correctly attributed validation metric shifts to specific code changes; a manual inspection of 20 lessons yielded 90% causal accuracy.
• Rule compliance and code originality. The official MLE-Bench log analysis tool (using gpt-4.1-mini) recorded a 0% violation rate across monitored dimensions ("Tried to access unauthorized resources", "Tried to call external LLM API service", "Manually-written submission"). Plagiarism detection based on Dolos found no submission exceeding a 60% similarity threshold against top public notebooks.
• Cost-controlled Pareto improvement. The 24-hour MARS cost $60.5 versus $39.0 for AIRA-dojo while nearly doubling Any Medal Rate (from 24.4% to 43.1%). A 4-hour version of MARS achieved a 28.4% Any Medal Rate at $9.6.
• Qualitative "Aha!" moment. On the iMet-2020-FGVC7 task, MARS progressed from a lightweight residual network to model ensemble techniques and reached a silver medal, where existing methods failed to reach medal-level performance.

Methodology in Plain English

The paper formalizes its task as a tuple of Instruction, Environment, and Objective, and asks the agent to maximize the Objective while staying under a cost budget. For MLE, the Environment is the provided datasets and the Objective is the metric on a held-out test set.

MARS then runs a loop with three parts:

• Planning. A cost-constrained MCTS navigates the search space using three actions: draft a new architecture, debug runtime errors, or improve a valid solution. Node selection uses UCT. The reward function normalizes performance relative to all previously explored nodes and multiplies it by a latency term [t(v)/L(v)]^w, so that two solutions with similar accuracy are ranked in favor of the faster one. Failed nodes enter an automatic debugging loop, and the number of improvement attempts per valid node is capped.
• Construction. Instead of writing one script, an Idea Generation Agent writes a natural-language plan, a Modular Agent decomposes it into independent functional modules, and a Coding Agent implements and validates each module before orchestrating them through a main script. Edits are applied as standardized diffs specifying the target file, the block to replace, and the new code.
• Memory. For successful solutions, an Empirical Analysis Agent extracts findings from logs and a Lesson Distillation Agent compares the current solution against the previous best to isolate causal changes, producing a lesson with the isolated change, a comparative impact analysis, and a generalized rule. For failures, a separate agent produces debugging lessons. A Review Agent filters redundant lessons, the K_m most recent lessons are kept in memory, and the agent is instructed to cite the lessons it applies.

MLE-specific components include multi-agent task preparation and metadata extraction, an Exploratory Data Analysis agent, and a curriculum-based idea generation strategy that moves from simple baselines to complex methods.

Evaluation followed the standard MLE-Bench protocol: a strict 24-hour wall-clock budget per competition on a node with one NVIDIA A100 GPU (40GB), 12 vCPUs, 220 GB of RAM, and 1 TB of SSD storage. Open-source baselines (AIDE and AIRA-dojo) were run under identical configurations with the same underlying LLMs. Results are the mean and standard error of the mean across three independent runs. Hyperparameters: K_m = 30, N_d = 10, N_i = 2, w = −0.07.

Why This Matters

The paper argues that MLE execution is the critical bottleneck for higher-level research automation: without robust autonomous MLE, more ambitious scientific automation cannot succeed. Its distinctive claim is that accounting for execution cost inside the search — rather than treating problem-solving as a purely code-based challenge — produces better discovery per unit of compute, and that analyzing deltas between solutions is a workable substitute for the isolated ablation experiments humans run.

Real-world applications suggested by the setting:

• Automated machine learning pipelines for tabular, vision, and language tasks under fixed time and cost budgets.
• Competition-style predictive modeling (the benchmark itself is drawn from 75 Kaggle competitions).
• Enterprise data science automation, where engineering teams need agent-generated repositories rather than single-file scripts, and where API spend matters.
• Foundation work toward automated research engineering and automated research replication, which the paper lists as adjacent domains for future evaluation.

Industry relevance comes from the cost argument: the paper reports that a 4-hour MARS configuration achieved a

Authors’ abstract

A critical bottleneck in automating AI research is the execution of complex machine learning engineering (MLE) tasks. MLE differs from general software engineering due to computationally expensive evaluation (e.g., model training) and opaque performance attribution. Current LLM-based agents struggle here, often generating monolithic scripts that ignore execution costs and causal factors. We introduce MARS (Modular Agent with Reflective Search), a framework optimized for autonomous AI research. MARS relies on three pillars: (1) Budget-Aware Planning via cost-constrained Monte Carlo Tree Search (MCTS) to explicitly balance performance with execution expense; (2) Modular Construction, employing a "Design-Decompose-Implement" pipeline to manage complex research repositories; and (3) Comparative Reflective Memory, which addresses credit assignment by analyzing solution differences to distill high-signal insights. MARS achieves state-of-the-art performance among open-source frameworks on MLE-Bench under comparable settings, maintaining competitiveness with the global leaderboard's top methods. Furthermore, the system exhibits qualitative "Aha!" moments, where 63% of all utilized lessons originate from cross-branch transfer, demonstrating that the agent effectively generalizes insights across search paths.