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

Authors: Daria Lioubashevski, Tomer M. Schlank, Gabriel Stanovsky, Ariel Goldstein
Primary area: interpretability and explainable AI
Keywords: mechanistic interpretability, transformer, large language model, efficient inference

Abstract:
Understanding the inner workings of Transformers is crucial for achieving more
accurate and efficient predictions. In this work, we analyze the computation performed
by Transformers in the layers after the top-1 prediction has become fixed, which has been
previously referred to as the “saturation event”. We expand the concept of saturation events
for top-k tokens, demonstrating that similar saturation events occur across language, vision, 
and speech models. We find that these saturation events happen in order of the 
corresponding tokens’ ranking, i.e., the model first decides on the top ranking token, then 
the second highest ranking token, and so on. This phenomenon seems intrinsic to the 
Transformer architecture, occurring across different architectural variants (decoder-only, 
encoder-only, and to a lesser extent full-Transformer), and even in untrained Transformers.
We propose an underlying mechanism of task transition for this sequential saturation, where 
task k corresponds to predicting the k-th most probable token, and the saturation events are 
in fact discrete transitions between the tasks. In support of this we show that it is possible to 
predict the current task from hidden layer embedding. Furthermore, using an intervention 
method we demonstrate that we can cause the model to switch from one task to the next. 
Finally, leveraging our findings, we introduce a novel token-level early-exit strategy, which 
surpasses existing methods in balancing performance and efficiency.

Decision: reject
Reviewer scores: 8, 5, 8, 8, 3

Reviews and discussion:
<thread>
creation date:
1730272521365

thread_type:
official_review

<thread_object>
creator: reviewer_u1oc

<content>
summary:
The paper explores hidden states of intermediate layers of transformer-based models by mapping them onto the vocabulary space to extract ranking over tokens and analyzing their changes over the layers. The authors claim to identify sequential saturation events for top-k tokens, a phenomenon observed across language, vision, and speech models. They propose a type of discrete transition occurs during decoding.

strengths:
Strengths:
The paper emphasizes the importance of investigating layer-wise states transition. The study tackles an important and interesting research question.

weaknesses:
The hypothesis of “saturation events” is taken as granted without sufficient evidence. The mechanism in producing top-1 tokens in LLMs is influenced by various factors, including LLM architecture, contexts, and how an LLM reacts to an input. Based on practice shared by many colleagues in probing LLMs, transformers often alternate their top-1 token predictions, even in later layers, which undermines the claimed findings by this study. 

The scope of the experiment is limited in terms of both the number and types of LLMs and tasks. This raises concerns about the generality of the findings. It is important to investigate both NLU and NLG tasks on a diverse range of LLMs under different regimes (e.g., pre-trained, post-trained, decoder-only, encoder-decoder) to ensure comprehensive insights. Instead of limiting to a small-scale LLM like GPT2-XL, the authors may consider exploring state-of-the-art (SOTA) LLMs, such as the LLAMA series, to broaden the scope of their research and enhance the robustness of the conclusions.

The use of ranks and central tendency in Figure 3 is confusing (i.e., why averaging ranks and std in figure 3). There are existing statistic rank tests the authors may leverage. The claimed sequential saturation pattern should be convincingly demonstrated.

The practical application of intervention is unclear, and the results shown in table 2 are not statistically significant. When you choose "Highest accuracy", how do you calculate it? Is it based on EM?

questions:
as above
</content>
</thread_object>

</thread>
<thread>
creation date:
1730596603109

thread_type:
official_review

<thread_object>
creator: reviewer_iroj

<content>
summary:
1) This study extends the concept of saturation from only the top-1 token to multiple top-ranked tokens, providing deeper insights into the computational processes within Transformer models.
2) Experiments across diverse domains, including text, vision, and speech, demonstrate the broad applicability of this sequential saturation mechanism.
3) Beyond observing this phenomenon, the authors propose an early-exit algorithm based on their findings, presenting a practical approach to enhance computational efficiency in Transformer models.

strengths:
The paper is written in a clear and understandable manner, with a well-defined approach and valuable findings that go beyond observation to demonstrate practical benefits. I believe this is an good research, and with a few additional points, it would be well-suited for acceptance at the conference.

weaknesses:
If the paper claims that these findings are unique to Transformers (as the title suggests), it should demonstrate that (1) this phenomenon does not occur in other architectures and (2) it consistently appears in recent, state-of-the-art models. Providing this evidence would make the paper much more logically sound and robust.

The practical applicability of the experiments in Sections 5.1 and 5.2 feels somewhat limited; further validation could enhance their real-world relevance. (For example, as shown in A.5, softmax and state cannot outperform ours even at the minimum speedup, which makes the comparison less compelling.)

questions:
(1) Table 1 shows different chance levels for each model (e.g., 25% for Whisper-large vs. 20% for GPT2-XL and ViT-L/16) Could the authors clarify the reason for these differences?

(2) In Figure 3, how did the authors handle cases where two top-ranked tokens reach saturation at the same layer?

(3) Following up, how were such cases handled during classifier training?

(4) In Table 1, how was accuracy calculated in these scenarios?

(5) typo: 400lines: l1(s1) < l2(s2) => l1(s1) < l1(s2)

(6) Could the authors provide a comparison with other early exit methods?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730598853969

thread_type:
official_review

<thread_object>
creator: reviewer_8pe6

<content>
summary:
This paper investigates the internal mechanisms of Transformers, particularly focusing on the computation performed after the top-1 prediction has been determined, a phenomenon referred to as the "saturation event." The authors expand the concept of saturation events to top-k tokens and demonstrate that these events occur in a ranked order across various modalities (language, vision, speech) and Transformer architectures. They propose a task-transition mechanism to explain this sequential saturation, where each task corresponds to predicting the k-th most probable token. The paper further shows that it is possible to predict the current task from hidden layer embeddings and to induce task transitions through interventions. Finally, the authors introduce a token-level early-exit strategy that improves performance and efficiency, outperforming existing methods.

strengths:
1. The paper provides a fresh perspective on how Transformers process predictions beyond the top-1 token, which is a significant contribution to the field of model interpretability.
2. The study's scope extends across multiple modalities, enhancing the generalizability of the findings and demonstrating the robustness of the proposed mechanisms.
3. The authors back their claims with extensive experiments and provide empirical evidence supporting the task-transition hypothesis.
4. The paper not only has theoretic analysis but also explores practical applications, such as the early-exit strategy, which has potential implications for improving model efficiency and accuracy.

weaknesses:
This is a suggestion rather than a weakness. 

It would be great to discuss whether this is a fundamental solution for early exit problem. After all, it is ridiculous for us to let problems with all level of difficulties to go through all the layers. There must be a solution that significantly improve the inference efficiency of the current LLM through sparsity or early exit. Yet the solution proposed by this paper give only marginal improvement.

questions:
Please see the weakness part. Highly appreciate if the author give some feedback.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703912038

thread_type:
official_review

<thread_object>
creator: reviewer_w56e

<content>
summary:
This paper is to understand the inner workings of transformers to achieve more accurate and efficient predictions. Authors expand the concept of saturation events for top-k tokens, demonstrating that similar saturation events occur across language, vision, and speech models. The experiments show that it is possible to predict the current task from hidden layer embedding. Furthermore, authors use an intervention method to cause the model to swithc from one task to the next.

strengths:
1. This paper is well written with clear illustrations.
2. There are extensive experiments across different modalities, such as vision, language, and speech.
3. Compared with two baseline models, the proposed method demonstrates improved performance.

weaknesses:
1. Compared with Line 275 and L 276, the proposed method seems to be model-specific. While the results in Table 1 show the layer embeddings contain information about the task number, such information may only work for language, but not vision and speech.

2. in Line 472, why the second token's saturation layer is at least "7" layers before the output? The hypothesis seems to be hand-crafted without any supportive evidence.

3. In Table 2, the comparisons of different strategies are a little bit weak to support the claim. Is 40% accuracy a significant improvement compared to the baselines 35.9% and 37.5%? As these numbers are pretty low to have any real applications.

questions:
See the above weakness. The main concern is that the proposed method is quite heuristic and hand-crafted without enough supportive evidence. While the authors demonstrate the results across different modalities, it is difficult to evaluate the effectiveness of this method given the relatively low numbers.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730749565977

thread_type:
official_review

<thread_object>
creator: reviewer_nrzw

<content>
summary:
This work expands the top-1 saturation event to top-k tokens and finds a significant phenomenon that there is an order law between the saturation events and the corresponding tokens' ranking. Additionally, such phenomenon is rooted in transformer architecture, which happens across all modality and even in untrained models. With the finding of such order law on the saturation events, an efficient decoding method is proposed for early-exit for next-token prediction.

strengths:
1. The concept and research problem in this work is very novel. It significantly extends the range of saturation event to the top-k token level. 

2. The findings are very brute-force, which can be easily used for early-exit decoding for next token prediction acceleration without introducing much additional computation. And such decoding strategy is general and easily applied across all transformer based LLMs.

weaknesses:
1. The evaluation of this method is limited in its size and targeted models. 60k tokens and 100 texts are a too small evaluation set size for a robust conclusion. GPT2-XL is also not well-trained and the experiments on Llama-3 is more recommended. If the saturation event for top-k tokens only happens at very late layers, then the acceleration ratio is not that ideal compared with GPT2-XL experiments. CNN summarization task is a very basic task for recent LLMs, while the benchmarks of MMLU, Hellaswag are more recommended for study and experiments.

questions:
1. I recommend to change the definition of "task" to other words like "process". The concept of "task transition" makes me feel like there is some generality in different tasks like summarization, QA, dialogues, etc.

2. The experiments on vision and speech transformers are great demonstrations to your claim on the generality of your findings but have no contribution to real applications like decoding acceleration. Maybe you can consider to remove them into appendix and enlarge your study on language model part.

3. I am very interested in whether such saturation event happens on all tokens in the context or just a few tokens? If it does not happen on all tokens, how you can deal with the prediction towards these non-saturation tokens?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731664787510

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank our reviewers for their insightful feedback and helpful suggestions.
 
In light of their comments, we’ve run additional experiments, which are presented in Appendix B:

1. To provide evidence establishing how common saturation events are in the data, we’ve added Section B.1 showing that in GPT2 model  over 80% of samples reach top-1 saturation. 

2. We have reproduced our results using Llama3-8B on MMLU and Hellaswag benchmarks to demonstrate the ordered saturation phenomenon in recent state-of-the-art models. 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 can 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.

3. In Section 5.2, the number 7 was chosen randomly as an 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 (see Section B.3).

We’ve also added some clarifications and statistical significance markers in the paper itself, color-coded in cyan for your convenience. 

Finally, we uploaded a Jupyter notebook to supplementary materials called “saturation_events.ipynb”, which allows the reviewers to verify the existence and frequency of the saturation events discussed in the paper on Llama3 model.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733312914021

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

Thank you for your insightful feedback and thoughtful engagement during the review process. 

To further support your evaluation, we have created an anonymous Colab notebook that allows you to verify the ordered saturation phenomenon of top-k tokens discussed in our paper. This notebook enables replication of results on the GPT-2 XL model with the MMLU dataset and is optimized to run on the free version of Colab using a T4 machine.

You can access the notebook here:
https://colab.research.google.com/drive/1qzhLcfeScbkTu99sq0SRWnlgKX9bBQ6h#scrollTo=Dcj0f03mLIiI

We hope this resource enhances your review experience and are grateful for the time and effort you’ve invested in reviewing our work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper receives the mixed ratings of 8, 5, 3, 8, 8. In this paper, the authors introduce a concept of "sequential saturation" in Transformers, analyzing how top-k tokens are fixed layer-by-layer beyond the top-1 prediction. This paper provides evidence supporting that saturation of top-ranking tokens occurs sequentially in various Transformer architectures and modalities. Furthermore, a task-transition mechanism is introduced to explain this phenomenon, validated through interventions and probing experiments.

Strengths:
- The idea of sequential saturation is intriguing, expanding current understanding of Transformer behavior.
- Extensive experiments across text, vision, and speech modalities strengthen the claim of generalizability.
- The proposed applications demonstrate potential practical utility in improving Transformer efficiency.

Area of improvements:
- The concept of task transitions, while innovative, lacks a clear connection to practical modeling scenarios. The assumption that each token represents a discrete task oversimplifies real-world complexities.
- Evidence from untrained models, while theoretically insightful, weakens applicability to real-world, pre-trained settings.
- Reviewer u1oC noted significant discrepancies when attempting to reproduce key results, raising concerns about the robustness of the experiments. The reliance on specific datasets and hyperparameters was not sufficiently detailed in the main paper, making reproducibility challenging.
- While statistical metrics like Kendall’s tau and probing accuracies are provided, the lack of clear visual/conceptual explanations for task transitions weaken the interpretability of the claims.
- Early-exit strategy, demonstrate only incremental gains over baseline methods.

The paper explores an interesting and promising direction, offering an innovative perspective on Transformer dynamics. However, concerns regarding the model's realism and the reproducibility of results persist despite the rebuttal and discussion. We encourage the authors to address these issues comprehensively, as doing so will significantly strengthen the work for a future resubmission.