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

TL;DR

SELA addresses suboptimal diversity and iteration in LLM-driven AutoML by introducing a tree-search framework that uses Monte Carlo Tree Search to explore a tree-structured pipeline space. It combines an insight proposer, an MCTS-based search module, and an experiment executor to iteratively propose, test, and refine ML configurations across five stages (EDA to evaluation). Empirical results on 20 AML datasets show SELA achieving higher average NS and best-rank metrics, beating traditional AutoML and existing LLM-based baselines, with ablations confirming the value of MCTS, rollouts, and LLM adaptability. The approach highlights the potential of structured search and iterative refinement for dynamic, human-like problem solving in AutoML and potentially other sequential decision tasks.

Abstract

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.

Figures (5)

• Figure 1: SELA's abstraction compared to other agent-based AutoML frameworks. There are two main types of agent-based approaches to AutoML problems. The first approach hong2024datainterpreterllmagent divides a machine learning task into multiple stages, proposing a plan for each stage, and generating and executing code step by step according to the plan, with no refinement after the solution is completed. The second Schmidt_Wu_Jiang generates the entire solution in one step and iteratively refines it as a whole. SELA integrates both approaches, enabling stage-wise planning while iteratively exploring better solutions at each stage level.
• Figure 2: SELA’s pipeline operates as follows: The system begins by inputting the problem description and dataset information into the LLM, which generates a search space of potential solutions, encompassing data preprocessing, feature engineering, and model training. The search module, powered by Monte Carlo Tree Search (MCTS), explores this space by selecting, expanding, and simulating potential configurations. The LLM agent then simulates the selected configuration by planning, coding, and executing the experiment. Feedback from the simulation is fed back into the search module, where it is used in the backpropagation step to refine future searches. This iterative process continues until a predefined stopping criterion is met, resulting in an optimized experimental pipeline.
• Figure 3: Rescaled NS of AutoML frameworks relative to SELA on tabular datasets. Points to the left of the vertical line indicate poorer predictions compared to SELA. Notably, SELA often occupies a leading position across the datasets.
• Figure 4: The average performance of SELA on six selected datasets with an increasing number of rollouts.
• Figure 5: Comparison of Normalized Scores between different base LLMs on six selected datasets.