Research map / Architectures and efficiency / Attention and transformers

Geometric Factual Recall in Transformers

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Transformer expressivity and learnability · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proves and empirically verifies that a single-layer transformer with learned embeddings can memorize NR random subject-relation-attribute bijections with embedding dimension O(R log N), storing subjects as linear superpositions of attribute vectors decoded by a ReLU MLP selector, and analyzes multi-hop capacity-depth tradeoffs with and without CoT.

Problem

How transformers memorize factual associations when embeddings are learned, and how much embedding dimension and MLP capacity multi-hop factual recall requires, especially with shared attribute codomains.

Approach

Learned subject embeddings encode linear superpositions of their R attribute vectors, and a small relation-conditioned ReLU MLP acts as a selector rather than an associative key-value store; this gives logarithmic embedding dimension, in contrast to near-orthogonal random embedding accounts requiring parameters linear in NR.

Key results

• Learned embeddings reach perfect memorization for all R when d >= 128, while frozen random embeddings need d >= 512 for R=16.
• Trainable-embedding minimal dimension grows as a + b log2 N (matching Theta(R log N)), whereas frozen embeddings grow as a power law in N with exponent about 1 or 1/2.
• Per-relation linear readouts from learned subject embeddings are near-perfect, and the frozen MLP achieves high zero-shot accuracy on new bijections with smart initialization of subject embeddings.

Surprising findings

• Frozen random-embedding MLP fails zero-shot transfer: the smart-init zero-shot column collapses, meaning the memorizer must be retrained for new facts.
• Best linear readout layer for real LMs lies at roughly 50-80% depth, and layer-0 signal is weaker than mid-late layers.

Limitations

• Theory and main experiments target a single-layer transformer with synthetic random bijections and arbitrary entity tokens, far from real factual knowledge.
• The lower bound leaves an intermediate regime (D between R and R^k) not tight.
• Real-LM probing evidence is correlational and uses a linear probe, not a causal test.

Experimental setup

Datasets: synthetic random bijections over 4096 subjects Claude-generated entity-attribute corpus (4,610 entities, 42 relations)

Benchmarks: synthetic shared-attribute factual recall multi-hop synthetic recall real LM linear relational probes across five categories

Metrics: accuracy linear readout accuracy selectivity zero-shot accuracy with smart initialization Hits@1 MRR

Baselines: frozen random embedding control

Builds on: associative memory (key-value view of FFN layers) fixed random spherical embedding factual recall construction ROME-style factual association localization

Closest papers by meaning

• Understanding Factual Recall in Transformers via Associative Memories (ICLR 2025 · similarity 0.89)
• Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformers (NeurIPS 2024 · similarity 0.88)
• Scaling Laws for Associative Memories (ICLR 2024 · similarity 0.87)
• Linearity of Relation Decoding in Transformer Language Models (ICLR 2024 · similarity 0.87)
• The Structure of Relation Decoding Linear Operators in Large Language Models (NeurIPS 2025 · similarity 0.87)
• A mathematical theory of balancing relational generalization and memorization (NeurIPS 2026 · similarity 0.86)