Source-linked AI summary
Layer by Layer: Uncovering Hidden Representations in Language Models
Oscar Skean, Md Rifat Arefin, Dan Zhao, Niket Patel, Jalal Naghiyev, Yann LeCun, Ravid Shwartz-Ziv
TL;DR
The paper asks whether final-layer embeddings are always the best representations for downstream tasks. It proposes a unified framework based on information theory, geometry, and invariance, and finds that intermediate layers often outperform final layers across architectures and modalities. The results motivate greater use of mid-layer representations while leaving the causes of intermediate-layer compression for future study.
Problem
The paper examines whether the conventional reliance on final-layer representations overlooks stronger intermediate-layer embeddings for downstream tasks.
Method
The paper combines information-theoretic, geometric, and invariance metrics in a unified framework for analyzing hidden representations across layers.
Results
Intermediate layers often surpass final layers in representation quality across 32 MTEB tasks, with improvements of up to 16% in downstream accuracy.
Takeaways & Limitations
Intermediate layers are potentially underutilized feature sources, and their quality can be analyzed through a common lens of compression, geometry, and invariance.
Takeaways & Limitations
The underlying causes of intermediate-layer compression and explicit finetuning to control it remain topics for future work.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper questions the assumption that final-layer representations are always best for downstream tasks. It introduces a unified analysis showing that intermediate layers can provide stronger representations across architectures, objectives, and domains.
- Key findings and contributions: Intermediate layers often surpass final layers by up to 16% in downstream accuracy across 32 MTEB embedding tasks.The evaluation spans transformers, state-space models, and encoder-based models from tens of millions to billions of parameters.
- A unified framework: The framework combines information-theoretic, geometric, and invariance perspectives to assess how layers compress information and preserve signal.These perspectives examine semantic information, embedding geometry, and robustness to input perturbations.
- Key findings and contributions: Intermediate layers consistently outperform final layers in both transformers and state-space models.This suggests the pattern extends beyond a single architecture family.
- Key findings and contributions: Autoregressive models show a pronounced mid-layer compression valley, while masked or bidirectional models exhibit milder intermediate changes.The contrast links representation dynamics to the training objective.
- Key findings and contributions: Autoregressive image transformers display the same mid-depth bottleneck as language models, indicating a training-objective effect across modalities.The result extends the observed phenomenon from language to vision.
- Key findings and contributions: CoT finetuning reshapes mid-layer entropy and preserves latent context for multi-step reasoning.The analysis identifies a representation change associated with chain-of-thought finetuning.
- Practical outlook: The results challenge default reliance on final-layer embeddings and identify intermediate layers as underutilized sources of meaningful features.The paper presents a unified framework and experiments in language and vision to analyze this pattern.
2. Related Work
Prior work studies representation dynamics through probing, similarity analysis, compression, geometry, and invariance metrics. This paper connects these strands while evaluating representations across architectures and training paradigms.
- Understanding Neural Representations: Earlier studies used linear probes and SVCCA to analyze how neural networks encode and organize information.Much of this work focused on vision models or shallow networks.
- Layer-wise Analysis in Language Models: Middle layers of BERT often encode linguistic features such as part-of-speech tags and semantic roles, while later work finds robust mid-depth features.These findings challenge exclusive emphasis on final-layer representations.
- Compression and Generalization: Research links compression and generalization, including evidence that discarding layers can improve downstream accuracy and that LLM embeddings often lie on low-dimensional manifolds.These results motivate examining compression within intermediate layers.
- Representation Quality Metrics: Representation-quality metrics fall into information-theoretic, geometric, and task-based or invariance categories.The categories cover information retention, high-dimensional structure, and robustness to perturbations or downstream goals.
- Representation Quality Metrics: The paper unifies seemingly distinct metric categories under a single lens to explain why some intermediate layers produce better downstream representations.Its analysis spans architectures and training paradigms rather than treating metrics independently.
3. A Unified Framework for Neural Representations
The framework treats representation quality as a unified problem involving compression, robustness, and the structure of hidden states. Matrix-based entropy provides the central theoretical lens connecting these properties.
- Unified framework: Matrix-based entropy unifies compression, geometry, and invariance as facets of representation quality.The framework uses this lens to analyze how representations distribute and preserve information.
- Motivation: Existing approaches assess internal representations through mutual information, geometric manifold structure, and invariance to augmentations.The paper consolidates these previously distinct ideas into one theoretical framework.
- Notation and motivation: The framework represents inputs x as hidden states Z, organized as a matrix of N samples or tokens in D dimensions.This notation provides the basis for analyzing internal representations mathematically.
- Motivation: The analysis asks how compressed representations are and how robust they remain to perturbations or augmentations.These questions define two central dimensions of representation quality in the framework.
3. How do they geometrically organize different inputs?
The framework organizes representations through matrix-based entropy, connecting compression, geometry, and augmentation invariance via the Gram matrix. These metrics reveal how layers preserve distinctions, compress redundancy, and structure local versus global information.
- Compression and diversity: A few dominant eigenvalues indicate low-rank, compressed representations, whereas a more uniform spectrum indicates higher-entropy, more diverse features.Entropy is small when only a few eigenvalues dominate and large when representations spread across principal directions.
- Local and global organization: Curvature distinguishes local token transitions from longer-range global structure: sharp turns correspond to skewed spectra, while smoother trajectories suggest global-level features.Higher curvature reflects abrupt directional shifts between consecutive token embeddings.
- Unified representation geometry: Matrix-based entropy applied to K = ZZ⊤ provides a unified lens for compression, geometric structure, and invariance.The framework relates all seven evaluation metrics to the Gram matrix.
- Levels of organization: Prompt entropy measures token-level diversity, while dataset entropy measures how distinctly embeddings separate different inputs.Higher prompt entropy indicates less redundant token features; dataset entropy aggregates mean token embeddings across prompts.
- Entropy and dimensionality: Effective rank lower-bounds exp(S1(Z)), linking representation dimensionality to entropy and showing how strong compression shrinks effective dimensionality.The paper connects this relationship to RankMe and LiDAR.
- Invariance under perturbations: Augmentation-invariance metrics assess whether perturbed versions of an input cluster together in embedding space, with lower InfoNCE loss indicating stronger invariance.LiDAR measures within-class versus between-class scatter, while DiME compares correct augmentation pairs with random pairings.
- Entropy and downstream structure: High token-level diversity encourages broader dataset-level diversity, whereas over-compression can limit separation among prompts.The paper connects these scaling behaviors to InfoNCE, higher entropy, robustness, and discriminability.
4. Empirical Results
Across architectures and tasks, intermediate layers often outperform final layers, with mid-depth representations balancing information retention and compression. Representation metrics track downstream utility and reveal architecture-, scale-, sub-layer-, and training-dependent differences.
- Downstream Task Performance: Intermediate layers consistently outperform final layers across transformers and state-space models on 32 MTEB embedding tasks.The evaluation spans classification, clustering, and reranking tasks.
- Downstream Task Performance: 2% to 16%: intermediate-layer improvements over final-layer performance across nearly every task, often peaking around network mid-depth.The pattern is reported across the evaluated architectures and benchmarks.
- Representation Quality: Intermediate-layer sweet spots balance sufficient information retention against low-level noise compression.The paper ties these sweet spots to how intermediate layers process information rather than treating them as random locations.
- Representation Quality: DiME, curvature, and InfoNCE show particularly high associations with downstream performance, with robustness across distance-correlation measures.The analysis uses distance correlations and reports consistent relationships across additional correlation measures.
- Architectural Differences: Autoregressive models show pronounced mid-layer compression valleys, while masked or bidirectional models change more mildly across depth.Pythia exhibits a strong entropy dip, BERT maintains higher entropy, and Mamba shows flatter curves.
- Scale Effects: Larger Pythia models display stronger intermediate compression, smoother token trajectories, and higher invariance.The reported scale effects are associated with more effective feature distillation and noise filtering.
- Transformer Sub-Components: Residual connections drive the observed midnetwork compression, whereas preceding attention and MLP sub-layers generally show only mild compression.Residual sub-layers also exhibit a pronounced entropy drop.
5. Extreme Input Conditions
Extreme inputs alter layer-wise prompt entropy in structured ways: repetition compresses intermediate representations, randomness inflates early-layer entropy, and longer prompts raise unnormalized entropy.
- Input Perturbations: Increasing token repetition sharply decreases prompt entropy in mid-depth layers.This suggests repetitive patterns are recognized and redundancy is discarded internally.
- Input Perturbations: Random tokens significantly increase entropy in early layers, while deeper layers are more robust.The result indicates greater early-layer sensitivity to input noise.
- Input Perturbations: Across perturbations, intermediate layers selectively compress repetitive patterns while retaining crucial distinctions.The findings characterize layer-wise information preservation and discarding under unusual inputs.
6. Comparison to Vision Transformers
Vision experiments show that autoregressive image transformers reproduce the mid-depth bottleneck seen in language models, whereas non-autoregressive vision models generally improve toward their final layers.
- Cross-Domain Comparison: Extreme-input analyses in Pythia show that repetition lowers intermediate-layer entropy while randomness raises entropy especially in early layers.These perturbation responses provide a language-model comparison for interpreting layer-wise representation changes.
- Autoregressive Vision Models: AIM shows an entropy valley and accuracy peak at intermediate layers, matching the behavior observed in autoregressive language models.AIM predicts image patches sequentially.
- Non-Autoregressive Vision Models: All evaluated vision models except AIM show downstream accuracy increasing toward the final layer.The passage contrasts these models with the intermediate-layer peak found in AIM.
- Training Objective: Autoregressive training, rather than text-versus-image modality, is identified as the factor associated with mid-layer information bottlenecks.The comparison covers supervised, self-supervised, and autoregressive vision approaches evaluated layer by layer on ImageNet-1k.
7. Discussion and Conclusion
The paper presents intermediate layers as important sources of representation quality across language and vision, while identifying further causal control of compression as future work.
- Conclusion: Intermediate layers often surpass final layers in representation quality across language and vision models.The conclusion frames this as an implication for feature relevance and extraction.
- Conclusion: DiME, curvature, and InfoNCE correlate well with downstream performance, linking representation metrics with generalizability.The conclusion reports this as a central empirical finding.
- Future Work: Future work should investigate the causes of intermediate-layer compression and explicitly finetune models to control it.This is the stated future-work boundary of the paper.
Impact Statement
Intermediate-layer representations may improve downstream performance and affect interpretability, robustness, efficiency, and fairness considerations. The paper analyzes these representations using entropy-based and related metrics across model architectures.
- Intermediate-layer representations can yield better performance across varied downstream tasks, with implications for interpretability, robustness, and efficiency.
- Leveraging intermediate representations may affect fairness and bias considerations in model evaluation and deployment.The paper warns that identifying latent features could amplify latent biases or contribute to unintended disparities.
- Prompt entropy measures the diversity of intermediate token representations.
- Matrix-based entropy uses a similarity kernel on sampled representations as a tractable surrogate for Rényi’s α-order entropy.The kernel can be Gaussian, linear, or Laplacian; this work restricts the analysis to the linear kernel.
- Low entropy corresponds to a heavy-tailed spectrum in which a few principal components dominate variance, whereas maximum entropy corresponds to uniformly distributed variance.
- Varying α recovers special cases including Shannon entropy as α approaches 1 and collision entropy at α = 2.The α = 2 case can be computed without explicit eigendecomposition, and Figure 7 examines α under a β-power-law eigenvalue spectrum.
C.1. Wikitext Dataset
The experiments use Wikitext prompts and 32 English MTEB tasks, while augmentation-invariance metrics are computed from paired prompt augmentations. The augmentations include word splitting, character noise, and keyboard-neighbor substitutions.
- The Wikitext dataset contains 100 million tokens from featured Wikipedia articles, filtered to prompts of at least 30 tokens excluding section headings.
- The benchmark comprises 32 English MTEB tasks spanning clustering, classification, reranking, and sentence-to-sentence evaluation.
- Augmentation-invariance metrics including InfoNCE, LiDAR, and DiME use NLPAug to generate paired prompt augmentations.The two augmented prompts, rather than one augmented prompt and the original, are used for these metrics; character-level changes may alter token length.
- SplitAug randomly divides words by inserting spaces, while RandomCharAug inserts, substitutes, swaps, or deletes characters.
- Keyboard augmentation replaces characters with characters one QWERTY-key distance away.For example, k may be replaced with i, l, m, or j.
E. Using Evaluation Metrics as a Performance Proxy
The paper uses unsupervised evaluation metrics as proxies for selecting high-performing layers without task-specific labels. It computes metrics across layers in one forward pass and reports a 3% average MTEB improvement from DiME-based selection on Pythia-410M.
- Unsupervised layer selection exploits correlations between evaluation metrics and downstream performance to choose high-performing layers without task-specific labels.
- For Pythia-410M and LLM2Vec-8B, DiME, InfoNCE, and dataset entropy are computed for every task across layers in a single forward pass.
- The method selects, for each task, the layer minimizing one metric because these metrics are negatively correlated with downstream performance.
- 3% average MTEB score improvement is obtained for Pythia-410M with DiME-based layer selection.The reported improvement requires no supervision.
- Token-level random noise replaces each token with probability p and preserves the prompt’s token length, unlike character-level random noise.
G. Theorems
The theorems connect matrix-based entropy with representation geometry, effective rank, dataset entropy, and InfoNCE. Together, they formalize how entropy and rank characterize compression and diversity in embeddings.
- Matrix-based entropy is Schur-concave for α > 0, linking eigenvalue concentration to representation entropy.
- The effective rank of an embedding matrix is bounded above by exp(S1(Z)), where S1 is matrix-based entropy at α = 1.
- Random unit vectors in high-dimensional space are nearly orthogonal with high probability, supporting geometric analyses of representation diversity.
- Maximum prompt entropy under orthogonal equivariance and isotropic Gaussian data yields a full-rank, highly diverse dataset representation with high probability.
- Minimal prompt entropy produces rank-one sample representations whose rows are identical, corresponding to strongly compressed representations.
- Dataset entropy bounds InfoNCE loss, while entropy is analogous to the dataset-level representation quantity used in the analysis.