Understanding Factual Recall in Transformers via Associative Memories
transformers;associative memories;factual recall;storage capacity;training dynamics
Yes
No ethics review needed.
(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). Indeed, there is a large debate in the community over the relationship between memorization and generalization. How could the memorization study of this work be scaled to study the *generalization* abilities of Transformers? I feel this theoretical framework could be used to tackle this problem.

(Q2) The paper only studies the causal-language modelling (i.e., GPT-style) objective. How would the theory extend to transformers trained on masked-token (i.e., BERT-style) tasks?

(Q3) Is $\sigma$ in Eq. (2) defined, or just an arbitrary non-linearity? The theory of [Dense Associative Memory](https://arxiv.org/abs/1606.01164) says that the choice of this non-linearity is critical to increasing the memory capacity of Associative Memories

(Q4) The experimental details around Fig 1 are incomplete -- I do not think I could replicate that experiment with the information included in the paper. Additionally, 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 (see (Q4)).
(S1) **Possibly high impact**. If the analyses performed in this paper are shown to hold for deeper transformers (with shared or unshared weights) using softmax activations in the attention and bigger datasets, the paper has the possibility to be high impact -- it would solidify the benefit of viewing transformers as formal Associative Memories, uniting the LLM community and the more niche physics- and math- communities studying associative memories. 

(S2) **Good, formal definition of Hallucination**. In a passing statement, this paper defines hallucination as the stage where [L455-460]

> ...the model ignores all other tokens in the sequence $z_{1:T}$ – including the useful subject token $s$ – and predicts based only on the relation r... the model is outputting a plausible, yet ultimately incorrect, answer to the prompt. 

As the associative memory perspective of transformers continues to gain traction and the theory of this paper extends to practically sized transformers and real NLP datasets, I can see this formal description of hallucination (i.e., the difference in the Cross-entropy loss of relationships and subjects) as game changing for the study of hallucination in LLMs. Related works do not seem to agree on a definition for "hallucination" except that it encompasses "undesirable predictions" from an LLM.
This paper is a natural extension to Cabannes et al. (2024) in three ways:

1. It extends the theory to a learnable task with trainable weights, not just studying random patterns
2. It considers the MLP and attention as contributions to the Associative Memory properties of transformers
3. It considers the (single-layer) transformer as a gradient flow

The goals are the same: to study the "memorization capacity" a.k.a. "scaling laws" of a transformer from the perspective of Associative Memory. The authors study this memorization capacity on a synthetic recall task consisting of memorizing subject-relationship correlations.
(W1) **Unclear application to real Transformers and datasets**. (dual to (S1)) The simplified Transformer architecture studied theoretically and empirically in this work is smaller than real Transformers, uses only a single update step for much of the analysis, only uses a synthetic dataset to validate the theory, and drops the softmax from the attention. It is unclear how the theory would generalize to unshared weights and real datasets where the task is to reproduce sequences of data.


(W2) **No comparison to the empirical scaling laws of modern LLMs**. Like (W1), this paper could be improved if the theoretical results are compared to the empirical results of e.g., the [Chinchilla scaling laws](https://arxiv.org/abs/2203.15556).

(W3) **Associative memories are single-step updates**. One limitation of this work is that the associative memories are not energy-based models where the memories are encoded as fixed points of energy descent dynamics. Indeed, the memory capacity of the linear associative memory model would not scale linearly if formulated as an energy-minimization problem. It would be nice to see memory capacity studied in the recurrent setting.

**Summary**

Despite these weaknesses, I believe that the contributions of this paper earn my vote for acceptance at ICLR. The paper is clearly presented, with testable claims, and sufficient empirical results. My concerns can be addressed in follow-up works by the community. 

**Note**

The MLP + attention architecture of Transformers has been explicitly formulated as an Associative Memory (specifically, a Hopfield Network with energy descent dynamics) in [Hoover et al. Energy Transformer (2023)](https://arxiv.org/abs/2302.07253). Please add this work to the related work section.
Yes
No ethics review needed.
1. 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. 
2. 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?
1. The theoretical contributions are solid
This papers show shallow transformers can use a combination of associative memories to obtain near optimal storage capacity. The verify this on a synthetic factual recall test.
1. The theoretical toy model is rather simple (MLP and one-layer transformer). It is unclear how it generalizes to multi-layer transformer.
2. If I understand correctly, the synthetic setting relies on the fact that noise tokens and subject tokens are disjoint. 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.
Yes
No ethics review needed.
- 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. 
- Section 4.2 appears to omit LayerNorm, which is part of the original Transformer layer. Would the theoretical results still hold if LayerNorm were considered?
- This work offers a theoretical understanding of how shallow Transformers handle factual recall, specifically revealing linear and MLP associative memory capacity in managing storage.
- The gradient flow analysis, particularly the identification of an intermediate "hallucination" phase, adds depth to the understanding of transformer training dynamics.
- The paper demonstrates mathematical rigor in its proofs regarding associative memory storage capacity
This paper examines factual recall in shallow Transformers, proposing that transformers can leverage associative memories for near-optimal storage of factual information. The authors demonstrate that linear and multi-layer perceptron (MLP) associative memories achieve storage capacity that scales linearly with the parameter count. They design a synthetic factual recall task to show that a single-layer Transformer can reliably store and retrieve facts using self-attention with or without an MLP. Additionally, their gradient dynamics analysis in a simplified linear attention model reveals a "hallucination" stage, where the model’s loss plateaus and predictions rely solely on relation information.
- The findings may be somewhat narrow, given the reliance on synthetic tasks. 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.
- Certain sections, particularly those describing empirical setups for synthetic tasks, could benefit from added detail. Greater specificity on initialization, parameter tuning, and model configuration across trials would improve clarity and reproducibility.
- More discussion on related work exploring neural networks and attention mechanisms as associative memories [1,2] would strengthen the paper and highlight the novelty of the insight. 

[1] Radhakrishnan, Adityanarayanan, Mikhail Belkin, and Caroline Uhler. "Overparameterized neural networks implement associative memory." Proceedings of the National Academy of Sciences 117, no. 44 (2020): 27162-27170.

[2] Le, Hung, Truyen Tran, and Svetha Venkatesh. "Self-attentive associative memory." In International conference on machine learning, pp. 5682-5691. PMLR, 2020.
@inproceedings{
anonymous2024understanding,
title={Understanding Factual Recall in Transformers via Associative Memories},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=hwSmPOAmhk},
note={under review}
}
Large language models have demonstrated an impressive ability to perform factual recall. Prior work has found that transformers trained on factual recall tasks can store information at a rate proportional to their parameter count. In our work, we show that shallow transformers can use a combination of associative memories to obtain such near optimal storage capacity. We begin by proving that the storage capacities of both linear and MLP associative memories scale linearly with parameter count. We next introduce a synthetic factual recall task, and prove that a transformer with a single layer of self-attention followed by an MLP can obtain 100\% accuracy on the task whenever either the total number of self-attention parameters or MLP parameters scales (up to log factors) linearly with the number of facts. In particular, the transformer can trade off between using the value matrices or the MLP as an associative memory to store the dataset of facts. We complement these expressivity results with an analysis of the gradient flow trajectory of a simplified linear attention model trained on our factual recall task, where we show that the model exhibits sequential learning behavior.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
transformers associative memories factual recall storage capacity training dynamics
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/4b2b9b6c4c0eb65bffcedd6dc6707944c4d110d0.pdf
learning theory
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Understanding Factual Recall in Transformers via Associative Memories
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission