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

Donya Jafari, Farzan Farnia|arXiv (Cornell University)|Mar 24, 2026

Advanced Bandit Algorithms Research0 citations

TL;DR

DAK-UCB introduces a diversity-aware contextual bandit method for online selection of generative models, balancing fidelity and output diversity via joint kernel scores. It also extends to prompt-conditioned mixtures to improve diversity.

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.

Motivation & Objective

• Motivate the need for prompt-aware model selection that accounts for diversity, not just fidelity.
• Propose a Diversity-Aware Kernelized UCB (DAK-UCB) framework that integrates diversity metrics into online model selection.
• Enable prompt-conditioned mixture strategies to further enhance diversity in generated outputs.
• Provide theoretical guarantees (regret bounds) for the proposed diversity-aware method.
• Demonstrate empirical gains in diversity and fidelity on text-to-image and language-model tasks.

Proposed method

• Extend KD and RKE to the conditional prompt–output setting using a joint product kernel.
• Define Joint Kernel Distance (JKD) and Joint-RKE (JRKE) as two-sample expectations over prompts and outputs.
• Integrate JKD/JRKE with fidelity scores (e.g., CLIP-Score) into a per-arm kernelized UCB objective with a diversity term.
• Estimate prompt-level targets via kernel ridge regression using single samples per prompt, enabling online confidence bounds.
• Provide a Mixture-DAK-UCB variant that optimizes prompt-dependent mixtures by solving a constrained quadratic program.
• Prove a phased regret bound showing the diversity-aware extension retains kernelized UCB guarantees.

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 t

Experimental results

Research questions

• RQ1 How can diversity be incorporated into online model selection for generative models driven by prompts?
• RQ2 Do joint kernel diversity scores (JKD, JRKE) improve model selection when balancing fidelity and diversity?
• RQ3 Can prompt-conditioned mixtures further enhance diversity beyond single-arm selections?
• RQ4 What theoretical guarantees (regret bounds) hold for diversity-aware contextual bandit approaches like DAK-UCB?
• RQ5 Do diversity-aware methods improve real-world tasks such as text-to-image generation and prompt-aware LLM routing?

Key findings

• DAK-UCB promotes diversity-aware model selection while maintaining fidelity compared to fidelity-only baselines.
• Joint JKD and Joint-RKE scores admit two-sample, prompt-aware forms enabling online estimation with a single sample per prompt.
• DAK-UCB achieves higher diversity metrics (Joint-RKE) and competitive fidelity (e.g., CLIP-Score) in MS-COCO based experiments.
• Mixture-DAK-UCB enables prompt-dependent mixtures that further improve diversity over single-arm selection.
• The approach provides a formal regret bound analogous to kernelized UCB under standard assumptions.

Figure 2: Performance comparison on JKD score and Joint-RKE for MS-COCO prompt clusters using Kandinsky, SDXL, and GigaGAN.