Probing classifiers

Definition & Problem Setup

Definition Probing Classifier

A probing classifier (or diagnostic classifier) is a supervised model—typically a simple linear classifier or shallow MLP—trained to predict an auxiliary task $\tau$ from the frozen representations $\mathbf{h}$ of a pre-trained model $f_\theta$ .

Setup:

Pre-trained model $f_\theta: \mathcal{X} \to \mathcal{Y}$ produces intermediate representations $\mathbf{h}^{(l)}(x)$ at layer $l$ .
Auxiliary dataset $\mathcal{D}_{\text{probe}} = \{(x_i, y_i^{\tau})\}$ with labels for task $\tau$ (e.g., part-of-speech tags, syntactic structure, semantic role).
Train probe $g_\phi: \mathbf{h}^{(l)} \to \hat{y}^{\tau}$ with frozen $f_\theta$ . Measure probe accuracy $A_\tau^{(l)}$ . Key principle: High $A_\tau^{(l)}$ implies layer $l$ representations linearly encode information about $\tau$ . Low accuracy means the info is absent or requires complex nonlinear extraction.

Why "frozen" representations? If we fine-tuned $f_\theta$ , the model could learn to encode $\tau$ during probing—we'd measure the model's capacity to learn $\tau$ , not whether it already learned it during pre-training.

Derivation: What Does Probe Accuracy Tell Us?

Let's formalize the intuition. Suppose $\mathbf{h}^{(l)} \in \mathbb{R}^d$ are layer- $l$ activations, and we probe with a linear classifier:

\hat{y}^{\tau} = \sigma(\mathbf{w}^\top \mathbf{h}^{(l)} + b)

where $\sigma$ is softmax (multiclass) or sigmoid (binary), $\mathbf{w} \in \mathbb{R}^d$ , $b \in \mathbb{R}$ .

What does high probe accuracy mean?

Information theoretically: $\mathbf{h}^{(l)}$ contains mutual information $I(\mathbf{h}^{(l)}; Y^\tau) > 0$ where $Y^\tau$ is the true label distribution for task $\tau$ . High accuracy $\implies$ high MI.
Geometrically: The classes of $\tau$ are linearly separable in the representation space. The decision boundary $\mathbf{w}^\top \mathbf{h}^{(l)} + b = 0$ cleanly partitions the space.
Causally: The pre-training objective induced features that correlate with $\tau$ . This happens if $\tau$ is a useful intermediate step for the main task, or if the data structure forces $\tau$ to be encoded.

Derivation of the separability condition:

For perfect linear separability, we need a hyperplane such that:

\begin{cases} \mathbf{w}^\top \mathbf{h}_i^{(l)} + b > 0 & \text{if } y_i^\tau = +1 \\ \mathbf{w}^\top \mathbf{h}_i^{(l)} + b < 0 & \text{if } y_i^\tau = -1 \end{cases}

Define the margin $\gamma$ :

\gamma = \min_{i} \frac{y_i^\tau (\mathbf{w}^\top \mathbf{h}_i^{(l)} + b)}{\|\mathbf{w}\|}

If $\gamma > 0$ , classes are linearly separable. Probe training (e.g., via logistic loss) finds $\mathbf{w}, b$ that maximize margin-like quantities.

Why this step? The margin quantifies how "easy" the linear separation is. Large margin $\implies$ robust linear encoding. Small margin $\implies$ the probe barely works, suggesting info is nonlinearly encoded or noisy.

Methodology: How to Run a Probing Study

Formula Probing Protocol (Step-by-Step)

1. Choose auxiliary tasks $\{\tau_1, \tau_2, \ldots\}$ :

Linguistic: POS tags, dependency parsing, named entities, semantic roles.
Vision: object presence, spatial relations, texture vs. shape.
Goal: tasks that test hypotheses about what the model learned.

2. Extract representations: For each layer $l \in \{1, 2, \ldots, L\}$ and each sample $x \in \mathcal{D}_{\text{probe}}$ : $\mathbf{h}^{(l)}(x) = f_\theta^{(l)}(x)$ Store these frozen activations.

3. Train probe $g_\phi^{(l)}$ for each layer $l$ and task $\tau$ : $\min_{\phi} \frac{1}{N} \sum_{i=1}^{N} \mathcal{L}(\hat{y}_i^\tau, y_i^\tau) + \lambda R(\phi)$ where $\mathcal{L}$ is cross-entropy, $R(\phi)$ is L2 regularization (controls probe complexity).

Why regularization? Prevents the probe from memorizing idiosyncrasies. We want to test if the info is easily accessible, not if a complex probe can extract it.

4. Evaluate: Compute accuracy $A_\tau^{(l)}$ on held-out test set. Compare across layers.

5. Interpret:

Layer-wise trends: Early layers encode low-level features (syntax), deeper layers encode semantics.
Task comparisons: High $A_{\text{POS}}$ but low $A_{\text{semantics}}$ $\implies$ syntax is explicit, semantics are entangled.

Example calculation: Suppose we probe BERT for POS tags.

Layer 0(word embeddings): $A_{\text{POS}}^{(0)} = 82\%$ (word identity strongly correlates with POS).
Layer 6 (middle): $A_{\text{POS}}^{(6)} = 97\%$ (contextualized representations make POS trivial).
Layer 12 (output): $A_{\text{POS}}^{(12)} = 95\%$ (slight drop; layer focuses on downstream task, not syntax).

Why this pattern? Middle layers balance local syntax (needed for MLM) and global context. Output layers specialize for the pre-training objective, potentially discarding task-irrelevant syntax.

Worked Example1: Probing for Subject-Verb Agreement

Task: Predict whether subject and verb agree in number (singular/plural) from a sentence.

Setup:

Model: GPT-2 (12 layers).
Dataset: 10k sentences with labels {agree, disagree}.
Probe: Logistic regression on activations at the verb token position.

Step 1: Extract representations

For sentence "The cat sleps", extract $\mathbf{h}^{(l)}_{\text{sleps}}$ at layers $l = 0, 1, \ldots, 11$ .

Step 2: Train probe

For layer $l=6$ :

\min_{\mathbf{w}, b} -\frac{1}{N} \sum_{i=1}^{N} \left[ y_i \log \sigma(\mathbf{w}^\top \mathbf{h}_i^{(6)} + b) + (1 - y_i) \log(1 - \sigma(\mathbf{w}^\top \mathbf{h}_i^{(6)} + b)) \right] + 0.01 \|\mathbf{w}\|^2

Why this step? Cross-entropy loss measures how well the linear projection $\mathbf{w}^\top \mathbf{h}^{(6)}$ separates the classes. L2 penalty $0.01 \|\mathbf{w}\|^2$ ensures $\mathbf{w}$ stays small, testing if the info is easily linear.

Result: $A_{\text{agreement}}^{(6)} = 94\%$ .

Interpretation: Layer 6 representations contain a linear subspace that encodes subject-verb agreement. GPT-2 learned this syntactic rule during language modeling, likely because predicting the verb requires tracking subject number.

Worked Example 2: Probing for Semantic Role (Agent vs. Patient)

Task: In sentences like "The dog chased the cat", identify whether a noun is the agent (doer) or patient (receiver).

Setup:

Model: BERT.
Dataset: 5k sentences with semantic role labels.
Probe:2-layer MLP (hidden size 100) on noun token activations.

Why MLP instead of linear? Semantic roles might require nonlinear combinations of features (e.g., position + grammatical relation + verb semantics).

Step 1: Extract $\mathbf{h}^{(l)}_{\text{noun}}$

Step 2: Train MLP probe

\hat{y} = \text{softmax}(\mathbf{W}_2 \, \text{ReLU}(\mathbf{W}_1 \mathbf{h}^{(l)} + \mathbf{b}_1) + \mathbf{b}_2)

Results:

Layer 8 $A_{\text{role}}^{(8)} = 78\%$ (MLP), $A_{\text{role}}^{(8)} = 62\%$ (linear probe).
Layer 10: $A_{\text{role}}^{(10)} = 89\%$ (MLP), $A_{\text{role}}^{(10)} = 71\%$ (linear).

Interpretation: Semantic roles are present in layer 8 but require nonlinear extraction (MLP much better than linear). By layer 10, roles are more linearly accessible (gap narows). This suggests deeper layers "untangle" semantic structure.

Why this step? Comparing linear vs. nonlinear probes tells us how the info is encoded. Linearly separable $\implies$ explicit representation. Nonlinearly separable $\implies$ entangled, requires disentanglement.

Common Mistakes & Steel-Man Fixes

Diagram: Layer-Wise Probing Accuracy Across Tasks

Figure Explanation: The diagram shows how probe accuracy for three tasks (POS tags, dependency parsing, semantic similarity) varies across layers in a 12-layer Transformer. POS tags peak early (layer 4-6), syntax peaks mid-network (layer 6-8), and semantics emerge later (layer 8-12). This visualizes the representational hierarchy: lower layers encode surface form, middle layers encode structure, upper layers encode meaning.

Connections to Other Concepts

Representation Learning: Probing measures what representations encode.
Attention Mechanisms: Probing attention weights reveals which tokens influence decisions.
Fine-tuning vs. Feature Extraction: Probing is feature extraction—we use frozen features for a new task.
Causal Intervention Methods: Complement to probing; interventions test usage, probing tests presence.
Adversarial Robustness: Probing adversarial representations reveals which features are robust vs. spurious.
Disentangled Representations: High linear probe accuracy $\implies$ the feature is disentangled (occupies an independent subspace).

Hidden Feynman Recall Block

Recall Explain Like I'm 12

Imagine you teach a robot to translate English to French. The robot has a "brain" (neural network) with many layers, like floors in a building. You wonder: Does the robot understand grammar? Does it know which words are nouns vs. verbs?

You can't just open up the robot's brain and read it—it's too complicated. Instead, you play a game: You show the robot some sentences and ask, "Tell me the part-of-speech tag for each word." But here's the trick—you freeze the robot's brain (don't let it learn anything new) and only train a tiny helper (the "probe") to guess the tags based on what the robot already learned.

If the helper guesses correctly 95% of the time, the robot's brain already contains the part-of-speech info, stored in a way that's easy to extract. If the helper only guesses 60%, the info is either missing or hidden in a complicated way.

Why does this matter? It tells us what the robot learned and where (which layer). Maybe the first few layers know spelling, the middle layers know grammar, and the last layers know meaning. This helps us trust the robot and fix it when it makes mistakes.

Mnemonic

Active Recall Flashcards

#flashcards/ai-ml

What is a probing classifier? :: A simple supervised model (e.g., logistic regression) trained on frozen pre-trained representations to predict an auxiliary task, measuring whether the representation linearly encodes task-relevant information.

Why do we freeze the pre-trained model during probing?

To test what the model already learned during pre-training, not what it can learn. Unfreezing would let the model adapt to encode the probe task, confounding the measurement.

What does high probe accuracy tell us?

The representation contains information about the auxiliary task in a linearly accessible form (geometrically: classes are linearly separable; information-theoretically: high mutual information).

Linear probe gets 65% accuracy, MLP probe gets 88%. What does this imply?

The information is present but encoded in an entangled, nonlinear way. The representation requires complex transformations to extract the feature.

What is the margin in linear separability?

\gamma = \min_{i} \frac{y_i (\mathbf{w}^\top \mathbf{h}_i + b)}{\|\mathbf{w}\|}

, the minimum distance from any point to the decision boundary. Positive margin implies linear separability.

Why regularize the probe during training?

To prevent overfitting and ensure we test easy accessibility. We want to know if the info is explicitly encoded, not if a complex probe can memorize the dataset.

Probing shows layer6 encodes POS tags at 95% accuracy. Does the model use POS tags for its task?

Not necessarily. Probing measures presence, not usage. The model might encode POS tags as a byproduct but not causally depend on them. Test with causal interventions.

What is the typical layer-wise pattern for linguistic probes in Transformers?

Early layers encode surface features (POS, morphology), middle layers encode syntax (dependencies, phrase structure), late layers encode semantics (entailment, paraphrase). This is the representational hierarchy.

How do you establish a baseline for probe accuracy?

Probe random representations (e.g., Gaussian noise) or shuffled labels. If the model's probe accuracy doesn't exceed this baseline, the feature is not encoded.

What is the difference between probing and causal intervention?

Probing measures if a feature is present in representations. Causal intervention (e.g., ablation) tests if the feature is used for the task. Both are needed for full interpretability.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Probing classifiers ka matlab kya hai? Jab ap kisi bade neural network ko train karte hain (jaise BERT ya GPT), toh woh andar se representations banata hai—matlab har layer mein activations jo model ne seekha hai. Ab sawal yeh hai: Kya model ne grammar seekha? POS tags? Syntax trees? Weights dekhne se pata nahi chalega, bahut complex hain. Toh hum ek trick karte hain: Probing. Hum model ko freeze kar dete hain (koi nayi chez mat seekho), aur uske representations parek chhota sa classifier train karte hain jo specific property predict kare—jaise part-of-speech tags ya dependency relations. Agar yeh chhota classifier achi accuracy deta hai (jaise 95%), toh iska matlab hai ki model ke representations mein woh information already encoded hai, aur woh bhi aise form mein jo easily extract ho sakti hai (linearly separable). Low accuracy matlab ya toh woh information hai hi nahi, ya bahut entangled form mein hai.

Kyun important hai yeh? Interpretability ke liye.Agar ap jaan lo ki model ne kya seekha (syntax, semantics, factual knowledge), kahan seekha (kaun se layers mein), aur kitna easily accessible hai (linear ya nonlinear), toh ap model par trust kar sakte ho, uske mistakes debug kar sakte ho, aur samajh sakte ho ki downstream tasks ke liye kaunsa layer best hoga. Probing essentially ek diagnostic tool hai—jaise doctor test karke dekhta hai ki patient ke body mein kya chal raha hai, waise hi hum model ke andar dekhte hain.

Medical diagnosis jaise example lo: CT scan se doctor ko pata chalta hai ki kaun se organs mein problem hai, lekin woh treatment nahi hai—sirf diagnosis hai. Similarly, probing se pata chalta hai model ne kya encode kiya, lekin yeh nahi bata ki model us info ka use kaise kar raha hai. Usage test karne ke liye causal interventions chahiye (jaise feature ko hata do aur dekho task performance girti hai ya nahi). Toh probing aur causal methods dono milke complete picture dete hain.

Definition & Problem Setup

Derivation: What Does Probe Accuracy Tell Us?

Methodology: How to Run a Probing Study

Worked Example1: Probing for Subject-Verb Agreement

Worked Example 2: Probing for Semantic Role (Agent vs. Patient)

Common Mistakes & Steel-Man Fixes

Diagram: Layer-Wise Probing Accuracy Across Tasks

Connections to Other Concepts

Hidden Feynman Recall Block

Mnemonic

Active Recall Flashcards

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Interpretability & Explainability

Connections