Probing Classifiers › Code Implementation
Language AI HandbookAll chapters Probing Classifiers › Code ImplementationOne section · Sections in this chapter8 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us build a complete probing experiment from scratch. We will probe BERT representations for part-of-speech information, implement control tasks, and measure selectivity. The experiment uses a small corpus to stay tractable, but the methodology scales to any size.
Setup and Data Preparation
We start with imports and a small POS-tagged dataset. In practice you would use a standard corpus like the Penn Treebank or Universal Dependencies; here we construct a compact synthetic dataset with varied grammatical structures.
import warnings
import numpy as np
import torch
warnings.filterwarnings("ignore")
# Fix random seeds for reproducibility
np.random.seed(42)
torch.manual_seed(42)# Small POS-tagged corpus (word, POS label pairs)
# In real research, you would use Universal Dependencies or Penn Treebank
labeled_sentences = [
[("The", "DET"), ("cat", "NOUN"), ("chases", "VERB"), ("mice", "NOUN")],
[("Birds", "NOUN"), ("fly", "VERB"), ("south", "ADV"), ("quickly", "ADV")],
[
("She", "PRON"),
("writes", "VERB"),
("clever", "ADJ"),
("stories", "NOUN"),
],
[("The", "DET"), ("old", "ADJ"), ("man", "NOUN"), ("walked", "VERB")],
[
("Dogs", "NOUN"),
("bark", "VERB"),
("loudly", "ADV"),
("sometimes", "ADV"),
],
[("A", "DET"), ("small", "ADJ"), ("child", "NOUN"), ("laughed", "VERB")],
[("He", "PRON"), ("runs", "VERB"), ("very", "ADV"), ("fast", "ADJ")],
[("Strong", "ADJ"), ("winds", "NOUN"), ("blow", "VERB"), ("here", "ADV")],
[("The", "DET"), ("red", "ADJ"), ("fox", "NOUN"), ("jumped", "VERB")],
[
("They", "PRON"),
("sing", "VERB"),
("beautiful", "ADJ"),
("songs", "NOUN"),
],
[("An", "DET"), ("eager", "ADJ"), ("student", "NOUN"), ("studied", "VERB")],
[("Rain", "NOUN"), ("falls", "VERB"), ("heavily", "ADV"), ("today", "ADV")],
[("The", "DET"), ("young", "ADJ"), ("dog", "NOUN"), ("sleeps", "VERB")],
[("We", "PRON"), ("read", "VERB"), ("long", "ADJ"), ("books", "NOUN")],
[
("Bright", "ADJ"),
("stars", "NOUN"),
("shine", "VERB"),
("nightly", "ADV"),
],
[("The", "DET"), ("happy", "ADJ"), ("child", "NOUN"), ("dances", "VERB")],
[("It", "PRON"), ("rains", "VERB"), ("quite", "ADV"), ("often", "ADV")],
[("Blue", "ADJ"), ("cars", "NOUN"), ("drive", "VERB"), ("fast", "ADJ")],
[("The", "DET"), ("tired", "ADJ"), ("worker", "NOUN"), ("rested", "VERB")],
[("She", "PRON"), ("speaks", "VERB"), ("fluently", "ADV"), ("now", "ADV")],
]Extracting BERT Representations
We load a pretrained BERT model and extract the hidden states from a middle layer (layer 6 is standard for syntactic tasks). Each word is represented by the hidden state of its first subword token.
import torch
from transformers import BertModel, BertTokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained(
"bert-base-uncased", output_hidden_states=True
)
model.eval()
PROBE_LAYER = (
6 # Layer 6 of BERT (0-indexed; BERT-base has 12 transformer layers)
)
def extract_word_representations(sentence_tokens, layer_idx):
"""
Extract BERT representation for each word in a sentence.
Returns (representations, labels) where representations has shape (n_words, 768).
Uses the first subword token for multi-subword words.
"""
words = [w for w, _ in sentence_tokens]
labels = [pos for _, pos in sentence_tokens]
text = " ".join(words)
encoding = tokenizer(text, return_tensors="pt", add_special_tokens=True)
word_ids = encoding.word_ids() # maps subword positions to word indices
with torch.no_grad():
outputs = model(**encoding)
hidden_states = (
outputs.hidden_states
) # tuple of (n_layers+1, 1, seq_len, 768)
# Layer layer_idx + 1 because hidden_states[0] is the embedding layer
layer_hidden = hidden_states[layer_idx + 1][0] # (seq_len, 768)
# Grab the first subword token for each original word
word_reps = []
seen_words = set()
for token_idx, word_idx in enumerate(word_ids):
if word_idx is None:
continue # [CLS] or [SEP]
if word_idx not in seen_words:
seen_words.add(word_idx)
word_reps.append(layer_hidden[token_idx].numpy())
return word_reps, labels# Extract representations for all sentences
all_representations = []
all_labels = []
for sentence in labeled_sentences:
reps, labels = extract_word_representations(sentence, PROBE_LAYER)
all_representations.extend(reps)
all_labels.extend(labels)
X = np.array(all_representations) # (n_tokens, 768)
y = np.array(all_labels)Total tokens extracted: 80
Representation dimensionality: 768
POS labels: [np.str_('ADJ'), np.str_('ADV'), np.str_('DET'), np.str_('NOUN'), np.str_('PRON'), np.str_('VERB')]
Label distribution:
ADJ: 15 (18.8%)
ADV: 13 (16.2%)
DET: 8 (10.0%)
NOUN: 18 (22.5%)
PRON: 6 (7.5%)
VERB: 20 (25.0%)Our dataset contains 80 token-level examples (20 sentences times 4 tokens each). Each example is a 768-dimensional BERT hidden-state vector. The labels span POS categories with a distribution shaped by our sentence construction. The high dimensionality relative to the number of examples is exactly the regime where control tasks become critical, since a linear classifier has more than enough parameters to memorize arbitrary labels.
Training the Linear Probe
With representations in hand, we split the data and train a logistic regression probe. The key is that the BERT representations are completely frozen; only the logistic regression parameters are trained.
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
# Encode string labels to integers
label_enc = LabelEncoder()
y_enc = label_enc.fit_transform(y)
# Train/test split stratified by label
X_train, X_test, y_train, y_test = train_test_split(
X, y_enc, test_size=0.3, random_state=42, stratify=y_enc
)
# Train linear probe (logistic regression with L2 regularization)
probe = LogisticRegression(
C=1.0, # Inverse regularization strength
max_iter=1000,
random_state=42,
solver="lbfgs",
)
probe.fit(X_train, y_train)
probe_accuracy = accuracy_score(y_test, probe.predict(X_test))Linear probe accuracy on POS tagging: 1.000 Chance baseline: 0.167 Accuracy above chance: 0.833
The accuracy above chance shows how much the BERT representations help, compared to predicting the most common label always. But this number alone is insufficient. We need to compare it against a control to understand whether the improvement reflects syntactic encoding or just the probe's ability to separate high-dimensional points.
Implementing Control Tasks
Now we implement the critical control: random label assignment. This tests whether the high accuracy is specific to POS structure or whether the probe could memorize arbitrary labels equally well.
def run_control_task(X_tr, X_te, y_tr, y_te, n_cls, n_trials=5):
"""
Run control task with random label assignment.
Returns mean and std accuracy across multiple random assignments.
"""
control_accuracies = []
for trial in range(n_trials):
rng = np.random.RandomState(trial)
y_random_train = rng.randint(0, n_cls, size=len(y_tr))
y_random_test = rng.randint(0, n_cls, size=len(y_te))
ctrl_probe = LogisticRegression(
C=1.0, max_iter=1000, random_state=trial, solver="lbfgs"
)
ctrl_probe.fit(X_tr, y_random_train)
ctrl_acc = accuracy_score(y_random_test, ctrl_probe.predict(X_te))
control_accuracies.append(ctrl_acc)
return float(np.mean(control_accuracies)), float(np.std(control_accuracies))
n_classes = len(label_enc.classes_)
control_mean, control_std = run_control_task(
X_train, X_test, y_train, y_test, n_classes
)
selectivity = probe_accuracy - control_meanReal task accuracy: 1.000 Control task accuracy: 0.225 +/- 0.090 Selectivity (gap): 0.775 Strong selectivity: probe finds genuine POS signal in BERT representations.
The selectivity score isolates what the real POS labels contribute over and above arbitrary label assignment. A large gap means the probe succeeded specifically because of the linguistic structure in the labels. A small gap means you cannot confidently attribute the probe's accuracy to POS encoding, regardless of how high the absolute accuracy is.
Note the subtlety in the control design. We assign random labels to both train and test independently, so the probe cannot memorize train-to-test mappings. The test accuracy for the control probe reflects how well a probe trained on random labels generalizes to new random labels, which, in 768 dimensions with regularization, is lower than its training accuracy. This is precisely the right comparison: we want to know whether the representations contain information about POS labels, not whether the probe can overfit on train.
Probing Multiple Layers
One of the most revealing analyses in probing research is comparing probe accuracy across all layers of a model. This produces a probing profile that shows where in the network different types of information are most linearly accessible.
def probe_single_layer(sentences, layer_idx, le):
"""Extract representations for a given layer and evaluate a linear probe."""
reps, lbls = [], []
for sentence in sentences:
r, l = extract_word_representations(sentence, layer_idx)
reps.extend(r)
lbls.extend(l)
X_l = np.array(reps)
y_l = le.transform(lbls)
X_tr, X_te, y_tr, y_te = train_test_split(
X_l, y_l, test_size=0.3, random_state=42, stratify=y_l
)
p = LogisticRegression(
C=1.0, max_iter=1000, random_state=42, solver="lbfgs"
)
p.fit(X_tr, y_tr)
return float(accuracy_score(y_te, p.predict(X_te))), X_tr, X_te, y_tr, y_te
# Probe all 12 transformer layers of BERT-base
layers_to_probe = list(range(12))
layer_accuracies = []
layer_ctrl_accs = []
for layer_idx in layers_to_probe:
acc, X_tr_l, X_te_l, y_tr_l, y_te_l = probe_single_layer(
labeled_sentences, layer_idx, label_enc
)
layer_accuracies.append(acc)
c_mean, _ = run_control_task(
X_tr_l, X_te_l, y_tr_l, y_te_l, n_classes, n_trials=3
)
layer_ctrl_accs.append(c_mean)Layer-by-layer probing accuracy (POS tagging): Layer Accuracy Control Selectivity -------------------------------------------- 0 0.958 0.153 0.806 ################### 1 0.958 0.125 0.833 ################### 2 0.958 0.139 0.819 ################### 3 1.000 0.125 0.875 #################### 4 1.000 0.153 0.847 #################### 5 1.000 0.167 0.833 #################### 6 1.000 0.153 0.847 #################### 7 1.000 0.167 0.833 #################### 8 1.000 0.153 0.847 #################### 9 1.000 0.153 0.847 #################### 10 1.000 0.139 0.861 #################### 11 1.000 0.139 0.861 #################### Best layer: 3 (accuracy: 1.000)
This layer analysis typically shows that syntactic information peaks in the middle layers of BERT (layers 4 through 8), while semantic information peaks in higher layers. This finding has been replicated across many probe tasks and model architectures. The pattern reflects the computational structure of the transformer: early layers process local token-level patterns, middle layers build syntactic structure through attention across the sentence, and late layers integrate meaning for the pretraining objective (predicting masked tokens requires semantic, not just syntactic, reasoning).
Visualization
The shaded region between the real-task curve and the control-task curve is the selectivity: the portion of probe accuracy that is attributable to linguistic structure in the representations. Layers where this gap is large are the layers most informative about POS structure. Notice that the control accuracy is not constant across layers; in higher layers, the representations tend to be more clustered in ways that make random-label memorization slightly easier, which is why you always need to measure selectivity rather than relying on raw accuracy alone.
Comparing Multiple Probing Tasks
It is also informative to compare multiple probing tasks at a fixed layer to understand what kinds of information that layer encodes most strongly. We compare POS tagging against two surface-level tasks: position (is the token in the first half of the sentence?) and token length (is the token short or long?).
def build_position_dataset(sentences, layer_idx):
"""Binary label: 1 if token is in first half of sentence, 0 otherwise."""
reps, lbls = [], []
for sentence in sentences:
n = len(sentence)
mid = n // 2
r, _ = extract_word_representations(sentence, layer_idx)
for i, rep in enumerate(r):
reps.append(rep)
lbls.append(1 if i < mid else 0)
return np.array(reps), np.array(lbls)
def build_length_dataset(sentences, layer_idx):
"""Binary label: 1 if token is short (<= 3 chars), 0 if long (4+ chars)."""
reps, lbls = [], []
for sentence in sentences:
r, _ = extract_word_representations(sentence, layer_idx)
for i, rep in enumerate(r):
word = sentence[i][0]
reps.append(rep)
lbls.append(1 if len(word) <= 3 else 0)
return np.array(reps), np.array(lbls)
def eval_probe(X_data, y_data, test_size=0.3):
"""Evaluate probe accuracy and selectivity."""
n_cls = len(np.unique(y_data))
X_tr, X_te, y_tr, y_te = train_test_split(
X_data, y_data, test_size=test_size, random_state=42
)
p = LogisticRegression(
C=1.0, max_iter=1000, random_state=42, solver="lbfgs"
)
p.fit(X_tr, y_tr)
task_acc = float(accuracy_score(y_te, p.predict(X_te)))
ctrl_mean, _ = run_control_task(X_tr, X_te, y_tr, y_te, n_cls, n_trials=5)
return task_acc, ctrl_mean, task_acc - ctrl_mean
target_layer = best_layer
X_position, y_position = build_position_dataset(labeled_sentences, target_layer)
X_length, y_length = build_length_dataset(labeled_sentences, target_layer)
pos_acc, pos_ctrl, pos_sel = probe_accuracy, control_mean, selectivity
position_acc, position_ctrl, position_sel = eval_probe(X_position, y_position)
length_acc, length_ctrl, length_sel = eval_probe(X_length, y_length)Task Accuracy Control Selectivity ------------------------------------------------------- POS tagging 1.000 0.225 0.775 Position 0.958 0.550 0.408 Token length 0.958 0.550 0.408
Interpreting the Results
The comparison across tasks reveals the layered nature of what BERT encodes.
Interpretation summary: POS selectivity: 0.775 Reflects genuine syntactic category structure in BERT representations. The probe succeeds because real POS labels partition the representation space in ways random labels cannot replicate. Position selectivity: 0.408 Positional information may be accessible due to BERT's positional embeddings, which encode token order directly. Length selectivity: 0.408 Short words (the, a, it) vs. long words differ in subword tokenization, so their embeddings may naturally cluster.
The key insight from multi-task probing is that different types of information have different accessibility profiles across layers. Surface features like token length tend to be decodable from shallow layers. Syntactic structure peaks in middle layers where attention patterns have had time to integrate context. Semantic properties often peak in higher layers where representations have been transformed by many attention and MLP operations.
The position probe result is particularly informative about BERT's architectural inductive biases. Because BERT adds sinusoidal positional embeddings to the input, position information is explicitly injected into the representation from the very first layer. It persists through the network because positional information is useful for many downstream tasks. The fact that a linear probe can decode position from BERT's representations is thus less surprising than the fact that it can decode POS, because position is explicitly encoded while POS is emergent.