Closing the Curious Case of Neural Text Degeneration
Authors: Matthew Finlayson
John Hewitt
Alexander Koller
Swabha Swayamdipta
Ashish Sabharwal
Venue: ICLR 2024 poster (2024)
Keywords: softmax
bottleneck
truncation
sampling
nucleus
top-k
theory
linear programming
linear algebra
decoding
generation
autoregressive
language model
NLP
open-ended generation
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We show that truncation sampling coarsely eliminates errors due to the softmax bottleneck, and develop a precise sampling method which outperforms truncation on open-ended text generation.

Abstract: Despite their ubiquity in language generation, it remains unknown why truncation sampling heuristics like nucleus sampling are so effective. We provide a theoretical explanation for the effectiveness of the truncation sampling by proving that truncation methods that discard tokens below some probability threshold (the most common type of truncation) can guarantee that all sampled tokens have nonzero true probability. However, thresholds are a coarse heuristic, and necessarily discard some tokens with nonzero true probability as well. In pursuit of a more precise sampling strategy, we show that we can leverage a known source of model errors, the softmax bottleneck, to prove that certain tokens have nonzero true probability, without relying on a threshold. Based on our findings, we develop an experimental truncation strategy and the present pilot studies demonstrating the promise of this type of algorithm. Our evaluations show that our method outperforms its threshold-based counterparts under automatic and human evaluation metrics for low-entropy (i.e., close to greedy) open-ended text generation. Our theoretical findings and pilot experiments provide both insight into why truncation sampling works, and make progress toward more expressive sampling algorithms that better surface the generative capabilities of large language models.

CLOSING THE CURIOUS CASE OF NEURAL TEXT DEGENERATION
Matthew Finlayson' University of Southern Californa
John Hewitt
Alexander Koller Saarland University
Stanford University
Swabha Swayamdipta University of Southern Californa
Ashish Sabharwal The Allen Institute for AI
ABSTRACT
Despite their ubiquity in language generation, it remains unknown why truncation sampling heuristics like nucleus sampling ical explanation for the effectiveness of the truncation sampling by proving that truncation methods that discard tokens below some probability threshold (the most common type of truncation) can guarantee that all sampled tokens have nonzero true probability. However, thresholds are a coarse heuristic, and necessarily discard some tokens with nonzero true probability as well. In pursuit of more precise sampling strategy, we show that we can leverage known source of model errors, the softmax bottleneck, to prove that certain tokens have nonzero true probability, without relying on a threshold. Based on our findings, we develop an experimental truncation strategy and the present pilot studies demonstrating the promise of this type of algorithm. Our evaluations show that our method outperforms its threshold- based counterparts under automatic and human evaluation metrics for low-entropy (i.e., close to greedy) open-ended text generation. Our theoretical findings and pilot experiments provide both insight into why truncation sampling works, and make progress toward more expressive sampling algorithms that better surface the generative capabilities of large language models.
effective. We provide theoret-
INTRODUCTION
Crucial to the remarkable generative capabilities of today's 's large language models (LLMs) (OpenAI, 2023; Touvron et al., 2023; Chowdhery et al., 2022) are the sampling algorithms responsible for selecting the next token at each timestep. The most common of these algorithms use a simple trunca- tion strategy: sample only the tokens that have probability greater than some threshold (Holtzman et al., 2020; Fan et al., 2018). In the quest for high-entropy generation wherein one wants to be able to generate multiple good completions, it has been empirically established that the search for the highest- lighest-iiklihhodd strings through e.g., beam search or greedy decoding led to low-quality generations (Hashimoto et al., 2019). Threshold-based truncation sampling presents compelling alternative: by avoiding the tokens at the tail end of the distribution which correspond to degenerate text it produces significantly more coherent generations (Ippolito et al., 2019; Holtzman et al., 2020; DeLucia et al., 2021). However, beyond the intuition that language models tend assign too much probability to tokens that should have 0 or near-O probability (akin to smoothing (Hewitt et al., 2022)), prior work has been limited in establishing why truncation sampling is so essential in autoregressive generation.
In this paper, we provide a precise mathematical explanation to elucidate the extraordinary success of threshold- based truncation sampling ($3). First, we prove via an argument about log-pro probability emrors that threshold sampling is guaranteed to only sample tokens in the support of the true distribution, so long as the chosen threshold is larger than some bound (Corollary 1). Next, we present method to more directly account for likely source of tail errors: the softmax bottleneck (Yang et al., 2018). which states that the low-rank softmax matrix used the output layer of language models causes probability errors in the model's output distribution ($4). Specifically, we show how to leverage
"Corresponding author mf ni mnlaysuscc.edu

---

Log probability
Swiff
nucleus
77
Accepted by ours
Not accepted by ours
had
Mart inez
Richardson
SW
Kits
Token (sorted by probability)
Figure 1: The next-token distribution from GPT-2 XL for the prefix "Tayler", with the tokens ordered by probability. Dashed vertical lines denote thresholds used to reject low-probability tokens, under various truncation strategies. Our masis-aware-threshold (BAT) sampling accepts tokens shown in blue and rejects those in orange. As evident, BAT rejects some implausible tokens assigned high probability under the model while accepting many plausible yet now-probaiility tokens- this not possible under truncation sampling. BAT uses the softmax matrix to find tokens that might have non-zero true probability, without relying on a threshold. See more examples in Fig. 4.
the restricted structure imposed by the softmax bottleneck to more precisely determine (relative to threshold-based truncation) which tokens are in the support of the true distribution (Theorem 2). At high level, the idea is to declare a token to be in the support if its probability is nonzero not only in the predicted distribution but also in all distributions that are "similar" to it (in a precise technical sense) from the perspective of the softmax matrix. This presents more nuanced strategy compared to threshold-bee based truncation sampling: our algorithm does not rely on a threshold, thereby allowing higher probability tokens to discarded while keeping some lower-probee power-probabilityy tokens.
We conduct pilot investigation ($5) to empirically evaluate this basis-aware truncation sampling approach. Our results shows improvements on an open-ended generation task via both automatic and human evaluation metrics under low-entropy generation (i.e., close to greedy). Figure illustrates our algorithm's 's more nuanced token selection strategy qualitatively (also see Figure 4). Unlike threshold- based truncation methods (each shown with a dotted vertical line), our method can selectively discard low-quality tokens while still keeping high-quality but lower-prodd wwer-prrbbbbllty tokens. This is accomplished by taking into account linear dependencies between token embeddings.
Overall our work provides theoretical insights which motivate a practical method and show how truncation sampling avoids errors in language model by mitigating the softmax bottleneck.
2 BACKGROUND
Autoregressive language models Autoregressive language models are trained next-word- predictors: given a prefix, the model assigns a probability to each token a vocabulary of size v as prediction of which token comes next. Given input prefix, model produces vector h E Rd. which we refer to as the hidden state, and hyperparameter d as the hidden size. The model then uses linear map with matrix W Ro to obtain logits Wh E R", to which it applies the softmax function to obtain a probability distribution over tokens in the vocabulary:
exp(Wh)
p softmax Wh)
E (h)),
The matrix W is commonly referred to the seftmax matrix because it is applied directly before the softmax, or the embedding matrix. Generally models are trained to output the á that minimizes the cross entropy with the conditional true distribution p" crossentropy/p2 .)) = 1 Pp log p:.
Generation via truncation sampling Language models can autoregressively generate text by sampling a token from P at each time step. Unfortunately, sampling directly from p. i.c., ancestral sampling, often leads to quality issues with unnatural, low-probability tokens. Truncation sampling aims to solve this issue post-hoc by choosing subset of the vocabulary to sample from, setting
com.coom/t.//ooooo t fin/bast s- re-are--hhesshh "I the case of natural language, it is not entirely clear what the "true" distribution pº means exactly. Nonetheless we use the distribution from which internet text is implicitly sampled as useful surrogate.
"Code for experiments: https: : / /gi thub. com
2

---

all other tokens have zero probability. Meister et al. (2023a) frame this strategy as reprioritizing precision over recall (i.e., removing some valid text from the distribution to avoid sampling unlikely text.) We focus a class of truncation methods that select tokens by choosing threshold at each timestep and truncating tokens with probability less than that threshold. This simple heuristic has been found to be effective and forms the basis of popular methods like nucleus (top-p) (Holtzman et al., 2020) and top-k (Fan et al., 2018) sampling.
Prior work has introduced several heuristics for choosing truncation thresholds. For instance, the threshold can be fixed constant as in sampling, or chosen dynamically across different distributions, as in 17. nucleus, top-k, and Mirostat sampling (Basu et al., 2021).² 17 sampling introduces the idea that the threshold should depend on the entropy of the distribution H (p) and sets the threshold" to min(7, (f)). In the latter three, the threshold is chosen implicitly rather than explicitly, for
instance, in nucleus sampling with parameter n, the threshold is min {p. E,,>,, P3 S =
In the extreme case, truncating all but the most likely token results in greedy decoding. Though this strategy makes it unlikely to sample a token outside the true support, often results in degenerative patterns like repetition (Holtzman al., 2020). Furthermore, even for modern language models that suffer less from greedy decoding traps, non-deterministic sample-baser decoding is useful for generating multiple completions and for more "creative" generations. Thus, the best choice of threshold must strike a balance between diversity (i.e., including as many tokens as possible in the set of candidates) and coherence (i.e., avoiding sampling tokens outside the true support).
The softmax bottleneck The sources of the probability overestimation errors are likely many, but one source of error is particularly compelling and well defined mathematically: the softmax bottleneck (Yang et al., 2018). The softmax bottleneck refers to the limited expressivity of models with a small hidden size and large vocabulary. Recalling the notation from Yang et al. (2018), let A R be the matrix where each entry Ar, log p"(i (i j) is the true prgobabbiillty of token i given a prefix j from some set of n > prefixes. Also, let W E Rexd be the softmax matrix for model, and H E IRdx be the matrix of model hidden states given each prefix. Finally, let J E Rexx be the all-ones matrix. The rank of the model's probabibitty matrix
A' = log softmax W H) = W H - Jdiag Σ eex((WH),)
is at most d+1 1 because W H has inner dimension d and therefore rank at most d, and the subtrahend has identical rows and therefore has rank at most The rank of A is at most v. If rank of A much larger than d, then A' can be at best low-rank approximation of A. From the Eokart-Young-Mit (EYM) theorem for low-rank approximations.
min MarankiA ( sdd+
A
A'
i=d+2
where I| denotes the Frobenius norm, and a the vector of singular values of A, ordered by decreasing size. Thus, there will always be some error in the model's log-probability estimations if there are more than d + 1 linearly independent columns in A. Yang et al. (2018) hypothesize that this is indeed the case.
Despite these theoretical shortcomings, language models still seem to perform quite well. We hypothesize that the reason for this is that default truncation sampling is sufficient to approximately mitigate errors from the softmax bottleneck. For a deeper discussion, see Appendix A.
3 A THEORETICAL EXPLANATION OF TRUNCATION SAMPLING
Given some textual context as input, let p" denote the true next-token distribution of the language and P the model's predicted next-token distribution. Intuitively, if the model's probability overestimation
Locally typical sampling (Meister et al., 2023b) truncates tokens whose log- log-probability diverges from the LM's 's conditional entropy. This may truncate poo-probayility tokens which likely have non-zero true probability. Hewitt al. (2022) instead s min (p)( for parameter We diverge for simplicity.
3

---

could be additively upper bounded, i.e., if we could show that P. S P. + for every token i, then this would yield natural way to avoid sampling tokens not in the support of p*: only sample tokens i with P; > 7 (which, along with the bound, would imply P. > 0). This is exactly what truncation sampling does. However, a difficulty in motivating truncation sampling via this argument is that is unclear how to derive such additive upper bound on probability overestimation.
Our key observation is that A' being a low-rank approximation A can be used to conclude that the model's log-probability underestimation non-zero but additively upper bounded. Indeed, assuming A' is reasonably good low-rank approximation of A, Equation 1 implies such an upper bound in the pog-pobabiiltyy space, which yields a multiplicative upper bound the probability space. We then combine this underestimation upper bound with basic properties of a probability distribution in order to derive the desired additive upper bound on the model's probability overestimation. Lastly, we show formally how this overestimation upper bound directly motivates truncation sampling.
3.1 BOUNDING LOG-PROBABILITY UNDERESTIMATION
We begin by proving bounds models' log- pe-probabilitty errors. Specifically, we find bounds on the maximum log-proo prepbabbbbiyyy underestimation error of the model, max(A A'). ) We focus exclusively on underestimation enrors because log- og-probability overestimation errors cannot be bounded above." 5
Maximum log-probability error upper bound We begin by boundinounding all model's log- probability underestimations. In particular, the underestimation errors A A' are upper-bouded by max(A A A') I max A min A' v min A', where the last inequality holds because max A is log- eee-prbbbbility and hence
by O. In other words, the negative minimum log- probbabbbi prediction min A' upper bounds all underestimation. As an example, a uniform predicted distribution underestimates the log- probbbbbbbii of a token by at most log(1/e).
Maximum log-probability error lower bound Next, we lower-bound maximum underestimation errors by showing that they strictly positive. We conjecture that this lower- bound error is loose, i.e. that the maximum error is bounded away from 0, depending on the singular values of A.
3.2 BOUNDING PROBABILITY OVERESTIMATION
Having established bounds on maximum log- probbabbbyilty underestimatiin, we now show that assuming such an upper bound implies an additive upper bound on maximum probability overestimation. As before, fix some imput textual context and let p' and p denote the true and model's predicted next-token distributions, respectively, for that context.
Theorem 1. If ffogg éú underestimates log pt by at most 8 for all tokens then p. overestimates PÊ by at most - exp( -8) for all tokens i.
See Appendix D for a proof. Note that the precondition log log p. < implies P4 2 p exp(-0) Intuitively, since is valid probability distribution summing to it cannot underestimate token probabilities beyond a factor of cxp(-5), then it also cannot overestimate other tokens' probabilities beyond certain additive factor. We compute this additive factor and find be 1 exp(-6).
3.3 EXPLAINING TRUNCATION SAMPLING
Recall that threshold-based truncation sampling works by only sampling tokens with probability greater than some threshold r. Sampling methods that choose a different 7 at every time step be viewed as additional heuristics for guessing when model outputs will have smaller errors. Theorem provides a direct explanation for why threshold-base truncation sampling might be successful:
Corollary 1 (Threshold based truncation works). Suppose log p underestimates log p" by at most 8. Then, for any threshold T 2 1- exp(-6), threshold based truncation sampling discards all tokens that not in the support of p"
'If the true distribution assigns zero probability to some tokens in some contexts (e.g.. p' ("ate" "I went to the") = 0), then the corresponding log-prohability is -oo Hence any finite proobbbbyyyyiyy es- timate will have infinite error.
4

---

softmax
00
R
1.0
0.5
0.0
ad
0.0
0.5
1.0
0
Figure 2: For toy model with hidden size 1, vocabulary size 3, and an embedding matrix W 33x1, W projects the space of possible hid- den states R into a I-dimensional subspace of the space of possible logits RR3. In turn, the softmax function projects this 1D logit subspace onto a ID) subspace of the space A3 of possible probability distributions over 3 tokens. Thus, our toy model can only output distributions within a ID sub- space of Ag, which is the image of softmax W.
Figure the model outputs (the blue dot) within the space of possible outputs (blue line), then each token might have zero true probability only if there is distribution p with p; = 0 that satisfies both the BA constraints (orange line) and the threshold constraints (orange area). For exam- ple, the orange line and area coincide at the green dot where P1 = O, therefore token I might have zero true probability. The other tokens must have nonzero probability since there other
such solutions.
Furthermore, based the above proof, present alternative formulation of truncation sampling. Corollary 2 (Threshold sampling reformulation). For model with maximum jog-probability un- derestimation error 5, if the model outputs p and there is no distribution p with ps such that Pj $ Py fxp( j E (1.2,.... v). then p > 0.
This follows directly from Equation (8) from the proof in the appendix, and is the contrapositive of the more straightforward statement that if p () then there exists distribution satisfying inequality conditions the corollary, namely p". One can check that only sampling tokens based on Corollary
yields the same candidate sets threshold sampling with 1 - (x) the parameter. This alternative formulation will become useful later on when we combine methods for proving certain tokens in the support.
DIRECTLY ADDRESSING ERRORS FROM THE SOFTMAX BOTTLENECK
The previous section demonstrates that we arrive truncation sampling by making assumption about the tog-probability errors, which allows us to prove that certain tokens have true probability greater than zero. However, truncating via threshold is inherently limited approach: if model assigns more probability to a "bad" (zero true probability) token than a "good" (nonzero true probability) token, then there is threshold that discards the bad token without discarding the good token. Naïvely, it would seem that this type of issue is unsolvable, however, it turns out that this error is caused by the softmax bottleneck, risking sampling the bad token. By exploiting W, the low-rank basis for the model's outputs, and we can deduce exactly which tokens may have errors due to the softmax bottleneck, regardless of their relative probability. In this section we show mathematically how we can extend threshold sampling take full advantage of our knowledge of the softmax bottleneck.
actually recover the good token without
4.1 BASIS- AWARE SAMPLING
At a high level, we will motivate this approach by showing that the function used to transform the hidden state to a probability distribution P restricts model's outputs to a subset of the possible probability distributions. When the true distribution p" lies outside of this set, then the model to output the P within the set that minimizes the model's ' training loss with respect to p=. We can exploit this property to identify the of distributions wherein the true distribution lies,
expect
5

---

namely the set of distributions that minimizes loss with. If no distributions within this set assign zero probability to a particular token, then that token must have nonzero probability.
To build intuition for how model's outputs are restricted. consider the toy model in Figure 2. We generalize this toy model to a model with hidden size d and vocabulary size v. Observe that the composed functions softmax W define linear map: first, the model's softmax matrix W E R defines a linear map Rd + R". Next, it is a lesser-known fact that the softmax function is a linear map from I + A.. where A, is the (v 1)-dimensional vector space of valid probability distributions over e variables (Aitchison, 1982) (see Appendix C for an explanation). Therefore, softmax W: Fd + Dg is a linear map from d-dimensional space (u -1) 1iddmennsilaal space, meaning the image of this function is an at-most d-dimensional subspace of A., In other words, the space of model outputs restricted to a subset of all possible probability distributions over the vocabulary.
What distribution should a model output, given that the true distribution p' may not lie in the subspace of possible outputs? Typically, language models are trained to minimize cross-entropy with the true distribution. Therefore, a well-trained model can be expected to output the distribution P within the image of softmax W that minimizes cross-entropy with p'. In other words, we assume that torseenttoo(ssootomnn
h), p²) is minimized. The key insight of our method is that if h does not minimize cross entropy with any distribution p such that pi = 0, then P # O, i.e., token i is in the true support.
the model will produce the hidden state h such that
Theorem 2 (Basis-aware sampling). If p the predicted distribution from a cross- caaasseettoo--inniii model with embedding matrix W, and if there is no valid probability distribution p such that pt = 0 and W T p Pp p. then the token's 's true probability Př is greater than O.
See proof in Appendix D. This gives new way to prove that tokens in the true support, similar to Corollary 2, but in a way that directly compensates for errors due to the softmax bottleneck.
4.2 COMBINING SAMPLING METHODS
Theorem 2 and Corollary equip us with methods for proving tokens are in the true support. By combining the constraints specified from each method we can create a hybrid proof strategy to take advantage of both methods' insights. In particular, if there does not exist a distribution p with pt = such that P5 S f. cxp(o) for all j (the truncation constraint) and W p W p (the basis- constraint), then pt > O.
This hybrid proof strategy naturally yields sampling method: sample only tokens that we can prove are the support. We call this method busis-aware threshold (BAT) sampling. Fortunately, both the threshold constraint and basis- (BA) constraints are linear, so can an off- fff-the-shelf linear programming optimizer to verify whether a token is in the support. Concretely, if the optimizer determines that there does not exist feasible solution p E IR" such that:
P 0,
Σe; Pj 1, Vj 0 S PS S i. exp(o), j=1
wpp WT p.
then pp >> Thus, our sampling strategy can be: sample a token i according to the model's output probabilities; if the optimizer finds a solution to (3), reject the token and re-sample; otherwise accept. See Algorithm I in the Appendix.
We expose 8 as parameter to tune the restrictiveness of the sampling method. For large 6, BAT becomes more like greedy sampling, and for small 8, more like ancestral sampling. The value of can be chosen sampling. Given a threshold 7 from the heuristic, set exp ő = 1/(1 - r). We call these variants of BAT sampling BA-c, BA- BA-nucleus sampling.
per-context basis using any threshold sampling heuristic, be it n, or nucleus
A toy example Suppose our model has hidden size 1, vocabulary size 3, and embedding ma- trix W [0.55 0.71 0.29). We employ the truncation sampling assumption that our model's output distributions are somewhat close to the true distribution by saying pf S P. exp ó and choosing
may correctly observe that the "bottleneck" is consequence of the linear map W, not the softmax function. We keep our notation for the sake of consistency with Yang et al. (2018).

---

5 = log 1.9 that pt I 1.9p; for all tokens i. Additionally, assume the model's outputs minimize cross-entropy with the true distribution, i.e., W p" W p for all p. Now suppose our model out- puts h = [2.55]. The output distribution is therefore p softmax (Wh) [0.33 0.50 0.17]
Our strategy only samples tokens for which we can prove that the true probability is positive. Referring to Figure 3, we see that there are no probability distributions p that satisfy our assumptions with P2 = 0 or P3 0. However, = [0 0.70 0.30] does satisfy our assumptions. Therefore, if we sample token we should reject it, as we only have evidence that pá A and p3 o 0. Notice that this strategy is non-monotonic: P1 > P3. but we only reject token 1, not token 3.
Basis-aware threshold sampling in practice The proposed implementation basis-aware sam- pling requires solving rather large linear programs, which tends to be too computationally expensive to be practical, even when using proprietary solvers. The long run times can mainly be attributed to the size of W. To make BAT feasible in practice, we approximate the full solution by replacing W with an much smaller matrix such that no additional tokens are accepted, and the set of rejected tokens minimally increases. More details are deferred to Appendix E. This shortens the time from over a minute on a proprietary solver to about a second. We further reduce the generation run time by observing that whenever token has probability greater than 1 - exp(-o) we can safely accept it without running the program, since the program will be infeasible. Since high byhh-probbbility tokens are most likely to be sampled, the program only needs to run once every few samples. The amortized cost of BAT' sampling comes to only about 0.1 seconds per token as the program typically runs every 10 samples.
PILOT EXPERIMENTS WITH BASIS- AWARE TRUNCATION
We conduct several evaluations with GPT-2 pilot BAT sampling viable alternative threshold sampling. While more powerful language models exist, these models suffice since we are primarily interested in testing the effect of the BAT sampling on performance under controlled settings.
As baseline methods for comparison, we select 7), e, and nucleus sampling (see $2). We also use and e as methods for selecting the d parameter at each time step for BAT sampling. In preliminary experiments, we also tried BA-nucleus, but found it to be significantly worse. One possible intuition for why is that the methods for choosing the threshold 6 and 1) are similar to the formulation of threshold sampling used to develop BAT. Nucleus sampling on the other hand determines the threshold using a function that is somewhat inconsistent with our framework.
We evaluate models on open-ended generation using both human annotators and automatic metrics. For each model and sampling setting, we generate completions for 5000 35-token prefixes taken from the Open Web Text (OWT) (Gokaslan et al., 2019). We use OWT because it comes from a similar distribution to GPT-2's 's training data. We report MAUVE (Pillutla et al.. 2021) similarity between human text and generated text for parameter selection and automatic evaluation.
Parameter selection and evaluation We perform a parameter sweep for nucleus, 7), and sampling and select the parameter that gives the highest MAUVE score the OWT validation set (see Table in the appendix). We control for the parameter choice in comparisons between BAT methods and their vanilla counterparts, by matching the parameters by selecting the BAT parameter that rejects the same proportion of tokens from corpus of human text as the vanilla method; see Appendix F for more details. Using these parameters, we generate completions on the OWT test set for automatic evaluation with MAUVE and human evaluation.
5.1 QUALITATIVE, AUTOMATIC, AND HUMAN EVALUATION
Qualitative analysis Figure 4 shows the effects of truncation methods on the next-token distribu- tions from 6 prefixes, drawn from Hewitt al. (2022). Unlike threshold sampling methods, BAT can reject low-quality high- probbbiibay tokens while accepting high-quality low-probability tokens.
BA-n outperforms all other methods for GPT-2-Large We compare the MAUVE scores on OWT for each method and model size in Figure 5. The results show that no single method consistently performs best, with BAT methods sometimes out-performing and sometimes under-performing their
7

---

My 100M
My name
The capisal of of the USA ie Mashinngton D.C, The capicaa of ladia is How Duibi. capstad sa Loadon. The capstel
67
Tichanli
130
10"1-
Stur
Chaza
fiweling
The feallin!
thwling!
Log
The feeling!
The fewling!
The cewling!
The fooling!
The focling!
107
Tresp
Naclos Enoobold
Ria thondaid
Kpelac tardoold
10-1
Ba-q accept
Bavq
GLapoE
10"4
624
Tobo (ondervel by prokodvility?
Figure 4: Additional qualitative examples, following the same setup as Figure 1.
95
BA-1
MAUVE
rj
90
Nucleus
I BA
85
Dle
Small
Medium
Large
XL
Model size
Figure 5: MAUVE scores for sampling methods on Open Web Text test set. No single sampling method consistently outperforms across sizes. BA-1) performs remarkably well for GPT-2-Large.
vanilla counterparts. We do, however, see that BA- outperforms 17 sampling for the two larger model sizes, and does particularly well against all methods for GPT-2-Large.
BA-r outperforms 1j sampling in low-entropy decoding across model sizes We compare BA-1 and 17 sampling across different 7) parameters, again matching our BA-1) parameter reject the same proportion of human text as the parameter. As shown in Figure 6, we find that for more restrictive sampling (i.e., larger 7, closer to greedy decoding), BA-7 consistently outperforms 1) sampling. To verify our results (since we know from Figure 5 that model size effects which method is best) we show in Table I that this pattern holds across all model sizes.
More constraints improves BAT Since reduce the number of constraints in the linear program to make it run quickly, we can add constraints back into to program to verify that the bwsre constraints are the reason for the gains in BAT sampling. We again adjust the BAT parameter match the proportion of rejected human text to control for the additional tokens added to the support from the new constraints. Figure 7 shows that adding more BA constraints indeed increases the MAUVE score for our method. This is direct evidence that controlling for the softmax bottleneck helps reduce errors in the model distribution.
Human annotators narrowly favor BA-7 and prefer coherence to diversity To support our automatic evaluations, we additionally use human annotators from Amazon Mechanical Turk to compare both methods. Annotators tasked with pairwise comparisons between generations from
8

---

94
92
MAUVE
92
MAUVE
90
90
88
BA-1)
88
86
86
10
10 10-3
10
10
20
30
7 closer to greedy)
Constraints
Figure 6: MAUVE scores for GPT-2-XL with BA-n and sampling for different v) Under low- entropy generation (i.e., closer to greedy), BA-1) consistently outperforms 7) sampling.
Figure 7: MAUVE scores for GPT-2-XL as the number of BA constraints varies. BAT sampling improves with more constraints.
Table 1: MAUVE scores for different GPT-2
Table 2: Pairwise human evaluation results. BA-
model sizes on low-entropy generation. BA-1) sampling outperforms each case.
? = x indicates the BA-r) parameter chosen to match 2.
Size Method
Small
Medium
Large
XL
Method
Method
wins
2 wins
Tic
BA-1 0.002
0.002
0.43
0.38
0.19
"
85.01.4
90.4o.1
86.0o.s
87.11.2
BA- 0.024
0.024
0.05
BA-n
87.81.0
92.20.6
88.40.s
89.60.4
BA-1 0,024
BA-1 0.001
0.50
0.42
0.08
each method and generated from the same prefix. See Appendix F.1 for more details. Table shows that, annotators narrowly prefer generations from BA-1 sampling to those from sampling. Furthermore we see that human annotators prefer lower entropy generations. This is likely because humans only see generation per method, making it impossible to assess diversity in the generations.
5.2 DISCUSSION
Overall, our results provide empirical evidence that the softmax bottleneck is responsible for signifi- cant errors in language model next-token distributions, and show that BAT sampling offers a viable method for mitigating those errors. Under low-entropy generation, BAT offers clear advantages to threshold sampling, where only few tokens are permissible.
Although our pilot study shows promising results for BA-1 sampling in low-entropy generation settings, there remain a number of limitations. For instance, as mentioned in $5, BAT does not pair well with nucleus sampling. Furthermore, sampling parameters, BA-c accepts no tokens. This is a non-issue for threshold sampling which can fall back greedy sampling, but because BAT relies rejection sampling, it is not known when to revert to greedy. Though it is possible to implement a max-retries guard, this remains computationally expensive and the generations themselves tend degrade.
find that for certain prefixes and sufficiently low-entropy
broader issue that BAT must deal with is the expensive computation associated with running the linear program. While this generally not issue for generation, certain tasks infeasible, such as finding the exact set of candidate tokens, which would require running the linear program on the full vocabulary. We remain optimistic that further optimizations the method be made allow this in future work, as well as enable BAT sampling with higher constraint counts.
CONCLUSION
Our work fills crucial gap in the theoretical understanding of truncation sampling methods and how they account for language model errors. These theoretical findings translate into a more direct method for mitigating errors due to the softmax bottleneck. As result, our BAT sampling method can discard higher-probale probability tokens while keeping higher-quality but lower probabiroba tokens. Lastly, our pilot study with BAT sampling shows promising results in low-entropy generation.

---

REFERENCES
J. Aitchison. The statistical analysis of compositional data. Journal of the Royal Statistical Society: Series (Methodological), 44(2):139-16)) 1982. doi: https//cch.og//0..111//.7777// tb01 URL https: http:/ssssooonnnnneeeeee.eeeeeee wwarrrar.wllll.... wttp://///////c.o/ooooo/ 2517- 6161. 1982. tb01195.x.
MOSEK ApS. MOSEK Optimizer API for Python 9.3.22. Version 10.0., 2023. URL https:
//docs mosek. com/ 9. httppp/tooappiinnneeee p
html.
Sourya Basu, Govardana Sachitanandam Ramachandran, Nitish Shirish Keskar, and Lav R. Varshney, Mirostat: A perplexity-controlled neural text decoding algorithm. In ICLR, 2021. URL https:
//openreview..om wtt.://ooooommmdd11111111
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: suite for analyzing large language models across training and scaling. In ICML, pp- 2397-2430, 2023. URL https: //aaxiiv.vrr. abs 2304 01373.
Haw- Haw-Shiuan Chang and Andrew McCallum. Softmax bottleneck makes language models un- able represent multi-mode word distributions. In ACL, volume 1, 2022. URL https: //aclantho ogy.
acl.:222cllllln.555551 554/-
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, et al. PaLM: Scaling language modeling with pathways, 2022. URL pttstts://arxxx org/abs/ 2204.222. 02311.
Alexandra Delucia, Aaron Mueller, Xiang Lisa Li, and João Sedoc. Decoding methods for neural narrative generation. In Proceedings of the Ist Workshop on Natural Language Generation, Evalua- tion, and Metrics (GEM 2021), ), pp. 166-185, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/201111110 1.16. URL https: htt.tttts//acaaae http:::ss//claahhologg.. ogg.:lgo2y...111.... gem-1.16.
David Demeter, Gregory Kimmel, and Doug Downey, Stolen probability: structural weakness of neural language models. In ACL, pp. 2191-2197, 2020. URL httpa////alaaaalloo.... org/2020 acl-main. 198/.
Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical neural story generation. In ACL, pp. 889-898, Melbourne, Australia, July 2018. doi: 10.18653/v1/P18 1082. URL https:// aclanthology.c wwtp:/yy.rr////0//00011101000010
Markus Freitag, Behrooz Ghorbani, and Patrick Fernandes. Epsilon sampling rocks: Investigating sampling strategies for minimum bayes risk decoding for machine translation, 2023. URL https:
/ //arxiv.. otgo abs 2305. 09860.
O. Ganea, S. Gelly, Gary Bécigneul, and Aliaksei Severyn. Breaking the softmax bottleneck via learnable monotonic pointwise non-linearities. In ICML, pp. 2073-2082, 2019.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Open Corpus, 2019. URL http://Skyliono/ http://kkkliin000..itt.t.... www.ggttthuugg.. ltt http:///....eeeeeeeeeetooooopppm
Andreas Grivas, Nikolay Bogoychev, and Adam Lopez. Low-rank softmax can have unargmax- able classes in theory but rarely practice. In ACL, pp. 6738-6758, 2022. URL https: / //aclanchology org. /tlgg//22..oo2 a0l-oncn....n 465.
Tatsunori B Hashimoto, Hugh Zhang, and Percy Liang. Unifying human and statistical evalua- tion for natural language generation. In NAACL-HLT, pp. 1689-1701, 2019. URL https:
/ //aclanthooooe wtt./o1--1119...//
John Hewitt, Christopher Manning, and Percy Liang. Truncation sampling language model desmoothing. In EMNLP, pp. 3414-3427, Abu Dhabi, United Arab Emirates. December 2022.
URL httttttp///aaae anthhln org. fi ndd indings-emn 249.
10

---

Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious of neural text de- generation. ICLR, 2020. URL
net/ ittp://wwweeeeeettoooooiiiooeeee
Daphne Ippolito, Reno Kriz, João Sedoc, Maria Kustikova, and Chris Callison- Burch. Comparison of diverse decoding methods from conditional language models. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 3752-3762, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/1//111-1365 URL https: http://aaaaaathlhoaeoooo -ttp://oooy.rr/11111111111 1365.
Fred Jelinek. Self-organized language modeling for speech recognition. Readings in speech recogni- tion, pp. 450-506, 1990.
Kalpesh Krishna, Yapei Chang, John Wieting, and Mohit lyyer. RankGen: Improving text generation with large ranking models. In EMNLP. pp. 199-232, Abu Dhabi, United Arab Emirates, December 2022. doi: 10. 18.18833/v wttp:////2222..mmmmaaaa 15. URL https: https///aaaaaaahoaaa.gee org/2022. emn1p-ma m 15.
Xiang Lisa Li, Ari Holtzman, Daniel Fried, Percy Liang, Jason Eisner, Tatsunori Hashimoto, Luke Zettlemoyer, and Mike Lewis. Contrastive decoding: Open-ended text generation as optimization. In ACL, pp. 12286-1231, Toronto, Canada, July 2023, doi: 10. 12118866/11222..eea https: http:::///aaaathhlloooo
URL
org/ 2023.023.2002 1.002.accl.. acl -ong- 687.
Clara Meister, Ryan Cotterell, and Tim Vieira. If beam search is the answer, what was the question? In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 2173-2185, Online, November 2020. Association for Computational Linguis- tics. doi: 10. 1111116133122200... http://222002mmmmmmaeeee ewn.e-mmin.nnnmnmm.eee-.--m. URL htl.httss///acaatttaae emm.pp ma in.inn... 170.
org/2020.
Clara Meister, Tiago Pimentel, L. Malagutti, Ethan Gotlieb Wilcox, and Ryan Cotterell. On the efficacy of sampling adapters. Annual Meeting of the Association Computational Linguistics, pp. 1437-1455, 2023a.
Clara Meister, Tiago Pimentel, Gian Wiher, and Ryan Cotterell. Locally Typical Sampling. TACL, 11: 102-121, 01 2023b. ISSN 2307-387X. doi: 10. 162/tacl URL https: http:/dooooooooooooooooooooe.oed orr// 10.1162/4 1t.1162/tacl... 00033.
OpenAl. Gpt-4 technical report, 2023. URL UtL.ttpsss//azwww.or/////s
Krishna Pillutla, Swabha Swayamdipta, Rowan Zellers, John Thickstun, Sean Welleck, Yejin Choi, and Zaid Harchaoui. MAUVE: Measuring the gap between neural text and human text using divergence frontiers. In A. Beygelzimer, Y. Dauphin, Liang, and Wortman Vaughan (eds.), NeurlPS, 2021. URI. https: htt.httts///oppeeeeeeee
net worumrrmmioo ddrrrmmmdddTqxnnnp77..
Raj Reddy. Speech understanding systems: summary of results of the five-year research effort at Carlegie-Mellon University. 1977. URL https://kilthut cmu emw..ed/artiillee
summary_maoooo inderstanding__sstteeee
hfror_fffott_ee results_of_tiee L ottpr//e__eeecrraaeeoooooo__..ar Uhiversity
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b- parameter open-access multilingual language model, 2022. URL https: /tt///axivv.///////////e///./aaa wwg//bs.......ooooo..oo...oo..b. 2211.05 05100.
Rico Sennrich, Jannis Vamvas, and Alireza Mohammadshahi. Mitigating hallucinations and off- target machine translation with source-contrastive and language-contrastive decoding, 2023. URL
https: //arxivaaviv//i//ia org/ 2339 ..0000000000 07098.
Felix Stahlberg and Bill Byrne. On NMT search emrors and model errors: Cat got your tongue? In
Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference Natural Language Processing (EMNLP-IJCNLP), pp. 3356 3356-3362, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10. 10.18653/v1 1/111/6//1/1111113/ 1331. URL htwpt:tpp:///e anthol org.ogggooooooooooooooo 9- 1331. 13331.1..11
11

---

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, et al. Llama 2: Open foundation and fine-tuned chat models, 2023. URL httpttt/a//aii..ortt abs/2307, 09288.
Zhilin Yang, Zihang Dai, Ruslan Salakhutdinov, and William W. Cohen. Breaking the softmax bottleneck: A high-rank RNN language model. In ICLR, 2018. URL httpp:://pperevv///
Zhilin Yang, Thang Luong, Russ R Salakhutdinov, and Quoc V Le. Mix- tape: Breaking the softmax bottleneck efficiently.
In NeurIPS, volume 32, URL https:: nips ctt.ccccppppppp papprrapppp.ee/ee files/ files/paper/201 9/ /
2019.
A FURTHER RELATED WORK
Generating from autoregressive language distributions. Generating strings from autoregressive, optionally conditional, generative models of language has a long history NLP; for decades, algorithms were developed for approximating the maximum-likelibee (Jelinek, 1990), e-g., beam search (Reddy, 1977), under the understanding that there is one best output for, e.g.. speech recognition. As Transformers became used for general-purposee and often high-entropy generation wherein one wants to be able to generate multiple good completions, it was found that search for the highest-likelihood strings led to low-quality generations (Fan al., 2018; Holtzman et al., 2020; Hashimoto et al., 2019; Stahlberg & Byrne, 2019; Meister et al., 2020). In developing algorithms for high-entropy generation, the afore- mentioned line of work attempts to maintain the learned distribution as much as possible (Holtzman er al., 2020; Hewitt et al., 2022); another significant design principle has related to the uniform information densily principle that humans are observed to obey, motivating Meister et al. (2023b). This algorithm intentionally deviates from the overall distribution more, by sometimes truncating high-probabile tokens in order never generate any tokens that are too high probability relative to the overall entropy. Krishna et al. (2022) show that language models do not effectively make use of long-term context, finding that an explicitly trained re-ranker can help. Li et al. (2023) hypothesizes that language model errors are distributed similarly small models as large, showing that taking the difference of their logits can help improve the large model's generations. Some ideas from high-entropy generation, including this, and the e-sampling algorithm, have shown to be useful even in low-entropy generation where previously algorithms like beam search have performed best (Freitag al., 2023; Sennrich et al., 2023).
string under the model
Did the softmax bottleneck turn out not to be a problem? After the demonstration of the softmax bottleneck by Yang al. (2018), various algorithms were proposed for efficiently learning high-rank language models (Yang et al., 2019; Ganea et al., 2019). Chang & McCallum (2022) showed that the softmax bottleneck makes certain multi-mode distributions difficult to model, while Demeter et al. (2020) demonstrated that the low-rank nature of language models means that it is possible for certain word tokens to be unable be the argmax, but Grivas al. (2022) demonstrated that this is rarely the case in practice. Overall, rank considerations have not been at the fore of language model development, as language models have scaled, their hidden state sizes have scaled as well, but stayed smaller than their vocabulary sizes (Scao et al., 2022; Biderman et al., 2023; Touvron et al., 2023). Throughout this time, when one generates from language models, one almost always lowers entropy and performs some kind of truncation sampling (or in the extreme, greedy decoding). Our results suggest that training high-rank language models may appear unnecessary because default truncation sampling mitigates errors stemming from the low-rank approximation.
B BAT SAMPLING ALGORITHM
Algorithm gives the procedure for BAT sampling.
12

---

Algorithm BAT sampling
1: procedure BAT(Threshold Next-token distribution p)
2: repeat
3: Sample
4: until ëp such that p. = 0 A W p= W
Use LP solver
5: return i
6:
end procedure
C THE SOFTMAX FUNCTION IS LINEAR
It is an unintuitive fact that the softmax function is linear map R2 + Aa. The key here is that addition and scalar multiplication are defined on Aa in a non-standard way. The elements of Ag tuples of length d whose entries sum one. Vector addition @ in As is defined as clementwise multiplication followed by normalization
p C q Σ P:q.
p
(4)
and multiplication @ by constant À R is elementwise exponentiation followed by normalization
A p 21=1
(5)
One can check that these operations satisfy the axioms of a vector space, and that the softmax function satisfies additivity and homogeneity under these operations, i.e..
softmax + v) softmax(()) soflmax(())
(6)
and
softmax(Au) = X & softmax(u).
D PROOFS
Proof of Theorem 1. By the precondition of the theorem, we have log P log for all i. It
follows that:
6; 2 p exp(-5).
(8)
Intuitively, since p is valid probability distribution summing to 1, if it cannot underestimate token probabilities beyond a factor of exp(-5). then also cannot overestimate other tokens' probabilities beyond certain factor; we will show that this factor is exp(-6).
To this end, we consider each token individually and calculate the maximum possible probability overestimation based on the maximum probability underestimation of the other tokens. Keeping in mind that any probability added to a token must be removed from other tokens to preserve valid probability distribution, the maximum probability added token is the sum of the maximum probabilities subtracted from the other tokens. This gives us that for all i:
P. P.
Σ PK
Σ Pk
(9)
Righi
lif
<
Pk
P% exp(-o)
From (8)
(10)
k
(1 exp(-6)) Σ Pk
Factor out pk
(11)
Aip (1 exp(-00 6)) -2,)
Probabilities sum
(12)
1 - exp(-o)
0 0 P < 1.
(13)
We thus have our desired probability overestimation bound, starting with the assumption of a log- probability underestimation bound.
13

---

Proof of Theorem 2. We begin by assuming that model has learned to minimize cross entropy with the true distribution, implying that
a
crossentropy (Wh, p")) 0.
(14)
Oh
Expanding and simplifying this equation, we can obtain
Σ pt log log(softmax(Wh))
cross entropy defn.
(15)
Ōh
Σ p Wh P log] Σ exp(Wh))
Log of softmax (16)
Ōh
Σ
p log
exp(Wh))
0h. Σ (WW);
Distribute
(17)
Oh
Oh
0
0 E (WW))
log
exp(Wh))
Σ Pt =1 (18)
0h
0h
E, exp(Wh)) 8 (Wh)) 0 on
8h Σ ((WW)):
Derivative (19)
exp(WWh),
0 Sh
(Wh) exp(Whh) 0 (Wh)"p' p'
Factor (20)
Σ, exp(Wh), Sh
a (Wh)'''
0h (Wh) (ith))sssotii softmax Wh)
0h. p'
Softmax defn.
(21)
W p W'p'' p'
Derivative
(22)
where P is the output distribution of the model. Thus, if there does not exist any valid probability distribution p such that p. = 0 and W p= W p. then pt o
E BWSIE THRESHOLD SAMPLING IN PRACTICE
Basis-aware sampling presents a number of practical challenges. Chief among them is the sheer size of the linear programs to be solved. These programs have v variables and d - 2u + 2 constraints. No open-source solver we tried was able to solve single problem in a reasonable amount of time, avoid hitting numerical errors, and solve within its default max-iteration limits. Proprietary solvers do better in some cases, but only the MOSEK solver (ApS, 2023) was able to solve the full problem in under minute. Even this relatively faster solving rate makes text generation at scale impractical.
To address this, we reduce the size of the linear program dramatically by discarding many constraints. While doing so, however, also aim to maintain much of the original solution space possible, so as to minimize the effect on the set of tokens discarded by basis-aware sampling."
In order reduce the number of constraints originating from the WT WT P term from d to c, can simply discard any d columns of W to obtain W Clearly, if p satisfies W p= W p. it will continue to also satisfy WeT p W p. Thus, if token was originally rejected by bottleneck-aware sampling, it would still be rejected, i.e., using We instead of W does not add new candidate tokens. may, however, remove some candidates, and we would like minimize this effect.
Suppose W has rank b < d. Then the set of probability distributions p satisfying WTp W forms a linear subspace implying the set of distributions p satisfying the relaxed condition W p W p forms linear superspace S of of dimension at least d min c}. Recall that the larger S is, the more candidate tokens will be removed by bottleneck sampling. Thus, to minimize candidate removal, we seek an 5e that of dimension exactly d min c). This can be achieved easily by keeping We any set of min (b, c} linearly independent columns of W. Note that if b S c, the use of such a We
C " of dimension d b. Further, W has rank most min(b,c c).
WWithout any constraints, hasis-aware threshold sampling reduces basic threshold sampling.
14

---

will, in fact, not remove any candidate, S will equal S. Otherwise S will be d dimensional superspace of S.
When b > c, however, this solution is still not optimal, which linearly independent columns of W we choose to keep in W determines how "close" S will be to the original solution space S. Intuitively,
would like to preserve S along dimensions that correspond to the largest eigenvalues of W. To accomplish this, we turn to singular value decomposition: find three matrices U E R Σ Rdxd and V Radd such that W = UEVI, then replace W with U R where U represents the first columns of U. Since U is simply a linear transformation of W, the solutions (in terms p) of U p U are precisely the subspace S of dimension b as before. Again, as before. replacing W with U does not add new tokens to the set of candidates, and may remove some candidates when b c. Importantly, when b c, U will intuitively be the "closest" possible approximation of W (capturing its e largest eigenvalues). Thus, S will form a desirable approximation of S.
The above SVD based approximation is what we use in practice. This reduces the number of constraints from (+s 700-1200 for our models) to e (typically 20), and shortens the run time from over minute proprietary solver to about second.
F PARAMETER SELECTION
Table 3: Parameter sweeps and chosen parameters for each method and size
Method
Swee
Small
Medium
Large
XL.
Nucleus
{0.89, 0.9.0.92 0.92, 0.95, 0,99)
0.92
0.89
0.92
0.95
(0.0003, 0.0006,, 0.0009, D.00I, O. D022
0.0009
0.0003
0.0009
0.0003
(0.0003, 0,0006, 0,0009, 0,002, 0.D04)
0.0009
0.0009
When comparing sampling methods, choice of parameters is very important, since each method has its own iiversity-coheremt trade-off characteristics. Without proper controls, impossible to tell whether the performance gap between two heuristics might be closed by simply adjusting the parameter of the worse-pent perssee-rfrrmingg method. To remedy this, we control for parameter choice by matching parameters of compared methods based on how conservative they are with respect to human text. In particular, for each vanilla threshold sampling method x, we choose the BA-c parameter that rejects the same proportion of tokens from human corpus. Table illustrates how measure this human-texi rejection rale (HRR). In our experiments, measure HRR by sampling 10,000 tokens with their prefixes from Open Web Text and calculating the proportion of the tokens that accepted sampling method with a given parameter.
Table 4: With hyperparameter of 0.002, e-sampling would have a human- rejection rate of 1/5 on this text.
Token Probability
I'm 0.02
the
problem, 0.01
it's
me.
0.3
0.001
0.3
Figure gives the sampling parameters as function of HRR. As HRR approaches zero, parameters become more permissive, i.e., nucleus approaches one, 1) and approach zero, in order to accept more tokens. We observe that as HRR increases, BAT parameters are consistently more conservative than their vanilla counterparts since BAT methods sample tokens beyond the threshold. In the case of BA-p, the parameter maxes out around 28% HRR, meaning that cannot reject more than 28% of human tokens.
F.I HUMAN EVALCATION
Annotators are paid $1 USD per annotation, and each annotation takes on average less than minutes. Figure 9 provides the exact instructions and layout given the annotators.
15

---

1.0
Epsilon
0.8
Eta
Hyperparameter
Nucleus
0.6
BA cpsilon
0.4
BA
BA nucleus
0.2
0.0
10
20
30
Rejection (%)
Figure 8: Truncation sampling parameters for various methods by HRR, the proportion of human-text the sampling methods reject.
an option
Text 1
Test As cokorful dicus of tenta Excrost Base Camp this spring, pa of
Which higher qualitty?
As colortul cireus of tants popa up at Everest Base Camp this spring a pair up camp the Moon They may be the frrst human belnge ever live a year in spece,
higher quality
studomss
feet
higher quality
peececal
2,400
They are equal quality
boore difoveet four look to pic the
the worldss largest human-poweer ponverd expedilng on camp. The team assembliig tbed by Inventorn
lever touas with grammas seseenc unen inss pooide that nets sesiies orehully and detoriine which in
On July 31, Astionaur Cameron...
Explain your cholos a line worcs.
may sefect weiter
te snerrshae tha one of the ne may
ancing straety
Figure An example of the interface and instructions shown to human annotators.
G TRUNCATED LANGUAGE MODEL DISTRIBUTIONS ARE HIGH-RANK
We motivated truncation sampling as helping to correctly discard tokens that are not in the support of the true distribution p" when those errors due the low-rank nature of language models' distributions. In this additional experiment, we show that the cost-truncation conditional distribution matrix A is high-rank relative to the pre-truncation distribution.
We run the GPT2-XL model on samples of OpenWebText concatenate the conditional log- distributions log P for each prefix, and compute the rank of the resulting matrix. This becomes a rather large matrix, since each log P is in R 50257 so we are limited in the number of prefixes we can consider. Since the number of prefixes upper-bounds the estimated rank, and we cannot run, e.g., 50257 prefixes, we plot the rank for various numbers of prefixes. We find that the GPT2-xl model, which has hidden dimensionality of 1600, has rank that saturates at 1600, as expected. For truncation sampling strategies nucleus, sammmig and e-sampling, find that the estimate of the rank continues to grow with the number of prefixes, far past 1600. See Table 10.
H MORE UNIT TESTS
We give the unit tests used Figures and tabular form (Tables 5-11).
16

---

GPT2-XL XL Rank Estimate as Function of Prefix Count
Nueleus
Elu
Epeilon
Estirtated Rank
Ancustaa (Unmmshileed)
3000
2000
Mosil Dinencieddits
2009
0000
5000
10665
Prefix Count for Estimate
Figure 10: The estimated rank of the log- probability distributions of model with truncation grow far past its hidden dimensionality; without truncation, the rank is constrained to the hidden dimensionality.
Table 5: A subset of the next-word distribution according GPT2 for the context "< endoftext >>Taylor.. The last four columns denote whether each token is in the sup- port of the titular strategies. Notice that BA-7 able accept good continuations like will' and Hanson' while excluding questionable continuations like Sw" which likely has higher probability because of its embedding alignment with Swift'.
Kank
Prob
Taken
BA Eio
Eta
Epsiion
Nucicus
4.3e-01
Swift"
True
2.1e-02
1s"
True
True
True
True
1.9e-02
Eail"
True
True
Fulse
True
2.2c-01
Sw"
True
False
True
2.2c-01 31
True
True
Fulse
True
2.0--11 35
Hiller"
True
True
False
Tre
1.9c-03
Thue
True
Fulse
Fals
39 1.9c
Fulse
True
False
False
40 1.9-00
eays' Lew'
True
True
Fulse
Falso
42 1.3e-0
False
False
Palse
43
1.7-00
True
Fulse
Fulse
False
44 1.7e-01
Eicks"
Thue
Fulse
Fulse
Fuls
1.6e-03 46
Fubs
Fulse
Fulse
False
47
1.fe-B
Mood"
Tiue
False
False
False
50
1.5c-0
Eoin'
True
False
False
60 1.2-00
Lee!
False
False
False
Palse
61 Lle-0
Kizz'
True
False
False
False
82
L.e-00
Hartin'
Fuls
Filse
Fulse
Falso
alao"
Table 6: <llenoftext >y name
Prob
Token
Fpsilon
Nucleus
9.6c
True
True
True
True
1.2e-00
was'
True
False
Palse
Palse
1,0e
Filse
Fulse
Fulse
Falso
7.7e-04
isn'
True
Fulse
Fulse
Faloc
4.le00
and'
Fabe
Fulse
Fulse
Palso
22
5.le-05
I5'
Fubse
Fulso
Fulse
Falso
17

---

Table 7: << endoftext
>yy name is'
Rank
Prob
Token
BA Eta
Eta
Epsilon
Nucleus
0
1.1c-02
David'
True
True
False
True
20
5.0e-03
Adam'
True
True
False
True
1156
1.3e-04
Ily'
True
True
False
False
1167
1.3e-04
Curt'
True
True
False
False
1168
1.3e-04
Sk'
False
True
False
Falsc
1169
1.3e-04
Stewart'
False
True
False
False
1170
1.3e-04
Avery'
True
True
False
False
1175
1.3e-04
Aud'
True
True
False
False
1176
1.3e-04
Eb'
False
True
False
False
1177
1.3e-04
Brock'
False
True
False
False
1178
1.3e-04
Franc'
True
True
False
False
1184
1.3e-04
Mercedes'
True
True
False
False
1185
1.3e-04
JJ
True
False
False
False
1194
1.2e-04
Sebast?
True
False
False
False
1195
1.2c-04
Di
False
False
False
False
1196
1.2c-04
Maxwell'
True
False
False
False
1205
1.2c-04
Mand"
True
False
Falsc
False
Table 8:
< eetof The capital of of the USA is Washi ngttn D.C.
The capital of India is New Delhi. The capital of the UK is Lendon. The capital of Ghana is'
Rank
Prob
Token
BA Eta
Epsilon
Nucleus
9.9e-01
Acc'
True
True
True
True
6.8c-03
Ab'
True
False
False
Falsc
2
6.7e-04
Kum"
True
False
False
False
3
2.6c-04
Con'
False
False
False
False
2.3e-04
Ghana'
True
False
False
False
1.4e-04
Abu'
False
False
False
False
22
1.2e-05
Tem'
False
False
False
False
Table 9: '</endoftext >Donald"
Rank
Prob
Token
BA Eta
Eta
Epsilon
Nucleus
0
9.4e-01
Trump'
True
True
True
True
1
1.e-02
J
True
False
False
False
2
6.2e-03
Glover'
True
False
False
False
1.9e-03
Sterling'
False
False
False
False
22
3.1e-04
Donald'
False
False
False
False
Table 10: *< </endoftext Thee'
Rank
Prob
Token
BA Eta
Eta
Epsilon
Nucleus
1.3e-02
first'
True
True
False
True
20
2.9c-03
Amcrican'
True
True
False
True
4338
3.5e-05
RNC"
True
True
False
False
4340
3.5e-05
poet'
True
True
False
False
4354
3.5e-05
True
True
False
False
4357
3.5e-05
inevitable'
True
True
False
False
4358
3.5e-05
hackers"
True
False
False
False
4359
3.5e-05
Bright'
True
False
False
False
5232
2.8e-05
PBS'
False
False
False
False
5233
2.8e-05
Granmy'
False
False
False
False
18

---

Table 11: '< endofftett |>The feeling! The feeling! The feeling! The feeling! The feeling! The feeling! The feeling! The feeling!
The feeling! The feeling! The feeling!"
Rank
Prob
Token BA Eta
Eta
Epsilon
Nucleus
9.5e-01
The'
True
True
True
True
2.3e-02
'nn'
True
False
True
False
2.6e-03
THE'
True
False
False
False
2.2c-03
"Annn'
False
False
Falsc
False
1.8c-03
'I'
False
False
False
False
9.9c-04
'The'
True
False
False
False
6.3e-04
-It'
False
False
False
False
22
1.3e-04
My"
False
False
False
False
19


Reviews:
[{"review_id": "K1uEotYN5O", "reviewer": "Reviewer_pgkn", "summary": "This paper aims to give a formal justification for why truncation based sampling approaches work well in language generation. They link this phenomenon to the softmax bottleneck—the problem that the final linear layer in a neural network often bottlenecks the expressivity of the model. Explicitly, given the difference between the hidden (embedding) dimension and the vocabulary size, the final linear transformation before the softmax projection can only perform a low-rank projection. The authors claim that the resulting approximation to the target distribution is likely the source of model errors that leads to the “unreliable tail” probabilities observed by prior work. The authors develop an algorithm for uncovering which tokens are necessarily in the support of the target distribution, and propose to use this algorithm as the basis for a truncation sampling method. They provide empirical results (including human evaluations) when using this method.", "questions": "* Since the matrix W is static, can (3) not be solved for all elements of the vocabulary that meet the desired constraint ahead of time?\n* Why is typical sampling omitted from the discussion of truncation-based sampling methods? It is arguably the most similar to the proposed method since it likewise “is able to discard higher-probability tokens while keeping… lower-probability tokens.” On a similar note, I don’t understand footnote 3; I don’t think it actually describes what is done by locally typical sampling\n* The phrasing of footnote 1 is strange. Specifically, the statement “setting p∗ to be the 1-hot vector indicating the gold token” is underspecified. I imagine that this is referring to the conditional distribution for a particular prefix\n* In theorem 1, these factors are the collective probability over/underestimation across all tokens, right? This then implies that no individual token probability can exceed these bounds. This logic should be made more explicit (the current notation is vague) \n* I don’t think that the softmax function satisfies the additivity property required of a “linear map.” Could you please elaborate on this claim at the top of page 6?\n* On page 5, what is the concrete distinction between low vs. high quality tokens? Is this another way of saying in vs. out of the true distribution? It would be helpful to change the language here to align with the other terminology used by the paper\n* The informal description about the practical implementation of basis-aware sampling is confusing. For example, what does “discarding the majority of the constraints” refer to?\n* Given the observation that BAT sampling performs more strongly in lower entropy settings, a logical next step would be to see how it performs in translation or summarization, where historically, sampling has not led to the best results.\n* It is perhaps more accurate to call the projection by W a linear layer instead of softmax layer, since the low-rank approximation is tied to this linear transformation, not the use of the softmax as a projection function. Further, are there insights into how the nature of these results will change with alternative projection functions, like the sparsemax?\n* Can these principles be used to explain the degeneracy that happens when selecting high probability tokens, i.e., during greedy decoding?\n* How do these results align with other work that has tried to explain why truncation methods work well in practice, such as [1]?\n\n[1] Meister et. al. 2023. On the Efficacy of Sampling Adapters.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The work offers a theoretical explanation for why certain ad-hoc methods used during language generator decoding work well. This is a valuable insight to the NLG community\n* The work then develops a sampling algorithm based on this theoretical explanation", "weaknesses": "* The theoretical portion of the paper is at times difficult to understand due to notational choices and lack of specificity (for example, switching between individual token probabilities ). This is particularly important since the theoretical portion is the main contribution of the work \n* The method does not appear to have empirical performance benefits and is computationally expensive, making it impractical\n* There lacks robust empirical justification of the hypothesis. Figure 7, which is intended to show that including more of the original optimization problem constraints lead to better results, only consists of 3 points, which hardly feels like enough evidence to claim a “trend”.\n* The terminology of the “true” distribution is perhaps misleading. I personally think that something like the “aggregate” distribution or the “data-generating” distribution would be more accurate\n* A small point: the intro of section 4 has some grammatical errors", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rrgW7V8GfO", "reviewer": "Reviewer_Hq1N", "summary": "The paper proposes a theoretical understanding of truncation sampling methods. The authors proceed to devise a novel sampling strategy, building upon the approximation error incurred by the softmax output bottleneck. Essentially, the idea is to assume that a token has a non-zero probability under the true distribution, not only if it has a non-zero probability under the predicted distribution (which is a source of overestimation errors) but also if it is non-zero under all distributions that \"map back\" to the hidden state by taking the transpose of the output embedding matrix. Intuitively, this measures whether a token has a non-zero probability by chance, i.e. if its information can be conveyed by any combination of other tokens while mapping back to the hidden state. This is formalized in the paper by assuming that the hidden state is the minimizer of the cross-entropy loss with respect to the true distribution. This insight serves to devise a novel sampling strategy that can sample low-probability next tokens, which differs from current approaches that rely on thresholding.", "questions": "Thank you for the efforts in writing a clear and enjoyable paper.\n\nMain questions:\n- Can you include a pseudo-code of the final BA implementation in the main paper?\n\n- Regarding Eq. 3: what I am going to propose is a bit dirty but would it be possible at each step to minimize (W^T p - W^T \\hat p)^2, wrt to p with a sparsity constraint (e.g. l1) and the range constraints, and reject all tokens for which |p| = 0? It might not derive from the theory but it might capture the overall idea? just wondering.\n\n- Concerning the results: there isn't much of a pattern in the MAUVE results if I look at the improvements of BA across model scales. Isn't  BA expected to help more with smaller scales given that the approximation error might be bigger?\n\n- Main problem: what happens with bigger models? given that the approximation error will be smaller, would your method still help?\n\nNitpicks:\n- It might be clearer to re-introduce the \\epsilon and \\eta baselines in the experiments. I struggled a bit to remember given that they are just introduced in the background section.\n- Eq. 10 in the appendix is missing a parenthesis.\n\nI would love to give a 7, but I can't (I have to choose between 6 and 8 now). I will give a 6 for now, and wait for authors responses with the will to increase my scores if more details / addition to the papers are given :-)", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Nice idea and analysis\n- Well written / clear\n- Shines new insights on a well-studied problem and could lead to more promising sampling methods", "weaknesses": "- Results are rather weak, efficacy of the method still remains to be demonstrated (minor)\n- An pseudo-code / algorithm box with the practical implementation of BA is needed in the main paper (minor)\n- Unclear whether the method will help for larger models or for models where the approximation errors (under-estimation / over-estimation) are small (kinda major).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0EYFL0hNoA", "reviewer": "Reviewer_8Qv3", "summary": "This paper provides a theoretical analysis of why decoding from language models with truncation sampling works well and provides a new decoding strategy called BAT-sampling. The gist is that if you truncate a sufficiently large portion of the distribution, then truncation sampling will avoid any tokens which are not in the support of the true language distribution. However, this may result in throwing out tokens which *are* in the support of the distribution. The paper then proposes BAT-sampling, which is an LP-based method that can determine which tokens are and are not in the support of the true distribution, even if they are \"out of order\" in terms of the probability assigned (i.e., the method can determine that lower-probability tokens are in the support of the distribution even when higher probability tokens are not). The paper concludes with a set of experiments, including an impressive discussion of speedups for BAT-sampling, as well as some (very slight) improvements over existing methods in certain conditions.", "questions": "1. Based on the figures (1,4), it seems like BAT is rejecting a lot of tokens corresponding to partial words. Out of curiosity: is this true, and do you have any insights into why this happens, or other qualitative insights into what tokens tend to get accepted/rejected?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This is a great analysis paper, providing an interesting explanation for why truncation sampling works so well in language model decoding. The paper's motivation is clear and well-written. The fact that BAT can determine that some tokens have nonzero true support, even though they are assigned less probability than others which are not in the support of the true distribution, is a surprising and compelling result. Leveraging the softmax bottleneck is a clever trick here and one that will be unexpected to most readers in NLP. \n\nI expected BAT to be computationally infeasible to run in practice due to its dependence on an LP-solver at each tilmestep of decoding. However, the speedups in the \"Basis-aware threshold sampling in practice\" (namely, using a decomposition of the softmax matrix and only relying on BAT when a token under the threshold probability is chosen) seem reasonable and compelling, and the amortized cost of 0.1s/token, while slow, is not infeasible for certain classes of applications.\n\nThe experiments, although not particularly compelling as a reason to start using BAT sampling in practice, seem reasonable and sufficiently thorough. In particular, the analysis of performance as more constraints are added back (after the SVD) is very clear. In contrast, I did not find the \"BAT outperforms all other methods for GPT-2-Large\" paragraph very compelling given that BAT is not the best-performing model on any other model size.", "weaknesses": "The primary weakness seems to be the performance of BAT compared to other methods. Despite its theoretical justification, it does not clearly outperform other sampling approaches (Figure 5). Although there is a preference for BAT to eta-sampling shown in Figure 6 and Table 1, this preference is very slight and the comparison is only between two sampling methods. However, I do not see this weakness as a legitimate reason to reject the paper, since its main contribution seems to be analysis and theoretical understanding of existing decoding algorithms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "U1ZdwK842i", "reviewer": "Reviewer_4bep", "summary": "This paper aims to develop a more precise sampling strategy for language models, specifically focusing on addressing errors arising from the softmax bottleneck. The authors establish two sufficient conditions, under certain assumptions, to ensure that sampled tokens belong to the true distribution's support. The first condition (Corollary 1) leads to the threshold sampling algorithm, providing a direct explanation for its success. Assuming the loss to be cross-entropy, the second condition imposes a linear constraint that can be solved with a linear programming optimizer.Combining both conditions, the proposed basis-aware threshold (BAT) sampling outperforms its threshold-based counterparts for low-entropy open-ended text generation.", "questions": "None", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed approach is innovative and theoretically-grounded. \n2. This paper brings new insights to the community. The theoretical concepts discussed in this paper are previously ignored but seems to be important.", "weaknesses": "1. Basis-aware sampling is specifically designed for models trained using cross-entropy loss. However, not all language models meet this criterion. For instance, LLMs fine-tuned with RLHF do not adhere to this condition.\n2. Theorem 2 and Corollary 2 provide sufficient but unnecessary conditions for proving tokens are in the true support. Therefore, the induced sampling algorithm may also discard tokens in the support of true distribution, leading to biased sampling.\n3. In Corollary 1, the statement \"threshold-based truncation sampling correctly discards all tokens that are not in the support of p*\" is not precise, as it may also incorrectly discard tokens that are in the support of p*. A more precise phrasing would be: \"all tokens that are not in the support of p* will be discarded by threshold-based truncation sampling.\"\n4. The relationship between the softmax bottleneck and text degeneration phenomena has not been verified. It remains unclear whether text degeneration is directly caused by the softmax bottleneck, or if increasing the dimensionality (d) beyond the vocabulary size (v) would effectively resolve the issue.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)