Probing Layers › Designing a Layer Probing Study

Running a layer probing study requires answering several design questions. Choices made here significantly affect the conclusions you can draw.

The most common choice is a linear probe: a single-layer classifier with no hidden layers. A linear probe tests whether the target property is linearly decodable from the representation, which is a strict test. If a non-linear probe does much better than a linear one, the property is encoded in a more complex geometry.

You should generally prefer linear probes for interpretability purposes. A complex probe might succeed by doing its own feature engineering, so a high accuracy number no longer tells you much about the representation; it tells you more about the probe's capacity. The canonical recommendation from Alain and Bengio (2017) is to keep probes simple enough that their success is attributable to the representation, not the classifier.

There is an important nuance here. A linear probe is not just simple; it makes a specific geometric claim. Linearly decodable means the property can be read off by a hyperplane in the representation space. This is the mildest possible version of "the information is there": if the property is present at all in any accessible form, a linear probe should be able to find it. If even a linear probe cannot, the property is either not encoded or is encoded in a fundamentally non-linear way that resists linear readout.

Single-hidden-layer MLPs with a small number of units (say, 50-200) are sometimes used when you want to test for slightly non-linear encodings without giving the probe too much freedom. The key constraint is regularization: strong L2 regularization or dropout prevents the MLP from memorizing the training set, which keeps the accuracy signal meaningful.

When probing at multiple layers with the same probe architecture, you should use the same hyperparameters and regularization at every layer. Varying regularization across layers would confound probe capacity differences with differences in how information is encoded.

A good probe task has the following properties:

• Well-defined labels: Each token or sentence should have a clear, unambiguous gold label.
• Linguistically interpretable: The task should correspond to a concept you care about (POS tag, NER label, syntactic role).
• Appropriate difficulty: Very easy tasks may show 99% accuracy at every layer, telling you nothing about localization.

Common probe tasks include:

One underappreciated issue is the relationship between class frequency and probe accuracy. POS tagging in English is dominated by a few frequent tags (NOUN, VERB, DET), and a probe that simply predicts the most common tag will achieve non-trivial accuracy by chance. Reporting only raw accuracy can therefore be misleading. Macro-averaged F1 across all classes, or accuracy stratified by tag frequency, gives a more reliable picture of what the probe has learned.

Another consideration is probe dataset size. Small datasets may not have enough examples of rare categories to produce reliable accuracy estimates. NER probing with 500 examples, for instance, might see very few ORGANIZATION entities and produce noisy accuracy for that class. Larger evaluation sets (ideally several thousand tokens per category) produce more reliable layer localization results.

To run a probing study, you need to extract the hidden states from a pre-trained model at every layer. With transformers from the Hugging Face library, this is straightforward. The critical keyword is output_hidden_states=True. When this is set, the model returns a tuple containing the hidden states at every layer, rather than only the final output. For a 12-layer BERT-base model, you get 13 tensors total: one for the embedding layer and one for each transformer block.

It is important to run the model in evaluation mode (model.eval()) and with torch.no_grad() to disable gradient computation. Since you are not fine-tuning the model, this is both faster and prevents any accidental weight updates. The representations you collect are the raw hidden states before any task-specific head, which is exactly what probing studies require: frozen, unmodified internal representations.

For sentence-level tasks (rather than token-level ones), you have additional aggregation choices. The [CLS] token representation is commonly used because BERT was trained to encode sentence-level information into it for the next-sentence prediction task. However, mean-pooling over all non-padding tokens often performs comparably or better on tasks that are not directly tied to the [CLS] training objective. The choice between CLS and mean-pooling can affect probe accuracy, so it is worth trying both and reporting results for both when presenting a probing study.

Another subtle choice is whether to extract representations before or after layer normalization. Most implementations return the post-normalization representations, which have been scaled to have consistent magnitude across the layer. This is usually the right choice for probing, since it ensures that probes at different layers receive comparably scaled input features. Probing pre-normalization representations would introduce confounds from scale differences across layers.

Transformers use subword tokenizers (WordPiece, BPE), which split words into pieces. The word "running" might become ["run", "##ning"]. If your probe targets word-level properties (like POS tags), you need to aggregate token-level representations. Common choices:

• First subword: Use the representation of the first subword token. Simple, fast, commonly used.
• Mean pooling: Average the representations of all subword tokens for a word. Tends to perform slightly better but requires token alignment.

The choice between these strategies matters more for morphologically rich languages, where a single word might be split into three or more subword pieces. In English, the difference is usually small, but in German or Finnish it can be substantial. For most English probing studies, the first-subword strategy is sufficient and is the standard baseline.

One subtle issue with subword alignment is that the tokenizer may insert special tokens ([CLS] at the start, [SEP] at the end) that shift all positions by 1. The Hugging Face word_ids() method returns None for these special tokens, making alignment straightforward: simply skip positions where word_id is None.

The aggregation question also matters for very long words. The word "nationalization" splits into ["national", "##ization"] and the two subword representations may encode different aspects of the word. The first subword "national" carries the stem meaning, while "##ization" carries the nominalization morphology. Mean-pooling integrates both signals, which can be beneficial for morphological probing tasks but may introduce noise for tasks where only the stem matters.