PaperlayerView exact revision on arXiv

Source-linked AI summary

AutoMind: Adaptive Knowledgeable Agent for Automated Data Science

Yixin Ou, Yujie Luo, Jingsheng Zheng, Lanning Wei, Zhuoyun Yu, Shuofei Qiao, Jintian Zhang, Da Zheng, Yuren Mao, Yunjun Gao, Huajun Chen, Ningyu Zhang

arXiv:2506.10974v3cs.CLcs.AIcs.HCcs.LGcs.MA

TL;DR

Data-science agents have limited empirical expertise and inflexible coding strategies, constraining their fit for complex tasks. AutoMind combines expert knowledge, tree search, and adaptive coding; it outperforms baselines on two benchmarks, with an 8.0% gain over prior SOTA on the MLE-Bench leaderboard. Its performance also depends on the coding capability of its backbone model.

  • Problem

    Data-science agents lack practitioners’ empirical expertise and often use inflexible coding strategies oriented toward simpler, classical tasks.

    Source: §1 · p. 1

  • Method

    AutoMind combines an expert knowledge base, agentic knowledgeable tree search, and self-adaptive coding to address automated data-science tasks.

    Source: §1 · pp. 1–2, §7 · p. 9

  • Results

    AutoMind outperforms baselines on both benchmarks and improves 8.0% over prior SOTA on the MLE-Bench leaderboard.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors characterize AutoMind as an efficient and robust step toward fully automated data science, reporting higher efficiency and lower token costs than prior SOTA.

    Source: §1 · p. 2, §7 · p. 9

    AutoMind’s performance is tightly coupled with backbone coding proficiency and may lag behind prior agents when that capability is insufficient for complex solutions.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Large Language Model (LLM) agents have shown great potential in addressing real-world data science problems. LLM-driven data science agents promise to automate the entire machine learning pipeline, yet their real-world effectiveness remains limited. Existing frameworks depend on rigid, pre-defined workflows and inflexible coding strategies; consequently, they excel only on relatively simple, classical problems and fail to capture the empirical expertise that human practitioners bring to complex, innovative tasks. In this work, we introduce AutoMind, an adaptive, knowledgeable LLM-agent framework that overcomes these deficiencies through three key advances: (1) a curated expert knowledge base that grounds the agent in domain expert knowledge, (2) an agentic knowledgeable tree search algorithm that strategically explores possible solutions, and (3) a self-adaptive coding strategy that dynamically tailors code generation to task complexity. Evaluations on two automated data science benchmarks demonstrate that AutoMind delivers superior performance versus state-of-the-art baselines. Additional analyses confirm favorable effectiveness, efficiency, and qualitative solution quality, highlighting AutoMind as an efficient and robust step toward fully automated data science. Code is at https://github.com/innovatingAI/AutoMind.

1 Introduction

AutoMind addresses data-science agents’ limited empirical expertise and inflexible coding with an expert knowledge base, knowledgeable tree search, and adaptive coding. It outperforms baselines on two benchmarks, including an 8.0% improvement over prior SOTA on the MLE-Bench leaderboard.

1 Introduction

  • Existing agents lack practitioners’ empirical expertise and use inflexible coding strategies suited mainly to simpler, classical tasks.The paper argues that complex, innovative problems call for context-aware code generation.Source: §1 · p. 1
  • AutoMind combines an expert-curated knowledge base, agentic knowledgeable tree search, and self-adaptive coding to address these limitations.The search explores multiple solution pathways, while coding adapts to task complexity.Source: §1 · pp. 1–2
  • AutoMind surpasses 39.9% of human participants on average on the MLE-Bench leaderboard, an 8.0% improvement over prior SOTA, and outperforms baselines on both benchmarks.The authors also report higher efficiency and lower token costs than prior SOTA.Source: §1 · p. 2

2 Preliminaries

The paper frames automated data science as an optimization problem over candidate solutions and uses tree search to find an optimal solution. This directly compares candidate solutions rather than modeling task solving as a long-horizon reward-maximizing process.

2 Preliminaries

  • AutoMind models automated data science as an optimization problem and applies tree search to find an optimal solution.The approach builds on recent work integrating tree search with LLM-agent workflows.Source: §2 · p. 2
  • A candidate solution is a tuple containing a textual plan, Python code, and the validation metric used to assess execution results.The paper denotes these components as p, σ, and η, respectively.Source: §2 · p. 2
  • Unlike general agents’ long-horizon reward maximization, this formulation directly evaluates and compares possible data-science solutions.General-agent examples cited by the paper include Yao et al. (2023), Hong et al. (2024b), Shinn et al. (2023), and Chen et al. (2023).Source: §2 · p. 2

3 AUTOMIND

AutoMind retrieves expert knowledge and iteratively explores solution alternatives through a tree-search loop, adapting its coding process to solution complexity. Its knowledge base combines research papers and competition expertise, while retrieval uses task labels to locate relevant material.

3 AUTOMIND

  • AutoMind retrieves relevant knowledge, iteratively selects and expands solution-tree nodes, and chooses the best node when its iteration limit or time budget is reached.Coding is invoked during each action to reconcile solution complexity with LLM coding capabilities.Source: §3 · pp. 2–3

3.1 Expert Knowledge for Data Science

  • The expert knowledge base combines domain-specific research with expert-curated insights from top-ranked competition solutions.The paper motivates this resource by noting that static pretrained knowledge may lack domain-specific or up-to-date insights.Source: §3.1 · p. 3

3.1.1 Knowledge Base Construction

  • The knowledge base retains 3,237 valid forum posts spanning 455 Kaggle competitions and preserves metadata and main content from recent papers in selected conferences and journals.The forum posts provide publicly shared competition solutions; paper sources include KDD, ICLR, NeurIPS, ICML, EMNLP, and Bioinformatics.Source: §3.1.1 · p. 3

3.1.2 Knowledge Retrieval

  • Because task descriptions correlate weakly with available technical approaches, AutoMind labels tasks and knowledge hierarchically, retrieves by label similarity, filters same-target solutions, and reranks by label priority.The trick taxonomy has 11 top-level categories and corresponding subcategories; paper summaries support retrieval from data, task, and technique perspectives.Source: §3.1.2 · p. 3

3.2 Agentic Knowledgeable Tree Search

  • The search policy prioritizes debugging buggy leaves, sometimes improves the best-metric node, and otherwise explores other valid nodes to reduce the risk of overlooking better solutions.It first ensures a maximum number of draft nodes to establish branches for further exploration.Source: §3.2 · p. 4

3.3 Self-Adpative Coding Strategy

  • AutoMind scores plan complexity: straightforward plans are coded in one pass, while complex plans are decomposed into tested substeps that use execution errors as feedback.Substeps receive AST checks and terminal execution; a retry limit can force abandonment of the current plan.Source: §3.3 · pp. 4–5

4 Experiments

The experiments compare AutoMind with baseline agents across MLE-Bench and two recent AI competitions. AutoMind outperforms baselines across benchmarks and improves on the prior state of the art on MLE-Bench.

4.1 Experimental Setup

  • Experiments test GPT-4o, o3-mini, and DeepSeek-V3; each task receives a 24-hour budget, with three runs summarized using best@3 and avg@3.The MLE-Bench baselines include prior-work results for MLAB, OpenHands, and AIDE with GPT-4o, plus reproduced AIDE runs using o3-mini and DeepSeek-V3.Source: §4.1 · p. 5

4.2 Benchmarks

  • MLE-Bench evaluation uses win rate, the proportion of human participants whose leaderboard scores agents surpass; the lite benchmark has 15 tasks across three difficulty tiers.The evaluation also includes the WhoIsWho-IND track of the 2024 OAG Challenge, scored by AUC, and the 2024 BELKA Challenge, scored by average precision.Source: §4.2 · pp. 5–6

4.3 Main Results

  • AUTOMIND beats all baselines across benchmarks; on MLE-Bench avg@3, o3-mini and DeepSeek-V3 surpass 38.7% and 41.2% of humans, gains of 11.0% and 5.2% over AIDE.Under best@3, win-rate improvements over prior SOTA are 20.2% with o3-mini and 10.6% with DeepSeek-V3; on BELKA, o3-mini reaches 0.44 average precision, 0.35 above prior SOTA.Source: §4.3 · p. 6

5 Analysis

Ablations associate both the expert knowledge base and self-adaptive coding with improved outcomes. AutoMind also reaches AIDE’s 24-hour performance sooner, uses fewer tokens, and illustrates knowledge retrieval in a BELKA case study.

5.1 Ablation Study

  • Without the expert knowledge base, win rate falls 11.8% and valid rate falls 5.6% on Medium MLE-Bench.The authors attribute the gains to human-validated knowledge constraining search and helping avoid rediscovering effective ideas within the time budget.Source: §5.1 · p. 7
  • Replacing self-adaptive coding with one-pass generation lowers win rate by 13.2% and valid rate by 27.6% on Medium MLE-Bench.The authors attribute the decline to limited one-pass coding capacity for complex plans; decomposition with execution feedback can limit error accumulation, while simpler plans may use one-pass coding.Source: §5.1 · p. 7

5.2 Efficiency Analysis

  • TOMIND achieves the prior SOTA’s 24-hour performance within an average of 15 hours, representing a 60% improvement in time efficiency.Source: §5.2 · p. 7
  • At the time it matches prior SOTA’s 24-hour performance, AUTOMIND reduces token costs by 9.6%.Table 2 reports input, output, and total token costs across MLE-Bench tasks in millions of tokens.Source: §5.2 · p. 7, Table 2 · p. 7

5.3 Case Study

  • On BELKA, AUTOMIND retrieves MolTrans and DeepDTA, derives frequent-subsequence mining with dual-CNN blocks, and contrasts with the no-knowledge variant’s simple MLPs and AIDE’s gradient boosting.The authors say the retrieved papers support a more expressive model for the complex task.Source: §5.3 · p. 7

6 Related Work

Prior LLM-agent work commonly relies on customized workflows or field-specific training, while data-science agents also face gaps in empirical expertise and task-adaptive strategies. AutoMind proposes combining an expert knowledge base with dynamic coding-strategy selection; Figure 4 compares its BELKA plans and code with AIDE’s.

6 Related Work

  • Figure 4 compares the solution plans and code implementations generated by AIDE and AutoMind on BELKA.Source: Figure 4 · p. 8
  • LLM-agent frameworks commonly use customized workflows or field-specific fine-tuning, with newer approaches also learning through self-exploration and rule-based rewards.The passage distinguishes training-free general architectures from fine-tuned models.Source: §6 · p. 8
  • Data-science agents commonly divide work into bounded subtasks and workflows, but still lack practitioners’ empirical expertise and often use uniform strategies.DS-Agent instead adapts legacy expert Kaggle solutions through case-based reasoning.Source: §6 · p. 8
  • AutoMind addresses these gaps with an expert knowledge base drawn from resources such as papers and Kaggle competitions, plus dynamic coding-strategy selection.The proposal follows prior work using expert Kaggle solutions and case-based reasoning.Source: §6 · pp. 8–9

7 Conclusion and Future Work

AutoMind combines an expert knowledge base, knowledgeable tree search, and self-adaptive coding for automated data science. The paper reports superior performance on two benchmarks and efficiency improvements, and envisions a continuously evolving knowledge ecosystem.

7 Conclusion and Future Work

  • AutoMind integrates an expert knowledge base, knowledgeable tree search, and self-adaptive coding, with superior performance versus state-of-the-art baselines on two benchmarks.The authors describe the framework as an efficient and robust step toward fully automated data science.Source: §7 · p. 9
  • The experiments validate AutoMind’s design and quantify efficiency improvements.Source: §7 · p. 9
  • The authors envision a continuously evolving knowledge ecosystem in which agents read and synthesize papers and code and generate novel insights.Source: §7 · p. 9

Limitations

The evaluation is limited to a resource-constrained subset of MLE-Bench, and AutoMind’s performance depends on its backbone’s coding proficiency.

Limitations

  • Limited computational resources led the experiments to use 15 representative tasks rather than all 75 MLE-Bench tasks.Source: Limitations · p. 9
  • AutoMind’s performance is tightly coupled to backbone coding proficiency; insufficient capability may leave it behind agents favoring simpler solutions.Source: Limitations · p. 9
  • The authors say the proprietary backbones used in the experiments can better reflect the method’s potential to implement complex, effective solutions.Source: Limitations · p. 9

A Search Policy

The section points to Algorithm 1 for a detailed illustration of AutoMind’s search policy π.

A Search Policy

B.1 MLE-Bench

The evaluation uses a 15-task lite version of MLE-Bench, filtered for valid prior-SOTA submissions and balanced across task categories. Tasks are grouped into Easy, Medium, and Hard tiers.

B.1 MLE-Bench

  • The lite MLE-Bench subset contains 15 tasks, with at least one and at most two tasks retained per category.Tasks without a valid submission from the prior state of the art were excluded first.Source: §B.1 · p. 12
  • The original MLE-Bench has 75 offline Kaggle competitions, and the evaluation subset is divided into Easy, Medium, and Hard tiers.The tiers are based on human experience and prior-SOTA performance.Source: §B.1 · p. 12
  • Table 3 lists the MLE-Bench tasks used for evaluation.Source: §B.1 · p. 12

B.2 Top AI Competitions

This subsection describes two competition datasets used in evaluation: BELKA for molecular binding-affinity prediction and OAG for detecting incorrect paper assignments. It also reports a metric label and points to the MLE-Bench task list in Table 3.

B.2 Top AI Competitions

  • BELKA provides 2.9B molecule–protein training pairs; the study samples 2.2M training rows and 590K test rows and evaluates with AP score.The samples follow the label distribution of the full training data.Source: §B.2 · p. 12
  • OAG detects incorrect paper assignments using paper and author metadata, including titles, abstracts, organizations, keywords, venues, and years.Source: §B.2 · p. 12
  • For OAG, the study samples 10k papers and retains their corresponding authors, then splits the authors into training and test data following the training-label distribution.Source: §B.2 · p. 12
  • Table 3 is identified as the full list of MLE-Bench tasks used for evaluation.Source: Table 3 · p. 13
  • The evaluation metric is identified as PM.Source: §B.2 · p. 14

C Runtime Environment

This section documents the agents’ execution environment and points to hyperparameter tables for AutoMind and AIDE. It also presents prompts used across AutoMind’s pipeline, covering complexity assessment, code generation and planning, implementation, runtime and submission constraints, debugging, and automated checks.

C Runtime Environment

  • Agents run in Ubuntu 20.04 Docker containers with prepared task data, standard Python ML packages, and access to a 48-vCPU, 448GB RAM node with 9.6TB SSD storage and one NVIDIA Tesla V100 32G GPU.The Anaconda environment includes packages such as PyTorch and scikit-learn.Source: C Runtime Environment · p. 14

D Hyperparameters

  • AutoMind and AIDE hyperparameters are listed separately in Tables 4 and 5, with AutoMind’s table defining draft count and probabilities for debugging, using tricks, and greedy selection.The AIDE table caption identifies draft count and debugging probability; both tables define &TARGET_MODEL as the evaluated foundation model.Source: D Hyperparameters · p. 14, Table 4 · p. 15, Table 5 · p. 15

E Prompts

  • The prompt appendix presents drafting and iterative-improvement prompts that use task descriptions, data analysis, prior plans or outputs, and—when applicable—memory and expert tricks to guide solution plans.Drafting instructions request an implementable natural-language plan; improvement variants address debugging, incorporating tricks, or making one atomic expert improvement, while retaining relevant prior-plan details.Source: E Prompts · pp. 14, 16, 17, 18
  • The prompt asks an engineer to assess task and proposed-solution complexity, briefly explain the assessment, and return one average floating-point score from 1 to 5 in 0.5-point increments, considering dataset size and implementation and execution resources.Source: E Prompts · pp. 21–22
  • The one-pass coding prompt requires a self-contained Python script that implements the proposed plan, evaluates on hold-out data, and saves test predictions and the metric.It also specifies package availability and execution constraints, including using PyTorch for neural networks and avoiding multiprocessing.Source: E Prompts · pp. 22–23
  • The stepwise-decomposition prompt requires a JSON sequence of detailed steps that preserves the proposed solution and its hyperparameters, with evaluation and submission requirements included where relevant.The steps’ order determines implementation order, and the prompt prohibits exploratory data analysis and requires a hold-out metric rather than a fabricated score.Source: E Prompts · pp. 23–24
  • The stepwise-coding prompt asks the engineer to implement only the current plan step using prior code and any failed attempt, without repeating earlier code.It directs saving predictions and any calculated validation metric, while requiring code to run in the __main__ module and prohibiting multiprocessing.Source: E Prompts · pp. 24–26
  • Debugging prompts prescribe exact SEARCH/REPLACE edits; stepwise debugging is limited to the current failed code, while one-pass debugging uses an improved solution plan.The stepwise prompt says not to modify previous-step code.Source: E Prompts · pp. 26–28
  • The prompts require runnable code to compute a genuine hold-out metric and save it with test predictions under ./submission/.They also require a self-contained single-file script with computational routines under an __main__ guard.Source: E Prompts · p. 26
  • Execution guidance directs code to use ./input, avoid unavailable ./data, multiprocessing, temporary serialization, progress bars, and exploratory analysis, while permitting GPU use.It also advises concise tensor-shape tracking for model debugging and consulting data analysis guidance for data errors.Source: E Prompts · pp. 23, 27
  • Output verification checks for bugs, overfitting, a correctly located CSV submission, a validation metric, and whether lower metric values are better.The verifier is instructed to report empirical findings or identify a bug or missing submission.Source: E Prompts · p. 29
Loading 2506.10974v3…