Interpretability Goals › Code: Exploring Model Internals
Language AI HandbookAll chapters Interpretability Goals › Code: Exploring Model InternalsOne section · Sections in this chapter7 / 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
To make interpretability concrete, let's implement a basic exploration of a small transformer model's internal representations. We will examine how activations change across layers for different types of inputs, a foundation for the probing and visualization techniques in later chapters. This code is exploratory rather than definitive: it demonstrates the setup and measurement tools, but four sentences is far too small a dataset to draw firm conclusions.
We start by loading a pretrained model and setting up extraction of intermediate activations.
import warnings
from transformers import AutoModel, AutoTokenizer
warnings.filterwarnings("ignore")
# Use a lightweight model for demonstration
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, output_hidden_states=True)
model.eval()# Define a set of sentences with different semantic properties
sentences = [
"The cat sat on the mat.", # simple, concrete
"The president signed the bill.", # factual, political
"The matrix was singular.", # technical, mathematical
"She felt deeply melancholic.", # emotional, abstract
]
# Tokenize all sentences
encoded = [
tokenizer(s, return_tensors="pt", padding=True, truncation=True)
for s in sentences
]import numpy as np
import torch
# Extract hidden states from all layers
all_hidden_states = []
with torch.no_grad():
for enc in encoded:
outputs = model(**enc)
# hidden_states: tuple of (num_layers+1) tensors of shape (1, seq_len, hidden_size)
# Index 0 is the embedding layer, rest are transformer layers
hidden_states = torch.stack(
outputs.hidden_states, dim=0
) # (layers, 1, seq, hidden)
# Take the [CLS] token representation for each layer
cls_states = hidden_states[:, 0, 0, :].numpy() # (layers, hidden_size)
all_hidden_states.append(cls_states)
all_hidden_states = np.array(
all_hidden_states
) # (sentences, layers, hidden_size)
num_layers = all_hidden_states.shape[1]
n_sentences = all_hidden_states.shape[0]Model: distilbert-base-uncased Hidden state shape: (sentences=4, layers=7, hidden_size=768) Sentences analyzed: 4 [0] The cat sat on the mat. [1] The president signed the bill. [2] The matrix was singular. [3] She felt deeply melancholic.
We now have hidden states for each sentence at each layer. A key interpretability question is how similar sentences are to each other as a function of layer depth. Do sentences become more or less similar as the model processes them? Intuitively, if the model is doing useful work, representations should become more differentiated and semantically meaningful as depth increases, diverging from the token-identity-driven representations at the input.
from sklearn.metrics.pairwise import cosine_similarity
# Compute pairwise cosine similarity between all sentences at each layer
layer_similarities = np.zeros((num_layers, n_sentences, n_sentences))
for layer in range(num_layers):
layer_reps = all_hidden_states[:, layer, :] # (sentences, hidden_size)
layer_similarities[layer] = cosine_similarity(layer_reps)
# Compute average off-diagonal similarity per layer
avg_sims = []
for layer in range(num_layers):
sim_matrix = layer_similarities[layer]
off_diag_indices = np.triu_indices(n_sentences, k=1)
off_diag = sim_matrix[off_diag_indices]
avg_sims.append(off_diag.mean())Average inter-sentence cosine similarity across layers: Layer Avg Similarity ---------------------------- Embed 1.0000 Layer 1 0.9876 Layer 2 0.9920 Layer 3 0.9457 Layer 4 0.9309 Layer 5 0.9021 Layer 6 0.9334
The similarity values reveal something characteristic of transformer models: representations tend to be similar at the embedding layer (shared token embeddings dominate) and diversify at intermediate layers as contextual information is mixed in. This layer-by-layer transformation is exactly what probing classifiers exploit. If you train a classifier on early-layer representations to predict semantic categories, it tends to fail; if you train it on late-layer representations, it tends to succeed.
The plot shows how representations evolve through the network. At the embedding layer, all sentences are represented primarily by their token identities, producing high similarity because different sentences share many common words ("The", "the", "a"). As processing proceeds, contextual information differentiates the representations. This is the transformation that makes deeper layers better sources for probing higher-level semantic and syntactic properties. A classifier trained on embedding-layer representations would find it hard to distinguish sentence topics, but a classifier trained on final-layer representations would have access to rich contextual signal.
Now let's look at whether different types of sentences cluster together in activation space by using PCA to project the high-dimensional representations to 2D.
from sklearn.decomposition import PCA
# We'll look at the last layer representations (most contextually rich)
last_layer_idx = num_layers - 1
last_layer_reps = all_hidden_states[:, last_layer_idx, :] # (sentences, hidden)
# Also extract first real transformer layer
first_layer_reps = all_hidden_states[:, 1, :] # (sentences, hidden)
# PCA to 2D - fit on combined data for consistent coordinate system
pca = PCA(n_components=2)
combined = np.vstack([first_layer_reps, last_layer_reps])
pca.fit(combined)
first_2d = pca.transform(first_layer_reps)
last_2d = pca.transform(last_layer_reps)
These projections illustrate one of the core findings of representation analysis: how sentences cluster in activation space changes across layers. While a six-layer model like DistilBERT-base shows modest differences with only four sentences, larger models like BERT-base and GPT-2 show dramatic stratification, with deeper layers encoding semantic neighborhoods that align closely with human-legible categories. With a larger dataset, you would observe that sentences with similar topics, styles, or syntactic structures group together in final-layer activation space, even when those sentences share no words in common.