Paper: Closing the Curious Case of Neural Text Degeneration (ICLR-2024)

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

Author response:
We chose to focus on cross-entropy loss since it is the most popular method used to pretrain base language models. We agree that it would be worthwhile to consider the implications of alternative loss functions, though these are not guaranteed to give the “nice” linear constraints that we derive from cross entropy.

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

Author response:
This is indeed what we meant. We believe the confusion may have come from the word “correctly” which we will remove.

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

Author response:
Prior work [1,2] offers evidence that the softmax bottleneck (SMB) indeed negatively impacts language model performance, both in theory and in practice, and it has been shown that “breaking” the SMB can lead to better performance. A question that remains is whether (and when) this effect will have an impact on generation, since heuristics like threshold sampling may already mitigate the SMB effects to some degree. Our pilot experiments with low-entropy decoding seem to suggest that the SMB may still have a negative impact on low-entropy decoding that can be mitigated by BAT sampling.

[1] Yang et al., Breaking the Softmax Bottleneck: A High-rank RNN Language Model.

[2] Yang et al., Mixtape: Breaking the Softmax Bottleneck Efficiently.

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

Author response:
We agree that our pilot experiments show only preliminary evidence for the effectiveness of our methods, and we are glad that you share our view that these results are supplementary, rather than the primary contribution of our paper, which is the theoretical explanation and mathematically grounded algorithmic proposal.

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

Author response:
We agree that our pilot experiments show only preliminary evidence for the effectiveness of our methods, and we are glad that you share our view that these results are supplementary, rather than the primary contribution of our paper, which is the theoretical explanation and mathematically grounded algorithmic proposal.

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

Author response:
This is a great observation! Our leading hypothesis for why these partial words are rejected (e.g., `Swift` accepted, but `S` and `Sw` rejected) is that they have high embedding similarity with the full word. Tokens whose embeddings (in $W$) have high dot-product similarity with the model’s hidden state receive the most probability mass. In our example, the model output a hidden state with high dot-product similarity with `Swift`, and consequently tokens like `Sw` that have similar embeddings to Swift also have high dot-product similarity with the hidden state, resulting in them getting probability mass.

Using our method, in order to identify a target token for truncation, it must be possible to transfer all the probability mass on the target token to other tokens while maintaining the constraints of the linear program. The $W^\top\hat{p}=W^\top p$ constraint specifies that the probability mass must be transferred to tokens with similar embeddings to the target token, and the $p_j\leq\hat{p}_j\exp(\delta)$ constraint limits the amount of mass that can be transferred to any one token, where the higher probability a token is, the more probability mass can be added to it. Therefore, a token is most likely to be truncated if there is another token with high probability and high embedding similarity to it.
In our example, the probability mass on `Sw` embedding can be transferred to `Swift` which has high probability and (likely) high embedding similarity.

Reviewer comment (Evaluation):
Results are rather weak, efficacy of the method still remains to be demonstrated (minor)

Author response:
We agree that our pilot experiment results are largely preliminary and will benefit from further investigation in future work. We take the view that our experimental findings are secondary and meant to supplement our primary theoretical results and method proposal.

Reviewer comment (Reproducibility):
An pseudo-code / algorithm box with the practical implementation of BA is needed in the main paper (minor)

Author response:
Thank you for pointing these out, we update accordingly in our revision.

Reviewer comment (Experiments):
Unclear whether the method will help for larger models or for models where the approximation errors (under-estimation / over-estimation) are small (kinda major)

Author response:
One important implication of our results is that larger models do not automatically avoid under/over-estimation errors, since the errors caused by the softmax bottleneck (SMB) depend not on the model size, but on the discrepancy between the hidden size and the vocabulary size. A larger model with a proportionally larger vocabulary size will suffer the same bottleneck effects. By the same argument, a model with a small vocabulary (i.e., close to equal hidden and vocab sizes) could avoid the SMB altogether, no matter the size. Unfortunately, we do not see many models of this type in common use.

One size-related investigation we would like to pursue in future work is the hypothesis that our method will become **more** effective for larger model sizes. We hypothesize that some of our results were weaker than expected since our method primarily addresses SMB-related errors. If the model predicts the next token incorrectly, none of the considered truncation methods (including ours) will be able to correct the error. This issue should be reduced with larger models, since larger models should have fewer modeling errors, leaving only SMB errors. We see some anecdotal evidence for this where, for instance, BA-$\eta$ becomes stronger compared to $\eta$ sampling as model size increases (Figure 5).

Reviewer comment (Reproducibility):
Can you include a pseudo-code of the final BA implementation in the main paper?

Author response:
Thank you for pointing these out, we update accordingly in our revision.

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

Author response:
See our response to your concern above: there may be multiple effects here, such as the rate at which the model makes prediction mistakes independent of the SMB, which would confound the trend because these types of mistakes would become less frequent at larger model sizes.

Reviewer comment (Experiments):
Main problem: what happens with bigger models? Given that the approximation error will be smaller, would your method still help?

Author response:
Again, we refer to our size-related responses above, with an additional perspective: we expect that small models will continue to play a role in language generation, given the ease of deployment in edge devices, as well as cheaper inference costs. Even if our method does not provide benefits for larger models, smaller models will continue to suffer from a SMB and our method (or some modification of it) will remain relevant.

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

Author response:
Thank you for pointing these out, we update accordingly in our revision.

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

Author response:
We strived to follow the ICLR style guidelines for notation, and also to be as consistent and specific as possible. We caught a few notational errors in our revision too. We would appreciate more feedback on which parts are confusing.

Reviewer comment (Evaluation):
The method does not appear to have empirical performance benefits and is computationally expensive, making it impractical.

Author response:
We address the overhead at the end of Sec. 4: the overall addition is $\approx0.1$ seconds per token, which is acceptable for many applications. Additionally, this method can be further optimized in the future.

Reviewer comment (Writing):
The terminology of the “true” distribution is perhaps misleading. The reviewer personally thinks that something like the “aggregate” distribution or the “data-generating” distribution would be more accurate.

Author response:
We address this in Footnote 1, where we refer to the distribution “from which internet text is implicitly sampled”, which we believe is equivalent to the “data-generating” distribution.

Reviewer comment (Writing):
A small point: the intro of section 4 has some grammatical errors.

Author response:
See revision.

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

Author response:
This would be desirable, but we cannot pre-solve since $\hat{\boldsymbol{p}}$ is different for each context. Pre-compilation of the program could speed up solving.

Reviewer comment (Novelty):
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, the reviewer does not understand footnote 3; they don't think it actually describes what is done by locally typical sampling.

Author response:
See the new Footnote 3. We focus less on typical sampling since it is markedly different from the other truncation schemes in that it does not attempt to recover the support of true distribution.

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

Author response:
“Discarding the majority of the constraints” means that we use an approximation instead of the full $\boldsymbol{W}$ matrix to make the linear program faster. See revision.

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

Author response:
This is a great future work direction, which we are pursuing!

Reviewer comment (Theory):
Can these principles be used to explain the degeneracy that happens when selecting high probability tokens, i.e., during greedy decoding?

Author response:
We suspect that degeneracy in greedy decoding stems from sampling from an unnatural distribution. One possible explanation for why our method does better in low-entropy (close-to-greedy) decoding may be that there are some non-greedy tokens that our method refuses to truncate ever. As an example, in the final frame of Figure 4 most threshold methods continue the repetition, but BAT allows other tokens to be sampled.

Reviewer comment (Novelty):
How do these results align with other work that has tried to explain why truncation methods work well in practice, such as [1]?

Author response:
Meister et al. (2023) approach the question of “why do truncation methods work” by hypothesizing that threshold sampling reprioritizes precision over recall. In some respects, our explanation presupposes the Meister et al. hypothesis then proves guarantees on the mechanism by which threshold sampling achieves high precision. Other work [2] has also made the assumption that threshold sampling aims to avoid zero-probability tokens. We add Meister et al. to our revision.