Research map / Architectures and efficiency / Attention and transformers
Looking Beyond the Top-1: Transformers Determine Top Tokens in Order
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that top-k tokens saturate in order of their rank across GPT2-XL, ViT-L/16 and Whisper-large (even untrained GPT2-XL), explains this as a discrete task-transition mechanism, and uses it for a token-level early-exit strategy.
Problem
Prior work found top-1 saturation events in Transformer layers, but what computation happens after the top-1 token has saturated, and whether lower-ranked top-k tokens also saturate in a structured order, was unexplored.
Approach
Extends saturation events to the k-th ranked token and finds that these saturation layers occur in ranking order; interprets this as discrete task transitions (task k = determining the k-th token), detectable by a linear probe and causally controllable by activation injection.
Key results
- Task-number logistic regression reaches 91.4% accuracy on pretrained GPT2-XL versus 20.6% with random embeddings (chance 20%).
- Token-level early exit reaches 40.0% accuracy at 1.185 speedup, versus 35.9% (softmax response, 1.126) and 37.5% (state saturation, 1.003).
- For incorrect top-1 predictions, the second-ranked token has 31.99% accuracy when it saturates at least 7 layers before output versus 17.14% when it only resolves in the final layer.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use task-transition classification from layer embeddings instead of random embeddings | GPT2-XL (pretrained) · CNN/DM | task-number classification accuracy: 20.6 ± 0.5 -> 91.4 ± 0.3 | Helped |
| Use task-transition classification from layer embeddings instead of random embeddings | GPT2-XL (random initialization) · CNN/DM | task-number classification accuracy: 32.7 ± 0.1 -> 86.1 ± 0.7 | Helped |
| Use task-transition classification from layer embeddings instead of random embeddings | ViT-L/16 · CIFAR-10 | task-number classification accuracy: 21.0 ± 0.5 -> 63.8 ± 0.1 | Helped |
| Use task-transition classification from layer embeddings instead of random embeddings | Whisper-large · LibriSpeech | task-number classification accuracy: 24.5 ± 0.4 -> 52.7 ± 0.1 | Helped |
| Exit based on the task-transition classifier rather than softmax response or state saturation | GPT2-XL · CNN/DM next-word prediction | accuracy: 35.9 ± 0.6; 37.5± 0.7 -> 40.0 ± 0.7 | Helped |
| Exit based on the task-transition classifier rather than softmax response or state saturation | GPT2-XL · CNN/DM next-word prediction | speedup ratio: 1.126 ± 0.004; 1.003 ± 0.001 -> 1.185 ± 0.009 | Helped |
| Use the second-ranked token after it reaches saturation rather than when it is determined only in the last layer | GPT2-XL · CNN/DM; top-1 prediction incorrect | second-ranked token accuracy: 17.14% -> 31.99% | Helped |
| Order top-token saturation by token rank | GPT2-XL (pretrained) · CNN/DM | strict Kendall's tau: 0.187 ± 0.004; p < 0.001 | Helped |
| Order top-token saturation by token rank | GPT2-XL (random initialization) · CNN/DM | strict Kendall's tau: 0.082 ± 0.009; < 0.001 | Helped |
| Order top-token saturation by token rank | ViT-L/16 · CIFAR-10 | strict Kendall's tau: 0.149 ± 0.007; < 0.001 | Helped |
| Order top-token saturation by token rank | Whisper-large · LibriSpeech | strict Kendall's tau: 0.210 ± 0.009; < 0.001 | Helped |
Surprising findings
- Untrained randomly initialized GPT2-XL also saturates top-k tokens in order (up to the 4th token), with task classifier accuracy 86.1%.
- Ordering is weaker for Whisper (only up to the third token) and for untrained models (lower stricter Kendall's tau).
Limitations
- Unclear which architectural components cause the ordering; no ablation studies.
- Did not consider whether evaluation data was seen during training.
- Only Transformer architectures were studied; RNNs untested.
Experimental setup
Datasets: CNN/DM CIFAR-10 LibriSpeech ImageNet-21k ImageNet 2012
Benchmarks: CNN/DM CIFAR-10 LibriSpeech
Metrics: average rank of saturation layer stricter Kendall's tau classifier accuracy ROC-AUC accuracy speedup ratio two-proportion z-test
Baselines: Softmax Response early exit Hidden-state saturation early exit
Builds on: logit lens top-1 saturation event tuned lens early exiting with softmax response and hidden-state saturation
Closest papers by meaning
- Transformers need glasses! Information over-squashing in language tasks (NeurIPS 2024 · similarity 0.89)
- Talking Heads: Understanding Inter-Layer Communication in Transformer Language Models (NeurIPS 2024 · similarity 0.89)
- Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models (NeurIPS 2025 · similarity 0.89)
- Dynamical Properties of Tokens in Self-Attention and Effects of Positional Encoding (NeurIPS 2025 · similarity 0.88)
- Transformer Block Coupling and its Correlation with Generalization in LLMs (ICLR 2025 · similarity 0.88)
- AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers (ICML 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.