Efficient Contextual LLM Cascades through Budget-Constrained Policy Learning

TL;DR

TREACLE is proposed, a reinforcement learning policy that jointly selects the model and prompting scheme while respecting the user's monetary cost and latency constraints, and provides the user with the ability to gracefully trade off accuracy for cost.

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.

Figures (17)

• Figure 1: ${\sf \small TREACLE}\xspace$ chooses LLMs to achieve high accuracy and $\sim$85% cost reduction, compared to individual LLMs.
• Figure 2: Overview of TREACLE framework. TREACLE decides on the next (LLM, prompt) to query in a context-aware fashion, summarized in the state variable. It can adapt to unseen tasks by projecting the new queries into the text embedding space.
• Figure 3: The performance of various methods for different cost functions and budget constraints. The dashed lines are methods that have ground knowledge, which is impractical but illustrates the best achievable performance.
• Figure 4: Number of times each model is re-queried.
• Figure 5: With and without re-querying. $\alpha=\frac{1}{20}$.

Theorems & Definitions (10)

• Proposition 1
• Definition 1
• Definition 2
• Proposition 2
• Lemma 1
• proof
• proof
• proof
• Proposition 3
• proof