Foundation Model Engineering

19.3 Probing Classifiers

The logit lens asks what a hidden state would predict through the model’s unembedding. A probe asks a different question: can a specified property be decoded from that state by a constrained auxiliary model?

That distinction sets the epistemic boundary for this chapter. A successful probe establishes decodability under an experimental protocol. It does not, by itself, show that the base model uses the information, stores it in a unique direction, or would expose the same feature under distribution shift.

1. The Measurement Object

For input xix_i, choose a model location and extract activation hi(l,p)∈Rdh_i^{(l,p)}\in\mathbb R^d, where ll is the layer and pp identifies the token position and stream location. Train a probe fθf_\theta to predict label yiy_i:

y^i=arg⁡max⁡c(Whi(l,p)+b)c.\hat y_i=\arg\max_c \left(Wh_i^{(l,p)}+b\right)_c.

A result is not reproducible unless the activation artifact records:

  • model checkpoint and tokenizer hashes, chat template, and input text;
  • layer, module, pre/post-normalization location, and token-position rule;
  • dtype, pooling, truncation and padding behavior;
  • dataset example ID, split, label provenance, and extraction code version.

“Layer 20 activation” is underspecified: residual-stream input, attention output, MLP output, and post-block residual are different tensors.

2. Why Use a Linear Probe?

A high-capacity nonlinear probe can create a useful decision surface even when the target is awkwardly represented. A linear probe deliberately asks the narrower question: is the label accessible by an affine boundary? The Linear Representation Hypothesis studies when high-level properties correspond to linear structure in representation space [1].

Low capacity improves interpretability but does not eliminate confounds. In a high-dimensional space, a linear classifier can memorize limited data, exploit lexical identity, or use a feature correlated with the label. Probe capacity, regularization, sample size, and representation dimension must be reported together.

3. A Minimal, Auditable Experiment

Split before fitting transformations. Related prompts, entities, paraphrases, and synthetic siblings belong in one split; otherwise the probe can recognize examples rather than generalize a property. Compute normalization statistics on training data only.

Scope of the code. This is executable for pre-extracted, fixed-length activations with class IDs 0,…,C−10,\ldots,C-1. Activation hooks, sequence pooling, class imbalance, hyperparameter selection, confidence intervals, and grouped splitting are experiment-specific and intentionally kept outside the helper.

import torch
import torch.nn as nn


class LinearProbe(nn.Module):
    def __init__(self, width: int, classes: int):
        super().__init__()
        self.readout = nn.Linear(width, classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.readout(x)


def fit_linear_probe(
    train_x: torch.Tensor,
    train_y: torch.Tensor,
    test_x: torch.Tensor,
    test_y: torch.Tensor,
    *,
    steps: int = 500,
    learning_rate: float = 3e-3,
    weight_decay: float = 1e-2,
) -> tuple[LinearProbe, float]:
    if train_x.ndim != 2 or test_x.shape[1] != train_x.shape[1]:
        raise ValueError("expected [examples, hidden_width] with matching widths")
    classes = int(train_y.max().item()) + 1
    if classes < 2:
        raise ValueError("training split must contain at least two classes")

    # Fit preprocessing on train only; keep it with the probe artifact.
    mean = train_x.mean(dim=0, keepdim=True)
    scale = train_x.std(dim=0, keepdim=True).clamp_min(1e-6)
    x_train = (train_x - mean) / scale
    x_test = (test_x - mean) / scale

    torch.manual_seed(0)
    probe = LinearProbe(train_x.shape[1], classes).to(train_x.device)
    optimizer = torch.optim.AdamW(
        probe.parameters(), lr=learning_rate, weight_decay=weight_decay
    )
    for _ in range(steps):
        optimizer.zero_grad(set_to_none=True)
        loss = nn.functional.cross_entropy(probe(x_train), train_y)
        loss.backward()
        optimizer.step()

    probe.eval()
    with torch.no_grad():
        accuracy = (probe(x_test).argmax(dim=-1) == test_y).float().mean().item()
    return probe, accuracy

In a real study, select hyperparameters on validation data, run multiple seeds, report uncertainty, and compare against a majority baseline, token or embedding baselines, and randomly initialized or shuffled-layer representations.

4. Controls and Selectivity

Hewitt and Liang’s control task assigns a random but deterministic label to each word type [2]. A probe that predicts both the linguistic task and this memorization-friendly control may be exploiting lexical identity or excess capacity. They define

selectivity=task performance−control performance.\text{selectivity}= \text{task performance}-\text{control performance}.

High selectivity is useful evidence that the chosen probe/data protocol favors the real structure over that control. It does not prove that the geometry “naturally represents” the human concept. The conclusion is conditional on the control construction, split, capacity measure, metric, and distribution.

A robust probe card includes:

  • label balance and grouped train/dev/test split;
  • layers and token positions chosen before inspecting the test set;
  • probe family, parameter count, regularization, seeds, and preprocessing;
  • in-distribution, paraphrase, entity-held-out, and template-held-out results;
  • control-task performance and simple representation baselines;
  • confidence intervals and all attempted hypotheses, not only the best layer.

5. From Decodability to Intervention

Decodability is correlational. Amnesic probing asks whether downstream behavior changes after removing probe-accessible information [3]. For a probe matrix WW, a single linear null-space projection can be written

P=I−W⊤(WW⊤)†W,h′=Ph,P=I-W^\top(WW^\top)^\dagger W,\qquad h'=Ph,

where †{}^\dagger is the Moore–Penrose pseudoinverse. INLP repeats probe fitting and projection until the selected linear probe family can no longer recover the label.

The careful interpretation is “the intervention reduced information recoverable by this probe family.” It did not erase every nonlinear encoding of the concept, and it may have removed unrelated information that shared the same subspace. Later work reports that INLP can introduce non-targeted changes and finds Mean Projection or LEACE more targeted in evaluated settings [5].

A causal claim needs controls:

  1. Verify completeness: can fresh probes still decode the target?
  2. Verify intervention selectivity: are matched non-target properties retained?
  3. Compare matched-rank random projections and magnitude-matched noise.
  4. Measure the predeclared downstream behavior, not only probe accuracy.
  5. Repeat across layers, seeds, removal methods, and held-out distributions.

Canby et al. formalize an empirical completeness–selectivity trade-off and report that no intervention family is perfect across their tests [6]. A downstream change after projection is therefore evidence for a causal contribution only to the extent that collateral damage has been ruled out.

6. Unsupervised Probing and CCS

Contrast-Consistent Search (CCS) uses pairs such as a statement and its negation, seeking a low-capacity direction whose predictions are confident and mutually consistent without using task labels [4]. A simplified consistency condition is

pθ(x)+pθ(¬x)≈1.p_\theta(x)+p_\theta(\neg x)\approx1.

This is an assumption and a training signal, not ground truth. Natural-language negation may not form a clean logical complement; both members can be ambiguous or false under different readings. The objective has a sign symmetry, so a learned axis still needs orientation, and label-free training does not make model selection label-free.

The CCS paper reports an average improvement over zero-shot accuracy across its evaluated datasets and models. It presents an initial method for eliciting latent knowledge—not a universal truth direction and not a detector that establishes whether a model is “lying.”

7. Product Debugging Without Overclaiming

Probes can be valuable instrumentation:

  • Over-refusal: test whether a risk-related signal separates benign security education from harmful requests across layers.
  • RAG grounding: ask whether answer-time states predict document-supported versus parametric answers on entity- and template-held-out data.
  • Tool readiness: compare pre-call states from successful and failed traces to find where action selection becomes decodable.

These are hypotheses, not production gates. Start in shadow mode. Calibrate false-positive and false-negative costs by slice, monitor drift after every model/template update, and require an output- or behavior-level fallback. Activation geometry and probe calibration can change when the base artifact changes.

Linear separability in a 2-D toy representation

Synthetic points, not model activations. This illustrates what a linear probe can measure; it is not evidence for a particular model.

Fixed linear decision boundary
Class A
Class B
Fixed-boundary accuracy: 92.5%. Increasing class signal makes the labels more linearly decodable. Distribution shift can still reduce test accuracy without changing the probe.

The interactive uses synthetic two-dimensional points. It demonstrates how a fixed linear boundary responds to signal and distribution shift; it does not simulate layer evolution or validate the LRH.

8. Takeaway

A probe is a measurement instrument whose validity comes from its controls. Linear accuracy supports a claim about decodability; selectivity narrows alternative explanations; a controlled intervention can support a causal contribution. None of these, alone, reveals a unique human-readable feature. The next section turns to sparse autoencoders, which replace a predeclared label with a learned feature dictionary—and introduce a different set of identifiability and evaluation problems.

Quizzes

Quiz 1: What is the strongest conclusion justified by high held-out linear-probe accuracy alone? Under the stated extraction, split, preprocessing, and probe family, the label is linearly decodable from those activations. Model use and causality require additional evidence.

Quiz 2: Why must normalization be fitted on the training split only? Using test statistics leaks information about the evaluation distribution into the probe pipeline and can bias the estimate.

Quiz 3: Why does a behavior drop after INLP not automatically prove targeted causal dependence? The projection may remove non-target information or move activations off distribution. Completeness, collateral-feature retention, and matched interventions are needed.

Quiz 4: What ambiguity remains after CCS learns a consistent statement–negation axis? Consistency does not identify which end is true, and natural-language pairs may violate the assumed complement relation. Orientation and external evaluation remain necessary.

Quiz 5: A probe performs well on random example splits but fails on entity-held-out data. What is the likely lesson? The random split allowed lexical or entity-specific shortcuts. The evidence supports memorization or narrow interpolation, not a general property representation.

References

  1. Park, K., et al. (2023). The Linear Representation Hypothesis and the Geometry of Large Language Models. arXiv:2311.03658.
  2. Hewitt, J., & Liang, P. (2019). Designing and Interpreting Probes with Control Tasks. ACL Anthology.
  3. Elazar, Y., et al. (2021). Amnesic Probing: Behavioral Explanation with Amnesic Counterfactuals. ACL Anthology.
  4. Burns, C., et al. (2022). Discovering Latent Knowledge in Language Models Without Supervision. arXiv:2212.03827.
  5. Dobrzeniecka, A., Fokkens, A., & Sommerauer, P. (2025). Improving Causal Interventions in Amnesic Probing with Mean Projection or LEACE. ACL Anthology.
  6. Canby, M. E., Davies, A., Rastogi, C., & Hockenmaier, J. (2025). How Reliable are Causal Probing Interventions? ACL Anthology.