Does Representation Matter? Exploring Intermediate Layers in Large Language Models

TL;DR

This work interrogates layer-wise representations in Transformer and State Space Model (SSM) large language models, showing that intermediate layers often yield more informative features for downstream tasks than final layers. It adapts matrix-based prompt entropy, curvature, and augmentation-invariance metrics (InfoNCE, DiME, LiDAR) to quantify representation quality across architectures, training stages, and prompt perturbations. Key findings include architecture-dependent dynamics ( Transformers exhibiting greater mid-layer changes and compression, SSMs being more stable), a pronounced impact of training progression on intermediate layers, and a surprising bimodal entropy pattern in Transformer middle layers whose cause remains unresolved. The results offer practical guidance for architectural optimization and training strategies that leverage intermediate representations for improved transfer and robustness.

Abstract

Understanding what defines a good representation in large language models (LLMs) is fundamental to both theoretical understanding and practical applications. In this paper, we investigate the quality of intermediate representations in various LLM architectures, including Transformers and State Space Models (SSMs). We find that intermediate layers often yield more informative representations for downstream tasks than the final layers. To measure the representation quality, we adapt and apply a suite of metrics - such as prompt entropy, curvature, and augmentation-invariance - originally proposed in other contexts. Our empirical study reveals significant architectural differences, how representations evolve throughout training, and how factors like input randomness and prompt length affect each layer. Notably, we observe a bimodal pattern in the entropy of some intermediate layers and consider potential explanations tied to training data. Overall, our results illuminate the internal mechanics of LLMs and guide strategies for architectural optimization and training.

Figures (8)

• Figure 1: Pythia’s intermediate layers show pronounced changes in representation quality metrics, while Mamba’s remain more stable. Representation evaluation metrics across layers in Pythia 410M and Mamba 370M architectures. The x-axis denotes model depth as a percentage, allowing fair comparison between models with different layer counts.
• Figure 2: Training effects are most pronounced in the intermediate layers. Representation metrics across layers at different training checkpoints (steps 1 to 143k). The x-axis is the depth percentage of the model, showing how training influences different layers, particularly those at intermediate depths.
• Figure 3: Prompt entropy across layers of Pythia 410M under various extreme input conditions. (a) Increasing token repetition leads to decreased entropy in intermediate layers. (b) Increasing token randomness results in higher entropy, especially in initial layers. (c) Unnormalized prompt entropy increases with prompt length due to the larger number of tokens. These results demonstrate how the model's internal representations adapt to different types of input perturbations.
• Figure 4: Bimodal distribution of prompt entropies observed in intermediate layers. The distributions of prompt entropies for WikiText and ai-medical-chatbot datasets are shown for Pythia, Mamba, and Llama3 models. The middle column highlights the layer with the highest Dip Test score hartigan1985dip, which measures the degree of multimodality in the entropy distribution.
• Figure 5: Entropy vs Accuracy of LLama3-8B on MMLU tasks. Each point represents a task in MMLU