How Many Features Can a Language Model Store Under the Linear Representation Hypothesis?

1. Problem Statement

The paper addresses a fundamental question in the theoretical understanding of Large Language Models (LLMs): Under the Linear Representation Hypothesis (LRH), how many features ($m$) can a single layer of $d$ neurons store and subsequently access?

The LRH posits that intermediate layers of LLMs store features linearly. However, previous literature often conflated two distinct concepts:

1. Linear Representation: Features are embedded linearly into neuron activations.
2. Linear Accessibility: Features can be decoded (extracted) using a linear probe.

The authors argue that while classical compressed sensing suggests $d = O(k \log(m/k))$ neurons suffice for $k$-sparse features (where $k$ is the number of active features per input), this assumes non-linear decoding (e.g., $\ell_1$-minimization). The paper investigates the stricter regime where decoding must also be linear, asking how this constraint affects the capacity for "superposition" (storing $m \gg d$ features).

2. Mathematical Framework

The authors formalize the problem using the following definitions:

• Activations: An input $\ell$ is mapped to a $d$-dimensional vector $f(\ell) = Az$, where $z \in \mathbb{R}^m$ is the feature vector and $A \in \mathbb{R}^{d \times m}$ is the representation matrix.
• Features: A feature $z_i$ is a function of the input. The set of features is assumed to be $k$-sparse (at most $k$ features are active for any given input).
• Linear Accessibility: A feature $z_i$ is linearly accessible if there exists a probe vector $b_i$ such that $\langle b_i, f(\ell) \rangle \approx z_i(\ell)$.
• The Goal: Find the minimum dimension $d$ required to find matrices $A$ (encoding) and $B$ (decoding/probing) such that for all $k$-sparse $z$:
  $$\| B^T A z - z \|_\infty < \epsilon$$
  This ensures that the linear probe $B^T$ can recover the original feature vector $z$ from the compressed representation $Az$ with error $\epsilon$.

3. Methodology

The paper employs tools from Compressed Sensing, Random Matrix Theory, and Extremal Graph Theory to derive tight bounds.

• Upper Bound Construction: The authors use standard random matrix constructions (specifically, matrices with approximately orthogonal columns, known as $\mu$-incoherent matrices). They show that if $A$ is a random matrix (e.g., scaled Rademacher or Gaussian), $A^T A$ approximates the identity matrix. By setting the coherence parameter $\mu$ appropriately relative to the sparsity $k$, they ensure that interference from other active features remains below the error threshold $\epsilon$.
• Lower Bound Proof: The proof is more complex and relies on a contradiction argument:
  1. Rank Constraints: They utilize a result by Alon (2003) regarding the rank of matrices with large diagonal entries and small off-diagonal entries.
  2. Submatrix Analysis: Instead of looking at the whole matrix, they apply rank bounds to principal submatrices of size $r$.
  3. Graph Reduction: They construct a graph where an edge exists between features $i$ and $j$ if their interference (inner product of probe $b_i$ and representation $a_j$) is high.
  4. Turán's Theorem: They apply Turán's theorem to show that if $d$ is too small, the graph must contain a large clique or lack a large independent set, forcing the existence of a "bad" feature pair where interference compounds beyond the allowed $\epsilon$.

4. Key Results

A. The Quantitative Gap (Main Contribution)

The paper establishes a significant quantitative gap between Classical Compressed Sensing (non-linear decoding) and Linear Compressed Sensing (linear decoding):

• Classical (Non-linear decoding): $d = O(k \log(m/k))$.
• Linear (Linear decoding): $d = \Omega_\epsilon\left( \frac{k^2}{\log k} \log \frac{m}{k} \right)$ (Lower Bound) and $d = O_\epsilon(k^2 \log m)$ (Upper Bound).

Implication: Linear accessibility is a strictly stronger and more expensive hypothesis than linear representation alone. The dimension $d$ must scale quadratically with sparsity $k$ in the linear setting, whereas it scales linearly in the non-linear setting.

B. Superposition Capacity

Despite the quadratic penalty on $k$, the upper bound $d = O(k^2 \log m)$ confirms that a layer can still store an exponential number of features ($m$) relative to the number of neurons ($d$), provided the features are sparse. This provides theoretical validation for the "superposition hypothesis" (Elhage et al., 2022).

C. Geometry of Feature Vectors

The paper challenges the intuition that feature representation vectors must be orthogonal.

• Proposition 9: It is possible to construct $A$ and $B$ where representation vectors are highly correlated (nearly parallel) and probe vectors are highly correlated, yet linear recovery is perfect. The only strict requirement is that a feature's probe vector must be orthogonal to other features' representation vectors.
• Proposition 11: If one imposes a constraint that representation and probe vectors must have unit norms (or bounded norms), then the geometry becomes "cleaner": representation and probe vectors align, and different features become approximately orthogonal.

D. Robustness to Non-linearities

The authors extend their lower bound to settings where the decoder includes a bias and a monotonic activation function (e.g., ReLU).

• Theorem 12 & Corollary 13: Adding a bias or a monotonic activation function to the linear probe does not asymptotically increase the storage capacity. The lower bound $d = \Omega(k^2 \log(m/k))$ remains valid even for binary features and non-linear decoders of the form $\sigma(W^T x + b)$.

5. Significance and Implications

1. Formalizing the LRH: The paper provides the first rigorous mathematical framework distinguishing between linear representation and linear accessibility, resolving ambiguities in previous literature.
2. Theoretical Justification for Superposition: It proves that LLMs can theoretically store exponentially more features than neurons, but only if the features are sparse and the model utilizes specific geometric arrangements (superposition).
3. Cost of Linearity: The quadratic dependence on $k$ ($d \propto k^2$) suggests a "cost" for linear accessibility. If a model needs to access many active features simultaneously ($k$ is large), the layer size must grow significantly. This offers a theoretical explanation for why deeper models or larger layers might be necessary for complex reasoning tasks involving many concurrent concepts.
4. Interpretability: The results clarify the role of Sparse Autoencoders (SAEs). The "true" direction of a feature is not necessarily the encoder weight; the representation vector and the probe vector can be distinct. However, for the probe to work, it must be orthogonal to the interference from other features.
5. Future Directions: The work opens avenues for studying multi-layer interactions, where subsequent layers might perform more complex operations than simple linear feature extraction, and exploring non-linear representation hypotheses.

In summary, the paper demonstrates that while the Linear Representation Hypothesis is powerful enough to explain the massive capacity of LLMs, the requirement for linear accessibility imposes strict geometric and dimensional constraints that distinguish it from general compressed sensing.