Technical Summary: Representational Capacity in Transformer Language Models

Problem Statement

The paper addresses a fundamental gap in understanding the geometric limits of feature representation within transformer-based language models. While the model dimension ($d_{model}$) is a primary hyperparameter controlling parameter count and latent space size, its specific role in determining the maximum number of representable features remains under-explored.

Existing theories, specifically the Linear Representation Hypothesis (features exist as directions in latent space) and the Superposition Hypothesis (models represent more concepts than dimensions by using near-orthogonal directions), suggest that the number of representable features is bounded by the geometry of the latent space. However, standard theoretical bounds, such as the Johnson-Lindenstrauss (JL) lemma, rely on assumptions of random projections. The paper argues that these standard bounds dramatically underestimate the packing efficiency of trained representations, failing to explain how models with dimensions in the thousands can support vocabularies and feature sets in the tens of thousands or millions.

Methodology

1. Estimating Accepted Deviation ($\epsilon$) via Embeddings

The authors propose using the embedding matrix as a measurable proxy for the near-orthogonality constraints of the entire latent space.

• Rationale: At initialization, random weights map orthogonal one-hot token vectors into $R^{d_{model}}$, resulting in near-orthogonal embeddings due to high-dimensional geometry. Training preserves this structure while developing a "right tail" in the cosine similarity distribution corresponding to meaningful lexical and semantic relationships.
• Metric Definition: The authors define $\epsilon$ (the accepted deviation from perfect orthogonality) as the boundary between incidental similarity and meaningful relationships. They estimate this threshold heuristically as $\mu + 2\sigma$ of the pairwise cosine similarity distribution of the embedding matrix, where $\mu$ and $\sigma$ are the mean and standard deviation.
• Classification: This metric is applied across dozens of open-source models to categorize them based on their adherence to near-orthogonality.

2. Deriving an Adjusted Capacity Formula

The paper critiques the standard JL lemma, which posits that the number of near-orthogonal vectors $k$ scales as $k \leq \exp(d \cdot \epsilon^2 / C)$.

• Empirical Failure: Applying the standard formula (even with empirically tightened constants) to trained models like Llama 2 7B yields capacity estimates (e.g., ~277 directions) that are orders of magnitude lower than the actual vocabulary size (32,000).
• Optimization Simulation: To characterize the packing efficiency of optimized vectors, the authors initialize random unit vectors and minimize a high-exponent penalty on the off-diagonal elements of the Gram matrix (simulating gradient descent).
• Formula Adjustment: By fitting the simulation data, the authors derive a modified relationship where capacity depends on the ratio of vectors to dimensions ($k/d$) rather than the raw count $k$ alone. The new relationship is:
  $$\epsilon = \sqrt{\frac{C \cdot \ln(k/d)}{d}} \implies k \leq d \cdot \exp\left(\frac{d \cdot \epsilon^2}{C}\right)$$
  This single-parameter modification reduces prediction error by two orders of magnitude compared to the standard JL bound.

3. Defining Representational Capacity

Combining the estimated $\epsilon$ from embeddings and the adjusted formula, the paper defines Representational Capacity as the quantitative upper bound on the number of distinguishable near-orthogonal directions available in the latent space for embeddings, unembeddings, and features.

Key Results

1. Two Distinct Model Classes

Analysis of dozens of models reveals a bifurcation into two classes based on $\epsilon$:

• Class 1 (High $\epsilon$): Models with lower $d_{model}$ and wide similarity distributions ($\epsilon > 0.2$). These models lack near-orthogonal structure at the embedding level and do not appear to leverage superposition in the same manner. Many in this class utilize tied embedding/unembedding matrices.
• Class 2 (Low $\epsilon$): Models with higher $d_{model}$ and tight distributions clustered near zero ($\epsilon < 0.1$). These models maintain strict near-orthogonality constraints, consistent with active use of superposition.

2. Capacity Estimates and Sensitivity

• Magnitude: For Class 2 models, representational capacity spans from $\sim 10^5$ to $\sim 10^{10}-10^{11}$. For example, Llama 2 7B ($d_{model}=4096, \epsilon=0.0645$) is estimated to have a capacity of roughly 40 million directions, comfortably accommodating its vocabulary and millions of features extracted via Sparse Autoencoders (SAEs).
• Sensitivity: Capacity is exponentially sensitive to $\epsilon$. A small increase in the tolerance for deviation (e.g., from 0.05 to 0.09) can result in a six-order-of-magnitude jump in capacity, making $\epsilon$ a more dominant lever than $d_{model}$ itself.
• Scaling Trends: Within Class 2, larger models tend to favor tighter orthogonality constraints (lower $\epsilon$) rather than maximizing raw capacity. This suggests a trade-off where larger models prioritize internal stability and reliable feature retrieval over raw geometric room.

Significance and Claims

The paper claims to provide the first quantitative framework for estimating the geometric limits of feature representation in transformers. Its primary contributions are:

1. Empirical Correction of Theoretical Bounds: It demonstrates that standard geometric bounds (JL lemma) are insufficient for trained models and provides an empirically adjusted formula that aligns with observed packing efficiency.
2. Quantification of a Shared Resource: It defines "Representational Capacity" as a finite geometric resource shared by embeddings, unembeddings, and features, governed by the model's tolerance for overlap ($\epsilon$).
3. Diagnostic Insight: The framework reveals that not all models utilize superposition equally; the distinction between high-$\epsilon$ and low-$\epsilon$ models suggests fundamental structural differences in how they organize latent space.
4. Design Implications: The authors suggest that $d_{model}$ selection, currently often heuristic (powers of 2), could theoretically be guided by capacity targets. They hypothesize a potential "ceiling" on useful capacity, beyond which increasing geometric room yields diminishing returns, and that larger models may naturally converge toward tighter orthogonality to maintain stability.

Limitations and Caveats

The paper explicitly notes several limitations:

• Hypothetical Generalization: The assumption that embedding geometry reflects the broader latent space (including intermediate features) is a motivated hypothesis, not a proven fact.
• Heuristic Threshold: The $\mu + 2\sigma$ threshold for $\epsilon$ is heuristic. Because capacity is exponentially sensitive to $\epsilon^2$, absolute capacity values are unreliable, though relative comparisons between models remain informative.
• Geometric vs. Computational: The derived capacity is a bound on geometric possibility, not on computational utilization. It does not account for whether the model's architecture (depth, attention heads) can effectively learn or process the maximum number of available directions.
• Data Scarcity: The observation that larger models favor tighter orthogonality is based on limited data points with $d_{model} > 6000$.