Probing Classifiers: Finding What a Layer Knows

A probing classifier is a small supervised model trained on a network’s hidden states to predict some property — part of speech, sentiment, whether the sentence is a question. High accuracy is taken to mean the property is encoded there. The problem is that a sufficiently expressive probe can learn the property from almost anything, so a probe without a control measures the probe.

The idea, and the confound in it

The method goes back to Alain and Bengio’s linear classifier probes in 2016 and became standard in NLP through the diagnostic-task work of Conneau and colleagues in 2018 and Hewitt and Manning’s structural probe for syntax trees in 2019. The setup is always the same: freeze the model, run your data through it, capture the hidden state at some layer and position, and fit a classifier from those vectors to your labels.

The confound is this. Suppose a probe recovers part-of-speech tags from layer four at 96% accuracy. Two explanations fit: the layer represents part of speech, or the layer represents word identity and the probe memorised the mapping from words to tags. The second requires nothing of the model at all — a lookup table gets there — and it is entirely consistent with the number.

Hewitt and Liang named the fix in 2019: a control task. Build a second labelling with the same structure as the real one but with the labels assigned arbitrarily — each word type gets a random tag, fixed across the dataset. A probe on the control task can only succeed by memorising word identities, because there is nothing else to learn. The difference between real-task accuracy and control-task accuracy is selectivity, and it is the number the paper should be reporting instead of accuracy.

Step 1: extract hidden states

Everything below runs on CPU with gpt2 and a few thousand tokens. Scale up once the harness works.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "gpt2"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name)
model.eval()

@torch.no_grad()
def hidden_states(text):
    """Returns (tokens, states) where states[layer] is (seq_len, d_model)."""
    ids = tok(text, return_tensors="pt")
    out = model(**ids, output_hidden_states=True)
    toks = [tok.decode([i]) for i in ids["input_ids"][0]]
    # out.hidden_states has n_layers + 1 entries; [0] is the embedding output
    states = [h[0] for h in out.hidden_states]
    return toks, states

toks, states = hidden_states("The cat sat on the mat.")
print(len(states), states[6].shape)   # 13  torch.Size([7, 768])

Two decisions are already made here and both matter. The layer index: hidden_states[0] is the embedding output, not layer one, so an off-by-one here shifts every conclusion you draw. And the position: for a token-level property you want the state at that token, but for a sentence-level property you need a pooling rule, and mean-pooling versus last-token gives materially different probes on a causal model where only the last position has seen the whole sentence.

Step 2: train the probe

Use a linear probe first. Not because linearity is principled — it is a choice about what counts as “encoded” — but because a linear probe with a control gives you an interpretable selectivity number, and an MLP probe usually does not.

import numpy as np
import torch.nn as nn

def train_probe(X_train, y_train, X_dev, y_dev, n_classes,
                epochs=200, lr=1e-2, weight_decay=1e-3, seed=0):
    torch.manual_seed(seed)
    probe = nn.Linear(X_train.shape[1], n_classes)
    opt = torch.optim.Adam(probe.parameters(), lr=lr, weight_decay=weight_decay)
    loss_fn = nn.CrossEntropyLoss()

    for _ in range(epochs):
        opt.zero_grad()
        loss = loss_fn(probe(X_train), y_train)
        loss.backward()
        opt.step()

    with torch.no_grad():
        acc = (probe(X_dev).argmax(-1) == y_dev).float().mean().item()
    return probe, acc

# X: (n_examples, d_model) float tensor of hidden states at one layer
# y: (n_examples,) long tensor of labels
probe, acc = train_probe(X_tr, y_tr, X_dev, y_dev, n_classes=12)
print(f"real task accuracy: {acc:.3f}")

Report the majority-class baseline alongside this. On a part-of-speech-shaped task the majority class alone can be well above 20%, and a probe that beats chance but not the majority baseline is not evidence of anything.

Step 3: the control task

This is the step that turns the exercise into an experiment.

import random

def control_labels(tokens, n_classes, seed=0):
    """Each word TYPE gets a fixed random label. Same label space, no structure."""
    rng = random.Random(seed)
    table = {}
    out = []
    for t in tokens:
        if t not in table:
            table[t] = rng.randrange(n_classes)
        out.append(table[t])
    return torch.tensor(out)

y_ctrl_tr  = control_labels(tokens_tr,  n_classes=12)
y_ctrl_dev = control_labels(tokens_dev, n_classes=12)   # same table, see below

_, ctrl_acc = train_probe(X_tr, y_ctrl_tr, X_dev, y_ctrl_dev, n_classes=12)
print(f"control accuracy: {ctrl_acc:.3f}")
print(f"selectivity:      {acc - ctrl_acc:.3f}")

Reading selectivity

Run the whole sweep across layers and plot selectivity, not accuracy, against layer index. The shape of that curve is the actual result, and it is what makes probing worth doing: properties tend to peak at different depths, and where they peak is a claim about the model that survives the confound.

Decodable is not used

A probe with excellent selectivity has established that the information is there. It has not established that the model uses it. This is a real distinction and there are documented cases where a highly decodable property turns out to be irrelevant to the model’s output.

The tools for the next question are causal. Amnesic probing (Elazar and colleagues, 2021) removes a property from a representation — typically by iterated nullspace projection, which repeatedly fits a linear classifier and projects out the direction it found — and then measures how much the model’s actual behaviour degrades. If you can strip part-of-speech information out entirely and language-modelling loss barely moves, the model was not relying on it, however cleanly it was encoded.

The other route is to intervene along the probe’s direction and see whether behaviour moves in the predicted way. That is exactly what steering vectors do, and a probe whose direction steers behaviour is a much stronger result than a probe that merely reads.

The four ways this goes wrong

1. Probe capacity is unreported. “A classifier” can mean logistic regression or a two-layer MLP, and the two support very different conclusions. State the architecture, the regularisation and the number of steps.
2. Subword alignment is silently wrong. Word-level labels against subword tokens need an explicit rule — first piece, last piece, or mean over pieces. Pick one, write it down, and check a few examples by hand. This is where most quiet failures are.
3. Train and test share sentences. Splitting at the token level rather than the document level puts near-duplicate contexts on both sides and inflates everything, including the control.
4. The result is stated causally. “Layer six encodes syntax” is supportable with a control. “Layer six computes syntax and passes it to layer seven” is not, and needs an intervention.