Geometric Factual Recall in Transformers

Abstract: How do transformer LLMs memorize factual associations? A common view casts internal weight matrices as associative memories over pairs of embeddings, requiring parameter counts that scale linearly with the number of facts. We develop a theoretical and empirical account of an alternative, \emph{geometric} form of memorization in which learned embeddings encode relational structure directly, and the MLP plays a qualitatively different role. In a controlled setting where a single-layer transformer must memorize random bijections from subjects to a shared attribute set, we prove that a logarithmic embedding dimension suffices: subject embeddings encode \emph{linear superpositions} of their associated attribute vectors, and a small MLP acts as a relation-conditioned selector that extracts the relevant attribute via ReLU gating, and not as an associative key-value mapping. We extend these results to the multi-hop setting -- chains of relational queries such as ``Who is the mother of the wife of $x$?'' -- providing constructions with and without chain-of-thought that exhibit a provable capacity-depth tradeoff, complemented by a matching information-theoretic lower bound. Empirically, gradient descent discovers solutions with precisely the predicted structure. Once trained, the MLP transfers zero-shot to entirely new bijections when subject embeddings are appropriately re-initialized, revealing that it has learned a generic selection mechanism rather than memorized any particular set of facts.

• The paper proposes a geometric mechanism for factual recall using transformers, achieving perfect single-hop accuracy with an embedding size of O(R log N).
• The relation-conditioned MLP selector effectively extracts fact components, and chain-of-thought reasoning circumvents exponential parameter growth in multi-hop settings.
• Empirical experiments on synthetic and real data validate the theory, demonstrating high linear decodability, accuracy, and successful zero-shot transfer.

Geometric Factual Recall in Transformers

Introduction and Motivation

The study addresses the fundamental mechanisms by which transformer-based LMs implement factual recall, particularly in multi-relational and multi-hop settings. Prior theoretical frameworks attributed memorization to associative memory in large parameter matrices and required embedding and parameter scaling proportional to the number of facts. This work departs from these algebraic accounts, proposing and validating a geometric mechanism where learned embeddings encode relational topology via linear superposition, leveraging a selection-based MLP architecture. The implications are non-trivial: embedding dimensionality scales merely logarithmically with subject set size, realizing efficient knowledge storage and recall.

Theoretical Analysis of Geometric Memorization

Single-Hop Recall and Logarithmic Embedding Dimension

The authors rigorously prove that, instead of encoding $N \times R$ facts through explicit key-value associations, a transformer with learned embeddings achieves perfect single-hop factual recall when the embedding dimension $d = O(R \log N)$. Each subject embedding is a concatenation (block vector) of its $R$ per-relation attribute vectors. At inference, the MLP serves as a relation-conditioned selector, extracting the desired component via ReLU gating. This capacity result, formalized in Theorem 1, marks a sharp distinction from prior random-embedding and associative memory analyses.

Moreover, the construction extends to a multi-head attention mechanism, where $R$ heads each implement selection over a block of the concatenated embedding (Figure 1 shows the effect of varying MLP width, highlighting minimal necessary width for successful recall).

Figure 1: Impact of MLP width on factual recall accuracy, demonstrating that small widths (e.g., $4d$) suffice when geometric superposition is utilized.

Parameter–Depth Tradeoff in Multi-Hop Recall

The analysis is extended to the $k$-hop case—answering composed queries via chains of relation functions. Here, the authors establish an information-theoretic lower bound (Theorem 2) on the tradeoff between embedding dimension and internal parameter count, scaling with the exponential of $k$ in the worst case without CoT. For instance, in direct multi-hop mapping, either the MLP width must be $\Omega(NR)$ or the embedding dimension must be $\tilde\Omega(R^k)$ to avoid information bottlenecks.

Importantly, they prove that inserting chain-of-thought (CoT) reasoning removes this exponential dependence: a single-layer transformer with $d = \tilde{O}(R + k)$ recovers full multi-hop capacity (Theorem 3). This highlights the computational and storage efficiency afforded by recursive, CoT-enriched inference, aligning with circuit complexity perspectives on autoregressive computation.

Empirical Verification

Comprehensive experiments on synthetic data validate the geometric memorization hypothesis. Systems with $d = O(R \log N)$0 subjects, $d = O(R \log N)$1 relations, and variable embedding dimensions are trained to recall attributes. The key outcomes are:

• Convergence: With learnable embeddings, perfect accuracy is achieved at embedding sizes matching the theoretical $d = O(R \log N)$2 threshold. The frozen-embedding controls only memorize when $d = O(R \log N)$3 is much higher, consistent with brute-force MLP lookups.
• Superposition Structure: Probing per-relation linear readout from subject embeddings shows near-perfect accuracy (Figure 2), confirming the theory’s prediction of linearly decodable superposition within embeddings.
• Relation-Conditioned Selector in MLP: Causal intervention experiments reveal that the MLP acts as a generic selector, with high intervention selectivity, verifying that the MLP executes relation-based extraction rather than acting as a memorizer (Figure 3).

  Figure 2: Linear readout accuracy per relation from subject embeddings and corresponding hidden states demonstrates robust linearly decodable superposition with learnable embeddings.

  

Figure 3: Intervention selectivity as a function of embedding dimension and relation count, indicating strong relation-specific selection by the MLP.

Further, the MLP, once trained, transfers zero-shot to new attribute bijections with high accuracy given smart re-initialization of the subject embeddings, isolating the MLP’s selection semantics from underlying factual associations.

Multi-Hop and CoT

In the multi-hop setting, experiments confirm the predicted expressivity gap: without CoT, performance degrades rapidly as $d = O(R \log N)$4 and $d = O(R \log N)$5 rise, but with CoT and geometric embeddings, perfect accuracy is maintained for all practical configurations (Figure 4). This supports the theoretical claim that CoT directly enables compositionality without exponential parameter growth.

Figure 4: Multi-hop recall accuracy across hops and relation count, showing CoT's effectiveness in eliminating exponential capacity blowup.

Probing Pretrained LMs

The geometric superposition structure was further probed in large pretrained LMs (Qwen, Llama, Phi) over real-world relational data. Layerwise probing demonstrates that subject hidden states encode linearly decodable output-embedding directions for attributes, with best-layer mean reciprocal rank (MRR) peaking in middle-to-high layers (Figure 5). This supports the generality of the geometric recall hypothesis, bridging synthetic constructs and real LMs.

Figure 5: Layerwise linear probe performance (MRR/Hits@1/10 per layer) across LMs on people-category relations, showing emergence and plateau of linearly-decodable structure mid-network.

Additionally, per-relation difficulty analysis exposes variance in decodability, indicating that specific relational topologies or surface complexity may modulate linear superposition fidelity (Figure 6).

Figure 6: Per-relation Hits@1 for Qwen3-14B on people-category, reflecting heterogeneity in factual relation encoding.

Contradictions with Prior Associative Memory Accounts

This work contradicts earlier models that necessitated large, often quadratic parameter counts for brute-force memorization of subject–relation–attribute mappings. By theoretically and empirically demonstrating efficient geometric solutions, it fundamentally shifts the consensus regarding necessary resource allocation for factual recall in neural sequence models.

Implications and Future Directions

The implications are broad for both AI theory and the practical deployment of large memory-augmented models:

• Resource Efficiency: The findings prescribe new design paradigms for embedding and parameter sizing in multi-relational models, supporting scalable integration of large fact bases with minimal overhead.
• Model Editing and Control: The identification of generic MLP selectors decoupled from factual content paves the way for systems where factual associations can be swapped or updated without retraining the full network.
• Compositionality and Reasoning: The capacity–depth tradeoff and CoT results demarcate the architectural boundaries for efficient multi-step reasoning and point to explicit chain-of-thought as a key mechanism for overcoming compositional inference bottlenecks.
• Deeper Models and Real Data: Observational alignment in real LMs suggests that these mechanisms extend beyond single-layer constructs, though further research is needed on datasets with rich semantic structure, correlated relations, and more complex relational graphs.

Conclusion

This work advances a precise geometric and mechanistic account of factual recall in transformers, shifting the focus from associative memorization to structured embedding superposition and selection. Both the theoretical bounds and experimental confirmations support a unified framework where relational knowledge is encoded efficiently, scalably, and with compositional flexibility. The implications extend to model architecture, editing, and our understanding of generalization and reasoning in emerging AI systems.