How Do LLMs Use Their Depth?

TL;DR

This work uncovers how LLMs utilize depth during inference by proposing a Guess-then-Refine framework and employing TunedLens to trace intermediate representations across multiple open-weight models. It demonstrates that early layers produce high-frequency token guesses, which are heavily refined in later layers as context accrues, revealing a dynamic, complexity-aware depth usage. Through POS analysis, multi-token fact recall, and downstream tasks, the study shows that easier subtasks are handled shallower while harder ones demand deeper processing, challenging assumptions about fixed-depth generation. The findings have implications for modeling efficiency and provide a principled lens to understand transformer computations, with TunedLens validated as a faithful probe for early-layer representations.

Abstract

Growing evidence suggests that large language models do not use their depth uniformly, yet we still lack a fine-grained understanding of their layer-wise prediction dynamics. In this paper, we trace the intermediate representations of several open-weight models during inference and reveal a structured and nuanced use of depth. Specifically, we propose a "Guess-then-Refine" framework that explains how LLMs internally structure their computations to make predictions. We first show that the top-ranked predictions in early LLM layers are composed primarily of high-frequency tokens, which act as statistical guesses proposed by the model early on due to the lack of appropriate contextual information. As contextual information develops deeper into the model, these initial guesses get refined into contextually appropriate tokens. Even high-frequency token predictions from early layers get refined >70% of the time, indicating that correct token prediction is not "one-and-done". We then go beyond frequency-based prediction to examine the dynamic usage of layer depth across three case studies. (i) Part-of-speech analysis shows that function words are, on average, the earliest to be predicted correctly. (ii) Fact recall task analysis shows that, in a multi-token answer, the first token requires more computational depth than the rest. (iii) Multiple-choice task analysis shows that the model identifies the format of the response within the first half of the layers, but finalizes its response only toward the end. Together, our results provide a detailed view of depth usage in LLMs, shedding light on the layer-by-layer computations that underlie successful predictions and providing insights for future works to improve computational efficiency in transformer-based models.

Figures (14)

• Figure 1: Illustrative depictions of our Guess-then-Refine framework and complexity-aware depth use in LLMs. (a) Early layers elevate high-frequency tokens as statistical guesses, which deeper layers refine into context-appropriate predictions (Frequency-Conditioned Onset). (b) Multi-token fact recall: the first token of the answer requires the largest depth on average, while subsequent tokens appear at much shallower depths. (c) Option-constrained downstream tasks (e.g., True/False): early layers collect valid options into the top ranks, while later layers reason between top choices to generate final answer. Overall: LLMs act as early statistical guessers and late contextual integrators, exhibiting complexity-driven depth use.
• Figure 2: Frequent tokens dominate early-layer predictions. At each layer, the layer’s top-1 ranked token is grouped by corpus-frequency buckets (Top1–10, Top11–100, Top101–1000, Top1000+). Early layers are dominated by Top1–10 bucket tokens, where LLMs make statistical guesses and propose high-frequency tokens as top-ranked candidates under limited information, while deeper layers increasingly replace them with rarer, context-appropriate tokens.
• Figure 3: Early-layer predictions are heavily revised. For each layer and frequency bucket, we show the fraction of top-1 ranked predictions that are overturned by the final layer (“decision flip rate”). We see that about 80% of layer-1 Top1–10 predictions change by the end, and this rises to 1̃00% for Top101–1000 and Top1000+ bucket tokens. As expected, flip rates drop with depth.
• Figure 4: Earliest crossing thresholds of predicted tokens by POS (top) and for multi-token facts (bottom). We determine which token the model predicts at the final layer and then track the earliest layer (x axis) at which the TunedLens rank of that token crosses a given threshold (y axis). (a, b) Case Study I: When predicted tokens are grouped by POS category, tokens corresponding to function words (DET, ADP) and punctuation appear much earlier than content words (ADJ, VERB, NOUN). (c, d) Case Study II: We plot tokenwise results for fact recall responses that consist of 1 token, 2 tokens, or 3 tokens. E.g., 3-token facts consist of 3 tokens at positions (pos) 1, 2, and 3. In both models, the first tokens of multi-token facts appear much later than subsequent tokens.
• Figure 5: Mean rank of answer tokens during constrained-choice downstream tasks. As the model performs tasks like MMLU (response options: ABCD) and sentiment analysis ( response options: positive/negative), we track the mean rank of the options through the layers. Early layers of the model promote valid options choices as top-ranked candidates, whereas later layers reason between option choices (shown as zoomed-in figures in the top-right corner of each plot). This figur shows results for Pythia-6.9B. Results for other models and tasks follow similar patterns and can be found in Appendix \ref{['appendix:downstream']}.