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. The code is available at https://github.com/Donya-Jafari/DAK-UCB.
Figures & tables
Figure 1: Comparison of baseline Kernelized-UCB model selection (CLIP-Score fidelity metric) ( Hu et al., 2025a ) vs. our proposed diversity-aware DAK-UCB over T=500 rounds. While the baseline Kernelized-UCB does not favor model G2 with higher diversity over model G1 , DAK-UCB selected the more diverse G2 more frequently.
Figure 2: Performance comparison on JKD score and Joint-RKE for MS-COCO prompt clusters using Kandinsky, SDXL, and GigaGAN.
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.
Figure 4: Expert Selection Ratio and Performance Comparison between DAK-UCB and baselines using the JKD score for diversity term in DAK-UCB.
Figure 5: Vendi diversity comparison across Llama3.2, Qwen2, Gemma3, and their mixture for prompts: “Generate a short sentence about a vibrant city in Northern America.” ,“Generate a short sentence about a vibrant city in Europe.” and “Generate a short sentence about a renowned celebrity.”
Appendix figures & tables22 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Results are averaged over 10 independent trials.
Figure 7: DAK-UCB selection ratio for LLM bias detection experiment.
Figure 8: Selection Ratios of biased arms and performance comparison on Joint-RKE.
Figure 9: Word Clouds of DeepSeek Responses Across Countries
Figure 10: Selection Ratios of image captioning models and performance comparison on KID using JKD metric. Results are averaged over 10 independent trials.
Figure 11: Dataset visualization for image captioning experiment.
Figure 12: Selection Ratios of image captioning models and performance comparison on Joint-RKE using I-JRKE metric.
Figure 13: Illustrative examples from the GPT-Generated prompt dataset and their SDXL-generated visualizations under varying guidance scales.
Figure 14: Cluster-conditioned expert selection ratios and performance comparison against baseline methods.
Prompt
DAK-UCB Selected Arm
Mixture-DAK-UCB Preference Vector
”an animal in the garden”
Arm 3 (unconditioned)
[0.218, 0.063, 0.719 ]
”an animal in the meadow”
Arm 3 (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”
Arm 3 (unconditioned)
[0.230, 0.082, 0.688 ]
”an animal in the desert”
Arm 3 (unconditioned)
[0.226, 0.092, 0.682 ]
Appendix
Table 1: DAK-UCB and Mixture-DAK-UCB Arm Preferences for Sample Prompts
Figure 15: Expert Selection Ratio and Performance Comparison on Joint-RKE.
Figure 16: Performance comparison on KID for MS-COCO prompt clusters using Kandinsky, SDXL, and GigaGAN. Results are averaged over 10 trials.
Figure 17: Left: CLIP-Score vs. noise level showing a monotonic fidelity drop. Right: Diversity in DINOv2 embedding vs. noise level showing reduced diversity under stronger corruption.
Diverse Arm (Blurred)
Limited-Diversity Arm
Selection Ratio
45.32%
54.68%
Appendix
Table 2: Final selection ratios when only the diverse arm is blurred (Radius = 20).
# RFF Features
Final Joint-RKE
Final CLIP
Final KD (× 10−3 )
Elapsed Time (× 103 )
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
Appendix
Table 3: Performance of RFF-based DAK-UCB compared to RBF-kernel DAK-UCB under the setup of Fig. 2 .
Figure 18: Adaptation behavior after introducing the third (diverse) arm at iteration 125. The algorithm subsequently converges to updated selection ratios that incorporate the new arm.
Figure 19: The average CLIP Score and Joint-RKE Diversity score over 10 trials using cosine similarity as kernel function.
Algorithm
Conditional Vendi Score
DAK-UCB
9.92
Mixture-DAK-UCB
12.53
PAK-UCB
8.75
Random
11.82
One-Arm Oracle
7.99
Appendix
Table 4: Comparison of Conditional Vendi scores for algorithms under kernel function sensitivity test.
λ=0.0
λ=0.1
λ=1.0
λ=10.0
PixArt
65%
43%
19%
12%
Stable Diffusion
35%
57%
81%
88%
Appendix
Table 5: Model selection ratio for different λ hyperparameter values in Mixture-DAK-UCB algorithm.
Metric
λ=0.0
λ=0.1
λ=1.0
λ=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
Appendix
Table 6: Performance metrics for different λ hyperparameter values in Mixture-DAK-UCB algorithm.
Figure 20: Clip, Vendi, and RKE scores for 100 ”image of an athlete” generations, comparing PixArt and Stable Diffusion.
Figure 21: Correlation among prompts in datasets used in Experiments.
Large language model routing selects models, tools, experts, or inference strategies to balance answer quality, cost, latency, and deployment constraints.
Efficient selection among multiple generative models is increasingly important in modern generative AI, where sampling from suboptimal models is costly. This problem can be viewed as a multi-armed bandit (MAB) task. Under diversity-aware evaluation scores, a non-degenerate mixture of generators can outperform any individual model, distinguishing this MAB setting from classical best-arm identification. Prior approaches incorporate an Upper Confidence Bound (UCB) exploration bonus into the mixture objective. However, across multiple datasets and evaluation metrics, we observe that the UCB term consistently slows convergence and reduces sample efficiency. In contrast, a simple Mixture-Greedy strategy without explicit UCB-type optimism converges faster and achieves even better performance, particularly for widely used metrics such as FID and Vendi where tight confidence bounds are difficult to construct. We provide theoretical insight explaining this behavior: under structural conditions, diversity-aware objectives induce implicit exploration by favoring interior mixtures, leading to sampling of all arms and sublinear regret guarantees for diversity-based objectives. These results suggest that in diversity-aware multi-armed bandits, e.g., for generative model selection, exploration can arise intrinsically from the objective's geometry.
Bahar Dibaei Nia, Farzan Farnia
Department of Computer Science and Engineering, Chinese University of Hong Kong.
When posed with prompts that permit a large number of valid answers, comprehensively generating them is the first step towards satisfying a wide range of users. In this paper, we study methods to elicit a comprehensive set of valid responses. To evaluate this, we introduce diversity coverage, a metric that measures the total quality scores assigned to each unique answer in the predicted answer set relative to the best possible answer set with the same number of answers. Using this metric, we evaluate 18 LLMs, finding no single model dominates at generating diverse responses to a wide range of open-ended prompts. Yet, per each prompt, there exists a model that outperforms all other models significantly at generating a diverse answer set. Motivated by this finding, we introduce a router that predicts the best model for each query. On NB-WildChat, our trained router outperforms the single best model baseline (26.3% vs 23.8%). We further show generalization to an out-of-domain dataset (NB-Curated) as well as different answer-generation prompting strategies. Our work lays foundation for studying generating comprehensive answers when we have access to a suite of models.
Yuhan Liu, Fangyuan Xu, Vishakh Padmakumar +2
New York University · Stanford University · Carnegie Mellon University
Generative models guided by text prompts are widely evaluated for fidelity and prompt alignment, yet their ability to produce diverse outputs remains underexplored. Existing diversity metrics such as Vendi and RKE, which are based on the von Neumann and Rényi entropies of kernel matrices, were developed for unconditional models and cannot distinguish prompt-induced from model-induced variability. We address this gap by introducing Conditional-Vendi and Conditional-RKE, diversity measures derived from the conditional entropy of positive semidefinite matrices. These scores isolate model-induced diversity in prompt-guided generation, with Conditional-RKE enjoying an O(1/n) convergence rate. For Conditional-Vendi, we introduce a truncated-spectrum approximation that yields scalable and consistent estimates. Experiments on text-to-image, image-captioning, and LLM tasks show that the conditional scores recover ground-truth diversity orderings and can also guide diffusion models toward more diverse samples. The codebase is available at https://github.com/mjalali/conditional-vendi.
Mohammad Jalali, Azim Ospanov, Amin Gohari +1
Department of Computer Science and Engineering, The Chinese University of Hong Kong · Department of Information Engineering, The Chinese University of Hong Kong