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

Authors: Yizhou Chi, Yizhang Lin, Duyi Pan, Bangbang Liu, Guanghao Mei, Yaying Fei, Tianqi Pang, Sirui Hong, Jacky Kwok, Bang Liu, Ceyao Zhang, Chenglin Wu
Primary area: applications to robotics, autonomy, planning
Keywords: AutoML, AutoDS, LLM, Agents, Tree Search, Planning

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-$\textbf{S}$earch $\textbf{E}$nhanced $\textbf{L}$LM $\textbf{A}$gents ($\textbf{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. The code will be open-sourced upon publication.

Decision: reject
Reviewer scores: 3, 5, 1, 5

Reviews and discussion:
<thread>
creation date:
1730463259940

thread_type:
official_review

<thread_object>
creator: reviewer_epej

<content>
summary:
The paper introduces SELA, a novel framework that integrates Monte Carlo Tree Search (MCTS) with Large Language Model (LLM) agents to enhance the process of Automated Machine Learning (AutoML). The authors claim that SELA improves the exploration and optimization of machine learning pipelines by leveraging MCTS to systematically explore potential configurations. The framework is evaluated on 20 datasets, demonstrating a win rate of 65% to 80% against traditional and agent-based AutoML methods.

strengths:
- The integration of MCTS with LLM agents is a creative approach to address the limitations of existing AutoML frameworks. By using MCTS, SELA aims to mimic the iterative and feedback-driven problem-solving approach of human experts, which is a commendable attempt to enhance the adaptability and effectiveness of AutoML systems.
- The experimental results indicate that SELA outperforms several baseline methods across a range of datasets. This suggests that the proposed method has potential advantages in terms of flexibility and performance, offering a fresh perspective on tackling complex machine learning challenges.
- The paper provides a detailed description of the SELA framework, including its components such as the insight proposer, experiment executor, and the MCTS-based search module. This thorough explanation helps in understanding the workflow and the role of each component in the overall system.

weaknesses:
- While the integration of MCTS with LLM agents is an interesting approach, the innovation appears limited. The application of MCTS in this context is relatively straightforward and does not introduce significant novel techniques beyond existing LLM-agent frameworks like Data Interpreter. The paper could benefit from a deeper exploration of how MCTS specifically enhances the LLM agent's capabilities in a unique way.
- The experimental evaluation, although showing promising results, is somewhat inadequate. The authors do not provide a clear rationale for the selection of the 20 datasets from the AutoML Benchmark (AMLB). Given the diversity of tasks available in AMLB, it would be beneficial to understand the principles behind the choice of these specific tasks to ensure a comprehensive assessment of SELA's capabilities.
- The paper lacks a detailed analysis of the results and the specific contributions of MCTS to the observed performance improvements. A more in-depth discussion on how MCTS influences the exploration and optimization process, compared to other search strategies (beyond random ones), would strengthen the paper's contributions. In particular, UCB typically requires a significant number of exploitation steps to accurately evaluate the value of each node. However, this would incur a huge workload in an AutoML scenario. It would be beneficial to discuss how the authors achieve a good trade-off between efficiency and effectiveness, and why limited exploitation of UCB suffices to gain a performance improvement.

questions:
- The experimental evaluation could be more comprehensive. Further exploration and justification of the chosen datasets, along with a deeper analysis of the results, would enhance the paper's impact and contribution to the field.
- Can you provide details on the size of the search pool for the MCTS procedure? Additionally, how many iterations does MCTS typically require to achieve good performance?

My score may be revised upon reviewing the authors' response to these limitations.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730581404308

thread_type:
official_review

<thread_object>
creator: reviewer_ppui

<content>
summary:
The work introduces an LLM based AutoML tool that uses MCTS to search for better performing configurations.

strengths:
The ablation study showing the strenghts of different base LLMs is interesting.

weaknesses:
The work largely seems to be unfamiliar with AutoML tools and the broader literature. This is particularly evident with the work repeatedly claiming falsehoods about AutoML. For example, early on it is said that AutoML tools typically only focus on the model training aspect, while ignoring feature engineering stages and so on. This claim is easily refuted as many of the cited tools do include data preprocessing and feature preprocessing into account and generally tend to the whole machine learning pipeline. It is often left to the user however what to include in the search space, thus a user might only want to focus on the model training element. Further, the work claims as a contribution that the search is inspired by how humans would conduct hyperparameter search/configuration. I.e., repeated experimentation while refining the search. This is exactly what nearly all AutoML tools do. Many are based on bayesian optimization which typically executes an initial design to probe the search space, then queries the most promising configuration(s) based on an internal model, based on the observation of the actual performance the model is refined and the search continues until a configuration budget is exhausted. Again, this is the standard approach in AutoML, making the claim in "contribution 1" not a contribution of this work. Further, contribution 2 hinges on the use of MCTS for AutoML (with an LLM base). However, MCTS has been already used and studied in the realm of AutoML (see the work by Rakotoarison https://inria.hal.science/hal-01966957/file/MOSAIC-wkpICML2018.pdf) strongly diminishing this claimed contribution. Further, the experiments do not consider an AutoML approach that uses MCTS for comparison.

I was already skeptical towards the experiments due to the untruths written about AutoML. I am further convinced that the evaluation of the aproaches is completely unfair. For AutoGluon and AutoSKLearn it is stated that they are simply used in their default configurations. However, it is not clear how much compute budget each method was allocated and that SELA very likely had a much larger configuration budget available than all other methods. Despite these weird setup issues, AutoGluon shows very strong performance, which is waved aside by the work. An important baseline that needs to be compared against in the AutoML realm is also missing. TabPFN uses a foundation model trained on synthetic data to learn to solve ML tasks and has shown highly competitive or even outperforming classical AutoML systems. This method thus has to be considered in a work as it is presented here. Overall the work can absolutely not be published. The work has to be largely rewritten to remove the false claims and the experiments likely need to be completely redone. I do not see immediate action items that would increase my trust in the work and would make me reconsider my scoring.

questions:
What is the configuration budget allocated to all methods? Since you mentioned AutoSKlearns 600seconds wall-clock time of budget, did you limit SELA to the same small budget?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730707814498

thread_type:
official_review

<thread_object>
creator: reviewer_u8rg

<content>
summary:
This paper introduces SELA (Tree-Search Enhanced LLM Agents), a novel approach that combines Large Language Model (LLM) agents with Monte Carlo Tree Search (MCTS) for automated machine learning (AutoML). The primary contribution is defining an AutoML pipeline using a tree structure and experimentally demonstrating its effectiveness.

strengths:
The paper presents an interesting pipeline for AutoML based on LLM agents, where the functional attributes of each node layer are intuitively and innovatively defined.

weaknesses:
1. The paper lacks sufficient detail and resources for reproducibility.
2. The study involves multiple modules and details; however, the results in Table 2 appear to show only marginal improvements.
3. Additional issues are mentioned in the "Questions" section.

questions:
1. Regarding the `simulate` phase of MCTS, I understand that your tree has a depth of 6 levels (including the root node). Consequently, $x_{\text{sample}}$ does not necessarily correspond to the final Model Evaluation node, so $c(x_\text{sample})$ does not represent a complete configuration file. How, then, can this be provided to the experimenter $E$ to obtain a score-evaluable code?

2. Regarding lines 205 and 323, when I expand a node at time $\tau_i$, I provide the problem description $p$, dataset information $d$, and $\lambda(x_i)$ (where $\lambda(x_i)$ might concatenate with $\lambda(x_{i-1})$) as the LLM’s prompt, generating $m$ values of $\lambda(x_{i+1})$, from which one is randomly sampled for simulation—is this correct?

3. Given that even the most SOTA models struggle to generate immediately executable code, especially for complex ML tasks, how do you handle cases where the code provided by the LLM fails to execute? Do you assign it a minimum score or apply a different approach?

4. In lines 242–258, could you clarify the distinction between Value and Simulation Score, as well as between Solution Code and Stage Code? My understanding is that Stage Code corresponds to the current node’s code, but is Solution Code a set of codes of child branches? Since the node may have multiple branches. 

5. In line 363, "Each framework utilizes the training and validation sets to train models," but I do not see any component in SELA that can be trained.

6. Regarding line 443, as a comparison, how exactly does SELA refine its Solution? I only understand that it finds the optimal solution through MCTS Search, but if a node's result is not good, it seems they won't refine it either, they just won't select it.

7. Line 393 mentions each AutoML rollout is conducted 10 times. Isn’t this count too low for a tree search, potentially resulting in a sparse tree structure?

8. Line 1160 states that each node corresponds to a single feature; however, combining multiple feature engineering techniques is standard practice. Does SELA not support such feature combinations?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721265473

thread_type:
official_review

<thread_object>
creator: reviewer_kfk4

<content>
summary:
The paper proposes a method (SELA) that combines LLMs with MCTS to solve AutoML tasks. More concretely, SELA uses LLMs to generate a search space of solutions (essentially generating different possibilities across different stages in the ML pipeline), which is then explored using MCTS. Through experiments on tasks from the AutoML benchmark and Kaggle competitions, the paper demonstrates SELA's efficacy against baseline methods.

strengths:
- The paper conveys the main ideas and results clearly.
- Handling the NP-hardness of the problem by incorporating LLM agents with a search procedure is intuitive and makes a lot of sense.

weaknesses:
My main concern with this paper is regarding the novelty of the proposed method. The authors note that prior work has already explored combining LLMs with tree search methods. This makes it unclear if the technical contribution of this paper is novel enough for ICLR. The primary contribution of applying LLMs with tree search to AutoML could be an interesting contribution had the paper shown that tree search cannot be trivially combined with LLMs to solve AutoML problems, and highlighted the key challenges faced.

questions:
Could the authors elaborate on some key challenges faced when integrating LLM agents with tree search for AutoML tasks, and how SELA addresses them? All in all, I am trying to understand if the application is straightforward.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732246789063

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Comments [1/2]

comment:
We really appreciate the feedback provided by all reviewers. To further help the audience understand our work, we provide concrete examples of insights and solutions generated by SELA in supplementary materials. In the meantime, we want to address some common issues raised by the reviewers and seek to clarify some points that might be misunderstood. 

# 1. Novelty comparing other AutoML frameworks

Given that the experimental results have demonstrated the superior performance of our method against the chosen baselines, we believe it is also beneficial to qualitatively elaborate on the unique values SELA brings about beyond the baselines. We summarize the key differences between SELA and other frameworks in the following table. 

|       | Dynamic Pipeline | Adaptive Feature Engineering    | Model Training    | Model Improvement          | Pipeline Optimization    | 
|---|---|---|---|---|---| 
| AutoGluon | ✗   | ✗    | Fixed models   | Multi-layer stacking + bagging | ✗       |
| AutoSklearn | ✗ | ✗    | Fixed models   | Bayes Opt. + meta-learning + ensemble | ✗   | 
| Data Interpreter | ✓   | Instinctive  | Instinctive | Instinctive   | ✗    | 
| AIDE  | ✓ | Instinctive | Dynamic & diverse| Dynamic & diverse | One-step refinement + LLM   | 
| **SELA**| ✓ |Dynamic & diverse | Dynamic & diverse | Dynamic & diverse |Stepwise MCTS + LLM|  

**Table**: Comparison of key capabilities across various AutoML methods. Dynamic indicates the system's ability to adjust workflows based on intermediate outcomes, allowing it to adapt as new information emerges. Diverse refers to employing multiple strategies or methods across tasks, capturing varied modeling needs. Instinctive means that the system directly relies on the decisions generated by an LLM and heavily depends on the model's inclination.

Characterized in the table, LLM Agents like Data Interpreter craft solutions dynamically but are unable to improve them. In contrast, traditional AutoML methods like AutoGluon and AutoSklearn employ optimization strategies but rely on static pipelines. We argue that our paper elegantly overcomes the limitation of both types of methods by mimicking human experts to dynamically construct pipelines and optimize them iteratively. Our contribution may be easier to understand with two detailed comparisons.

**Comparing with existing LLM Agents:** 

While SELA leverages Data Interpreter [1] as its Experimenter, their focuses have fundamental differences. Data Interpreter is an LLM agent aiming to produce functional and performant pipelines for ML problems. Its main techniques such as Hierarchical Graph Modeling and Programmable Node Generation are proposed to guarantee successful construction of the pipeline. However, Data Interpreter is unable to further improve the pipeline after the first construction, even given multiple trials. SELA, on the other hand, abstracts away the pipeline-building process assuming a capable Experimenter (section 3.2) and places a crucial emphasis on optimization. The optimization consists of an insight proposer for diverse configurations and a stepwise MCTS process for efficient exploration, allowing continuous discovery of better pipelines, which Data Interpreter falls short of. In fact, when presented with a Data-Interpreter-like agent, it is not obviously clear how we can equip it with an optimization capability. We see limited literature discussing such strategies for ML agents (AIDE [2] counts as one, and we include it in our evaluation comparison and show its performance can be surpassed). By introducing a structural and systematic search process, we provide an intuitive yet highly effective methodology to fill the blank.

**Comparison with traditional AutoML frameworks:** 

Our method shares the same important idea of searching with traditional AutoML frameworks, introducing systematic and feedback-driven optimization procedures crucial to acquiring competitive solutions for ML problems. However, our method is distinctive in two ways. First, SELA proposes a dynamic and task-aware search space consisting of natural language insights specifically based on the datasets. For example, when the dataset requires house price prediction, our Insight Proposer may put forward options such as "Create a feature by calculating the age of the house at the time of sale", a semantically meaningful and potentially useful feature operation that we cannot expect from frameworks such as AutoGluon and AutoSklearn. We have included more examples of insights in the supplementary materials for demonstration. Additionally, traditional AutoML frameworks face significant challenges in performing end-to-end ML tasks starting from natural language task descriptions and dataset information to producing final solutions. In contrast, SELA seamlessly executes the entire pipeline: interpreting the task description and dataset, building and optimizing the ML pipelines, and delivering final predictions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes to do a specific variant of MCTS for AutoML tasks, where the possible actions consist of standard data-science actions, e.g. data processing, feature engineering, proposing models, etc.

Comparisons are made against standard AutoML packages such as AutoGluon, AutoSklearn, with generally competitive results.

## Strengths: 
* The main message of the paper, from my perspective, is showing that LLMs can perform agentic approaches for AutoML problems, which has the potential for future work in e.g. data science, scientific discoveries, AI for AI, etc.
* If we assume that the experiments were rigorously done, then results seem comparable to SOTA AutoML platforms (AutoGluon, etc.)

## Weaknesses:
* There was an overall sentiment of "It's just another LLM can do XYZ" paper.
* In my opinion, the writing style is pretty opaque - Section 3 I think has a poor combination of notation and wordiness, that made it difficult to read.

Unfortunately, none of the reviewers wishes to champion this paper, as the highest score as a 5. The overall message "LLM Agents can do AutoML" wasn't well-received, as most of the negative feedback came from novelty and experimental rigorousness.

In general, this paper might not be suitable for ICLR, and I actually recommend the authors send this paper to a LLM-focused conference, provided that they rewrite the paper to focus more on the overall agentic message. Otherwise, I can forsee it having the same feedback issues if submitted to NeurIPS/ICML (lack of novelty, etc.). This is unless, somehow the proposed method would lead to significantly better experimental outcomes.