MARS: Modular Agent with Reflective Search for Automated AI Research

Jiefeng Chen, Bhavana Dalvi Mishra, Jaehyun Nam, Rui Meng, Tomas Pfister, Jinsung Yoon

arXiv:2602.02660v3 · This workspace stays on this version. Check arXiv for newer versions

Affiliations and author notes

Google Cloud AI Research

Some source details need review 29 to check

Original formulas, table values and figures are preserved. Comparison can miss differences; use the original PDF to check details.

Reading view checked: 234 math expressions, 0 rendering failures; 45 of 45 images available.

Compared 166 of 174 source passages against 44 of 66 PDF pages.

The HTML source contains failed formatting commands; these were omitted from the reading view. The original sources are retained. (1)

Some graphics are laid out differently in the HTML source than in the paper (text or formulas outside their frame); check the original PDF. (27)
A8_p2_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p3_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p4_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p5_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p6_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p7_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p8_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p9_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p10_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p11_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p12_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p13_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p14_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p15_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p16_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p17_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p18_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p19_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A8_p20_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A9_p2_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A9_p3_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A9_p4_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A9_p5_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A9_p6_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A9_p7_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A9_p8_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
A9_p9_pic1: the arXiv HTML draws text or formulas of this graphic outside its frame and it could not be corrected safely (too-small); the image shows the HTML layout, which differs from the PDF.
This long paper was only partially compared; the original sources remain available.

45 of 45 figure, table and graphic images captured. Open original PDF

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.

Figure 1: The “Aha!” moment of MARS on the challenging iMet-2020-FGVC7 task. The visualization tracks validation performance gains triggered by specific strategic lessons. While existing methods fail to reach medal-level performance, MARS progressively refines its strategy – evolving from a lightweight residual network to model ensemble techniques – to ultimately achieve a silver medal.

The integration of Large Language Models (LLMs) into software engineering has fundamentally transformed code generation, evolving from simple auto-completion to autonomous agents capable of resolving GitHub issues Jimenez et al. (2023); Yang et al. (2024) and generating functional scripts (Li et al., 2022; Wang et al., 2025; Jiang et al., 2025). However, while current agents excel at general software maintenance tasks – such as patching bugs or writing unit tests – they face significant hurdles when applied to the domain of Automating AI Research (Chan et al., 2025; Tian et al., 2024; Wijk et al., 2025; Yamada et al., 2025; Starace et al., 2025). A critical bottleneck in this domain is the execution of complex machine learning engineering (MLE) tasks. Unlike standard software development, where correctness is often binary and verification is computationally cheap, MLE is a probabilistic, resource-intensive endeavor. It requires not only coding intelligence but also the strategic foresight to navigate a landscape defined by computationally expensive evaluations, opaque performance attribution, and high architectural complexity.

Existing agentic frameworks, designed primarily for monolithic code generation, struggle to adapt to these constraints. First, they typically view problem-solving as a purely code-based challenge (Huang et al., 2024; Jiang et al., 2025; Toledo et al., 2026), ignoring the economic reality of research: model training and data processing consume vast computational resources. An agent that improves model accuracy by 0.1% but increases training time from one hour to ten hours is often practically useless, yet standard search algorithms would prioritize it. Second, the monolithic and unstructured scripts often produced by previous LLM agents are fragile and ill-suited for the modular complexity required in research repositories, where data loading, model architecture, and training loops must interact seamlessly. Finally, research progress is iterative and opaque; when a new experiment yields better results, it is difficult to isolate the causal factor. Standard memory-based agents (Packer et al., 2023; Shinn et al., 2023; Xu et al., 2026; Ouyang et al., 2026) lack the mechanism to solve this credit assignment problem, often failing to learn effectively from past trials.

To bridge this gap, we introduce MARS (Modular Agent with Reflective Search), a framework explicitly optimized for the distinct constraints of autonomous scientific discovery. MARS reformulates the research process as a search for an optimal software repository, governed by three core pillars. To address the high cost of evaluation, we employ Budget-Aware Planning via a cost-constrained Monte Carlo Tree Search (MCTS). Unlike general search algorithms, our method explicitly balances performance maximization with execution expense, prioritizing efficient solutions – such as favoring a 1-hour training run over a 4-hour run if performance is comparable – to optimize the discovery rate within a fixed budget. To manage architectural complexity, we replace fragile scripting with a Modular “Design-Decompose-Implement” pipeline. This structure employs specialized agents to architect solutions into independent, testable modules. Finally, to resolve the credit assignment problem, we introduce Comparative Reflective Memory. By analyzing the differences between the current solution and the best-known solution, the agent distills high-signal, causal insights, isolating the specific factors driving performance shifts in a way that standard memory mechanisms cannot. As illustrated in Figure 1, these pillars allow MARS to experience “Aha!” moments during long-horizon exploration, successfully navigating complex optimization landscapes where baselines fail.

Our contributions are summarized as follows:

•
We introduce MARS, a framework aimed at automating AI research by optimizing for complex MLE tasks, featuring a novel combination of Budget-Aware MCTS, a modular implementation pipeline, and Comparative Reflective Memory.
•
We perform extensive evaluation on the MLE-Bench benchmark, where MARS achieves state-of-the-art performance among open-source frameworks under comparable settings. Ablation studies further validate the necessity of each proposed mechanism.
•
We provide qualitative analyses of how MARS drives long-horizon exploration. To facilitate future research, we release prompts in Appendix H, and MARS generated code, trajectories in https://github.com/jfc43/MARS.

Recent advancements in LLMs have driven ambitious efforts to automate complex, long-horizon AI research (Lu et al., 2024; Yamada et al., 2025; Weng et al., 2025; Tang et al., 2026; Schmidgall et al., 2025; Pu et al., 2025; Starace et al., 2025). Central to these efforts is the realization that Machine Learning Engineering (MLE) serves as a critical execution bottleneck (Chan et al., 2025; Wijk et al., 2025); without robust autonomous MLE capabilities, higher-level research automation cannot succeed. While numerous agentic frameworks have been proposed to address these challenges (Jiang et al., 2025; Toledo et al., 2026; Yang et al., 2025; Liu et al., 2025b; Team et al., 2025; Li et al., 2025; Nam et al., 2026; Zhu et al., 2026), existing systems predominantly operate under a monolithic paradigm, generating expansive, single-file scripts. This approach typically results in fragile codebases that lack the modularity essential for rigorous engineering. MARS departs from this by enforcing a repository-level paradigm that systematically decomposes tasks into distinct, testable, and maintainable modules, mirroring the professional software architecture required for scalable AI research.

Search Algorithms in Code Generation. Solving long-horizon AI research problems, where code execution is resource-intensive, necessitates effective search strategies for code optimization. While various algorithms have been adapted for these systems – including greedy search (Jiang et al., 2025), Monte Carlo Tree Search (MCTS) (Kocsis and Szepesvári, 2006; Liu et al., 2025b), and Evolutionary search (Team et al., 2025) – they typically optimize solely for task performance, neglecting computational cost. While recent work has introduced “budget awareness” for tool-augmented agents via external plug-ins (Liu et al., 2025a), such methods are primarily designed for discrete actions like web search. In contrast, we introduce Budget-aware MCTS, which integrates an efficiency-guided reward function directly into the search tree. This allows MARS to balance the exploitation of high-performing strategies with the exploration of novel ideas, penalizing computationally expensive solutions to ensure both performance and efficiency.

Reflective Learning and Memory. Enabling agents to improve iteratively through environmental interaction is a rapidly evolving research area. Early approaches, such as Reflexion (Shinn et al., 2023), focus on self-correction using verbal reinforcement derived from prior mistakes. Recently, the field has expanded to include sophisticated memory and distillation frameworks. For instance, several works explore advanced memory architectures for personalization and associative retrieval Tan et al. (2025); Xu et al. (2026), while others focus on refining process memory and distilling execution traces into stable strategies Forouzandeh et al. (2025); Ouyang et al. (2026); Zhu et al. (2026). Closer to our domain, Zhao et al. identify “good practices” by comparing high-level action trajectories, and Jansen et al. cache useful code blocks for future reuse. MARS advances this paradigm by introducing “Lesson Learning”. Unlike prior methods that summarize execution logs to debug errors or track binary success-failure trajectories, our method explicitly analyzes the causal link between code changes and performance variations. This comparative analysis isolates effective algorithmic changes from confounding factors, distilling high-value insights into a lesson pool to guide future exploration.

Table 1 summarizes the key differences between MARS and existing MLE agent frameworks.

We first formalize the general problem of Long-Horizon Agentic Problem Solving, where an autonomous agent is tasked with constructing a complex artifacts (e.g., a software system) to satisfy a set of requirements within a constrained budget. Let $\mathcal{P}$ denote a problem instance defined by the tuple $\mathcal{P}=(\mathcal{I},\mathcal{E},\mathcal{O})$ , where: (1) $\mathcal{I}$ represents the Instruction or requirements provided in natural language. (2) $\mathcal{E}$ denotes the Environment with which the agent interacts to validate its solutions. This can be a compiler, a simulator, or a dataset depending on the task scenario. (3) $\mathcal{O}$ is the Objective function that quantifies the quality of the solution.

The goal is to find a solution $s^{*}$ that maximizes $\mathcal{O}$ by interacting with $\mathcal{E}$ , subject to a cost constraint $B$ (e.g., time budget or monetary cost):

\displaystyle s^{*}=\argmax_{s}\quad\mathcal{O}(s,\mathcal{E}),\quad s.t.\quad\text{Cost}(s)\leq B

where the search space for $s$ is often vast and unstructured (e.g., the space of all possible Python programs).

Machine Learning Engineering (MLE) is a representative and challenging instantiation of this general problem class. MLE requires the agent to engineer a full pipeline that processes data, trains models, and validates results. In this scenario, $\mathcal{E}$ consists of the provided datasets while $\mathcal{O}$ is the performance metric (e.g. accuracy) on the held-out test set. Refer to Appendix B for details.

Figure 2: Overview of the MARS Framework. MARS reformulates long-horizon coding as a search for an optimal software repository. (1) Task Preparation: The agent grounds the abstract problem (Instruction, Environment, Objective) tuple by exploratory analysis of the given dataset and metadata. (2) The MARS Loop: The agent iteratively evolves solutions through three synergistic modules: **(A) Resource-Aware Planning:** A Budget-Aware MCTS strategically navigates the search space by selecting actions from {Draft new architecture, Debug runtime errors, Improve a valid solution}. It optimizes an efficiency-guided reward that explicitly balances performance maximization with the penalty of high execution costs. **(B) Modular Decomposition:** To replace fragile monolithic scripting, the system employs a “Design-Decompose-Implement” pipeline. Specialized {Idea, Modular, Coding} agents architect the solution into independent, testable modules. This structure enables precise Diff-Based Refinement, allowing the agent to update specific logic blocks without regenerating the entire codebase. **(C) Reflective Memory:** This module distills raw execution logs into structured Debugging and Solution Lessons to proactively prevent error repetition and accelerate convergence in later iterations.

We propose MARS, a general agent scaffolding framework designed to enable autonomous agents to solve long-horizon AI Research problems, as illustrated in Figure 2. Formally, we define the problem as a tuple $\mathcal{P}=(\mathcal{I},\mathcal{E},\mathcal{O})$ , where the agent must follow the instruction $\mathcal{I}$ within an environment $\mathcal{E}$ to maximize an objective $\mathcal{O}$ under a cost budget $B$ . To address the core challenges of exploration complexity, context management, and solution robustness in this setting, our framework integrates three key capabilities:

•
Modular Construction Strategy: Instead of generating monolithic scripts, we enforce a structured, repository-level software architecture. This paradigm allows for handing complex logic with greater accuracy, efficient code reuse, and improving testability.
•
Comparative Reflective Memory: To overcome context window limitations, we introduce a “Lesson Learning” mechanism that distills high-value, causal insights from past interactions (both successes and failures) into a compact, retrievable knowledge base.
•
Resource-Aware Planning: We employ a budget-aware Monte Carlo Tree Search (MCTS) algorithm to systematically explore the solution space. This allows the system to balance the exploitation of promising candidates with the exploration of novel ideas, preventing local optima, and penalize solutions that are costly.

A primary contribution of this work is the strategic shift from generating monolithic scripts to a Modular Implementation paradigm. This paradigm addresses several inherent limitations of LLM-based coding. First, it bypasses token output limits by distributing code across multiple files. Second, it enhances precision; by focusing on smaller logical units, the agent encounters less context noise and can handle complex logic with greater accuracy. Third, it enables efficiency via caching, as validated modules can be reused without regeneration. Finally, it significantly improves testability, as debugging is localized to specific files rather than requiring full-script diagnosis.

We define a node solution $s_{n}$ as a tuple comprising a set of $l$ independent modules and one orchestration script:

\displaystyle s_{n}=\langle\{\mathcal{M}_{j}\}_{j=1}^{l},\pi_{main}\rangle

Each module $\mathcal{M}_{j}$ encapsulates a specific sub-task (e.g., data preprocessing, configuration), while the main script $\pi_{\text{main}}$ orchestrates the end-to-end pipeline.

To instantiate this structure, we employ a three-stage “Design-Decompose-Implement” workflow:

•
Idea Generation: An Idea Generation Agent articulates a comprehensive natural language plan covering various aspects of the solution.
•
Module Decomposition: A Modular Agent parses the plan and decomposes the solution into logical, independent functional modules.
•
Component Implementation and Debugging: A Coding Agent sequentially implements each module $\mathcal{M}_{j}$ , employing a validation script to debug and verify functionality. Once validated, the agent orchestrates the modules via the main script $\pi_{\text{main}}$ .

To prevent wasteful full-repository regeneration, we adopt a Diff-Based Editing mechanism. Code modifications are structured in a standardized diff format, specifying the target file, the block to replace, and the new code. This enables atomic, multi-file updates in a single inference step.

Solving complex tasks requires long-horizon exploration, generating extensive interaction trajectories that often exceed context window constraints. More importantly, research progress is inherently iterative and opaque; when a new experiment yields improved results, isolating the specific causal factors remains a challenge. Standard memory-based agents often lack the mechanisms to solve this credit assignment problem, failing to learn effectively from past trials. To address this, we propose Comparative Reflective Memory, a mechanism designed to distill high-signal, causal insights from the exploration process into a compact lesson pool.

We employ a two-stage process to resolve the credit assignment problem by synthesizing lessons from valid solutions. First, an Empirical Analysis Agent reviews execution logs to extract objective findings (e.g., metric trends). Subsequently, a Lesson Distillation Agent performs a comparative reflection by analyzing the delta between the current solution and the previous best-known solution. This isolates the specific algorithmic changes driving performance shifts, resulting in a structured lesson containing: (1) The isolated causal change, (2) A comparative impact analysis, and (3) A generalized rule for future iterations.

For failed executions, a dedicated agent analyzes the buggy code, error logs, and the applied fix. It outputs a lesson confirming the fix’s efficacy, explaining the failure logic, and providing guidelines to preemptively identify similar errors.

To maintain a high-signal lesson pool, a Review Agent evaluates new lessons against the existing pool through LLM-based reasoning, filtering out redundant insights to ensure the retrieved context remains diverse and relevant.

When executing solution improvement or debugging actions, the agent utilizes relevant knowledge from the corresponding lesson categories. We retain the $K_{m}$ most recent lessons in the agent’s memory to manage context. To ensure interpretability, the agent is instructed to explicitly cite specific lessons whenever they are applied.

We adopt the Monte Carlo Tree Search (MCTS) framework to explore the solution space, which iterates through four phases: Selection, Expansion, Simulation, and Backpropagation. In this section, we detail our domain-specific modifications: (1) specialized expansion operators, (2) a coherent node selection strategy, and (3) an Efficiency-Guided Reward Function that balances performance with cost. Appendix C provides a review of standard MCTS principles.

We define three distinct operators to transform a parent state $s_{parent}$ into a child solution $s_{new}$ :

•
Drafting (Root Expansion): Generates a completely new solution $s_{new}$ from scratch.
•
Improvement: Applied to valid, executable nodes. The agent modifies the modules and the main script from $s_{parent}$ to maximize the objective $\mathcal{O}$ .
•
Debugging: Applied to nodes where execution failed. The agent inherits the solution structure from $s_{parent}$ but modifies specific modules or the orchestration script to resolve runtime errors. Buggy children enter an automatic debugging loop with up to $N_{d}$ debugging actions to fix the errors.

We employ the Upper Confidence Bound for Trees (UCT) algorithm to navigate the solution space, balancing the exploitation of high-performing solutions with the exploration of new solutions.

The selection phase begins at the root node. In each step, we select the child node that maximizes the UCT value. This traversal continues recursively until we identify a candidate node, defined as a node that is not yet “fully expanded”.

The root node is set fully expanded unless any of the follow condition occurs: (1) It does not have any children; (2) the best solution has not been improved after implementing $n_{s}$ valid nodes.

If the traversal reaches a leaf node that is already fully expanded (it implies that no further debugging or improvement is permitted for that branch), then the root node is re-activated to allow for new drafts.

The buggy nodes are always set fully expanded. The valid nodes are set fully expanded if they have $\geq N_{i}$ children (attempts to improve).

To guide the search efficiently, we design a reward function $R(v)$ that rewards performance gains and penalizes long execution time. Let $M(v)$ denote the performance metric of a node $v$ , and let $t(v)$ and $L(v)$ represent its execution time and time limit, respectively. We first normalize the performance metric relative to the history of explored nodes $\mathcal{V}$ . Let $M_{max}=\max_{v^{\prime}\in\mathcal{V}}M(v^{\prime})$ and $M_{min}=\min_{v^{\prime}\in\mathcal{V}}M(v^{\prime})$ . We define the global normalized score $G(v)$ as:

\displaystyle G(v):=\begin{cases}0.5\quad\text{if }M_{max}=M_{min},\\ \frac{M(v)-M_{min}}{M_{max}-M_{min}}\quad\text{otherwise}\end{cases}

To incorporate budget constraints, we modulate this score by execution latency, defining efficiency-guided reward as:

\displaystyle R(v):=G(v)\cdot[t(v)/L(v)]^{w}

Where $w$ is a penalty weight hyperparameter. A similar function has been proposed in Tan et al. (2019).

Table 2: Performance comparison on MLE-Bench. Results are reported as mean $\pm$ SEM across three independent runs. All values are in percentages (%). **Bold** and underlined values denote the best and second-best performance, respectively. Refer to Appendix E for a detailed comparison of evaluation setups.

While MARS is a general framework, its application requires task-specific components. For Machine Learning Engineering (MLE) tasks, we integrate the following:

Task preparation. We employ a multi-agent system to extract task metadata, formalizing the optimization objective and preparing training, validation, and test datasets.

Data analysis. We employ an agent to perform Exploratory Data Analysis (EDA) to generate a report that guides downstream feature engineering.

Curriculum-Based Exploration. We implement a curriculum-based idea generation strategy that progressively explores simple baselines to complex methods.

Refer to the Appendix D for the details.

Datasets. We evaluate our agent on MLE-Bench (Chan et al., 2025), which consists of 75 challenging competitions from Kaggle, forming a diverse collection of tasks covering natural language processing, computer vision, and tabular data analysis.

Environments. We adhere to the standard MLE-Bench protocol, where agents are allocated a strict 24-hour wall-clock time budget per competition. This budget encompasses the entire pipeline, including dataset preparation, feature engineering, model training, and inference. The experiment for each agent on each competition is conducted on a standard node equipped with one NVIDIA A100 GPU (40GB), 12 vCPUs, 220 GB of RAM, and 1 TB of SSD storage. This setup simulates a realistic, resource-constrained machine learning engineering environment.

Baselines. We compare our method to the agents in the MLE-Bench leaderboard $^{1}$ and two state-of-the-art open-source agents: AIDE (Jiang et al., 2025) and AIRA (Toledo et al., 2026). For open-source baselines, we ensure a strictly fair comparison by running them under identical environment configurations and using the same underlying LLMs.

Metrics. Following the standard MLE-Bench evaluation protocol, we report the mean and standard error of the mean (SEM) across three independent runs. Our evaluation focuses on three primary metrics: Above Median Rate (percentage of runs outperforming the median participant), Any Medal Rate (percentage achieving at least a Bronze medal), and Gold Medal Rate (percentage securing a Gold medal).

Hyperparameters. We set the maximum number of lessons in the agent’s memory to $K_{m}=30$ to maintain relevant context without context window overflow. We allow up to $N_{d}=10$ debugging actions per failure to resolve runtime errors effectively. The branching factor for valid nodes is set to $N_{i}=2$ , balancing exploration breadth with depth. We set $w=-0.07$ in the reward function (4) following Tan et al. (2019) to penalize excessive execution time.

We compare MARS against state-of-the-art baselines in Table 2. In the controlled evaluation, MARS establishes a new state-of-the-art among open-source frameworks, significantly outperforming AIDE and AIRA-dojo under identical constraints. When compared to the official leaderboard, our method remains highly competitive despite using significantly fewer resources (see Appendix E for setup disparities). Notably, the standard MARS achieves the highest Gold Medal rate (31.1%) among all reported agents. To assess scalability, we evaluate MARS+, a variant configured to execute two concurrent search trees with increased compute (2 $\times$ H100 GPUs and 48 vCPUs). This scaled approach achieves the highest Above Median rate (74.2%), Gold Medal rate (33.8%), and Any Medal rate (62.7%), outperforming strong competitors like ML-Master 2.0. Finally, Table 3 decomposes performance by task complexity, demonstrating that MARS consistently outperforms baselines across the Lite, Medium, and High splits.

Original Table 3 — Table 3: Controlled evaluation in our environment across different splits of MLE-Bench. Results are reported as mean $\pm$ SEM across three independent runs. The best performance is highlighted in **bold**, and the second-best is underlined. The complete results including leaderboard results and other metrics are in Appendix F.1.

Table 3
Agent	Model	Any Medal
Lite (%)	Medium (%)	High (%)
AIDE	Gemini-2.5-Pro	$36.4\pm 4.5$	$18.4\pm 2.6$	$15.6\pm 2.2$
Gemini-3-Pro-Prev.	$53.0\pm 6.1$	$26.3\pm 3.0$	$17.8\pm 2.2$
AIRA	Gemini-2.5-Pro	$40.9\pm 2.6$	$16.7\pm 3.5$	$20.0\pm 0.0$
-dojo	Gemini-3-Pro-Prev.	$56.1\pm 1.5$	$29.8\pm 3.8$	$\underline{31.1}\pm 4.4$
MARS	Gemini-2.5-Pro	$\underline{68.2}\pm 2.6$	$\underline{33.3}\pm 1.8$	$\underline{31.1}\pm 2.2$
(ours)	Gemini-3-Pro-Prev.	$\textbf{74.2}\pm 1.5$	$\textbf{52.6}\pm 3.0$	$\textbf{37.8}\pm 2.2$

Figure 3: Ablation study results on MLE-Bench Lite. Results for the Medium and High splits are provided in the Appendix F.3.

We evaluate the individual components of MARS on 22 competitions from the MLE-Bench Lite dataset (Figure 3).

Modular Decomposition. As illustrated in Figure 3(a), removing the modular decomposition component degrades the agent’s overall success, validating its ability to reduce context noise and improve testability across independent functional modules.

Memory Mechanisms. Figure 3(b) highlights the critical importance of our memory architecture. Complete removal of memory mechanisms results in a drastic performance drop. Furthermore, to isolate the specific benefits of comparative distillation, we tested a variant of our method against a baseline utilizing only empirical analysis to distill lessons, without comparing them to the best-known solution. The results demonstrate that the comparative component provides a consistent performance boost. This confirms that code-level comparisons are essential for isolating causal factors. Without this “delta” analysis, the agent often over-generalizes from noisy logs. In contrast, comparative reflective memory allows it to distill precise, actionable heuristics, effectively mirroring human-led ablation experiments.

Tree Search Strategies. Figure 3(c) compares different tree search algorithms for MARS. While Greedy Search selects the node with the best validation metric for expansion at each step, Vanilla MCTS operates as a variant of Budget-Aware MCTS where $w=0$ (Eq. 4). The results indicate that the proposed Budget-Aware MCTS consistently yields superior performance over time by effectively balancing exploration with resource constraints.

Penalty Weight $w$ . Figure 3(d) evaluates Budget-Aware MCTS across varying penalty weights. The default setting of $w=-0.07$ proves optimal. Removing the penalty entirely ( $w=0$ ) causes performance degradation, underscoring the importance of penalizing long execution times alongside rewarding performance. Conversely, a stronger penalty ( $w=-0.15$ ) excessively biases the reward toward latency, causing the search to prioritize trivial, fast nodes over high-performing ones.

We investigate whether Modular Decomposition facilitates the construction of complex solutions for each task. Table 4 compares the repository statistics of MARS with and without modular decomposition for the best solution. The results show that the modular approach encourages the generation of more extensive and structured codebases (measured by lines of code and number of files in the best solution). To illustrate this structural adaptability, Table 5 enumerates the specific modules synthesized for five representative competitions. The diversity of these modules – tailored to specific sub-tasks such as preprocessing and model architecture – demonstrates the agent’s ability to decompose intricate problems into logical components. This capacity to architect organized, repository-level solutions closely mirrors professional software engineering workflows.

Original Table 4 — Table 4: Comparison of repository statistics between MARS and the variant without Modular Decomposition on MLE-Bench Lite.

Table 4
Metric	MARS without Modular	MARS
Lines of Code	$474.8\pm 13.5$	$1103.9\pm 35.9$
Number of Files	$1.0\pm 0.0$	$6.7\pm 0.1$

Table 5: Modules generated by MARS on challenging competitions.

Figure 4: Reward modulation: Budget-aware MCTS assigns higher rewards to faster candidates when performance is comparable.

We examine whether Budget-aware MCTS discovers high-quality solutions more frequently than the Vanilla MCTS. We define the effective solution rate as the proportion of explored solutions that improve upon the current best validation metric per task. Empirically, Budget-aware MCTS achieves an effective solution rate of $19.5\%\pm 1.5\%$ , notably higher than the $16.1\%\pm 1.3\%$ observed with Vanilla MCTS. This suggests that the latency penalty acts as a heuristic to prune inefficient trajectories. As illustrated in Figure 4, when the agent encounters solutions with comparable accuracy but differing costs, our efficiency-guided reward favors the faster candidate. This bias directs computational resources toward efficient nodes, accelerating the discovery of optimal solutions within the time limit.

We examine the role of Lesson Learning in guiding the agent’s solution exploration. Figure 1 illustrates an example where the agent formulates lessons from early failures or partial successes and applies them to refine subsequent solutions. To quantify this behavior, we introduce two metrics: the lesson-utilization rate (the proportion of solutions that incorporate existing lessons) and the lesson-transfer rate (the proportion of utilized solution lessons originating from a different tree branch). MARS achieves a lesson-utilization rate of $65.8\%\pm 1.1\%$ and a lesson-transfer rate of $63.0\%\pm 1.8\%$ on MLE-Bench. These results demonstrate that the agent actively leverages learned knowledge and cross-branch transfer to steer the search toward high-quality strategies.

Figure 5: Distribution of maximum code similarity scores for medal-winning submissions from AIRA-dojo and MARS, compared against top public Kaggle notebooks.

To verify compliance, we employ the official MLE-Bench log analysis tool, which utilizes gpt-4.1-mini to audit the logs and code outputs of all medal-winning submissions. The evaluation confirms that MARS strictly adheres to the protocol, registering a 0% violation rate across all monitored dimensions, including “Tried to access unauthorized resources”, “Tried to call external LLM API service”, and “Manually-written submission”. Furthermore, we assess code originality using the provided plagiarism detection tool based on Dolos (Maertens et al., 2024). We calculate the maximum similarity score between our agent’s code – concatenated into a single file for multi-module repositories – and the top public notebooks for each competition. As shown in Figure 5, the similarity distribution of MARS mirrors that of the baseline AIRA-dojo. Crucially, no submission exceeds a 60% similarity threshold, demonstrating that MARS generates distinct, original solutions rather than reproducing existing public code.

To verify whether MARS performs genuine credit assignment, we conducted a systematic evaluation of its distilled lessons. First, we utilized Claude 4.6 Sonnet to audit all 3,611 solution lessons generated by Gemini-3-Pro-Preview for causal accuracy. This audit confirmed that 88.34% of the lessons correctly attributed validation metric shifts to specific code changes, rather than relying on hallucinated narratives. Second, we manually inspected a representative subset of 20 lessons. This human evaluation yielded a 90% causal accuracy rate, confirming that the LLM-based audit serves as a reliable proxy for expert judgment. Together, this analysis provides concrete evidence that our method effectively resolves the credit assignment problem. Qualitative examples demonstrating this are provided in the Appendix G.

As detailed in Appendix F.4, maintaining a comprehensive memory context results in a higher cost for the 24-hour version of MARS compared to AIRA-dojo ($60.5 vs. $39.0). However, this investment nearly doubles the Any Medal Rate (from 24.4% to 43.1%). More importantly, a cost-controlled evaluation demonstrates that a 4-hour version of MARS achieves a higher Any Medal Rate (28.4%) while spending significantly less ($9.6) than the 24-hour baselines. This Pareto improvement confirms that our performance gains stem from architectural efficiency rather than simply a larger resource budget.

In this work, we addressed the limitations of current autonomous agents in Long-Horizon AI Research by introducing MARS. Unlike traditional code generation approaches, our framework treats research as a rigorous, repository-level engineering challenge. By integrating Resource-Aware Planning via Budget-Aware MCTS, Modular Construction, and Comparative Reflective Memory, MARS effectively resolves the credit assignment problem while balancing exploration with computational efficiency. Our extensive evaluation on MLE-Bench demonstrates that this structured approach – mimicking the strategic foresight of human engineers – enables state-of-the-art performance in complex Machine Learning Engineering tasks. Future work will focus on extending MARS to broader scientific discovery domains and optimizing the framework’s economic viability through advanced context caching and early stopping mechanisms.

MARS contributes to the advancement of autonomous AI agents. While our work aims to enhance the reliability and efficiency of automated software engineering, we acknowledge potential broader impacts. The deployment of LLM-based agents involves risks related to the generation of incorrect or hallucinatory code; we mitigate this through iterative self-correction with code execution feedback. We do not foresee immediate negative societal consequences beyond those generally associated with the advancement of generative AI.

References

Chan et al. (2025) J. S. Chan, N. Chowdhury, O. Jaffe, J. Aung, D. Sherburn, E. Mays, G. Starace, K. Liu, L. Maksin, T. Patwardhan, et al. MLE-bench: evaluating machine learning agents on machine learning engineering. In International Conference on Learning Representations, Vol. 2025, pp. 50466–50494. Cited by: §1, §2, §5.1.
Forouzandeh et al. (2025) S. Forouzandeh, W. Peng, P. Moradi, X. Yu, and M. Jalili Learning hierarchical procedural memory for llm agents through bayesian selection and contrastive refinement. arXiv preprint arXiv:2512.18950. Cited by: §2.
Huang et al. (2024) Q. Huang, J. Vora, P. Liang, and J. Leskovec MLAgentBench: evaluating language agents on machine learning experimentation. In International Conference on Machine Learning, pp. 20271–20309. Cited by: §1.
Jansen et al. (2025) P. Jansen, O. Tafjord, M. Radensky, P. Siangliulue, T. Hope, B. Dalvi Mishra, B. P. Majumder, D. S. Weld, and P. Clark CodeScientist: end-to-end semi-automated scientific discovery with code-based experimentation. In Findings of the Association for Computational Linguistics: ACL 2025, W. Che, J. Nabende, E. Shutova, and M. T. Pilehvar (Eds.), Vienna, Austria. Cited by: §2.
Jiang et al. (2025) Z. Jiang, D. Schmidt, D. Srikanth, D. Xu, I. Kaplan, D. Jacenko, and Y. Wu Aide: ai-driven exploration in the space of code. arXiv preprint arXiv:2502.13138. Cited by: Table 6, Table 7, Table 8, Table 9, §1, §1, §2, §2, Table 1, Table 2, §5.1.
Jimenez et al. (2023) C. E. Jimenez, J. Yang, A. Wettig, S. Yao, K. Pei, O. Press, and K. Narasimhan SWE-bench: can language models resolve real-world github issues?. CoRR abs/2310.06770. External Links: Link Cited by: §1.Source
Kocsis and Szepesvári (2006) L. Kocsis and C. Szepesvári Bandit based monte-carlo planning. In European conference on machine learning, pp. 282–293. Cited by: item 1, §2.
Li et al. (2025) A. Li, C. Wu, Z. Ge, Y. H. Chong, Z. Hou, L. Cao, C. Ju, J. Wu, H. Li, H. Zhang, et al. The fm agent. arXiv preprint arXiv:2510.26144. Cited by: Table 6, Table 7, Table 8, Table 9, §2, Table 2.
Li et al. (2022) Y. Li, D. Choi, J. Chung, N. Kushman, J. Schrittwieser, R. Leblond, T. Eccles, J. Keeling, F. Gimeno, A. Dal Lago, et al. Competition-level code generation with alphacode. Science 378 (6624), pp. 1092–1097. Cited by: §1.
Liu et al. (2025a) T. Liu, Z. Wang, J. Miao, I. Hsu, J. Yan, J. Chen, R. Han, F. Xu, Y. Chen, K. Jiang, et al. Budget-aware tool-use enables effective agent scaling. arXiv preprint arXiv:2511.17006. Cited by: §2.
Liu et al. (2025b) Z. Liu, Y. Cai, X. Zhu, Y. Zheng, R. Chen, Y. Wen, Y. Wang, S. Chen, et al. ML-master: towards ai-for-ai via integration of exploration and reasoning. arXiv preprint arXiv:2506.16499. Cited by: Table 6, Table 7, Table 8, Table 9, §2, §2, Table 2.
Lu et al. (2024) C. Lu, C. Lu, R. T. Lange, J. Foerster, J. Clune, and D. Ha The AI scientist: towards fully automated open-ended scientific discovery. arXiv preprint arXiv:2408.06292. Cited by: §2.
Maertens et al. (2024) R. Maertens, M. Van Neyghem, M. Geldhof, C. Van Petegem, N. Strijbol, P. Dawyndt, and B. Mesuere Discovering and exploring cases of educational source code plagiarism with dolos. SoftwareX 26, pp. 101755. Cited by: §6.
Nadaf et al. (2025) A. Nadaf, A. Mohammadshahi, M. Yazdani, and Leeroo Coding Agent Kapso: a knowledge-grounded framework for autonomous program synthesis and optimization. External Links: Link Cited by: Table 6, Table 7, Table 8, Table 9, Table 2.Source
Nam et al. (2026) J. Nam, J. Yoon, J. Chen, J. Shin, S. Arik, and T. Pfister Mle-star: machine learning engineering agent via search and targeted refinement. Advances in Neural Information Processing Systems 38, pp. 116692–116712. Cited by: §2, Table 1.
Ouyang et al. (2026) S. Ouyang, J. Yan, I. Hsu, Y. Chen, K. Jiang, Z. Wang, R. Han, L. Le, S. Daruki, X. Tang, V. Tirumalashetty, G. Lee, M. Rofouei, H. Lin, J. Han, C. Lee, and T. Pfister ReasoningBank: scaling agent self-evolving with reasoning memory. In The Fourteenth International Conference on Learning Representations, Cited by: §1, §2.
Packer et al. (2023) C. Packer, S. Wooders, K. Lin, V. Fang, S. G. Patil, I. Stoica, and J. E. Gonzalez MemGPT: towards llms as operating systems. arXiv preprint arXiv:2310.08560. Cited by: §1.
Pu et al. (2025) Y. Pu, T. Lin, and H. Chen PiFlow: principle-aware scientific discovery with multi-agent collaboration. arXiv preprint arXiv:2505.15047. Cited by: §2.
Schmidgall et al. (2025) S. Schmidgall, Y. Su, Z. Wang, X. Sun, J. Wu, X. Yu, J. Liu, M. Moor, Z. Liu, and E. Barsoum Agent laboratory: using LLM agents as research assistants. Findings of the Association for Computational Linguistics: EMNLP 2025, pp. 5977–6043. Cited by: §2.
Shinn et al. (2023) N. Shinn, F. Cassano, A. Gopinath, K. Narasimhan, and S. Yao Reflexion: language agents with verbal reinforcement learning. Advances in Neural Information Processing Systems 36, pp. 8634–8652. Cited by: §1, §2.
Starace et al. (2025) G. Starace, O. Jaffe, D. Sherburn, J. Aung, J. S. Chan, L. Maksin, R. Dias, E. Mays, B. Kinsella, W. Thompson, J. Heidecke, A. Glaese, and T. Patwardhan PaperBench: evaluating AI’s ability to replicate AI research. In Forty-second International Conference on Machine Learning, Cited by: Appendix A, §1, §2.
Tan et al. (2019) M. Tan, B. Chen, R. Pang, V. Vasudevan, M. Sandler, A. Howard, and Q. V. Le Mnasnet: platform-aware neural architecture search for mobile. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 2820–2828. Cited by: §4.4.3, §5.1.
Tan et al. (2025) Z. Tan, J. Yan, I. Hsu, R. Han, Z. Wang, L. Le, Y. Song, Y. Chen, H. Palangi, G. Lee, et al. In prospect and retrospect: reflective memory management for long-term personalized dialogue agents. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 8416–8439. Cited by: §2.
Tang et al. (2026) J. Tang, L. Xia, Z. Li, and C. Huang Ai-researcher: autonomous scientific innovation. Advances in Neural Information Processing Systems 38, pp. 9481–9520. Cited by: §2.
Team et al. (2025) N. Team, B. Zhang, S. Feng, X. Yan, J. Yuan, Z. Yu, X. He, S. Huang, S. Hou, Z. Nie, et al. NovelSeek: when agent becomes the scientist–building closed-loop system from hypothesis to verification. arXiv preprint arXiv:2505.16938. Cited by: Table 6, Table 7, Table 8, Table 9, §2, §2, Table 2.
Tian et al. (2024) M. Tian, L. Gao, S. Zhang, X. Chen, C. Fan, X. Guo, R. Haas, P. Ji, K. Krongchon, Y. Li, et al. Scicode: a research coding benchmark curated by scientists. Advances in Neural Information Processing Systems 37, pp. 30624–30650. Cited by: §1.
Toledo et al. (2026) E. Toledo, K. Hambardzumyan, M. Josifoski, R. Hazra, N. Baldwin, A. Audran-Reiss, M. Kuchnik, D. Magka, M. Jiang, A. Lupidi, et al. Ai research agents for machine learning: search, exploration, and generalization in mle-bench. Advances in Neural Information Processing Systems 38, pp. 35309–35348. Cited by: Table 6, Table 7, Table 8, Table 9, §1, §2, Table 1, Table 2, §5.1.
Wang et al. (2025) X. Wang, B. Li, Y. Song, F. F. Xu, X. Tang, M. Zhuge, J. Pan, Y. Song, B. Li, J. Singh, et al. Openhands: an open platform for ai software developers as generalist agents. In International Conference on Learning Representations, Vol. 2025, pp. 65882–65919. Cited by: §1.
Weng et al. (2025) Y. Weng, M. Zhu, Q. Xie, Q. Sun, Z. Lin, S. Liu, and Y. Zhang DeepScientist: advancing frontier-pushing scientific findings progressively. arXiv preprint arXiv:2509.26603. Cited by: §2.
Wijk et al. (2025) H. Wijk, T. R. Lin, J. Becker, S. Jawhar, N. Parikh, T. Broadley, L. Chan, M. Chen, J. M. Clymer, J. Dhyani, E. Ericheva, K. Garcia, B. Goodrich, N. Jurkovic, M. Kinniment, A. Lajko, S. Nix, L. J. K. Sato, W. Saunders, M. Taran, B. West, and E. Barnes RE-bench: evaluating frontier AI r&d capabilities of language model agents against human experts. In Forty-second International Conference on Machine Learning, Cited by: Appendix A, §1, §2.
Xu et al. (2026) W. Xu, Z. Liang, K. Mei, H. Gao, J. Tan, and Y. Zhang A-mem: agentic memory for llm agents. Advances in Neural Information Processing Systems 38, pp. 17577–17604. Cited by: §1, §2.
Yamada et al. (2025) Y. Yamada, R. T. Lange, C. Lu, S. Hu, C. Lu, J. Foerster, J. Clune, and D. Ha The ai scientist-v2: workshop-level automated scientific discovery via agentic tree search. arXiv preprint arXiv:2504.08066. Cited by: §1, §2.
Yang et al. (2024) J. Yang, C. E. Jimenez, A. Wettig, K. Lieret, S. Yao, K. Narasimhan, and O. Press SWE-agent: agent-computer interfaces enable automated software engineering. Advances in Neural Information Processing Systems 37, pp. 50528–50652. Cited by: §1.
Yang et al. (2025) X. Yang, X. Yang, S. Fang, Y. Zhang, J. Wang, B. Xian, Q. Li, J. Li, M. Xu, Y. Li, et al. R&D-agent: an llm-agent framework towards autonomous data science. arXiv preprint arXiv:2505.14738. Cited by: Table 6, Table 7, Table 8, Table 9, §2, Table 1, Table 2.
Zhao et al. (2024) A. Zhao, D. Huang, Q. Xu, M. Lin, Y. Liu, and G. Huang Expel: llm agents are experiential learners. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 38, pp. 19632–19642. Cited by: §2.
Zhu et al. (2026) X. Zhu, Y. Cai, Z. Liu, B. Zheng, C. Wang, R. Ye, J. Chen, H. Wang, W. Wang, Y. Zhang, et al. Toward ultra-long-horizon agentic science: cognitive accumulation for machine learning engineering. arXiv preprint arXiv:2601.10402. Cited by: Table 6, Table 7, Table 8, Table 9, §2, §2, Table 1, Table 2.

Appendix

This Appendix is organized as follows: Appendix A discusses the limitations of our work. We then provide background on the MLE task scenario and the standard MCTS algorithm in Appendices B and C, respectively. Next, Appendix D details the instantiation of MARS for MLE tasks, and Appendix E contrasts our evaluation setup with those of prior agents. Finally, we present additional experimental results (Appendix F), qualitative examples of causally correct distilled lessons (Appendix G), comprehensive agent prompts (Appendix H), and representative code examples generated by our system (Appendix I).

While MARS demonstrates significant improvements in autonomous machine learning engineering, we acknowledge several limitations that present important avenues for future research.

Scope of Evaluation. Our evaluation is currently conducted exclusively on the MLE-Bench dataset. While MLE-Bench is a rigorous and comprehensive benchmark for applied ML engineering, it does not encompass the full spectrum of open-ended tasks implied by general “Automated AI Research,” such as novel algorithm formulation, literature synthesis, or deployment in unconstrained real-world environments. Future work should evaluate the framework across a broader range of domains, including research engineering Wijk et al. (2025) and automated research replication Starace et al. (2025), to fully substantiate its broader applicability.

Computational Cost. The current pipeline incurs a notable computational expense. This cost primarily stems from the extensive API calls required for iterative tree search and comparative distillation, which may serve as a barrier to widespread adoption by independent researchers. To mitigate this, future iterations must prioritize inference efficiency. Potential cost-reduction strategies include implementing KV-cache and prompt caching for unchanged modular environments, introducing early-stopping mechanisms within the Budget-Aware MCTS when performance improvements plateau, and routing simpler implementation sub-tasks to smaller, more cost-effective models.

Reliability of Distilled Lessons. Finally, our approach relies on LLMs to deduce causal relationships between code modifications and performance metrics. Consequently, the distilled lessons in the memory pool remain susceptible to hallucinations or misattributions, where the model might incorrectly credit an incidental code change for a performance gain. Although our comparative “delta” analysis significantly mitigates this risk by grounding reflections in specific code diffs, it does not entirely eliminate it. Future frameworks could improve lesson reliability by integrating formal program analysis, cross-model verification, or lesson-confidence scoring based on repeated empirical validation across multiple runs.

Machine Learning Engineering (MLE) is a representative and challenging instantiation of this general problem class. MLE requires the agent not just to write a snippet of code, but to engineer a full pipeline that processes data, trains models, and validates results.

We map the general problem $\mathcal{P}=(\mathcal{I},\mathcal{E},\mathcal{O})$ to an MLE task $Q=(I,D,M)$ , where:

•
$I$ corresponds to the natural language task description ( $\mathcal{I}$ ).
•
$D$ represents the datasets ( $D=\{D_{dev},D_{test}\}$ ) which form the data environment ( $\mathcal{E}$ ).
•
$M$ is the evaluation metric (e.g., Accuracy, F1-score) defining the objective ( $\mathcal{O}$ ). Without loss of generality, we treat the optimization of $M$ as a maximization problem.

If a pre-defined validation set is not provided in the development set $D_{dev}$ , the agent must partition $D_{dev}$ to create a validation set $D_{val}$ for internal evaluation, as the test set $D_{test}$ is strictly hidden.

We aim to build an MLE agent $\mathcal{A}$ that explores a space of possible solutions and outputs a final executable solution $s$ . We define the solution $s$ as a structured software repository comprising the distinct code modules, dependencies, and entry points required to orchestrate the end-to-end pipeline.

The performance of a solution $s$ is quantified by the metric function $f(s,D,M)\in\mathbb{R}$ . While the ultimate goal is to maximize performance on the unseen test set $D_{test}$ , the agent must rely on a proxy objective using the validation set $D_{val}$ . The optimization objective becomes:

\displaystyle s^{*}=\argmax_{s\in\mathcal{S}_{\mathcal{A}}}\quad f(s,D_{test},M),\quad s.t.\quad C(s)\leq T

where $T$ is the wall-clock time budget, $\mathcal{S}_{\mathcal{A}}$ is the set of candidate solutions generated by agent $\mathcal{A}$ given task $Q$ , and $C(s)$ denotes the total wall-clock time consumed by the agent to search for the solution $s$ . Since $D_{test}$ is unobservable, the agent optimizes via $f(s,D_{val},M)$ .

Monte Carlo Tree Search (MCTS) is a heuristic search algorithm for decision processes, most notably employed in game play. The algorithm builds a search tree where each node $v$ represents a state $s$ , and each edge represents an action $a$ leading to a new state. The value of a state is estimated by simulating outcomes from that state. As shown in Algorithm 1, each MCTS iteration consists of four distinct phases:

1.
Selection: Starting from the root node $v_{0}$ , the algorithm recursively traverses down the tree by selecting child nodes according to a selection policy, typically aiming to balance exploration and exploitation. a common strategy is the Upper Confidence Bound for Trees (UCT) Kocsis and Szepesvári (2006):
$\displaystyle a^{*}=\argmax_{a\in\mathcal{A}(s)}\left(Q(s,a)+c_{uct}\sqrt{\frac{\ln N(s)}{N(s,a)}}\right)$
where $Q(s,a)$ is the estimated value of taking action $a$ in state $s$ , $N(s)$ is the total visit count of state $s$ , $N(s,a)$ is the number of times action $a$ has been selected from $s$ , and $c_{uct}$ is a constant controlling the exploration weight.
2.
Expansion: Once a leaf node $v_{l}$ is reached (or a node with unexplored actions), one or more child nodes are added to the tree, representing reachable states from standard actions.
3.
Simulation: From the newly expanded node, a rollout policy (often random or heuristic-based) is executed to simulate a sequence of actions until a terminal state is reached or a resource limit is met. This produces a reward $R$ .
4.
Backpropagation: The reward $R$ obtained from the simulation is propagated back up the tree from the leaf to the root. For each node $(s,a)$ traversed during the selection phase, we update the visit count and value estimate as follows:
$\begin{aligned}\displaystyle N(s,a) & \displaystyle\leftarrow N(s,a)+1\end{aligned}$
$\begin{aligned}\displaystyle Q(s,a) & \displaystyle\leftarrow Q(s,a)+\frac{R-Q(s,a)}{N(s,a)}\end{aligned}$

In our MARS framework, we adapt MCTS to the space of automated AI Research. A state $s$ corresponds to a partial or complete solution $s_{n}$ , and actions correspond to modification operators (Drafting, Improvement, Debugging). The reward is derived from the efficiency-guided validation performance.

Algorithm 1 Monte Carlo Tree Search (MCTS)

1: Input: Task

\mathcal{P}

, Time Budget

T

2: Output: Best Solution Node

v^{*}

3: Initialize root node

v_{0}

with empty solution

v^{*}\leftarrow v_{0}

5: while Time used

<T

v_{l}\leftarrow\text{SELECT}(v_{0})

{Tree Traversal using UCT}

v_{new}\leftarrow\text{EXPAND}(v_{l})

{Apply Drafting/Improvement/Debugging}

R\leftarrow\text{SIMULATE}(v_{new})

{Execute and Evaluate Solution}

\text{BACKPROPAGATE}(v_{new},R)

{Update

Q

and

N

values}

10: if

ValMetric(v_{new})>ValMetric(v^{*})

then

11:

v^{*}\leftarrow v_{new}

12: end if

13: end while

14: return

v^{*}

In this section, we detail the instantiation of MARS for Machine Learning Engineering (MLE) tasks. The comprehensive procedure is formalized in Algorithm 2. Corresponding instruction prompts for the agents involved are provided in Appendix H.

The workflow initiates by formalizing the optimization objective through task metadata extraction. A Metric Extraction Agent parses the natural language task description $\mathcal{I}$ to identify the primary evaluation metric $M$ and the optimization direction $d\in\{\text{maximize},\text{minimize}\}$ .

Simultaneously, a Multi-Agent Subsystem processes the raw data to generate metadata descriptors (e.g., sample IDs) for the training ( $D_{train}$ ), validation ( $D_{val}$ ), and test ( $D_{test}$ ) sets. These metadata descriptors are saved to files for later usage.

To ensure robust evaluation, we employ a strict protocol:

•
Validation Dataset Creation: If a pre-defined validation set is not provided, the agent performs a stratified or group-based split (defaulting to a 80:20 ratio) on $D_{dev}$ to create $D_{train}$ and $D_{val}$ . This ensures that $D_{val}$ maintains a distribution $P(D_{val})\approx P(D_{train})$ , enabling reliable proxy evaluation.
•
Verification & Documentation: Distinct agents perform key integrity checks (e.g., no leakage between splits) and generate comprehensive documentation describing the data schema and split logic.

Following preparation, a Data Analysis Agent performs Exploratory Data Analysis (EDA) on $D_{train}$ . This agent generates a detailed report highlighting data distributions and potential correlations, which serves as a critical reference for feature engineering during the solution exploration. Furthermore, a Search Agent identifies $K_{a}$ candidate model architectures across diverse algorithmic families (e.g., gradient-boosted trees, deep neural networks) using web search tools.

Once initialized, MARS enters an iterative Tree Search Stage. In each iteration, a node $v$ is selected via the Upper Confidence Bound for Trees (UCT) formula. If the root node is selected, the system enters the Draft Phase; otherwise, it proceeds to the Improvement Phase. Following code generation, a Debugging Loop is triggered to resolve execution errors, after which the results are reviewed, lessons are distilled, and rewards are backpropagated.

This phase initializes new branches of the search tree using a curriculum-based strategy that progresses from simple baselines to sophisticated ensembles.

•
Initial Seed: When the solution lesson pool $\mathcal{L}_{solution}$ is empty, an Initial Idea Generation Agent proposes a solution based on the most lightweight model from the $K_{a}$ candidates.
•
Evolutionary Growth: As lessons accumulate, an Idea Improvement Agent formulates advanced proposals by integrating insights from $\mathcal{L}_{solution}$ .
•
Modular Implementation: A Modular Agent decomposes the proposed idea into independent functional units, which are implemented and unit-tested by a Coding Agent before being orchestrated into a final execution script $\pi_{\text{main}}$ .

This phase focuses on local optimization. An agent analyzes the current solution and its performance metrics to propose targeted, ablation-style modifications. By leveraging the learned lessons in $\mathcal{L}_{solution}$ , the agent avoids previously identified pitfalls and focuses on high-impact refinements (e.g., hyperparameter tuning or feature engineering).

If a candidate node $v_{new}$ fails execution, the system enters a debugging loop (up to $K_{debug}$ attempts). We maintain a dedicated debugging lesson pool $\mathcal{L}_{debug}$ to store error-correction patterns. This prevents the agent from repeating previous mistakes in subsequent iterations.

Algorithm 2 MARS for MLE Tasks

1: Input: Task Description

I

, Raw Dataset

D

, Time Limit

T

2: Output: Optimized solution code repository

s^{*}

d\leftarrow\text{MetricParsing}(I)

{Extract optimization objective and direction}

\mathcal{C}_{meta}\leftarrow\text{Preprocess}(D)

{Generate metadata and stratified splits}

\mathcal{C}_{eda}\leftarrow\text{Analyze}(I,D,\mathcal{C}_{meta})

{Perform EDA and statistical profiling}

\mathcal{C}_{model}\leftarrow\text{SearchArchitectures}(I)

{Retrieve SOTA model candidates via search}

\mathcal{C}=\{\mathcal{C}_{meta},\mathcal{C}_{eda},\mathcal{C}_{model}\}

v_{root}\leftarrow\text{InitializeTree}()

v^{*}\leftarrow\text{None}

10:

\mathcal{L}_{solution}\leftarrow\emptyset

{Solution Lesson Pool}

11:

\mathcal{L}_{debug}\leftarrow\emptyset

{Debug Lesson Pool}

12:

\mathcal{Y}\leftarrow\emptyset

{Explored Ideas}

13: while

\text{Time used}<T

14:

v\leftarrow\text{SelectNode}(v_{root})

{Using UCT selection}

15: if

v

v_{root}

then

16:

Y\leftarrow\text{ProposeIdea}(I,\mathcal{Y},\mathcal{C},\mathcal{L}_{solution})

{Curriculum-based idea generation}

17:

Z\leftarrow\text{ProposeModules}(I,Y)

{Decompose idea into functional modules}

18:

\{\mathcal{M}_{j}\}_{j=1}^{l}\leftarrow\text{ImplementModules}(I,Y,Z)

{Implement modular components}

19:

\{\mathcal{M}_{j}\}_{j=1}^{l}\leftarrow\text{DebugModules}(I,\{\mathcal{M}_{j}\}_{j=1}^{l})

{Unit-test modules}

20:

\pi_{\text{main}}\leftarrow\text{ImplementMainScript}(I,Y,\{\mathcal{M}_{j}\}_{j=1}^{l})

{Orchestrate pipeline}

21:

v_{new}\leftarrow\text{DraftNode}(\{\mathcal{M}_{j}\}_{j=1}^{l},\pi_{\text{main}})

22:

\mathcal{Y}\leftarrow\mathcal{Y}\cup\{Y\}

23: else

24:

v_{new}\leftarrow\text{ImproveNode}(v,\mathcal{L}_{solution})

{Ablation-style local optimization}

25: end if

26:

k\leftarrow 0

27: while

\text{IsBuggy}(v_{new})

and

k<N_{d}

28:

v_{new}\leftarrow\text{DebugNode}(v_{new},\mathcal{L}_{debug})

{Apply

\mathcal{L}_{debug}

for debugging and then update

\mathcal{L}_{debug}

}

29:

k\leftarrow k+1

30: end while

31:

r_{e}\leftarrow\text{ExecuteAndReview}(v_{new})

{Execute code and review execution results}

32:

\text{ExtractLesson}(v_{new},r_{e},\mathcal{L}_{solution})

{Distill lessons from results}

33:

\text{Backpropagate}(v_{new},r_{e})

{Update tree statistics with rewards}

34: if

v^{*}\text{is None}

\text{IsImprovedMetric}(r_{e},v^{*},d)

then

35:

v^{*}\leftarrow v_{new}

36: end if

37: end while

38:

s^{*}\leftarrow\text{GetRepoCode}(v^{*})

39: return

s^{*}

Since MLE-Bench allows for open-ended submissions with varying computational budgets and system architectures, direct comparisons on the official leaderboard can be influenced by hardware disparities. To ensure a fair assessment, we detail the specific hardware, time limits, and auxiliary resources used by top-performing leaderboard agents alongside our own in Table 6. In our Controlled Evaluation (AIDE, AIRA-dojo, and MARS), we standardize the environment to a single A100 GPU node with no external knowledge bases to isolate algorithmic effectiveness from resource scaling.

Table 6: Comparison of leaderboard agents’ setup and our agent’s setup.

This section presents a comprehensive evaluation across the various subsets of MLE-Bench. Detailed performance metrics for the Lite, Medium, and High splits are provided in Tables 7, 8, and 9, respectively.

Table 7: Performance comparison on MLE-Bench Lite. Results are reported as mean $\pm$ SEM across three independent runs. All values are in percentages (%). The best performance is highlighted in **bold**, and the second-best is underlined.

Table 8: Performance comparison on MLE-Bench Medium. Results are reported as mean $\pm$ SEM across three independent runs. All values are in percentages (%). The best performance is highlighted in **bold**, and the second-best is underlined.

Table 9: Performance comparison on MLE-Bench High. Results are reported as mean $\pm$ SEM across three independent runs. All values are in percentages (%). The best performance is highlighted in **bold**, and the second-best is underlined.

To demonstrate that the success of MARS stems from its architectural robustness rather than being overly fitted to a specific language model, we conducted a cross-model evaluation utilizing Claude 4.6 Sonnet. The results in Table 10 confirm that MARS maintains a substantial performance advantage over the baseline across all difficulty splits, verifying its efficacy independent of the underlying reasoning engine.

Original Table 10 — Table 10: Model generalizability on MLE-Bench using Claude 4.6 Sonnet. Values represent Any Medal Rate (Mean $\pm$ SEM) over three independent runs, evaluated on a random sample ( $n=10$ per split) of Lite, Medium, and High tasks.

Table 10
Method	Lite	Medium	High	All
AIRA-dojo	40.0 $\pm$ 5.8	6.7 $\pm$ 3.3	13.3 $\pm$ 3.3	20.0 $\pm$ 1.9
MARS (ours)	80.0 $\pm$ 0.0	43.3 $\pm$ 3.3	43.3 $\pm$ 3.3	55.6 $\pm$ 1.1

Figure 6: Ablation study results on a random sample of 10 tasks from the MLE-Bench Medium split.

Figure 7: Ablation study results on a random sample of 10 tasks from the MLE-Bench High split.

Building upon the ablation study conducted on the MLE-Bench Lite split (Section 5.3), we extend our evaluation to include a random sample of 10 tasks each from the Medium and High splits. As illustrated in Figures 6 and 7, the performance advantages provided by our core modules remain highly consistent across these more challenging tasks. Furthermore, the performance gaps between our full method and the ablated variants widen on the harder splits, demonstrating that components like modular decomposition and comparative memory distillation become increasingly essential as task complexity scales.

In this section, we analyze the computational costs associated with our framework compared to the baselines (Table 11). When operating under the same 24-hour time limit as the baselines, MARS exhibits a distinct resource profile. It achieves the lowest generation volume, producing fewer output tokens ( $6.0\times 10^{5}$ ) than both baselines and making significantly fewer API calls than AIRA-dojo (594.8 vs. 867.0). However, its input consumption is substantial ( $286.6\times 10^{5}$ tokens) – approximately 2.2 times that of AIRA-dojo. This increase is largely due to maintaining a comprehensive memory context filled with learned lessons and modular structures. Because the Gemini-2.5-Pro model applies premium pricing for long-context prompts (over 200k tokens), the 24-hour version of MARS incurs a higher total cost ($60.5) than AIRA-dojo ($39.0). Crucially, this investment yields substantial returns: the Any Medal Rate increases from 24.4% to 43.1%, justifying the higher expense through superior efficacy.

To demonstrate that the performance of MARS is driven by architectural efficiency rather than simply a larger financial budget, we also conducted a cost-controlled evaluation. As shown in Table 11, MARS running under a restricted 4-hour time limit achieves a higher Any Medal Rate (28.4%) while spending less ($9.6) than both 24-hour baselines (AIDE and AIRA-dojo). This Pareto improvement – achieving better results with a smaller budget – directly substantiates our claim that the specific architectural choices of MARS are responsible for the performance gains, not merely an increased allocation of resources.

In this section, we provide two qualitative examples demonstrating the causal accuracy of the lessons distilled by our framework.

MARS diagnosed a 0.28% AUC performance drop in a Dual-Stream architecture compared to an Early Fusion baseline. By correlating architectural diffs with execution-log overfitting, the agent synthesized a transferable lesson: Early Fusion is superior for heterogeneous tabular data because it preserves low-level cross-modal interactions – such as category-specific continuous thresholds – that independent projection streams discard as an information bottleneck.

MARS diagnosed why a multi-model ensemble underperformed (0.976 AUC) against its own single-model component (0.982 AUC). By cross-referencing code with execution logs, it identified a convergence gap: the ConvNeXt model matured significantly faster than the EfficientNet model.The agent’s causal analysis determined that unweighted soft voting diluted the high-quality signals of the lead model with the under-optimized predictions of the weaker one. Instead of simply discarding the result, the agent synthesized a negative constraint: avoid a naive ensemble of disparate performers. This demonstrates the system’s ability to prune the search space by extracting principled heuristics from failed experiments.

This section provides the full suite of instruction prompts utilized by MARS to orchestrate the various agents involved in solving MLE tasks.

Graphic from the paper; its text follows

In this section, we showcase the final solution generated for the iMet-2020-fgvc7 competition. This repository was produced by MARS utilizing Gemini-3-Pro-Preview and achieved the highest validation performance during the search process.

MARS: Modular Agent with Reflective Search for Automated AI Research

Abstract

References

Meow summary

The problem MARS tackles

How the method works

What the tests show

Where the evidence stops

Make notes of your own

Abstract

1 Introduction

2 Related Work

Automated AI Research & MLE Bottleneck.

3 Problem

MLE Task Scenario.

4 Method

4.1 Overall Framework

4.2 Modular Decomposition

4.3 Comparative Reflective Memory

Solution Improvement via Comparative Reflection.

Debugging Lessons.

Lesson Management.

Lesson Utilization.

4.4 Budget-Aware MCTS

4.4.1 Actions and Expansion

4.4.2 Node Selection

4.4.3 Efficiency-Guided Reward Function

4.5 Task Specific Components

5 Experiment

5.1 Setup

5.2 Main Results

5.3 Ablation Study

6 Discussions

How does Modular Decomposition impact solution complexity?

Does Budget-aware MCTS improve exploration?

How lessons guide the evolution process?

Does MARS follow the MLE-Bench rules?

Are the Distilled Lessons Causally Correct?

Cost Analysis.

7 Conclusion

Impact Statement

References

Appendix A Limitations

Appendix B MLE Task Scenario

Appendix C Monte Carlo Tree Search (MCTS)

Appendix D MARS for MLE Tasks

Drafting Phase.

Improvement Phase.

Debugging Phase.

Appendix E Setup for Leaderboard Methods vs. Our Setup

Appendix F Additional Results

F.1 Evaluation across Different Splits of MLE-Bench

F.2 Generalization to Other Models

F.3 Ablation Study Results on MLE-Bench Medium and High Splits

F.4 Cost-Performance Trade-off

Appendix G Qualitative Examples of Causally Correct Distilled Lessons

Example 1: Tabular Playground Series (May 2022).

Example 2: Histopathologic Cancer Detection.

Appendix H Prompts

Appendix I Example Solution Code

Meow summary

The problem MARS tackles

How the method works

What the tests show

Where the evidence stops

Make notes of your own