# Unraveling Token Prediction Refinement and Identifying Essential Layers in Language Models

arXiv: https://arxiv.org/abs/2501.15054 · HTML: https://mlsanity.com/paper/2501.15054/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 44 references that are in the map.
- First posted on arXiv 25 Jan 2025.

## Builds on (most cited first, 24 of 24)

- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [Sparse Autoencoders Find Highly Interpretable Features in Language Models](https://mlsanity.com/paper/2309.08600/) (2023)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [Representation Engineering: A Top-Down Approach to AI Transparency](https://mlsanity.com/paper/2310.01405/) (2023)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Discovering Latent Knowledge in Language Models Without Supervision](https://mlsanity.com/paper/2212.03827/) (2022)
- [Designing and Interpreting Probes with Control Tasks](https://mlsanity.com/paper/1909.03368/) (2019)
- [Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task](https://mlsanity.com/paper/2210.13382/) (2022)
- [Mechanistic Interpretability for AI Safety -- A Review](https://mlsanity.com/paper/2404.14082/) (2024)
- [A Survey of Machine Unlearning](https://mlsanity.com/paper/2209.02299/) (2022)
- [A Survey on Neural Network Interpretability](https://mlsanity.com/paper/2012.14261/) (2020)
- [AI Alignment: A Comprehensive Survey](https://mlsanity.com/paper/2310.19852/) (2023)
- [An Overview of Catastrophic AI Risks](https://mlsanity.com/paper/2306.12001/) (2023)
- [Localizing Model Behavior with Path Patching](https://mlsanity.com/paper/2304.05969/) (2023)
- [AI Deception: A Survey of Examples, Risks, and Potential Solutions](https://mlsanity.com/paper/2308.14752/) (2023)
- [Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit](https://mlsanity.com/paper/2207.08799/) (2022)
- [Still No Lie Detector for Language Models: Probing Empirical and Conceptual Roadblocks](https://mlsanity.com/paper/2307.00175/) (2023)
- [Efficient Long-Text Understanding with Short-Text Models](https://mlsanity.com/paper/2208.00748/) (2022)
- [X-Risk Analysis for AI Research](https://mlsanity.com/paper/2206.05862/) (2022)
- [An overview of 11 proposals for building safe advanced AI](https://mlsanity.com/paper/2012.07532/) (2020)
- [Circumventing interpretability: How to defeat mind-readers](https://mlsanity.com/paper/2212.11415/) (2022)

## Built on by (most cited first, 3 of 3)

- [LogitLens4LLMs: Extending Logit Lens Analysis to Modern Large Language Models](https://mlsanity.com/paper/2503.11667/) (2025)
- [Layer-wise Positional Bias in Short-Context Language Modeling](https://mlsanity.com/paper/2601.04098/) (2026)
- [Off-Axis, On Purpose: Where a Transformer Computes Concepts and Why it Does So](https://mlsanity.com/paper/2608.10251/) (2026)

## Abstract (arXiv)

This research aims to unravel how large language models (LLMs) iteratively refine token predictions through internal processing. We utilized a logit lens technique to analyze the model's token predictions derived from intermediate representations. Specifically, we focused on (1) how LLMs access and utilize information from input contexts, and (2) how positioning of relevant information affects the model's token prediction refinement process. On a multi-document question answering task with varying input context lengths, we found that the depth of prediction refinement (defined as the number of intermediate layers an LLM uses to transition from an initial correct token prediction to its final, stable correct output), as a function of the position of relevant information, exhibits an approximately inverted U-shaped curve. We also found that the gap between these two layers, on average, diminishes when relevant information is positioned at the beginning or end of the input context. This suggested that the model requires more refinements when processing longer contexts with relevant information situated in the middle. Furthermore, our findings indicate that not all layers are equally essential for determining final correct outputs. Our analysis provides insights into how token predictions are distributed across different conditions, and establishes important connections to existing hypotheses and previous findings in AI safety research and development.