Understanding Factual Recall in Transformers via Associative Memories

Authors: Eshaan Nichani, Jason D. Lee, Alberto Bietti
Primary area: learning theory
Keywords: transformers, associative memories, factual recall, storage capacity, training dynamics

Abstract:
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.

Decision: accept (spotlight)
Reviewer scores: 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730088674641

thread_type:
official_review

<thread_object>
creator: reviewer_fxtn

<content>
summary:
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.

strengths:
- 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

weaknesses:
- 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.

questions:
- 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730587397133

thread_type:
official_review

<thread_object>
creator: reviewer_kbfg

<content>
summary:
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.

strengths:
1. The theoretical contributions are solid

weaknesses:
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.

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730743983823

thread_type:
official_review

<thread_object>
creator: reviewer_wfua

<content>
summary:
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.

strengths:
(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.

weaknesses:
(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.

questions:
(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)).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732286263193

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Comment

comment:
Thank you to all of the reviewers for their detailed feedback and the overall positive evaluation of our paper.

## Re: Application of our theory to larger models and more realistic datasets.

The reviewers point out that one limitation of our work is that our theory focuses only on a single-layer transformer. We do acknowledge this point. We would like to additionally remark that our work is motivated by the prior work Allen-Zhu & Li (2024), which shows empirically that larger scale (GPT2 size) transformers trained on a dataset consisting of synthetic biographies can store an amount of information proportional to the model’s parameter count. One of our contributions is that a shallow transformer can provably obtain such linear in number of parameters scaling using associative memories to memorize the factual information, providing evidence that associative memories may be used to store information for large-scale models and more complex factual recall tasks. Understanding how our mechanism generalizes to larger models and more realistic datasets is indeed an interesting direction of future work.

We respond to individual comments below.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper offers compelling theoretical insights into transformers' factual recall capabilities through associative memories, demonstrating significant contributions to understanding information storage and retrieval. Despite initial concerns about synthetic tasks, reviewers praised the mathematical rigor, with multiple reviewers supporting acceptance and highlighting the novel analysis of gradient flow dynamics. The authors effectively contextualized their work within broader empirical findings, showing how their shallow transformer model provides evidence for associative memory mechanisms in larger-scale models. Comprehensive responses to reviewer questions and the potential for future research extensions make this paper a strong candidate for acceptance.