Conformal Nucleus Sampling (2023)

Abstract: Language models generate text based on successively sampling the next word. A decoding procedure based on nucleus (top- p ) sampling chooses from the smallest possible set of words whose cumulative probability exceeds the probability p . In this work, we assess whether a top- p set is indeed aligned with its probabilistic meaning in various linguistic contexts.We employ conformal prediction, a calibration procedure that focuses on the construction of minimal prediction sets according to a desired confidence level, to calibrate the parameter p as a function of the entropy of the next word distribution. We find that OPT models are overconfident, and that calibration shows a moderate inverse scaling with model size.

## introduction modern language generation methods are all based on computing the conditional next-word distribution. however, there is still considerable debate about the best way to extract the next word from that distribution. most current text generation methods employ one of a handful of standard decoding strategies, which are characterized as either deterministic or stochastic in nature. a greedy search strategy selects the word with the highest probability at each timestep. the greedy method and its beam search variations work remarkably well for machine translation but outside of this context, tend to return dull text or degenerate text @xcite . @xcite argued that high-quality human language does not follow a pattern of highestprobability next words, as humans expect the generated text to not be repetitive or boring. the same problem occurs with beam search. direct sampling from the next-word distribution computed by the model often generates incoherent gibberish text. temperature sampling @xcite is a word sampling approach based on rescaling logit scores before applying the softmax function to compute the word distribution. other methods limit the sampling space to a small prediction set to avoid the "unreliable tail" @xcite . in top-k sampling @xcite , we sample only from the top-k most likely words. instead of sampling only from the most likely k words, top-p (nucleus) sampling chooses from the smallest possible set of words whose cumulative probability exceeds the probability p @xcite . top-p sampling enables a dynamically sized window of words, unlike top-k which fixes the size of k for every step. finally, locally typical sampling @xcite and truncation sampling @xcite are recent variants of top-p that aim to make it more suitable for language generation. the top-p prediction set has a concrete probabilistic interpretation. here we examine whether the probability that the "correct" word belongs to the set of words produced by the top-p algorithm is indeed p. more generally we expect that the nextword prediction would be calibrated, meaning that the output of the next-word softmax layer would accurately reflect the true word distribution. parametric calibration methods, such as temperature scaling @xcite , which adjust the confidence of the most probable word, are not suitable for adjusting the size of the prediction set. conformal prediction (cp) @xcite @xcite is a non-parametric calibration method that, given a value p, aims to build a prediction set with a guarantee that the probability that the correct word is within this set is indeed p. note that this notion of calibration, which is distinct from the way calibration is usually formulated in language mod-eling settings, exactly coincides with the goal of the top-p prediction model. the model-agnostic and distribution-free nature of cp makes it particularly suitable for large neural network models. we thus applied cp analysis to asses whether the top-p procedure is calibrated and, if needed, tune it to have the desired probabilistic interpretation. we find that opt models of different sizes @xcite are not calibrated according to the conformal prediction theory, and that calibration shows moderate inverse scaling. additionally, we show that the degree of calibration varies significantly with the entropy of the model's distribution over the vocabulary. we thus propose a new conformal top-p decoding algorithm, which ensures that the top-p sampling has a meaningful probabilistic interpretation. ## cp for language generation in this section, we briefly review the split conformal prediction algorithm @xcite and discuss its relevance to language generation models. consider a network that classifies an input x into k pre-defined classes. the network (softmax layer) output has the mathematical form of a distribution. however, this does not necessarily mean that it accurately reflects the true class distribution. let (x, y) be a test instance and its corresponding class. we want to find a small subset of classes (a prediction set) c(x) ⊂ {1, ..., k} such that let @xmath0, 1] be the smallest threshold value that p(@xmath1. if @xmath2. if @xmath3. if the model is not calibrated, we can calibrate it using a labeled validation set (x 1 , y 1 ), ..., (x n , y n ). denote p t (i) = p(y @xmath4; θ). define the conformal scores to be: this cp score is known as the adaptive prediction sets (aps) score, and was first introduced in @xcite . note that y @xmath5. we next look for a minimal threshold q such that the correct label y t is included in the prediction set c q(x t ) for at least (1-α)n points of the validation set. in other words, q calibrates the top-(1-α) prediction-set on the validation set. we can easily find q by first sorting the n scores s 1 , ..., s n and then q is the (1-α)-quantile of the validation-set scores. once the network is calibrated, if we want to form a prediction set for a new test sample x, that contains the true class with probability (1-α), we use c q(x). the cp calibration procedure for calibrating the top-p word decoding is summarized in algorithm 1. the conformal prediction theory provides the following guarantee on the threshold q @xcite . theorem: assume a test point (x, y) and the n validation points are independent and identically distributed (or at least exchangeable). let q be the ⌈(n + 1)(1α)/n⌉-quantile of the validation set scores. then note that this is a marginal probability over all the test points and is not conditioned on a given input. exchangeability means that the sequence distribution is not altered by permuting the order of the random variables. in this study, we aim to apply the conformal prediction framework to language generation models to analyze the prediction sets used for sampling the next word. the joint distribution of words in a text is neither iid nor exchangeable, since the words are correlated and the order of the words in a sentence is significant. a recent study @xcite showed that applying the usual cp algorithm to a stationary β-mixing process (rather than an exchangeable one) results in a guaranteed coverage level of 1-α-η, where η depends on the mixing properties of the process and is theoretically hard to know, or bound. roughly speaking, β-mixing processes are stochastic processes in which far-away ## experiments in this section, we apply the conformal prediction calibration method to analyze the calibration status of the top-p nucleus sampling. setup. we experimented with variants-from 125m parameters up to 30b parameters-of opt @xcite , a left-to-right language model. we ran the models on 10,000 english wikipedia ## conclusions to conclude, in this study we apply the notion of calibration by conformal prediction to calibrate the top-p nucleus sampling as a function of the next word distribution entropy and thus made the top-p decoding policy consistent. the same analysis and calibration can also be applied to other commonly used decoding methods, such as variants of topp @xcite and truncation sampling @xcite .