Language Model Cascades: Token-level uncertainty and beyond

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Gupta, Neha, Narasimhan, Harikrishna, Jitkrittum, Wittawat, Rawat, Ankit Singh, Menon, Aditya Krishna, Kumar, Sanjiv
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917641122217984
author Gupta, Neha
Narasimhan, Harikrishna
Jitkrittum, Wittawat
Rawat, Ankit Singh
Menon, Aditya Krishna
Kumar, Sanjiv
author_facet Gupta, Neha
Narasimhan, Harikrishna
Jitkrittum, Wittawat
Rawat, Ankit Singh
Menon, Aditya Krishna
Kumar, Sanjiv
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2404_10136
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Language Model Cascades: Token-level uncertainty and beyond
Gupta, Neha
Narasimhan, Harikrishna
Jitkrittum, Wittawat
Rawat, Ankit Singh
Menon, Aditya Krishna
Kumar, Sanjiv
Computation and Language
Artificial Intelligence
Machine Learning
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.
title Language Model Cascades: Token-level uncertainty and beyond
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2404.10136