# PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 51 references that are in the map.
- First posted on arXiv 17 Oct 2024.

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

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [BERTScore: Evaluating Text Generation with BERT](https://mlsanity.com/paper/1904.09675/) (2019)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [CLIPScore: A Reference-free Evaluation Metric for Image Captioning](https://mlsanity.com/paper/2104.08718/) (2021)
- [VBench: Comprehensive Benchmark Suite for Video Generative Models](https://mlsanity.com/paper/2311.17982/) (2023)
- [Demystifying MMD GANs](https://mlsanity.com/paper/1801.01401/) (2018)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2305.01569/) (2023)
- [BigCodeBench: Benchmarking Code Generation with Diverse Function Calls and Complex Instructions](https://mlsanity.com/paper/2406.15877/) (2024)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [Sora: A Review on Background, Technology, Limitations, and Opportunities of Large Vision Models](https://mlsanity.com/paper/2402.17177/) (2024)
- [CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X](https://mlsanity.com/paper/2303.17568/) (2023)
- [One Transformer Fits All Distributions in Multi-Modal Diffusion at Scale](https://mlsanity.com/paper/2303.06555/) (2023)
- [Finite-Time Analysis of Kernelised Contextual Bandits](https://mlsanity.com/paper/1309.6869/) (2013)
- [DreamBench++: A Human-Aligned Benchmark for Personalized Image Generation](https://mlsanity.com/paper/2406.16855/) (2024)
- [Practical Contextual Bandits with Regression Oracles](https://mlsanity.com/paper/1803.01088/) (2018)
- [On the Error of Random Fourier Features](https://mlsanity.com/paper/1506.02785/) (2015)
- [Random Fourier Features for Kernel Ridge Regression: Approximation Bounds and Statistical Guarantees](https://mlsanity.com/paper/1804.09893/) (2018)
- [Consistency-diversity-realism Pareto fronts of conditional image generative models](https://mlsanity.com/paper/2406.10429/) (2024)
- [Mutual Information Divergence: A Unified Metric for Multimodal Generative Models](https://mlsanity.com/paper/2205.13445/) (2022)
- [Towards a Scalable Reference-Free Evaluation of Generative Models](https://mlsanity.com/paper/2407.02961/) (2024)
- [EVALALIGN: Supervised Fine-Tuning Multimodal LLMs with Human-Aligned Data for Evaluating Text-to-Image Models](https://mlsanity.com/paper/2406.16562/) (2024)
- [Contextual-Bandit Based Personalized Recommendation with Time-Varying User Interests](https://mlsanity.com/paper/2003.00359/) (2020)
- [Stylus: Automatic Adapter Selection for Diffusion Models](https://mlsanity.com/paper/2404.18928/) (2024)
- [Near-linear Time Gaussian Process Optimization with Adaptive Batching and Resparsification](https://mlsanity.com/paper/2002.09954/) (2020)
- [A Multi-Armed Bandit Approach to Online Selection and Evaluation of Generative Models](https://mlsanity.com/paper/2406.07451/) (2024)
- [Analyzing Quality, Bias, and Performance in Text-to-Image Generative Models](https://mlsanity.com/paper/2407.00138/) (2024)
- [Stochastic Bandits for Egalitarian Assignment](https://mlsanity.com/paper/2410.05856/) (2024)

## Built on by (most cited first, 12 of 12)

- [SPARKE: Scalable Prompt-Aware Diversity and Novelty Guidance in Diffusion Models via RKE Score](https://mlsanity.com/paper/2506.10173/) (2025)
- [Do Vendi Scores Converge with Finite Samples? Truncated Vendi Score for Finite-Sample Convergence Guarantees](https://mlsanity.com/paper/2410.21719/) (2024)
- [Towards an Explainable Comparison and Alignment of Feature Embeddings](https://mlsanity.com/paper/2506.06231/) (2025)
- [Conditional Vendi Score: Prompt-Aware Diversity Evaluation for Generative AI Models and LLMs](https://mlsanity.com/paper/2411.02817/) (2024)
- [MMD Guidance: Training-Free Distribution Adaptation for Diffusion Models via Maximum Mean Discrepancy Guidance](https://mlsanity.com/paper/2601.08379/) (2026)
- [The Maximum von Neumann Entropy Principle: Theory and Applications in Machine Learning](https://mlsanity.com/paper/2602.02117/) (2026)
- [Efficient Adversarial Attacks on High-dimensional Offline Bandits](https://mlsanity.com/paper/2602.01658/) (2026)
- [DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative Models](https://mlsanity.com/paper/2603.23140/) (2026)
- [KODA: Contrastive Representation Comparison and Alignment for Vision-Language Foundation Models](https://mlsanity.com/paper/2606.04180/) (2026)
- [Mixture-Greedy for Online Generative Model Selection: Is UCB Necessary in Diversity-Aware Multi-Armed Bandits?](https://mlsanity.com/paper/2603.21716/) (2026)
- [Exploring More to Solve More: Boosting Diversity in Text Diffusion Models via Entropy-Based Guidance](https://mlsanity.com/paper/2608.00024/) (2026)
- [TasteRoute: Personalized Routing for Video Generation](https://mlsanity.com/paper/2610.05896/) (2026)

## Abstract (arXiv)

Selecting a sample generation scheme from multiple prompt-based generative models, including large language models (LLMs) and prompt-guided image and video generation models, is typically addressed by choosing the model that maximizes an averaged evaluation score. However, this score-based selection overlooks the possibility that different models achieve the best generation performance for different types of text prompts. An online identification of the best generation model for various input prompts can reduce the costs associated with querying sub-optimal models. In this work, we explore the possibility of varying rankings of text-based generative models for different text prompts and propose an online learning framework to predict the best data generation model for a given input prompt. The proposed PAK-UCB algorithm addresses a contextual bandit (CB) setting with shared context variables across the arms, utilizing the generated data to update kernel-based functions that predict the score of each model available for unseen text prompts. Additionally, we leverage random Fourier features (RFF) to accelerate the online learning process of PAK-UCB. Our numerical experiments on real and simulated text-to-image and image-to-text generative models show that RFF-UCB performs successfully in identifying the best generation model across different sample types. The code is available at: github.com/yannxiaoyanhu/dgm-online-select.