PaperlayerView exact revision on arXiv

Source-linked AI summary

Geometric Factual Recall in Transformers

Shauli Ravfogel, Gilad Yehudai, Joan Bruna, Alberto Bietti

arXiv:2605.12426v1cs.CL

TL;DR

Transformers can memorize many relational facts, but how their representations support factual recall remains unclear. This paper develops geometric constructions and experiments showing that logarithmic-dimensional embeddings can encode relational superpositions, while a small MLP selects queried attributes and extends to multi-hop recall with a capacity–depth tradeoff.

  • Problem

    The paper asks how transformers memorize factual associations and what latent representations support recalling and composing facts.

  • Method

    The paper analyzes single- and multi-hop factual recall using geometric embedding constructions, capacity bounds, lower bounds, and controlled synthetic experiments.

  • Results

    Embedding dimension d = O(R log N) suffices for single-hop recall, while multi-hop recall exhibits a capacity–depth tradeoff that chain-of-thought can reduce to the single-hop case.

  • Takeaways & Limitations

    Learned subject embeddings encode relational superpositions and a relation-conditioned MLP selects attributes, with gradient descent recovering this structure and transferring to new bijections without retraining.

  • Takeaways & Limitations

    The setting uses opaque relations and arbitrary attributes, leaving extension to semantically structured data and deep models for future work.

Abstract

from arXiv · show

How do transformer language models 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.

1 Introduction

The paper argues that transformers can memorize relational facts geometrically, encoding structure in embeddings rather than relying on unstructured key-value associations. It develops logarithmic-capacity constructions, multi-hop capacity-depth results, and experiments showing gradient descent discovers the predicted mechanisms.

  • Motivation: Transformers must store and compose richly structured facts involving multiple relations over the same entities.The motivating examples include querying an entity’s birthplace or occupation and composing relations such as a book author’s mother.
  • Geometric memorization: Existing theories model memorization as storing input-output embedding pairs, whereas recent evidence points to geometric representations encoding relational topology.The contrasted algebraic accounts use unstructured, nearly orthogonal vectors, while learned embeddings appear to reflect stored relational structure.
  • Main contributions: A single-layer transformer memorizes N subjects across R relations with d = O(R log N) using superposed attribute vectors and relation-conditioned ReLU decoding.Each (subject, relation) pair maps to one attribute, and the construction directly encodes relational structure in embedding geometry.
  • Empirical validation: Experiments confirm the predicted capacity threshold d = Θ(R log N) across embedding dimensions and numbers of relations.The synthetic shared-attribute experiments test whether gradient descent discovers the theoretically predicted structure.
  • Multi-hop recall: Multi-hop recall exhibits a capacity–depth tradeoff, with separate constructions for solving k-hop queries with and without CoT.The introduction identifies this tradeoff as a central theoretical result for multi-hop relational queries.

2 Preliminaries and Problem Setup

The paper formulates multi-hop factual recall as composing relation-specific bijections over a shared subject and answer set. It studies autoregressive transformers whose layers combine multi-head self-attention, residual connections, and position-wise ReLU MLPs.

  • Problem Setup: Each relation r ∈ [R] is a ground-truth bijection g_r: [N] → [N], and the target answer composes k relations applied to the initial subject.Training sequences have the form (s_0, r_1, …, r_k, y), with y = (g_rk ◦ ··· ◦ g_r1)(s_0).
  • Problem Setup: At inference, the model predicts the final token by maximizing its conditional probability given the subject and relation prefix.The prediction is ŷ = arg max_v∈[N] p_θ(v | s_0, r_1, …, r_k).
  • Problem Setup: The single-hop case k = 1 generalizes factual recall with a shared codomain, unlike frameworks using mutually disjoint attribute codomains.The shared codomain [N] induces a non-trivial extension relative to the cited framework.
  • Transformer Architecture: Each transformer layer applies multi-head self-attention to token representations, adds a residual connection, and then applies a position-wise MLP.The MLP acts independently on each token, uses hidden dimension m, and consists of two linear transformations with a ReLU activation.
  • Transformer Architecture: Theoretical constructions may use a two-hidden-layer MLP, while the analysis omits normalization and experiments use standard normalized transformer layers.Degenerate identity normalization layers can implement the theoretical omission without limiting generality.

3 Related Work

Prior work often models factual recall as component-localized or associative key–value memory, while this paper shows that learnable embeddings enable logarithmic dimension through superposition. It also situates the work among memorization bounds and circuit-complexity results on capacity–depth tradeoffs with and without chain-of-thought.

  • Component-localized factual recall: Prior studies localize factual recall to FFN or MLP components, with attention transporting subjects and targeted weight updates enabling factual editing.These works variously interpret FFN layers as key–value memories, identify mid-layer MLPs as dominant stores, and trace recall to MLP lookup.
  • Associative-memory models: Associative-memory accounts represent transformer weights as outer products over near-orthogonal embedding pairs and analyze the resulting matrix capacity.This line has roots in neural computation and includes capacity analyses in simplified settings.
  • Geometric memorization: Theorems 4.1 and F.1 show that learnable embeddings reduce dimension to logarithmic scale without inflating the MLP, using superposition rather than near-orthogonal random codes.This directly addresses open questions about whether learnable embeddings reduce d and what structure the resulting solution adopts.
  • Memorization bounds: Worst-case memorization bounds depend on example count and sequence length, whereas Dugan et al. store f : [K] → [V] in Θ(K log V) parameters, matching the information-theoretic lower bound.The cited bounds do not exploit relational structure; the closest construction uses an encoder–decoder MLP.
  • Capacity–depth tradeoff: With CoT, d = O(R log N) suffices for any number of hops; without CoT, either dimension or MLP width grows exponentially or linearly in k.Theorem 4.3 and Theorem 4.4 establish this capacity–depth tradeoff, paralleling circuit-complexity results on bounded-depth transformers and CoT.

4 Factual Recall with Learned Representations

Learned embeddings enable geometric factual recall: a logarithmic-dimensional embedding can store relational structure while a small MLP selects the relevant attribute. For multi-hop recall, zero-shot capacity depends on a trade-off between embedding dimension, parameter capacity, and the use of chain-of-thought.

  • 1-hop factual recall: d = 4R log(N) + 1 suffices for a 1-layer transformer with a 3-layer MLP of width R to solve single-hop recall.An alternative uses R attention heads with head dimension dh = 4 log(N) and d = 4R log(N) + 4 log(R) + 1.
  • 1-hop factual recall: Subject embeddings stack representations of all R possible relation answers, while the MLP acts as a relation-conditioned selector rather than a lookup table.This achieves d = O(log N) with MLP width independent of N in the shared-attribute regime.
  • Multi-hop lower bound: D < R implies W ≥Ω(N log N), whereas D > Rk makes both global and local capacity constraints vanish.Small embeddings force the model to store NR facts in its weights; sufficiently large embeddings can encode each subject’s complete k-hop evaluation tree.
  • Multi-hop constructions: Key-Value Memory uses d = ˜O(k) and MLP width O(N · R), while Embedding Pre-computation uses d = ˜O(Rk) and MLP width ˜O(Rk).These constructions demonstrate exponential blowup in either embedding dimension or MLP width, up to suppressed logarithmic factors.
  • Chain-of-thought: CoT reduces the required depth from k to 1 and yields a 1-layer solution with d = ˜O(R + k) and MLP width ˜O(R).Autoregressive intermediate outputs let the model solve the 1-hop problem locally at each generation step.

5 Experiments

Experiments show that gradient descent discovers the predicted geometric recall structure: learned embeddings linearly superpose relation-specific attributes, while the MLP acts as a generic relation-conditioned selector. This structure enables perfect memorization at lower dimension, transfers to new bijections without retraining, and appears as relational signal in pretrained LMs.

  • Experimental setup: With N = 4096 subjects, random bijections, R ∈ {2, 4, 8, 10, 12, 14, 16}, and d ∈ {32, 64, 128, 256, 512, 768}, the model is trained on all N × R pairs.Accuracy is defined as the fraction of subject–relation pairs classified correctly.
  • Scaling and memorization: d ≥128 yields perfect memorization for every R with learned embeddings, whereas frozen embeddings require d ≥512 for R = 16.The frozen-embedding result is consistent with the associative-memory parameter-count bound ˜Ω(NR).
  • Superposition structure: Trainable embeddings achieve near-perfect per-relation linear readout accuracy, showing that each relation’s attribute is encoded as a linearly decodable component of the subject embedding.Random frozen embeddings instead collapse to chance for subject-embedding readouts because they cannot encode relation-specific attribute information.
  • MLP as a generic selector: Across sufficient-capacity configurations, the MLP has high selectivity and follows counterfactual relation changes while leaving other relations unaffected.This matches the relation-specific ReLU-gating selector predicted by the construction.
  • MLP as a generic selector: After reinitializing embeddings for fresh random bijections, the frozen MLP achieves substantial accuracy without retraining, demonstrating generic relation-conditioned selection rather than memorization of specific associations.The new embeddings encode the fresh attributes using the smart initialization.
  • Pretrained language models: Across five pretrained LMs and six categories, linearly probed relational signal exceeds random-vocabulary chance and a constant majority baseline, with best-layer MRR reaching 0.69 on people and 0.44 on films.Intermediate values are 0.58 for companies, 0.56 for buildings, and 0.55 for programming languages.

6 Discussion

The paper presents geometric factual memorization, showing that subject embeddings can encode relational structure while a relation-conditioned ReLU MLP selects relevant attributes. It extends this account to multi-hop reasoning and identifies limitations of the controlled setting for natural-language knowledge.

  • Geometric factual memorization: Dimension d = O(R log N) suffices for memorizing random subject–attribute bijections, with subject embeddings storing linear superpositions and a ReLU MLP selecting attributes.The MLP is relation-conditioned rather than an associative key-value mapping.
  • Multi-hop reasoning: The geometric account extends from single-hop memorization to multi-hop reasoning with a provable capacity–depth tradeoff without chain-of-thought.The supplied passage indicates this extension but is truncated before stating the full multi-hop construction.
  • Limitations and future directions: The controlled setting omits semantic regularities found in real factual knowledge, including correlated relations, typed attributes, polysemy, multi-token entities, and long-tailed frequencies.The passage motivates extending the theory so embedding geometry reflects these regularities rather than encoding facts independently.

Acknoweldgements … B.1 Architecture

The appendix situates the paper’s capacity claims within prior work on factual storage, tensor decompositions, and neural-network memorization, then specifies the shared transformer architecture and experimental setup. It also records the paper’s acknowledgement and computational hardware.

  • Acknoweldgements: The authors thank Yanai Elazar for valuable comments.
  • A Additional Related Work: Closed-book QA and the Physics of Language Models series provide empirical precedents for studying factual-storage capacity and extractability.Prior work reports that storable knowledge scales with model size and that manipulating stored facts can require more capacity than retrieval.
  • A Additional Related Work: Tensor-decomposition methods model multi-relation fact storage by factorizing an N × R × N relation tensor.For random bijections, each relation slice is a permutation matrix of rank N, while worst-case CP rank for the full tensor is O(NR).
  • A Additional Related Work: Prior memorization theory establishes bounds for ReLU networks and transformers, including logarithmic width in robust regimes and universal approximation via quantize-and-memorize.These bounds depend on the number of examples and sequence length and do not directly capture the paper’s geometric formulation.
  • B Experimental Details: All experiments were conducted on a single H200 node.
  • B.1 Architecture: Single-hop and multi-hop experiments use a single-layer, one-head transformer with a two-layer GELU MLP, expansion ratio 4, Pre-LayerNorm, and RMSNorm.Subjects and attributes come from one size-N embedding pool, with independent input and output projections indexed by the shared entity vocabulary.

B.2 Single-Hop Training · B.3 Analysis Methodology · C Additional Experimental Results

The paper specifies the single-hop training protocol, defines post-training tests for superposition and relation-specific selection, and reports final accuracy for models with frozen or trainable embeddings. These analyses include linear readouts, causal interventions, and zero-shot transfer under smart reinitialization.

  • B.2 Single-Hop Training: Models use AdamW for 15,000 steps with batch size 1024, gradient clipping at norm 1, and early stopping below 10−4.Training uses learning rate 1.0, weight decay 0.1, and three seeds per (d, R) configuration, with mean ± standard deviation reported.
  • B.2 Single-Hop Training: The frozen-embedding control holds the input entity embedding table at random initialization while training attention, MLP, and output-projection weights.All other settings remain identical to the trainable-embedding condition.
  • B.3 Analysis Methodology: Linear readouts fit one ridge-regression map W_r per relation and evaluate held-out-subject classification to test whether embeddings encode superpositions of attribute vectors.High per-relation accuracy would support recovery of each attribute from the predicted superposition.
  • B.3 Analysis Methodology: Causal interventions swap a queried relation’s attribute using minimum-norm embedding perturbations derived from the readout pseudoinverse.The perturbation uses the difference between substituted and original attribute vectors and a rank-truncated Moore–Penrose pseudoinverse.
  • B.3 Analysis Methodology: Selectivity combines whether the queried relation changes and whether all other relations remain stable, using their geometric mean across rank-k truncations.The best selectivity per cell is reported because the pseudoinverse is sensitive to small singular values.
  • B.3 Analysis Methodology: The MLP freeze experiment tests zero-shot transfer by freezing the trained MLP and reinitializing subject embeddings for fresh random bijections.Subject embeddings are reinitialized either randomly as the baseline or through a smart initialization that solves for the new superposition encoding.
  • B.3 Analysis Methodology: Smart initialization should yield high zero-shot accuracy without retraining if the MLP learned a generic relation-conditioned selector; associative storage instead predicts chance accuracy.This contrasts selector behavior with the hypothesis that the MLP stores subject–attribute mappings.
  • C Additional Experimental Results: Figure 4 reports final accuracy for models trained with frozen or trainable embeddings.The passage identifies the comparison but does not provide numerical accuracy values.

C.1 Multi-Hop Experiments

Multi-hop experiments test the predicted capacity–depth tradeoff in a single-layer transformer. With chain-of-thought and learned embeddings, the model reaches near-perfect accuracy across all tested configurations, while performance without chain-of-thought degrades as hops and relations increase.

  • Capacity–depth tradeoff: Without chain-of-thought, k-hop queries require either embeddings of dimension d = O(Rk log N) or an MLP of width O(N · R), whereas chain-of-thought needs d = O(R log N).These constructions predict that chain-of-thought removes the dependence of embedding dimension on the number of hops.
  • Experimental setup: The experiments fix N = 2048 and d = 256, sweeping k ∈ {1, 2, 4} hops and R ∈ {2, 4, 8, 12, 16} relations across four chain-of-thought and embedding-freezing conditions.The design compares learned versus randomly initialized frozen embeddings, with one learned attention head and learned positional embeddings.
  • Training protocol: Models use AdamW with learning rate 10^-2, weight decay 0.1, batch size 1024, gradient clipping at norm 1, and up to 15,000 training steps.Three seeds are run per condition, and results are reported as mean ± std.
  • Results: With chain-of-thought and learned embeddings, the model achieves near-perfect accuracy across all configurations, including k = 4 hops with R = 16 relations.This result is reported for the top-right panel of Fig. 5.
  • Results: Without chain-of-thought and with learned embeddings, performance degrades sharply as k and R increase, despite d = 256.The passage attributes this degradation to insufficient embedding dimension for the increasing multi-hop and relational complexity.

D Frozen-Embedding Control: Full Results … E.2 Probe training details

Frozen embeddings force substantially larger dimensions and eliminate the geometric readout and zero-shot transfer mechanisms, while the real-LM probing appendix specifies confidence-filtered data generation and a fixed-rank linear-probe training recipe.

  • D Frozen-Embedding Control: Full Results: D Frozen-Embedding Control: Full Results: R = 16 relations require d ≥512 with frozen embeddings, versus d ≥128 when embeddings are learned.The frozen model still achieves perfect memorization, but at a substantially higher embedding dimension.
  • D Frozen-Embedding Control: Full Results: D Frozen-Embedding Control: Full Results: Subject-embedding linear readouts remain at chance, and post-attention hidden-state readouts stay low across the grid.Random fixed embeddings do not encode relation-specific attribute information by construction.
  • D Frozen-Embedding Control: Full Results: D Frozen-Embedding Control: Full Results: Smart-initialization zero-shot accuracy on new bijections g′ collapses to near chance without learned per-relation readouts.Retraining subject embeddings can recover high accuracy by directly re-encoding g′ in the embedding table.
  • E Real-LM Probing — Details: E Real-LM Probing — Details: The appendix documents subject–attribute data generation, probe training, and full per-(model, category) results.These details support the experiment in Section 5.3.
  • E.1 Data generation pipeline: E.1 Data generation pipeline: Six categories are assembled by prompting Claude in two stages, applying a confidence-or-skip filter at the entity level.The categories are people, companies, films, species, buildings, and programming languages.
  • E.1 Data generation pipeline: E.1 Data generation pipeline: Batches contain BATCH_SIZE=50 entities and continue until fewer than EARLY_STOP_NEW_UNIQUE = 8 new uniques appear or approximately 2000 entities are reached.Each category defines an entity kind and relation-specific field descriptions.
  • E.1 Data generation pipeline: E.1 Data generation pipeline: Entities are included only when every attribute is certain, with canonical real entities preferred and fewer than the batch ceiling allowed.The pipeline forbids guessing, placeholders, and fictional or composite entities.
  • E.2 Probe training details: E.2 Probe training details: Each category–relation pair uses a biased LinearProbe(d_in →rank → d_out) with rank k = 512, AdamW at 10−3 learning rate and weight decay, and 200 epochs.Training uses patience-5 early stopping and the cosine objective L = 1 −cos(Wrsℓ+ br, D[to]).

E.3 Full results

Across all evaluated language models, factual-recall performance improves from token-embedding lookup through the middle layers, then plateaus in the upper stack. At the best layer, performance is substantially higher than at layer 0, with relation difficulty varying across attributes.

  • E.3 Full results: Best-layer MRR is significantly higher than layer-0 MRR across every evaluated model and category.Layer 0 corresponds to token-embedding lookup, while the best layer lies in the latter half of the network.
  • E.3 Full results: Across all five models, MRR, Hits@1, and Hits@10 rise through the middle layers and plateau in the upper layers.Layer 0 carries non-trivial signal, and some models show a slight decline at the very top.
  • E.3 Full results: At Qwen3-14B’s best layer, gender and country_of_birth are easier relations than religion and occupation.The per-relation breakdown identifies substantial variation in recall difficulty across attributes.

E.4 Training Transformer Language Model on Natural Language Relational Data · F Disjoint Attributes · G Proofs from Section 4

The natural-language Transformer memorizes its corpus while encoding relational answers linearly, with performance depending on layer and tokenization. In the disjoint-attribute setting, logarithmic-dimensional constructions follow from low-correlation sign vectors and uniform-attention composition.

  • E.4 Training Transformer Language Model on Natural Language Relational Data: A 1-layer GPT-style decoder with d=256 and ∼7.3M parameters effectively memorizes the natural-language corpus, reaching 99.4% greedy Hits@1.The model uses one attention head, Pre-RMSNorm, untied input/output embeddings, learned absolute position embeddings, and MLP width 4d.
  • E.4 Training Transformer Language Model on Natural Language Relational Data: 0.71 vs. 0.52 mean full-vocabulary Hits@1 shows sharper linear decoding from layer 0 than layer 1.Per-category averages are 0.85/0.57 for people, 0.80/0.63 for programming_languages, and 0.73/0.68 for companies.
  • E.4 Training Transformer Language Model on Natural Language Relational Data: Over 97% linear readout occurs when each entity receives a single token, exceeding the multi-token word-level setup.The passage identifies tokenization as a factor and leaves its effects on linear relational encoding for future study.
  • F Disjoint Attributes: d = 4 log(NR) suffices for a 1-layer transformer to solve factual recall when each relation’s attribute set is disjoint.The setting has |A| = R · N and requires disjointness between attribute sets associated with distinct relations.
  • F Disjoint Attributes: The construction encodes subjects and relations with sign vectors and outputs the attribute gj(i) on the second token.Subject vectors use dimension 4 log(N), relation vectors use dimension 4 log(R), and the combined embedding dimension is their sum.
  • F Disjoint Attributes: Uniform attention with WV = 2 · Id averages the two input tokens, after which the output matrix recovers the exact attribute token.The input tokens are xi and rj, and the second-token output represents gj(i).
  • G Proofs from Section 4: For any k ≥2, there exist {±1}4 log(k) vectors whose pairwise inner products satisfy |⟨vi, vj⟩| ≤3 log(k).Random sampling, Hoeffding’s inequality, and a union bound establish a non-zero probability of obtaining these vectors.

G.1 Proof of Theorem 4.1 … G.4 Proof of Theorem 4.4

The proofs construct geometric solutions for one-hop and multi-hop relational recall, using compact embeddings with MLP or attention-based selectors. They also establish storage and decoding lower bounds and give explicit chain-of-thought and no-chain-of-thought constructions.

  • G.1 Proof of Theorem 4.1: 4 log(N) coordinates encode attribute vectors, while a 3-layer ReLU MLP with O(R log(N)) neurons selects the relation-matched vector.The construction uses embedding dimension d = 4R log(N) + 1 and outputs the correct vector gj(i).
  • G.1 Proof of Theorem 4.1: R attention heads replace the MLP selector: each relation head extracts the correct attribute vector, with value/output head dimension 4 log(N) and internal key/query dimension one.For sufficiently large β, the relation token attends to the subject token for its relation and otherwise attends to itself.
  • G.2 Proof of Theorem 4.2: The lower bound combines global storage and local decoding constraints for zero-error k-hop evaluation across all R-regular relation graphs.The global argument counts (N!)R possibilities modulo right-translation, while the local argument extracts disjoint k-hop trees.
  • G.3 Proof of Theorem 4.3: D = 1 + dS + 2dR + k supports a chain-of-thought construction whose MLP width is O(N · R) for one-step subject-relation lookup.The changing token extracts relations sequentially, and the MLP projects a one-hot subject-relation representation to the next subject.
  • G.3 Proof of Theorem 4.3: D = Rk · dS + dR + k + 1 gives a no-CoT construction that stores the entire k-hop answer tree in the initial subject embedding.Each layer fetches the preceding subtree and uses an MLP of width O(Rk ·dS) to select one of R branches.
  • G.4 Proof of Theorem 4.4: D = 1 + dS + dR + R · dS + (2k + 1) supports autoregressive chain-of-thought generation by embedding each subject’s complete 1-hop topology.A single attention head routes the needed subject and relation information, while an O(R · dS)-width Boolean selector computes each hop.
  • G.4 Proof of Theorem 4.4: After k autoregressive steps, highest-correlation dictionary lookup emits the correct chain-of-thought subjects and final answer sk.The generated token sequence proceeds through s1, …, sk, with each next token selected from its subject coordinates.
Loading 2605.12426v1…