Source-linked AI summary
Comparing Linear Probes with Mahalanobis Cosine Similarity
Zhuofan Josh Ying, Peter Hase, Nikolaus Kriegeskorte
TL;DR
Linear probes can behave unreliably under distribution shift, and prior MCS evidence was limited to one model and truthfulness datasets. This paper combines broad empirical evaluation with a closed-form theory of MCS–AUROC linearity, finding strong cross-setting agreement and identifying conditions where it fails.
Problem
Linear probes often degrade under distribution shift, while prior MCS–OOD AUROC evidence was limited to one model and truthfulness datasets.
Method
The paper evaluates MCS across models, layers, and concept domains and derives its relationship with OOD AUROC under balanced classes and Gaussian projections.
Results
R^2>0.93 holds across all tested models, layers, and datasets, while ECS drops to R^2=0.06 in the worst condition.
Takeaways & Limitations
MCS is a theoretically grounded, task-aware alternative to Euclidean cosine similarity for comparing linear probes.
Takeaways & Limitations
The evaluation is limited to binary classification on residual-stream features of autoregressive LLMs, and formal guarantees assume per-class projection Gaussianity.
Abstract
from arXiv · showhide
Linear probes are widely used in interpretability research and often compared by cosine similarity. The Mahalanobis cosine similarity (MCS) between two directions, which reweights the inner product by test data covariance, is a natural task-aware refinement. Ying et al. (2026) report that a probe's MCS to a reference probe trained on the out-of-distribution (OOD) data near-perfectly linearly predicts the probe's OOD AUROC (R^2 = 0.98). Here, we extend this empirical finding across models, layers, and concept domains, and prove this general phenomenon in closed form: For balanced classes whose projections are Gaussian, OOD AUROC and MCS to the reference probe are linear because both are sigmoid-shaped functions of the probe's signal-to-noise ratio (SNR) on the test data. The theory also predicts when this linearity fails, which we verify empirically. MCS offers a theoretically grounded and empirically effective alternative to Euclidean cosine similarity for comparing linear probes.
1 Introduction
Linear probes can generalize poorly under distribution shift, motivating task-aware comparison. This paper extends MCS evidence across settings and explains its near-linearity with OOD AUROC theoretically.
- Linear probes often fail to generalize across closely related datasets and subtle shifts such as negation or prompt-format changes.
- Prior work found MCS between ID and OOD probes linearly related to OOD AUROC, with R^2=0.98, but only on one model and truthfulness datasets.
- R^2>0.93 holds across three models, layers 20–65, and 24 datasets spanning truthfulness, gender classification, and general NLP benchmarks; ECS falls to R^2=0.06.
- Under balanced classes and Gaussian projections, OOD AUROC and MCS are S-shaped functions of the probe’s OOD signal-to-noise ratio.
- Composing these functions yields near-linearity, while small Fisher distance, class imbalance, non-Fisher references, or pooled covariance predict failures.
- The paper’s novelty is composing covariance-projected cosine machinery into a closed-form law relating MCS and OOD AUROC.
- MCS provides a theoretically grounded, task-aware alternative for comparing linear probes.
2 Empirical evidence
The empirical study evaluates logistic-regression probes across models, layers, and concept domains using held-out OOD data. MCS closely tracks AUROC across all tested conditions.
- Experiments use Llama-3.3-70B at layer 33, additional layers and models, and 24 datasets across truthfulness, gender classification, and general NLP.
- Logistic-regression probes are trained for each task, with ID and OOD directions denoted w_id and w_ood.
- MCS uses the full-sample OOD covariance, while pooled within-class covariance is evaluated as a contrast.
- OOD data are split into disjoint train and test halves to train probes, estimate task quantities, and compute held-out AUROC without overfitting bias.
- R^2≥0.93 for held-out AUROC against MCS in every tested condition, whereas ECS reaches R^2=0.06 in the worst condition.
3 Theory
Under balanced classes with Gaussian projections, AUROC and Mahalanobis cosine similarity are both governed by probe SNR. Their sigmoid and softsign forms compose into an approximately linear AUROC–MCS relationship, especially for large Fisher distance.
- Setup: The analysis assumes binary labels, Gaussian class-conditional projections, and a fixed candidate direction evaluated on OOD data.Projection Gaussianity is weaker than joint Gaussianity and is retained as a simplifying assumption.
- Setup: MCS uses covariance weighting appropriate for comparing probe directions, transforming data and probes inversely so whitened projections are preserved.This contrasts with the inverse-covariance weighting used for Mahalanobis distances between data points.
- Linearity mechanism: For fixed Fisher distance, AUROC and MCS are monotone functions of SNR, tracing a parametric curve whose curvature largely cancels.AUROC follows the binormal Gaussian-CDF relationship, while MCS approaches a bounded softsign shape as Fisher distance grows.
- Closed form: Theorem 1 expresses MCS between a candidate direction and the Fisher direction using only candidate SNR and task Fisher distance.The resulting identity is obtained by computing the relevant quadratic forms from the covariance decomposition.
- Linearity mechanism: The cancellation is approximate: near |MCS| = 1, the local slope drops toward 0, producing the observed top-right flattening.Figure 2 reports empirical points largely following the parameter-free theory and a near-straight AUROC–MCS curve that bends in the top-right corner.
- Linearity mechanism: At Fisher distance greater than 20, the task-dependent factor is within 0.5% of its limiting value, with central slope 1/√π.Empirical slopes are slightly lower because saturation-tail sampling reduces the global slope.
4 When does the linearity break down
The AUROC–MCS linearity depends on total-covariance weighting, large Fisher distance, balanced classes, and Fisher-style reference probes. Violating these conditions produces characteristic deviations and lower linear-fit R^2.
- Conditions: The law requires MCS computed with total covariance, large Fisher distance, roughly balanced classes, and OOD probes near the Fisher direction.Figure 3 tests these assumptions using common axes and linear fits of OOD AUROC against MCS.
- Wrong covariance: Using pooled instead of total covariance makes MCS equal s/zmax, turning AUROC into a sigmoid in MCS and reducing R2 from 0.98 to 0.83.This is the wrong-covariance failure mode on the main experiment data.
- Non-Fisher probe: Difference-of-means probes yield R2 = 0.79 on the same empirical data, delimiting the prediction to Fisher-style probes such as LR, LDA, and shrinkage variants.The theory does not hold when the OOD reference probe deviates too far from the optimal Fisher direction.
- Small Fisher distance: Small Fisher distances create a fan of task-specific slopes, and synthetic data at zmax ∈ {0.1, 0.5, 1, 2} reduces the combined linear-fit R2 to 0.666.The LLM experiments are unaffected because every task has zmax > 20.
- Class imbalance: Class imbalance steepens the AUROC–MCS slope; synthetic settings with π ∈ {0.5, 0.1, 0.02, 0.004} reduce R2 to 0.84.The balanced-class factor generalizes to π(1 −π), weakening saturation cancellation as class proportions diverge.
5 Discussion
The results support MCS as a theoretically sound, task-aware alternative to Euclidean cosine similarity. The authors suggest extending this covariance-aware comparison to other interpretability methods.
- Discussion: MCS is presented as a theoretically sound alternative to standard Euclidean cosine similarity for comparing linear probes.The proposed broader direction includes steering-vector comparison, SAE feature alignment, concept-direction clustering, data filtering, and data attribution.
Limitations
The study requires labeled OOD moments, is theoretically calibrated to Fisher-style probes and a restricted model setting, and provides only approximate linearity under Gaussian projection assumptions.
- Data requirements: Computing MCS with total covariance requires estimating the OOD probe and covariance from labeled OOD data.The candidate probe is training-free, but the method is not label-free for the target distribution.
- Probe scope: The closed-form theory is calibrated to Fisher-style probes, while difference-of-means probes show marked degradation and non-Fisher references remain analytically open.This limits direct theoretical coverage beyond LR, LDA, and related Fisher-style variants.
- Evaluation scope: Evaluation is restricted to binary classification on residual-stream features of autoregressive LLMs, leaving multiclass probes, attention and MLP features, and non-LLM architectures untested.The experiments use several LLMs and layers but do not establish coverage outside this setting.
- Approximation and assumptions: The linearity law is approximate: per-task slopes approach 1/√π, local slopes go to 0 at the extreme, and heavy-tailed projections lack formal guarantees.Empirical projections are often approximately Gaussian and simulations show linearity can survive clear non-Gaussianity, but the stated guarantee remains limited.
B.3 Dataset details
The study uses datasets spanning truthfulness, gender classification, and general NLP, with large datasets subsampled in some cases. The datasets are publicly available for noncommercial academic evaluation, and all data is English-language.
- Truthfulness datasets: The truthfulness collection contains ten datasets covering fundamental truth types, sycophantic and expectation-inverted lying, and on-policy deception.The collection includes definitional, empirical, logical, fictional, and ethical truth, plus roleplay, insider trading, and sandbagging datasets.
- Gender classification datasets: The gender-classification collection contains six datasets, with datasets larger than 4,000 samples randomly subsampled.The listed datasets include PAN17, CrowdFlower Twitter, HappyDB, BiosBias, WinoBias, and GAP.
- General NLP benchmarks: The general NLP collection covers sentiment, topic classification, natural-language inference, question answering, and commonsense reasoning.Large datasets are also subsampled to 4,000 samples.
- Licensing and language: The datasets are publicly available and used only for noncommercial academic evaluation under their stated licenses or research-use terms.The passage lists MIT, Apache 2.0, CC BY-SA 3.0, Academic Free License v3.0, and CC BY 4.0 among the applicable licenses.
- Licensing and language: All data used in the work is English-language.
C Additional empirical results
Additional experiments show that AUROC–MCS linearity is robust across conditions and covariance estimators, while its predicted boundaries and Gaussianity assumptions can be tested empirically. The results also identify covariance structure and Fisher distance as important determinants of the relationship.
- MCS and ECS against AUROC: The empirical AUROC–MCS relationship remains strong across all eight conditions, whereas ECS is substantially weaker.The study reports near-linear fits across models, layers, and concept domains, with MCS outperforming ECS.
- Verification of the theory: Theoretical predictions track empirical data very well across all eight conditions.
- Empirical Gaussianity: 73% of empirical skewness values lie within ±0.5 and 79% of kurtosis values within ±1 across conditions, directions, and test distributions.Most empirical projections are therefore reported as having small skewness and kurtosis.
- Non-Gaussian simulations: The AUROC–MCS R2 remains extremely high for deliberately constructed non-Gaussian distributions that violate projection-Gaussianity.This result comes from simulations designed to break the theoretical Gaussianity assumption.
- Empirical Fisher distance: All empirical zmax values exceed 20, where per-task AUROC–MCS slopes lie within 0.5% of their limiting values.This saturation supports a single strong linear fit across tasks and a universal slope approximation.
- Robustness to covariance estimation: Ledoit–Wolf covariance matches the full covariance, while the diagonal approximation degrades sharply, reaching R2=0.002 versus 0.936 on general NLP.The comparison indicates that off-diagonal covariance structure is important for MCS performance.
E Proofs of background results
The proofs derive covariance identities, AUROC from Gaussian class projections, and the Fisher-optimal probe by reducing the optimization to a constrained inner-product problem. They then express the relevant quantities through scale-invariant quadratic forms.
- Covariance decomposition: The law of total covariance decomposes total covariance into pooled within-class covariance and between-class covariance.The between-class contribution is represented by a rank-one term proportional to δδ⊤.
- AUROC under Gaussian projections: Under Gaussian class projections, AUROC is the probability that the probe scores a class-1 sample above an independent class-0 sample.The score variances combine to 2w⊤Σ_poolw under the stated independence assumption.
- SNR optimization: The probe SNR is scale-invariant, so maximizing it reduces to maximizing (w⊤δ)^2 subject to w⊤Σ_poolw = 1.
- Whitened optimization: Whitening transforms the constrained optimization into maximizing (u⊤q)^2 over unit-norm u, whose maximum follows from Cauchy–Schwarz.The resulting maximum is expressed through δ⊤Σ_pool^-1δ.
- Closed-form quantities: The proof computes the quadratic forms needed for the closed-form MCS expression using the relation between total and pooled covariance.The derivation also identifies zmax as the Fisher distance.
G Slope along the AUROC–MCS curve
The theory factorizes the AUROC–MCS slope into task and shape factors, explaining near-universal empirical slopes when tasks operate in the saturation regime. It also identifies finite-saturation, small-distance, and class-imbalance conditions that reduce linearity.
- Slope factorization: The slope factorizes into a task-independent shape factor and a task-dependent factor determined by zmax.This separates within-task curve flatness from variation across tasks.
- Saturation regime: Every empirical task has zmax > 20, placing its AUROC–MCS slope within 0.5% of the saturation-limit curve.The task factor approaches its limit with deviations falling off as zmax^-2.
- Saturation cancellation: The softsign and sigmoid saturation effects largely cancel, keeping the local slope near its central value across most MCSΣtot values.The Gaussian derivative shrinks while the MCS derivative grows as |s| increases.
- Saturation cancellation: The cancellation is imperfect: the slope peaks near s = ±2 at approximately 0.629 before decaying toward zero near the MCS boundaries.This finite-tail decay explains why empirical AUROC–MCS fits have R2 below 1.
- Empirical validation: Empirical slopes across eight conditions agree with the theoretical central slope 1/√π ≈0.564.Figure 9 reports linear-fit slopes with 95% bootstrap confidence intervals.
- Failure modes: Heterogeneous small-zmax tasks produce a fan rather than a single line, with synthetic global linear R2 dropping to 0.666.Strong class imbalance also weakens cancellation and steepens the AUROC–MCS slope; experiments use approximately balanced datasets.
I Empirical alignment of LR and LDA reference directions
The empirical analysis substitutes an OOD-trained logistic-regression direction for the Fisher reference direction because balanced, well-separated Gaussian tasks make them nearly proportional. Across models, layers, and concept domains, the resulting MCS measures remain strongly aligned and preserve the observed linearity.
- Experimental setup: The experiments compute Fisher and logistic-regression directions on the train half of stratified 50/50 splits, then compare their MCS values.The Fisher direction uses pooled covariance with λ = 10^-6 regularization.
- Empirical alignment: Across all models, layers, and concept domains, pairwise MCS values for logistic-regression and Fisher references correlate strongly.This alignment supports substituting the logistic-regression direction in the empirical analysis.
- Empirical alignment: Replacing the Fisher reference with logistic regression changes the headline linear-fit R2 by less than 1%.The two reference-direction MCS measures therefore preserve the observed linearity across the eight conditions.
- Reference-direction relationship: For balanced binary Gaussian classification with shared covariance, logistic-regression and Fisher directions coincide up to scaling.The paper notes that real activations only approximately satisfy these assumptions.
- Scope: The substitution may be less reliable on harder, less-separable tasks, where the residual mismatch between directions is expected to be larger.All probe datasets in these experiments have ID AUROC above 0.95.