MARS: Modular Agent with Reflective Search for Automated AI Research

TL;DR

MARS reframes automated AI research as a constrained, long-horizon software-repository problem and introduces three core innovations: Budget-Aware Planning via cost-constrained Monte Carlo Tree Search (MCTS), Modular Construction through a Design-Decompose-Implement pipeline, and Comparative Reflective Memory to distill high-signal lessons. The framework explicitly balances performance gains against execution costs, manages architectural complexity with modularization, and resolves credit attribution by analyzing differences between solutions. Across the MLE-Bench suite, MARS achieves state-of-the-art results among open-source approaches and demonstrates strong cross-branch knowledge transfer, validating its capacity to generalize insights across search paths. The work also provides ablations, cost analyses, and prompts to facilitate reproducibility and future research in autonomous AI discovery.

Abstract

Automating AI research 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.

Figures (7)

• 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.
• 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.
• Figure 3: Impact of Modular Decomposition and Lesson Learning.
• Figure 4: Comparison of tree search strategies for MARS.
• Figure 5: Reward modulation: Budget-aware MCTS assigns higher rewards to faster candidates when performance is comparable.