In-Context Decision Making for Optimizing Complex AutoML Pipelines
This addresses the need for more efficient AutoML systems in heterogeneous ML pipelines, though it appears incremental as an extension of existing CASH and bandit methods.
The paper tackles the problem of optimizing complex AutoML pipelines that involve fine-tuning, ensembling, and adaptation techniques, extending the CASH framework to handle modern ML workflows. It proposes PS-PFN, a method using posterior sampling and prior-data fitted networks for in-context learning, and demonstrates superior performance on benchmark tasks compared to other approaches.
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.