Paper: Scaling Laws for Associative Memories (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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).

Reviewer comment (Evaluation):
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.

Author response:
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.

Reviewer comment (Theory):
Is it possible to extend the analysis to probabilistic associations?

Author response:
Yes, it is possible, although it would require lengthier derivations that are beyond the scope of this paper.

Reviewer comment (Theory):
Is it possible to analyze the softmax model, or is the hard softmax assumption central to the tractability of the model?

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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".

Reviewer comment (Theory):
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.

Author response:
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).

Reviewer comment (Writing):
In the abstract, the sentence 'We derive precise scaling laws with respect to sample size and parameter size,' seems to be missing a subject, for example: "We derive precise scaling laws of quantity XY with respect to ...."

Author response:
The *scaling laws* are the scaling laws of the generalization error (we will write it explicitly).

Reviewer comment (Writing):
In Section 1, it would be nice to give an example of a 'behaviour' of models that can be accessed with scaling laws.

Author response:
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.

Reviewer comment (Writing):
In Section 1, what is the criterion to qualify a scaling law as 'improved'?

Author response:
*"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).

Reviewer comment (Writing):
In Section 1, what is exactly meant by a 'statistical rate' in the present context?

Author response:
*"Statistical rates" means the average (over inherited randomness) rates of convergence of the generalization error towards zero.

Reviewer comment (Presentation):
In Figure 8, clarify whether it is SignGD, Adam, or SGD in the plots.

Author response:
*Figure 5* represents Adam with $\beta_1 = \beta_2 = 0$, which is equivalent to SignSGD (https://openreview.net/forum?id=S1EwLkW0W).

Reviewer comment (Evaluation):
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).

Author response:
*"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.

Reviewer comment (Presentation):
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.

Author response:
Thank you for your feedback, we will add a diagram to ease readability.

Reviewer comment (Theory):
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").

Author response:
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.