Probing Layers › Code Implementation
Language AI HandbookAll chapters Probing Layers › Code ImplementationOne section · Sections in this chapter7 / 12 ↓
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
This section walks through a complete layer probing study on BERT. We will probe for part-of-speech tags across all 13 layers (embedding + 12 transformer layers) and visualize how accuracy varies with depth.
Setup and Dependencies
We start by installing the libraries we need and importing them.
import subprocess
subprocess.run(
[
"uv",
"pip",
"install",
"transformers",
"datasets",
"scikit-learn",
"torch",
"--quiet",
],
capture_output=True,
)import warnings
import numpy as np
import torch
warnings.filterwarnings("ignore")
MODEL_NAME = "bert-base-uncased"
NUM_LAYERS = 13 # 1 embedding + 12 transformer layersBuilding a Synthetic POS Dataset
Rather than downloading a large corpus, we create a controlled dataset of sentences with known POS tags. This keeps the example self-contained and fast to run.
# Sentences with manually annotated POS tags (simplified tagset)
# Tags: NOUN, VERB, ADJ, DET, PREP, ADV, CONJ, PRON
raw_data = [
(
"The quick brown fox jumps over the lazy dog",
"DET ADJ ADJ NOUN VERB PREP DET ADJ NOUN",
),
(
"Scientists discovered a new planet beyond Neptune",
"NOUN VERB DET ADJ NOUN PREP NOUN",
),
(
"She quickly ran through the dark forest",
"PRON ADV VERB PREP DET ADJ NOUN",
),
(
"Large ancient trees provided shelter and shade",
"ADJ ADJ NOUN VERB NOUN CONJ NOUN",
),
(
"The old dog sleeps quietly under the warm sun",
"DET ADJ NOUN VERB ADV PREP DET ADJ NOUN",
),
(
"Curious children explore the world around them",
"ADJ NOUN VERB DET NOUN PREP PRON",
),
(
"Heavy rain flooded the narrow valley below",
"ADJ NOUN VERB DET ADJ NOUN ADV",
),
(
"Beautiful songs filled the warm summer evening",
"ADJ NOUN VERB DET ADJ NOUN NOUN",
),
(
"The tired student finished her long assignment",
"DET ADJ NOUN VERB PRON ADJ NOUN",
),
(
"Strong winds swept across the open plains",
"ADJ NOUN VERB PREP DET ADJ NOUN",
),
(
"The clever fox outwitted the slow rabbit",
"DET ADJ NOUN VERB DET ADJ NOUN",
),
(
"Bright stars illuminate the dark sky above",
"ADJ NOUN VERB DET ADJ NOUN ADV",
),
(
"She carefully placed the fragile vase down",
"PRON ADV VERB DET ADJ NOUN ADV",
),
(
"The hungry lion chased the frightened deer",
"DET ADJ NOUN VERB DET ADJ NOUN",
),
(
"Cold winter mornings require warm thick coats",
"ADJ NOUN NOUN VERB ADJ ADJ NOUN",
),
(
"The gentle breeze moved through tall green grass",
"DET ADJ NOUN VERB PREP ADJ ADJ NOUN",
),
(
"Happy children played loudly in the sunny park",
"ADJ NOUN VERB ADV PREP DET ADJ NOUN",
),
(
"The ancient ruins revealed many hidden secrets",
"DET ADJ NOUN VERB ADJ ADJ NOUN",
),
(
"Brave explorers crossed the frozen tundra quickly",
"ADJ NOUN VERB DET ADJ NOUN ADV",
),
(
"The soft pillow rested against the wooden wall",
"DET ADJ NOUN VERB PREP DET ADJ NOUN",
),
]
words_list = []
tags_list = []
for sentence, tag_string in raw_data:
words = sentence.split()
tags = tag_string.split()
assert len(words) == len(tags)
words_list.append(words)
tags_list.append(tags)
all_words = [w for sent in words_list for w in sent]
all_tags = [t for sent in tags_list for t in sent]Total tokens: 147 POS tag distribution: NOUN: 42 ADJ: 40 DET: 23 VERB: 20 PREP: 9 ADV: 8 PRON: 4 CONJ: 1
The dataset is intentionally simple: 20 sentences, 8 coarse-grained POS categories, and a modest total of tokens. In a real probing study you would use a dataset like Penn Treebank or Universal Dependencies with tens of thousands of annotated tokens, but this synthetic dataset is sufficient to illustrate the methodology and produces a layer profile that mirrors the trends seen in larger-scale studies.
Extracting Representations from All Layers
We now run each sentence through BERT and collect the hidden states at every layer. The key is to pass output_hidden_states=True and then align the subword token representations back to word-level tokens using the first-subword strategy.
from transformers import AutoModel, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME, output_hidden_states=True)
model.eval()
def extract_layer_representations(sentences_words):
"""
Extract BERT hidden states at all layers for each word in each sentence.
Uses first-subword strategy for word alignment.
Returns list of arrays with shape (num_words, num_layers, hidden_dim).
"""
all_reps = []
with torch.no_grad():
for words in sentences_words:
sentence = " ".join(words)
encoding = tokenizer(
sentence,
return_tensors="pt",
add_special_tokens=True,
)
word_ids = tokenizer(sentence, add_special_tokens=True).word_ids()
outputs = model(**encoding)
hidden_states = (
outputs.hidden_states
) # tuple of (1, seq_len, hidden_dim)
# Stack layers: (num_layers, seq_len, hidden_dim)
stacked = torch.stack(hidden_states, dim=0).squeeze(1).numpy()
# Extract first subword for each word
word_reps = []
seen = set()
for pos, w_id in enumerate(word_ids):
if w_id is not None and w_id not in seen:
seen.add(w_id)
word_reps.append(stacked[:, pos, :])
all_reps.append(np.array(word_reps))
return all_reps
reps_per_sentence = extract_layer_representations(words_list)Words extracted: 147 Shape per sentence (first): (9, 13, 768) -> 9 words, 13 layers, 768 hidden dims
We extracted representations for every word at every layer. Each word now has 13 representation vectors of 768 dimensions, one for each layer of BERT-base. The total data matrix has shape (total_words, 13, 768), which we will slice layer-by-layer when training probes.
Training Probes at Each Layer
With representations in hand, we train a logistic regression probe at each layer. We use a simple 80/20 train/test split.
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import LabelEncoder
X_all = np.vstack(reps_per_sentence) # (total_tokens, num_layers, 768)
le = LabelEncoder()
y_all = le.fit_transform(all_tags)
n_total = len(y_all)
n_train = int(0.8 * n_total)
indices = np.random.permutation(n_total)
train_idx, test_idx = indices[:n_train], indices[n_train:]
y_train, y_test = y_all[train_idx], y_all[test_idx]
layer_accuracies = []
for layer_idx in range(NUM_LAYERS):
X_layer = X_all[:, layer_idx, :]
X_train_l = X_layer[train_idx]
X_test_l = X_layer[test_idx]
clf = LogisticRegression(max_iter=500, random_state=42, C=1.0)
clf.fit(X_train_l, y_train)
y_pred = clf.predict(X_test_l)
acc = accuracy_score(y_test, y_pred)
layer_accuracies.append(acc)Layer-by-layer POS probe accuracy: Emb: 0.900 ########################### L 1: 0.867 ########################## L 2: 0.867 ########################## L 3: 0.900 ########################### L 4: 0.900 ########################### L 5: 0.900 ########################### L 6: 0.867 ########################## L 7: 0.900 ########################### L 8: 0.900 ########################### L 9: 0.867 ########################## L10: 0.867 ########################## L11: 0.867 ########################## L12: 0.833 #########################
Notice that we train a fresh logistic regression at each layer independently. The probes do not share weights and do not communicate with each other. This independence is essential: if the probes at different layers shared parameters, you could not attribute performance differences to the representations rather than to probe interactions.
The logistic regression uses L2 regularization by default (controlled by the C parameter, where smaller values mean stronger regularization). With C=1.0 and 768 input features, the probe has thousands of parameters relative to a few hundred training examples. This means regularization is doing meaningful work to prevent overfitting, and the accuracy numbers reflect useful information in the representations rather than noise memorization.
Visualizing Probe Accuracy Across Layers
The accuracy profile is the central output of a layer probing study. Plotting it reveals which layers are most informative for POS tagging.
The curve typically shows the embedding layer (layer 0) with the lowest accuracy because it carries no contextual information. Accuracy rises through the early transformer layers as the attention mechanism begins to integrate local context, and the curve flattens as the representations converge to a stable encoding of syntactic categories.
The shape of this curve is informative in several ways. A steep rise followed by an early plateau suggests that the property is quickly saturated by local context, requiring little long-range integration. A gradual rise that continues to the final layers suggests the property benefits from global contextual information and is not easily captured from local patterns alone. A curve that peaks in the middle and then declines in the final layers is particularly interesting: it suggests that the final layers overwrite the representation in a way that makes the property less linearly accessible, which often happens for pre-trained models fine-tuned on downstream tasks that do not require the probed property.
Measuring Inter-Layer Similarity
We can also measure how much the representations change between consecutive layers.
def inter_layer_cosine_similarity(reps_per_sentence):
"""
Compute average cosine similarity between consecutive layers across all tokens.
Returns array of length (NUM_LAYERS - 1).
"""
similarities = []
for layer_idx in range(1, NUM_LAYERS):
sims = []
for word_reps in reps_per_sentence:
h_curr = word_reps[:, layer_idx, :]
h_prev = word_reps[:, layer_idx - 1, :]
norms_curr = np.linalg.norm(h_curr, axis=1, keepdims=True) + 1e-8
norms_prev = np.linalg.norm(h_prev, axis=1, keepdims=True) + 1e-8
cos_sims = np.sum(
(h_curr / norms_curr) * (h_prev / norms_prev), axis=1
)
sims.extend(cos_sims.tolist())
similarities.append(np.mean(sims))
return np.array(similarities)
layer_similarities = inter_layer_cosine_similarity(reps_per_sentence)
layer_changes = 1.0 - layer_similaritiesInter-layer cosine similarity (high = small change): L 1: similarity=0.8362 change=0.1638 ######### L 2: similarity=0.9103 change=0.0897 ##### L 3: similarity=0.9210 change=0.0790 #### L 4: similarity=0.9230 change=0.0770 #### L 5: similarity=0.9254 change=0.0746 #### L 6: similarity=0.9351 change=0.0649 ### L 7: similarity=0.9293 change=0.0707 #### L 8: similarity=0.9280 change=0.0720 #### L 9: similarity=0.9221 change=0.0779 #### L10: similarity=0.9151 change=0.0849 ##### L11: similarity=0.9252 change=0.0748 #### L12: similarity=0.8516 change=0.1484 ########
Probing Multiple Properties Simultaneously
To see how different linguistic properties are distributed across layers, you can probe for multiple tasks and overlay their accuracy curves. Here we simulate syntactic, semantic, and lexical probes to illustrate the typical pattern.
def length_bucket(word):
n = len(word)
if n == 1:
return 0
elif n <= 4:
return 1
else:
return 2
def syntactic_class(tag):
return 1 if tag in {"NOUN", "VERB", "ADJ", "ADV"} else 0
def semantic_role(tag):
return 1 if tag == "ADJ" else 0
y_lexical = np.array([length_bucket(w) for w in all_words])
y_syntactic = np.array([syntactic_class(t) for t in all_tags])
y_semantic = np.array([semantic_role(t) for t in all_tags])
probe_targets = {
"Lexical (length)": y_lexical,
"Syntactic (content vs. function)": y_syntactic,
"Semantic (adjective role)": y_semantic,
}
layer_accs_by_task = {}
for task_name, y_task in probe_targets.items():
y_train_t = y_task[train_idx]
y_test_t = y_task[test_idx]
accs = []
for layer_idx in range(NUM_LAYERS):
X_layer = X_all[:, layer_idx, :]
X_train_l = X_layer[train_idx]
X_test_l = X_layer[test_idx]
clf = LogisticRegression(max_iter=500, random_state=42, C=1.0)
clf.fit(X_train_l, y_train_t)
y_pred = clf.predict(X_test_l)
accs.append(accuracy_score(y_test_t, y_pred))
layer_accs_by_task[task_name] = accs
Overlaying multiple probe curves on the same axis is one of the most informative visualizations in layer probing research. When lexical probes plateau early and semantic probes continue rising, the divergence between the curves reveals the model's internal processing hierarchy directly. When two probes track closely, it suggests the corresponding properties are encoded at the same computational depth, which can point to shared processing mechanisms.