Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory

We thank the reviewer for the careful reading of the paper and constructive suggestions and feedback. We provide our response to the reviewer’s individual comments here, and strongly encourage the reviewer to check our global response for updates regarding associative memory formulation, which is also relevant to the reviewer’s comments. Please find our response to the questions below:

A slight nitpick, but on line 060 the authors state that the Hopfield network has only been recently generalized to continuous values. However, the continuous value Hopfield network was introduced only shortly after the 1982 paper. See Hopfield, J. "Neurons with graded response have collective computational properties like those of two-state neurons", 1984. The Dense Associative Memory (Modern Hopfield Network), introduced by Krotov and Hopfield 2016, was generalized to continuous values by Ramsauer et al. 2020, so perhaps this is what the authors intended? Reading on, this is discussed more in Section 2, subsection on Hopfield models, so my apologies for citing literature the authors are aware of! A slight rephrasing of line 060 may still be prudent.

We appreciate the suggestion. Indeed, we recognize that the original phrasing could be misleading, and it has been corrected in our revised manuscript (lines 60-61). Additionally, we have expanded the "Related Work" section to incorporate more recent literature.

The discussion on Hopfield models in Section 2 is excellent, but towards the end (e.g. works on the Energy Transformer and U-Hop networks) are seemingly only tangentially related to the main text. These could be left in for completion, but are they necessary for a reader's understanding of the main text? Perhaps with more explicit ties in later sections these would make more sense for a reader.

We appreciate the suggestion. In response, we have relocated certain related work to the appendices, thereby creating more space for the mathematical components of the paper.

Section 3 lays out the model architecture. Section 3.1 discusses the associative memory formulation and how sequences of tokens may be interpreted as patterns in associative memories. Since each prediction will have access to a different number of previous tokens (e.g. predicting token 5 will have access to 4 previous tokens, predicting token 6 will have access to the previous 5, and so on) does this mean the associative memory here must be capable of handling retrieval for patterns of different lengths? Or are all patterns to be padded to a uniform length, such as in a traditional transformer, in which was the associative memory would have a much lower proportion of initial contents for subsequences that are early in the task (e.g. predicting the 5th token compared to predicting the 500th)? Perhaps I am misunderstanding the associative memory structure or pattern structure here? The initial paragraph of Section 3 seems to imply the latter, that all patterns are equal to the maximum token length. Is this correct? If so, how does this affect the retrieval of patterns that have only a small number of tokens available?

We appreciate the insightful question from the reviewer. It is true that transformer-based models operate within a predefined maximum context length, and during pre-training, input texts are indeed segmented into sequences of this maximum token length. You are correct in your understanding that in Section 3, we intend to pad the sequences to a uniform length. In the context of associative memory, when only a small number of tokens are provided, the model updates by aligning the padded input with the nearest stored pattern. This ensures that the associative memory can effectively handle retrieval for patterns of varying lengths in the latent space. Based on your insights, we have refined our assumptions. We have expanded upon this discussion in the global response to provide a more detailed explanation of how our model handles patterns in the latent space.

We appreciate the reviewer's time and constructive feedback. Your insightful suggestions have been invaluable in enhancing our manuscript. We provide our response to the reviewer’s individual comments here, and strongly encourage the reviewer to check our global response for updates which is also relevant to the reviewer’s comments. Below, we provide a detailed response to each of the comments.

The authors might not be aware of this paper due to its title. But the proposed energy function (eq. 3) in the paper was recently introduced/used in Saha et al. (2023), with the exception that they use $\beta$ in their formulation and descriptions of the function's properties. However, I do not think this fact compromises the novelty of the work.

We appreciate the reviewer's recommendation. Indeed, we were previously unaware of this research. We have now expanded our literature review in the revision. Specifically, we have added a reference (Lines 137-140) to the independent introduction of the layer-wise energy function equation (3) by Saha et al. (2023), which was presented within the context of clustering. Your suggestion also inspired us to discuss the collective attraction mechanism in Remark 1 (Lines 186-189).

Much of the current analyses of storage capacity in Hopfield models relies heavily on uncorrelated patterns. Consequently, the claim of exponential capacity does not necessarily hold true for natural (or correlated) data. Similarly, a weakness of this work is assumption 3, since the work focuses on text data. The authors can instead make a claim that the Transformer blocks are storing the latent data, obtained from passing the text data to an embedding layer ---which should be much more uncorrelated--- to reinforce their narrative.

We sincerely appreciate the excellent suggestion. In response, we have revised our assumptions to incorporate the recommended claim. Please also refer to "Regarding the assumptions" in our global rebuttal.

The writing for the experimentation is unclear, especially to unfamiliar readers. For example, I think the authors should elaborate more on the findings of the mean and median nearest neighbor distances, and why such results corroborate eq. 12. Finally, I'm not sure what section 6.3 is trying to demonstrate. How was $\sigma$ chosen? To find $D^*$, you would have to train multiple models across a variety of data sizes $D^{(k)}$ for $k=1,\ldots,M$ sizes. Perhaps, there should be an elaboration on the selection of the size of different models and each chosen data sizes in relation to $N$ and $D=T_{max}n.$ The inclusion of these details should fortify the reasonings behind these experiments.

We appreciate your suggestions. We have elaborated our method in lines 477-479. In Section 6.3, our objective is to investigate the relationship between the number of transformer parameters $N$ and the data size $D$. To find $D^*$, we trained multiple models across a variety of data sizes $D^{(k)}.$ Due to space constraints, the data sizes $D^{(k)}$ and their corresponding results are presented in Table 3 in the appendix. Both the data sizes $D^{(k)}$ and the threshold $\sigma$ were empirically selected based on our preliminary experiments, as depicted in Appendix F.3, where a grid search was conducted over $D$ and $\sigma$.

Is it possible to collect the activation outputs up to 5% of the documents? If it's not possible due to limited compute, I understand. However, I think collecting activation outputs from 1% of the documents might be too limited.

Thank you for the suggestion. We have managed to updated our experimental result using 1.3% of the documents. Unfortunately, the scale of our experiments is constrained by our computational resources. Please also refer to our responses on empirical results in the global response.

I believe you forgot to include the citation for RMSNorm, see the paragraph which describes the hypersphere volume on page 8 of the manuscript and where you mentioned LayerNorm.

We appreciate the reviewer's observation. We have added the citation in our revision in line 407.

In the description for the cross entropy loss $L$ in Proposition 4, you should perhaps replace that equation with eq. 13 in appendix D, since there is a mention of $c$ which is not found in the equation in the main text.

Thank you for the suggestion. The variable $c$ in equation 13 is now mentioned in equation 10 (originally equation 9) in section 5.

There is an inconsistency in labeling the equations in the main text. Please organize them better since the mathematical descriptions provided in this paper are great. For example, I understand that $N$ and $D$ are mentioned in the prior sections of the manuscript. It is still certainly an important equation to be labeled/numbered.

Thank you for the suggestion. We have numbered these equations (line 152), and we have further provided a summary of notations in Appendix A.

Thank you to the authors for their detailed response. After reading over the revised manuscript, I am happy to say that I am considering raising my score, but there are a couple of more things that should be addressed before I do so. I understand this might be over the top but for the sake of improving the quality of the manuscript, it is essential for us (authors and reviewers) to do so.

My first qualm, that has not been addressed, is related to section 4 of the manuscript. Firstly, thank you to the authors for mentioning the work Saha et al. (2023) in the revised version. However, it is clear that eq. 4 is not a new energy function by any means, even though it does not include the inverse temperature as in Saha et al. (2023); also, I am not claiming that their work is the first to introduced this energy function. Instead, I believe the authors should rely heavily on eq. 7 as their novel energy function, which is a global energy function that captures the probability density of all transformer blocks. I believe all of the formulations, done in the work, which lead up to eq. 7 is fantastic and it should be the focus of section 4. Similarly, the same can be said with Millidge et al. (2022) as I did with Saha et al. (2023).

Secondly, I somewhat understand the experimentation, but what I do not get is the significance of the results after reading the manuscript. Specifically, why do these results tie up with the theoretical contributions and why are they significant to the claims made in the abstract. I do not think that the authors have made it very clear to me and possibly, to the other reviewers too. I suggest that the authors should have a separate discussion section where they discuss the significance of their results. Also, for context, after reading lines 477-479, they did not give me that sense anyhow.

Thirdly, in the Hopfield section of the related work section, the authors should use Millidge et al. (2022) and Saha et al. (2023) instead of Millidge et al. and Saha et al.. The same can be said in Remark 1. This is simply a stylistic comment.

We thank the reviewer for their positive endorsement of our work. We have improved our work based on the suggestions made by the reviewer. The answers to the specific queries of the reviewer are provided below:

The section on cross entropy loss is a little less clear than the rest of the work: Maybe you could add another ‘proposition’ or a statement that summarizes the results of the second part of the section?

We appreciate your suggestion. In response, we have introduced Proposition 5 (Lines 421-423) to summarize the findings from the second part of the section. This proposition essentially establishes that the optimal relationship between the model size $N$ and data size $D$ is given by $N \propto O(D^2)$ under idealized conditions where well-separated patterns are learned from the data and memorized by the model.

Scaling laws in the related works: N and D are never introduced.

We appreciate your attention to detail. We have revised the manuscript to include the definitions of N and D in the related works (Line 108).

“It has been shown that the feed-forward layers can be interpreted as persistent memory vectors, and the weights of the attention and feed-forward sub-layers can be merged into a new all-attention layer without compromising the model’s performance” This part is not well explained in the introduction in my opinion: could the authors dedicate to this a couple more sentences? Or, remove it and leave it as it is later in the paper?

Thank you for your valuable feedback regarding the clarity of our introduction. We have removed the mentioned sentence concerning the interpretation of feed-forward layers and the merging of weights into a new all-attention layer. We agree that introducing this complex concept early on might overshadow the primary focus of the introduction and could potentially confuse the reader.

The provided energy function is not ‘local’, as it is in other energy based models such as standard Hopfield networks, predictive coding, etc. This is quite interesting. However, I feel there is little discussion on the consequences of this. Why is this important? What does it tell us that is different from the previous literature?

We appreciate the reviewer's keen observation regarding the non-local nature of the energy function in our model, which indeed differs from traditional energy-based models such as standard Hopfield networks and predictive coding. Your comment has sparked a crucial discussion on the implications of this distinction. While the existing literature, such as Hierarchical Associative Memory (Krotov, 2021), employs a system of differential equations to design a global energy function that can encompass feedback connections, our proposed model tackles the global energy specific to feedforward architectures. Predictive coding networks, as referenced in (Tang et al., 2023) and (Li et al., 2023), incorporate recurrent connections; however, their dynamics are focused on minimizing the total squared prediction errors, which results in a limited connection to the attention mechanism. By framing the information retrieval properties of the Transformer as a sequence of associative memories, our model seeks to establish correlations between model size and the memorization of training data, all within the context of statistical physics. We have integrated this discussion into our revised manuscript in the new paragraph "Relationship to Prior Work" in Lines 829-844.

Using the euclidean distance inside an Hopfield networks: This has already been done, and it has shown good performance in simple image retrieval tasks [1]. Hence, interestingly your approach could be seen as using universal Hopfield networks to model the energy function of the transformer model.

We appreciate the reviewer's recommendation. We were previously unaware of the universal Hopfield network, which is a valuable addition to our literature review. Recognizing the connection, we can indeed consider our proposed energy function as a special case within the broader framework of universal Hopfield networks. To reflect this insight, we have expanded our "Related Work" section to incorporate this perspective (Lines 139-140) and modified "Our Contribution" (Lines 89-90) accordingly.

Thank you for your time and suggestions. We address each of your comments below.

• (W1) We sincerely appreciate the reviewer's observation regarding the inconsistency in our notation. We acknowledge that we adopted the MCHN definition from Ramsauer et al. (2020) without thoroughly aligning the notations. In our revised manuscript, we have included the definition of $M$ and replaced $x$ with the boldface $\mathbf{x}$ to clearly denote the vector form. The following theoretical results remain unaffected by these notational discrepancies.
• (W2) We apologize for any confusion caused by our previous notation. We would like to clarify that the matrices denoted by $Q, K$ and $V$ are indeed weight matrices. To avoid any ambiguity, we have revised the introduction of our attention mechanism section by replacing these notations with $W_Q, W_K$ and $W_V$ respectively. Again, these do not affect the following theoretical results.
• (W3) We appreciate the reviewer's question. For a comprehensive response addressing the assumptions, please refer to the global response.
• (W4) In Propositions 1, 2, and 3, we have established that both the proposed energy and the MCHN energy, previously demonstrated to be closely related to the attention mechanism, can be effectively approximated through an argmin search. Moreover, the MCHN update has been shown to converge in a single step. Given our focus on the collective behavior of attention layers, we consider the current approximation to be reasonable.
• (W5) Although our paper primarily focuses on the theoretical analysis of the transformer's behavior, rather than an empirical examination of scaling laws, we recognize the constraints of our current experimental setup. Regarding the theoretical results in [L449], please refer to our detailed responses on empirical results in the global response.
• (W6) Thank you for the suggestion. To improve the clarity, we have added a summary of notations in Appendix A.
• (Q1) We appreciate the reviewer's inquiry. Initially, we differentiated between the two sets in Assumption 2, which stated that held-out samples are stored in a manner analogous to the training set. Our revision now articulates Assumption 2 to emphasize the consistency of latent patterns. We invite the reviewer to review our responses regarding the assumptions in the global rebuttal.
• (Q2) The paper is structured into two parts. The earlier sections, including Section 4, are dedicated to presenting our model and introducing a novel energy function for the transformer-based networks using associative memory. Section 5, in particular, builds upon this foundation by applying our model to analyze the cross-entropy loss within the context of the pre-training. This analysis not only validates the theoretical framework but also elucidates its practical relevance to the transformer model.
• (Q3) Thank you for your question. It appears there may be a misunderstanding regarding the role of the softmax function in both the attention mechanism and the CE loss within the training process. To clarify, within the attention mechanism, the softmax function ensures that the model assigns appropriate importance to different elements within the input sequence. In the final layer of the transformer before calculating the CE loss, the softmax function converts the output logits from the transformer into a probability distribution, which represents the model's predictions for the next token in the sequence. The CE loss then measures the difference between these predicted probabilities and the actual target distribution, providing a measure of how well the model's predictions match the true labels. The attention softmax influences the model's ability to understand and process the input data, which in turn affects the output probabilities that are used to calculate the CE loss. In essence, the attention softmax indirectly influences the loss by shaping the model's predictions. We have made revisions to the manuscript to ensure that this relationship is more clearly articulated.
• Typos and Improvements: Thank you. We have corrected the identified typos and made revisions to the manuscript in accordance with your suggestions (I1) and (I2) as requested.
• (C1) Thank you. Our intuition is that as the model scales, the attention and FF layers, being stacked, constitute the majority of the model's parameters. However, taking your suggestions into consideration, we have further clarified that the fundamental operations of the Transformer are indeed the attention and FF layers.
• (C2) We appreciate the reviewer's feedback. As discussed in the paper (L406-408), the stability of the probabilities is attributed to the application of normalization operators which regulate the distribution of activations. In response to the specific point raised about our description at lines L419-420, our intention was to emphasize that our exploration is narrow in scope without considering other potential modifications.

Once again, we thank the reviewers for continuing to provide valuable suggestions. In response to the concerns raised by Reviewers Xt6m and zNBU regarding the significance of the experimental results in our theory-focused paper, we have made substantial revisions. Specifically, we have relocated all experimental sections to Appendix F, where we have also included additional descriptions of their significance. This change allows us to allocate more space in the main text for the mathematical derivations that are central to our study. In the main content, we now offer a concise summary of the experiments and their significance in Section 5.3. Consequently, we have removed all summary statements related to the experimental contributions that were previously present in the abstract, "our contribution", and the conclusion.

This revision emphasizes the theoretical findings of our research. However, we remain open to the reviewers' preferences. If they believe that retaining the empirical results within the main body of the paper is more appropriate, we are willing to revert the changes and incorporate those sections back into the main text.