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

We appreciate the reviewer's helpful feedback.

Methods And Evaluation Criteria

A simpler and more direct option could be reporting the percentage of inputs where the order of the saturation layers matches the order of the tokens.

We calculated this percentage for Llama3-8B on 1K MMLU questions for top-k tokens, where top-1 refers to cases where the 1st tokens reach saturation in order, top-2 includes the 1st and 2nd tokens, and so on:

These results further validate the ordered saturation phenomenon, and we will include this analysis in the camera-ready version.

it’s unclear whether the data used in the experiments was part of the pre-training set.

To mitigate the concerns of data pollution, we reproduced our findings on datasets which came out after the pretraining, and hence less likely to be contaminated. For text, in Appendix A.4, we report results for GPT2-XL model on the MMLU dataset—a benchmark introduced after the model’s release.

For vision, we reproduced our results using ViT on 1K images from ImageNet-D[1], a newer synthetic benchmark. The average rank of the k-th saturation layer increases monotonically (up to k=3), with statistically significant differences between each two consecutive token ranks (p < 0.001, average ranks: 1.0, 2.0, 2.39).

For speech, we reproduced our results using Whisper on 200 English audios from the newer Emilia [2] benchmark. Here we also found the average k-th saturation layer rank to increase with k (up to k=3), with the differences being statistically significant (average saturation ranks in order: 1.01, 1.94, 2.42). Furthermore, since we demonstrate this order phenomenon in a randomly initialized Transformer, our findings suggest robustness to both training and dataset choice.

For the early-exit strategy .. it may be more appropriate to use a metric that assesses the model's alignment with the language distribution rather than just the top token prediction/accuracy.

Following your suggestion, we calculated the perplexity metric for each early-exit strategy using probabilities from the early-exit layer. For Llama3-8B and 500 CNN/DM texts, our method significantly outperforms Softmax Response in speedup-performance tradeoff and performs similarly to State Saturation, both of which are considered SOTA. See results in table below:

Weakness

If this implies that the architecture plays a major role in unlocking this phenomenon (line 289), then it would be useful to verify this observation across other architectures.

In this work we are the first to notice this phenomenon, and focus our work on developing measures of quantifying its extent. We agree that other architectures can also lead to ordered saturation. We will refine this claim in the paper to make it clear that this is an hypothesis raised by our findings.

Questions:

1. We extract embeddings for training only when saturation layers are in order, as our goal is to explain this mechanism.
2. When evaluating our early-exit strategy we also used cases where this order doesn’t hold, i.e., we exit when the first token saturates in all cases. This demonstrates the robustness of our probing classifier and method.
3. Chance level varies as it is inversely proportional to the number of tasks/classes, which is set to ensure data balancing. We will clarify this in the table caption.
4. i. Saturation events occur fairly often for other ranks. In LLaMA3-8B on 1,000 MMLU questions: top-1 saturates in 65.7% of tokens (77K), top-2 in 28.3% (33K), top-3 in 15.2% (18K), top-4 in 9.9% (12K), and top-5 in 7.2% (8.5K). ii. The strict Kendall’s coefficient values reported in the paper, suggesting moderate agreement between token rank and saturation order, are calculated only on the subset of samples where top-1 reaches saturation.
5. As noted in Section 5.1, dynamic decoding is not our focus, so we propagate states from layers after early-exit as in regular inference. This allows later tokens to use these computations if needed, and the effect of a state-copying mechanism on performance is negligible [3].
6. Yes, in LLaMA3-8B on 500 MMLU questions, the first saturation layer ranges from 2 to 30. Analyzing its correlation with token position, we found a weak but significant negative correlation (r = -0.1, p < 0.001), suggesting later tokens tend to saturate earlier. We will add this to the camera-ready version.

We welcome further discussion if needed. If our clarifications are satisfactory, we kindly ask the reviewer to consider raising their rating.

[1] Zhang et al. 2024

[2] He et al. 2024

[3] Schuster et al. 2022

Thank you for your thoughtful and positive assessment of our work! We appreciate your recognition of the novelty of the ordered saturation phenomenon and its implications for understanding Transformer models. Your acknowledgment of both the theoretical contributions and practical applications, such as our early exit strategy, is encouraging.

Other Strengths And Weaknesses:

However, the applicability of complex tasks/models is still unclear, such as some of the more advanced LLMs, Qwen Audio series, LLaVA series, and it is not clear whether they are universal for these complex tasks/models

We agree that exploring the applicability of our findings to more complex tasks and advanced models is an exciting and important direction. Following your suggestion we replicated our results using LLaVA-1.5-7B model on 1K questions from MMLU dataset treating it as a language model. We found the average rank of the k-th saturation layer for each k increases monotonically with k (up to k=5) and the difference between each two consecutive token ranks is statistically significant with p < 0.001 based on a pairwise independent samples t-test (average saturation ranks in order: 1.01, 2.14, 2.66, 2.87, 3.00). We also trained a classifier to predict task number from hidden layer representations, and achieved an 82.4% average accuracy over 5 classes (chance level accuracy being 20%).

We then re-run the analysis on 500 questions from MMMU dataset [1], utilizing this model’s multi-modal capabilities on the more complex task of Visual Question Answering, where the input contains both image and text. We found the average rank of the k-th saturation layer for each k increases monotonically with (k up to k=5), and the difference between each two consecutive token ranks is statistically significant with p < 0.001 based on a pairwise independent samples t-test (average saturation ranks in order: 1.01, 2.38, 2.91 3.14, 3.33). We also trained a classifier to predict task number from hidden layer representations, and achieved an 82.4% average accuracy over 5 classes (chance level accuracy being 20%).

Additionally, we ran Qwen Audio model on 500 audio samples from LibriSpeech dataset and found that the average rank of the k-th saturation layer for each k increases monotonically with k (up to k=4) and the difference between each two consecutive token ranks is statistically significant with p < 0.001 based on a pairwise independent samples t-test (average saturation ranks in order: 1.01, 1.85, 2.73, 2.93). We also trained a classifier to predict task number from hidden layer representation and achieved an 85.5 average accuracy over 4 classes.

These results further emphasize the robustness of the ordered saturation phenomenon over models, tasks and modalities.

Thank you again for taking the time to review the paper and providing helpful feedback! Do the above actions address your concerns with the paper? If not, what further clarification or modifications could we make to improve your score?

---

[1] Yue, Xiang, et al. "Mmmu: A massive multi-discipline multimodal understanding and reasoning benchmark for expert agi." IEEE/CVF 2024.

Thank you for your insightful comments and helpful suggestions. We appreciate that you found the sequential determination of top-k tokens to be an interesting and well-supported phenomenon. We're also glad that you recognized the broader relevance of our findings in relation to existing literature on token determination in Transformers.

Claims And Evidence:

I believe the obvious order of token determination doesn't necessarily result the task transition. There is possibly another possible modeling: if the weight of all top tokens is increasing together but with different scale, we can also observe sequential saturation of top tokens starting from random weights.

We agree that the sequential saturation of top tokens could emerge from a different underlying mechanism, such as the one suggested—where all top token probabilities increase together but with different scales due to residual connections. This is an interesting alternative perspective, which we will add to the paper.

To clarify, our proposed interpretation is not that the first task corresponds strictly to processing only the first token, the second to the second, and so on. Rather, what we can say with confidence is that by the time the model transitions to the second task, it is no longer focusing on the first token in the same way. This distinction is supported by the probing classifier results, which show high accuracy in distinguishing the tasks as we define them, indicating that these transitions are not simply a gradual rescaling but rather discrete shifts in model behavior.

Additionally, the intervention experiments provide further support for our interpretation. Specifically, when we inject layer activations from later stages, the model abruptly shifts its behavior, ceasing to focus on the top-1 token and instead progressing to the next stage. This causal evidence suggests that the observed transitions are more than just a continuous adjustment of weights. That said, we agree that further investigation is needed to rigorously rule out alternative explanations, and we appreciate the reviewer’s suggestion to refine the claim for greater clarity.

Given the opportunity we will incorporate those points into the discussion section in the camera-ready version.

Other Comments Or Suggestions:

We will follow your suggestion to expand on previous work on extracting predictions from intermediate layer in the introduction section in the camera ready version of our paper.

Questions For Authors:

I think this phenomenon will be quite related to the existence of residual connection, because this results that every layer will have minor modification to the previous layer, i.e., for x_{n+1} = x_n + f(x_n), actually many works show that |x_n| >> |f(x_n)|, I think this may be related?

You raise a very interesting question about the relevance of residual connections to the ordered saturation phenomenon. We address this indirectly in our Related Work section when considering the iterative inference hypothesis. This approach interprets each layer as an iteration from an iterative and convergent process [1], suggesting that each layer incrementally refines the hidden representation by gradually shaping the next token prediction [2-4].

At first glance, the existence of residual connections should make it more difficult to distinguish layer-wise task transitions, as they enforce strong similarity between consecutive layers. However, our results show that saturation events still occur in a structured and discrete manner, rather than as a continuous drift. This suggests a phase transition-like dynamic: while each layer makes small refinements, at a certain point, a critical threshold is reached, leading to a sharp shift in representation and task focus.

More work is needed to determine which components in the Transformer architecture give rise to this phenomenon, and the residual connections is certainly one to examine alongside self-attention and MLP layers. One possible way to do so is through various ablation studies.

Thank you again for taking the time to review the paper and providing helpful feedback! Do the above actions address your concerns with the paper? If not, what further clarification or modifications could we make to improve your score?

---

[1] Simoulin, Antoine, and Benoit Crabbé. "How many layers and why? An analysis of the model depth in transformers."

[2] Geva, Mor, et al. "Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space."

[3] Belrose, Nora, et al. "Eliciting latent predictions from transformers with the tuned lens."

[4] Rushing, Cody, and Neel Nanda. "Explorations of self-repair in language models."

Thank you for your thorough review.

Claims And Evidence

I don't believe results provided in Figure 3 .. are convincing enough beyond the 3rd top token.

Beyond the top-3 tokens, saturation order is indeed less consistent in some models. However, our results are statistically significant across models and modalities.This variability may be partly due to model depth—fewer layers may limit lower-ranked tokens from reaching saturation. For instance, GPT-2 XL (48 layers) maintains order up to the 5th token, unlike Llama3-8B (32 layers) (Appendix A.4). We will add this observation to the paper.

With respect to the results on classifier accuracy… I'm afraid the same can't be said about the Image and speech transformers.

While the classifier accuracy in Vision and Speech Transformers is lower than in text models, it remains statistically significant (Table 1) and over twice of chance level. Though this may limit practical applications, this experiment validates the task transition mechanism.

I believe the model's accuracy on specific k-top layer is probably more informative than the overall accuracy.

Table 8 shows no significant accuracy differences across tasks, with ViT even performing slightly better for Task 5 than Task 1. The classifier was trained on balanced data representing all tasks. We didn't investigate these differences, but agree this is an interesting avenue for future research.

With regards to the layer activation injection, I have some questions on what the activation layer is.

We inject the first saturation layer’s output from sample s1 into sample s2’s subsequent layer and observe its effect on s2’s first saturation layer. As a control, we repeat this with activations from three layers before and after. Please see Appendix A.8 for details. We will clarify this in the camera-ready version.

However, just extending the results to 2-tokens might not be significant enough for the paper.

Following this comment we analyzed the 3rd and 4th top tokens using Llama3-8B on 500 CNN/DM texts, assuming higher-ranked tokens were incorrect, and found the results generalize.

For the 3rd top token, accuracy increased from 17.01% (no saturation) to 23.3%, 26.9%, 28.9%, 29.2%, and 31.0% when saturation occurred 2 to 6 layers before output, respectively.

For the 4th top token, accuracy increased from 13.4% (no saturation) to 16.8%, 19.7%, 20.45%, 19.5%, and 18.3% when saturation occurred 2 to 6 layers before output.

A two-proportion z-test confirms statistical significance (p < 0.001) for 2 ≤ i ≤ 6 in the 3rd token and 2 ≤ i ≤ 5 in the 4th token. We will include this analysis in the camera-ready version.

Questions

Q1: We report in our response to Reviewer J6pR that other Vision and Speech models achieve high accuracy, comparable to text models. In Whisper, we speculate that order deteriorates in later tokens because all decoder layers are conditioned on the encoder’s last layer, potentially blurring task boundaries. In contrast, in Qwen-Audio, where the encoder only generates input features for the decoder, our task classifier achieves high accuracy.

Q2: Non-linear probing could hurt our claim, as powerful probes may learn arbitrary mappings rather than reveal the model’s inherent representations [1, 2].

Q3: We computed standard error as std/sqrt(n) [3]; large n results in small values .

Q4: Omitted values indicate that the probing classifier was trained on a different number of tasks per model. For Llama3-8B for example, it was trained on three tasks, so accuracies for Tasks 4 and 5 are not reported. We will clarify this in the caption.

Q5: We did not explicitly test task classifier transferability across models, but we indirectly demonstrate cross-dataset transferability by applying a classifier trained on MMLU to CNN/DM in our early-exit experiment (Section 5.1).

Q6: In our analysis of the ordered saturation of top-k tokens, we set k = 5 to ensure consistency across models while maintaining sufficient examples (>4% of inputs for all models). As k increases, saturation cases decline. The 5% value was a typo. We will clarify this and add detailed statistics in the appendix.

Q7: As noted in Appendix A.3, stricter Kendall’s tau p-value was computed via a permutation test of saturation layer sequences (1K random shuffles per instance). In all cases, real tau values exceeded shuffled results (p < 0.001). We will further clarify this in the paper.

Q8: Table 3 analyzes the second token’s accuracy when the top-1 token is incorrect, depending on its saturation layer (2 ≤ i ≤ 6 layers before output). Variability in saturation layer is an open question, but our main goal was to show a statistically significant difference between saturation and no-saturation cases.

We welcome further discussion if needed. If our clarifications are satisfactory, we kindly ask the reviewer to consider raising their rating.

---

[1] Hewitt et al., 2019

[2] Belinkov, 2022

[3] Casella, G., & Berger, R. L. 2002