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.

• The paper introduces DAK-UCB, a contextual-bandit router that combines fidelity with joint kernel distance or inverse Rényi kernel entropy to optimize prompt-relevant diversity from one sample per interaction.
• The method and its mixture extension favor diverse models without sacrificing correctness, achieving the strongest Joint-RKE and kernel-distance results in text-to-image tests and improving Conditional-Vendi diversity for complementary LLMs.
• The phased algorithm has a sublinear regret guarantee under RKHS assumptions, while practical limitations include restricted diversity metrics, quadratic model-pair costs for mixtures, and the absence of guarantees for fully adaptive routing.

Motivation and problem statement

Existing prompt-aware model selection methods for generative AI—both offline selectors such as DiffusionGPT, FrugalGPT, and Prompt-to-Leaderboard, and the online PAK-UCB contextual bandit of Hu et al.—route each prompt to a candidate model based solely on per-sample fidelity scores (e.g., CLIP-Score in text-to-image generation). The authors argue that this is structurally inadequate: fidelity-based rewards are means of sample-level scores, whereas diversity is a group-level property determined by the relative positioning of multiple outputs. Consequently, a selector can converge to models whose individual outputs align well with prompts but which collectively collapse onto narrow modes—a failure with practical consequences for representation of sensitive attributes such as gender or ethnicity.

The paper's motivating example makes this concrete: over $T=500$ MS-COCO prompts, the fidelity-only kernelized UCB baseline selects between an SD-XL instance conditioned on "young male" and an unconditioned SD-XL instance at roughly equal rates, ignoring that the latter produces far more diverse outputs. The paper proposes DAK-UCB (Diversity-Aware Kernelized UCB), a contextual bandit algorithm that augments the standard kernelized UCB objective with a diversity term, together with a mixture extension, Mixture-DAK-UCB, that assigns prompt-dependent probability distributions over models rather than deterministic selections.

Joint kernel scores: JKD and JRKE

The technical core is a family of prompt-conditional diversity scores defined via product kernels $$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$ over the joint (prompt, output) space:

• Joint Kernel Distance (JKD): the kernel distance (squared MMD) between $P_T \cdot P_{X|T}$ and $P_T \cdot Q_{X|T}$, serving as a distribution-matching/correctness score against a reference model or dataset.
• Joint RKE (JRKE): the Rényi kernel entropy of the joint distribution, with the paper optimizing its inverse, I-JRKE, as a diversity penalty.

The key structural result (Proposition 1) is that both scores decompose into expectations of prompt-level functions of a single generated sample—for example, I-JRKE equals $$\mathbb{E}_{t \sim P_T,\, x \sim P_{X|T=t}}[\phi_{\text{I-JRKE}}(t,x)]$$. This two-sample decomposition is what makes the scores compatible with online learning: each interaction round yields an unbiased stochastic label for the diversity function, so kernel ridge regression (KRR) can be run on these labels exactly as it is on fidelity labels, producing confidence bounds of the same form. The authors note explicitly that this structure is specific to JRKE and JKD; arbitrary group-level diversity metrics would not admit this treatment.

Algorithm design

At each round, DAK-UCB treats the embedded prompt as context and compares arms via

$$\widehat{J}_g^{\mathrm{UCB}}(t) = \big(\widehat{s}_g(t) + \beta^{(s)}\widehat{\sigma}_g^{(s)}(t)\big) + \lambda\big(\widehat{D}_g(t) - \beta^{(D)}\widehat{\sigma}_g^{(D)}(t)\big),$$

where $s_g$ is a fidelity score (CLIP-Score in experiments), $D_g$ is the signed diversity reward (negative I-JRKE or negative JKD penalty), and optimism is applied to fidelity while pessimism is applied to the diversity penalty. A single generated sample per round suffices to update both KRR estimators.

For Mixture-DAK-UCB, the I-JRKE of a prompt-dependent mixture $\boldsymbol{\alpha}(t)$ admits the quadratic form $$\mathbb{E}_t[\boldsymbol{\alpha}(t)^\top M(t)\boldsymbol{\alpha}(t)]$$, where $$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$0 collects cross-kernel expectations across model pairs. Mixtures are restricted to a kernel-Lipschitz class $$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$1, under which a proxy objective incurs only an $$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$2 approximation error (Proposition 2 in the appendix). The per-prompt decision reduces to a concave quadratic maximization over the simplex using UCB estimates of fidelity and a PSD-projected estimate of $$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$3, with optional panel sampling (probability $$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$4) to obtain unbiased cross-model labels.

Regret analysis

The regret guarantee applies to a phased variant, Sup-DAK-UCB, following the standard device of staged analysis used for Sup-Kernelized-UCB and Sup-PAK-UCB. Under normalized kernels, sub-Gaussian noise, and RKHS boundedness assumptions on both $$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$5 and $$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$6, the bound is

$$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$7

with $$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$8 information-gain terms for the fidelity and diversity targets. Two caveats deserve emphasis. First, the analysis covers only the phased surrogate, not the fully adaptive DAK-UCB used in practice—the same limitation inherited from prior kernelized bandit work due to statistically correlated selections across rounds. Second, the diversity labels are constructed from stage-frozen archives to preserve conditional independence, meaning new samples are not usable for diversity estimation within the same stage. A corresponding mixture regret bound (Theorem 6 in the appendix) adds a panel-sampling term scaling as $$k_{\text{joint}}([t,x],[t',x']) = k_{\mathcal{T}}(t,t')\cdot k_{\mathcal{X}}(x,x')$$9 plus the $P_T \cdot P_{X|T}$0 Lipschitz-approximation cost, though this portion appears in commented-out material in the source and should be treated as provisional.

Empirical results

Experiments use CLIP embeddings for text and DINOv2 for images, with baselines comprising a one-arm oracle, uniform random selection, and PAK-UCB.

• Text-to-image selection (MS-COCO): across 2000 iterations over ten prompt clusters with Kandinsky, SDXL, and GigaGAN as arms, Mixture-DAK-UCB achieved the highest Joint-RKE diversity score and the best KD score against MS-COCO references.
• Simulated arms with controlled diversity: with three "animal" arms (two SD-XL instances conditioned on "cat" and "dog", one unconditioned over ten animals), DAK-UCB preferentially selected the diverse third arm, whereas CLIP-Score-based PAK-UCB favored a less-diverse arm.
• Prompt-relevant diversity: in a four-expert setup where non-expert arms return irrelevant content, both JKD-based and CLIP+I-JRKE-based DAK-UCB avoided trading correctness for spurious diversity.
• LLM diversity collapse: Llama3.2, Qwen2, and Gemma3 each exhibited persistent, distinct geographic biases when asked about "a vibrant city in North America" (New Orleans, New York City, and Chicago respectively). Because the collapse modes are complementary, the Mixture-DAK-UCB mixture attained substantially higher Conditional-Vendi diversity than any single model—an argument that the optimal diversity-aware policy can be a non-degenerate mixture, consistent with prior findings in the unconditional setting.
• Image captioning: routing among LLaVA, InstructBLIP, and BLIP-2, JKD-based selection improved KID correctness, while I-JRKE-based selection improved Joint-RKE diversity (favoring InstructBLIP over the more faithful LLaVA).

Ablations support robustness: CLIP-Score decreases monotonically under image blur while the DINOv2 diversity metric does not inflate under corruption; asymmetrically blurring only the diverse arm drops its selection ratio to 45.32%, confirming fidelity sensitivity; Random Fourier Feature approximation matches exact RBF-kernel performance with near-linear per-round cost; and the algorithm adapts when a new arm is introduced mid-run. Sweeping the trade-off parameter $P_T \cdot P_{X|T}$1 from 0 to 10 shifts selection from PixArt (higher fidelity) toward Stable Diffusion (higher diversity), raising the Vendi score from 15.25 to 35.04 while CLIP-Score declines modestly from 26.8 to 25.55—and incidentally correcting a gender bias in which PixArt generated male athletes and Stable Diffusion female ones.

Limitations and open questions

Several limitations are acknowledged or implicit. The regret guarantee holds only for the phased Sup-DAK-UCB variant under RKHS regularity assumptions on both target functions, leaving the practical single-pass algorithm without formal guarantees. The diversity score family compatible with the framework is limited to those admitting two-sample expectation decompositions; whether other prompt-aware metrics (e.g., Scendi or Conditional Vendi, used here only for evaluation) can be incorporated remains open. The mixture analysis relies on the kernel-Lipschitz competitor class, introducing an irreducible $P_T \cdot P_{X|T}$2 bias term, and the matrix-learning cost scales with $P_T \cdot P_{X|T}$3 pairs, which may limit scalability to large model pools. Evaluation relies partly on GPT-4o-generated synthetic prompts, and the LLM experiments involve short-form generation tasks whose diversity structure may not transfer to more complex workloads. The authors themselves identify extension to protein, molecular, and graph generative models, and application of the scores in general bandit settings beyond model selection, as unresolved directions.

Conclusion

This paper extends prompt-aware contextual bandit routing of generative models from pure fidelity optimization to explicit fidelity–diversity trade-offs. Its central contribution is identifying the two-sample expectation structure of joint kernel distance and Rényi kernel entropy scores, which permits unbiased per-round diversity labels and thus integration with kernelized UCB confidence bounds, along with a regret guarantee for a phased variant. The Mixture-DAK-UCB extension demonstrates empirically that prompt-dependent mixtures of individually collapsed models can substantially outperform any single model on diversity, a finding relevant wherever multiple generative services exhibit complementary failure modes.