Paper: Memorization Capacity of Multi-Head Attention in Transformers (ICLR-2024)

Reviewer comment (Writing):
One of my main concern is the illustration or definition of "memorization" in this paper. The inputs of attention include both the key matrix and the query vector. In a common understanding, attention plays a role to capture knowledge from the context according to the "attention" on other tokens for each token. So what does attention memorize? I think the paper should make it clearer before or after the theorectical analysis, or even verify the memorized knowledge with some visualization.

Author response:
**Q1: One of my main concern is the illustration or definition of "memorization" in this paper. The inputs of attention include both the key matrix and the query vector. In a common understanding, attention plays a role to capture knowledge from the context according to the "attention" on other tokens for each token. So what does attention memorize? I think the paper should make it clearer before or after the theorectical analysis, or even verify the memorized knowledge with some visualization.**

In response to the concern about the definition of "memorization", we adopt a classical definition applicable to any parameterizable function $f$ and a set of $T$ input/output pairs $x_i$, $y_i$. This definition aligns with established literature on both ReLU networks and Transformers [3, 4, 5, 6, 7]. Our focus in this work is on Attention-only models, isolating them from the well-studied ReLU networks to explore their capabilities in the absence of ReLU activation.

While the common belief attributes the success of attention to its ability to capture knowledge from various parts of the input, we argue that this is not the sole reason behind its efficacy. We explicitly recognize the utility of attention in capturing input context information in Section 4.4 and Proposition 4, where we demonstrate its effectiveness in token-mixing (or contextualization). Moreover, in our proof technique, we show that we can partition input data into subsets and make each head of the Multi-Head attention to be responsible for memorizing only one subset. This responsibility is not in the form of "capturing knowledge from the context", but rather, exploiting the context to build a desired convex combination output vector (see Appendix B for complete proof). Moreover, we do not impose any explicit assumption on the relationship between the input and label, and our theorem works for any arbitrary label assigned to an input datapoint, therefore the notion of "knowledge" is not well-defined. Finally, there have been recent empirical observations that show merely extending the transformer's input context with *constant* tokens can enhance the performance of transformers (e.g., [1, 2]). Incorporating such constant tokens enhances the performance without adding any "knowledge" to the context.

Reviewer comment (Novelty):
Analysis and comparison with existing works are insufficient, for example, <<Transformer Feed-Forward Layers Are Key-Value Memories>>. How are the memories different between this paper and the reference? There are also some works about transformer interpretability and model editing. Do the observations in these related works support the findings in this paper? I think the authors should give more analysis, rather than demonstrating their propsed theorem only.

Author response:
**Q2: Analysis and comparison with existing works are insufficient, for example, <<Transformer Feed-Forward Layers Are Key-Value Memories>>. How are the memories different between this paper and the reference? There are also some works about transformer interpretability and model editing. Do the observations in these related works support the findings in this paper? I think the authors should give more analysis, rather than demonstrating the their propsed theorem only.**

Thanks for bringing this paper to our attention. First, we would like to highlight the difference between "neural memory" definition, and the "memorization" definition used in our work. In neural memory, typically the differentiable weights are seen as memory units, and the neural network is seen as a function that utilizes/modifies these weights to produce an answer. In our work, however, memorization is defined as the network being able to memorize (i.e., interpolate) the output $y$ for each corresponding input $x$. In memory units, the network has a specific structure such that some weights have specific interpretations as memory units, while in memorization, the goal is to achieve an output given input with the smallest possible parameter count. Second, the work you mentioned focuses on the fully-connected ReLU part of a Transformer, and interprets its first layer weights as "memory keys", and its second layer as "memory values". **While it is interesting to study the role of ReLU FCN in transformers, we take an orthogonal approach and focus on isolating the attention part of the transformer.** As a result, the focus and scope of our paper differ from the work you mentioned, making direct comparisons challenging. In the new revision, we updated our manuscript with a footnote to clarify the distinction between concepts.

If you have additional specific papers in mind that you believe would enhance the discussion, we are open to incorporating them into our analysis to further enhance the exposition of our work and its relation to existing literature. 
    
[1] S. Goyal et al. Think before you speak: Training Language Models With Pause Tokens, 2023
    
[2] T. Darcet et al. Vision Transformers Need Registers, 2023
    
[3] R. Vershynin. Memory capacity of neural networks with threshold and rectified linear unit activations. SIAM J., 2020.

[4] G. Vardi, et al. On the optimal memorization power of reLU neural networks. ICLR, 2022.

[5] C. Yun, et al. Small relu networks are powerful memorizers: a tight analysis of memorization capacity. NeurIPS, 2019.

[6] S. Bubeck, et al. Network size and size of the weights in memorization with two-layers neural networks, NeurIPS, 2020.
    
[7] J. Kim et al. Provable Memorization Capacity of Transformers, ICLR 2023.

Reviewer comment (Experiments):
It might be significantly different between the image patch tokens (ViT) and the language tokens. Can the author's experimental verification of those assumptions be verified on NLP tasks?

Author response:
As mentioned in Section 4, our linear independence assumptions are primarily best-suited for continuous input data such as images, hence our focus on ViT and ImageNet. For completeness, we have included experiments on GPT-2 and BERT in Appendix J. We find that randomly initialized GPT-2 satisfies both assumptions and Theorem 1 is applicable in this setting. Moreover, a randomly initialized BERT satisfies Assumption 2. For BERT, while it appears to violate Assumption 1, our theory provides a (weaker) bound on memorization as discussed in Remark 1.

Reviewer comment (Theory):
The authors demonstrate that Assumption 2 typically holds in practice due to positional encoding. Does the assumption still hold when the positional encoding is learnable?

Author response:
Yes. When positional encodings are learnable, they become a part of the designed parameter set for memorization. One simple such design is setting the positional encodings randomly from a continuous distribution, such as uniform or Gaussian. This also agrees with practice when initializing neural network weights. This random initialization ensures that each entry is independent. Given that the number of positions ($n$) is less than the dimensionality of the encoding ($d$), the positional encoding vectors $\text{pos}_1, \text{pos}_2, \dots, \text{pos}_n \in \mathbb{R}^d$ are almost surely linearly independent (i.e., are linearly independent with probability one). This is because for each $\text{pos}_i$, the set of possible vectors that make $\text{pos}_i$ a linear combination of $\text{pos}_1, \text{pos}_2, ..., \text{pos}_j$ (for $j<i$) has measure zero. Consequently, these learnable positional encodings provably help satisfy Assumption 2 (like the sinusoidal positional encodings). However, for a trained network, no definitive statement can be made. Learnable positional encodings are subject to change during training, and their linear independence may evolve in either direction. Therefore, adherence to Assumption 2 cannot be theoretically guaranteed once the network has undergone training. But, again, since memorization capacity asks for *existence* of parameters memorizing the data, positional encoding becomes a part of the whole set of parameters, which can be "designed" by our proof.

Reviewer comment (Experiments):
Limited Empirical Testing: While the paper includes synthetic experiments, real-world data experiments might be needed to fully understand the practical implications of the findings.

Author response:
We would like to point out that we have indeed verified our assumptions on real-world models and datasets, as discussed in Section 5.1 of the paper. For testing memorization capacity on real-world data, the issue lies in the nature of memorization question, which only asks for the **existence** of a solution. There is no apriori easy way to check this. As we discussed in Appendix A, the nature of gradient-based optimization (e.g., using Adam optimizer) as a proxy for quantifying memorization presents a significant challenge in terms of computational resources. For instance, even for a small-scale model and synthetic dataset with a dimension of 64, we had to run for 300,000 optimization steps to mitigate optimization-related confounding limiting factors in our plot. In conclusion, given the existential nature of the question of memorization capacity, the more relevant question is that of verifying the validity of our data assumptions. As mentioned, we have done so in real-world experiments as suggested by the reviewer.

Reviewer comment (Experiments):
Focus on Single-Layer MHA Module: The study primarily focuses on a single-layer Multi-head Attention (MHA) module. Expanding the analysis to multi-layered architectures could provide more comprehensive insights.

Author response:
We would like to draw attention to the progression of results for memorization in fully-connected networks. The first results appeared in 1988[1] for only a single hidden layer, and with a strong “general position” assumption (which, in our paper, we showed is highly unrealistic in our setup and offered a relaxed version based on Kruskal Rank). Only decades later, results for multi-layer ReLU networks have been proposed and various assumptions lifted[2, 3, 4], yet, the exact memorization capacity of MLPs of depth $\geq 3$ is still not fully known (only known up to logarithmic and constant factor).

Here we take an analogous first step for attention models. Thus, we expect our analysis to inspire further investigations and extensions to multi-layer architectures.

Reviewer comment (Theory):
Could you explain the rationale behind the specific assumptions made regarding the linear independence of input data? How do these assumptions align with real-world data scenarios in transformer applications?

Author response:
Our Assumption 1 is inspired by the "General Position" assumption found in several prior works on ReLU FCN [1, 6]. However, as our experiments indicate (Table 1), this assumption is often not met in practice. This led us to propose the Kruskal rank assumption as a relaxation to the general position (i.e., it subsumes the general position assumption as a special case), which we also found to better-suite practical settings (specifically kurskal rank larger than context size $n$ in Table 2). Moreover, as mentioned in section 4.4, the positional encoding used in transformers is the primary motivation behind Assumption 2.

Reviewer comment (Evaluation):
How do your findings contribute to the ongoing discussion in machine learning between model generalization and memorization, particularly in the context of overfitting?

Author response:
While our work does not directly consider generalization, we believe precise quantification of the memorization ability of models on real-world datasets provides significant insight into generalization (e.g., measuring how much of the success of foundation models can be attributed to memorization rather than generalization). We take a step towards better such quantification by (1) refining common assumptions on the data towards being more realistic (Assumptions 1 and 2) and (2) proposing a lower bound on the memorization with such assumptions (Theorem 1). Please also see our response to Q3.