I-MCTS: Enhancing Agentic AutoML via Introspective Monte Carlo Tree Search

TL;DR

This work addresses limitations of current LLM-based AutoML agents, notably limited thought diversity and suboptimal code generation. It introduces Introspective Monte Carlo Tree Search (I-MCTS), which couples introspective node expansion with a hybrid reward mechanism driven by an LLM value model to guide search more effectively, including pre-rollout node evaluation. The approach yields a 4% absolute improvement in performance over strong AutoML baselines across 20 tabular datasets, with ablations confirming the contributions of introspection and hybrid rewards. The findings suggest that integrating reflective reasoning into tree-search can enhance agentic AutoML systems, potentially enabling more efficient and robust automated ML workflows in practice.

Abstract

Recent advancements in large language models (LLMs) have shown remarkable potential in automating machine learning tasks. However, existing LLM-based agents often struggle with low-diversity and suboptimal code generation. While recent work has introduced Monte Carlo Tree Search (MCTS) to address these issues, limitations persist in the quality and diversity of thoughts generated, as well as in the scalar value feedback mechanisms used for node selection. In this study, we introduce Introspective Monte Carlo Tree Search (I-MCTS), a novel approach that iteratively expands tree nodes through an introspective process that meticulously analyzes solutions and results from parent and sibling nodes. This facilitates a continuous refinement of the node in the search tree, thereby enhancing the overall decision-making process. Furthermore, we integrate a Large Language Model (LLM)-based value model to facilitate direct evaluation of each node's solution prior to conducting comprehensive computational rollouts. A hybrid rewarding mechanism is implemented to seamlessly transition the Q-value from LLM-estimated scores to actual performance scores. This allows higher-quality nodes to be traversed earlier. Applied to the various ML tasks, our approach demonstrates a 4% absolute improvement in performance compared to the strong open-source AutoML agents, showcasing its effectiveness in enhancing agentic AutoML systems. Resource available at https://github.com/jokieleung/I-MCTS

Figures (1)

• Figure 1: I-MCTS architecture featuring (a) Introspective node expansion through parent/sibling analysis, (b) Hybrid reward calculation combining LLM predictions and empirical scores. The red arrows indicate the introspective feedback loop that continuously improves node quality. For the full prompt design, please refer to Appendix \ref{['app:prompt']}.