In-Context Decision Making for Optimizing Complex AutoML Pipelines

TL;DR

This paper introduces PS-PFN, a novel AutoML approach that extends CASH to modern ML pipelines using in-context learning and bandit algorithms, demonstrating superior performance on benchmarks.

Contribution

It extends the CASH framework with PS-PFN, leveraging posterior sampling and PFNs for efficient pipeline selection and adaptation in complex ML workflows.

Findings

PS-PFN outperforms other bandit and AutoML strategies on benchmark tasks.

The method effectively models reward distributions with PFNs for pipeline optimization.

Experimental results show improved efficiency and accuracy in pipeline selection.

Abstract

Combined Algorithm Selection and Hyperparameter Optimization (CASH) has been fundamental to traditional AutoML systems. However, with the advancements of pre-trained models, modern ML workflows go beyond hyperparameter optimization and often require fine-tuning, ensembling, and other adaptation techniques. While the core challenge of identifying the best-performing model for a downstream task remains, the increasing heterogeneity of ML pipelines demands novel AutoML approaches. This work extends the CASH framework to select and adapt modern ML pipelines. We propose PS-PFN to efficiently explore and exploit adapting ML pipelines by extending Posterior Sampling (PS) to the max k-armed bandit problem setup. PS-PFN leverages prior-data fitted networks (PFNs) to efficiently estimate the posterior distribution of the maximal value via in-context learning. We show how to extend this method to consider varying costs of pulling arms and to use different PFNs to model reward distributions individually per arm. Experimental results on one novel and two existing standard benchmark tasks demonstrate the superior performance of PS-PFN compared to other bandit and AutoML strategies. We make our code and data available at https://github.com/amirbalef/CASHPlus.