On the Optimal Memorization Capacity of Transformers
Authors: Tokio Kajitsuka
Issei Sato
Venue: ICLR 2025 Poster (2025)
Keywords: Transformer
Self-Attention
Memorization
Contextual Mapping
Permutation Equivariance
Primary area: learning theory
TL;DR: We prove that the memorization capacity of Transformers is optimal up to logarithmic factors.

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.

Published as a conference paper at ICLR 2025
ON THE OPTIMAL MEMORIZATION TRANSFORMERS
CAPACITY OF
Tokio Kajitsuka & Issei Sato
Department of Computer Science
The University of Tokyo (kajitsuka-tokii sato ecc. uot.........................
ABSTRACT
Recent research in the field of machine learning has increasingly focused the memorization capacity of Transformers, but how efficient they are not yet well understood. We demonstrate that Transformers can memorize labels with Ö ((vN) parameters in next-token prediction setting for N input sequences of length which is proved to be optimal up to logarithmic factors. This indicates that Trans- formers can efficiently perform memorization with little influence from the input length n owing to the benefit of parameter sharing. We also analyze the mem- orization capacity in the sequence-to-seqqenee setting, and find that Ö(v parameters are not only sufficient, but also necessary at least for Transformers with hardmax. These results suggest that while self-attention mechanisms can ef- ficiently identify input sequences, the feed-forwart network becomes a bottleneck when associating label to each token.
INTRODUCTION
In recent years, the Transformer architecture (Vaswani al. 2017) has played pivotal role in the field of machine learning, becoming indispensable for variety of models in the community. In addition to the original breakthroughs in natural language processing, such as the GPT series (Brown et al. 2020 Radford et al., 2018: 2019), it has been observed that in numerous applications, higher accuracy can be achieved by replacing existing models with Transformers. such the Vision Transformer (Dosovitskiy et al. 2021) in image processing and the Diffusion Transformer (Peebles & Xie 2023) in generative tasks have demonstrated exceptional performances in wide variety of tasks. These examples demonstrate how effective and versatile Transformers for a diverse range purposes.
particular, models
Although the high performance of Transformers has led their widespread use in practice, there ongoing attempts to theoretically analyze what exactly contributes to their superior performance. In particular, important aspect of Transformers is their representational capabilities. Previous stud- ies have explored from variety of angles why Transformers have high expressive capacity and can memorize vast amounts of data (Edelman et 2022 Gurevych et 2022 Takakura Suzuki 2023). For example, it has been shown that Transformers are universal approximators (having the ability to approximate arbitrary functions) (Yun al. 2019) that a particular Transformer config- uration can memorize a given set of data (Kim et al.
2023 Kajitsuka & Sato 2023 Mahdavi et al.,
2023 Madden et al. 2024).
Nevertheless, while various studies have suggested that Transformers are indeed capable of mem- orizing data, our understanding of how efficiently they can do so remains limited. Specifically, not yet fully clear how certain characteristics of Transformers, such as parameter sharing, influ- ence the reduction of model parameters and overall efficiency with respect to their memorization capacity, the minimum size of networks required for memorizing any sequence of a given number of data.
There several key advantages to investigating whether Transformer efficiently memorize data, such as the possibility of gaining a better understanding of Transformer's strengths and pro- viding useful insights for model design and selection. In addition, knowledge of memorization

---

Published as a conference paper at ICLR 2025
efficiency can provide important information for evaluating generalization enror (Belkin al. 2019 Nakkiran et al., 2021). Alternatively, if it turns out that Transformers do not offer a sigmificant efficiency advattage over feed- forward networks, it may suggest that currently widely used Trans- formers may in fact be substitutable for feed-forward networks.
This paper investigates the efficiency of Transformers in achieving data memorization by analyzing the necessary and sufficient model complexity for this task. To be more precise, we establish both upper and lower bounds on the number of parameters needed for memorization in the next-token prediction setting and demonstrate that they are of the same order up to logarithmic factors, thereby showing that Transformers can achieve data memorization with nearly optimal efficiency.
Furthermore, the upper bound on memorization capacity in the next-token prediction setting can be naturally extended to the soeenceet--seueneee setting. This upper bound also proved to be optimal up to logarithmic factors in the sequence-to-sequems setting, at least for Transformers with the hardmax function.
2 RELATED WORK
MEMORIZATION CAPACITY
Research memorization capacity began at least as late the 1960s (Cover 1965 Nilsson 1965
Minsky & Papert 1969). Specifically,NNilssse Nilsson (965) showed that one- one-hidden- hhdden-llyer neural networks with N 1 nodes is able compute any label assignments for N data points. Later, |Bam (1988) exhibited that /ddd] neurons are sufficient for one-hidden-layer neural networks with threshold units to memorize any set of N input-label pairs with the input dimension d, and Huang & Babri [1988 Zhang et al. (2021) extended the results to more general activation functions.
The analysis of memorization capacity is closely linked the concept of the Vapnik-Chervonenkis (VC) dimension. While the memorization capacity of model refers to the minimum size of the model required for memorizing any tuple of N input-label pairs for some N, the VC dimension considers whether the model is capable of shattering, that is, memorizing any possible label assign- ments for some set of N input points, which in turn provides lower bound on the memorization capacity. For example, Goldberg & Jerrum (1995) estimated that the VC dimension of feed- forward network with ReLU activation functions and parameters is at most O(W?) by reducing the network to boolean formula. From this upper bound, it can be inferred that feed-forward network with ReLU activation functions requires least 2?( ((VN) parameters memorize arbitrary N data points. Bartlett et al. (2019) further refined this analysis by examining the behavior of the network function of its parameters and analyzing it layer by layer, and demonstrated that the VC dimension of a ReI.U network with width W and depth T. is O(W log W).
Remarkably, Park et al. (002) proposed construction method under the assumption that the data points are separated by at least 6, showing that a feed-forward network using sigmoid or ReLU activation functions with sub-linear parameter order (
2/3 + log 5) can memorize N data points. Later, Vardi et al. (2022) demonstrated that, under similar assumptions, a ReLU network with O(VN log N) parameters suffices for memorizing arbitrary N data points. This result is opti- mal up to logarithmic factors, as it matches the lower bound 2 ((vN) implied by the VC dimension discussed above. Note that the assumption that data points are well separated is crucial achieve sub-lined memorization capacity; in fact, it has been shown that at least (N 1)/2 parameters are required to memorize arbitrary N distinct data points without such separation (Sontag, 1997). Ad- ditionally, (Segell (2024) proved that ((N) parameters are necessary for memorizing N data points when the separation between data points is exponentially small with respect N.
Memorization capacity is not only theoretically intriguing but also practically significant. As the model size increases, classical learning theory predicts that the training error decreases while the generalization error follows a U-shaped curve. However, recent observations of the double descent phenomenon (Belkin et al. 2019 Nakkiran et al. 202I) revealed that after achieving zero training loss. the generalization error begins to decrease again. Analyzing memorization capacity helps iden- tify the critical model size at which this shift occurs, providing valuable insights into the dynamics of model performance.
2

---

Published as a conference paper at ICLR 2025
Table 1: Comparisons between results and related work regarding the memorization capacity of Trans- formers. The variable w in bounds presented by Madden al. (2024) represents the vocabulary size, or the number of distinct word vectors that appear in input sequences.
Paper
Setting
Input
#layers
Upper bound
Lower bound
Kim et al.
(2023)
seq-to-seq
(r, -seearrad
token-wise
Ö(n + VnN)
Ö(n + VnN)
Mahdavi
linearly
al. (2023)
next-token
independent
1
O(d2N/n)
Kajitsuka &
token-wise
Sato (2023)
seq-to-seq
(r, 5--sparaatedd
1
O(dnN + d
Madden
with positional
next-token
1
O(WN)
(((NN)
et al. 2024)
encoding
Ours
next-token
token-wise
Ö(VN)
Ö(VN)
(((NN)
(r, 8) seeparttd
seq-to-seq
(r, 8) -spparaeed
token-wise
Ö(VnN)
Ö(VnN) 5 (v log(nN)
EXPRESSIVITY OF TRANSFORMERS
One of the foundational studies the representation power of afaansffommrss is the work by Yun al. (201)), who demonstrated that Transformers are universal approximators. Their proof already incor- porates the idea of constructing contextual mapping from data points to contexts and linking these context ids to labels. Kim et al. (2023), whose work is most closely related to our work, improved their contextual mapping approach and demonstrated that this mapping, constructed using 2n layers of self-attention for N input sequences of length n, allows for memorization with Ó( + vnN) parameters under the same assumption that data points are well separated as in Park et al. (2021); Vardi al. (2022). Later, Kajitsuka & Sato (2023) showed that a single-layer, single-head Trans- former already possesses memorization capacity under the same assumption, while self-attention with hardmax does not. In contrast to the studies mentioned above, Mahdavi et al. (2023) demon- strated that under the assumption that data points are linearly independent, multi-head attention with H heads and embedding dimension > n. can memorize (2(n)) data points in a next-token prediction like setting. Madden et al. (2024) proved upper and lower bounds on the memorization capacity of one-layer Transformers with parameters of infinite precision in the next-token predic- tion setting. Chen & Zou (2024) investigated the behavior of Transformers with varying depths, and specifically demonstrated that a single-layer Transformer can achieve memorization input se- quences are sufficiently gero-padded. However, they noted that their objective was not explore efficient constructions. The comparisons between our results and related work are summarized in Table T Note that all the papers listed here that investigate single-layer Transformers assume either infinite parameter precision or do not consider the bit-length required to represent parameters.
In addition to memorization capacity, there are studies highlighting other perspectives on Transform- ers, including their function approximation capacity (Gurevych et al.] [022] Takakura & Suzuki
2023
Jiang & Li| 2024), and their ability to efficiently represent sparse functions (Edelman et al,
2022
Bhattamishra et al. 2023 Sanford et al. 2023
lrrauger& & Tewari 2024 Wang et al. 2024b).
PRELIMINARIES
3.1 NOTATION
We denote vectors and matrices by bold lowercase and uppercase letters, respectively. Given vector v, we denote its i-th element as v, Given a matrix A, we denote its i-th row as A its j-th column A., and the element at position (i,j) Aij. For natural number E N+ we use [m] to denote the set f1....,m, m). In the context of the self-attention mechanism, we use 0g to represent the column-wise softmax function. Specifically, for matrix A xb [A] Ra is calculated by a [s[A), = exp(Ar,)) // Ei= exp(Ar.)). Likewise, we use Hr to denote the column- wise hardmax function. Note that if there are multiple values in column, its outputs normalized
3

---

Published as a conference paper at ICLR 2025
so that they sum up to 1. Mathematically, for matrix A Raxb [H[A] R is calculated as follows.
[1/12] if A., = max Akj.
Th [n[Akk =
(1)
(0 otherwise,
where 1 == arg max; Ak. = (k' E [a] Axr. maxi Aky) for any j E [b]. We use on to denote the ReLU activation function, that is, on(z) == max(0, 2). Unlike as and Øn, on always applied element-wise. regardless of whether the input is vector or a matrix. For any natural number x E I. 1. (x) E N represents the sequence of bits from the i-th bit to the j-th bit (counting from the left) of z, interpreted
natural number. For vector RR, the L norm of is denoted by ||ll2 = Σi=1 of. We use standard asymptotic notation. Specifically, f(n) (g(n)) indicates that the function f grows most fast as g for sufficiently large and f(n) (g(n)) represents
that I grows at most as fast as 9, up to logarithmic factors. Likewise, f(n) = [2(g()) means that the function f grows least as fast g for sufficiently large n. ((n E g(n) means that there exists a positive constant e such that f(m) S cg(n) holds.
In this paper, we basically use n to denote the length of an input sequence, N denote the number of input sequences, C' to denote the number of classes, and d to denote the dimensionality of each token. Additionally, index is typically used to refer to the position of input sequences, while index k is used to refer to the position of the token within an input sequence.
3.2 TRANSFORMER BLOCK
In this subsection, we introduce the architecture of Transformers (Vaswani et al. 2017). We ba- sically follow the notations by Kim et al. (2023). Transformers are defined by stacking multiple Transformer blocks, each of which consists self-attention layer and a feed-forward layer.
Self-attention layer: Given input sequence Z _m=xn, the output of a self-attention layer (((A). m + mm at block E [T.] is calculated by
((A)) (Z) = Z
w (0) w Zas (wix))))
(V)
(Wil))
2)
E I m
where (V)) w(K). (( Rm and (0) Rmxs are value, key, query and projection
matrices at head h E [H] with head size s, respectively.
Feed-forward layer: The output H mmxx of the self-attention layer block / is then passed to the feed-forward layer, which performs the following loken-wise operation:
J(PEF ( H. + W! OR W H + b69] + b1th T E [n]).
(3)
where (1)) E Roxm and W(22) R are weight matrices with hidden dimension q, and b(1)
RS and b," 1m are bias terms.
Using the self-attention layer and the feed-forward layer, the Transformer block Fi Rmxn
Rmxa at block 1 [T.] defined as composition of these two layers, that is, Fi == ((FF) o ((SA).
and the whole architecture of the Transformer. N pdxa + R Ixn is expressed by
N == Eout FL ... Fi Ein.
(4)
where Enn Rg + m2 and Eout g + _ x are token-wise linear mappings.
In a Transformer, the width is determined by the combination of self-attention layers and feed- forward layers. According to the definition proposed by Kim et al. (2023), the width of the Trans- former model is defined as max(m, sH, q). We define the depth Transformer by the number of blocks L.
Remark 3.1. The use of in/out token-wiss linear mappings comes from the fact that Transformer blocks by definition have the same input and output dimensions. The token-wise linear mappings can be removed at the cost of a linear dependence of the number of parameters required for memo- rization the embedding dimension d.
4

---

Published as a conference paper at ICLR 2025
3.3 BIT COMPLEXITY
In this paper, consider not only the number of parameters but also the number of bits required to represent the model. Specifically, we adopt the definition bit complexity proposed by Vardi al. (2022). According this definition, the bit complexity of parameter is defined as the number of bits needed to represent that parameter. The bit complexity of a model is then defined as the maximum bit complexity among its individual parameters. It important to that by multiplying the bit complexity of the model by the number of parameters, we can estimate the total number of bits required represent the entire model.
MEMORIZATION CAPACITY OF TRANSFORMERS
In this section, we state the main theorems of this paper regarding the optimal memorization capacity of Transformers. Section 4.1 defines the memorization capacity of Transformers and discuss the main challenge behind this concept. In Sections 4.2 and 4.3) we provide upper and lower bounds on the number of parameters required for Transformers to achieve memorization the next-token prediction setting and the sequence-to-sequee prediction setting, respectively.
4.1 PROBLEM SETTING
The aim of this study is to analyze the memorization capacity of Transformers. Informally, mem- orization capacity refers to the minimum size model that can memorize a specific number of arbitrary data points. To be more precise. let X' and y be input space and output space, respectively. Then, given N input-label pairs (x(1), ( (X(M), (()) X x V we are interested in the model complexity of a model f X' + y such that ((x y( holds for any i E [N]. In the case of Transformers, the input space consists input sequences made up of n tokens, each of which is d-dimensional vector. Hence, we define the input space X' as X' == Rdxn.
Without any assumptions on the input data, has been shown by |Sontag (1997). that a linear order of parameters is required to memorize arbitrary N data points. To achieve sub-linear memorization capacity, in this paper, we assume that the data points are well separated, a common assumption in prior work PPar et al. 2021 Vardi al. 2022, Kim et al. 2023 Kajitsuka & Sato 2023 Siegel 2024). In the case of Transformers, this concept is formalized as token-wise (r, d) o---s-paratednes Kim et al. 2023
Kajitsuka & Sato 2023).
(NN)) E Ruxx be N input sequences, Assumption 4.1 (Token-wise separatedness). Let ( X(),...,,
each of which consists of n word vectors with its dimension d. Then, we say that ( (I)).... (N)
token-wise (r, -ssepartte for some r, if the following two conditions are satisfied:
1. for every i [M] and [n), X (6) la S holds.
2. for every i.j [N] and k, [n], either x xx or 1x x holds.
The notion of tosen-wise (r, 6) separateednss ensures that the word vectors appearing in the input sequences have an r norm of at most r, and are separated by at least 5 1,2 norm from each other.
The main difficulty of memorization with Transformers, compared with in the fact that tokens with identical values do not necessarily correspond to the same label. Instead, is crucial to capture the context in which each token appears within the entire input sequence. In Transformers, while feed-forward layers operate on individual tokens, self-attention layers are the only place that enables interactions between tokens within the input sequence. Therefore, the central question we consider in this paper is:
networks, lies
how efficiently can self-attention layers capture the context of tokens?
To explore this issue, analyze both upper and lower bounds on the number of parameters quired for memorization with Transformers two settings: next-token prediction and sequence-to- sequence prediction.
5

---

Published as a conference paper at ICLR 2025
4.2 NEXT-TOKEN PREDICTION SETTING
4.2.1 UPPER BOUND
First, given N input sequences of length consider the problem setting in which Transformer memorizes labels corresponding to the n-th token of input sequences. We call this task next- token prediction setting. In this problem setting, how many parameters does Transformer archi- tecture require? Surprisingly, O(vN) is sufficient, that is, the input length has almost no effect on the number of parameters required for memorization, as the following theorem states.
In the next theorem, F(FF) and ((F)) represent feed-forward networks of arbitrary depth, unlike eq. (3). which is limited two layers. Note that deep feed-forward networks can also be imple- mented with standard Transformers, by setting the projection matrix of the self-attention layer in each block to zero. Furthermore, the assumption of consistency on labels in Theorem is a neces- sary requirement to perform memorization with Transformer, due to its permutation equivariance.
Theorem 4.1 (Next-token prediction). Let (()),, (X()), y()) E Rdxx x [C] be sequence of input-label pairs such that
1. (x ) ..... (X(N),
,y(N) are consistently labeled, in the sense that for any i,j
[N]. we have ()) =z(") if
X (8) X (
and X x() x(() up to permutations.
(5)
2. X(I)
(N) token-wise (r, separated for some 2 1 and 0 <
Then, there exists Transformer N [Rdx Rn with width 14 and depth Ö(VN) that memorizes the dataset, that is,
N ((x1) = Eout ((FP) (SA) F((F) Ei (x(()), -0(()
(6)
holds for every i [N], as long n, C,r8-1 NO N
The formal statement of Theorem 4.1 and its proof can be found in Appendix B.1.1
Remark 4.1 (Deep sets), In fact, Theorem 4.1 can extended to deep sets (aheer et al. 2017),
which is a popular architecture model mapping from sets to labels. For details on this result, see Appendix D
Remark 4.2 (Embedding layer). A similar result holds for Transformer with an embedding layer. However, in this case, the presence of an embedding layer introduces a dependency on the size of the vocabulary, which may result in non-optimal order of parameters in the worst-case scenario. Details regarding this discussion can be found in Appendix E
Remark 4.3 (Dependence on d). The Transformer architecture defined by eq. (4) includes token-
wise linear muppings Ein R2d + I.m and Evut. I - Rd, leuding to Ö(d + VNN parameters for a Transformer with depth Ö(VN) and width 14. As noted by Vardi et al. 2022) and Kim et al. 2223), this dependence the dimension d is unavoidable preserve the information of the imput tokens.
Theorem 4.1 demonstrates that as long as the dimension is of the order d - =Ö(VNN, the Trans- former with a single self-attention layer can memorize N input sequences and their labels for next- token prediction with Ö(vN) parameters, showing negligible dependence on the input length n. In contrast, accomplish the same task with a feed-forward network. necessary use parameters to retain the information of the input sequence in Rd This illustrates a significant efficiency advantage of Transformers over forw-fowwar networks, thanks to parameter sharing.
4.2.2 PROOF OUTLINE OF THEOREM 4.1
Here provide outline of the proof of Theorem 4.1 See Appendix B.I.I for its full proof.
The proof strategy is to construct a contextual mapping as in Yun et al. (2019) Kim et al. (2023) and Kajitsuka & Sato (2023), and then construct mapping from the context id to the label. Here,

---

Published as a conference paper at ICLR 2025
contextual mapping is function used to distinguish tokens in each input sequence with the following properties:
Definition 4.1 (Contextual mapping). Let x ....... X Rdxn be input sequences. Then, map CM Rdxa + R" is called an (r, -conttextl mapping if the following two conditions hold:
1. For any e [N] and & [n], () S holds.
2. For any í, E [N] and k, E [n] such that x * X or X + x up to permuta- tions, [CM( CM ) 2 8 holds.
In partieular, CM(( (+)) is called a context id of the k-th token in X
Intuitively, the two conditions above ensure that the contextual mapping is injective from "distinct" data points to scalars. If such mapping can be constructed, then mapping from context ids to labels can be realized using feed-forward network with Ö(vN) parameters, as shown by Vardi et al. (2022). In particular, if we can associate each distinct input sequence with a unique value, referred to as a sequence id, then the context id of, for example, the k-th token in X can be constructed from the sequence id of (+) and the token vector X Therefore, the primary focus of our proof is on how to construct mapping from each input sequence to its sequence id using feed-forwe forward network and single self-attention layer.
From a high-level perspective, our goal is construct feed-forward network Φ Rd R with Ö(VN) parameters such that the sums
Σ
x (N)
Rm
are well-separated
The sum En of (i E [N]) is then used as the sequence id of (1)
Crucial observations for constructing @ with Ô( parameters
follows.
1. To distinguish input sequences,
sufficient to focus on at most distinct word vectors. More precisely, given N input sequences, there are most N distinct word vectors such that the input sequences can be identified by counting occurrences of these N words UN C Rd (Lemma B.1.
A {.......
2. Although a feed-forward network requires 2( ((VN) parameters to memorize N data points and their labels (Goldberg & Jerrum 1995). network that outputs zero for additional data points not among N data points can he constructed without significantly affecting the order of the parameter count (Lemma C.1). Together with the first observation, all we need is to construct feed-forward network Ø Rd + E such that
Σ
Σ
(8)
1= EA
1= EA
are well-separated.
3. The final key observation is that rather than directly constructing o, we first consider the dinensiinn representation. Concretely, given arbitrary bijection g A + [N].
map each input sequence (+) to high-dimensional vector X IR follows:
"(X!2)
where ,(x())) E 1N is a one-hot vector with only in the (X ) position. While x (N) distinct from the first observation and suitable candidates for sequence
For simplicity, here that the input sequences (1) (N) distinct up permutations.

---

Published as a conference paper at ICLR 2025
ids, it requires (2)N parameters express these N-dimensional vectors with ferd-forward networks. This problem can be circumvented by compressing the nigh-dimensionll vectors
into scalars using adequate vector and we define ø by o(x) == u ((())
To ensure that a feed-forward network with Ö( parameters can indeed implement the function s we need to carefully analyze how separated the compressed versions of the high-dimensional rep- resentations (1) (N) are. Detailed proof of this implementation is provided in Lemma B.3
4.2.3 LOWER BOUND
In this subsection, we evaluate the minimal model complexity required for memorization with Trans- formers in the next-token prediction setting determine how close to optimal Theorem 4.1 is.
First, notice that the model obtained in Theorem 4.1 is optimal, in terms of bit counts. Remark 4.4 (Optimality in terms of bit counts). As previously discussed in Remark the Trans- former model obtained in Theorem 4.1 has O(VN) parameters long as ÖV.N). On the other hand, the bit complexity of the model Ó(log d VN) (see the formal statement in Ap- pendix B.I.I.. Therefore, if d = Ö(.N, the total number of bits required to represent the model Ö(N). Given that there are 2 possible label assignments for N distinct data points with binary labels, Ó(N) bits are optimal up to logarithmic factors for this setting. The more general case where bit complexity is restricted to ÓN*) for some e E 1/2] is discussed in Appendix
Having established the optimality in terms of bit counts, we now turn to evaluating how efficient the number of parameters the Transformer model considered in Theorem 4.1 is. The next theorem provides lower bound on the number of parameters required for memorization in the next-token prediction setting.
Theorem 4.2 (Lower bound). Suppose Transformer N R + IR" defined by eq. (4 can
shatter a set of N distinct input sequences (1), ,.... X Rdxx with x(9 =x for any i E [N] and ls E [n]. in the sense that for any label assignments y(), (".,,..((N),, E {0,1), there parameters with which N(X =)()) holds for any E [N). Then, the Transformer N has at least (VVN) parameters.
The proof of this theorem can be found in Appendix B.1.2 This result indicates that the Transformer model described in Theorem 4.1 is also optimal in terms the number of parameters. Specifically, since memorization in the next-token prediction setting requires the ability to distinguish N input sequences, this result provides the following crucial insight.
Transformer with a single layer of self-attention already possesses necessary and sufficient expressive capacity to identify input sequences.
In fact, as indicated in the proof outline Section 4.2.2 we only employ the self-attention layer as an averaging operation in the model obtained by Theorem 4.1 The observation that simple averag- ing provides sufficient representational power has been confirmed experimentally by Yu et al. (2022) with their PoolFormer architecture. In this paper, we provide theoretical support by demonstrating that a Transformer with just a simple averaging operation already has optimal memorization capac- ity. We also conducted experiments on two real-world datasets and a randomly generated dataset, confirming that even a single layer of self-attention, as averaging, possesses sufficient representa- tional capacity for memorization. For further details, please refer to Appendix H
4.3
COQURENEETO-SEQUFEEE PREDICTION SETTING
Next, consider the problem setting in which each token in input sequence is assigned some label and a Transformer memorizes them all. We call this task a sequence-to-sequeene prediction setting, or seq-to-seq prediction for short.
is readily apparent that the seq-to-seq prediction can be regarded as rearranging the input sequence so that each token is placed the end of the sequence, and then performing next-token prediction on n.N input sequences obtained in this way. From this observation, we have the following corollary from Theorem 4.1
8

---

Published as a conference paper at ICLR 2025
Corollary 4.1 (Seq-to-seq prediction). Let (x(1), y (x(Ny(N) Raxn [C]" be sequence of input-label label pairs such that
1. (x(), )) (X(M), y) consistently labeled, in the sense that for any i.j
[N] and k. I E [n], we have v/2
= yy if
and X x = =xxi up to permutations.
(10)
2. ( X (N) are token-wise (r, separated for some r 2 1 and 0 < 1.
Then, there exists a Transformer N R - R with width 14 and depth Ö(vnN) that memo- rizes the dataset, that is,
N (x (xx())=== Eout F(FF) F(AA) F(FF) Ei (x(1)
Ukk
(6)
(11)
holds for every i [N] and [n), as long as C, -- (nN)o()) as n.V - CC.
Remark 4.5 (Sparse Transformers). While Corollary 4.1 demonstrates that Transformer with single-layer self-attention can achieve memorization in the seq-to-seq prediction setting, it inevitably requires (nn²) computational complexity due to the self-attention mechanism. In line with recent efforts to improve the scalability of Transformers by making attention maps sparse (Zaheer et al. 20200 Yun et al. 2020). using two self celf-attention layers and appending additional token to the input sequence allows us to achieve the same behavior with an (n)) connections without affecting the order of parameter counts. This idea of aggregating global information into the additional token has gained interest in recent studies (Darcet et al. 2023: Wang et al. 2024a).
This corollary shows that at least ÖVvnN) parameters with bit complexity Ö(vnN) are enough to memorize N input sequences of input length n. The next question is: is this order optimal for the seq-to-seq prediction setting? As in the case of next-token prediction setting (Remark 4.4). we can leverage a similar argument to show that this is optimal, least in terms of bit counts.
Remark 4.6 (Optimality in terms of bit counts). Ifd = ÖVVnN), the construction by Corollary 4.1 uses Ö( Ö(VnN) parameters with bit complexity Ö(vnN) to memorize N input sequences of inpus length n, which amounts to Ö(NN) bits. If all word vectors in input sequences are different, there are 2nN binary label patterns. Therefore, to memorize such patterns, the number of states of the model must be at least 2nN, which means that log 2 = nN bits are required.
Unlike the next-token prediction setting, it challenging analyze the optimal lower bound on the number of parameters necessary to memorize N input sequences with input length for the seq-to- seq prediction setting, mainly due to the presence of the softmax function. However, we partially answer this question by considering a Transformer that uses not the softmax function, but instead the hardmax function, often viewed as an approximation of the softmax.
More rigorously, we introduce the following self-attention layer with the hardmax function, which we call the hard attention layer. For each block E [L] and its input Z E mx the hard attention layer at block / calculates
H J(HA) (Z) = Z Σ
( Zoh
( kl
(w hi
z) T (W/l (w(()2)
2)
RmxX
(12)
A=1
1 for its definition), and hl w((k) w(( Raxm and hi mRmxx are value, key, query and projection matrices head [H] with head size 8, respectively. worth noting that simple averaging operation can also be implemented using hard attention layer by setting key and query matrices to zero.
where 0H Raxx_ [0, 1 |mxx the column-wise hardmax function (see eq.
With this definition, we demonstrate that the number of parameters by Corollary 4. is actually optimal up to logarithmic factors, at least for Transformers with the hardmax function. To state the theorem, let W be the number of parameters and 0 R be vector of all parameters of the Transformer. We also denote by No the Transformer to emphasize the presence of the parameter vector 0.

---

Published as a conference paper at ICLR 2025
Theorem 4.3 (Lower bound). Ne R - RR" be Transformer defined by eq. 4 with self- attention layers replaced with hard attention layers (eq. (12). In addition, suppose No can shatter a set of N imput sequences X(1) ,..,..., x.,x
E IRdxn with X + I. for any i,j E [N] and k,l E [n]] (i + j or lo * D, in the sense that for any label assignments y yy.) E 10,1)" there is parameter vector 0 E IRW such that
No y
(13)
for any i [N]. Then, the Transformer has at least W = S2 (V(V aN log(nNN
parameters.
The proof of Theorem builds on the approach used by |Brrlett et al. (219 to evaluate a lower bound on the VC dimension of feed-forward networks. Specifically, considering Transformer as a function in variable its parameter vector, we partition the parameter space of the Transformer in such a way that, within each cell of this partition, the function can be expressed as a polynomial in terms of its parameters, and then evaluate the number of cells and the properties of the polynomials within those cells.
The key novelty of the proof lies in the analysis of how parameter sharing and the hardmax function affect the memorization capacity of Transformers. Parameter sharing in Transformers allows the model to effectively behave like a network with its width scaled by the number of tokens, without actually increasing the number of parameters. However, the proof shows that merely increasing the width by factor of does not lead to fundamental improvement in the memorization capacity of the Transformer. The full proof of Theorem 4.3 can be found in Appendix B.2
Theorem 4.3 demonstrates that the number of parameters in the model from Corollary 4.1 is within logarithmic factors of the optimal lower bound. addition, it provides another crucial insight. As shown in the next-token prediction setting, Transformers can identify N input sequences with Ö(VN) parameters and single self-attention layer, which implies that they are capable of capturing the context of each token. In contrast, the memorization capacity in the seq-to-seq setting is prov- ably lownrebounded by ((VnN), which includes an additional /n factor compared to the Ö( bound in the next-token prediction setting. Therefore, in the seq-to-seq prediction setting, the pri- mary bottleneck is not the contextual mapping of tokens, but rather the feed-forward layers' capacity to map this token-level contextual information to labels.
We conclude this section by leaving an open problem. Based on Theorem 4.3 for Transformer to memorize N sequences of length n with of Vn.) parameters, it is necessary to exploit the unique characteristics of the softmax function, rather than using it as approximation of hardmax.
Open Problem. Does a Transformer using the softmax function require Q(vnN) parameters to memorize N input-label pairs (X11, ())),...,,(X!) ()) E Rdxn x [C]"? Alternatively, is possible to construct a Transformer with of o(vnN) parameters that can shatter arbitrary N token- wise (r, 0) -spparadd input sequences in the seq-to-seq setting?
5 CONCLUSIONS
In this paper, showed that in the next-token prediction setting, Transformer with Ö (( parameters can memorize N input sequences of length n and their labels, which we showed to be optimal up to logarithmic factors. This result indicates that Transformers can perform next-token prediction with almost no impact from the length of the input sequence. Notably, its proof indicates that even single self-attention layer used as an averaging operation possesses sufficient expressive seruenee power to distinguish between input sequences efficiently. Furthermore, we demonstrated that in the
seq-to-seg prediction setting, Ö(VRN) parameters are also sufficient, and we proved that this is optimal up to logarithmic factors, at least for Transformers using hardmax. These findings highlight that the main bottleneck in seq-to-seq prediction tasks lies in the feed- feed-forward layers' capacity to map each token to the corresponding label.
Given that single layer of self-attention as an averaging operation suffices for distinguishing input sequences from memorization perspective, our results suggest that the advantages of using self- attention might rather lie in the perspectives of optimization and generalization.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by JSPS KAKENHI Grant Number JP24H00709. We would like to thank all the collaborators and anonymous reviewers for constructive discussions.
REFERENCES
Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama, Optuna: A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, 2019.
Peter L. Bartlett, Nick Harvey, Christopher Liaw, and Abbas Mehrabian. Nearly-tight ve-dimension and pseudodimension bounds for piecewise linear neural networks. Journal of Machine Learning Research, 20(63):1-17, 2019. URL http: http://aa...oooppppppe/////////1
Eric B Baum. On the capabilities of multilayer perceptrons. Journal of Complexity, 4(3):193-
215, September 1988. ISSN 0885-064X. doi: 10. 10.10160088500444X88900000000000 URL https:
/www. www.sscciencediret.tcommm www ott.aaaccc.............aaaaaaaaa
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine-
learning practice and the classical bias variance trade-off. Proceedings of the National Academy
of Sciences, 116(32):15849 -5811158999-154,,,.1- August 2019. doi: 0..1073/pnas.190300 1903070116. URL / // /ww. pnas. pogg/di. full/lll/000 1.00..1100 1073/ 1073/pnas: pnas. 1903070116 of the National Academy of Sciences.
Publisher: Proceedings
Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity Bias in Trans- formers and their Ability to Learn Sparse Boolean Functions. In Proceedings of the 61st An- nual Meeting of the Association for Computational Linguistics (Volume 1: Longg Papers), pp- 5767-5791, Toronto, Canada, July 2023. Association for Computational Linguistics. URL https://acce
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Merbert-Vosss. Gretchen Krueger. Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Rad- ford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners In Ad- vances in Neural Information Processing Systems 2020, volume 33, pp. 1877-1901. Cur- ran Associates, Inc., 2020. URL htttps://papers..rtt 1 4457c0d6bfcb4565 41 81 wwt.ffa8ccc220 Absttraaaaabbbrrttt..... html
Xingwu Chen and Difan Zou. What Can Transformer Learn with Varying Depth? Case Studies on Sequence Learning Tasks. In Proceedings of the 41st International Conference on Machine Learning, pp- 7972- PMLR, July 2024. https::::: http:pss/pppcoeeeinngee .mlr warrr....///// ISSN: 2640-3498.
v235/chen24bp html
Thomas M. Cover. Geometrical and statistical properties of systems of linear inequalities with applications pattern recognition. JEEE Trans. Electron. Comput., 14:326-331 1965. URL
https: /pppi sewant htt.://cc.....aa.oooooooooaa.aa.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision Transformers Need Registers. In The Twelfth International Conference on Learning Representations, October 2023.
URL httpss://oppeeeeiieeeeee wwt.:iiwww.w.eeeeeee.emeeeeeeeee http://pppaareeeeeeeeeeeeeorrree 2andddinnnnnniinn
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In The Ninth International Conference on Learning Representations, January 2021. URL https: /ttp:::///perrvvvve lew net nwt. hett:/orrumiiie>.crrdd.

---

Published as a conference paper at ICLR 2025
Benjamin L. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive Biases and Variable Creation in Self-Attention Mechanisms. In Proceddings of the 39th International Conference on Machine Lecamingg pp. 5793-58311 PMLR. June 2022. URL htllaataa///prcceennnee .mlr.
htmml
ISSN: 2640-3498.
Paul W. Goldberg and Mark R. Jerrum. Bounding the vapnik-chervonenkis dimension of concept
classes parameterized by real numbers. Machine Learning, 1802)11111148 1995. ISSN 1573-0565. doi: 10. N..107/BF00993408.
Iryna Gurevych, Michael Kohler, and Gözde Gül Şahin. On the Rate of Convergence of Clas- sifier Based on a Transformer Encoder. IEEE Transactions on Information Theory, 68(12); 8139-8155, February 2022. ISSN 1557-9654, doi: 10.10099TT22022..1197777 URL [RLhtsps:
//ieeexplore leee. argg//ocument/9857855 Conference Name: IEEE Transactions on Information Theory.
Guang-Bin Huang and H.A. Babri. Upper bounds on the number of hidden neurons in feedforward networks with arbitrary bounded nonlinear activation functions. IEEE Transactions on Neural Nenworks, 91)::224-23 22,,,, 1998. doi: 10.1 10.1109772.655445.
Haotian Jiang and Qianxiao Li. Approximation Rate of the Transformer Architecture for Sequence Modeling, February 2024. URL http://arxixs http::://ax.v..r//abbbbss 18475 18475
arXiv:2305.1847, [es].
Tokio Kajitsuka and Issei Sato. Are Transformers with One Layer Self-Attention Using Low- Rank Weight Matrices Universal Approximators? The Twelfth International Conference on Learning Representations, October 2023. URL hhttp:://openeeet
net. net/forum?id
hJnky5K944
Junghwan Kim, Michelle Kim, and Barzan Mozafari. Provable Memorization Capacity of Trans- formers. In the Eleventh International Conference on Learning Representations, February 2023.
URL htttps://ppenreu//
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https::/
www.eee/ffooooe formm/ ddd.dkk-eee Bkg (kkgRiiyY)
Andrew L. Maas, Raymond E. Daly, Peter Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Dekang Lin, Yuji Matsumoto, and Rada Mihalcea (eds.), Proceedings of the 49th Annual Meeting of the Association for Compu- fational Linguistics: Human Language Technologies, pp. 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL http:://ccaatthllogg org/
P11 P11-1015
Liam Madden, Curtis Fox, and Christos Thrampoulidis. Next-token prediction capacity: general upper bounds and a lower bound for transformers, September 2024. URL http:://rxxixv
org/abs/ 2405.4455..0 13718 arXiv:2405.13718 [cs, math].
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization Capacity of Multi-Head Attention in Transformers. In The Twelfth International Conference Learning Representations, October 2023. URL httttttpss:/peeeeeteeee/eeeeeeee
.nt wteee//ooumiidddM33333
Marvin Minsky and Seymour Papert. Perceptrons: An Introduction to Computational Geometry. MIT Press, Cambridge, MA, USA, 1969.
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: where bigger models and more data hurt. Journal of Statistical Mechunics: Theory and Experiment, 2021(12):124003 February 2021. ISSN 1742-5468, doi: 10. 1088/ htlLttts::1 lisher: IOP Publishing and SISSA.
7742-5468/ac3a744
doi Dtt.:ptt.//.a.......o......././. 020744422546/2222
Pub-
Nils J. Nilsson. Learning Machines. McGraw-Hill, New York, 1965.
12

---

Published as a conference paper at ICLR 2025
Sejun Park, Jaeho Lee, Chulhee Yun, and Jinwoo Shin. Provable Memorization via Deep Neural Networks using Sub-linear Parameters. In Proceedings of Thirty Fourth Conference on Learning Theory, pp. 3627-3661. PMLR, July 2021. URL https: http:::://proceecings. mmlr. . press/
v1 v134/park211 html
ISSN: 2640-3498.
William Peebles and Saining Xie. Scalable Diffusion Models with Transformers. In-
ternational Conference on Computer Vision, pp. 4195-4205, 2023. URL https:
//openaccess. thecvf. www. http:///c.cccoee////////////////
_esls___caalll__aaa___a__bbaaa_s with wwt.wwihh__h_aeeeeeee Twans cateeeesrseccCC ICCV Saceccvv_22___ppee__ 2023_paper hamperr.mmll .hmml
Diffusion_Monn
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving Language Under- standing by Generative Pre- Training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language Models are Unsupervised Multitask Learners. 2019.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational Strengths and Lim- itations of Transformers, June 2023. URL http:///rxxt..rr//bb/// 02896
arXiv:2306.02896 [cs, stat].
Jonathan W. Siegel. Sharp Lower Bounds on Interpolation by Deep ReLU Neural Networks at Irregularly Spaced Data, February 2024. URL httt:://axii...e../////// http///xxvv..ogbb////aaa/a 00834
arXiv:2302.00834 [cs, stat].
Eduardo D. Sontag.
Shattering All Sets of 'K' Points in "General Position" Requires ( (k-1) 1)/2 Parameters. Neural Computation, ((2:3337-348, February 1997. ISSN 0899-7667, 1530-888X. doi: 10.1162/ncco.19997 2/337-340/6035
URL https: htmpss//ddrrec wwteeetttiitt.........
wwtptreettttttttteeeeeeeeeeeee edu. www.ee//ecc/aaiceeeeeee/eee
Shokichi Takakura and Taiji Suzuki. Approximation and Estimation Ability of Transformers for Scquence-to-SSquence Functions with Infinite Dimensional Input. In Proceedings of the 40th
International Conference on Machine Learning, pp. 33416-33447. PMLR, July 2023. URL
mlr www.ppeessvvl0 http./ppee//////taaaaauate. tttpt/taaaaaaaaaaaa html
ISSN: 2640-3498.
Jacob Trauger and Ambuj Tewari. Sequence Length Independent Norm-Based Generalization Bounds for Transformers. In Proceedingy of The 27th International Conference Artificial In- telligence and Statistiics, pp. 14405-1413 PMLR, April 2024. URL /tttttps//prroeeiinggss mlr. wttp:p//s///////reree2222
html
ISSN: 2640-3498.
Gal Vardi, Gilad Yehudai, and Ohad Shamir. On the Optimal Memorization Power of ReLU Neural Networks. In The Tenth International Conference on Learning Representations, January 2022.
URL
www.:/etttttteeeeeeeeee forum? id MdMMMTttjeeMTTT
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances Neural Information Processing Systems, volume 30. Curran Asso-
http:://ppaaa...oooo....aer///// ciates. Inc., 2017. URL hatpos://papes..ni nips. cc httpppper/21///aas. html
Edward Wagstaff, Fabian B. Fuchs, Martin Engelcke, Michael A. Osbore, and Ingmar Posner. Universal Approximation of Functions on Sets. Jourmal of Machine Learning Research, 23(151); 1-56, 2022. ISSN 1533-7928. URL UtRL:tt::////0 wttp/pappss//////111111--
html
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding. September 2018. URL httpttss//ppenneeiiieeeee
net htt://ormmmmmm fttt
Feng Wang, Jiahao Wang, Sucheng Ren, Guoyizhe Wei, Jieru Mei, Wei Shao, Yuyin Zhou, Alan Yuille, and Cihang Xie. Mamba-R. Vision Mamba ALSO Needs Registers, May 2024a. URL
http://arrii..oorr/a//////////// htrgg/bb//244m 14858 arXiv:2405. 14858 [cs].
13

---

Published as a conference paper at ICLR 2025
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D. Lee. Transformers Provably Learn Sparse Token Selection While Fully-Conmected Nets Cannot. In Proceedings of the 41st Interna- tional Conference Machine Learning, pp. 118844551122 PMLR, July 2024b. URL https:
//proceedin ngs. ml www.prres//ee...
saress/v2355 v235/ 20Tess//2335wwag222cc.m html ISSN: 2640-3498.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for Marilyn Walker, Heng Ji, and Amanda Stent (eds.), Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long Papers), pp- 1112-1122, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 1/16533/1/118-10111111 JRL/https: https///aaaaatthoogooo 1tt.///ooooooo..oo///1/1100000
sentence understanding through inference.
Weihao Yu, Mi Luo. Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang. Jiashi Feng, and Shuicheng Yan. MetaFormer is Actually What You Need for Vision. In 2022 IEEE/CVF Con- ference on Computer Vision and Pattern Recognition (CVPR), pp. N0.1109/CVPR52288.2222
June 2022. doi: ieee org/ document/
UtLhhttpss//ieeexpp...
9879612 2
ISSN: 2575-7075
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are Transformers universal approximators sequence-to-sseuunec functions? In The Eighth International Conference on Learning Representations, December 2019. URL https://
Chulhee Yun, Yin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and San- jiv Kumar. O(n) Connections are Expressive Enough: Universal Approximability of Sparse Trans- formers. In Advances in Neural Information Processing Systems, volume 33, pp. 13783-13794 Curran Associates, Inc., 2020. URL
hww.::eiiiinnnnnnnneeeeeee. neurips cc cc/paper/ htm_
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep Sets. In Advances in Neural Information Processing Systems, vol- ume 30. Curran Associates, Inc., 2017. URL httptttttt///// //apeers///ppprsseee nips wttpscccppppperfiiiiiii files/ Abstrectbbtttt......m .htll
7da 74a
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Al- berti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big Bird: Transformers for Longer Sequences. In Advances in Neu- ral Information Processing Systems, volume 33, pp. 17283-17297. Curran Asso- ciates.
Inc.. 2020. URL https::/papers nips. httpp/sccccaapee22///aaaaa html
c851 2d1 4 a
Chiyuan Zhang, Samy Bengio, Moritz Hardt. Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Commun. ACM, 647111177115,, feb 2021. ISSN 0001-0782. doi: 10.1 145/3446776. URL htt.:httt://///t htttpp:///bl..rr/1..111.2. /t1.:////11111//4477777
NOTATION TABLE
Numbers and Arrays
A scalar
A vector
A
A matrix
The length of an input sequence
N
The number of input sequences
C
The number of output classes
d
Embedding dimension
i-th input sequence, consisting of n tokens of embedding dimension
Sets
14

---

Published as a conference paper at ICLR 2025
{{...}}
Multiset (see Definition A.1
R
Set of real numbers
N
Set of all multisets over the domain X'
[m]
Set of all integers from to
Indexing
a,
Element of vector a, with indexing starting at
Aij
Element i.j of matrix A
A.
Column i of matrix A
A Row of matrix A
Functions
|
L norm of
supp(m) Support of m (see Definition A.2
ByNay(1) (x) The sequence of bits from the i-th bit the j-th bit (counting from the left) of
as
Softmax function
Hardmax function
OR
ReLU activation function
((RA) Hardmax-based self-attention mechanism with a skip- ckip-connection
((SA)
Softmax-based self-attention mechanism with a skip-connection
((F))
Feed-forward neural network with skip-connection
No
Transformer with parameter vector 0
Asymptotics
f(n) O(g(n)) f grows most as fast for sufficiently large
f(n) = Ö(g(n)) f grows at most as fast as for sufficiently large n, up to logarithmic factors
f(m = ((g(n)) grows at least as fast as g for sufficiently large
f There exists positive constant e such that f < cg holds
A DEFINITION OF MULTISETS
A multiset is generalization of a set whose elements are allowed to be duplicated.
Definition A.1 (Multiset). A multiset over the domain X' is identified by function X - N. which indicates the multiplicity mn rn(x) of each element x E the multiset. The set of multisets over the domain X' is denoted by N
Definition A.2. The support of multiset E N is defined by supp(m) = far X m(x) > 0).
In addition, the cardinality of a multiset m N is defined by
Em)m m(x) if supp(m)) 00, otherwise,
|m|
(14)
and the multiset m called finite if |ml] < oc.
In this paper, we only consider finite multisets, and in an abuse notation we sometimes denote fintie multiset No by ((z1)...,? N where ; F[m] E X' are possibly dupli- cated elements.
The following assumption guarantees that each value of the multiset is separated by certain amount, and the coken-wise separatedness in the analysis of Transformer's 's memorization can be translated into this assumption.
15

---

Published as a conference paper at ICLR 2025
Assumption A.1 (Element-wis wise separatedness). Let X' R and m ,...,, (N) E N be (5) (m(0]] (on) for each i [N]. Then, we say that
sequence of finite multisets with m) ({x)
m!)), ....., .....,.N) m satisfied:
are element-wise (r, -separaae for some r, 5 the following two conditions
1. for every i E [N] and & E [[m((]] k 112 S holds.
2. for every i,j E [N] and ls E [[. E [m either di 2i or 1| -x0()) ||2 2 holds.
B PROOF OF MAIN RESULTS
In the following. will extensively use the concept of multisets. For the definition of multisets and the notation used in this paper, refer to Appendix A We also use the operator X frequently. See
Section 3.1 for its definition.
To prove Theorem 4.1 we present several lemmas. Lemma B.I shows that to distinguish between N distinct multisets, suffices focus on the occurrence counts of at most V values. Lemma B.2 establishes the existence of function that computes sequence ids used to distinguish between N different multisets. Finally, Lemma B.3 states that the function obtained from Lemma B.2 with Lemma B.1 can be implemented using ferwaro network with OVVN) parameters.
Definition B.1. Let C X' and m E N a multiset. Then, we define the restriction of m to A by
(m(x) i A, m|a(x)
(15)
otherwise.
Lemma B.1. Let m!))...,,
(N) E Nº be a sequence of distinct multisets. Then, there exists a subset A C X with its cardinality at most N such that m!!) A....,m!N)] distinct.
Proof. We prove the lemma by induction. The base of N is obvious.
Suppose that the lemma is correct for the case N k; and we prove the case for N
Let (1), ....,, (...,,m(+))) (k+1) E N be a sequence of distinct multisets. Then, by applying the assumption the first multisets (1)
(A) E N we have subset C X with its cardinality at most (x) A distinct. If m ( are distinet, there is nothing prove. So assume that [ coincides with (+) a for some i [k)]. Notice that for any j E [K] with + m l and m(+1)) LA are distinct by the assumption.
k such that mL))
Since ( and m (k++1 are distinct, there is element X' \ A such that m (x)) * m (x). Then, the subset A' C. X' defined by A' AU {x} is the desired set for the case N = k +
In the next lemma, we say that scalars aj..... am are (r, 6) -epparated if [a,l S r for all i e [m]] and la: a,l 2 d for all i [m] with # aj.
Lemma B.2. Let m1!)) ..... ......!!! M! E N be a sequence of finite and distinct multisets with [N]. Furthermore, let S X be the union of all supports; that is,
mm() < M for every i
= U (8).
Then, there exists a function f:S [[4N2] [
such that
(=))())))...... Σ #Esupp(m(1)
Σ
m (N) (x)
(16)
(Esupp(m(AN))
(AM V SSeparspar
16

---

Published as a conference paper at ICLR 2025
Proof. Let - ||||]
....... ( be arbitrary bijective function. For each multiset m) with = ,,....A. N,, we define its -ige--imnniional representation by
Σ
m (1)eg))) E NNS!,
(17)
x Esupp(m())
where (g(x) E {0, 16.1111111 is a one-hot vector with in the g(z)-th position. Since m ...., m (N) are
distinct with Im ( s .M for every E [NV] we have
m
m m0l
Σ
-1lm)mmm00 (x)
1,
(18)
FES
for any i,j E [N] with i + j. and the norm of each mn (8) is upper-bounded by
m (0)
Σ
(x)
||a0t))12 |m S M.
(19)
#Esupp(m()))
By applying Lemma
G.1
m
(1)
m (N),
there is unit vector v E IR 1S such that
1
8
m
m 0
|| S T (
ml)) s m (6) m (
(20)
N2
S
holds for any j.j E [N]. Let h be the function : S Z, + [Na S SITT0 Hereafter, we that this function has the desired properties.
Let = /Toil [ /TEN 1) z i... the vector approximating N2|S|/Te
with integers. The approximation error estimated as follows:
[s|
N
(N'IS [N s 1S).
(21)
which means that | -112 S ST. Notice that
Σ
(x)h(())
Σ
(x) [N31S|
(Exupp(m())
(E)upp(m(9)
Σ
m0 (x)
TT eg(m)
#Esupp(m(9)
m
(22)
holds for every E [N]. Then, the absolute value of the left-hand side is upper-bounded by
(x))(())
T mn
#E)upp(m(9)
|| (
2
(23)
since the norm upper-bounded by
| <
||
| -
VIS
(24)
17

---

Published as a conference paper at ICLR 2025
On the other hand, for any i,j E [N] with i + J. we have
Σ (2)h())
E
(x)h())
<Esupp(m(9)
=Esupp(m(())
π. m
m
(m) N
T
m (6)
(N - v) (i - mm
N21SS//T
(m m (1)
/22IIVV - 8112 |m m7 (1)
2/15|
m ( m
(3)
v|s|
m m (1)
2 viSh
(25)
since eq. (20 implies
N NT (ml) - -m(())
= =N
(m
mm
1
8
2 N
m (1)
m
N2
T|S|
> 2
m
ml)
(26)
Finally, the output of the function is always bounded by
(Vr E S). (x)) = |[N < N S+
(27)
Thus, by setting f(z) h(x) |2N /T] we have desired function.
Lemma B.3 (Separation of multisets). Let X = Rd and m!)) .!."....,m""
m (N) E N be sequence of ".ri"..... m(a }} for each i [N]. Suppose that m ..... (N) satisfy
multisets with ()) (1 (6) (x)
the following three conditions:
1. (1) (N) are finite multisets whose cardinalities are at most M.
2. m m (N) are distinct.
3. m1)) ,...,, (1)....,,.()) (N) element-wise separatee for some 2 2
Let C. = -INN and Rs = 20r(NN / Then, there exists a neural network à
Rd + with width 12, depth
N
NVllgge
max(log Rag, log Co
(28)
log N
(for the definition of S. Section 3.1
and bit complexity bounded by
log d +
max Rss, log C
(29)
log N
such that ((x) E [[4N3 VT]| U {0) holds for any x E U (f)) and
(
mm [M)
(a
(1)
Σ
H(
(N1)
(30)
(4MN3 V 1)-separated.
18

---

Published as a conference paper at ICLR 2025
..... m m(,, we have finite subset A C IRd with |A| < N such that ml1) ( A...., ........./l la are distinct. Then, according to Lemma B.2
Proof. By applying Lemma B.1 to the sequence of distinct multisets m!).....
there exists function f A [[4N2] | 7 such that
m aa)))(())
Σ m m ((()f(())
(31)
#Esupp(m (a)
(esupp(mttlaa)
are (4MN Hereafter, we consider a function Φ F - R such that
h -eeearrpe, and in partieular (4MN' V., seeaprateed..
f(x)
i E A,
(xx) = 0 otherwise,
(32)
and simulate ø by neural network. Notice that the possible number of inputs for the function ø is most MN, and all outputs are natural numbers equal less than [4N]|A|/ I [4N3 /T. We define constants Rgg. and Co by
Co = [4N3V/]
(33)
RS = 20r
(34)
Then, Lemma
C.I guarantees the existence of the feed-forward network ò with width 12, depth
N
Nog
V
log N
max(log RS, log Co):
(35)
and bit complexity bounded by
log d +
max(log Rs, log C
(36)
log N
such that for any E [N] with (+) [{( |mttss }} and any E [[m(1], we have
[a
(1) if o E otherwise.
((m(0)
(37)
Thus, the outputs of à coincide with those of d for all inputs 2. with i [M] and [lm()]]-
Finally, we verify that the neural network Ø actually satisfies the desired property. For any E [N].
we have
Im
(1)
Σ
(x))))))))))))))))))))))))))))))
#Esupp(m(t)
Σ
m ( (x){(x)
rEsupp(m()))
Σ
ml ( (x(x)(xx)
(38)
(Esupp(m(e)a)
Thus, eq.
31
implies that Σ mm
1 )......
k=1
((
Dimi
(m(N))
((x(( are (4MN3 sesparred
B.
NEXT-TOKEN PREDICTION SETTING
B. UPPER BOUND
Here we state the complete statement of Theorem 4.1 with its bit complexity. Before moving on to the theorem, we introduce a uniform attention layer; that is, self-attention layer with the
While the upper bounds provided in Theorem B.1 in the form D(VNlog logn), these upper bounds can actually be reduced to O log(nNV) by modifying the data point partitioning Stage of Lemma C.1 to use /N log(rNN) subsets, each containing
N elements.
log(nNy
19

---

Published as a conference paper at ICLR 2025
softmax function replaced by simple averaging. For an input Z Rmx" layer calculates
the uniform attention
(2) 2+ w W Σ Z.k (1....,1) E R ERIXa
(39)
where W(V) E R and W E Rmxs are value and projection matrices with head size 8, respectively. A uniform attention layer is subset of self-attention layer it can be implemented using a aele-tttention layer by setting key or query matrices to zero.
In the next theorem, J(FF) and ((FF) represent feed-forward networks of arbitrary depth, unlike eq. (3, which is limited to two layers.
Theorem B.1 (Next-token prediction). Let (x0)),0)) (x((), ((o)) E dd x [C] be sequence of input-late label pairs such that
1. (0))))0)
(NNNN,,)) ,(()) are consistently labeled, in the sense that for any i, E
[ we have (( = y if
x ( x and x = -x() up to permutations.
(40)
(NN token-wise (r, 8) separated for some 2 1 and 0 < 5 1. 2. X
Let R 400 n.. Then, there exists a Transformer N Rdxn + R with width 14, depth
N log N
VVllo +
maag R. log C},
(41)
(for the definition of S. see Section 3.1 and bit complexity bounded by
N
log d V log N
maxxlog log C}
(42)
that memorizes the dataset, i.e.,
N (x()) = Eout (FF)) F(UA) o (FF) Ein (x(1), - holds for every i E [N].
(43)
Proof. For simplicity, we assume in this proof that there is no skip-connection in feed-forward layers, the modification for networks with skip-comeections is straightforward. For details on implementing the memorization results for feed-forward networks in Transformers, refer to Kim
et al. (023).
For each input sequence x (8) with [VV. we define its multiset expression m.)) E N((*) by
m (i) gd + N. { E [m] x (1)
x)
(44)
The cardinality of 6) for each [N] is at most n, and the token-wise (r, -eeparrttess of (1)) ..... (N) implies that m1)) , (N) are element-wise (r, -separteed In addition, the consistency the labels are rephrased follows: for any i,j [N], we have y(6) y(u) if
x((i x and (+) = m) hold.
Construction of ((FF), Applying Lemma B.3 sequence of all distinct multisets which appear (1))) ....,m. m (N)), we have feed-forward network ò gd with width 12, depth
N
S / VVl N
+
log max max(log R1, log C1
(45)
20

---

Published as a conference paper at ICLR 2025
with C1 == [4N2 π and R1 == 20r(22))8-'vv and bit complexity bounded by
N
x log d
max(log R1, C13
(46)
log N
such that o(
/ E [[4N8
holds for any i E [M] and l: E [nn]. and
o(z)
Σ
p(x)
Σ
21
(47)
#Esupp(m())
#E
holds for any i,j [N] such that m())) o m())
We extend the ferd-forward network T to retain the information of the input token. Let V be a set of all input tokens, that is, V = {X
i E
[N], h E [x]]- Since the input sequences are token-wise
(r, -separaae,, by applying Lemma C.2 to V. we have a feed-forward network F IR - R with
width 1, depth 2 and bit complexity log Bdr(nN)2 -=) such that
0 F(X
10r(nN) Add
(48)
for every E [N] and & E [n]. and
F( F 22
(49)
for every i,j E [N] and k E [n] with X : + x() Notice that the depth of the feed-forward network Φ is at least 2. Thus, it is possible to parallelly attach the above 2-layer network F to the first 2-layer of o, and extend the hidden dimension of the remaining layers of o by one to propagate the value of F the last layer. Furthermore, we augment the output dimension by one more and pad
by 0, which is used to store the average value of Φ. Let / ((FF) RRd + IR be the network obtained by the above procedure, that is, for any Rd the output of ((FF)
{(PF)
(z) = ((x)),
(50)
Then, the width of ((FF) is that of ē plus two, which is 14. The depth and the bit complexity of f(FF), on the other hand, remain the same, because the depth and the bit complexity of F is smaller than those of Φ. We also define a token-wiss operation ((FF) Rdxx + p3xn by
((PF) (X)& /(PF) (X..) (k: = 1.....m.)) n).
(51)
Construction of the self-attention layer: Let W e 33x3 and WO E R3x3 be any value matrix and projection matrix such that their multiplication is
0 0
w (V)
0
(52)
The output, which we denote by % R of the self-attention layer with the value matrix (V)) and projection matrix (O) for the input ( at index l: E [ri] is calculated as
((UA) ((FF) (FF) (1())
n ( w (V) ((F)
(x(!) + (PF) (x(1)
1=1
x(x
((x(1))
0
1
0000
F(x!!)
+
0
x
F(X!9)
(53)
1 ΣΣ 0(x(1))
21

---

Published as a conference paper at ICLR 2025
We verify that the right- hand side is context id, in the sense of Definition 4.1 Fix any i.j [NT]
If X +
(i) hi then according to eq. (49), we have F(X F(X..
2 2. On the other hand,
if X (1) are not permutation of X i... ml) # ml)), then eq. 47 implies that
(54)
k=1
Therefore, the difference of any two n-th outputs of the self-attention layer is lower-bounded by
x
)(x!!)) (1)
//
(3)
F F(X
F(X..)
11
Σk= o(x
1 -= 8(x
[(x()) F(X
Σ 3
min
IV
(55)
for any i,j E [N] such that either X x or m m holds. As for the magnitude of each output of the self-attention layer, it is upper-bounded by
8 (
F(X
2
x(x
2 Φ 90) + F(X (1)
+
[4N3 / + 10r(n)) 1 Rd + [4N3 VT
20rn N
(56)
where we used the assumption r 2 and å < 1 in the last line.
Construction of ((FF); What remains to do is construct a network f2 associates outputs of the self-attention layer with their corresponding labels. Specifically, since
(FF)
R3 + R which
we know from eqs. (5)) and (56) that the sequence of unique elements in ((1)) ..... 8n are (20rn? 2 VTdd 1/n) .spaarree by applying Lemma C.1 to N inputs ,..., labels y(),..., ")...,y(N) we have a ford-forward network f/
(N)
(N) and their (PF) R3 + R with width 12, depth
N
S N log N +
maxx R2, log C}
(57)
log N
with R2 20 20rn N35 V N bounded by
v 400v3dd"r and bit complexity
N
max R2, log C}
(58)
log N
such that f (1 = (8) for every E [N], In particular, this means that by defining a token-wiss wise
(FFF)
operation F (FF).
1x3xn + RR as
F FF)
X k: = 1
(FF) (X..)) (k 1.......))
(59)
we have
F (FF)
(FF)
o F(UA)
x(FF) (x(i) = ())
(60)
22

---

Published as a conference paper at ICLR 2025
for every i E [M].
Model complexity: The width the Transformer F2 }(FF) F(UA) J(FF) is the maximum of widths of w(FF) F(UA) and x(FF) which is max(14, 3, 12) 14. The depth is upper-bounded by the addition of depths of F( and F plus one, which implies that the depth is
N
N log
max R1, log C1}
log N
N log N max R2, log C}
v Nlog N +
V
N log N max R. log C}
V N Nolog N
(61)
with ==
= R2 - 400 30000/dd///0 N
2 max{log R1, log C₁ log Ra}. Likewise, the bit complex-
ity is
N
< log d +
max R1, log C1, log R2, log C}
log N
N log N
S log d +
max{log R. log C)-
(62)
B.1.2 LOWER BOUND
For convenience, we restate the statement of Theorem 4.2 below.
Theorem B.2. Suppose a Transformer RRdxn + IR defined by eq. (4) can shatter a set of N distinct input sequences xx),..,,X"" x Rudxa with x- x( for any i E [N] and k: E [m], in the sense that for any label assignments y(1),,.. "i.),
((N),,,))))))))))))))))) E {0, 13, there are parameters with which ( = y holds for any E [N]. Then, the Transformer N has at least M(vN) parameters.
Proof. We denote by L the depth of the Transformer F, and a feature matrix at block 1 = 1, 1,...,II L by
hy x = F (FF) F(SA) o ... ((FF) ((A)) a Ein E Rmxn
(63)
with (X) (X) and (XX) = Eout ohg (XX for any input X E Rddx. Then, the permutation
equivariance of Transformers implies that the feature matrix h at block 1 = 1,..., L satisfies
(x (0) ).. .... xx(x(()) (1))
(64)
for each E [N]. Thus, the self-attention layer block 1,..., ,,,,,,, can be calculated by
((A)) (xx1)()(10)
k
= (x(1),
wlo (0) (V
( ( - (w!l)), (x (())
)
hd hil
H
-1-1(x(1)) x+ +
w h.
w h (x(
H1
I+ =11+
W (0) (1)
h (x
(65)
b.t
1=1
23

---

Published as a conference paper at ICLR 2025
where w W hy and (Q) with h H are weight matrices for the self-attention at block i, and I E Rm is the identity matrix. This observation indicates that calculations of self-attention layers for inputs X ((()) reduces to linear transformations, which in turn implies that the behavior of the Transformer N inputs ( ..... ( can be simulated by ferd-forward network with equal or fewer parameters, and with inputs .......,, x. Since is known that the VC dimension of ReLU-basse
worward networks with W parameters is at most O(W2) (Goldeerg & Jerrum 1995), the Transformer N must have at least D (VVN) parameters. I1
B.2 SEQUENCE-TO TOOQUUCCETOSSEQUENCE SETTING LOWER BOUND
Before proceeding to the proof of Theorem function:
4.3
we cite the following lemma. Here sgn is the sign
1
if x > 0,
sgn
0
if x 0,
(66)
-
if O.
Lemma B.4 (Goldberg Jerrum (1995)) Suppose W S M and let Pi,..., PM be polynomials of degree at most D in W variables. Define
K = (sgn (gg)),...,ss((PM((()) a aER").
(67)
ie., K is the number of possible sign vectors attained by the polynomials. Then we have K I (8eM D/W)W.
Hereafter, let W be the nubmer of parameters and 0 E R be a vector of all parameters of a Transformer. We also denote by No the Transformer to emphasize the presence of the parameter vector 0. For convenience, we present the statement of Theorem 4.3 below.
Theorem B.3 (Lower bound). Let No Bdxn - R" be Transformer defined by eq. (4) with self- attention layers replaced with hard attention layers (eq. (12). In addition, suppose No can shatter set of N imput sequences
x Tgdxn with x( * x(i for any i,j E [N] and k E [n]] (i # j or lo + 1), in the sense that for any label assignments y......, (N) E {0,1}",
there is parameter vector O E R.W such that
(o( (b)
(68)
for any i e [N]. Then, the Transformer has at least W = Ω
ON Tog(nN)
parameters.
Proof. Recall that the Transformer No R - Rn is defined as No = Eout o Fr, .... Fi Ein.
(69)
where the 1-th block F; : Rmxxn Rmxn is composed of self-attention layer and a ferd-forward layer. For the Transformer No to memorize all label assignments for given N input sequences with length n, the number of possible sign assignments for outputs of the Transformer must be least p equal to or more than 2n.V that is,
K{
H(000 (No(X!)))a
2nN N S KK=
KE(N] O
MER' RW
(70)
kEjn]
must hold. We estimate the upper- upper-bound on the right- of the above inequality.
Our strategy is to partition the set of parameters inductively with respect to the layers, so that each cell the output of the Transformer can be expressed by some polynomial function on the parameters. To be more precise, construct sequence of partitions So, Si,...,& S.....,,. St. (PR such that
1. for each I 0, 1....,,.,, St is partition of the set of parameters, that is,
S, S, = (VS, Sy e S with S. # St) and =S=R".
(71)
SES,
and is also a refinement of S when 2 1, in the sense that for every cell S E S,, there is cell S' E S with S C S'.
24

---

Published as a conference paper at ICLR 2025
2. for each I [L], the number of cells in St satisfies
|
88 88c-n H HN-. (81 (81- (8)-4) 4)
Wi-1+2maH
8
Wi-1 2msH
Se ngg. 41
Wi-1 11-+t2ms++mm-1
(72)
Wi-1 2msH + (m 1)g
where Wi-1 is the number of parameters up to the (1 1)-th block, with Wo dm, the number of parameters in Ein.
3. for each I 0, 1, ................................ L, outputs of the 1-th block for imput X on each cell S Si
(1) k.s (0) = Fio Fi Em (X
Ju,k (u E [m], kt [n])
(73)
are polynomial functions in variable O₁ of degree at most 41 as long as 0 varies within the cell S. Here 0 E Rw; a part of & corresponding to parameters up to the 1-th block, with On defined by parameter vector of Ein-
First, set Sp (R"), }. polynomial functions in variable Og of degree 1, because En Rdxx - Rmxnn is token-wise linear mapping.
Notice that outputs Em ( ( u,k for all [m]. k [n] and i e [M]
Next, suppose sequence of partitions So..... ....... S for I [L- 1] already given, and we construct a partition Si from them. Specifically, we subdivide each cell $ E S to create new partition S, By ussumption, on each cell S the inputs of the (1 1)-th block Fi-1 for the input (
( Pi- Pi-L,,,.s(01-11)) Fi-1 ....o Fi o En (+) u.k
(u [m], ks E [n]),
(74)
are polynomial functions in variable a parameter vector 0 of degree no more than 4(1 - 1) + 1, as long as 0 varies in the cell S.
Sell-attention subblock: Recall that the self-attention layer with the hardmax function in the (-the block for the input sequence X (8) is calculated as follows.
(HH)) x(HA))Z(()) = 2
H
w
(0)
(V) M
Z0On Z(1l))z)) (w z )
(75)
h=1
where (0) E Rmxn is the input of the self-attention layer for input X (9) In particular, when varies in cell S E S1 2 for each [m], k [n] can be expressed by the polynomial
function (8) ki (0:-) of degree 4(1 1) 1.
Hereafter, subdivide each cell S S to construct refinement S(SA) of S so that the hardmax patterns
(w/(k) z!0) (W 200) hl
E Ru
(Ah H)
(76)
remain the same on each cell S' E S(SA)_ The (k, k')tth element of the attention matrix at head [H] can be written by
(3)
W (K) W
(4)
w 2(0)
(W(O) (Q) z(i)
i,ke
m
(K)
w
(6) Pi
(9 (0
M (0-1),
bl
u,k==1
(77)
from which we see that each element of the attention matrix is a polynomial function in variables (K) w(Q) and O,-1, of degree most 8(1 1) 4, long as 9 varies in the cell S E S
25

---

Published as a conference paper at ICLR 2025
We define a partition P(AA) of S based the hardmax patterns, that is, the minimal partition of S such that each cell, all outputs of the hardmax function remain the same. To estimate the size of P(SA), we instead consider sign patterns of polynomials
(0) L.A.A. - Of OLA.
(
E [N], he [H], K,K, [nl].
(78)
because whenever sign patterns of the above set of polynomials do change subset of the parameter space, the hardmax patterns must also remain the same. Applying Lemma B.4 to the above collection of polynomials, the size of the partition LLs is upper-bounded by
Se: n HNN- (81 - 4)
Wi-1+2mst
(79)
Wi-1 + 2msH
of S by subdividing each cell S E S in this way, and its size is We define a refinement (SA) upper-bounded by
(SA) $ ||(111
(8e Wi-1 2msH
(Be. 88 Bn (81 - 4)
W- W
(80)
On each cell S' S(SA) the hardmax patterns remain unchanged, which implies that
ES
Z Z
oH
( (K) z(0)
W
2
(u [m], Is E [ni], h e [H]) u.k
(81)
polynomial functions in variable 0:-1 of degree at most 4(1 1) 1, as long as 9 moves in the cell ES(SA) S(SA) This further means that each element of the output F. (HA Z (0) is a polynomial
function in variables w(O) (i() with h E [H] and O--1, of degree at most 4(1 + 3 on each
cell S' E (SA)
Fordardrward subblock: As for feed-forward layers, we follow the analysis given by Bartlett al. 2019). On each cell S" E S s(SA) the hidden layer at the k-th token for input X (8) is
F
(HA)
) + b1" E R?,
(82)
whose e-th element is a polynomial function in variables w-(o) (V with h E [H], w( and 0111 of degree at most 4(1 1) Notice that sign patterns of polynomials
(
W
(1) J HA) Z ()).. + E [N], u E 14]. k ke E
(83)
completely determine whether or not the ReLU activation function in the middle layer fires. There- fore, by defining 1.5' same on each cell, the size of (( is upper-bounded by Lemma B.4 as
(FF)) as the minimal partition of S' such that the activation pattern remains the
|(SF)
SE Be.ngN. .11 Wi1+ + 2msH + (m + 1)q
#m (m++1g
(84)
We define refinement (FF)) ((SA) by subdividing each cell S' E E (SA)) into ((FF) Then, outputs of the feed-forwer ferwadd layer
Pi..k.s (0:) = (HA) (Z()) Ju,k W L.a
(D J. (Z ) + bi [10]] + (2)
(85)
for any u E [m],k k E [n]] and [N] are polynomial functions in variable 0, of degree at most 4(1 1) + = 41 1, as long as the parameter vector varies in the cell S" E ((PF)
Finally, we set S, as S = F!FF Then, from the above observations we know that outputs of the 1-th block are polynomial functions in variables W, of degree at most 41 + 1 as long as O moves
26

---

Published as a conference paper at ICLR 2025
within each cell S" St, desired. In addition, we have
|
8e nnnH&N HN (81 4)
W1-1+2meH
S
Wi-1 + 2msH
Wi-1 Wi--+meeH+++tll
Se 12-N 41
(86)
W-1+ + 2msH + (m + 1)g
which satisfies the second property. In this way, have desired sequence of partitions Su. SL.
Outputs of the L-th block for input X (0) E [N])
F...
Ei (x Ju, (u [m]] k [n]]
(87)
are polynomial functions in variable 9,, of degree at most 4L + 1 as long as 0 varies in each cell of St. which in turn implies that final outputs of the Transformer
1 Pk.s (0) = No (0)) )a- = Enut Fr F.... Fi Ein ( ())x (k E [n])
(88)
polynomial functions in variable 0 of degree at most 4L if the parameter vector 0 moves within each cell S E SL-
Applying Lemma B.4 to the set of polynomials (0),s(0)) on each cell St allows us to upper-bound K as tollows.
{
K
sgn
Noo (oo()(0))))
iEIN kE(m
0 R
Σ
sgn (No(x(())))
o
iE(N) kE(n
SESL
88 n.. (4I+ +
W
$ |
(89)
W
Since |1t] |8o1 II | S -11
and [So]
1, the right-hand side is further expanded as
nN (4L + 2)
w
S
k
W
S
8e nN (4L 2)
W
W
8e n HH (81 4)
Wi 2mmH
8e ngNN -11
2m
Wi-1 + 2msH
W1 + 2msH + (m 1)g
1=1
Se nN. (4L 2 + Σt [8 n HN- (81 (8)-4) + Se ngN 41] W
(90)
W
where we used the weighted arithmetic geomeerie inequality in the last line, with W defined by
W +
[w + 2msH + Wi + 2msH + (m + 1)g]
(91)
1=1
Notice that W; for each / [L] is the number of parameters up to the (-th block, which indicates
Wi = md +
Amss + 2(m + 1)g]
1=1
= md + [4msH + 2(m + 1)q]
2 4.H 21g
(92)
27

---

Published as a conference paper at ICLR 2025
with W Wy + md. With this observation, the numerator on the right-hand side of eq. (90 is upper-bounded by
Se n. (4L + Σ 88 n HN (81 4) 8e nqN 41
S 88 nN (4L+2) + + N Σ
(SIH + 4lg)
l=1
< 88 NN (4L - 2) + Se - N Σ
2W
1=1
S 48c n W
(93)
where used 4T. 2 < 4W and Σt 2W, S 2W the right-hand side of eq. 90 is opper-bounded by
Σt= 2 2W in the last line. Thus,
W
K 48 W W
(4 W
(94)
Recall that in order to memorize all label assignments for N input sequences with length n, K is at least equal to more than 2n. which gives us an upper-bound of n.V:
(48en* N) W
nN & log2
W logg (48en N)
3W loga (48enN).
(95)
Here we evaluate a crude upper-bound of W with respect to the number W of parameters as follows.
W
W
2W 4msH (m 1)g]
L
<
W
W
1=1
< W
2LW & 3W?
(96)
which implies nN < 3W log, (48en.) < 9W/2 lo/g (48enN). Therefore, the Transformer has at
least W = Ω (V (Vogg())) (og(nN)
nN
parameters.
C MEMORIZATION OF FEED-FORWARD NETWORKS
In this section, we extend the result on the optimal memorization of feed-forward networks proved by Vardi et al. (2022). Specifically, the following lemma states that we can freely add data points without severely affecting the memorization capacity of feed-forward networks, as long as their labels are zero. We would like to note that Vardi et al. (202) implicitly used this result to show the memorization capacity of feed-forward networks with a bounded depth. Thus, our aim here is to explicitly state the result and provide rigorous proof.
Lemma C.1 (Extension of Vardi al. (2022)). Let N, N,V,d.C C E N with N < V, and
1,0
5 < 1. Let (L)), ((V) E [C] be sel of labels and (11) (..,),,.,..... xv(V) E Rdd be set of inputs
such that [0x & r for every E [V] and [|x](0) - x05i 2 for every i,j j E [V] with i + j. Denote
R = 228 / nd. Then, there exists neural network F Rdd + IR with width 12, depth
N
$ / Nlog
*V log log N
max max(log R, C)g
(97)
28

---

Published as a conference paper at ICLR 2025
(for the definition of S. see Section 3.1 and bit complexity
N
S log d + V log max R, log O}
(98)
such that F(x(()) w for every i [M] and F(xx()) 0 for every E [V] [N].
Proof. The proof goes basically the same as was done in the proof of the original theorem by Vardi ctt al. (2022): we construct a three sub-networks F1, F2 and F'3 with width at most 12, and then concatenate those networks to create the final network F F3g F22 o F1. The only architectural difference lies in the construction of F1, and the rest of the proof is dedicated to verifying that the resulting network F satisfies ((x(()) = for E [V] [N].
STAGE I: PROJECTING ONTO A ONE-DIMENSIONAL SUBSPACE
In this stage, we construct a sub-network Fi TR. - IR, which projects each input onto the line TR while approximately retaining their distance. We use the following lemma from Vardi et al. al. (2022).
Lemma C.2 (aadd et al. (2022). Let (1). a..,,,..,,a" (N) E Rdd with ||,( S for every i E [N]
and ((x(() x!''|| 2 8 for every i,j E [N] with i + j. Then, there exists neural network
F : Rd - TR with width 1, depth 2 and bit complexity log(3dr VTo-1), N22 such that 0 I F(x(1)) s
10 Vrd every E [N] and [F(x!)) F(x())) 2 for every i,j E [N] with i A.
Instead of applying the above lemma to the set of N inputs x(1), x (N), (N),, here we apply it to the Then, we obtain a neural network Fi Rd + R with width 1, depth 2 and bit complexity log 1og(3rrV2 VT==1) such that 0 S (z((()) I 10rV28-1 for every
set of V inputs (11), .....
xV)) (
i [V] and (x)()) Fi (xx()) 2 2 for every i,j E [V] with i j.
By slight modification the bias term, we may construct neural network Fi R R such that S Fi (z) < R == 20r VTd without affecting its width, depth and bit-complexity. We adopt Fi as the first sub-network.
STAGE II: FINDING THE RIGHT SUBSET
In this stage, we adopt the same construction strategy for the second sub-network F22 : IR R as was done in the proof of Vardi et al. (2022). We use Lemma G.3 whose statement is the strengthened version of the one by Vardi
(2022).
We denote the outputs Ii..... .y. In addition, by rearranging labels, we assume without loss of generality that the first N outputs T1, ...,....2....................... ZN are in an increasing order, that is, 21
(x())) of the first sub-network F1 for (1)),... (V)) by
IN.
Let m = VN log N. and wi...., WN No and uI..... " V be two sets of log N
log C- log R-bit sequences, respectively, such that for every i [N], let
N
bit sequences and
N
V
log N
j
log N
E [m],k = mod
N log N then wi...., NNolg and un,..., VNoloN
are defined by identities
BIN BIN lo C411(++)))og (ww)) (0())
(99)
BIN BINg-tog R+1:(k+1) ne(u;) [x,],
(100)
where we used the fact that the outputs of the first sub-network F, are non-negative and upper- bounded by R 20.
Next, by applying Lemma G.3 to wi...., W. W and u1...., / log N. respectively, we obtain
two networks F.r R RR and Fr R - IR with width 4, depth 3 N log N + and bit complexity
29

---

Published as a conference paper at ICLR 2025
at most V
N Nog max{logs C, log R) + [log R| such that for every i E [N],
F [z)) = Wj, and (2x) uj.
(101)
hold with ji i. log N
By concatenating these two networks F2" and F2. construct
second sub-network F2 IR
R such that for any i [N] we have
F2(x,)
W
(102)
Uj.
As for the outputs of F2 for #N+1,-... IV, since the construction of the first sub-network Fj assures that jr, - syyl 2 2 for every i.j [V] have
with + j. Lemma G.3 indicates that for any [V] [N].
2i
F2(1;)
(103)
where (resp. u) is either 0 or wy (resp. uj) for some j [m]]
STAGE III: BIT EXTRACTION FROM THE CRAFTED WEIGHTS
As in the previous stage, we follow the same construction strategy as is done in Vardi et al. (2022) However, here we inspect the behavior the third sub-network for :
We the function obtained by Lemma G.6 p log C,
log and log R the third
sub-network F3 1 with width 12, depth 3/ log N max R, log C} + 22
N log N
+ 2 and
bit complexity log max floglog R. log C} + 2. Then, we construct the final network F RRd + IR
by setting F F3 F2 Fi.
VERIFICATION OF BEHAVIOR AND MODEL COMPLEXITY
Hereafter, we check that the configured network F F3 F2 Fi correctly outputs the desired values, that is, for any i E [N] we have
F(x)())) = (0)
(104)
and for any i [M]][N] [N]
()) = O.
(105)
Fix i [N] with je
1. log
The output of F2 Fi for x00 is
F2 o (x)
Wjj Ubjs
(106)
Since [x:] BIN (1+ (Uz:) with : mod
V
.N log
by definition, Lemma G.6 implies
Fg F2 o F ()) = BIN p. (IN (0) as desired.
On the other hand, for any E [V] \ [N], the output of F2 o Fj for x0 is
Pg o Fi (x)
w
(107)
where wo (resp. u) is either 0 or wy (resp. uy) for some j E [m]. If u = 0, then xy satisfies
[r 1/2 BLV p. p1pp) = jz; 1/2) 1,
(108)
30

---

Published as a conference paper at ICLR 2025
because the construction of the first sub-network Fi guarantees that 21..... 2v 2. Thus, Lemma G.6 implies that F(xx()) = F3 F2 Fi (xx())) = 0 desired. On the other hand, "j for some j [m]. at should satisfy [ 1/2 BIN, (++))p((++))()))) > for any
k E {0, {O..... log 1}. This is because for each k, BIN p+++1:p(++)) equals [xi]] for some [M] by definition, which together with the separatedness of T1. .......,, By implies
N
I - 1/2 BIN (o+
= i. 1/2 [x
> rr x1| - | 1/2 [
2 2 - 1/2 1.
(109)
Therefore, the output of F = F3 F2 Fi for x; in this case is again 0.
The width of F is the maximal width of its sub-networks which corresponds the width of F3, i.e., 12. The depth of F is the sum of the depths of F1, F2 and F3, which is estimated as
N
N
2 + 3 N N + +2+3
log max R, log C} + 2/ log N
+2
N
2 Nog +
max R, log C}.
(110)
log
The bit complexity of F is the maximal bit complexity of its sub-networks, which is upper-bounded by
N max C, log R} [log R],
max
log(3drV2 VTo
log N
N log max R, log C} +
V
N V log N
S log d +
max,log C, log R}.
(11I)
D MEMORIZATION CAPACITY OF DEEP SETS
Refer to Appendix A for the definition of multiset and the notation in this paper.
Deep set (Zaheer al. 2017 is well-known architecture used for modeling functions that take set, or more generally a multiset as input. The architecture is stated in a very general form, and is known (Wagstaff al., 2022) that any permutation invariant function for multisets over countable domain X' can be decomposed by appropriate functions á and as follows:
(d.p)(m)
Σ ((zk))
with m ff21 Init E N
(112)
In this paper, define deep set tuple (0, where á and p are feed-forward networks. In addition, the width of the deep set (ó, p) is defined as the maximum of the widths of ó and and the depth of (Ø, p) as the addition of the depths of @ and p.
(M) (()) N x [C] be a sequence of input-late label pairs such that ml)) ,..., (N) satisfy the following three conditions:
Theorem D.1 (Memorization of deep sets). Let X = Rd and (m()), ())),...
1. (1) "..., (1),..., (N) are finite multisets whose cardinalities are at most M.
2. m11) 3.r.. m are distinct.
31

---

Published as a conference paper at ICLR 2025
3. m()) m are element-wise (r, -seaarraed for some 1, 0 < d S (Assump- tion A.7.
Let R = 8 o d Then, there exists a deep set (Ø. p) with width 12, depth
S Nog +
max R. log C}
(113)
log
(for the definition of S. see Section 3.1 and bit complexity bounded by
N
2 llog d + V log
max{log R, log C}
(114)
which memorizes the dataset, that is,
(ø. (m. =
Σ o(af
-y
(115)
holds for every i E [N] with ml() = {(m))
imn m
Proof of Theorem D.7 Let C. [4N and Rs 20r(N Then, applying
Lemma B.3 readily implies that there exists a neural network @ R2 - with width 12, depth
N
$ Noo.
+ V log N max(log RS log Co):
(116)
and bit complexity bounded by
N
log d +
max max(log Rs. log C
(117)
log N
such that d(x) [IAANVVT] (0) holds for any U (u)p(m(()). and
(A)
Im (m)
pl
(1)
(N))
(118)
are 3 VT, 1)) separratd.
Since the correspondence Σ ((xx(² rization of N input-label pairs
Σk=1
(m (0)
to the label (+)) is injective, we can consider the memo-
m
Im (m))
g()))
Σ
ER x [C]
(119)
with feed-for ferward networks. Specifically, let Rp be
Rp = 20 4 VT 2 VT 80MN
(120)
Then, according to Lemma we have feed-forward network P R - R with width 12, depth
N log N
S / log
max/log Ro log C},
(121)
and bit complexity bounded by
N
max(log Rys log C)
(122)
log N
32

---

Published as a conference paper at ICLR 2025
such that for any E [N] we have
63 6(2)())
p
- y0(
(123)
as desired.
Model complexity. With the configurations defined above, the deep set (Ø. p) provably memorizes the dataset. Lastly, check its model complexities, that is, width, depth and bit complexity.
The width of both ø and is 12, and thus the width of the deep set (Ø, p) is also 12. As for depth and bit complexity, we define R by
R N 'vVd.
(124)
Notice that Rg, Cs and R2 are all upper-bounded by R, because of the assumption 2 1 and < I 1. The depth of the deep set (Φ) p) the addition of the depth of each feed-forward network, and thus upper-boune bounded by
N log N max R, log Cj.
S V N Nog
(125)
Likewise, the bit complexity of the deep set () is the maximum of the bit complexity of each feed-forward network, which is upper-bounded by
N
logd d +
max R. log C).
(126)
log N
E TRANSFORMERS WITH EMBEDDING LAYER
In this section, we examine the memorization capacity of Transformers equipped with an embedding layer. When considering an embedding layer, input sequences consist of sequence of token ids, rather than sequence of word vectors. Mathematically, N input sequences we consider in this section are expressed by N vectors
x
(N) E [w]",
(127)
where represents the vocabulary size, i.e., the number of distinct token ids that occur in the input sequence. Then, the embedding layer (FM)) operation
[w]" mm is defined by the token-wise
(EM = (EM) Ezy
IR (k E [n]).
(128)
with (EM) E R.mxw the embedding matrix used as a lookup table, and ez, E f0, 1} one-hot vector with in the rą-th position.
Given the input sequences and the embedding layer defined in this way, we now state the theorem on the memorization capacity of Transformers with the embedding layer. As Theorem B.1 (FFF) in the next theorem represents a token-wise
network arbitrary depth, and (UA) is
uniform attention layer (see eq. (39) for its definition). Remarkably, the number of parameters now depends on w, which is unavoidable due to the use of the embedding layer.
Theorem E.1. Let (2), ())),, (xo(N), ,y(NN) E [W]" [C] be a sequence of input-label pairs that are consistently labeled, in the sense that for any i,j E [N], we have y xydd) if
x = ()) and ( = (5) up to permutations.
(129)
Let R. 2000/322r'' wn. Then, there exists a Transformer with the embedding layer (EM)). W/" + RD with the number of parameters
N
S w NN+NNooo V + log N
inx{flog R, log C),
(130)
33

---

Published as a conference paper at ICLR 2025
and the bit complexity
N log.N N
x log +
max (log R, log C}
(131)
that memorizes the dataset, i.e.,
(xEM)(x Eout o F((F) F(UA) o (x
(132)
holds for every i [N].
B.1 that the role previously performed by the feed-forward network (PF) is now implemented by the embedding layer F(EM), Specifically,
Proof. The only difference the proof from Theorem
for each input sequence x with i E [N]. we define its multiset expression m E M144 by
m (1) [w]] - N, > 1(k ok = x]|-
(133)
The cardinality of (1 for each E [N] is
n, and the consistency the labels are rephrased
as follows: for any i,j E [V], we have ( = y( if ( = xهاا and m m hold.
According to Lemma B. there exists subset AA [w] with its cardinality at most min N such that (1) A..... m
are distinct. Then, by applying Lemma B.2 to (1) lA...., m (N) we have a
function f 4N2 V min(w,N]. such that
Σ
(x)f())
Σ
m (N) (x)f(x)
(134)
(Esupp(m(*)a)
*Eaupp(m(<0))
are (4nN2 VT min(w, NN, () sepsararte We directly implement the function f in the embedding layer. Namely, we define the embedding matrix (EE)) e Raxw the embedding layer F(EM) with = by
W(EM) 1
(f( 20,
i E
(135)
otherwise,
for each 2 [w]. The bit complexity of the embedding layer is most logl|4N2 V Vi-min(w, NN F1++ log w, and since the construction of the remaining parts of the Transformer can be carried out simi- larly Theorem B.1 we omit those details here.
F MEMORIZATION CAPACITY WITH LIMITED BIT COMPLEXITY
In this section, consider how the memorization capacity of networks changes when the number of bits available for each parameter of the network is bounded. The following lemma extends Theorem 6.2 from Vardi al. (202)) with the only difference that also explicitly supports additional data points with zero labels.
Lemma F.1 (Extension of Vardi et al. (2022)). Let N, V.V.CCC E N with N < V, and r 2 1,0 < 8 < 1. Let (())
such that [| ...., ...,y()) [[C] be set of N labels and (1 ,..., (v Rd be set of V imputs
r for every i E [V] and [|a,(0] x11i]| 8 for every i,j [V] with # j. Denote R 10r and let [VN]. Then, there exists neural network F ; Rdd R
with width 13, depth
N B
N
max{log R, log C),
(136)
B
B B B
(for the definition of S, see Section
3.1
and bit complexity
S log + Vlog B
max R, log C}
(137)
such that F(xx(0) () forevery every [N] and F(x((9) 0 for every i [V] [N].
34

---

Published as a conference paper at ICLR 2025
Proof. The proof idea is the same as the one for Theorem 6.2 from Vardi al. (2022): : N struct a B² + sub-networks
with width at most 13, and then concatenate those networks to create the final network F. For the first sub-network Fi, work as in the proof of Lemma C-1 which projects the inputs ,..)..., xx(,), (N++1) ".,,"" xV(V) into scalars 21,.-.; #N++ +1 -..,, By while approximately keeping distance between them.
the same net-
Next, partition 21,...; EN into N subsets each containing B data points. For each sub- set 1(j-1))B8+1 Z2.22 (i E B D. we use Lemma C.1 with zero labels at other data points "...), and 42,, TV to obtain sub-networks F F..,FN///241 with width 12, depth
I Bv Bo B
Vlg mmx R, log C
(138)
and bit complexity
B
S
Vlog B max R, log C).
(139)
Finally, we extend the widths of Ez...., BN/Ba by one to create sub-networks Fz...., FNy such that
Fi(x) ) (++))))+())))
F r )==-(
(== 2,....
(140)
By concatenating all sub-networks and one projection layer 6(x, y) y to obtain the final network
F = Φ F .... F2 Fi, whose depth is upper-bounded by
B
B B B
+
maxflog R, log C7)
B
Vog B
N /oo B
N
max RG log Cj.
(141)
B
Bylog
For each i = [N]. there is unique index j E (,...., N/B²) such that Fa));) = (()) holds and at the same time have (x;))) = 0 for any ho N/...,,/BB) with is j. Thus, the output of F for (+) is calculated as
F(x()) = (
) ,
Iz
(142)
On the other hand, for any E {N + 1....,, always zero, which implies F(x(1) desired.
V} and j E
N/B*), (2.....N/B)), the output of (z,)) is
By replacing feed-forward networks used in the proof of Theorem 4.1 with Lemma F.1 obtain the upper bound on the memorization capacity of Transformers with limited bit complexity. Notably, the following theorem shows that a Transformer with Ö(N'- parameters can memorize data points in the next-token prediction setting when each parameter is restricted to Ô(N) bits for some C [0, 1/2], under the condition that n. C, C5-51 NO1) and d = Ó(N==)) as N cc.
Theorem F.1 (Next-token prediction with limited bits). Let (x(1), ()..... (X(N), (())
Bdxx x [C] be sequence of imput-label pairs such that
1. (x(), ())) ..... ",.)),..,""""..... (X(N), ,y(NN) are consistently labeled, in the sense that for any i, [N], we have ((6) = if
x = x and x = x(0) up to permutations.
(143)
2. (1). ..... (N) token-wise (r, 8) sepaaradd for some 2 1 and 0 <
35

---

Published as a conference paper at ICLR 2025
Let R 4000 T and B [VM] Then, there exists a Transformer N R + R with width 15, depth
Nv B
N
Bylog max R, log C),
(144)
B
(for the definition of S. see Section
3.1
and bit complexity bounded by
B
log d +
max max R, log C}
(145)
Vlog B
that memorizes the dataset, i.e.,
N (x(1) = Eout F2 (FF))
o F(UA) o F Em (xi()), = =
(146)
holds for every i [N].
Proof. The proof goes basically the same as is done in the proof of Theorem 4.1 but this time feed- forward networks with limited bit complexity (Lemma F.I replace two token-wise feed-forward networks in the proof of Theorem 4.1 namely, F(FF) and (FFF)
(FF)
The first token-wise feed-forward network F(FP) defined by eq. 51 is essentially composed of é R + R obtained from Lemma B.3 which in turn is constructed from - R defined by
eq. 32 using Lemma C.I Therefore, let us consider representing á with feed-forward network with limited bit complexity using Lemma F.1 and the resulting network Ø R - R has width 13, depth
Nvlog B
N
B Bylog B max{log Rig log Co
(147)
B
where Cs 3 VT and Rs 20r(N (N nd, and bit complexity
B
logg d
max max flog Rog log Co
(148)
Vlog B
Then, the first token-wise feed-forward network F is defined using á in the same manner as in eq.
FF)
pdx= - 3x with limited bit complexity
Similarly, the second token-wise feed-forward network F (PF) is detined using Lemma C. to asso- ciate // (N²) defined by eq. (53
with labels y(),..., (N), Thus, this time we use Lemmal F.1
construct a heed-forward network F2
(F))
R with width 13, depth
N Aog B
N
B Bog B max(log R2,logg,,110 log C},
(149)
B
with Rg = 20 2 nd bounded by
N
n V /d0 and bit complexity
B max (log R2, log C}
(150)
v Vlog
such that S (8n = (t) for every i E [N], which induces the second token-wise feed-forward
(PF)
(1)
network F (FF) p3xn Ra with limited bit complexity.
G TECHNICAL LEMMAS
This section summarizes various technical lemmas. In this section, LEN(m) E for any n e N represents the number of bits in its binary representation.
36

---

Published as a conference paper at ICLR 2025
Lemma G.1 Park al. (2021)). Let d E N. Then, for any finite subset X C R there exists unit vector v E R such that
8 nd
x w' 1l2 S 10 (æ: (e=) S <1 æ'll2
0/0000
(151)
X
holds for any x, x' E X.
Lemma G.2 (Vardi et al. (2022)). ). Let a, N with < b. Then, there exists neural network F with depth 2, width 2 and bit complexity LEN(b) such that F(x) = for x (a, bj and F(x)
for x > b + or I < a
Lemma G.3. Let 21 < ... <. EN < R with R > O and [x; 2z1 2 2 for every i,j E [N] with
i + j. Let m E N with m < and Wi...., Wm E NN where LEN(w) S b for every j E [m].
Let = M1 Then, there exists a neural network F R - with width 4, depth 3m + and bit complexity b [log R. such that F satisfies
1. for every i E F(x,) WT1
2. for every IR with |x xz 2 for all [M], the output F(x) is either w, for some j E [m].
Most of the proof the same as Lemma A.4. from Vardi et al. 2022). and the only difference that we now examine how the function behaves outside of I1, ........2 ZN.
Proof of Lemma G.3
[zy. + 1] to construct feed-forward network F3 IR + R such that (x) = 1 for any x E [[=j Lay. + 1J], and
For any j E [m.] we use Lemma G.2 with a [ and
2. In particular, this means that 1.j k). Here j -- k may become bigger than N, and in such case j- l is replaced with N. Then, we define feed-forwart network Fj IR + Ry
Fj (a:) = 0 for any > [ary-e + 1J + or < [
Fy (2x) for any []
Bg (( ))
y y + wy ( - ()())
(152)
RR + R by F(x) (1))
F o Fi ((
)) For the
and the whole network F
verification of the correct behavior of the function F for inputs F1, ........... NN. and the analysis of its model complexity, we refer the reader to the proof by Vardi et al. (2022). Instead, we check the
output of for inputs outside of x; with i = I,..., N. For any input such that | 24 2 for all E []] there are two situations.
1. x [|æ Lzj- + 1J] for some j [m]: in this case, only (x) outputs and other sub-network Fg (z) with j' + output which results in F(x) Wj.
2. for any [m]] Irj for j E [m] and thus F(x) = O.
or x [z - holds: in this case, (x(x)
that the output F(x) for every & E R with Ja: xz (Vi = 1,..., N) is or wy for some j E [m].
Putting the above two together,
Lemma G.4 (Vardi et al. (002)) Let n E and let i,j E N with
< S n. Denote Tel- garsky's triangle function by (2)) (2(oR(22) - oR(42 4z 2)). Then, there exists a neural network F R² + R3 with width 5, depth 3 - 1), and bit complexity n + 2, such that 15(i-1)
+ JOFT
for any x. E N with LEN(x) & if the imput of is (1--11)
then it outputs:
2014
(b))
(b))
SHIZ
(2))
(x)
37

---

Published as a conference paper at ICLR 2025
Lemma G.5 Vardi al. 2.22)). There exists network F R² + R with width 2, depth 2 and bit complexity 2 such that
) if E 1v. + 1] and F (( = 0 if x > 3 or
x:
The following lemma is extension of the lemma by Vardi al. (2022), in that the outputs for unexpected inputs are also considered.
Lemma G.6. Let P, n, e E N. Let E N with (EN(u) = p and let w E N with LEN(w) = c- nh Assume that for any l, k (0.1,..., 13 with E k we have that BIN p-++1:p-(+11)() Then, there exists a network F TR3
BIN p-k+1:p-((k41 (u)
with width 12, depth 3n max (,, 6J + 2n + and bit complexity n max(p, c} + 2, such that for every x > there exist where [x] = BIN piNppptp(( then:
j E {0, 1...,..,
2 BIN ()++)p-(4)(())
W
(153)
u
In addition, if x satisfies (xt 1/2 - BIN p)tz)(+))(u)) for any j E (0...., n 11. then
F
-.
(154)
Proof. We follow exactly the same construction of a neural network by Vardi et al. (022). As such, for a detailed analysis of the depth and bit complexity of each network defined here, we refer the reader to the original paper and omit it here.
For each = O,..., n 1, we construct a network F; as follows, First, we use Lemma G.4 for and respectively to obtain two networks F and F" with width 5, depth at most 3 max p. c} and bit complexity at most n max(p, c} + such that
(il(it))))) 1)(()))))))
a
F
p)
2010 +
+ 20
2 20
2
2n 2mo01
u
2n
(155)
yan
BIN (p+1
()((++1)-e)
(b(ire)
(b-))
2weet
(((+1))-))
(156)
20007
BIN
1( (w)
Next, Lemma G.5 with inputs and BIN, (( to obtain the neural network F with width depth 2 and bit complexity at most P such that
F ((
(III - if BIN, PLN SI BIN pt + 1, if 1/2 - BIN +(N 1.
(157)
In addition, construct 1-layer feed-forward network F./ by
F.y )
on (x. nx(
20+1 20+1 + v).
(158)
38

---

Published as a conference paper at ICLR 2025
Putting the networks defined above and trivial modifications together, we define neural network F; such that F3 satisfies
vb(tm)
pb(bp)
Fg
ailb-e)
+ 2
gb(bre)
+ 201013
2n
(1((+))))))
1
200p
20
((+1)))
20
2
(1++1)))
w 2a.e
cm.
1p ga 2u
y + oR P. ((
(BBIN
x
20+1 20+1 BIN ((+)))
(159)
Finally, we concatenate Fg for each = 0,.............
n ,, F= G Fa-1 .... Fo H
to construct a network F IR3 - by
(160)
where G and are additional 1-layer feed-forward networks such that
T
2
H R
R3
2
(161)
U
2
2"
20
20
and G IR$ I outputs the fifth coordinate of the input. Note that with these configurations, it can be proved by induction that inputs F; for each = O,..., n -
are always of the form eq. (159).
Hereafter, we verify that the network F actually satisfies the desired behavior. Notice that the output of F' is expressed as
)
FP (( )) 2 20+1 ((
F
OR
==
(162)
there exist j E (0.1....,.
1} with | BBN, pi+tep((t))()))) the right-hand side becomes
4
w = aR (1 2c+1 20+1 BIN BIN)e+11(+))))(()))
n-1
Σ oR (0-2ct1 2c+1 20+1 BIN BIN ((eDee)
BIN (e
(163)
because BIN (1N 2c+1 holds for any i ....., 02...,
On the other hand, if x satisfies æ 1/2 - BIN (+)ll > for any j E (o..... 1}, the output of F becomes
)
F
w
Σ
on 2 BIN B
20+1
1=0
0,
(164)
desired.
n
39

---

Published as a conference paper at ICLR 2025
H EXPERIMENTS
In this section, we empirically investigate whether the memorization capacity real-world Trans- formers aligns with the behavior predicted by our theoretical analysis when varying the size of the dataset and the length of input sequences.
H.1 SETUP
We trained Transformers the next-token prediction setting on two real-world datasets and one randomly generated dataset of various sizes and evaluated the minimum network size required to memorize each dataset, plotting the results to examine the correlation between dataset size and network size. To validate our theoretical analysis, the architecture of the Transformer used in our experiments followed the same structure as the model described in Theorem B.I To be more precise, consider the following architecture:
N x Eout FFFF) F(UA), J(FF) Ein (x(1),:
(165)
where F(FF) and F/2 F.FF) are token-wise feed-forward layers (eq. (3)) stacked for #blocks blocks with the hidden dimension q 4m and embedding dimension m 2 Since the number of parameters in the model is approximately proportional to #blocks, we use #blocks as proxy for memorization capacity in our experiments by varying it to evaluate the minimum network size re- quired for memorization. The model was trained using the AdamW optimizer (Loshchilov & Hutter, 2019) with full-batch updates. To focus on the representational capacity of models and minimize the influence of optimization, we tuned hyperparameters such as a learning rate and warmup interval using Optuna (Akiba et al. 2019).
H.2 RESULTS
Validation of memorization with Transformers using single uniform-atttenion: We first validate that a single layer of uniform attention actually suffices for memorization. Specifically, we trained simplified Transformer defined in eq. (165). consisting of one uniform attention layer and two token- wise feed-forward networks, on two real-world datasets: MultiNLI dataset (Williams et al., 2018) from GLUE benchmark (Wang et al. 2018) and IMDb dataset (Maas et al. 201 l). For both datasets, the length of input sequences was truncated to 8, and outputs at the 0-th token were compared with labels using cross-entropy loss. While this setup does not correspond to next- teke-token prediction in the traditional sense, it aligns with the next-token prediction setting considered in this paper.
The results are summarized in figure for MultiNLI dataset and figure for IMDb dataset. Overall, experiments confirmed that as the number of blocks increases, the training loss can be reduced to nearly zero, and the accuracy tends to approach one. This outcome aligns with the predictions of Theorem B.I demonstrating that even single layer of uniform attention, when paired with an appropriate number of feed-forward networks, is sufficient for memorization.
Varying the dataset size while the sequence length is fixed: We next examined how the memo- rization capacity of Transformers changes when varying the dataset size while keeping the sequence length fixed. Specifically, we trained Transformers on datasets sampled from the MultiNLI dataset, where the sequence length was fixed at = 8 and the dataset size N ranged from 600 to 1700 in increments of 100. For each dataset, we determined the minimum number of #blocks required for the network to memorize the data. Here, a network was considered to have successfully memorized the dataset when the training error fell below a threshold of 0.01.
Figure 3 summarizes the evaluation of the memorization capacity of Transformers on MultiLII datasets of varying sizes. From this figure, we can observe the following two points.
1. Square-root scaling for small datasets: For smaller dataset sizes, particularly in the range of N = 600 to N = 1400. the memorization capacity of the Transformer scales approx- imately as VN. This behavior aligns well with the theoretical prediction of Theorem 4.1
The emhedding dimension was set to 2 so that it becomes difficult for models to memorize even small datascts.
40

---

Published as a conference paper at ICLR 2025
0.30
1,000
0:975
0.350
9.20
Training Loss
0.923
0:35
0.900
0.875
2.10
9.05
9.00
0.800
290
400
600
800
1000
200
400
600
800
1000
Takr ing Epocts
Tarring Epocts
(a) Training loss on MultiNLI dataset
(b) Accuracy on MultiNLI dataset
Figure 1: Training losses and accuracies of Transformers with #blocks = 1, 10, 20, 30, 40, 50 on a dataset of size N = 2000 with input sequence length n 8 sampled from MultiNLI dataset. Each model was trained using full-hatch gradient descent for 1000 epochs, and the best-performing model was selected after running two trials of hyperparameter tuning with Optuna.
0.25
1.00
#bacces 11
Rolocks =1
#blocks 10
Abtocks
0.20
#blocss 20
9.58
Wntockes 20
#bioces 30
#bioces 40
#olocdss -10
Training Loss
0.35
0.10
0..4
0.05
0.00
0.50
"Priring Eppccts
(a) Training loss on IMDb dataset
(b) Accuracy on IMDb dataset
Figure 2: Training losses and accuracies of Transformers with #blocks 1, 10, 20, 30, 40 dataset of size N 3000 with input sequence length n = sampled from IMDb dataset. Each model was trained using full-batch gradient descent for 1000 epochs, and the best pestt-erforming model was selected after running two trials of hyperparameter tuning with Optuna.
and Theorem which suggests that the memorization capacity of Transformers in the next-token prediction setting scales OvN.
2. Rapid increase for larger datasets: Beyond N = 1400, the memorization capacity ex- hibits a sharp increase, deviating from the earlier VN scaling. This phenomenon has also been observed in the experiments conducted by Kim et al. (2023). plausible explana- tion is that the bit-length of each parameter in the network is fixed during the experiments. As the dataset size grows, the precision of the parameters becomes insufficient for optimal memorization. Under this regime, the analysis of Transformers with limited bit complexity, as discussed in Appendix [F] becomes applicable, predicting that the memorization capacity scales linearly with the dataset size N.
Varying the sequence length while the dataset size is fixed: We also investigated how the mem- orization capacity changes when the size of a randomly generated dataset is fixed at N 500 and the input sequence length n is varied across 10. 100, 1000 and 10000. In this experiment, each word token is a 6-dimensional vector, with each element sampled independently from the uniform dis- tribution over the interval [0, 1). Similarly, each label is either +1 or -1, sampled independently from the Rademacher distribution. Using the mean squared error as the loss function, a network was considered have successfully memorized the dataset when the training error fell below thresh- old of e = 0.01. Surprisingly. the results showed that, for all sequence lengths, a Transformer with #blocks = 4 the smallest model capable of achieving memorization. An insight from this
41

---

Published as a conference paper at ICLR 2025
16
12
#blocks
10
600 800
1000
1200
1400
1600
Dataset Size
Figure Memorization capacity, that is, the minimum size Transformers required for memorizing MultiNLJ dataset with size N = 600..... 1700 in increments of 100. In this figure, the depth #blocks of the two token-wise ferd-forward networks J(FF) and ((FF) in eq. (165 is used as the variable on the vertical axis to control the size of the network. Each model was trained using full- batch gradient descent for 1000 epochs, and the best-performing model was selected after running trials hyperparameter tuning with Optuna.
perimental result is that, while the upper bound of the memorization capacity given by Theorem 4.1 has a gap of (log n) compared the lower bound in Theorem 4.2 real-world Transformers appear to align more closely with the lower bound of Theorem that is, the memorization capacity might be nearly independent of the input sequence length n.
42


Reviews:
[{"review_id": "oVRILdytKq", "reviewer": "Reviewer_VZKs", "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.", "questions": "See weaknesses above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper is well-motivated and clearly presented, making it easy to follow.\n- 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.\n- 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. \n- 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?\n - 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "guGD9Pmpkl", "reviewer": "Reviewer_ZZQ4", "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.", "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?\n\n- 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?\n\n- 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]$?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The proof sketch of Theorem 4.1 is clear and well motivated\n- 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\n- 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.\n\n- **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.\n\n- minor typo on line 462: viewd -> viewed", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "52f1DRhGYy", "reviewer": "Reviewer_6ZoB", "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.", "questions": "Following the weaknesses part, I have the following questions.\n\n1.  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?\n\n2.1 Is there any examples of real-world datasets that satisfy or approximately satisfy this assumption?\n2.2 Can authors discuss how the results might change if this assumption is relaxed?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The results seem promising and strong, because both upper and lower bounds are provided.\n2. 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.\n2. The paper can have more discussion on the separatedness assumption. For example, how realistic is the assumption?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fcc21EUg5I", "reviewer": "Reviewer_hFhi", "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.", "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.\n\n2. If the feed-forward layer is identified as the bottleneck, is there an approach to mitigate this through architectural changes or modifications?\n\n3. 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.\n\n4. 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.\n\nI am open to increasing the score for the paper based on the authors' responses and feedback from other reviewers.\n\nReferences:\n\na) Rajaraman, Nived, et al. \"Transformers on Markov data: Constant depth suffices.\" arXiv preprint arXiv:2407.17686 (2024).  \n\nb) Makkuva, Ashok Vardhan, et al. \"Local to Global: Learning Dynamics and Effect of Initialization for Transformers.\" arXiv preprint arXiv:2406.03072 (2024).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper addresses the important problem of quantifying the memorization capacity of transformers.\n2. 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.\n2. While the paper is theoretical, and experiments may not be essential, including experiments to support the claims would strengthen the work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)