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

Paper: 2603.23140 Authors: Donya Jafari, Farzan Farnia Categories: cs.LG Venue: ICLR 2026

Abstract

Existing model selection and prompt routing methods for LLMs and generative models focus primarily on fidelity scores, overlooking the diversity of generated outputs. This paper introduces DAK-UCB (Diversity-Aware Kernelized Upper Confidence Bound), a contextual bandit algorithm for online model selection that jointly considers both fidelity and diversity. The method incorporates prompt-aware diversity score functions that decompose into two-sample expectations over prompt-output pairs, applying joint kernel distance and kernel entropy measures. By extending the kernelized UCB framework with a diversity-oriented term, DAK-UCB addresses the critical fairness concern that diversity-unaware selection can lead to restricted representation of sensitive attributes such as gender and ethnicity in generated content. Code is available at github.com/Donya-Jafari/DAK-UCB.

Key Contributions

• Diversity-aware model selection: Formulates the prompt routing problem as a contextual bandit that optimizes for both output fidelity and diversity, moving beyond fidelity-only approaches
• Prompt-aware diversity scores: Introduces diversity score functions that decompose into two-sample expectations over prompt-output pairs, enabling tractable diversity estimation
• Kernel-based diversity measures: Applies joint kernel distance and kernel entropy as principled diversity metrics within the selection framework
• Fairness-motivated design: Demonstrates that diversity-unaware routing can systematically underrepresent sensitive attributes (gender, ethnicity), framing diversity as a fairness requirement
• Extended kernelized UCB: Augments the standard kernelized UCB bandit framework with a diversity-oriented exploration term

Methodology

DAK-UCB extends the kernelized UCB framework for contextual bandits to incorporate diversity considerations:

1. Problem formulation: Online model selection is cast as a contextual bandit problem where the context is the input prompt and the arms are available models/configurations
2. Fidelity-diversity objective: The reward function combines a fidelity score (measuring output quality) with a diversity score (measuring variety across generated outputs)
3. Prompt-aware diversity decomposition: Diversity scores are decomposed into two-sample expectations over prompt-output pairs, making them amenable to kernel-based estimation
4. Joint kernel measures: Two kernel-based diversity metrics are employed — kernel distance (measuring distributional spread) and kernel entropy (measuring information-theoretic diversity)
5. UCB exploration: The upper confidence bound naturally balances exploitation (selecting high-fidelity models) with exploration (trying models that may improve diversity)

Results

• Diversity-unaware model selection methods can lead to restricted representation of sensitive attributes in generated outputs, confirming the fairness motivation
• DAK-UCB effectively balances fidelity and diversity in online model selection, outperforming fidelity-only baselines on combined metrics
• The kernel-based diversity measures provide a principled and computationally tractable way to quantify output variety
• The contextual bandit formulation enables adaptive routing that improves over time as more prompt-output observations are collected

Takeaways

• Fidelity alone is insufficient for model selection in generative systems — diversity of outputs matters for both quality and fairness
• The fairness dimension is underappreciated: routing that maximizes only quality scores can systematically produce outputs with narrow demographic representation
• Kernel methods provide a natural framework for measuring diversity in high-dimensional output spaces without requiring explicit attribute labels
• The contextual bandit formulation is well-suited for online prompt routing, as it handles the exploration-exploitation tradeoff inherent in learning which models produce diverse, high-quality outputs
• DAK-UCB is a modular extension to existing UCB-based routing — the diversity term can be integrated without redesigning the entire selection pipeline