Linear Classifier Probes

A from-scratch implementation of the linear probing technique from Alain & Bengio (2016), applied to GPT-2 using TransformerLens. Linear probes reveal what information each layer of a neural network has learned to encode by training simple linear classifiers on frozen intermediate representations.

Background

Alain & Bengio (2016) proposed attaching independent linear classifiers ("probes") to the output of each layer in a deep network. The key insight: if a linear classifier achieves high accuracy on a task using only a given layer's representations, then that layer must encode the task-relevant features in a linearly separable way. By measuring probe accuracy across all layers, the technique reveals how information emerges and is refined through the depth of the network.

The original paper observed that linear separability increases monotonically with depth — early layers produce noisy, entangled representations, while later layers progressively disentangle features into linearly decodable formats.

How It Works

At each layer l, a linear probe computes:

P(y = 1 | h_l) = sigmoid(w^T h_l + b)

where h_l is the frozen residual stream activation at layer l, and w, b are the probe's learned parameters. The probe is trained with binary cross-entropy loss via Adam, entirely independently of the base model. Only the probe parameters are updated — the model's weights stay frozen.

This is the simplest possible classifier. If it works, the representation itself must contain the signal; the probe cannot manufacture features that aren't already there.

Implementation

This project probes GPT-2 (12 transformer layers, d_model=768) for binary sentiment classification:

1. Dataset construction: 40 positive and 40 negative sentiment words are each embedded into 5 carrier sentence templates (e.g., "The movie was absolutely {wonderful/terrible}."), producing 400 labeled examples.

2. Representation extraction: Each sentence is tokenized and run through GPT-2 with TransformerLens caching. The residual stream vector at the sentiment word's token position is extracted from all 13 points (embedding + 12 post-block layers).

3. Word-level train/test split: The split is done by unique words (75/25), not by sentences. This forces probes to generalize to sentiment words unseen during training, preventing memorization of word-template combinations.

4. Probe training: At each layer, an independent nn.Linear(768, 1) probe is trained for 200 epochs with Adam (lr=0.01) on the frozen representations.

5. Visualization: A layer-vs-accuracy plot shows how sentiment information emerges across depth.

Usage

uv sync
uv run python main.py

Output

Dataset: 400 sentences (200 positive, 200 negative)
Loading GPT-2...
Extracting representations at every layer...
Split: 300 train, 100 test (by word, not sentence)
  Embed:  train_acc=0.xxx  test_acc=0.xxx
  Layer  1:  train_acc=0.xxx  test_acc=0.xxx
  ...
  Layer 12:  train_acc=0.xxx  test_acc=0.xxx

Best layer: XX (test_acc=0.xxx)

Generates probe_accuracy.png — a plot of probe accuracy across layers with training loss curves.

Dependencies

• Python 3.12+
• PyTorch
• TransformerLens
• matplotlib

Key Concepts Demonstrated

• Linear probes as interpretability tools: measuring what information is linearly decodable at each layer
• Frozen representations: the base model is never updated — only the probe parameters change
• Linear separability vs. depth: confirming that deeper layers produce more linearly separable representations
• Proper evaluation: word-level train/test splits to prevent data leakage from shared sentence templates
• Connection to mechanistic interpretability: linear probes complement techniques like the logit lens and tuned lens by testing specific hypotheses about what each layer encodes

Reference

Alain, G. & Bengio, Y. (2016). Understanding intermediate layers using linear classifier probes. arXiv:1610.01644.