Layer by Layer: Uncovering Hidden Representations in Language Models

TL;DR

The work challenges the standard emphasis on final-layer embeddings by showing that intermediate layers often deliver stronger representations across diverse architectures and tasks. It introduces a unified framework grounded in matrix-based entropy, geometry, and augmentation invariance to diagnose hidden-layer quality and links these metrics to downstream performance. Extensive experiments across transformers, state-space models, and vision tasks reveal mid-depth bottlenecks and robust correlations between representation metrics and task accuracy, with implications for layer selection, model design, and fine-tuning. The findings suggest practical strategies for exploiting mid-layer representations to improve robustness and efficiency while offering a theory-driven lens for cross-architecture analysis.

Abstract

From extracting features to generating text, the outputs of large language models (LLMs) typically rely on the final layers, following the conventional wisdom that earlier layers capture only low-level cues. However, our analysis shows that intermediate layers can encode even richer representations, often improving performance on a range of downstream tasks. To explain and quantify these hidden-layer properties, we propose a unified framework of representation quality metrics based on information theory, geometry, and invariance to input perturbations. Our framework highlights how each layer balances information compression and signal preservation, revealing why mid-depth embeddings can exceed the last layer's performance. Through extensive experiments on 32 text-embedding tasks across various architectures (transformers, state-space models) and domains (language, vision), we demonstrate that intermediate layers consistently provide stronger features, challenging the standard view on final-layer embeddings and opening new directions on using mid-layer representations for more robust and accurate representations.

Figures (15)

• Figure 1: Intermediate layers consistently outperform final layers on downstream tasks. The average score of 32 MTEB tasks using the outputs of every model layer as embeddings for three different model architectures. The x-axis is the depth percentage of the layer, rather than the layer number which varies across models.
• Figure 2: Pythia and Mamba's intermediate layers show pronounced changes in representation quality metrics, while BERT’s remain more stable. Three representation evaluation metrics calculated on the wikitext dataset for every layer in Pythia-410M, Mamba 370M, and BERT-base architectures. The x-axis denotes layer depth as a percentage, allowing fair comparison between models with different layer counts.
• Figure 3: Relationship between representation metrics and task performance averaged across layers for Pythia 410M. Using distance correlation (dCor), we see strong associative relationships across the board with DiME exhibiting the strongest relationship with downstream performance. We use dCor due to its robustness and ability to measure both linear and non-linear relationships (dCor $\in [0,1]$ with 0 indicating statistical independence and 1 indicating strong dependency). We defer additional results to the Appendix.
• Figure 4: Strong trends in intermediate behavior emerge during training Representation evaluation metrics across layers at various Pythia-410M training checkpoints, ranging from step 1 to the final step at 143k. The x-axis is the model layer, showing how training affects different layers, while the colors are different checkpoints during training.
• Figure 5: Token-level prompt entropy across sequence lengths for Qwen 2.5 and Qwen 2.5-Math. The base model (Qwen 2.5) exhibits greater prompt compression, while the finetuned (Qwen 2.5-Math) has higher entropy, indicating more information retention.

Theorems & Definitions (16)

• Theorem 1: Lower Bound via Effective Rank
• Theorem 2: Informal
• Theorem 3: Dataset Entropy Bounds InfoNCE
• Definition 1
• Definition 2
• Lemma 1
• Theorem 4
• proof
• Proposition 1
• proof