Library

Efficient Contextual LLM Cascades through Budget-Constrained Policy Learning

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.

cs.CLcs.AIcs.LGAnomaly Detection Techniques and ApplicationsTime Series Analysis and ForecastingMachine Learning and Data ClassificationComputer ScienceEconomics
Source

Paper

Similar papers

Pay for Hints, Not Answers: LLM Shepherding for Cost-Efficient InferenceZiming Dong, Hardik Sharma, E. O’Toole et al.2026 · arXiv.org · 0 citations In Library
PickLLM: Context-Aware RL-Assisted Large Language Model RoutingDimitrios Sikeridis, Dennis Ramdass, P. Pareek2024 · arXiv (Cornell University) · 14 citations In Library
C3PO: Optimized Large Language Model Cascades with Probabilistic Cost Constraints for ReasoningAntonios Valkanas, Soumyasundar Pal, P. Rumiantsev et al.2025 · Neural Information Processing Systems · 5 citations In Library
Dynamic Policy Induction for Adaptive Prompt Optimization: Bridging the Efficiency-Accuracy Gap via Lightweight Reinforcement LearningJie Xu2025 · arXiv.org · 0 citations In Library
Online Multi-LLM Selection via Contextual Bandits Under Unstructured Context EvolutionManhin Poon, Xiangxiang Dai, Xutong Liu et al.2026 · Proceedings of the AAAI Conference on Artificial Intelligence In Library
Online Multi-LLM Selection via Contextual Bandits under Unstructured Context EvolutionM. Poon, Xiangxiang Dai, Xutong Liu et al.2025 · AAAI Conference on Artificial Intelligence · 16 citations In Library
LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM InferenceQichen Fu, Minsik Cho, Thomas Merth et al.2024 · arXiv (Cornell University) · 91 citations In Library
Conformal Constrained Policy Optimization for Cost-Effective LLM AgentsWenwen Si, Sooyong Jang, Insup Lee et al.2026 · Proceedings of the AAAI Conference on Artificial Intelligence In Library
Conformal Constrained Policy Optimization for Cost-Effective LLM AgentsWenwen Si, Sooyong Jang, Insup Lee et al.2025 · AAAI Conference on Artificial Intelligence · 3 citations In Library
Large Language Models Are Multitask Chain-of-Thought Prompting OptimizersFeihu Jin, Ying Tan2025 · IEEE Transactions on Neural Networks and Learning Systems · 0 citations In Library

Peer review

Accept (poster)NeurIPS 2024 · via OpenReview
Reviewer gLy55/10 · confidence 2/52024-06-15

Summary

The paper introduces TREACLE, a reinforcement learning policy designed to optimize the selection of LLMs and prompting schemes based on a user's budget constraints in terms of cost and latency.

Strengths

1. TREACLE enables substantial cost savings compared to existing methods by intelligently choosing among various LLMs and prompts based on the monetary cost, latency, and accuracy. 2. By considering the context of the question, including embeddings and past interactions, TREACLE customizes the prompting strategy to balance accuracy against cost, often using advanced prompting strategies like Chain-of-Thought to improve answer quality at a controlled cost. 3. The system dynamically decides whether to re-query a model based on the current response's consistency and the remaining budget, which helps in refining the answers further without exceeding budget constraints.

Weaknesses

1. Dynamic selection of models and re-querying could lead to increased computational costs and delays, especially in scenarios requiring high real-time performance. Although the system is designed to save costs, frequent model switching and complex queries might backfire. 2. The reward mechanism mentioned in the text depends on accurate answer feedback to adjust strategies, but in practical applications, users may not always provide clear or consistent feedback. This could lead to instability during the learning process and inaccuracies in reward distribution. 3. My biggest concern is that the architecture of the LLM itself has not been changed. It merely adds additional reinforcement learning, which seems overly reliant on data, and might perform poorly on new types of questions or unseen data, limiting the model's generalizability.

Questions

1. Is the setting in lines 103-107 somewhat far from reality? The same question is asked many times; wouldn't that mean, in practice, there is an endless number of questions that need to be learned, and indeed, the same question could be asked in many different ways. Also, users might not always provide feedback on a question, so how would one obtain the reward? 2. How large is the actual action space, particularly with the three different prompt strategies (standard, domain expert, and CoT)? 3. Optimizing the selection of LLMs and prompting schemes with constraints is an intuitive idea. There are already some similar works: (1)- "Which LLM to Play? Convergence-Aware Online Model Selection with Time-Increasing Bandits" (2)- "Cost-Effective Online Multi-LLM Selection with Versatile Reward Models" (3)- "Best Arm Identification for Prompt Learning under a Limited Budget" So, the Table 1 listed by the authors is not comprehensive. If the authors do not have time to conduct new experiments for comparison, you could provide a textual description comparing these works, highlighting the advantages of your own work. 4. Where are "finetune" and "scratch" mentioned in Figure 7?

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

As mentioned by the authors, the RL policy’s budget does not account for the cost of collecting the training data.

Authorsrebuttal2024-08-12

We thank the reviewer for reassessing our work and raising their rating. Regarding Reviewer LPQa's concern: We realized that Figure 1 may not reflect the full picture due to the log-scale of the x-axis. Kindly consider the equivalent table provided below which demonstrates that TREACLE with penalty facilitates significant cost reduction. | Budget | Total Cost (No Penalty) | Accuracy (No Penalty) | Total Cost (With Penalty) | Accuracy (With Penalty) | |----------|---------------------------|-------------------|--------------------|----------------------------| | 0.05 | 0.050 | 0.273 | 0.047 | 0.268 | | 0.15 | 0.150 | 0.799 | 0.149 | 0.791 | | 0.3 | 0.298 | 0.843 | 0.297 | 0.842 | | 0.6 | 0.599 | 0.891 | 0.596 | 0.890 | | 0.9 | 0.884 | 0.909 | 0.876 | 0.907 | | 1.5 | 1.471 | 0.918 | 1.261 | 0.911 | | 3 | 2.826 | 0.920 | 2.471 | 0.917 | | 6 | 5.951 | 0.938 | 4.517 | 0.926 | Above, the last data point (budget of $6) shows a **24.09% reduction in cost** compared to the original (from 5.951 to 4.517), with a corresponding 1.279% decrease in accuracy (from 93.8% to 92.6%). At a budget of $3, the improvement is also evident with **12.56%** lower costs and accuracy drop of only 0.3%. Finally, this table also demonstrates that TREACLE works as intended: - Without penalty, the method fully utilizes the available budget i.e. the actual cost is very close to the max available budget. This is in line with the goal of maximizing accuracy subject to the budget constraint. - With penalty, it utilizes the budget more efficiently (e.g. 24% cheaper) at the cost of slightly reduced accuracy. Note that such an accuracy-cost tradeoff is fundamental and not really avoidable (there is no free lunch; we cannot perfectly know which are the easy questions to save on). Also note that by increasing the penalty parameter that trades off between cost and accuracy, one can achieve a more drastic cost reduction at the expense of (slightly) worse accuracy. We are happy to provide more results with different penalty parameter $\lambda$ choices. In summary, we hope this clarifies that if we wish to minimize actual costs as pointed out by the reviewer, the model with penalty can efficiently allocate resources to avoid additional spending, when accuracy gains start to diminish (using tradeoff parameter $\lambda$). We thank the reviewer again for raising the actual spent cost as an important consideration.

Reviewer LPQa4/10 · confidence 5/52024-07-13

Summary

The paper proposes a reinforcement learning method to select the model and prompting. It combines with monetary cost and latency constraints. The design of the features contains question text embeddings and response history. Experiments studies the cost savings.

Strengths

1, Important problems. 2, Interesting algorithm design 3, Many experiments

Weaknesses

The main weaknesses of the paper are that the method overlooks the actual cost and the experiments are not convincing enough. 1, Insufficient to differentiate simple and difficult questions in model selection. A good model selection mechanism should choose low-cost models for simple questions and more powerful models to improve accuracy for difficult questions. The selected models should differ between simple and complex questions. From Figure 4a, it can be seen that when the budget is low (0.05), the selected models are limited to llama2-7b and llama2-13b, failing to call powerful models for solving difficult problems. However, when the budget is high (10), the selected models no longer include llama2-7b, thus missing the opportunity to use low-cost models to reduce expenditure. 2, Lack minimizing actual cost in method design. The method proposed in this paper is designed to only consider staying within the long-term (maximum) budget, without optimizing for minimal actual cost, which can lead to cost inefficiencies. This may derive from overlooking individual cost optimization. 3, Unnecessary high actual costs on simple questions. In simple questions, the method proposed in this paper shows similar accuracy to existing methods, but it consumes an additional 19% of the actual cost. Figure 8b shows that when the total (maximum) budget is \$1, the method proposed in this paper (TREACLE) performs similarly to Single model and Calibrated cascade in terms of accuracy. However, according to Figure 8c, Single model/Calibrated cascade actually costs 0.69/0.702, while TREACLE costs 0.84 (an additional 19%).

Questions

Why do the authors summarize in Table 1 that MoT, a consistency-based method, cannot be robust to new models? I believe this statement is incorrect because this method is training-free. If it is due to limitations in various model capabilities, I suggest introducing the option of 'partially limited'.

Rating

4

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

It is recommended to apply the method proposed in this paper to more tasks. The tasks in this paper are too limited.

Reviewer EMUg5/10 · confidence 2/52024-07-13

Summary

This paper presents a framework for managing different budgets—such as accuracy, cost, and latency—when utilizing Large Language Models (LLMs) for reasoning tasks. Recognizing that reasoning tasks can be broken down into a series of question-and-answer interactions, the authors propose a method to allocate models of varying sizes to handle these multi-round interactions effectively. To achieve this, they use a reinforcement learning approach to train a model that can estimate the budget space. The efficacy of the proposed method is validated through evaluations on standard benchmarks, including GSM8k, CSQA, and LLC.

Strengths

* Introduces a novel framework for modeling multi-turn reasoning sequences, taking into account the holistic aspects of LLM cost, latency, and other factors. * Empirically evaluates the reinforcement learning-based training and execution of the algorithm in realistic settings using popular datasets. * Provides intriguing real-world observations (Section 5.2.1) regarding the impact of pricing changes and the introduction of new models.

Weaknesses

* The current framework is inadequate if a capable model can plan ahead by considering multiple questions or a trajectory of questions in advance, even while using various models for the answers. * The state vector is limited, as text embeddings alone may not fully capture the complete characteristics of the prompt. * There is a need for a reliable method to estimate the quality of questions and answers generated by the model.

Questions

* Many figures require a color printer and are difficult to read due to solid and translucent lines. The authors should consider improving the visuals to be more readable across different printing methods. * Please also refer to the weaknesses section for additional feedback.

Rating

5

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

Ok.

Reviewer kukn7/10 · confidence 4/52024-07-16

Summary

This paper aims to solve the problem that LLMs can be costly, in particular using technologies such as COT. It proposes to apply RL to select the model and prompting scheme. Experimental results show that the proposed method can maintain the model performance while saving up to 85% costs.

Strengths

1. This paper is generally well-written and is well motivated. 2. The experiments are solid with great cost savings. 3. The idea of using RL to reason over question and response embeddings is interesting.

Weaknesses

No great weakness.

Questions

NA

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Reviewer LPQa2024-08-11

Thank you for the detailed response. I appreciate the design regarding the penalty. However, based on the newly added experiments, Figure 1 in the attachment shows that the cost has not been significantly reduced, and the accuracy has decreased. Therefore, I will maintain the original score.

Authorsrebuttal2024-08-12

Thank you for your time and feedback. We realized that Figure 1 may not reflect the full picture due to the log-scale of the x-axis. Kindly consider the equivalent table provided below which demonstrates that TREACLE with penalty facilitates significant cost reduction. | Budget | Total Cost (No Penalty) | Accuracy (No Penalty) | Total Cost (With Penalty) | Accuracy (With Penalty) | |----------|---------------------------|-------------------|--------------------|----------------------------| | 0.05 | 0.050 | 0.273 | 0.047 | 0.268 | | 0.15 | 0.150 | 0.799 | 0.149 | 0.791 | | 0.3 | 0.298 | 0.843 | 0.297 | 0.842 | | 0.6 | 0.599 | 0.891 | 0.596 | 0.890 | | 0.9 | 0.884 | 0.909 | 0.876 | 0.907 | | 1.5 | 1.471 | 0.918 | 1.261 | 0.911 | | 3 | 2.826 | 0.920 | 2.471 | 0.917 | | 6 | 5.951 | 0.938 | 4.517 | 0.926 | Above, the last data point (budget of $6) shows a **24.09% reduction in cost** compared to the original (from 5.951 to 4.517), with a corresponding 1.279% decrease in accuracy (from 93.8% to 92.6%). At a budget of $3, the improvement is also evident with **12.56%** lower costs and accuracy drop of only 0.3%. Finally, this table also demonstrates that TREACLE works as intended: - Without penalty, the method fully utilizes the available budget i.e. the actual cost is very close to the max available budget. This is in line with the goal of maximizing accuracy subject to the budget constraint. - With penalty, it utilizes the budget more efficiently (e.g. 24% cheaper) at the cost of slightly reduced accuracy. Note that such an accuracy-cost tradeoff is fundamental and not really avoidable (there is no free lunch; we cannot perfectly know which are the easy questions to save on). Also note that by increasing the penalty parameter that trades off between cost and accuracy, one can achieve a more drastic cost reduction at the expense of (slightly) worse accuracy. We are happy to provide more results with different penalty parameter $\lambda$ choices. In summary, we hope this clarifies that if we wish to minimize actual costs as pointed out by the reviewer, the model with penalty can efficiently allocate resources to avoid additional spending, when accuracy gains start to diminish (using tradeoff parameter $\lambda$). We thank the reviewer again for raising the actual spent cost as an important consideration.

Reviewer gLy52024-08-12

Thank you for your response, which has addressed my concerns. However, as Reviewer LPQa mentioned, the main weaknesses of the paper are that the method overlooks the actual cost, and the supplementary experiment TREACLE+penalty does not show a significant effect. Overall, I choose to increase the score by one point, but my confidence is not high.

Reviewer EMUg2024-08-13

Thank you for your rebuttal. However, I still have some concerns that remain unaddressed. The responses provided were not entirely satisfactory. I also share the concerns raised by other reviewers (LPQa, gLy5) regarding the absence of direct cost minimization in the modeling process. The two lines in Figure 1 of the rebuttal PDF (with and without the cost penalty) looks nearly identical. It suggests that the cost penalty term is not working effectively. Even in the table the authors provided in the comment, the cost and accuracy trade-off is not great. e.g. compare a) budget of 1.5 gives the original method 0.918 accuracy with cost of 1.471. b) budget of 3 with the cost penalty method gives 0.917 accuracy with the cost of 2.471. Given these remaining concerns, I will be adjusting my score accordingly.

Authorsrebuttal2024-08-13

Thank you for your valuable feedback. With the cost-penalized formulation, tuning $\lambda$ can further adjust the cost-accuracy tradeoff. For example, with regards to the reviewer's example, for a different parameter setting ($\lambda=100, \text{budget}= 3$; previous figures were for fixed $\lambda=10$), the accuracy is 0.918 and the cost is 1.454. This is comparable to TREACLE without penalty and \$1.5 budget, which has the same accuracy of 0.918 and cost 1.471. In other words, with a larger budget and by adding the cost penalty, we can more finely control the actual cost. Overall, many papers use the budget-constrained setting [1,2,3,4]. We explored the cost penalty based on reviewer feedback and found that the TREACLE framework flexibly extended to such a formulation. The results suggest that the cost penalty formulation can perform well with some parameter tuning, which is a general disadvantage of cost-penalty formulations. The advantage of our original formulation is that it does not require parameter tuning, only a simple total budget setting, which easier for practitioners. We would like to add the cost-penalty results as a subsection to the paper further exploring this alternative formulation and different settings of $\lambda$. [1] Bai, Fan, Alan Ritter, and Wei Xu. "Pre-train or Annotate? Domain Adaptation with a Constrained Budget." Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. 2021. [2] Hoffmann, Jordan, et al. "Training compute-optimal large language models." Proceedings of the 36th International Conference on Neural Information Processing Systems. 2022. [3] Chen, Lingjiao, Matei Zaharia, and James Zou. "Frugalgpt: How to use large language models while reducing cost and improving performance." arXiv preprint arXiv:2305.05176 (2023). [4] Shi, Chengshuai, et al. "Best arm identification for prompt learning under a limited budget." arXiv preprint arXiv:2402.09723 (2024).

Authorsrebuttal2024-08-13

We sincerely thank all reviewers for their detailed and constructive feedback. Since there is a shared concern, we would like to clarify that **the two curves in Figure 1 being close is the expected and the desired behavior**. The constrained and penalized optimization are dual forms, both aiming to allocate resources optimally to maximize accuracy. Thus, for the same level of actual cost, we expect that constrained and penalized forms achieve (approximately) the same accuracy. This is why the **actual-cost vs accuracy pareto-front in Figure 1 are essentially the same for the penalized and constrained forms**. This also means that **our base method can minimize the actual cost by adjusting the budget to match the desired cost level**. In our initial response to Reviewer LPQa, we viewed the total budget as a fixed user-specified quantity. This was to highlight the fact that, by adding penalty, one can improve the utilization of this fixed budget (rather than maxing it out to increase accuracy). However, following the reviewer discussion and potential confusion, we realize that the correct viewpoint is that by adjusting the budget, the algorithm indeed minimizes the actual cost. We hope this clarifies our contribution and Figure 1. Authors

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC