Linear Probes

The problem, from first principles

Does a model's hidden state contain some property, say whether a sentence is in the past tense or whether a number is even? A linear probe answers by training the simplest possible classifier, a logistic regression, on frozen activations. If a linear map can read the property off the hidden state, the model represents it in an easily accessible way. A probe is only as informative as it is simple: a powerful probe could learn the property itself rather than read it, so interpretability work keeps probes linear and compares against a control. You will build the probe from scratch with gradient descent and an L2 penalty.

From theory to code

Implement train_linear_probe and probe_accuracy.

Constraints

• X has shape (n, d) and y holds 0/1 labels. The probe is p = sigmoid(X @ w + b).
• train_linear_probe(X, y, steps, lr, l2) starts from w = 0 and b = 0 and performs steps full-batch gradient-descent updates on the mean cross-entropy plus 0.5 * l2 * ||w||^2 (the bias is not penalized). Return (w, b).
• The gradient of the data term is X.T @ (p - y) / n for w and mean(p - y) for b. The penalty adds l2 * w to the weight gradient.
• probe_accuracy(X, y, w, b) is the fraction of examples where (X @ w + b > 0) equals y.
• Use a stable sigmoid.

Hints