Understanding Factual Recall in Transformers via Associative Memories
Authors: Eshaan Nichani
Jason D. Lee
Alberto Bietti
Venue: ICLR 2025 Spotlight (2025)
Keywords: transformers
associative memories
factual recall
storage capacity
training dynamics
Primary area: learning theory

Abstract: Large language models have demonstrated an impressive ability to perform factual recall. Prior work has found that transformers trained on factual recall tasks can store information at a rate proportional to their parameter count. In our work, we show that shallow transformers can use a combination of associative memories to obtain such near optimal storage capacity. We begin by proving that the storage capacities of both linear and MLP associative memories scale linearly with parameter count. We next introduce a synthetic factual recall task, and prove that a transformer with a single layer of self-attention followed by an MLP can obtain 100\% accuracy on the task whenever either the total number of self-attention parameters or MLP parameters scales (up to log factors) linearly with the number of facts. In particular, the transformer can trade off between using the value matrices or the MLP as an associative memory to store the dataset of facts. We complement these expressivity results with an analysis of the gradient flow trajectory of a simplified linear attention model trained on our factual recall task, where we show that the model exhibits sequential learning behavior.

Published as a conference paper at ICLR 2025
UNDERSTANDING FACTUAL RECALL IN TRANSFORMERS VIA ASSOCIATIVE MEMORIES
Eshaan Nichani
Jason D. Lee Princeton University
Alberto Bietti Flatiron Institute
Princeton University
ABSTRACT
Large language models have demonstrated impressive ability to perform fac- tual recall. Prior work has found that transformers trained on factual recall tasks information
rate proportional their parameter count. In work, we show that shallow transformers can use a combination of associative memories obtain such near optimal storage capacity. We begin by proving that the stor- age capacities of both linear and MLP associative memories scale linearly with parameter count. We next introduce a synthetic factual recall task, and prove that transformer with single layer of self-attention followed by an MLP can obtain 100% aceuracy on the task whenever either the total number of self-attention pa- rameters MLP parameters scales (up log factors) linearly with the number of facts. In particular, the transformer can trade off between using the value matrices the MLP ment these expressivity results with an analysis of the gradient flow trajectory of a simplified linear attention model trained on our factual recall task, where show that the model exhibits sequential learning behavior.
associative memory to store the dataset of facts. We comple-
INTRODUCTION
One hallmark capability of transformer-based large language models (LLMs) is factual re- call (Petroni et al.] 2019; Diang et al., 2020) Roberts et al. (020). Given a prompt of the form "In what year was George Washington born?" an LLM will correctly respond with "1732." Lan- guage models thus act as databases, storing somewhere their parameters mappings of the form (George Washington, birth year) (1732) which can be easily accessed during inference time.
Prior work (Allen-Zhu & Li 2024) has observed that transformers trained on factual recall tasks can store information at rate proportional to their parameter count. Other studies (c.g.. Meng jet al. 2022, Geva et al.) 2023 Nanda et al. 2023 Ly et al. 2024) have sought to understand the specific mechanism by which transformers implement factual recall, probing models to understand specifically which transformer blocks "contain" certain facts. However, these studies do not consider the memorization capacity of such constructions, and it is thus open question to understand how transformers optimally encode such factual information within their weights.
In this work, we show that shallow transformers can use combination of associative memories to obtain near-optimal storage capacity for factual recall tasks. Associative memories store pairs of input-output embeddings through their outer products, and thus well-suited for modeling the weight matrices of a transformer. Prior work (Bietti et al., 2023) has shown that this associative memory model is a key primitive towards understanding both the representational capacity and optimization dynamics of transformers on synthetic tasks.
Our specific contributions are as follows: Section we begin by studying the ability of linear and MLP associative memory models to store associations between discrete vocabularies. We prove that when the embeddings are sampled randomly over the sphere, these models can store a number of associations proportional to their parameter count, significantly improving over the case where the embeddings are orthogonal. In Section we introduce a synthetic next-token prediction task which models factual recall. The data distribution consists of prompts containing subject token and relation token r hidden amongst a set of noise tokens, which the learner must map to a ground truth answer a" rs,r). Our main theorem is that transformer consisting of single multi-head

---

Published as a conference paper at ICLR 2025
self-attention layer followed by an MLP can obtain 100% accuracy when either the number of self- attention parameters or MLP parameters scales (up to logs) proportionally with the dataset size. In Section [5] study the gradient descent dynamics of single linear self-attention head trained on the synthetie task. We prove that the model undergoes a sequential learning dynamics, consisting of "hallucination" stage where the model outputs the conditional distribution for the answer based on only the relation. Finally, in Section 6 we complement our constructions with lower bounds, showing that they are optimal up to logarithmic factors. Overall, our work makes progress towards understanding the mechanism by which transformers learn and store factual information.
2 RELATED WORK
Associative memories. Associative memories have a long history in the neural computation litera- ture (Hopfield 1982 Kohonen 1972. Willshaw et al. 1969). More recently there has been renewed interest in extensions of such models with larger capacity (Krotov
Demircigil et al. 2017 Lucibello & Mézard 2024). These have been linked to the attention blocks in Trans- formers (Ramsauer et al.) 2020; Schlag et al., (02T), with (Le et al., 2020 Hoover et al. 2023) in particular using the connection between self-attention and associative memories to design new variants of the attention module. (Radhakrishnan et al. 2020) show that overparameterized autoen- coders can also behave as associative memories. However, these connections diffeers from our work, where we consider instead the role of both self-attention and MLP weights as associative memories, Cabannes et al. 2024
Hopfilld 2016
in similar vein to (Bietti al. 2023
Memorization and factual recall. Large language models are known to store vast amounts of factual knowledge in their weights (Jiang et al. 2020 Roberts et al. 2020 Geva et al. 1021). Sev-
eral recent works in the mechanistic interpretability literature have attempted to understand how
transformers store facts (Meng al.) [2022] Geva al. 2023 Nanda et al. 2023 Lv et al.) 2024).
Allen-Zhu & Li (024) empirically studied the memorization capacity for Transformer language models of different sizes trained on synthetic factual recall tasks, and observed near-linear scaling with the number of parameters. Jiang et al.] (2024) demonstrate how shallow transformers can solve a related latent concept association task by viewing the weight matrices as associative memories. At more basic level, several works have studied the memorization capacity of neural networks, using constructions that differ from our associative memory approach, both in the context of re- gression (Bubeck et al., 20000 Vardi et al. [202]] Madden & Thrampoulidss] [024] and (next) token prediction (Mahdavi et al. 2023; Kajitsuka & Sato 2023 2024 Madden ct al. 2024).
Gradient dynamics. Training dynamics of transformer models on various tasks has been a pop- ular recent line of research (Jelassi et al. [0022] Snell et al. 2021 Li et al.) 2023 Bietti et al.)
2023 Tian et al. [223]] Nichani et al.) 2024). Zhang et al. (2024); Mahankali et al. (2024) studied training dynamics of transformers with linear attention on in-context learning tasks. Ghosal et al. 2024) studied the fine-tuning dynamics a similar factual recall task, showing how training
lesser-known facts may hurt performance. Our emphasis differs in that we consider non- non-orthogonal embeddings, and require the model to additionally filter out the relevant subject and relation tokens from the noise tokens, which requires learning of the key and query matrices.
ASSOCIATIVE MEMORIES
In this section, we show that associative memories have storage capacity the order of the number of parameters (up to logarithmic factors), which is near-optimal (as we show in Section 6).
Setup. Our setting follows that of Cabannes et al. 2024). Let [N] be the set of input tokens, and [M] be the set of output tokens. Our goal is to store a set of associations given by the function f" [N] [M- For each input token E [N] we assign corresponding embedding vector e, E Rd, and likewise for each output token y E [M] we associate an unembedding vector uy E Rd. We primarily focus on the setting where the embeddings (euhne(N) and {uy)eejM, are drawn i..d uniformly from the sphere of radius 1. Let F Idd + Rd be our model which "stores" the associations f Given such an F, the prediction fxx for (xx) is given by the arg-max decoding f(x) == arg ma u.] F(ez).
2

---

Published as a conference paper at ICLR 2025
3.1 LINEAR ASSOCIATIVE MEMORIES
We first consider the case where F is a linear map F(ex) Wez.
Theorem 1. Assume that f injective. If d2 Z N poly log N, then with high probability over the draw of the embeddings, there exists a W such that
arg max u, We, (x) for all x E [M]. Ve[ad]
This capacity is obtained by the construction W ExE(M (x) Furthermore, if W is restricted be a rank m matrix, then such W exists when md 2 N poly log N; this construction is W ExE(N] u..- (a) Di=1 U,D. where E Rdd are drawn ii.d from the standard Gaussian.
Since W has de parameters, Theorem shows that the number of associations that can be stored scales (up to log factors) linearly in the number of parameters. We note that in this linear case, injectivity assumption on f* is important, as otherwise the capacity may as low as d, as in (Cabannes et al.) 22024). Additionally,
remark that these constructions by gradient descent; the general W construction corresponds to one-step of GD on the correlation Σx u
easily obtained
((x) We,, while the low-rank construction corresponds to parameterizing W = UV T for U, V E Rdxm, and taking one step of GD on U while V is fixed to random Gaussian initialization. The proof of Theorem iis deferred to Appendix
loss L(W)
Remarks. Our setting bears similarity to associative Hopfield networks (Hopfield 1982), yet dif- fers in that we decode a fixed discrete set of output tokens [M] rather than exactly matching the target output. This more closely resembles the language modeling framework, and allows us to improve the memorization capacity from da (McEliece et al. 1987). Next, we note that non-orthogonality of the embeddings is necessary for Thcorem []] as the optimal storage capacity for one-hot embeddings is only = d. Since our constructions are in the regime N > d, the associative memory W is superposition (Elhage et al. 2022) of the outer products (=)) Finally, we remark that the random, rather than tramable, embeddings setting was also studied in Cabannes et al. (2024). The embeddings
be viewed global quantities in larger network, of which the associative memory is implementing some subtask, and is thus not able to optimize these embeddings in order to solve its specific task.
3.2 MLP ASSOCIATIVE MEMORIES
Next, we consider the case where F is a two-layer neural network with hidden width m; that is F(e,)= VTo(We, e for V. W E mmxx Theorem 2 (Informal). If md Z N poly log N, then with high probability over the draw of the embeddings, there exists V, W such that
arg max u V (WWe)) (x) for all [N]. yE[M]
Since the MLP has 2md parameters, Theorem 2 shows that the MLP associative memory scheme has storage capacity which is (nearly) linear in the parameter count.
Proof Sketch. The construction for Theorem mimics that of Theorem after an appropriate ran- dom feature transformation. First, sample the rows of W from the standard Gaussian. Then, set each
v. = m 1 Σ.. (x) (uw,,, ew))), where hy is the kth Hermite polynomial (see Appendix F.I
We then see that
1 m m
F(e.)
Σ [M
ur-Geryha(w,, ex)) (w)) ex:))) e,)) se
Σ (rr ex z'E(N]
(3)
for sufficiently large m. Such polynomial associative memory is reminiscent of that Krotov & Hopfield (2016), and can store many more associations for large k. By choosing & x loga and appropriately dealing with concentration, reasons, we must also use the Neural Tangent Kernel (Jacot et al. 2018) rather the random feature model). The full proof of Theorem 2 is deferred to Appendix
obtain the O(md) storage capacity (for technical
3

---

Published as a conference paper at ICLR 2025
Linear Associative Memory
MLP Associative Memory
MlogN dd
10
m/d
10a
m/d 4
mid 8
Nogg x
102
104
10a
10
10'
102
10a
10
d2
md
Figure I: We train linear and MLP associative memories to store (*)x== Left: Linear associa-
tive memories require df N log parameters store N associations. Right: MLP associative memories require md x N log N parameters to store N associations, as predicted by Theorem 2
On Optimal Storage Capacity. Prior works (Bubeck et al. 2020 Madden & Thrampoulidis, 2024 Madden et al. 2024) studying the memorization capacity of neural networks focus on regression setting, and thus do not directly apply to our setup with multiple outputs and discrete
set of output tokens. Other works (Vardi et al. 2021 Kajitsuka & Sato 2023; 2024) show that one can memorize N arbitrary labels with O( O(vN) parameters, at the expense using a bit complexity of S2(VN). Such networks still require ((N) bits, which matches our lower bounds in Section 6 These constructions, however, unwieldy, and not learnable if restrict the precision to be poly log N. Instead, our constructions are learnable - the linear construction results from one step of GD, while the ReLU construction uses the NTK and can thus be learned via GD on convex loss. In Corollary 2 we show that a quantized version of the construction from Theorem succeeds with bit precision Ó(1), and thus more accurately captures realistic training regimes where models do seem to succeed with low precision (Dettmers et al. 2022 Allen-Zhu & Li 2024).
[ indeed
Empirical Validation. In Figure I we train both linear and MLP associative memories to store the association (x)x) x. Given a fixed model size (d, m), we fit datasets with increasing values of N using the cross entropy loss, and plot the largest value of N for which we can obtain at least 99% accuracy. We observe that the linear associative memory can store θ(d²) associations, while the MIP associative memory can store (mdd) associations. See Appendix A for additional experiments where the number of output tokens M does not scale with N.
A SYNTHETIC TASK FOR FACTUAL RECALL
In this section introduce synthetie factual recall task, and show that one-layer transformers constructed via associative memories can store a number of facts proportional to parameter count.
4.1 THE TASK
We first define global dictionary of facts. Let S be a of subject tokens and R be set of relation tokens, where S = |S] R |R]. Let A be the set of answer tokens. We let a" S x R + A be the ground truth association function, which maps subject-relation tuples (s. r) to their corresponding answer a" a""ss r A similar such task was considered in Petroni et al. (2019); Ghosal et al. (024)).
Define A, be the set of answers corresponding relation r, i.e A, = (a)(s)r) Define A, := (a"(s,r) r R) analogously. We assume that each relation corresponds to disjoint set of answers:
S).
Assumption 1. A, A, 0 forr, E R with fr'. Furthermore, define D maxreR |Ar].
For example, S could be the set of all countries. while R could be [papsident.capitta] in this case, the set of all presidents and set of all capitals are disjoint.
Wee focus on one-to-one relations, where each (s, r) pair corresponds to a unique a'. This is in contrast to the one-to-many setting, where each (s, r) maps to many possible answers (for example, s = "France," "city." a" ("Paris", "Toulouse", ....
4

---

Published as a conference paper at ICLR 2025
S++
BOS
Figure 2: diagram of the synthetic factual recall task.
We next define our data distribution D over sequences. Let T > 0 the context length. Let N be set of noise tokens, and define the vocabulary to be V == S U R AUN (EOS), where EOS is a special "end-of-sequence" token. The data distribution is over length T sequences
21:741 (=1; #2..... - FT; 2T+1) E yT+1, generated via the following procedure:
First, sample a subject and relation tuple (s, r) from some distribution p over S x R.
2. Next, sample two distinct indices i, [T 1). Set 2; and 2j
3. For the remainder of tokens where ko E [ 1] {i,j}, draw 2μ uniformly at random from the noise tokens N.
4. Set 27 FOS..
5. Finally, set 2T+1 = a" (s,
The goal of this task is to predict 2 from ZT). A model which can successfully do so must (1) be able to isolate the relevant subject and relation tokens from the noise tokens and (2) store all of the associations (s,r) a" c*(s,r). See Figure 2 for diagram of the task.
4.2 THE MODEL: ONE-LAYER TRANSFORMER
Our learner for the task a single layer of multi-head self attention followed by an MLP. Define d to be the embedding dimension. The input to the transformer is sequence of vectors X (..,...., 2rT' E IRT'xd Each self attention head is parameterized by the key, query, and value matrices WK. WR Wo, Wv Rd where d2 is the head dimension. The self attention head is then map attn( WK, WQ. Wv) Rxx - Rd,, which operates as
attn X: WK. Wo. Wv) = W. S(XW WoRr),
(4)
exp(x,) where S(z); = Σ, exp(=,) the softmax operator.
A multi-head self-attention layer with H heads is parameterized by H different key, query, and value matrices, along with H output matrices. Let == {(w w w((h) w ) K
(h)
where (h(h) (h()) (h()), wi(h) E xd, A multi-head self-attention layer is then a map FMHSA( ;0) RTxd + Rd given by
FMHSA (X:0) ( = Σ whOh (X: W K
(h))
(5)
NE[H]
Finally, single-layer transformer combines multi-head self-attention layer with an MLP. Let m be the MLP width. Let V. W E Rm be the MLP parameters, and define OTF := 0 U (V, W). Then, single-layer transformer is the map FIe( (rrF): Txx Rd given by
Fre X: θTP) FMISA X: () V a W MIISSA (X;0)).
(6)
single-layer transformer parameterized by the tuple of hyperparameters (d, H, dr, m). The model has 4Hdd,, self-attention parameters, and 2md MIP parameters.
4.3 OAYER TRANSFORMERS HAVE (ALMOST) LINEAR STORAGE CAPACITY
We next characterize how large single-layer transformer must be in order to obtain 100% ac- curacy the synthetic task. For each token V, sample its embedding vectors (p(2) E Rd iid uniformly over the sphere of radius An input sequence (21,..., 2T) gets embedded as X (((2z))-..,((()))) We use argmax decoding predict the next token; that is,
f(x1.) = arg max ((2) Fir (x;Orr..
(7)
=EV
5

---

Published as a conference paper at ICLR 2025
Our first result is that there exists attention-only single-layer transformer that obtain 100% accu- racy on the factual recall task, as long as the total number of self-attention parameters dHd,, scales (up to logarithmic factors) linearly with the dataset size SR.
Theorem (Attention-only, informal). Assume that Z St(max(R, D)) and Hd, 2 Ô(S + R). With high probability the embeddings, there exists single-layer attention-only transformer Frr( (rr) with embedding dimension d, number of heads H and head dimension dy such that
IP PELTIID arg max (p()) FTP(X;;OF) = 2T+1 1.
(apta//())P(PPPPPPPPPPPP =EV
We next show that single-layed transformer with an MLP can obtain 100% accuracy the factual recall task, if the number of MLP parameters md scales linearly with the dataset size:
Theorem (Attention + MLP, informal), Assume that a is a polynomial of sufficiently large degree. Define C(a) [(s,r): a" (s. r) = a.}. Let (d, H, du, m) satisfy d 2 (11)
Hd, 2 +R(S + R)
m 2 ((max((a))
md 2 2Q(S)).
(9)
Then with high probability over the embeddings there exists single-layer transformer Fre( : θTF) with embedding dimension d, number of heads H, head dimension dh and MLP width m such that
=EV IP ParriiD arg max (p)) ((2) Fir(X;Orr) 2T+1
(10)
The proofs of Theorem 3 and Theorem 4 are deferred Appendix C
Remarks. Theorems aand have two main constraints the size of the architecture needed to obtain 100% accuracy. First, the quantity Hd, must be larger than S + R. This corresponds to self-attention having sufficient capacity to filter out the tokens S U R from the noise tokens N. For the attention-only architecture, we additionally require d (2mmx D)). When R > D, the total number of parameters Hdd, is (up to logs) at least the total number of facts SR. For the MLP construction, the second condition is that the number of MLP parameters, md, scales nearly linearly with the number of facts SR. As such, as long either the total number of self- attention parameters or the total number of MI.P parameters large enough, 100% accuracy can be obtained. The single-layer transformer can thus trade off MIP and self-attention parameters while still maintaining perfect accuracy. This phenomenon is reflected in the experiments in Section 4.4 We remark that it is straightforward extend our construction to the case where we only need to store size M subset of S x R, where the constraints now become Hdd,, md = â²(M).
Proof Sketch. Theorems 3 and 4 both utilize the associative memory framework of Section 3 First, the key and query matrices of each self-attention head act as a denoiser, selecting the relevant subject and relation tokens in &1-7 while ignoring the noise tokens. To the hth attention head, associate subset ((h) C S U R of subject and relation tokens. Then, setting
w(() K
with
R Σ ((=){(EOS)
(11)
ZESUM)
for large constant B, that the /th head will only attend to the tokens in the subset S(A), We remark that since the embeddings are d-dimensional, at most d/ poly log(d) embeddings can be in superposition, and thus we must have |s()|| < d/ poly (o).
For the attention-only construction, the output-value matrix (())) W acts a linear associative memory, mapping each in ((h) a superposition of all possible answers associated with the subject/relation 2. Letting PA be projection onto a random d, -imeensional subspace of Rd we set
w T
( x Σ Σ ((a){p(2) Pr.
(12)
2E A.A.

---

Published as a conference paper at ICLR 2025
Attention-oni Construction
Head lamends so "axe
41XW, Miae?
a
Head it to
Attention MLP Construction
ddid W'X
wgw.
Hest of heads attend BOS
MLP Ascociatie Memoory X VTa(Mo)
4IXW,
didd,ltedds [heads
Figure 3: Both the Attention-only and Attention+MLP constructions for the factual recall task.
In Lemma 2 we show that this construction stores at most da tokens per head (i.e |s(h| S dx),, and requires the dimension to scale with the number of elements in superposition (i.e |AAi| S d). Since | < R + D, and the (th) partition S U R, it suffices to take d ž R + D and Hd, Z S + R.
For the MLP construction, we instead associate the subset ( with [d/da] attention heads. This is equivalent to having single full-rank attention head per subset. We set the aggregate output-value matrix to the identity, so that the output of the self-attention layer is F'MHSA (X;0) ( - (r)) Finally, the MLP layer for each (s, r) pair. Via a similar computation to Theorem 2 it suffices to make the total number of
an MLP associative memory, mapping (p(s) (p()) to pp (8, r)
parameters mn be md (2SR). Since the (h) partition S U R. it suffices to take Hd, 2 S + R as well. See Figure 3 for diagram describing both constructions.
4.4 EMPIRICAL VALIDATION
We next empirically validate the claims of Theorems 3 and that 100% accuracy can be obtained as long as either the total number of self-attention MLP parameters scales with SR. We further observe that 100% accuracy can be achieved as long as the total number of parameters scales with SR, providing evidence that the model can simultaneously
attention and the MLP store facts.
In Figure we train a wide range of models of various "shapes" on datasets of varying sizes. A model shape is defined by the tuple (a,3, . and corresponds to the family of models satisfying Hd, = ad and m = 3d. The total number of model parameters is 4 HHd,dd 2md = (4a + 2,3 2 which can thus be varied by increasing d. For fixed model size (d, H, dr, m.), binary search the largest dataset size that can be memorized. Specifically, we fix D = 8 and vary S, R jointly as S R. Experiments with different scalings are considered in Appendix A For each (S, R. D), the fact dataset is generated at random by selecting [A| = RD, W = R, and for each s sampling a" [s. r) uniformly random from {(r 1)D 1, 1...., ...). We say the dataset was successfully memorized, and as such SR facts were stored, if the model can obtain an accuracy of at least 99%.
On the left panel Figure we observe that, across different model shapes, the maximum number of facts stored scales linearly with the total number of parameters. On the right panel, we consider a specific dataset with S 32, 32, D 8, and plot the accuracy as the number of parameters vary. We observe that the model can trade off MLP parameters for self-attention parameters, while still maintaining an accuracy of near However, we do still require the total number of attention parameters to be large enough; this corresponds to the Hd, = Ộ + R) constraint.
Thhe "standard" transformer scaling takes 1,8 4.

---

Published as a conference paper at ICLR 2025
42092
104
1590
21
CA,
2503
36
Number k.0
k276 Lo
2003
1100
1000
50ts 25 cmh
412
101
10
8278 A25
06b 46
4000
5000
26
Number Parameters
Fact Peice
Figure 4: (Left) The number of facts stored scales linearly with the total number of parameters, for wide range of model sizes. (Right) For a fixed dataset, the model can trade off MLP parameters for attention parameters obtain 100% aceuracy. The heatmap color corresponds model accuracy.
OPTIMIZATION DYNAMICS
We next study the optimization dynamics of the factual recall task. To simplify the model, consider a linear attention transformer (i.e., the softmax is replaced with the identity map) with orthogonal embeddings. We set d |V/, and let the embedding vectors {((=))=ev satisfy
(p(()),p(22) 62=2' Such linear attention or orthogonal embeddings assumptions are common prior works studying the gradient descent dynamics of transformers (Li al. 2023 Von Oswald
ct al. 2023 Ahn et al. 2024; Mahankali et al. 2024 Zhang et al. 2024
Nichani et al. 2024).
The linear attention model is given by
Fiim X: 0) = Wov XWKQIT,
(13)
where we set d, = d and let Wov = Wo Wr. WXQ = W Wo denote the non-factorized output-value and key-query matrices. Let p(- ZI.T) E Aa be the predicted next token distribution on an input sequence z1.T, i.e
exp ((p(()) Fin (X;0))
p(a (1TT) EveA exp ((y((a)) Fin (X; 0))))
(14)
One can then rewrite the cross entropy loss as
L( = lo p00g(((22) 21:r)).
(15)
We would like to characterize the output of rumning gradient flow, (i.e θ = -V(()) with respect to the non- non-factorized parameters 0 : (Wov, WxQ] on the cross-entropy loss (15). For notational
convenience, we denote Wov (a, 2) == cp(a) Wovvp(z), WKQ(=) == ()=) W (FOS). and note that by isometry gradient flow on 0 is equivalent to gradient flow on these quantities.
Let us assume that we start from the following "balanced" initialization.
Assumption 2. Given initialization scale
0, set Wor(a,
and WKQ(2)
a//Al + for each A, E
Our first result is that the gradient flow indeed converges to zero loss. As a consequence, the pre- dicted next token probabilities p(2T+1 21:T) converge to 1(2741 as(s)r)). where s,r r are the subject and relation contained in the sequnece 21.7-
Theorem (Global Convergence). For 2 let O(t) be the output of running gradient flow for time. For any 8 > O, there exists time 1g such that for 2 1,. L(O(()) < 8.
We next show that the model undergoes a sequential learning dynamics. Let us assume that the number of subjects S is much greater than the number of facts R. We show that during the first stage of training only the Wov (a, r) and (r)(r) components grow for relations r E R, while the As such, the model gets close outputting the best predictor based on just the relation token r. Define p" r) to be the conditional distribution of the answer, given the relation r, i.e p" (a r) E.es p(s r) a" (s)r)).
remainder of the parameters stay close
8

---

Published as a conference paper at ICLR 2025
3.5
3,0
Stage 12 Init
Stage 2: Halludination
Stage Convergence
3.0
Stage 1:
Stage 2: Hallucination
Stage Convergence
Int
2.5
2.5
2.0
Reaten--cy only Subyect only
Relatio onl Subject onl
1.5
1.5
0.5
0.5
2000
40000
6000
8000
1000
2000
3000
4000
Steps of GD
Steps of GD
Figure 5: (Left) Loss of the linear attention model with orthogonal embeddings. There is interme- diate hallucination stage where the loss plateaus and the model predicts based on only the relation. (Right) Loss of the softmax attention model with random embeddings. We again observe an inter- mediate hallucination stage, where the relation-on -only loss is zero but the total loss is still large.
Theorem (Sequential Learning). Assume that S 2 8R/2D, and [N] 2 4R/2DT. Let p(s,)) SR Pick O. There exists runtime T and initialization scale (both depending such that:
1. For all I & T and E SUNN A, we have WWv (a, =)1, [K]]WKQ [=(()] S 2
2. There exists t I T such that, for any input sequence 21-r containing a relation r, Exe (p" ()) i(a =1:T))² S
Proofs of Theorems and 6 are deferred Appendix D
Remarks. Theorem tells that at some intermediate time, the prediction of the model ý(- 21:T) is approximately equal to (- r), the conditional distribution of the answer given the relation At this stage, the model ignores all other tokens in the sequence 21:7 including the useful subject token s and predicts based only on the relation r. For example, S is the set of all countries and is the relation "capital," then on the prompt "What is the capital of France?" the model will output a random countries' capital. We view this as an instance hallucination: the model is outputting plausible, yet ultimately incorrect, answer to the prompt. We remark that without the assumption that S >> R, it is possible for this intermediate hallucination stage to exhibit different behavior.
first train the linear attention model with orthogonal embeddings (15) with 5 = 16. R 4 and D = 8, and plot the loss over time. the left pane of Figure 5 we observe three distinct stages. At the start training, the prediction is close to uniform over all possible answers, and the model obtains loss of log [A]. Next, the loss plateaus at log D, and the model outputs the conditional distribution a given the relation Finally, training continues, the model escapes the plateau and converges to zero loss. We include the "relation-only loss" in the plot, defined as Ex ExyTTI log (EaeA. p(a 21:r))] where any probability mass assigned be correct; the subject-only loss is defined analogously.
Empirical Validation. We next empirically verify Theorems 5 and
answer which is valid for the relation is considered to
In the right pane of Figure 5 we plot the loss of single softmax attention head with random embeddings trained on the same factual recall task. We observe similar phenomenology as for linear attention, and identify an intermediate "hallucination" stage where the relation-only loss drops to zero, but the subject-only loss is still far from zero.
LOWER BOUNDS
In this section, we argue via information -teoorrttttttthccctcoo arguments that the results from Sections are optimal to logarithmic factors. Proofs are deferred to Appendix E
Associative Memories. Let [N] and [M] be the input and output vocabularies, respectively. To establish lower bound, must consider distribution over association functions f* For each x E [N]. assume that the output f* (x) is sampled independently from the uniform distribution over [M]. We model the learning protocol follows. At train time, the learner observes the randomly

---

Published as a conference paper at ICLR 2025
sampled ground truth f", and writes down B bit model F. At test time, the learner generates set of predictions f TR.N from F, where f(x) E Am is the prediction for f Both the mappings f - and F + f can be randomized. Let p be a probability distribution over the input
space [N]; assume WLOG that p(1) 2 p(2) 2 ... p(N). The goal of the learner is to minimize
the cross entropy loss I(f)
EPE[NT p(z) log f()))-())-
Theorem 7. The expected loss of the learner can be lower bounded by
B y [(]]]
2 log M
E>> B
p(xx).
(16)
We thus see that in order to obtain zero loss, the learner must use 2 N log M bits; this matches the construction from Theorem 2 up to log factors. As corollary of Theorem 7 we can obtain scaling law lower bounds with respect to model size.
Corollary 1. Assume that is power law, ie p(x) x x for > 1. Then B [M]]]] 2 Bl-a.
This lower bound is obtained by the MLP associative memory by storing the most probable Ō (B) associations. This matches the scaling law with respect to model size considered in Michaud et al.)
2023
Cabannes et al. (2024), which also considered storing the ((B) most frequent associations.
The constructions in Section require storing Ö(N) network parameters, along with input and output embeddings. We view F in Theorem ass containing only the network parameters, while the embeddings are "global" quantities, independent of the ground truth f", used to compute the predictions f. This matches our interpretation of the embeddings as fixed global quantities which cannot be modified by the associative memory. We remark that the associative memory constructions from Section 3 match the lower bound, since they hold for ((1)-bit precision (Corollary 2).
Remark.
Factual Recall. We next prove a lower bound for the factual recall task; a similar bound was proven in Allen-Zhu & Li
(024). Let S and R be the fixed set of subjects and relations and V be sampled
SS) [R). The association function a" : S x R -> randomly as follows. First, for each relation E R, the answer set A, is chosen to be a uniformly random size D subset of V, conditional on all subsets A,,. being disjoint. For each s E S, the answer a" (s) is sampled uniformly at random from A,. The learner sees the association a", writes down a B bit model F, and from F generates set of predictions ff E Rx where f(s, r) Ay is the prediction for a" rs. We lower bound L, the expected cross entropy loss with respect to distribution p(s, r) over S x RR defined as follows:
the full vocabulary, where V >>
L = a [L(I]] Ea..f [-E,r
E.,, p(s, r),, log f(s, ((s,)) (em))
(17)
Theorem 8. Assume that |>] 2 2RD and S 2 CD log (2 log IV]) for sufficiently large constant
C. There exists a constant e E (0, 1) such that, if L = 0. the number of bits B must satisfy
B 2 SR log D + (1 c) ((VI/D)
(18)
We thus see that Ö(SR) parameters are needed to achieve loss of zero. For this lower bound, the learner knows the sets S and R and does not have to distinguish them from the noise tokens N,, making it a strictly easier problem than the factual recall task in Section
7 DISCUSSION
In this work, we showed that shallow transformers can use associative memories to obtain near optimal storage capacity for factual recall tasks. Furthermore, by studying the optimization dynam- ies of simplified model, we also showed that transformers undergo an intermediate hallucination stage. One interesting direction of future work is to better understand the role of the embeddings, and whether there exists an optimal choice of (non-random) embeddings leading to more efficient constructions. Another important direction is to understand the implication of our results towards understanding empirical LLM scaling laws (Hoffmann et al. 2002). In particular, does there exist scaling law lower bound for the factual recall task? Finally. it would be very interesting to under- stand the extent to which larger models utilize similar associative memory constructions, and if one can probe whether specific "facts" are stored in either the self-attention matrices or the MI.P.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This work was conducted while EN was an intern at the Flatiron Institute. JDL acknowledges support of NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award, and NSF CAREER Award 2144994, This work was conducted in part at the Simons Institute.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.3, knowledge capacity scaling laws. arXiv preprint 44X5::2404005405, 2024.
William Beckner. Inequalities in fourier analysis. Annals of Mathematics, 102(1):159-182, 1975. ISSN 0003486X, 19398980. URL (tLhtt::///wee http:/.ww....oo................. jstor. wtt.:sstttttttttteeeeeeeeeeeeeee
William Beckner. Sobolev inequalities, the poisson semigroup, and analysis on the sphere sn. Pro- ceedings the National Academy of Sciences the United States of America, 19111)4881-4481, 1992. ISSN 00278424, 10916490. URL BRLhttp:/ /ttk::://www... www. jstor wtggsttabeee23555555
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, and Dan Mikulincer. Network size and size of the weights in memorization with two-layers neural networks. Advances Neural Information Processing Systems, 2020.
Vivien Cabannes, Elvis Dohmatob, and Alberto Bietti. Scaling laws for associative memories. In International Conference on Learning Representations (ICLR), 2024.
Mete Demircigil, Judith Heusel, Matthias Löwe. Sven Upgang, and Franck Vermet. On model of associative memory with huge storage capacity. Journal of Statistical Physics, 168:288-299, 2017.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3. int8 (): 8-bit matrix multiplication for transformers at scale. Advances Neural Information Processing Systems, 35: 30318-30332 2022.
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac HatfieldDDodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy models of super- position. Transformer Circuits Thread, 2022. URL httt::://trassfooomee:
htrotooooooeeeeeeeereeeeeeeeeeee html
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Conference on Empirical Methods in Natural Language Processing (EMNLP),
Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in auto-regred neto-reegeessvee language models. In Conference Empirical Methods in Natural Language Processing (EMNLP), 2023.
Gaurav Ghosal, Tatsunori Hashimoto, and Aditi Raghunathan. Understanding finetuning for factual knowledge extraction. arXiv preprint arXiv:2406. 14785, 2024.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
11

---

Published as a conference paper at ICLR 2025
Benjamin Hoover, Yuchen Liang. Bao Pham, Rameswar Panda, Hendrik Strobelt, Duen Homn Chau, Mohammed Zaki, and Dmitry Krotov. Energy transformer. Advances Neural Infor- mation Processing Systems, 36, 2023.
John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(8):2554-2558, 1982.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Samy Jelassi, Clara Mohri, David Brandfonbrener, Alex Gu, Nikhil Vyas, Nikhil Anand, David Alvarez-Melis, Yuanzhi Li, Sham M Kakade, and Eran Malach. Mixture of parrots: Experts improve memorization more than reasoning. arXiv preprint PrXiv:2410.19033 2024.
Yibo Jiang, Goutham Rajendran, Pradeep Ravikumar, and Bryon Aragam. Do llms dream of ele- phants (when told not to)? latent concept association and associative memory in transformers. arXiv preprint arXiv:2406. 18400, 2024.
Zhengbao Jiang, Frank F Xu, Jun Araki, and Graham Neubig. How can we know what language models know? Transactions of the Association for Computational Linguistics, 8:423-438, 2020.
Tokio Kajitsuka and Issei Sato. Are transformers with one layer self-attention using low-rank weight matrices universal approximators? arXiv preprint arXiv:2307.14000 2023.
Tokio Kajitsuka and Issei Sato. Optimal memorization capacity of transformers. arXiv preprint arXiv:2409. 17677 2024.
Teuvo Kohonen. Correlation matrix memories. IEEE Transactions on Computers, 1972.
Dmitry Krotov and John J Hopfield. Dense associative memory for pattern recognition. Advances in neural information processing systems, 29, 2016.
Hung Le, Truyen Tran, and Svetha Venkatesh. Self-attentive associative memory. In International conference on machine learning, pp. 5682-5691. PMLR, 2020.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In Proceedings of the International Conference on Machine Learning (ICML), 2023.
Carlo Lucibello and Mare Mézard. Exponential capacity of dense associative memories. Physical Review Letters, 132(7):07731,, 2024.
Ang Lv, Kaiyi Zhang, Yuhan Chen, Yulong Wang, Lifeng Liu, Ji-Rong Wen, Jian Xie, and Rui Yan. Interpreting key mechanisms of factual recall in transformer-basee language models. arXiv preprint arviv:2403. /9521, 2024.
Liam Madden and Christos Thrampoulidis. Memory capacity of two layer neural networks with smooth activations. SIAM Journal on Mathematics of Data Science, 6(3):679-702, 2024.
Liam Madden, Curtis Fox, and Christos Thrampoulidis. Upper and lower memory capacity bounds of transformers for next-token prediction. arXiv preprint arXiv:2405. 13718, 2024.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent prov- ably the optimal in-context learner with one layer of linear self-attention. In Proceedings of the International Conference on Learning Representations (ICLR), 2024.
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi-head attention in transformers. arXiv preprint arXiv:2306.n 02010. 2023.
R. McEliece, E. Posner, E. Rodemich, and S. Venkatesh. The capacity of the hopfield associative memory. IEEE Transactions on Information Theory, 33(4):461-48) 1987. doi: 1091TT./TIT 1987.1057325
12

---

Published as a conference paper at ICLR 2025
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quantization model of neural scaling.
Advances in Neural Information Processing Systems, 36, 2023.
Ashley Montanaro. Some applications of hypercontractive inequalities in quantum information the- ory. Journal of Mathematical Physics, 53(12), 2012.
Neel Nanda, S Rajamanoharan, J Kramár, and R Shah. Fact finding: Attempt- ing to reverse-engineer factual recall on the neuron level. Al Alignment Forum,
2023. URL http:: al otte http://sttt///w//HHo/ 3DF3333 6DFY3sKB/ fact
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. In Proceedings of the International Conference on Machine Learning (ICML), 2024.
Fabio Petroni, Tim Rocktäschel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, Alexander H Miller, and Sebastian Riedel. Language models as knowledge bases? Conference on Empirical Methods in Natural Language Processing (EMNLP). 2019.
Adityanarayanan Radhakrishnan, Mikhail Belkin, and Caroline Uhler. Overparameterrized neural networks implement associative memory. Proceedings of the National Academy of Sciences, 117 (44):27162-27170, 2020.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.0222 2020.
Adam Roberts, Colin Raffel, and Noam Shazeer. How much knowledge can you pack into the parameters of language model? In Conference Empirical Methods in Natural Language Processing (EMNLP), 2020.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Proceedings of the International Conference on Machine Learning (ICML), 2021.
Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.07601, 2021.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding train- ing dynamics and token composition in 1-layer transformer. In Advances in Neural Information Processing Systems (NewrIPS), 2023.
Gal Vardi, Gilad Yehudai, and Ohad Shamir. On the optimal memorization power of relu neural networks. arXiv preprint arXiv:2110.03188 2021. 2021.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174, PMLR, 2023.
David J Willshaw, O Peter Buneman, and Hugh Christopher Longuet-Higgins. Non-holographic associative memory. Nature, 922(5197)9900-96 1969.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-555 2024.
13

---

Published as a conference paper at ICLR 2025
A ADDITIONAL EXPERIMENTS
A.1 MLP ASSOCIATIVE MEMORY
In Figure 6 train MLP associative memories store the association (x) mod M. We fix M 32 throughout. In this case, we see that the number of associations N which can be stored by a model with md parameters scales N md. This linear scaling the absence logarithmic factors is due to the fact that the number of output tokens M is a constant, and does not scale with the number of input tokens N.
MLP Associative Memory
m/d =
m/d =
10*
m/d
m/d
103
md
102
102
103
104
md
Figure 6: We train an MLP associative memory to store the association f*(x) = mod 32, Em- pirically, md x N parameters are required store N associations.
A.2 FACTUAL RECALL
In Figures and scalings of (S, R, D). In all plots, we observe the general trend that the number facts stored scales proportionally the number of model parameters.
repeat the experiment in the left pane of Figure 4 for different choices of H
16,0-
H -
10
10'
100
10
125
122
10
325
Number Paramenen
A22
Number
We repeat the experiment in Figure varying the number of head to be 4 (Left) or 16 (Right). In both cases, we observe that the number of facts stored scales linear with the parameter count.
Figure 7:
A.3 EXPERIMENTAL DETAILS
Figures 1, 6: For the MLP associative memory experiments, for each choice m, d, N, we first sample random embeddings (e>)=E]N), [ i.i.d uniformly over the sphere. We train a two- layer neural network on the cioss entropy loss to predict the association (x) We standard parameterization and initialization, and the activation a ReLU. The network is trained using ADAM with learning rate of 10-2 for 24 steps. We compute the maximum accuracy the network
14

---

Published as a conference paper at ICLR 2025
H 8,
H --
10
10f
Numbe Fadts Storid
1.00
10'
Numbe
10
80th s4k 6.5
10
125
10
109 Number of Parsmntee
16
104
10a
101
5545 125 069
125 W5--5
Number of Parameters
Faas Jornn
Fect x -arrra
Figure 8: We repeat Figure 4 on factual recall tasks where each subject and relation map distinct answer (i.e D = S).
R²
R²
10
10f
Numbo- 2313 Storio 1
of Faac
1.0
100
36-22
5.2 99
30
3,2
10 Number Paramerers
Mumber Pammetee
Farana
Figure 9: We repeat Figure on factual recall tasks where the number of subjects much larger than the number of relations; specifically, we take S = R2.
achieves over the training run, and say that the network has "stored" the dataset if the highest accu- racy is at least 99%. We repeat this procedure binary search over N, find the largest value of N such that the network achieves accuracy of at least 99%. Error bars are shown for 5 random seeds.
Figures 4, 7, 8, 9: We consider a fixed prompt length of T 32, and train the models via online batch gradient descent with batch size 1024 on the population loss (i.e we sample an independent batch each timestep). We use standard parameterization and initialization for both self-attention and the MI.P For fixed model size, we binary search over the maximum value of S SR such that the model achieves an accuracy of least 99%. All models were trained using ADAM for 214 steps, with sweep over learning rates in {.001, .00.000..000 .01} (where we consider the best performing model over all learning rates).
In the left pane we train a linear attention head with orthogonal embeddings. The weights are all initialized to be equal to 10-5, In the right plot, we train a softmax attention head with random embeddings, which are fixed throughout training.
Figure 5:
Code for all the experiments can be found at http:://gittub..comeeeiiio pectual-recall iccr
B PROOFS FOR SECTION 3
Proof of Theorem 7 Let us set W = E=EIN U)=(=)e] For # (x), define the quantity Tay by
Yay = (uf-(z) uy)) We.. TWe..
15

---

Published as a conference paper at ICLR 2025
We first see that (where the expectation is taken the randomness of the embedding vectors)
E Elmsyl = Σ B(( -uy) Ble]
xE|M]
= E (u)=(s)) - uy)"uyua (x)
= 1.
We can next compute the second moment of Yay- Since the e, are drawn uniformly on the sphere, the e_ Cr terms for 2 # x are independent and mean zero. Therefore
Σ E ((u))(x)) - -.y) (ug(())) 2 (e]
ME|N]
E
(1
u
T-(e)) uy )/
+
E -uy) (
1
1(f*()) (1+2) ((f(())=) =)) v)
Σ
d z/
1 1
2
2(N 2)
1+
d²
Therefore
2(N 2)
Var(Tay) =
+
+
d²
d
d
d
Let 5 be fixed failure probability, and 8' NM-
Observe that Yzy a degree polynomial. By CN lo (1/8') for a sufficiently large constant d 2 C Colog" (1/8)) and 2
Lemma 14 by choosing
C, we have that 16e
log"(1 (188(11))aar1a)) < and thus P(7zz S 0) s 8'.
(E>r,)*
Therefore union bounding over all (x, 3) pairs with # (I), we have that
P(3 Tay S 0) < N(M 1)6' < 6.
Thus with probability - 8, Yag >
for all x and y + (z),, and on this event
arg max u. We, = fxx for all x E [N].
For Theorem 2 need the following assumption the activation
Assumption 3. is polynomial of degree Furthermore, ifa(z) Eh (zhx(2) is the Hermite decomposition of a, then ej + 0 for all
q.
We prove the following formal version of Theorem
Theorem 9. Let 10,1) be fixed constant. Assume that d > N and N > Ci()), where Ci(c) a constant depending only on Assume that q in Assumption satisfies q C2 for some C2 > 2.
Then, if md. Z (C3 log( (/MM) Caje, with probability 1 - 8 over the draw of the embeddings, there exists W such that
V arg max u. V (Wex)=f'(a yE[M]
(19)
for all x E [N).
Proof of Theorem 2 Let us consider the lincarization, or Neural Tangent Kernel, of
FNTK(*) v (((("()))))))))))) (W )
Σ 0. <)) ( wp.2) 2).
mm)
16

---

Published as a conference paper at ICLR 2025
where 8 is the initialization which are linearizing with respect to. By rescaling the parameters as W 0 + <(W W and V t eV we see that F - FNIK as e + 0. thus suffices to work with FNTK instead of F. For case of notation, we redefine W as W, and W W as Q, so that
F(z)= = Σ ()))
RE[m]
Let be even integer, to be chosen later. Assume without loss of generality that ck+1 > 0 iif negative, we can simply negate all the in the construction below). Set
qi
Σ =E[M]
(ne(,, un))) ur ex)
m
where h2 is the kth Hermite polynomial. Then
F(e,)) = m
Σ
( us ( ex) hx((), wi)) ez)
iE iE|m),aE[M)
As in the proof of Theorem define the margin between x and some y + (x) as
Yay = (Ufx(() (a)
Myy (TFFF(e))))
Σ
(0s, (x)) Uy) (=s) ((ws, exr) ((e,, wi) ez)-
m
We will show that, with high probability over the draw of the embeddings over the sphere, and the 17 u. independently from the standard Gaussian, that Yzy > 0 for all * (*)).
The expectation of the margin is
E Eaayy]
Σ LE (V, (x)())) uy)(o,, (') ex)) wi)) ex)]
EE (uf - uy (ex)) k411
C++1
Ck+1+
We next compute the variance. Define wiz zy as
W
(v; (z) - uyy (us, (x) (ww;, ez)) (e ez),
so that Yay 1 Er, wiyh First, observe that when # 2', have E w w even. For i + j. we have that
0. since A is
E T[ WT] = 2 (( uy u)())) Bi B ex) 2(k+1))
First, by Lemma 12 we have that
(e ez) (2(++))]
1
x=2
(2k + 1
x
Next, we see that
[1+1
(())=()))))))
B[(()((=)) My [²) =
otherwise
d
Finally, we have that
E my 20 = LE (v,, Ux)(x)) 2(; u (z) (ow,, hulle, we)) ez)2]
=E (vi, Up-(x) u (vi )))())) E o'ow,,, hu w:) (e,, (er,
17

---

Published as a conference paper at ICLR 2025
The first quantity can be bounded as
I (v,, (x) ( (x)) - [[,, (x) - [1 (x))
(x)) 1/2
2 v3- v3 =6
The second term is bounded as
((w;, ez)) [[w
ex)²] & E[o [hal(e,, [[e ez)
By Gaussian hypercontractivity (Lemma
13p.
E e()(((1,,,)))))))4 11o'11 89110'1122 S 89,
and likewise
E[hx( w,)) wx))) I gh
Finally, E (( e2) s 4d a # 2, and I otherwise. Altogether,
Bl W 2 (d + 1(x 2))
Altogether, we get that
2 Yayl
m 1
E [ [ + Σ
Elw Blw
The first quantity is
1+ d ch. )ohn + (
1)
2N.
The second quantity
Σ E Blw W 201 S 230+3k 1 N
(1+a)
wiz
Therefore
1 (2k)
23 N
Var(Tny) S
d
d++1
md
Choose t: = [] then
d
N
2
21
(2k) N
(4/e)2+2/e
for N 2 (c), and so
1 S d
23 N
Var(Tay)
+
md
Observe that Yzy is degree 2q 2k: + 4 4g + 4 polynomial. If md 2 N C2 40 (1/8') for unspecified constant Ca. then 2' (1/""(111""))) Var(tuu) 1, and thus by Lemma 14 have (B>r,)"
40+4
that (P < < 8' Choosing 8' MN and union bounding over all (z, v) pairs with f*(x) yields the desired result.
18

---

Published as a conference paper at ICLR 2025
B.1 BOUNDED BIT COMPLEXITY
Corollary 2. Under the setting of Theorem 1 if df Z N Npoly log N, then with high probability
there exists a quantized weight matrix W, where each weight requires Oflog d) bits to store, such that
arg max uy We, = (x) for all x E [N]. yE[M]]
(20)
Proof. One sees from the proof of Theorem I that, with high probability over the embeddings, the weight matrix W = E (=) has margin Yay satisfies Yay 2 for all y & f*(x). Each
entry of W lies in the interval [- N, Nj. For some 6 > O, define W by rounding each entry of W to the nearest multiple of e. By definition,
W w
S e. We also see that
(W e)
x
W
w
e 111 S de.
Thus choosing e < 82. the margin of the quantized network satisfies
Pay == (u) - -uy) We,
(u)-(a)- uuy) We, - <We)
uy) (uy
uuy) (W W)e,
IV
-2de
> 0.
Finally, the number of bits required to store each weight is log(2N//)) = log(16NN)) = Olog d).
We remark that similar quantization argument was proven in Jelassi et al. (2024).
C PROOFS FOR SECTION
Lemma 1. Ler ((h) C R. Assume that d 2 log(([V/))) Define v == Esevon (0(2) + (E(EO)). Then, with probability 1 8 over the draw of the embeddings.
(u, > (w, (E(SS) + > (w, (o())))))) +
for any E V(b) and z' g D(A)
Proof. Define 7/2 as
(v,)(z)) -1 ((A) (v, ((EO)) (v, spx))
Yo =
EOS
2 & (hh)
We first see that EB/s]] =
Next, observe that
Σ Everin (0(2)), 72 Eveemn (p())
2 Φ v(h)
(())
Since each of the (4((2), (p(2)) are independent subGaussian variables with variance proxy 1/d, by Hoeffding's 's inequality we have that, with probability 8'.
|c0) log(1/8')
Irsl S
d
Setting 8' = 5/ and union bounding over all V yields the desired result.
19

---

Published as a conference paper at ICLR 2025
C.I CONSTRUCTION VIA SELF-ATTENTION
Lemma 2. Let V(h) C V. and for each E V(h, let A² C V. Assume that 2 max:Ev(M) |A4 logoo"()) (121/8) and da 2 [ (oolloo'(()) (1V)6)) Define
d E p(a))(=)) w.a.
W =EVO aE
dh
where w; chosen uniformiy the sphere of radius 1. conditioned on being orthogonal to (pFOOS). Then, with probability 5 over the draw of the embeddings and the wi,
(0(a) T Wcp(=) 1(a E A
for all z E V(h), a E V.
Proof. Define
Yaz == s0(a) T Wcp(z).
We first that
E B/yas]
E
Σ
E
d-1 Po(Ros)(((())
1( E A²).
='VV(A)
a'EA"
We next compute the variance. For a & A
de di
E
[
Σ
Σ
Σ (p(a), ((a)))(p(=), (1) ww, sp(2))
('VV()) d'EAA'
Define ware = (c(a), sp(aa)) (cp(z), w.) ((=))) We see that E ElWasnnwanze)) nonzero only
a1 ag and 2 = 22. For + j. we have that
E(warssatt = E (s)() sp 2 (v(2), wi) ((=)))))(=)))
d 3 (pa')) " (s((2), ((()))
< d Pan' Pzd'
== Also,
where Pij
d=
*
R[w2 = E (sp(a):
Since E| N/w0001 1](d+T) Sym (P (Po))))))
3
| (EOS) (0()) 19]
E [[ S (c2 - 1) [(((( (((())))) ((a)))?
IR
+ 2F (<p(=): Pol Pe(Bos)s(())) (o(2)))]]
d Pan' Paul (1 )p=z).
Altogether,
d² d
Σ
Σ
1os
z'E((N) d'E t==1
Σ
Σ
dp
1 Paa (1 + 2pzx')
Padt
Past
dr
dn
3'EV(N) d'EAA'
d, +2
d d+ +1
Σ
Σ Pag'
Pan' +
dn
ddp,
d'EA"
z'EIN a'EAA"
dn +2)
(((EEA)))
14:1)
d + de +1
Σ
dn
d
ddp
#EVINN
20

---

Published as a conference paper at ICLR 2025
and thus
d-1
Var(nas) = 18 [azz)
1(a E A+)
1.4
1(aE A=
dp
3'EV(N)
max 2EVON)
AA*1
D(N))
d
dp
since d 2 dz | Next, since Yas is degree 6 polynomial, with probability 1
8
have that
Yas - 1(a A
V Var(Tes) log (11/5)
max |A= log (1V168)
[(h) log (121/0)
24
d
dp
5.
Union bounding over all E 7(h), V yields the desired result.
Let us state the formal version of Theorem 3 which we aim to prove:
Theorem 10. Assume that d Z max( R, D) log (1V))/ and Hd, Z Slgg (IV/R)/). Then, with probability 1-6 8, there exists a single-layer attention-only transformer FTE( :0TF) with embedding dimension d, number of heads H and head dimension dysuch that
P PRITHNDD
arg max =EV
F PTP(X;OTF) = ZT+1
Remark.
When 2 D. one can obtain an accuracy of 100% whenever the total parameter count
is
Hdda 2 PRly (Rg(O))SRR)))
Proof of Theorem 10
Partition S into the sets S(1) .....(ss) and R. into the sets R(1), R(N,..,,NN
(Nn),,
such that [80) RD))
< M and Ns 1, [R] 1.
Let us choose M so that 2 d, 2 M log looo ((V1/5). The total number of attention heads is then
H No + NN Slog (10/0
dn
For each i [Ns], we construct the attention head as follows. First, let
i w = Σ
(B=)) + 3 + ((OS)) T
K
SES())
for a large constant B. Next, set
dp
W (1)
d
Σ
Σ
(())p(2)) w.w.
dn
ZES0) aEA,
for w, sampled uniformly on the sphere, orthogonal (EOSS)
Consider input sequence (z1....,27) and let be the subject token in this sequence. On the event that Lemma hollds if s & S(i), then 21 ((i), and thus
((x))
W (1) K
( (EOSS) < (BOSS w() (i) w (1) (BOS)
3 2
21

---

Published as a conference paper at ICLR 2025
for all t < T. As 3 + oo, the self-attention module fully attends to the EOS token. On the other
hand, if 8 Sti), then if zς- we have
W ((OOS) .3
((+))
((ROS)) < (((+))
W K
K
2
for all t + S. Likewise, as 3 - oo, the softmax converges to hardmax on the 2,- token. Alto- gether, we get that
[p(EOS))
S(1)
S Xw
whe (6) ØT)
((((s)
E 8(i)
Next, the event that Lemma
holds, since d 2 Rlog (1V)/0)
|As 14llog9((V1/)) we have that
(p(a) T W W (ps)) - 1(a E As)
5
for st) Defining attn; == w(()T ( 8
XW T
w 2T
have that
K
{0}
s S
p(a) T attn; E
1 5:
S S A. S A,
By an identical construction, for each j [Nn]. with probability 1 28 we can construct the attention head Ns + such that
{0}
r & R r E R() Ar re ( a Ar
((a) T attnns++
long as Z log (1V1/))) 2 |A, log (MV18). Therefore by union bound, with probability
1 2SR have that (where S( and 8
No+Na
<0(a) T FMHSA (X:0)
Σ <p(a) attn
h=1
= =((aa atin; +ap(a) allnNgti
If (8,r), then E A, O Ar, and thus
<pa) T FMHSA (X:
5
Otherwise, either ( attn, or ppaa) attnns+> is and thus
1
<pa) T FMHSA (X; ク
Therefore arg e vp(a) FMHSA (x:) a" (8, Replacing 2S with yields the desired result.
C.2 CONSTRUCTION VIA MLP
Lemma 3. Let e be a fixed constant. Assume that in Assumption 3 satisfies = C2 for some
C2 > 2. Assume that d 2 Se. R Define C(a) = 1((s.r) a" (s, a)l-
Let d be odd, and let P,Q be orthogonal [d/2] dimensional subspaces of Rdd. Define cp(s) psss(s) ((r) (r))
There exists universal constants Cs, C4 such that if
d (Ca log( /1/
m Z log( V1/8)) max ((a)
md 2 (Ca log( V1/8) SR,
22

---

Published as a conference paper at ICLR 2025
then with probability 1 - over the draw of the embeddings there exists two-layer neural network F(z) = -ie((m) (w) 2) of width m satisfying
arg max ppla T F({((s) +(()) = as,))
aEV
for all 8 E S, E R.
Proof. For odd integers p. k to be determined later, let us set
m Σ
(Fs) to (p(r) (s,)", He p+b + (w)))
v;
(Rd)@(pth) is the Hermite tensor of degree p ++ k (see Appendix F.1 Assume without loss generality that cp+b == [[o(ott) (z)] the (p + k)th Hermite coefficient of o is positive (the negative case can be handled by negating all the v; in the construction)
where He
Rd
For some (s, r). the margin for some a² (s,r)
Ysra (pp(a"(s (s,r)) p(o), (()))
1 Σ Σ (wwi, (
Hew.He (())) & (0r(')o0)))
kE|m] 8' (4p(a (d' (d), (()))) p((a).
We first see that
E [sraa
Σ Blo((ws, + (H))
(I(a"(s,)) a - 1(a -==== a
(sw, ((8) EEBolotb) (w,
+ ((r)) (()))
00()) (())))))
8'r'
a'(, - -a= -(a=)))))
If either s' * or + r', we see that conditioned on (s)), ((r), the quantity (ss')) @ (p(")oo is mean zero. Therefore the only nonzero term in the sum is when (s,r) = (s', r'), and so
B[Tara] = EzmN((,.))
o ((tt)
(zv ||(())]1 + [I((())))
||(
The quantities [1)11 - 2
[|(((1))1 - are subexponential random variables with Orlicz norm
1/d, and therefore bound
E [(Yrra] Cpt22 d
We next compute the variance. Define wher by
Wist ((w,, (ss) + ((rr)))
( & ((r))88)) (s((a"(5',))))
sp(a)
We first observe that E[weerr, Wjezral is zero, unless expectation of Wis'rr, conditioned on the embeddings (i.e with respect to the randomness w.):
and = r2. Next, we compute the
E BlWist [pl
((p(s) ((++)) (()) ss (( Eus; (oioth) ((s), + (())))
(slla r'),, pplaa(aasa (s.r)) - ppa))
(oansoooon ((ott)) z ( (1 +I(())))) (0(6)) (())))))(())))))))).)))))))))")) ((1))) (p(("(6',))) ((),)))))((((())))))))))))))))))
(x,r)) -p(a))
23

---

Published as a conference paper at ICLR 2025
Therefore for i # j.
BluirwWjsrr B Wuis sp]²
S & (f((s), ()) (e[[)))((())))))))))))))())))))) (((())))))))) (()) ,<(a)))).
When (s, = (s', r'), then
B[Wianijjs)] EE Ez-N((.1) (((+h)) [V(1@())11+= + (1((0))))++((())))))))))) [++()))))) (1+d-1)
= C 2p-24 2k + O(1/d)
Next, define the quantities
Pus {(()()))<<'')))))
Prr (
Pan ([((p(())))(a))))
so that
T3] Fldis S 2 SppalsssPPr Per (Pan)(s) + (or)))))
We that for + sr + r # a',
Pes' S (2)){dd2)) (4p)
Prr S (2))((//)==<444))
Pao' < d
Next, see that
E
(He [[o(ww,, ((s) + ()r))) (He (He (w:))
((s)))))(((s)))))))))))))))))))) ((5)) (0 (((((()))) (
[[o((us, ((s) + (1 /B (He)+)(w;))).....)) (88))8p s (())) 1/2 (poa + (ar(aa))
S $2t0 + Pa (s*r))
where we have applied Lemma
13 to the first two expectations. Altogether, we have that
1
E [w]
E Tsral
E WWjer
m
m
E Wjirs + m
1
Σ
1 Σ m 8'
E E[wler]]
m
(s.r))
and thus
Var(Tsra) = E EE(rra
-2
cpt
+ 24g m
S E
Σ
Pas' Per (Paa (8.r) + (arra
(Paa (d(rr)) + (arraa))
o'
For the first sum, we can bound
Peo Per (e + (o)( S
Σ Per
(H
(d'r")
$ SR. ( +R(k)*&
For the second sum, we get that
Σ (Pan)()tsr)) + (e-(a))))(()))))) s C(()))((()r))) 2SR
24

---

Published as a conference paper at ICLR 2025
Altogether,
Var(Tisra
S(4p)
R(4k)* R(4k)* dk
S(4p)
2* + C(a'(s,)))
SR md p+h
B[Yera]
dh
mc
Let 8' be fixed failure probability. We see that, for p 2 +1
2 log
log 4 (1/8))S(4p))
2 4Ca log
4C₂ +2 (1/8')
2 1
whenever d 2 (C' log(1/8)))) Cafe for appropriately chosen constants C3. C4. Likewise, setting k 1 + 1, we get that
24 log
47+ (1/8") 2
dk
maxo ((a) and md 2 2 log 2 (e Next, setting m 2 2 2og 40 (1/8)) yields
2 2 log 10 (1/8') (Ca() C(a*(s,))))
2
mc² mC
24
2 2 to 2
S
mde
p+b ptd
Altogether, by choosing constants appropriately, we get that 2 log +7++ (1/8")e- Var(Ysrra < 1. B[ysra] ]
Therefore by Lemma 14 with probability 1 8' we have that Yara 0. Union bounding over all 8, r, a and setting 8' - yields the desired result.
We state the formal version of Theorem which wish prove:
Theorem 11. Let be fixed constant. Assume that is degree q polynomial, where q C1/e for some C₁ > Assume that d 2 S RR. Define ((a) = [(s,r) (s, = ax}.
Let (d, H, dn m) satisfy
d 2 (C2, (og()))
Hd, (S R) log (IV18)
m 2 (C2 log( log(()/ max C(a)
md 2 (C, log(ll((()/))) log() SR.
Then, with probability 1 - 8, there exists single-layer transformer Frr( 0rr)) with embedding dimension number of heads H, head dimension dn, and MLP width m such that
IP ParrAND
arg max sEV (p2) T FIr(XX;rr)) 2T+1
Remark. Ignoring polylog factors, and treating as a constant, the constraints on the architecture size become
Hd, 2 S R and m 2 C(a) and md Z SR.
We first note that C(a) < S, and so m Z is sufficient. It is possible for C(a) to be much smaller; average expect C(a) S/ D, and also note that it is possible for C(a) 1. The main constraint is that md 2 SR, ie that the number of MLP parameters scales linearly with the number of facts that need be stored.
25

---

Published as a conference paper at ICLR 2025
Proof of Theorem 77
Partition S into the sets ((1) ..... S(Ns) and R into the sets R(1),
..... R(
such that (5() R
R I M and Ng = M 1.NNR = M
Assume that = ( log((v)/8'))
Let H d. For each [Ns], we construct the H' attention heads corresponding to h {(i 1)H + .....,,.fH.. } follows. First, for all such h, let
3 ((=)))(OOS) + (BOS)))(BOSS
(h)
W
(h)
3 E
K
SESO)
for a large constant 3. By an identical argument to as in Theorem g on the event that Lemma holds we have that
(EES) (p(s)
E 8 S(1)
X
XW (h) K
W
(h)
zr
The total contribution from these attention heads is then
W Σ
(h)
attn x; (tt Wh
w (h)
w =
Σ
wh)
w
(EOS)))))))
s s S
Q
(p(s))
8 8 S())
11
Since H'd, 2 d, we can let E (i- 11----11100 1)1+1
'H'
( ( be a projection onto [d/2] dimensional
subspace P, orthogonal to (FOOS, and thus
iH'
w (h) attn( X; Wi(h
(h)
(h) w()) Q
0
$ s S()
(Tpp(s))
8 8 S()
Altogether, if the sequence (ZI, contains the subject 8, then
N'Ns
(h)
attn(X; W (h) w(()) w((h) Isp)(()
(h)
Similarly, if we let Q be [d/2] dimensional subspace orthogonal to and (EOSS) then we can construct the attention heads h E (H'Ns 1,..., H'Ns H' Nn) such that
H' H'N
Σ
w(th attn (etn((;;W)) W (h) K
(h) w(()) NQp(()),
A=
where r is the relation the sequence (21...., 2T).
Such a construction exists with probability
1 -Ns Nn)86. The total number of heads
d(S + R)
(S R) log (118))
H H'N + H'
M
dp
The output of the
component is then
FMHSA [X; 0) = 1lp((()) + llQ((()) = ((s) + ((r).
On the event that Lemma 3 holds, we have that there exists two-layer neural network F(z) Eiemm) (w. 2) of width m such that
arg max sp(a) (((())(((()))) + ((r)) = as)))))
Scaling V by large enough constant ensures that
arg max (px) Fre X: 0rr) = a" rs)
Union bounding over all the high probability events and setting 88/(Ng + Nn + 1) yields the desired result.
I
26

---

Published as a conference paper at ICLR 2025
D PROOFS FOR SECTION 5
D.1 PRELIMINARIES
Recall that the parameters are 0 (Wor(1 (Wkg(z))zev, and that the cross entropy loss is
L( == B 11
(1)(2T+1), (X; + log
Σ aEA
exp Fiin (X; (X:
where
T
((aa) Pri (x;0) = Σ
Wor(a, 21
We consider running gradient flow:
θ = V
from the initialization (a, a WkQ(=) av A for some a > 0.
We also define by
(a, (
and remark that the loss L is convex in O
Lemma 4 (Balancedness). Let C(21r, 2) denote the number of tokens in 21:T equal to 2. The loss gradients are given by
Owvo(a:))L(0) WkQ(z) Enr (C(TT, (1(a = a² (Z1T)) - p(a 21:T))l (wrr(z)
Σ (o 2) Danl(C(an:) (1(a a" (2(11))) p(a 2x)))]
As such, the quantity
WKQ(2)² - Σ Wvola,2)
aEAA
is constant throughout the gradient flow trajectory.
Proof. We first see that
@Wvola.>((p(a') Fin (X;0)) 1(a =') a') C(EIT,2) WkQ(2),
Similarly,
( (X;(()))) C(a.T, 2) Wor(d,z).
Therefore
L(v
WKKQ2)) 1 (2r+1 = a) C(a.r, 2) + Ense exp ((s((a)) Fin (x; 1(a d') C(ZIT,2)
Edes exp Fin (x;;(0;))
-=)Q(>)) E E (C(I,, (I(a = a²(x)))) M(a =1.T)))
By similar computation,
3wxe(s)TN I(0)
E EEL
Wor(2r41>2) C(ZIT,=) + Σ p(a Wor(aa) C(E121
Exr
C(aIT)2)- (oollat(iT))) + Σ p(a (or
Wor (a, 2). Eser EanrIC (C(z))))) (1(a a² (au))) - p(a 2x])
27

---

Published as a conference paper at ICLR 2025
Under gradient flow, that
Wkooz)²²
Σ
Wvola,z)3
2 di
aEA
WKKoz)) WKolz)
Σ Wvo(a, Wvo(a,z)
dt
dt
aEd
=-Wxo(2) -wx((s)) +
Σ Wvo(a,z) (Wro(a.))
aEA
Wko(z)
Σ Wov(a, EST (C(=)T,))) (1(a (1(a= a'l==a((=xi)))) p(a =1))
Σ Wor(a;;)Wko(() PelrCC(ai (1 a (zI.)) p(a =(.))
aEA
Corollary 3. Throughout the gradient flow trajectory, WKQ(2) 2 a
Proof. At initialization, WxQ(2)² - EaEA Wvo(a, z)² = a². Since this quantity is an invariant
of gradient flow, it is impossible for WkQ(=) = 0, and thus Wxo(z) > 0 throughout the entire trajectory. Furthermore,
WKo(2)² Σ aeA
Wvola, + 2
and thus WKQ(=> 2 a.
D.2 PROOF OF THEOREM 5
Proof of Theorem
Let select
2 rap(s (T-33 - 1
< min
There exists time T. such that for all Te, (] < Let us set to = Te. Now, consider some iterate 0 (t) fort 2 tg.
First, see that for E S.
(wov( L(0) WKQ(s) E., [C(an, 2) (1(a (21r)) p(a 21:)))
WKe@ss) p(s) E [1 a" (ZIT)) i(a auT) |s zuT]].
Consider some & A.. Then E.., [1a a*(x1))) 8 E æ 0, and thus Owor(a.) WkQ(s) p(s) E.gr [p(a m s E 21:7] WKR(s)) Σ p(s,r) E.. [pa (1.)) s, 2mrl
As such, since [wwor S.
Es,r [p(a 21:T) s, E æ S ea pp(ss)
By an identical argument, since [awor (()) s then for & A,
E [Ā(a 21:1) s, murl & ca pps,r)
For any a # a" rs, either a & A, or & Ar. Therefore E (ī(a air) 8, E 2ur] S
ea 1o-1 p((,r)) for all a (s,r), and thus
E,,r(p((" (s,r) 2xT) 8, E 2uतe
ca -a--a-pppss "p(ssr)~111
28

---

Published as a conference paper at ICLR 2025
There are most T sequences 21.7 containing (s,r), each of which occurs with equal probability. Therefore
(s, 21.T) 2 1 1 1
ca - p
for all such 21.7. Then, bounding log(1 - =) < 2z for 2 E [0. 2
B1TT 8, 2u:T| 2 2 f(aa' (s.r) =1T) 8, anr]
v 2ca - [|A]
Altogether, the loss is
E B1-logØ(2T+1 21:T]] Ep(s,r) log log (8,r) (ITT) 8, E zmml
2.r
< 2ea
!|AISR
S
as desired.
D.3 SEQUENTIAL LEARNING
The goal of this section is to show that the model learns sequentially; first, the relation components grow, then the subject components grow. This is given formally by Theorem
We first prove that weights corresponding to the subject and noise tokens stay bounded during the beginning of the trajectory.
Lemma 5. For s E S.
WkQ(2) &
Likewise, for N.
WxQ(=) & exp(2Tt/[W] av +1
Proof. Recall that the update for WkQ(s) is
WKQ(s) =
EZar [p( ZITT) s E aur))
S < 2p
s) - -E,. P0. (1T) I ærrll
S Wko(s)
Therefore by Gronwall's inequality,
WKQ(s) S exp(2p(s)t) -xp(2(())))-a//11++ +
Similarly, the update for (xQ(2) for E
WxQ(2) = (Wor(,2). (C(r,,2)) a - p(. (xr))))
(or(; =( (((:)) pl. =x))
< WWv ((C,,)))))
2T
W (x)(2).
W
Again by Gronwall's inequality,
WkQ(2) exp(2Tl/WW) ay +I
The following lemma is our key result, and shows that, assuming that the subject and noise weights stay bounded, the relation weights grow until the output the model approximates the best relation- only prediction.
29

---

Published as a conference paper at ICLR 2025
Lemma 6. Let aams
0 be arbitrary parameters satisfying
e²
OK 9m
T
log
a²( + +1))
min (-1r)) - -poll-
150
1
2² 50[[A] + 1)
min IIp" (- (-(r)) -pll
For a target accuracy Emin > 0, define T by
T = (1p 1 Poll poll log
(vvia++1) 100(|AA) + log |AAcc 3c mmm
Assume that for E S UN that WxQ(2) S aam- Then, there exists S such that
( - 10( =1r) =1Tlll? S emin.
Proof. The proof proceeds in three stages. First, bound the time required for the relation weights to escape the origin. Next, we prove that the relation weights stay large. Finally, we show conver- gence.
Stage 1: Escaping the origin. The gradient flow update on Wov (a.r)
(a) WkQ(r)) p(r)(p" (a r) EE (i z1r)) |r 2)rl)
We thus have
Wor(.r) WkQ(r)
(xq(()-p(r)( p())() Po(a)) WkQ(r) E DP z1:) r E 21:rl - Pu
1 1A. Observe that Define Po [A
|| pp( (ITT) r [1:]] - poll S Baar [|][(a aIT) poll r r 21rl
< B
W [ a:) rE 21
W (or(sr)l + Ta WxQ(r)² + Ta
Thus
Wor(sr)- WkQ(r) (i ((1))
5 (WW)))o((((<<<(1) Ta??)
Likewise,
WkQ(T)= = ( E Ip(a 21T) r 21Tl))
and thus
W
pol) < < PP( a1T) r ø1.7] Poll
((r)) Ta
Define the vector 2 by
WKQ(r)
(Wor(.r)
(-(-)r) Po (-r)))))p
We that
0 1
- ---p(r)lp"( (:1 -poll
S 2pp(() <2p()WWec(((W<xQ11> +Ta2
1
0
30

---

Published as a conference paper at ICLR 2025
Therefore
d di (1lul()) s 2(ū, u)
- 2p(r))lp" (:1 ()) poll|u/| + u (Wxp:)/ + +T
- 2p(r))lp"( (1Ir) Po |ull|l1|
1 4p(r))l [1u1] + +Ta
- 2p(r) (Ilp" (-1r) Poll 2T || 4p(r))
where the last inequality bounds k (r) S ||]l
Define "r = (1p r) - poll + 2m By Lemma we have that for
t < v/ log
+ +1)
4pol)]|2ol1 7r. 1ll exp(
||ul|2 S 7r + 4p(r)3(1 - exp((t,t)
Let T. be the first time that ||u/|| 2 If < 7,-1 log
Th 1p(r)
|u0 1| + 11) then
20110o11 exp(7xT.)
+ + 1) exp(7, T.)
2 |ull
5
7r + 4p(r)2(1 (xxpr))((-eep(7xT3))
7r + 4p(r) exp(xxT.))
Therefore
Eiy 40(A) +
2
T. 7, log (IAA + 1)7 4p(r) ((AA +1)
7r log
20t2 (IA + 1)
for 22 S
On this assumption, 20-((++1]
2
S 4p(r)lluo1** and thus we always have T. 2
2r
4p(r)
7,- log [1<(T1+
Define Lp by
exp exp ((p(a), (in(a))) Ire 21 00))
(0) == =p(r)
Fin
(x; + log
Σ
A.EA
Let us define the relation-only model
((a) Fxi(X;)) Wor(er))
where Z1:7- We see that
( (X: Finn (X:0) S (T 1)a2
Define RIA + IR by g(z) = log (Ea exp(2x)). We see that V=9(2) = S(z), where S is the softmax, and thus
(g(21) 9 S sup V [=g(2)11ı /1((2))) - g(22)ll S [199(2)) g(22)ll
Therefore defining the relation-only loss L, as
exp ((p(a), (r; 0))) ]
I, (0) : E, - Frei(r; 0)) + log
Σ NEEA
=-p(r) Σ p(a (s.)) + p(r) log
Σ
exp ()
REA
31

---

Published as a conference paper at ICLR 2025
see that
- s 2(T -1)a
Since log-sum-exp I-strongly-convex, recalling that ( r) = r)
log
Σ exp(O(a.))) $ log(1AA]) + Σ
1
- (a. r) + 1
| (..r)
A
2
Therefore
L, log ||| -()) (-1 - pu) + 2 ()(r))((.))))
Lr. - -(r- -pol ur + ((r))(((,))1
We next track the evolution of Uz22
dl u2) = ú1 + u1 2ي
2 p(r)lle' (: (1)) - poll|u||² (((r)>/u12((11 (T - 1)
2 p(r) [Ip" () - -oll - -/||ul12 - 4(T 1)a ||
2 p(r) (Ilp" (- r) - poll - -4Taam) |ell? - 4p(r)) 2/11
for I < Tr. Since
||||||
S e, this is increasing in ||2e|/-
We first have the bound |//l| > 2 2 a² Next, we have the bound |uull 2 2u, u2 Pick
some time Te. Define y, (p(r)(1pp"()))
poll a. We see that
(u)u2)(T) 2 3:p=2=-p())a++ )- 3./=a²
1
Next, by Lemma 7 for t < Te we have
75 (u (t (t -7))
(2) 2
7f Sp(r)((r)s)()) exp( (t (t-)))
Plugging in 6 1 1==-=110 log
2 2a<[[A|+T]
(T) (t (T) 2 7 a e'xxxx(-77 eexp
1r log
2a2([AA 1)
7
Selecting T 11/5 we get
a
(: (r) -poll a sm [pp(:
é²
(x))))() (1)
>
4e
exp
poll
+ 27 o2
log
20 20a( + 1)
a²
é²
6T 2
2²
4e
2a
[|A[ +1)
eexp
[1p" r) Po || + 27
sm
am
log
2a [A] +
a²
2²
11 (1-
aT km
2²
log
4e
1p-( ((Ir) poll
+ 2Ta2
2a2([AA + 1)
22
50([A] + 1)
whenever [-r()))ppo 27 mm
2
log (
< 150
32

---

Published as a conference paper at ICLR 2025
Therefore
2²
(u) 2 50( +1 1
1
p(r)
50 [[
2
50[ + 1) po A[+1)
4e2
2
100(|AA +1)
Altogether, we get that the loss is
c²
L S Ly,0 p(r)lpp"( (.1r) Poll 100([AA) + 1) + =2((r)e*
62
L p(r)llp"((()))))))))))))p)) Poll
200([A4 + 1)
whenever 22 5
(-1r)- poll [A][11)
Stage 2: Norm stays large Next, we want to show that WKQ(r) stays large. We first show that the relation-only loss is decreasing, We can compute that
dt = =Vool Vol.)
Define p(: ()) by
p(a b(alr) () = exp(Wx(()Woov(a.)) En cop (a',r)
We observe that
Owo (x (-1r) p((1r)
@workijh Wxo(r)p(r) (-10 (-Ir) EET (p( [1:T) r ((t))
and thus
@wor L < Wx 1 (I;) TE (1:) - x(- |r)
< (x
Likewise,
[Owwols) - Ow (Wov (.WWoo. Pp 21T) r 21rl - p( [r))
S (tr) (. (1 p0. (ITT 21:T p(: [r]]
S W
Therefore
dt (,) 2 V / 1oL V
Assume that & (.) 0. Then
V2Wxo(T)(()TTa.. > Vol. (xQ(()(()ll0((11) PC (r)),,
ie
(p-(-1r)- |r)l V2T22
33

---

Published as a conference paper at ICLR 2025
Assuming that a 28 since p'(a r) > for * (a r) 0 we have that
(a r)) log pa(a r)
log (1+ p(a (r) ((alr)-prr -r 1
= p" (a r)) log
((((alr)- ("
P(a r)
pa(a r))
Therefore I,, p-rrp(() r)) - 2p(r))le(-1 (r) r) --pp- -r)-1 S V2Ta sm D. As such, we have 2
that L, stays below Ln0 - p(r))lp" (-1 - poll 200[[[[+1] for the remainder of the gradient flow trajectory.
By convexity in @ space,
Lr.0 - (.) (0))
$ [()
s p(r) (:1r) - [[)
Therefore
(+r)) 2 100(|AA 1
Stage 3: Convergence.
Next, we can bound the loss decrease by
d dt T(00) 1|
S Σ [1)
Σ Ln)
Σ [iolr))po)>l'' (-1 [r) E.. [p( 21:r)
r e srrlll
Σ (1067))1177)111 (r) - -EZ. P0 21:T)
r m ?
2²
Σ p(r))1lb 11 (1r) Ea,r I6( (IT)
zur]|l
200 1)
Since L(O) < L(O)) = log AA), and L())= 2 0, there exists some s max, t" + 100([4) +
1) log AAl -2 mm. such that
Ee(r) E.,.. D6 ( r e 2utlll < 2
desired.
To conclude, we must set at, Orems T* appropriately in terms of in order to apply Lemma
Proof of Theorem o Let Cmin
100(|A|+) -inn I|p (- r) - Poll be the target accuracy. Let
1
us choose the initialization a so that log
-viAi=I) = 4, where is chosen so that
2 100(|A] + 1) log LAC (1(r"(())) Poll-
In this case, we see that
T < 2amax (p(r)~2110* |r) -oll poll--))
34

---

Published as a conference paper at ICLR 2025
Since p (r) is supported at most D elements, and |A] 2D, we have llp" (r)rr)
-oll -rol- 1
v2D. Therefore T < 2RVD Let compute Otym-
We see that, since $ 2 8Rv2D,
4RvD S
WKQ(s) S exp
+v
S exp
|
+1
22
VEa
(IA) + 1)1/4
S Va
Similarly, since N 2 4Rr2DT,
ARVDT
WKQ(s) S exp
.aVIA++1
W
$ exp
2" avVAA++
$ Va
Therefore the assumption holds for a'm Va. To conclude, we must verify that
To S 300 min Ilp"( (r) Pall-
the RHS scales with em and the RHS scales with and thus the condition can be obtained for choosing & sufficiently large.
But since A
Under the setting of paramters the conditions of Lerm satisfied, and thus the claim holds.
D.4 HELPER LEMMA
Lemma 7, Let z(1) 2 0 satisfy
8 < Az + Bz?.
for positive constants A, B. Then
Az(0)e Ar
=(t) A + Bz(0)(1 e.
Furthermore, if
2 Az 1322
and 2 E [0, A on the interval [0, T1, then
Az
=(1) 2
A + (z 11
for all t E [0,
Both claims follow from the Bihari-LaSalle inequality.
E PROOFS FROM SECTION 6
E.1 ASSOCIATIVE MEMORIES
Proof of Theorem 7 l + F f is Markov chain, so by the data processing inequality.
I(f";f) I(f"FF)
35

---

Published as a conference paper at ICLR 2025
Also, by definition of mutual information
I) (F)) S B,
where the last inequality follows since F is an B-bit message. Thus T((" B.
Let (- f) be the conditional distribution of (x) given f. Consider some fixed f. (xx is also probability distribution over [M]] and thus by Gibbs' inequality
E g (1)) [- log f(())) L (- [- log (v lo [f).
Therefore, letting be the marginal distribution over f and the marginal over I (z),
E Ey log f(x))((0)
E (-1))
- log [(x))]]
2 y B (-1)
log Qu (y []]]
llgg (f(()
y - llog (x(f"xx),f) log a (f*
{(f)g((f(())
--(( (x);; log M.
where in the last step we use the fact that Oz is uniform over [M] and plug in the definition of mutual information. The total loss is thus
L() 2 Σ p(x) p(a)((1(5(00)))) log M).
(21)
=[[N]
Since the yi independent.
BZ B21 E 10(()); DEIN
Also, 0 < (( S H()(()) log M. Therefore equation 21 minimized when
(x); f) = log for the B/ log M most frequent tokens. Altogether,
E Ey. [w]] 2 log
Σ log KT
p(x).
Proof of Corollary 7
Let p(x) Zax where Za EmE(N x can bound
Σ
p(z)
z= x 11-0
2=1
x
z=k
Therefore
B l-a
E E [()]] 2 log M
Σ p(x) R log M
B Bl-a
log M
Tox
E.2 FACTUAL RECALL.
Proof. Define ((s,r) == Ep.r|-ton log (s, that
L p(s,r) ((s,r).
36

---

Published as a conference paper at ICLR 2025
Let us define the expanded dataset D hAr U (a* We observe that D F + is a Markov chain, and thus by the data processing inequality
B I)D:f).
Next, by the chain rule, we can decompose
I(D: 1 ((A
ART f) + I(a': AI,...,AR)
2 T(AI,
An; Σ (8"(s))))) A AR)
= =I(A
Iaaa"ss)): Ar).
8,5
where the first inequality uses the fact that the as(s r) are conditionally independent given the Ar, and the second uses that a"(s, r) is independent of A, given Ar, for + r'.
We can decompose the first mutual information term, using the fact that the A, are nearly indepen- dent:
Lemma 8. Assume that M 2RD. Then
2R [
I(A AR: I((AnA))
We next relate I(Ar; to the loss. The intuition for this lemma is that for fixed r the quantity Σ, r() is small, then the predictor f must contain information about the answer set A,.
Lemma 9, Assume that [V] 2 2D. Define n = CV D log log [l) for a sufficiently large
constant C, and assume that 1j < 1. Then
D
M
2D log VI
2D
I(Ar:) 111+1)
Σ +E[S]
(s,rr) Dlog
nD
D
V
V
lower order term
Finally, we relate rs, A,) to the loss. Similarly, the intuition for this lemma is that if the loss r(),r) is small, then f must contain information about the true association a (s,
Lemma 10. For all s.r,
I(a (s,r): Ar) log D E(s
The proof's for Lemmas Combining Lemmas
deferred to Appendix B33
to 10 we get
B> 2 I(A AR:3 f) Σ I(a (s.r)s) (A,)
D = = s 8,7
Σ
((s + RD log
M
2RD log VI
2RD
RRD -
2R D
D
|
M
lower order term
+ SRlog D Σ ((s,r)
= --(1--) n) +1
Σ ((,r) + SRlog D + RD log
V
Stats
D
8,7"
where Etot
2RD lox V
2RD M
+ nRD + R+ DR
RD log is a lower order term.
[V]
[
37

---

Published as a conference paper at ICLR 2025
Altogether, that in order for all the losses ((s,r) to equal zero, require
[
M
B2 SRlog R.D log
Elot 2 SRlog D + (1 - c)
Furthermore, when p(s, r) RS. then L
SR En,r E(s, r), and the bound becomes
B2 ((1 + RD + RS) L + SR log RDD log
Eta
V
- ((1++) + RS). L SRlog (1 c) RD log
D
E.3 AUXILIARY LEMMAS
Lemma 11. For random variables X, Y,Y, Z,
[(X, Y; 2) 2 ((x;))+((Y22) 1(xY)
Proof. By standard properties of mutual information:
I(x,Y; -
-) - H(X,Y Ø) - H(X) + H(X Ø) H(Y + H(Y (2)
=({;Y ((;Y)
2 -)(;))).
Proof of Lemma
g
By Lemma 11
I[A ....... Agi f) Σ I(Arff) Σ A...AAA
Σ((Ani)) (II(A.) + +H(AA......I)))) H1 HAA,,
((.)........)) A.)
ΣnAAn H((A) H(AI AR).
Since each A, is uniformly random subset of V, we have H (A,) = loog (1))
Also, we can bound
H(AI. .R......))))..). = =log (()) (1) (P/DD ... (M)) - (R ( (1-1))) 1)2)
D
D
Thus
(b)
ΣH(A) (004)-H((AA.... AR) & Rlog
(MI-(R-1)2)
| RD)!
Rlog
([V)
D)L([VI - - 1RD
[V]
RDD log |V]
RD
2R D2
[V]
where used the bound log - 2x (0, ))) Plugging in yields the desired bound.
n
38

---

Published as a conference paper at ICLR 2025
Proof of Lemma 9
Let (21,...,, 2D) be random permutation of A,, We first aim to relate [ f) [ (245 f). By the data processing inequality,
I(A;f)) 2 1 (2:
By Lemma 11
2D: 2
Σ T(2sf)) Σ 22
- ((: HH(z) (2),....,
The tuple (21,...,,3) 2p) is chosen uniformly at random from DD, conditioned on all the 2; being
distinct. Therefore (z:) log [V. and (21....., log ((V) ... (IV) D 1)). Thus
1210
Σ H (z,) - (z1,, =D) log
-..... ((V) D+
[V
D log
[ D
2D |V]
Altogether,
2D |V
T(Arf) 2 Σ Tzs
Next, using the definition of mutual information and Gibbs' inequality,
PP#:: f) P((2))
E B..f log {(=: f)
I(2: f) E E
log
P(()))
for any probability distribution q. Let us define & as follows. First, define f(s,r) == (1 (1-c)((s,)) r)+ M Av, for small constant to be chosen later. Next, define
q(= f) S E0,r))
Plugging in, and observing that IF(2)) = [V] we get that
[[z55 Egg log 1/5
((.)))))))
]
+ log \//.
Define Nº. : {s a" (s,r) 2). Let E be the event that [M11 2 M for all A. On the event E, we can bound
log 5 Σ f(s,T)
2 log
5 Σ f(s))))((a)) sENs,
W
log
Σ
f(s, )ar(())))
+ log
W BEN
Wea
Σ
log [(8, (a)ssr)) + log
W
S
sEN.
M log ((, + log S
M s
39

---

Published as a conference paper at ICLR 2025
Thus
1
Σ log
Σ r))
2
log ())(a))) + D Dlog s
8
M
[E(D]
sENV,,
M
Σ
log (s. -log
M
#E[S]
Se
M log VI + D Dlog
log (8,
M
M
S
SE/S]
On E, we have the naive bound
Σ xE(D)
log
1 8
5 F(srr))))
2 D log |
Altogether, we have
Σ (zx;.)
xE(D)
Σ Σ M(rr))a) )
2 B
Σ log LE(D)
+ Doog |
]
1-6 Σ
Se
log M Dlog M
2 j
1(8).
log Ts,
M
M
sE(S)
+ +IE(E) Dlog
D Dog |
2 M
1
Bznff
Σ sE|S]
log f(s)
Se
log M Dlog M
P(E) Dlog | + D log |
M
Σ #E|S]
((8,r)
Sc
log M1 + D log
M + P(E) Dlog
c + D Dlog |V]
M
M
5
By Bernstein's 's incquality and a union bound (a similar such concentration argument was used in the lower bound proof in Allen-Zhu & Li 2024 ), there exists a constant C' such that P(E) 2 1 for
S
C /V: log(D
M
D
as long as S 2 Dlog(D/5)
Set c =
M =
2D log M and define =
20
/
D 3
22 log (V) < 1. We have that
M
1
D (Dg(
1 - (1-
1
Co
S D
and thus
S =
11/2 < D(1 + m)
M
Therefore
D
log Vl
V
Σ ((xa)) 2 -11 1)
Σ ((s,T) - -(1 + -1)
D log
+ Dooog(1 1/2) 1
[V]
D
[E[D]
$E|S
D
Vl
2D log V
2 - + 11) 5
((S,T) + -log
nD
M
+E(S]
40

---

Published as a conference paper at ICLR 2025
Altogether, have
I(Ar:f) 2 E f( 2D
[
VI
2D log V
2D
2 - (1+ S
Σ sE|S]
((,,r)
Dog
nD 1,
D Vl
[V]
as desired.
Proof of Lemma 10 By the definition of mutual information and Gibbs' inequality,
r)) log P(a (A)
(s) Ar) = EA,
E (art)).
= EA, [a
log P(g (.AA)) ]
+ log D
2 EA, [
log (asa(aa(s,))) (8,r) S.
+ +lg D
where q(. (() is any distribution over V. Let us define be
a(a (.A.) & ((, 1(a E A.) Since a" (8, Ar always, have that a(a (s, A,) 2 f(s, (s.r,, and thus
((aa(s,r); Ar) 1A) EA. E Ba log f(s, ()))))
+ +og D
Ep log rs, + log 1 log D f(s,r).
F TECHNICAL LEMMAS
Lemma 12. Let u, be drawn uniformly over the d-dimensional sphere of radius 1. Then
EE [u, 2| < Pd
Lemma 13 (Hypercontractivtty for product distributions). Let f (8t-1)+ R. RR be a poly- nomial of total degree at most p. Then
(mm)) (
where va is the uniform distribution over the sphere Sd-1, and pm is the standard Gaussian in m dimensions.
Hypercontracctvityy for the Boolean hypercube (which implics hy hypercontractivity for Gaussian space)
and for the sphere are consequences of [Bekner (1975 1992). To show Lemma 13, one can use similar techniques to the proof of Corollary 12 in Montanaro (2012).
Lemma 14. Let f (gd-1)* TM R be polynomial of total degree most p. Assume that Soh Eʃ 2 O, where the expectation is laken with respect to v Mmn Then if
2. log" (ar()) -1,
(Ef))
with probability most 8.
41

---

Published as a conference paper at ICLR 2025
Proof. By Markov's inequality,
PUS S P(S Rʃ 2 ES)
P(S R 2 (B.)))
f E
(E,f)
Since f is degree polynomial, by Lemma 13 we have that
Elf Bf Bf s (ar
Therefore
q'Var(s) 4/2 (Ef)2
P(S 50)
Setting q 1og((1/)), we see that whenever
2 log"
S
(B.))²
we have
( (EEf}
9/2
P( 50) 0)
S
as desired.
F.1 HERMITE POLYNOMTAT.
Let j be the standard Gaussian in dimension, and let (p)) be the function space of square- integrable functions with respect to this Gaussian measure. The Hermite polynomials (hal220 form an orthonormal basis of I. (1)). In particular, h2 is a degree polynomial, satisfying
(he by) (2()) dig-
One useful property Hermite polynomials the following: Lemma 15. Let u, Rd with ||a|||
||/|| and let N Ia). Then Eg [he ((u, x) (ww (w,x)) (u, w)*
Next, let na be the standard Gaussian in d dimensions. The function space (2((a)) has an orthonor- mal basis of Hermite tensors (Hex k2z. where Heg RT > (Rd) 8t
Definition 1. Let the kth Hermite tensor Hee : Rmm + (Rd) so be defined as
Hel((a)) (-1)* Vh ua(ar)
pa(x))
where (x) = (2m)-a/2 exp (xp (==1l)) is the Gaussian density. We remark that each entry of He, (x) is a degree l polynomial in x, and
The Hermite tensors satisfy the following useful properties:
Lemma 16 (Properties of Hermite Tensors).
(Connection to Hermite Polynomials) If E Rd ||w||| = 1, then
he((w, x)) = (Hew (x:), 2000) (Stein's 's Lemma) For x NoO0, Ia), f E I.²(Ha))
Ex Hew (x)] Exe | fx))]
42


Reviews:
[{"review_id": "BM3RBytPKL", "reviewer": "Reviewer_Wfua", "summary": "This paper is a natural extension to Cabannes et al. (2024) in three ways:\n\n1. It extends the theory to a learnable task with trainable weights, not just studying random patterns\n2. It considers the MLP and attention as contributions to the Associative Memory properties of transformers\n3. It considers the (single-layer) transformer as a gradient flow\n\nThe goals are the same: to study the \"memorization capacity\" a.k.a. \"scaling laws\" of a transformer from the perspective of Associative Memory. The authors study this memorization capacity on a synthetic recall task consisting of memorizing subject-relationship correlations.", "questions": "(Q1) Standard Transformers typically output a probability distribution over all output tokens. The theory of Eq. (7) only considers the most-likely token (equivalent to setting the sampling temperature of LLMs -> close to 0). Indeed, there is a large debate in the community over the relationship between memorization and generalization. How could the memorization study of this work be scaled to study the *generalization* abilities of Transformers? I feel this theoretical framework could be used to tackle this problem.\n\n(Q2) The paper only studies the causal-language modelling (i.e., GPT-style) objective. How would the theory extend to transformers trained on masked-token (i.e., BERT-style) tasks?\n\n(Q3) Is $\\sigma$ in Eq. (2) defined, or just an arbitrary non-linearity? The theory of [Dense Associative Memory](https://arxiv.org/abs/1606.01164) says that the choice of this non-linearity is critical to increasing the memory capacity of Associative Memories\n\n(Q4) The experimental details around Fig 1 are incomplete -- I do not think I could replicate that experiment with the information included in the paper. Additionally, Fig 1 could be improved with error bars and notable \"failure cases\" e.g., the task of storing one-hot vectors [L130] or under different choices of $\\sigma$ in the MLP (see (Q4)).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "(S1) **Possibly high impact**. If the analyses performed in this paper are shown to hold for deeper transformers (with shared or unshared weights) using softmax activations in the attention and bigger datasets, the paper has the possibility to be high impact -- it would solidify the benefit of viewing transformers as formal Associative Memories, uniting the LLM community and the more niche physics- and math- communities studying associative memories. \n\n(S2) **Good, formal definition of Hallucination**. In a passing statement, this paper defines hallucination as the stage where [L455-460]\n\n> ...the model ignores all other tokens in the sequence $z_{1:T}$ – including the useful subject token $s$ – and predicts based only on the relation r... the model is outputting a plausible, yet ultimately incorrect, answer to the prompt. \n\nAs the associative memory perspective of transformers continues to gain traction and the theory of this paper extends to practically sized transformers and real NLP datasets, I can see this formal description of hallucination (i.e., the difference in the Cross-entropy loss of relationships and subjects) as game changing for the study of hallucination in LLMs. Related works do not seem to agree on a definition for \"hallucination\" except that it encompasses \"undesirable predictions\" from an LLM.", "weaknesses": "(W1) **Unclear application to real Transformers and datasets**. (dual to (S1)) The simplified Transformer architecture studied theoretically and empirically in this work is smaller than real Transformers, uses only a single update step for much of the analysis, only uses a synthetic dataset to validate the theory, and drops the softmax from the attention. It is unclear how the theory would generalize to unshared weights and real datasets where the task is to reproduce sequences of data.\n\n\n(W2) **No comparison to the empirical scaling laws of modern LLMs**. Like (W1), this paper could be improved if the theoretical results are compared to the empirical results of e.g., the [Chinchilla scaling laws](https://arxiv.org/abs/2203.15556).\n\n(W3) **Associative memories are single-step updates**. One limitation of this work is that the associative memories are not energy-based models where the memories are encoded as fixed points of energy descent dynamics. Indeed, the memory capacity of the linear associative memory model would not scale linearly if formulated as an energy-minimization problem. It would be nice to see memory capacity studied in the recurrent setting.\n\n**Summary**\n\nDespite these weaknesses, I believe that the contributions of this paper earn my vote for acceptance at ICLR. The paper is clearly presented, with testable claims, and sufficient empirical results. My concerns can be addressed in follow-up works by the community. \n\n**Note**\n\nThe MLP + attention architecture of Transformers has been explicitly formulated as an Associative Memory (specifically, a Hopfield Network with energy descent dynamics) in [Hoover et al. Energy Transformer (2023)](https://arxiv.org/abs/2302.07253). Please add this work to the related work section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eBUBkjHefH", "reviewer": "Reviewer_kBfg", "summary": "This papers show shallow transformers can use a combination of associative memories to obtain near optimal storage capacity. The verify this on a synthetic factual recall test.", "questions": "1. How did you get the plot for Figure 5 (left)? It seems a bit too smooth. Also, why didn’t the first 2000 GD change the loss? It seems counter-intuitive. \n2. Could you give an intuitive explanation for why the model, in theorem 6, learns conditional distribution of relations as opposed to conditional distribution of subject?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The theoretical contributions are solid", "weaknesses": "1. The theoretical toy model is rather simple (MLP and one-layer transformer). It is unclear how it generalizes to multi-layer transformer.\n2. If I understand correctly, the synthetic setting relies on the fact that noise tokens and subject tokens are disjoint. In reality, usually the last token (not eos in practice) should be somewhat relevant for token selection (for example, “in” as the last token would look for locations). It doesn’t make too much sense to let model do next token prediction after some random tokens.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "npWxIKCFtz", "reviewer": "Reviewer_fXTn", "summary": "This paper examines factual recall in shallow Transformers, proposing that transformers can leverage associative memories for near-optimal storage of factual information. The authors demonstrate that linear and multi-layer perceptron (MLP) associative memories achieve storage capacity that scales linearly with the parameter count. They design a synthetic factual recall task to show that a single-layer Transformer can reliably store and retrieve facts using self-attention with or without an MLP. Additionally, their gradient dynamics analysis in a simplified linear attention model reveals a \"hallucination\" stage, where the model’s loss plateaus and predictions rely solely on relation information.", "questions": "- The statement on line 117—that the number of associations scales linearly with the number of parameters (up to log factors)—is somewhat unclear. Without injective assumption, it looks like the number of tokens, not associations, scales linearly. Please elaborate on that point. \n- Section 4.2 appears to omit LayerNorm, which is part of the original Transformer layer. Would the theoretical results still hold if LayerNorm were considered?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- This work offers a theoretical understanding of how shallow Transformers handle factual recall, specifically revealing linear and MLP associative memory capacity in managing storage.\n- The gradient flow analysis, particularly the identification of an intermediate \"hallucination\" phase, adds depth to the understanding of transformer training dynamics.\n- The paper demonstrates mathematical rigor in its proofs regarding associative memory storage capacity", "weaknesses": "- The findings may be somewhat narrow, given the reliance on synthetic tasks. While the theoretical insights are valuable, it remains uncertain how well these translate to complex, real-world scenarios involving non-random and interdependent factual associations.\n- The study is centered on shallow models, raising questions about the applicability of these findings to deeper, large-scale transformers commonly used in research and industry.\n- Certain sections, particularly those describing empirical setups for synthetic tasks, could benefit from added detail. Greater specificity on initialization, parameter tuning, and model configuration across trials would improve clarity and reproducibility.\n- More discussion on related work exploring neural networks and attention mechanisms as associative memories [1,2] would strengthen the paper and highlight the novelty of the insight. \n\n[1] Radhakrishnan, Adityanarayanan, Mikhail Belkin, and Caroline Uhler. \"Overparameterized neural networks implement associative memory.\" Proceedings of the National Academy of Sciences 117, no. 44 (2020): 27162-27170.\n\n[2] Le, Hung, Truyen Tran, and Svetha Venkatesh. \"Self-attentive associative memory.\" In International conference on machine learning, pp. 5682-5691. PMLR, 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)