Paper: Language Model Cascades: Token-Level Uncertainty And Beyond (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
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).

Reviewer comment (Evaluation):
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?

Author response:
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,

Reviewer comment (Evaluation):
The author claims that Chow-Sum is overly biased towards deferring longer predictions. However, from Figure3(a) we can obverse 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.

Author response:
**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.**

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Experiments):
What is the performance of Post-Hoc-Embed-1?

Author response:
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.
| Dataset | Post-Hoc-Embed-1 AUC-DF (% improvement over random) |
|----------|-------------------|
|anli-r1 |0.546 (+ 8.80)|
|anli-r2 | 0.450 (+ 2.21) |
|anli-r3 |0.426 (+ 6.07)|
| boolq | 0.840 (+ 2.96)|
| imdb |0.964 (+ 1.23) |
|lambada |0.693 (+ 2.07) |
|mnli |0.719 (+ 11.82) |
|piqa |0.718 (+ 2.17)|
|squad|0.408 (+ 1.65)|
|triviaqa |0.099 (+ 17.65)|
|tydiqa-fi|0.333 (+ 4.47)|
|tydiqa-id|0.242 (+ 4.62)|
|tydiqa-sw|0.162 (+ 5.47)|
|winogrande|0.644 (+ 1.02)|
|wmt-de-fr| 0.404 (+ 4.23)|
|wmt-en-fr| 0.444 (+ 3.40)|
|wmt-fr-en| 0.618 (+ 2.99)|

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.

Reviewer comment (Reproducibility):
How is the $\Phi(x)$ computed in detail?

Author response:
We will clarify how $\Phi(x)$ is computed for each of the Post-Hoc methods with the help of an example. 

Let us say that the example $x$ outputs a prediction which has 5 output tokens. Let the output log probability associated with each output token be $p_i(x)$, where $0 <= i < 4$. Then, the Chow-Sum score is $s_{\rm sum}(x) = \sum_i p_i(x)$ and the Chow-Average score is $s_{\rm avg}(x) =\frac{1}{5} \sum_i p_i(x)$. Let the various quantiles be $s_{\rm quant}(x,\alpha_i)$ where $\alpha_i$ is in [0, 0.01, 0.02,... 0.1, 0.2, 0.3, …, 1.0]. We take 20 quantiles. 

Let the decoder embedding from model 1’s final layer (averaged over all output tokens) be denoted as ${\rm Emb}_{\rm dec,final,avg,m1}(x) \in \mathbb{R}^{n_1}$. 

Let the decoder embedding from model 2’s first layer (for the first output token) be denoted as ${\rm Emb}_{\rm dec,int,first,m2}(x) \in \mathbb{R}^{n_2}$.

**Post-Hoc-Quantile**

For the method Post-Hoc-Quantile, we concatenate the quantiles, Chow-Sum and Chow-Average scores:

$$ \Phi(x) = [ s_{\rm sum}(x), s_{\rm avg}(x), \{ s_{\rm quant}(x,\alpha_i) \} ] \in \mathbb{R}^{22} $$ 

**Post-Hoc-Embed-1**

For the method Post-Hoc-Embed-1, we concatenate the quantiles, Chow-Sum, Chow-Average scores and model 1’s decoder embeddings:

$$ \Phi(x) = [ s_{\rm sum}(x), s_{\rm avg}(x), \{ s_{\rm quant}(x,\alpha_i) \}, {\rm Emb}_{\rm dec,final,avg,m1}(x) ] \in \mathbb{R}^{22 + n_1} $$ 

**Post-Hoc-Embed-1+2**

For the method Post-Hoc-Embed-1+2, we concatenate the quantiles, Chow-Sum, Chow-Average scores, model 1’s decoder embeddings and model 2’s intermediate decoder embeddings:

$$ \Phi(x) = [ s_{\rm sum}(x), s_{\rm avg}(x), \{ s_{\rm quant}(x,\alpha_i) \}, {\rm Emb_{\rm dec,final,avg,m1}}(x), {\rm Emb_{\rm dec,int,first,m2}}(x) ] \in \mathbb{R}^{22 + n_1 + n_2} $$

We have also added this example in Appendix F.5.

Reviewer comment (Evaluation):
What is the performance of the post-hoc-quantile when predicting the golden deferring label?

Author response:
Thanks for the interesting question. The final performance of a cascade is typically measured by the accuracy of the interleaved predictions from the small and large model. Nonetheless, the reviewer’s suggestion is indeed a useful diagnostic. 


To that end, the table below shows the AUC-ROC score for predicting the golden binary deferring labels, which is 1 if model 2 is better than model 1, and 0 otherwise. We see that Chow-Sum and Chow-Average tend to be poorly predictive of this golden label. The recommended Post-Hoc deferral approaches yield significant gains in the AUC-ROC. We also added the table in Appendix F.6.

| Method |MNLI | TriviaQA | WMT DE -> FR|
|-|-|-|-|
|Chow-Sum|0.46|0.31|0.51|
|Chow-Average|0.53|0.53|0.56|
|Chow-Quantile-Best|0.73|0.52|0.62|
|Post-Hoc-Quantile|0.86|0.68|0.62|
|Post-Hoc-Embed-1|0.92|0.70|0.63|
|Post-Hoc-Embed-1+2|0.95|0.71|0.62|

Reviewer comment (Evaluation):
In figure2, figure3 and figure5, all figures use the deferring rate as the x-axis. I am curious about whether we could use the inference time cost as the x-axis.

Author response:
Thanks for the suggestion. Note that the two are linearly related: for a deferral rate of $d%%$, the inference cost is $c_{\rm S} \cdot \frac{100 - d}{100} + (c_{\rm S} + c_{\rm L}) \cdot \frac{d}{100}$, where $c_{\rm S}, c_{\rm L}$ are the costs of the small and large model respectively.
We have also added a plot in Appendix F.2 where we plot the relative latency on the x-axis. In our experiments, we estimate the base model latency to be 36% that of the large model. (Please note that the latency estimate may vary depending on the dataset, because the number of output tokens varies. We estimate the latency based on 500 random examples on XSum dataset; this is intended to serve as a rough estimate of the relative speedup from cascading.)

Reviewer comment (Experiments):
Motivation for the lack of comparison with other uncertainty estimation methods.

Author response:
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.

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

Author response:
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’$.

Reviewer comment (Experiments):
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).)

Author response:
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!

Reviewer comment (Experiments):
For the machine translation evaluation: Is the output generated by beam search? Please speculate on the effect of the hyperparameters used on the generation, do they have an effect on the output length?

Author response:
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.

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.

Reviewer comment (Evaluation):
Please elaborate on the use of output probabilities from the LLMs as uncertainty estimates compared to other methods (e.g. deep ensembles, MC dropout)? (e.g. Baan, Joris et al. “Uncertainty in Natural Language Generation: From Theory to Applications.” ArXiv abs/2307.15703 (2023): n. pag.)

Author response:
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!

Reviewer comment (Experiments):
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. Some of the conclusions of the experiments might be artifacts of these hallucinations.

Author response:
**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.

Reviewer comment (Evaluation):
There was a negative correlation between translation quality and sequence length; was this because many of the long sequences were hallucinations?

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

Reviewer comment (Evaluation):
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.

Author response:
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.

Reviewer comment (Experiments):
Which WMT set was used? It needs to be identified and cited.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.