Understanding Factual Recall in Transformers via Associative Memories

Thank you to the reviewer for your thoughtful and detailed review. We respond to your comments below.

“While the theoretical insights are valuable, it remains uncertain how well these translate to complex, real-world scenarios involving non-random and interdependent factual associations.”

“The study is centered on shallow models, raising questions about the applicability of these findings to deeper, large-scale transformers commonly used in research and industry.”

We agree with these comments, but would like to also remark that Allen-Zhu & Li (2024) show empirically that larger-scale models can store an amount of information proportional to their parameter count. Our work demonstrates that a shallow transformer can use associative memories to obtain such linear scaling. Please refer to the global comment.

“Greater specificity on initialization, parameter tuning, and model configuration across trials would improve clarity and reproducibility.”

Thank you for the suggestion. In Appendix A, we have added additional experimental details.

“The statement on line 117—that the number of associations scales linearly with the number of parameters (up to log factors)—is somewhat unclear. Without injective assumption, it looks like the number of tokens, not associations, scales linearly. Please elaborate on that point.”

We say that the model can “store” the association $(x, f^*(x))$ if $\arg\max_y u_y^TF(e_x) = f^*(x)$. Theorem 1 states that if $N = \tilde O(d^2)$, then with high probability there exists a model that can store all $N$ associations $\{(x, f^*(x)\}_{x \in [N]}$, as long as $f^*$ is injective. We remark that in the MLP associative memory construction in Theorem 2, we no longer need the injectivity assumption, and thus we can always store $N = \tilde O(md)$ associations regardless of the choice of $f^*$.

“Would the theoretical results still hold if LayerNorm were considered?”

We expect the theoretical results to still hold with layer norm. On a prompt containing subject $s$ and relation $r$, the transformer will output a vector which is close to $\phi(a^*(s, r))$. Since the embeddings are drawn uniformly on the sphere, as $d$ grows the mean will be close to zero and the variance will concentrate, and as such layer norm would have the same effect on each token.

Thank you for the additional references, we have added discussion on these points to the related work.

Thank you to the reviewer for your thoughtful and detailed review. We respond to your comments below.

“The theoretical toy model is rather simple (MLP and one-layer transformer). It is unclear how it generalizes to multi-layer transformer.”

We agree with this comment, but would like to also remark that Allen-Zhu & Li (2024) show empirically that larger-scale models can store an amount of information proportional to their parameter count. Our work demonstrates that a shallow transformer can use associative memories to obtain such linear scaling. Please refer to the global comment.

“in reality, usually the last token (not eos in practice) should be somewhat relevant for token selection (for example, “in” as the last token would look for locations). It doesn’t make too much sense to let model do next token prediction after some random tokens”

Our task can be thought of modeling the scenario where the prompt is a question, and the final token in the sequence is a question mark. As you suggest, one alternate task that is also reasonable would be to have the final token in the sequence either be the relation or be a relevant token for the relation (i.e in your example, the token “in” signals that the relation is “location”). We hypothesize that a minor modification of our construction in Section 4.3 succeeds on this task with parameter count remaining proportional to number of facts. A brief sketch is as follows: Each head is associated with a subset $S^{(h)}$ of subjects and a subset $R^{(h)}$ of relations. Setting $W^{(h)}\_{KQ} \propto \sum\_{s \in S^{(h)}, r \in R^{(h)}} \phi(s)\phi(r)^T$, this head attends only to tokens in $S^{(h)}$ assuming that a relation in $R^{(h)}$ is present. We can add a formal description of this construction to the next revision.

“How did you get the plot for Figure 5 (left)? It seems a bit too smooth. Also, why didn’t the first 2000 GD change the loss? It seems counter-intuitive.”

Figure 5 (left) is the loss when training the linear attention model with orthogonal embeddings. Since the embeddings are orthogonal, the dynamics for each subject and relation are decoupled, and thus better behaved than the non-orthogonal case (right). The loss does not decrease for the first 2000 steps due to the small initialization, and the existence of a saddle around 0. While the exact saddle only exists for the linear attention model, we do see a similar phenomenology in the softmax attention model (Figure 5; right) where the loss decreases slowly for the first few steps of GD.

“Could you give an intuitive explanation for why the model, in theorem 6, learns conditional distribution of relations as opposed to conditional distribution of subject?”

In theorem 6, we assume that the number of subjects S is much larger than the number of relations R. Since there are fewer relations, predicting using only the relation is more “useful” than predicting using only the subject. We expect that the reverse order would hold if there are more relations than subjects.

Thank you to the reviewer for your thoughtful and detailed review. We respond to your comments below

“W1: Unclear application to real Transformers and datasets.”

We agree with this comment, but would like to also remark that Allen-Zhu & Li (2024) show empirically that larger-scale models can store an amount of information proportional to their parameter count. Our work demonstrates that a shallow transformer can use associative memories to obtain such linear scaling. Please refer to the global comment.

“W2: No comparison to the empirical scaling laws of modern LLMs.”

We briefly discuss one connection to scaling laws in Section 6. For the basic associative memory task, we prove that if the token frequencies are distributed as a power law ($p(x) \propto x^{-\alpha}$, then any model with $B$ parameters must incur a loss of at least $B^{1 - \alpha}$. This lower bound is obtained by an MLP associative memory which stores the $\Theta(B)$ most probable associations, and matches the scaling law with respect to model size obtained in prior works. We agree that it would be an interesting direction of future work to understand what scaling laws are achieved if the subjects and relations in our factual recall task are also distributed according to a power law.

“W3: Associative memories are single-step updates.”

We remark that the weight matrices in both self-attention and an MLP can be interpreted as these “single-step update” associative memories, as demonstrated in the prior works (Bietti et al. (2023), Cabannes et al. (2024)), which is why we adopt this viewpoint in our work.

“(Q1) Standard Transformers typically output a probability distribution over all output tokens. The theory of Eq. (7) only considers the most-likely token (equivalent to setting the sampling temperature of LLMs -> close to 0)”

We remark that even though we focus on arg-max prediction, this can be achieved for any fixed sampling temperature by scaling up the magnitude of our constructions to infinity, so that the next token distribution concentrates on the most likely token. We do agree that it would be interesting to design a synthetic task where there are multiple possible valid responses for the next token.

“How could the memorization study of this work be scaled to study the generalization abilities of Transformers?”

Our factual recall task can be viewed as one instance of generalization, where the model must ignore the noise tokens while storing (s, r, a) tuples; in this case, the model may be able to generalize to sequences with different noise tokens or where the subject and relation are in different positions. In this case, it would be very interesting to prove a finite-sample bound on the number of such sequences the transformer must see in order to obtain such generalization.

“(Q2)...How would the theory extend to transformers trained on masked-token (i.e., BERT-style) tasks?”

It would be interesting to understand how our theory extends to the masked language modeling setting. One hypothesis is that if the token containing the answer is the one that is masked out, then an identical construction to that in Section 4.3, with $\phi(EOS)$ replaced by $\phi(MASK)$, would succeed.

“Q3: Is $\sigma$ in Eq. (2) defined, or just an arbitrary non-linearity?”

We require $\sigma$ to be a polynomial of sufficiently high degree. This coincides with the theory for Dense Associative Memories, which require the energy function to also be a polynomial of high degree in order to store many more associations. We expect that this polynomial assumption can be relaxed to instead assume that the Hermite coefficients of $\sigma$ are nonzero, which is a common assumption in prior deep learning theory work and includes standard activations such as ReLU with random bias.

“(Q4) The experimental details around Fig 1 are incomplete…Fig 1 could be improved with error bars and notable "failure cases" e.g., the task of storing one-hot vectors [L130] or under different choices of $\sigma$ in the MLP."

Thank you for this suggestion. We have added additional experimental details to Appendix A. We will regenerate Figure 1 with more random seeds and error bars and add to the next revision. We remark that in our experiments, the scaling was relatively consistent for different choices of $\sigma$.

Thank you for suggesting the additional reference, we have added it to our related work section.