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

TL;DR

This paper introduces PAK-UCB, an online learning algorithm for selecting the most suitable generative model for a given prompt, improving efficiency by adapting to prompt-specific model performance.

Contribution

It proposes a novel online learning framework with a kernel-based PAK-UCB algorithm and RFF acceleration for prompt-aware model selection in generative tasks.

Findings

RFF-UCB effectively identifies the best model for different prompts.

The approach reduces costs by avoiding sub-optimal model queries.

Experimental results demonstrate high accuracy in model selection across tasks.

Abstract

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.