Scaling Laws for Associative Memories
Authors: Vivien Cabannes
Elvis Dohmatob
Alberto Bietti
Venue: ICLR 2024 spotlight (2024)
Keywords: scaling law
associative memory
mechanistic interpretability
Hopfield network
Primary area: learning theory
TL;DR: Mechanistic interpretability of memorization (modeling transformer intermediate layers)

Abstract: Learning arguably involves the discovery and memorization of abstract rules. The aim of this paper is to study associative memory mechanisms. Our model is based on high-dimensional matrices consisting of outer products of embeddings, which relates to the inner layers of transformer language models. We derive precise scaling laws with respect to sample size and parameter size, and discuss the statistical efficiency of different estimators, including optimization-based algorithms. We provide extensive numerical experiments to validate and interpret theoretical results, including fine-grained visualizations of the stored memory associations.

Published as a conference paper at ICLR 2024
SCALING LAWS FOR ASSOCIATIVE MEMORIES
Vivien Cabannes FAIR, Meta
Elvis Dohmatob FAIR, Meta
Alberto Bietti Flatiron Institute
ABSTRACT
Learning arguably involves the discovery and memorization of abstract rules. The aim of this paper is to study associative memory mechanisms. Our model is based high-dimensional matrices consisting of outer products of embeddings, which relates to the inner layers of transformer language models. We derive precise scaling laws with respect to sample size and parameter size, and discuss the sta- algorithms.
tistical efficiency of different estimators, including optimization-base We provide extensive numerical experiments validate and interpret theoretical results, including fine-grained visualizations of the stored memory associations.
INTRODUCTION
As the scale of large language models (LLMs) keeps increasing, scaling laws have become a crucial tool to empirically assess and predict the behavior of these models when varying the number of parameters and training data (Kaplan et all 2020; Hoffmann et all, 2022).. Despite their practical impact, the underlying phenomena leading to such scaling laws remain poorly understood. A better understanding of such phenomena could guide researchers towards improved models, algorithmss, and datasets which may lead to improved scaling laws.
Our study focuses on simple model that aims to be representative of LLMs in two ways. First, we focus on heavy-tailed data distributions over discrete tokens, a natural assumption for text data (Pr antadosi 204). Second, consider associative memory models that store input-output pairs through outer-products of finite-dimensional embeddings. and can be seen as proxy of the in- termediate layers of transformers. Indeed, some transformer layers have been found to behave as key-value memories (Geva et all 2021; Meng et al, 20222), and more generally outer-product as- sociative memory matrices arise naturally from training dynamics on intermediate weights (Biel al 2023). Beyond simple associative recall, the combination of multiple such associative rules at different layers may lead to certain circuits with rich "reasoning" behaviors based on context (Elhage etall 2021 Brett et all 2023; Michaud ala 2023). For example, an intermediate layer input token may encode for the topic "linux", leading to an output token that will trigger specific behavior in the transformer's 's following layers when processing the token "terminal".
Our contributions are as follows:
We provide precise statistical rates for outer outer-product memories with random embeddings, and compare different memory storage schemes in the context of Zipf-distributed data. We compare theoretical schemes to the weights learned by various optimization algorithms used in practice, and illustrate the role of different design choices with numerical experiments.
Related work. Associative memory models have a long history in the literature on neural com- putation (Stembuch 1961; Willshaw et all 1999;; Bonguct-Higging et all, 1970, Kohonen, 1972: Amari, 1972 Little 1974: Hopheld, 1982: Smolensky, 1990 Schlag et all, 2021: Valle Lishoa all 20239, though the statistical insights we provide for continuous-valle heavy-tailed tokens distributions are new, to the best of our knowledge. Memorization behaviors have drawn a lot of attention recently, and are believed to be an important notion to understand the learning happening in deep neural network (e.g., Sukhbaatar et all 20 [Eeldman, 2020: Ealdman & Zhang, 2020 Geva all 2021; Wi etal 2022). Building on memorization and heavy-tailes discrete data, our model bears similarities to the ones of Hutter (2021). Michaud et all (2(23) or Debowski (2023), although we focus on practical models with finite capacity. The discrete nature of
random embeddings and

---

Published as a conference paper at ICLR 2024
10-1
Error
10-1
Error
T = 10
d = 50
T 10
d 100
10
10a
10-2
200
10
1000
10
10a
10
10
10a
10s
10
d
T
Figure I: Scaling laws with respect to model capacity d (left), respectively the number of data seen T (right), for various numbers of dataset size T, respectively various model capacity d. This plots validates empirically theory developed the paper that proves scaling laws {(fa) x d= + 1 (dashed lines) under our setting with a = (ID) (1), (E), and the association scheme (12) with p = and P = d/8/8. The experiments averaged over 100 runs, standard deviations are shown with solid color.
tokens contrasts with other recent works on scaling laws that have focused on continuous Gaussian inputs (e.g., BBahi et all, 2021: Maloney et all, 2022: Sorscher et all, 2022).
2 MODEL FOR ASSOCIATIVE MEMORY
The data. In the following, we consider a joint distribution p E [[N]x[M] on inputs x E [N] and outputs y [M]. The inputs and outputs are respectively assumed to solely take N and M discrete values respectively. For example. N could be the number of potential sequences of fixed word length in the English language, while M would be all the potential words to complete the sequence. Abstractly, x and will be referred to as tokens. To simplify the study, we assume for now that deterministic function of x, i.e., there no noise in the labels. consistency with language modeling, we equally assume that p(z) follows Zipf law. Formally, there exists an parameter
0, a normalizing constant C.. permutation a G, and a function f- [N] + [M] such that p(a(z)) =
VH,y E [N] [M],
p(y/z) f=
(1)
The distribution p is not known, but has generated T known independent samples (11,, ye [E(]] p. For readability sake, we will assume without restriction that a is the identity (so that p is decreasing).
The model, and the loss. The input tokens are embedded into space Rd of dimension d through an embedding map e : [N] - Te.. This space is used for computation purposes. In particular, we focus the linear transformation parameterized by a matrix W E Rdxd mapping W We(r). This latter vector is mapped back to the output space through an unembedding map u [M] - I2d and the decoding rule
Sw (w) arg max u. Wez; yE[M]
W W
where ε, and by are abbreviations for exx) and u(y). The model (α) can be seen as analogous an attention layer where keys e, are tested against queries uy through a matrix W before going through a softmax layer, which, when the attention is peaky, identifies to an argmax. It also resembles next- token prediction from an intermediate representation Wez, which may itself be the output of an attention block that attends to a token æ. The matrices W will be expressed as associative memories. Memory of an observed pair (x, v) is represented as outer product u Remembering those with respect to a probability q E A [N][[M] leads to the matrix
W
Σ
((I,y)u,eT
9 E [[NTX[M)
[M]
This representation (3) is justified as the predictions (2) are insensitive to modifications of M outside the span of (Myee Jz.g- (z.g.)) In our deterministic setting (II) where one only observes pairs (x)),, we shall consider the simpler model where
W
Σ ((E)USLEEC ZE[N]
q
A []
(4)
It should he noted that the proof techniques hehind Theorem II do not break when considering 9 g( v): both models would lead to similar results, with the case q = q(2, y) being simpler comprehend.
2

---

Published as a conference paper at ICLR 2024
Table 1: Summary of key elements in the study. We are given discrete tokens y with deterministic relation y f-rx). We emhed tokens in RR", acts as "model capacity" parameter. We store association y the matrix W through scheme q and recall them through the decoding fq. We will first study scaling law of the generalization error & as a function of the number of data '', and the model capacity for different schemes q. We will later study the scheme found by optimization- based algorithms.
Tokens
Embeddings Ex Uy E 12d ~ 1)
Model
Scaling
Vi, = (x() E (1,2,...,
E, ( f((x) arg max, My Wen
W
E(g) E (1 E(g) - ((d,T:g)
Table 2: Some insightful provable scaling laws with respect to the model capacity d. and the number of data T, two schemes that store associations as (L) and random embeddings.
Model
Error scaling
Comment
((x) p(z)
d 1ll11)/20 +T T-11/
Found with large batches in one step
g(x) l dati T-1+1/a
Optimal scaling with random embeddings
To simplify notations, we will write Sn for Sw, (2). The model Sg is seen as superposing memories since all associations are mixed together in single matrix. The quality of mapping / quantified through the generalization eπror
{f) = Exx,))
f [M] + [M].
(5)
Which questions are we interested in? Several questions naturally arise from our model. The first ones are related scaling laws: how does the error depend on T, the number of data? How does it scale with d that encodes for model capacity? The second ones relate to the model itself: how does the error behave for different g? What about optimization based algorithms?
Arguably, the model (2) lays out simple model to study memorization, which could easily be ex- tended to model more intricate memorization and training behaviors inside a transformer language model. Indeed, memories of the form (g) were found to accurately model the behavior of weight matrices in multi-layer transformers trained by gradient methods on certain tasks (Bieth 2023).
Hence, expect our study to be generalizable to more complex mechanisms in transformers, re- sulting in rich token interactions to predict the next token in a sequence.
SCALING LAWS WITH RANDOM EMBEDDINGS
Why do we make errors? With simple deterministic model, one may wonder how can not learn perfectly the mapping f.. There are two sources of error. One is due to not having enough data to see all the potential association (Ir, (z)), and has already been studied by Hutter (2021)). The other one is due to the limited memory capacity of our model, which we illustrate in Figure D.
Proposition 1 (Finite data, infinite memory). Consider infinite memory model f. which at time T predicts correctly all x that were seen in the past training, i.e., x E {XthtetT) where the (X,, Y.) were drawn independently at random from distribution generalization error reads, with respect to the random dataset Dr = (X,, Y)le(T),
A [[]] [M] Under the data model the
EDy (E()) x 1-1///a.
(6)
Here, the notation b means that there exist two constants с₁ and Co such that c, S a c2b.
3.1 TIGHT ERROR CIARACTERIZATTION
The case where one has infinite data but finite memory is intrinsically a deterministic problem. How- ever, characterizing interferences between embeddings and the corresponding generalization error is combinatorial in nature, and is hard to study without specific assumptions on the embeddings and A natural choice is to consider them to be random, is the case at initialization.
Theorem (Infinite data, finite memory). Let M 2 4 and d > 8 logg( (.). For any memory weight scheme q : [N] + R, when the embeddings e, are independent random variables e, NO I), and the unembeddings are taken uniformly at random on the sphere.
Ee,u [E, inf (f [N] dd(x)) 16c, (0xx
8c,, 11111 }))
(7)
d
3

---

Published as a conference paper at ICLR 2024
1/3
1/2
62
Figure 2: Error due to finite memory capacity: the stacking of associative memories in matrix W may exhibit pattern W "r ff. (=)e, where three inputs mapped three different outputs interact in such way that 2z Wei 62 e1 + M2 Mire3 2 + uI uzez in Wei, so that fw(2 1) = 2 # = f. 1). In other terms, memory interference may lead wrong prediction, illustrating the finite capacity the model fw (2) store all data associations.
where Qx == maxy Ea:f-(a)=y (())3, c log(M) log(d), and p(X) ExEx denotes the probability of belong X c [M]. In terms of lower bound,
B.,. (E(G)
20 p(fr E [N]]3(d 1)((2)² Qx0)).
88)
Theorem M Willustrates how the error made by scheme g at the input x relates to the ratio between the signal dg(x)), provided by the associative memory (f) e and the noise Qoe. which corresponds to the signal provided by the most competitive class for E [M]. This is true up to a higher term in ||q 2 which corresponds to a class y = fx) competing against itself when the random embeddings exr for z' such that f- (x') ) point in the opposite direction of ez. When is large and p is regular, c, lell/| (12/ /d will be dominated by Qx and the cut-off of g(x)² /Q 32c./ / will behave similarly to a cut-off at 1 1/d up to logarithmic terms. Moreover, when g is chosen independently of p(yl d one can expect Qo p=l19ll? where p. = ma ply)) a consequence, up to constants and logarithmic term, we get
B[E(fA) p(Iz [N]] p-llgl1
(9)
3.2 MEMORY SCHEMES
Let us now discuss several natural choices for and compare their corresponding performance. The first naive choice consists in storing all the data seen at time T in memory. It reads
4o (x) 1 1x
go(x)= 1
(10)
Here, îo corresponds to the learned weighted scheme based on the T data, while denotes an idealized limit when one has infinite data. In the idealized setting (go Np. where p. max p(y). From Theorem II. we deduce that & {fw.,) will follow two regimes: an overflow regime where 3(d + 1) < Np. and in essence the memory W., is too full to recover any signal in it, and (ff. > 1/20 (B); infinite memory regime where d 2 N and all associations e, r. (x) can be stored orthogonally to another, and the error (fw...) quantifies the tiny probability that some random inputs embeddings appear to be too correlated.
Equipped with the knowledge that associative memory model (2) has finite capacity, may weight memories according to their frequencies, leading to the scheme, for 2 0
(x) = ( Σ Σ xx=x)
(p(x) = p())
(11)
1E[T]
better option consists in explicitly limiting the storage of our model with a simple thresholding algorithm
(x)
(x) = (ecee)==pp)n)acepp((((C ((=))))((()))))r)))
(p)1r(())= = p(x)11zeim)
(12)
where t(p)((x:)) denotes the set made the P most frequent inputs in the data (æx).
"To more precise, one should actually choose g(w) class dependent so cram memory as many as possible for each different class = (z), ensuring that y + Ex))(r))pp/ (x))² constant with respect to y. For simplicity, we will discuss this behavior that does not change the hig picture beyond our exposition.
4

---

Published as a conference paper at ICLR 2024
10
10
0.6
10
10a
10
Error
10a
10a
0.3
S 5/4/
10
5
10
10
10
10a
10a
10
d
d
Figure 3: Generalization error (N) as function of cl and T for the model (H) averaged over 100 runs. The data follows Zipf law with a 0.5, N = 100, M = 5 and f. (1) = mod. M. Left: enror for go (IT), either d too small and there will be memory overflow leading to large error (red area), either is big cnough and with enough data, the error will be null (blue area). Middle: error for Qa (II), for small and big T. avoid memory overflow allowing a smaller error then qui however for big d does not allocated enough memory to rare association, leading to a bigger error. Those results can be interpreted mechanistically by looking at the corresponding memory matrices (see Figure ItO). Right: Generalization error when T +00, N 100 and error decreases fast to zero (in blue); scheme 9i leads to an
2: the scheme go leads
zero-one type of plot where if
< the error is high, and d > d 1/4 as predicted by theory (in orange); the scheme Q (12) with P = =d/8,, decreases in d' (a-1) = d' until reaching the tipping point when d/8 > N (in green).
error decreasing d' da(a-1)/2
Proposition 2 (Without thresholding). Let p be o-Zipf distribution (.). For p >> O, the perfor- mance of Sp = Jap (II) is, up to poly-logarithm factors and constants that depends on both p and a,
-aa-1)/2pa
if >
E(pp) (log)
where p(N)= =
(N) N 2pix
2f =
(13)
((NN)
if <
In particular, when p E,, Ee (fo) scales in (a-a//) In the limit where p Ee, can be understood as (d/N)->0 which will go to zero if and only if is bigger than N.
Proposition 3 (With thresholding). Assume that p(x) follows a a-Zipp law (I) with N = +00. For 2 0, setting = d1/(200+1), the evror made by the memory scheme (12) scales as
(log) d-(a-1)((2pat1)
Eef)
(14)
In particular, when p == 0 and P~ one gets a scaling in d d++ which is actually optimal. The fact that this maximum is reached for P = d is reminiscent of Hopfield networks (Hopfield, (182) which can only store d/ log(d) patterns with d by matrix. Similarly, our model stores most associations, which, when in presence
Zipf law, leads to an error scaling in d (a-11)
Theorem 2 (Minimax performance). Assume that p(x) follows a a-Zipf law (I) with N +00. For any weighting scheme q, and p- E 1). there exists conditional distribution p(ulz) with p. max, ply) such that the error made for the distribution p lower bounded by
Cap.
E 2 (d
where
Co. = 20(a + 3
Moreover, this performance is reached (up logarithms factor) by the thresholding algorithm (l2) with P = d/ log(d) and p=
Finally, we prove that the scaling laws proved for d when T = +oo and for T when d = +o0 appears jointly when both and are finite.
Proposition 4 (Finite data and finite memory). For the previous bound with respect to d, Proposi-
tion a and Proposition B, considering finite data simply adds term T-1+1/a (up to constants and logarithmic terms), matching the optimal bound of Proposition D. In particular, with p = 0 and P ~ d/ log(d) reaches the optimal scaling in
T .e.. E(T)) x 1+1// d=
(15)
The optimal scaling (L5) recovers the law of Hutter (2021) with respect to T, and the one of Michaud all (2023) with respect to d. This is intuitive, since Hitter (202T) assumes memorizing exactly
5

---

Published as a conference paper at ICLR 2024
109
real approx
real
real
Error
Error
approx
Error
approx
10-1
10-1
10
0
500
1000
500
1000
500
1000
Epochs
Epochs
Epochs
Figure 4: Comparison between the error found by optimizing W (C) with SGD on the cross coos--y loss,
and its approximation with g(x)) (H) and the approximate update rule (20). We consider N = 100, M 5.
f=( mod. M, 2, and hatch size equals one. Left: One with d = N 100 with 7 10.
Middle: Average over 100 runs with
N 100 with 7 Right: Average when = N/11 = 10
with y = 1. which implies that our approximation is not valid anymore. The same results can be obtained for bigger batch sizes shown in Figure D
10
10
0.15
10-1
0.10
0.05
10
0.00
10-3
10
100
50
Figure 5: Theoretical approximation of the association scheme found with stochastic gradient descent with batch size equals one and fixed learning rates. Left: Plot of f0) a function of where is effect of one gradient update on ((x) (20).. Right: Plot
resulting 9- (z) when ney x p(r) x (ar + 3)-" " with a = and nu = 1. In dashed, we represent dg (II) for 0.05, p = 0.35 and 1. Those curves map well 4, for 10, 10 and = 10 respectively.
all previously seen data, while each memory could be seen as specifying a "quantum of knowledge" as modeled in Michard et all (2023), with d d corresponding to the risk (5) of only storing the
most frequent d tokens. However, associative memories can be understood at different level of granularity, and while one may argue that transformer acts as a big associative memory machine and derives LLMs scaling laws approximations as corollaries, we prefer to understand a transformer as combination of hidden associative memories suggested by Sukhbaatar et atal (20T9); Geva atll (2021)); Wu et all (20022); [Bietti et all (2023) among others.
4 OPTIMIZATTION-ASSHH MEMORIZAATION
This section studies memory schemes privileged by optimization the training dynamics behind memorization. In terms of relevance, argue that model (2) is proxy for the inner layers of transformer that memorize patterns before matching them against new data at inference time. As such, a transformer influence storage our memory model.
algorithms, digging into
understand how different key elements in the training of
Gradient updates. We consider the cross entropy loss as a surrogate objective minimize, and study the form of gradient updates on batches of data. Formally, the matrix W E Rdxd in (2) is optimized to minimize the loss
L(W)- Exx,Y) ((x, ys W)).
((x, y: W) - We, + log( Σ exp(u Wez)). (16)
zE(M]
The gradient of this loss with respect to W takes the following form, as detailed in Appendix A.l.:
Vwl(x, W) = -(1 pw (w (u) - e)eT with = Σ
pw(2/u: y)u..
(17)
xE[M]
where (w(y)æ) x exp(u,) We,) ez) are model predictions for the current W. For a batch of n data
B = [mi,"" ... In], a gradient update with step size 7/t updates W, as
W W4+1=
W, Yt
Vwl(x, f.):
(18)
ZEB

---

Published as a conference paper at ICLR 2024
Figure 6: Gradient descent dynamics from perspective of the matrix (u,, Wies)y E RMNN with N = 10, M = 5, = 1.5, f-(x) = 2 mod. 5, and d = < N. A lighter color in the square (y, x) means higher value of u, We.. The optimal W corresponds to two diagonal strips of yellow boxes (see Figure DT. The matrix W. is updated with stochastic gradient descent with batch size equal to one. From time to time, stochastic gradient descent will hit an association that is not properly stored in memory yet (the red boxes). It will consequently update the weight matrix Wt + We+1 (side by side pairs) to store (TE). Left pair: update with a big learning rate Y 10, whose risk is to erase previous memories (the light colored boxes), similarly to qu (IT). Right pair: update with a small learning rate = 10 which will not store rare memory, similarly to 9, (III) with large p.
Approximation of the updates. When pw (z(x) does not change much for all 2 * f. (2). since Ur were sampled at random in Sd, we expect E (lZ)) to concentrate around zero with ||=|| = 1/M, hence to be negligible in front of (x). As a consequence,
Vw (1:); W) R (1 pw (f)))))))))
(19)
This is notably the case for W 0, random W. if W only stores pairs fz. (x)) with d > N. With the update model above (I9) T' steps of SGD with batch size one lead to an association scheme of the form (HI) with (see Appendix [IIII)
q- (x)) = TT(z) (0) fofo.... fo))
where
f:z
(20)
1 M exp(x))
Tp(z) times
This equation tells us what form to expect for for optimization schemes with different hyperparam- eters. This approximation shown in Figure D, and is validated empirically in Figure g1
Step size effect. When d > N, the updates approximation (20) and the resulting show how large learning rate y is beneficial for our problem, in particular when using SGD with batch size one. Interestingly, the same behavior holds in the presence of limited capacity, i.e., d < N, although (I9). In those settings, we resort to numerical simulation to study how optimization manages to rearrange memories. Figure E showcases two types behaviors depending on the size 7. (i) When the learning rate 7 large, associations will be stored easily in memory. but will tend to overwrite previous storage. (ii) When the learning rate 7 is small, associations need be seen often to build up in the matrix W (eI which will take more time, but will not erase memory. This provides another intuition explanation for why a bigger step size leads to better results on the left Figure Ø. The previous considcrations also explain the usefulness of scheduling in our simple model, which we illustrate Figure I[ using large learning rate enables us to store associations while there still memory space, while reducing later in training avoids overwriting previous storage unless an association is highly frequent.
interferences between embeddings (Figure D) break our approximation
Batch size effect. Table 2 recalls how storing associations with g = under the model (e) is better than storing them with q p. As such, it suggests that, when processing finite number of data T. smaller batch size is preferable. Intuitively, processing an input æ in a batch will reweight it by its frequency p(z). while processing it by itself will update W similarly to setting (x) 1 if has not been already seen. Indeed, in the large batch limit where [B] + +o0, one batch update corresponds population gradient update, which when Pw < assimilates to Vw L(W) x -. (w)e' This contrasts with many small batch updates that rather lead to an association scheme akin to (eI with y = 1. support of this line of reasoning, Figure (middle) illustrates the benefits of splitting the descent with many steps, with small batch size and large step size.
4.1 PRACTICAL CONSIDERATIONS
In order to optimize our simple model the fastest, we have seen the usefulness of large step size and small batch size. However, for large transformers such design choices are impractical. First, large step sizes may lead to instability in realistic models (Gimeer al. 2021). Second, in order to reduce training time and improve hardware efficiency, one should process large batches (Smith all 2018).
7

---

Published as a conference paper at ICLR 2024
SGD. [8]] T-10240
SCD, 7-1024
1022,
10240
16.7-1
1064
Etror
Error
Error
10-1
10-2-
SGD
Adamm
BGB--122
Adao-EN
10a
10 10 10
10 10
10
10a
10a
Figure 7: Effect of step size, batch size, layer-norm and Adam (with 3: = B2 0, which corresponds
to SignGD). All the experiments are conducted with N = 100, M = 5, = 2, fx) z mod M, runs. We initialized parameters and rescale learning rates updates, as explained in Appendix B.I. To avoid confounders, we scale ry on the middle plot for the variance of the gradient updates to be independent of the batch size.
averaged
maximal feature
Adam. We have seen before how the update of SGD with large batch can be approximated with -L-
7' (We+1-W.-Jj)
pw (
F
Those naive updates would lead to model that resembles (2) with 9 pf for p 11=111e In con- cordance with previous research on the matter (Zhang et all 2020, Kunstner et all 2023), we found Adam be helpful in our setup as well, see Figure 7 (right). In first order approximation, Adam is approximated as signSGD (Balles & Hennig, 2018). Arguably, this introduces a normalization effect to the gradient, helping to reach the saturation phase of + f (2L) shown on Figure E homogenizing the resulting matrix W to behave similarly to Q1 = 1, therefore optimizing memory capacity. Experiments to underpin this intuition are reported in Figures LLS and IGG in Appendix [B
Layer normalization. Minimizing the cross-entropy loss implies setting pw (ylæ:) 1, which will lead to W diverging to infinity and unstable loss gradients. In order to ensure numerical stability, is natural to rescale the vector We, RR, especially since what matters for the final prediction fw is only its direction. This is precisely what layer-norm does, introducing the logit score
We
9y (x)))) (uy,
instead of (T(x)) Ugy We
|Weal
This leads to an added projection on the gradients in (L), as detailed in Appendix denoting W- W/llWezll
WwelN (z, y. W) = Vwl(s, W)
(w
(21)
[Wea
We recognize a projection that kills the signal that already aligns with Wez. We conjecture that this introduces a clipping effect on the corresponding q(x), optimizing for memory storage, and explaining the good performance observed in the right of Figure Ø.
4.2 THE BENEFITS OF LEARNING THE EMBEDDINGS
Taking a step back, Theorem II implies that our model with d2 parameters, the matrix W E dgddd H. only memorize about d/ (og(d) associations (ez; Myy E (I2d) of size 2d. Intriguingly, Lemma below states that an exponential number of quasi-orthogonal elements can be put in R4, an event that actually holds with high probability when embeddings are random, showcasing intrinsic limitations of our "linear" model (2).
Definition 1 (Qussi-orhoooonality) The family (u=)ze(p) with us E Rdd is n-quasi orthogonal if
and Lemma 1. For any d E N and P 2 3, there exists an embedding u [P] + Rd such that the family (uz)zejr) is = / (oo(())-quas))) orthogonal.
A v(2,2'} C [P], Kuz, Uz>) S n,
||u=|| 1.
(22)
As a consequence of Lemma MI the following model
fi (x) arg max uy
Σ ZE[[
urcero((ebee -));
(23)
y
8

---

Published as a conference paper at ICLR 2024
Figure 8: Experiments with learned emheddings when = 2, N 100 and M = 5 with = (z) x mod. M and d = 2. Left: level lines of the function IR² - [5]; + arg max (e(s) My u with uy the learned unembedding. Middle: seatter plot of the learned input embeddings ez E R for x E [N] colored accordingly to fz(z) for ez. It illustrates how the input embeddings match with the output ones, similarly to ( and Proposition & Right: learned input embeddings obtained with M = 10, and allowing again zero generalization emor. Reaching
with d' = 2 greatly contrasts with the condition 2 needed to get to zero generalization enror when the embeddings are random.
where o(x)) 24 is the ReLU function, can fit P exp(n) elements in memory, leading to a scaling in E(f)) x exp(-(a - 1_) d when p(z) follows a a-Zipf law. Similarly, one could consider higher moments of ex which has been the basis for modern Hopfield networks (Krotov & Hopfield, 2016; |Ramsauer et all, 202I). However, implementing the model (23) requires keep track of each of the P vectors ez R&d leading to Pd parameters, in order to only store associations of size d, needing compute that scales with Pd at inference time, rather than just d,
We also note that when embeddings are learned, it actually possible to store as many memories as desired, which can be seen from the fact that
W I, Vy E [M] Uy E &S = (f)(x)
fx(x) arg max uy Wez,
(24)
In particular, Figure N illustrates the solution found when d = 2 by optimization-baser algorithms in order to get zero generalization error on the task of Figure where M = 5. Optimizing token embeddings is probably an important element to increase memorization capacity in transformers, although enforcing E₂ = u. is unrealistic when embeddings are shared over different heads, and the input/output relationships to be learned differ across heads.
5 CONCLUSION
This work considers a simple model to study memorization in transformers. Here, memorization seen a valuable behavior, the network memorizing useful patterns and association rules. We derive precise scaling laws with respect to both the number of data, and the model size, which plays the role of model capacity. We quantify the effect of different memorization schemes, illustrating the benefits of uniformly weighted outer products. We leverage these theoretical results to study how different optimization algorithms commonly used for transformers may lead to more efficient memorization. In particular, we showcase the efficacy of small batches and large learning rates, and, under the design constraints resulting from efficient hardware utilization and training stability, the usefulness of Adam and layer normalization.
While our study focuses on simple memorization schemes, it opens up many possible new direc- tions. This includes extending study to richer models that are closer to transformers, where embeddings, attention and feed-forward layers are trained. This could allow models scaling laws that capture interactions between tokens, well as hierarchical behaviors that require multiple lay- ers. We would equally like to leverage our framework for assessing memorization and generalization through clear metrics, and eventually automatically adapt the learning rates function of the "free" memory capacity left in layer.
The authors would like to thank Léon Bottou as well Hervé Jégou for many fruitful discussions on memorization mechanisms in transformer language models.
Acknowledgements.
'This result follows directly from two facts. When input embeddings are chosen at random, the probability that they are not 7-quasi orthogonal is bounded by pd exp( dn² /2). When input embeddings are r--uasi orthogonal, fz) fx) for any E (P).

---

Published as a conference paper at ICLR 2024
REFERENCES
Shun-Ichi Amari. Learning patterns and pattern sequences by self-organizing nets of threshold elements. IEEE Transactions on Computers, 1972.
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jachoon Lec, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXiv: 102.6010,,, 2021.
Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance of stochastic gradients. In ICML, 2018
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: memory viewpoint. In NeurIPS, 2023.
Thomas Cover and Joy Thomas. Elements of Information Theory. Wiley, 1991.
Lukasz Debowski. A simplistic model of neural scaling laws: Multiperiodic santa fe processes. arXiv preprint arXiv:2302.09049 2023.
Ian Dinwoodie. Mesures dominantes et théorème de Sanov. Annales de I'Institut Henri Poincare, 1992.
Nelson Elhage. Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-oodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Ka- mal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Technical report, Anthropic, 2021.
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-oodde Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy models superpo- sition. Technical report, Anthropic, 2022.
Vitaly Feldman. Does learning require memorization? a short tale about a long tail. In STOC, 2020.
Vitaly Feldman and Chiyuan Zhang. What neural networks memorize and why: Discovering the long tail via influence estimation. In NeurIPS, 2020.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In EMNLP, 2021.
Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Car- doze, George Edward Dahl, Zachary Nado, and Orhan Firat. loss curvature perspective on training instabilities of deep learning models. In International Conference on Learning Represen- tations, 2021.
Wassily Hoeffding. Probability inequalities for sums of bounded random variables. Journal of the American Statistical Association, 1963.
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. In NeurIPS, 2022.
John Hopfield. Neural networks and physical systems with emergent collective computational abili- ties. Proceedings of the National Academy of Sciences of the United States of America, 1982.
Mareus Hutter. Learning curve theory. arXiv preprint arXiv:2102.04077 2021.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836/ 2020.
Teuvo Kohonen. Correlation matrix memories. IEEE Transactions on Computers, 1972.
10

---

Published as a conference paper at ICLR 2024
Dmitry Krotov and John Hopfield. Dense associative memory for pattern recognition. In NeurlPS, 2016.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. In ICLR, 2023.
William Little. The existence persistent states in the brain. Mathematical Biosciences, 1974.
Christopher Longuet-Higgins, David. Willshaw, and Peter Buneman. Theories of associative recall. Quarterly Reviews of Biophysics, 1970.
Alexander Maloney, Daniel Roberts, and James Sully. A solvable model of neural scaling laws. arXiv preprint arXiv:2210. 2022.
Kevin Meng. David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in GPT. In NeurlPS, 2022.
Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quamtization model of neural scaling. arXiv preprint arXiv:2303.13506 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Ed- ward Yang. Zach Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, nigh-performance deep learning library. In NeurlPS, 2019.
Steven Piantadosi. Zipfs word frequency law in natural language: A critical review and future directions. Psychonomic Bulletin and Review, 2014.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlovi, Geir Kjetil Sandve, Victor Greiff, David Kreil, Michael Kopp, Güitter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hop- field networks is all you need. In ICLR, 2021.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In ICML, 2021.
Samuel Smith, Pieter-Jan Kindermans, Chris Ying, and Quoe V. Le. Don't decay the learning rate, increase the batch size. In ICLR, 2018.
Paul Smolensky. Tensor product variable binding and the representation of symbolic structures in connectionist systems. Artifical Intelligence, 1990.
Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. In NeurlPS, 2022.
Karl Steinbuch. Die Lermmatrix. Kybernetik, 1961.
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Aug- menting self-attention with persistent memory. arXiv preprint 1007:0190000141 2019.
Juan Valle-Lishoo,, Andrés Pomi, and Eduardo Mizraji. Multiplicative processing in the modeling of cognitive activities large neural networks. Biophysical Reviews, 2023.
David Willshaw, Peter Buneman, and Christopher Longuet-Higgins. Non-holographic associative memory, Nature, 1969.
Yuhuai Wu, Markus Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In ICLR, 2022.
Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the infinite width limit. arXiv preprint 238.v22308.08144 2023.
11

---

Published as a conference paper at ICLR 2024
Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. In NeurlPS, 2021.
Jingzhao Zhang, Sairaneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? In NeurIPS, 2020.
12

---

Published as a conference paper at ICLR 2024
A PROOFS
A.1 FINITE DATA PROOF OF PROPOSITION II
Let us consider the infinite memory model, where an LLM can store memory all previously seen associations (z, y). At each time f, random positive integer is drawn from some fixed probability distribution. At time T, the LLM would have seen T1,... 2T and the associated fe (Zz), where each 2, is random positive integer drawn independently from p. As such, the LLM would have learned map f, that only miscorrects the inputs which are different from all the 24 for & [7]. The generalization error reads, with respect to the random dataset Dr = (Xt, Yi)eEm
ED, = PX,D, (X & (Xxhemm) = Σ
p(x) & (XoJuerr))
Σ ME[N]
(()(1-p(x)))
=E[N]
Using that (1 a)²² exp(Tloo(( - a)) and log S log(1 a) for any -1/2, get
Σ 1/())/1/2 p(x) exp( (2g(-=)))())((( S Σ 2p(e))xp((2loog()bb xE|M]
$ LD, Σ e(x))xxp(p(=))) =E[M]
Relating this series to the corresponding integral, we have
JEELLN p(x) ((t)cx((22lgg//) 2log(2)p(z)))dl 1/T
J2c(en 11/TH) - 1)-))p((22ll(02ppp 1)T) de
p(x)) p(o)exp(-2log(2))())...)
< Σ p(x) (200-222gg)))()111))) ED, Ifl
Σ ((=)exp(-p(a)T)
2=2
xE[N]
1/T
ZE[1,N]
Letting N goes to infinity, we get the scaling
EDr x f p(z)e Tp(z) dr + 1/T.
(25)
Assuming that p(x) Cf(x)) for some constant C, and smooth strongly decreasing function f: R4 + R4 such that lim f(z) = too, one may consider the change of variable u = f(x),
ie., x = (u). If so,
du d = d(f = fo f (2)
Hence it holds that
Ep
100
2
uf du.
(26)
fl
This relates to the Laplace transform of the function inside the integrand. In particular, one can work out that when p(z) x Ca fu) 2- -1/ from which one can deduce that
exp(-Txx )) dx
[
which recovers result of Huftei (2021).
13

---

Published as a conference paper at ICLR 2024
A.2 MEMORY CAPACITY PROOF OF LEMMA II
The proof of Lemma DI concerning quasi orthogonal embeddings can be done through a reasoning on random embeddings. Let (X;) be P independent identically distributed random variables. We are interested in the event where the normalized (X;) are 77-qasi orthogonal.
P( X3Xl S x||xxIIIIXx;I]] = (P [I,,, X 2 all x (II,II)
2 1 P(P - P( ž X||
2
If this event can happen, it means that there exists such 7--qasi orthogonal samples. As a conse- quence, we are looking to maximize 7 such that
P(KXI, Xa)1 Z 7 xXIxXXI <
(27)
1)
Let us consider (X:) to be distributed accordingly to a rotation-invariant probability. By symmetry, we have, with fi denoting the first vector of the canonical basis in R.
X [xIl||| (2l) = P( fi)l 2 x|| = P(K || fill 2 7)
P([(X,, Z n/ X1
(28)
By symmetry, the vector X is uniform on the sphere. 2P( fi) > 7)) and
Using that P([(X, fill > 7) =
2
/
P(((,, fill 2n) Vol(Sd-1)
1 10,2n dr
wEgd-1
2
Voll / - xg. 2 dxi dr1
Vol(Sd-1)
Ja1=
21 +1) /'
(1 (1- + dt
(1 (1-t)) -1- dt.
Vol(S*-1)
v + }) |
To upper bound this probability, we proceed with
P(KXAN)22
2 + 2
(1
- dt
2 + 1)1/2
t (1- + dt
( c + )
2
L=y
Vi
L=n 17
_ + 1)
v2
exp( _1dd
d+1 722)
2
1/(d +
1)
V V
2
The last inequality follows from the fact that
(d + 2) d+1+1 1 1+ म (d + 1)2
dd
= 1+ d
d+1 d+1
and that for any x E (-1, 1), the concavity of the logarithm mean that log(1 + x) S x hence that (1 + z)" exp(n log(1 + x)) S exp(nz).
This leads to the following series of implications
1
pdd
112
17
2
3 (X;) quasi orthogonal
expl
>
VT
2
2
p2
pdd
1/2
17'd
p2
expl
2
2
2
2/TT
r
11d
p2
>1
and
expl
2
2V7
rfd 2 2log(P)
- logg2 2
2
nfd
1 log (og((//)) T))
2 log(P) >
4
2
Finally, we have proven the existence of a 1u-quasi orthogonal family for
17 2 / (llo(P)d-)
long as
P 3.
(29)
14

---

Published as a conference paper at ICLR 2024
A.3 GENERIC ERROR DECOMPOSITION
The error made by fw relates to the ordering between the signals "f-(x) Wey and the noises fa (a) My W Cr.
Let fg be defined as in the main text. We have the following sequence of equivalence, assuming uniqueness of the argument of the maximum for simplicity.
(z0) # fx(0)
arg max DE[Mi
Σ =EIN]
q(x)e Gzo 220 (,)) + fx(o
ezO U (1()) Σ ((x)e
Σ
g(x)e (
max VE[M] xEIN]
#E[N]
max VE[M]
Σ
(x))ez
Ezo U1 (. (xo)) >
=E
As consequence,
E(fa)
Σ p(ao)1 &S.
zoE|M]
= Σ p(zo)1 mmxx, Exe(Ny g(x)ee exg ". (=))))))))))
500
ZoE[N]
In other terms, we have proven the following characterization, which holds for any q, even if derived from a finite number of data,
E(fa) = p((x [N]] max
Σ e'E[N]
-(,,a)) > 0f).
(30)
A.4 RANDOM EMBEDDINGS PROOF OF THEOREM II
Let us introduce randomness the model. If each e ~ (0,, is actually independent random Gaussian vector in Rd, we continue our derivation with
E. [E(f(]
Σ p(Zo) (oz))rro (fg( + (IO) exo)
Z.E(NI
= Σ
p(ro) Eesp [P (Pee)aazo
(max Σ
q(z)c_ Eza u ( (uy - u))(2))) > 0 ezo
ZE[N]
Σ )(Z0)
Beep [P (ex))at (max Zy 0 (x)))))
2.E[N]
Here, we have introduced the random variables Z, for # fx(o)), inheriting their randomness from (elexo), and defined by
Z, Σ q(x)e exo - .E[M]
(31)
Those are projections of Gaussian variables, hence are Gaussian. Using the fact that E[eg] = 0, their
mean is
My == E = q /1. (mo) (uy-uuy (x)()))) Those variables correlated. Using the characterization of the mean, deduce that their covari- ance reads
(32)
Eyr. B[( ElZn, - E(2y
Σ (())g(2)) Elee exp enp e.o (u ( (f.(
$200
-Un Us (so))
Σ
(xx)2 Te Eleze 11) - us,
=un, - ((oo))
Σ
((0)² e.. 1 - (f)(
15

---

Published as a conference paper at ICLR 2024
Finally, obtain the following covariance
Ey.y |lezo (uy "fa()o))
Σ zzzo
(33)
We are left with the computation of the probability that the maximum of the correlated, non- centered, exchangeable, Gaussian variables (Zy) is bigger than zero.
Generic upper bound. Since we do not care about the scaling with respect to M, we proceed with
max P(Zy S & (mmmax& Zy 50) & yE M
Σ
P(Zy 0) S
max P(Zy S 50) (34)
yE[M]
vE[MS]
which leads to P ((e))adeD (max
Σ
q(a)ee Cwo "o (.x)( us. > 0fe(sro)
DE(N]
Myy
exp( 10
25
(
ezz 1| g (1 (2o), Uy -
Σ (x))
exp(- 11-1 (,
2
Σxtro (2) us-(w0)
Finally, recognizing variable with d degrees of freedom, for any a > 0.
Ejexp(-allez,) ||3]] (1 +1+22a 2a) d/2 exp( 1xxp(-220p) d log +a)).
This leads to the final bound, with Xu,x = minye[M]
"10(=))00- -400))000
E. ( S E p(x) ppp))mmn()) m .E[N]
Σ (1+
g(z) (u. uy -
us.(ar)) -
4 Xux }. (35)
En'pa g(x (Uf.(x)), (x)
This holds for any unembedding u and associative weight scheme g. In the following, will assume that the unembedding u are such that Xu,x = 1, which is notably the case when the My are normalized (i.e., Uy E Id-1).
Matching lower bound.
Going back to (34), one can get matching lower bound.
Eg 1 N
Σ B(r)Ee, max P(Zy S O)ex)] V#I-(x)
DE[N]
2 Σ p(x) max E., Plyy I 0lez)]
v=f=(x)
=E[N]
1
Σ ((z)(1
max Ez, [ers( w7/f-(x)
Py
)))
2
V /
2E[N]
To conclude, need an inequality of anti-concentration for Gaussian variables. In essence, should distinguish two type of inputs x E [N]:
the ones where My Eg.y will be large enough store the association ( which will
lead to an error decreasing exponentially fast;
the ones where the same ratio is too small and that should count in the lower bound.
16

---

Published as a conference paper at ICLR 2024
Following this split, one go for the simple "survival" lower bound
1 - erf(t)
E. [E(a)l 2 sup
p(zo) max Beeg [1 1235220,.0 (
1>0
2
M0E[M]
1- crf(t)
Σ p( max
= sup t>0
2
V=S-(wo)
[M]
Pezo (Ilezol 11' (20) (zo) "y us.()m)) $ 21 E
- -)1(20)))))
crf(t)
2 sup t,s20
Σ
p(zo)]Fero (llexo 5(
max yfS.
2
2go) (00))) ".. (og) 221 Eiggion (w)
Without optimizing for constants, taking = 1/v2 and d, we get the simple "survival bound" that there exists a constant e such that
E. (E(a) 2
Σ #E(N]
a(2 (1 (on un.00)
(36)
The constant can be computed explicitly
er[(1/V2), P([ 1 S d) 0.158 1/2 0.079,
where we have used that leex 1|2 as x² variable with mean d hence smaller median, which implies that ([|ee 1| d) > 1/2.
Quasi-orthogonal output embeddings. Let us consider [M] Rdd such that (uy)eE(M) is 7-quasi orthogonal.
Upper bound. Going back (355), we can work lower bound with
((00) ((1.)) My (x)())))
E Eo (())2 Uyy - u (mn))
Entor (10)2 (1,)) + 112 + (200) -)) n)
(z0= 1 +1
((10) n)
2 4 Eayer ((x) (1 (x()== + (m)= ==((0))((o
(1-0))
= (,00 + 772) g(00)2
(000) (1)
1 (1- 172) Em.ts (x)) g(20)²
((20) (1- 1)
= 4 1 (1- (Q2 + (z.()) q((x)²²
Here, we have used that for the numerator
(u (2o): (xo); My
"f. (za)) )2 = ((1s, (mn)) Uy) - (1-
and the same for the term in the denominator (since their ratio cancels out), well as
(uys (xa))) < (1 ?
(ug, 2 < (2n)²
Moreover, we have introduced
Q
Σ d(a')²
(37)
17

---

Published as a conference paper at ICLR 2024
Using the fact that (1 + 2² = exp(dlog(1+ x)) z exp(dr), an upper bound directly follows from those derivations,
p(ao) min M [ep( d(1 (1)
g(20)²
[.(E(G)]]
Σ
}
(38)
4nj 1101 +0020
=oE[N]
where
Qx
max Qy
Σ
g(x)²
(39)
DC[M]
DE
yE(M
z:f-(x)=y
Matching lower bound. Similarly, one can work out lower bound with q(01)² (u uy - "))((a))? Enjrr (()) (U),(x)) Uy U (no)) Enyan ((0) (1)))(1,,00==(1 17)²
( +)²
17)2
q(50)2
14 Qp,(a)) 7(00)2
Combining this with (3G), we get the lower bound, with .079.
Ex (E())
Σ p(x)1 1(xx1 (4D =E[N]
(40)
Remark that in the previous lower bound, we have dropped the previous factor 7 19/l that appears in the upper bound. We expect this term to actually be present in a tighter error characterization. In essence, we expect the embeddings to fill the full space Sd-1 so that most of the difference (1)) Ml (zo)) typically behave 172. However, quantifying this precisely is beyond the scope of this paper.
Random output embeddings. In the case where the output embeddings random, we can dis- tinguish two cases. The cases where the embeddings are r--uaai orthogonal, where one can retake the previous derivations, and the case where they are not, which will have small probability if is large enough.
Consider be random embeddings taking uniformly on the unit sphere. Let us introduce the event E, {u is quuus orthogonal].
We have seen in the proof of Lemma that
M 2
1d
P(F,)) S
exp(
(41)
2 r7ad
2
For any random variable Z that is bounded by one, we have the bounds
P(E)E[Z|E) < E[Z] = (1 - (()) + P(E)E[Z
+ E]
(42)
The upper bound of Theorem M1 directly follows from plugging (L and (II) into this last equation
Ee,u (efu(E((f)]
/ ped exp(
M
2
n2d
Σ
Σ (1+
(1 n)² q(x)))
(43)
2
4
#E(N]
Since this is true for any 7), one can consider the infimum in the upper bound.
In term of lower bound, retaking HI
2 2
n2d
Ee. (E(g)) 2 c(1
exp(
Σ p(x)
(0+10
(44)
2/
dd
n20
xE[N]
In particular, when d > 8 log( M) one can consider
< 1/2 such that r/d > Mog( which leads
(17 1)/()) + 1) > 1 1/3, and, if M 2
2
rd
1
1
exp(
2/77 π / MM) 2/3.
2 2/ π 2d
All together we have proven that, as long as M 2 4 and d 2 8 Blog(M with c > .079 2/3 .052
and 02 > 1/3,
E (E()))
Σ
(45)
N
18

---

Published as a conference paper at ICLR 2024
Writing upper bounds survival bounds. Until now, have written the upper bounds as the sum of exponential (EX) and the lower bounds as a sum of missed associations (95), which we called "survival" bound. In order to best read how tight
characterization is, one rewrite the
upper bounds as survival bounds. In particular, as we did in the lower bound, we will dissociate x
corresponding to small exponential and the other ones. Using the fact that the p(x) sum one, get, when the output embeddings are 17-qasii orthogonal,
d(1 7)
((50)2 4n
[e(E(fa)]
Σ IoE|N]
p(zo) min M exp(
4 + 2Q )
2
t(1 7)?
Σ 2oE|N]
p(xo)
inf M exp( 130
4
((1-n)
inf exp
loggMM) +
Σ p(o))aagoo3322//1113-00020 zE[N]
120
4
To simplify the bound, consider the constraints
11² - /
and
" < 1/2.
(46)
get, using& t 16(log(M) + 7og((d)) df) for O. we get
1(1 7)2
Ee (f(g) I inf exp(
+ log(M)) + Σ
p(x) (401
130
ZE|M]
16
S inf exp( 130
p(x) dgd)))(())))31
16
s eexp(-7log(d)) +
Σ
*E[N]
Finally, when the output embedding 17-quasi orthogonal with satisfying (EG), get Ex(E(fa)] s inf d=>
Σ #()))laapp222sa1(0(
(47)
730
ME[N]
When the unembeddings are chosen at random, when d > 8log(M). one can choose nj < 1/2, and is cast as, chosen dr1² 4 log(MM 2y log(d),.
M
2
n'd
E (E(Ja) < inf
exp( exp(-
n.
2/T
1
2
+
+
Σ p(x) hatc))/soolooAAAAooooooooooheeee
[E[N]
d
inf
2 log M) + log(d)
+ d + Σ
1 (00 logg Myy logid))
[N(N]
< 2nf
Σ
p(x)1
+7 log(())) log( togtif)
l+ll++Q=""
=<[N]
Finally, we have shown that when the embeddings are taken at random
Be < inf 2d +
Σ p(x) 16()) a log(d) [1+l)+Q.w'* (48)
A.5 PROOF OF PROPOSITION
When p(z) ~ - x= = hence,
p(fx E [N] dg(x)² < p. = [M] || x < (d]lall 2211/22) ~ (d|lgl) (dl_all-))(a--112200 We are left with the computation of ((N) == 1 J q(x)² dx S 2pc dx. When 2pa > this integral reads 1 2apt1 which is bounded by one.
19

---

Published as a conference paper at ICLR 2024
A.6 PROOF OF PROPOSITION E
When p(x) = x g(x) - 1etp)p(()) = 12E[Pp0 pa, we get
= p((2 [N] ddqx))² s p p((x E [P]]d s sp.llal?)) p((x PJ)
d (PP)
) 1a-1p/2pa
+ P p
The optimal threshold P is set by equalizing the two terms, which we compute as
(a- 1)/2pa = p-at1
((P)
a +1 1 log(d) -a +1
log(P) (P)( (-a 1) log(
2pa
2pa
log(d) - log P)) = 2pa ()
p= d2/(2pa+1).
This choice of P leads to scaling in, with So.[P] = Sam
Ee.u(E(fo.r)) (log) p(( [N] dq(x)) S p.lllll p- (a-11)
A.7 PROOF OF THEOREM P
The lower bound directly follows from (5) together with Qx P- 1lll and the fact that q is invariant to rescaling, so the best we can do is fit as much memories as we can until reaching 3(d+1) = p.P leading to scaling in Jr p(x) dr C. C-at1 ((a + 1).
A.8 PROOF OF PROPOSITION B
In order to get scaling with both finite data and finite memory simultaneously, we used a simple strategy:
With high probability 1 - 1+1// for some constant c, is similar q-
When ÿ is similar to q. the scaling with derived from Theorem M is left unchanged by substituting g by q.
Rather than using a uniform concentration inequality on the full q. we will proceed individually on each q(x). Denoting by Dr the random dataset of T data, for any sequence of set (Ex) (E))2E[ -typically will choose Eg- (a(x) g(x)/2).
Bue, D.r (E(fa]= = 1101->ne p(x)F u. pe) (f(x) + fa(1))
((x))Pa..r7 ( & E.)) + P # fa(x) q E Ex).
The second term has been worked out before, using that Qxx S ||g||2
I inf 2d 7 + (dg(x)2 < 16c, 81/11
Pu, Pu,r)(((()) # fa(x) Ex)
) }1g E Ex).
d
< inf 7 Pr (dg(z)² < <,16111 q E E.z)
?
where e, = 16c, 336=,,,000 + Sc,
Without thresholding.
Let us first start with the scheme (LI), with p > 0
1 -
((x) =
g(z) = p(x)..
[E[T]
Using a simplification of Chernoff bound for Bernoulli variables (see e.g., Hoeflding 1963), we get the probability bound (the randomness being due to the data),
g(x)
< p(z) ) I
(T((()) < 21/2
= Pr(((()
2
20

---

Published as a conference paper at ICLR 2024
As consequence, reusing the proof of Proposition M. when follows Zipf law (I),
[[E([a)] = Σ +
+f(())
$ >E p(x) -()o)))p(-T)(()18) + Σ + f.(x) a(x)/21/0)
ST E + {.(x)
We are left with the computation of the second term, denote cp 1/0, we have
R Pr.(P)((() # fa) ((x) cpa(x)) & inf 2d 7 (r(((())) = < (2))/2).
?
By definition of G, together with Jensen's inequality when 5 1/2
Σ =EIN]
(a(x)))1/20 z 1112)
hence l1q N1-2p, When o > 1/2, the worst value of ||4|l is when all the mass is concentrated on one q(x'), in which case ||q|| < 1. With the corresponding (NN), we get
PT((p) + f. (x) l((x) > cp((x)) S inf 2d ? + (NC
Finally, reusing the proof of Proposition A. and hiding logarithmic factors,
B[E(fa) = Σ p(=))(f(=) fz)))
$ T 1+1/a + inf 2d p(tz de>g(z)² S (()(()))
S TT-1+1/0 ((N)
(k-1)/22
The case p 0, can be easily treated by considering an error if and only if the number of elements [ 1 E [T]H is smaller than
With thresholding.
Let us now consider the thresholding scheme (12), with P E and p 2
((=))p(((=)))))))) ((tpp
((x: ((=))
<(x)) = p(z)
We basically proceed with the same technique but with the event Es the probability that belongs to the top P of the empirical frequencies. When dealing with a binomial distribution, one can enumerate all possible outcomes for the empirical frequencies. For a template E AWN-. we said that sequence (x:)) is of type if its empirical frequency is equal to a,
T(a) ((x)) e [N]]]
Aa [N),
Σ 1 -a{x)) [E[T]
Some enumeration arguments that can be found in Cover & Thomas (1991. Chapter 11) leads to
Pd, ((x,) E T(a)) = [Taa) exp(-T((H() (xu(allp)) I exp( (al(l))))))
Hence, the probability that does not belong the top P of the empirical frequencies of (Ix) is bounded by
(D topp(xx) T(a)) < Σ Eesp(TT-Dxu(llpp))
aeA
where the set of all templates where æ is not in the top of (a(="))=<(N)) With T' samples over elements there are at most (N + 1)" different type templates, hence
exp(ca T) (T 1) sup exp(ca T) sup exp(c. T NloggT (T + 1)).
aEA
aEA
As consequence,
PDr (x & topp(x:) E T(a)) S sup exp(-I - DKL(alpp)) + N log(T + 1)) aeA
21

---

Published as a conference paper at ICLR 2024
Now, it is actually to possible to remove the N log(T 1) in the exponential and extends this type of result to generic Polish spaces (see, c.g. Dinwoodie (992)..
PDr (x & topp(r.) E T())) & sup exp(-I- DKL(allp))
acA
We are left with the computation of the "information projection distance" between p and the set of distribution where does not belong to the top P. In order to get out of the top P of p one should switch p(z) with pPP) which leads to (without caring for exact constants)
DR. (p' (I)) = p(a) log(p(z)/p(P))) + p(P) log(p(P)//(x))) (p(x) p(P)) p()))ooo((()))
When considering P/2 and following a Zipf law we get
DKL(P' (1n) 2 (p(x) p(2x)) (og(p((/2))/((((c 2 Ca "a - - cap(x)
where ca ca 2 2a log(2), As a consequence, for any P E PEN.
Ed, ()(a)] S co P-
p(z)P(f(z) + fz(x)). PE/P/2
Σ
S cpp-at1 at +
Σ
(ex) T (d)(())) +(f(()) + fx(x) topp((s:)))
xE|P/2
I Scp atl + exp(p2271
2 p
Σ p(x) * fx)) x topp((>)))- xE|
When p 0, setting p min(e; T 1/a log(T)) with e₁ chosen so that all x stored in memory lead to f.) f(x) gives to the right scaling with both T and d: up to logarithmic factors,
[[E(f+)] & d= + T-11///
+ exp(-c3 log(
Because a >> the last term decreases faster than any polynomial power of T, hence ends up being negligible in front of -1+11/a
For the case E (0, 1] one can dissociate two events: the event where x belongs to the top P/2 empirical frequencies; the event where p(x) > p(x)/2; and conclude with similar derivations as precedently
B[E(fa)] < coP-atl + ex 1+11/
Σ p(x) # fa(x) | (opp((x))) > p(x)/2). xE(P/2]
Retaking previous arguments leads to the same scalings as the ones of Proposition B with respect to and scaling in 1-1+//a with respect T. This ends the proof of the mixed scaling with both finite data and finite memory capacity.
A.9 LEARNING THE INPUTS EMBEDDINGS
In instances where the embeddings are learned within the linear model (2), one may optimize them by merging all input token embeddings that are associated with the same output, which what we actually observed in practice in Figure N. Proposition E captures the resulting theoretical perfor- mance.
Proposition 5 (Improvement for learned inputs embeddings). Let the input embeddings be set to ea = (z)(z)) Assume without restrictions that p(y) is decreasing with Consider the unembeddings where (My)ye(r) are n-quasi orthogonal, and uyy = 0 if is not among the P-th most frequent classes. Let E RVV, and set E Re as q(y)= (x) 90(x), then
E(fw.,) ) S p(tx 1 (=)) < 27/ + 2 11l1113)
(49)
In particular, it is possible to consider a thresholding associative scheme qa such that, if y follows a Zipf law plu) Cay- E(Jw... O(d/ (o)
22

---

Published as a conference paper at ICLR 2024
Proposition B shows that when learning the input embeddings in d that depends on the law of æ, by a scaling that depends on the law of y. It illustrates the usefulness to learn embeddings when the law of z is well factored by the law of y. This is typically the case when x are news articles associated with a few topics y.
expect to replace the scaling
Proof. When e can be optimized, is natural set ez to be a constant for all x that are associated with the same output. Let E [ be an associative scheme,
fw (wra # (w6) max (mo) Eee(s) ( sx (ov=)))_)))))00
In order to lower the probability, wants minimize the left expression, which leads to the will to maximize exo ex f. (x) (f.(z)) This can be done by setting
Va, [M], efe = (1,())
(50)
Such an isometry can be built by setting 6z = (x)) leading to the new characterization
fwap (=oo)) maxxttv Eec(se) (o())a M.nw (u, Mxg
where yo (20) and
((v) = E go(x).
(51)
an 2:f-(x)=y
When 17-quasi orthogonal for its first P values and set to otherwise, have
Σ (())u u (u, - (yy))= = q(v0) "y - 1) + ((u)(" "y - (u ),,))
*E|M]
Σ
( uz uu uy u Ugo)
x [
I <9((8) Ia(10)l + lg(v)ln
Σ
la(2)ln(1 + n)
*
S <-9(10) 2n sup |a(2)1 + 21)
Σ |x(2)
=#00
zE|M
- 2 + 2n7 12/211
As consequence, we get
E(fwa)) - Σ plx))1((f(())))241lll. +22l1l
xE[N]
Using that, for any A: [M] - Rdd (.EE(())AACC (x)) = E(, = E(y
E(fw,,) ) S Σ L( +2122=
(52)
ME[M]
Note that when the embeddings u are chosen uniformly at random on the sphere, and > log( similar bound will hold up to extra higher-order term as seen in the proof of Theorem II
When u is defined to be zero on [M] \ [P]. and only 7-quasi orthogonal for (Uy)yE[P), the same characterization holds with
E(Sw..))
p(v) +
Σ p(u) plu) (00)s) sE[P!
(53)
vE[M[N[
Finally, 7 is set to 1/4P, and = [yj( we get the upper bound
E(Jw,)) ) S p( (f))))) P)).
The best P that one can consider that such d/4P r2=d log(P). Setting P (o16)o)(d). and bounding E < Jp 1-8 dt ends the proof.
17
23

---

Published as a conference paper at ICLR 2024
A.9.1 DISCUSSION ON COMPENSATION MECHANISMS
When optimizing the embeddings, one may turn the negative interference mechanisms illustrated in Figure into positive ones.
Assume that Ex (z)(x), model (HI) become, denoting (f.(z) ug for simplicity,
f(x)
arg max " W Wuo:
W
Σ
uug.
(54)
sE[M]
y'E[M]]
Similarly as before an error is made when
max VEIM
Σ M'EIM
(lo') uo) Uy T > 0.
(55)
When the output embeddings are learned,
optimize them induce compensation mecha- nisms. For example, when M = 3, and y1 is competing when fx) as the argmax of (54)) due large storage of ((v1) compared (Bo). one could benefit of (322 ensure that
(u) uo 1) 9(31) (1 u, uo)u un + (12) uo) uguu ug < 0 < ( uo 1) + (1 - u un uu ug.
In this situation, the score u6 We, e, of yo would be higher then " We, ensuring that we not make an error when predicting f(x (544).
We refer the interested reader to Elhage al (222)) for related investigation.
A. LOsS GRADIENT
The cross-entropy loss is written as
exp(u, Wex) [((r, y, W) - log( ExeM exp(u.. Wex) )= -1 - We, + log( Σ exp(u
(xe()]Wex))) Wez)).
sE[M]
Hence stochastic gradient descent will update the matrix W by adding terms of the form Exe(an exp(u W We.)u.e
Owe((z, v), W) = -uee T. +
EverM exp(u, Wew)
= u- Σ Pw Pw (2)
=E[M]
-(1 - pw (u e T + Σ Pw (2
Pw (2|z))
(1 pw (ylx
Σ 2
1 Pw Pw (y|æ)
Note that Pw (w(22x)) (2/x)//11 (1 (y)=)) corresponds the the probability of the z conditioned with respect under the event that is not y, formally
pw (2(a) = p(2222 v)-
Pw (y/x)
Finally,
Owl((2,8),,)) W) = (1 pw (w(y))))(une) (w v)use
#V
-1 Pw (y)=))(uge - Ez-pw lugl,, y| y)
While, it is clear that the model ( does not describe the solution found by cross entropy, might hope that the term E|u will somewhat cancel themselves out and be an order of magnitude smaller than the leading term u
24

---

Published as a conference paper at ICLR 2024
A.11 APPROXIMATE UPDATES
The formula (201) is justified by the fact that a matrix W. - Wg. will lead to an update (E) at time I according to the rule (I9), assuming exp(u, We,) ex) 1 for any (x)),
at+1 (2:) ge (x) = (1 - pw,, (z)) = 1+ (M - 1) 1
exp()(x)))
together with the fact that will be seen Tp(x) times average in T samples.
Similarly, very large batch size = [B] and T/b update steps, each x will appear in each batch about bp(x) times, which leads the rough approximation
pbp(a) 1 M- exp(x)) (56)
= f = foofo......... ((),
where f:I+a+
T/b times
In practice, we can approximate the effect of a batch by counting how many times x was in this batch and setting bp(x) be the exact count, which will lead to tighter approximation. This is this approximation that we plot on Figure IK
A.12 GRADIENT FOR LAYER NORM
Let [N], y E [M] and W e Rdxd When processing the input æ, layer norm adds normaliza-
tion layer
W
f: W
Well
Using the chain rule, with D denoting the Jacobian operator,
Vwl(r, y: f(W)) = (Dw f(W)) V s(w)((E, y: f(W)) = (Dw f(W)) (WW))TWw Tw W). We are left with the computation of the Jacobian. We proceed with chain rule
S(W) = ((2((a((W)) W, A:LER-C .2: + [Te|l, Sa W E Raxd + We..
Dw D(W) Vw(hoft f2 fs)(W) + (ffa((f(W))) I
-w (W)
I= fss(W) fs(W)) WT IIl
WT
Weel I
fe(f
Wezll
-e W
= ||
|We -1 ||Wex|l
(1 We
This proves the formula written in the main text.
B EXPERIMENTAL DETAILS
B. MAXIMAL PARAMETERS UPDATES
In order carefully choose step-sizes that scale well with width d in optimization algorithms, we fol- low Yang et al. (202l) and consider learning rates consistent with maximal feature learning updates. Here we consider the following initializations:
Wi is initialized as a Gaussian random matrix with (O 1) entries.
Input embeddings e, and output embeddings uy are initialized as either random on the unit- sphere dimensions, with Gaussian (O, 1) entries. In both cases, every embedding has norm R 1.
Updates to W. The updates to the matrix W look as follows:
SGD with step-size Iw:
W' W + nwoW, 8W
Σ aj anuuy Gays
25

---

Published as a conference paper at ICLR 2024
with aj Oa()), and dimension
dimension-indeeendeen independent number of elements in the sum. Choos- ing 7j = ((1) then ensures that for any input embedding 6z, we have e = ((1) desired.
Adam (idealized here as signSGD) with step-size vy:
8W
W W + '/w sign(ōW), sign(öW)s = [l
The coordinates of sign (W) are now ((1) instead of ((1/d), thus the step-size needs to be taken Ijw (1/d)) in order satisfy ||"e (()) (see (Yang et al, 2021: Yang & Littwin, 2023) for more details)
Updates to embeddings.
The updates to embeddings look as follows:
SGD updates:
u/y = My + nuduy
Juy = Σ aayWesy
c = + noder, Se, = Ea' Bgy
with a, = ((1) and a dimension-indeeendeee number of js. Since the algorithm en- sures [We,, (1) and || Uy = ((1) throughout training, choosing The Tle (1) ensures that these conditions continue to hold after each update. Adam/signSGD updates:
(Suy)i
uny = uy + the sign(ōuy),
(sign(6uy)):
((Ouy)sl'
(ōez):
e = ex + ne sign (deg), (sign(dez)),
((ōex)ll )al
Since the updates have coordinates of order ((1), in order to ensure that embeddings re- main of norm (1) after each update, we thus need 7a, Te = (11/Vd).
B.2 ADDITIONAT. FIGURES
Our theory predicted optimal scaling laws d-1-a However, there are some catches behind the proof:
The lower bound is true when d < N = 100, otherwise the error can actually reach zero when d becomes larger than tipping number de which compares N. This fact was illustrated on Figure E Increasing N augments the tipping point de, rectifying the learning curve as illustrated
Figure
This was proven for models where q(z, y) (zx), and where q(z) is not optimized with respect to fz).. As such, is not clear if those lower bounds hold for optimization- algorithms, although we argue that we do not expect different mechanisms to take place in the proofs. We illustrate this empirically in the left of Figure 121
Similarly, the unreasonable effect of learning the embeddings would be highly disappointing if those were hard to optimize in practice. The right of Figure I illustrates how with a few steps, one can achieve zero generalization error when learning the embeddings.
In order to better understand gradient updates, Figure 14 shows the dynamic of the association memory W updated with SGD and a large step size. To validate the approximation (20), Figure plots the generalization error associated with SGD and its theoretical approximation, while Figure illustrates the idealized association scheme 4, associated with step size y, batch size one and Zipf law on x E [N].
In order to understand the effect of Adam, compare it with plain SGD and SGD with rescaled variance on population data. That is, we consider gradient descent with V CWC(W) (IC). The rescale variance SGD, consists in dividing the gradient by the variance of V V (x();; W) (l7) when
26

---

Published as a conference paper at ICLR 2024
100
10
Error
=
10 10-2
1544/8
10
10a
10a
Figure 9: Same figure right one of Figure B yet with a bigger here 1000. The dashed curves represent & = .35 dt 1/4 (orange) and & = 3.5 d' (green). They validate the scaling predicted by theory where used N +000 get tight polynomial scalings of ((S with respect to d.
Figure 10: Representation of the weight matrix (u,, We y. MM for N 10, M 5 (x)
mod. M. The data r follows Zipf-law with a = and 7 = 10ah The matrix W is obtained according to (E) together with the scheme ([I) Left: p (II), d = 10, there is not enough memory capacity, and memories, leading
large generalization (I)), = 50, there is enough memory capacity, we learn right association y = x mod. Middle right: (II)) d = 10, the weighting allows to store the most important memories beside having small memory capacity. Right: p (L.) d = 50, the weighting too strong which does not allow to memory associated with rare association (bottom of the matrix).
the model does succeed
Middle left: p
X ~ p (I)) For simplicity, we consider Adam with 31 = B2 = 0. in which case, equates sign SGD, i.c. SGD when considering the sign of each entries of V C(W) in the updates W. + W++1- Figures TS and IG underpins our intuition that the usefulness of Adam lies in its ability to rescale gradient updates, an effect that could equally be obtained by tuning the learning rate.
27

---

Published as a conference paper at ICLR 2024
100
-6
T
10
7 = StepLR
Error
Error
4
10-1
10
.1
.2
StepI.R
101
10a
10a
T
Figure 11: Learning curve of the generalization enror (N) with respect to the number of data processed by stochastic gradient descent in the setting Figure E Left: comparison size allows store more memory the risk overwriting past association, which explains the higher variance of the blue curve but its overall better performance. A small step size will avoid loss spikes due to memory overwriting, but will take more time to store rare associations, leading to worse performance. By decreasing the learning rates along training, e.g.. with the "StepLR" scheduler (Paszke etal. (DTS), can get the best both world, i.c., store memories fast the heginning of training when storage capacity underused, while heing more cautious the end training when there is no more "free" memory space. Right: Similar plot with N = 30 averaged over one hundred runs.
single run. A big step
1100, |B|=16, T=10240
10"2
Eroe
Brrot
Leuming
S0D 100
SGD
10
10a
10
10a
10
10P
10a
Figure 12: Scalings with respect to d for optimization-based algorithms, in the setting of Figure B. Left: optimization- based algorithms beat the best algorithm designed by hands with g(z, the
2(x). Note how d=++1 (the left part of the figure show similar
have the same optimal exponent slopes curves) yet smaller constant in front, leading to earlier typing point before reaching zero generalization emror due to full storage of all the associations. Middle: Comparison of learning the sole matrix W (blue), or learning the embeddings e and (orange), together with the possibility to use non-linear model uy ReLU(e..) with and learned (green). All curves obtained after 10 updates with batch size 10 Right: Comparison with same setting Figure Learning embeddings or going non-linear allows to impressively optimize memory storage, leading to better exponent with respect to and earlier tipping point for fixed number of updates.
10
x 10
x 10
6
mel aррооx
Broo
x 10
10-1
roal
approx
Tamor
10-1
real
2x
6 10
Appprox
50
100
50
100
30
100
Epochs
Epochs
Epochs
Figure Same as Figure @ yet with batch size equals one thousands [B] = 103.
28

---

Published as a conference paper at ICLR 2024
Figure 14: Gradient descent dynamics similar Figure i with d 10 and fixed step size y 10. From time to time, we represent here E {0, 4, 6, 8,9, 11,30, 30, 32, 37, 49, 62, 75, .0), stochastic gradient descent will hit an association that is not properly stored in memory yet (the red hoxes). It will consequently update the weight matrix W, -> W.i1 (side by side pairs) to store it. When d is big enough, here d = 10, W will end by storing correctly all associations, leading perfect generalization for future examples.
29

---

Published as a conference paper at ICLR 2024
Figure 15: Comparison between SGD, signSGD and SGD with normalized variance population gra- dient seen from the association matrix W. at different times in the setting of Figure I The different correspond to the matrices W. time & (1,2,3,7, 100). Left: Plain SGD. Middle; Adam with 81 = A3 = 0, i.c., SignSGD. Right: SGD with normalized variance.
10
Gradient variance
SCD
Emor
sign SGD SGD
10-3
10-2
50
100
50
100
Epoch
Epoch
Generalization error in the setting of Figure Observe how SGD with rescaled variance (in green), an effect that can be done with SGD after adapting the learning rate, actually performs better than sign SGD (i.c., Adam with 31 82 0). Right: Variance of SGD along the training. As the training goes, SGD is losing momentum due to smaller gradient variances, henee smaller updates.
Figure 16:
Left:
30


Reviews:
[{"review_id": "Oc77EHhYtM", "reviewer": "Reviewer_cFvv", "summary": "This paper performs a study of the scaling laws from the perspective of associative memory, studying the phenomena by formalizing a highly controlled experimental setting. They test this phenomena across the amount of training data and the embedding (memory) dimension of that data and observe scaling law trends that resemble those of LLMs, indicating that the conclusions drawn in this paper will likely extrapolate beyond the scope of small associative memories.", "questions": "1. Sec 2 par 1: \n\n> \"For example, $N$ could be the number of potential sequences of fixed word length in the English language, while $M$ would be all the potential words to complete the sequence\" \n\nUnfortunately, there is no modern model that actually treats all possible sequences of a fixed word length as a single token. But a recently proposed method derives the Transformer as an Associative Memory (see [Energy Transformer](https://arxiv.org/abs/2302.07253)). Could you explain how the experimental setup could be adapted to more advanced associative memory structures that contain multiple weight matrices and allow token-token interaction?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "## Admirably formalizes the scaling laws in Transformers as a memorization/memory retrieval task in Associative Memories\n\n- (+ +) The paper clearly and thoroughly defines a \"sandbox\" problem setting where we can study scaling laws (of discrete data domains, like the vocabulary tokens in NLP) using principles of Associative Memory\n- (+) The paper includes experiments using the Associative Memory sandbox to draw conclusions about good optimizers, learning rates, and batch sizes in larger models.\n- (+) Supplementary includes complete and well-organized code for all experiments in the paper.\n\n**Originality**: I am not aware of existing works studying the scaling laws from the perspective of Associative Memory.\n\n**Quality**: The paper is of high quality, though I did not read the Appendix.\n\n**Clarity**: The paper is very clear, though accessibility to the average reader could be improved.\n\n**Significance**: Medium-High -- this is another work drawing formal connections between foundation models and associative memory, providing theoretical structure to a field designed primarily by empirical results.", "weaknesses": "## Experiments not able to scale to large models\n\n1. (-) It took several readings to understand the experimental setup. The clarity of the paper would be improved with a small architectural diagram describing the setting.\n2. (-) To my understanding, the proposed method can only study Transformer blocks individually, not the entire Transformer as a whole (This is my understanding of Sec 4 paragraph 1: \"our model is a proxy for the inner layers of a transformer\")\n3. (-) Like 2., the proposed method does not allow words in an input sequence the ability to talk to each other, which is how the attention mechanism in Transformers actually works (see Question 1). Thus, the sandbox is a very limited tool to study larger language models.", "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": "dArbaocGYT", "reviewer": "Reviewer_qN1v", "summary": "The authors explore the behavior of a simple model for associative memory as a weighted sum of outer products of query and key vectors for tokens (matrix). In particular they provide bounds for its generalization error as the number of tokens and the encoding vector size varies and for different choices of the weights in the sum (memory scheme). \n\nMore specifically they provide scaling laws in the case memory or data is infinite (respectively for finite data or memory) and memory performance characteristics for weights that are constant or seen-data specific (frequencies). They also study optimization based learning of memorization and how training choices and hyperparameters as in transformers affect its characteristics.", "questions": "- For ranges of values for T and d  for data distributions that could map to / feed actual transformers what would be the recommended memory scheme to try in order to minimize generalization error? (This could be a high level and practical direction to the reader who seeks a brief takeaway message).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The presentation is excellently organized, the notations, definitions and associated propositions and theorems are carefully stated and accompanied by clean supporting simulation plots, the cases explored make up a comprehensive and complete narrative for this interesting theoretical work.", "weaknesses": "- The current setup is synthetic/artificial: it is a drastic simplification of configurations found in practice, e.g. for real transformers. Although there are clear notes in the text for the potential deviations of this simplified model to a real one, it remains to be seen how well analogies hold. To this end, perhaps crisper (albeit riskier) predictions of how some of these results would translate/map to tangible observations in a real transformer would help the reader better appreciate the implications of the theoretical results.", "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": "Db1Hewjxcl", "reviewer": "Reviewer_fGhC", "summary": "The authors provide scaling laws for the error of a specific model of associatice memory (it takes inputs x and predicts outputs y, which deterministically depends on x) in terms of the strategy of the construction of this models parameters, the number of input-output pairs seen, and the distribution of the input tokens.\nThey further experiment how optimized, rather than prescribed weights, relate to the error scaling recovered in the theory. They investigate how several specific architectural and optimization choices affect this error in practice.\n\n---- Update ----\nThanks to the authors clarifications during the rebuttal, my confusion got cleared up. I now understand the paper to be not only an interesting theoretical contribution, but also one that belongs in this context. During our discussion we converged on the points that lead to my misunderstanding, and the authors intend to improve some aspects in the camera-ready version. In the light of this, I improved my score and think that modulo the changes the authors promise, the paper should be accepted.", "questions": "Abstract\n- 'We derive precise scaling laws with respect to sample size and parameter size,' -> it seems there is a subejct missing \"We derive precise scalings laws of quantity XY with respect to ....\"\n\nSection 1\n- It would be nice to give an example of a 'behaviour' of models that can be accessed with scaling laws.\n- what is the criterion to qualify a scaling law as 'improved'?\n- what is exactly meant by a 'statistical rate' in the present context?\n- 'theoretical schemes' -> theoretical predictions?\n- 'based on specific' -> 'for specific'?\n\nSection 2\n- 'number of data *samples* '\n- 'The first/second ones' seems like a wrong english construction of mixing singular and plural.\nSection 3\n-  as is the case at initialization -> of a neural network/transformer?\nSection 4\n- m = 5. -> M = 5?\n\nFigures\nFig 5 batch one -? batch size one?\nFig 8 is it SignGD, Adam, or SGD in the plots?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The derived results look interesting in the context of the chosen model and construction of its parameters.\n- Looking at discrete data with a real-word-like distribution is a promising idea.\n- I think re-framed in the correct context, the result could add nice insights to scaling laws, even though in their current presentation they are more confusing than insightful.", "weaknesses": "- While the introduction and title suggest that the paper considers the memory capacity of associative memories, it seems that in fact it is investigating the error scaling laws of a specific learning problem, where a discrete input x determines an output y. The suspicion that this is learning, is corroborated by the  fact that giving more data for a fixed dimension (e.g. Fig.3 right) improves the error. If the model was truly memorizing, eventually there would be a cut-off and no new data could be taken up by the model for a fixed dimension d (as there is in Hopefield networks, the original 'associative memories'). Under the present title and introduction I would expect scaling laws of the memory capacity in terms of the input parameters, and this is not what the paper is giving. This is the main weakness of the paper; that the motivation, theory and experiments do not form a coherent line of arguments which improve understanding of associative memories and their memorization capacity.\n- I find it difficult to comment on the results of the paper in the light of this mismatch, for me, the stated goal to investigate \"[...] how different key elements in the training of a transformer influence storage in our memory model.\" which motivates the experimental section, is not answered at all. \n- I want to note that I would be happy to read a rebuttal about why the authors believe their theoretical and empirical analysis is connected to memory capacity as discussed in Figure 2 - it could be that I am missing a piece. Otherwise, I think the results, stated differently, could still be useful to the community, but this would require a complete revision of the paper's motivation and contextualization.", "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": "80BB2VsQYE", "reviewer": "Reviewer_fK8b", "summary": "The paper presents an extensive rigorous theoretical analysis of the associative memory capacity of simplified transformer layers. Associative memories are formalized as cross products of input and output tokens, which are assumed to be associated deterministically in the analysis. These associations are combined in the key matrix of an attention layer with hard argmax attention. The analysis is carried out assuming that the input tokens follow a Zipp law, which is commonly observed in naturalistic data. \n\nFinally, the paper contains a numerical analysis of SGD leaning in these layers and it provides some recommendations for Transformers training.", "questions": "- Is it possible to extend the analysis to probabilistic associations?\n- Is it possible to analyze the softmax model, or is the hard softmax assumption central to the tractability of the model?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "I must start by stating that I am not very familiar with the kind of proofs given in the paper and I did not have the time to study them in detail. Therefore, my judgment is conditional on the validity of the statements.\n\nI found the analysis to be useful as it offers a detailed theoretical view on an essential component of modern language models. While the paper makes several simplifications, I do think that the resulting model captures several of the main components of commonly used attention layers.\n\nAll in all, I do think that research of this kind is highly needed toi bridge the gap between our understanding of language models and our ability to use them. While this paper is just a small step in this direction, I do think that it is a much needed one. In particular, I highly appreciated the focus on token memorization as the phenomenon seems to be behind most of the capabilities of generative models.", "weaknesses": "- Some of the assumptions are rather strong and it is therefore unclear if the insights will generalize to more realistic scenarios. In particular, deterministic associations are rare in real data.\n\n- There is some evidence on the importance of lower weighted components of the attention blocks in the performance of Transformers, which are entirely ignored in the hard argmax model.\n\n- While I do think that the theoretical analysis is insightful, I am not sure that the result of the SGD experiments on the simplified model can cast much insight on actual Transformer training. In fact, the recommendation of small batches and larger step sizes seem to be the opposite of what is known to work in large architectures.", "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": "OX3YBzduQI", "reviewer": "Reviewer_vzgF", "summary": "The paper studies scaling laws in a simple, linear associative memory model. The model is aimed at capturing trends in LLMs, which use similar mechanisms. The authors derive scaling laws for the error of the model under varying amounts of data and memory capacity. The results reveal an optimal method of storing information in the model to minimize error. Next, the authors demonstrate that memorizing via gradient updates can be modeled in the framework derived earlier, and show trends of error with respect to learning rate and batch size. The authors finally discuss some additional considerations related to optimization, layer normalization and learned embeddings.", "questions": "What are the key differences between this work and related work? What are the scaling results for similar memory models that have been previously proposed?\n\nWhat is the practical significance of having discrete input values? How does this affect how one may consider generalization to unseen inputs?\n\nWhat key gaps remain between the model considered and practical LLMs?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "**Originality**\nThe approach taken by this paper is novel, and contrasts with perspectives in prior papers that focus on continuous inputs. What is particularly notable is that the paper considers both errors arising from finite data size and finite capacity and derives relatively explicit expressions for each of these.\n\n**Quality**\nAlthough the proposed model is relatively simple, the theoretical analysis in the paper is extensive as evidenced by Appendix A. The theoretical statements are also backed up by experiments where applicable.\n\n**Clarity**\nThe paper is mostly adequately well-written. Figures are well illustrated. The notation is well chosen.\n\n**Significance**\nThe paper appears relatively significant to the field of associative memory. Although it considers a relatively simple memory model, the analysis is quite extensive and could be used in future work. Moreover, connections are drawn to practical LLMs which greatly improves the paper's relevance.", "weaknesses": "In my view, the main weaknesses of the paper are related to its clarity. The paper is quite dense with theoretical results, which is good in that the authors provide many contributions. On the other hand, it makes it difficult to interpret and contextualize the results. I would suggest that the authors use more space discussing their results and interpretation, and move some theoretical results to the supplement. One possibility to consider might be adding an extended discussion subsection at the end of each of sections 3 and 4.\n\nAnother point of weakness is the description of related work; it would be ideal to significantly expand this section, particularly with respect to the theory on associative memory models. It may be helpful to highlight key results in the associative memory and scaling laws in the related work section (e.g. results on the capacity of other associative memory models, scaling laws for LLMs). This is important to establish the significance of the results in this paper relative to prior work.\n\nOne key assumption in the paper is that inputs take discrete values, and that unseen input values lead to errors. It would be helpful to further discuss the realism of the assumption. In particular, when inputs are continuous-valued, we may expect generalization to unseen input values that are similar to previously seen values. When is it (or is not) reasonable to expect this kind of generalization in the discrete setting?\n\nFinally, it would be worth discussing in further detail what the key gaps remain from using the theory developed in this paper to explain scaling in actual, practical LLMs (e.g. what remaining architectural features of LLMs prevent the theory from applying to them).\n\n**Minor Comments**\n\nThe placement of figures is sometimes far from where they are referenced in the text\n\nIt is unclear what the error margins in figures 3 and 4 represent\n\nThe trends in Figure 7 are difficult to interpret due to the variation- it would be ideal to plot an average of many trials\n\nAdding some additional models to Table 1 could be helpful; it might not be worth having a table here if there are only two rows\n\nThe log scaling symbol in equation 9 is not formally defined in the main text", "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)