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

arXiv: https://arxiv.org/abs/2505.18901 · HTML: https://mlsanity.com/paper/2505.18901/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 39 references that are in the map.
- First posted on arXiv 24 May 2025.

## Builds on (most cited first, 21 of 21)

- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [CodeBERT: A Pre-Trained Model for Programming and Natural Languages](https://mlsanity.com/paper/2002.08155/) (2020)
- [BigCodeBench: Benchmarking Code Generation with Diverse Function Calls and Complex Instructions](https://mlsanity.com/paper/2406.15877/) (2024)
- [CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X](https://mlsanity.com/paper/2303.17568/) (2023)
- [Nash Learning from Human Feedback](https://mlsanity.com/paper/2312.00886/) (2023)
- [Training GANs with Optimism](https://mlsanity.com/paper/1711.00141/) (2017)
- [Finite-Time Analysis of Kernelised Contextual Bandits](https://mlsanity.com/paper/1309.6869/) (2013)
- [Improved Optimistic Algorithms for Logistic Bandits](https://mlsanity.com/paper/2002.07530/) (2020)
- [Reward-Directed Conditional Diffusion: Provable Distribution Estimation and Reward Improvement](https://mlsanity.com/paper/2307.07055/) (2023)
- [On the Performance of Thompson Sampling on Logistic Bandits](https://mlsanity.com/paper/1905.04654/) (2019)
- [Prompt-to-Leaderboard](https://mlsanity.com/paper/2502.14855/) (2025)
- [DiffusionAgent: Navigating Expert Models for Agentic Image Generation](https://mlsanity.com/paper/2401.10061/) (2024)
- [The Multi-fidelity Multi-armed Bandit](https://mlsanity.com/paper/1610.09726/) (2016)
- [An Online Learning Approach to Generative Adversarial Networks](https://mlsanity.com/paper/1706.03269/) (2017)
- [Cost Aware Best Arm Identification](https://mlsanity.com/paper/2402.16710/) (2024)
- [Multi-Fidelity Multi-Armed Bandits Revisited](https://mlsanity.com/paper/2306.07761/) (2023)

## Built on by (most cited first, 3 of 3)

- [Towards an Explainable Comparison and Alignment of Feature Embeddings](https://mlsanity.com/paper/2506.06231/) (2025)
- [MMD Guidance: Training-Free Distribution Adaptation for Diffusion Models via Maximum Mean Discrepancy Guidance](https://mlsanity.com/paper/2601.08379/) (2026)
- [DAK-UCB: Diversity-Aware Prompt Routing for LLMs and Generative Models](https://mlsanity.com/paper/2603.23140/) (2026)

## Abstract (arXiv)

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.