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

2026-03-24

AI Summary

This work addresses the limitation of existing generative model selection methods, which predominantly prioritize fidelity while neglecting output diversity, thereby failing to meet users’ demand for varied responses. To bridge this gap, we propose DAK-UCB, the first approach to explicitly incorporate diversity into online generative model selection within a contextual bandit framework, jointly optimizing both fidelity and diversity. Our method introduces a prompt-aware diversity scoring function, leveraging kernel distance and kernel entropy to construct a two-sample expected diversity metric, and employs a kernelized upper confidence bound (UCB) strategy for dynamic model routing. Experimental results demonstrate that DAK-UCB significantly enhances output diversity while maintaining high fidelity—as measured by metrics such as CLIP-Score—enabling efficient and diverse model orchestration.

Technology Category

• Machine Learning: Online Learning & Bandits
• Natural Language Processing: Generation
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Fairness, privacy, and diversity in economic environments
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• diversity
• generative models
• prompt routing
• LLMs
• model selection

Innovation

• diversity-aware model selection
• contextual bandits
• kernelized UCB
• generative AI routing
• prompt-based diversity