AutoMLGen: Navigating Fine-Grained Optimization for Coding Agents

TL;DR

AutoMLGen addresses the bottleneck of expert intervention in end-to-end ML engineering by coupling a curated ML domain knowledge base with Monte Carlo Graph Search (MCGS) to enable trajectory recall, cross-branch knowledge transfer, and multi-solution fusion. The graph-based exploration broadens search diversity beyond traditional tree-based MCTS, while fine-grained operators stabilize execution and improve the quality of generated ML pipelines. Across the MLE-Bench, AutoMLGen achieves state-of-the-art performance with a 12-hour budget, notably a 36.4% average medal rate and 18.7% gold, along with high valid submission rates and competitive efficiency. Ablation studies validate the contributions of the knowledge base and MCGS components, and results generalize across different foundation models, underscoring practical impact for autonomous ML engineering tasks. The work points to future directions including broader benchmarks and decomposed multi-step code generation for even more complex AI tasks.

Abstract

Large language models (LLMs) have shown impressive performance in general programming tasks. However, in Machine Learning Engineering (MLE) scenarios such as AutoML and Kaggle competitions, achieving high performance depends heavily on expert intervention and repeated adjustments rather than simply generating correct code. When applied directly to these tasks, LLMs often lack fine-grained domain priors, and existing MLE approaches that use linear or tree-structured searches limit knowledge transfer to adjacent hierarchical links. As a result, they cannot leverage past full trajectories or share information across branches, limiting self-evolving ability and search space diversity. To address these limitations, we introduce AutoMLGen, an LLM-based coding agent that integrates a domain knowledge base for high-quality prior guidance and Monte Carlo Graph Search (MCGS) for efficient exploration. MCGS retains the tree-guided exploration of MCTS while embedding a graph structure into the expansion stage to enable dynamic path reorganization, historical trajectory reuse, and multi-solution fusion to support both self-evolution and collaborative learning. Combined with fine-grained operator sets, this design improves stability and accelerates convergence. Evaluation on the MLE-Bench shows that AutoMLGen achieves state-of-the-art performance in numerous dimensions, such as the average medal rate and the valid submission rate, under a 12-hour budget (half the standard runtime). The code is available at https://github.com/Alpha-Innovator/InternAgent.

Figures (9)

• Figure 1: The comparison across various methods on MLE-Bench. Our AutoMLGen wins the championship within a 12-hour budget.
• Figure 2: The framework of AutoMLGen. It consists of two main modules: (i) a curated ML domain knowledge base, and (ii) the MCGS module, which integrates graph-based exploration with a set of fine-grained operators. Detailed descriptions are provided in Section \ref{['sec:method']}.
• Figure 3: MLE-Bench results of AutoMLGen and other methods. It is noticeable that AutoMLGen performs better at all these dimensions with the shortest run time.
• Figure 4: Performance comparison on all high-level tasks of MLE-Bench. Win means the average (3 seeds) test score of the task is better, so on for equal and lose. It can be seen that AutoMLGen achieves a better win rate against different baselines.
• Figure 5: (a) The comparison of different models by task type. AutoMLGen is tested with different LLMs (DeepSeek-R1, o4-mini, and Gemini-2.5-pro) across image, text, and tabular tasks. (b) The evolution of beat ratio over time. This figure shows how AutoMLGen compares with the baseline under a 12-hour budget, where AutoMLGen consistently achieves higher leaderboard standings as search progresses.