DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative Models

Donya Jafari; Farzan Farnia

cs.LGMar 24, 2026

DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative Models

Authors: Donya Jafari, Farzan Farnia

Organizations: Sharif University of Technology · The Chinese University of Hong Kong

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. 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 $G_{2}$ with higher diversity over model $G_{1}$ , DAK-UCB selected the more diverse $G_{2}$ 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

Panel of Figure 15 — 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.

Panel of Figure 17 — 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 (× $10^{3}$ )
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.

	$\lambda=0.0$	$\lambda=0.1$	$\lambda=1.0$	$\lambda=10.0$
PixArt	65%	43%	19%	12%
Stable Diffusion	35%	57%	81%	88%

Appendix

Table 5: Model selection ratio for different

\lambda

hyperparameter values in Mixture-DAK-UCB algorithm.

Metric	$\lambda=0.0$	$\lambda=0.1$	$\lambda=1.0$	$\lambda=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

\lambda

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.

Panel of Figure 21 — Figure 21: Correlation among prompts in datasets used in Experiments.

LLM Routing

Large language model routing selects models, tools, experts, or inference strategies to balance answer quality, cost, latency, and deployment constraints.

See all 345 papers

Mixture-Greedy for Online Generative Model Selection: Is UCB Necessary in Diversity-Aware Multi-Armed Bandits?
Mar 23, 2026Bahar Dibaei Nia, Farzan FarniaMulti-Armed Bandits Exploration-Exploitation Tradeoff
Department of Computer Science and Engineering, Chinese University of Hong Kong.
No Single Best Model for Diversity: Learning a Router for Sample Diversity
Apr 2, 2026Yuhan Liu, Fangyuan Xu, Vishakh Padmakumar +2Diverse Text Generation LLM Routing
New York University · Stanford University · Carnegie Mellon University
Conditional Vendi Score: Prompt-Aware Diversity Evaluation for Generative AI Models and LLMs
Nov 5, 2024Mohammad Jalali, Azim Ospanov, Amin Gohari +1Generative AI Evaluation Language Model Generation Evaluation
Department of Computer Science and Engineering, The Chinese University of Hong Kong · Department of Information Engineering, The Chinese University of Hong Kong

Mar 23, 2026J/K move · Enter open · S save

Mixture-Greedy for Online Generative Model Selection: Is UCB Necessary in Diversity-Aware Multi-Armed Bandits?

Bahar Dibaei Nia, Farzan Farnia

Department of Computer Science and Engineering, Chinese University of Hong Kong.

Multi-Armed Bandits Exploration-Exploitation Tradeoff Generative Modeling UCB Algorithms

Open paper AskPDF arXiv

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.

Abstract

LLM Routing

Explore similar work

Mixture-Greedy for Online Generative Model Selection: Is UCB Necessary in Diversity-Aware Multi-Armed Bandits?

No Single Best Model for Diversity: Learning a Router for Sample Diversity

Conditional Vendi Score: Prompt-Aware Diversity Evaluation for Generative AI Models and LLMs