Looking Beyond the Top-1: Transformers Determine Top Tokens in Order

We’re glad that you find our research question important and interesting.

Although, to the best of our knowledge, the existence of “saturation events” is not controversial in the literature [1,2,3], we agree that we should have provided some evidence demonstrating their frequency in the data. For this reason, we’ve added a section (currently in Appendix B.1) showing that in GPT2, saturation events are common for top-k tokens, with over 80% of samples reaching top-1 saturation. Even if we consider only cases where top-1 saturation happens in the first 85% layers of the model, as we do in all of our experiments, we find that this includes 51.5% of all samples.

To facilitate your evaluation, we have also uploaded a Jupyter notebook to the supplementary materials called “saturation_events.ipynb”. This notebook contains code which allows you to verify the existence of saturation events as described in the paper on Llama3 model.

Regarding the robustness of our findings, we have reproduced our results in different regimes and architecture types and even across modalities - pre-trained and untrained (randomly initialized) GPT2 which is decoder-only, encode-only pre-trained ViT, and encoder-decoder pre-trained Whisper (see Sections 3.2 and 3.3).

To address your concern of our results generalizing to SOTA models and different NLG tasks, we have reproduced our order saturation and task transition probing experiments using Llama3-8B model on MMLU and Hellaswag benchmarks. We show that Llama3’s top-k tokens reach saturation in order of their ranking up to (and including) the 4th ranking token, and that we’re able to predict the task number from layer embeddings with 88% accuracy. Please see Section B.2 in the appendix for more details.

We show that the phenomenon of ordered saturation is statistically significant using a stricter version of a well known ranking metric Kendall’s tau. If this is not sufficient, could you please suggest which other statistical rank tests we should run?

The goal of the intervention experiment is not to demonstrate a practical application for our findings, but rather to show causality and further validate our proposed task transition mechanism.

Regarding Table 2, thank you for bringing this to our attention. After running an independent samples t-test we find that the difference in accuracy between our method and each of the alternatives is statistically significant, with p < 0.001 and make note of it in Section 5.1.

In Section 5.1 we report the highest accuracy achieved across all tested thresholds for Softmax Response and State Saturation methods, which is calculated for the next-word prediction task as the percent of samples where the predicted top-1 token matches exactly the next token in the text. Please see Section A.7 in the Appendix for further details.

If there are specific models, benchmarks, or analyses you would like us to explore to strengthen the work further, please do not hesitate to let us know.

1 - Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. (https://arxiv.org/abs/2203.14680)

2 - Confident adaptive language modeling. (https://proceedings.neurips.cc/paper_files/paper/2022/hash/6fac9e316a4ae75ea244ddcef1982c71-Abstract-Conference.html)

3 - Language Models Implement Simple Word2Vec-style Vector Arithmetic (https://arxiv.org/abs/2305.16130)

I appreciate the authors' persistency and the code. However, I'd like to reiterate that my concerns are not addressed by the empirical/experiential findings of this submission. I'd like to highlight that my own observations of alternative layer-wise outputs (when probing LLMs with prompts in zero-shot and few-shot settings) also constitute empirical evidence. I ran the script (the logit lens) provided in the supplementary material and obtained the output. Contrary to the claimed "saturated events", the results show that neither the top 1 nor top K tokens maintain their positions across layers. In fact, even the top 1 token fails to hold its position, let alone a consistent picture for ordered top K saturated events (LLM: Meta-Llama-3-8B). This raises concerns about the robustness of the findings.

(Pdb) text = "What is the location of China's capital?"

(Pdb) logit_ret = logit_lens(tokenizer, model, text, top_k=5)

(Pdb) print_logit_lens(logit_ret)

Layer 1, top_5 logit lens is ['ToPoint', 'ĠKERNEL', 'amac', 'vell', 'Ġà¤ħà¤Ĺ']

Layer 2, top_5 logit lens is ['ĠPitt', 'Ġchin', 'Ġperms', 'agli', 'semb']

Layer 3, top_5 logit lens is ['Ġperms', 'ĠCarson', 'âĢĮاÙĨ', 'BOOLE', 'Ġciv']

Layer 4, top_5 logit lens is ['ĠPitt', 'Ġperms', 'pun', 'ONGL', 'olin']

Layer 5, top_5 logit lens is ['Ġchin', 'NECT', 'ĶåĽŀ', 'Ġredistribute', 'Äįet']

Layer 6, top_5 logit lens is ['#ab', '#ac', 'Sharper', 'TRGL', 'à¹Ħลà¸Ļ']

Layer 7, top_5 logit lens is ["'gc", 'akah', 'Äįet', 'Enlarge', 'Ġcue']

Layer 8, top_5 logit lens is ['#ab', '#ac', 'Ġdisp', "'gc", 'RetVal']

Layer 9, top_5 logit lens is ['#ab', 'RetVal', "'gc", 'λικ', 'eph']

Layer 10, top_5 logit lens is ['eph', '#ab', '\model', '675', 'Ĥ¨']

Layer 11, top_5 logit lens is ['ÑģÑĤе', '485', '#ab', 'óst', 'habit']

Layer 12, top_5 logit lens is ['#ab', '.Xaml', 'taboola', 'VISIBLE', 'RIPT']

Layer 13, top_5 logit lens is ['#ab', 'STYPE', '#ac', '#ad', 'ystate']

Layer 14, top_5 logit lens is ['STYPE', 'HCI', '#ab', '#ad', '#ac']

Layer 15, top_5 logit lens is ['#ad', '#ac', '#ab', 'modifiable', 'STYPE']

Layer 16, top_5 logit lens is ['Ġanswer', 'λικ', 'ĶåĽŀ', 'ĠHint', '#ab']

Layer 17, top_5 logit lens is ['Ġanswer', 'ĠAnswer', 'ĠHint', 'ModifiedDate', 'weit']

Layer 18, top_5 logit lens is ['Ġanswer', 'ĠAnswer', 'answer', 'ĠANSW', 'Ġanswers']

Layer 19, top_5 logit lens is ['Ġanswer', 'answer', 'ĠAnswer', 'Ġanswered', 'ĠANSW']

Layer 20, top_5 logit lens is ['Ġanswer', 'ĠAnswer', 'answer', 'Ġanswered', 'ĠANSW']

Layer 21, top_5 logit lens is ['Ġanswer', 'ĠAnswer', 'çŃĶæ¡Ī', 'answer', 'Ġanswered']

Layer 22, top_5 logit lens is ['Ġanswer', 'ĠAnswer', 'çŃĶæ¡Ī', 'answer', 'Ġanswered']

Layer 23, top_5 logit lens is ['Ġanswer', 'answer', 'ĠAnswer', 'çŃĶæ¡Ī', 'ĠWhich']

Layer 24, top_5 logit lens is ['Ġanswer', 'Ġwhere', 'Ġwhat', 'ĠWhere', 'Ġwhy']

Layer 25, top_5 logit lens is ['ĠBeijing', 'Ġanswer', 'Ġwhere', 'ĠWhere', 'Ġwhat']

Layer 26, top_5 logit lens is ['ĠBeijing', 'Ġwhere', 'Ġwhat', 'ĠWhere', 'ĠChina']

Layer 27, top_5 logit lens is ['ĠBeijing', 'Ġwhere', 'Ġwhat', 'ĠWhere', 'Ġhow']

Layer 28, top_5 logit lens is ['ĠWhere', 'Ġwhere', 'ĠBeijing', 'ĠWhich', 'Ġwhat']

Layer 29, top_5 logit lens is ['Ġwhere', 'Ġwhat', 'ĠBeijing', 'ĠWhere', 'where']

Layer 30, top_5 logit lens is ['ĠA', 'Ġwhere', 'ĠWhere', 'ĠBeijing', 'ÂłA']

Layer 31, top_5 logit lens is ['ĠWhere', 'ĠA', 'ĠBeijing', 'Ġwhere', 'ĠWhat']

Layer 32, top_5 logit lens is ['ĠBeijing', 'ĠWhat', 'ĠA', 'ĠThe', 'ĠWhere']

To verify the "ordered state saturation" phenomenon under more real-world conditions, I propose the following experiments:

1. In zero-shot settings, use paraphrased prompts (e.g., 5-fold for each prompt) to record changes in top N outputs, simulating real-world prompt variability.
2. In few-shot settings, experiment with appended or preceded contexts (including negative examples acting as interference) to examine the impact of in-context learning on the phenomenon.
3. When projecting hidden states into the vocabulary space, consider the potential suppression of meaning from other features (BTW "Polysemanticity" is a well-accepted concept in the community of mechanistic interpretability). To better understand the inner mechanism, I suggest exploring SAE-based analytic approaches to uncover potential footprints of "ordered state saturation" in neuron or feature circuits.

We deeply appreciate the reviewer's willingness to run our code and engage in further discussion. We apologize if we’ve made it seem like saturation events occur for all tokens in all prompts, and acknowledge that although they are prevalent in the data it may not be evident without analyzing a lot of examples.

We note that the reviewer's output was not produced by our code provided in the supplementary materials. We tried our best to remain as true as possible to the reviewer’s points (see code in next comment) and ran the same exact prompts.

We find that in the prompt “What is the location of China's capital?” the following token position predictions reach top-1 saturation at least 3 layers before output:

• Position 4 (prediction for next-word after ‘location’) - top-1 saturates at layer 24
• Position 6 (prediction for next-word after ‘China’) - top-1 saturates at layer 26.

And in the prompt “Can you tell me where the capital of China is?” the following token position predictions reach top-1 saturation at least 3 layers before output:

• Position 3 (prediction for next-word after ‘tell’) - top-1 saturates at layer 16
• Position 9 (prediction for next-word after ‘China’) - top-1 saturates at layer 27
• Position 11 (prediction for next-word after ‘?’) - top-1 saturates at layer 25

At position 3 we can even see some evidence for ordered saturation as the 2nd ranking token reaches saturation at layer 28, and the 4th token reaches saturation at layer 31 (3d token doesn’t undergo saturation). Considering that both prompts are very short (10 and 11 tokens, respectively), this is a significant percentage of tokens that undergo top-1 saturation.

Using the same prompts as the reviewer we got different results to their output, which clearly demonstrate the existence of saturation events (at the very least for the top-1 token). The disparity might be due to 3 factors:

1. Quantization - as we noted in the appendix and can be seen in our code we use 8-bit quantized version of Llama3 with the following parameters: BitsAndBytesConfig(load_in_8bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
2. LayerNorm - as is the standard when using logit lens, we apply layer norm to the outputs of each hidden layer before projecting onto the vocabulary space. This is done to ensure consistency with the last layer.
3. We look at the model predictions for all token positions in the prompt, and not just the last one.

We reiterate that we do not claim that all predictions undergo top-k saturation, just that enough undergo top-1 saturation for this to be a relevant phenomenon worth exploring. As can be seen by running our code on 100 questions randomly sampled from the Hellaswag benchmark, almost 47% of predictions reach top-1 saturation at least 3 layers before output. When looking at the next highest ranking tokens, although the numbers get lower they are still a significant fraction of the data (Over 11% of samples for the 2nd highest ranking token, almost 4% of samples for the 3d ranking token, and around 2% of samples for the 4th ranking token).

Regarding the experiments proposed by the reviewer:

1. To simulate real-world prompt variability we followed the reviewer’s suggestion by:

   • randomly sampling 200 questions from MMLU dataset,
   • automatically creating 5 paraphrases for each question with gpt3.5-turbo OpenAI model using the following prompt (resulting in 1K texts): "Paraphrase the following question into 5 different versions:\nQuestion: <Question>\nParaphrases:"
   • re-running our analysis to validate the existence of saturation events.

   We find that the overall saturation statistics remain almost the same with over 46% of samples reaching saturation at least 3 layers before the output for the top-1 token.

2. Can the reviewer explain what is meant here by negative examples? Additionally, as we’ve run our experiments on multiple text datasets (CNN/DM, MMLU and Hellaswag) we would argue that we’ve checked many different contexts, and even shown this phenomenon is content independent by reproducing it in a randomly initialized Transformer (for which any context is meaningless) and Transformers from other domains (vision and audio).

3. We first note that investigating the inner mechanism that produces the ordered saturation phenomenon implies its existence. In our work, we focus on establishing the generality of this newly discovered phenomenon and propose a high level task-transition mechanism for it, which we validate using both probing and intervention procedures, and even show how it can be leveraged for practical applications. We agree that given the establishment of our claims, a great direction for future work to explore is which feature circuits and specific neurons give rise to this phenomenon, but this is beyond the scope of our paper. We will expand on this topic in the discussion.

We're glad you found our paper clear and understandable, with a well-defined approach and practical contributions.

We do not claim that the ordered saturation phenomenon is unique to Transformers, which, as we mention in the “Conclusion and Future Work” section, is a question worth exploring, particularly in RNNs. We argue that it is inherent to the architecture and support our claim by reproducing the findings in an untrained randomly initialized Transformer model (Section 4.1).

Following your comments and to demonstrate that this phenomenon appears in recent, state-of-the-art models, we have reproduced our results using Llama3-8B on MMLU and Hellaswag benchmarks. We find that Llama3’s top-k tokens reach saturation in order of their ranking up to (and including) the fourth-ranking token and that we’re able to predict the task number from layer embeddings with 88% accuracy. Please see Figure 7 and Table 5 in Section B.2 of the appendix for more details.

You noted that the practical applications require further validation to enhance their real-world relevance. Could you please explain what type of validation would be helpful here?

Regarding your questions:

1. The difference in chance level for each model follows from the different number of tasks used when training the probing classifier for each model. The number of tasks is determined per model to be the maximum number for which, after balancing the data, there are at least 10 embeddings from each layer in each class from at least 4 different layers. To clarify this point, we’ve added the relevant footnote to Section 2.3.

2. In cases where two top-ranked tokens reached saturation as the same layer, they both get the same ranking which works against our claim of ordered saturation. Despite this, we see a statistically significant difference between the average ranks of adjacent top tokens.

3. When creating the data for the task number classifier, if two adjacent top tokens reached saturation at the same layer (for example 3d and 4th) then these would be labeled as belonging to the task of the smaller token between them (in this case 3d task). This adds some noise to the data, but does not prevent the classifier from achieving high accuracy.

4. Classifier accuracy was calculated as the percentage of samples where the predicted label matches exactly that of the true label.

5. Fixed typo.

6. We provided comparison with the only known token level early exit methods known to us based on the cited survey [1] . If you can suggest others, we would be happy to compare our strategy against them as well.

We will add clarifications to all of these questions in the paper.

If there are specific models, benchmarks, or analyses you would like us to explore to strengthen the work further, please do not hesitate to let us know.

[1] - A survey on efficient inference for large language models. (https://arxiv.org/abs/2404.14294)

Thank you to the authors for engaging in discussion. I appreciate the time that put in to address my concerns and feel they were sufficiently addressed, so I raised my score.

We are pleased that you appreciated our contributions to advancing the understanding of Transformer models. It is encouraging to see that our efforts to establish generalizability across modalities and provide empirical backing for the task-transition hypothesis were well received.

Although we provide a new early-exit strategy as a practical application for our findings, this is not the main focus of our paper. Furthermore, our proposed strategy focuses only on the top-1 as is the standard in most efficiency focused works, but as we argue in the Related Works section, this does not take into account the other top-k tokens which are often used in common generation sampling methods. For this reason, to fundamentally improve inference efficiency as you suggest, future work should consider the trade-off between speedup and top-k accuracy.

If there are specific models, benchmarks, or analyses you would like us to explore to strengthen the work further, please do not hesitate to let us know

Thank you for the opportunity to elaborate on this point. When we mention the trade-off between speedup and top-k accuracy, we are emphasizing the importance of considering how early-exit strategies affect more than just top-1 predictions. As discussed in Section 6 of our paper, although greedy decoding (which uses top-1 predictions) is considered suboptimal for many language generation tasks compared to stochastic sampling [1], it remains the standard approach in early-exit evaluations, including in recent works such as [2]. Even early-exit papers that use top-k or top-p sampling techniques during evaluation such as [3,4] treat them as a technical detail. Based on our work it is possible to establish which top-k tokens reach saturation and could be relevant predictions. Future research should explore these sampling strategies more thoroughly, considering them as key hyperparameters to strike a balance between efficiency and maintaining the quality of generated outputs.

We appreciate your thoughtful question, as it underscores an important implication of our work. We will incorporate this discussion into the conclusion of our paper to emphasize its relevance to the field.

[1] A Thorough Examination of Decoding Methods in the Era of LLMs - https://arxiv.org/abs/2402.06925

[2] LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding - https://arxiv.org/pdf/2404.16710

[3] SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference - https://arxiv.org/pdf/2307.02628

[4] ConsistentEE: A Consistent and Hardness-Guided Early Exiting Method for Accelerating Language Models Inference - https://ojs.aaai.org/index.php/AAAI/article/download/29922/31612

We are glad you found our writing and illustrations clear and noted the extensive experiments across different modalities.

Answers to points raised:

1. Although the classifiers for vision and speech models achieve lower accuracy than that of the classifier trained on the language model, the results are still statistically significant with p < 0.00001 based on a Binomial Distribution probability test. The difference in accuracy might be due to the different number of layers between the models, where GPT2 has 48 layers, ViT and Whisper have 32. The larger number of layers lends itself to more layers per task and better separation between them.

2. The number 7 was chosen arbitrarily as a hyperparameter; to demonstrate the robustness of the effect, we re-ran the experiment with values in the range 2-10 and indeed observed the same phenomenon, indicating that there’s nothing unique about the number 7: the accuracy of 2nd tokens which reach saturation i layers before output (with 2 <=i <=10) is significantly higher than of those that are determined only at the last layer. Please see Section B.3 in the Appendix for more details.

3. Thank you for bringing this to our attention. After running an independent samples t-test, we found that the difference in accuracy between our method and each of the alternatives is statistically significant, with p < 0.001 and made note of it in Table 2. We would like to emphasize that the practical applications are not the main focus of our paper, and are meant to showcase how our theoretical insights can be used.

You mentioned that the proposed method seems heuristic and hand-crafted, and that there is not enough supportive evidence. Can you please expand on what you mean by that? Which method used in the paper are you referring to and what additional supportive evidence could we provide?

If there are specific models, benchmarks, or analyses you would like us to explore to strengthen the work further, please do not hesitate to let us know.

We’re glad that you found our concept novel and see the potential of our proposed practical applications.

Per your suggestion we re-ran our order saturation and task transition probing experiments using Llama3-8B model on 1K randomly chosen samples from MMLU and Hellaswag benchmarks. We show that Llama3’s top-k tokens reach saturation in order of their ranking up to (and including) the 4th ranking token, and that we’re able to predict the task number from layer embeddings with 88% accuracy. Please see Section B.2 in the appendix for more details, And let us know if you feel that more samples or additional datasets are needed.

Regarding your questions:

1. We appreciate your suggestion, and acknowledge that in the NLP field the word “task” has other more downstream connotations. Yet we feel that the word “process” here would be misleading, and that “task” better encapsulates the discrete nature of the proposed mechanism, where the model determines each token in sequence.

2. We recognize that our practical applications sections apply only to language models, but we feel that the main focus of our paper is the discovery of the ordered saturation phenomenon and the proposed task transition mechanism, which as we argue are inherent to Transformer architecture regardless of modality. For this reason we think that focusing only on language models would be doing a disservice to the generality of our findings.

3. To address the question of which tokens undergo saturation we’ve added a section to the Appendix (B.1) where we show that in GPT2 saturation events are common for top-k tokens, with over 80% of samples reaching top-1 saturation. Additionally, these samples belong to different parts of speech (POS), and are not just function words for example, with over 27% of them being nouns. Our findings do not apply to samples which do not reach top-1 saturation, but our results in Section 5.2 suggest that these predictions night be less accurate in terms of language modeling.

If there are specific models, benchmarks, or analyses you would like us to explore to strengthen the work further, please do not hesitate to let us know.

In addition, we’ve reproduced our preliminary practical applications results showing that our insights can be leveraged for better efficiency and performance in LLMs using Llama3, Mistral and Falcon models (Sections B.2.1, B.4.1 and B.4.2 respectively in the Appendix) to further show the generality of our findings:

• We show that our novel early-exit strategy based on our task-transition classifier outperforms SOTA softmax-response and hidden-state saturation methods in Llama3 and Mistral, and outperforms softmax-response in Falcon in terms of the efficiency and accuracy tradeoff.

• We show that the saturation of the second highest ranking token affects the accuracy of it’s language modeling (when the top-1 token is incorrect), so that the difference in next-word prediction accuracy between 2nd token achieving saturation i layers before final layer (with 2 <= i <= 6) condition and the condition of the 2nd token being determined only in the last layer is statistically significant with p < 0.001 even when correcting for multiple comparisons.