Unraveling Token Prediction Refinement and Identifying Essential Layers in Language Models

TL;DR

This paper investigates how large language models refine token predictions internally, revealing the influence of input position on prediction refinement depth and identifying which layers are crucial for accurate outputs.

Contribution

It introduces a logit lens analysis to understand token prediction refinement and identifies essential layers, highlighting the impact of input positioning on model behavior.

Findings

Prediction refinement depth varies with input position, forming an inverted U-shape.

Relevant information at input edges reduces the need for multiple refinements.

Not all layers are equally important for final token prediction accuracy.

Abstract

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.