Layer by Layer: Uncovering Hidden Representations in Language Models

Layer by Layer: Uncovering Hidden Representations in Language Models

Speakers: Oscar Skean, PhD Student, University of Kentucky; Md Rifat Arefin; Dan Zhao; Niket Patel; Jalal Naghiyev; Yann LeCun; Ravid Shwartz-Ziv

Conference: ICML 2025

YouTube: https://slideslive.com/39043900

Overview

In the rapidly evolving landscape of large language models (LLMs), a set of deeply ingrained assumptions often guides their application: that the final layers yield the most optimal embeddings for downstream tasks, or that intermediate layers are largely unhelpful. This talk, "Layer by Layer: Uncovering Hidden Representations in Language Models," delivered by Oscar Skean, a PhD student at the University of Kentucky, boldly challenges this conventional wisdom. Skean, presenting on behalf of a diverse, cross-institutional team including prominent researchers like Yann LeCun, unveils compelling empirical evidence demonstrating that intermediate layers frequently surpass the performance of final layers across a wide array of models, scales, and modalities.

The research not only empirically validates this phenomenon but also introduces a comprehensive theoretical toolkit of evaluation metrics designed to demystify the internal workings of these complex models. By exploring information-theoretic, augmentation invariance, and geometric properties of hidden representations, the team provides crucial insights into why these intermediate layers exhibit such strength. This work is profoundly significant for the broader ML community, offering immediate practical benefits—such as easy performance boosts and reduced inference times—and fostering a deeper scientific understanding of how information is processed and transformed within neural networks. It encourages practitioners to critically re-evaluate their approaches to leveraging pre-trained models, moving beyond default settings to unlock hidden potential.

Background

▶ Watch: Challenging common ML assumptions about model layers (0:00)

The pervasive use of large language models has led to the emergence of certain "folklore" within the machine learning community. A common practice involves taking an off-the-shelf pre-trained model and directly utilizing its final layer outputs as embeddings for various downstream tasks. This approach is often predicated on the assumption that as information flows through successive layers, it becomes progressively refined and abstracted, culminating in the most task-agnostic and robust representations at the network's deepest point. Consequently, intermediate layers are frequently overlooked, deemed either too raw or too specialized to be universally useful. This implicit bias is evident in how many benchmarking leaderboards report scores, almost exclusively focusing on the performance derived from the final layer.

However, this widely adopted methodology often leaves significant performance on the table. The problem stems from a lack of granular understanding of how information is encoded, transformed, and potentially compressed or expanded across different layers within a deep neural network. While a model's ultimate objective (e.g., next token prediction for autoregressive models) dictates the utility of its final layer for that specific task, it doesn't necessarily imply optimal general-purpose representations for a diverse set of other downstream applications. Prior work has largely focused on the model's output, with less emphasis on a systematic, layer-by-layer analysis of representation quality across varied tasks and model architectures. This gap has led to a suboptimal utilization of powerful pre-trained models, hindering practitioners from extracting maximum value without resorting to extensive fine-tuning or the acquisition of labeled data. The work presented by Skean and his co-authors directly addresses this fundamental challenge, aiming to demystify the internal dynamics and provide actionable insights for better model deployment.

Key Findings

▶ Watch: Introducing theoretical toolkit: information, invariance, and geometric metrics (3:00)

The central discovery of this research is a profound challenge to the prevailing assumption that the final layers of large language models consistently produce the best embeddings for downstream tasks. Through rigorous empirical testing across diverse model architectures, scales, tasks, and even modalities, the team unequivocally demonstrates that intermediate layers often significantly outperform final layers. Specifically, they found that middle layers can yield a substantial performance improvement, often by 5-10%, on various benchmarks.

To explain this phenomenon, the researchers introduced a novel theoretical toolkit comprising three families of evaluation metrics: information theoretic, augmentation invariance, and geometric metrics. Among these, prompt entropy emerged as a particularly insightful information-theoretic metric, serving as a central theoretical link between the different families. Prompt entropy quantifies the compression level of representations at any given layer.

A key finding related to prompt entropy is its distinct behavior across different model architectures:

• Autoregressive models (e.g., Pythia, Llama, Qwen) exhibit a characteristic "compression bottleneck." This is observed as a sharp dip in prompt entropy in the middle layers, indicating that the token embeddings become highly compressed and low-rank before expanding again towards the final layers. This bottleneck is interpreted as the model transitioning from a "token space" to an "abstract idea space" and then back to "token space" for output generation. Crucially, for autoregressive transformers, low prompt entropy in these intermediate layers is strongly and negatively correlated with high downstream task performance.
• This compression bottleneck is not an artifact of specific training stages but emerges rapidly during training (around 10,000 steps for Pythia, which trains for 143,000 steps) and deepens with increased model scale (from 14 million to 1 billion parameters in Pythia).
• Masked language models like BERT demonstrate a different pattern, with a much flatter prompt entropy profile across layers, suggesting less internal compression. This difference is attributed to BERT's bidirectional attention, which provides more flexibility compared to the causal attention constraints of autoregressive models.
• The phenomenon of strong intermediate performance and the associated compression bottleneck is not limited to language models but extends to other modalities, specifically autoregressive vision models like AIMv1. This suggests that autoregressive training itself, rather than the modality, is a key driver for these internal representation dynamics.

Practically, these findings translate into an easily achievable performance boost. For autoregressive transformers, simply identifying the layer with minimum prompt entropy and using its embeddings for downstream tasks can yield a 5-10% improvement without any additional training or labeled data. This offers a straightforward, no-cost method to enhance model performance and efficiency.

Technical Deep Dive

▶ Watch: Prompt entropy: measuring representation compression (effective rank) (4:00)

The methodology employed in this research involved a comprehensive, layer-wise probing of diverse language models and, to a lesser extent, vision models. The core idea was to treat the hidden states from every layer of a pre-trained model as potential embeddings for downstream tasks and evaluate their efficacy.

The empirical evaluation leveraged the Massive Text Embedding Benchmark from Hugging Face, a robust benchmark comprising hundreds of tasks. The researchers focused on 32 diverse tasks spanning five different task domains to ensure broad applicability of their findings. For each task, several models were selected, and the hidden states from every layer were used as input to a simple probe designed to solve the task. The primary objective was to identify which layer's representations yielded the optimal performance for a given task.

To move beyond mere empirical observation and understand why intermediate layers perform so well, the team developed a theoretical toolkit of evaluation metrics, categorized into three families:

1. Information Theoretic Metrics: These metrics aim to quantify how much information is preserved or compressed within the embeddings at different layers. The flagship metric in this family is prompt entropy.

• Prompt Entropy Calculation: Given a prompt of N tokens, the model generates an N x D matrix of token embeddings at layer K. From this matrix, the covariance matrix (Σ) is computed. Prompt entropy is then defined as the Shannon entropy of the eigenvalues of Σ, which is also known as the effective rank.
• Interpretation: A high prompt entropy indicates that token embeddings are high-rank, widely spread out, and diverse, suggesting a rich amount of information. Conversely, low prompt entropy signifies highly compressed, low-rank token representations.

1. Augmentation Invariance Metrics: Inspired by self-supervised learning techniques in computer vision (e.g., joint embedding SSL algorithms), these metrics assess the stability of representations when the input is subtly perturbed. The process involves performing a forward pass with a clean input, then augmenting the input (e.g., adding typos, random words), and observing how much the internal representations change. The goal is to understand how invariant the model's internal representations are to minor input variations. The talk noted that these metrics showed even more diverse trends across model families than prompt entropy.

1. Geometric Metrics: This family of metrics investigates the "shape" of the data manifold within the embedding space. Specifically, the researchers looked at curvature, which helps understand the sharpness of transitions between causal tokens (e.g., from token I to token I+1). Like augmentation invariance, geometric metrics revealed unique "fingerprints" for each model family's layer-wise behavior.

The analysis of prompt entropy across different model architectures revealed distinct internal behaviors:

• Autoregressive Transformers (e.g., Pythia, Llama, Qwen): These models exhibited a strikingly consistent pattern. The first few layers maintain a relatively stable, high prompt entropy. Subsequently, there's a sharp, significant dip in entropy, indicating a strong compression of representations. This is followed by a gradual increase in entropy towards the final layers. This "implicit bottleneck" suggests a phase where the model transforms high-dimensional token-level information into a more abstract, compressed conceptual space, before re-expanding it for the final token prediction task.
• Masked Language Models (e.g., BERT): In contrast, BERT-like models showed a much flatter prompt entropy profile across layers. This difference is attributed to their bidirectional attention mechanism, which allows them to leverage context from both past and future tokens. This bidirectional capacity reduces the pressure to compress information into an "abstract idea space" as intensely as autoregressive models, which are constrained by causal attention masks to only look backward.
• Autoregressive State Space Models (e.g., Mamba): While distinct from transformers, Mamba also exhibited a different prompt entropy trend, though the talk emphasized the clear autoregressive transformer pattern.

Further technical insights came from analyzing the training dynamics and scaling effects using Pythia, a model family known for releasing all its training checkpoints. The compression bottleneck was observed to emerge rapidly, solidifying around 10,000 training steps (out of 143,000 total steps) and remaining stable thereafter. Moreover, scaling up Pythia models from 14 million to 1 billion parameters consistently led to deeper and more profound compression bottlenecks, suggesting that this internal compression mechanism is an inherent property that intensifies with model capacity.

The critical theoretical link established is the negative correlation between low prompt entropy and high downstream task performance specifically in autoregressive transformers. This means that the layers where representations are most compressed (lowest prompt entropy) are often the ones that yield the best results for various downstream tasks, offering a straightforward heuristic for layer selection.

Experimental Setup & Results

▶ Watch: Emergence of the entropy bottleneck during model training (6:00)

The empirical backbone of this research relied on a systematic evaluation across a broad spectrum of models and tasks.

Datasets and Benchmarks:

The primary benchmark utilized was the Massive Text Embedding Benchmark from Hugging Face. This comprehensive benchmark contains hundreds of tasks, but for this study, the researchers focused on a curated subset of 32 diverse tasks spanning five different task domains. This selection ensured that the findings were generalizable and not specific to a narrow task type.

Models Tested:

The study evaluated models from various architectural families and scales:

• Autoregressive Transformers: Pythia (used as a representative, with models ranging from 14 million to 1 billion parameters), Llama, and Qwen.
• Masked Language Models: BERT.
• Autoregressive State Space Models: Mamba.
• Vision Models: AIMv1 (an autoregressive vision model).

Experimental Procedure:

For each of the 32 tasks and every model, the researchers adopted a "probing" approach. This involved:

1. Passing input prompts through the model.
2. Extracting the hidden states (embeddings) from every single layer of the model.
3. Using these layer-wise hidden states as input to a simple downstream probe (e.g., a linear classifier or regression model) designed to solve the specific task.
4. Measuring the performance of each layer's embeddings on the task to identify which layer produced the best results.

Headline Results:

The results were "astonishing," demonstrating that intermediate layers consistently outperformed the final layers across the tested autoregressive models. The typical scores reported on leaderboards, which use the final layer (depth 100 on their normalized axis), were significantly surpassed by embeddings from layers typically found in the "middle" (often between 20% and 70% of the total depth). For autoregressive models, this translated into an easy 5 to 10 percentage point performance improvement on the Massive Text Embedding Benchmark tasks, simply by selecting the optimal intermediate layer.

Ablation Studies and Key Observations:

• Prompt Entropy Correlation: A crucial finding was the strong negative correlation between prompt entropy and downstream task performance in autoregressive transformer models. Layers exhibiting the lowest prompt entropy (highest compression) were found to yield the best task performance. This provides a direct, measurable heuristic for identifying optimal layers.
• Training Dynamics: Analyzing Pythia's training checkpoints revealed that the "compression bottleneck" (the sharp dip in prompt entropy) emerged rapidly during training, stabilizing early in the training process (around 10,000 steps for Pythia, which trains for 143,000 steps). This indicates it's an intrinsic property developed quickly.
• Scaling Effects: As Pythia models scaled up from 14 million to 1 billion parameters, the compression bottleneck became deeper and more profound, suggesting that larger models learn to compress information more effectively in their intermediate layers.
• Cross-Modality Validation: The findings were not limited to language. The team extended their investigation to the vision domain, specifically examining AIMv1, an autoregressive vision model. They observed the same pattern of superior intermediate accuracy and a corresponding compression bottleneck in terms of prompt entropy. This critical result suggests that autoregressive training, rather than the input modality itself, is a key factor driving these internal representation dynamics.
• BERT's Distinct Behavior: In contrast to autoregressive models, BERT (a masked language model) exhibited a much flatter prompt entropy profile, with less pronounced compression. This difference was attributed to BERT's bidirectional attention, which provides greater flexibility in processing information compared to the causal attention constraints of autoregressive models.

In summary, the experimental results provided robust empirical evidence for the superior performance of intermediate layers in autoregressive models, both language and vision, and offered a quantifiable metric (prompt entropy) to explain and predict this behavior.

Practical Implications

▶ Watch: Model scaling increases representation compression across layers (6:40)

The findings presented in "Layer by Layer: Uncovering Hidden Representations in Language Models" carry profound practical implications for anyone working with pre-trained AI/ML models, from individual practitioners to large infrastructure teams. They challenge established norms and offer concrete strategies for immediate performance enhancements and resource optimization.

1. Easy Performance Boost Without Additional Training:

Perhaps the most immediate and impactful implication is the ability to achieve a 5-10% performance improvement on downstream tasks without any fine-tuning, labeled data, or additional computational training. For autoregressive transformers specifically, practitioners can simply:

• Access the hidden states from all layers of a Hugging Face model (typically a one-line code change: output_hidden_states=True).
• Empirically test layers, focusing on depths between 20% and 70% of the total model layers.
• Alternatively, calculate prompt entropy for each layer and select the layer with the minimum entropy, as this is strongly correlated with optimal downstream performance in autoregressive models.

This provides a low-effort, high-reward strategy to enhance existing deployments or prototypes.

2. Reduced Inference Time and Resource Optimization:

By identifying that optimal performance often resides in intermediate layers, practitioners can significantly reduce inference costs and latency. If only a subset of layers (e.g., the first 50 out of 100 layers) is required to achieve superior embeddings, then:

• Less memory is needed to load the model, as only the necessary layers need to be resident.
• Fewer computations are performed during the forward pass, leading to faster inference times.
• This enables the use of larger batch sizes, further improving throughput and efficiency.

These benefits are critical for deploying LLMs in resource-constrained environments or high-throughput applications where every millisecond and byte counts.

3. Enhanced Understanding of Internal Model Behavior:

The introduced toolkit of evaluation metrics—prompt entropy, augmentation invariance, and geometric metrics—provides powerful lenses through which to examine and understand the black box of deep neural networks.

• Model Builders and Researchers can use these metrics to gain insights into how different architectures (e.g., autoregressive vs. masked) and pretext tasks influence the internal representation learning.
• Understanding phenomena like the "compression bottleneck" helps in designing more effective training regimes or even novel architectures. It sheds light on how models transform raw input into abstract concepts and back, offering a more nuanced view than just observing input-output pairs.

4. Guiding Model Improvement and Fine-tuning:

The utility of these metrics extends beyond mere evaluation. The talk mentioned follow-up work where prompt entropy metrics have already been used to improve math reasoning capabilities on benchmarks like GSM8K. This suggests that:

• These metrics can serve as diagnostic tools during training or fine-tuning, helping identify layers that are underperforming or exhibiting undesirable representation characteristics.
• They could potentially be incorporated into loss functions or used for layer pruning/selection strategies to optimize models for specific downstream tasks more effectively.

Tradeoffs and Limitations:

While the benefits are substantial, it's crucial to acknowledge certain tradeoffs and limitations:

• Architecture Dependence: The strong negative correlation between low prompt entropy and high performance is most pronounced and consistent in autoregressive transformers. For masked language models (like BERT) or state-space models (like Mamba), the prompt entropy profile is flatter, and empirical testing across layers is still recommended, as the direct heuristic might not apply as strongly.
• Task Specificity: While the Massive Text Embedding Benchmark covers diverse tasks, the optimal layer might still vary slightly depending on the specific downstream task. Therefore, a quick empirical check remains a robust strategy.
• "One Line of Code" Simplicity: While accessing hidden states is simple, implementing the prompt entropy calculation might require a few more lines of code. However, given the potential performance uplift, this is a minor overhead.

In essence, this research empowers practitioners to move beyond default settings, encouraging a more informed and efficient use of powerful pre-trained models. It's a call to "check your middle layers" and unlock the hidden potential within.

Key Takeaways

• Challenging ML Folklore: Intermediate layers of large language models, particularly autoregressive transformers, often significantly outperform final layers for various downstream tasks, contrary to common assumptions.
• The "Compression Bottleneck": Autoregressive models (e.g., Pythia, Llama) exhibit a distinct "compression bottleneck" in their middle layers, characterized by a sharp dip in prompt entropy. This indicates a phase where information is highly compressed into abstract representations.
• Prompt Entropy as a Predictor: For autoregressive transformers, there's a strong negative correlation between low prompt entropy (high compression) and high downstream task performance, offering a direct heuristic for identifying optimal layers.
• Cross-Modality Phenomenon: This behavior of superior intermediate performance and the associated compression bottleneck is not limited to language models but also observed in autoregressive vision models (e.g., AIMv1), suggesting it's a fundamental property of autoregressive training.
• Easy Performance & Efficiency Gains: Practitioners can achieve a 5-10% performance boost by simply selecting optimal intermediate layers (empirically or via minimum prompt entropy for autoregressive models), leading to reduced inference time, lower memory usage, and larger batch sizes without additional training.
• Enhanced Model Understanding: The introduced toolkit of information theoretic, augmentation invariance, and geometric metrics provides valuable tools for researchers and developers to gain deeper insights into the internal workings and representation learning dynamics of diverse model architectures.

About the Speaker(s)

The work "Layer by Layer: Uncovering Hidden Representations in Language Models" was presented by Oscar Skean, a PhD student at the University of Kentucky. Skean served as the lead presenter, articulating the motivation, methodology, and key findings of this extensive research. The project itself was a highly collaborative, cross-institutional effort, involving numerous co-authors spread across the globe who contributed through remote collaboration channels like Discord and Zoom. The distinguished list of co-authors includes Md Rifat Arefin, Dan Zhao, Niket Patel, Jalal Naghiyev, Yann LeCun, and Ravid Shwartz-Ziv, highlighting the broad expertise and collaborative spirit underpinning this significant contribution to the field of machine learning.