Probing Transformer Representations Layer by Layer and Module by Module
Summary:A study of vision foundation models and large language models evaluates more than 144,000 linear probing configurations. The central finding is that representation quality cannot be characterized by network depth alone: the choice of module and the layer containing it must be optimized jointly. The best layer-module pair substantially outperforms probing the residual stream at the best layer. The authors also identify simple fixed alternatives with limited engineering overhead: the final-layer
The central finding
A common way to study representations in a Transformer foundation model is to freeze the network and train a linear probe on the final-layer residual stream. After more than 144,000 linear probing evaluations across vision foundation models and large language models, this study shows that the default is not consistently optimal.
Useful representations cannot be characterized by depth alone. The module selected within a layer can vary substantially in performance, and jointly choosing the best layer and module produces larger gains than keeping the residual stream fixed while optimizing only the layer.
In practical terms, the representation search space should be treated as a two-dimensional grid of layers and internal modules, rather than a one-dimensional list of network depths.
What the study probes
Instead of treating each Transformer block as an atomic unit, the researchers extract and evaluate multiple internal signals, including:
- Residual-stream states, which act as the main information highway through the network;
- Attention outputs, which summarize interactions between tokens or image patches;
- Feed-forward activations, which encode nonlinear transformations produced by the MLP;
- Normalization-layer states, such as LN1 and LN2 at different depths.
A linear classifier is then trained on each candidate representation. Because the probe remains linear, the comparison is intended to isolate the usefulness of the underlying representation rather than relying on a large nonlinear head to compensate for a weak feature space.
Architecturally, the residual stream is a shared bus into which attention and MLP blocks continually write information. It therefore accumulates multiple overlapping abstractions. A module-local activation instead reflects one particular computation and may expose a cleaner encoding of semantics, geometry, syntax, or downstream task structure. This helps explain why the best layer-module pairing can outperform the residual stream even at the best layer.
Practical fixed probing strategies
Exhaustively searching every layer-module combination can be expensive. The paper therefore evaluates two simple fixed strategies:
- Act@100 for vision: the feed-forward activation from the final layer;
- LN2@60 for language: the second normalization layer at approximately 60% of model depth.
Both strategies outperform the standard final-layer residual-stream probe in most evaluations. Relative to the oracle layer-module pairing, Act@100 recovers an average of 42% of the achievable gain in vision, while LN2@60 recovers 86% in language.
For vision models, the benefit also increases as transfer between pretraining and downstream data becomes more difficult. This suggests that intermediate or module-specific features may become particularly valuable when the deployment distribution differs from the pretraining distribution, whereas the final residual stream may be more specialized to the original training regime.
How to interpret the benchmark
The scale of the evaluation supports the main empirical conclusion, but the abstract does not expose every experimental control. Linear-probe performance can be influenced not only by representation quality but also by feature dimensionality, probe capacity, optimizer settings, and training-set size. A larger activation vector may improve downstream scores simply because the probe has more parameters.
A rigorous reproduction should therefore add dimension-matched probes, stronger regularization controls, and repeated cross-validation. It should also compare against simple alternatives such as PCA or random projection, ensuring that the gains are not merely due to increased input size.
The fixed configurations should also be interpreted as empirical defaults rather than universal laws. Activation placement, normalization order, residual architecture, model scale, and pretraining recipe can all change the best probe location. A small validation sweep remains safer than hard-coding Act@100 or LN2@60 across unrelated models.
Finally, predictability is not causality. A linear probe can exploit features correlated with the target without identifying the mechanism the model uses to solve the task. The findings therefore support representation extraction and diagnostic tooling, but do not by themselves prove that a particular module contains the model's true reasoning computation.
Practical guidance
Researchers and engineers who need access to internal model states can adopt the following workflow:
- Scan modules as well as depths. Include the residual stream, attention output, MLP activations, and multiple normalization positions.
- Control probe capacity. Normalize dimensions, match parameter counts, and use cross-validation to avoid confusing larger representations with better ones.
- Start with low-cost fixed locations. Use Act@100 for vision and LN2@60 for language as initial candidates, then refine them for the target architecture.
- Account for inference overhead. Reading intermediate states requires running the model to the relevant depth and managing hooks and memory. Where possible, truncate later computation or use forward caching.
- Use probes for diagnosis. Compare layer-module combinations under in-distribution, shifted, and transfer tasks to identify where different kinds of information are encoded.
The main contribution is methodological rather than architectural: Transformer internals should be indexed by both depth and module type. For interpretability research, linear evaluation, feature extraction, and lightweight adaptation, that shift can materially improve probe quality while reducing the need for exhaustive end-to-end fine-tuning.
— END —Source: Hugging Face Daily Papers (2026-11-08)
Tags: #Transformer #Linear Probing #Internal Representations #Interpretability #Foundation Models
Editorial & Fact-Checking Note: This article is compiled from primary research, official release documentation, and source papers by the ZICQ Newsroom pipeline with automated entity verification and human editorial review. If you notice any technical inaccuracy, please submit a correction via our corrections policy or email our editorial desk directly.
Community Comments