Efficient Contextual LLM Cascades through Budget-Constrained Policy Learning

W1: Dynamic selection of models and re-querying could lead to increased computational costs and delays

As the reviewer points out, latency plays a crucial role in real-time settings (e.g., voice assistants). To address scenarios requiring real-time performance, we incorporated latency as a component of the cost constraint (see L142). TREACLE allows users to choose the trade-off between monetary cost and latency by adjusting a trade-off coefficient $\beta$, where $cost=\text{latency} + \beta * \text{monetary price}$. This enables the users to choose the balance between monetary cost and latency according to their specific needs. The rightmost two subfigures in Figure 3 show the performance when latency is included in the cost constraint, using real end-to-end latency values (including computation and communication) that we measured by querying Llama and GPT models (see Figure 13 in the Appendix). Further experiments regarding latency are described in Appendices C.1 and C.4. We conducted experiments when the API query latency varied over time; these results are summarized in Table 2 in Appendix C.1 and show that TREACLE can adapt to API querying latency that differs across models and time. Figure 15 in Appendix C.4 shows results with latency in the cost function for more datasets and $\beta$ values. Finally, an additional component of computation latency is the RL model itself. Our RL model is a two-layer neural network, so the inference time is very fast (on the order of ms) and hence negligible.

W2: Users may not always provide clear or consistent feedback.

To clarify, we do not rely on user feedback. Rather, during the training phase, the RL learns a policy to maximize rewards using the correct answers from the datasets, and executes that trained policy during the test phase (see L197). The main form of “feedback” during the test phase is the consistency of the responses returned by the LLMs, as in prior work [20]. However, we found that response consistency alone is insufficient feedback, so TREACLE combines response consistency with prompt text embedding, for the first time, to enhance the overall effectiveness. Adding user feedback, along the lines of active learning, could be an interesting future extension for our framework.

W3: The architecture of the LLM itself has not been changed

We believe that adding reinforcement learning on top of existing LLMs is a strength of our framework, since the modular design enables incorporating new LLMs that are constantly emerging. TREACLE’s framework is generalizable and we investigated its adaptability to new types of questions (Section 5.2.3), unseen data with harder questions (Section 5.2.2), and new LLMs (Section 5.2.1). For example, to understand whether TREACLE can adapt to new types of questions, we conducted experiments (L384-394). The base model is trained using commonsense reasoning questions (CSQA dataset), and the new unseen question type are math problems (GSM8K dataset). The results in Figure 10b show that TREACLE can achieve high accuracy on new question types with only a minimal amount of extra training (i.e., “Fine-tune on 200 GSM8K” is close to “Train on GSM8K”). This minimal amount of extra training is done by freezing the base RL policy trained on CSQA and fine-tuning the “text embedding” feature in the state vector using GSM8K.

Q1: Is the setting in lines 103-107 somewhat far from reality?

The setting of re-querying the same question multiple times has been established in the literature [16,20]. In practice, there are a limited number of unique combinations of language models and prompts. In our experiments, there are 6 possible combinations. The reinforcement learning (RL) model is designed to be general; it does not need to be trained with every possible question variation and answer to choose the best action.

Q2: How large is the actual action space?

There are 3 possible actions in the action space, no matter which prompt or model was used on the previous query: Return the current response, re-query with the same model-prompt pair, or select a new model-prompt pair (from the next option in the cascade, not choosing from all possible model-prompt combinations). We will elaborate on this in L188.

Q3: Related work

Compared to our work, firstly, we are the only work that considers latency constraints. Some scenarios require high real-time performance, as mentioned by Reviewer gLy5. Second, we take into account both LLMs and prompts, since performance and cost are influenced by both factors. The related papers only focus on one of them. Third, we show generalization to unseen tasks and that our method is pretty sample efficient. This is related to the online decision exploration cost mentioned in (2). Fourth, unlike the mentioned papers, we treat the LLMs and prompt pair as a black box, without training or fine-tuning them. Finally, we determine different model and prompt pairs for each sample, rather than for the entire task. We will add these papers to the table in the related work section, and hope to add them as baselines once their code is released.

Q4: Where are "finetune" and "scratch" mentioned in Figure 7?

“finetune” (solid blue/orange lines) means that we start with a model that was trained using the old API prices and language models (LLMs), and then fine-tune it with new state-action trajectories collected using the new API prices and LLMs. In contrast, “scratch” (dashed blue/orange lines) means that we train from scratch, i.e., we initialize the reinforcement learning (RL) model randomly and train it directly with the newly collected trajectories. We will add these clarifications to the text of the paper.

L1: Cost of collecting training data

As mentioned in the paper, we plan to release the training datasets and code for reproducibility so that others can avoid the cost of collecting training data and adapt the framework to other LLMs and tasks.

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.

Above, the last data point (budget of $6) shows a **24.09$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.

We greatly appreciate the reviewer’s detailed and constructive feedback.

Weakness 2: Main weaknesses of the paper are that the method overlooks the actual cost.

We acknowledge the reviewer's concern and agree that our method uses a total budget constraint which minimizes the individual query costs in an indirect fashion. In general, the choice of cost function (e.g. the total budget constraint or individual costs or both) is a design decision. The constraint budget formulation we adopt is widely utilized in resource allocation problems and is typically more difficult to solve/enforce compared to unconstrained RL. That said, our approach is flexible and can also directly minimize the actual cost through a penalized reward formulation. Concretely, following the reviewer’s suggestion, we examined a penalized variation of TREACLE where the reward formulation is (see PDF in rebuttal for all reviewers for precise form) $\mathbb{E}[query\\_acc - \lambda\cdot query\\_cost]~~~subject~to~~~total\\_cost \leq B$ Here, $\lambda\cdot query\\_cost$ is the penalty term we incorporate. This penalized form brings us to Weakness 3 discussed below.

Weakness 3: Unnecessary high actual costs on simple questions.

We have conducted experiments on TREACLE+penalty which are provided in Figures 1 and 2 of the PDF.

• Over the set of easy questions, penalty promotes the use of cheaper models and matches the single model baseline without loss of accuracy. This highlights that penalty aids the efficient use of budget in line with the reviewer’s intuition. The TREACLE+penalty model tends to solve easy questions with less powerful LLMs and the budget spent decreases to 0.717 which is almost the same as the single model baseline while also achieving on-par accuracy, 0.987 (TREACLE+penalty) v.s. 0.986 (single model).

• In general, penalty term will result in a budget-accuracy tradeoff. This is because we don’t exactly know the optimal model for solving a particular question. When evaluating TREACLE+penalty over all queries, we find that penalty improves the budget utilization of TREACLE (actual cost decrease from 0.97 to 0.94) however it also results in (minor) accuracy degradation (accuracy decrease from 0.92 to 0.90).

Weakness 1: Insufficient to differentiate simple and difficult questions in model selection.

The results actually show that our method makes intelligent choices in a budget-aware fashion. Concretely:

• When the budget is low (0.05), the model correctly prioritizes easy questions, as calling an expensive model for one difficult question would leave more than 20 other questions unanswered, harming the overall performance. For instance, with the settings of Figure 4a, if we pick 5 difficult questions and solve them with GPT-4 (CoT), and use the remaining budget for the rest, the total accuracy drops from 0.31 to 0.20.

• When the budget is sufficiently high, the model opts for using powerful models for all questions to maximize acuracy. This is because relying on smaller models or difficulty estimation stage can lead to errors. For example, the average accuracy of Llama-2-7b (CoT) is only 23.65%. Forcing the model to use Llama-2-7b (CoT) decreases the performance from 92.47% to 92.20%, and for the Majority Voting baseline, from 83.62% to 83.30%. As discussed under Weakness 2 and 3, when we use TREACLE-penalty, the model starts prioritizing Llama-2 even with a high budget.

Question 1: Why MoT cannot be robust to new models

We agree that MoT is training-free; however, the original MoT paper does not address the inclusion of new models. We believe it has two limitations when incorporating new models: Firstly, MoT only allows for two models, a weak and a strong one, and there is no clear way of adding more models. Secondly, MoT uses a threshold to decide whether an answer is accepted or not. There is no guarantee that a fixed threshold will work across multiple distinct models or how to find/optimize such thresholds efficiently. For instance, our calibrated cascade approach is reminiscent of MoT and, empirically, we find that it is not robust to distribution change.

Limitations: It is recommended to apply the method proposed in this paper to more tasks.

Thank you for your suggestion. We plan to include additional tasks in future work. Currently, we have evaluated our approach using three representative datasets. Additionally, we demonstrate that our model can be seamlessly adapted to unseen tasks (L384-395, Fig 10b) and that a single model can effectively handle multiple types of tasks while maintaining a shared budget constraint (L374-383, Fig 10a).

We genuinely appreciate the reviewer’s excellent points which motivated us to study penalized TREACLE. We will incorporate these discussions and evaluations and revise the manuscript accordingly. We hope that this response has addressed their concerns. We would be happy to engage further during the discussion week.

W1: 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.

An approach of considering a batch of questions at once could certainly work, whereas in our framework we consider questions one-by-one. We do this mainly for scalability, as considering multiple questions into the future would greatly increase the dimension of the state and action space and hence the amount of training needed. We believe that our framework can extend to considering batches of questions in the future and it would be interesting to explore.

W2: The state vector is limited, as text embeddings alone may not fully capture the complete characteristics of the prompt.

We agree that text embeddings alone may not capture all characteristics of the prompt. Therefore, our state vector includes both text embeddings and response consistency, which is an indirect measure of another prompt characteristic, its difficulty (for example, difficult prompts tend to have less consistent answers [20]). We are the first to implement this combination of state features, resulting in notable performance improvements. Also, our framework is quite flexible, so additional features relating to prompt characteristics can easily be added.

In greater detail, including both text embeddings and response consistency in the state vector improves performance. One baseline method, FrugalGPT, only uses text embeddings. Another baseline method, Majority Voting, only uses response consistency. We have demonstrated significant improvements of TREACLE over both baselines. We also provide theoretical justification supporting why policies considering response consistency perform well (L216-239).

W3: There is a need for a reliable method to estimate the quality of questions and answers generated by the model.

Indeed, we initially had the same thought, and hence we developed the Calibrated Cascade baseline, which explicitly estimates answer quality (using the same state vector as TREACLE as input). It then uses the estimated answer quality to decide whether to query another LLM. While this baseline generally outperforms the other baselines, it is not as robust as TREACLE particularly when there are shifts in question difficulty, because it does not carefully consider the remaining budget when answering easy vs hard questions. In contrast, TREACLE implicitly estimates answer quality by combining previous answers with text embeddings of the question, in order to decide which LLMs to query.

Q1: improving the visuals to be more readable across different printing methods

We will improve the figures to enhance readability for different printing methods, ensuring they are clear even in black-and-white print.

Thank you for your positive feedback on the manuscript. We appreciate your recognition of the reinforcement-learning based approach to achieve significant cost savings when querying LLMs.