Probing Transformer Representations Layer by Layer, Module by Module

Probing Transformer Representations Layer by Layer, Module by Module. By Ambroise Odonnat and four others, from Noah's Ark Lab, Inria, and five other institutions. Published on the eighth of November.

Probing a transformer usually means picking a layer. Recent work has shown that the standard choice, the final layer, is often suboptimal, because intermediate layers can provide better features. Even so, representations are still mostly studied along depth alone, with each transformer block treated as one atomic unit. The claim here is that depth is only half the question. Which module inside a block gets probed matters too, and layer and module should be chosen jointly.

The evidence comes from more than one hundred and forty-four thousand linear probing evaluations, run on vision foundation models and on large language models. Two comparisons carry it. Choosing the best layer while keeping the residual stream fixed gives some gain over probing the final-layer residual stream. Choosing the best layer and module together gives a substantially larger one. Module performance also varies substantially from one layer to the next, which is why useful representations cannot be characterized by depth alone.

From there the authors test simple fixed strategies. For vision, probe the feedforward activation of the final layer, which they call Act at one hundred. For language, probe the second normalization layer at sixty percent depth, called LN two at sixty. These fixed choices improve over the probing baseline in a majority of evaluations. On average they recover forty-two percent of the gain from probing the best layer-module pair for vision, and eighty-six percent of it for language.

Several limits sit inside those numbers. The fixed strategies help in a majority of evaluations, not in all of them. The share of the available gain they recover differs sharply by modality: under half for vision, most of it for language. For vision foundation models the gains also depend on the task, growing with the transfer difficulty between pretraining data and downstream data. And everything here is measured with linear probes, on the two families of model the study covers.

The practical change is to the search space. Asking only which layer to probe treats a block as a single point. Asking which layer and which module turns every block into several candidates, and the study finds better representations among them, with a simple fixed default for each modality when a full search is not on the table.