DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative Models
Authors: Donya Jafari
Farzan Farnia
Venue: ICLR 2026 Poster (2026)
Keywords: Generative AI Models
LLMs
Prompt Routing
Contextual Bandits
Diversity
Primary area: reinforcement learning
TL;DR: We introduce DAK-UCB, a diversity-aware contextual bandit framework for online prompt routing across LLMs and generative AI models.

Abstract: The expansion of generative AI and LLM services underscores the growing need for adaptive mechanisms to select an appropriate available model to respond to a user's prompts. Recent works have proposed offline and online learning formulations to identify the optimal generative AI model for an input prompt, based solely on maximizing prompt-based fidelity evaluation scores, e.g., CLIP-Score in text-to-image generation. However, such fidelity-based selection methods overlook the diversity of generated outputs, and hence, they can fail to address potential diversity shortcomings in the generated responses. In this paper, we introduce the *Diversity-Aware Kernelized Upper Confidence Bound (DAK-UCB)* method as a contextual bandit algorithm for the online selection of generative models with diversity considerations. The proposed DAK-UCB method incorporates both fidelity and diversity-related metrics into the selection process. We design this framework based on prompt-aware diversity score functions that decompose to a two-sample-based expectation over prompt-output pairs in the previous generation rounds. Specifically, we illustrate the application of our framework using joint kernel distance and kernel entropy measures. Our experimental results demonstrate the effectiveness of DAK-UCB in promoting diversity-aware model selection while maintaining fidelity in the generations for a sequence of prompts.

Published as a conference paper at ICLR 2026
DAK-UCB: DIVERSITY-AWARE PROMPT ROUTING FOR LLMS AND GENERATIVE MODELS
Donya Jafari
Farzan Farnia
Sharif University of Technology donya. janyaajafariii jafariilliii Feffrrrl1@@sh weeili1@sarff....m
The Chinese University of Hong Kong farni a@cse. cuhk.edu hk
ABSTRACT
The expansion of generative Al and LLM services underscores the growing need for adaptive mechanisms to select an appropriate available model to respond to a user's prompts. Recent works have proposed offline and online learning formula- tions to identify the optimal generative AI model for an input prompt, based solely on maximizing prompt-based fidelity evaluation scores, c.g., CLIP-Score in text- to-image generation. However, such fidelity-based selection methods overlook the diversity of generated outputs, and hence, they can fail to address potential diversity shorteomings in the generated responses. In this paper, introduce the Diversity-Aware Kernelized Upper Confidence Bound (DAK-UCB) method as contextual bandit algorithm for the online selection of generative models with diversity considerations. The proposed DAK-UCB method incorporates both fi- delity and diversity-related metrics into the selection process. We design this framework based on prompt-aware diversity score functions that decompose to a two-sample-bast expectation over prompt-output pairs in the previous gen- eration rounds. Specifically, we illustrate the application of our framework us- ing joint kernel distance and kernel entropy measures. Our experimental results demonstrate the effectiveness of DAK-UCB in promoting diversity-aware model selection while maintaining fidelity in the generations for a sequence of prompts.
INTRODUCTION
The past few years have witnessed rapid surge in generative Al services capable of addressing a wide array of tasks, ranging from large language models (LLMs) answering arbitrary questions to text-to-image and video models generating visual content guided by user prompts. Given the growing number of available generative models, key challenge is how to effectively select suitable generative models for a sequence of user- ssev-provided prompts. A conventional approach is to com- pute overall evaluation score for each candidate generative AI model and subsequently select the model with the highest aggregate score to address all future prompts. However, this approach implicitly assumes that single model consistently outperforms the other models across all possible prompts. This assumption has been demonstrated to be untrue in realistic scenarios where different models may excel on different topics or prompt categories (Hu et 2025a Frick et al. 2025
To address this limitation, recent literature has introduced prompt-aware model selection mecha- nisms. These methods include offline learning algorithms (Qin et al. 2024 Chen et al.) 2024 Frick
et al.) 2025), which train selector model using batch of pre-collected responses of models to prompts as training data. Also, Hu et al. (2025a) propose the online learning PAK-UCB method, formulating the selection task as contextual multi-armed bandit problem to utilize the user's ob- served model performances in the previous rounds.
Despite the development of several model selection approaches, the existing methods focus only on the fidelity scores data generation, while overlooking the diversity of generated samples. For example, in text-to-image generation tasks, existing frameworks evaluate models based on the align- ment of the input prompt and generated image, without considering diversity in the image outputs. selection lead to output samples that, although individually aligned
Such diversity-uman well with prompts, collectively lack diversity. addition, overlooking the diversity of output data potentially lead to restricted representation of sensitive attributes, such gender ethnicity,

---

Published as a conference paper at ICLR 2026
Candidate Generative Models
Text to Image Generative Model1
Text to Image Generative Model
Conditioned on "Young Male" (Lower Diversity)
Unconditioned generation (Higher Diversity)
An individual jogging in the park.
An Individual dancing the rain
Individual jogging in the park.
An individual dancing the rain.
individual riding
individual painting portrait
individual bicycle.
individuall painting portrait,
Online Selection Strategy
(Diverstty-nnwware) Kernelized UCB
Diversity-Aware Kernelized UCB (Ours)
selection ratio
selection ratio
51.2%
48.8%
21.8%
78.2%
Generative Model
Generative Model 2
Generative Model
Generative Model
Figure Comparison baseline Kernelized-UCH model selection (CLIP-Score fidelity metric) (Hu et al. 2025a) our proposed diversity-awarr DAK-UCB over T = 500 rounds. While the baseline Kernelized- UCB does not favor model G2 with higher diversity over model G₁, DAK-UCB selected the more diverse G2 more frequently.
in generated datasets. Figure I displays an example, where the diversity- baseline kernel- ized UCB selection algorithm in (Hu et al., (025a)) chooses the less and more diverse generative models (conditioned to "young male" and none) with similar frequencies, ignoring the diversity factor in the selection process. This limitation arises because PAK-UCB, and standard contextual bandit methods more broadly, compute rewards using the mean of sample-level scores. Diversity, however, is group-level property determined by the relative positioning of multiple samples, which cannot be expressed through simple average of individual rewards.
These limitations highlight the importance of diversity-aware selection methods, which explicitly incorporate considerations of output diversity into the model selection process. In this work, we specifically focus the online selection task, proposing an algorithm designed to leverage previ- ously generated data to select generative models that achieve an optimal balance between fidelity and diversity. Our proposed method, which we call Diversiry-Awar Kernelized Upper Confidence Bound (DAK-UCH extends the kernelized UCB framework (Valko et al. 2013 Hu et al. (025a)) by integrating diversity-orient term, the form of the expectation of two-sample (prompt, random variable, into the contextual bandit objective.
A key challenge designing DAK-UCB is to determine which diversity scores are compatible with the contextual bandit selection framework. Specifically, we identify family of joint kernel scores ioass-iiniiinngg prompt-conditiona extensions of kernel distance (Binkowski et al. 2018). RKE (Jlali et al., 2023), and MMD (Gretton et al. 2012) that be expressed as expectations of two- sample quadratic forms over prompts and outputs. This structure central to DAK-UCB, it enables fast-converging estimation from streaming data via kernel ridge regression, yielding principled con-
2

---

Published as a conference paper at ICLR 2026
fidence bounds in the UCB process. Moreover, by combining joint kernel scores with task-specific fidelity metrics (c.g.. CLIPScore in text-to-image generation), DAK-UCB can be extended to pro- vide unified approach to prompt-adaptive selection that balances the fidelity and diversity factors.
Figure | shows an application DAK-UCB for a diversity waree generative model selection in response to T 500 prompts of MS-COCO dataset (Lin et al. 2014) on generating numan-related scenes. The candidate model G represents the Stable Diffusion XL (SD-XL) (Stability-A] (023) model conditioned to "young male individual", whereas candidate model G2 outputs the SD-XL outputs without any conditioning. While the baseline kernelized UCB fidelity-based selection did not favor the more diverse model G2, generating samples from both models with equal probabilities, the DAK- model selection with the joint RKE diversity score chose model G1 more often over the 500 online selection iterations.
Beyond deterministic prompt-to-model assignment at every iteration of DAK-UCB, DAK-UCB can be adapted to assign the model an input prompt based on non-degenerate mixture of the models. selection
As noted by Rezaei et al. (2025) in the unconditional setting, the optimal diversity strategy can itself be non-degenerate mixture of models. Extending this insight to the conditional, prompt-aware setting, mixture-based selector effectively rolls biased vn-sided die to determine which of the m models is queried for a given prompt. We introduce the Mixture-DAk UCB method to realize this idea: online algorithm that optimizes prompt-dependent mixture probabilities. Kacture-DAK-UGG
generalizes the prompt-free Mixture-UCB framework of |Rezaei et al. (2025) to the conditional case, enabling diversity-enhancinn mixtures tailored to incoming prompts and yielding further improvements in diversity metrics.
We empirically evaluate different variants of the proposed DAK-UCB and gorithms on text-to-image and language model generation tasks. Our results demonstrate improve- ments in diversity and overall correctness metries relative existing contextual bandit algorithms, such as Kernelized UCB. PAK-UCB, and randomized selection strategies. We also validated the defined Joint-RKE and Joint-KD measures for capturing diversity and distributional matching char- acteristics of pooppt-gudded generative models. Here we summarize the work's 's main contributions:
al-
Studying the role of diversity in prompt-aware selection of generative AI models, Introducing the Diversity-Aware Kernelized UCB (DAK-UCB) algorithm, contextual bandit approach explicitly accounting for the diversity factor in model selection,
Extending deterministic DAK-UCB selection to prompt-conditioed mixture selection,
Demonstrating numerical effectiveness of DAK-UCB several text-to-image generation tasks.
2 RELATED WORKS
Contextual bandits (CB) extend the multi-armed bandit (MAB) framework by incorporating the context variable to guide the selection process (Langford & Zhang 2007 Foster et al. 2018). A widely-studied CB is the linear CB, which assumes that the expected reward of each arm linear function of context (Li al., 2210 Chu et al.) 2011). Kernelized CBs gener- alize to non-linear reward models by using kernel methods to capture more complex dependencies
Contextual Bandits.
between contexts and rewards (Valko et al. | (0013). Due to the computational cost of kernel methods, recent works have explored approximations using relevant assumptions on the kernel (Calandriello
et al.
2019 2020 Zenati et al. 2022).
To address exploration in linear CBs more effectively. (Abbasi- Yadkori et al., 2011) propose tighter confidence sets using martingale inequalities, leading to stronger theoretical guarantees and im- proved empirical performance. Moving beyond linearity, (uu et al.) (025b) introduce PromptWise, a multi-iteration-per--ooo
cost-aware contextual bandit for prompt routing in I.I.Ms and generative models. Also, Kveton et al. (2020) propose two randomized exploration algorithms for generalized linear bandits, which leverage Laplace approximations and perturbations of past data to efficiently explore under non-linear models. However, the above CB methodologies do not target diversity awareness in the online learning setting.
Diversity/Novelty Evaluation Scores and Guidance in Generative Models. Several methods have been proposed for evaluating and improving the diversity of generative and diffusion models. On the diversity evaluation, the metrics Recall (Saijadi et al. 2018 Kynkäänniemi et al. 2019), Coverage
3

---

Published as a conference paper at ICLR 2026
Nacem al. 2020). Vendi (Dan Friedman & Dieng 2023 Ospanov et al. 2024 Ospanov & Farnia 2025). and RKE (Jalali et al. 2023) have been proposed for unconditional (prompt-free) sample generation, and Conditional Vendi/RKE (Jalali et al. 2026 2025a) and Scendi (Ospanov et al. 2025) have been suggested for prompt-aware diversity measurement. We note that (Zhang et al,,
2024 2025) propose entropy-based measures for novelty of generative models and their comparison, and (Jalali et al. 2025b Gong et al. 2025) study kernel-based comparison of embeddings.
For guiding sample generation, (Miao et al. 2024) employed reinforcement learning with a diver- sity reward function in the generation process. (Sehwag et al. (2022) proposed sampling from low- density regions of the data manifold to encourage diverse outputs. (Corso et al. (202] introduced a particle-based potential function that explicitly maximizes pairwise dissimilarity. [Sadat et al. (2024) explored the addition of Gaussian noise to conditioning inputs during inference to promote variabil-
ity. Lu et al. (2024) developed ProCreate, a distance- based guidance technique. [Askari Hemmat al. (2024): Jalali et al. (2025a) proposed Vendi/Conditional-RK. Score Guidance, which incor- porates diversity score guidance in diffusion models. Similarly, [Sani et al. (2026) propose MMD guidance to align the diffusion model to a target distribution by minimizing the MMD distance. We highlight that these works aim improve the diversity and alignment over the sample generation process, unlike our work on the diversity- aware online selection of pre-trained models.
Multi-Armee Bandit for diversity-based selection. In related work, Rezaei et al. (2025) propose Mixture- UCB, a bandit algorithm for selecting mixtures of generative models to maximize diversity, while their proposed approach is not prompt-aware and therefore not applicable to prompt-guided sample generation. (Chen et al. 22025). (Yang et al. 2024), and (Hou et al. (2024) improve the best arm identification by multi-objective optimization, regret minimization, and sample efficiency. (Sani et al. (2012) introduce a framework for risk-averse pacision-making in bandit problems by inte- grating variance- senssiivee utility functions into exploration strategies. Weinberger & Yemini| (2023) study bandits with self- information- based rewards, proposing algorithms that leverage information- theoretic concepts to balance exploration and exploitation. //hu & Tan (2020) develop Thompson Sampling algorithms for veannnariance bandits, optimizing both expected returns and reward vari- ability. We note that our work focuses on diversity in a contextual bandit setting, where the prompt plays the role of the context, which is not the case in the context-free MAB setting of these works.
3 PRELIMINARIES
3.1 NOTATIONS AND DEFINITIONS
Throughout the paper, we define conditional generative model & a conditional distribution Pg(x)() where E X' is the generated data variable conditioned to the randomly-observed prompt t e T. Following this definition, every sample generation of model is conditioned on a user's provided prompt T f and then drawing sample from the conditioned distribution Pa(x[T' t).
3.2 BERNEL-BASER SCORES FOR GENERATIVE MODELS
In a sample space X, we call k X X + R a kemel function if there exists a feature map Φb : X' -> He such that for every 2,, X' we have k(x, x²) (6(2)), d(x²)) where (;-) denotes the inner product in the Hilbert space H4 of kernel function kt Examples kernel functions include the degree-r polynomial kernel kpaly(r) (x. y) (1+ xxx,)))))))))))))))))))>)))))+ with parameter y > and the RBF (Gaussian) kernel with parameter a defined as:
kigusssion( (o) (I., y) = cxp ( || |m-vvl? y/1
20
Given a kernel function k. we can define the n X n kernel matrix K = [k(2,, x. Missssn for n samples Zi,..., In X. Note that every valid kernel function will result in positive semi-definite (PSD) kernel matrix for every set of samples. our analysis, we use the following kernel-based scores and their variants in the online selection process:
Maximum Mean Diserepancy (MMD) and Kernel Distance (KD): For two probability dis- tributions P. Q sample space X. (Biñkowski et al.; 2018) consider the kernel distance (KD) between P and Q as the square of the maximum mean discrepancy (MMD) Gretton et al. (012)) i.e.,
KD(P,QM Q) = E
R [k(x, (,t,, x')] + u. B. [k(,,'')] - 2. Pxx [k(2,)))
(1)
4

---

Published as a conference paper at ICLR 2026
In the above definition, the samples x, Px.y.v/ Qx are drawn independently according to the specified distributions.
Rényi Kernel Entropy (RKE): For probability model P on space X. the Rényi kernel entropy (RKE) (Jalali et al. 2023) is defined as the order-2 Rényi entropy of the normalized population
kernel matrix, which reduces to
RKE(Px) - 1/18 x. Ik, x 2²)²
Considering the empirical samples Z P. the empirical RKE score reduces to RKE(E)...., =I -| K|
DIVERSITY-AWARE KERNELIZED UPPER-CONFIDENCCE BOUND
To develop a diversity-awe aware online selection of conditional generative models, we first propose sasple-basedd extensions of the KD and RKE scores to the conditional sample generation case. Subsequently, we extend the standard Kemnliized-U online learning framework by including an upper confidence bound of the joint proposed score functions.
4.1 EXTENSION OF KD AND RKE SCORES TO CONDITIONAL GENERATIVE MODELS
We propose the following extensions of the KD in equation 1 and RKE in equation 2 the con- ditional sample generation task. Both the extensions in the following apply the original scores to the joint (prompt t,data x)) variable, by using the product kernel function kyoint (It, x). [t' x²))
kiext (t, (') kdatt (z, x'). As demonstrated by Bamberger et al. (0222) ) Wu et al. (0225), the product kernel function corresponds to the Hilbert space of the tensor product of the (embedded) prompt and data vectors, effectively capturing the clusters in the dataset of the joint prompt,data vectors.
Joint Kernel Distance (JKD) distribution matching score. We propose the following extension of the marginal (prompt-unaware) kernel distance in cquation[Tto the prompt-aware kernel distance, which we call Joint Kernel Distance (JKD), for two conditional distributions Pxir and QXIT:
JKD(PxIT, Qx(r) = KD((r Pxpr, Pr Qx(r)
(3)
= E Pr. wny wy Pxir=e Px Qx11 (T(((()
(x((xx))) + kx 9') kx(x,))) -x kx(x'x0) wi)]
x X' R denote the kernel functions for the input prompt t and output x, and Pr is reference distribution on the input variable T (i.e., prompt) over space T. Importantly, the empirical estimation of the expectation in equation 3an be performed by accessing only one sample generated by Pxjr for each input prompt t.
where ky T T R and hx x
Joint RKE (JRKE) diversity score. Similarly, we propose the following definition for the joint (prompt-dar RKE score, which we call Joint-RKE (JRKE) score. JRKE is defined to be the RKE score of the joint sample (T, ~ Pr Pxtr given a reference prompt distribution Pr:
JRKE(PxT)) = RREEEE PXIT)
(4)
-1 -1/B B. Pre Pxirer Pxyt- [ko (t, t') kx x')]
This score varies monotonically with its inverse, i.c, Inverse-JRKE score denoted by I-JRKE:
I-JRKE(Pxw) t, Pr 3 Te x KT (t, (xx,,
(5)
Similar to the JKD score, the expectation the diversity-basee Inverse-JRKE: score can be estimated using a single output X ~ Pxit=: for every prompt 7 f.
4.2 DIVERSITY-AWARE ONLINE LEARNING VIA DAK-UCB
To propose a diversity- aware online selection framework, we leverage our proposed conditional diversity scores in Equations 3 and within the contextual bandit framework. The prompt & serves
5

---

Published as a conference paper at ICLR 2026
as the context, and we seek policy II T - [G] that balances fidelity and diversity objectives. A key feature of the introduced diversity scores is that they both decompose into expectations of prompt-level functions, enabling online estimation with single sample per prompt. The following proposition highlights this property of the JKD and Inverse-JRKE scores.
Proposition 1. For conditional distributions Pxyr, Qxir and reference distribution Pri
(a) The Inverse- JRKKE admits the decomposition:
I-JRKE(PxI)) = R Pr a~ (1,
(6)
where (I.ske((tx = Ee [KT
2 xx,,
(b) The JKD for comparing model against reference Q admits:
JKD Q) Eg Pr. and (11)) ( (f, where (9) (t, x)) By (L) [xxx (2r, - (ox) v)D)
Proposition ! highlightss a crucial structural property of the proposed diversity scores: both I-JRKE and JKD admit two-sample expectation form, in which the overall metric decomposes into the expectation a prompt-level function of a single generated sample. This is important in the online setting, because ensures that each round of interaction with model provides an unbiased stochas- tic label for the corresponding diversity function, even though the original metric is defined in terms of expectations over pairs of prompts and outputs. Therefore, the two-sample form makes these scores applicable to the kernelized UCB algorithm, as we can run kernel ridge regression (KRR) on the stochastic labels and obtain confidence bounds that are comparable those for the fidelity score.
Based on this decomposition, we define for each model g and prompt t prompt-level target functions: xr; (1) := Pg [onaetiyy (1, D1(() == Eg~Pa((lo) [6 x1; Hc)]
(8)
Here Qridelit (1, x) denotes fidelity score of prompt-output pair, instantiated in our experiments
as the CLIP-Score between text prompt and generated image The function d 21 H4, is per-sample diversity score, whose expectation recovers the desired diversity metric in Proposition]] The history H, is only used to instantiate reference expectations over past outputs.
At each round f, DAK-UCB treats the prompt pt as context in the per-arm kernelized contextual bandit process (Hu et al. 2025a). and compares arms via per-arm UCB on the combined objective (t) = (t) + xD,(tt)
where (t) (e.g. CLIP-Score in our experiments) the fidelity score and (1) is defined with d2g instantiated either the (negative) I-JRKE score or the (negative) JKD
score as in Proposition After observing a single sample x; ~ Pg. (- 4)), we form unbiased labels (x)
yj Φna (f,, 2T)) and y; (D) vg (t,, z4; H.). update per-arm KRR models for sg and Dg. and select the next arm using an optimistic estimate
(t,) (t, + (()))(()))) (t.)) + A(D,(tt) (DD)
i.., an upper bound for sy and a lower bound for D, (since Dg is a signed diversity reward, equal to the negative of the underlying penalty). Confidence radii 3(s), g(D) follow the standard KRR-UCB form as detailed in Algorithm
we establish a regret bound for a phased variant of our algorithm. Sup-DAK UCB. This result shows that the known regret guarantees of kernelized UCB methods (Chu et al. 2011 Valko et al. 2013 Hu et al. 2025a) can be systematically extended to our diversity-aware objec- tive. key technical component of this analysis that the JRKE and JKD metries admit the two- sample expectation structure, thereby enabling integration with kernelized- UCB confidence bounds. This structural property is specific to JRKE and JKD and allows us to obtain regret guarantees for diversity-aware model selection. The following provides informal statement of the resulting gret bound, and the proof is deferred to Appendix B
In Appendix B
Theorem 1 (Informal regret bound for DAK-UCB, Under Assumptions 3 in Appendix B (nor- malized kernels, sub- Gaussian noise, and RKHS regularity for sy and Dz). the phased variant SUP-
DAK-UCB algorithm satisfies the following regret bound where the information-gan -ain r(.)) and effective-dimension D terms are defined in Appendix
Regret(T) = Ö(VGTT + GTG

---

Published as a conference paper at ICLR 2026
Algorithm/ 1: Diversity-Aware Kernelized UCB (DAK-UCB)
Input: generative models, horizon T, prompt distribution P. trade-off X, diversity score v E {-{ {-I-JRKE, -JKD}
Output: 7 generated outputs
Initialize per-arm KRR estimators for fidelity 8g and diversity Dg
2 Sample prompt 4 ~ P
for g to G do
Predict fidelity and diversity with KRR: (3) og (t)).. (D, (ti)::
Form UCB score: Jy
(t+) JUCB((t) + (Gg(t)) B6 ( (t.) (t:) + (t))
Select model 9i < arg maxy Jg (t,)
Generate output 2; ~ p. t2
Form labels y!"'! = Φna (na vz). Vs (D) = th (ti, T3; H4)
Update KRR models of 9: with (ti, yi ( and (ti, y/ (D) Update history Hi+1 +- Hi U {14,, (1,
4.3 PROMPT-AWARE MIXTURE SELECTION VIA QUADRATIC OPTIMIZATION
While DAK-UCB selects a single model per prompt, maximizing diversity can require prompt- dependent mixtures of the available models, where we denote the model mixture probability values of prompt & with notation axtt) Åc. Therefore, for G conditional generation mod- els in {P)(-)t)) we consider a prompt-aware mixture at E Ac. yielding Pa(-It) = Eg Og ( It). We focus here the I-JRKE diversity penalty; the analogous construction A Using the product kernel, the I-JRKE admits
for JKD is deferred to Proposition 2 in Appendix
the quadratic form
I-JRKE(Pa)) = [a(t) (M())))))))))MMM(((((((())
where (t) E [0, 1jxx collects cross-kernel expectations across models. To ensure stability across prompts, restrict mixtures to a kernel-Lipschie competitor set
A. = (a: T AG Vt, [t, ||a(t) - a(t)), sc). sc)
which guarantees that nearby prompts yield similar mixtures and incurs only an (cc) approximation error. In the Appendix [ we discuss that, under the above mixture feasible set, approximate solution follows solving the following problem where at each prompt I, the decision rule reduces to the concave quadratic maximization
a argmax (a, (..JUG SUCB xx MUB
aεAc
where 3UCB(1) are fidelity UCB estimates and MUCB(()
(1) the projection of the kernelized UCB estimation of (t) onto the PSD matrices by zeroing its negative eigenvalues. We call the resulting mixture-model selection method Mixture-DAK-UCB, as detailed Algorithm 2] at Appendix A
NUMERICAL RESULTS
We numerically evaluated the proposed DAK-UCB and its mixture variant, Mixture-DAK-UCB in several experiments. In our numerical experiments on text and image data, we used the CLIP encoder (Radford et al., 202T) as the backbone text embedding and DINOv2 (Oquab et al. (023) as the image embedding as suggested by Stein et al. (2023). We considered the following online model selection baselines our evaluation of UCK- and Mixture-DA--UCB:
The one-arm oracle baseline has knowledge of the evaluation scores of each individual generative model (aggregated over the validation prompt set). This baseline universally selects the individual model with the best aggregate score to handle all the prompts.
One Arm Oracle:
Random Selection: This baseline randomly selects an arm for an input prompt, where each arm is selected uniformly with equal probability, and the selection across prompts are independently.
7

---

Published as a conference paper at ICLR 2026
180
6.5
DAGUIE
32.2
Joint:KKE Score
160
KD Score (x<30-1)
6.0
Random
32.0
One Arm Gradle
5.5
PAK-UCB
DAGUCS
5.0
31.6
DACUCS
4.5-
31.4
Mixture DAKUCB
Mixture DAKUCB
One Am Cracle
4.0
One Am Cracle
100
500
1000
1500
500
1000
1500
31.0
500
1000
1500
Number Iteration
Number iteration
Number Iteration
Figure 2: Performance comparison JKD score and Joint-RKE for MS-COCO prompt clusters using Kandinsky, SDXL, and GigaGAN.
Candidate simulated text to image generative models
arm I'ast" condibonec)
amm 2 'ddog' canditione))
am 3 (uncondtioned]
Comparison:
Pilurs DMALO inlect
Nurbbar beracin
Cond-Vend 14.24 Joint-RKE = 21.07
Cond-Veedd 17.35
Lond-Vendii22 : 20
Cuanergenct:
Joint-RKE 36.25
Joint-RKE 51.01
Prompts:
DAK-UCB (Ours):
PAK JCB:
ein animal street animal in the forest
Selecctoon kasts
Sellerton Rath
6.8%
25.0
enn onimal in the park
505
200
animet the river
Cond-Vendi 23.87 Jaint RKE = 53.85
Cond-Vendi 20.13
an smirell the animal the beach
Joint-RKE = 39.31
Figure 3: Visualization of simulated generative models with less-diverse Models 1,2 and more- diverse Model 3. DAK-UCB and PAK-UCB selection ratios,scores over 500 rounds are reported.
PAK-UCB:
This baseline is a iiversity-unawarn contextual bandit algorithm (embedded prompt is the context variable), selecting the model only based the CLIP-score fidelity score in text-to- image generation.
DAK-UCB applied to diversity-aware text-to-image model selection on MS-COCO prompts. We considered the prompts in the MS-COCO (Lin al. 2014) validation subset. We uniformly sampled a thousand prompts containing the words: cat, dog, car, cake, bowl, bike, tree, airplane, park, and elephant. Three generative models were used candidate text-to-imaae generation mod- els in the experiment: Kandinsky (Arkhipkin et al. (024)), SDXL (Stability-Al] (203)), and Giga- GAN Kang al. 2023). The experiment ran for 2000 iterations, where at each iteration a random prompt from a random cluster was chosen, and our objective selected the best arm that balanced both diversity and fidelity. The results are averaged over 10 trials reduce noise from random prompt selection. Figure 2 shows that Mxiture-DAKe DAK-UCB could achieve the highest diversity Joint-RKE score. We also used MS-COCO test samples as the reference dataset and report the KD scores with Mixture-DAK-UBB obtaining the best score. Note that the KD metric evaluates both diversity and quality factors.
Experiment on simulated text-to-image models with varying diversity in "animal" image gen- eration. In this experiment, we simulated three animal image generation arms, where the first two arms (less-diverse)
outputting the SD-XL generated data conditioned on "cat" and "dog" sam- ples, respectively. On the other hand, the third model (more-diverse) generates the picture of an animal uniformly selected from list of 10 animals. To run the experiment, we used GPT-40 to
8

---

Published as a conference paper at ICLR 2026
Ing
Expert Selection Ratio
Cat promptss (Cluster 0)
Expert Selection Ratio
Dog promets (Cluster 1)
1.00
X
X
X
0.75
0.4
DeGaca sin N6-C3
Wodi 1 coet
X
0.50
0.25
Motur
0.2
wode
X
X
X
0.00
500
1000
Oit Ohodle
1500
0.0
1000
1500
X
Ind pace iage
Ageiling
Iteration
Iteration
Expert Expert Selection Ratio
0.6
Car prompts (Cluster 2)
Expert Selection Ratio
Cake prompts (Cluster
0.6
CLIP Score
DeGucc Molax
Hizaie DICGER Orada
0.4
0.2
0-2
Ded
0.0
0.0
500
1000
1500
0
500
1000
1500
500
1000 Number Iteration
Iteration
Iteration
Figure 4: Expert Selection Ratio and Performance Comparison between DAK-UCB and baselines
using the JKD score for diversity term in DAK-UCB.
generate 200 independent prompts about "an animal" in different scenes, with sample prompts pro- vided in Figure Figure3]shows shows the conditional- -ennii and Joint-RKE scores for each of the three simulated arms, which indicate that the "cat" and "dog" simulated arms were less diverse than the third simulated arm. We ran the DAK-UCB algorithm for 500 iterations, where at each step a ran- dom prompt was chosen and the output from the algorithm's 's selected arm was observed. The results demonstrate that the DAK-UCB tended to generate images from the more diverse third arm, while the TL.I--Score-basee more frequently.
baseline generated samples from the less-diverse second model
Identification of prompt-relevant diversity via DAK-UCB. In the experiment, we tested DAK- UCB in outputting samples with prompt-relevant diversity. We used these four prompt groups of the experiment of Figure [ from MS-COCO validation set: "Cat", "Dog", "Car", and "Cake". We designed four arms, each acting an expert on one of these clusters, where the expert arm on each subject generates aligned samples with the prompt of the same type, while it generates images of a mandomly-selected incorrect type for the remaining three subjects For example, ARtt expert) generates an image using SDXL when given a prompt from the Cat cluster; otherwise, it samples a random irrelevant prompt from other types and generates in Figure we report the expert arm selection ratio for each prompt category and the average CLIP- Scores iterations for each baseline. As suggested by the results in Figure 4 both the JKD-based and Clip-Score++-JJRR -eereee---------------- based DAK-UCB (Appendix. Figure 15) methods could avoid generating prompt-irrelevant output and did not attempt to increase diversity by generating unrelated content.
image for that prompt using SDXL. In
Despite their strong capabilities, all three models suffered from diversity collapse, but crucially, their collapse modes were complementary rather than identical. This directly motivates mixture- based selection: although each arm exhibits low diversity on its own, their differing failure modes allow a mixture to achieve significantly higher output diversity. Using the Cond- Vendi metric, we show that the mixture selected by MixtureDDA-UCC attains substantially higher diversity than any single model. We repeat this experiment for two additional prompts- a vibrant city in Europe" and renowned aclebrity"-an the visualizations in Figure F consistently demonstrate the diversity gains unique to model mixtures.
also report the numerical results of applying DAK-UCB for iiversity-aware model selection in the tasks of prompt-awar selection of simulated LLMs with different diversity scores and the image-to-model assignments of image captioning mod- els. We also present the results of the ablation study for testing the effect of the choice of image embedding and the coefficient of DAK DAK-UCB objective's diversity term.
Additional Numerical Results. In Appendix

---

Published as a conference paper at ICLR 2026
"rompt="Generate "Generate short sentence about vibrant city in Northern America."
Uama322
68% '2%
Vendi
2.39
Prompt: "Generate short sentence about vibrant city in Europe."
11ff/111
Prompt: "Generate short sentence about renowned celebrity."
Vendi
Figure 5: Vendi diversity comparison across LJama3.2, Qwen2, Gemma3, and their mixture for prompts: "Generate short sentence about vibrant city in Northern America." "Generate short sentence about vibrant city in Europe." and "Generate a short sentence about a renowned celebrity."
6 CONCLUSION
In this work, we proposed an online learning framework for a diversity-aware prompt-based se- lection of multiple generative models. Our proposed DAK-UCB can be applied using the defined I-JRKE diversity and JKD correctness scores for improving the diversity factor in sample genera- tion. Notably, the proposed scores reduce the two-sample expectation over the observed samples, which can be estimated using only one generated sample for an input prompt. In addition, we intro- duced the Mixture-DAK-U DAK-UCK extension, which enables optimized prompt-dependent mixtures of generative models and further improves diversity-aware selection.
Beyond text-to-image generation, we also demonstrated applications DAK-UCB multi-LLM prompt assignment and image captioning models in Appendix C, illustrating its applicability in broader generative settings. Extending the framework to additional modalities, such as protein, molecular, and graph generative models, is relevant future direction. The extension of the proposed scores for evaluating and guiding prompt-aware diversity and correctness in data generation will be relevant for future exploration. Also, studying the application of the scores in general bandit settings beyond generative model selection problems is another related future direction.
ACKNOWLEDGMENTS
The work of Farzan Farnia is partially supported by a grant from the Research Grants Council the Hong Kong Special Administrative Region, China, Project 14210725, and is partially supported by CUHK Direct Research Grants with CUHK Project No. 4055164 and 4937054. The work is also supported by a grant under 1+1+1 OOHK-CUHK(SZ-GDSTT Joint Collaboration Fund. Finally, the authors thank the anonymous reviewers and metareviewer for their constructive feedback and suggestions.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We have taken several steps for the reproducibility of our work. The proposed DAK-UCB and Mixture-DAK-UCh algorithms are fully specified in Section 4 and Appendix A, with pseudocode provided in Algorithms 1-3. Theoretical results are stated with assumptions and proofs in Appendix B. The datasets used in our experiments include standard publicly available benchmarks including MS-COCO as well as synthetically generated prompt sets by the specified GPT-4o model. The de- tails of data selection, prompt construction, model candidates, and evaluation metries are described in Section 5 and Appendix C. We provide ablation studies in Appendix C.4 to clarify the sensitivity of our results to hyperparameters and embedding choices. An anonymous implementation of our method will be released as supplementary material.
REFERENCES
Yasin Abbasi- Abbasi-Yadkori, Dávid Pál, and Csaba Szepesvári. Improved algorithms for linear stochastic bandits. Advances in Neural Information Processing Systems, volume 24, pp. 2312-2320 2011.
Meta AI. Llama 3.2 series (Ib, 3b, 8b, 70b). 2024.
ama
Vladimir Arkhipkin, Viacheslav Vasilev, Andrei Filatov, Igor Pavlov, Julia Agafonova, Nikolai Gerasimenko, Anna Averchenkova, Evelina Mironova, Anton Bukashkin, Konstantin Kulikov, et Kandinsky 3: Text-o-immae synthesis for multifunctional generative framework. In Pro- ceedings of the 2024 Conference Empirical Methods in Natural Language Processing: System Demonstrations, pp. 475-445555 2024.
Reyhane Askari Hemmat, Melissa Hall, Alicia Sun, Candace Ross, Michal Drozdzal, and Adriana -Sorian.. Improving geo-diversity of generated images with contextualized vendi score guidance. In European Conference on Computer Vision, pp. 213-229. Springer, 2024.
Stefan Bamberger, Felix Krahmer, and Rachel Ward. Johnson-lindenstinn necker structure. SIAM Journal Matrix Analysis and Applications, 43(4):1806- 1850, 2022.
embeddings with kro-
Mikotaj Binkowski, Danica Sutherland, Michael Arbel, and Arthur Gretton. Demystifying MMD) GANs. In International Conference on Learning Representations, 2018.
Daniele Calandriello, Luigi Carratino, Alessandro Lazaric, Michal Valko, and Lorenzo Rosasco. Gaussian process optimization with adaptive sketching: Scalable and no regret. In Alina Beygelz- imer and Daniel Hsu (eds.), Proceedings of the Thirty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pp. 533-557. PMLR, 25-28 Jun 2019.
Daniele Calandriello, Luigi Carratino, Alessandro Lazaric, Michal Valko, and Lorenzo Rosasco. Near-linear time Gaussian process optimization with adaptive batching and resparsification. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Ma- chine Learning. volume 119 of Proceedings of Machine Learning Research, pp. 1295-1305. PMLR, 13-18 Jul 2020.
Lingjiao Chen, Matei Zaharia, and James Zou. FrugalGPT: How to use large language models while reducing cost and improving performance. Transactions on Machine Learning Reseurch, 2024. ISSN 2835-8856. URI httth::ts:///gtg /ttp::/:///eeeeviieeeeee ew. net / forum? fideiid=s53 mkw.5
Zhirui Chen, P. N. Karthik, Yeow Meng Chee, and Vincent Y. F. Tan. Optimal multi-objectio best arm identification with fixed confidence. arXiv preprint 1xXiv:2501.13600 2025, doi: 10.48550/ arXiv.2501. 13607. Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics (AISTATS).
Wei Chu, Lihong Li, Lev Reyzin, and Robert Schapire. Contextual bandits with linear payoff func- tions. Geoffrey Gordon, David Dunson, and Miroslav Dudík (eds.), Proceedings of the Four- teenth International Conference Artificial Intelligence and Statistics, volume 15 of Proceed- ings of Machine Learning Research, pp. Fort Lauderdale, FL, USA, 1 11-13 Apr 2011. PMLR.
11

---

Published as a conference paper at ICLR 2026
Gabriele Corso, Yilun Xu, Valentin De Bortoli, Regina Barzilay, and Tommi S. Jaakkola. Parti- cle Guidance: non-I.I.D diverse sampling with diffusion models. In The Twelfth International Conference on Learning Representations, 2024.
Dan Dan Friedman and Adji Bousso Dieng. The vendi score: A diversity evaluation metric for machine learning. Transactions machine learning research, 2023.
Google DeepMind. Gemma Open weights family (Ib, 4b, 12b, 27b).
https: //
mwnggingface.co/gooooo e e gemma 3-10-it it 2024.
DeepSeek-AL.
bttp:://uuginngfacc..o//// comtooothh
DeepSeek-Rl-Distil -wee---wwen--- 5GUFGUF) 2024.
Jianfeng Dong, Xirong Li, Chaoyou Xu, Xun Yang, Gang Yang, Xun Wang, Yuan Meng, Qiang Li, Feng Zheng, Xiaogang Zhang, al. Instructblip: Towards general-purpose vision-language models with instruction tuning, 2023.
Dylan Foster, Alekh Agarwal, Miroslav Dudik, Haipeng Luo, and Robert Schapire. Practical con- textual bandits with regression oracles. In Jennifer Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1539-1548. PMLR, 10-15 Jul 2018.
Evan Frick, Connor Chen, Joseph Tennyson, Tianle Li, Wei-Lin Chiang, Anastasios N Angelopou- los, and Ion Stoica. Prompt-to-lcaderbbard. arXiv preprint arXiv:2502. 14855, 2025.
Shizhan Gong, Yankai Jiang, Qi Dou, and Farzan Farnia. Kernel-based unsupervised embedding alignment for enhanced visual representation in vision-language models. In International Con- ference on Machine Learning, pp. 19912-19931. PMLR, 2025.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
Yunlong Hou, Vincent Y. F. Tan, and Zixin Zhong. Almost minimax optimal best arm identification in piecewise stationary linear bandits. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
Xiaoyan Hu, Ho-fung Leung, and Farzan Farnia. Pak-uch contextual bandit: An online learning approach to prompt-aware selection of generative models and llms. In Proceedings of the 26th International Conference on Machine Learning (ICML), 2025a.
Xiaoyan Hu, Lauren Pick, Ho-fung Leung, and Farzan Farnia. Promptwise: Online learning for cost-aware prompt assignment in generative models. arXiv preprint arXiv:2505. 18901, 2025b.
Mohammad Jalali, Cheuk Ting Li, and Farzan Farnia. An information-theorrtie evaluation of gen- erative models in learning multi-modal distributions. Advances Neural Information Processing Systems, 36:9931-9943, 2023.
Mohammad Jalali, Haoyu Lei, Amin Gohari, and Farzan Farnia. Sparke: Scalable prompt-aware diversity guidance in diffusion models via rke scores. Advances in Neural Information Processing Systems, 2025a.
Mohammad Jalali, Bahar Dibaei Nia, and Farzan Farnia. Towards explainable comparison and alignment of feature embeddings. Forty-ecoon International Conference on MachineLearn- ing, 2025b. URL httl:ttts////peeeeee
net http://nett/ettteeeeeeee www.moorriiooooooooo Fow.ommmmmoooo000000
Mohammad Jalali, Azim Ospanov, Amin Gohari, and Farzan Farnia. Conditional vendi score: Prompt-aware diversity evaluation for text-guided generative AI models. In The 29th Interna- tional Conference on Artificial Intelligence and Statistics, 2026.
Minguk Kang, Jun Yan Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung Park. Scaling up gans for text-to-image synthesis. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2023.
12

---

Published as a conference paper at ICLR 2026
Branislav Kveton, Manzil Zaheer, Csaba Szepesvári, Lihong Li, Mohammad Ghavamzadeh, and Craig Boutilier. Randomized exploration in generalized linear bandits. In Proceedings of the 23rd International Conference Artificial Intelligence and Statistics (AISTATS), 2020.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in neural information processing systems, 32, 2019.
John Langford and Tong Zhang. The epoch-greeds algorithm for multi-armed bandits with side information. In Platt, D. Koller, Y. Singer, and S. Roweis (eds.), Advances in Neural Information Processing Systems, volume 20, Curran Associates, Inc., 2007.
Junnan Li, Dongxu Li, Silvio Savarese, and Steven Hoi. Blip-2: Bootstrapping language-imas -magg pre-training with frozen image encoders and large language models. 2023.
Lihong Li, Wei Chu, John Langford, and Robert E. Schapire. A contextual-bandit approach to per- sonalized news article recommendation. In Proceedings of the 19th International Conference World Wide Web, WWW 10. pp. 661-670, New York, NY, USA, 2010. Association for Comput- ing Machinery. ISBN 9781605587998.
doi: (o.114//779900.177227
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft COCO: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 740-755. Springer, 2014.
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Llava: Large language and vision assistant, 2023.
Jack Lu, Ryan Teehan, and Mengye Ren. ProCreate, Don't Reproduce! propulsive energy diffusion for creative generation. Computer Vision ECCV 2024 18th European Conference, Milano, Italy, September 29- October 27, 2024, 2024.
Zichen Miao, Jiang Wang, Ze Wang, Zhengyuan Yang. Lijuan Wang, Qiang Qiu, and Zicheng Liu. Training diffuuson models towards diverse image generation with reinforcement learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),
pp. 10844-10853, June 2024.
Muhammad Ferjad Nacem, Scong Joon Oh, Youngjung Uh, Yunjey Choi, and Jacjun Yoo. Reliable fidelity and diversity metrics for generative models. In International conference on machine learning, pp. 7176-7185 PMLR, 2020.
Maxime Oquab, Timothée Darcet, Theo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint 3rXiv:2304.07799,, 2023.
Azim Ospanov and Farzan Farnia. Do vendi scores converge with finite samples? truncated vendi score for finite-sample convergence guarantees. In The 4/st Conference Uncertainty in Artifi- cial Intelligence, 2025.
Azim Ospanov, Jingwei Zhang, Mohammad Jalali, Xuenan Cao, Andrej Bogdanov, and Farzan Farnia. Towards scalable reference-free evaluation of generative models. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Azim Ospanov, Mohammad Jalali, and Farzan Farnia. Scendi score: Prompt-aware diversity evalu- ation via schur complement of clip embeddings. In Proceedings of the IEEF/CVF International Conference Computer Vision (ICCV), pp. 16927-16937, October 2025.
PixArt-alpha.
2xxrttx1-2-512x5512
https://haggiigeeecco/otttttttt
PixArt 21.-2-51 2x212 2024.
Jie Qin. Jie Wu, Weifeng Chen, Yuxi Ren, Huixia Li, Hefeng Wu, Xuefeng Xiao, Rui Wang, and Shilei Wen. Diffusiongpt: Lim-driven text-to-image generation system. arXiv preprint arXiv:2401. 1006/, 2024.
13

---

Published as a conference paper at ICLR 2026
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Phil Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. International Conference Machine Learning, pp. 8748- 8763, 2021.
Parham Rezaci, Farzan Farnia, and Cheuk Ting Li. Be more diverse than the most diverse: Optimal mixtures of generative models via mixture-UCB bandit algorithms. In The Thirteenth Interna tional Conference on Learning Representations, 2025. URL httpps://peereeiee..rr
Forum?1d=22hkk5Yee
Runway-ML. Stable diffusion v1-5 model card. https::/huugingaaace wwccceee.ccruie wwwpuwwwwm
2023. Version 1.5 the model, released in February 2023.
Seyedmorteza Sadat, Jakob Buhmann, Derek Bradley, Otmar Hilliges, and Romann M. Weber. CADS: Unleashing the diversity of diffusion models through condition-anneale sampling. In The Twelfth International Conference on Learning Representations, 2024.
Mehdi SM Sajjadi, Olivier Bachem, Mario Lucic, Olivier Bousquet, and Sylvain Gelly, Assessing generative models via precision and recall. Advances in neural information processing systems, 31, 2018.
Amir Sani, Alessandro Lazaric, and Rémi Munos. Risk-aversion in multi- bandits. In F. Pereira, C.J. Burges, L. Bottou, and K.Q. Weinberger (eds.), Advances in Neurul Information Processing Systems, volume 25, 2012.
Matina Mahdizadeh Sani, Nima Jamali, Mohammad Jalali, and Farzan Farnia. Training-free distri- bution adaptation for diffusion models via maximum mean discrepancy guidance. arXiv preprint arXiv:2601. .0.326 2026. /tK/:httpp.://rrv...rr/////
08379
Vikash Sehwag, Caner Hazirbas, Albert Gordo, Firat Ozgenel, and Cristian Canton Ferrer. Gen- erating high fidelity data from low-density regions using diffusion models. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition CCpiiiionnCCPP (CVPR), pp. 11482 1 11491, 2022. doi: 0T0119//CVVPR222822222..11111
Stability-AI. Stable diffusion xl 1.0. stable-diffusion-x//
2023.
George Stein, Jesse Cresswell, Rasa Hosseinzadeh, Yi Sui, Brendan Ross, Valentin Villecroze, Zhaoyan Liu, Anthony L Caterini, Eric Taylor, and Gabriel Loaiza-Ganena Exposing flaws of generative model evaluation metrics and their unfair treatment of diffusion models. Advances in Neural Information Processing Systems, 36:3732-3788 2023.
Michal Valko, Nathaniel Korda, Remi Munos, Ilias Flaounas, and Nelo Cristianini. Finite-time analysis of kernelised contextual bandits, 2013.
Nir Weinberger and Michal Yemini. Multi-armed bandits with self-information rewards. IEEE Transactions on Information Theory, 69(11):7100-7884 2023. doi: 10.1 109/TTT.200233329944))
Yougi Wu, Jingwei Zhang, and Farzan Farnia. Fusing cross and uni uni-modal representations: A kronecker product approach. In Advances in Neural Information Processing Systems (NeurIPS), 2025.
Junwen Yang, Vincent Y. F. Tan, and Tianyuan Jin. Best arm identification with minimal regret. arXiv preprint 200 2024. doi: 200550/axXiv.249.11599
Houssam Zenati, Alberto Bietti, Eustache Diemert, Julien Mairal, Matthieu Martin, and Pierre Gail- lard. Efficient kernelized ucb for contextual bandits. In Gustau Camps- Valls, Francisco J. R. Ruiz, and Isabel Valera (eds.), Proceedings of The 25th International Conference Artificial Intelii- gence and Statistics, volume 151 of Proceedings of Machine Learning Research, pp. 5689-5720 PMLR, M.LR, 28-30 Mar 2022.
14

---

Published as a conference paper at ICLR 2026
Jingwei Zhang, Cheuk Ting Li, and Farzan Farnia. An interpretable evaluation of entropy-based novelty of generative models. In Proceedings of the 41st International Conference on Machine Learning. volume 235 of Proceedings of Machine Learning Research, pp. 59148--59172 PMLR, 21-27 Jul 2024.
Jingwei Zhang, Mohammad Jalali, Cheuk Ting Li, and Farzan Farnia. Unveiling differences in generative models: A scalable differential clustering approach. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), 2025.
Qiuyu Zhu and Vincent Tan. Thompson sampling algorithms for mean-variance bandits. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Ma- chine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 11599-11608. PMLR. 13-18 Jul 2020.
APPENDIX A DERIVATION OF THE MIXTURE-DAK-UCB PROXY OBJECTIVE FUNCTION
A.1 UCB FORMULATION OF MIXTURE- CANUURE-DAK-UB
Building on the mixture objective in the main text, we here prove the approximation guarantee for Mixture-DAK DAK-UCB, and then provide the corresponding UCB formulation.
Proposition 2. Assume the kernel functions
normalized and satisfy kT (t, 1 and
(x(x,x)) 1 for all t,t', x'. For every mixture weight A. {a T + Ag V, (L,(L,)) (|a(()) - ai(')),, S the following hold:
(a) For the I-JRKE score of the mixture Pa defined as I-JRKE(Pa) Ei~Pr (a(P)(((t) (t))
the proxy I-JRKE score I-JRKEapprod (Pa)= = E. Eg Og (B) (t) results in an c- bounded error:
I-JRKE(Pa) BURKEupporr (Pa) Se
(b) For the JKD score of the mixture Pa defined as JKD Q) Eg kr t') Eg (t)) (t. the proxy JKD score JKD (Pamme(Par (asy Q) =
(t), Egy (t, Eg (g(t)ag (ID (t, (t, results in an e-bounded error:
JKD Q) JKD aapppruee (Pa 2) Sc.
Proof. Proof for (a). Considering the definitions, we have
[-JRKE(Pa) I-JRKE upprox (Pa)l
Er, (t,t'))) Σ (g(t) (t') RK (t, E, Σ
(g(())ag (t MRKE (t)
9.9'
9.9"
= BA (t, Σ (t)aa - 0gg ( (t, r2)
99'
9..
For fixed E, t', we bound the inner sum as
Σ (1))))as ag (1) 98 9.8
(1, S Σ (t)la) ag (4)
3.9
< ((a((()) - a<t))
where the last line uses E, (t()) = 1. From the Lipschitz condition a E A::
(1, 1')l [|a((t)) ax(1)) Sc.
15

---

Published as a conference paper at ICLR 2026
Since tt, & (t(tt)))r(t = 1, we have (t, z (ty(t, and then: (t, (la((t))))))))) a(t)ll - (t, 10) ||a(tt) a((1]] sc..
Therefore, we can write
[-JRKEE(Pa) I-JRKE upprox (Pa) S Ex (t,t') [|a(tt) (t)))a)
I Et
c.
Proof for (b). Using the definitions,
JKD( Q) - JJKD
JKDapprou upprox (Pa
1 (1, Σ
ag(t)ag (1' K2p
ng (1, (1,
- E,
Σ (g(()
on"
9.0"
9.9
Euy (t,tt,t)
Σ
(tlt ag (t)] IKD
(t, e
9:9"
For fixed t, t', we can bound the inner sum as follows
Σ ag (g((a(ag())) ag (4) (1, (1,))) S Σ Eag())agg() ag)
9.9'
9.9
$ (lattt) - a(tt)
From the Lipschitz condition and noting that (t, S (t](t
(t, ||a(t') a(t)ll1 S (t, -
s.
As result, the following holds
JJKD(Pa Q) JKD approx (Pa. Q)1 E, -()-t((()()))))))))))))))()))){) (aa(a)la((ttt))tt a(t)lla)
< E,eld
Therefore, to formulte the UCB formulation of Mixture-DAK-UCB, at each round i for prompt 4, the learner maintains UCB predictors for
SUCB (t.)
E R
MUCB E RGG MUCB (t.) 2
where SUCB (t+) collects fidelity UCB scores for each model and MUC PSD estimate of the cross-model diversity matrix M (t:). The per-prompt mixture decision then follows the concave quadratic program
(a, SCCB - A M (UCB(L)
a =rg max atAg
This UCB objective ensures optimism for fidelity while pessimistically accounting for diversity penalties. The chosen a specifies sampling distribution over models, from which the algorithm draws 9: ~ Multinomial(a2)) and obtains x; ~ Pg. (-14). The resulting procedure, which extends the single-model DAK-UCB mixture assignments, summarized in Algorithm
APPENDIX B REGRET ANALYSIS OF SUP-DAK-UCB (PHASED VARIANT OF DAK-UCB)
As noted in the literature, the theoretical analysis of the standard kernelized UCB method faces the challenge of potentially statistically correlated model selection at different rounds, which renders
16

---

Published as a conference paper at ICLR 2026
Algorithm 2: Mixture Diversity-Awand Kerneiized UCB (MITURRE-DA-UCCH
Input: models; horizon T: prompt distribution P; trade-off A; diversity primitive
& E (HRKE, JKD); (optional) panel rate E [0.
Output: generated outputs
Initialize per-model KRR for fidelity (5 and matrix-valued KRR for diversity {Ma ba
for i to T do
Sample prompt t2 ~ P;
For all g: (ty(t.). (6) (t:)) -RRR-PREDICT,(g t;); (M)
For all (g.g'): (Myy' (pi), 99" (t:)) t KRR-PREDICTMM(O.C tc);
Set 5
SUCB((t) (t:) (t4))) + //5)0g (t,) for all g:
Set 29' (t:) Mg (4) 3 (M)
(t:); project PSD if needed;
"g"
a; arg max T g ST (t:) Xa M (LLCB(t:): acAc
Sample gi ~ ai draw a; ~ Pa (:));
Vi (+) (ii 24);
10
11
if panel step with prob. then
12
for g 1 to do
draw 20 P- (-1 ti) (reuse x. a(00)= =I;)
UPR-EEDAATE (fz, v. ))): if panel step then
14
15
16
Update (May') with cross-kernel labels built from { 1g=13
}g=1"
17
else
Update }Mar'} using the available pairs;
standard concentration analysis by means of independent observations inapplicable. To circumvent this challenge, adopt the standard approach of analyzing staged variant of the proposed DAK- UCB algorithm, which we call Sup-DAK-UCB. The same technique of analyzing Sup-Kerenlzies UCB and Sup-PAK-UCB have been applied in the related works (Chu al., 2011 2013 Hu et al. 2025a).
Valko et al.
In the phased variant of SUP-DAK-UCB, within each urm-stage-target triple, the data used by kernel ridge regression (KRR) are independent. This enables the analytical derivations of confi- dence bounds and proper regret decomposition, similar to the analysis in (Valko al. 2013 Hu et al. 2025a)). In the following, we first state the updated setting and assumptions in our theoretical analysis, then present the phased algorithmic structure, followed by the theorems and their proofs.
B.I ASSUMPTIONS IN THEORETICAL ANALYSIS OF DAK-UCB
Let G (1,..,.., be G generative models, T prompt space with i.i.d. prompts Pr, and X' the output space. At round I, the algorithm selects single model 9t- Note that the objective function is the following for parameter
0: (t) == sy (D,
Assumption (Normalized Prompt and Data Kernel Functions). The prompt kernel KT TTX + [-1, 1] and the output kernel xx X 1-11,11, 1] are positive definite with kr lx (z, x) 1 for all t E T, x E X. For g.g E G. define
K Kgg' (p.p')
E [0.
Assumption "Sub-Gaussian noise in kernel regression). All sealar observations are conditionally o- o-sub-Caussian Gaussian
given the history: Efexp(A=) 744 I cxp(X2o2 /2) for all X R.
Assumption (RKHS boundedness for single-model case). Let HT be the RKHS of ky. Assume 8g HT with ||s सt S B, and Dg HT with [| || BD for all g G.
17

---

Published as a conference paper at ICLR 2026
B.2 INTRODUCING PHASED -UP-DAK-UCB ALGORITHM
As mentioned earlier, we analyze a staged variant of DAK-UCB, called Sup-DAK-UCB, possessing M [log2 7 stages. In Sup-DAK-UCH for each generative model-stage pair (g, m) and target type r E {s, D). we maintain frozen index set 4g To guarantee independence for diversity labels, we additionally maintain stage snapshot of the archive of past (1, pairs for each arm: the beginning of stage m, we freeze D and use this snapshot to build diversity labels throughout stage m. We note that new pairs collected during stage m are not used to form diversity labels in the same stage; They will be only available from stage m + 1 onward.
To explain the steps of Sup-DAK-UCB in Algorithm note that iteration in stage m with candidate set 9Th we perform the following.
1. For each g E 9Th, we perform KRR predictors (gy.. ag m.())) and (Drs om.(D)) based on (g and ()) respectively.
g,55 gg
We let 17 == a /2 (2GMT/6) and set B(0) == B.Va + foa
3(D) = =BDVa+ go.
We define the optimistic score and the width as
J'm 3.3 = (sg), m())) AD g.a ((D)) 5.2
w 9. = (l0))m(()) + A0(D))m.())
g.
43
3. The stage selection rule in Sup-DAK-UCB as follows:
If max, ugght 1/12 exploit: gi arg max gax Jm
Else if max, w I 21-"", we eliminate {g maxy' Jw. jmm > 22-m] and set m +
+ 1.
Else (explore), we pick every with wh, > 21 21-m and append to (m) and m.))D)
4. Feedback: For single-model selection, we draw x; ~ Pg. (-It,), observe y (0) and build a stage-frozen, unbiased, bounded diversity statistic using Di
sg. (ti) + =
d, max(1, D
Σ
(t,,),))) 2x
E dg t., DT Dg. (ty).
(M,")ED
We define the zero-mean diversity noise E ( = d Dg. (t,). Finally, update the archive
Dr + Dy u ((t,, x;)), which will only be snapped at the next stage.
B.3 KRR NOTATION AND INFORMATION MEASURES
For an index set V, let Φ = [o(t:) [o(t:)Tkev. K2 = 2y T.. (t) = [tt(t, and Ay =
d v D + al. The KRR predictor and posterior deviation at are
7et: V) = Ky (1) (K + "n = <8() T A x Tywe
()11;)= =()) A A() d(1).
We use the shorthand na (8) o /2logg2//6 Also, use the following complexity measures:
T(T) = log de + a -I Kv), TT = max 7(T).
V: ST
B.4 SINGLE SELECTION SUP-DAK-UCH REGRET BOUNDS
Lemma 1, Consider arm 9. stage m, and target T E D). Consider the sequence of time indices
(i) that get appended (T by the stage rule. Conditional on the prompt sequence (tx) and the stage-froze frozen archive snapshot Dmm (for T = =D), the random variables {v)" (r) Neegitt are mutually independent and satisfy Ely(()) t,, D f(²) (tt).
(t)). where f(a) s and f(D) D. Moreover,
( d. -- -D2((+) is conditionally 1/2-sub-Gaussian.
18

---

Published as a conference paper at ICLR 2026
Algorithm 3: Sup-DAK-UCB
Input: models, rounds, prompt dist. P, trade-off X, kernels ky, kx, ridge at confidence
Output: 7 generated outputs
Set number of stages M < [log2 T];
Initialize per-stage
(pr,((s) (m. 0 and archives D, +- 0:
for i == to T do
Sample prompt 1 P; set 1. 9 [G]; Freeze stage snapshots D & Dg for all gi
repeat
for 9 E 5Th do
Compute fidelity and diversity predictions by KRR: (8% Og om.(e)), (D og .o..(D))
Form optimistic score Jgn and width wy:
if maxy "0 S T-1/2 then
9. + arg max, Jm, break;
11
12
13
else
if maxy wy S 2l then
Eliminate: gm+ t (g: Jm 2 max Jm - 22-m);
15
16
m + m + 1; freeze new snapshots;
17
else
Explore: choose 9i with w > 2 21-m.
Append i to um (e) wm.(D)
Vg
20
break;
21
until;
Generate x; ~ Pg. (-lt:), observe (*) and stage-frozen y!?
22
Update live archive Dg, Dg. U + ::)]:
Proof. For r = yi sp(p:) + with et conditionally independent across given (Fs-1,t)) t;) by Assumption Thus, {uj" (0) immm. are mutually independent given (4:).
For = D, by construction we use the stage-frozen archive snapshot D Given t; and Dm we draw a fresh x, ~ P. t,) independent of (z,))<s. Each summand inside d, equals
Z, = kT t') kx x')² [0,
and these are i.i.d. across (t',') E Dit conditional on (t,, (t). Hence,
B dg 4, D
Σ
kr( 1')² (ao Dy(ts).
max ma (D
(U,+))D
Define (D) = -p.)),(i)). Since d, E [0, and Dy(4:) E [0, 1J. Hoeffding's lemma implies e! (D) conditionally 4/2-sub-Gaussiae The mutual independence {v) (D)) over E wm holds because each y! (P) is function of (f;, 2;) and the fixed snapshot Di and (æ:) are independent across i.
(pi).
Lemma 2. Fix (g, m, and a nonempty index set Vp = V (.n.(t) that satisfies Lemma Let f ( E HT with || < Br, and write yo f( tew, where Ex zero-mean, conditionally independent and --sub-Gaussia (with a = for fulelity and = 1/2 for diversity). Then, for any t and any 8' (0, 1). with probability at least -
[att v) -(()) < To (8'))(( o + Va (BB
naisee
strinkage
19

---

Published as a conference paper at ICLR 2026
where no (8') 2 (/8"). Equivalently, with 3(7) == B. no(0')
8 + 3())g())
2, δ D-3 S.
Proof. Let A
4 ty + al. Using the identity Po + +al) = ( + a0)-1186
A T. we can write
[ v) s() -()) (()) ( vv) 5(() wy
where Wy Hy is the (unique) RKHS representer of f, i.. f(.) = (w o(:)), with |
Ilfll S B,, Inserting and subtracting <(t) 1 + w wy gives
A(t: v) 4
- f(t) = (tt) A (t)
A 1 2 D 1 wy
Term ())
Term
Bounding Term (I): Let A E RI By independence and sub- Lub-Gaussianity of s and the standard bound for linear forms sub-Gaussian vectors, for any 8' E (0, 1).
P((a [ - n) [|alla) 2 1 8' no = o (log((/8))
We bound ||ll2 by noting that
la12 (
A 2 4 0(t) s - A = 6²;
since T6 3 A implies A $ 2 1 1- A-1 Hence, with probability at least 1 - 8',
[Ø(t) ( A 4 TT s No ()))
Bounding Term (II): Since A = $ Φv + al, we have A A- tv I= -aA-! (II) becomes equal to ((t) T wy. By Cauchy-Schhieee we can write
(-t))T--'wf = ||A-1/26(t)l -|a-)))/_l1--/3 ||A-1/2w,ll, S vo(I) A
||--1/2
| ||-
=((122
-a-1/2
Thus (o(t)
(A-10 (
- 1)wsl s \( S va B, Ōt: v) Combining (I) and (II)
completes the proof.
Lemma 3. Fix any round and stage m. Suppose the confidence radii in Lemma 2 hold for all arms 9th (with the choice 8' (/(G) and union bound across i, 9, m.). Then the following hold simultaneously with probability at least 1 8:
S - sg((t) s(a) gi and 5Tm g
Dg(t,)l (D gá for each E 9m.
2. The true per-prompt maximizer 9f E 9m after any elimination step.
3. For every E T"h, J9g (tx) (y((t) < 23-
Proof. Item follows directly from the application of Lemma to s and / with union bound over i, m.
Regarding Items and 3, consider a step in which the algorithm chooses either to exploit or elim- inate. Throughout this step, by definition of wit and the trigger max max.53m wigh I 21-m have
for every g 9mh.
g
< 2l-m
A xg(D)am 21-m
(10)
9. Using Lemma 2 we have the following to hold for every component:
gi
8g (14)) 8%h + g()))m.(c)) 9.
Dy(1,) 2 Dm
3(D)
20

---

Published as a conference paper at ICLR 2026
Therefore,
Jn(44) = -bD(()))))(())
S gm 9.
+ (
-> (D.r, - ((D)omm.0))
Jm
=22-m.
g.
9,
3.a
S g.. +-23mm
2 B 22
Similarly,
Jg(t,) 2 Jm 22-m
9.2
In particular, for the optimal arm 9f and any 9 E 9m.
Jm g. 2ot (t:) - Jg (t:) - -22-m 22- -23
Hence Imm is at most 2 m below any Jm. During elimination, the algorithm removes only arms 99. whose J is more than 22-m below the current maximum. Since 23-m > 2 is the two-sided tolerance and 9f attains the one-sided tolerance bound, of cannot be eliminated: this proves (2). Finally, any survivor satisfies
J0p (1:) Jo(()) S (J - - -Jen -22-m) < 22 + 22- 23-m
which proves (3).
Lemma 4. Fix (g, m, T) and let (1r)jEv be the prompts indexed by V V. in the sorted order. Define Ag = aT and A₃ = A5-1 + ((t)) Then, the following holds
Σ E Σ o(t)) (7))(1)))))))) (7()), JEW
jeW
Also, the application of the Cauchy-Schwart inequality implies that
E S
Σ
EE((iVe))) S (2)
jEV
jEV
Proof. For the determinant of matrices, we can write the following:
det(A,) det(AA-1) (1+ oCtz) ( Af)1o(t)))
Therefore,
log det(A,) log (1 02 (tz; V7 <i))
dct(A,-)
Summing the above from j = 1o = | shows that
log + (1:50 V v)) = log det(A[w1) = log det (1 del(An)
Σ JEW
= (7V))
Since 0 - (t,; V2 x)) < al (because A;-1 2 al), we have 22 when a 2 1. Also, note that for every x E [0, 1). < 2 log(1 + x) holds, which implies that E, @25 V<i) S 2, log(1 + (t; V 22)) (V). Finally, the linear-sum bound follows directly from the application of the Cauchy-Schwer inequality.
Theorem (DAK-UCB regret). Define instance-wise regret be b; b Jty (t.) (tt). Under As- 3 and with g(s) = aa g(D) = Bpya+ where a / / 22gg/22GTT)))
sumptions
the phased DAK-UCB satisfies, with probability at least 1
Regret(T) ((VTTT+) x GVTTT)))
r6 by an effective-dimension proxy where p(r) = maxg.m 7 ( Equivalently. replacing
d' the bound Ō( voly (3 (0 voly dD GT).
21

---

Published as a conference paper at ICLR 2026
Proof. We partition rounds into To (exploitation: max; ugg 1 and Ti (exploration: width value greater than 21-m).. On To. by Lemma Part 2 and the exploitation trigger, 7; S 2w ge S 27-1/2 and therefore ΣieTo S 2,T. Subsequently over T1, we group considering (g. m)) as
M Σ
Σ RETI
Σ
23m Σ 23-m (ymm(())
3=1 g=1==l
9.m
For each such i, the append rule implies wigh ming over Vg (.) and applying Lemma and Schwarz to the two targets separately yields
> 21-m iee, 3 91
+
AB(D))
21-m Sum-
2 [mm & g(a)
Σ
g(a) git
+ AB(D)
((D)
SEW (4))
g(0)
/
2/
v
7g
mild)
+ (3
23 (m) m.(D)
the above, 79 m.(F)) = > V (T)) Taking the summation over (g, m) pairs and noting [loge T'] completes the proof of the regret bound.
APPENDIX C ADDITIONAL NUMERICAL RESULTS
C.I DAK-UCB APPLIED TO LLM SELECTION PROBLEM
DAK-UCB for Diversity-Aware LLM Selection Using Synthetic Prompts: In this experiment, we asked GPT-4o to provide five words as categories: temple, painting, market, horse, and farm. We then selected three LLMs: DeepSeek (DeepSeek-Al) 2024), Gemma (DeepMind) (024), and Llama (Al 2024 At each iteration, a random cluster from these five was selected, and a prompt of the form "Describe scene containing a [cluster]." designed the input to the LLM. In Figure you can see the performance comparison and selection ratios of these models.
DAK-UCB Selection Ratios
DeepSeek
Gemma
Joint-RKE Score
4.8
0.8
Llama
DAK-UCB
Random
0.6
4.6-
Mixture- DAK-UCB
One Arm Oracle
0.4
PAK-UCB
4.4
0.2
0.0
100
200
300
400
0
100
200
300 400
Number of Iteration
Number of Iteration
Figure 6: Results are averaged over independent trials.
Detecting Bias in LLMs Using the I-JRKE Diversity Metric: In this experiment, we used (DeepSeek-AI 2024 to generate sentences about cities in the US Canada, China, and England. One arm biased toward the capitals of countries, while the other shown in Figure 7 our algorithm DAK-UP -CB preferred the unbiased arm, as it resulted in greater diversity across the arms. The results reported after 200 iterations.
unconstrained. As
22

---

Published as a conference paper at ICLR 2026
China
Canada
England
US
China city
Serani conic Vibrantain
Vancouve
London
known
city
Unconditioned LLM
historyCanad
New York
SUBel
city ronawand
known
known
City ty Canada known Beg attawa capitall vibrant
China
Ontario
Capita tata
saly Kro
Conditioned LLM
rish
London
Unire
(Conditioned capitals)
Vibrant
capital
diver'e city.
VINTONT BBSD LDS
Egtoo.
DAK-UCB Selection Ratio
28.5%
71.5%
Figure 7: DAK-UCB selection ratio for LLM bias detection experiment.
Enhancing LLM Diversity via I-JRKE: In our experiment, we implemented a four-arm setup us- ing DeepSeek, where each arm was coollasty-oollapsee by country-specing biasing through modified prompts of the form: "Describe one of the famous cities in [Country] in no more than 20 words." (with Japan, France, Brazil, and Egypt as respective biases). Results from 10 runs demonstrated that mixture effectively enhanced diversity. Scores and selection ratios of each arm are presented in Figure 8 The observed decay in DAK-UCB is explainable: With only one cluster, the algorithm sticks to one arm after some point. This naturally converges to a single-arm oracle scenario.
1.0
MK-UBB Selection Ratios
Japan
France
Joint-RKE Score
2.5
Brazil
Egypt
0.6
2.0
DAK-UCB
0.4
Random
1.5
Mixture- Mixture-DAK-UCB
One Arm Oracle
PAK-UCB
1.0
0.0
500
1000
1500
0
500
1000
1500
Number of Iteration
Number of Iteration
Figure 8: Selection Ratios of biased arms and performance comparison on Joint-RKE.
capital
Egypt
one
Brazil
Rio de
cultura
dknown
ncn history
city
alpan
JOKYO
CapitalFrant
Paraiss
Known
France e
Egypt
Cairostruc
Aworld
history
Sarurssructure
known a
GIZa
eancient
ture famqus
ncteet
Great prr
ramid
Eiffel Tower
BOOE Known
enthlechue
deilaneiro
famours
Pyramid
massive
Brazil
Japan
France
Egypt
Figure 9: Word Clouds of DeepSeek Responses Across Countries
23

---

Published as a conference paper at ICLR 2026
C.2 DAK-UCB APPLIED IMAGE- CAGGICNATIONIIG MODEL SELECTION TASK
Improving Correctness in Image Captioning via JKD: We evaluated three state-of-the-ar image captioning models arms: LLAVA (Liu al. 2023), INSTRUCTBLIP (Dong et al.) 2023), and BLIP-2 Li et al. 2023). At each iteration, we sampled an image-caption pair from a thousand MS-COCO instances, distributed across 10 previously mentioned clusters. By minimizing the Joint Kernel Distance (JKD) between generated captions and MS-COCO reference captions, we aimed enhance captioning correctness through dynamic model selection. Figure [T0] demonstrates the evolution of KII) scores across iterations compared to baseline methods, along with the empirical selection ratios for each captioning model. Figure [1 visualizes the dataset.
1.0
0.045
DAK-UCB
MK-UCB Selection Ratios
Blip2
Random
InstructBlip
0.040
Mixtue
0.8
Llava
KID Score
One Arm Oracle
0.035
0.6
0.030
0.4
0.025
0.2
0.020
0.0
500
1000
1500
500
1000
1500
Number of Iteration
Number of Iteration
Figure 10: Selection Ratios of image captioning models and performance comparison on KID using JKD metric. Results averaged over 10 independent trials.
24

---

Published as a conference paper at ICLR 2026
Cluster
Image
Captions
MS-COCO: of dog laying the ground.
Blip2 A dog sleeping on the sidewalk.
Dog
InstructBlip a white deg leying on the ground next to bicycle, indicatinng that the dog is sleeping. dog sleeping
sidewall..
MS-COCO close of child next to cake with balloons.
Blip2 A baby eating cake.
Cake
cake,, Liava A baby with cake his face.
InstructBlip bait sitting
birthday present.
Blip2 i cat sitting on couch with remote control
Cat
InstructBllp: cat laying on loather chair with two remote controls its mouth.
sitting on
MS MS-COCO-1 A rear view mirror bike is reflecting poster building.
Blip2 A reflection bicycle in side mirror of motorcycle
Bike
InstructBlip: motorcycle
side reflecting sign. Llava A reflection of sign motorcycle mirror.
running gress Blip2 i zebra large, fast-running black and white horse
Park
InstructBllp: zebra running through a grassy fleld with treas in the background. walking on grass.
MS-COCO:
Two beautiful young ladies giving an elephant a bath a river.
young girl bathing elephant
Elephant
InstructBlip: two woree washing elephent in water, which be practice parts the world.
girls washing an elephant mvvr.
MS DBS-COCO:: plane parked in large sirplene hanger.
Bllp2 A large airpiane parked in hangar
Airplane
InstructBlip: large airplane parked hangar, several people working around
Liava
parked inside hangar.
COC.: A large white many green apples.
green apples
Bowl
InstructSlip: bowl of green apples stacked a white bowl. The bowl is placed table, which likelyy
tabletop, and the bowl is filled with green apples.
green apples.
MS MSs-CCCO Large tree with woaden picnic bench in wilderness area Blip2 A tree and bench
Tree
InstructBlip: large, green
gressy field,
nearby,
Lieva :A pienic table in field.
like moves through Blip2 small boet is treveling the water
InstructBllp: small motor boat travaling on body of water. The boat white color and has small attached
faster afficiently.
Liava small boat water.
Figure 11: Dataset visualization for image captioning experiment.
Improving Diversity in Image Captioning via I-JRKE: We repeated the exact same experiment on image captioning, replacing the JKD objective with I-JRKE to observe our algorithm can enhance diversity. As shown in Figure [2] algorithm demonstrated superior performance compared to the basclines. While LLaVA worked best in terms of correctness, our algorithm tends to prefer InstructBIIIP for diversity.
25

---

Published as a conference paper at ICLR 2026
MK-UCB Selection Ratios
1.0
Blip2 InstructBlip
Joint RKE Score
220
0.8
Llava
0.6
200
DAK-UCB
0.4
180
Random
Mixtue-DAK-UCA
0.2
160
PAK-UCB
0.0
500
1000
1500
500
1000
1500
Number of Iteration
Number of Iteration
Figure 12: Selection Ratios of image captioning models and performance comparison on Joint-RKE using I-JRKE metric.
C.3 ADDITIONAL NUMERICAL RESULTS ON APPLYING DAK-UCB TO THE TEXT-TO-IMAGE GENERATION TASK
Conditional Expert Selection via I-JRKE Using Classifier-Free Guidance: We used the dataset from (Rezaci et al.) 2025). The dataset was generated by selecting four categories: dog, river, airplane, and building. For each category, GPT was asked to generate 10 adjectives, 10 activities, and 10 places. By mixing these with category, 1000 prompts were obtained. Some samples of the dataset can be seen in Figure 1 We used SDXL to generate images with classifier-fres guidance scales and 30 for each cluster. We designed the arms as follows: ARMI: Generates images of dogs and rivers with classifier-free guidance scale of 2.0 (less guided, more diverse) and images of airplanes and buildings with a classifier-free guidance scale of 30.0 (more guided, less diverse). expert in the dog and river clusters. ARM2: Does the opposite, making
Thus, this arm acts it an expert in the building and airplane clusters. In this experiment, our objective was solely to minimize I-JRKE, and CLIP was involved in the optimization. As shown Figure I4 (averaged over 10 trials), the expert for each cluster was successfully detected. However, we observed that in some clusters, the selection ratios were very close. This is acceptable, as we know that mixture can enhance diversity.
Prompts
Classifler-Free Guidance 2.0
Classifier-Free Guidance 30,0
happp dog running the park happy
happy nunntng backyand. happy dog
wide river flowing through the valley. A wide rheer flowing in forest. wide river moxtin by the mountxine. A wide rierra floving near the stlinge.
faat sirptane flying the douts. A fast airplane flying the rumwey. A twes sirplane flying over the coean.
the
tall building standing in eity. tall building' standing by the rives. call building is standing is the coumtyside. till building is standing on she hill.
Figure 13: Ilustrative examples from the GPT-Generated prompt dataset and their SDXL-generated visualizations under varying guidance scales.
26

---

Published as a conference paper at ICLR 2026
Dog prompts
River prompts
1.00
SoXI (CFGm2)
Selection Ratio
1.00
SDXL(CFG=2)
0.75
30X
0.75
SDXU(CFG=30)
alection
0.50
0.25
0.25
0.00
0.00
500
1000
1500
500
1000
1500
Joint-RKE
Iteration
Iteration
Airplanc prompts
Building prompts
DNG-IC
Rendom
1.00
(CF
1.00
(CFG()) S0XUCTG=101
Mishar-DAAA DAKAII
Oracle
20.75
SoxLIcrG=301
2.75
60.50
Selection
0.50
Selection
0.50
500
Number Iteration
1000
1500
0.25
0.00
0.00
0
500
1000
1500
500
1000
1500
Iteration
Iteration
Figure 14: Cluster-conditioned expert selection ratios and performance comparison against baseline methods.
Diversity Identification in Text-to-Image Generative Models: An Offline Study on Varying Diversity in "animal" Generation: In this experiment, modified the setup described in Fig- ure 3 and considered an offline setting. To achieve this, we evaluated the reward with respect to all available offline data and removed the UCB radius term, since the entire dataset was accessible and no exploration was required. Figure I presents sample prompts along with the arm-selection preferences of DAK-UCB and Mixture -Akture-DAK-UCB....
Table 1: DAK-UCB and Mixture-AKK-UCB Arm Preferences for Sample Prompts
Prompt
DAK-UCB Selected Arm
Mixture-DAK-UCR Preference Vector
"an animal the garden"
Arm (unconditioned)
[0.218, 0.063, 0.719]
"an animal in the meadow"
Arm (unconditioned)
[0.218, 0.064, 0.718]
"an animal near the lake" Arm 3 (unconditioned)
[0.225, 0.069, 0.706]
'an animal in the jungle" Anm (unconditioned)
[0.230, 0.082, 0.688)
"n animal in the desert"
Arm (unconditioned)
[0.226, 0.092, 0.682]
Conditional Expert Selection via I-JRKE Using Synthetic Diversity Control: In this experiment, we utilize four clusters from the MS-COCO dataset: Bike, Car, Bowl, and Airplane, each comprising hundred prompts. The prompts partitioned into ten groups via K-means clustering applied to their CLIP-embedded vector representations. We designate four specialized experts corresponding each cluster. When the selected expert matches the revealed prompt's 's cluster, it generates the appropriate SDXL output for that prompt. For mismatched cases, we employ a diversity-limitin strategy by outputting the image generated for a group representative prompt rather than the actual prompt. Specifically, we assign representative prompt for each of the ten groups, and for any prompt within a group, the system outputs the SDXL-generated image corresponding to the group's representative. This experiment was conducted over 2000 iterations across 20 independent runs. The resulting expert selection ratios and Joint-RKE scores are presented in Figure 15 The observation that PAK-UCB works competitively well is that expert selection for individual prompts -whnpe where each image is generated specifically for its prompt rather than for group representatives intresentttvess-iprooosss CLIP score alongside diversity.
27

---

Published as a conference paper at ICLR 2026
Expert Selection Rotio
Bike (Cluster
Airplane (Cluster
1.0
0.88
Che are unie
Expert Sellection
0.2
75
DAC
70
Randum
0.2
DACUCE
500
1500
Joint-RKE
Itoration
Iterati
60
Expert Selection Ratio
Car prompts (Cluster 2)
Expert Selection Ratto
Bowl promots (Cusser 31
155
0.4
50
0.2
500
1000
1500
Number Iteration
0.0
500
1000
1000
1500
Itera
Iterati
Figure 15: Expert Selection Ratio and Performance Comparison on Joint-RKE.
Performance Comparison by KID via JKD: In this experiment, we used the same ten clusters of prompts and three generative models as in the first experiment. We employed MS-COCO images as reference and optimized solely on the JKD score defined earlier. The performance was evaluated using KID to measure how close each baseline gets the reference. The results are shown in Figure T6 We observed that Mixtue-DAL Mastue-DAK-UCCE and DAK-UCB achieve close performance this setup, which suggests the optimal solution per cluster is not mixture.
0.016
DAK
Random
0.014
Misture
Orade
0.012
PAK-ICS
KID
0.010
0.008
0.006
500
1000
1500
Number of Iteration
Figure 16:
Performance comparison KID for MS-COCO prompt clusters using Kandinsky, SDXL, and GigaGAN. Results are averaged over 10 trials.
C.4 ABLATION STUDIES
Testing the Robustness of Fidelity and Diversity Scores Under Noise: In this experiment, our ob- jective was to evaluate how reliably our fidelity and diversity metrics behave when the input images undergo controlled degradation. Using 1000 samples from the MS-COCO validation set, we pro- gressively increased the level of blur applied to the images and measured both CLIP-Score fidelity and DINOv2-based diversity. The CLIP-Score decreases monotonically with increased corruption, demonstrating its sensitivity to image quality. Likewise, the diversity metric also decreases, indi- cating that it does not mistakenly interpret noise as meaningful variation. The two plots in Fig. 17 summarize these trends.
28

---

Published as a conference paper at ICLR 2026
CLIP Blur Radius
Joint-RKE Radius
0.25
15
1ont-RME
12-
0
10
12
16
10
Blur Redius
Blur Radius
Figure 17: Left: CLIP-Score vs. noise level showing a monotonic fidelity drop. Right: Diversity in DINOv2 embedding noise level showing reduced diversity under stronger corruption.
Awdel Behavior in Asymmetric Degradation: To verify that algorithm does not blindly prioritize diversity at the cost of fidelity, we revisited the introductory two-arm experiment shown in Fig. [ This time, however, we introduced strong blur (radius 20) exclusively to the diverse arm while keeping the limited-diversity arm intact. As expected, the selection ratio of the degraded arm dropped sharply, demonstrating that our method correctly down-weight samples whose fidelity deteriorates, even if they originate from a high- nigh-diversity source. The final selection ratios are reported in Table showing behavior consistent with our fidelity-sensitive design.
Table 2: Final selection ratios when only the diverse arm blurred (Radius 20).
Diverse Arm (Blurred) Limited-Diversity Arm
Selection Ratio
45.32%
54.68%
Efficiency Random Fourier Feature Approximation: To address the quadratic growth of ker- nel computations in DAK-UCB, we adopt Random Fourier Features (RFFs) to approximate the RBF kernel, reducing the per-round computational cost to (ddt) for an RFF dimension Since the proxy embedding produced by RFFs is d-dimensional, this approximation offers a scalable alter- native while preserving the behavior of the original kernelized method. In the experimental setting of Fig. 2 verified that based DAK-UCB closely matches the performance of the exact RBF- kernel version across all metrics. As shown in Table 3 increasing the number of random features improves stability while maintaining nearly identical scores.
Table 3: Performance of RFF-based DAK-UCB compared to RBF-kernel DAK-UCB under the
setup of Fig. 2
# RFF Features
Final Joint-RKE
Final CLIP
Final KD (x10-3)
Elapsed Time (x10*)
32
158.4
29.80
7.20
4.947
64
160.5
30.40
6.60
4.994
128
162.3
31.00
5.80
5,006
256
172.65
31.59
4.78
5.013
512
182.79
32.17
3.60
5.019
Adaptation to the Introduction of a New Arm We repeated the experiment shown in Figure with a modified setup in order to evaluate the algorithm's ability to adapt when a new generative model is introduced partway through the process. Specifically, we began with two arms and intro- duced the third (and most diverse) arm at iteration 125. Once the new arm appeared, the DAK-UCB algorithm gradually adjusted its selection behavior. As shown in Figure 18 the algorithm success- fully adapted to the presence of the newly added arm, redistributing selection ratios and eventually converging to the appropriate mixture for this expanded arm set.
29

---

Published as a conference paper at ICLR 2026
DCK-UBB Selection Ratios
1.0
animal
cat
0.6
0.2
100
200
300
400
Number Iteration
Figure 18: Adaptation behavior after introducing the third (diverse) arm iteration 125. The algo- rithm subsequently converges to updated selection ratios that incorporate the new arm.
Sensitivity to Kernel Function and Embedding Choice: We repeated the initial experiment (re- sults shown in Figure 2 with slight modification: used CLIP embeddings for the images and employed cosine similarity as the kernel function. Figure 19 presents a comparison to the baselines in terms of Joint-RKE and CLIP-score, while Table 4 reports the final cond-vendi scores.
[9
12.5
32.25
12.0
Joint-RKE Score
32.00
11.5
CLIP Score
31.75
11.0
31.50
DAK-UCB
DAK-UCB
10.5
Random
31.25
... Random
Mixture-DAK-UCB
Mixture-DAK-UCE
10.0
One Arm Oracle
31.00
One Arm Oracle
PAK-UCB
PAK-UCB
30.75
9.5
500
1000
1500
500
1000
1500
Number of Iteration
Number of Iteration
Figure 19: The average CLIP Score and Joint-RKE Diversity score over trials using cosine similarity as kernel function.
Table 4: Comparison Conditional Vendi scores for algorithms under kernel function sensitivity test.
Algorithm
Conditional Vendi Score
DAK-UCB
9.92
Sixture-DAK-UGG PAK-UCB
12.53
8.75
Random
11.82
One-Arm Oracle
7.99
Sensitivity to the Diversity Term Scalar Hyperparameter: In this experiment, we used the models from (PixArt-alpha) 2024) and (Runway-ML) 2023) to generate images an athlete. We generated hundred images from each model. As suggested by the CLIP, RKE, and Vendi scores in Figure 20 PixArt generated images with higher fidelity, while Stable Diffusion produced more diverse images. deteet this trade-off and monitor the selection ratios and final metrics for diversity and fidelity under different values (the diversity term multiplier). we report the results in Tables and 6
To determine if algorithm
30

---

Published as a conference paper at ICLR 2026
Table 5: Model selection ratio for different À hyperparameter values in Mixture-AK--CB algo- rithm.
A= 0.0
À == 0.1
== 1.0
10.0
PixArt
65%
43%
19%
12%
Stable Diffusion
35%
57%
81%
88%
Table 6: Performance metrics for different hyperparameter values in MixtureDAA--UCB algo- rithm.
Metric
À = 0.0
À = 0.1
À 1.0
A
10.0
Vendi Score
15.25
22.41
33.39
35,04
RKE Score
4.03
6.07
9.14
9.39
CLIP Score
26.8
26.25
25.89
25.55
PixArt
Stable Diffusion
GARDEN 20
Clip Score:
27.42
Clip Score:
25.29
Score:
RKE Score:
9.22
Vendi Score:
5.99
Vendi Score:
37.00
Figure 20: Clip, Vendi, and RKE scores 100 "image of an athlete" generations, comparing PixArt and Stable Diffusion.
As we increased X, the algorithm preferred Stable Diffusion to maximize diversity metrics, which resulted in corresponding decrease in fidelity metrics. This trade-off is clearly observable in our results. We also observed a gender bias: PixArt tended to generate male athletes, while Stable Diffusion tended to generate female athletes. Our algorithm accounted for gender fairness while enhancing diversity.
31

---

Published as a conference paper at ICLR 2026
Visualization of Prompt Correlations:
(a) Kernel matrix for Experiment
(b) Kerel matrix for Experiment 13
Figure 21: Correlation among prompts in datasets used in Experiments.
APPENDIX D STATEMENT ON THE USE OF LARGE LANGUAGE MODELS (LLMs)
LLMs used solely for proofreading and polishing the language of this manuscript, well as for generating the prompts of the numerical experiments for the prompt-aware model selection task. All technical content developed entirely by the authors.
32


Reviews:
[{"review_id": "GrDwzWbOS8", "reviewer": "Reviewer_4Gv8", "summary": "The paper studies the online selection of generative models for given prompts, with explicit consideration of both fidelity and diversity in the generated outputs. \nThe existing prompt-aware selection methods only focus on fidelity (e.g., CLIP-Score), often overlooking diversity, which can lead to repetitive or biased outputs. To overcome this, the authors propose a method, Diversity-Aware Kernelized Upper Confidence Bound (DAK-UCB), that extends the kernelized UCB framework by incorporating diversity metrics \n(joint kernel distance (JKD) and joint Rényi kernel entropy (JRKE)) into the selection process to balance fidelity and diversity. The authors also propose a mixture variant (Mixture-DAK-UCB) that allows for prompt-conditioned mixtures of models. The approach is validated on text-to-image and language model generation tasks, showing improvements in both diversity and fidelity metrics compared to baselines.", "questions": "Please address the weaknesses of the paper.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "**The following are the strengths of the paper:**\n1. This paper considers the online selection of generative models for given prompts that explicitly consider both fidelity and diversity in the generated outputs, overcoming repetitive or biased outputs.\n\n2. The authors propose a new algorithm, Diversity-Aware Kernelized Upper Confidence Bound (DAK-UCB), that incorporates diversity metrics into the selection process to balance fidelity and diversity, while having theoretical guarantees on the regret upper bound (better bound for Sup-DAK-UCB).\n\n3. The authors empirically validated the proposed algorithm on text-to-image and language model generation tasks, showing improvements in both diversity and fidelity metrics compared to baselines.", "weaknesses": "**The following are the weaknesses of the paper:**\n1. The kernel methods and the estimators may lead to high computational overhead (due to kernel matrices), especially when the number of prompts (or models) increases. There should be more details on how the computational cost varies with the increase in samples (number of iterations)\n\n2. It is unclear why the paper only focuses on JKD and JRKE, as there are other diversity metrics.\n\n3. There is not much novelty in the regret bounds, as the derivation depends on existing results, which make many assumptions (normalized kernels, sub-Gaussian noise, RKHS boundedness), which may not hold in all practical settings, especially for LLMs.\n\n4. The diversity-aware selection may degrade fidelity, and there needs to be a discussion about this trade-off. Furthermore, there should be a discussion about the choice of kernel functions, as it may depend on the prompts and models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "3UsEFp8kIL", "reviewer": "Reviewer_rNZi", "summary": "This paper recasts prompt‑conditioned model routing as a contextual bandit that balances generation quality and diversity when choosing among generative models. Building on Upper Confidence Bound, which optimizes only fidelity, the proposed DAK‑UCB jointly estimates prompt fidelity and diversity using separate regressors and an optimism bonus. The method also extends to prompt‑conditioned mixture selection that learns routing probabilities. The authors evaluate DAK‑UCB on the MS‑COCO validation set and on a synthetic dataset created with GPT‑4o to assess its effectiveness.", "questions": "See “Weakness”", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper presents a complete and clearly derived theory. It details how KD and RKE are extended to metrics suitable for joint variables and provides the theoretical basis for a diversity‑aware extension of UCB.\n2. DAK‑UCB is applied to image generation, text generation, and image captioning, and the experiments validate the approach across these tasks.", "weaknesses": "1. The only major difference from PAK‑UCB is the shift from prompt‑aware selection alone to a joint selection based on both prompt fidelity and diversity. Many other aspects, including the theoretical tools and the choice of backbone models in experiments, overlap substantially, which reduces the perceived novelty.\n2. The DAK‑UCB score relies on CLIP text embeddings and DINOv2 image embeddings as the algorithm’s inputs, and the evaluation uses the same embeddings. Prior work indicates that CLIP and DINOv2 can encode biases stemming from imbalances in their training data across geography and race. Using biased embeddings may affect distances between text‑image vector pairs and thereby influence routing, yet the paper does not analyze whether this impacts DAK‑UCB’s outcomes.\n3. In this work, fidelity refers only to alignment between an image and its prompt, not to the realism of the image itself. It therefore remains unclear whether DAK‑UCB can select diffusion models that produce more realistic images, for example better FID.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "kVUOLpg0pd", "reviewer": "Reviewer_a2p3", "summary": "The paper proposes DAK-UCB, a contextual-bandit algorithm for the online selection of generative models with diversity considerations. It incorporates fidelity and diversity-related metrics into the selection process. The experimental results demonstrate considerable improvements.", "questions": "See the weakness.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1) The topic of this paper is meaningful. The new method of delicate online model selection is valuable for the community.\n2) The experimental results on benchmarks demonstrate considerable improvements.", "weaknesses": "1) I do wander the cost of Mixture-DAK-UCB: it solves quadratic program each round and stores matrices, will this bring scalability concerns? \n2) More details and analysis of the bias evaluation beyond gender is necessary.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ZY03vYv9Ce", "reviewer": "Reviewer_E2Mq", "summary": "The authors propose DAK-UCB a diversity-aware bandit algorithm for selecting the best generative model for a sequence of prompts. The main contribution lies in adding diversity scores to the bandit as well as allowing for a mixture of generative models, instead of selecting a single model per prompt. The experiments show that, compared to simple baselines, the proposed approach 1) achieves better diversity-aware scores, 2) is able to discover the most diverse generative model and 3) does not fall for the trap of selecting generative models that increase diversity in images but based on irrelevant prompts.", "questions": "- I am not convinced by the practical use case of the method. If we are interested in diversity along a certain dimension, could we not specify an appropriate prior distribution of prompts, e.g., a 50/50 gender distribution, use existing algorithms to select the most appropriate generative model, and then conditionally sample 50% male and 50% female images? What would be the limits of this approach?\n- When do I need to rerun the algorithm? For instance, is a re-training required whenever the user prompt or prompt distribution changes? Could any information from the previous model be salvaged for a warm start?\n- In Equation 6 and 7, $t'$ is just another prompt drawn from $P_t$. Should this step not appear in the algorithm?\n- In Figure 3, it seems unnatural to have models conditioned on cat or dog classes but keep the prompt about a general \"animal\". This seems rather detached from how the conditioning is applied in practice. Would you expect different results when being more specific in the prompt?\n- In Figure 4, why does the mixture DAK-UCB lead to a worse CLIP score than DAK-UCB, should the former not be nested in the latter?\n- How to interpret the expert selection ratio in Figure 4? Should we not expect to select each cluster 25% of the time (assuming the cat, dog, car and cake clusters have equal size)?\n- Can the model be used in other domains such as protein generation, where diversity has a clear benefit?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Using UCB or a mixture UCB for the selection of a generative model is not new. However, the authors add diversity scores as an additional criterion, which is an original idea. Going beyond the images themselves and incorporating further available information is important in making a more informed choice between generative models.\n- The work is grounded in a thorough mathematical argument that is worth sharing with the community and may prompt further research.\n- The authors conduct many experiments trying to showcase the strengths of their method in various different contexts.\n- The experimental results look promising: Making UCB diversity-aware appears to work as expected, both on (conditional) image generation and text / caption generation tasks.\n- The appendix includes ablation studies on the sensitivity to hyperparameter choices, such as the Kernel function.", "weaknesses": "- The motivation is weak. For example, to get diverse data with respect to gender, I can just define an appropriate prior distribution of prompts. Hence, the practical benefit over existing models is not always clear.\n- $D_g$ is defined twice (in Equation 8 and in line 312 with different notation). $s_g$ is defined twice as well (in Equation 8 and in line 310 with different notation). In general this part is rather inconsistent in notation. Similarly, first $t$ is defined as a text prompt (line 186) and later as a round (line 309).\n- I would have expected a comparison to the already existing, and cited, mixture UCB to show the value of adding diversity scores.\n- The presented metrics in Figure 2 are the same ones used for the optimization. Therefore, it is no surprise that the proposed models perform best here. Possibly using additional metrics that are not directly part of the algorithm could strengthen the argument.\n- Some implementation details remain unclear. For example, after equation 6 and 7 the prompt $t'$ is suddenly no longer considered.\n- The method balances fidelity and diversity but the costs in terms of fidelity have not been directly addressed, e.g., there is no CLIP Score in Figure 2.\n- The actual costs associated with using this approach for increasing diversity remain unclear. No running times for the experiments are mentioned.  I expect the exploration phase to incur monetary costs that may not be trivial, depending the number of generative models, the complexity of the prompt, etc. While regret bounds are provided, the actual incurred costs for the experiments would be interesting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)