PaperlayerView exact revision on arXiv

Source-linked AI summary

SELA: Tree-Search Enhanced LLM Agents for Automated Machine Learning

Yizhou Chi, Yizhang Lin, Sirui Hong, Duyi Pan, Yaying Fei, Guanghao Mei, Bangbang Liu, Tianqi Pang, Jacky Kwok, Ceyao Zhang, Bang Liu, Chenglin Wu

arXiv:2410.17238v1cs.AIcs.CLcs.LGcs.SE

TL;DR

Traditional AutoML uses fixed search routines, while LLM-based agents can struggle to explore alternatives and refine solutions. SELA combines stage-wise planning with MCTS-guided experimentation, and competing methods lose to SELA on 65–80% of datasets.

  • Problem

    Traditional AutoML relies on predefined search spaces, while LLM-based solutions often have limited exploration and iterative refinement, leaving workflows suboptimal.

    Source: §1 · pp. 1–2

  • Method

    SELA combines stage-wise LLM planning with MCTS over a tree-structured space, using experiment feedback to guide subsequent pipeline configurations.

    Source: §1 · p. 2, §3.2 · p. 5, §3.3.2 · p. 6

  • Results

    65–80% of datasets: the “Losses” column shows each competing method falling short of SELA.

    Source: §4.2 · pp. 8–9

  • Takeaways & Limitations

    SELA’s tree-search approach may extend to sequential decision problems representable as trees with scalar rewards at their leaf nodes.

    Source: §5 · p. 10

    Beyond machine learning, the proposed adaptation is scoped to sequential decision problems representable as trees with scalar leaf rewards.

    Source: §5 · p. 10

Abstract

from arXiv · show

Automated Machine Learning (AutoML) approaches encompass traditional methods that optimize fixed pipelines for model selection and ensembling, as well as newer LLM-based frameworks that autonomously build pipelines. While LLM-based agents have shown promise in automating machine learning tasks, they often generate low-diversity and suboptimal code, even after multiple iterations. To overcome these limitations, we introduce Tree-Search Enhanced LLM Agents (SELA), an innovative agent-based system that leverages Monte Carlo Tree Search (MCTS) to optimize the AutoML process. By representing pipeline configurations as trees, our framework enables agents to conduct experiments intelligently and iteratively refine their strategies, facilitating a more effective exploration of the machine learning solution space. This novel approach allows SELA to discover optimal pathways based on experimental feedback, improving the overall quality of the solutions. In an extensive evaluation across 20 machine learning datasets, we compare the performance of traditional and agent-based AutoML methods, demonstrating that SELA achieves a win rate of 65% to 80% against each baseline across all datasets. These results underscore the significant potential of agent-based strategies in AutoML, offering a fresh perspective on tackling complex machine learning challenges.

1 INTRODUCTION

SELA combines stage-wise planning with iterative refinement, using experimental feedback to search for better machine-learning configurations. In evaluation across 20 diverse datasets, it outperformed traditional and agent-based AutoML systems.

1 INTRODUCTION

  • Traditional AutoML relies on predefined search spaces and routines, while LLM agents can produce low-diversity, suboptimal workflows with limited iterative exploration.Traditional systems also leave preprocessing and feature engineering underexplored.Source: §1 · pp. 1–2
  • SELA combines stage-wise planning with iterative refinement, exploring better solutions at each stage.Its stages include Exploratory Data Analysis, Data Preprocessing, Feature Engineering, and Model Training.Source: §1 · p. 2, Figure 1 · p. 2
  • MCTS searches a tree of potential configurations, balancing exploration of new strategies with exploitation of promising ones.Experimental results feed back into the search so the agent can select subsequent configurations.Source: §1 · pp. 2–3
  • SELA consistently outperformed traditional AutoML systems and agent-based approaches across 20 diverse datasets.The evaluation used datasets from the AutoML Benchmark.Source: §1 · p. 2

2 RELATED WORKS

Prior work spans traditional AutoML, LLM-enhanced pipeline components, and dynamic LLM agents; tree search has also been applied across several AI domains. SELA combines stage-wise planning and iterative refinement, with a pipeline that uses MCTS to explore configurations and return experiment feedback to the search.

2 RELATED WORKS

  • Tree search, including MCTS, has been used in robotics, chemistry, gaming, and LLM-based exploration and planning.Prior work also adapted AlphaZero-style tree search to LLM tasks.Source: §2 · p. 3
  • Traditional AutoML advanced through hyperparameter optimization, model selection, stacking, ensembling, and meta-learning strategies.Extensions have broadened AutoML to multimodal data settings.Source: §2 · p. 3
  • LLMs have been applied to feature engineering, hyperparameter tuning, and multiple workflow stages, while traditional systems remain constrained by rigid pipelines.These approaches address pipeline flexibility for unique datasets and task requirements.Source: §2 · p. 3
  • Dynamic LLM agents have used hierarchical graph modeling, programmable node generation, structured inputs, past experience, and case-based reasoning to build adaptable pipelines.Guo et al. (2024) face challenges generating solutions from scratch because of reliance on existing codebases; Schmidt et al. (2024) iteratively refines a whole pipeline generated in one step.Source: §2 · pp. 3–4
  • SELA’s pipeline has an LLM insight proposer define a search space, MCTS explore it as a tree, and an LLM experiment executor run selected configurations and return results.The feedback is used to refine future searches.Source: Figure 2 · p. 4
  • SELA integrates stage-wise planning with iterative refinement, enabling autonomous exploration and generation of machine-learning solutions from the ground up.The paper describes this as providing greater flexibility and control during search, with optimized solutions generated at each stage.Source: §2 · p. 4

3 METHOD

SELA combines stage-wise LLM planning and pipeline execution with MCTS over experiment configurations. Experiments provide feedback that guides subsequent search, while stage-level code reuse supports efficient execution.

3 METHOD

  • SELA has an LLM insight proposer, an MCTS search module, and an LLM experiment executor that feeds experiment results into future searches.The cycle continues until its termination criterion is met.Source: §3 · p. 4

3.1 INSIGHT PROPOSAL AND SEARCH SPACE CREATION

  • The insight proposer uses the problem description and dataset information to generate stage-specific insights that form SELA’s search space.The five stages are exploratory data analysis, data preprocessing, feature engineering, model training, and model evaluation.Source: §3.1 · p. 4

3.2 PIPELINE EXECUTION AND CODE GENERATION

  • For each configuration, the experiment executor plans stage-specific instructions, writes and runs code for each task, and concatenates the stage code into a full solution.Execution produces a final score.Source: §3.2 · p. 5

3.3 TREE SEARCH IN MACHINE LEARNING EXPERIMENTS

  • SELA models machine-learning experiment configurations as paths through a hierarchical tree so tree-search algorithms can explore them.The search process uses repeated rollouts to seek the best solution.Source: §3.3 · p. 5, §3.3.2 · p. 6

3.3.1 EXPERIMENT NODE

  • The search space is a hierarchical tree in which each node represents an insight and each root-to-node path defines an experiment configuration.Nodes track simulation scores, visit counts, and solution or stage code.Source: §3.3.1 · p. 5

3.3.2 TREE SEARCH FOR ML EXPERIMENTS

  • MCTS selects and expands paths, simulates sampled configurations, and backpropagates scores and solution code to guide later rollouts.Its depth-preferred UCT selection prioritizes deeper nodes early, while an unvisited-node preference balances exploration and computational efficiency.Source: §3.3.2 · pp. 6–7

3.3.3 EXPERIMENT STATE SAVING AND LOADING

  • SELA caches stage-level code and reruns saved code when configurations share components, conserving resources and addressing variation from LLM non-determinism.Saved stage code is loaded when a node is selected for execution.Source: §3.3.3 · p. 7

4 EXPERIMENTS

SELA was evaluated against traditional and agent-based AutoML on 20 datasets, with ablations testing search strategy, rollout count, and base-LLM variation. It achieved the highest average normalized score and best rank, although AutoGluon had a marginally higher average rank.

4.1 EXPERIMENTAL SETUP

  • Evaluation covered 20 datasets—13 classification and 7 regression tasks—with four baselines and 6:2:2 training, validation, and test splits.LLM-based methods shared an initial prompt and used DeepSeek v2.5; AIDE ran 10 iterations and SELA 10 rollouts.Source: §4.1 · pp. 7–8

4.2 RESULTS

  • SELA had the highest average NS and best rank, the most Top 1 predictions, and baseline losses of 65–80%; Figure 3 shows it often leading, though AutoGluon’s average rank was slightly higher.Table 2 reports the framework results across 20 datasets.Source: Table 2 · p. 8, §4.2 · pp. 8–9, Figure 3 · p. 8

4.3 ABLATION STUDY

  • MCTS achieved higher average NS and a better overall ranking than random search; even random-search SELA outperformed Data Interpreter.The authors also report that their insight proposer supplies useful insights regardless of the selection method.Source: Table 3 · p. 9, §4.3 · p. 9
  • SELA delivers similar results across different LLMs on six selected datasets, indicating flexibility with various models depending on user preference and availability; numeric results are reported in Appendix C.2.Source: Figure 5 · p. 10, §4.3 · p. 10
  • SELA’s normalized score improved 4.7% after 10 rollouts and 6.4% after 20, compared with the initial rollout.Figure 4 shows average performance across six selected datasets as permitted rollouts increase.Source: §4.3 · p. 9, Figure 4 · p. 10

5 CONCLUSION

SELA combines LLM-based agents with Monte Carlo Tree Search to automate machine-learning workflows. The paper reports evaluation on 20 datasets and identifies broader applications and search improvements as future directions.

5 CONCLUSION

  • SELA integrates LLM-based agents with Monte Carlo Tree Search to automate machine-learning workflows.The conclusion presents this combination as the paper's framework.Source: §5 · p. 10
  • The authors report experiments on 20 machine-learning datasets and advantages over traditional AutoML frameworks and existing LLM-based approaches.The conclusion does not give metric-level comparisons here.Source: §5 · p. 10
  • SELA could be adapted to other sequential decision problems when they use tree structures and scalar rewards from leaf nodes.Software engineering, scientific discovery, game playing, and robotics are named as possible domains.Source: §5 · p. 10
  • The conclusion identifies search efficiency and scalability for larger solution spaces, and interpretable explanations, as areas for future work.The authors also point to transparency and trustworthiness as motivations for interpretability work.Source: §5 · p. 10

A DATASETS

Table 4 summarizes the datasets used in the evaluation. Its caption notes that OpenML datasets have dataset IDs and Kaggle datasets are accessible through competition sources.

A DATASETS

  • Table 4 is the dataset-summary table for the evaluation.The caption describes it as outlining detailed dataset information.Source: Table 4 · p. 13, A DATASETS · p. 13
  • The caption says OpenML datasets can be accessed using their respective dataset IDs.Source: Table 4 · p. 13
  • The caption points to Kaggle competition sources for access to the Kaggle datasets.It gives a competition-source URL pattern rather than listing dataset names in the passage.Source: Table 4 · p. 13

B PROMPTS

The appendix describes prompts for dataset tasks, framework instructions, and generating stage-specific insights. Together, the prompts specify the task inputs and expected outputs and ask for diverse, implementable insights.

B.1 TASK PROMPT

  • The base task prompt supplies the dataset name, target label, metric, file path, and metadata path to LLM-based methods.Source: §B.1 · p. 14
  • The task template asks the model to analyze, preprocess, engineer features, and model the dataset to predict the target.It provides train and development data with labels, and a test set without labels.Source: §B.1 · p. 14
  • For AIDE, the original training and validation sets are combined as input, with one-fold validation configured to align more closely with the other setup.Both setups give frameworks access to labels for the training and validation sets.Source: §B.1 · p. 14

B.2 INSTRUCTION PROMPT

  • The instruction prompt directs frameworks to save the final prediction file for evaluation.Source: §B.2 · p. 14
  • The instruction template prohibits target leakage, requires dev and test prediction files, and asks frameworks to print train and dev performance.It also specifies prediction formatting and coordinated transformations across train, dev, and test sets.Source: §B.2 · p. 14

B.3 INSIGHT PROPOSAL PROMPT

  • The Insight Proposer prompt uses dataset information to generate stage-specific insights intended to improve model performance.Source: §B.3 · p. 15
  • The insight prompt requests at least five diverse, separately implementable insights per task type, with specific model, ensemble, tuning, preprocessing, and feature-engineering choices.The template lists EDA, Data Preprocessing, Feature Engineering, and Model Training as task types.Source: §B.3 · p. 15

C.1 MAIN RESULTS

The main-results tables present methods' rankings and Normalized Score percentages across tabular datasets. The supplied captions identify their scope but do not state which methods perform best.

C.1 MAIN RESULTS

  • Table 5 presents methods' rankings for each tabular dataset.Source: Table 5 · p. 16
  • Together, the tables cover both per-dataset method rankings and per-dataset NS percentages.Source: Table 5 · p. 16, Table 6 · p. 16
  • Table 6 presents methods' Normalized Score percentages for each tabular dataset.Source: Table 6 · p. 16

C.2 PERFORMANCE USING DIFFERENT LLMS

Table 7 reports SELA results with different base LLMs on selected tabular datasets.

C.2 PERFORMANCE USING DIFFERENT LLMS

  • Table 7 compares SELA using different base LLMs on selected tabular datasets.The supplied passage gives the table's scope but no cell values.Source: Table 7 · p. 17

D COST-EFFECTIVENESS ANALYSIS

The cost analysis estimates per-task agent-framework costs over iterations or rollouts and describes mechanisms associated with lower execution costs.

D COST-EFFECTIVENESS ANALYSIS

  • SELA is less costly than Data Interpreter, its base experimenter, because state saving and loading reduces repeated task and code generation.The comparison is reported for execution cost; Table 8 covers estimated costs on a single dataset over k iterations or rollouts.Source: D COST-EFFECTIVENESS ANALYSIS · p. 17, Table 8 · p. 17
  • AIDE has the lowest execution cost among the compared methods, attributed to fewer token generations without stage-wise planning.Source: D COST-EFFECTIVENESS ANALYSIS · p. 17
  • The analysis estimates average running costs for LLM-based baselines using multiple execution trials.Source: D COST-EFFECTIVENESS ANALYSIS · p. 17

E.1 OVERVIEW OF SELA’S SEARCH PROCESS

SELA uses MCTS to search a staged machine-learning workflow through repeated selection, expansion, simulation, and feedback updates. In the case study, the best-performing node combines time-based features with an RBF-kernel SVM.

E.1 OVERVIEW OF SELA’S SEARCH PROCESS

  • The case study shows search over alternative feature strategies and model choices, including a weighted-ensemble branch.The trace reports 29 unique codes generated.Source: §E.1 · pp. 18, 20, 21
  • SELA structures the workflow into stages, initializes candidate insights, and expands a selected time-feature node into alternative model strategies.The example alternatives include Random Forests, SVMs, Neural Networks, Gradient Boosting, and Grid Search.Source: §E.1 · p. 21
  • Each search cycle simulates an expanded candidate, evaluates it, and backpropagates its score to parent nodes before the cycle repeats to a stopping condition.Source: §E.1 · p. 21
  • Node 0-1-1, using an SVM with an RBF kernel, achieved the highest score in the current dataset, with test_score 0.6928451194338062 and dev_score 0.6944266833187726.Source: §E.1 · pp. 19, 21

Related papers

Loading 2410.17238v1…