DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative Models
Donya Jafari, Farzan Farnia
TL;DR
This paper introduces DAK-UCB, a novel online model selection algorithm for generative AI that balances prompt fidelity and output diversity, improving response variety without sacrificing quality.
Contribution
It presents a diversity-aware contextual bandit framework for model selection, integrating fidelity and diversity metrics, and demonstrates its effectiveness in promoting diverse, high-quality generations.
Findings
DAK-UCB effectively balances diversity and fidelity in model selection.
The method outperforms fidelity-only approaches in promoting output diversity.
Experimental results validate the framework's ability to generate diverse responses.
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…
Peer Reviews
Decision·ICLR 2026 Poster
**The following are the strengths of the paper:** 1. This paper considers the online selection of generative models for given prompts that explicitly consider both fidelity and diversity in the generated outputs, overcoming repetitive or biased outputs. 2. The authors propose a new algorithm, Diversity-Aware Kernelized Upper Confidence Bound (DAK-UCB), that incorporates diversity metrics into the selection process to balance fidelity and diversity, while having theoretical guarantees on the reg
**The following are the weaknesses of the paper:** 1. The kernel methods and the estimators may lead to high computational overhead (due to kernel matrices), especially when the number of prompts (or models) increases. There should be more details on how the computational cost varies with the increase in samples (number of iterations) 2. It is unclear why the paper only focuses on JKD and JRKE, as there are other diversity metrics. 3. There is not much novelty in the regret bounds, as the deri
1. The paper presents a complete and clearly derived theory. It details how KD and RKE are extended to metrics suitable for joint variables and provides the theoretical basis for a diversity‑aware extension of UCB. 2. DAK‑UCB is applied to image generation, text generation, and image captioning, and the experiments validate the approach across these tasks.
1. The only major difference from PAK‑UCB is the shift from prompt‑aware selection alone to a joint selection based on both prompt fidelity and diversity. Many other aspects, including the theoretical tools and the choice of backbone models in experiments, overlap substantially, which reduces the perceived novelty. 2. The DAK‑UCB score relies on CLIP text embeddings and DINOv2 image embeddings as the algorithm’s inputs, and the evaluation uses the same embeddings. Prior work indicates that CLIP
1) The topic of this paper is meaningful. The new method of delicate online model selection is valuable for the community. 2) The experimental results on benchmarks demonstrate considerable improvements.
1) I do wander the cost of Mixture-DAK-UCB: it solves quadratic program each round and stores matrices, will this bring scalability concerns? 2) More details and analysis of the bias evaluation beyond gender is necessary.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsAdvanced Bandit Algorithms Research · Artificial Intelligence in Healthcare and Education · Explainable Artificial Intelligence (XAI)