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

Computer Science > Machine Learning arXiv:2603.23140 (cs) [Submitted on 24 Mar 2026] Title:DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative Models Authors:Donya Jafari, Farzan Farnia View a PDF of the paper titled DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative Models, by Donya Jafari and 1 other authors View PDF HTML (experimental) 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. Specific...