PaperlayerView exact revision on arXiv

Source-linked AI summary

MLE-STAR: Machine Learning Engineering Agent via Search and Targeted Refinement

Jaehyun Nam, Jinsung Yoon, Jiefeng Chen, Jinwoo Shin, Sercan Ö. Arık, Tomas Pfister

arXiv:2506.15692v3cs.LG

TL;DR

Existing MLE agents may rely on familiar internal knowledge and coarse code-wide changes, limiting task-specific model selection and deep component exploration. MLE-STAR uses web search to seed a solution and refines targeted components, including an agent-proposed ensemble strategy. It achieves medals in 64% of MLE-bench Lite Kaggle competitions.

  • Problem

    Existing MLE agents can favor familiar methods and change whole code structures at once, limiting task-specific model selection and deep exploration of individual pipeline components.

    Source: §1 · p. 1

  • Method

    MLE-STAR searches the web for effective models, uses ablations to guide targeted refinement of pipeline components, and iteratively improves an ensemble strategy proposed by the agent.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    MLE-STAR achieved medals in 64% of MLE-bench Lite Kaggle competitions, with 36% of those competitions yielding gold medals.

    Source: §6 · p. 11

  • Takeaways & Limitations

    With minimal human effort, MLE-STAR significantly outperforms previous methods, including approaches requiring manual collection of Kaggle strategies.

    Source: §1 · p. 2

    Because Kaggle competitions are publicly accessible, LLMs may have been trained on relevant challenge discussions, although the authors report sufficient solution novelty.

    Source: §6 · p. 11

Abstract

from arXiv · show

Agents based on large language models (LLMs) for machine learning engineering (MLE) can automatically implement ML models via code generation. However, existing approaches to build such agents often rely heavily on inherent LLM knowledge and employ coarse exploration strategies that modify the entire code structure at once. This limits their ability to select effective task-specific models and perform deep exploration within specific components, such as experimenting extensively with feature engineering options. To overcome these, we propose MLE-STAR, a novel approach to build MLE agents. MLE-STAR first leverages external knowledge by using a search engine to retrieve effective models from the web, forming an initial solution, then iteratively refines it by exploring various strategies targeting specific ML components. This exploration is guided by ablation studies analyzing the impact of individual code blocks. Furthermore, we introduce a novel ensembling method using an effective strategy suggested by MLE-STAR. Our experimental results show that MLE-STAR achieves medals in 64% of the Kaggle competitions on the MLE-bench Lite, significantly outperforming the best alternative.

1. Introduction

MLE engineering agents automate parts of a labor-intensive process, but existing agents may favor familiar methods and shift between pipeline components before exploring them deeply. MLE-STAR combines web search with targeted component refinement and achieves a higher medal rate than the top-performing baseline.

1. Introduction

  • Existing MLE agents can favor familiar methods and modify whole code structures at once, limiting deep exploration of individual components.Source: §1 · p. 1
  • MLE-STAR uses Google Search to find relevant models, then refines specific pipeline components using ablation studies to identify blocks with the greatest performance impact.It also proposes candidate solutions and iteratively refines an agent-suggested ensemble strategy.Source: §1 · pp. 1–2
  • MLE-STAR raised medal achievement from 25.8% to 63.6% compared with the top-performing baseline.Source: §1 · p. 2

2. Related work

Prior work includes general LLM agents, AutoML systems that search predefined spaces, and specialized LLM-based MLE agents. MLE-STAR is positioned as code-space exploration that uses search and targeted component refinement rather than relying on internal knowledge or a fixed case bank.

2. Related work

  • General-purpose agents use external tools, while specialized agents such as Voyager and AlphaCode use execution feedback to improve domain-specific approaches.Source: §2 · p. 2
  • AutoML explores predefined model, hyperparameter, and specialized search spaces, which often require domain expertise to define.Source: §2 · pp. 2–3
  • AIDE may favor outdated or simple models and shift refinement prematurely, while DS-Agent relies on a manually curated case bank that can limit scalability and novel-task applicability.MLE-STAR instead searches for task-relevant models and explores options within specific pipeline components.Source: §2 · p. 3

3. MLE-STAR

MLE-STAR builds an initial solution using web-retrieved models, then uses ablation-guided nested loops to refine individual code blocks. It also explores ensemble strategies and includes modules for debugging and checking data use and leakage.

3. MLE-STAR

  • MLE-STAR combines web search, targeted code-block refinement, and an LLM-proposed ensembling method in a multi-agent framework for machine-learning tasks.The framework takes datasets and a task description as input and uses agents with specific functions.Source: §3 · p. 3

3.1. Generating an initial solution using web search as a tool

  • Web search retrieves task-specific models and example code to guide the LLM in generating candidate solutions rather than relying only on its model suggestions.The retrieved model description and code examples provide guidance for generating executable code.Source: §3.1 · p. 4
  • After evaluating retrieved scripts, MLE-STAR starts with the top performer and tries sequential simple-average merges, stopping when the validation score no longer improves.Each remaining script is merged into the current initial solution in turn.Source: §3.1 · p. 5

3.2. Refining a code block for solution improvement

  • Ablation studies identify code blocks whose modification most affects performance, while prior ablations and previously targeted blocks guide selection toward different pipeline parts.The extractor also produces an initial refinement plan for the selected block.Source: §3.2 · p. 5
  • For a selected block, the inner loop tests refinement plans using prior attempts as feedback, retains the best-performing candidate, and updates the solution only if it improves performance.The outer loop repeats this process for a predetermined number of steps.Source: §3.2 · pp. 5–6

3.3. Further improvement by exploring ensemble strategies

  • MLE-STAR combines distinct solutions, including parallel runs, by iteratively proposing ensemble plans based on earlier plans and their performance, then selecting the best-performing ensemble.A simple initial plan can average the solutions’ final predictions.Source: §3.3 · p. 6

3.4. Additional modules for robust MLE agents

  • When a script errors, the debugging agent attempts corrections until execution succeeds or a preset round limit is reached, then proceeds with the latest executable version if unresolved.The agent uses the error record, such as a traceback, to guide debugging.Source: §3.4 · p. 7
  • A leakage checker examines preprocessing code before execution, while a data-usage checker reviews the initial solution against the task description and revises it if relevant data is underused.The leakage checker leaves the script unchanged if it detects no leakage.Source: §3.4 · p. 7

4. Experiments

Across 22 MLE-bench Lite Kaggle competitions, MLE-STAR is evaluated against existing agents and variants, including different LLMs and an ensemble strategy. The experiments report stronger medal results, compatibility across LLM types, and benefits from ensembling.

4. Experiments

  • Experiments cover 22 Kaggle competitions with three seeds and Gemini-2.0-Flash by default, within a 24-hour limit.MLE-STAR retrieves four model candidates, uses four inner and four outer refinement loops, and explores ensemble strategies for five rounds.Source: §4 · p. 7

4.1. Main results

  • MLE-STAR with Gemini-2.0-Flash raises AIDE’s any-medal rate from 25.8% to 43.9% and its above-median rate from 39.4% to 63.6%.With Gemini-2.5-Pro, MLE-STAR achieves a medal rate above 60%.Source: Table 1 · p. 8, §4.1 · p. 8
  • On four tabular tasks over five seeds, MLE-STAR with Gemini-2.0-Flash significantly outperforms DS-Agent without human effort.The comparison uses the same tasks employed during DS-Agent’s development; the paper notes that a broader direct comparison is infeasible because DS-Agent requires a manually curated case bank.Source: Table 2 · p. 8, §4.1 · p. 8

4.2. Ablation studies

  • Ensembling multiple solution candidates improves any-medal and gold-medal success rates and outperforms simpler ensemble strategies.The authors also report a larger margin over median human experts with ensembling than without it.Source: Table 4 · p. 9, §4.2 · p. 9
  • With Gemini-2.5-Pro, MLE-STAR earns two gold medals and one silver on denoising-dirty-documents, while Gemini-2.0-Flash earns none.The authors present this as evidence that MLE-STAR can harness advances in reasoning-based LLMs.Source: §4.2 · p. 9
  • Across four competitions spanning image-to-image, image, text, and audio classification, Claude-Sonnet-4 results support compatibility across LLM types.Each competition was run for three seeds.Source: Table 3 · p. 8, §4.2 · p. 9

5. Discussion

The discussion illustrates MLE-STAR’s model choices, data checks, and refinement process through examples. It also describes how manual input can add models that search did not retrieve.

Qualitative observations on selected models

  • MLE-STAR uses newer EfficientNet or ViT models more often than AIDE’s ResNet-focused approach, winning 37% of image-classification medals versus AIDE’s 26%.Figure 4 reports model usage percentages and omits other models, which account for 11.7% and are used by MLE-STAR.Source: §5 · p. 9, Figure 4 · p. 10

Human intervention

  • Manually adding a RealMLP model description lets MLE-STAR integrate its training even though search had not retrieved the model.Users can also replace the ablation summary with instructions to specify target code blocks.Source: Figure 5 · p. 10, §5 · p. 10

Progressive improvement via MLE-STAR refinement

  • Across 22 MLE-bench Lite challenges, refinement consistently reduces average relative error, with notable gains in early steps that target influential code blocks.Each refinement step focuses on one code block; the authors attribute the early gains to ablation identifying influential blocks first.Source: Figure 8 · p. 11, §5 · p. 11

Misbehavior of LLMs and corrections

  • When feature engineering uses the unavailable-at-test-time target Transported, validation improves but test accuracy drops, illustrating data leakage.MLE-STAR’s leakage checker instead uses training statistics to preprocess test data.Source: Figure 6 · p. 10, §5 · p. 11
  • MLE-STAR’s data-usage checker reexamines task descriptions and can incorporate overlooked sources, as with geometry.xyz in nomad2018-predicting.The paper reports that performance significantly improves in this example.Source: Figure 7 · p. 10, §5 · p. 11

6. Conclusion

MLE-STAR retrieves effective models with a search engine, then refines targeted pipeline components and explores ensemble strategies. It wins medals in 64% of MLE-bench Lite Kaggle competitions.

6. Conclusion

  • 64% of MLE-bench Lite Kaggle competitions yielded medals for MLE-STAR, including 36% gold medals.The conclusion also notes a potential risk that LLMs encountered relevant Kaggle discussions during training; it reports checking solution novelty against those discussions with an LLM judge.Source: §6 · p. 11

A.1. Retriever agent; A.2. Candidate evaluation agent

  • The retriever returns task-specific model-and-code pairs, and candidate evaluation runs a simple validation script, using training subsets when samples are numerous.The retriever returns M pairs; the evaluation script computes results on a hold-out validation set.Source: §A.1 · p. 15, §A.2 · p. 16

A.3. Merging agent

  • The merger sequentially incorporates the best unmerged reference model code into the current base script to produce the next candidate solution.The process continues until the retrieved models have been merged.Source: §A.3 · p. 17

A.4. Ablation study agent; A.6. Extractor

  • Ablation studies modify or disable pipeline components; their results guide extraction of the highest-impact code block and its initial refinement plan.The extractor also considers the current solution and previously refined blocks.Source: §A.6 · p. 20

A.7. Coder; A.8. Planner

  • The coder applies a refinement plan to the extracted block, while the planner proposes further plans using that block and previous attempts.The planner’s examples explore alternative feature-engineering strategies.Source: §A.7 · p. 21, §A.8 · p. 22

A.9. Ensemble strategy planner; A.10. Ensembler

  • The ensemble planner uses prior plans and their performance as feedback, and the ensembler applies a plan to two final solutions generated in parallel.The resulting merged script is among the candidates from which the final solution is selected.Source: §A.9 · p. 24, §A.10 · p. 27

A.11. Debugging agent; A.12. Data leakage checker; A.13. Data usage checker

  • When execution errors occur, a debugger attempts correction; separate checks inspect preprocessing for leakage and revise the initial script to use relevant data.The leakage checker extracts the preprocessing block and corrects it if it detects leakage.Source: §A.11 · p. 28, §A.12 · p. 29, §A.13 · p. 30

A.4. Ablation study agent

MLE-STAR uses ablation studies to assess pipeline components and guide targeted code-block refinement; supplementary algorithms describe initial-solution generation, refinement, and ensembling.

A.4. Ablation study agent

  • Aabl generates code that modifies or disables specific components, with prior ablation summaries supplied as feedback to encourage exploration of different pipeline parts.The generated script performs an ablation study on the current solution.Source: §A.4 · p. 18

A.5. Ablation study summarization agent

  • Asummarize uses the ablation script and its execution output to parse the component-level results, filtering out unrelated output such as training-epoch losses.The ablation script is included because it describes the modification being evaluated.Source: §A.5 · p. 19

B.1. Algorithm for generating an initial solution

  • Algorithm 1 takes the task description, datasets, score function, and number of retrieved models as inputs for generating an initial solution.Source: §B.1 · p. 31

B.2. Algorithm for refining a code block for solution improvement

  • Algorithm 2 takes an initial solution and outer- and inner-loop step counts as inputs for refining a code block.Source: §B.2 · p. 32

B.3. Algorithm for further improvement by exploring ensemble strategies

  • Algorithm 3 covers ensembling final solutions and takes candidate final solutions as input.Source: §B.3 · p. 33

C.1. Generated code for ablation study

  • The supplementary material provides an example of code generated by Aabl for an ablation study.Source: §C.1 · p. 34

C.2. Raw output of ablation study after execution

Ablation execution output can include irrelevant training information, so MLE-STAR summarizes results before using them to select a code block. The appendix also identifies the benchmark scope and a qualitative solution comparison.

C.2. Raw output of ablation study after execution

  • The example raw ablation output can contain unrelated information, such as LightGBM training details, so MLE-STAR parses it with Asummarize before extraction.Figure 23 shows an example execution output.Source: §C.2 · p. 34, Figure 23 · p. 34

C.3. Summary of ablation study

  • Asummarize organizes component-impact information from raw ablation output, and its summary is passed to Aextractor to identify the block with the greatest performance impact.Figure 24 shows an example summarized result.Source: §C.3 · p. 35, Figure 24 · p. 35

D. Qualitative comparison

  • The qualitative comparison provides final solution code for MLE-STAR and AIDE, with both agents using Gemini-2.0-Flash as the base LLM.Source: D. Qualitative comparison · p. 35

E.1. MLE-bench Lite

  • MLE-bench Lite is the low-complexity split of MLE-bench, comprising 22 competitions across problem types including tabular, text, and image prediction.The full benchmark consists of 75 offline Kaggle competitions.Source: §E.1 · p. 36

E.2. Tabular tasks from DS-Agent

The appendix lists tabular competitions used in DS-Agent’s development phase and describes how MLE-STAR generates submissions using test samples.

E.2. Tabular tasks from DS-Agent

E.3. Generating submission file

  • At submission time, Atest takes the task description and final solution as input, then outputs code to load test samples and create a submission file.Source: §E.3 · p. 37
  • MLE-STAR removes the subsampling used to speed refinement by extracting that code block and modifying it to use all provided samples when creating the submission.Source: §E.3 · p. 38

F. Experimental setup

With four retrieved models, four inner loops, four outer loops, and five ensemble rounds, MLE-STAR took 14.1 hours per solution on average, compared with 15.4 hours for AIDE.

F. Experimental setup

  • Experiments mainly used 96 vCPUs, 360 GB of memory, and eight NVIDIA V100 GPUs with 16 GB of memory each.Source: F. Experimental setup · p. 38
  • 14.1 hours per solution on average for MLE-STAR, versus 15.4 hours for AIDE, under a 24-hour limit for both methods.The average covers 22 tasks and three random trials (66 experiments), using four retrieved models, four inner loops, four outer loops, and five ensemble rounds.Source: F. Experimental setup · p. 38

G. Additional quantitative results

MLE-STAR consistently outperformed DS-Agent with Gemini-2.0-Flash and surpassed AutoGluon by a high margin on three tabular tasks, although AutoGluon is designed specifically for tabular data.

G. Additional quantitative results

H. Analysis on data contamination

Because public Kaggle discussions may enter LLM training data, the authors assess solution novelty against those discussions. They also describe potential broader impacts and contrast data-science agents' evaluation focus with specific engineering challenges.

H. Analysis on data contamination

  • All Gemini-2.0-Flash MLE-STAR solutions were judged sufficiently novel against top discussions across 75 pairs from 25 posts and seven competitions.Gemini-2.5-Pro served as the novelty judge.Source: H. Analysis on data contamination · p. 40

I. Broader impacts

  • MLE-STAR could lower barriers to using ML, and its web search may let generated solutions benefit as state-of-the-art models improve.Source: §I · p. 40

J. Related works on data science agents

Loading 2506.15692v3…