MARS: Modular Agent with Reflective Search for Automated AI Research
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.
Lay Summary
AI agents are getting better at writing standard software, but they still struggle with the complex task of researching and building new AI models — a process known as Machine Learning Engineering (MLE). This is because training AI is expensive and requires highly organized code. Current AI assistants usually generate massive, inefficient scripts without considering the high computing costs or learning from their successes and mistakes. To solve this, we built MARS, an AI system designed specifically to tackle MLE tasks and act like an autonomous researcher. MARS works on three principles: it carefully manages its "computing budget" to find the most cost-effective ways to train models; it builds code in organized, bite-sized pieces rather than fragile chunks; and it uses a unique memory system to compare experiments and figure out exactly why a new idea succeeded. MARS achieves state-of-the-art performance among open-source frameworks, successfully experiencing "Aha!" moments during its problem-solving process. This work brings us closer to fully automating scientific discovery in AI, making it significantly faster and more cost-effective to build the next generation of machine learning models.