# In-Context Decision Making for Optimizing Complex AutoML Pipelines

arXiv: https://arxiv.org/abs/2508.13657 · HTML: https://mlsanity.com/paper/2508.13657/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 37 of its 62 references that are in the map.
- First posted on arXiv 19 Aug 2025.

## Builds on (most cited first, 30 of 37)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Learning to Optimize Via Posterior Sampling](https://mlsanity.com/paper/1301.2609/) (2013)
- [Auto-Sklearn 2.0: Hands-free AutoML via Meta-Learning](https://mlsanity.com/paper/2007.04074/) (2020)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [Large Language Models for Automated Data Science: Introducing CAAFE for Context-Aware Automated Feature Engineering](https://mlsanity.com/paper/2305.03403/) (2023)
- [FLAML: A Fast and Lightweight AutoML Library](https://mlsanity.com/paper/1911.04706/) (2019)
- [Better by Default: Strong Pre-Tuned MLPs and Boosted Trees on Tabular Data](https://mlsanity.com/paper/2407.04491/) (2024)
- [HEBO Pushing The Limits of Sample-Efficient Hyperparameter Optimisation](https://mlsanity.com/paper/2012.03826/) (2020)
- [Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining](https://mlsanity.com/paper/2310.08566/) (2023)
- [XTab: Cross-table Pretraining for Tabular Transformers](https://mlsanity.com/paper/2305.06090/) (2023)
- [PFNs4BO: In-Context Learning for Bayesian Optimization](https://mlsanity.com/paper/2305.17535/) (2023)
- [YAHPO Gym -- An Efficient Multi-Objective Multi-Fidelity Benchmark for Hyperparameter Optimization](https://mlsanity.com/paper/2109.03670/) (2021)
- [TuneTables: Context Optimization for Scalable Prior-Data Fitted Networks](https://mlsanity.com/paper/2402.11137/) (2024)
- [Cost-aware Bayesian Optimization](https://mlsanity.com/paper/2003.10870/) (2020)
- [Efficient Bayesian Learning Curve Extrapolation using Prior-Data Fitted Networks](https://mlsanity.com/paper/2310.20447/) (2023)
- [TabRepo: A Large Scale Repository of Tabular Model Evaluations and its AutoML Applications](https://mlsanity.com/paper/2311.02971/) (2023)
- [Efficient Automatic CASH via Rising Bandits](https://mlsanity.com/paper/2012.04371/) (2020)
- [Thompson Sampling for Budgeted Multi-armed Bandits](https://mlsanity.com/paper/1505.00146/) (2015)
- [Fine-tuned In-Context Learning Transformers are Excellent Tabular Data Classifiers](https://mlsanity.com/paper/2405.13396/) (2024)
- [Non-Stationary Bandit Learning via Predictive Sampling](https://mlsanity.com/paper/2205.01970/) (2022)
- [Quick-Tune: Quickly Learning Which Pretrained Model to Finetune and How](https://mlsanity.com/paper/2306.03828/) (2023)
- [Cost-aware Bayesian Optimization via the Pandora's Box Gittins Index](https://mlsanity.com/paper/2406.20062/) (2024)
- [Tabular Data Augmentation for Machine Learning: Progress and Prospects of Embracing Generative AI](https://mlsanity.com/paper/2407.21523/) (2024)

## Abstract (arXiv)

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.