PaperlayerView exact revision on arXiv

Source-linked AI summary

Language Model Cascades: Token-level uncertainty and beyond

Neha Gupta, Harikrishna Narasimhan, Wittawat Jitkrittum, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar

arXiv:2404.10136v1cs.CLcs.AIcs.LG

TL;DR

Generative-LM cascades lack a well-understood deferral rule, and simple sequence confidence can be biased by variable output lengths. The paper studies token-level uncertainty quantiles and learned post-hoc rules; these rules perform consistently well against aggregation baselines, with embeddings providing further improvement.

  • Problem

    For generative LMs, it is unclear how to turn token-level uncertainties across variable-length outputs into a deferral score.

    Source: §1 · p. 2

  • Method

    The paper uses token-level uncertainty quantiles as features for learned post-hoc deferral rules, optionally adding model embeddings.

    Source: §1 · p. 2, §4.1 · p. 10, §4.2 · p. 10

  • Results

    Across NLP benchmarks with FLAN-T5, learned quantile-based deferral is the most consistently performant against natural aggregation baselines, and intermediate large-model embeddings further improve performance.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Token-level uncertainty and intermediate embeddings can improve cost-quality tradeoffs beyond simple sequence-level uncertainty aggregation.

    Source: §5 · p. 11

    Summing or averaging token probabilities can mask highly uncertain individual tokens, including entity tokens relevant to fact-answering.

    Source: §3.2 · p. 6

Abstract

from arXiv · show

Recent advances in language models (LMs) have led to significant improvements in quality on complex NLP tasks, but at the expense of increased inference costs. Cascading offers a simple strategy to achieve more favorable cost-quality tradeoffs: here, a small model is invoked for most "easy" instances, while a few "hard" instances are deferred to the large model. While the principles underpinning cascading are well-studied for classification tasks - with deferral based on predicted class uncertainty favored theoretically and practically - a similar understanding is lacking for generative LM tasks. In this work, we initiate a systematic study of deferral rules for LM cascades. We begin by examining the natural extension of predicted class uncertainty to generative LM tasks, namely, the predicted sequence uncertainty. We show that this measure suffers from the length bias problem, either over- or under-emphasizing outputs based on their lengths. This is because LMs produce a sequence of uncertainty values, one for each output token; and moreover, the number of output tokens is variable across examples. To mitigate this issue, we propose to exploit the richer token-level uncertainty information implicit in generative LMs. We argue that naive predicted sequence uncertainty corresponds to a simple aggregation of these uncertainties. By contrast, we show that incorporating token-level uncertainty through learned post-hoc deferral rules can significantly outperform such simple aggregation strategies, via experiments on a range of natural language benchmarks with FLAN-T5 models. We further show that incorporating embeddings from the smaller model and intermediate layers of the larger model can give an additional boost in the overall cost-quality tradeoff.

1 Introduction

The introduction frames generative-LM cascades as a deferral problem complicated by token sequences of variable length. It motivates token-level quantile features and learned post-hoc rules as alternatives to simple sequence-level confidence aggregation.

1 Introduction

  • Cascades route easy instances to a small model and defer hard instances to a larger model, making the deferral rule central to cost-quality tradeoffs.In classification, predicted class confidence is a standard deferral signal, but its generative-LM extension faces the variable-length uncertainty problem.Source: §1 · p. 2
  • Generative-LM deferral must summarize uncertainty across multiple token distributions, with output token counts varying between examples.This makes a single deferral score harder to construct than in classification, where confidence comes from one distribution over labels.Source: §1 · p. 2
  • Token-level uncertainty quantiles offer a finer-grained alternative to simple aggregation and can capture complementary uncertainty information.The introduction reports that quantiles outperform simple sum and average schemes, but no single quantile works across all datasets.Source: §1 · pp. 2–3
  • Post-hoc deferral trained on quantile features is the most consistently performant strategy against natural aggregation baselines on NLP benchmarks with FLAN-T5; large-model intermediate embeddings further improve performance.The paper also reports that quantiles outperform simple sum and average aggregation, but no fixed quantile works well across all settings.Source: §1 · p. 3
  • Sequence-level confidence measures can produce strongly sub-optimal cost-quality tradeoffs because they are affected by output-length bias.Taking the product of token probabilities tends to defer longer predictions, while length normalization can overcorrect and defer shorter ones.Source: §1 · pp. 2–3

2 Background and Problem Setup

This section sets up cascades as adaptive inference under a cost constraint and extends the setup from classifiers to generative language models. It defines deferral rules that select which model supplies each prediction.

2 Background and Problem Setup

  • An LM defines token probabilities from context, which compose into sequence probabilities by the chain rule.The section describes both classification, which selects among predefined choices, and generation, which samples a response sequence.Source: §2 · p. 3
  • A cascade first runs the lowest-cost model and uses a deferral rule to decide whether to stop or invoke a more costly model.Average inference cost can fall when only a small fraction of inputs are deferred.Source: §2 · pp. 3–4
  • The setup seeks a rule that maximizes expected prediction quality subject to a budget on average inference cost.The average cost depends on the deferral rate and the costs of the small and large models.Source: §2 · p. 4
  • For generative LMs, the cascade selects between two models’ response distributions using a deferral rule based on the input prompt.The quality measure evaluates a model’s response distribution for a prompt and target sequence.Source: §2 · p. 4
  • Prior learned routers may use prompts or generated text, while some do not use model confidence or assume models of increasing capacity.FrugalGPT learns a deferral score using the prompt and generated output text, rather than token-level uncertainty.Source: §2 · p. 4

3 Confidence Measures for Language Model Cascades

This section develops deferral rules for language-model cascades, moving from sequence-level probability and length-normalized averages to token-level quantiles. Across evaluated tasks, Chow-Quantile consistently outperforms Chow-Sum and Chow-Average, though no single quantile works well across all tasks.

3 Confidence Measures for Language Model Cascades

  • The section examines how different deferral functions behave on NLP tasks, extending confidence-based deferral from classification to generative language models.The starting point is the predicted class probability used in Chow’s rule for classification.Source: §3 · p. 5, §3.1 · p. 5

3.1 Chow-Sum and Chow-Average

  • Chow-Sum adds per-token log-probabilities, while Chow-Average uses their length-normalized average, also known as log-perplexity.The generated sequence is the prediction whose confidence is aggregated for deferral.Source: §3.1 · p. 5

3.2 Limitations of Chow-Sum and Chow-Average

  • Chow-Sum can favor longer outputs, while Chow-Average can over-correct and favor shorter ones; both may mask highly uncertain individual tokens.A few uncertain tokens can leave a sequence with a relatively high Chow-Sum score when many other tokens are predictable.Source: §3.2 · p. 6, Figure 2 · p. 6

3.3 Beyond Chow-Sum and Chow-Average: Chow-Quantile

  • Chow-Quantile converts the per-token log-probabilities into a deferral score using a quantile, retaining richer token-level uncertainty information than a sum or average.At α = 0, the score corresponds to the minimum per-token log-probability.Source: §3.3 · pp. 6–7

3.4 Experimental Setup

  • The evaluation uses few-shot FLAN-T5 Base and Large models with greedy decoding, and reports deferral-rate versus cascade-quality curves and AUC-DF.The datasets span classification, question answering, reading comprehension, translation, and commonsense reasoning; higher AUC-DF indicates better deferral curves, though values vary across datasets.Source: §3.4 · p. 7

3.5 Evaluating Confidence Measures for Cascades

  • Chow-Quantile consistently outperforms Chow-Sum and Chow-Average on three evaluated datasets, but no single quantile performs well across all tasks.The deferral curves cover MNLI, TriviaQA, and WMT DE→FR; the broader AUC-DF results show that the best quantile varies by task.Source: §3.5 · p. 7, Figure 3 · p. 8
  • Chow-Sum tends to defer longer outputs and Chow-Average shorter ones, although longer WMT outputs can contain repetitions and be good candidates for deferral.Thus length carries some signal, but the two aggregation rules can bias deferral toward one output length.Source: §3.5 · p. 8
  • On WMT FR→EN, Chow-Quantile-0 identifies repetitions and “??” unknown tokens, while Chow-Quantile-0.8 better captures uncertainty in shorter predictions; the quantiles provide complementary uncertainty measures.For FLAN-T5 Base on WMT EN→FR, average token probability increases with token index, and the model tends to become more confident.Source: §3.5 · p. 9

4 Post-hoc Deferral Rules

The section develops learned post-hoc deferral rules from token-probability quantiles, optionally enriched with model embeddings. Across datasets, post-hoc routing is generally on par with or better than non-learned baselines, and larger-model intermediate embeddings can improve performance.

4.1 Post-hoc deferral rule training

  • Post-Hoc-Quantile trains a scorer on small-model token-probability quantiles, Chow-Sum and Chow-Average scores, and targets reflecting whether deferral would benefit the example.Accuracy tasks use binary logistic loss with a target indicating when the large model is correct and the small model is incorrect; translation uses squared loss on the BLEURT-score difference.Source: §4.1 · p. 10

4.2 Leveraging intermediate embeddings

  • The embedding variants use averaged final decoder embeddings from the small model, then optionally add the large model’s first decoder-layer first-token embedding.The added large-model information is available only where intermediate embeddings can be extracted without querying the full large model; doing so can still involve non-trivial cost.Source: §4.2 · p. 10

4.3 How well does post-hoc deferral work?

  • Across datasets, post-hoc deferral is generally on par with or better than non-learned rules, while larger-model intermediate embeddings yield further gains.The comparison summarizes AUC-DF across datasets and full deferral curves on MNLI, TriviaQA, and WMT DE →FR.Source: §4.3 · p. 10, Table 1 · p. 11, Figure 6 · p. 11
  • Post-Hoc-Embed-1 slightly improves on Post-Hoc-Quantile and Chow-Quantile methods, while Post-Hoc-Embed-1+2 performs better than Post-Hoc-Embed-1.Larger-model intermediate embeddings can yield significant gains across deferral rates, particularly for classification; the first token may be less informative for longer generation and translation outputs.Source: §4.3 · pp. 10–11

5 Discussion and Future Work

The discussion highlights token-level uncertainty and intermediate embeddings as useful ingredients for LM cascades. It identifies model architectures, scorer design, and fine-tuning as directions for further evaluation.

5 Discussion and Future Work

  • The paper concludes that token-level uncertainty can improve on simple sequence-level uncertainty and that larger-model intermediate embeddings can further boost performance.These are the discussion’s stated findings about deferral rules for cascades.Source: §5 · p. 11
  • Future work could test whether a 1-layer Transformer better exploits token-probability sequence structure than the 5-layer MLP used here.The proposed Transformer is a direction for evaluation, not a tested alternative in this work.Source: §5 · pp. 11–12
  • The experiments focus on FLAN-T5 instruction-tuned encoder-decoder models; decoder-only architectures and the effects of fine-tuning, including RLHF, remain future evaluation questions.The paper suggests its methods may generalize to decoder-only models but does not evaluate them.Source: §5 · p. 12

6 Uncertainty Quantification for LMs

The section reviews uncertainty methods for language models, including sequence sampling, ensembles, calibration, direct confidence probing, and routing approaches. It notes trade-offs in computational cost and in how uncertainty is represented.

6 Uncertainty Quantification for LMs

  • Consensus-based methods sample multiple outputs and assess agreement, helping distinguish locally diffuse predictions from distributions peaked around one candidate [Wang et al., 2023; Xiao et al., 2021; Chen et al., 2023c].Yue et al. (2023) used answer consistency to construct cascades.Source: §6 · p. 12
  • Ensembles and dropout estimate uncertainty through model disagreement, but require additional computation compared with a single model’s softmax probabilities [Van Landeghem et al., 2022; Wang et al., 2019b; Gleave & Irving, 2022].The passage describes that extra computation as less appealing when the goal of a cascade is efficiency.Source: §6 · p. 12
  • Calibration work addresses answer and length biases [Zhao et al., 2021; Holtzman et al., 2021; Murray & Chiang, 2018; Jiang et al., 2021], while semantic entropy groups sampled, semantically equivalent outputs [Kuhn et al., 2023a].The passage proposes evaluating semantic entropy for cascading.Source: §6 · p. 12
  • Generative uncertainty methods directly probe an LM for confidence or train a confidence head, while reported linguistic confidence estimates can be worse than softmax-based estimates [Shrivastava et al., 2023].The cited work also reports that linguistic and softmax-based uncertainty sources can be complementary [Shrivastava et al., 2023].Source: §6 · pp. 12–13
  • Other approaches include sequence-level calibration [Zhao et al., 2023], ambiguity detection with clarification, external quality-estimation routing [Hendy et al., 2023], and meta-model selection among models [Šakota et al., 2023].Fadeeva et al. (2023) provide an experimental analysis of uncertainty methods.Source: §6 · p. 13

A Experimental details

The appendix reports training settings for the post-hoc rules and describes the dataset information tabulated for the experiments.

A Experimental details

  • Quantile-based deferral rules use a 5-layer, 32-hidden-unit-per-dimension MLP, while embedding-based rules use a 2-layer, 8-hidden-unit MLP to limit overfitting.Training uses 200 epochs with early stopping, ADAM at 10^-5, batch size 16, and five random runs; datasets are capped at 10,000 examples per split.Source: A Experimental details · p. 26
  • Table 2 records each dataset’s shots, training and testing splits, example counts per split, and average target-token count in the training split.The training set is divided into 80% training and 20% validation where datasets lack all three splits.Source: A Experimental details · p. 26, Table 2 · p. 27

B.1 Ablations with probability vectors as features

Quantile-based token uncertainty features outperform padded or sorted probability-vector alternatives, while output length also correlates with BLEURT scores. The analyses show that useful quantiles vary by dataset.

B.1 Ablations with probability vectors as features

  • Quantile features outperform padded probability vectors, and sorting probabilities also performs worse than quantiles.The probability-vector comparisons account for variable-length sequences by padding with zeros.Source: §B.1 · p. 28, Figure 7 · p. 29
  • Token uncertainties are not aligned across variable-length sequences, making padded vectors harder to learn from; quantiles better align uncertainty scores.Source: §B.1 · p. 28

B.2 Output length and BLEURT correlation

  • BLEURT scores have a non-zero correlation with prediction lengths across FLAN-T5 model sizes, providing some signal for deferral.Table 3 reports the correlation analysis.Source: §B.2 · p. 29, Table 3 · p. 30
  • Different quantile scores have different length correlations and deferral-curve AUC-DF; quantiles 0.0 and 0.8 work best on the illustrated dataset.Source: Figure 8 · p. 30

B.3 Additional results for various models and datasets

The appendix extends cascade-curve analyses across model pairs and datasets, and examines which output lengths different rules defer. It also presents curves for a cascade with relative latency.

B.3 Additional results for various models and datasets

B.4 Different version of length plot

  • For WMT FR→EN, Oracle defers a larger fraction of shorter predictions, while Chow-Sum favors longer predictions.Figure 12 presents the comparison at different deferral rates; Figure 13 reports the same length preference.Source: Figure 12 · p. 34, Figure 13 · p. 34

B.5 Cascade curves with relative latency

  • Figure 14 presents relative-latency deferral curves for a Base→Large cascade on MNLI, TriviaQA, and WMT DE→FR.Source: Figure 14 · p. 35

B.6 Results with encoder and decoder embeddings

The embedding experiments compare variants that add encoder and decoder representations to deferral methods. Including encoder embeddings yields no additional gains in these methods.

B.6 Results with encoder and decoder embeddings

  • Post-Hoc-Embed-Dec-1 adds model 1 decoder embeddings, while Post-Hoc-Embed-Enc+Dec-1 also includes model 1 encoder embeddings.Source: §B.6 · p. 36
  • Post-Hoc-Embed-Dec-1+2 adds model 1 decoder and intermediate decoder embeddings; its Enc+Dec variant additionally includes model 2 encoder embeddings.Source: §B.6 · p. 36
  • Including encoder embeddings provides no additional gains in these methods.Figure 15 shows cascade curves for the Base→Large cascade with additional encoder embeddings.Source: §B.6 · p. 36, Figure 15 · p. 36

B.7 Results with beam search decoding

With beam-search decoding, the appendix presents cascade curves and examines output-length preferences. It also reports that post-hoc approaches predict golden deferral labels better than Chow-Sum and Chow-Average.

B.7 Results with beam search decoding

B.8 Performance of methods in predicting golden deferral label

  • For WMT FR→EN under beam search, Chow-Sum tends to defer longer prompts, Chow-Average over-corrects toward shorter prompts, and Oracle's preference for longer outputs is more moderate.Figure 18 also shows the corresponding deferral curves.Source: Figure 18 · p. 38
  • Post-hoc deferral approaches yield significant AUC-ROC gains over Chow-Sum and Chow-Average, which are poorly predictive of golden deferral labels.The golden label marks cases where FLAN-T5 Large is better than FLAN-T5 Base.Source: Table 4 · p. 38

C How is Φ(·) computed?

Φ is computed as a feature vector for Post-Hoc methods, combining uncertainty scores and, in some variants, decoder embeddings. The variants differ in which quantities they concatenate and in the resulting feature dimensions.

C How is Φ(·) computed?

  • The example defines token log probabilities for a five-token prediction and uses 20 quantiles for the quantile scores.The specified quantile levels range from 0 to 1.Source: C How is Φ(·) computed? · p. 39
  • Post-Hoc-Quantile concatenates Chow-Sum and Chow-Average scores into a 22-dimensional feature.The quantile method is introduced as a Post-Hoc method variant.Source: C How is Φ(·) computed? · p. 39
  • Post-Hoc-Embed-1 concatenates Chow-Sum and Chow-Average scores with model 1’s decoder embeddings, yielding a 22 + n1-dimensional feature.Model 1’s decoder embedding is averaged over the output tokens.Source: C How is Φ(·) computed? · p. 39
  • Post-Hoc-Embed-1+2 concatenates Chow-Sum and Chow-Average scores with embeddings from both models, yielding a 22 + n1 + n2-dimensional feature.The embeddings are the smaller model’s decoder embedding and the larger model’s intermediate decoder embedding.Source: C How is Φ(·) computed? · p. 39
Loading 2404.10136v1…