Optimising Calls to Large Language Models with Uncertainty-Based Two-Tier Selection

Authors: guillem  ramírez
alexandra  birch
ivan  titov
anna  funder
nick  payne
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-22 19:23:58

Abstract: Researchers and practitioners operating on a limited budget face the well-known cost-performance trade-off dilemma. The challenging decision often centers on whether to use a large LLM with better performance or a smaller one with reduced costs. This has motivated recent research in the optimisation of LLM calls. Either a cascading strategy is used, where a smaller LLM or both are called causally, or a routing strategy is used, where only one model is ever called. This is dependent on a decision criterion in both scenarios which is typically an auxiliary neural model. In this work, we propose a cost-effective solution; we use only the uncertainty of the generations of the small LLM as the decision criterion. We compare our approach with both cascading and routing strategies using three different pairs of pre-trained small and large LLMs, on nine different tasks and against approaches that require an additional neural model. Our experiments reveal this simple solution optimally balances cost and performance, outperforming existing methods on 25 out of 27 experimental setups.

Main paper: 
 introduction: 
 Large Language Models (LLMs) offer a high performance for a wide range of text tasks. Their widespread popularity both in research and industrial applications necessitates an understanding on how to optimally use them. Bigger models tend to have better performance, while smaller models are faster and cheaper to run. Deciding which model to use is a common dilemma for many researchers and practitioners with limited budgets, time-constraints or environmental concerns.Recent works attempt to optimise calls to a set of LLMs. In this work we consider the set-up with two LLMs, where one is more expensive with greater performance than the other. In this scenario, there are two main strategies (Figure 1 ): routing, where a query from a user is directed to only one model based on a decision criterion; cascading, where the query always goes to the cheaper model and may subsequently go to the more expensive model depending on the cheaper model's output. These previous studies use one of these calling strategies, and involve either using an auxiliary model to score an LLM output (Chen et al., 2023; Sakota et al., 2023; Ding et al., 2024; Madaan et al., 2023) or using repeated calls to the small cheaper LLM (Yue et al., 2024; Madaan et al., 2023) .Studies that use an auxiliary model introduce further complexity in the optimisation approach, and it remains unclear when practitioners should rely on these auxiliary models. Not only do they require additional training, but they also usually require specific training data and the auxiliary models may not generalise to other tasks. For studies that rely on repeated calls to the small LLM, this approach can become expensive, undermining the original practical motivation for its use. It is perhaps for these reasons, that neither approach has gained traction among practitioners. However, we question whether these additional models or the repeated calls are required to optimise LLM calls. We hypothesise that they may be unnecessary in many use cases, as we can extract confidence measures from the generations of the small model. To investigate this, Margin Figure 1: Routing (left) attempts to select the LLM with the best cost-accuracy trade-off given an incoming query. In cascading (right), all queries are passed through the small model, and depending on its output, the large LLM is consulted. We propose using a cascading approach that uses the margin of the generations to score outputs from the small LLM.we propose a cascading policy that uses a simple measure of confidence of the small LLM to decide whether the large LLM needs to be called. We evaluate our policy on nine tasks, and for three different pairs of pre-trained small and large LLMs. We focus specifically on shortgeneration tasks, which are widely used and versatile, and do not include long-generation tasks in this study due to the additional complexity they introduce for evaluation. In our experiments on classification, multiple-choice and Question-Answering (QA) tasks, this policy outperforms methods that require additional training.We believe our findings suggest that the optimisation of LLM calls should gravitate towards understanding their existing available signals, as opposed to training and running further auxiliary models. Moreover, simple and cheap policies have a lower entry cost for researchers and practitioners, which has an impact on their adoption. 1The key contributions of this work are as follows:• We propose using the margin of the generations of LLMs to optimise LLM calls. The approach is simple, non-parametric, and does not require data or running multiple calls to the small model. • We test our policy on nine tasks with three different pairs of small and large LLMs, and against relevant neural methods that use auxiliary models. We find that Margin Sampling outperforms the other methods, despite using fewer resources. In addition, we further test our policy in a multi-task set-up and obtain similar results.  
related work llm uncertainty: 
 The margin between the two most likely classes has been widely regarded as an uncertainty measure adopted from early Active Learning literature (Scheffer et al., 2001; Luo et al., 2004) and considered more recently in the context of Knowledge Distillation from LLMs (Baykal et al., 2023; Ramírez et al., 2023) . Other measures of uncertainty exist for LLMs (Baan et al., 2023; Huang et al., 2023) . The primary challenge in text generation lies in the difficulty of distinguishing between various forms of uncertainty. Specifically, when evaluating the confidence level of a text generator, we would ideally want to concentrate on uncertainties related to the change in meaning. This entails distinguishing uncertainties that affect the intended meaning from meaning-preserving variations in the generated text. In this study, we address this challenge by concentrating on tasks that require generating very short sequences of tokens. We find that even a basic approach to uncertainty can be advantageous. Additionally, while it might seem logical to assess uncertainty through multiple text outputs from a model, this would lead to significant computational costs, which are impractical for our purposes. Therefore, we employ a straightforward method that does not necessitate generating multiple text samples.Several methods have been proposed to improve the latency of LLMs, such as speculative decoding (Leviathan et al., 2023) , knowledge distillation (Bucila et al., 2006; Hinton et al., 2015) and model quantisation (Jacob et al., 2018) . However, operational costs rather than latency are the focus of this work. Our method only requires access to the logits of a small pre-trained LLM with no need for its parameters.Recent work deals with the problem of optimising calls to a pool of LLMs (Wang et al., 2024) . Sakota et al. (2023) and Lu et al. (2023) propose to train an auxiliary model that predicts the success of calling each LLM. Similarly, Ding et al. (2024) developed a model that predicts the quantitative benefit of utilizing a smaller language model over a larger one. Chen et al. (2023) proposed cascading: using an auxiliary model to predict the accuracy of the small LLM's output. Madaan et al. (2023) and Zhang et al. (2023) used cascading in conjunction with multiple calls to the small model. Finally, Yue et al. (2024) propose a cascading approach that requires repeated calls of the small LLM for reasoning tasks; we show that this can be simplified to just looking at the most likely tokens for short-generation tasks. Ramírez et al. (2023) showed that the margin on a knowledge-distilled model could optimise calls to the larger LLM. This work was limited to one pre-trained large LLM and a smaller fine-tuned local model. However, many practitioners and researchers may not be able to fine-tune their own models due to budget constraints. We extend these findings to pairs of both small and large pre-trained LLMs, as well as to other generation tasks including a multi-task setup.Concurrent work (Gupta et al., 2024) showed that the uncertainty at the token level can be used to effectively leverage a smaller LLM in a cascade setup. Their focus is on the combination of the multiple tokens involved in longer generation, and propose a supervised method; we propose using the margin, a simple rule that optimises short-generation tasks without needing additional data. In addition, we show its applicability in the multi-task setup.Our proposed method differs from previous studies as it does not require previously annotated data for the task and does not require repeated calls to the small LLM. Finally, it does not require training and deploying an auxiliary model. 
problem definition: 
 In this work, we predict mappings between elements in the input space, X , and the corresponding labels in the output space, Y. We have access to the small and the large LLMs that we can prompt to become predictors f s , f l : X → ∆(Y ), where ∆(Y ) denotes the class of probability distributions over Y. We simulate the online setting, where users send queries sequentially. We have q queries (x 1 , . . . , x q )iid ∼ X and we predict the corresponding labels (y 1 , . . . , y q ). For each incoming query x i , we decide whether to call an LLM based on a calling strategy (see below), and incur a given cost c s (x i ) or c l (x i ) respectively. The average cost of the queries by both models is then given by ĉs = 1 q ∑ i c s (x i ); ĉl = 1 q ∑ i c l (x i ). We assume that ĉs < ĉl .For strategies that require training an additional model, we use a train dataset X train , Y train , and use either the corresponding labels from the small LLM only, f s (X train ), or from both LLMs, f s (X train ) and f l (X train ), depending on the strategy. For the training of the auxiliary models, we follow the original papers as much as possible and perform a hyperparameter search where values are omitted. See Appendix A for a detailed explanation of the training process.Since only the small or the large LLM is called in routing strategies, then for a given target average cost per query c ( ĉs ≤ c ≤ ĉl ), we call the large LLM with a probability p r calculated from the re-arranged form of Equation 1.c = (1 -p r ) ĉs + p r ĉl (1)Random routing For every incoming query we call the large LLM with the probability p r .We train a meta-model to predict the performance of the small LLM only, given an incoming query. If this prediction is below a threshold value related to the probability p r , it indicates the small LLM's performance is insufficient and thus we must call the large LLM.The performance of both the small and large LLMs are modelled in this strategy. We train a meta-model to predict if an incoming query is likely to be better solved by the small LLM than by the large LLM. As in Routing above, if this prediction is below a threshold value related to the probability p r , we call the large LLM, otherwise the small LLM.Since the small LLM is always called, then for a given target average cost per query, c, we call the large LLM with a probability p c calculated from the re-arranged form of Equation 2.c = ĉs + p c ĉl (2)FrugalGPT (Chen et al., 2023) We train a model that, given a query and a candidate answer, predicts if the latter is correct. If this prediction is below a threshold value related to the probability p c , we call the large LLM.We suggest using the uncertainty of the output, namely the margin (Scheffer et al., 2001; Luo et al., 2004) , defined by:EQUATIONwhere k t=1 1 and k t=1 2 are the first and second most likely tokens, respectively, according to the distribution of f s for the first predicted token position, t = 1. One advantage of this approach is that it does not require generating a full sequence to be able to compute uncertainty. Moreover, for the tasks we consider there is generally more uncertainty in the first token. If the margin is below a threshold value related to the probability p c , we call the large LLM.All the investigated strategies require setting a threshold for the decision criterion, and we select a dynamic threshold in this work. An initial threshold is calculated using the first 20 queries. We do not evaluate whether to call the large LLM for these 20 queries, we only obtain outputs from the auxiliary models, or the margin value for Margin Sampling. This may require calling the small LLM depending on the strategy. We then use this distribution to calculate an initial p r or p c -th percentile value. For all subsequent queries, the decision to call the large LLM is made, and the threshold is dynamically updated based on all past queries. We test the effect of the dynamic threshold in additional experiments (Table 10 ) and conclude it does not result in major performance inefficiencies. 
llms: 
 We study three pairs of small and large LLMs in our experiments: Mistral 7B (Jiang et al., 2023) and Mixtral 8x7B (Jiang et al., 2024) ; Llama-2 of size 13B and Llama-2 of size 70B (Touvron et al., 2023) ; GPT-3 (Brown et al., 2020) and GPT-4 (OpenAI, 2023) . We selected these pre-trained LLMs because of their popularity among practitioners, as well as to show robustness across different scales. We have arranged the pairs this way to keep the family of LLMs similar; however, this is not a requirement for any of the calling strategies.For the open-source families (Mistral, Llama-2), all our experiments are done locally in one NVIDIA A100 GPU (80 GB), after applying a 4-bit quantisation. Section C contains more details about the LLMs used.We draw inspiration from Liang et al. (2022) and Ramírez et al. (2023) and choose a wide range of tasks, showcasing different difficulties. The sizes of the datasets, along with the label distribution and accuracy of the LLMs, can be found in Appendix (Section B, Section C).For emotion classification we use ISEAR (Shao et al., 2015) ; for factchecking we use FEVER (Thorne et al., 2018) ; for sentiment analysis we use RT-Polarity (Pang & Lee, 2005) , CR (Ni et al., 2019) , and SST-2 (Socher et al., 2013) . All of these datasets are balanced.We use Openbook (Mihaylov et al., 2018) , a popular multiple-choice dataset that involves common knowledge of the world.We use NaturalQuestions (Kwiatkowski et al., 2019) , that contains real questions from human users; we use Wikifact (Petroni et al., 2019; Goodrich et al., 2019) , which consists of a knowledge base completion problem to test factual knowledge; we use bAbI (Weston et al., 2016) , that tests language understanding and reasoning.For each dataset we set aside 1, 000 data-points for a train dataset that is used only to train the auxiliary models. The remainder is used for the online test set. We use n = 500 data-points from the train dataset in the Routing, HybridLLM and FrugalGPT strategies when training the auxiliary model, unless stated otherwise. For the auxiliary models, we fine-tune DistilBERT (Sanh et al., 2019) Unless stated otherwise, our results assume a simple cost scheme with c s (x i ) = 1 and c l (x i ) = 10, consistent with the pricing of commercial APIs (Wang et al., 2024) and similar to the cost schemes of related work (Madaan et al., 2023; Chen et al., 2023; Lu et al., 2023; Sakota et al., 2023) . We do not take into account in our experiments the latency of running DistilBERT, which we deem negligible compared to running the LLMs. To evaluate accuracy across budgets, we report Area Under the Curve (AUC) of the accuracy divided by ĉlĉs . Bolded results mark best performance, and underlined results mark second-best. We run our experiments with three random seeds, and report average results.Figure 2 : Accuracy curve with respect to budgets. We have averaged results for all the tasks. 
comparison of calling strategies: 
 Table 1 shows the AUC for all five evaluated calling strategies, across the nine tasks and for the three pairs of LLMs. Averaged across all nine tasks (the final column), we see that Margin Sampling outperforms all strategies for all LLM pairs. Across individual tasks, we see that it performs best or second best for seven of the nine tasks consistently across all three LLM pairs. Of all the nine tasks and three LLM pairs, it does not achieve best or second-best performance for only two of the 27 combinations: for ISEAR and for NaturalQuestions for the Mistral 7B and Mixtral 8x7B LLM pair. This is likely due to the poor accuracy of the small LLM Mistral 7B on these tasks (Appendix C). Figure 2 shows how performance changes with budget. We see that Margin Sampling dominates across all budgets, despite not using any training data.Indeed, the performance of Margin Sampling seems to improve as the performance of the cheaper LLM improves, which is to be expected; it achieves its best results when applied on top of GPT-3.FrugalGPT is on average the second-best performing strategy. This is due to its good performance on classification tasks, which is expected as it uses a classifier trained on the task as the auxiliary model. However, it performs worse than the random baseline on the more 1 : Accuracy (AUC) for the three LLM model pairs across the five classification tasks (columns 2-6), the multiple-choice task (column 7) and the three generation tasks (columns 8-10).challenging multiple-choice task, Openbook. FrugalGPT also performs inconsistently for QA tasks; we conclude that FrugalGPT may be satisfactory on relatively easy classification tasks and struggle with harder generation tasks.Finally, Routing and HybridLLM seem to have a good performance in QA tasks while having a worse performance in classification tasks. We note that HybridLLM on average has the same performance as random for the OpenAI models, which is a surprising finding.We have not shown comparisons to approaches that require repeated calls to a small LLM throughout this work, as in preliminary experiments we found they do not perform well (see Section D.2).LLMs are often used to handle various tasks simultaneously. To simulate this scenario, we create an artificial multi-task setting by merging the datasets from all nine tasks. We then sample 10,000 data-points. We split this dataset into a 10% train set (n = 1, 000 data-points) and a 90% online test set.Figure 3 and Table 2 show the results for our multi-task experiments. We see again that Margin Sampling has the best performance. This shows the versatility of this method, that it can be applied across tasks with ease. In contrast, HybridLLM has a poor performance both for Llama-2 and OpenAI models. We found these results still hold when using 5,000 data-points as training data (Appendix D.1).To ensure that auxiliary models are not being unfairly handicapped by a low-data setting, we train them with double the data points as before (n = 1, 000), in spite of this being a possibly infeasible amount of data for many researchers and practitioners. FrugalGPT have been trained with n = 1, 000 data-points.Figure 3 : Accuracy curve with respect to budgets, in the multi-task setting.Table 3 shows the averaged AUC results across tasks for the three LLM pairs. The additional data generally improves the performance of the auxiliary models, however, Margin Sampling still performs competitively, achieving the best performance for two of the three pairs. Table 3 also shows that the gap with FrugalGPT closes and that HybridLLM has again a limited performance on the OpenAI models. Additionally, we run experiments in settings with more data (n = 5, 000), and confirm that Margin Sampling still performs competitively (Table 11 ).While previous studies typically only study a single cost setting (Sakota et al., 2023; Chen et al., 2023; Madaan et al., 2023; Ding et al., 2024) , this could influence findings. We carry out further experiments with alternative cost schemes for the three LLM pairs across all tasks. 2 The cost ratio r cost = c l /c s is sufficient to parameterise the effect of cost. We need only vary c l , and maintain c s = 1. We study values of c l = 2, 5 and 20. Table 4 shows the results of these experiments. Independently of the cost, Margin Sampling appears the best cascading strategy (against FrugalGPT). In addition, we observe that Margin Sampling is the best overall strategy for r cost ≥ 5; for r cost = 2, routing strategies could be preferred. Intuitively, cascading needs the cost of the small LLM to be relatively cheap enough to not sacrifice too many calls to the large LLM (p c < p r ).c l =2 c l =5 c l =20 c l =2 c l =5 c l =20 c l =2 c l =5 c l =20Margin Sampling can also be applied when the two LLMs are from different families; we further test its robustness with pairs Llama-2 13B -Mixtral 8x7B, Mistral 7B -Llama-2 70B. Our findings confirm the effectiveness of Margin Sampling (Table 5 ). 
discussion: 
 In this paper, we have used a simple measure of confidence of the small LLM, known as Margin Sampling, to estimate the uncertainty of an LLM output in short-generation tasks, and we leave for future work the generalisation of this approach to long-generation tasks. Our experiments show that Margin Sampling performs consistently well on a range of short-generation tasks relevant to researchers and practitioners, such as QA and multiplechoice/classification tasks.Existing approaches on the optimisation of LLM calls require training an auxiliary model. We hypothesised that these approaches introduce additional complexity that obfuscates their understanding, and our findings that HybridLLM performs poorly with the OpenAI models lends credibility to this hypothesis. In contrast, we propose moving towards research that leverages information within the LLM itself. Classic notions of the uncertainty of the LLM's generation, such as perplexity, margin or entropy, could be relevant signals to help optimise LLM calls. A natural extension of our work is to generalise it for an arbitrary task length, for which it may be that a global notion of uncertainty also heavily depends on the first token. Our preliminary results with Machine Translation (Appendix D.5) indicate that some further adaptation of our method may be required for tasks where the first token may not be directly or partially the answer.It was beyond the scope of this work to investigate cascading and routing strategies of three or more LLMs. However, we hypothesise our approach may still perform well, as Margin Sampling has shown to be robust to different-sized LLMs. 
conclusions: 
 We have proposed a method for LLM call optimisation that achieves superior performance without the need of an auxiliary model. To the best of our knowledge, the field of LLM call optimisation has not yet gained widespread adoption among practitioners. This may be due to the limited versatility and increased complexity of previous solutions. In contrast, our proposed simple approach can be easily and quickly implemented with most commercial LLMs. We believe that our findings could encourage a new direction in the research of LLM call optimisations. FrugalGPT (QA) We train a binary classifier that predicts if an answer is correct. To do so, we use as the positive class the gold labels. We generate answers with Llama-2 13B and tag them as either positive or negative class depending on whether they match the gold labels.• Input: 'Who wrote 'If There Is I Haven't Found It Yet? ANSWER: Anna Funder'• Output: '0'• Explanation: Anna Funder is an incorrect answer. 
 
 Appendix: Table 6 contains some statistics on the datasets used. All the classification datasets are uniformly distributed. For our experiments, we have reserved 1,000 datapoints for training the scorers (Router, HybridLLM and FrugalGPT) and used the remaining of the datasets for online inference.We load the open-source LLMs with Huggingface's AutoModelForCausalLM.from pretrained, activating load in 4bit. We use models meta-llama/Llama-2-13b-hf, meta-llama/Llama-2-70b-hf, mistralai/Mistral-7B-Instruct-v0.2 and mistralai/Mixtral-8x7B-Instruct-v0.1. We set the temperature to 0 and look at the most likely token.For ISEAR, RT-Polarity, FEVER, OpenbookQA, SST-2 and CR, we use the prompts from Ramírez et al. (2023) (0-shot) . For Wikifact, NaturalQuestions and bAbI we use the prompts from the HELM benchmark (Liang et al., 2022) .For the OpenAI models, we use davinci-002 (GPT-3) and gpt-4. Annotating all the datasets has a cost of around $180. We make our code publicly available: https://github.com/guillemram97/margin llmsWe note that, at the time of writing, the pricing ratio for the OpenAI models that we used (GPT-3 and GPT-4) is $30/million tokens $2/million tokens = 15 


Decision: Pros:

* The proposed method aims at optimizing LLM inference calls by using a novel approach to cascading
* The paper does a thorough evaluation on several models and the results show that the method is an improvement over existing approaches.

Cons: 
* Adding an analysis of long form generation would significantly improve the paper

[At least one review was discounted during the decision process due to quality]

Reviews: review: paper_summary: paper_summary: This work aims at optimizing the LLM inference cost by running user queries through small models and switching to larger models when the performance is not satisfying. 
In contrast to other work that trains router models to predict the poor small model performance, the paper proposes to use the margin token prediction as a metric to decide when to switch between small and large models.
The proposed method is compared to several supervised router approaches including Routing, FrugalGPT, Hybrid LLM, and random. The proposed margin sampling approach outperforms the other methods on most tasks.
strength_weakness: strength_weakness: - The supervised models trained for comparison seem rather weak. Only 500-100 examples are used for training which does not sound optimal. Moreover, the various models are trained on top of distilbert which is a very small model. Since the training cost is negligible, it makes sense to train the compared models on more data and with a much larger model. Running BERT or Roberta for classification would be much smaller than using the actual model. 
- The evaluation uses only short context tasks and it is not clear how the proposed approach would work on a longer context. One concern is that the margin sampling might not be super efficient in larger contexts.
questions: questions: What is the maximum number of examples that can be  used for training the rankers?
paper_summary: paper_summary: The paper proposes a strategy for optimisation of LLM calls - a practical problem when making multiple API calls for solving a problem. There are two main strategies - a cascading strategy and a routing strategy. The strategy choice depends on a decision criterion in an auxiliary neural model.
strength_weakness: strength_weakness: The strategies can be viewed as descriptions of practices that worked for a set of pairs and might not be sufficient for a research paper.
questions: questions: Both small and large models contain billions of parameters. What would constitute the threshold for categorising a model as small or large?
paper_summary: paper_summary: This paper tackles the challenge of optimizing LLM calls to two LLMs that gap in parameter size with margin sampling for a better cost-performance trade-off. Specifically, the proposed method leverages uncertainty, defined as the probability difference of the most and the second most likely tokens, of the small LLM's first predicted token position with a statistically obtained dynamic threshold to decide if a secondary call of the large LLM is necessary. The evaluation conducted on a comprehensive set of short-generation tasks and the comparison with a comprehensive set of strong baseline methods warrant the effectiveness of a margin sampling based calling strategy.
strength_weakness: strength_weakness: - The authors only conducted experiments on short-generation tasks citing the evaluation complexity of long-generation tasks, which leaves a common use-case of LLMs unevaluated. 

- It seems that the threshold is a key factor of the performance, but this paper only conducted experiments with the dynamic thresholding strategy that relies on stats of seen samples. The paper also lacks adequate analysis and assessment of the dynamic thresholding strategy.

- This paper is a perfect fit for an industry track, but it might not be of very significant interest to the research community as evidenced by the rejection reason of the FrugalGPT paper on OpenReview. I'm leaning towards an open mind on research topic but would like to keep a note here for the meta-reviewer's reference.
questions: questions: - Could you provide the variance or standard deviation for the reported results?

- How would the order of samples affect the effectiveness of dynamic thresholding?
paper_summary: paper_summary: The paper addresses the optimization of calls to Language Model (LLMs) within the context of having access to both a smaller and larger LLM. The primary research objective is to determine when to utilize the larger LLM for queries that cannot be adequately resolved using the smaller one.

Main contribution of this work is the validation of Margin Sampling (Scheffer et al. 2001, Luo et al. 2004, Ramirez et al. 2023) as a criterion for selecting the larger LLM. Margin Sampling operates by assessing the difference in probabilities between the most likely and second most likely tokens according to the probability distribution of the smaller LLM for the first predicted token position. The paper further refines this threshold on a per-dataset basis with minimal examples.

The experimental evaluation focuses on various tasks including classification, multiple-choice selection, and short-text generation. Results from these experiments demonstrate the efficacy of margin sampling compared to several baseline approaches.
strength_weakness: strength_weakness: 1. Lack of comparison with the related baseline such as speculative decoding (which also requires access to the token probability distribution as current method).

2. All the evaluation is done for short-form generation tasks, what happens when the method is used for long-form generation such as paraphrase, translation, etc. Even if the results are negative, the discussion on long-form generation is crucial to understand the limitation of proposed approach.

3. Limited novelty as the idea of using Margin sampling is proposed in the similar context by Ramirez et al. 2023.
questions: questions: 1. How does the method perform on long-form tasks? A study on this limitation would be appreciated.

2. What happens when LLMs of different families are used for margin sampling? Is it still a reliable estimator for routing?

3. Since the threshold for routing is dynamic, what is the trend of accurate predictions in the beginning and the end? Does it improve with better thresholds from the previous examples? How crucial is the threshold?

4. What is the standard size of dataset used to train the routing module as used in FrugalGPT and other baselines?

Comments: title: Response to authors' rebuttal
comment: Thank you for providing additional information. I'm convinced of the method's effectiveness and the appropriateness of using dynamic sampling. I have slightly increased the rating assuming the authors will _incorporate additional results and a discussion on long-form generation in the final version of the paper_.
Thanks for the rebuttal. I have raised my rating.