MARS: Modular Agent with Reflective Search for Automated AI Research

• On iMet-2020-FGVC7, MARS progressively refines its strategy from a lightweight residual network to model ensemble techniques and ultimately achieves a silver medal, while existing methods fail to reach medal-level performance.Source: Figure 1 · p. 1
• MLE tasks require strategic foresight because evaluations are computationally expensive, performance attribution is opaque, and architectural complexity is high.Unlike standard software development, MLE is probabilistic, resource-intensive, and not cheaply verified.Source: §1 · pp. 1–2
• Existing agents often ignore execution costs, produce fragile monolithic repositories, and fail to isolate causal factors behind performance changes.Standard search may favor a 0.1% accuracy gain that increases training time from one hour to ten, while memory-based agents struggle with credit assignment.Source: §1 · p. 2
• MARS combines cost-constrained MCTS, a modular Design-Decompose-Implement pipeline, and Comparative Reflective Memory to optimize research under fixed computational budgets.These mechanisms balance performance against execution expense, organize independent testable modules, and extract causal insights by comparing current and best-known solutions.Source: §1 · p. 2
• MARS achieves state-of-the-art performance among open-source frameworks on the MLE-Bench benchmark under comparable settings.Ablation studies further validate the necessity of each proposed mechanism.Source: §1 · p. 2

Comparison of MLE Agents

• Table 1 summarizes how MARS and existing MLE agents differ in modular code generation, runtime or budget awareness, and memory mechanisms.The comparison organizes the frameworks along architecture, search-cost consideration, and memory dimensions.Source: Table 1 · p. 3, §2 · p. 3

Automated AI Research & MLE Bottleneck

• Autonomous MLE is a critical execution bottleneck, motivating MARS’s repository-level decomposition into distinct, testable, and maintainable modules instead of monolithic scripts.Existing agentic frameworks predominantly generate expansive single-file scripts that produce fragile codebases lacking modularity for rigorous engineering.Source: §2 · p. 3

Search Algorithms in Code Generation

• Existing code-search methods optimize task performance while neglecting computational cost; MARS’s Budget-aware MCTS integrates efficiency-guided rewards to balance exploration, exploitation, and expense.The method penalizes computationally expensive solutions directly within the search tree, unlike budget-awareness approaches designed primarily for discrete tool actions.Source: §2 · p. 3

Reflective Learning and Memory

• Prior reflective methods emphasize self-correction, retrieval, process-memory distillation, trajectory comparison, or code reuse, whereas MARS introduces Lesson Learning for causal lesson extraction.Its comparative analysis links code changes to performance variations, isolating effective algorithmic changes from confounding factors and storing them in a lesson pool.Source: §2 · p. 3

• The problem is to construct a solution s* that maximizes objective O through interaction with environment E under cost constraint B.A problem instance is defined as P = (I, E, O), where I contains natural-language requirements, E validates solutions, and O quantifies solution quality.Source: §3 · p. 4
• The search space is often vast and unstructured, including the space of all possible Python programs.Source: §3 · p. 4

MLE Task Scenario

• MLE instantiates this framework by requiring a full pipeline that processes data, trains models, and validates results using held-out performance metrics.In this scenario, E consists of the provided datasets and O is a performance metric such as accuracy on the held-out test set.Source: §3 · p. 4

4.1. Overall Framework

• MARS formulates long-horizon AI research as maximizing objective O under budget B while integrating modular construction, reflective memory, and resource-aware planning.The framework addresses exploration complexity, context management, and solution robustness through these three capabilities.Source: §4.1 · p. 4

4.2. Modular Decomposition

• MARS replaces monolithic scripts with a Design-Decompose-Implement pipeline that creates independent, testable modules and applies atomic diff-based repository updates.Idea, Modular, and Coding agents plan, decompose, implement, validate, and orchestrate modules through a main script.Source: §4.2 · pp. 4–6, Figure 2 · p. 5

4.3. Comparative Reflective Memory

• Comparative Reflective Memory extracts causal lessons by contrasting current solutions with previous best solutions and maintains a filtered, diverse lesson pool.Lessons capture the causal change, comparative impact, and generalized rule; failed executions additionally produce debugging lessons.Source: §4.3 · p. 6

4.4. Budget-Aware MCTS

• Budget-Aware MCTS explores solution repositories through specialized expansion, UCT-based selection, and an efficiency-guided reward balancing performance against execution cost.The search iterates through Selection, Expansion, Simulation, and Backpropagation.Source: §4.4 · p. 6

4.4.1. Actions and Expansion

• MARS expands search nodes by drafting new solutions, improving valid solutions, or debugging failed ones, with buggy children entering an automatic loop of up to N_d debugging actions.Improvement and debugging modify modules or the orchestration script while inheriting the parent structure.Source: §4.4.1 · p. 7

4.4.2. Node Selection

• UCT selection traverses the search tree by balancing exploitation of high-performing solutions with exploration of new solutions.Traversal continues until an incompletely expanded candidate node is identified.Source: §4.4.2 · p. 7

4.4.3. Efficiency-Guided Reward Function

• The efficiency-guided reward normalizes node performance against explored-history extrema and modulates it by execution latency using a penalty weight w.This reward is designed to favor performance gains while penalizing long execution time under budget constraints.Source: §4.4.3 · p. 7

4.5. Task Specific Components

• For MLE tasks, MARS adds multi-agent task preparation, EDA-driven data analysis, and curriculum-based exploration from simple baselines to complex methods.Task preparation formalizes objectives and prepares training, validation, and test datasets, while EDA guides downstream feature engineering.Source: §4.5 · p. 8

5.1. Setup

• MARS is evaluated across 75 Kaggle competitions spanning natural language processing, computer vision, and tabular analysis under the standard MLE-Bench protocol.Each competition receives a strict 24-hour wall-clock budget on one NVIDIA A100 GPU, 12 vCPUs, 220 GB RAM, and 1 TB SSD.Source: §5.1 · p. 8
• Performance is measured by Above Median, Any Medal, and Gold Medal rates, reported as mean ± SEM across three independent runs against leaderboard agents, AIDE (Jiang et al., 2025), and AIRA (Toledo et al., 2026).Open-source baselines use identical environment configurations and underlying LLMs for fairness.Source: §5.1 · p. 9

5.2. Main Results

• 74.2% Above Median, 33.8% Gold Medal, and 62.7% Any Medal rates make MARS+ the strongest reported scaled variant, while standard MARS leads Gold Medal rate at 31.1%.MARS+ uses two concurrent search trees with 2×H100 GPUs and 48 vCPUs; MARS remains highly competitive with the official leaderboard and significantly outperforms AIDE and AIRA-dojo under identical constraints.Source: §5.2 · p. 9, Table 3 · p. 9

5.3. Ablation Study

• Removing modular decomposition reduces overall success, while removing memory causes a drastic performance drop and comparative distillation consistently improves over empirical-only lesson analysis.Comparative reflective memory isolates causal factors through code-level differences and produces more actionable heuristics than noisy-log generalization.Source: §5.3 · p. 10, Figure 3 · p. 10
• Budget-Aware MCTS consistently outperforms Greedy Search and Vanilla MCTS over time by balancing exploration with resource constraints.Vanilla MCTS corresponds to the zero-penalty setting, w = 0.Source: §5.3 · pp. 10–11, Figure 3 · p. 10
• The default penalty w = -0.07 is optimal; w = 0 degrades performance, whereas w = -0.15 over-penalizes latency and favors trivial fast nodes.The penalty balances execution time against performance in the search reward.Source: Figure 3 · p. 10, §5.3 · p. 11

Modular Decomposition

• Modular Decomposition produces more extensive, structured codebases and task-specific modules for preprocessing, model architecture, and other subproblems.The generated repository-level organization mirrors professional software-engineering workflows.Source: §6 · p. 11, Table 4 · p. 11, Table 5 · p. 11

Budget-aware MCTS

• Budget-aware MCTS raises the effective solution rate to 19.5% ± 1.5% from Vanilla MCTS’s 16.1% ± 1.3% by favoring faster comparable candidates.Its latency penalty prunes inefficient trajectories and directs computation toward efficient nodes within the time limit.Source: §6 · p. 11, Figure 4 · p. 12

Lesson Learning

• Lesson Learning achieves 65.8% ± 1.1% utilization and 63.0% ± 1.8% cross-branch transfer, steering exploration with knowledge from prior solutions.Lessons arise from early failures or partial successes and refine subsequent solutions toward higher-quality strategies.Source: §6 · pp. 11–12

Compliance and Originality

• MARS records a 0% rule-violation rate, and no submission exceeds 60% code similarity, indicating protocol compliance and distinct solutions.Similarity distributions mirror AIRA-dojo while remaining below the stated threshold against top public notebooks.Source: §6 · p. 12, Figure 5 · p. 13

Causal Correctness

• 88.34% of 3,611 distilled lessons correctly attribute validation shifts to code changes, while manual inspection finds 90% causal accuracy.The audit used Claude 4.6 Sonnet, and the manual evaluation covered 20 representative lessons.Source: §6 · pp. 12–13

Cost Analysis

• A 4-hour MARS achieves a higher Any Medal Rate of 28.4% while spending $9.6, outperforming 24-hour baselines more cheaply.The 24-hour version costs $60.5 versus AIRA-dojo’s $39.0, while nearly doubling Any Medal Rate from 24.4% to 43.1%.Source: §6 · p. 13

Conclusion

• MARS combines Budget-Aware MCTS, Modular Construction, and Comparative Reflective Memory to balance computational efficiency with exploration and resolve credit assignment in repository-level AI research.The framework treats research as a rigorous engineering challenge rather than traditional code generation.Source: §7 · p. 13
• MARS achieves state-of-the-art performance on MLE-Bench for complex machine learning engineering tasks.The evaluation supports the framework’s structured approach to long-horizon AI research.Source: §7 · p. 13
• Future work will extend MARS to broader scientific discovery domains and improve economic viability through context caching and early stopping.Source: §7 · p. 13

Impact Statement

• MARS targets more reliable and efficient automated software engineering as a contribution to autonomous AI agents.Source: Impact Statement · p. 14
• Incorrect or hallucinatory code is a recognized deployment risk for LLM-based agents.Source: Impact Statement · p. 14
• Iterative self-correction using code-execution feedback is the stated mitigation for incorrect or hallucinatory code.Source: Impact Statement · p. 14
• The authors foresee no immediate negative societal consequences beyond those generally associated with generative AI.Source: Impact Statement · p. 14

Appendix

• Appendices A–E cover limitations, the MLE task scenario, standard MCTS, MARS’s MLE instantiation, and comparisons with prior agents’ evaluation setups.Source: Appendix · p. 17
• Appendices F–I provide additional experiments, causally correct distilled-lesson examples, comprehensive prompts, and representative generated-code examples.Source: Appendix · p. 17

Scope of Evaluation

• Evaluation is limited to MLE-Bench, excluding novel algorithm formulation, literature synthesis, and unconstrained deployment; broader testing should include research engineering Wijk et al. (2025) and replication Starace et al. (2025).MLE-Bench is rigorous but does not cover the full spectrum of open-ended automated AI research.Source: A. Limitations · p. 17

Computational Cost

• Extensive API calls for iterative tree search and comparative distillation create substantial computational costs that may hinder independent researchers’ adoption.Proposed mitigations include KV-cache and prompt caching, early stopping when MCTS improvements plateau, and routing simpler subtasks to smaller models.Source: A. Limitations · p. 17

Reliability of Distilled Lessons

• LLM-generated lessons can hallucinate or misattribute causal effects despite comparative delta analysis, motivating formal program analysis, cross-model verification, and confidence scoring from repeated runs.Delta analysis grounds reflections in code differences but does not eliminate reliability risks.Source: A. Limitations · p. 17

B. MLE Task Scenario

• MLE requires engineering a complete pipeline that processes data, trains models, and validates results rather than writing an isolated code snippet.Source: B. MLE Task Scenario · p. 17
• The MLE task is represented as Q = (I, D, M), where the natural-language description, datasets, and evaluation metric define the problem.Source: B. MLE Task Scenario · pp. 17–18
• Because the test set is hidden, the agent creates D_val from D_dev when necessary and optimizes validation performance as a proxy for unseen-test performance.Source: B. MLE Task Scenario · p. 18
• The agent outputs an executable solution s as a structured repository containing the modules, dependencies, and entry points for an end-to-end pipeline.Source: B. MLE Task Scenario · p. 18
• The MLE agent searches candidate solutions under a wall-clock budget T, with C(s) denoting the total time consumed while searching for solution s.Source: B. MLE Task Scenario · p. 18

C. Monte Carlo Tree Search (MCTS)

• MCTS builds a tree of states and actions, estimating state values through simulated outcomes across selection, expansion, simulation, and backpropagation phases.Source: C. Monte Carlo Tree Search (MCTS) · p. 18
• In MARS, MCTS states are partial or complete solutions, actions are drafting, improvement, or debugging operators, and rewards derive from efficiency-guided validation.Source: C. Monte Carlo Tree Search (MCTS) · p. 19

Task and data preparation

• MARS extracts the evaluation metric and optimization direction, generates dataset metadata, and creates verified stratified or group-based splits when validation data are unavailable.The default split is 80:20, with checks for leakage and documentation of schema and split logic.Source: D. MARS for MLE Tasks · pp. 19, 21

Analysis and architecture search

• Data analysis produces distribution and correlation reports for feature engineering, while search identifies K_a candidate architectures across diverse algorithmic families.The candidates include gradient-boosted trees and deep neural networks and are retrieved through search tools.Source: D. MARS for MLE Tasks · pp. 19, 21

Iterative tree search

• MARS iteratively selects tree nodes with UCT, enters drafting or improvement, debugs generated code, reviews results, distills lessons, and backpropagates rewards.Root selection triggers drafting; non-root selection triggers improvement.Source: D. MARS for MLE Tasks · pp. 20–21

Drafting phase

• Drafting grows branches from lightweight baselines to ensembles, using accumulated solution lessons to generate ideas and modular agents to implement and unit-test functional components.The components are orchestrated into a final execution script π_main.Source: D. MARS for MLE Tasks · pp. 20–21

Improvement and debugging phases

• Improvement performs targeted ablation-style refinements using learned lessons, while debugging applies stored error-correction patterns for up to K_debug attempts to prevent repeated failures.Typical refinements include hyperparameter tuning and feature engineering.Source: D. MARS for MLE Tasks · pp. 20–21

E. Setup for Leaderboard Methods vs. Our Setup

• MARS’s controlled comparison uses a single A100 GPU node without external knowledge bases to isolate algorithmic effectiveness from resource scaling.The setup also documents the hardware, time limits, and auxiliary resources used by leaderboard agents and MARS in Table 6.Source: E. Setup for Leaderboard Methods vs.… · p. 22

F.1. Evaluation across Different Splits of MLE-Bench

• The evaluation covers Lite, Medium, and High subsets of MLE-Bench, with detailed results reported in Tables 7, 8, and 9.Results are reported as mean ± SEM across three independent runs, with values expressed as percentages.Source: §F.1 · p. 23, Table 7 · p. 23, Table 8 · p. 23
• Table 7 compares performance on the MLE-Bench Lite split, highlighting the best result and underlining the second-best.Values are percentages reported as mean ± SEM across three independent runs.Source: Table 7 · p. 23
• Table 8 compares performance on the MLE-Bench Medium split, highlighting the best result and underlining the second-best.Values are percentages reported as mean ± SEM across three independent runs.Source: Table 8 · p. 23
• Table 9 compares performance on the MLE-Bench High split, highlighting the best result and underlining the second-best.Values are percentages reported as mean ± SEM across three independent runs.Source: Table 9 · p. 24

F.2. Generalization to Other Models

• MARS maintains a substantial performance advantage over the baseline across all difficulty splits when evaluated with Claude 4.6 Sonnet.This cross-model evaluation is intended to show architectural robustness rather than overfitting to a specific language model, and Table 10 evaluates Any Medal Rate on random samples of 10 tasks per split across three independent runs.Source: §F.2 · p. 24, Table 10 · p. 24

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

• On random samples of 10 Medium-split tasks, MARS’s core modules retain consistent performance advantages over ablated variants.Figure 6 reports the Medium-split ablation results.Source: §F.3 · p. 24, Figure 6 · p. 25
• On random samples of 10 High-split tasks, MARS’s core modules likewise retain their performance advantages over ablated variants.Figure 7 reports the High-split ablation results.Source: §F.3 · p. 24, Figure 7 · p. 26
• The performance gaps between full MARS and ablated variants widen on harder splits, making modular decomposition and comparative memory distillation increasingly essential as task complexity scales.This finding extends the earlier ablation evaluation on the MLE-Bench Lite split to Medium and High samples.Source: §F.3 · p. 24

F.4. Cost-Performance Trade-off

• Under the 24-hour limit, MARS produces 6.0×10^5 output tokens and 594.8 API calls, versus 867.0 calls for AIRA-dojo, but consumes 286.6×10^5 input tokens.Its substantial input consumption is linked to maintaining comprehensive memory and modular structures.Source: §F.4 · p. 25
• MARS’s Any Medal Rate rises from 24.4% to 43.1% but costs $60.5 versus AIRA-dojo’s $39.0 under the same 24-hour limit.The higher expense is attributed to Gemini-2.5-Pro’s premium pricing for long-context prompts over 200k tokens.Source: §F.4 · p. 25, Table 11 · p. 25
• 28.4% Any Medal Rate at $9.6 under a 4-hour limit exceeds both 24-hour baselines while costing less, demonstrating a Pareto improvement from MARS’s architecture.The result supports architectural efficiency rather than gains from simply allocating a larger financial budget.Source: Table 11 · p. 25, §F.4 · p. 26

G. Qualitative Examples of Causally Correct Distilled Lessons

• A 0.28% AUC drop revealed that Early Fusion outperformed Dual-Stream architecture on heterogeneous tabular data by preserving cross-modal interactions discarded by independent projections.MARS linked the architectural difference to execution-log overfitting and distilled a transferable causal lesson.Source: G. Qualitative Examples of Causally… · p. 26
• The 0.976-AUC multi-model ensemble underperformed its 0.982-AUC single-model component because unweighted voting diluted the faster-converging ConvNeXt model with weaker EfficientNet predictions.MARS converted this failure into a negative constraint against naive ensembles of disparate performers, pruning the search space.Source: G. Qualitative Examples of Causally… · p. 27

H. Prompts

• MARS’s prompts coordinate metric-direction parsing, metadata-file generation, validation-set handling, execution checks, and structured review or documentation outputs.The prompts require explicit JSON responses, isolated metadata generation, reproducible representative splits when needed, and failure checks rather than suppressed errors.Source: H. Prompts · pp. 27–29
• The EDA prompt requires leakage-free, modality-specific analysis of targets, inputs, missingness, signal relationships, and feature importance to guide preprocessing and feature engineering.It covers tabular, image, audio, and text diagnostics, including distributions, correlations, metadata relationships, and top-five Random Forest features.Source: H. Prompts · pp. 30–31
• Architecture prompts require diverse algorithmic families, data-modality alignment, lightweight designs, and explicit reasoning about efficiency, data constraints, and overfitting.Candidate models must include reasoning and technical descriptions in a specified JSON format.Source: H. Prompts · pp. 31–32
• Initial idea prompts demand novel, simple, fast baselines that prioritize speed over accuracy and exclude resource-intensive techniques such as heavy augmentation and ensembles.Responses must describe the model, data preprocessing, training, and evaluation in natural language without code.Source: H. Prompts · p. 32
• Idea-improvement prompts require structural innovation grounded in lessons, retention of successful components, and feasibility under standard training budgets rather than trivial hyperparameter tuning.Applied lessons must be cited inline using their lesson identifiers.Source: H. Prompts · p. 33
• Modular decomposition prompts require logical modules, a topologically ordered JSON repository description, and a main entry point coordinating the end-to-end pipeline.Each module description must specify its purpose, classes or functions, implementation logic, and interactions with other modules.Source: H. Prompts · p. 34
• Implementation prompts enforce reuse of provided code, configurable dataset and training limits, metadata-based splits, safe caching without pickle, full-precision metrics, early stopping, and standardized submission files.Caching recomputes and stores data when loading fails or caching is disabled, while training modules print key metrics and submission modules write predictions to ./submission/submission.csv.Source: H. Prompts · pp. 34–35
• The end-to-end solution draft must run a fast baseline that trains, validates on the metadata-defined hold-out set, performs error-feature correlation analysis, and generates the complete test submission.It must print the full-precision final validation metric as `Final Validation Metric: < value>` and save predictions to ./submission/submission.csv.Source: H. Prompts · p. 36
• The solution-improvement prompt requires targeted ablations, lesson citations, hyperparameter optimization, preserved core logic, validation reporting, failure analysis, and test-set submission generation.It explicitly balances predictive performance with execution speed while preventing wholesale rewrites.Source: H. Prompts · p. 37
• Debugging prompts require root-cause fixes, targeted changes that preserve intended functionality, lesson citations, and no error suppression through try...except blocks.A separate distillation prompt determines whether fixes succeed, explains the causal mechanism, and records detection signals for future debugging.Source: H. Prompts · p. 38
• Solution lessons are distilled through comparative or empirical analysis, causal tracing, generalizable algorithmic reasoning, performance trade-offs, and actionable key lessons, then deduplicated by semantic overlap.The execution-review prompt separately requires JSON fields for findings, validation metric, and metric validity, including checks for leakage and incomplete hold-out evaluation.Source: H. Prompts · pp. 40–41

I. Example Solution Code

• The showcased iMet-2020-fgvc7 repository was generated by MARS with Gemini-3-Pro-Preview and achieved the highest validation performance during the search process.The example begins with configuration code defining working, submission, and metadata paths for train, validation, and test files.Source: §I · p. 42
• The configuration fixes reproducibility and resource-aware training choices, including seed 42, 12 epochs, batch size 48, 352x352 inputs, and A100 40GB tuning.The code also sets deterministic cuDNN behavior and selects CUDA when available.Source: §I · pp. 43–44
• The pipeline loads and caches processed metadata, parses multi-label attribute IDs, handles invalid images, and constructs train, validation, and test DataLoaders.It supports debug subsampling and configurable batch size, worker count, and cached dataframes.Source: §I · pp. 44–48
• Inference applies horizontal-flip test-time augmentation by averaging original and flipped sigmoid probabilities, with cached predictions reused when available.Prediction caching is validated against the DataLoader dataset size before inference is rerun.Source: §I · pp. 49–50
• The submission workflow predicts with two models, ensembles their probabilities, optimizes a validation threshold, applies it to test predictions, and saves submission.csv.The documented sequence uses Model A and Model B with TTA before threshold calibration and test-time application.Source: §I · pp. 51–52
• The inference pipeline ensembles validation and test predictions, optimizes a validation threshold, and writes thresholded attribute IDs to a submission CSV.Validation uses the optimized Micro-F1 threshold, which is then applied during test submission generation.Source: §I · pp. 52–53
• The implemented Asymmetric Loss separates positive and negative components, down-weights easy negatives, and hard-thresholds negatives below the clipping margin.The loss uses separate focusing parameters for positive and negative samples and clamps logarithm inputs for numerical stability.Source: §I · pp. 53–54
• The model factory uses timm to instantiate pretrained architectures such as ResNet101d or ConvNeXt-Base and replaces their classifier heads for the dataset’s class count.The architecture follows global average pooling with a linear projection through timm’s configurable num_classes argument.Source: §I · pp. 55–56
• Training combines mixed-precision optimization, OneCycleLR scheduling, optimized-F1 checkpointing, and patience-based early stopping.Validation reports standard and optimized F1, while the optimized score determines checkpoint improvement and threshold calibration.Source: §I · pp. 57–59
• The code trains two models for five epochs within a four-hour limit, applying lessons on longer schedules and diverse-model ensembling.It explicitly cites Lesson 00003 for convergence and Lesson 00018 for ensemble diversity.Source: §I · pp. 63–64
• The validation pipeline averages predictions from both models, optimizes the decision threshold for Micro F1, and prints the final validation metric.The ensemble uses simple averaging before threshold optimization.Source: §I · p. 64
• Failure analysis estimates per-sample error magnitude from F1 scores and relates it to label cardinality and image brightness using sampled validation images.Image feature extraction is limited to at most 2,000 validation samples, and the code reports both correlations when sufficient data exist.Source: §I · pp. 64–65
• Submission generation is gated by whether the optimized validation score exceeds a predefined threshold; otherwise, submission creation is skipped.The code reports the threshold check and generates fresh test predictions only when the condition is met.Source: §I · p. 66