Learning to Recall with Transformers Beyond Orthogonal Embeddings

Nuri Mert Vural; Alberto Bietti; Mahdi Soltanolkotabi; Denny Wu

stat.MLMar 16, 2026

Learning to Recall with Transformers Beyond Orthogonal Embeddings

Authors: Nuri Mert Vural, Alberto Bietti, Mahdi Soltanolkotabi, Denny Wu

Organizations: University of Toronto and Vector Institute · Flatiron Institute · University of Southern California · New York University and Flatiron Institute

Abstract

Modern large language models (LLMs) excel at tasks that require storing and retrieving knowledge, such as factual recall and question answering. Transformers are central to this capability because they can encode information during training and retrieve it at inference. Existing theoretical analyses typically study transformers under idealized assumptions such as infinite data or orthogonal embeddings. In realistic settings, however, models are trained on finite datasets with non-orthogonal (random) embeddings. We address this gap by analyzing a single-layer transformer with random embeddings trained with (empirical) gradient descent on a simple token-retrieval task, where the model must identify an informative token within a length- $L$ sequence and learn a one-to-one mapping from tokens to labels. Our analysis tracks the ``early phase'' of gradient descent and yields explicit formulas for the model's storage capacity -- revealing a multiplicative dependence between sample size $N$ , embedding dimension $d$ , and sequence length $L$ . We validate these scalings numerically and further complement them with a lower bound for the underlying statistical problem, demonstrating that this multiplicative scaling is intrinsic under non-orthogonal embeddings.

Figures & tables

Panel of Figure 1 — Figure 1: (a) Empirical scaling of the parameter size required for a GD-trained one-layer transformer to learn factual recall, where we use $m\asymp d^{2}$ (see Section 4.2 for details). For small $L$ , the trained model achieves the optimal capacity $V\asymp md$ (purple line). As the sequence length $L$ increases, the scaling changes, suggesting a multiplicative rate (blue and red lines). (b) Phase diagram for the theoretical scaling of the parameter count given in Corollary 2 . Each region corresponds to a regime where a particular noise term in Theorem 1 is dominant. The parameter-size condition ( $md$ ) in each region is given in Corollary 2 .

Panel of Figure 2 — Figure 2: Empirical scaling of parameter size via three-step GD for the Attention-only model. In (a) and (b), top-left and top-right use $L\asymp V^{0.5}$ ; bottom-left and bottom-right use $L\asymp V$ . In (c), we compare the parameter counts from (a) and (b) for the $L\asymp V$ case under two sample-size regimes, $N\asymp V\log V$ and $N\asymp V^{1.5}$ . Line fitting: We identify in the heatmaps the smallest embedding dimension that achieves accuracies $\{0.1,0.125,0.15\}$ and perform a least squares fit. The slopes of the fitted lines and their theoretical counterparts are reported on the heatmaps.

Panel of Figure 3 — Figure 3: Empirical scaling of parameter size for the Attention-MLP model under two sample size regimes, $N\asymp V\log V$ and $N\asymp V^{1.5}$ . In (a) and (b), top-row uses $L\asymp V^{0.25}$ ; bottom-row uses $L\asymp V^{0.5}$ . In (c), we compare the parameter counts from (a) and (b) for the $L\asymp V^{0.5}$ case under both sample-size regimes.

Panel of Figure 4 — Figure 4: Empirical scaling of embedding parameter size for the Attention-MLP model under two width regimes, $m\asymp d^{2}$ and $m\asymp d^{3}$ . In (a) and (b), top-row uses $L\asymp V^{0.25}$ ; bottom-row uses $L\asymp V^{0.5}$ . In (c), we compare the parameter counts from (a) and (b) for the $L\asymp V^{0.5}$ case under both width regimes.

Panel of Figure 5 — Figure 5: Empirical scaling of the parameter size for the Attention-only model under two sample size regimes. Top row (a–d): $N\asymp V\log V$ . Bottom row (e–h): $N\asymp V^{1.5}$ . The model uses $L\asymp V$ and is trained using Adam over $16$ epochs. We observe that the capacity improves as the number of epochs increases.

Panel of Figure 6 — Figure 6: Empirical scaling of the parameter size for the Attention-only model with $N\asymp V\log V$ and $L\asymp V^{0.85}$ . The model is trained using Adam over $16$ epochs. The slopes are smaller than in Figure 5 , which is consistent with the shorter sequence length.

Panel of Figure 7 — Figure 7: Empirical scaling of the parameter size for the Attention-only model ( $L\asymp V$ ), trained with Adam using a batch size of $\lfloor N/16\rfloor$ . Top row (a–d): $N\asymp V\log V$ . Bottom row (e–h): $N\asymp V^{1.5}$ .

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 528 papers

Geometric Factual Recall in Transformers
May 12, 2026Shauli Ravfogel, Gilad Yehudai, Joan Bruna +1Neural Representation Geometry Memorization in Language Models
1New York University · 2Flatiron Institute
MLPs are Hebbians: Constructing Efficient Fact-Storing MLPs for Transformers
Jul 10, 2026Roberto Garcia, Jerry Liu, Ronny Junkins +3Transformer Memorization in Language Models
Institute for Computational & Mathematical Engineering, Stanford University · Department of Computer Science, Stanford University · Department of Computer Science and Engineering, University at Buffalo
Mixture of Chapters: Scaling Learnt Memory in Transformers
Mar 22, 2026Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda +2Memory-Augmented Language Models Persistent Memory for Language Models
IIT Kharagpur · Fractal AI

May 12, 2026J/K move · Enter open · S save

Geometric Factual Recall in Transformers

Shauli Ravfogel, Gilad Yehudai, Joan Bruna, Alberto Bietti

1New York University · 2Flatiron Institute

Neural Representation Geometry Memorization in Language Models Neural Network Memorization Memorization in Generative Models Factual Knowledge in Language Models

Open paper AskPDF arXiv

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, 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 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.

Abstract

Transformer

Explore similar work

Geometric Factual Recall in Transformers

MLPs are Hebbians: Constructing Efficient Fact-Storing MLPs for Transformers

Mixture of Chapters: Scaling Learnt Memory in Transformers