[논문 리뷰] DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative Models

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

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.

연구 동기 및 목표

• 충실도뿐만 아니라 다양성도 고려하는 프롬프트 인식형 모델 선별의 필요성을 동기 부여합니다.
• 다양성 메트릭을 온라인 모델 선별에 통합하는 다양성 인식 커널화된 UCB(DAK-UCB) 프레임워크를 제안합니다.
• 생성 출력의 다양성을 더욱 향상시키기 위한 프롬프트-조건 혼합 전략을 가능하게 합니다.
• 제안된 다양성 인식 방법에 대한 이론적 보장(후퇴 한계)을 제공합니다.
• 텍스트-투-이미지 및 언어 모델 작업에서 다양성과 충실도에 대한 실증적 향상을 시연합니다.

제안 방법

• 결합 곱 커널을 사용하여 조건부 프롬프트–출력 설정으로 KD와 RKE를 확장합니다.
• 프롬프트와 출력에 대한 두 샘플의 기대값으로 JKD와 JRKE를 정의합니다.
• 충실도 점수(예: CLIP-Score)와 JKD/JRKE를 각 팔의 커널화된 UCB 목표에 다양성 항과 함께 통합합니다.
• 단일 프롬프트당 샘플을 사용하여 커널 릿지 회교로 프롬프트 수준 목표를 추정하고 온라인 신뢰 구간을 가능하게 합니다.
• 제약된 이차 계획법을 해결하여 프롬프트 의존 혼합을 최적화하는 Mixture-DAK-UCB 변형을 제공합니다.
• 다양성 인식 확장이 커널화된 UCB 보장을 유지함을 보이는 단계별 후퇴 경계(bound)를 입증합니다.

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는 충실도만을 기준으로 한 기반선에 비해 다양성 인식 모델 선정을 촉진합니다.
• Joint JKD 및 Joint-RKE 점수는 두 샘플의 프롬프트 인식 형태를 허용하여 프롬프트당 단일 샘플로 온라인 추정을 가능하게 합니다.
• MS-COCO 기반 실험에서 DAK-UCB는 더 높은 다양성 지표(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.