PaperlayerView exact revision on arXiv

Source-linked AI summary

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.02660v3cs.AI

TL;DR

Automating AI research is limited by complex MLE tasks involving expensive evaluation, monolithic code, and difficult credit assignment. MARS combines budget-aware MCTS, modular repository construction, and comparative reflective memory to address these constraints. It achieves state-of-the-art performance among open-source frameworks on MLE-Bench under comparable settings, while its impact statement acknowledges risks from incorrect or hallucinatory code.

  • Problem

    Current agents struggle with MLE because evaluation is computationally expensive, research repositories are architecturally complex, and performance attribution is opaque.

  • Method

    MARS combines cost-constrained MCTS, a Design-Decompose-Implement modular pipeline, and Comparative Reflective Memory for autonomous AI research.

  • Results

    MARS achieves state-of-the-art performance among open-source frameworks on MLE-Bench under comparable settings.

  • Takeaways & Limitations

    MARS treats autonomous research as a repository-level engineering problem requiring resource-aware planning, modular construction, and reflective learning.

  • Takeaways & Limitations

    LLM-based agents can generate incorrect or hallucinatory code, which MARS mitigates through iterative self-correction with execution feedback.

Abstract

from arXiv · show

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.

1. Introduction

Automating AI research is bottlenecked by complex MLE tasks whose expensive evaluation, architectural complexity, and opaque attribution challenge current agents. MARS addresses these constraints through budget-aware search, modular construction, and comparative reflective memory.

  • MLE is probabilistic and resource-intensive, requiring strategic planning around expensive evaluations, opaque performance attribution, and high architectural complexity.
  • Existing agents often treat research as code generation, neglecting execution costs and producing fragile monolithic scripts for complex repositories.
  • MARS uses cost-constrained MCTS to balance performance gains against execution expense within a fixed budget.
  • MARS uses a Design-Decompose-Implement pipeline to organize research repositories into independent, testable modules.
  • Comparative Reflective Memory analyzes differences from the best-known solution to distill high-signal insights for credit assignment.
  • MARS achieves state-of-the-art performance among open-source frameworks on MLE-Bench under comparable settings.

2. Related Work

Prior MLE agents commonly use monolithic code and performance-focused search, while reflective systems emphasize self-correction or memory. MARS combines repository-level modularity, budget-aware search, and comparative analysis of code changes and performance.

  • Table 1 compares MLE agents by modular code generation, runtime or budget awareness, and memory mechanisms.
  • Automated AI Research & MLE Bottleneck: Existing systems predominantly generate expansive single-file scripts that lack modularity for rigorous engineering.
  • Search Algorithms in Code Generation: Prior search algorithms typically optimize task performance while neglecting computational cost.
  • Reflective Learning and Memory: Reflective learning systems include self-correction, associative memory, process-memory distillation, trajectory comparison, and reusable code blocks.
  • Reflective Learning and Memory: MARS analyzes causal links between code changes and performance variations to isolate effective algorithmic changes.

3. Problem

The paper formalizes long-horizon agentic problem solving as constructing an artifact that satisfies instructions and maximizes an objective through environmental interaction under a cost budget. MLE instantiates this framework with data processing, model training, and validation.

  • An agent constructs a complex artifact to satisfy requirements within a constrained budget.
  • A problem instance is represented as P = (I, E, O), where I is the instruction, E the validation environment, and O the objective function.
  • The goal is to find s* maximizing O through interaction with E subject to cost constraint B.
  • The solution space can be vast and unstructured, including the space of all possible Python programs.
  • MLE Task Scenario: MLE requires engineering a full pipeline that processes data, trains models, and validates results.

4. Method

MARS treats long-horizon AI research as repository-level search, combining budget-aware planning, modular construction, and reflective memory to manage costly, complex MLE tasks.

  • MARS formulates autonomous AI research as maximizing objective O under cost budget B in environment E, using a repository-level search process.
  • Modular construction supports targeted repository edits, code reuse, and localized debugging by updating specific files or logic blocks rather than regenerating the full codebase.Diff-Based Editing specifies the target file, replacement block, and new code for atomic multi-file updates.
  • MARS replaces monolithic scripts with a Design-Decompose-Implement workflow that organizes solutions into independent modules and an orchestration script.Idea, Modular, and Coding agents plan, decompose, implement, and validate components before orchestration.
  • Comparative Reflective Memory addresses opaque credit assignment by comparing current and previous best solutions to isolate changes associated with performance shifts.Separate agents distill solution and debugging lessons, while a review agent filters redundant insights before retrieval.
  • MARS uses specialized MCTS expansion, node selection, and efficiency-guided rewards to iteratively evolve solutions while restricting further actions on fully expanded branches.Valid nodes become fully expanded after sufficient improvement attempts, whereas failed nodes are always fully expanded.
  • Budget-Aware MCTS balances performance gains against execution expense while exploring solution states through drafting, improvement, and debugging operators.The search uses UCT to balance exploitation and exploration, with an efficiency-guided reward that penalizes long execution time.

5. Experiment

MARS is evaluated on MLE-Bench under controlled resource and comparison settings, with results showing strong performance and ablations supporting its modular, reflective, and budget-aware components.

  • Evaluation Setup: The evaluation uses 75 Kaggle competitions spanning natural language processing, computer vision, and tabular data analysis.
  • Evaluation Setup: Each competition receives a strict 24-hour end-to-end budget on a node with one NVIDIA A100 GPU, 12 vCPUs, 220 GB RAM, and 1 TB SSD storage.
  • Evaluation Setup: MARS is compared with leaderboard agents, AIDE, and AIRA under identical environments and underlying LLMs for open-source baselines.
  • Main Results: MARS establishes a new state-of-the-art among open-source frameworks and consistently outperforms baselines across Lite, Medium, and High task-complexity splits.
  • Main Results: 31.1% is MARS's highest reported Gold Medal rate, while MARS+ reaches 74.2% Above Median, 33.8% Gold Medal, and 62.7% Any Medal rates with increased compute.
  • Ablation Study: Removing modular decomposition degrades overall success, while removing memory causes a drastic performance drop and comparative distillation provides a consistent boost.
  • Ablation Study: Budget-Aware MCTS outperforms alternative search strategies over time by balancing exploration with resource constraints.
  • Ablation Study: The default penalty weight w = -0.07 is optimal; no penalty degrades performance, whereas w = -0.15 over-prioritizes fast but trivial nodes.

6. Discussions

The discussion evaluates MARS’s modularity, budget-aware exploration, lesson transfer, compliance, originality, causal attribution, and cost efficiency. Results support structured repository construction, efficient search, cross-branch learning, and lower-cost performance gains.

  • Modular Construction: Modular decomposition produces more extensive and structured repositories than the non-modular variant.Repository complexity is measured by lines of code and number of files, with representative modules enumerated for five competitions.
  • Budget-aware Exploration: 19.5% ± 1.5% of explored solutions improve the current best validation metric with Budget-aware MCTS, versus 16.1% ± 1.3% with Vanilla MCTS.The latency penalty favors faster candidates with comparable accuracy and prunes inefficient trajectories.
  • Lesson Learning: 65.8% ± 1.1% of solutions utilize existing lessons, and 63.0% ± 1.8% of utilized lessons originate from different search-tree branches.Lessons formed from early failures or partial successes guide subsequent solutions and transfer knowledge across branches.
  • Compliance and Originality: MARS records a 0% violation rate across monitored MLE-Bench protocol dimensions, and no submission exceeds 60% code similarity to public notebooks.The similarity distribution mirrors AIRA-dojo while remaining below the stated threshold.
  • Credit Assignment: 88.34% of 3,611 audited lessons correctly attribute validation-metric shifts to specific code changes.The audit evaluates whether distilled lessons perform genuine credit assignment rather than relying on hallucinated narratives.
  • Cost Analysis: The 4-hour MARS version achieves a 28.4% Any Medal Rate for $9.6, while the 24-hour version costs $60.5 and reaches 43.1%.The cost-controlled comparison is described as a Pareto improvement, supporting gains from architectural efficiency rather than only a larger resource budget.

7. Conclusion

MARS addresses limitations in autonomous long-horizon AI research by treating research as repository-level engineering. Its integrated planning, modular construction, and comparative memory achieve state-of-the-art performance on complex MLE tasks while balancing exploration and computational efficiency.

  • Conclusion: MARS integrates Budget-Aware MCTS, Modular Construction, and Comparative Reflective Memory for repository-level autonomous AI research.The framework is presented as addressing credit assignment while balancing exploration with computational efficiency.
  • Conclusion: MARS achieves state-of-the-art performance on complex Machine Learning Engineering tasks.The conclusion attributes this outcome to a structured approach that mimics strategic foresight in human engineering.

Impact Statement

The impact statement acknowledges risks from incorrect or hallucinatory code generated by LLM-based agents. It describes iterative self-correction with execution feedback as a mitigation and reports no immediate additional societal consequences beyond generative AI’s general risks.

  • Impact Statement: LLM-based agents may generate incorrect or hallucinatory code, which MARS mitigates through iterative self-correction with code execution feedback.The statement frames this as a broader impact risk associated with deploying LLM-based agents.
  • Impact Statement: The authors foresee no immediate negative societal consequences beyond those generally associated with generative AI.This scope is limited to the impact statement’s stated assessment.

A. Limitations

MARS’s evaluation is limited in scope, incurs substantial computational expense, and relies on LLM-derived lessons that can remain unreliable despite comparative analysis.

  • Scope of Evaluation: Evaluation is confined to MLE-Bench, which does not cover the full spectrum of open-ended automated AI research tasks.Excluded areas include novel algorithm formulation, literature synthesis, and unconstrained real-world deployment.
  • Computational Cost: The pipeline’s extensive API calls for iterative tree search and comparative distillation create notable computational costs that may hinder independent researchers.Proposed mitigations include caching, early stopping, and routing simpler subtasks to smaller models.
  • Reliability of Distilled Lessons: Distilled lessons may hallucinate or misattribute performance gains because LLMs infer causal relationships between code changes and metrics.Comparative delta analysis reduces but does not eliminate this risk; formal analysis, cross-model verification, and repeated validation are suggested.

B. MLE Task Scenario

The MLE scenario treats research automation as constructing and searching over executable software repositories under hidden-test and wall-clock constraints.

  • MLE Task Definition: MLE requires engineering a full pipeline that processes data, trains models, and validates results rather than writing an isolated code snippet.The solution includes distinct modules, dependencies, and entry points for end-to-end orchestration.
  • MLE Task Definition: An MLE task is represented as Q = (I, D, M), combining a natural-language description, datasets, and an evaluation metric.The metric defines the objective and is treated as a maximization target without loss of generality.
  • Evaluation Setting: When no validation set exists, the agent partitions D_dev to create D_val because D_test remains strictly hidden.Validation therefore provides the observable basis for internal evaluation.
  • Evaluation Setting: The agent searches candidate repositories under a wall-clock budget, optimizing validation performance while accounting for search cost C(s).The objective seeks a high-performing executable solution within the available time budget T.
  • Search Procedure: MCTS selects actions by balancing estimated value with exploration, then expands nodes, simulates outcomes, and updates tree statistics from rewards.UCT uses state-action values, visit counts, and an exploration constant.

D. MARS for MLE Tasks

MARS operationalizes MLE research as an iterative modular tree-search workflow that parses tasks, prepares data, generates and improves solutions, debugs execution, and distills lessons.

  • Initialization: MARS begins by extracting the optimization metric and direction, preprocessing datasets, performing EDA, and searching for candidate architectures.These steps create metadata, statistical profiles, and model choices for subsequent exploration.
  • Data Preparation: The workflow creates and verifies training, validation, and test metadata, using stratified or group-based splits when a predefined validation set is unavailable.The default split is 80:20, with integrity checks intended to prevent leakage.
  • Iterative Search: MARS uses tree search to choose either curriculum-based drafting or ablation-style improvement, followed by debugging when execution fails.Drafting progresses from simple baselines to sophisticated ensembles, while improvement targets high-impact refinements.
  • Modular Construction: The framework decomposes proposed ideas into independent functional modules that a coding agent implements and unit-tests before orchestration.Modularity supports complex research repositories by separating functional units.
  • Reflection and Feedback: After execution, MARS reviews results, extracts lessons into a solution pool, and backpropagates rewards to update search statistics.Debugging uses a dedicated lesson pool to avoid repeating prior errors.

E. Setup for Leaderboard Methods vs. Our Setup

The controlled evaluation standardizes hardware and resources to isolate algorithmic effectiveness, while reporting performance across MLE-Bench difficulty splits and a cross-model test.

  • Controlled Setup: Controlled comparisons use one A100 GPU node and no external knowledge bases to isolate algorithmic effectiveness from resource scaling.The setup is intended to address hardware disparities affecting open-ended leaderboard comparisons.
  • Evaluation Protocol: MLE-Bench results are reported separately for Lite, Medium, and High subsets using three independent runs and mean ± SEM percentages.Tables 7–9 provide the corresponding performance comparisons.
  • Model Generalizability: MARS retains a substantial performance advantage over the baseline across all difficulty splits when evaluated with Claude 4.6 Sonnet.This cross-model evaluation is presented as evidence that performance is not specific to the underlying reasoning engine.

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

On harder MLE-Bench Medium and High tasks, MARS’s core modules retain consistent advantages, with larger gaps on the harder splits. A cost-controlled evaluation further attributes gains to architectural efficiency rather than simply greater resource allocation.

  • Ablation results: Performance advantages from MARS’s core modules remain consistent across random samples of 10 Medium and 10 High-split tasks.The gaps between the full method and ablated variants widen on the harder splits.
  • Ablation results: The performance gaps between MARS and ablated variants widen on the harder MLE-Bench splits.This result supports the increasing importance of modules such as modular decomposition as task difficulty rises.
  • Cost-performance trade-off: 28.4% Any Medal Rate is achieved by MARS under a restricted 4-hour limit while spending $9.6, outperforming both 24-hour baselines.The comparison is against AIDE and AIRA-dojo under the cost-controlled evaluation.
  • Cost-performance trade-off: The cost-controlled results attribute MARS’s performance gains to architectural choices rather than merely increased resource allocation.MARS achieves a Pareto improvement by obtaining better results with a smaller budget.

G. Qualitative Examples of Causally Correct Distilled Lessons

MARS demonstrates causally informed lesson distillation through two examples in which code differences and execution behavior explain why alternative architectures or ensembles underperformed. These analyses yield transferable design lessons for subsequent research solutions.

  • Cross-example lesson: The examples illustrate how MARS uses code and execution-log differences to produce causal, transferable lessons rather than merely empirical observations.The section presents these cases as demonstrations of the causal accuracy of distilled lessons.
  • Example 1: Tabular Playground Series (May 2022): MARS diagnosed a 0.28% AUC drop in Dual-Stream versus Early Fusion and attributed it to overfitting and discarded cross-modal interactions.It distilled the lesson that Early Fusion is superior for heterogeneous tabular data in this example.
  • Example 2: Histopathologic Cancer Detection: MARS found that an ensemble reached 0.976 AUC, below its single-model component at 0.982 AUC, because unweighted voting diluted the stronger model’s signal.Execution logs showed ConvNeXt converged faster than EfficientNet, creating a convergence gap.
Loading 2602.02660v3…