Is a Linear Probe Evidence of a Linear Representation?

AI summary

Overview

Research area: Representation learning and interpretability — specifically, whether linear probing of pretrained vision transformers provides evidence for the "linear representation hypothesis" (that generative factors are encoded as linear directions in feature space).

Technical level: Advanced. The paper uses covariance algebra, random-features/RKHS theory, neural tangent kernel arguments, and spectral analysis; the core intuition is accessible, but the formal machinery (Definition 1, Propositions 1–3) requires comfort with linear algebra and statistics.

Scope (one sentence): The paper formalizes a strict "linear world model" definition, proposes three asymmetric diagnostics to test it, validates them in a controlled simulation with known ground truth, and applies them to six pretrained vision transformers on a new two-style benchmark (SVG-World) plus CelebA, dSprites, 3D Shapes and MPI3D.

What This Paper Is About

A standard practice in representation learning is to train a linear probe on a model's frozen features and declare a latent factor (e.g., object position) "represented" if probe accuracy is high. The authors show this test is too permissive: random-features theory implies that any sufficiently rich feature map — even a randomly initialized, untrained network — admits a linear readout of almost any smooth latent once the feature dimension is large enough. The goal is therefore to distinguish two very different situations that both produce high probe accuracy: a genuine linear world model, where the latents dominate the representation's variance, versus a kernel machine, where the latents are merely one linear readout from a much larger curved feature manifold.

Key Contributions

1. A formal definition of a linear world model. Definition 1 states that a feature map f implements a linear world model of latents z with residual level σ² if f is, up to a bounded, z-orthogonal residual, an affine map of the whitened latents — i.e., the latents carry most of f's variance. This is strictly stronger than "z is linearly decodable from f."

2. Three asymmetric diagnostics plus a ceiling. The authors propose (D1) reverse predictivity R²_rev — how much of the feature variance an affine function of the latents explains; (D2) spectral concentration — whether the latent-aligned variance lives in the leading principal components; and (D3) out-of-distribution generalization — whether probe directions transfer across conditions. They also derive a ceiling C_r (the variance share of the top r principal components, where r = rank of Cov(z)), which makes reverse predictivity comparable across encoders with different spectra.

3. Controlled validation in a fully observed simulation. Two-dimensional latents are lifted by a frozen random Fourier map into a two-layer ReLU MLP, and trained versus randomly initialized versions are compared at matched hidden width, confirming that forward decoding alone cannot separate the regimes while the asymmetric diagnostics can.

4. SVG-World, a deterministic two-style benchmark. A 32-dimensional latent layout (20 active dimensions: 4 people each with (x, z) position, facing direction and skin hue; 2 animals each with (x, z) position) rendered in two visual styles (AquaWorld and WestWorld) that share the layout but differ in rendering, with N = 10,000 scenes per split at 224×224. This enables the cross-style transfer test (D3), which existing disentanglement benchmarks do not support.

Main Findings

• Forward decoding is asymptotically free. In the simulation, random-init forward R² rises smoothly with width and reaches approximately 1 by h ≈ 100 — matching the random-features prediction of Proposition 1 — so forward R² alone cannot distinguish a trained from an untrained feature map beyond very small widths.

• The asymmetric diagnostics open a robust gap in simulation. At h = 1024 with 5 seeds, reverse R² is 0.9935 with weight decay λ = 1 versus 0.7841 with λ = 0, at an identical forward R² of 1.000. With λ = 1, all diagnostics sit flat near 1 across widths; random init sits well below on reverse R² and top-2 variance share. Random init's reverse R² plateaus near 0.2 rather than decaying to 0 as Proposition 2 would predict asymptotically — attributed to finite-width ReLU MLPs not realizing strict z-orthogonality.

• Linear world models are visible as geometry. At h = 1024, the trained model's top-2 principal components carry 99.57% of the variance (shares 0.5657 and 0.4300, against 0.1384 and 0.1151 at random init), and its top-2 projection is an axis-aligned image of the latent square; the random model's spectrum decays smoothly over hundreds of components and folds the latent square onto itself.

• Satisfying Definition 1 is a property of training pressure, not just training itself. Without weight decay the model decodes as well as the regularized one on forward and OOD metrics but sags to roughly 0.7 on reverse R² and top-2 share at intermediate widths. The authors hypothesize that penalizing the readout norm pushes the latent-aligned subspace toward the leading directions.

• Pretrained DINOv3-B/16 decodes SVG-World latents far better than random init. Averaged over both worlds on CLS features, pretrained reaches mean ρ ≈ 0.94 (R² = 0.88), versus 0.55 (R² = 0.33) for random init (Δρ = 0.38, rising to 0.58 on animals). Mean-pooled patches give 0.91 (R² = 0.82) for pretrained. Random init partially decodes person positions (R² = 0.41) but loses animals almost completely (R² = 0.17); with a cross-validated ridge penalty random init reaches R² = 0.66, but the animal gap remains.

• The reverse gap is much larger than the forward gap. On raw CLS features, pretrained reverse R² is 0.18 versus 0.02 for random init; on mean-pooled patches, 0.28 versus 0.02. Over five random initializations the CLS gap is 0.187 [0.182, 0.191] in AquaWorld and 0.143 [0.139, 0.148] in WestWorld (95% paired scene-bootstrap intervals). The pretrained-vs-random ratio is roughly 10×, against roughly 1.4× for forward decoding with a cross-validated penalty.

• Against the ceiling, pretrained models reach about a third at best. On CLS features in AquaWorld (R²_rev = 0.212, C₁₂ = 0.609), pretrained reaches 0.35 of the attainable share and random init 0.05. Across all encoders and datasets the probed factors take at most about 0.53 of what any linear image of them could reach.

• Latent-aligned variance is spectrally concentrated but at a low absolute level. Pretrained R²_rev peaks at k ≈ 32 (0.17 on CLS, 0.19 on patches) when features are PCA-whitened, falls below 0.04 at k = 256 and to 0.01 at full whitening (k = 768); random init is flat at about 0.02. Pretrained DINOv3-B/16 needs its top 17 principal components to collect 90% of its latent-predictable variance, against 55–59 for its random initializations.

• Probe directions transfer across visual style. On mean-pooled patch features, pretrained averages ρ = 0.91 within-domain versus 0.51 for random (Δρ = 0.41); cross-domain the absolute values fall to 0.65 and 0.23 but the gap does not shrink (Δρ = 0.42). Animal positions transfer best (pretrained cross-domain ρ ≈ 0.76, roughly 0.68 to 0.84), while random init falls to zero or negative. With probe outputs rescaled on half the target style, transfer R² is 0.16 and 0.31 for pretrained against 0.01 and 0.02 for random.

• Results extend across architectures but with a size paradox. All six transformers (DINOv3 S/16, B/16, L/16; DINOv2-B/14; DeiT-B/16; CLIP-B/16) exceed their random initializations on reverse R² by 0.145 to 0.235; a ResNet-50 shows a gap of 0.105. Within DINOv3, larger models decode better but carry less latent-aligned variance: from S/16 to L/16 forward R² rises from 0.855 to 0.914 while reverse R² falls from 0.257 to 0.169.

• On standard benchmarks the picture is mixed. On dSprites and MPI3D the pretrained encoder has higher forward R² than its random init (0.42 vs 0.16 and 0.85 vs 0.61) but a lower reverse R²; on 3D Shapes forward R² is tied. Random initializations there are nearly low-dimensional (participation ratio at most 3.8, against 8.7 to 16.6 for pretrained), so their ceiling is close to 1. Four of six encoder–benchmark pairs keep the sign of the gap under all three factor encodings tested; two do not.

• On real photographs the SVG-World pattern holds. On CelebA, pretrained models reach 0.28 to 0.36 of their ceiling against 0.07 to 0.11 at random init, and spectral concentration is 0.15 to 0.17 versus 0.02 to 0.03. CLIP's raw reverse R² (0.192) exceeds DINOv3 CLS (0.170), but both reach 0.28 of the ceiling.

• Overall conclusion. Pretraining produces structure well beyond a random kernel but falls well short of a linear world model. The diagnostics are necessary conditions: low scores rule out a linear world model, but high scores do not establish one.

Methodology in Plain English

The authors build a ruler and then measure models with it.

First, they state precisely what it would mean for a network to "really" represent a factor linearly: the factor should account for most of the network's feature variance, with only a small leftover that is unrelated to it. That is Definition 1.

Second, they argue that a normal linear probe cannot test this, because a wide enough random feature map can fit almost any smooth target. So they design three tests that exploit asymmetries:

• Reverse predictivity: instead of predicting the latent from features, predict the features from the latent. If the latent really dominates the representation, most feature variance is explained; if the latent is a minor readout from a big curved manifold, little is explained.
• Spectral concentration: compress the features to their top few principal components (equalizing their variances) and check whether the latents are linearly recoverable from just those. A linear world model keeps its latent information in the very top directions; a kernel machine buries it.
• Out-of-distribution transfer: train the probe in one visual condition and test it in another with the same underlying latents. A factor learned as a genuine axis should transfer; a factor learned as a style-specific pattern should not.

They also derive a ceiling on reverse predictivity from the feature covariance spectrum, so numbers can be reported as a fraction of what is attainable rather than as raw values.

Third, they check that these tests behave as claimed in a simulation where the true latents exist and are known: 2D uniform latents are lifted by 256 frozen random Fourier features (cosines of random projections) and then passed through a two-layer ReLU MLP, compared at matched hidden width either untrained or trained end-to-end with AdamW (learning rate 2·10⁻³, 2000 steps, batch 256) with weight decay of 0 or 1. They verify the network is genuinely nonlinear — at h = 1024 it visits 1988 ReLU activation regions on average and its Jacobian varies across inputs.

Finally, they apply the diagnostics to real models, using SVG-World (which supplies paired visual styles) for the transfer test and CelebA plus three standard disentanglement benchmarks for generality, always comparing each pretrained encoder against a random initialization of the same architecture and probing protocol.

Why This Matters

Impact on research: A large body of interpretability and "world model" claims rests on linear probe accuracy. This paper shows that such accuracy is, in the limit, a property of feature richness rather than of learned structure, and supplies concrete diagnostics — reverse predictivity, spectral concentration, and cross-condition transfer — that researchers can adopt. It also reframes the binary question "is this factor represented?" as a graded one: how much of the representation does a factor set occupy, and where in the spectrum.

Real-world applications:

• Safety and auditing of deployed vision models: knowing whether an internal direction is a genuine axis of the model's representation, rather than a decodable but peripheral signal, bears on whether editing or monitoring that direction will behave predictably.
• Robustness and OOD engineering: the paper argues that the curved structure orthogonal to the probe is what matters for out-of-distribution generalization and compositional transfer, so probe-only evaluations can overstate reliability under domain shift.
• Model selection for downstream transfer: the finding that larger DINOv3 models decode better while devoting less variance to the factors suggests probe accuracy and representation alignment can point in opposite directions when choosing an encoder.
• Benchmark design for disentanglement: SVG-World, with its shared latent layout across two renderings, offers a template for testing whether learned factor directions are style-invariant, something existing benchmarks do not test.

Industry relevance: teams that use frozen foundation-model features for regression or control tasks (robotics, medical imaging, geospatial, content moderation) routinely validate with linear probes. The diagnostics give a cheap additional check — reverse predictivity and spectral concentration require only the same features and labels — for deciding whether a feature extractor's internal organization will hold up under distribution shift, and the participation-ratio and ceiling diagnostics are useful for comparing encoders with very different spectra.

Future Directions

• Which objectives produce more linear-world-model-like encoders? The paper's discussion points toward objectives that explicitly favor compressed or bottlenecked codes as the natural place to look; this remains untested here.
• Can the weight-decay effect be reproduced in self-supervised training? The simulation shows that penalizing the readout norm moves a network toward Definition 1, but the authors note that self-distillation losses have no obvious analogous pressure. Whether a practical regularization can induce the same effect at scale is open.
• Reconciling the encoder-size paradox. Larger DINOv3 models decode factors better (forward R² 0.855 to 0.914 from S/16 to L/16) while carrying less latent-aligned variance (reverse R² 0.257 to 0.169). Understanding why capacity trades off against alignment is unresolved.
• Extending past linear diagnostics. The authors note their diagnostics are themselves linear, so a factor encoded faithfully along a curved manifold will also score low; nonlinear or manifold-aware versions of the tests, and an account of superposition, are natural extensions. The Discussion section in the provided text is truncated, so the paper's full closing research agenda is not reported here.

Target Audience

Researchers and graduate students in representation learning, mechanistic interpretability, and self-supervised vision; practitioners who rely on linear probing to justify claims about what a model "knows"; and benchmarking researchers interested in disentanglement evaluation. Readers evaluating whether foundation-model features are suitable for a downstream task that requires out-of-distribution robustness will find the diagnostic toolkit directly actionable, even if they skip the proofs.

Authors’ abstract

The linear representation hypothesis holds that pretrained vision models encode generative factors as linearly decodable directions in feature space, and comes with a standard methodology: train a linear probe, declare a factor represented if accuracy is high. We show this test is too permissive. Asymptotically, any sufficiently rich feature map, including a random one, admits a linear readout, so high probe accuracy is consistent with two very different geometries: a true linear world model, where the latent-aligned subspace dominates the representation, or a kernel machine, where the latents lie linearly inside a much larger curved manifold. This distinction matters because the curved structure orthogonal to the probe is relevant for out-of-distribution generalization, robustness, and compositional transfer. We propose three asymmetric diagnostics that distinguish the two cases: reverse predictivity, spectral concentration, and out-of-distribution generalization. We validate them in a fully observed simulation where ground truth is known. Applied to six pretrained vision transformers on SVG-World, our deterministic generative benchmark with paired visual styles, and to two of them on CelebA, they show pretraining produces structure beyond a random kernel but well short of a linear world model. On dSprites, 3D Shapes and MPI3D, forward decodability and latent-aligned variance dissociate. Thus, forward decoding alone does not show that the probed factors align with the dominant geometry of the representation. In summary, our work cautions about the success of linear probes as opposed to truly linear and low dimensional representations.