Gen AI

Probing Classifiers

How do we know if a language model actually learned grammar, or if it just memorized text? We train a tiny 'probe' model on its hidden states to see if we can extract grammatical concepts.

A probing classifier tests if a large model has secretly learned a concept by trying to extract that concept from its hidden states.

Why Does This Exist?

When you train a Large Language Model (LLM) to predict the next word, it gets incredibly good at it. But how is it predicting the next word?

Did the model blindly memorize 50 million sentences from Wikipedia? Or did it actually learn the underlying rules of human language—things like nouns, verbs, plurality, and sentence structure?

Because the inside of an LLM is a giant matrix of billions of unreadable numbers (the hidden states), you cannot simply "look" to see if a concept like "Grammar" is stored in there. Mechanistic Interpretability is the field of trying to reverse-engineer the brain of the AI to prove what it actually learned. Probing Classifiers are one of the oldest and most fundamental tools in this field. They exist to chemically test the hidden states for specific human concepts.

Think of It Like This

Think of an LLM's hidden state like a highly compressed zip file.

You know the zip file contains a massive amount of data, but you can't read it directly. You want to know if there is a picture of a dog inside it.

A probing classifier is like a tiny, specialized search script you write. You run the script on the zip file. If the script easily finds a picture of a dog, you have proven that the zip file contains dog data. If the script fails, the zip file probably doesn't contain dog data.

How It Actually Works

Probing is an experiment in two parts. You freeze the massive LLM, and you train a tiny new model on top of it.

1. Extract the Hidden States

First, you define the concept you want to test. Let's say you want to prove the LLM learned Parts of Speech (POS). You get a dataset of sentences where every word is labeled by a human (e.g., "The [Determiner] dog [Noun] ran [Verb]").

You feed these sentences into the LLM. You do not care what the LLM predicts. Instead, you pause the LLM at the middle layer (e.g., Layer 12) and extract the hidden state vector for the word "dog." This vector is just a list of 4,096 floating-point numbers.

2. Train the Probe

You now have a dataset of hidden state vectors paired with human labels (e.g., [0.4, -1.2, 3.1...] -> Noun).

You train a Probing Classifier on this data. This classifier must be extremely simple—usually just a single-layer linear regression. Why? Because if you use a massive, complex neural network for your probe, the probe itself might learn the concept of grammar from scratch! You want the probe to be stupid. It should only succeed if the grammar concept is already perfectly organized and clearly exposed inside the LLM's hidden state.

3. Analyze the Results

If your tiny, stupid linear probe achieves 99% accuracy in predicting Parts of Speech using the LLM's vectors, you have achieved your goal. You have mathematically proven that the LLM internally organized its data by grammar, even though it was never explicitly told to do so.

Show Me the Code

Here is a conceptual pipeline using PyTorch to extract a hidden state and train a linear probe.

1import torch2import torch.nn as nn3from sklearn.linear_model import LogisticRegression4
5# 1. Freeze the massive LLM6llm.eval()7for param in llm.parameters():8    param.requires_grad = False9
10hidden_states = []11labels = []12
13# 2. Extract hidden states for your specific dataset14for text, pos_label in dataset:15    # Run the text through the LLM16    outputs = llm(text, output_hidden_states=True)17    18    # Extract the vector from the middle layer (e.g., layer 12)19    # for the specific token we care about20    vector = outputs.hidden_states[12][0, target_token_idx, :].numpy()21    22    hidden_states.append(vector)23    labels.append(pos_label)24
25# 3. Train a tiny, dumb linear probe26probe = LogisticRegression(max_iter=1000)27probe.fit(hidden_states, labels)28
29# 4. Check accuracy30accuracy = probe.score(hidden_states, labels)31print(f"Probe Accuracy: {accuracy * 100:.2f}%")32# If accuracy is very high, the LLM learned Parts of Speech!

Watch Out For

The Probe Memorization Problem

The biggest debate in probing is: Did the probe find the concept, or did the probe learn the concept? If you make your probing classifier too complex (e.g., a multi-layer deep network), it might just memorize the dataset you gave it. The LLM might not know anything about grammar, but your deep probe figured it out anyway. To solve this, always use the simplest possible probe (a linear classifier) and rely heavily on holdout test sets.

The Quick Version

Mechanistic interpretability tries to prove what an LLM learned internally, rather than just looking at its outputs.
A probing classifier is a tiny, simple model trained to predict a specific human concept (like grammar or sentiment) using the LLM's internal hidden states.
If a simple linear probe achieves high accuracy, it proves that the LLM already did the hard work of learning and organizing that concept internally.
Probes must be kept as simple as possible to ensure they are merely "reading" the LLM's knowledge, not "learning" the knowledge themselves.