Kiru Lab  /  Mechanistic Interpretability  /  Causal Methods

Probing: What a Classifier on Activations Can and Cannot Tell You

Training a probe shows information is present. It does not show the model uses it.

Concept  ·  about 35 minutes

A probe is a small classifier — usually logistic regression, from track three — trained on a model's internal activations to predict some property. High probe accuracy shows that the property is linearly recoverable from that layer. That is a real and useful finding, and it is narrower than it is usually reported to be.

The core confound

Information being present is not the same as information being used. A probe can read a property out of a layer that no downstream component ever attends to. A sufficiently expressive probe can also extract structure that the model itself could not use — at which point you are measuring your probe, not the model.

  • Keep probes linear unless you have a specific reason not to. A deep probe learns the task itself.
  • Always include a control task: a probe trained on random labels establishes what accuracy your setup produces from nothing.
  • Report selectivity — real-task accuracy minus control-task accuracy — rather than raw accuracy.
  • Follow a probing result with a causal test. If the direction the probe found matters, intervening along it should change behavior.
From correlation to causation
from sklearn.linear_model import LogisticRegression
import numpy as np

probe = LogisticRegression(max_iter=2000).fit(activations, labels)
direction = probe.coef_[0] / np.linalg.norm(probe.coef_[0])

# The claim: this direction carries the property.
# The test: push activations along it and see whether the output moves.
def steer(activation, alpha):
    return activation + alpha * direction

# If large alpha does not change the model's behavior on the property,
# the probe found something readable that the model does not read.
The through-line

This is the same warning as the logistic coefficient in track three and the attention weight in track five, now stated in its final form. Every convenient number that looks like an explanation must be converted into a prediction and tested by intervention before it counts. That conversion is the discipline this entire curriculum exists to teach.

Hold on to

  • Probes establish availability of information, not use of it
  • Control tasks and selectivity are mandatory, not optional
  • A probing result becomes a claim only after a steering or patching test

Work through

Try each one before opening the solution. Getting it wrong first is most of where the learning happens.

  1. Train probes for one property at every layer and plot accuracy against depth. Add the random-label control on the same axes.
    Hint

    Train the same probe at every layer, and a second probe on shuffled labels as the control.

    Solution

    Real-task accuracy typically rises through the early-to-middle layers and plateaus or falls near the output. The control probe sits near chance if your setup is sound — if it does not, your probe is memorizing and every number above it is inflated. Plot them together and report the gap, not the raw accuracy; that gap is selectivity and it is the only part of the result that means anything.

    import numpy as np
from sklearn.linear_model import LogisticRegression

real, control = [], []
shuffled = np.random.default_rng(0).permutation(labels)

for layer in range(n_layers):
    A = activations_at(layer)                       # (n_samples, d_model)
    real.append(LogisticRegression(max_iter=2000).fit(A, labels).score(A_test, labels_test))
    control.append(LogisticRegression(max_iter=2000).fit(A, shuffled).score(A_test, shuffled_test))

import matplotlib.pyplot as plt
plt.plot(real, label="task"); plt.plot(control, "--", label="random labels")
plt.xlabel("layer"); plt.ylabel("probe accuracy"); plt.legend(); plt.show()
  2. Take your best probe direction and steer along it. Report whether behavior changed, and what you conclude either way.
    Hint

    Add alpha times the unit direction to the residual stream and sweep alpha.

    Solution

    Two outcomes and both are publishable. If behavior shifts systematically with alpha, the direction is used by downstream computation and the probe found something causal. If large alpha changes nothing while the probe still reads 95% accuracy, the information is present but unused — the probe found a readable correlate, not a mechanism. Report which one you got. The second result is the one most people quietly leave out, and it is the more informative of the two.

    import torch

def steer(model, ids, layer, direction, alpha):
    d = torch.as_tensor(direction, dtype=torch.float32)
    d = d / d.norm()
    handle = model.blocks[layer].register_forward_hook(
        lambda m, i, o: o + alpha * d)
    with torch.no_grad():
        out = model(ids)
    handle.remove()
    return out

for alpha in (-20, -5, 0, 5, 20):
    logits = steer(model, ids, LAYER, probe.coef_[0], alpha)
    print(alpha, measure_property(logits))

Sign in to track your progress through the lab.

← Ablation and Lesion Studies Back to track →