[論文レビュー] DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative Models

Donya Jafari, Farzan Farnia|arXiv (Cornell University)|Mar 24, 2026

Advanced Bandit Algorithms Research被引用数 0

ひとこと要約

DAK-UCBはオンラインで生成モデルを選択する際の多様性を考慮したコンテキストバンディット手法を導入し、結合カーネルスコアを用いて忠実度と出力の多様性をバランスさせる。さらに、プロンプト条件付きのミキシングへ拡張し多様性を向上させる。

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.

研究の動機と目的

• デ diversity を考慮したプロンプト対応のモデル選択の必要性を示す。
• Diversity-Aware Kernelized UCB (DAK-UCB) フレームワークを提案し、多様性指標をオンラインのモデル選択へ組み込む。
• 生成出力の多様性をさらに高めるためのプロンプト条件付きミキシング戦略を実装。
• 提案された多様性を意識した手法の理論保証（後悔界限）を提供。
• テキストから画像へ・言語モデルタスクにおける多様性と忠実度の実証的改善を示す。

提案手法

• 共分散プロンプト–出力設定へKDとRKEを拡張するため、結合積カーネルを用いる。
• Joint Kernel Distance (JKD) と Joint-RKE (JRKE) を、プロンプトと出力の2サンプル期待値として定義する。
• 忠実度スコア（例：CLIP-Score）と多様性項を組み込んだ、アームごとにカーネル化されたUCB目的関数へJKD/JRKEを統合する。
• オンライン信頼境界を可能にする、各プロンプトにつき単一サンプルを用いたカーネルリッジ回帰によるプロンプトレベルターゲットの推定。
• Mixture-DAK-UCB バリアントを提供し、制約付き二次計画問題を解くことでプロンプト依存のミキシングを最適化。
• 多様性を意識した拡張がカーネル化UCBの保証を保持することを示す段階的後悔界限を証明。

Figure 1: Comparison of baseline Kernelized-UCB model selection (CLIP-Score fidelity metric) (Hu et al. , 2025a ) 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 t

実験結果

リサーチクエスチョン

• RQ1 プロンプト駆動の生成モデルのオンライン選択において、多様性をどのように組み込むことができるか。
• RQ2 結合カーネル多様性スコア（JKD, JRKE）は、忠実度と多様性のバランスを取る際にモデル選択を改善するか。
• RQ3 プロンプト条件付きミキシングは、単一アーム選択を超えて多様性をさらに高められるか。
• RQ4 DAK-UCB のような多様性を意識した文脈バンディットアプローチにどのような理論的保証（後悔界限）が成立するか。
• RQ5 多様性を意識した手法は、テキストから画像生成やプロンプト対応のLLMルーティングといった実世界のタスクを改善するか。

主な発見

• DAK-UCBは忠実度のみを重視するベースラインと比較して、多様性を意識したモデル選択を促進する。
• 結合JKDとJoint-RKEは、プロンプト対応の2サンプル形式を許容し、1サンプルずつのオンライン推定を可能にする。
• DAK-UCBはMS-COCOベースの実験で多様性指標（Joint-RKE）を高く、忠実度（例：CLIP-Score）も競争力を示す。
• Mixture-DAK-UCBはプロンプト依存のミキシングを可能にし、単一アーム選択よりも多様性をさらに改善する。
• このアプローチは、標準的な仮定の下でカーネル化UCBに類似した正式な後悔界限を提供する。

Figure 2: Performance comparison on JKD score and Joint-RKE for MS-COCO prompt clusters using Kandinsky, SDXL, and GigaGAN.