Language Model Cascades: Token-Level Uncertainty And Beyond

Thanks for the positive feedback and thoughtful comments.

The authors only work with a single base model: FLAN-T5. It's not clear how well these results generalize.

Thanks for the suggestion. Our focus on FLAN-T5 models was motivated by their versatile performance on a range of downstream tasks (owing to instruction tuning), as well as them offering models of different sizes.

There are other methods of confidence estimation beyond logprobs (see, e.g. https://direct.mit.edu/tacl/article/doi/10.1162/tacl_a_00598/117737/Calibrated-Interpretation-Confidence-Estimation-in) -- would like to see more analysis here.

Thanks for this reference. From our reading, the paper considers min aggregation versus mean aggregation for confidence measure (Section 4.4). The paper finds that min aggregation leads to better calibration which we have already included in this work.

We note also that Appendix D covers some other approaches to estimating uncertainty. While some of these may be computationally prohibitive for a cascade setting, it would be of interest to see if others offer value over the information in softmax probabilities. This could be a good direction for future work (as noted in Section 5).

In the "Intermediate embeddings" approach, only decoder representations are used. However, wouldn't it potentially be useful to characterize aspects of the input? I could see that some inputs might be more reasonable for a smaller model; others might have complexities that are more suited to a larger model. Even input length could potentially be useful. Do you have empirical experimentation here?

Yes, we had experimented with both output only, as well as input and output combined representations. But, adding input representations did not seem to help. We have added results for 3 datasets with input representations also in Appendix F.3.

We really appreciate your time and effort in giving detailed comments and feedback. We have updated our paper based on this and have answered your questions above. Please let us know if you have any further questions or comments.

As another line of work, speculative decoding also aims at the trade-off balance between efficiency and performance and I think the authors should discuss or compare the difference between these two lines of work.

Thanks for raising this interesting point. Indeed, the recent line of work on speculative decoding is conceptually related to cascading: both approaches involve orchestrating between a small and large model to improve inference efficiency.

However, they have slightly different use-cases: speculative decoding critically assumes that it is feasible to use the large “verifier” model to score predictions from the small “drafter” model. Note that such scoring needs to be done for every test example. However, this may not always be feasible: e.g., consider a setting where the large model has >100B parameters, which could result in a high latency even when scoring sequences. By contrast, cascading can still be applicable in such cases, since the large model needs only be invoked on samples where the deferral mechanism predicts to forward.

On the flip side, an appealing characteristic of speculative decoding is that it is provably quality-neutral. On the other hand, cascading offers a tradeoff between quality and cost. (Typically, it is empirically observed that the cascade can also be quality neutral at moderate deferral rate to the large model.) It is of considerable interest to study mechanisms that combine the strengths of cascading and speculative decoding; however, this would be worthy of a separate paper in itself.

What is the performance of Post-Hoc-Embed-1?

Thanks for pointing this out. We did not include this in the main paper for the sake of clarity and readability. Nonetheless, here is the performance comparison, summarized via the AUC values for each dataset, and percentage improvement over random.

We see that Post-Hoc-Embed-1 is able to improve upon Post-Hoc-Quantile and Chow-Quantile methods slightly, but is slightly inferior compared to the Post-Hoc-Embed-1+2 method. This intuitively makes sense as this has more information compared to the Post-Hoc-Quantile method but still does not include any information about model 2. We have also added this text in Section 4.3 in the draft.

Thanks for the detailed feedback and insightful comments.

Compared with simple averaging the log probability, the major advantage of quantile is that it reflects more about the overall log probability distribution of the sequence and is more robust to outliers. To highlight the motivation of the proposal, more evidence for the existence of the outlier is expected and the showcase in Figure 1 is not sufficient.

We wish to emphasize that our goal in leveraging probability quantiles is not only to achieve robustness to outliers: it is to obtain a richer view of the distribution than a single summary statistic (like the sum or mean). For example, while the sum of log-probabilities may be low, if even a couple of tokens have high uncertainty, that could indicate a prediction is uncertain (per Figure 1).

Our experimental results confirm this goal is worthwhile: by learning a post-hoc deferral rule on top of quantiles, one can significantly improve the cost-quality tradeoff compared to simply using a single summary statistic (see Table 1).

In Figure 2 and Figure 3, it seems that the best generation performance is obtained in the middle of the curve, other than the endpoint where all examples are deferred. Does this mean that sometimes smaller LM outperform larger ones?

Thanks for raising this: indeed, this indicates that the small model may be superior to the larger one on certain examples. For example, on TriviaQA, we find that on 1.1% of test examples, the FLAN-T5 Base model is correct while the FLAN-T5-Large model is incorrect. This “non-monotonicity” of the performance of models of different sizes has been observed previously, e.g.,

Narayan et al., Predicting on the Edge: Identifying Where a Larger Model Does Better. 2022. Kim et al., Speculative Decoding with Big Little Decoder, 2023,

The author claims that Chow-Sum is overly biased towards deferring longer predictions. However, from Figure3(a) we can observe that when the output length (y-axis) is between 150 words to 250 words, the score of the oracle deferring strategy is smaller than the Chow-Sum, which says the opposite: the chow-sum isn’t biased towards deferring longer predictions when compared with the oracle.

We believe there is a misunderstanding in the interpretation of Figure 3(a). The x-axis denotes the score quantile for different methods. At a given score quantile (say, x=0.2), we consider the behavior of each method when the corresponding fraction of lowest scoring samples are deferred (e.g., samples with lowest 20% scores are deferred). The y-axis represents the average output length of the deferred samples.

At x=0.2, the Chow-Sum method has a y value of 250, whereas the Oracle method has a y value of around 125. So, this means that the first 20% of samples which Chow-Sum defers have an average length of 250, whereas the Oracle method’s deferred samples have an average length of 125. Thus, Chow-Sum prefers to defer longer predictions.

To further illustrate this, we have also added more plots in Appendix F.1 which shows the histogram and CDFs of lengths of outputs deferred by each method at different deferral rates. Each plot considers the samples which are deferred by each method at the mentioned deferral rate. Figure 11 and 12 shows that the Oracle method has a higher fraction of points deferred with lower prediction lengths as compared to the Chow-Sum method at different deferral rates. These plots conclusively provide evidence that Chow-Sum is overly biased towards deferring longer predictions as compared to the Oracle method.

We really appreciate your time and effort in giving detailed comments and feedback. We have updated our paper based on this and have answered your questions above. Please let us know if you have any further questions or comments.

Please speculate for the use of the proposed method for robustness to OOD in MT. different domains can be used to evaluate a change in distribution, is the uncertainty estimate robust to such change?

Thanks for the interesting suggestion. We note that the distinction between OOD and ID data is not immediately clear in our setting of instruction tuned LMs. Specifically, we are operating in a setting where there is a single LM that is first pre-trained, and then fine-tuned on a given instruction-tuning set (the FLAN mixture for FLAN-T5). The latter implicitly contains multiple tasks, and aims to help the model generalize to new tasks. One narrow definition of “OOD” in this setting could be based on whether or not a task is included in the instruction-tuning mixture. For example, the TyDiQA dataset is not included in the FLAN mixture. One might thus consider TyDiQA as, in a narrow sense, “OOD” for FLAN-T5 models. We have included results for TyDiQA in our paper. On one of the three languages for TyDiQA, our method outperforms other methods. For OOD in machine translation, if the reviewer has a specific task in mind, we would be happy to consider experimenting with it.

Thanks for the detailed feedback and encouraging comments!

Please elaborate on the relation/difference of the Chow estimates with proper scoring rules (e.g. NLL, Brier score).

This is a good question. A proper loss (or negative proper scoring rule) is typically defined for a predicted probability and a ground truth label; e.g., we may compute the log-loss $-log \hat{p}(y \mid x)$ for a predicted distribution $\hat{p}( \cdot \mid x )$ and label $y$.

In our setting, we only have the predicted distribution at inference time; the ground truth label $y$ is unknown, and so cannot be used to compute a confidence measure for deferral. With Chow’s rule, we compute the maximum value of $\hat{p}( y’ \mid x )$ over all possible labels $y’$.

Could the proposed estimates be directly compared/evaluated to estimates from deep ensembles instead of a cascade or even jointly? (Lakshminarayanan, Balaji et al. “Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles.” Neural Information Processing Systems (2016).)... Please elaborate on the use of output probabilities from the LLMs as uncertainty estimates compared to other methods (e.g. deep enembles, MC dropout)? (e.g. Baan, Joris et al. “Uncertainty in Natural Language Generation: From Theory to Applications.” ArXiv abs/2307.15703 (2023): n. pag.)

Thanks for the interesting suggestion. First, we note that cascades and ensembles solve slightly different problems. Cascading is often used as a way of improving the prediction efficiency, by selectively leveraging one model from a given family. Ensembling is often used as a way of improving the prediction quality, by suitably leveraging multiple models from a given family.

Deep ensembles are generally understood to provide more reliable uncertainty estimates than a single model. However, this comes at the expense of a higher compute cost than invoking a single model. For cascading, efficiency is paramount: such a higher cost would defeat any gains obtained by the improved uncertainty estimates. A similar concern would arise in approaches based on averaging multiple models using dropout.

One can in fact combine cascading and ensembling, per Wang et al., “Wisdom of Committees: An Overlooked Approach To Faster and More Accurate Models”. However, such approaches have not (to our knowledge) been explored in the language model domain. We agree this could be extremely interesting for future work!

Motivation for the lack of comparison with other uncertainty estimation methods.

As discussed above, many uncertainty estimation procedures involve additional computation (e.g., multiple inferences with a single model in dropout-based approaches and single inference with multiple models in ensemble-based approaches) compared to simply using softmax probability outputs from a single network. Such approaches are less appealing for use in cascades, where the primary goal is to improve efficiency. We have also updated the text in Appendix D to include this discussion.

For the machine translation evaluation: Is the output generated by beam search?

The outputs were generated by greedy decoding. To verify that the results are not an artifact of this decoding strategy, we have also added results for beam search in Appendix F.4 in Figures 15, 16 and 17. The output lengths are on average lower in beam search as compared to greedy decoding. We observe that the same findings still hold in this case.

Please speculate on the effect of the hyperparameters used on the generation, do they have an effect on the output length?

Thanks for raising this important point. We added a plot (Figure 18) in Appendix F.4 which compares the distribution of output lengths using greedy decoding versus beam search (beam size 10). As our intuition suggests (and also pointed out by Reviewer 7G2G), we see that outputs generated by beam search have lower lengths on average as compared to the outputs generated by greedy decoding because they are less prone to get stuck in repetitions. But, as the figure suggests, there is still a considerable probability mass on longer lengths and hence, they are still likely to generate longer outputs with repetitive text. Thus, as the rest of the figures in Appendix F.4 suggest, our insights and methods still generalize to outputs generated using beam search.

Thanks for the detailed feedback and encouraging comments!

Additionally, there are places where the experimental set-up seems to be suboptimal: greedy decoding was used, but this is more prone to hallucination than beam search.

We have added cascade plots with beam search for three datasets (TriviaQA, WMT DE->FR and WMT FR->EN; see Appendix F.4, Figures 15, 16), and a length analysis plot for WMT FR->EN (Appendix F.4, Figure 17). We used a beam size of 10 for this experiment. The output lengths are on average lower in beam search as compared to greedy decoding (as shown in Figure 18 in Appendix F.4 as suggested by reviewer WdbA) which gives more support to the reviewer’s intuition that beam search is less likely to get caught in repetitive hallucinations as compared to greedy decoding but still gets affected by this problem. Thus, we observed similar conclusions to our main results. Thus, our results also extend to decoding with beam search, and are not purely an artifact of greedy decoding.

the Chow-average model favored longer sequences, which there is no intuitive reason for. Could it be because hallucinations often get trapped in loops where the same short phrases get repeated with high probability? This would drive up the average.

We believe the reviewer has the correct intuition. As noted, Chow-average preferred keeping longer sequences and deferring shorter sequences (see (0, 0.2) quantile score range of Figure 3a). Longer sequences got a higher Chow-average score on average for two reasons: (a) the later tokens tend to have higher probability as per Figure 3c, and (b) as the reviewer suggests, there are loops where the same short phrase gets repeated. That is why, after normalization by the sequence length, Chow-average prefers these longer sequences.

there was a negative correlation between translation quality and sequence length; was this because many of the long sequences were hallucinations?

Per discussion in Section 3.5 (see also Figure 4), longer predictions tend to have repetitions, which results in lower translation quality.

--Which WMT set was used? It needs to be identified and cited.

Thanks for pointing this out. We have used WMT 14 FR-> EN, WMT 14 EN-> FR and WMT 19 DE-> FR. We also updated the draft to include these citations in Section 3.4.

--I don't understand how the MNLI experiments work. It is noted in Section 3.4 that MNLI is a multi-class classification problem, but the techniques proposed in this paper are not for classification.

Sorry for the confusion. Following the T5 and FLAN-T5 papers, we treat all problems as finding a text to text mapping. So for MNLI, we encode the classes as strings, namely, “entailment”, “neutral”, “contradiction”. We take the model’s output text and perform a string comparison to the label. We have updated the draft to clarify this.