On the Optimal Memorization Capacity of Transformers

Q1. I don't completely follow the reasoning behind the statement: "the primary bottleneck is not the contextual mapping of tokens, but rather the feed-forward layers’ capacity to map this token-level contextual information to labels." Clarification on this point would be helpful.

Thank you for pointing this out. To provide some background, prior work [1] suggested that self-attention mechanism performs contextual mapping, while feed-forward layers are responsible for the actual memorization, that is, mapping token-level contextual information to labels. The statement "the primary bottleneck is..." was written with this context in mind and reflects how our results align with this background.

Here is a more detailed explanation: as shown in Theorem 4.1, the upper bound for the next-token prediction setting already allows for distinguishing between $N$ input sequences using $\tilde{O}(\sqrt{N})$ parameters. Given that the ability to distinguish input sequences, in addition to token-level information, implies that the context of each token can also be identified, this implies that contextual mapping (in the sense of Definition 4.1) is already achieved within this parameter bound. The additional $\sqrt{n}$ factor in the seq-to-seq prediction setting reflects the extra effort required to map this contextual information to labels. In this sense, we mentioned that the feed-forward network is the bottleneck, as it is primarily responsible for this context-to-label mapping.

Q2. If the feed-forward layer is identified as the bottleneck, is there an approach to mitigate this through architectural changes or modifications?

Addressing the bottleneck in the feed-forward network would essentially require designing a model that performs memorization more efficiently than a feed-forward network. At present, this appears to be a challenging task. One potential research direction could involve exploring architectures that are better suited to real-world datasets, as memorization capacity is a worst-case analysis. For instance, just as mixture of experts can serve as an alternative to pure feed-forward layers, there may be architectures that are more effective for specific data distributions. In this context, advancing the study of memorization capacity of feed-forward networks in more developed settings could be a valuable avenue for future research.

Q3. Are the data assumptions too restrictive? It feels surprising that the results hold if the attention layer primarily averages, especially when compared to findings in studies such as (a) and (b), despite the differences in setting.

We would like to note that in the study of memorization capacity, $(r,\delta)$-separatedness is considered a standard assumption [1, 2, 3, 4, 5, 6, 7]. Since memorization capacity is the notion for the framework of finite dataset, and any finite discrete dataset (for example, any text dataset) is inherently $(r,\delta)$-separated to some extent, this assumption is relatively loose. As a next step, it would indeed be valuable to consider more sophisticated data assumptions in future research.

Regarding the cited work, we believe the differences arise primarily from the problem settings, not the data assumption. The paper (a) investigates how Transformers can represent conditional $k$-th order Markov processes. Here, "conditional" refers to computing an empirical $k$-th order Markov distribution in the in-contextual way, given a specific input sequence, and as such this is a distinct problem setting from our work. In such settings, uniform attention may indeed be insufficient, and fully leveraging the keys and queries in self-attention could be essential for capturing these contextual relationships.

The paper (b), on the other hand, analyzes Transformers trained on data generated by a first order Markov process. Their results effectively disregard self-attention mechanism (this conclusion is natural, considering the fact that a first-order Markov process depends only on the current token).

References

[1] Kajitsuka and Sato, Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? ICLR 2024.

[2] Hardt and Ma. Identity Matters in Deep Learning. ICLR 2017.

[3] Vershynin, Memory Capacity of Neural Networks with Threshold and ReLU Activations. arXiv 2020.

[4] Park et al., Provable Memorization via Deep Neural Networks using Sub-linear Parameters. COLT 2021.

[5] Vardi et al., On the Optimal Memorization Capacity of ReLU Neural Networks. ICLR 2021.

[6] Kim et al., Provable Memorization Capacity of Transformers. ICLR 2023.

[7] Siegel, Sharp Lower Bounds on Interpolation by Deep ReLU Neural Networks at Irregularly Spaced Data. arXiv 2024.

Thank you for taking the time to review our paper and sharing your insightful comments. We will address your questions below.

W1. The sketches of the proofs in the paper were somewhat difficult to follow without referring to the full proofs. Improving the presentation could help a wider audience better understand and appreciate the work.

We have added more detailed explanations about the high-dimensional representations $\tilde{X}^{(1)},\dots,\tilde{X}^{(N)}$ used to construct the feed-forward network $\phi$ in the proof sketch of Theorem 4.1, particularly around eq. (9). We hope this provides additional clarity and aids in understanding. If there are still unclear parts, please let us know, and we would be happy to provide further explanations.

W2. While the paper is theoretical, and experiments may not be essential, including experiments to support the claims would strengthen the work.

In response to your suggestion, we have included three types of experiments in the revised version of the paper to verify our theoretical analysis (details are provided in Appendix H). First, we conducted an experiment to verify that the model described in Theorem 4.1 can successfully perform memorization using two real-world datasets. Specifically, we trained a Transformer consisting of two feed-forward networks and a single layer of uniform attention in between on MNLI and IMDb datasets. The experimental results are summarized in the table below, where \\#\mathrm{blocks} indicates the depth of each of the two token-wise feed-forward networks in the model.

For MultiNLI dataset:

For IMDb dataset:

Overall, the results confirm that, given adequately deep feed-forward networks, the training error can be reduced to near zero, which validates that uniform attention actually possesses sufficient representational capacity for memorization.

In addition, we evaluated the behavior of the memorization capacity, that is, the minimum number of \\#\mathrm{blocks} required for memorizing real-world dataset (MultiNLI) as the dataset size varies.

The experimental results are summarized in the table below.

The results indicate that for smaller datasets, the memorization capacity scales as $\sqrt{N}$, consistent with the predictions of Theorems 4.1 and 4.2. For larger datasets, however, the memorization capacity increases rapidly.

While the influence of optimization in real-world experiments may result in sub-optimal models, a plausible explanation for this phenomenon is that the fixed bit-length of model parameters becomes a bottleneck as the dataset size grows. We have added a discussion in Appendix F on the memorization capacity of Transformers in scenarios where the bit-length of parameters is insufficient. Specifically, the result from Appendix F suggests that the memorization capacity scales linearly with $N$ when the bit-length of parameters are limited (up to logarithmic factors).

We also trained the model on randomly generated datasets with varying input sequence length and fixed dataset size. The dataset size was fixed at $N=500$, and we considered four input sequence lengths: $n=10,100,1000,10000$. For all cases, we verified that a Transformer combining a single uniform attention layer with $4$-block feed-forward networks achieved a training error below the threshold $ϵ=0.01$. This result not only demonstrates that memorization is possible even with a single uniform attention layer, but also confirms the prediction made in Theorem 4.1 that the memorization capacity in the next-token prediction setting is nearly independent of the input sequence length $n$.

The experiments we have conducted significantly deepened our understanding of the memorization capacity of Transformers. It has also allowed us to experimentally confirm the occurrence of the phase transition in memorization capacity. Thank you for your valuable feedback!

Thank you for taking the time to review our paper and sharing your insightful comments. We will address your questions below.

W1. bit complexity: As the authors point out, the transformer needs at least $\Omega(N)$ bits, so to succeed with $\sqrt{N}$ parameters it needs a bit complexity of $\sqrt{N}$. However, this is very unrealistic as these models are trained in finite precisions, most often 16 bit precision. It is therefore unclear if these results say much about transformers ability to memorize in practice. This also affects the motivation, for example I would expected double descent (lines 102-107) to occur at $\Theta(N)$ parameters, not $\Theta(\sqrt{N})$ parameters.

We have included an analysis of the memorization capacity of Transformers with limited bit complexity in Appendix F. As you mentioned in your question 3, Theorem F.1 demonstrates that a Transformer constrained to a bit complexity of $\tilde{O}(N^{1 - \alpha})$ for some $\alpha \in [1/2,1]$ can memorize $N$ data points in the next-token prediction setting using $\tilde{O}(N^{\alpha})$ parameters.

W2. experiments: Related to the above point, the paper would benefit from some simple experiments that demonstrate whether such constructions are learnable or whether the empirical limit is $\Theta(N)$ or $\Theta(\sqrt{N})$ parameters. It could also be interesting to test this with various levels of precisions (FP16/32/64) to see whether higher precision really does let the transformer memorize with fewer parameters.

In response to your feedback, we evaluated the behavior of the memorization capacity, that is, the minimum size of Transformers required for memorizing real-world dataset (MultiNLI) as the dataset size varies (details are provided in Appendix H). The experimental results are summarized in the table below, where \\#\mathrm{blocks} indicates the depth of each of the two token-wise feed-forward networks in the model and was used as a proxy for the model size.

Overall, the results indicate that for smaller datasets, the memorization capacity scales as $\sqrt{N}$, consistent with the predictions of Theorems 4.1 and 4.2. For larger datasets, however, the memorization capacity increases rapidly.

While the influence of optimization in real-world experiments may result in sub-optimal models, a plausible explanation for this phenomenon, as you suggested, is that the fixed bit-length of model parameters becomes a bottleneck as the dataset size grows. In such a regime, the results from Appendix F on Transformers with limited bit complexity suggest that the memorization capacity scales linearly with $N$.

Your feedback has greatly enhanced our understanding of the memorization capacity in real-world models. It has also allowed us to experimentally confirm the occurrence of this phase transition in memorization capacity, thereby broadening the scope of this work. We sincerely appreciate your insightful comment.

W3. minor typo on line 462: viewd -> viewed

Thank you for pointing this out. We have corrected the typo in the revised paper.

Thank you for taking the time to review our paper and sharing your insightful comments. We will address your questions below.

W1. Token-wise $(r,\delta)$-separated sequences may not be an expressive enough model to fully capture the memorization capacity of Transformers. It appears that the authors construct a Transformer that captures contextual information by 'counting' the occurrences of tokens in a subset. However, Transformers typically capture contextual information from the mutual relationships among tokens via the attention mechanism.

Regarding the use of $(r,\epsilon)$-separated sequences, we note that this is a standard assumption in the study of memorization capacity [1, 2, 3, 4, 5, 6, 7]. We understand the reviewer’s point about Transformers capturing contextual information through the mutual relationships among tokens via the self-attention mechanism. One of the main messages of our work is that to fully understand how keys and queries are utilized in real-world self-attention, it would be better to analyze Transformers from the perspectives of optimization or generalization, rather than solely focusing on memorization capacity.

W2. In the proof of Theorem 4.1, the attention layer is set to uniform attention, functioning as a column-wise average operator. This over-simplification of the attention mechanism limits the ability of the work to reveal how attention contributes to memorization, as the attention mechanism should gather context-aware representations of tokens based on their relationships with others in the sequence. Therefore, how do the authors understand the role that the attention mechanism plays in the memorization capacity of the Transformer within their framework?

We would like to clarify that we did not oversimplify the problem setting by replacing self-attention with uniform attention. Rather, our analysis of the memorization capacity of a Transformer with general version of self-attention revealed that uniform attention already achieves the sufficient representational capacity. Moreover, prior research [8, 9] has indicated that even with uniform attention, Transformers can achieve sufficient accuracy for certain classification tasks. In this sense, our work provides theoretical support for these empirically observed phenomena, reinforcing the understanding of uniform attention's effectiveness in specific contexts, including memorization.

In response to your feedback, we conducted experiments to demonstrate that the model described in Theorem 4.1 can indeed achieve memorization on two real-world datasets (details are provided in Appendix H). Specifically, we trained a Transformer consisting of a single uniform attention layer and two token-wise feed-forward networks on MultiNLI dataset and IMDb dataset. The experimental results are summarized in the table below, where \\# \mathrm{blocks} indicates the depth of each of the two token-wise feed-forward networks in the model.

For MultiNLI dataset:

For IMDb dataset:

Overall, the results confirm that, given adequately deep feed-forward networks, the training error can be reduced to near zero, which validates that uniform attention actually possesses sufficient representational capacity for memorization.

References

[1] Hardt and Ma. Identity Matters in Deep Learning. ICLR 2017.

[2] Vershynin, Memory Capacity of Neural Networks with Threshold and ReLU Activations. arXiv 2020.

[3] Park et al., Provable Memorization via Deep Neural Networks using Sub-linear Parameters. COLT 2021.

[4] Vardi et al., On the Optimal Memorization Capacity of ReLU Neural Networks. ICLR 2021.

[5] Kim et al., Provable Memorization Capacity of Transformers. ICLR 2023.

[6] Kajitsuka and Sato, Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? ICLR 2024.

[7] Siegel, Sharp Lower Bounds on Interpolation by Deep ReLU Neural Networks at Irregularly Spaced Data. arXiv 2024.

[8] Yu et al., MetaFormer Is Actually What You Need for Vision. CVPR 2022.

[9] Hyeon-Woo et al., Scratching Visual Transformer's Back with Uniform Attention. ICCV 2023.

Thank you for taking the time to review our paper and sharing your insightful comments. We will address your questions below.

W1. The title wording could be improved for accuracy, I would suggest 'near-optimal...' instead of 'optimal...' since there is a logarithm term mismatch between lower bound and upper bound.

Thank you for your suggestion. Following the title of prior work [1], which established upper and lower bounds on the memorization capacity of feed-forward networks with a logarithmic factor gap, we have revised the title to "On the" Optimal Memorization Capacity of Transformers.

W2. The paper can have more discussion on the separatedness assumption. For example, how realistic is the assumption?

In the study of memorization capacity, the $(r,\delta)$-separatedness is a standard assumption [1, 2, 3, 4, 5, 6, 7], and any finite discrete dataset (for example, any text dataset) is inherently $(r,\delta)$-separated to some degree. More specifically, what really matters is the values of $r$ and $\delta$ in the $(r,\delta)$-separatedness. Prior work [7] has shown that when $r\delta^{-1} = \Omega(\exp(N))$, $O(\sqrt{N})$ parameters are no longer sufficient, and $\Omega(N)$ parameters are required for memorization. Given this, the assumption $r\delta^{-1} = N^{O(1)}$ in Theorem 4.1 is sufficiently loose to accommodate cover scenarios outside the range identified in prior work [7]. Since we demonstrated that optimal memorization is achievable under standard $(r,\delta)$-separated datasets in this work, it would indeed be valuable to explore more advanced data assumptions in future research.

Q1. What is the significance of the logarithmic gap in the paper. For instance, do authors believe this gap can be closed with further analysis, or is it likely to be fundamental?

In our current lower bound analysis, we did not explicitly incorporate the $(r,\delta)$-separatedness of the data into the proof. However, we believe that by integrating this assumption into the analysis, it might be possible to eliminate the logarithmic gap. In fact, it is already known [7] that when $r\delta^{-1} = \Omega(\exp(N))$, $O(\sqrt{N})$ parameters are not sufficient, and the lower bound on the memorization capacity becomes $\Omega(N)$. This suggests that the separation of the data is inherently linked to the lower bound of the memorization capacity.

Q2. Is there any examples of real-world datasets that satisfy or approximately satisfy this assumption? 2.2 Can authors discuss how the results might change if this assumption is relaxed?

Memorization capacity is the notion for the framework of finite datasets, and any finite, discrete dataset (for example, any text dataset) is inherently token-wise separated to some degree. We hope this answers your question.

References

[1] Vardi et al., On the Optimal Memorization Capacity of ReLU Neural Networks. ICLR 2021.

[2] Hardt and Ma. Identity Matters in Deep Learning. ICLR 2017.

[3] Vershynin, Memory Capacity of Neural Networks with Threshold and ReLU Activations. arXiv 2020.

[4] Park et al., Provable Memorization via Deep Neural Networks using Sub-linear Parameters. COLT 2021.

[5] Kim et al., Provable Memorization Capacity of Transformers. ICLR 2023.

[6] Kajitsuka and Sato, Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? ICLR 2024.

[7] Siegel, Sharp Lower Bounds on Interpolation by Deep ReLU Neural Networks at Irregularly Spaced Data. arXiv 2024.