Scaling Laws for Associative Memories

Thank you for your positive review, we appreciate the time you took to help us get the most out of our work.

Your questions relate to the weaknesses you mentioned, and that we will comment now.

• "In my view, [...] sections 3 and 4": Thank you for your suggestions to improve the clarity of our work. We will do our best to better contextualize the different results with some extended discussion at the end of Sections 3 and 4.

• "Another point [...] prior work": Regarding related work, Hutter derived Proposition 1, and our Theorem 1 is to put in perspective with traditional Hopfield networks that can store $d / \log(d)$ spin patterns, although our setup is slightly different (working with random embeddings). Up to our knowledge, no paper has provided a clear understanding of the sample complexity for associative memory, at least not in the realm of LLMs.

• "One key assumption [...] discrete setting?": The memorization phenomenon described in the paper is intrinsically of discrete nature. It is true that in LLMs, tokens are not completely disjoint unity. E.g., if your tokenizer tokenizes "happy" and "happily" as two tokens, their embeddings are probably going to be close, and learning that "happy" is associated with positive sentiments will transfer to "happily". This is notably why we were keen to add a paragraph on learning the embeddings (see also Appendix A.8). We believe that those phenomena will somehow decrease the "effective number" of tokens, but that the described memorization mechanisms will stay the same (similarly to how the study of sample complexity for linear regression can be extended to non-parametric estimation through the notion of effective dimension in Hilbert spaces).

• "Finally, [...] applying to them)": We see our work as a module to be reused to comprehend more finely more complex architecture. There are three components that seem important to us to generalize this work.

1. Hierarchy: what happens when we put many associative memories one after the other (e.g., the first ones may memorize low-level semantics, while the last ones may perform more abstract pattern matching)?

2. Superposition: what happens when those modules are put in parallel (e.g., multi-head) and extract patterns that can be reused differently later by subsequent layers?

3. Entanglement: how to disentangle transformer representation to be able to track pattern matching mechanism / associative memory? Can we easily erase facts from LLMs memory, like someone's birthdate, or the knowledge of Harry Potter, based on manual interventions of the weights?

Thank you for your other comments:

• We will try to remedy the situation regarding the placement of figures.
• In Figure 3, the randomness comes from the random embeddings, in Figure 4, it also comes from the random data. We will clarify it.
• The right of Figure 7 is averaged over several runs in order to showcase the trends.
• We were hesitant for Table 1, some of us believed that it makes sense to keep it really simple as a minimal form of take-home message.
• Thanks for pointing out the undefined symbol. We will define it.

Thank you very much for the great suggestions, and the many comments that will help us improve our work.

Thank you for your appreciation of our work, and your great questions to extend our first steps toward a better understanding of modern language models.

Let us first address your comments on the weaknesses.

• "Some of the assumptions [...] in real data"" We agree that deterministic associations are rare in real data. Happily, the proof technique behind our main theorem, Theorem 1, still applies to noisy associations, although the derivations would be more lengthy.

• "There is some evidence [...] model": There are definitely second-order effects that we did not succeed to quantify in this paper. Typically, lower weighted components could introduce "positive interferences" that improve the performance of the model, but our analysis only considers "worse-cases" when dealing with interferences between memories (see, e.g., paragraph "upper bound" on page 16, or Appendix A.8.1).

• "While I do think [...] large architectures." We understand your concern about take-away generality. We hope that thinking in terms of memorization capacity could help design optimal optimization parameters for more complex models, although we agree that the conclusions are likely to vary from our single-layer model. Nonetheless, note that our conclusions regarding small batches and large learning rate are about SGD, while Adam (the method of choice in practice) does work well with large batches as commonly used in practice. The non-linear nature of the Adam updates makes it more difficult to study precisely, but we hypothesize that its benefits for large batches are due to its pre-conditioning, clipping, and normalization effects, as discussed in Section 4.1.

Regarding your questions:

• "Is it [...] probabilistic associations?" Yes, it is possible, although it would require lengthier derivations that are beyond the scope of this paper.

• "Is it [...] the model?" Yes, it is possible to analyze the softmax layer. Indeed, in our simple model replacing the hard max by a softmax convexifies the problem, and one can use tools from convex analysis to do so. This is a choice we initially considered, although we decided to focus on scaling laws with respect to the 0-1 classification loss, which implies the use of a hard max.

Thank you very much for your review. We are happy that you find our results insightful and interesting. We are sorry to hear that you did not relate with the framing of our results in terms of statistical learning bounds, instead of memory capacity results more common in the neural computation literature. We will do our best to clarify how these viewpoints relate to each other. Your comments will help us improve our work, and increase its reach.

Let us address the three weaknesses you are mentioning. We hope that our answer will help you better appreciate our work.

-"I want [...] contextualization": Figure 2 was actually useful for us to engineer the proof of Theorem 1 (see for example Eq. (30) page 14 in the supplementary materials). The mechanism regarding Hopfield networks you are mentioning is exactly the one we are deriving: when $d$ is smaller than the number of associations to memorize, there are interferences between memories, and you can not store all of them. Theorem 1 was useful to later predict the empirical observations in Section 4 (as we will detail in point 2).

-"While the introduction [...] capacity": It seems that there is some confusion regarding the right of Figure 3. The x-axis is the embedding dimension (i.e., the memory capacity), not the number of samples, it shows that as the memory capacity increases, the error decreases. The memory cut-off can be seen on the left and middle of Figure 3: as T, the number of data augments, the level lines of the generalization error do not change (the color do not change as we go up) after a certain threshold with respect to $T$ that is linear in $d$. This cut off can also be seen on the right of Figure 1. Note that in addition to the number of data $T$, the scaling laws are given as a function of $d$ which parameterized the sole parameter of the model, which is the $d$-by-$d$ matrix $W$ Eq. (2) (itself defined through $q$, Eq. (4)). Hence one simply has to take $d = \theta^{1/2}$ to cast our scaling into one with respect to the number $\theta\in\N$ of parameters of the model. We will make those points clearer in the next revision.

-"I find it difficult to comment [...], at all": In order to understand how the key elements in the training of a transformer influence storage in our memory model, we perform some gradient update approximations (which were controlled by experiments on Figure 4), and derive the memorization scheme $q$ that arises from different choice of optimization hyperparameters, plugging this into the theorems of Section 3 allows us to predict the behavior observed on Figure 8: e.g., large learning rates are better, small batch size is better; as well as Figure 7: scheduling is useful when reaching memory overflow. Those insights are what we meant by "illustrating how different key elements in the training of a transformer influence storage in our memory model".

Thank you for pointing many typos:

• The scaling laws are the scaling laws of the generalization error (we will write it explicitly).
• By "behavior" we mean the generalization error, and to a certain extent "emerging properties" such as mastering grammar, syntax, translation, reasoning... (lower generalization error in language modeling means better next tokens predictions, which implies the understanding of more phenomenon, and is associated with "emerging properties"). We will make it clearer.
• "Improved scaling law" means that the generalization error is decreasing faster to zero (which can be read in a bigger exponent in the power law of the generalization error with respect to the scaling parameters, i.e. the number of data, or the embedding dimension).
• *"Statistical rates" means the average (over inherited randomness) rates of convergence of the generalization error towards zero.
• Figure 5 represents Adam with $\beta_1 = \beta_2 = 0$, which is equivalent to SignSGD (https://openreview.net/forum?id=S1EwLkW0W).
• Thank you for spotting all the other typos!

We hope that our answer will help you better appreciate our work, and stay at your disposal to answer further concerns.

Thank you for your appreciation of our work. We are happy to hear that you found our study complete and thorough. We appreciate your will to validate our theory with concrete recommendations and predictions of tangible observations.

"The current setup [...]. To this end, [...] theoretical results.": At the time of writing, we cannot provide you one crisp prediction that is easy to test, but can provide some additional discussion on the matter. Here are some practical directions for the reader who wants to take-away our perspective to enlighten their experiments.

1. Associative memory could be understood with different granularity: e.g., a transformer is seen abstractly as a big associative memory machine, and tokens represent abstract "pieces of knowledge"; but layers can also be seen as associative memories working at finer levels, e.g., remembering attributes of people, such as birth dates.

2. Knowledge can be entangled. For example, to recognize the language of a document, it is probably enough to average all tokens of this document and learn a simple classification rule between this average (the input) and language (the class/output). However, some of our experiments suggest that some "interpretable behaviors" (e.g., "people birth dates") can be entangled/distributed across heads or layers. Being able to train while enforcing "layer disentanglement" would solve many issues in LLMs.

3. Once you have a layer working for a sub-task you do not need much to learn to solve new instances of this sub-task. E.g., if you have learned to recognize languages (the sub-task of recognizing languages is useful for next token prediction) through tokens averages (an attention head does the average, and the MLP layer do the association "tokens average"->language); it will be easy learn to recognize new languages with a few gradient update on the MLP to add the new association new average->new language.

Let us now turn those guidelines into a concrete prediction on the minimal curated data and the (curriculum) learning process needed to see emergent reasoning behavior. To ground the discussion, let us consider a "piece of knowledge" that corresponds to the ability to prove statements by induction. Roughly speaking we can assume that a type of reasoning should have been seen about 30 times for it to be learnable. This particularly true when the learner has already been tuned/wired to be able to do variables binding and recognize reasoning patterns, i.e., there is an (entangled) inner transformer layer that has been learned and acts as an associative memory of reasoning patterns (similarly to the language recognition layer described above). Hence, we could assume that with proper curriculum learning, 30 proofs by inductions (that are various enough to abstract the induction mechanism, and avoid prediction based on confounding factors) should be enough to learn to prove things by induction (this is one concrete prediction).

"For ranges of [...] takeaway message)": To answer your specific question, in general the best memory scheme is the thresholding one where one remembers all the "pieces of knowledge" $x$ that have been seen frequently enough (say more than 30 times to average out noise - in our noiseless setting 1 time is enough). When in the presence of limited memory, best is to memorize the most frequent ones. Concretely, when having $T$ data, we want to remember all $x$ such that the number of time $x$ has been seen is greater than 30, i.e., $T p(x) > 30$, which can be done by considering a memory capacity $d$ (which, in our paper, is the embedding dimension - but in practice could be defined as some effective dimension) such that $30 = T p(d) \simeq T d^{-1/\alpha}$ (under the Zipf law hypothesis). This retrieves the optimal scaling $d \simeq T^{1/\alpha}$ found by equating the two terms in Eq. (15). Although in this example, $T$, $d$ and $x$ corresponds to abstraction to be determined in a large corpus of texts processed by a transformer.

We appreciate your will to get the best out of our work, notably through the search of its practical implications, and hope that our response brings more answers than questions, and convince you of the valuable perspectives brought by associative memories to comprehend LLMs.

Thank you for your appreciation of our work. We appreciate your curiosity towards the next steps to be taken to perform a more comprehensive study.

• "It took several [...] setting": Thank you for your feedback, we will add a diagram to ease readability.

• "To my understanding, [...] transformer)" There are two ways to use our result. Either one sees a transformer as a big associative memory machine and derives LLMs scaling laws approximation from ours (this is what Hutter, 2021 or Maloney et al., 2022 do, as well as the energy transformer you point out). Either one sees our result as describing the behavior of some blocks (which is the future use we argue for). We will make this dichotomy clearer.

• "Like 2., [...] models." To see the relevance of our model for studying interactions between input tokens (as opposed to input-output associations), we refer to Bietti et al. 2023, where it is shown that gradient dynamics on key-query matrices in a transformer lead to outer-product associative memories of a similar form to the ones we study. In that work, the weighting of different outer products elements also seems related to token frequencies (see e.g., Figure 7 in their paper, which shows how different tokens are more or less prominent in key-query matrices). Studying precise scaling laws for such multi-layer transformers is beyond the scope of the present work, but we expect our work to be generalizable to that setting.

Thank you for the time you devoted to review and help us improve our work.