Efficient Contextual LLM Cascades through Budget-Constrained Policy Learning

TL;DR

This paper introduces TREACLE, a reinforcement learning policy that optimally selects among multiple large language models and prompts to minimize costs while meeting accuracy, latency, and budget constraints in question-answering tasks.

Contribution

The paper presents TREACLE, a novel RL-based approach for dynamic model and prompt selection that considers context and constraints, improving cost efficiency in LLM inference.

Findings

Achieves up to 85% cost savings compared to baselines.

Maintains high accuracy while reducing costs.

Enables flexible trade-offs between accuracy and expense.

Abstract

Recent successes in natural language processing have led to the proliferation of large language models (LLMs) by multiple providers. Each LLM offering has different inference accuracy, monetary cost, and latency, and their accuracy further depends on the exact wording of the question (i.e., the specific prompt). At the same time, users often have a limit on monetary budget and latency to answer all their questions, and they do not know which LLMs to choose for each question to meet their accuracy and long term budget requirements. To navigate this rich design space, we propose TREACLE ($\underline{T}$hrifty $\underline{Rea}$soning via $\underline{C}$ontext-Aware $\underline{L}$LM and Prompt S$\underline{e}$lection), a reinforcement learning policy that jointly selects the model and prompting scheme while respecting the user's monetary cost and latency constraints. TREACLE uses the problem context, including question text embeddings (reflecting the type or difficulty of a query) and the response history (reflecting the consistency of previous responses) to make smart decisions. Our evaluations on standard reasoning datasets (GSM8K, CSQA, and LLC) with various LLMs and prompts show that TREACLE enables cost savings of up to 85% compared to baselines, while maintaining high accuracy. Importantly, it provides the user with the ability to gracefully trade off accuracy for cost.