Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Donya Jafari, Farzan Farnia

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.

DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative 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. The code is available at https://github.com/Donya-Jafari/DAK-UCB.
Paper Structure (24 sections, 8 theorems, 48 equations, 21 figures, 6 tables, 3 algorithms)

This paper contains 24 sections, 8 theorems, 48 equations, 21 figures, 6 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Preliminaries
  4. Notations and Definitions
  5. Kernel-based Scores for Generative Models
  6. Diversity-Aware Kernelized Upper-Confidence Bound
  7. Extension of KD and RKE Scores to Conditional Generative Models
  8. Diversity-Aware Online Learning via DAK-UCB
  9. Prompt-Aware Mixture Selection via Quadratic Optimization
  10. Numerical Results
  11. Conclusion
  12. Derivation of the Mixture-DAK-UCB Proxy Objective Function
  13. UCB Formulation of Mixture-DAK-UCB
  14. Regret Analysis of Sup-DAK-UCB (Phased Variant of DAK-UCB)
  15. Assumptions in Theoretical Analysis of DAK-UCB
  16. ...and 9 more sections

Key Result

Proposition 1

For conditional distributions $P_{X|T}, Q_{X|T}$ and reference distribution $P_T$:

Figures (21)

  • Figure 1: Comparison of baseline Kernelized-UCB model selection (CLIP-Score fidelity metric) hu2025online 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.’’
  • ...and 16 more figures

Theorems & Definitions (14)

  • Proposition 1
  • Theorem 1: Informal regret bound for DAK-UCB
  • Proposition 2
  • proof
  • Lemma 1
  • proof
  • Lemma 2
  • proof
  • Lemma 3
  • proof
  • ...and 4 more