Stop Probing, Start Coding: Why Linear Probes and Sparse Autoencoders Fail at Compositional Generalisation

Abstract: The linear representation hypothesis states that neural network activations encode high-level concepts as linear mixtures. However, under superposition, this encoding is a projection from a higher-dimensional concept space into a lower-dimensional activation space, and a linear decision boundary in the concept space need not remain linear after projection. In this setting, classical sparse coding methods with per-sample iterative inference leverage compressed sensing guarantees to recover latent factors. Sparse autoencoders (SAEs), on the other hand, amortise sparse inference into a fixed encoder, introducing a systematic gap. We show this amortisation gap persists across training set sizes, latent dimensions, and sparsity levels, causing SAEs to fail under out-of-distribution (OOD) compositional shifts. Through controlled experiments that decompose the failure, we identify dictionary learning -- not the inference procedure -- as the binding constraint: SAE-learned dictionaries point in substantially wrong directions, and replacing the encoder with per-sample FISTA on the same dictionary does not close the gap. An oracle baseline proves the problem is solvable with a good dictionary at all scales tested. Our results reframe the SAE failure as a dictionary learning challenge, not an amortisation problem, and point to scalable dictionary learning as the key open problem for sparse inference under superposition.

• The paper demonstrates that linear probes and SAEs fail out-of-distribution generalization due to nonlinear distortions from latent superposition.
• It reveals that per-sample nonlinear sparse coding outperforms amortized inference, underscoring a persistent amortisation gap.
• The study identifies that inadequate dictionary learning, rather than encoder design, is the primary bottleneck hindering true compositional recovery.

Failure Modes of Linear Probes and Sparse Autoencoders in Compositional Generalization

Introduction and Problem Setting

Understanding the failure mechanisms of compositional generalization in neural network interpretability pipelines is essential, especially as advances in LLMs and other high-dimensional models outpace our ability to extract reliable, robust representations. The linear representation hypothesis (LRH) posits that model activations are linear combinations of abstract high-level concepts. While this provides a foundation for approaches such as linear probing and sparse autoencoding, it is often confounded with a stronger notion: that these concepts are linearly accessible, i.e., can be recovered via linear decoding regardless of architectural superposition effects. The current work rigorously exposes the circumstances under which this conflation fails, and dissects the underlying causes of observed OOD generalization collapse in two widely-used families of methods: linear probes and sparse autoencoders (SAEs) (2603.28744).

The fundamental setting is that of superposition: more independent latent factors are encoded than there are activation dimensions ($d_z > d_y$), so activation observations are projections of sparse, high-dimensional codes. This linear bottleneck is underdetermined; activation space does not support unique, linear invertibility to recover the true generative factors. The compressed sensing framework stipulates that, provided codes are sparse ($k \ll d_z$), unique recovery is theoretically possible, but only through nonlinear inference procedures sensitive to the combinatorics of sparse selection and dictionary geometry.

Compositional Generalization and the Linear Representation Hypothesis

A core claim is that, under superposition, the critical challenge is not simply that concepts are encoded as linear mixtures, but rather that the mapping back from activations to concepts is nonlinear and task-dependent. Concretely, even if a downstream classification task is linearly separable in latent space, the corresponding decision boundary can become highly nonlinear after projection into activation space. This failure mode is exacerbated in OOD scenarios involving novel recombinations of latent factors—the precise locus of compositional generalization.

Figure 1: OOD evaluation is performed by withholding novel combinations of latent factors during training; projections into activation space fold these combinatorial patterns nonlinearly.

Linear probes—supervised, task-specific decoders fit atop neural activations—succeed only insofar as decision boundaries remain linearly separable after projection. As shown, even perfect linear probes trained ID can catastrophically fail OOD due to the geometric nonlinearity induced by high overcompleteness ratio ($d_z \gg d_y$).

Figure 2: Linear probes fit ID decision boundaries but suffer catastrophic OOD collapse under superposition, reflecting the nonlinearity introduced by projection from latent to activation space.

Sparse Coding, Autoencoders, and the Amortisation Gap

Classical sparse coding achieves latent recovery (and thus compositional generalization) using iterative, per-sample nonlinear programs (e.g., Lasso, ISTA/FISTA) that exploit compressed sensing guarantees; recovery succeeds when the dictionary satisfies the restricted isometry property (RIP) at sufficiently sparse levels (i.e., $d_y \gtrsim O(k \log(d_z/k))$). SAEs, in contrast, amortize this inference by training a fixed encoder, often a shallow neural net, to map activations to sparse codes in a single learned step.

Amortisation introduces a systematic, data-dependent generalization gap: the encoder is optimized only for training-time support structures. When code supports shift compositionally OOD, the fixed encoder is nonadaptive and generalizes poorly. The paper provides both controlled experiments and theoretical justification for this amortization gap, and demonstrates its persistence even when scaling up training data, latent dimensionality, and other key parameters.

Figure 3: Toy example demonstrating that SAEs under superposition fail to recover ground-truth latent factors, while per-sample sparse coding recovers all supports up to scale.

Figure 4: The amortisation gap persists across undersampling ratios; per-sample FISTA achieves a sharp phase transition (near-perfect recovery above the compressed sensing threshold), while SAEs plateau at suboptimal recovery regardless of $\delta = d_y/d_z$.

Dictionary Learning as the Primary Bottleneck

A central finding is that the bottleneck in compositional generalization is not the inference procedure (encoder) per se, but the quality of the learned dictionary. Extensive ablations show that performing per-sample (FISTA) inference with a dictionary learned by the SAE yields no improvement in recovery—i.e., the failure propagates from the choice of unmixing weights, not from encoder parameterization.

Figure 5: Swapping out the SAE encoder for per-sample FISTA on the same dictionary achieves no gain; only replacing the dictionary via alternating minimization (DL-FISTA) recovers performance. The oracle baseline (FISTA with ground-truth dictionary) closes the gap entirely.

This points to a universal failure mechanism: as $d_z$ increases, both SAEs and alternating-minimization dictionary learners degrade, but the amortized encoders degrade faster and saturate; the dominant source of error is that the learned dictionaries do not align with the true latent directions (as measured by cosine similarity and support diagnostics).

Scaling Laws, Data Regimes, and the Limits of Amortisation

The authors systematically explore scaling in latent dimension, sparsity, and sample count. Increasing the number of training samples benefits dictionary learning in classical alternating minimization (DL-FISTA), but not in SAEs, which plateau far below compressed sensing optimal recovery. Remarkably, DL-FISTA dominates linear probes OOD when it does not fail to learn the dictionary ($d_z \leq 1000$), but both methods collapse at large scale due to the difficulties of high-dimensional nonconvex dictionary learning.

Figure 6: Scaling latent dimension does not close the compositional gap. DL-FISTA degrades due to dictionary learning difficulty but maintains a small ID--OOD gap. SAEs saturate at low recovery even as the number of dimensions grows.

Figure 7: Increasing training set size closes the gap for DL-FISTA but not for SAEs. OOD MCC for SAEs plateaus or degrades; DL-FISTA and oracle approach perfect recovery with sufficient samples.

Moreover, SAE-learnt dictionaries provide some utility as warm starts for dictionary learning at small dimension, but this advantage vanishes at scale, reinforcing the criticality of the dictionary learning algorithm rather than initialisation.

Figure 8: DL-FISTA initialised from the SAE decoder converges more rapidly at small scale ($d_z=5000$), but both warm and random starts converge to the same optimum.

Theoretical Implications and Forward Directions

The results clarify several key theoretical points for the practice of mechanistic interpretability and compressed-sensing-based sparse recovery in modern neural systems:

• Linear probes are fundamentally limited under superposition: They can only reliably separate concepts OOD if supports do not recombine in novel, unexpected ways. Highly overcomplete activation spaces invalidate the direct use of linear accessibility.
• Amortizing inference does not circumvent the nonconvexity of dictionary learning: SAE encoders cannot generalize supports OOD because the learned dictionary basis is aligned with spurious training co-occurrences, not abstract latent directions.
• Scalable, robust dictionary learning is the central open question: Theoretical conditions (RIP, spark) guarantee identifiability only for sufficiently sparse and diverse data, but practical stochastic optimizers and deep SAE encoders do not reliably find the ideal solution at relevant scale.

Conclusion

This paper provides an authoritative analysis of compositional generalization failure in linear probes and SAEs under latent superposition. It demonstrates that while classical (per-sample) sparse coding is theoretically sufficient for recovery—even OOD—the bottleneck is dictionary learning: the directionality of SAE-learned dictionaries limits both identifiability and downstream task performance. Amortized inference, by design, overfits the training distributions of support co-occurrences and exhibits a persistent amortisation gap even with ample data. Future progress in robust AI interpretability thus depends on scalable, high-fidelity dictionary learning algorithms capable of exploiting the sparsity inherent in large model activations, with implications for both mechanistic interpretability and high-stakes application robustness.