PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Authors: Xiaoyan Hu
Ho-fung Leung
Farzan Farnia
Venue: ICML 2025 poster (2025)
Keywords: Online Learning
Large Language Models
Evaluation and Selection of Generative Models
Contextual Bandits
Primary area: reinforcement_learning
TL;DR: We propose an online learning method to adaptively select the best generative model or LLM for each input prompt.

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](github.com/yannxiaoyanhu/dgm-online-select).

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Xiaoyan Hu'
Ho-fung Leung² Farzan Farnia
Abstract
1. Introduction
Large language models (LLMs) and text-guided generative Al models have found numerous applications in various engineering tasks. A prompt-guided generative Al repre- sents a conditional generative model that produces output samples given an input text prompt. Over the past few years, several frameworks have been developed to train generative models and perform text-guided sample generation in var- ious domains. including text, image, and video (Touvron et al., 2023; Gemini-Tean et al., 2025; OpenAI et al., 2024; DeepSeek-AL. et al., 2025; Rombach et al., 2022; Bao et al., 2023b; Podell et al., 2024; Singer et al., 2022; Bai et al., 2023; Liu et al., 2024). The great number of available LLMs and soompt-gudded generative models has led to significant interest in developing evaluation mechanisms to rank the ex- isting models and find the best model from a group of avail- able ones. To address this task, several evaluation metrics have been proposed to quantify the fidelity of samples gen- erated by prompt- -bsed generative models, such as CLIP. Score (Hessel et al., 2021) and PickScore (Kirstain et al., 2023) for prompt-guided image and video generation, and BLEU score (Papineni et al., 2002) and BERTscore (Zhang et al., 2020) for LLMs.
Selecting sample generation scheme from mul- tiple prompt-based generative models, includ- ing large language models (LLMs) and prompt- guided image and video generation models, is typ- ically addressed by choosing the model that maxi- mizes 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 mod- els. In this work, we explore the possibility of varying rankings of text-based generative mod- els for different text prompts and propose an on- line learning framework to predict the best data generation model for a given input prompt. The proposed PAK-UCB algorithm addresses con- textual bandit (CB) setting with shared context variables across the arms, utilizing the generated data to update hernel-based functions that pre- dict the score of each model available for unseen text prompts. Additionally, we leverage random Fourier features (RFF) 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/ yannxi maanxiaoyannhu/eao doppp///gmmonno ne- ect.
The existing model selection methodologies commonly aim to identify the generative model with the highest averaged fidelity score over the distribution of text prompts, produc- ing samples that, on average, align the most with input text prompts. A well-known example is the averaged CLIPScore for text-to-image models, measuring the expected alignment between the input text and output image of the model using the CLIP embedding (Radford et al., 2021b). While the best-averaged score model selection strategy has been fre- quently utilized in generative Al applications, this approach does not consider the possibility that the involved models can perform differently across text prompts.
However, we highlight the possibility that one model out- performs another model responding text prompts from certain categories, while that model performs suboptimally in responding to prompts from other categories. Figure 1 shows examples of both L.LMs and text-to-image models, where two widely-used models exhibit different rankings across the text prompts from different categories (code gen-
'Department of Computer Science and Engineering, The Chi- nese University of Hong Kong 'Independent Researcher. Corre- spondence to: Xiaoyan Hu <xyhu21 www..cch..coo edu.hk>, Farzan Farnia <famia wwweeeeweeeeoooo.eeo www.@eeeee..cceeeeee cuhk.edu.hk>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Example on
Example on LLMs
Stable Diffusion v1.5
PixArt-a
Gemini-2.55Flasst
Qwen-Plus
T2I models
Instroot
Instruct
Prompts of Type "dog"
Code Completion
Pass rate (%)
92.32 (+0.43)
86.46 (+0.7
avg. CLIPScore
36.37 (=0.13)
37.24 (=0.09)
-Jav
Sava
Code Translation
Prompts of
Type "car"
Mrongi
Cocrwct!
Pass (%)
48.78 (+0.52)
82.26
(=1.95)
avg. CLIPScore
36.10 (+0.06) 35.68 (+0.15)
(a). Example 1: Gemini-2.5-Fii attains higher pass rate (Python) code completion on the HumanEval benchmark (92.3% versus 86.5%) while underperforms for Java-to-C+ translation on the HumanEval-X benchmark 44.88 versus 82.3%).
(b). Example 2: Stable Diffusion v1.5 attains higher CLIPScore in generating MS-COCO prompts with term "car" (36. 10 versus 35.68) while underperforms for MS-COCO prompts with term "dog" (36.37 versus 37.24).
Figure 1. Widely-used LLMs and text-to-image models could exhibit different rankings across the text prompts with different categories.
eration and translation tasks for LLMs, and MS-COCO prompts with terms "dog""'car"). In general, the different training sets and model architectures of L.MSs and prompt- guided models can result in the models' varying perfor- mance in response to different text prompts, which is an important consideration in assigning prompts of various types to a group of prompt-input generative models.'
We highlight that the described online learning task can be viewed contextual bandit (CB) problem widely stud- ied in the bandit literature (Langford & Zhang, 2007; Li et al., 2010; Chu et al., 2011; Agrawal & Goyal, 2013). In the CB task, the online learner observes a single context variable (the text prompt in our setting) and predicts the best arm for the current input context. Specifically, we focus the kernel-based method to predict the score of each available model for incoming prompt. As the text prompt is shared across the models, the CB task may be suboptimally addressed by the common-weight formulation of linear CB (Chu et al., 201 1) and kernelized CB (Valko et al., 2013) algorithms, as these implementations of the algorithms learn a single shared weight vector to predict the reward functions for all the arms (i.e., generative models in our setting). To relax this constraint, we propose Per- Arm Kernelized UCB (PAK-UCB) to address the online prompt-basee based selection of generative models. According to the PAK-UCB approach, the learner utilizes the computed UCB-scores of arm-specific kernel- based prediction func- tions to choose the generative model for the incoming text prompt and subsequently update the kernel-based prediction rule based the generated data for the future rounds. We attempt to theoretically analyze the proposed PAK- and show that its variant could achieve regret bound of Ö(vT) over T rounds.
In this work, we aim develop a learning algorithm to identify the best generative model for given input prompt, using observed prompt/generated samples collected from the models in the previous sample generation rounds (Fig- ure 2). Since the goal of prompt-based model selection is to avoid sample generation queries from suboptimal generative models for given text prompt, we view the model selec- tion task as an online learning problem, where after each data generation, the learner updates the prediction on which generative model performs the best in response to input text prompts. Here, the objective of the online learner is to uti- lize the previously generated samples to accurately guess the generation model with the best performance for the incom- ing text prompt. An optimal online model selection method will result in bounded regret value, measured in compar- ison to the sample generation from the groundtruth-ber model for the text prompts.
'The full models' response codes in the figure's experiment are provided at: https: / / //ggtttubbb com/ yannxi waaanniiiayannuuu cgm-onninee-nn select.
Since the user applying the CB-based model selection ap-
2

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Candidate Models
Single Model Selection
Offline Search for the Model Maximizing the Averaged CLIPScore
PixArt
PixArt
PixArt- or
PixArt- a
PixArt-
(Conventional approach)
39.13
31.93
31.33
43.18
31.24
CLiPSSCre
35.36
Stable Diffusion v1.5
Contextual Bandit- Based Selector
(Our proposed approach)
PixArt-a
PixArt-a
415
Prompt-based
(PAK-UCB)
PixArt-a
Model Selection
39.13
39.00
36.90
43.18
42.34
Avg.CLiPScore: 40.11
Figure 2. Selection of text-to-image (T2I) models based on the CLIPScore: The conventional best-score model selection assigns all prompts the single model with the best averaged CLIPScore. In contrast, our proposed online approach performs prompt-level model selection by leveraging the UCB contextual bandit algorithm to assign the model to the prompt, learning from previous data generations.
proach may have limited compute power and not be able to afford growing computational costs in the online learning process, we propose applying the random Fourier features (RFF) framework (Rahimi & Recht, 2007b) to reduce the computational load of PAK-UCB. We discuss that in the pro- posed PAK-UCB, the computational cost per iteration will grow cubically (tt²) with iteration f. We demonstrate that after leveraging the RFF method, our developed proxy RFF-UCB algorithm can effectively approximate the solu- tion to PAK-UCB with the computational costs increasing only linearly (t).
Developing the contextual bandit algorithms of PAK- UCB and RFF-UCB algorithms for the online selection of prompt-based generative models.
Theoretically analyzing the regret and computational costs of PAK-UCB and RFF-UCB online learning methods
Presenting numerical results on the prompt-based selec- tion of generative models using PAK-UCB and RFF-UCB.
2. Related Work
Finally, we present the results several numerical experi- ments show the efficacy of our proposed PAK-UCB and RFF-UCB in the online selection of prompt-guided genera- tive models and LLMs. In our experiments, we test several pre-trained and simulated text-to-image, image-captioning (image-to-text), and language models, where different mod- els could lead to different rankings across various sample types. Our numerical results suggest fast convergence of the proposed online learning algorithms the best prompt- based selection of the available models in response dif- ferent prompt types. In experiments, the proposed PAK- UCB and RFF-UCB could perform better than several base- line methods, including the common-weight implementation of LinUCB (Chu et al., 201 1) and KernelUBB (Valko et al., 2013) methods. The following a summary of this work's ss contributions:
(Automatic) Evaluation of conditional generative mod- els. Evaluating the conditional generative models has been studied extensively in the literature. For text-to-image (T2I) generation, carlier methods primarily rely on the Incep- tion score (Salimans et al., 2016) and Fréchet inception distance (Heusel al., 2017). More recent works propose reference-free metrics for robust automatic evaluation of T21 and image captioning, with notable examples being CLIP- Score (Hessel et al., 2021) and PickScore (Kirstain et al., 2023). Kim et al. (2022) propose intuaal-iforrmation based metric, which attains consistency across benchmarks, sample parsimony, and robustness. To provide a holistic evaluation of T2I models, several works focus on multi- objective evaluation. Astolfi al. (2024) propose to evalu- ate conditional image generation in terms of prompt-sample consistency, sample diversity, and fidelity. Kannen et al. (2024) introduce a framework to evaluate T2I models re- garding cultural awareness and cultural diversity. Mas- rourisaadat et al. (2024) examine the performance of several T2I models in generating images such human faces and
Studying the prompt-based selection of prompt-guide generative models to improve the prompt-level perfor- mance scores,
3

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
groups and present a social bias analysis. Another line of study explores evaluation approaches using large language models (LLMs). Tan et al. (2024) develop LLM-based evaluation protocols that focus on the faithfulness and text- image alignment. Peng et al. (2024) introduce a GPT-based benchmark for evaluating personalized image generation. In addition, a line of study leverages human feedback for scoring/ranking generated images and improving T2I mod- els (Xu et al., 2023). For evaluation of text-to-video (T2V) generation, Huang et al. (2024) introduce VBench as com- prehensive evaluation of T2V models in terms of quality and consistency. Finally, we note the prompt-aware diversity scores, Conditional- Vendi (Jalali et al., 2023) and Schur- (Ospanov et al., 2024a), developed for prompt-guided guided generative models.
tion of image-output generative models based on the FID score (Heusel et al., 2017). Rezaei et al. (2025) show that mixture of several unconditional generative models can attain higher diversity scores than each individual model, and propose the Mixture-UCB MAB framework to find the best mixture of the generative models in terms of the kernel distance (Bińkowski et al., 2018) and Renyi kernel entropy (Jalali et al., 2023; Ospanov et al., 2024b). We note that the mentioned papers study the selection of unconditional generative models, different from our problem setting of selecting the conditional prompt-guided generative models.
3. Preliminaries
3.1. Kernel Methods and Random Fourier Features
Let Φ : Rd - H be a feature map from the ambient space R2d to the (possibly infiniteddimensional) space cor- responding to the associated reproducing kernel Hilbert space (RKHS) H. The kernel function for RKHS H is defined as the inner product of the representations of the inputs as k(y, v/) == (o(v), ,(v')) = <(v) TTv')) for every 3/ E R2, where (,, .) denotes the standard inner product.
(Kernelized) Contextual bandits. The contextual ban- dits (CB) is an efficient framework for online decision- making with side information (Langford & Zhang, 2007: Foster et al., 2018), which is widely adopted in domains such as recommendation system and online advertisement (Li et al., 2010) and resource allocations (Lim et al., 2024). A key to its formulation is the relationship between the con- text (vector) and the expected reward. In linear CB, the reward is assumed to be linear to the context vector (Li et al., 2010; Chu et al., 2011). To incorporate non-linearity, Valko et al. (2013) propose kernelized CB, which assumes the rewards are realaalee in a reproducing kernel Hilbert space (RKHS). To address the growing computation, recent work leverages the assumption that the kernel matrix is often approximately low-rank and uses Nyström approxi- mations (Calandriello et al., 2019; 2020; Zenati et al., 2022).
Every kernel function k will satisfy the positive-defimitnnee condition, that means Σx= Σ}= Geyk(Us Vi) 2 0 holds for every integer n E N+. Vi, , Ur E Rdd and C1, Cn E R. Furthermore, we call a kernel function A shift invariant if k(y. s/) == k(y- yy) (lyy-') A well-known example the radial basis function (RBF) kernel with bandwidth parameter o > 0:
for every y. v E R
1l - 211
KRBF(Y, 3/) = exp
20
Prompt-based Selection of LLMs and Generative Mod- els. We note that a line of study proposes offline methods for learning a universal prompt-to-mont model assignment rule by training a neural network model over a dataset of ranked responses of the models to large set of prompts (Luo et al., 2024; Qin et al., 2024; Frick et al., 2025). Unlike these methods, our proposed approach follows an online learning strategy which can significantly lower the computational and statistical costs of finding a satisfactory assignment rule. In addition, as we later discuss in the numerical results, the online nature of our proposed algorithm can improve the adaptability of the model to the changes in the user's prompt distribution and models' behavior over the iterations of evaluating the models in real time.
Kernel ridge regression (KRR). Given empirical labeled samples (11, si),..., (Na, 8n), where {yi E Rd are feature variables and {s; E [R] are target variables, respectively, the kernel ridge regression assumes that for a w E H we have E[s; [vs] = <dv)) T "n for borany =1 = =....... n..
For regularization parameter a 2 0, KRR estimator
SKRR (3) = kn (K + ala)
for every y E Rd, where K = [k(2), y3 V)) denotes the kernel matrix, V := [s1, [s,,... , sn]] hy = [k( V), ... k(yn, v)] E R". The KRR estimator can be interpreted as the solution w* to the ridge regression:
E R
R". and
Multi-armed bandit selection of generative models. We note that the references (Hu et al., 2025a; Rezaei et al., 2025) propose multi multi-armed bandit (MAB) approaches to the selection of unconditional generative models without an input text prompt. Specifically, Hu et al. (2025a) propose an online learning framework for selecting generative models and develop the FID-UCB algorithm for the online selec-
W = arg min E(((vi)T -w- - si)² + +all a0 we H.
1=1
where || := Vw Touu w for any E H. and then making the prediction SKRR = (d())) @.
Random Fourier features (RFF). To reduce the compu- tational costs of kernel methods for shift-invariant kernels,
4

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Rahimi & Recht (2007a) proposes the random Fourier fea- tures framework. Specifically, Bochner's Theorem (Rudin, 2017) shows that for every shift-invariant kernel k s/) = k(y-y') satisfying x(0) = 1. the Fourier transform of k pro- vides a valid probability density function, which we denote by p E A(RR), such that k(y.y') = Ew-ple" [empll"""" (v-x'), with being the imaginary unit. Following this property, the RFF approach independently samples w1,'". , DDD ~ P and proposes the empirical mean proxy for k(y. y/) as Di Ee=1 etu] (v-v"), Due to the real Fourier transform of the even function K, we can simplify the equation as
4. Prompt-Based Model Selection as a Contextual Bandit Problem
In this section, we introduce the framework of online prompt-based selection of generative models, which is given in Protocol 1. Let [N] = (I...., ,N for any positive integer N E N We denote by g == [G] the set of (prompt-based) generative models. The evaluation proceeds in T N+ iterations.
1
At each iteration & [7]. prompt ye E 3 is drawn from fixed distribution p E ())) the prompt space y C gd-1, e.g., (the normalized embedding of) picture in image cap- tioning or a paragraph in rext-to-imaeeviideo generation. Based on prompt yi (and previous observation sequence), an algorithm A picks model gr E 9 and samples rexponse Iy ~ Pa. (-1yc), where P, () E (X is the conditional distribution of the generated response from any model g E G. The quality of response 2, is evaluated using a score func- tion y x X' - [-1, 1]. which assigns scores(y), The algorithm A aims to minimize the regret
1 ((v)= = VD
[ws (u)Ty y). sin(w] v), cos(wDy) sin sin(wDD))
(1)
where fuyl Lid and 3') is approximated by the inner product ((y) "vy).. The resulting approximate KRR estimator is
(y) == ( t + al2D) "'''
where T- == (p(v)) E RTx2D and we denote by
T
Regret = Σ (s.(v)) - Sg. (2t)),
hht its conjugate transpose, can be computed in (nD²)
(3)
time and (nD) memory, giving substantial computational savings if D < n. For the RBF kernel k (#1, 222) =
where we denote by sy()) expected score of any model g E G and s,) := maxgeg sy(v) the optimal expected score, both conditioned on the prompt y.
(enr,,,(w xg)ه the
exp(( _ |1 x2 22213), the PDF PRBF will be the multivari- ate Gaussian (0, 0² Ia).²
3.2. CLIPScore for Evaluating Text-to-Image Models
Protocol Online Prompt-based Selection of LLMs and Prompt-guided Generative Models
CLIPScore (Hessel et al., 2021) has been widely-used to evaluate the alignment of samples generated by text-to- image/video (T2I/V) and image captioning models. Let (v,z) E y x X be text-image pair. We denote by Gy E gd-1 = {= E Rd ||ı||2 = 1} and ny E sd-1 the (normalized) embeddings of text y E y and image X E X, respectively, both extracted by CLIP (Radford et al., 2021a).
Require: total iterations T E N4, set of generators G [G]. prompt distribution A()), score function
y x + [-1, 1], algorithm A : (Y' x G x R)* x y +
Agg)
Initialize: observation sequence D +- 2 , ,T do Prompt y/ ~ p is revealed. 3: Algorithm .A picks model ge ~ A(-ID, y1) and sam-
The CLIPScore (Hessel et al., 2021) is given by
1: for iteration = 1, 1..,....
CLIPSSore T21 (y, x² == max (0, 100 cos(U G,),, (2)
ples response ,, ~ Pg. (-1vt).
where cos(ez, Cgy) = (U., c2) for the Ey-normalized CLIP-
4: Score + s(81, zx) is assigned.
embedded Ug Cy. Extending the definition to (text, video)
5: Update observation D t- D { St, (11,,tt,,))) 8t)).
pairs, for a video X == (z()) / consisting of L frames, where (!!) is the (-th frame, the score is the averaged frame- level CLIPScore:
6:
end for
5. An Optimism-based Approach for Prompt-based Selection
CLIPScore T2V (y. = CLIPScore T21 (2l x ()).
1 L 1=1
Under the online prompt-based selection setting, a key chal- lenge is to learn the relationship between the prompt and the score of each model. In this paper, we consider kernel methods for score prediction. Specifically, for each model g E G. we assume the existence of (possibly non- non-linear and
We note that an alternative form of RFF is given by: ppv) V 3 [cas wIT (y+b)).... coo((wb++b))
where (w,) p
p and [b:).p j lid Unif([O, 2m)). For the RBF kernel, it can be shown that (1) has uniformly lower variance (Sutherland & Schneider, 2015).
feature map Φ and a weight w
5

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
in an RKHS, such that the score sy(y) conditioned on the
Algorithm 2 Per-Arm Kernelized UCB (PAK-UCB)
prompt y is given by (ø(y)) Twg Notably, the weight vector n is arm-specific and can vary across the models, which is
Require: iteration T, generators g [G] prompt distribu-
tion p. score function y x X + [-1,1], positive definite kernel k, regularization and exploration param-
stated in the following assumption.
Assumption 1 (Realizabilitt), There exists a mapping o: R + H. and weight wn E H. such that score sg (V) = ky. w u for any prompt vector y R and model g G. Further, it holds that |!w6g| I 1, and k(y, v) S x² and
eters 7 2 0
Initialize: observation sequence D + Ø and index set
V. k Ø for all g E
1: for iteration = 1,2, , do
||(())|| S 1 for any y E Vi where k Rd x Rd + R. is the kernel function of the mapping db.
Prompt yt ~ p is revealed.
3: Compute (f.g. 0,)
COMPUTT__CCB(2) Vc: V/)
and set $ + Ml +jōg for each g G.
Remark (Comparison to linear and kernelized CB). We note that Assumption I differs from the standard settings of linear and kernelized contextual bandit (CB) (Chu et al., 2011; Valko et al., 2013; Zenati et al., 2022). Specifically, in the settings of these references, a potentially different context variable is considered for each arm, and the same weight vector is applied across all arms. In the prompt- based model selection setting, the text prompt (i.e., the context variable) is shared among arms (i.e., the gener- ative models) and can result in different performance scores across the models, which cannot be captured by the standard kernelized CB approach. On the other hand, we note that the existing arm-specific bandit algorithms in (Li et al., 2010; Xu et al., 2020) are designed considering linear pay-off structure, which is generally different from the kernel-based formulation in our work.
4: Pick model gt. <- arg max (8g)-
5: Sample 2t- ~ Pg. (-lea) and receive score St.
6: Update D + DU {(1,, $t)} and V. + V U {t).
7:
end for
8:
function COMPUTE. CUTITEEUUBB.B.. y, Vo)
9:
if 1/ y is empty then
10:
Big +- +00,0g + +00.
11:
else
12:
Set K + [K(v, Wilkk,,,le,, + [iilew, and
ky
Ik My..i))
13:
my + ky (K + +al) '.
3 + a : /klu.v)) v) kj (K + at) -ky-
14:
15:
end if
16:
return (ig, o₂).
17:
end function
5.1. The PAK-UCB Algorithm
In this section, we present PAK-UCB in Algorithm 2, an online learning approach to prompt-based model selection. For each incoming prompt, PAK-UCB first estimates the performance scores via kernel ridge regression (KRR) and then picks the model with the highest estimated score. Un- like conventional formulations of LinUCB (Chu et al., 2011) and KerneUUCB (Valko et al., 2013) algorithms, which learn a single reward function shared across all arms, PAK-UCB learns arm-specific functions to predict the score of each model.
Theorem 1 (Regret, informal). With probability of at least 8, the regret of running a variant of PAK-UCB (stated in Algorithmm 3 in the Appendix) for T iterations is bounded by Ö( O(VGT).
5.2. PAK-UCB with Random Fourier Features
The PAK-UCB solves a KRR for each model at an itera- tion to estimate the scores, which can be expensive in both computation and memory for large number of iterations. To address this problem, we leverage the random Fourier features (RFF) sampling (Rahimi & Recht, 2007a) for shift- invariant kernel functions, e.g., the RBF (Gaussian) kernel. At a high level, RFF maps the input data, e.g., the prompt (vector) in our setting, to a randomized low-dimensional feature space and then learns a linear model the result- ing random space by solving a standard linear regression problem. Note that following the design of RFF, the inner product of the projected randomized features will be an unbiased estimation of the original kernel inner product.
The key component in PAK-UCB is the function COMPUTE.UCB (lines 8-17), which outputs both the KRR estimator P. and an uncertainty quantifier o As the weight vector w can vary across the arms, PAK-UCB constructs the KRR dataset using prompt-score pairs from iterations where model g is chosen, with the corresponding indices stored in the set My (line 6). The estimated score is then computed by S, = Fh + nō, (line 4), which is initially set to +0o to ensure each model is picked at least once (lines 9-) To provide theoretical justification for our method, we show that a variant of PAK-UCB, which is illustrated in Algorithm 3 in the Appendix, can attain a squared-root regret bound. The formal statement and the proof can be found in Appendix A.
We present RFF-UCB as an RFF implementation of PAK- UCB (Algorithm in the Appendix). Particularly, RFF. UCB leverages an baF-based approach to compute the mean
6

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
0.70
0.6
0.4
0.65
0.2
PAK-UCB-paly3
0.60
PAK-UCB p3ly-
RFF-UCB-RBF
RFF-UCB-ABF
0.0
KemelUCB poly3
0.55
KernelUCB paly3
One-arm Oracle
One-arm Oracle
-0.
Random
0.50
Random
0
1000
2000
3000
4000
5000
1000
2000 3000
4000
5000
(a) Outscore-the best (OtB)
(b) optimal-pick-ane (OPR)
Figure 3. Prompt-based selection between Stable Diffusion v1.5 and PixArt-a (Figure (a)): Results are averaged over 20 trials.
and uncertainty quantifier, which is stated in Algorithm 5 in the Appendix. Upon receiving the regression dataset con- sisting of prompt-score pairs, COMPUTE_UCB BFFFF first projects each d-dimensional prompt vector to a randomized 2.D-dimensional feature space according to Equation (1) and then solves a linear ridge regression estimate the mean and uncertainty. To see why RFF reduce the computation, note that the size of the (regularized) Gram matrix (0, g 9, al2p) in Line 8 is fixed to be 2D in the whole process, while the size of (K + al) in line 13 of Algorithm 2 scales with |Vyl and can grow linearly over iterations. Particularly, the following lemma shows that COMPUTE, UCB RFF can reduce the time and space by an order of O(e) and (t)), respectively.
prompt-based selection among standard text-to-image (T2I) models: Stable Diffusion v1.5², FixAAtto-XXL-2-51xx55 UniDiffuser", and DeepFloyd IF-1-M-v1.0.. For the LLM experiments, we provide numerical results for prompt-based selection of the following large language models (LLMs): Gemini-2.5 2.5-Flash-ppeeee Qwen-Plus. In the Appendix, report additional results image captioning (image-to-text) task and video data under synthetic setups. Additional results can be found in Appendix C.
03-mini, Deepseek- Chepseek-Chat, and
We compare PAK-UCB-poly3 and RFF-UCB- RBF with three baselines, including 1) KernelUCB-polyy: standard KemnlUUB (Valko et al., 2013) using the klo kernel, 2) One-arm Oracle: always picking the model with the maximum averaged score, and 3) Random: selecting an model uniformly randomly. In the Appendix, we re- port results for three additional baselines: 4) PAK-UCB-lin: PAK-UCB with linear kernel, i.e., Amm (21, = z 22, which does not incorporate non-linearity in score estima- tion, 5) LinUCB: standard LinUCB (Chu et al., 2011). and 6) Naive-KRR: PAK-UCB-polyy without exploration, i.e.. parameter O, which selects the model with the highest estimated mean conditioned to the prompt.
Baselines.
Lemma 1 (Time and space complexity). At any iteration t E [T). COMPUTE_UCE UCB (Lines 8-17 of Algorithm 2) re- quires Ot /G²) time and O(t² /G) space, while CoM- PUTE UCB RFF with random features of size N (AI- gorithm 5) requires O(tD²) time and Ott)) space, where G is the number of generators. See Appendix B.1 for details.
To provide theoretical justification to the RFF approach, we show that the implementation of PAK-UCT with RFF can attain the exact same regret bound with a carefully chosen feature sizes. The formal statement and the proof can be found in Appendix B.1.
Performance metries. For each experiment, we report two performance metrics: (i) outscore- thereessheeet (OtB): the difference between the CLIPScore attained by the algorithm and the highest average CLIPScore attained by any single model, and (ii) optimal-picc--auie (OPR): the overall ratio that the algorithm picks the best generator conditioned to the prompt. Specifically, to determine the best model for each specific prompt, we generate five responses from every model (to that prompt), and the best model in the OPR calculation the model with the highest averaged score
Theorem 2 (Regret of RFF implementation, informal). With probability at least - ô, the regret of running variant of RFF-UCB (stated in Algorithm 6) for T iterations is bounded by O(VGT).
6. Numerical Results
In this section, we present numerical results for the proposed 1) PAK-UCB UCB--pll3: PAK-UCB using a polynomial kemnel
hattes::/hoggingfeeecc///////// en http://pppulnesetttaaeeeeeeaee/s "ttpper/hagginggaee.c///aaaaaa
with degree 3, i.e., 1-7 a ( 222 (1 + y ar )²23 and 2) RFF-UCB-RBER PAK-UCB with RFF implementation of
the RBF kernel, i.e., Agg [ x2)= ) = cxp( - ||ا -22
Implementation details and the choice for hyperparameters
can be found in Appendix C.4. Our primary focus on

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
0.25
0.6
0.8
0.00
0.4
0.6
-0.25
0.2
0.4
PAK-UUB-ppR poly3
PAK-UCB poly3
UAK- poly3 RFF-UCB- RBF
0.0
RFF- OFF-UCB-RBS RBF
0.2
RFF-UCB-RB RBF
0.50
KernelUCB-palyy
KernelUCB-poll Banely
KernelUBB-pely: One-arm Oracle Random
-0.2
Om Oracle
0.0
One-arm Oracle Random
-075
Random
-0.4
-02
-1.00
0
1000
2000
3000
4000
5000
1000
2000 3000
4000
5000
1000
2000
3000
4000 5000
(a) UniDiffuser and Prx-A (Figure 9).
(b) Stable Diffusion v1.5 and UniDif- fuser (Figure 11)
(c) Stable Diffusion v1.5, Pix.Art-a, and DeepFloyd
Figure 4. Prompt-based selection among several standard T2I models: Outscore-the-best (OtB) is reported. The selection of our proposed PAK-UCB-pon algorithm converges quickly to the optimal model for incoming text prompts. Full results can be found in Figures 10, 12, and 13 in the Appendix. Results are averaged over 20 trials.
over the 5 runs. Then, OPR is defined as the fraction that the algorithm's 's selected model matches the best model.
"elephant' /" f'ieehrrnnt (Figure 4(b)), and "carrot'/'bww' (Figure 4(c)). The input of the PAK-UCB, Naive-KRR, and RFF-UCB-RBE methods is the embedded prompt that is output by the pretrained model from the open_elip repository". For LinUCB and Ker- neIUCB baselines, we also concatenate the one-hot encoded vector of the model index to the CLIP-embedded prompt. Full results can be found in Figure 8 in the Appendix.
6.1. Summary of the Numerical Results
D---ViT-B-32-laioni
The results of our numerical experiments indicate the im- provement of the proposed PAK-UCB algorithm over the one-arm oracle baseline that is aware of the single gen- erative model with the maximized averaged scores over the prompt distribution. This result suggests that the online learning algorithm could outperform a selector with the side- knowledge of the single best-performing model. We note that this improvement follows the prompt-based selection design of our proposed PAK-UCB algorithm. Moreover, our numerical results indicate that PAK-UCB can signif- icantly improve upon the online learning baselines with shared weights, including LinUCB (Chu et al., 2011) and KernelUCB (Valko et al., 2013) methods, indicating the effectiveness of the arm-specific design of PAK-UCB in the prompt-based model selection setting. Finally, in our experiments, the proposed RFF-UCB-RBF variant could reduce the computational costs of the general PAK-UCB algorithm.
Setup 2. Prompt-based selection of LLMs. In these ex- periments, we focus on prompt-based based selection of large lan- guage models (LLMs) on various tasks, including Sudoku- solving and code generation (code completion and trans- lation). In the first setup, we select LLMs to solve code generation tasks. The model given a Python code com- pletion problem (sampled from the first 164 tasks in Big- CodeBench (Zhuo al., 2024)) or C++-to-Jva code trans- lation problem (sampled from the HumanFxal-X bench- mark (Zheng et al., 2023)). Sample prompts can be found in Figure 18 and 19 in the Appendix. Particularly, the Gemini- 2.5 ---FFa---preiiew model attains a higher pass @ rate on code completion (55.91% % versus 13.84%) while underper- forms for C++-to-Java translation (24.39% versus 40.91% %) compared to the Qwen-Plus model. The rewards are bi- nary, indicating whether the generated code can pass all the test cases. The input of the PAK-UCB, Naive- and RFF-UCB-RBE methods is the 768-dimensional embedded prompt that is output by the RoBERTa model (for code com- pletion) or the CodeBERT model (for code translation). For LinUCB and KernelUCB baselines, we also concate- nate the one-hot encoded vector of the model index the embedded prompt. The results (Figures 6(a) and 20) show that the proposed RFF-UCB- BBF method queries from the better model conditioned to the task category (Figure 5),
6.2. Experiment Settings
Setup 1. Prompt-based selection among standard text- to-image models. The first set of experiments are on the setup illustrated in Figure 1(a), where we generate images from Stable Diffusion v1.3 and PixArt-a. The results show that PAK-UCB-poly3 outperforms the baseline algorithm and attains a high optimal-pick-ration which shows that it can identify the optimal model conditioned the prompt (see Figure 3). Additionally, we provide numerical results various T21 generative models, including UniDiffuser and DeepFloyd (see Figure 4). Prompts are uniformly randomly selected from the MS-COCO dataset under two categories: 'dog'' "oo'/'car' (Figure 3), 'train') baseeball-bat (Figure 4(a)),
codebert-base
8

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
0.75
0.70
0.75
Pek-UCB-poly3 RFF-UCB-RBF
0.65
RFF-UCB-ABF
KernelUCB pely3
0.60
KernelUCB paly
One-arm Oracle
One-arm Oracle
0.55
Random
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
(a) Code generation: Claude- 3.5- and Gemini 225 2.5-Flash- ---llash-rrevieww
(b) Sudoku Deepseek-Chie and o3- mini
Figure 5. Prompt-basee selection of LLMs (Setup 2): Optimal-pick-ratte (OPR) is reported.
Testruct
UCB-RBF algorithm outperforms all the baseline methods.
Setup 3. Adaptation to new prompt types and models. We consider scenarios where new text-to-image models or prompt types are introduced after the initial deployment. At the beginning of the first experiment, Stable Diffusion v1.5 and PixArt-ar are available, and UniDiffuser introduced after 2500 iterations. Prompts are uniformly randomly se- lected from the MS-COCO dataset under categories 'train' and "baseball-bat' In the second experiment, generate samples from both Art-rt and UniDiffuser. In the first Ik iterations, the prompts are uniformly randomly selected from a pool that initially includes categories 'person' and 'bicycle' in the MS-COCO dataset. Then, categories 'air- plane', 'bus", "train', and 'truck' are added the pool after each Ik iterations. The results show that PAK-UCB-poly3 well adapt to new prompt types and generators (see Figure 14). The input context follows Setup 1.
PY
Contextual Based Selector (RFF-UCB B- RBF)
Code Completion
Genini25
Java
Code Translation
Figure 6. Prompt-based selection of LLMs on code generation task (Setup 2): The Gemini-2.5-4 2.a.-2.55--aaa--eeiiiee model attains higher pass rate on the code completion task while underperforms on the code translation task. Our proposed SFF-UCB-RRB algorithm queries solutions from the hetter model conditioned to the task category.
Setup 4. Synthetic experiments on other conditional generation tasks. We provide numerical results on image- captioning (image-to-text) and text-to-video (T2V) task un- der synthetic setups in Appendix C.3.
In the second setup, we utilize two LLMs, including o3- mini and Deepseek-Chat, to solve 9 9 Sudoku problems, which contains 5-30 blank entries. Sample prompt can be found in Figure 16 the Appendix. The 03-mini model at- tains a higher success rate on harder Sudoku problems with more than 15 blank entries (70.82% versus 50.42%), while slightly underperforms easier Sudoku problems (86.85% versus 93.6%). The rewards are binary, indicating the solu- tion is correct or wrong. The input of the PAK-UCB, Naive- KRR, and RFF-UCB-RBF methods is the 768-dimensional embedded prompt that is output by the RoBERTa model". For LinUCB and KernelUCB baselines, we also concate- nate the one-hot encoded vector of the model index to the RoBERTa--mbedden prompt. The results summarized in Figures 5(b) and 17, which show that the proposed RFF-
7. Conclusion
In this work, investigated prompt-based selection of gen- erative models using contextual bandit algorithm, which can identify the best available generative model for a given text prompt. We propose two kernel-based algorithms, in- cluding PAK-UCB and RFF-UCB, to perform this selec- tion task. Our numerical results on LLMs (text-to-text), text-to-image, text-to-video, and image-captioning tasks demonstrate the effectiveness of the proposed framework in scenarios where the available generative models have vary- ing performance rankings depending on the type of prompt. A potential future direction is to extend the online selection to capture the diversity and novelty of generated data. Also, developing cost-aware online selection methods, similar to the cost-aware CB algorithm in (Hu al., 2025b), will be relevant direction for future studies.
Phttps://hughts waw.ggiigafccc. cocccs/ccccc////coooo/o trans formerfrres///eo /tdd/////////////////////// (acddel_ooooaee roberta

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Acknowledgements
technical report, 2023, URL https://:/arrii abs (bb/2239.16600)
The work of Farzan Farnia partially supported by a grant from the Research Grants Council of the Hong Kong Spe- cial Administrative Region, China, Project 14209920, and is partially supported by CUHK Direct Research Grants with CUHK Project No. 4055164 and 4937054. Also, the authors would like to thank the anonymous reviewers for their constructive feedback and suggestions.
Bao, F,, Nie, S., Xue, K., Cao. Li, C., C., Su, H., and Zhu, J. All are worth words: A vit backbone for diffusion models. In CVPR, 2023a.
Bao, F., Nie, S., Xue, K., Li, C., Pu, S., Wang, Y., Yue, G., Cao, Y., Su, H.. and Zhu, One trans- former fits all distributions in multi-modal diffusion at scale. In Krause, A., Brunskill, E., Cho, K., En- gelhardt, B., Sabato, S., and Scarleu, J. (eds.), Pro- ceedings of the 40th International Conference on Ma- chine Learning, volume 202 of Proceedings of Ma- chine Learning Research, pp. 1692-1717. PMLR, 23- 29 Jul 2023b. URL https: http://oooeednggeee www.gm.............mm..........m .mlr. press/ /V0/s//v02//2 wwt.//aaaaaaaaaaae/..ooea.aaaa.. html.
Impact Statement
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
References
Binkowski, M., Sutherland, D. J., Arbel, M., and Gret- ton, A. Demystifying mmd gans. arXiv preprint arXiv: arXiv:/801.0144 /801.01401, 2018.
Agrawal, S. and Goyal, N. Thompson sampling for con- textual bandits with linear payoffs. In Dasgupta, S. and McAllester, D. (eds.), Proceedings of the 30th Interna- tional Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pp. 127- 135, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL hwtphttps:// /ptppppeeenngs...//ppppp/p www.mmrrrr.ess..e..m press/v28/ agrawall3.m html.
Calandriello, D., Carratino, L., Lazaric, A., Valko, M., and Rosasco, L. Gaussian process optimization with adaptive sketching: Scalable and no regret. In Beygelzimer, A. and Hsu, D. (eds.), Proceedings of the Thirty-Second Confer- ence on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pp. 533-557. PMLR, 25- 28 Jun 2019. URL https: pto.eeiicceeenns.ooooeo// .mlr. press / 9 ca lalarreellllll htm.....
Astolfi, P., Careil, M., Hall, M., Mañas, O., Muckley, M.. Verbeek, J., Soriano, A. R., and Drozdzal, M. pareto fronts of condi- tional image generative models, 2024. URL https: //arxiv. .rgg abaaa 10429.
Consistency-diveriityy
Calandriello, D., Carratino, L., Lazaric, A., Valko, M., and Rosasco, L. Near-lincar time Gaussian process optimization with adaptive batching and resparsifica- tion. In III, H. D. and Singh, A. (eds.), Proceed- ings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1295-1308 PMLR, 13-18 Jul 2020. htwphtttss//ppproee v119 / www.aaaaaaaeeeee00 html.
Auer, P. Using confidence bounds for exploitation- exploration trade-offs. J. Mach. Learn. Res., 3(null): 397-422, mar 2003. ISSN 1532-4435.
Avron, H., Kapralov, M., Musco, C., Musco, C., Velingker, A., and Zandieh, A. Random Fourier features for kernel ridge regression: Approximation bounds and statistical guarantees. Precup. D. and Teh, Y. W. (eds.), Proceed- ings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learn- ing Research, pp. 253-262. PMLR, 06-11 Aug 2017. httpt pesttps://procee http://poccedeeege/ mlr press/v70/ avron 7a html.
press/ pr.ppssss//
ml
Chu, W., Li, L.,, Reyzin, I.., and Schapire, R. Contextual ban- dits with linear payoff functions. In Gordon, G., Dunson, D., and Dudík, M. (eds.), Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, volume 15 of Proceedings of Machine Learning Research, pp. 208 208-214 214, Fort Lauderdale, FL, USA, 11- 13 Apr 2011 PMLR. URL https: wwtpttt/tt/pooeeeeeeee ppw.eeddios..s... mir www.spressseeeee httpc///hhhhheae ch..la.hllll..
Bai, J., Bai. S., Chu, Y.. Cui, Z., Dang, K., Deng, X., Fan, Y., Ge. W., Han, Y., Huang, F., Hui, B., Ji, L., Li, M., Lin, J., Lin, R., Liu, D., Liu, G., Lu, C., Lu, K., Ma, J., Men, R., Ren, X., Ren, X., Tan, C., Tan, S., Tu, J., Wang, P.. Wang. S., Wang, W., Wu, S., Xu, B., Xu, J., Yang,
DeepSeek-AI, Liu, A., and al. Deepseek-v3 technical re- port, 2025. URL https:::: //tpxiva.vvvv///////aa/ org/ abs/ 2412. 19437.
A., Yang, H., Yang, J., Yang. S., Yao, Y., Yu, B., Yuan, H., Yuan, Z., Zhang. J., Zhang, X., Zhang, Y., Zhang, Z., Zhou, C., Zhou, J., Zhou, X., and Zhu, Qwen
Foster, D., Agarwal, A., Dudik, M., Luo, H., and Schapire, R. Practical contextual bandits with regression ora- cles. In Dy, J. and Krause, A. (eds.), Proceedings of
10

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
the 35th International Conference on Machine Learn- ing, volume 80 of Proceedings of Machine Learn- ing Research, pp. 1539-1548. PMLR, 10-15 Jul 2018. https:/ httppp:://ppoceeeee /ttpp//rcednggss//// mlr. press/v80/ cester18a.hie html.
Kim, J.-H., Kim, Y., Lee, J., Yoo, K. M., and Lee, S.-W. Mutual information divergence: A unified metric for mul- timodal generative models. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https: / //openreviww.oo net wtw.mmmdddddddddtesddsssee
Frick, E., Chen, C., Tennyson, J., Li, T., Chiang. W.- L., Angelopoulos, A. N., and Stoica, I. Prompt-to- leaderboard, 2025, URL https://arxiv..r abs/2502. 14855.
Kirstain, Y., Polyak, A., Singer, U., Matiana, S., Penna, J., and Levy, O. -ic:- An open dataset of user pref- erences for text-to-image generation. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https:////// /tppneeeeeeeeeeeeeeeeeeeeeeeeeee net/ forum? id-G5RwHpBUvC.
Gemini-Team, Anil, R., and et al. Gemini: family of highly capable multimodal models, 2025. URL https: // //arxiv. org/ abs/ 22122.21222222 11805.
Langford, J. and Zhang, T. The epoch-greedy algo- rithm for multi-armed bandits with side information. Platt, J., Koller, D., Singer, Y., and Roweis, (eds.), Advances in Neural Information Process- ing Systems, volume 20, Curran Associates, Inc.. 2007. URL https::/prooeeedings new.oeeeeeednn.....e..........p cc/ hter///iles//pper///////eee (aaee@bbaa111ce33ffeee3 pdf.
Hessel, J., Holtzman, A., Forbes, M., Le Bras, R., and Choi, Y. CLIPScore: A reference-free evaluation metric for image captioning. In Moens, M.-F., Huang, X., Specia, L., and Yih, S. W.-t. (eds.), Proceedings of the 2021 Con- ference on Empirical Methods Natural Language Pro- cessing, pp. 7514-7528, Online and Punta Cana, Domini- can Republic, November 2021. Association for Computa- tional Linguistics. doi: 10. 595. URL https:/ttttt/....o //t:///cc a http://aaaaaalhllloooooooooooooo ogy. or2/2021 ctmommmin.5990
Paper.
Li, L., Chu, W., Langford, J., and Schapire, R. E. A contextual-bandi approach to personalized news arti- cle recommendation. In Proceedings of the 19th In- ternational Conference on World Wide Web, WWW "00 pp. 661-670, New York, NY, USA, 2010. Associa- tion for Computing Machinery. ISBN 9781 17858877899 doi: 10. 10.1145/117726601777777 org htgg/1..115//// 114111155//1 /1144////77999../ 1772758,
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Hu, X., fung Leung, H., and Farnia, F. A multi-armed bandit approach online selection and evaluation of generative models, 2025a. URL https: /ttp/vaxivv////////////////a.ava org/ abs/ 2406. 07451.
URL https: /oii
Lim, E., Tan, V. Y. F., and Soh, H. Stochastic bandits for egalitarian assignment. Transactions on Machine Learn- ing Research, 2024. ISSN 2835-888 8856. URL https: / //opperevv net. forum? d= 2J
Hu, X., Pick, L., fung Leung, H., and Farnia, F. Prompt- wise: Online learning for cost-aware prompt assignment in generative models. 2025b.
Liu, Y., Zhang, K., Li, Y., Yan, Z., Gao, C., Chen, R., Yuan, Z., Huang, Y., Sun, H., Gao, J., He, L., and Sun, L. Sora: A review on background, technology, limitations, and opportunities of large vision models, 2024. URL https: www.aiii.oooo.ooo wtwgg/ab////ee http//bs//eee /tw./////.2.... 140..7777.
Huang, Z., He, Y., Yu, J., Zhang, F., Si, C., Jiang, Y., Zhang, Y., Wu, T., Jin, Q., Chanpaisit, N., Wang, Y., Chen, X., Wang, L., Lin, D., Qiao, Y., and Liu, Z. VBench: Compre- hensive benchmark suite for video generative models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pottern Recognition, 2024.
Luo, M., Wong, J., Trabucco, B., Huang. Y., Gonzalez, J. E., Chen, Z.. Salakhutdinov, R., and Stoica, Stylus: Automatic adapter selection for diffusion models, 2024. URL https: org abs 18928.8888...
Jalali, M., Li, C. T., and Farnia, F. An information-theorne evaluation of generative models in learning multi-modal distributions. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Masrourisaadat, N., Sedaghatkish, N., Sarshartehrani, F., and Fox, F. A. Analyzing quality, bias, and performance in text-to-image generative models, 2024, URL https: /atx/vv.og//bb//240..
Kannen, N., Ahmad, A., Andreetto, M., Prabhakaran, V., Prabhu, U., Dieng, A. B., Bhattacharyya, P., and Dave, S. Beyond aesthetics: Cultural competence in text-to- image models, 2024. URL https: //ap:ss:///xx..../// /aw.xa...voorooe abs/2407. 06633.
OpenAl, Achiam, J., and et al. Gpt-4 technical report, 2024. URL /Rt.hhcpps//aarrroooro/////////s

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Ospanov, A., Jalali, M., and Farnia, F. Dissecting clip: Decomposition with a schur complement-based approach. arXiv preprint arXiv:2412. 2024a.
cc//apper_illee/ppppr//////// 01 0.330066f03ddcc3322fffeee -55 pdf.
Ospanov, A., Zhang. J., Jalali, M., Cao, X., Bogdanov, A., and Farnia, F. Towards scalable reference- free evaluation of generative models. arXiv preprint arXiv:2407.02966 2276199,7,,,,,, 2024b.
Rahimi, A. and Recht, B. Random features for large-scale kemell machines. Advances in neural information pro- cessing systems, 20, 2007b.
Rezaei, P., Farnia, F., and Li, C. T. Be more diverse than the most diverse: Optimal mixtures of generative models via mixture-uch bandit algorithms, 2025. URL https: / //arxiv. org. abs 2212.
Papineni, K. Roukos, S., Ward, T., and Zhu, W.-J Bleu: a method for automatic evaluation of machine translation. In Isabelle, P.. Charniak, E., and Lin, D. (eds.), Proceed- ings of the 40th Annual Meeting of the Association for Computational Linguistics, pp. 311-318, Philadelphia, Pennsylvania, USA, July 2002. Association for Computa- tional Linguistics. doi: 10.31155/10730033007733 1073135. URL. ht https: / /ta///aaaanthllcce wwwp:y...rrggge... wwg/po2 1040/.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with la- tent diffusion models. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10684-10695, June 2022.
Peng, Y., Cui, Y., Tang, H., Qi, 7., Dong, R., Bai, J., Han, C., Ge, Z., Zhang, X., and Xia, S.-T. Dreambench++: A human-aligned benchmark for personalized image genera- tion, 2024. URL https: ////arxiiivv..a www.iiivorrgaaaao aass///b///// 2206 16855.
Rudin, W. Fourier analysis on groups. Courier Dover Publications, 2017.
Salimans, T., Goodfellow, Zaremba, W., Cheung, V., Rad- ford, A., Chen, X., and Chen, X. Improved techniques for training GANs. In Lee, D.,, Sugiyama, M., Luxburg, U.. Guyon, I., and Garnett, R. (eds.), Advances in Neu- ral Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
Podell, D., English, Z., Lacey, K., Blattmann, A., Dock- horn, T., Müller, J.. Penna, J., and Rombach, R. SDXL: Improving latent diffusion models for rigo-resolution image synthesis. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: //openreviex wwt.:/iwwwweeeeeeeeemmmmeeeeee htrum/orrmmmmooo wwt.mmiiddddddddde.....
Singer, U., Polyak. A., Hayes, T., Yin, X., An, J., Zhang, S., Hu, Q. Yang, H., Ashual, O., Gafni, O., Parikh, D., Gupta, S., and Taigman, Y. Make-a-video: Text-to-video generation without vixt- data, 2022. URL https: //arxiv. org/ abs//aa////// 2209 14792.
Qin, J., Wu, J., Chen, W., Ren, Y., Li, H., Wu, H., Xiao, X., Wang, R., and Wen, Diffusiongpt: Llm-driven text-to-image generation system, 2024. URL https: //aaxiv. org/ abaa/aaa.11.. 10061.
Sutherland, D. J. and Schneider, J. On the error of random fourier features. In Proceedings of the Thirty-First Confer- ence on Uncertainty Artificial Intelligence, UAI' 15, pp. 862-871, Arlington, Virginia, USA, 2015. AUAI Press. ISBN 9780996643108.
Radford, A., Kim, J. W., Hallacy C., Ramesh, A., Goh, G. Agarwal, S.. Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever,
Learning transferable visual models from natural language su- pervision. In Meila, M. and Zhang, T. (eds.), Pro- ceedings of the 38th International Conference on Ma- chine Learning, volume 139 of Proceedings of Ma- chine Learning Research, pp. 8748-8763, PMLR, 18- 24 Jul 2021a. URL https: htttts::/pproeeeeeeeeee mmr. press/v133 39/ htt99//ddffrf2d// html.
Tan, Z., Yang, X., Qin, L., Yang, M., Zhang, C., and Li, Evalalign: Supervised fine-tuning multimodal Ilms with human-alignn data for evaluating text-to-image mod- els, 2024. URL https::///// /taa/aiivvv. org/ httg//bb//406 16562.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. Llama: Open and efficient foundation lan- guage models, 2023. URL. https: http::ss///xxvv.ooo///o abs/2302.1397.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, Learning transferable visual models from natural language supervision, 2021b.
Rahimi, A. and Recht, B. Random features for large-scale kernel machines. In Platt, J., Koller, D., Singer, Y., and Roweis, S. (eds.), Advances in Neural Information Processing Systems, volume 20. Curran Associates, Inc., 2007a. URL https: http///oceedddeeec ngs. neurips.
Valko, M., Korda, N., Munos, R., Flaounas, I., and Cristian- ini, N. Finite-time analysis of kernelised contextual ban- dits, 2013. URL httss: : /arxivv. org/ org/abs// abs/ 1309. 6869.
12

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Xu, J. Mei, T., Yao, T., and Rui, Y. Msr-vtt: large video description dataset for bridging video and language. In 2016 IEEE Conference Computer Vision and Pattern Recognition (CVPR), pp. 528855288299 5296, 2016. doi: 10.1109) CVPR.2016.571.
Xu, J., Liu, X., Wu, Y., Tong, Y., Li, Q., Ding, M., Tang. J.. and Dong, Y. Imagereward: Learning and evaluat- ing human preferences for text-to-image generation. In Thirty-seventh Conference on Neural Information Pro- cessing Systems, 2023. URL https::tp... http:::://ppeeeviiic net/ forum? -irrumiid--Vze0EExxx6
Xu, X., Dong, F., Li, Y., He, S., and Li, X. Contextual- bandit based personalized recommendation with time- varying user interests, 2020. URL https://arxiv. org/abs// 200///bb//00....// 00359.
Zenati, H., Bietti, A., Diemert, E., Mairal, J., Martin, M., and Gaillard, P. Efficient kernelized ueb for con- textual bandits. In Hamps-Valls, G., Ruiz, F. J. R., and Valera, I. (eds.), Proceedings of The 25th Interna- tional Conference on Artificial Intelligence and Statis- tics, volume 151 of Proceedings of Machine Learn- ing Research, pp. 6689-5720. PMLR, 28-30 Mar 2022. httpst / http://proeeeeeeeee/eeeee mlr.press/ v151/zenati221 html.
Zhang, T., Kishore, V., Wu, F,, Weinberger, K. Q., and Artzi, Y. Bertscore: Evaluating text generation with bert, 2020.
URL TtlLhhtss://axxtt v. org/ abs/ 904. 09675.
Zheng, Q., Xia, X., Zou, X., Dong, Y., Wang, S., Xue, Y., Wang, Z., Shen, L., Wang, A., Li, Y., Su, T., Yang, Z., and Tang, J. Codegeex: A pre-trained model for code gener- ation with multilingual benchmarking on humaneval-x. In Proceedings of the 29th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 5673-5684, 2023.
Zhuo, Y., Vu, M. C., Chim, J., Hu, H., Yu, W., Widyasari, R., Yusuf, I. N. B., Zhan, H., He, J., Paul, I., et al. Big- codebench: Benchmarking code generation with diverse function calls and complex instructions. arXiv preprint arXiv:2406. 15877, 2024.
13

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
A. Proof in Section 5.1
The technical challenge analyzing PAK-UCB is that predictions later iterations make use previous outcomes. Hence, the rewards {s; Nev, are not independent if the index set Dg updated each time when model g is chosen (Line 6 of Algorithm 2). To address this problem, we leverage a standard approach used prior works (Auer, 2003; Chu al., 2011; Valko et al., 2013) and present variant of PAK-UCB in Algorithm 3, which is called Sup-PAK-UCB.
Algorithm 3 Sup-PAK-UCH
Require: total iterations I' E N+, set of generators G [G], prompt distribution p E ( score function s y X'
[], positive definite kernel k x y R, regularization and exploration parameters a, 2 0, function
COMPUTE_UCB in Algorithm 2
Initialize: observation sequence D <- Ø and index sets { +- 2) for all g E 9, where M <- log T
1: for iteration / = 1, 1....... , do
2: Prompt 8/2 ~ p is revealed.
3: Set stage
and 9 4- g.
4: repeat
5:
Compute (1ng o0g < COMPUTE (CB yt. VT ge ge@m.
6: Set 8g (vr) t Ag + nōg" and sy (V+) + Fg + (2n vajōg for all g Gm.
7:
if (17 + vajōm < 1/vT for all E Gmh then
8:
Pick model gt < arg max 9 (y).
9:
else (2n + 0a) s 2mmm for all g E GTL then
gm+I1 9 E 5m (ve) 2 max gE (yr) 22-m)
10:
Set stage m m +
12:
else
Pick Sr -LH such that (2n+ va)on" > 2
13:
14:
Update Tm < Tm U {t}.
15:
end
16:
until a model gr selected
Sample an answer H, ~ Pg, (-1y)) and compute the score s, s(y:: 22).
Update D t- D ((),, $t)).
19:
end for
A.1. Regret Analysis of Sup-PAK-UCB
Theorem (Regret of Sup-PAK-UCB), Under Assumption 1, with probability least 1 d, the regret of Sup-PAK-UCB with 1 /2 log(2GT/8) is bounded
Regret(T) Sō (a+ 5 (u+vay vayy dotf |+ (1+1 GT
(4)
where dett is data- data-dependent quantity defined in Lemma 5 and logarithmic factors are hidden in the notation Ö(-)-
Notations. To facilitate the analysis, add an subscript all the notations in Algorithm 3 indicate they are quantities computed at the t-th iteration, i.e., Pyn and g.t (first appeared Line 5),3 and ay, (first appeared Line 6), I.m. and Tim. In addition, we define g.,, arg maxgeg (y)) as the optimal model for prompt yı and g.,. := arg max 8g. is optimistic model at stage m. 10
g.t"
Th Proof of Theorem 3. Let Ti == Ume[M]].ed) 9.T41 and To := Tz. Note that To and Ti are sets of iterations such that the model is selected Lines and 13 of Algorithm 3, respectively.
Note that Line 14 of Algorithm is rewritten "mm for any g o gr and m [M]" In addition, we set qo
51.41
qm 40,.41 U (t), vm v..441 me for any m' o m, and vp upmm
9.441 Tmm
g.t41 f- Waa41 for all g E & and m E [M] Line
14

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Regret incurred in To- For any & [T), let my denote the stage that model Dt is picked the /-th iteration. We have that
Σ KETO
(v+) - Sy. (v:)) = Σ (V+)
Sy. (v) + 8 gg,t, (U)
8g (Vt)
+8
(ve) Sg. (36))
To
S 5y Line
<(n + va)
Jm + (3r) va)
m ged
(5)
tETo
1ETo
<O(vT),
where the first inequality holds by the fact that In E g.me
Regret incurred in 71.
Note that
Σ
(s. - (Br)) = Σ
Σ
(sa(Be)
Sg. (vr))
RE
gEg mE[M
LE Wim
g.
(6)
23-m
gEg ME[M]
where the inequality holds by the last statement Lemma 3. It remains to bound Im we have that
g.r+1 First note that for any m E [M].
Σ
(2n va)a 21- 1 9.7+1
TE D.T
from Line 13 of Algorithm 3. In addition, by a similar statement of (Valko et al., 2013, Lemma 4), which is stated in Lemma 5, we have that
S Ō 56 (V/ca (1+ (( Im WTT+1
Σ
m 3.t
where der is defined therein and logarithmic factors hidden in the notation Ö(.). Plugging in Equation (6) results in
(s)(v)) - Sz. (Ui) =0
Σ Σ
V den (1+ 1
Imm
g
IE
MEB mE(Md]
VGM
deff (1+
Σ {
Σ
N
g.T+1
(8)
9E9 ME[M]
(vaar
ddar
1
GT
where the second inequality holds by Cauchy- Schwarz inequality.
3. Putting everything together. Combining Inequalities (5) and (8) leads to
(s) - (0,)) sp. s Ô (a+ va) va) dott (1+ GT
(T)
Σ
+
Σ tET,
GT
KETo
which concludes the proof.
A.2. Auxiliary Lemmas
Lemma 2 (Auer (2003), Lemma 14). For any iteration t E [T'), model E G, and stage m E [M], the set of rewards (sihe are independent random variables such that B[ss] = sg (y:). Lemma 3. With probability at least (MGT)). for any iteration E [T] and stage m E [M] in Algorithm 3, the following statements hold:
15

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs Iige sy(Ve)] S (2n + yo for any g 9r
arg mggga 8y (y1) 9.mh and
s. - - (y6) 2 for any g 9tn
Proof. The first statement holds by both Lemma and Lemma and uniform bound over all f E [T],g e g. and m [M]
To show the second statement, first note that g..t E T.. Assume g..t E 9IT for some m E [M 1]. Let 5g.
F (2n + Va)ag, (computed in Line 10 in the algorithm). Then, by the first statement, obtain that
8m
(y1) +8x(2))
max {Fg, sz(Ue) 8y (v:)) gEgN"
9g9t
g..el
gE
2 max
(8y,e- 88g(1)}
gEst
22n + valo dit 22 2 (-21-m) -22-m
where the last inequality holds by Line 9 of the algorithm. To show the last statement, note that for any E 3.Th. holds that
(-(1 sg(1e) (v)
3m
(8g(ve) 5g.t + 331,)
<2(2 + va mm + 13
8%
gr.ett
<2 2 1 2-- < 23-m
which concludes the proof.
Lemma 4 (Optimism). Let V C. [T] be an index set such that the set of scores {st : E V/ are independent random variables. Then, under Assumption 1, with probability at least 8, the quantity Bg computed in function COMPUTE.UCB(D vn satisfies that
(5-5,(u)1= s (2m + va)ō,g
(9)
where 1) = 2 (2/6). Hence, it holds that So = u. + (2n + vajo, 2 sy()).
Proof. We rewrite the proof using the notations in Section 5. Obviously, Equation (9) holds when the index set V is empty. In the following, we consider non-empty Mg Let 42 == (o(v:) "Iiew Note that ky = [k(y, y:lleew, = =b (o(r)) and
K (1y,,,))le.ee = We have
Mg sy sy( =((())) Φ (K (K + al) (ő(v)) "
=(o(y)) T ( o + al) _ (0())) (Φ o + - Tb + al)
(10)
=d(v)) (( J 1 + al) $
" (v 1 w a((())) (0) o al) w
(2)
(6)
where the second equation holds by the positive definiteness of both matrices (K + al) and (Ø) 2 al) and hence
$ ( + al) (0) 2 al) -
(11)
= Bounding Term (a). Note that the scores {sr E , are independent by the construction of o and [[ss]] (wx)Te ((y)), we have that
(6()))((1))) (0] 4y + +al) T -byw)) = > ((((v) [ p, + +al) a wyle
16

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
the summation of mean independent random variables, where denote by Els the i-th element of a vector. Further, each variable satisfies that
< (0, gl: 1v 2 ugll SV (d(y))
((o(t)) +al)
1|0 2 wiltl
( 2, + a 2 ( 2 + (o(x) (1+ + (l))))
~25g.
where the last incquality holds by || I and the second inequality holds by
6 =a + Vk(y.V) [K +
-a + (oll)) (b(w)) T T (K + (
=a + (o(u)) (1 (1- ( + aa))
_ 2b) (0()))
(d(U)) ((()))((()))) [ ( $ + (b))),
Then, by Azuma-Houuoaoo Hoeffding inequality, it holds that
P ((((())) ( o + +al) '2,
1gw > 2nτ²,)
<2exp
(12)
2
<2exp(-ni//2)
Solving the above inequality to be smaller than leads to = / 12g(2/6).
Bounding Term (b). By the Cauchy-Schward inequality, holds that
(((()) T (6 T T + a <|Iugll (1(Ø(v)) (0) 2, + +al)-111
-||wا|| ((o(y)) 2, + tal)-ta
+ al)-
(13)
Sa 1/2
1 (o(x)) (Ø] 4 + +al) aal)-1(07) 2 + aa) a
-1/20
where the second inequality holds by the positive definiteness of $ Dy
Putting everything together. Plugging (12) and (13) in (10) and setting 2 exp(- 2 /2) concludes the proof.
Lemma 5 (Valko al. (2013), Lemma 4). For any model g E G and stage m [M], let X²n, 2 A31 2 ... denote the
eigenvalues (in the decreasing order) of the matrix (Tm)) T. + +al, where g = (p(x)) lie Then, for any iteration 1 [T], it holds that
am 2ō ( (1+
Im 1yutall
TEW
where deft := me(g min{j N4 ja log T > And and Agg = Eiss Ag
is the effective dimension.
17

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
B. Missing Algorithms and Proofs in Section 5.2
B.1. The RFF-UCB Algorithm
Algorithm RFF-UCB: PAK-UCH with RFF implementation
Require: iteration T, generators G = [G], prompt distribution p. score function : 3 x & - [-1, positive definite kernel k, regularization and exploration parameters ax, 1j 2 0, function COMPUTE_UCB RFF (Algorithm 5) Initialize: observation sequence D <- and index set V f for all g E g
1: for iteration f 1.2..... do
2: Prompt 1/+ p revealed.
3: Compute (Fy 0g)) + COMPUTE UCB RFF(D, gg) and set 8gg + Mg + 170g for each g E G.
4: Pick model + arg max
5: Sample Hę ~ Pg. (-1yx) and receive score St-
6: Update D <- D U ((1t,1)) and v f- Vn (1).
end for
Algorithm 5 Compute UCB with Random Fourier Features
Require: the Fourier transform of a positive definite exploration parameters o,, 2 0
kernel k(v,8') = k(y v²), regularization and
Initialize: number of features D
1: function COMPUTE_ UCB CCB_RFF(DD2) v if L empty then
3: 1g +00, og
+00
4: else
Draw WT w iid.
5:
6:
Define mapping pp(') & D [cos(w] v), sin(wT v).... c sin(wTy)]1 for any s e Rd.
7:
Set Φ + (p())) liev, and [silee,
8:
Bg - ( 0 +alnn) 1 y
9:
g - a (1 ( + alan) (v((()))
10: end if
11: return (7y, kg).
12: end function
Analysis of Lemma 1. Solving KRR with n regression data requires (nr²) time and O space. Hence, by the convexity of the cubic and quadratic functions, the time for COMPUTE UCB scales with O geg ng) O /G)). and the space scales with O m²) O /G), where nh == [ is the visitation to any model g up to iteration t, and we have Egeg Rg t. On the other hand, solving KRR with regression data and random features of size requires On)D²) time and (nD) space. Therefore, the time for COMPUTE.UCH RFF scales with O(Egeg n, D²) O(ID"), and the space scales with (Egeg D) = O(tD).
B.2, Regret analysis of Sup-RFF-UCB
Algorithm description. We present Sup-RFF-UCB Algorithm 6, where we utilize function COMPUTE_ CCB.RFFFR to compute the UCB values. achieve the regret bound (4), an important problem is to design (adaptive) error thresholds, i.e., <RFF and AREE, when computing UCB each stage and iteration I. We prove the regret bound in the following theorem.
18

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Algorithm 6 Sup-RFF-UCB
Require: total iterations T' E N+, set of generators G [G], prompt distribution p A(V), score function s y X +
-1, 1]. positive definite kernel k y y R. regularization and exploration parameters a, 2 0, function COMPUTE_UCH RFF in Algorithm 5, parameters D B g.1, and g22,,
Initialize: observation sequence D <- and index sets { Ø } for all g E G. where M log
1: for iteration & 1. T do
2: Prompt J/2 p revealed.
3: Set stage +- and 9 -- G.
4: repeat
5:
Compute { gg) &- COMPUTE CCMPUTE_UCB.AARR UCF.UT Vt, DD")
geg with RFF dimension D
6: Set 8 = Pg + noo n for all g E 9mh
7: Set 5g t Fg gg1 + (2n + vallōg + 2,,,)) for all g e 9mh.
8:
if (1 + vajōg" I 1/VT for all E Gmh then
9:
Pick model 9r t arg max gc $g
10:
else if B 21 2 + (2n + Vā) + Bm g22,) S 2 21 for all g E 9mh then
11:
Gm+ {9 E 5m sy
f - 22-m)
12:
Set stage m t m + 1.
13:
else
Pick 9t E 9m such that B + (2n + (a) + 1 21-m.
Update 4m - vm (t). 9,
end
17:
until model 9r is selected Sample an answer ~ Pg. (-lvs) and compute the score + s(yi,
18:
19:
Update D + D U {.),,
20:
end for
Theorem 4 (Regret of -FF-RFF-- With probability at least 1 - 6, the regret of running Sup-RFF-UCE with exploration parameter 17 /2 /2log' log( (8) and regularization S 1or T iterations is bounded by (4). Specifically, at stage m [M] at iteration t E [T), the parameters satisfy
Dm max
[(d + SERFE
2
log
Op
+ log
3d
8|W6l (A)FF
2
324 (K)
(14)
(ERIF t
1+3
2
ERFF
m g.t
3x
d
8m
m
ERIF g.t +a - 1m
(m g.tl
+ aa);
(15)
TT 9.2.t (at) 1/2 + (a ARFFg g.
+ a)
2 ERFF
(16)
where la is defined in Lemma 10, Ç and on are defined in Lemma 11, and the error thresholds are given by
CRFT S 2
ARFF at 3100 g. (110)++a) +
(17)
Proof. Note that 1Th S O(t and 1m I ()). The proof is based on Lemma 6. For iterations in To (model gr is
picked in Line 8 of Algorithm 6), we have
8g. (wr))
To
Σ
(vr)
gime ( + 5Th (yy)..(y) Sot +8mh 9g. (3re)
To
5 Line
8 ( + +s
8gn (1)
8
sen (yr) ge.t
3m (yr + (yx) - 8g. (V+))
1oTo
< by Lemma
definitions
19

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
B
ge,1,t11 + (n + vajō gnitt + 2B 9r
+ 2(2n + vajom.. med + (2n + (B) s.,t, + 225m. ge.t,2)
To
<4
(1) + vo: (o g.a,t + ggii)
B sneed.1 + g0.t11 2(22 va) br. + g2.1.2)
LE7o
2 by Line
<ō
VT+ Σ
5 (1+ v+
TETo
<ō
(1
where the last inequality holds by the fact that each E [T) appears in at most one index set. Further, for iterations in Ti (model is picked in Line 13 of Algorithm 3), the third statement in Lemma 6 ensures that
Σ ss. (16) Sge (3e))
23-m.
Im g
SEG [M[M]
Σ
(B + (2n + vα) (ot), + B.m,./1
BEG mc [M]
2.111
<
Σ Σ
(8%).. 1. + (2n + va) (og,, + 25
SEG mE[M]
RE m g.7+1
ΣΣ
(BTh. g1.1tt + 2(2 + va)B 91,2,1) + +21 va) Σ am
ME(M]
LEW 2,711
LE LE-VMM Mo
9.T11
Σ
B 2(21) va)B mi.2.))
+4 +
(2n+
Σ m LE
m.2.1)
ME[M]
MEW 9,741
SEC mE[M]
g.
Note that the upper bound of the second term has been derived in Equation (7). It remains to bound the first term. Essentially, we will find a sequence of eтror thresholds, and hence the number of features defined in Inequality (20), such that the first term is bounded by Õ((1 + a --=/)/T)
For any iteration & [7), model g E G. and stage m E [M]. we define Km 9.8 First, by the definition of By in Equation (15), have
($m T where t = [ Tieem.. == 2.. (1 (B)).)
Σ
Σ
8 9.141
gr,1,
9Eg geg.me[ MM LE
Σ
(N ERFF +a +| n. ARFF gf
(( mi + +a)
gEG.mE M] M1
(18)
(1-+
50 (log
where the first inequality holds by the fact that each i E [T] appears in at most one index set. On the other hand, by the
20

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
definition of g.2 in Equation (16), derive
Σ
gEGmmE M] te 9iT+1
Σ
((at) $ - +a-
(R aa-+taa-aaAaaa + a) + ERFF
9 ME M]
tE 9.741
(19)
T
4a - 6-
t=1
< a vT)
Therefore, we conclude the proof.
B.3. Auxiliary Lemmas
Lemma 6. With probability at least (2MGT)6, for any iteration E [T] and stage m E [M], the following hold:
1 - sy(vi)l & BT g.1.t + (2m + va) (o g.t + +B, 9.2,,) for any G...
arg gag 8y(11) Gim and
- 23-m for any T.
where the first statement is guaranteed by Theorem 7, and Bm g.1.t and B 9/2,t are given in (15) and (16), respectively.
Proof. The proof follows the exact same analysis of Lemma 3.
Lemma 7 (Optimistic KRR estimators with RFF). Assume the error thresholds input Algorithm 5 satisfy that ARFF, ERFT S 1/2. Let regularization < 1. With probubility at least 1 - 26, the quantity Py computed by function COMPUTE UCB RFF satisfies that
F.g. - Nyg-sy(00 (v) < B, + (2n + (a) - -By22),
with the number of features satisfying Inequality (20) and bonus terms BB,1 and Bg,2 given by Equations (15) and (16). where 1) = / (2/()) Hence, holds that ky = 7o B + (2n + (ā) + Bg22) 2 sy(y)
Proof. The proof based on the following two lemmas, which analyze the concentration error of the quantities jig and ōg.
Finally, combining Lemmas 8, and derive that
[ig sy(v)) S Tig - pg + 18g sy(')l
SB,,1 (2n + vajō,
BB,.1 (2n + Vā) Bg,2),
which concludes the proof.
Lemma 8 (Concentration of mean using RFF). Let ARTT, CRFF S 1/2 and the regularization parameter S 1. Let U C [T] be an index set such that the set of scores {st :t E V are independent random variables. Then, with probability at least 8, the quantity Bg computed by function COMPUTE.UCB. RRF satisfies that
F.g p. M CRFF +a - | ARFF( V a:)
with the number of features
[4dd + 2)
2 log Op
Ba 8
8/
(32 (K)
D > max
+ log
ARFF -2 log
(20)
ERFF
1+
2
ERFF
3a
where T. (p(1)) T K + aa) -1 and os Somp 3a and ta (.) are quantities defined in Lemmas 10 and 11, respectively.
21

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Proof. For convenience, define ky (v(((())) E RI Q (K +a) E R and := (K + al)- R v xx Wal, where K- == - T Using the same notations in the proof of Lemma 4, we obtain that
11g - Bgl
(o(u)) $. (K + +al)='o v- - ky (K + aal)
Ry Qu kyy Qu
(21)
R ( - Q)u + (Ty - kiy) QqQ
where we use Equation (11) to derive F.g = (p(1)) (6) 4 + all - (p(u)) all) in the first equation.
Bounding [k3 TQul. Note that y C gd-1, We evoke (Sutherland & Schneider, 2015, Proposition 1), which rewritten in Lemma using our notations. For a desired threshold GRIF > set
4(d + 2
2
log Op + log
B8 5
D
"RIF"
I+3
2 d
CRFF
Then, with probability at least 1- 6. it holds that w.yeeyy ((p(1)) p(()) -kl// 3' S CRFF, and hence 1 kyylo S CRFF. Therefore, obtain
(kh, Ty)TQri S 1k0 kull2 l1Q112 |lul2a S CRFF / (1 a1+ /lWgll //l s / CRIF.
(22)
where the second inequality holds by ||Ql22 Ami (K + af) s a1+ and |10llcc S 1.
2. Bounding [R. (Q - Q)ul- Note that
Ki ( Q)ol 511 12 |1 Ql1 ||w||2 - / V 11 QQ/2 V
where the first inequality holds by the fact that ((())) (vx) S 1. To bound |1 -Q12. we evoke (Avron et al., 2017, Theorem 7), which is rewritten in Lemma 10. For a desired threshold ARFF I 1/2, the following inequality holds with probability at least 1 1- 6.
2
(1 - ARF) + al) K aa
for D > SlVel
ARFF 2 "log(32.. (32c.. (K). By therman-Morrioon-oooddee formula, i.c., A - B-1 -B A)B
30
where A and B are invertible, we derive
I| Ql1 =I((K al) (K +al+ 1|2
<I[(K + al)- [[ + al) - (K + al)12 [[K + al)-
(23)
<a DRFF (| + a) < a v +a),
where the second inequality holds by the fact that [/]12 2 1[ + al) a 2 a 1|Q112 II(K + a S (1 a1+a)-1
and [K + al) - (K + a2) $ ARFF K + ad)ll2 < ARF ( ||2 + a)., and the last inequality holds by the fact that |||| < | for any shift-invariant and PSD kernel.
3. Putting everything together.
Combining Equations (22) and (23), with probability least 1 5, it holds that
Fig - pal < [b CRIF +o - ARTT( [T + a:)
when
D2 max
[4dd + Srum
2 log
Op
+ log
3a
8 g.
ARFF -2
log
32t (K)
"RFF"
1 2
ERFF
5
3a
which concludes the proof.
22

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
Lemma 9 (Concentration of variance using RFF). Conditioned on the successful events in Lemma 8, the quanfity ōg computed by function COMPUTE, UCB RFF satisfies that
[0g - ووو S max (al) -a- ARFF( +a) + 2 ERFF)
(24)
where o = a / - k (K + al)-
Proof. We use the same notations in the proof of Lemma 8, Note that if o o S (at) 3, we have 1o, - -āgl s (at) On the other hand, if one of og. 2 (at)- 2, have
1
=a
2
1 (p(()) $5 (K + al) al)-11 /1 -1- kJ(KK
(K + aa)-
SV Svla
(v(()) T T5 (K + at) aa) ((p(() kiy (K + a
=Via- 4
ky Qky Qky
(25)
Svla= $ (R. (0 Q) + (ky - kin) TQkyy + k (T kb)))
SVta
[1111811@ Ql1 + 1k - K ll K 12l1Q112 + 11 -
<iaa=
(a-1 ARFF / + a) + ERFF),
where we utilize Inequality (23) to obtain the last inequality. We conclude the proof.
Lemma 10 (Avron al. (2017), Theorem 7). Let K [k(y 3 Js)le. denote the Gram matrix of {ys E R where shift-invariant kernel function. Let A 1/2 and ō E (0, 1). Assume that [||]2 2 a. If we use D z 8n A 1og(16ta(K)/0 random Fourier features, then with probability at least 1 - 8. it holds that
(1 - (1- + al) = K al al (K +
where to (K) == Tr[(K + = and we denote by K- = [(((y))) (s((e)))) the approximated Gram matrix using E N., random Fourier features, where :: pdd R2 is the feature mapping.
Lemma 11 (Adapted from (Sutherland & Schneider, 2015, Proposition 1)). Let k be a continuous shift-invariant positive-
definite function k(y.v') k(y - -') defined on y S Rd, with k(0) = 1 and such that V²k(0) exists. Suppose y is compact, with diameter di) Let (p()) be as in Equation (1). For any c > 0, let
Se = -immin sup + k(2,, 2n/ - A(y,8') + 1
1
)
ey 2
v.
Ba = प ) +
d 12x () 2
(
arي2 ककर
Then, assuming only for the second statement that S didmm()))
(0
2 exp
De 4(d + 2)
IP sup
((((())) ((v) k(.. S 66 ea
where of E.p [|]] is the second moment of the Fourier transform of k. Thus, we can achieve an embedding with pointwise error more than with probability at least - 8 as long as
4(d + 2){e
2
log dzdiam())
Bd log
D
2
+
E
8
For the RBF kernel with parameter o², i.e., keg v/)
exp(
1 lly - of 1112), we have o
20
23

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
C. Additional Experimental Details and Results
C.1. Results on Setups and 2
Varying rankings of text-to-image models. We provide more examples showing that prompt-based generative models can outperform for text prompts from certain categories while underperforming for other text categories (see Figures 7, and 11).
Stable Diffusion v1.5
PixArt-a
Examples (clockwise)
woman sitting with her dog and a toy" "the little dog stands near the toilet in small bathroom"
Prompts of Type "dog"
3. "a dog on leash drinking water from a water bottle" " "a dog laying on the floor next door"
avg. CLIPScore
36.37 (+0.13)
37.24 (=0.9
"a motorcycle is on the road behind car" "two cars and a motorcycle road being crossed by herd of elephants"
Prompts of Type "car"
3. "a car that had run
over a red fire hydrant"
"a taxi driving down city street below tall white buildings"
avg.CIe CLIPScore
36.10 (+0006)
35.68 (+0.15)
Figure 7. Prompt-based generated images from Stable Diffusion v1.5 and PixArt-o: Stable Diffusion v1.5 attains higher CLIPScore in generating type-2 prompts (36.10 versus 35.68) while underperforms for type-1 prompts (36.37 versus 37.24).
0.70
0.6
PAK-UCB-poly3
0.65
PAK-UCB-poly3
0.4
RFF-UCB-RBF
RFF-UCB-RBF
LinUCB
LinUCB
0.2
KernelUCB-poly3
0.60
KernelUCB-poly3
One-arm Oracle
One-arm Oracle
0.0
PAK-UCB-lin
0.55
PAK-UCB-lin
Naive-KRR
Naive-KRR
-0.2
Random
0.50
Random
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
(a) Outscore-the-bes (OIB)
(b) Opcimal-pick-ratio (OPR)
Figure Prompt-hased selection between Stable Diffusion v1.5 and PixArt-a (see Figure 7): Results are averaged over 20 trials.
Prompt-based selection over three T2I models. See Figure 13.
24

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
UniDiffuser
PixArt-a
Examples (clockwise)
"the rusted out remains of a small railway line"
2. "a skier stands in the snow outside of train
Prompts of Type "train"
3. "train cars parked on train track near pile of construction material"
4. "several people on horses with a train car in the background"
avg. CLIPScore
35.29 (+000
34.25 (#0.12)
1. "a man in red shirt holding a baseball bat"
Prompts of Type "baseball bat"
2. "a woman holding buseball but with her head resting on it"
3. "baseball player in the batter's 's box hitting baseball
4. "hind view of a baseball player, an umpire, and catcher"
avg. CLIPScore
32.51 (+0.05)
34.30 (+0.04)
Figure 9. Prompt-bas based generated images from UniDiffuser (Bao al., 2023a) and PixArt-a: UniDiffuser attains higher CLIPScore in generating type "train" prompts (35.29 versus 34.25) while underperforms for type "baseball bat" prompts (32.51 versus 34.30).
0.70
0.50
PAK-UCB-poly3
0.65
PAK-UCB-poly3
0.25
RFF-UCB-RBF
RFF-UCB-RBF
LinUCB
LinUCB
0.60
KernelUCB-poly3
KernelUCB-poly?
0.00
One-arm Oracle
One-arm Oracle
0.55
PAK-UCB-lin
PAK-UCB-lin
-0.25
Naive-KRR
Naive-KRR
0.50
Random
Random
1000 2000 3000
4000
5000
0
1000
2000
3000
4000
5000
(a) Outscore-the-bes. (OIB)
(b) Opcimal-pick-raiio (OPR)
Figure 10. Prompt-based selection between UniDiffuser PixArt-ce (see Figure 9): Results are averaged over trials.
Adaptation new models and prompts.
See Figures 14 and 15.
25

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
UniDiffuser
Stable Diffusion v1.5
Examples (clockwise)
an elephant is carrying people across a forested area"
Prompts of Type "elephant"
2. elephant has a sheet its back
3. "a small gray elephant standing on a beach next to lake"
4. "a large elephant in an open field approaching a vehicle"
avg. CLIPScore
36.67 (+0005)
35.08 (+0006)
1. "a fire hydrant in clump of flowering bushes"
Prompts of Type "fire hydrant"
2. "a fire hydrant on a gravel ground with a fence behind it"
3. "a fire hydrant that is in the grass"
4. "a toy Ford truck next to a fire hydrant"
CLIPScore
35.11 (+0.14)
37.23 (+.00))
Figure 11. Prompt-based generated images from UniDiffuser and Stable Diffusion v1.5: UniDiffuser attains a higher CLIPScore in generating type "elephant" prompts (36.67 versus 35.08) while underperforms for type "fire hydrant" prompts (35.11 37.23).
0.8
0.75
0.70
0.6
PAK-UCB-poly3
PAK-UCB-poly3
RFF-UCB-RBF
RFF-UCB-REE
0.65
0.4
LinUCB
LinUCB
KernelUCB-polly
BerneiUCB-polyy
0.60
0.2
One-arm Oracle
One-arm Oracle
PAK-UCB-lin
0.55
PAK-UCB-lin
0.0
Naive-KRR KRR
Naive-KRR
Random
0.50
Random
-0.2
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
(a) Outscore-the-hest (OtB)
(b) Optimal-pick-rane (OPR)
Figure 12. Prompt-based selection between UniDiffuser and Stable Diffusion v1.5 (see Figure 11): Results are averaged over 20 trials.
C.2. Results on LLM Selection
C.3. Results on Other Conditional Generation Tasks
Text-to-Image (T2I). In this setup, we synthesize five T2I generators based on Stable Diffusion 12, where each generator an "expert" in generating images corresponding to prompt type. The prompts are captions in the MS-COCO dataset from five categories: dog, car, carrot, cake, and bowl. At each iteration, caption is drawn from (random) category.
lerbbbllliie e-df ffusion-2
26

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
0.25
0.55
0.00
0.50
PAK-UCB-polyy 3
PAK-UCC-poly3
RFF-UCB-REE
RFF-UCB-RBF
-0.25
0.45
LinUCB
LinUCB
KernelUCB poly3
MernelUc--poly:
-0.50
0.40
One-arm Oracle
One-arm Oracle
PAK-UCB-lin
PAK-UCB-lin
-0.75
Naive-KRR
0.35
Naive-KRR
Random
Random
0.30
-1.00
1000 2000
3000
4000
5000
0
1000
2000
3000
4000
5000
(a) Outscore-the- the- (OIB)
(b) Optimal-pick-ratoo pick-ratio (OPR)
Figure 13. Prompt-bar based selection among Stable Diffusion v1.3, PixArt-a, and DeepFloyd: Prompts are drawn from types "carrot" and "bowl" in the MS-COCO dataset. Results averaged over 20 trials.
0.4
0.6
PAK-UCB-poly3
PAK-UCB-poly3
0.2
RFF-UCB-RBF
RFF-UCB-RBF
0.5
LinUCB
LinUCB
0.0
KernelUCB-pole
KernelUCB-por -ply33
One-arm Oracle
One-arm Oracle PAK-UCB-lin
PAK-UCB-lin
0.4
-0.2
Naive-KRR
Naive-KRR
Random
Random
1000 2000
3000
4000
5000
1000
2000
3000
4000
5000
(a) Outscore- seusore--hhe-eee best (OIB)
(b)
(OPR)
Figure 14. T2I generation with newly-introduced models: Initially, Stable Diffusson v1.5 and PixArt-cr available. UniDiffuser introduced after 2500 iterations. Results are averaged over 20 trials.
and an image is generated from Stable Diffusion 2. If the learner does not select the expert generator, then we add Gaussian noise to the generated image. Examples are visualized Figure 21.
this setup, the images are chosen from the MS-COCO dataset from five categories: dog, car, carrot, cake, and bowl. We synthesize five expert generators based on the vit-gpt2 model the Transformers repository If non-expert generator is chosen, then the caption is generated from the noisy image perturbed by Gaussian noises. Examples are visualized in Figure 23. The numerical results are summarized in Figure 24.
Image Captioning.
3. Text-to-Vided (T2V). We provide numerical results on synthetic T2V setting. Specifically, both the captions and videos are randomly selected from the following five categories of the MSR-VTT dataset (Xu et al., 2016): sports/action, movic/comedy, vehicles/autos, music, and food/drink. Each of the five synthetic arms corresponds to an expert in "generating" videos from a single category. Gaussian noises are applied to the video for non-experts. The results are summarized in Figure 25.
27

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
0.4
0.60
0.55
PAK-UCB-poly3
PAK-UCB-polys
0.2
RFF-UCB-RBF
0.50
RFF-UCB-RBB
LinUCB
LinUCB
KernelUCB-poly3
0.45
KernelUCB-poll:
0.0
One-arm Oracle PAK-UCB-lin
One-arm Oracle
0.40
PAK-UCB-lin
Naive-KRR
Naive-KRR
-0.2
Random
0.35
Random
1000 2000 3000
4000
5000
0
1000 2000
3000
4000
5000
(a) Outscore-the-best the (OIB)
(b) Optimal-pick-ratie (OPR)
Figure 15. T21 generation with newly-introduse prompt types: Prompts are drawn from two categories in the MS-COCO datasct for the first Ik iterations. After that, additional prompt category is added after each IK iterations. Images generated from PixArt-a and UniDiffuser. Results are averaged over 20 trials.
C.4. Implementation Details and Ablation Study
1. Hyperparameters and Implementation. For the PAK-UCB-poly3 algorithm, we utilize the k 222 (1 ar (x2)3 kernel with 7 5.0, 10.0. For RFF-UCB-RBE, we utilize the (1 .22) = exp(-22g 200-dimensional RFF (i.e., D 200), where we 7 0.5 for LLM-related experiments (Setup 2) and 7 5.0 for the rest of experiments. In both algorithms, the regularization parameter is set to be parameter 1) / (og which follows our regret analysis in Sections 5.1 and 5.2. We choose the standard 0.05.
20- ||>| x2 (12) kernel with
In addition,
the exploration
2. Ablation study on hyperparameters. We conduct ablation studies the selections of parameter a in the RBF kernel function and the number of features in RFF-UCB-RBE The results are summarized in Figures 26 and 27, respectively, We select a 1,3,5 5, and 7, and the number of features varying between 25, 50, 75, and 100. Results show that the RFF-UCB- REF-UC--BBB-A algorithm can attain consistent performance. Additionally, we test the LK--UCB-pooyy algorithm with ? == 1,3,5, 5, and 7 in the polynomial kernel and regularization parameter a 0.5, 1.0, and 1.5. The results are summarized in Figures 28 and 29, respectively.
Comparison of PAK-UCB-REE and RFF-UCB-RBE. We compare the performance of PAK-UCB-RBF and its RFF counterpart, i.e., the SFF-UCB-RBL algorithm. The results show that the RFF method attains similar performance with the original PAK-UCB-RBEE algorithm (Figures 30 and 31). Moreover, the RFF-UCB-RII algorithm can significantly speed up the computation (Figure 32).
28

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
follosing by filling enpty jnarked ".)
A4d Puzzle Detaila;
Roard Size: 9 rows ool urne
Rulex: Standard Sudoku rulex apply.
Btater
1
9
4
7
2
8
User
2
8
2
6
8
Requirenents:
Provide complete solved in
format as above, reglacing all 'V" with the corzect valces
not include additional explanation or commentary
Nointain
zos/colunn
your solution
speai Fied rules
#48 Your Solution rconplate board only):
Figure 16. Sample prompt for random Sudoku-solving task (Setup 2)
0.04
0.75
0.02
PAK-UCB-polyy
PAK-UCB-poly3
RFF-UCB-RBF
0.70
RFF-UCB-RBF
0.00
LinUCB
LinUCB
KerrelUCB-polk
0.65
KernelUCB-poly3
One-arm Oracle
One-arm Oracle
-0.02
PAK-UCB-lin
0.60
PAK-UCB-lin
Naive-KRR
Naive-KRR
-0.04
Random
Random
0.55
1000 2000 3000
4000
5000
0
1000
2000
3000
4000
5000
(a) Outscore-the-bet (OIB)
(b) Oicimal-pick-raiio (OPR)
Figure 17. Random Sudoku-solving task (Setup 2): Deepseek-Chee o3-mini, Results are averaged over trials.
29

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
coxplete the Python function below. Follow these roles strictly:
1. Provide OKLY the code solution (ao explamations, markdom,, Haintain
textl.
indentation spaces for Pythos).
Exanple:
Prongt: def return]():
Corplstion:
import
subprocesa check outpust -If
Now complete this Eusction:
User
from typing ingort Liat.
has alarenta (nunberas (loat],,
shreahold: floas) => bool: <**Cheok it in given list of numbers, are any numbers closer to wach other than given threshold.
haa_close, www__coo_oeeeeee close_eienonta(11.) 2.0, 3.01, 0.5)
has close, elersnta((110) 2.8, 3.0, 4.0, 5.0, 2.0], 0.31 True
Figure 18. Sample prompt for Python code completion task (Setup 2)
Pranslate the following Ce+ code Javáz
Declaration einclode<stdio.bs #inolode-nnstb.le using namespace std; ponollddeaagorithae sanolude<atddib
truncate, junber nunber)
Completion the sourcs codes zetarn panber- ink (urber)) AB)
User
Follow
should nainsain
functionalit the souzes code.
2, only outpot completion part vithout
declaration part and any additional explanation comments.
Heintain indertation.
complete tzanclationn
import java.util.*;
inport java.iang.f;
clasa Bolution public double tzuncatolumber (dobble numberi
Figure 19, Sample prompt for C++tto-ava code translation task (Setup 2)
30

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
0.025
0.80
PAK-UCB-poly3
0.000
RFF-UCB-RBF
LinUCB
-0.025
KernelUCB-polyy
0.75
PAK-UCB-poly3
One-arm Oracle
RFF-UCB-RBF
-0.050
PAK-UCB-lin
KernelUCB-poNy
-0.075
Naive-KRR
0.70
One-arm Oracle
Random
Random
1000 2000 3000
4000
5000
0
1000
2000
3000
4000
5000
(a) Outscore- the best (OtB)
(b) Optimal pickal--ii---iiiiii (OPR)
Figure 20. Code generation task (Setup 2): Claude-3.5 Haiku and Semmnn-2.55FFlsh--reeeee Results averaged over 20 trials.
Generator
5
Type 1
Type 2
Type 3
Type 4
Type 5
Figure 21. Generated images with noise perturbations: Each and column display the generated images from synthetic generator according to one single type of prompts. Images generated by the expert models are framedd by green boxes. Gaussian noises are applied non-expert models.
31

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
0.35
2
PAK-UCB-poly3
PAK-UCB-poly3
0.30
RFF-UCB-RBF
RFF-UCB-RBF
1
LinUCB
LinUCB
KernelUCB-pone 3
0.25
KernelUCB-polly
One-arm Oracle
One-arm Oracle
PAK-UCB-lin
0.20
PAK-UCB-lin
Naive KRR
Naive-KRR
-1
Random
0.15
Random
500
1000
1500
2000
500
1000
1500
2000
(a) Outscore- the best (OtB)
(b) Optimal-pick-ratte pick-ratio (OPR)
Figure Synthetic expert model rext-to-image task: The prompts are uniformly randomly selected from MS-COCO dataset under categories 'dog', "car', 'carrot", "cake", and "bowl", Results are averaged over 20 trials.
Example
Example 2
Example 3
"a person a surfboard in the air above the water" 28.57
Clean
"a standing in kitchen with dog"
"a bowl filled with ice cream and strawberries
40.53
27.98
"a blurry photo of a skateboarder flying through the air"
"a cat that is standing in the grass" 13.27
"a blue and white bowl filled with water"
Noisy
24.72
25.83
Figure Generated captions for the clean noise- pertsereeturbedd images from vit-gpt2 the corresponding CLIPScore.
2
0.35
PAK-UCB-poly3
0.30
PAK-UCB-poly3
RFF-UCB-RBF
RFF-UCB-RBF
0
LinUCB
LinUCB
0.25
KernelUCB-poly3
KernelUCB-polyy
One-arm Oracle
One-arm Oracle
0.20
PAK-UCB-lin
PAK-UCB-lin
-2
Naive-KRR
Naive-KRR
0.15
Random
Random
0
500
1000
1500
2000
0
500
1000
1500
2000
(a) Outscore-the the best (OtB)
(b) Optimal-pick-ratie rack-ratio (OPR)
Figure 24. Synthetic expert model on image captioning task: The prompts are uniformly randomly selected from the MS-COCO dataset under categories 'dog", 'car", "carrot', "cake", and 'bowl'. Results averaged over 20 trials.
32

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
PAK-UCB-polys
0.6
PAK-UCB-poly3
RFF-UCB-RBF
RFF-UCB-RBF
LinUCB
LinUCB
KernelUCB-poly3
0.4
KemelUCB-polyy
One-arm Oracle
One-arm Oracle
PAK-UCB-lin
PAK-UCB-lin
Naive-KRR
Naive-KRR
0.2
-2
Random
Random
0
500
1000
1500
2000
0
500
1000
1500
2000
(a) Outscore-the-best (OtB)
(b) Optimal-pick-raiio (OPR)
Figure Synthetic expert model on text-to-video task: captions are uniformly randomly selected MSR-VTT dataset under categories 'sports/action', 'movie/comedy", 'vchicles/auto', 'music', and 'food/drink'. Results are averaged over 20 trials.
4
0.4
2
0.3
0
0.2
-2
3
5
5
7
0.1
7
4
0
500
1000
1500
2000
0
500
1000
1500
2000
(a) Outscore-the-hest (OtB)
(b) Optimal- (OPR)
Figure 26.
Parameter a the RBF keree function: Results are averaged over 20 trials.
3
0.40
2
0.35
1
0.30
25
25
0.25
50
50
-1
75
75
0.20
100
100
-2
0
500
1000 Iteration
1500
2000
0
500
1000 Iteration
1500
2000
(a) Outscore-the-best (OtB)
Optimal-pick-ratio (OPR)
Figure 27. Number of random features in RFF-UCB-RBE: Results are averaged over 20 trials.
33

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
0.75
0.7
0.50
0.25
0.6
0.00
1
0.5
-0.25
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
Iteration
Iteration
(a) Outscore-the-best (OtB)
(b) Optimal-pick-ratio (OPR)
Figure 28, Parameter yy the polynomial kernel function: Results are averaged over trials.
1
0.30
0.25
-1
0.5
0.20
0.5
-2
1.0
1.0
1.5
1.5
500
1000 Iteration
1500
2000
0
500
1000 Iteration
1500
2000
Outscore-the-best (OtB)
Optimal-pick-ratio (OPR)
Figure 29. Regularization parameter in KRR: Results are averaged over 20 trials.
0.02
0.75
0.00
-0.02
0.70
PAK-UCB-RBF RFF-UCB-RBF
PAK-UCB-RBF
-0.04
RFF-UCB-RBF
1000
2000 3000
4000
5000
0
1000
2000
3000
4000
5000
(a) Outscore-the-bes (OIB)
(b) Optimal-pick-ratio (OPR)
Figure Performance PAK-UCB-RUE and RFF-UCC-RBE Results are reported Sudoku-sol solving task (Setup 2). Results are
averaged over 20 trials.
34

---

PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs
0.04
0.02
0.00
-0.02
0.75
-0.04
PAK-UCB-REE
PAK-UCB-RBF
-0.06
RFF-UCB-RBF
0.70
RFF-UCB-RBF
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
(a) Outscore-the-bett (OIB)
(b)
Optimal-pick-rann (OPR)
Figure 31. Performance of PAK-UCB-RBE and REF-UCB-RBE: Results are reported on the code generation task (Setup 2). Results are averaged over 20 trials.
10
PAK-UCB-RBF RFF-UCB-RBF
Time/mins
0
500
1000 Iterations
1500
2000
Figure 32, Running time: The execution time of PAK-UCB-RBF (PAK-UCB using the RBF kernel) and REF-UCB-RBF on Setup 4. PAK-UCB-RBE takes around 10 minutes to finish 2,000 iterations of model selection, while RFF-UCB-RIE uses less than minutes. Results are averaged over 20 trials.
35


Reviews:
[{"review_id": "pZSnD3yYfK", "reviewer": "Reviewer_MgCP", "summary": "This study focuses on the task of selecting the generative model that achieves the highest reward for a given input prompt. The authors formulate this task as a contextual bandit (CB) problem, treating it as an online learning problem where past records are used to update the predictive model dynamically. They explore a UCB-based approach to solve the CB problem, specifically introducing PAK-UCB, which employs a kernel-based prediction function for each model (arm). Additionally, they present RFF-UCB, which reduces computational burden at the cost of some performance loss. The proposed methods were validated on the text-to-image generation task using models such as Stable Diffusion v1.5 and PixArt-alpha, demonstrating superior performance compared to other UCB-based approaches.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer MgCP for the thoughtful feedback on our work. Below is our answer to the reviewer's comments and questions.\n\n**1. Arm-specific reward model**\n\nWe thank the reviewer for pointing out the bandit algorithm in [1] that utilizes arm-specific prediction functions. To the best of our knowledge, the arm-specific bandit algorithms in the literature (References [1-3]) consider a linear pay-off. However, in our experiments, we observed that the PAK-UCB algorithm performs better when non-linear kernel functions (e.g., polynomial kernel and RBF kernel) are applied. We think this phenomenon is due to the normalized output of the CLIP embedding that is on the surface of the unit sphere, where a linear classification could be sub-optimal compared to a non-linear rule provided by the kernel-based approach. In the revision, we will discuss the existing arm-specific implementations of LinUCB and the necessity of including non-linear kernel functions to obtain better results in the case of text-to-image model selection.\n\n[1] Li et al. \"A contextual-bandit approach to personalized news article recommendation.\" WWW 2010.\n\n[2] Fang et al. \"Networked bandits with disjoint linear payoffs.\" KDD 2014.\n\n[3] Xu et al. \"Contextual-bandit based personalized recommendation with time-varying user interests.\" AAAI 2020.\n\n\n**2. Performance on ImageReward DB dataset**\n\nThank you for the suggestion. Due to the limited time, we report preliminary results on the ImageRewardDB ReFL dataset [1]. In the experiment, the algorithm selects among three T2I models considered in our paper: StableDiffusion v1.5, UniDiffuser, and PixArt-$\\alpha$, which attain a CLIPScore of 35.54, 34.40, and 37.20 on (a subset of) the dataset, respectively. After 5k iterations, our proposed PAK-UCB-poly3 and RFF-UCB algorithms attain an OPR (ratio of picking the best model PixArt-$\\alpha$) of 63.76% and 42.94%. On the other hand, the KernelUCB-poly3 baseline achieves only 34.18%. We will cite and discuss the dataset in the revised paper.\n\n[1] Xu et al. \"Imagereward: Learning and evaluating human preferences for text-to-image generation.\" NeurIPS 2023.\n\n\n**3. Comparison with the DiffusionGPT framework [1]**\n\nWe thank the reviewer for pointing out the alternative DiffusionGPT method in [1], which leverages an LLM agent for model selection. Our proposed PAK-UCB approach can be viewed as complementary to the DifussionGPT method, which employs the UCB bandit approach to address this task. We will cite and discuss the potential combination of the PAK-UCB and DiffusionGPT approaches as a future direction in the revised text.\n\n[1] Qin et al. \"DiffusionGPT: LLM-driven text-to-image generation system.\"\n\n\n**4. Handling unseen categories**\n\nWe agree with the reviewer that adaptation to unseen prompt categories would be challenging if the new prompts were *fully orthogonal* to previous prompts. On the other hand, in practice, it is often the case that the incoming prompt has *some correlation* with some previous prompts. For example, a user who has generated images of cats and dogs is likely to generate images of other pets or animals in the future. Assuming that the optimal model choice changes continuously with input prompts, the PAK-UCB online learning approach would be capable of predicting the optimal arm after observing previously correlated prompts. \n\nTo test this hypothesis, we conducted a numerical experiment to predict the CLIPScore of StableDiffusion v1.5. We train a poly3-based prediction model on $n=1,2,3$ categories and compute the prediction MSE on a new category after observing 10 samples in this new category. The results show that the prediction function can generalize effectively to unseen but correlated prompt categories.\n\n|Trained categories|bird|bird, horse|bird, horse, sheep|\n|-|-|-|-| \n|**New category**|horse|sheep|cow| \n|**Prediction MSE**|$0.22$|$0.19$|$0.08$|\n\n\n**5. Performance on a large number of arms**\n\nWe appreciate the reviewer's comment on setups with a large number of arms. Please note that the primary goal of the online selection task is to minimize the regret. Our regret bound in Theorems 1 and 2 are on the order of $\\widetilde{O}(\\sqrt{GT})$ for $G$ arms and time horizon $T$. In practice, the effective number of arms will be lower if the performance scores change more smoothly across arms.\n\nTo numerically test the effect of a larger number of arms, we conducted an experiment to evaluate PAK-UCB-poly3 and RFF-UCB in a setting with ten arms: We included five more arms to the synthetic T2I task (Setup 3), which generate a clean image less frequently (only $10\\%$ of the time). The results show that PAK-UCB-poly3  and RFF-UCB can outperform the best single arm and the KernelUCB baseline.\n\n|Metric (After 2k iterations)|PAK-UCB-poly3|RFF-UCB|KernelUCB-poly3|\n|-|-|-|-|\n|**O2B**|0.84|2.05|0.22|\n|**OPR**|0.13|0.16|0.11|", "claims_and_evidence": "This study argues that the performance of generative models can vary significantly across different prompt categories. This claim is supported by the analysis in Figure 1 and further reinforced by the main experimental results, which show that approaches selecting the optimal model for each prompt outperform the strategy of consistently choosing a generally high-performing model (One-arm Oracle).", "methods_and_evaluation_criteria": "In this study, the authors formulate the given task as a contextual bandit (CB) problem, which I believe is a valid problem formulation from the perspective of online model selection. Additionally, their proposed method (PAK-UCB), which defines a different kernel function for each arm, is also a reasonable approach within the CB framework.\n\nHowever, regarding this methodology, the authors emphasize the use of an \"arm-specific\" prediction function as a key distinguishing feature compared to other UCB-based approaches. However, based on my understanding, using an arm-specific prediction function is a commonly adopted approach in previous CB-related studies [1], making this claim somewhat overemphasized. Therefore, it would be beneficial to compare their method with other UCB approaches that also utilize arm-specific functions and discuss the differences between them.\n\n[1] A Contextual-Bandit Approach to Personalized News Article Recommendation", "theoretical_claims": "In this study, the authors theoretically derive the regret bound for PAK-UCB. Since they appropriately extend the proofs used in UCB-based approaches, I find this approach valid and well-founded.", "experimental_designs_or_analyses": "I believe that the overall evaluation setting, including the metrics and baselines, is well-designed. However, it is unclear which prompt set was used for the main experiments presented in the paper. Based on inferences from Appendix Section 7, Figure 10, and Figure 12, it seems that the experiments were conducted using only two categories from the MS-COCO dataset. If that is the case, it would be beneficial to validate the approach not only on such a clearly defined and limited category set but also on more general prompt sets, such as ImageRewardDB[1], for a more comprehensive evaluation.\n\n[1] ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation", "relation_to_broader_scientific_literature": "The \"generative model selection based on prompts\" emphasized in this study is a valuable research topic that could be further explored in future work. Additionally, this study has identified a meaningful domain where UCB literature can be effectively applied.", "essential_references_not_discussed": "This study explores UCB-based approaches for the task of model selection. However, there may be alternative methodologies for this task beyond UCB. Notably, as presented in [1], an Agent AI approach is also a viable option. It would be beneficial to include a discussion on alternative methodologies beyond UCB as well.\n\n[1] DiffusionGPT: LLM-Driven Text-to-Image Generation System", "other_strengths_and_weaknesses": "Strengths: Identifying the importance of generative model selection based on prompts and formulating it as a contextual bandit problem is novel and interesting. Additionally, the presentation is clear and well-structured.\n\nWeaknesses: Achieving meaningful performance requires multiple iterations, and even after several iterations, the model seems to work effectively only when the prompts fall within previously seen categories. While the appendix (Fig 12) demonstrates adaptation ability through additional training when a new category is introduced, handling unseen categories still necessitates substantial additional learning, which may limit the practical usability of the proposed approach.", "other_comments_or_suggestions": "While this study focuses on selecting a model based on the input prompt, another practical approach involves selecting an appropriate fine-tuned adapter for each prompt, as proposed in [1]. I am curious whether the proposed PAK-UCB method could also be applied to such scenarios, where adapters are chosen based on the prompt. Unlike models, adapters are generally more abundant (typically at least 10 or more), so demonstrating the effectiveness of the proposed method in cases where the contextual bandit (CB) has more than three arms could further highlight its generalizability.\n\nAdditionally, if the concerns mentioned above are addressed, I would be willing to increase the rating.\n\n[1] Stylus: Automatic Adapter Selection for Diffusion Models", "questions_for_authors": "As mentioned in the \"Experimental Designs or Analyses\" section, it would be helpful to specify the exact prompt set used in the experiments presented in the main paper.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "z2Wc2YHCUe", "reviewer": "Reviewer_3Dpj", "summary": "Generative models are increasingly being used in numerous applications. Evaluation scores are typically used when selecting a sample generation from multiple models. The drawback of evaluation scores is that different models perform better under different text prompts. The paper proposes a method to address this issue by learning the ranking of generative models for a given prompt. \n\nThe proposed method goes beyond standard LinUCB and KernelUCB. Specifically, it introduces PAK-UCB, which learns an arm-specific function to predict the score of each model. Furthermore, the paper seeks to reduce expensive computation and memory overhead by incorporating Random Fourier Features into PAK-UCB. The proposed algorithm is evaluated against several baselines for prompt-based selection of text-to-image models and wa shown to outperform all of them.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer 3Dpj for the thoughtful feedback on our work. Below is our answer to the reviewer's comments and questions.\n\n\n**1. Details of the datasets**\n\nWe note that the details of experiment settings are discussed in Appendix D. The following is a summary of the details asked by Reviewer 3Dpj, which we will include in the revised main text to improve the clarity:  \n\n* **Setup 1:** Prompts are uniformly randomly selected from the MS-COCO dataset under two categories: 'dog'/'car' (Fig.2), 'train'/'baseball-bat' (Fig.3a),  'elephant'/'fire-hydrant' (Fig.3b), and 'carrot'/'bowl' (Fig.3c).\n* **Adaptation to new models (Fig.4a):** Prompts are uniformly randomly selected from the MS-COCO dataset under categories 'train' and 'baseball-bat',\n* **Adaptation to new prompt types (Fig.4b):** In the first 1k iterations, the prompts are uniformly randomly selected from a pool that initially includes categories 'person' and 'bicycle' in the MS-COCO dataset. Then, categories 'airplane', 'bus', 'train', and 'truck' are added to the pool after each 1k iterations,\n* **Synthetic T2I and image-captioning experiments (Fig.19 and Fig.21):** The prompts are uniformly randomly selected from the MS-COCO dataset under categories 'dog', 'car', 'carrot', 'cake', and 'bowl',\n* **Synthetic T2V task (Fig.22):** The captions are uniformly randomly selected from the MSR-VTT dataset under categories 'sports/action', 'movie/comedy', 'vehicles/autos', 'music', and 'food/drink'.\n\n\n**2. Evaluation scores**\n\nIn the numerical experiments of this paper, we primarily focus on text-to-image generation tasks and use CLIPScore as the evaluation score. We note that the online selection framework can be applied to other prompt-guided generation tasks as long as we know the score values assigned to the generated samples.  \n\n\n**3. Features for kernel methods**\n\nThe input of the PAK-UCB method and other baselines is the embedded prompt that is output by the pretrained CLIP-ViT-B-32-laion2B-e16 model from the open_clip repository (https://github.com/mlfoundations/open_clip/tree/main). Only for LinUCB and KernelizedUCB baselines, we also concatenate the one-hot encoded vector of the model index to the CLIP-embedded prompt.\n\n**4. Performance of the random selection strategy**\n\nWe thank the reviewer for the suggestion of including the random selection strategy as a baseline. The random selection strategy would be expected to underperform for prompt-based model selection. Below, we provide the results in Setup 1 (Figure 3b) where we use PAK-UCB-poly3 as the competing strategy.\n\n| Metric (After 5k iterations) | Random | PAK-UCB-poly3 | \n| - | - | - |\n| **O2B** | -0.13 | **0.86** |\n| **OPR** |  0.50 | **0.77** |\n\n**5. Expressivity of score estimation functions**\n\nThe use of a polynomial kernel with degree 3 is inspired by the kernel Inception distance (KID) in the literature of generative models [1, 2]. Please note that a higher degree can lead to higher expressivity while increasing the risk of overfitting the data. Below, we conduct an ablation study to test the effect of degree on the performance of the PAK-UCB algorithm using a polynomial. We observe that a degree of 3 can achieve a better tradeoff between expressivity and generalization.\n\n|Metric (After 5k iterations)|poly1|poly2|poly3|poly 4|\n|-|-|-|-|-|\n| **O2B** | 0.13 | 0.30 | **0.70** | 0.39 |\n| **OPR** | 0.54 | 0.58 | **0.71** | 0.61 |\n\n[1] Stein et al. \"Exposing flaws of generative model evaluation metrics and their unfair treatment of diffusion models.\" NuerIPS 2023.\n\n[2] Bińkowski et al. \"Demystifying mmd gans.\" ICLR 2018.\n\n**6. Comparison to KernelUCB**\n\nIn the following, we will compare the process in PAK-UCB and KernelUCB side-by-side to highlight their differences:\n\n- **Problem Setting of PAK-UCB (ours):** We have $G$ arms where each arm represents a *fixed generative model* that remains unchanged across rounds: for example, Arm 1 represents the Stable Diffusion model in all rounds. At each round, the arms observe *one shared* context variable (i.e., the text prompt). We learn $N$ *separate kernel-based models* with *different weights* to predict the CLIPScore of a shared incoming prompt (context) for the $G$ fixed generative models.\n- **Problem Setting of KernelUCB [1]:** At every round, we have $N$ arms where the expected reward of each arm is fully characterized by its context variable. The arms have *different* context variables at each round. We learn *one shared* set of weights to predict the expected reward for the $N$ observed contexts (i.e., arms) in the next round.\n\nAs explained above, in the KernelUCB [1] setting, there is not a fixed model corresponding to one arm across iterations, and the arms will perform independently across iterations depending on their context. However, in the setting of PAK-UCB (our method), each arm will represent one fixed generative model in all the learning rounds. \n\n[1] Valko, et al. \"Finite-time analysis of kernelised contextual bandits\", UAI 2013.", "claims_and_evidence": "The paper claims that different models perform differently depending on the text prompts provided, and evaluation scores do not capture this inherent limitation. The authors present evidence of this through their experiments, which is well-known in the literature. They propose to address this issue by converting the problem of finding which model responds best to which prompt—maximizing evaluation scores—into a contextual bandit problem.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation make sense for the problem and application.", "theoretical_claims": "The paper claims that the proposed algorithm achieves a \\bigO\\sqrt{GT} regret bound. I briefly reviewed the proof for correctness.", "experimental_designs_or_analyses": "The paper conducts experiments using state-of-the-art models: UniDiffuser, Stable Diffusion, PixelArt, and DeepFloyd. Specifically, it evaluates the performance of these models using two metrics: (i) outscore-the-best (02B) and (ii) optimal-pick-ratio (OPR). The main results of the paper are sound, and the experimental setup is well-defined. However, one issue is that while the models used were provided, the paper does not mention any of the datasets that were used in the experiments.\n\nIn addition to the main results, the paper includes two additional ablation studies—one for adaptation to new prompts and the other for synthetic experiments. Similar to the main results, the authors do not clarify how the new prompts were selected or what the original datasets were that the models were trained on. Furthermore, the paper does not explain the overlap between the new prompts and the training prompts. The synthetic experiments also lack detail, making it unclear how these experiments were conducted.", "relation_to_broader_scientific_literature": "The authors are attempting to address an ongoing and challenging problem: models are sensitive to prompts, and small changes in the prompts can have significant effects on the model's decision-making. In particular, the authors are trying to improve the selection criteria for determining which model should be used for a given prompt in order to maximize performance.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "None", "other_comments_or_suggestions": "None", "questions_for_authors": "1) What dataset did you perform your experiments on?\n\n2) What evaluation score (e.g., ClipScore) did you use in your main experiments? \n\n3) What features are you using for your kernel methods?\n\n4) How does randomly selecting a model per prompt perform?\n\n5) Are non-polynomial degree 3 algorithms suffering from the model not being expressive enough?\n\n6) Could you elaborate on the difference between PAK-UCB and KernelUCB?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "None"}}, {"review_id": "U6MXOPPJzo", "reviewer": "Reviewer_5QGu", "summary": "- This paper frames the task of optimally routing prompt to a data generation model as a contextual bandit problem \n-  In doing so, the authors design a contextual bandit algorithm called PAK-UCB and prove uapper bounds on its expected regret\n-  The overcome the computational overhead of PAK-UCB, the authors propose a variant, called RFF-UCB, based on the random Fourier Features framework that approximates PAK-UCB. They prove that RFF-UCB is efficient and obtains expected regret that is not too much larger than that of PAK-UCB. \n- Finally, the authors present experimental results showcasing the performance of PAK-UCB and RFF-UCB for text-to-image and image-captioning tasks.\n\n## Update after rebuttal\nI thank the authors for their response. As they have satisfactorily addressed my questions and concerns, I will maintain my positive score for this paper.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer 5QGu for the thoughtful feedback on our work. Below is our answer to the reviewer's comments and questions.\n\n\n**1. Regret and complexity of PAK-UCB and RFF-UCB**\n\nFirst, we would like to clarify that the numerical results in Section 6 are reported for PAK-UCB (Alg.2) and RFF-UCB (Alg.3) in the main text. \n\nAlso, as we stated in Theorem 1, the regret bound is shown for the variant of PAK-UCB in the Appendix, which we titled Sup-PAK-UCB (Alg.4). We note that our analysis of PAK-UCB and Sup-PAK-UCB parallels the analysis of LinUCB [1] and KernelUCB [2] references. Specifically, [1] introduces two versions of the LinUCB algorithm: \"LinUCB\" (Alg.1 on p.g. 3 in [1]) is recommended for usage in practice, whereas [1]'s' regret analysis is performed for \"SupLinUCB\" (Alg.3 on p.g. 4). Similarly, Reference [2] also introduces two versions of KernelUCB, \"KernelUCB\" (Alg.1 on p.g. 5) and \"SupKernelUCB\" (Alg.2 on p.g. 5), where the numerical application and theoretical analysis are aimed for the algorithms. Our analysis also follows a similar approach. We will include this clarification in the revised paper.\n\n**2. Regarding Lemma 2**\n\nWe would like to clarify that Lemma 2 still holds for the variant of RFF-UCB analyzed in Theorem 2 (described in Appendix B.2). Note that Sup-PAK-UCB (Alg.4) computes the UCB values on the mutually exclusive subsets $\\\\{\\Psi\\_g^m\\\\}\\_{m \\in [M]}$, which satisfies $\\sum_{m,g}|\\Psi_g^m| \\le t$ at the $(t+1)$-th iteration. Therefore, Sup-PAK-UCB using Compute_UCB_RFF (Alg.3) requires time at most $\\Theta(\\sum_{m,g}|\\Psi_g^m| s^2) = O(ts^2)$ and space $\\Theta(\\sum_{m,g} |\\Psi_g^m| s) = O(ts)$. We will add this clarification to the revised paper.\n\n**3. Performance of RFF-UCB**\n\nAs pointed out by the reviewer, there could be a performance gap between PAK-UCB-poly3 and RFF-UCB (e.g. in Figure 2). Note that this is not in contradiction with the regret bound of Theorem 2, because PAK-UCB-poly3 and RFF-UCB use *different kernel functions* to predict the scores: RFF-UCB uses the Gaussian kernel, whereas PAK-UCB-poly3 uses a polynomial kernel with degree 3. Therefore, RFF-UCB is *not* supposed to match the result with PAK-UCB-poly3.\n\n\n**3. Discussing Remark 1 in the Related work**\n\nThank you for the suggestion. We will update the related work with a summary of Remark 1.\n\n\n**4. Typo in lines 212-213**\n\nWe thank the reviewer for pointing this out. We would like to clarify that the tilde in the notation $\\widetilde{\\Phi}$ is a typo, which we will correct in the revision. In the equation, $\\widetilde{\\Phi}$ should be replaced with $\\Phi$.\n\n\n**5. Results on regret**\n\nWe appreciate the reviewer's question about the numerical performance in terms of regret values. We note that the evaluations based on O2B and the (average) regret are equivalent. This is because the *average regret* is computed as $\\text{Avg.Regret}(T) := \\frac{1}{T} \\sum_{t=1}^T (s_\\star (y_t) - s_{g_t}(y_t))$, while *O2B* is computed as $$\\text{O2B}(T) := \\frac{1}{T}\\sum_{t=1}^T (s_{g_t}(y_t) - s_{g^\\star}(y_t)) = -\\text{Avg.Regret}(T) + C,$$ where $g^\\star$ is the best single model with the highest expected score. Therefore, the O2B scores of different policies are all shifted with the same constant $C := \\frac{1}{T}\\sum_{t=1}^T (s_\\star (y_t) - s_{g^\\star}(y_t))$ compared to their average regret. As a result, the O2B rankings of the policies in the plots are *identical* to the regret-based rankings. \n\n\n**6. Regret lower bound**\n\nFirst, we note that the regret (upper) bound derived in Theorem 1 matches the regret of the LinUCB [1] and KernelUCB [2] algorithms up to a factor of $\\sqrt{G}$, where $G$ is the number of models. On the other hand, we anticipate that the regret lower bound could scale with $\\Omega(\\sqrt{dGT})$ for a kernel function with a finite dimension $d$ (e.g., by slight modification of [Theorem 2, 1] for linear bandits without arm-specification). Formally proving a regret lower bound for the arm-specific setting is an interesting future direction for our work. We will discuss this in the revised conclusion.\n\n\n[1] Chu, et al. \"Contextual bandits with linear payoff functions.\" AISTATS 2011.\n\n[2] Valko, et al. \"Finite-time analysis of kernelised contextual bandits\", UAI 2013.", "claims_and_evidence": "Yes, the claims made in the submissions are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Yes, the proposed methods and/or evaluation criteria make sense for the problem at hand.", "theoretical_claims": "As no proofs were provided in the main text, I did not verify the correctness of the theoretical claims.", "experimental_designs_or_analyses": "Yes, I checked the experimental results in the main text and the Appendix.", "relation_to_broader_scientific_literature": "This paper empirically finds that different models perform better on different prompts. This motivates the design of routing mechanisms which can sequentially learn to route prompts to model so that prompts get answered by their \"optimal\" model. The authors frame this problem through as a contextual bandit problem, where now the context is the prompt, and the arms are the different models that the prompt can be routed to. Unlike the standard contextual bandit setting where a different context vector is variable is observed for each arm, but a single weight vector is applied to all arms, in the authors setting, the flip is true: there is a single context that is shared across all arms, but now each arm has a fixed weight vector.  As far as I can tell (I am not an expert in this area), both the framing of prompt-based selection of generative models as a variant of the CB problem and the particular variant of the CB problem are novel, and worth futher investigations.  As such, a key contribution of this paper is the introduction of a new variant of the CB problem and its applications to modern-day generative machine learning.", "essential_references_not_discussed": "From what I can tell, the authors adaquetly cite existing work and make sure to distinguish their settings with existing work.", "other_strengths_and_weaknesses": "**Strengths**: \n- This paper is well written and easy to follow\n- The problem studied is well-motivated and the novel variant of the CB problem is interesting to me\n- The experimental results for PAK-UCB seem fairly strong to me\n\n**Weaknesses**: \n- Clarity.  The regret guarantees of PAK-UCB are with respect to a variant of the PAK-UCB algorithm presented by the authors in the main text, with no discussion on what this variant is and why it is needed. Moreover, it now becomes unclear whether the xperimental results for PAk-UCB in Section 6 are with respect to Algorithm 2 or its variant in Appendix A.1. The same can be said about RFF-UCB. In fact, its unclear whether Lemma 2 still holds the variant of RFF-UCB that satisfies Theorem 2. \n- Mismatch between theory and practice. According to Theorem 1 and Theorem 2, the reader gets the sense that PAK-UCB and RFF-UCB have the same regret bound and thus should perform empirically similary. In fact, the authors state \"It can be shown that the implementation of PAK-UCB with RFF attains the exact same regret guarantees for adaptively selected feature sizes.\" However, the experiments tell a different story, with PAK-UCB consistently and significantly outperforming RFF-UCB across all experimental setups. Moreover, in many of these experiments, RFF-UCB doesn't seem to significantly outperform the baselines. It would be nice if the authors can provide some reasonable justification about why RFF-UCB doesn't do as well as PAK-UCB despite having similar regret guarantees.", "other_comments_or_suggestions": "- I think the authors should move Remark 1 into the Contextual bandits section of the Related works to futher drive home the differences between their setup and the standard CB setup\n- In the Equation in Lines 212-213, $\\tilde{\\Phi}$  and $\\tilde{\\Phi}^*$ are not defined any where. \n- The authors provide theoretical guarantees on regret, but the experiments are evaluted with respect to O2B and OPR. I would be interested in seeing plots on regret as well.", "questions_for_authors": "(1) What are lower bounds on expected regret in your variant of the CB problem? Are your algorithms optimal with respect to the relevant problem parameters?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper formulates model selection as a bandit problem, specifically designing a contextual bandit algorithm named PAK-UCB and proving upper bounds on its expected regret. Additionally, the authors introduce an approximate version of PAK-UCB tailored for practical applications. The effectiveness of the proposed algorithms is verified through applications in text-to-image and image captioning tasks.

The problem addressed is well-motivated, and the novel variant of the contextual bandit problem is intriguing. The experimental results are also robust. However, the connection between the proposed bandit formulation and the demonstrated application domains (mainly in image generation) is not entirely clear. For instance, extending the experiments to include other generative models, such as LLMs, could enhance the paper’s impact. Nonetheless, I believe that accepting this work would be a valuable contribution to the community. Therefore, I recommend a weak accept.
Decision: Accept (poster)