🤖 machine learning

Comparing Linear Probes with Mahalanobis Cosine Similarity

This paper theoretically proves and empirically validates that Mahalanobis cosine similarity to an out-of-distribution reference probe linearly predicts a linear probe's OOD AUROC under Gaussian assumptions, offering a superior alternative to standard cosine similarity for comparing probes.

Original authors: Zhuofan Josh Ying, Peter Hase, Nikolaus Kriegeskorte

Published 2026-06-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhuofan Josh Ying, Peter Hase, Nikolaus Kriegeskorte

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Why Do Some "Probes" Work and Others Don't?

Imagine you have a giant, complex machine (like a Large Language Model) that thinks in a language of numbers and patterns we can't easily see. Researchers use "linear probes" to peek inside this machine. Think of a probe as a flashlight beam. You shine it in a specific direction to see if the machine is "thinking" about a specific concept, like "Is this sentence true?" or "Is this person male?"

The problem is that these flashlights are fragile. A beam that works perfectly in one room (the training data) might go dark the moment you walk into a slightly different room (new data). Researchers want a way to predict if a flashlight will work in a new room just by looking at how it was built.

The Old Way vs. The New Way

To compare two flashlights (probes), researchers usually measure the angle between them using Euclidean Cosine Similarity.

The Analogy: Imagine you are comparing two compass needles. The old method just looks at the angle between the needles, treating every direction on the map as equally important.
The Flaw: This is like comparing compasses on a map where some directions are "sticky" (hard to move through) and others are "slippery" (easy to move through). If you ignore the sticky parts, your comparison is wrong.

The authors propose a new method called Mahalanobis Cosine Similarity (MCS).

The Analogy: This is like comparing compass needles while wearing special glasses that account for the "stickiness" of the terrain. It knows that some directions in the data are more important (where the data actually varies) and others are just noise. It weighs the comparison based on the shape of the data itself.

The Big Discovery: A Perfect Line

The authors found something surprising:
When they compared a probe trained on old data to a probe trained on new data using this new "terrain-aware" method (MCS), the result was a perfectly straight line.

The Analogy: Imagine you have a ruler (MCS) that measures how similar two flashlights are. If you use this ruler, you can predict exactly how well the flashlight will perform in the new room (its "Generalization Performance").
The Result: The paper shows that if the MCS score is high, the performance is high. If the MCS score is low, the performance is low. This relationship is so strong that it looks like a straight line drawn with a ruler (mathematically, an $R^2$ of 0.98, which is nearly perfect).
The Contrast: If they used the old "angle-only" method (Euclidean), the points would be scattered all over the place, like a messy cloud, making it impossible to predict performance.

Why Does This Happen? (The Magic of "S-Curves")

The paper explains why this straight line exists using a bit of math, but here is the simple version:

The Signal-to-Noise Ratio (SNR): Think of this as how loud the "truth" is compared to the "noise" in the data.
Two Different Shapes:
- The Performance (how well the probe works) follows an S-shape (a sigmoid curve). It starts slow, speeds up, and then levels off.
- The MCS Score also follows an S-shape, but it's a slightly different kind of curve.
The Cancellation: When you combine these two S-shapes, they cancel each other out! The "bend" in one curve perfectly straightens out the "bend" in the other.
- The Analogy: Imagine you are walking up a hill that curves to the left, but you are also walking on a conveyor belt that curves to the right at the exact same rate. The result is that you move in a perfectly straight line.

Because of this mathematical "cancellation," the relationship between the similarity score and the actual performance becomes a straight line.

When Does the Rule Break?

The authors also tested when this perfect line breaks, which helps us understand the limits of the rule:

Wrong Map (Wrong Covariance): If you use the old "angle-only" method (ignoring the terrain), the line disappears. You must use the "terrain-aware" map.
Bad Compass (Non-Fisher Probes): If the probe isn't built using the optimal mathematical method (like a simple "difference of means" instead of a more sophisticated one), the line gets messy.
Tiny Signal (Small Fisher Distance): If the signal is so weak that the "S-curves" haven't started to bend yet, the relationship isn't linear. (However, in the models they tested, the signal was strong enough that this didn't happen).
Unbalanced Teams: If one group of data is huge and the other is tiny (like 99% cats and 1% dogs), the line tilts and breaks.

Summary

The Problem: It's hard to tell if a tool (probe) will work on new data just by looking at it.
The Solution: Use Mahalanobis Cosine Similarity, which accounts for the shape and "stickiness" of the data.
The Result: This method creates a near-perfect straight line between the similarity score and the actual performance.
The Takeaway: You don't need to test the probe on new data to know if it will work. You just need to measure its similarity to a reference probe using this specific "terrain-aware" ruler.

The paper proves this mathematically and shows it works across different AI models, different layers inside those models, and many different types of tasks (like checking truthfulness or gender). It turns a messy, unpredictable problem into a clean, predictable line.

Technical Summary: Comparing Linear Probes with Mahalanobis Cosine Similarity

Problem Statement

Linear probes are standard tools in interpretability research for extracting concepts from neural network representations. However, their transfer behavior is notoriously brittle; probes trained on in-distribution (ID) data often fail to generalize to out-of-distribution (OOD) scenarios, even under subtle shifts like negation or prompt formatting changes. A critical challenge is determining which probes will generalize and understanding why probes that appear similar on one distribution diverge on another.

Current methods for comparing probes typically rely on Euclidean cosine similarity (ECS). However, ECS treats all dimensions equally, ignoring the fact that only dimensions where the data actually varies significantly impact probe performance. While Ying et al. (2026) empirically observed that the Mahalanobis cosine similarity (MCS)—which reweights the inner product by OOD data covariance—linearly predicts OOD AUROC ( $R^2=0.98$ ) for truthfulness datasets, this finding lacked a theoretical explanation and had not been validated across diverse models, layers, or concept domains.

Methodology

The authors propose a theoretical and empirical framework to explain and validate the relationship between MCS and OOD generalization performance.

Theoretical Framework

The paper derives a closed-form relationship between the OOD AUROC of a probe and its MCS to a reference probe trained on OOD data. The derivation relies on three standard statistical assumptions:

Balanced Classes: $Pr(Y=0) = Pr(Y=1) = 0.5$.
Projection Gaussianity: The projection of the feature vector onto any probe direction follows a Gaussian distribution for each class.
Covariance Structure: The total covariance ( $\Sigma_{tot}$ ) is decomposed into the pooled within-class covariance ( $\Sigma_{pool}$ ) and the between-class component.

Under these assumptions, the authors prove that:

AUROC is a sigmoid-shaped function of the probe's signal-to-noise ratio (SNR), $s$ : $AUROC = \Phi(s/\sqrt{2})$ .
MCS (specifically against the Fisher discriminant direction) is a soft-sign shaped function of $s$ : $MCS \approx s/\sqrt{4+s^2}$ (in the limit of large Fisher distance $z_{max}$ ).

The core theoretical insight is that composing these two functions cancels out their respective non-linearities (the sigmoid and the soft-sign), resulting in a near-linear relationship between AUROC and MCS. The slope of this line is theoretically predicted to be approximately $1/\sqrt{\pi}$ when the Fisher distance is large.

Empirical Evaluation

The authors extend the empirical validation of this relationship across:

Models: Llama-3.3-70B, Llama-3.1-8B, and Qwen2.5-7B.
Layers: Various layers (e.g., 20, 33, 50, 65) within Llama-70B.
Domains: 24 datasets spanning truthfulness, gender classification, and general NLP benchmarks.
Baselines: Comparison against standard Euclidean Cosine Similarity (ECS) and using different covariance estimators (total vs. pooled).

Key Contributions

Theoretical Proof of Linearity: The paper provides a closed-form proof demonstrating that for balanced, Gaussian-projected classes, OOD AUROC and MCS are linearly related. This explains the near-perfect empirical correlation ( $R^2 \approx 0.98$ ) previously observed.
Characterization of Failure Modes: The theory explicitly predicts conditions under which this linearity breaks down, which are verified empirically:
- Using pooled covariance ( $\Sigma_{pool}$ ) instead of total covariance ( $\Sigma_{tot}$ ) destroys linearity.
- Small Fisher distances ( $z_{max}$ ) result in task-specific slopes rather than a universal line.
- Class imbalance steepens the slope.
- Probes that deviate significantly from the Fisher direction (e.g., simple difference-of-means probes) show degraded linearity.
Empirical Generalization: The study validates the MCS-AUROC linearity across a wide range of models, layers, and concept domains, achieving $R^2 > 0.93$ in all tested conditions. In contrast, ECS performance drops significantly, with $R^2$ as low as 0.06 in general NLP benchmarks.
Robustness to Estimators: The authors demonstrate that the result holds for full sample covariance, Ledoit-Wolf shrinkage, but fails with diagonal approximations, confirming the necessity of off-diagonal covariance structure.

Results

High Predictive Power: Across all 8 experimental conditions (varying models, layers, and domains), the linear fit of OOD AUROC against MCS yields $R^2 \geq 0.93$ .
Superiority over ECS: ECS consistently underperforms, with $R^2$ values ranging from 0.057 (General NLP) to 0.723 (Gender), whereas MCS maintains high predictive power.
Theoretical Alignment: The empirical slopes of the AUROC-MCS curves align closely with the theoretical central slope of $1/\sqrt{\pi} \approx 0.564$ . Deviations occur primarily at the saturation tails (high AUROC/MCS values), consistent with the theory.
Failure Mode Verification: Simulations and real-data experiments confirm that linearity degrades when assumptions are violated (e.g., using $\Sigma_{pool}$ reduces $R^2$ from 0.98 to 0.83; using difference-of-means probes reduces it to 0.79).

Significance and Claims

The paper claims that Mahalanobis Cosine Similarity offers a theoretically grounded, task-aware alternative to Euclidean cosine similarity for comparing linear probes.

Mechanism: Unlike ECS, which treats all dimensions equally, MCS reweights the inner product by the OOD data covariance, effectively comparing probes in the subspace that actually matters for the specific task.
Predictive Utility: The linear relationship allows researchers to predict the generalization performance (OOD AUROC) of a probe trained on ID data simply by measuring its MCS to a reference probe trained on OOD data (or a proxy for the Fisher direction).
Broader Implications: The authors suggest that this approach may benefit other interpretability methods currently reliant on cosine similarity, such as steering-vector comparison, SAE feature alignment, and concept-direction clustering.
Limitations: The authors note that the method requires labeled OOD data to estimate the covariance and reference direction (it is not label-free for the target distribution). Furthermore, the theory is calibrated for Fisher-style probes (Logistic Regression, LDA); performance degrades for difference-of-means probes. The evaluation is currently confined to binary classification on residual-stream features of autoregressive LLMs.

In summary, the work establishes that the success of MCS in predicting probe generalization is not a coincidence but a consequence of the geometric interplay between signal-to-noise ratios, Gaussian projections, and covariance-aware similarity metrics.