Probing Internal Representations

Visualizing attention patterns provides a way to understand a model's focus, but this doesn't directly tell us what kind of linguistic or semantic information is encoded within the high-dimensional hidden state vectors produced by each layer. These vectors, often having thousands of dimensions, are the internal currency of the Transformer, carrying information from one layer to the next. To understand what these representations capture, we turn to a technique known as probing.

Probing involves training simple, auxiliary models, called probes, to predict specific properties of interest directly from the LLM's internal representations. The core idea is: if a simple probe can accurately predict a property (like part-of-speech tags or dependency relations) using only the hidden state vector from a particular layer as input, then that information is likely explicitly encoded or at least linearly separable within that representation. We are less interested in building the best possible predictor for the property itself; rather, we use the probe's performance as a diagnostic tool for the LLM's representation quality.

The Probing Methodology

The typical workflow for probing involves several steps:

1. Define the Task and Obtain Data: Select a property you want to investigate (e.g., identifying the grammatical number of nouns, predicting syntactic heads). You need a dataset where text inputs are annotated with these properties. For linguistic properties, resources like Universal Dependencies or the Penn Treebank are often used.
2. Extract Representations: Process the annotated dataset using the pre-trained LLM you want to analyze. Crucially, the LLM's weights are kept frozen during this process. For each token in the input sequences, extract the hidden state vector from the specific layer(s) you are interested in probing.
3. Train the Probe: Train a simple classifier or regressor using the extracted hidden states as input features and the annotated properties as target labels. Common choices for probes include linear classifiers (logistic regression, linear SVM) or very shallow multi-layer perceptrons (MLPs). The simplicity is important: a complex probe might learn the task independently rather than just "reading out" information already present in the representation.
4. Evaluate the Probe: Assess the probe's performance on a held-out test set using relevant metrics (e.g., accuracy for classification, mean squared error for regression).
5. Interpret Results: High performance suggests the representation at that layer contains easily accessible information about the probed property. Comparing performance across different layers can reveal how information is transformed or refined as it passes through the model. Comparing against baseline probes (e.g., trained on random vectors or static word embeddings) helps establish the significance of the results.

Example: Probing for Part-of-Speech Tags

Let's consider probing a pre-trained Transformer model (like BERT or a GPT variant) for Part-of-Speech (POS) information.

1. Data: We need a corpus annotated with POS tags (e.g., the Universal Dependencies English Web Treebank). 2. Representations: We feed sentences from the corpus into our frozen LLM and collect the hidden state vectors for each token from, say, layers 6, 12, and 18. 3. Probe: We choose a simple linear classifier. 4. Training & Evaluation: For each layer, we train a separate linear probe to predict the POS tag for each token based solely on its hidden state vector from that layer.

Here's a PyTorch snippet illustrating representation extraction and probe definition:

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer

# Load pre-trained model and tokenizer
model_name = "bert-base-uncased" # Or any other model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
model.eval() # Set model to evaluation mode
# Freeze model parameters
for param in model.parameters():
    param.requires_grad = False

# Example sentence and POS tags (replace with actual dataset loading)
sentence = "Probing helps analyze model representations."
# Assume tags are ['NOUN', 'VERB', 'VERB', 'NOUN', 'NOUN', 'PUNCT']
# In practice, align tokenization with tags carefully

inputs = tokenizer(sentence, return_tensors="pt")

# Extract hidden states from a specific layer (e.g., layer 8)
target_layer = 8
with torch.no_grad():
    outputs = model(**inputs, output_hidden_states=True)
    # hidden_states is a tuple: (embedding_layer, layer_1, ..., layer_N)
    layer_representations = outputs.hidden_states[target_layer]
    # Shape: [batch_size, sequence_length, hidden_size]

# Assume we have extracted representations and corresponding
# POS tag IDs for many examples
# representations_tensor: [num_examples, hidden_size]
# labels_tensor: [num_examples]

# Define a simple linear probe
hidden_size = layer_representations.shape[-1]
num_pos_tags = 17 # Example number of unique POS tags in UD EWT
probe_classifier = nn.Linear(hidden_size, num_pos_tags)

# --- Training Loop ---
# Standard PyTorch training loop here:
# - Define loss (e.g., CrossEntropyLoss)
# - Define optimizer (e.g., AdamW, only optimizing
#   probe_classifier.parameters())
# - Iterate over batches of (representations_tensor, labels_tensor)
# - Calculate loss, backpropagate, update probe weights
# - Evaluate on a validation set
# --------------------

# After training, evaluate probe_classifier on a test set of representations.

Interpreting Probe Results

The results of probing experiments can be quite revealing:

• Performance Magnitude: High accuracy (e.g., >90% POS tagging accuracy with a linear probe) strongly indicates that the specific layer encodes this information in a readily accessible format.
• Layer-wise Comparison: Often, performance on syntactic tasks like POS tagging or dependency parsing peaks in the middle layers of a Transformer, while tasks requiring more semantic integration might peak in later layers. This suggests a progression from lower-level linguistic features to higher-level abstractions.
• Control Tasks: It's good practice to run control tasks. For instance, train a probe to predict random labels assigned to the tokens. If this probe performs significantly above chance, it might indicate biases or issues with the experimental setup. Another control is selectivity: ensuring the probe trained for POS tagging performs poorly on, say, predicting dependency relations, and vice-versa. This confirms the probe is specifically leveraging the target information.

Comparison showing POS tagging accuracy peaking earlier than dependency relation accuracy across model layers.

Scope of Probing Tasks

Probing can be applied to a wide range of linguistic and semantic phenomena:

• Morpho-syntax: Tense, number, case, POS tags.
• Syntax: Syntactic constituent boundaries, dependency relations, grammatical role identification.
• Semantics: Semantic role labeling (SRL), named entity recognition (NER), coreference resolution, relation extraction.
• Knowledge: While harder to probe linearly, some studies attempt to probe for factual knowledge embedded within representations.

Limitations and Challenges

Probing is a powerful analysis tool, but it's important to be aware of its limitations:

• Probe Simplicity vs. Capacity: The probe must be simple enough that we can reasonably attribute high performance to the representation itself, not the probe's learning capacity. There's always a trade-off; a slightly more complex probe might reveal information not linearly separable but still easily extractable.
• Correlation, Not Causation: Finding that a layer encodes POS information doesn't necessarily mean the model uses this information explicitly for its primary objective (e.g., next-token prediction) in the way we might expect. It indicates the information is present and accessible.
• Data Dependency: The quality and nature of the annotated dataset used for probing are significant.
• Tokenization Alignment: Careful alignment between the LLM's tokenization and the linguistic annotations (which might be word-based) is needed.

Despite these points, probing offers valuable insights into the internal knowledge structures learned by large language models. By systematically examining how different types of information are represented across layers, we gain a better understanding of how these models process language, which can inform debugging, model improvement, and efforts to build more reliable and interpretable systems.