PromptWise: Online Learning for Cost-Aware Prompt Assignment in Generative Models

TL;DR

PromptWise is an online learning framework that efficiently assigns prompts to generative models by balancing performance and cost, significantly reducing service expenses while maintaining output quality.

Contribution

It introduces a cost-aware bandit approach for prompt-model assignment, addressing the gap of cost consideration in existing selection methods.

Findings

Achieves comparable performance to baseline methods

Reduces total service cost substantially

Effective in code generation and translation tasks

Abstract

The rapid advancement of generative AI has provided users with a wide range of well-trained models to address diverse prompts. When selecting a model for a given prompt, users should weigh not only its performance but also its service cost. However, existing model-selection methods typically emphasize performance while overlooking cost differences. In this paper, we introduce PromptWise, an online learning framework that assigns prompts to generative models in a cost-aware manner. PromptWise estimates prompt-model compatibility to select the least expensive model expected to deliver satisfactory outputs. Unlike standard contextual bandits that make a one-shot decision per prompt, PromptWise employs a cost-aware bandit structure that allows sequential model assignments per prompt to reduce total service cost. Through numerical experiments on tasks such as code generation and translation, we demonstrate that PromptWise can achieve performance comparable to baseline selection methods while incurring substantially lower costs. The code is available at: github.com/yannxiaoyanhu/PromptWise.