On the Optimal Memorization Capacity of Transformers

Authors: Tokio Kajitsuka, Issei Sato
Primary area: learning theory
Keywords: Transformer, Self-Attention, Memorization, Contextual Mapping, Permutation Equivariance

Abstract:
Recent research in the field of machine learning has increasingly focused on the memorization capacity of Transformers, but how efficient they are is not yet well understood.
We demonstrate that Transformers can memorize labels with $\tilde{O}(\sqrt{N})$ parameters in a next-token prediction setting for $N$ input sequences of length $n$, which is proved to be optimal up to logarithmic factors.
This indicates that Transformers can efficiently perform memorization with little influence from the input length $n$ owing to the benefit of parameter sharing.
We also analyze the memorization capacity in the sequence-to-sequence setting, and find that $\tilde{O}(\sqrt{nN})$ parameters are not only sufficient, but also necessary at least for Transformers with hardmax.
These results suggest that while self-attention mechanisms can efficiently identify input sequences, the feed-forward network becomes a bottleneck when associating a label to each token.

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

Reviews and discussion:
<thread>
creation date:
1730501809401

thread_type:
official_review

<thread_object>
creator: reviewer_hfhi

<content>
summary:
The paper examines the memorization capacity of Transformers. The authors claim that Transformers can memorize labels with $O(\sqrt{nN})$ parameters in next-token prediction for sequences. They argue that parameter sharing reduces input length sensitivity. In sequence-to-sequence tasks, $ O(\sqrt{nN})$ parameters are shown to be necessary and sufficient for hardmax transformers. While self-attention identifies input sequences effectively, the feed-forward network seems to be a limiting factor.

strengths:
1. The paper addresses the important problem of quantifying the memorization capacity of transformers.
2. The authors present theorems that cover a wide range of tasks for which transformers are used, making the analysis comprehensive.

weaknesses:
1. 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.
2. While the paper is theoretical, and experiments may not be essential, including experiments to support the claims would strengthen the work.

questions:
1. 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.

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

3. 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.

4. Would it be possible to include experiments, even toy examples, to support some of the paper’s claims? Such experiments could significantly enhance the paper's impact.

I am open to increasing the score for the paper based on the authors' responses and feedback from other reviewers.

References:

a) Rajaraman, Nived, et al. "Transformers on Markov data: Constant depth suffices." arXiv preprint arXiv:2407.17686 (2024).  

b) Makkuva, Ashok Vardhan, et al. "Local to Global: Learning Dynamics and Effect of Initialization for Transformers." arXiv preprint arXiv:2406.03072 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730657928141

thread_type:
official_review

<thread_object>
creator: reviewer_6zob

<content>
summary:
This paper studied the memorization capacity of transformers in two tasks, namely next-token prediction and sequence-to-sequence prediction. The paper showed that, by construction, a transformer with $\tilde{O}(\sqrt{N}log n)$ layers can memorize $N$ data points of length $n$. The paper further showed that a transformer with $\tilde{O}(\sqrt{nN})$ layers can memorize $N$ data points in the sequence-to-sequence prediction problem. A lower bound is provided to demonstrate the near-optimality of the upper bound.

strengths:
1. The results seem promising and strong, because both upper and lower bounds are provided.
2. The studied problem and method provide interesting insight into transformers.

weaknesses:
1. 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.
2. The paper can have more discussion on the separatedness assumption. For example, how realistic is the assumption?

questions:
Following the weaknesses part, I have the following questions.

1.  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?

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

</thread>
<thread>
creation date:
1730675740876

thread_type:
official_review

<thread_object>
creator: reviewer_zzq4

<content>
summary:
This paper studies the number of parameters and the bit complexity required by a transformer to memorize a dataset of size $N$. In the context of next token prediction, the paper proves that for any dataset with well separated tokens, there exists a transformer with $\tilde O(\sqrt{N})$ parameters+bit-complexity which interpolates it, along with a matching lower bound that $\Omega(\sqrt{N})$ parameters are necessary. In the sequence to sequence setting, the paper proves a similar upper bound with $\tilde O(\sqrt{nN})$ parameters+bit complexity where $n$ is the sequence length, and they prove a matching lower bound for hardmax transformers.

strengths:
- The proof sketch of Theorem 4.1 is clear and well motivated
- In the next-token setting, the paper proves matching upper and lower bounds, and in the sequence to sequence setting they match for hardmax transformers
- The authors provide clear comparisons with prior work throughout the paper

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

- **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.

- minor typo on line 462: viewd -> viewed

questions:
- If I understand correctly, neither of the setting in the paper correspond directly to the standard next-token prediction training paradigm in which the transformer needs to simultaneously predict the $i$th token given $X_{:,< i}$ for all $i$. Corollary 4.1 does not directly apply because of the varying sequence lengths, but it suggests a construction with $\tilde O(N)$ parameters+bit complexity. However, it's not clear if this is optimal since Theorem 4.3 is lower bound against the more general sequence-to-sequence task?

- Could you explain footnote 2 (line 1025)? I assume it's not possible to remove the $\log(r/\delta)$ dependence so is the goal to remove the $\log d$ or $\log C$ dependencies? If so, what is the idea?

- Is it possible to achieve the full tradeoff between bit-complexity and number of parameters? E.g. $N^{\alpha}$ parameters and $N^{1-\alpha}$ bit complexity for any $\alpha \in [1/2,1]$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730724901138

thread_type:
official_review

<thread_object>
creator: reviewer_vzks

<content>
summary:
The paper focuses on the memorization capacity of Transformers in next-token prediction and sequence-to-sequence settings, and shows the scaling in the number of input sequences and their lengths.

strengths:
- This paper is well-motivated and clearly presented, making it easy to follow.
- This paper is technically solid; the construction of a single-layer Transformer that can identify identical sequence ids from token-wise \((r,\delta)\)-separated sequences with near-optimal parameter order is novel. Additionally, the authors provide a parameter lower bound for the seq2seq case when the self-attention layer uses hard attention.
- Furthermore, the paper considers bit complexity, which is often overlooked in related work.

weaknesses:
- 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. 
- 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?
 - The proof for the parameter lower bound in seq2seq cases relies on hardmax activation, deviating from the real-world scenario where the attention layer adopts softmax activation. Is it possible to extend the proof to the case with softmax activation?

questions:
See weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732697143426

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Updates Based on Reviewer Feedback

comment:
Thank you so much for taking the time to review our paper and for your helpful comments. Your feedback has really helped us make the paper clearer and better.
Below, we summarize the revisions and additions made during the rebuttal period based on your responses:
1. **Title Revision**: To enhance precision, we have updated the paper's title to "On the" Optimal Memorization Capacity of Transformers, 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
2. **Proof Sketch of Main Theorem**: We have included more detailed explanations regarding the high-dimensional representations $\tilde{X}^{(1)},\dots,\tilde{X}^{(N)}$ used to construct the feed-forward network in the proof sketch of Theorem 4.1, particularly around eq. (9).
3. **Memorization Capacity with Limited Bit Complexity**: We have added a discussion in Appendix F about cases where the bit complexity of parameters is insufficient for optimal memorization. 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.
4. **Experiments**: To support our theoretical analysis, we conducted experiments on two real-world datasets and one randomly generated datase (details are provided in Appendix H):
    - **Verification of Uniform Attention's Expressive Power**: We trained the model described in Theorem 4.1, that is, a Transformer consisting of single uniform attention layer sandwiched between two token-wise feed-forward networks, on the MNLI and IMDb datasets. 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.
    - **Analysis of Memorization Capacity**: Using the same model, we trained it on MNLI datasets of varying sizes and plotted the minimum model size required for memorization. The results show 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. A plausible explanation for this phenomenon 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$.
    - **Effect of Input Sequence Length on Memorization Capacity**: We trained the model on randomly generated datasets with fixed data sizes but varying input sequence lengths. The results indicate that the minimum model size required for memorization remains unchanged regardless of input sequence length, consistent with the predictions of Theorem 4.1 that the memorization capacity in the next-token prediction setting is nearly independent of the input sequence length $n$.

Your thoughtful feedback not only helped us address the concerns you raised but also deepened our understanding of the memorization capacity of real-world Transformer models. We are deeply grateful for your valuable insights and for giving us the opportunity to refine and strengthen our work.

Thank you once again for your constructive comments.

**Reference**  
[1] Vardi et al., On the Optimal Memorization Capacity of ReLU Neural Networks. ICLR 2021.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies the number of parameters and the bit complexity required by a transformer to memorize a dataset of size. In the context of next token prediction, the paper proves that for any dataset with well separated tokens, there exists a transformer optimal 
 parameters+bit-complexity (both upper and lower bounds.) In the sequence to sequence setting, the paper proves a similar upper bound with, and they prove a matching lower bound for hardmax transformers. All reviewers believe this paper is technically strong, and the AC agrees and recommends acceptance.