Provable Memorization Capacity of Transformers
Authors: Junghwan Kim
Michelle Kim
Barzan Mozafari
Venue: ICLR 2023 poster (2023)
Keywords: Transformer
Expressivness
Memorization
Deep learning theory
contextual mapping
permutation equivariance

Abstract: Quantifying memorization capacity is essential for understanding the expressiveness and generalizability of deep learning model architectures. However, the memorization capacity of the Transformer architecture has yet to be explored. In this work, we present the first study of the memorization capacity of the Transformer architecture. We prove that Transformers are capable of memorizing $N$ sequence-to-sequence mappings of length $n$ with $d$-dimensional input tokens using $\tilde{O}(d + n + \sqrt{nN})$ parameters. Our theory supports memorization both with and without permutation equivariance, utilizing positional encodings in the latter case. Building on our theory, we also analyze the memorization capacity of Transformers in the sequence classification and language modeling tasks. To verify these theoretical findings, we conduct experiments analyzing the memorization capacity of Transformers in the natural language domain.

Published as a conference paper at ICLR 2023
PROVABLE MEMORIZATION CAPACITY OF TRANS- FORMERS
Junghwan Kim
Michelle YoungJin Kim CSE Department Michigan State University East Lansing, MI kimmi cl 60msu. edu
Barzan Mozafari
CSE Department University of Michigan Ann Arbor, MI kimjhje @ umi ch. edu
CSE Department University of Michigan Ann Arbor, MI mozafari ch, edu
ABSTRACT
Quantifying memorization capacity is essential for understanding the expressive- ness and generalizability of deep learning model architectures. However, the memorization capacity of the Transformer architecture has yet to be explored. In this work, we present the first study of the memorization capacity of the Trans- former architecture. We prove that Transformers are capable of memorizing N sequence-to-sequenec mappings of length n with d-dimensional input tokens us- ing O(d + VnN) parameters. Our theory supports memorization both with and without permutation equivariance, utilizing positional encodings in the lat- Building
theory, also analyze the memorization capacity of Transformers in the sequence classification and language modeling tasks. To ver- ify these theoretical findings, conduct experiments analyzing the memorization capacity of Transformers in the natural language domain.
INTRODUCTION
Transformer networks Vaswani al. 2017) have shown tremendous success in natural language processing tasks (Devlin et al.) 2019: Yang et al., 2019 Brown et al. 2020; Fedus et al. 2022) rapidly becoming the standard architecture for natural language modeling. The success of Trans- formers has also transferred to various other sequence and set modeling tasks, including image recognition (Parmar et al. 2018 Dosovitskiy al. 202T), semantic segmentation (Zheng et al.)
2021). video understanding (Akbari et al.] 2021 Bertasius et al.) (021), reinforcement learn- ing (Parisotto et al. 2020 Chen et al. 2021 Janner al. (022)) 3D point cloud processing (Zhao 2002)) protein structure prediction (Jumper et al. (2021) and automatic theorem proving (Polu & Sutskever 2020). Despite this success across various areas, the theoretical understanding of Transformers lags behind that of standard fully-connected networks.
The major strength of Transformers in their efficient scaling, which enabled through parallel token processing with parameter sharing and simple dot-product-based token interaction. Surpris- ingly, even though the parameter sharing and simple token interaction impose constraints on the function space of Transformers, Yun et al. (2020a) show that Transformers can approximate any continuous function from input to output sequences. However, their result focuses on the function approximation capacity with infinite precision, leaving the finite sample memorization capacity with finite precision unexplored. We note that universal function approximation does not automatically imply efficient memorization in terms of the number of parameters. Generalizing infinite precision results to the finite precision case is not straightforward and may not be possible in some cases. For example, Transformers are Turing complete only with infinite precision (Pérez al. 2019), not with finite precision (Dehghani al. 2019).
Understanding the memorization capacity of a model critical for choosing an appropriate model size. Practitioners often choose model size with enough representation capacity to achieve zero training loss (i.e., a size larger than memorization capacity). Moreover, the memorization capac- ity has generalization implications, 2019 Nakkiran et al. 202T). As the network size increases, generalization performance exhibits bias- hias-variance tradeoff until the memorization is possible and then improves monotonically after-
observed in the double descent phenomena (Belkin et al.)

---

Published as a conference paper at ICLR 2023
ward. Understanding the memorization capacity of Transformers requires answers to the following questions: How large should the size and precision of the Transformer architecture be to enable memorization of any given number of input-output sequence pairs? How does the memorization capacity of Transformers differ for various problem settings practical application scenarios?
In this paper, answer these questions by proving that Transformers memorize N sequences of d-dimensional tokens with length using O(d + + VnN) parameters. Our proof constructs permutation equivariant Transformers that can memorize all permutations of N input sequences. We extend this construction to the memorization without permutation equivariance by adding positional encodings. In addition, we derive the memorization capacity for sequence classification task from our proposed theory.
The key technical component of our construction is efficient contextual mapping, which requires only n self self-attention layers. Our contextual mapping also applies aparse-attention Transformers making fewer assumptions on sparsity patterns than Yun et al. (2020b) and[Zaheer et al. (2020). Fur- thermore, we present the generalization of contextual mapping to function approximation settings, vastly improving the parameter efficiency of attention layers compared to the selective-shifting based contextual mapping in Yun et al. (2020a).
Our main contributions are summarized as follows:
We prove the memorization capacity of Transformers for sequence-fo-sequencee mappings with and without permutation equivariance. We analyze the memorization capacity in other standard task settings, such sequence classification and language modeling.
We show that the efficient contextual mapping presented in our theoretical analysis extends sparse attention settings and improves the function approximation results.
We provide experiments validating the memorization capacity of Transformers for token classification and sequence classification tasks.
1.1 RELATED WORKS
Memorization capacity. Characterizing the memorization capacity of neural networks has been
an active research area with a long history (Baum) 1988 Sontag 1997, Huang Babri 1998]
Huang 2003 Zhang et al. 2017 1. Yun et al.] 2019 Bubeck et al. 2020: Vershynin 2020 Rajput
ct al. 2021 Park et al. 2021 Vardi et al. 2022). Recently, Park et al. (2021) constructed neural
networks with O O22
parameters to memorize N data points. They bypass the ((NN) lower
bound in Sontag 1997 by assuming simple separation (i.e., ||x; xj 2 8, Vi + j). Vardi
al. 2022 improve further, showing that ((N1/2) parameters sufficient. They also prove the matching lower bound of (((1122) through the VC-dimension analysis.
Inspired by Park et al. 2021) and Vardi et al. (2022), our construction assumes a similar separation, but between pairs of distinct tokens, not the whole sequence pairs. (See Definition 3.] and the discussion that follows the definition.) In addition, our construction uses the same pipeline of projection, string matching, and bit extraction as in Vardi et al. (2022). However, we introduce an additional critical step: efficient contextual mapping to complement projection by summarizing all token information via self-attention layers.
In contrast to the extensive results fully-connected networks, there are few studies on the memo- rization capacity of specific modern architectures. Hardt & Ma show that the residual network with ReLU activation and O(N) hidden neurons memorize N data points under the separation assumption. Nguyen & Hein (2018) show that the convolutional network with O(N) hidden neu- memorize N data points. To the best of our knowledge, there is existing literature the memorization capacity of the Transformer architecture.
Transformer expressivity. Given the recent empirical success of Transformers observed across multiple areas, several papers have studied the expressivity Transformers. Yun et al. (2020a) establish the first universal approximation theorem for Transformers, and the result later extended to sparse-attention Transformers (Yun et al. 20206) Zaheer et al.) (020)) and Transformers with
hard constraints (Kratsios et al. 2022 All these results study the function approximation but not
2

---

Published as a conference paper at ICLR 2023
the finite sample memorization in our paper. We also note that construction reduce the number of self-attention layers in the function approximation setting.
There are other lines of studies focusing on different aspects of the representation capacity of Trans- formers. Some papers aim to characterize the representation capacity of a single self-attention layer. Bhojanapalli et al. 2020) suggest that the small size of attention heads limits the rank of self- attention matrix. Dong et al. (2021) show that the rank of self-attention decays exponentially when self-attention layers are composed without skip-connections or feedforward layers. Likhosherstov et al. (02T) show that a fixed self-attention module can approximate any sparsity pattern.
Other papers investigate tradeoff between width and depth since it is crucial issue when scaling Transformers, Levine et al. (2020) demonstrate the depth efficiency in modeling feature interaction through the separation rank analysis of Transformers. [Wies al. (202I) identify the rank of the input embedding matrix as bottleneck for the network width's ' contribution to expressivity.
Finally, Vuckovie al. (2020) and Kim al. (202T) study the Lipschitz smoothness of attention operations. Edelman et al. (2022) derive the norm- norm-based generalization bounds from the Lipschitz smoothness of norm-bounded attention layers to analyze the inductive bias of attention layers. Wei et al. (2021) study the expressivity of Transformers under the constraint of statistical learnability. Our memorization study provides a complementary understanding of the capacity of Transformers.
PRELIMINARIES
This section establishes the notation and defines the Transformer architecture.
2.1 NOTATION
Denote the number of input-output pairs as N, the number of output classes as C, the token em- bedding dimension as d, and the sequence length as n. We use Ö() to hide logarithmic factors and O()) to hide constant factors. We use on to represent the ReLU activation function. We let as and oh be the softmax and hardmax operators, respectively. These operators take a matrix as an input, apply softmax/hardmax columnwise, and output a column stochastic matrix of the same size.
For m N, we define [m] = {1, m.. We use [S] to denote the number of elements in a set S. We denote set function
upper case caligraphic letter, matrix by upper case bold letter and a vector by a lower case bold letter. We denote the standard unit vector with all but i-th coordinates 0 as e,, the dimensional vector with all coordinates 1 1m and the m dimensional vector with all coordinates O as O,- For a vector æ we represent its i-th entry as æ[i] and its Euclidean norm ||az|]. For matrix X, we use X [i, j.. X ii, and {(: represent (i,j) j-th entry, i-th row, and j-th column, respectively. We use XX FF to represent the Frobenius norm of the matrix X.
2.2 TRANSFORMER ARCHITECTURE
We define a Transformer N dR - 1 with input and output embedding mappings:
x of depth I. as a composition of Transformer blocks
N Eout o Ft. ... o F2 E,n.
where each Transformer block Fi x + mm is sequence- sequence-to-sequene cooennce-ttssepeee function consisting of two subblocks: self-attention subblock and a tokenwise feedforward subblock. The input em- bedding block m,, dx + mm and the output embedding block Eout m - R I-layer tokenwise linear mappings.
The self-attention subblock represents the interaction among tokens. Formally, given input Z Rmxn, the self-attentio subblock F ((A) with h heads and head size & computes
((AA) Z) = Z
(W!.))
1, z) os
(V)
1,3 (w
W 1,2
2) 2)
where W(OO) E ma and ((()) ((K)) the self-attention subblock. We include skip-connection in the self-attention subblock.
(K) (! R are the weight matrices parametrizing
3

---

Published as a conference paper at ICLR 2023
The feedforward subblock processes each token independently in parallel by applying two feed- forward layers. Given an input H Rmmx,, the feedforward subblock F(FF) with dimension computes
( = = W on (W H b + b IIT n (
where w/22) E Rm b1²) E RT, w( E R and b1"h E I" parametrize the feedforward subblock. The feedforward subblock also includes skip-connection.
Finally, the Transformer block composes two subblocks as
Fi(Z) {
(Z)
Unlike the original formulation in Vaswani al. (2017). our definition excludes layer normalization as Yun et al.J (2020a) to simplify our analysis. Since layer normalizations mainly contribute to optimization without much effect on expressivity, our definition still captures the representation aspect of the Transformer architecture.
Since each Transformer block consists of a fixed number of layers even in the most fine-grained sense, we use the number of blocks L as the depth of the network. We define the width of the network as max(m, kh, q
| The number of parameters is the number of non-zero weights in our network. We note that the single parameter is reused times for sequence length n. but still is counted as one parameter. The bis complexity of the network is the maximum bit complexity of its weights, where the bit complexity adopt these definitions of the number of parameters and the bit complexity from the convention in the VC dimension literature (Bartlett al., (019) and the recent paper the optimal memorization capacity of fully-connected networks (Vardi et al. 2022).
weight is the number of bits required to represent it. We
MEMORIZATION CAPACITY OF TRANSFORMERS
In this section, we describe the problem setting and present main theorem the memorization capacity of the Transformer architecture. Then, we sketch the proof for our main theorem and discuss the memorization capacity of Transformers in other standard task settings.
3.1 PROBLEM SETTING
We consider the memorization of N input-output sequence pairs (X(1), y (X(M, (N)) where each input ( E Rdx= is a sequence of n token vectors in dimension d. Each output ( [[]] a sequence labels where each label [11 k] assigned a token X ( k).
We define the context of each input sequence ( as y(t) = {o Rd v = x (;) k] for some e [n]]. We define the vocabulary V UiE[N] ((0) as the set of all tokens appearing in the input sequences. Note that [V] < n.V.
As pointed out in Park al. (2021) and Vardi al. (2022), we must assume some conditions on the dataset to bypass the lower bound in Sontag 1997) and memorize N data points with o(N) parameters. We present a natural generalization the separation condition defined in Vardi et al. 2022) to sequence modeling settings.
Definition 3.1. Let 2 1,0 < Let X(), ........ (N) E d be N input sequences with vocabulary V. We say that (1) .... X (N) are tokenwise (r, -esepaaate if
(NN)
1. |||| < r for all v E V and
2. 0 llll 2 8 for all v, e' with + v'.
This condition requires (1) each token to have a bounded norm and (2) each pair of distinct tokens to be separated. We note that the tokenwise separation is stronger condition than the separation of input s However, the condition better captures many practical settings where the number of tokens in the vocabulary is much smaller than the number of input sequences.
We recall that m is the embedding dimension, is the number of attention heads, k is the attention head size, and y is the feedforward dimension.
WWhe (
x (N) E d are distinct and (r, 6)-eparrtte,, then (1) ||x (Xx S Sr/n for [N]
and (2) (x) x [ || 2 5 for j E [N] with
4

---

Published as a conference paper at ICLR 2023
For the permutation equivariant mappings, need the following label consistency condition. Definition 3.2. Let (X(1), y()),.. .... (X(N), ((D)) Rdxa [C]1*a be N input-output pairs of sequences. We say that (X y (X(N) Y(N) consistently labeled if
x 1 1 RJ = x (:., implies (1 (1,
for every E [N] and k, 1 [n].
We emphasize that we impose this condition only on the memorization with permutation equivari- ance but not on the memorization without permutation equivariance The condition implies that two identical tokens appearing in the same context should have the same label. Consider permu- tation equivariant mapping IRdxn + Rix= and an input sequence X Rdxn, Define X' by swapping two tokens [:; ki] and X [., I] in X. Then, we have [(X)(1, kj [1,) due to the permutation equivariance. However, if two tokens X|: k] and [[:,4] were identical, then X = X and consequently F(X) kj = [1X)
3.2 MAIN RESULTS
We now present our main theorem on the memorization capacity of Transformers.
Theorem 3.1. Let N, d, n, E N and r 2 1,0 < ô < 1. Let (X y """""""""""""" (X(M), y())
Rdxa x [C be input-output sequence pairs of sequences where input sequences are distinct and tokenwise (r, 8)-spparated.
1. (With permutation equivariance) Suppose that contexts v(i) are distinct and sequences are consistently labeled. Then, there exists a Transformer network N Rdxx - Blxn such that
N
(0) = -
for every i [N] and for every permalation matrix P Raxn.
2. (Without permutation equivariance) There exists Transformer network N Rdxn >> IR Ixx and positional encoding E E dRd such that
N(X + +))= (1)
for every i E [N]. In both cases, the Transformer N has width 16 (m 8, h k = and 16), depth
o n nN log(nN)+
nN log(nN) max(log C, log R}
and bit complexity bounded by
nN max(log Co R}
o
log d + V log(nN)
where denote R 8000 2
Theorem 3.1 shows that O(d - + vnN) parameters are enough to memorize N sequence-to- sequence mapping of length n with token dimension d since the initial embedding layer has pa- rameters and the rest layers have constant number of parameters. We provide the proof sketch in Section 3.3 and the full proof in Appendix
Remark 3.2. Extensions to real vector outputs. Some application scenarios of Transformers quire real vector outputs. As proposed in Park et al. (021 and Vardi et al. (002), extension to real scalar values easily achievable by using O(2) classes when the output has bounded range. More concretely, we partition the output range into 6-length intervals and match each class to one partition to perform regression with error per token. This replaces log Theorem 3.1 with log( !
Thhe general gequence-to-squueee mappings do not satisfy the condition. For example, the same word appearing multiple times in sentence may have different meanings.
5

---

Published as a conference paper at ICLR 2023
Similarly, extension to vector outputs is possible when the output has a bounded domain. Suppose that we aim to minimize the tokenwise L2 distances in dimension p. We partition the output range
into yp -ength cubes and match each class to one cube. Then, we use O (vP) classes to perform
regression with emror c per token. This construction replaces log C' in Theorem 3.1 with plog(()) ")-
Remark 3.3. Large width and fixed bit complexity. Theorem 3.1 uses fixed width and large bit complexity to minimize the number of parameters. However, common approach to scaling Trans- formers to increase width (Levine et al. 2020) while using the same number of bits per param- eter. Using similar argument as Vardi et al. (2022), extend Theorem 3.1 to the cases with larger width and with bounded bit complexity. When the larger width allowed, Transformers of width Dn)N/L²), depth On + L) and bit complexity ((L) memorize the same dataset for some I. S VnN. When the bit complexity is bounded, Transformers of width O(1), depth Ö(n O(n+nN// and bit complexity (BB) memorize the same dataset for some B S V VnN. We provide the formal theorem and the proof in Appendix
Remark 3.4. Tightness in the order of bit counts. Suppose the token dimension d and the sequence length n. both O(N). Theorem 3.1 shows that Transformer memorizes N input-output sequueeeeessuuueee pairs using Ö(vnN) parameters of bit complexity Ö(vn.N) which sum up to (nn.) bits. Without any additional assumption on dataset, representing models that memorize all CRN possible label for N input sequences requires (2nNV) bits. Thus, Theorem 3.1 1 is tight upto logarithmic factors in the order of bit counts.
Remark 3.5. Comparison against fully--nnncee ReLU networks. With a slight modification of results in Vardi et al. (2022 fully-connected ReLU networks require Ö(dn + vnN) parameters.
The dependence on the number of data points is the same as Ö(Vn.N. However, for the permutation equivariant case, Transformers memorize all permutations of each input sequence. That is, Transformers are capable of memorizing at most n! times more data points the cost of reusing each parameter n. time
Our construction has a better dependence on d and n: O(d+n) n)) for Transformers and O(d)) (dd) for fully-connected ReLU networks. Transformers exploit the structure of sequence data through parameter sharing. As a result, Transformers do not need O(d) parameters to read all dn values in the input sequence.
We note that this comparison is not completely fair because (1) our result makes a slightly stronger assumption of the separation between tokens than the separation between whole input; and (2) fully- connected networks are not designed for permutation equivariant datasets. However, the difference in assumption affects the final bound only the logarithmic terms. Furthermore, there is no known result on the memorization capacity of other permutation equivariant architectures.
3.3 PROOF SKETCH
We outline the proof of Theorem? 3.1 A more formal statement of each stage with detailed proof is in Appendix A Our proof adopts the approach from Vardi et al. (2022) and shares similar steps. We discuss technical novelty after sketching the main ideas of our proof.
Our proof constructs a Transformer code as a "contextual token id" that identify each token in the different context of sequence. We ensure that the "contextual token id" is permutation equivariant and that the "contextual token id" is uniquely assigned each token each context. Then. the last two stages map each "contextual ken id" to the corresponding label using the bit-extraction network adopted from Vardi et al. (2022). We describe key ideas of each stage.
stages. The first two stages assemble input values and en-
"There are C' possible lahels for n tokens in sequences. 'We consider inputs tteddimesiionl flattened vectors and outputs values from [C"]. We choose different balancing point between stages and 3 in their construction balance the extra factors.
A similar benefit of Transformers has been previously observed in the function approximation Yun et al. 2020ay: the number of parameters is reduced by (n 1). times to successfully approximate permutation equivariant functions.

---

Published as a conference paper at ICLR 2023
Tokenwise Projection. We project each token vector to scalar token id while keeping distinct tokens well separated.
Stage 1.
Stage 2. Contextual Mapping. We compose sequence id linear combination of token ids. Weights of each id in the linear combination depend on the order of token id within each sequence. The resulting sequence ids are permutation invariant. We concatenate the token id and the sequence id to obtain a contextual token id.
Stage String Lookup. We partition all n. contextual token ids into intervals, each con- taining the same number of ids. We construct two encoding numbers for each interval by concatenating all corresponding contextual token ids and token labels. Then, we find which group the contextual token id falls into and retrieve the corresponding encoding numbers.
Stage 4. Bit Extraction. We extract each contextual token id and token label from the encoding numbers. If the extracted contextual token id agrees with the one composed in stage 2, then output the corresponding token label.
Our technical novelty in this proof is in (1) the implementation of the contextual mapping in stage using self-attention subblocks and (2) generalizing stages 3 and 4 to incorporate skip connections. Remark 3.6. Contextual mapping. The number of self-attention subblocks that proof uses is proportional to the length of the sequence n but independent of the number data points N. This requirement is in striking contrast the
swwree-shiiiinnbsssee contextual mapping from Yun et al.
2020a). which requires (4) dm layers for shifting each grid cell of side length ò. In the memorization setting, may remove layers for grid cells without any data point. But the selettie-shifftin--aase contextual mapping would still need n. self-attention subblocks, which alone is already larger than the number of required layers our result. In contrast, our efficient contextual mapping uses n layers, improving all the above when 6, N > 1. We showcase the benefit of our contextual mapping in Appendix C Specifically, we show that our contextual mapping are capable of incorporating sparse self attenttnnn settings with minimal parameter overhead. Moreover, we show that the same idea of our contextual mapping can be applied in the function approximation setting.
Remark 3.7. Parameter contribution. In our construction, self-attention subblocks contribute only (n) parameters while feedforward subblocks contribute O OvvnN). When n < N. most of the pa- rameter count comes from feedforward subblocks. Although the self-attention layers play a critical role in contextual mapping, we do not need many of them. This explains the model design in prac- tice, where more than half of the parameters are in the tokenwise feedforward subblocks (Vaswani et al. 2017: Devlin et al., 2019, Brown et al. 20200 Geva et al.) [021). Moreover, Mandava et al. (2020) observes that the self-attention layers can be further reduced without much performance loss.
MEMORIZATION CAPACITY IN OTHER TASKS
In this section, we generalize Theorem 3.1 to analyze the memorization capacity for other standard task settings. We consider sequence classification in Theorem 4.1 masked language modeling in Theorem D. and autoregressive language modeling in Theorem D.2 Due to the limited space, for- mal results on language modeling tasks and the proofs of the theorems are provided in Appendix D
4.1 SEQUENCE CLASSIFICATION
In sequence classification, assign single label ())) [C] for each input sequence ( We present our theorem for sequence classification task.
Theorem 4.1. Let N,d,n,
n. C e and r 2 1,0 < 8 < 1. Let ()),
(N), g())
Raxn x xC] be N input-output pairs of sequences where input sequences are distinct and tokenwise (r, --eeparrtt
1. (With permutation invariance) Suppose that contexts (i)) distinct. Then, there exists Transformer network N dx + *R Rlxn such that
7 appropriate choice of positional encoding hypasses permutation invariance hy collecting token the position order, See the second part of Theorem I for the result and Section A.5 for the details. "The most practical settings fall into this regime.
7

---

Published as a conference paper at ICLR 2023
for every i [N], E [n] for every permutation matrix P E R
2. (Without permutation invariance) There exists Transformer network N: R2 + Rixa and positional encoding E Raxx such that
N k]] yy)
for every i E IN].EEl [n].
In both cases, the Transformer N has width 16 (m 6, h and 16), depth
0 n 1 v Nog +
N log
max[log log R8}
V
and bit complexity bounded by
N log N
o
log d +
max C, log R}
where denote R 80000~8- 2 N
Theorem 4.1 shows that Öd + n- v N) parameters are enough to memorize N sequence classifica- tion examples of length n. with token dimension d. Compared to the eqquence-to-squeene mapping, there is v factor of saving in the last term.
4.2 LANGUAGE MODELING
We consider two language modeling tasks commonly used for pre-training Transformers: masked language modeling and autoregressive language modeling. We consider the memorization of all possible length n sequences obtainable from the given corpus of length T. The input is embedded in the d-dimensional space while the output is mapped to one V tokens in the dictionary.
The memorization of masked language modeling requires Ö + vn parameters (The- orem D.I while the memorization of autoregressive language modeling requires Ö + n + VT) parameters (Theorem D.22). Compared the autoregressive language modeling, the masked lan- guage modeling has the additional factor n m that comes from memorizing all masking patterns separately and the factor n that comes from memorizing all masked tokens instead of one next token. For more details on the settings and formal statement of the result, we refer to Appendix D
5
EXPERIMENTS
5.1 EXPERIMENTAL SETUP
We complement our theory with experiments on real-world dataset. We train encoder-only Trans- former models (Vaswani et al.) 2017) on token classification task where each token is assigned label as Theorem 3.1 and sequence classification task where each sequence assigned a label as in Theorem 4.1 We study the relationship between the memorized dataset size and the model size.
For token classification, we use 14,000 randomly selected examples among 14,041 training exam- ples in the named entity recognition dataset from CoNLL--2003 (Tjong Kim Sang & De Meulder, 2003). For sequence classification, we use 50,000 randomly selected examples among 392,702 training examples in the MNLI dataset from GLUE benchmark (Wang et al. 2019).
We vary the model size through the embedding size m while fixing the number of layers as 6. We fix the number of attention head as h = 12, the embedding head size ratio m/k h 12 and the feedforward to embedding size ratio as q = as commonly done in practice. More details experiments are in the Appendix F
5.2
RESULTS
Figure shows heatmaps of training errors as the dataset size and the model size vary. There is clear trend that the training emoor is smaller (darker in color) for the smaller dataset size and the
8

---

Published as a conference paper at ICLR 2023
0.08
14000
0.40
50000
12900
0.07
45000
11200
0.30
9000
8400
0.08
35003
0.25
Dalasel
0.04
7000
25000
5600
20000
1.15
4200
0.00
15300
0.10
2600
0.01
10300
5000
408
S16
92
860
888
708
96
192
288
384
460
578
672
788
0.00
Model Embedding Size
Model Embedding Size
(a) Token Classification
(b) Sequence Classification
Figure 1: Heatmaps of training errors. We show the color-coded training errors as the dataset size and the model size vary. The dataset size represented in the Y-axis while the model size is represented in the X-axis the embedding size. The training error tends to get better for the smaller dataset size and the larger model size.
4.0
2.5
1.5
2000
10000
12000
14020
10000
20200
30003
40000
50000
Dataset Size
Dataset Size
(a) Token Classification
(b) Sequence Classification
Figure 2: The number of parameters required for memorization. We show the number of pa- rameters (Y-axis) of the smallest model that achieves less than 0.005 training error for each dataset size (X-axis). We observe linear trend in the number of parameters with a subtle concavity as our theory predicts, but only at the low data regime.
larger model size. To the clearer relationship between model size and the dataset size, plot in Figure the minimum model size that achieves less than 0.005 training error for each dataset size. In general, there is linear increase in the model size the memorized dataset size increases. We also observe a slight downward curvature (concavity) as our theory predicts, but only at the low dataset size. We conjecture that the linear trend may be due the fixed depth and bit complexity during the experiments. See Remark 3.3 for the discussion of this bounded depth and bit complexity regime and see Appendix B for the formal result in these regimes. Indeed, Theorem B.1 and Theorem B.2 predicts linear dependence of the model size in the dataset size.
CONCLUSIONS
In this paper, prove that Transformers are capable of memorizing N length-n sequence-to- sequence mappings with O(d+ n + VnN) parameters. We extend our theory to analyze the mem- orization capacity of Transformers in other standard task settings. Our proof constructs a contex- tual mapping with (n)) self-attention layers, which significantly improves the previously proposed selective-shee selective-shifting baastiin---aeddddddd contextual mapping in terms of parameter efficiency. Finally, we provide experimental results that verify theory.

---

Published as a conference paper at ICLR 2023
REFERENCES
Hassan Akbari, Liangzhe Yuan, Rui Qian, Wei-Hong Chuang, Shih-Fu Chang, Yin Cui, and Boqing Gong. Vatt: Transformers for multimodal self-supervised learning from raw video, audio and text. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 24206-24221. Curran As- sociates, Inc., 2021. /RLp: http:/eepss..ppeeee///////// pt/pppppr/222///11eee 25222366644322333666666666333333
Paper. pdf
Peter L. Bartlett, Nick Harvey, Christopher Liaw, and Abbas Mehrabian. Nearly-tight vc-dimension and pseudodimension bounds for piecewise linear neural networks. Journal of Machine Learning
Research, 20(63):1-17, 2019. URL
htm
Eric B Baum. On the capabilities of multilayer perceptron............ Journal of Complexity, 4(3):193-215, 1988. ISSN 0885-064X. doi: http://dooaogg//l.oooooosssssss2 www. www.sciencedirectt
URL https:
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconeiling modern machine-
learning practice and the classical bias-variance trade-off. Proceedings of the National Academy
of Sciences, 1 116(32):15849 15854, 2019. doi: 10.1073/paas.1990070111606 1903070116.
URL https:
//ww www. pnas htlg://o//aaae /ato/////////// 1073 ppa../a....... 1903070112
Gedas Bertasius, Heng Wang, and Lorenzo Torresani. Is space-time attention all you need for video understanding? In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 324. PMLR, 18-24 Jul 2021. URL https: /tt.o//nnnntteeeoeeoeeeeeennnnnn bertas us2la att.smaa..mml.....mm
mlr. wttpsmess/.pr..e3....
Srinadh Bhojanapalli, Chulhee Yun, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Low- rank bottleneck in multi-head attention models. In Hal Daumé III and Aarti Singh (eds.), Pro- ceedings of the 37th International Conference on Machine Learning, volume 119 of Proceed- ings of Machine Learning Research, pp. 864-873. PMLR, 13-18 Jul 2020. URL https: /pproceedings.o www.:sas.........mmmmmm.mm.m..em svieess/v199 janapa 1200...... html
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Ma- teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Asso- ciates, Inc., 2020. URL 1457c0d6bf cb4 1964996440 we1.18fbacc12fff6 Pawe:e22oeaaaaar..ppepe Paper. pdf
neurips wttpccpppp///////......
Sebastien Bubeck, Ronen Eldan, Yin Tat Lee, and Dan Mikulincer. Network size and size of the weights in memorization with two-layers neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 4977-4986. Curran Asso- ciates, Inc., 2020. URL
http::occceeiino..eeiiii.e.eeese ce/ ctt.cc//ppp/p/ee ptcp/pppe//220//i1ee ppdf
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 15084-15097. Cur- ran Associates, Inc., 2021. URL httpttteep/ppoooeeiiiitetooo ngs neurips cttpcc//pppaaa pccpppper/2211 fi
02 o2w2b5231 057 00wp Paper .pdf
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2019. URL https:
//openreviee ew. www. ddw.mmdddd-oyoodoodooo R9Y7
10

---

Published as a conference paper at ICLR 2023
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Technologies, Volume (Lng and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 18853/v1/N119-142)) URL https:
Yihe Dong, Jean- cann-Batistte Cordonnier, and Andreas Loukas. Attention is not all you need: pure attention loses rank doubly exponentially with depth. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceed- ings of Machine Learning Research, pp. 2793-2803. PMLR, 18-24 Jul 2021. URL https:
//proceedings..o http://wwww........eeeeeese.eeee
www.ssevl33e http:/33//ongna.aheeeca html
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 recogni- tion at scale. In International Conference on Learning Representations, 2021. URL https: //openzeview feree//ooumiid wow.aaidd-Y.ccaaaa
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and vari- able creation in self-attention mechanisms. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Con- ference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 5793-5831. PMLR, 17-23 Jul 2022. URL https: http://ppooeeengg...oe mlr...mmrr.. www.s/wwessl2
edelman22a. html
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1
2022. URL CtRphhtt::///00 http://mmmr.rrrpppppe/////////// 20998981119911.11001101 html
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 5484-5845 5495, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. 186518/6/1111 10.18653/1122221.co hwwp///..mmmmmmii...4..mmm URL https://ach antho w
oro.or/////02 /2021 www.lemlppemme. in. 446
Moritz Hardt and Tengyu Ma. Identity matters in deep leamning.
In International Confer- net/ forum?id
ence on Learning Representations, 2017. URL
EXXBORLXX
Guang-Bin Huang. Learning capability and storage capacity of two-hidden-layer feedforward net- works. IEEE Transactions on Neural Networks, 14(2):274-281, 2003. doi: 10. 1009 809401.
Guang- -Bn 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 Networks, 91)::224-22 229, 1998. doi: 10. (0.100972.6550455
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big se- quence modeling problem. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 1273-1286. Curran Associates, Inc., 2021. URL http::ttt///t/ h/tpcceee///rcceeeeeeee ngs www.innn..nnnnnnnneeeeeee www.ris...... cc/ Scc/paper// (apperr/2221/ /21///22211111/e1 filele////e htt.///ee999fffeee/ wttp//99ffebbbb.. 442000044442233.. 444444422233444553344444443
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(78733:583-558) 2021.
Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 5572-5571 mlrr. http://wwww....../.............. wwwps//vwe//........e
PMLR. 18-24 Jul 2021. URL
html

---

Published as a conference paper at ICLR 2023
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
orq/ aa////aa/1112 6980
Anastasis Kratsios, Behnoosh Zamanlooy, Tianlin Liu, and Ivan Dokmanić. Universal approxima- tion under constraints is possible with transformers. In International Conference Learning Representations, 2022. URL
Yoav Levine, Noam Wies, Or Sharir, Hofit Bata, and Amnon Shashua. Limits to depth efficien- cies of self-attention. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 22640-22653 Curran As- sociates, Inc., 2020. URL
Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self- attention matrices. CoRR, abs/2106.037664 2021. URL https//caaaerr/2///1eeee
Paper.
03764]
Swetha Mandava, Szymon Migacz, and Alex Fit-Florea. Pay attention when required. CoRR, 0454,, 2020. (RIpps: wrw. wtwpgrggg//ggoo../oo//o 04534
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 Mechanics: Theory and Experiment, 2021(12):124003, dec 2021. doi:
URL
https://doo www.orgg111. 1001111888--488888881 ac/344
Quynh Nguyen and Matthias Hein. Optimization landscape and expressivity of deep CNNs. In Jennifer Dy and Andreas Krause (eds.), Pruceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Reseurch, pp. 3730-3739. PMLR, 10-15 Jul 2018. URL htthttss:/pproceeiinngg /ttp:///rrroceeeeeeooeeeeeeee/// wtt.ppers///////////// 18a html
Emilio Parisotto, Francis Song. Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphaël Lopez Kaufman, Aidan Clark, Seb Noury, Matthew Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learning. In Hal Daumé III and Aarti Singh (eds.). Proceedings of the 37th International Conference on Machine Learning, volume 1 119 of Proceedings of Machine Learning Research, pp. 7487-7498. PMLR, 13-18 Jul 2020. htt.:tttps.//prceeed.aie
www.srppress/vieee pttp:///m..ppeeeee/////a//r.reee wotp:///o.aaotooooaaaao html
Sejun Park, Jaeho Lee, Chulhee Yun, and Jinwoo Shin. Provable memorization via deep neural net- works using sub-linear parameters. In Mikhail Belkin and Samory Kpotufe (eds.), Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pp. 3627-3661. PMLR, 15-19 Aug 2021. URL httppss/pprceeeinggg.. press/v134/parccec
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 4055 4064. PMLR, 10-15 Jul 2018. URL mlr wwt.ppres/v//ppaa. http://a.maaaaa......... 8tmh
Stanislas Polu and Ilya Sutskever. Generative language modeling for automated theorem proving. CoRR, wtw/2009.0033999 2020. Stl.hhtts:///rxiiss htt.s/orrg/bb///0000 htt.tppp://arra..orrr///////3 03393
Jorge Pérez, Javier Marinković, and Pablo Barceló. On the turing completeness of modern neural network architectures. In International Conference on Learning Representations, 2019. URL
https://apenrrviiee net htt.pp/aeeeeeeeeeeoeeeeooooroaaa
Shashank Rajput, Kartik Sreenivasan, Dimitris Papailiopoulos, and Amin Karbasi. An expo- nential improvement on the memorization capacity of deep threshold networks. In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34. pp. 12674 12685. Curran Asso- ciates. Inc., 2021. / nww.iiissnnn.nnnn..enn...iniiiie http://cecc.......e.oooo///////a 69dd2ef 26aa
pd
12

---

Published as a conference paper at ICLR 2023
Eduardo D. Sontag. Shattering All Sets of 'k' Points in "General Position" Requires (k 1)/2 Parameters. Neural Computation, ((2):337-348, 02 1997. ISSN 0899-7667. doi: 10.1162/neco. 1997.9.2.337. URL https://doi.ccccmc /ttp://....rr/11111112//..o..... http://ecco..99....o
Erik Tjong Kim Sang and Fien De Meulder. Introduction to the CoNLL-2003 shared task: Language-independen named entity recognition. In Proceedings of the Seventh Conference on Natural Language Learning at HLTTNAACI. 2003, pp. 142-147, 2003. URL https:
Gal Vardi, Gilad Yehudai, and Ohad Shamir. On the optimal memorization power of reLU neural networks. In International Conference on Learming Representations, 2022. URL https://
openreview net/ forum? MKIMiidMMMMMeee
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL
neurips.....u......o............ Paper. ppf
Roman Vershynin. Memory capacity of neural networks with threshold and rectified linear unit activations. SIAM Journal on Mathematics of Data Science, 2(4): 1033, 2020. doi: 10. 1137/20M1314884 URL /dopi.::////..o http::::///....oo/11..... 1ttt:/100.13//2MMM1 314884 314884]
James Vuckovic, Aristide Baratin, and Remi Tachet des Combes. A mathematical theory of atten- tion. CoRR, hins2207.002876 2020. btt.htttps://arx wwtpiiiiiorggaasaaooeeoooo atrrr//bb/2000..o. 02876
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: multi-task benchmark and analysis platform for natural language understanding, In International Conference on Learning Representations, 2019. htLhttp::
/ opp.reviee..e...
net/Forum?id=r:JA
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. CoRR, abs/210..11666,,, 2021. URL https:
/aaxxvv.org/bblll. 21t.b////111111..
Noam Wies, Yoav Levine, Daniel Jannai, and Amnon Shashua. Which transformer architecture fits my data? a vocabulary bottleneck in self-attention. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 1110-111181 PMLR, 18-24 Jul 2021. URL mlr.
https://proceedingg wwt.mmm.........1...
wies2la..e.. htmll
Zhilin Yang, Zihang Dai, Yiming Yang. Jaime Carbonell, Russ Salakhutdinov, and Quoe V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bue, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32, Curran Asso- ciates, Inc., 2019. URL https: http:////occeeeeesseeee pdf
http.//00/..///ttp///////oodddd neurssn wttp:rrpss...pppperr//0..... cctppppee//19///ieeeeeee dc6a7e655d7e5844 0wt.eee666633999776666665
Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Small relu networks are powerful memorizers: a tight analysis of memorization capacity. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett (eds.), Advances Neural Information Processing Systems, vol- ume 32. Curran Associates, Inc., 2019. URL https: http:s/:////cceeed/e/// neurips... cc/
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal upproximators of soquence-to-sequenee functions? In International Con- ference on Learning Representations, 2020a. URL
net forum?
id=ByxRMONEvr
Chulhee Yun, Yin- Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. O(n) connections are expressive enough: Universal approximability of sparse
13

---

Published as a conference paper at ICLR 2023
transformers. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. 1779443379. Curran Asso- ciates, Inc., 2020b. URL https::/pproceedigg// neuripss cc paper/2020/File// 9ed27554 & 2aw.aaaaadd555522233555-------a- .ddf
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, and Amr Ahmed. Big bird: Trans- formers for longer sequences. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances Neural Information Processing Systems, volume 33, pp. 17283-17297. Cur- ran Associates, Inc., 2020. URL https: http::/ee///cceeedeeeoooo wgg.:eeeeddnn.g......oo.. nwwrips https//nnrrpsssccpppp///////
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Rep- resentations, 2017. URL petttss:/openreee
net. ottppp/www........e.e..oooooaaaa
Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip H.S. Tomr, and Vladlen Koltun. Point transformer. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 16259- 16268, October 2021.
Sixiao Zheng, Jiachen Lu, Hengshuang Zhao, Xiatian Zhu, Zekun Luo, Yabiao Wang, Yanwei Fu, Jianfeng Feng, Tao Xiang, Philip H.S. Torr, and Li Zhang. Rethinking semantic segmentation from geguence-to-sequeee perspective with transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 6881-6890, June 2021.
14

---

Published as a conference paper at ICLR 2023
A PROOF OF THEOREM 3.1
Our proof of Theorem 3.1 consists of four stages as described in Section 3.3 We state and prove the main lemma for each stage in the following subsections. Then, combine all stages in the end. The lemmas in each stage assumes permutation equivariance. We analyze how to circumvent permutation equivariance through the positional encoding in separate subsection.
A.1 STAGE 1: TOKENWISE PROJECTION
The main lemma for stage 1 is stated below.
Lemma A.1. Let N, d, N and 2 1,0 < 8 1. Let ,..,.... E Rdxx be a set of N in- put sequences that are distinct and tokenwise (r, o.separaad Denote r' = 2 V V
Then, there exists network N1 Rddx + _xx= consisting of the input embedding block and one tokenwise feedforward subblock with feedforward dimension ( separated. Moreover, for i, E [N] and k. E [n].
1 and bit complexity such that (1 (x)), [N] are non-negative and tokenwise (2r',2) 2)-
(x [1, Nx x [1,4] if and only if [1 = [1.
Proof. Our proof defines vector Rdd such that sequences (f) u² ( E RIx= i E [N] tokenwise (r', s)-searateed and satisfy that, for i E [N] and k; E [n).
1 [1, [1, and only x (c, k] = x (3) [:
Then, we construct the network Ag with the required size and bit complexity that computes (G( = u +r
Construction of u. Recall the definition of the vocabulary V = Uke[N] V( = {u E Rd
X [t, k) for some E [N], [n]]. Note that |V] I nN. We use Lemma E.1 on find such that
w'll
r'll S E (v w')l s 1 o'll Rd v
N
V12
for every v. E V.
Let i Rd be vector with each coordinate being the first Ilog(n?N) 2 T
bits of the corre-
Va sponding coordinate of u. We note that || ú piogin NAav))
n'N²
Rd-
Define u Sú with [2n N2, 1 2 2. We now check that (+) u X i [N] are
tokenwise (r',, seeppartte with r' 25[r]. Let i,j E [N],k, e [n] with x * X X [
and v, o E V with ( k)), v' ( (:, Then, we have
x
[1, k|
1 x (1) l:. k]
|
S 1 2
+ |(u ú)/
1
S S
||a||
nd ||a||
nz No
25||ea||
25r <r.
15

---

Published as a conference paper at ICLR 2023
We also have
[1
x0 [1.1
u (X (X | []]] (=))=.
u (o -')l
2 S (v 2') - (ā ū))² (v - (u)o'))
28
N2 V
10 1'll N2 Ta 11 2'1l
nd
S
0 l'll
n1
N2
Rd
2
Rd 2 2,
N2
which also implies
1 x0 if and only if ( [g,k] [:.0.
Construction of Ni. We construct Mi as composition of the input embedding block &n. Raxd + Rixn and a tokenwise feedforward block F(FF), RIx= - 1xx with skip-connection. We define { (X) a X and
F(FF) = =2 2
Then, we have
Ni(X (Em))(
= + - 11)on(u x + (r + 2
= + (S 115-117 -1)((i) X x 2 + 2
=Sa"X +25Tri1 u X +r'1
where we removed the ReLU activation because n X 2 have all positive values.
It is straightforward from the definition of ((FF) that the feedforward dimension of the network 1. Moreover, we can represent each coordinate of i with (N(n)))d)/x)) bits, the weights in F(FF) with [log S7 = [log(2n² /dd8-1d bits, and the biases in FFFF) with [log 2r] bits. Thus, the bit complexity of the network M1 is [log(2rn' 2 T8 ')..
A.2 STAGE 2: CONTEXTUAL MAPPING
Before proceeding to the stage 2, we reindex the tokens of each output from Lemma A. Due to the permutation cquivariance the Transformer architecture, we may reorder tokens in any order we want without loss of generality.
For each E [N], we consider (i) R as vector and reindex as follows. Suppose that there are unique tokens in (+), Then, there are also unique values in (+) We assign indices for tokens so that the first n, tokens have unique values of (i) in descending order:
x (x x ( [2] > ... x [nx]-
Then, we index the remaining redundant tokens in descending order:
[n, 1| [14 + > (0)) [m].
We note that the resulting vector (i) is uniquely defined. That is, x corresponds to one and only valid resulting vector (+).
We need the following definition in this section.
Definition A.1. Let N, n E N and 2 1,0 < < 1. Let (,() We say that (1) .* (N) are (r, (separred if
(N) E R. be data instances.
16

---

Published as a conference paper at ICLR 2023
|| for all i [N] and ||x0)) xا01 2 8 for all i E [A] with # j.
We now state the main lemma for stage 2.
Lemma A.2. Let N, n' M. Let x!1)... .... (N) R be of N input sequences that non-negative and tokenwise (2r', 2)-separated. Denote R' = 4 VTn [r][/
Then, there exists network N2 R3xn 33xa consisting 2n Transformer blocks with the and bit complexity
number of head h = 1, head size k = 1, feedforward dimension 1log(2r'nN2 /+)] such that
x 0
0 a 0
N2
(1) 1
where 2(0) ERiE [N] are (R' + +1,2))separated
Proof. Our proof defines vector E R such that values g(t) = w n.1" (1 nn E R,ic E [N] are (R', 4))-epprratte where we denote the number of unique values in (6) as Then, we construct the network N2 with the required size and bit complexity that computes
(1)2 o 0
0
N2
0
1
where 200) approximates g00
within 1 as
200 2(1)
s 1.
Since have
200) I (00
+ 200 z(0)
S R' +1
for E [N] and
2(1) 200 2 1200
z00|
=00
=0) -500
2 24-1-1=2
for i,j E [N] with * j. we conclude that 2 E RRiE [N] are (R' + 1,22 separaee..
Construction of w. For [N], define a(1) E R vector having the same values as (+)
in the first n, coordinates and 0 in the rest. We use Lemma E.I on (11) (2) ..... (NN) to find i such that
| (
TT (z(0)
x(:) S || - 2(11
Tm
for every i,j E [N].
Let V E d be a vector with each coordinate being the first [log(nN2 [T]] coordinate of w. We note that 20
bits of the corresponding
<
T
1
2
No
17

---

Published as a conference paper at ICLR 2023
Define Pwù with P 12N2 V w² x (+), LiE E [M] are (R', -)-eparatedd. Let i,j [N] with i j. Then, we have
We check that 20) w(1 ngl" "1"x"" (x( m.]
w xx
P w x(0
+ (z 20) T
< P
||
+
|
N2
nn
S 1
2P, Σ (0))
< 4Pr Vn S R.
and
200 z00
w (x() - -
P w (z(1) m(01
P
1 (1) xb)))
(1 10)² ( 2011))
8
N2 1 (00 x(71|)
P
5601
N
m
P
N2
P N
22 4.
Construction of 22- We construct N2 = Fzn o F2n-1 .... Fi in n steps. Each step [n] consists of 2 Transformer blocks 12 F21 33xn + 3x
First, we use Lemma E.2 obtain self-attention module ((SA) Blxn -> Rlxa that computes
a vector with all coordinates holding 2P-4 approximaainnnnn &max of 2max = max:e]m] []]]] given
R". We extend
((A)
to define valid self-attention subblock ((8A) 2/-1 123x + pax
x
x
x 2 max 17
((AA) 21-1
0
0
+
((SA)
(wT)
for x, E R.. The subblock F(SA)) finds the 2P approximate maximum token id among the first coordinate values and outputs the second coordinate.
Next, we use Lemma E.3 to obtain a tokenwise feedforward module ((FF) : R22nn R that 21-1
outputs r' for tokens with the same value in two coordinates. We define a tokenwise feedforward subblock F(FM)
g + 323a by extending ((PP) that, for æ. y, 1²
21-1
2 (PP)
([x, y
""
(FFF) 21-1
y"
y T
0
y T
T
z
0
z
where
2:
2rd
if
if
y
18

---

Published as a conference paper at ICLR 2023
The subblock r(FF) compares the first rows of the input and subtracts 2r' from the first row if two values are not separated. Since ar is bounded above by 2r', the subtracted entries become negative.
Since do not need the self-attention subblock from F21, set all weights of F(SA) 21 to zero. We
define F(PP) 1g + R3xa as
x
0
x²
((FF)
y²
y"
+
0
I/P
P
OR
Pe
y"
2T
0
Wk
2
x
0
x T Pyj
y
0
1/P P
OR
0
Wk
x. T
+ on(nmit
xT)
(
wu P
R((
RR( + WAOR(Y")
where P > 0 and = 2x P are defined earlier. When 2 O,, we can further simplify as
x
GR(x²) o n wk y
F(FF)
y"
21
T
( o
Let ((.)) FFx o F21-1 oF.
R
pan the output of the 1-th step when
0.
input is ( for for , n. We show inductively that
(1,0 [x] if < n, and [x) < (0
0
otherwise
(1,))2 k] = O,
(1)
min((,ne)
26.03, k] w(1 min n4 m.111 min(1,n,)] n;
Σ
wli [
for E [N], [n], where a00 is 2P./
approximation of (+)
(1)77
For = 0, the conditions hold for the input
O
Suppose that the conditions hold for
0
1 - 1. When i > Bhy the induction hypothesis implies that Z (0. [1,:] Of Thus, obtain
19

---

Published as a conference paper at ICLR 2023
the output
0
F(FF)
F(FF)
F(SA) 2ं-1
(z(l-11)
(((FF)
((FF)
F(AA)
21
21-1
21
2
21-1
Z
[3,
T
FF(PP)
F((F))
0
2f
21-1
Z (11.7-1) [3,
22
F(FF)
0
(1,, [3,
0
0
Z [3 :]
When i S n,, the induction hypothesis implies that Z (1.7 1) [1, has non-zero entries among which the largest value is (0) Then, it follows that
z (1,
F(FF) o
F(FF) 21
F(SA) (z(r-1)) = =F((F)) o o((FF) o
F(SA)
z z 12,
21-1
13. =1
22(ii11) [1,1]
((PP)
(FF) (SA)
0
21
2-1
2--1
(1 =1
z 2 [1,
F(FF)
(FF)
$01
21
2
1(1.-1))
(FF)
21
1)
0
(1, 1) [3, :] + w [1 0 1
where 2000 is 2P
1 approximation of x (1) Here, x (1.i-1) 1) denotes
7 ((( [1, 2r
if (1. [1, 2000 3
a [K]]
20-ii [1, [1
if 2 (1,T [1, [000 >
so
10
if z [1, k, 0 < 1 if (0 -1,] 20 >
On(i(1/-11 (1))) =
Z ( 1) [1, k]
Since distinct tokens are separated by 2, 2P apprximatiioss them are separated by 1. On the other hand, since 2P-a < 1.2 the 2P/m approximation of the maximum element stays closer than }-
Consequently, (=((i.11)) is the same as Z (i.0 [1, Thus, the conditions hold and we conclude
our induction proof.
0
In the end of n steps, the output is ((.)) =
0
with 20 = T (+) where each entry of (+)
z 201
2P approximates the corresponding entries of a We check that (0)) approximates 2(0) within
20

---

Published as a conference paper at ICLR 2023
1 as
20(0) -=00
< 10" t==a 26)) - -w="a
|/ ||(((0) =
VT
Pl
2P
(||zo| || + ||20 - u Il)
2
1 N
S 1.
2
Our construction of N2 involves 2n Transformer blocks. From Lemma E.2 the subblock ((SA)
2/-1
uses 1 head (h 1) with head size (k 1) and bit complexity [log (log(8r)) 1/2 3/' p)))
[log (1'
From Lemma E.3 the subblock F(FF) uses feedforward dimen-
sion 4 with bit complexity 1lo22'1. The subblock F
F(FF) uses feedforward dimension 2. All weights in F22 F((F)) are either 0, P, or Wn. Since we represent each coordinate of wn using ( bits, the bit complexity of ((FF) is max{[log P1. [log(n.22 [T)]} [log( (2n VT) Thus, the network N₂ has 1 head
1) with the head size (k 1). the feedforward dimension 4 and the bit complexity [log(2r'n 22 non'nN22 vz)].
A.3 STAGE 3: STRING LOOKUP
In stages 3 and 4, we map each token X (;) lt] to the corresponding label using both the token id
x [*]k from stage I and the sequence id ( from stage 2. Now that the token id x [x( and the sequence id ((+) hold enough information to identify the label, we process each token independently
using tokenwise feedforward blocks from now on. We first combine two ids into a single "contextual
token" id and map to the corresponding token label in the last two stages. We adopt stages 2 and 3
in Vardi al. (2022) to our architecture that involves skip-connections.
We can extend stage 2 by using extra dimension pass (+) from stage 1 without any additional parameter. Then, after stage 2, the k-th token in the i-th sequence contains 2(() and [1).. which are enough to identify the corresponding label [k).. We note that 0 S (*] < 2r' and |2 <
R' + with r', RR 6. We define a [k) ([z(0) R' + 1)(2r' + 1) + [k]] + be the unique integer id for each token in each sequence. Then, we have
1 a [k] < (2R' + 3)r +1 < '''R' for E [N] and
-aa0 a [1 2 2 for i.j E [1] k, E [m]] with V(t)
* vGD or X X [:AA+ + [:
We denote R = 9 Now, our goal is to map a [ E [R] to y [ E [C] using tokenwise feedforward subblocks. We denote the number of distinct a (1)) [k]'s 's as N' < nN. We reindex each
of unique tokens and the corresponding label as a and j(() for i E [N'], respectively. Without loss
of generality, we supppose that 1 S a(1) <... G(N²) < R.
Then, we partition N' contextual token ids into groups of ids. For each group g E [A], we construct two strings Uy and Wg. The binary string ug is concatenation of B ids in the group g. Each id is represented as an integer of p = [log R bits. The binary string w, is a concatenation of B labels corresponding to B ids in the group g. Each label is represented as an integer of y [log C bits. Thus, My and Wy are strings of length B and yB, respectively.
We now state the main lemma for stage 3.
Lemma A.3. Let N', R E N and I (1) .... < @(N)) R to be distinct integer ids that identify each token in each sequence. We suppose that |āl0) - -(b) 2 for i,j E [N'] with # j. Let
A, B, b E I with < N' and B = A Let w1, ... W E N with the number of bits in their binary representation most b.
21

---

Published as a conference paper at ICLR 2023
Then, there exists network N3 R - R consisting of A feedforward blocks with skip- connections every layers, feedforward dimension 9 = 4 and bit complexity b b+ [2o + 1)] such that
Na (0) = (al), w w1+1)
Proof. We construct F1ooo-100--- of as a composition of feedforward blocks Fi IR² R E [A] where each feedforward block contains single hidden layer and skip-connection.
Let 1 E [A]. We use Lemma E,4 define Fi
R22 [all-I)-B+1), aul-B)] and (z) of [all-11)-++)
R such that Fi(x) wy if
au-BB B) b where regard
1. B > be N We extend Fr to define a valid feedforward block Fr as
Fi (z) = (x, (0. = (x.v+ F(()))
Throughout the computation of N3, the first coordinate is the same across blocks. For E (N',. ( only activates one of Fi, E [A]. In particular, Fi(a(() = wi = + and Fi(a()) = otherwise. Thus, we get
Na (2(9,0) = (a()) , wrr+1)
Our construction involves A feedforward blocks with skip-connections every 2 layers. From Lemma E.4 feedforward dimension and the bit complexity is [log(2R + 1)].
Remark A.4. We can extend the construction from Lemma A.3 to output multiple strings corre- sponding to the same range without additional feedforward dimension. particular, the first layer in the constraction of Lemma E.4 does not depend Thus, we reuse this four units with different output weights to output additional strings corresponding to the same range. In stage 4, we need 2 strings for each range.
A.4 STAGE 4: BIT EXTRACTION
We first define (n) to be the substring of n with bits in places from j inclusive for n, E N with i j. We now state the main lemma for stage 4.
Lemma A.5. Let P, E N and u, N. Suppose that the number of bits in binary representa- tion of u and w are pB and BB respectively. We assume that
BIN (- (1 BIN (1N
(w) 22
for i,j E [B] with i i j.
Then, there exists a network N4
3 + I consisting of output embedding block and (max(p, 7 + B + feedforward blocks with skip-connections every layers, feedforward di- mension q 16 and bit complexity 2 max{p, TBB such that
(I, BIN
p. (()
if there exist E [B] such that BIN BINN-ii-i1 (- ((-
Proof. We construct N1 Fpost FB FB-1 Fi Fpre in B steps with pre-processing F Fpre: 3 + R2 and post-processing F post g R. We define the pre- preproccessing network as
Fpre(x, u. w. = xr 2
"
+
n + 208
O, 27 + B7 n
w
w
+ B2 0.0)
1
and the post-processing network as
F Fpost (21, 22, 266 2z.
22

---

Published as a conference paper at ICLR 2023
In each step [B], Fi R - R first implements two sub-networks J7
R
R and
F : R + R in parallel on the middle coordinates. Two sub-ne -ebworksssks uses Lemma E.5 to compute
b(e(l-1)) ((ell-1))
Tolat) siee)
2
2
F
20B
2 + 2014 BIN
p(l 1)+11:
and
[ b(1
w
22 + 2-g-1) d(a)) 27 + 2- (1-1))
2
2-8-1
Fl
BIN
Then, Fi combines the result using two additional feedforward blocks as
T
glel)
s(al)
(lel)
20
2015
"
+
2
202
s(o))
2
2
20B 2P1F2
BIN
pll
1)-1:pt(
F (x,
BIN
p(
(s(l)
27
10
270-1)
pbot)
w 2
22
d(l))
2
270+2)
xbx))
w BIN
2
27
BIN (1-11)11)1)
(1 1)
olal)
2 2011
+ +
clat)
i
d(11)
w +
w
+
0
F7 BIN
(1
where Fi uses Lemma E.3 compute
(z,0)) = (7
if r - vl < if r yl >
(0
and F? computes, for 0 03y522 2
() 27 +2) = y if = 2
if 0
The resulting value in the last coordinate is (IN if [x BIN - 1:pt (w) < } and 0 if |a BIN ((-1))++00 oz (w) > 1. Therefore, the last coordinate throughout each step of N₁ keeps the value 0 until it finds = BIN BIN 1(11))+ When such / is found, the value is updated to BIN,((1)+1->> as the requirement.
Finally, the parallel implementation of F7" and F7" requires max(p. 7 feedforward blocks, feed- forward dimension 16 = 8 8 and bit complexity 2 max)(, Moreover, each of Fi and F? requires feedforward block and bit complexity 7. The feedforward dimension of Fi is where F7 incurs 4 from Lemma E.3 and additional unit wipes out the (positive) carried value in 4-th coordinate from the skip-connection. The feedforward dimension of F is 3 where 877 incurs I and 2 additional units wipe out the carried values in 4-th and 7-th coordinates from the skip-connection. In total, each step / [B] consists of max(p, 7 + 2 feedforward blocks with skip-connections, feedforward dimension 16 and bit complexity 2 max(p, B
The pre- processennn network Fpre and the post-processing network Fpost can be implemented with feedforward block, feedforward dimension at most 2 (to carry u and w or 2g) and the bit complexity most max(p. 7)
Thus, Nu consists (max(p, 2) feedforward blocks with saip-connections, feedforward dimension 16 and bit complexity max 7
23

---

Published as a conference paper at ICLR 2023
A.5
POSITIONAL ENCODING
Consider the token id u² X (x) in stage 1. If we define the positional encoding as
E r r'uln 1, ,,,,,
then we have
+ E) u X (n) r'1 1, 2, 1,
Since every element u² (m) bounded above by r', positional encoding enforces the decreas- ing order used in stage as the usual sequential order. Thus, the sequence id is no longer permutation equivariant. The upper bound on the magnitude of token ids increases by factor of n, but it only affects parameter complexity and bit complexity logarithmically through R.
A.6 PROOF OF THEOREM 3.1
The final Transformer network combines 4 stages as N = N4 o N3 o N22 o Ni where A/3 and Na applies the same function to each token independently. The mismatch in the embedding and feedforward dimensions is easily resolved by using the maximum dimension required and set all weights in the unused dimension zero. We modify N3 output both " + in the Remark A.4 We summarize all stages:
and w 1 mentioned
1. Mi projects the input sequence X(1 _ the token ids x()) R" This stage con- sists of feedforward block of dimension 1 and bit complexity
S
log(r' dd).
2. N2 further projects the token ids xx() IR the sequence id 20)) TR permutation equivariantly. This stage consists of 2n Transformer blocks of attention dimension 1, feed- forward dimension 4 and bit complexity [log(2r'n.3) N22 /n)l S log(R').
3. Na combines a token id and sequence id to obtain the contextual token id and finds two group strings. N' < n. memorized contextual token ids are partitioned into groups of B ids so that N' < AB. Two group strings are crafted as concatenations of contextual token ids and corresponding labels in the group. This stage consists of A feedforward blocks of dimension 8 and bit complexity [max(log R. log C}] [log R] +1
4. Na extracts the correct label from the crafted strings. This stage consists of (max p. 7 + 2) + 2 feedforward blocks dimension and bit complexity 2[max (log Rt log C
In total, the network V uses 1 2n A (max > + + O(n A max B) Trans- former blocks of dimension 16 and bit complexity log(r' V + log R') mmx (log log C) B+ [log R| + 1 + 2 max log R, log B = Ollog d + [max(log R, log CJ We note that
R 9r Rl < 1 < 8000, 2 2
Finally, we balance and B as A = v nN log(nN) and B Theorem
and conclude the proof of log(nN)
nN
B LARGE WIDTH AND FIXED BIT COMPLEXITY
In this section, formally study the generalization of Theorem
3.1 in the case of large width
(Section B.I) and fixed bit complexity (Section B.2) from Remark 3.3 To simplify the argument, we state and prove the result only for the case without permutation equivariance. The theorem for the permutation equivariance case is straightforwardly obtained from our stated theorem.
B.1 LARGE WIDTH
We state and prove the theorem for large width.
Theorem B.1. Assume the same setting as in Theorem 3.1 Let S. VnN. There exists a Trans-
former network N: _xdx + glxn and positional encoding E IRdxn such that
(i E (+)
24

---

Published as a conference paper at ICLR 2023
for every [M] In both cases, the Transformer N has width 16nN/ (m 6nN/L k= and 9 = 1inNinn/// / depth
o n Lo/((L + V log(T)
1
maxg C, log Rt}
and bit complexity bounded by
log d + V log(L) max max C, log R}
where we denote R 30000~8--20 2dn No
Proof. Stages and 2 are the same as the proof of Theorem 3.1 For stages 3 and 4, instead of directly memorizing n.V contextual token ids directly, we construct nN/T.2 subnetworks where each memorize I,2 contextual token ids. By stacking subnetworks horizontally across width, we obtain the result. We remark that the width is not increased for the self-attention layers which are not used in the parallelized stages and
Theorem B.1 shows that, if depth is bounded above by Ö(L) with >> then Ö(d + + nN/L) parameters are enough memorize N sequence classification examples of length n with token dimension d. We count the number of parameters as linear in width instead of quadratic in width because our construction uses parallel LL2 subnetworks without interaction among them.
B.2 FIXED BIT COMPLEXITY
We state and prove the theorem for fixed bit complexity.
Theorem B.2, Assume the same setting as in Theorem 3.1 Let < VuN. There exists a Trans- former network N: Rdxx Blxn and positional encoding E Rdxx such that
N(X (8) + =E)
for every E [N]. In both cases, the Transformer N has width 16 (m = 6, h = ks and q = 16), depth
n aN
BB + nN
max C, log R}
B
B
log(B)
and bit complexity bounded by
o
log B,
max(log C, log R}
log(B)
where we denote R 8 No
Proof. Stage and 2 are the the proof of Theorem 3.1 For stage 3 and 4, instead of directly memorizing n. contextual token ids directly, we construct nN/B B2 subnetworks where each memorize B² contextual token ids. By stacking subnetworks vertically across depth, we obtain the result.
Theorem B.2 shows that, bit complexity is bounded above by Ö(B), then O(d + n + nN/B) parameters are enough to memorize N sequence classification examples of length with token dimension d. We count the number parameters as linear in width instead of quadratic in width because construction uses parallel n./T.2 subnetworks without interaction among them.
25

---

Published as a conference paper at ICLR 2023
C CONTEXTUAL MAPPING IMPLICATIONS
C.I SPARSE ATTENTION TRANSFORMERS
This section shows how our result generalizes sparse-attention Transformers, Sparse-attention Transformers replaces self-attention subblocks with sparse counterparts. Let A [m] be the I- th sparsity pattern of k-th token where E [n], E [p]. Given an input Rmxn, the sparse self-attention subblock F/ (SSA) with h heads and head size & computes
((SSA) Z Z
W (0)
(Wil) 22s os (wl.
(K) (w 2 (w/)
where 4zz R |A'l denote the submatrix consisting of columns of Z in the index Ai. 1, E Rkxm are the weight matrices parametrizing
Again, W(()) 1,f E m and w w:(k), w!( the sparse self-attention subblock.
We make the following assumption the sparsity pattern, which is the last condition among three conditions in Assumption in Yun et al. (2020b). Thus, our assumption is strictly weaker.
Assumption C.1. (Relaxed version of Assumption Yun al. (2020b)) Define sequence of he as
s = A St =
U
sj-
We assume that the sparsity patterns { satisfy that there exists a finite s such that
minfulSt [n] for all & [n]].
We provide the sparse-attention version of our main result.
Theorem C.1. Let N, d, n. C, E N and T
1,0
8 S
1.
Let
(X ( (x(N), Y Y(M) E Rdxn [C]1xa be N input-ouiput pairs of sequences that satisfies. Denote = 80 2n Then, there exists a Transformer network gdx= R xn with width 16 (m 6, h
and 16), depth
nN
O
VnNlog
log N max maxg C, log Rt
and bit complexity bounded by
o log d +
n.N log N
max{log log Rt}
such that ((X(0) ( for every i [N] and for every permutation matrix P paxn
Proof. The stage 3, and 4 are the same the proof of Theorem 3.1 In stage 2, each step of NG in Lemma A.2 computes the maximum token id over the whole sequence using self-attention layer. Under Assumption C. instead compute the maximum token id over the allowed sparsity pattern. Since the whole sequence covered within a recursion consecutive sparsity patterns and taking maximum is associative, repeating self-attention layer will give the desired maximum token id over the whole sequence again. Now, the other component in stage 2 works as before and the resulting memorization is achieved in the same way. The only overhead in this approach is the times larger number of self-attention layers.
The simplicity of our contextual mapping enables easy generalization to the sparse attention. Since the sparse attention Transformer only constrains the sparsity pattern in the self-attention subblock, stages 1, and in our construction works without any modification. For the contextual mapping in stage 2, we achieve the same memorization capacity with only times more self-attention layers. The number of parameters is Ö(d + VN).
26

---

Published as a conference paper at ICLR 2023
C.2 IMPROVED CONTEXTUAL MAPPING FOR FUNCTION APPROXIMATION
Our idea can be used to reduce the number of self-attention layers for the contextual mapping in function approximation settings, too. We recall the formal definition of the contextual mapping as follows.
Definition C.2. (Definition 3.1 in Yun et al. (2020a), Contextual Mapping) Consider a finite set C Rdxx, A map q L + R Rlxn is contextual mapping if the map satisfies the following:
1. For any L L, the entries in (L) are distinct.
2. For any L, I' E with L + I', all entries of ((L) and ((')) are distinct.
We state our improvement of the contextual mapping.
Theorem C.2. (Improved version of Lemma 6 in Yun al. (2020al) Consider the following subset
of Go = {0,6, ... .1- 8) :
Go = IL E GolL,, # L.g for all i + j}.
Assume that n 2 2 and 8-1 2 2. Then, there exist a function ge R4xn - RR4x composed of
3n Transformer blocks with h 1, k and q 4 that employ the hardmax operator, vectors
o w E Rd u E RA, constants 41, br E IR (0 < 4 < (.), such that g(L) == uu go(w" L, w L, 0, 0) satisfies the following propeties:
1. For any L E 98, the entries of (()) are all distinct.
2. For any L, E Gs such that L is not permutation of L', all entries of ((D), q(L') are distinct.
3. For any L E Gs, all the entries of q(D) are in [4, [r].
4. For any L E 93 Gs. all the entries of ((L) are outside [tz, tr]
Proof. Since token embeddings are 6-discretized, we can concatenate each coordinate to obtain the token id as in Yun et al. (2020a). Let be the vector that represents such concatenation linear operation. As in stage 2 of our proof for Theorem 3.I we concatenate the token in the decreasing obtain the
order of magnitude obtain the sequence id in steps. Then, "contextual token id" or the concatenation of token id and sequence id in (L). Then, the first three conditions are easy to check. The first condition is trivially true because distinct token will have distinct token id and consequently distinct contextual token id. The second condition is also true because is not a permutation of L', then their sequence ids should differ. The third condition is true since linear function in compact region is bounded.
appropriately
Consider the final condition. any step of our efficient contextual mapping, when the maximum token id is zero, there must be duplicate tokens in the input sequence. Conversely, if there are duplicate tokens in the input sequence, the maximum token id should be zero some steps of our efficient contextual mapping. We may use one more feedforward block in each step of the efficient contextual mapping to subtract
sequence by if such zero maximum token id is observed. Here, M is the maximum value possible for the sequence id. Then, the sequence id will still be negative in the end so qLL) will also negative. Thus, the final condition also truc.
We highlight the difference in the architecture. The major difference is the number of Transformer blocks used. We use 3n layers (linear in the sequence length) while Yun et al. (2020a) use 8-4 layers (exponential in the embedding dimension). Since the sequence length and the embedding dimension are of the same order in practice, our construction exponentially improves Lemma 6 in Yun al. (2020a). The minor difference in the architecture is the intermediate embedding dimension (4 in ours and in Yun et al.) (2020a)) and the number of attention heads (1 in ours and 2 in Yun al. (0200a))
27

---

Published as a conference paper at ICLR 2023
D OTHER TASKS
We provide the proof for Thcorem 4.1 [Section[D.] and formal results on language modeling tasks. We consider two language modeling tasks that commonly used pre-training Transformers: masked language modeling (Section D.2 and autoregressive language modeling (Section D.3
D.1 PROOF OF THEOREM
Stages and 2 are the id, we can directly classify sequence ids in stages 3 and 4. Since there are N possible sequence id, replace n.V with N in the parameter complexity from Theorem 3.1
the proof of Theorem 3.1 Instead of classifying the contextual token
D.2 MASKED LANGUAGE MODELING
Let X E RRdx and Y E [V]" be corpus data of T tokens represented as embedding vectors and token ids, respectively. The corpus data is divided into P = T - n sequences X .... (P) e Rdxa and y
(()) E [V]" of length n by taking (6) X it-) and y = Y it n]. Here, we use the slice index notation i j for items from (inclusive) to j (exclusive). We mask m of n tokens in the sequence with
(m) possible masking patterns M (
{0, 1)". We define the masked sequences M o x()) extracted from as the sequence x with
columns masked according to the pattern M (i)
We say that Transformer network N Rdxn memorizes masked language modeling of X, x,, if
Rixa and positional encoding E E Rdx=
N(M X (0) E (
for every E [P].] E [Q].
Theorem D.1. Let T, d, n, m. V E N and 1 Let X E BdXT, y E [V] be corpus data T tokens represented as embedding vectors and taken ids, respectively. Suppose that the masked sequences extracted from X are distinct and tokenwise (r, 8--sparrated.
Then, there exists Transformer network N RRdx RIln and positional encoding E Rdxx that memorizes masked language modeling of X.Y. The Transformer N has width 16 (m = 6, k and 16), depth
o 0 n nnQ log(nQQ) +
nPQ log(nPQ) max Vo R3
and bit complexity bounded by
o (10g d V log(nnPQ) max{log V, R)
nPQ
where we denote P -
= (m) and R = 8000, P6
Proof. We apply Theorem 3.1 memorize PQ masked sequences M x))
D.3 AUTOREGRESSIVE LANGUAGE MODELING
Let X E Rdxx and Y E [V] be a corpus data of T tokens represented as embedding vectors and
token ids, respectively. The corpus data is divided into P = T n sequences X (1) (P) E
kdd= and y()) ..... .y e [V] of length n by taking (6) X i n) and ( Y ii + n]. We call (1) as the input sequences extracted from X.
We say that Transformer network N dxx Rixn and positional encoding E E d memorizes autoregressive language modeling of X.Y if
(1 E) = (0)
for every E [P].
28

---

Published as a conference paper at ICLR 2023
Let X B Y [V] be corpus data Tf tokens represented as embedding vectors and token ids, respectively. Suppose that the input sequences extracted from X are distinct and tokenwise (r, -eparaatd
Theorem D.2. Let T,d, m. N and
Then, there exists a Transformer network N
RR - 12 and positional encoding E R that memorizes autoregressive language modeling of X,Y. Y. The Transformer N has width 16 (m = = and = 16), depth
o
n +
Plog P
+
V log P
P
maxflog V, log R3)
and bit complexity bounded by
o (10g [log& d
P
max{log V, log R}
log P
where denote P T - and R = 2n
Proof. We apply Theorem 4.1 to memorize P inppu sequences X
E TECHNICAL LEMMAS
Here, state technical lemmas that used in our proofs.
Lemma F.1. (Lemma 13 from Park et al. exists a unit vector u E Rdd such that N² nd (+) for every j.j [N].
2022 Ler N, d and x:11, .... .N) E Rd. Then, there xb0i s Iu (x(()) x(11) s ||x((0) x011
Lemma E.2. Let n E IN and r',P > 1. Then, there exists a neural network F R RIXB consisting of a single softmax self-attention with head, head size 1 and bit complexity [log 3/2 satisfies
such that F(x")) = el with Fmax 2P I c S mmax whenever a E TR.
| 2r for [n]] and
< Tmaa - for E [m] with a[i] F Fmaxo
where we denote Imax (maxigjn) #[())
Proof. When we have hardmax activation on the attention matrix, is easy to construct the network that satisfies the condition. Consider the following self-attention
F(x) = 1- (1- Z)OH [(1 x) (0 + + = æo x 11
1
max
/E(n)
Indeed, the output exactly Tmax and the bit complexity is since all weights are either 0 or
To approximate hardmax with softmax, we introduce some large factor t >> 0 in the attention matrix. Consider the following self- attention:
F(x) = 1- (1 x))s (tt T() 1,)1 wos [xTT,,7 1 =c1,
where
exp(tx il) =
c =
exp(1xl)])
Since c is a convex combination of x x Js,, it is easy to see that 2max upper bounds c. suffices to find / that satisfies the lower bound condition.
29

---

Published as a conference paper at ICLR 2023
Choose
[ log((n ///r/////)))) We lower bound the softmax weights on Fmax
Σ Eiarli =x =mxx
exp(Laoli)
Pmax
1= exp( 12: [i]
Linli (xp(tx(li)
E Sisel muraax exp (ta) [i] + Σ in / exp(txl(])
Eisrl =xmax cxp(txmaxx
E =izeli]=mmas expl tx max)
+ Σ (100
Suae exp(t(Fmmx 2))
Ramax
mmax + (n mmax expl- -21)
1)
Rimax
1
1
12 Ramax
1)
8
1 + Sv'P
where Timax == Hi x1
Tmax] H. Now, we can lower bound as
c 2 Tmax
2r' - Pmax)
Tmax
(Fmax + 21')(1 - Pinax)
Imax - 4r (1 - Pmax)
4r'
Tmax
1 1 8r PVn.
2Pvn
-mmxx
1 1
Tmax
2P
This self-attention module only has 1 head with head size 1. All weights are either 0, 1 or t so the bit complexity is [logt] t S [log (log(8n3/2//P))
Lemma E.3. Let r' E N. Then, there exists neural network F IR2 - IR with I hidden layer, width 4 and bit complexity [log 2r" such that
F(x, %
o y if
0
if
y
Proof. Consider the following neural network:
2
2
2 -2
1
F(x,y) = 1r' - M|
+
2 -2
y
1
2 2
<o'(on(22 - 2y + 22(2x - 2y + Rn(2x - 2y - + o - 2y - 2))
straightforward to see that the network computes the desired function. This network has hidden layer and width 4. All parameters either +1, =2 or tr' so the bit complexity is [log 2r'
Lemma E.4. Let a, b, w N with b. Then, there exists a neural network F R. - IR with hidden layer, width 4 and bit complexity flog w| + [log(2b + 1)] such that
w
if x [a,
F(x)
0
if x $ la 1. 1
30

---

Published as a conference paper at ICLR 2023
Proof. Consider the following neural network:
2a
2
2a
F(I, v) = [w0 -w -w
w OR
x
2b
2
2b +
w ( (2r 2a + 1)
R 2x - 2a) oR(2x - 26) + Rn - 2b - 1)).
It is straightforward to that the network FF computes the desired function. This network has hidden layer and width 4. All parameters are either tw, 2 or 2a 1, 2a, 2b, 2b + 1 so the bit complexity is [log w| [log(2b + 1)].
Lemma E.5. Let i.J.n, E N with w < n. Denote 6(2) = oR(oR(22) o (4z 2)). Then, there exists neural network F RR R$ skip-commections every layers, feedforward dimension 8 and bit complexity 2n such that
consisting of i i feedforward blocks with
() 1) 1(i-1)
x 2
Φ
+
F
2
2 Bg (2)
2013
Proof. We construct F;
F,, as a composition of j i+ feedforward blocks Fi : IR2 + R2 < S j where each feedforward block contains single feedforward layer with hidden units and a skip-connection. In step I, the block Fi computes
coll1) 1) 6
+
2n+2
(IN
Since we have
BIN Σ 2' (BN,
the block computation in Equation ensures that computes the desired function.
Construction of Fi. connection
To obtain Equation
2 we define the feedforward block Fi with a skip-
Æ1
[2]
T6()) - 21
Fi
22
=
D2
+
((22)
22 =
((22)
23.
fi(E1, 22)
23 + fi(21, 22)
where
f((1, 22)=2=20 = 20-1 (20+1-10(2)) ((tt))-2+t1-16(60) + 1 )/
We first note that the triangle function & can be implemented with one hidden layer
d(²) R22) a 2) + aRR22 2).
Also, we can implement identity function as
1 2 n2(2)) 2 -22).
Thus, the following 8 hidden units
oR (221) on (42) - 2) (22) on (-221) σR (2x2) OR (4zg 2) RR (2z2 2) OR (-2z2)
enough represent Fi.
31

---

Published as a conference paper at ICLR 2023
We check Equation
2 as follows:
Φ
(ol-1))
2
Fi
6
2
y fi (+)(-1)
2
20+1
y 2
(BIN(
because
BIN 22-1- (100)
x
2n+1-(ol1)
x +
1
2n
2
2"
2
2
Finally, since all parameters are have the form 2 for some -1 S h S 2n, the bit complexity is 2n.
F EXPERIMENTAL SETUP
We use HuggingFace Py Torch implementation of the BERT model for our experiments. All ex- periments are conducted on an Nvidia Quatro RTX 5000, 16 GB memory GPU in a machine with Intel(() Xe) Silver 4214 CPU @ 2.20GHz.
As mentioned in the main paper, we use 14,000 random samples in the named entity recognition dataset from CoNLL-2003 (Tjong Kim Sang & De Meulder (0033 for token classification and 50,000 random samples in the MNLI dataset from GLUE benchmark (Wang et al., (019) for se- quence classification. For token classification, the task is to classify the named entity type for each token among 9 possible classes. The sequence classification dataset aims to classify the relationship between sentence pairs as 3 classes: entailment, contradiction, and neutral. We vary the dataset size by randomly order examples and picking first p% for p 10, 20. 100.
We vary the model size through the embedding size m, which is varied by 12 and 96 for token and sequence classification tasks, respectively. As mentioned in the main paper, fix the number of layers as = 6, the number attention head as h = 12, the embedding to head size ratio as m/k h 12 and the feedforward embedding size ratio g/m 4, commonly done in practice.
We optimize using Adam optimizer (Kingma & Ba 2015) with learning rate 0.00002, batch size 32 and dropout rate 10%. We train our models for .500 and 7,500 steps for token and sequence classification, respectively. We determine the above number of steps to ensure that the training error does not improve at least for the last 3 epochs.
choose the minimum size memorizing model the smallest model that reaches the For Figure 2]
training error 0.005. The maximum training errors of selected models are 0.00499 and 0.00450 for token and sequence classification tasks, respectively. The average training errors of selected models are 0.00464 and 0.00259 for token and sequence classification tasks, respectively.
htttps///uugiingffac..oo
32


Reviews:
[{"review_id": "xc6oe_TnwY", "reviewer": "Reviewer_SL1m", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proves the memorization capacity of Transformer. The main technical theory is: given $N$ input-output pairs of sequence, this paper constructs a transformer that memorizes them with permutation equivariance, under slightly stronger assumptions. The mathematical tools used in this paper adopts the approach from Vardi et al., but with its own technical novelties. Finally, experiments with BERT validates the proposed theory. ", "summary_of_the_review": "The paper makes a solid contribution to first provide a memorization capacity result for transformers. ", "strength_and_weaknesses": "Strength: the memorization capacity of transformers is an interesting and important, though not surprising, topic to study. This paper is the first to prove the memorization capacity of transformers. The paper is well-written overall, with minor typos. \n\nWeakness:  Two minor weakness are as follows. The technical novelty seems to be limited, as it mainly follows from Vardi et al. Also, there are some typos in the paper, e.g., in section 3.1 should $Y[:, k]$ and $Y[1, k]$ be using the same notation?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "KZICagDwll", "reviewer": "Reviewer_YzgV", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper gives an upper bound on the number of parameters required to remember $N$ length-$n$ sequences, where each token is of dimension $d$.\nThe parameter upper bound is $\\tilde{O}(d + n + \\sqrt{nN})$. Together with bit complexity $\\tilde{O}(\\sqrt{nN})$, the upper bound is optimal up to log factors in terms of bit counts.\n\nThe proof technique is closely related to that of Vardi et al. 2022, with novelty lying in the use of self-attention, which improves parameter efficiency. As a comparison, directly applying the result in Vardi et al. (which is for fully-connected ReLU nets) will give a parameter count of $\\tilde{O}(dn + \\sqrt{nN})$, i.e. there is a factor of $dn$ as opposed to $d+n$ as in this paper.\n\nThe paper also provides a similar upper bound for sequence classification setup (i.e. the label is a scalar, rather than a sequence as in the main result), as well as empirical evidence that the parameter counts roughly follow the trend predicted by theory.", "summary_of_the_review": "To my knowledge, this paper is the first to provide memorization capacity of Transformers with finite precision. The paper is well written and the results are of interests to the community, so I recommend an accept.", "strength_and_weaknesses": "In addition to the upper bound in main result, the paper provides remarks on:\n- extending the results to real-valued / vector-valued output;\n- allowing wider width (the main result has width=16) or bounded bit complexity (the main result has bit complexity growing roughly as $O(\\sqrt{nN})$);\n- connection to fully-connected ReLU net, i.e. the bound if we apply Vardi et al. 2022 directly.\n\nI have no major complaint about the technical part of the paper. The main novelty over Vardi et al. 2022 comes from the use of self-attention blocks, which allows parameter sharing across tokens in a sequence, saving a factor linear in the sequence length.\nI should note though that even though I read Vardi et al., I'm not well versed with the literature so may miss some subtleties.\n\nSome minor comments about the experiments:\n- Fig 1: how many runs are performed for each cell of the heatmap? Could you increase the granularity of Fig 1(b) to show a clearer trend?\n- Fig 2: do the points correspond to models that have reached 0 (or near 0) training error? How many runs are there for each point? If there's only 1 run, then the trends may not be trustworthy. Please consider using multiple runs and show the median and standard error.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "Gv-gLDvMQl", "reviewer": "Reviewer_Nzx7", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors construct a transformer that can memorize $N$ sequences of length $n$ and dimension $d$ that have permutation equivariance, with a transformer that has width 16 and depth $\\tilde{\\mathcal{O}}(n + \\sqrt{N})$, with bit complexity $\\tilde{\\mathcal{O}}(\\sqrt{N})$. Furthermore, their theory can show memorization for sequences without permutation equivariance using positional embeddings.", "summary_of_the_review": "Overall, my scores are slightly on the negative side, mainly because (a) I would like to see how the attention module has been used in the proof sketch, and (b) I would like to see how the authors managed to improve upon Yun et al (2020a)'s construction. I would be very happy to discuss with the authors in the rebuttal period on the above questions.", "strength_and_weaknesses": "The major strength of the paper is the parameter efficient transformer (which is nearly optimal in the number of bit counts) constructed for memorizing $N$ sequences. Furthermore, the authors show experiments verifying some of their claims.\n\nHowever, I feel that the use of the self-attention module in the construction hasn't been highlighted in the paper and is difficult to read from the proof in the appendix (please see more questions below). One of the major questions that I have is whether we can use the result of [1] in two steps: first, we project each token in the $n$ length sequence into a single bit, and then further project the $n$ dimensional representation into a single bit.  \n\nA section similar to 3.1 in [1] would have helped a lot.\n\nMoreover, I have the following questions:\n\n(a) Can the authors explain the role of the self-attention module inside the construction? Can one simply use a feed-forward network in place of the self-attention module? My reasoning is that after stage 1, you only need to store a single bit for each token in the sequence. Hence, is it possible to use a feedforward layer of dimension $n$ that \"approximates maximum token id among the\nfirst coordinate values and outputs in the second coordinate\"? \n\nIf there is a difference, this can be used to compare attention models with MLP-only models introduced in [2, 3].\n\n(b) Can the authors provide a discussion on how the contextual mapping used in Stage 2 improves over the selective shifting-based contextual mapping from Yun et al. (2020a)? Does the pre-processing in Stage 1 help reduce the overhead in Stage 2 for contextual mapping? Also, can the authors provide some discussion on how their contextual mapping is capable of \"incorporating sparse self-attention settings with minimal parameter overhead\"?\n\n(c) In figure 2, how do the authors observe a square root relation between training data size and model size? The plots seem very linear. In plot 2(b), if we remove the point (30000, 2.5e7) which could have been a noisy point, the plot will become more linear.\n\n(d) The experiments have been conducted with a fixed hyperparameter set (batch size, learning rate, etc.). That implies the plots in 2(a) and (b) could be highly dependent on the hyperparameters used. How do the plots look with a different set of hyperparameters?\n\n(e) Is there a relation between the entropy observed in the experiments in Figure 3 and the number of sequences $N$? (By theorem 3.1, there should be a square root dependence between bit complexity and $N$). \n\n(f) How do the plots in Figures 2 and 3 change with changes in the length of the sequences? \n\n(g) The construction doesn't require the number of attention heads to change with a change in the number of sequences. Do the authors observe the same in their experiments, i.e. the behavior of the training loss doesn't depend on the number of attention heads used?\n\n\n1: On the Optimal Memorization Power of ReLU Neural Networks. Gal Vardi, Gilad Yehudai, and Ohad Shamir\n\n2: MLP-Mixer: An all-MLP Architecture for Vision. Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, Alexey Dosovitskiy\n\n3: Pay Attention to MLPs. Hanxiao Liu, Zihang Dai, David R. So, Quoc V. Le\n \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "BEGImD4xyW", "reviewer": "Reviewer_XAtG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper conducts the theoretical analysis of the memorization capacity of the Transformer model with finite machine precision and shows $\\tilde{O}(d +n \\sqrt{nN})$ model parameters are enough to memorize $N$ seq2seq samples with/without permutation equivariance assumption, where $d$ is the embedding dimension/model size and $n$ is the sequence length, which improves the results in previous literature (e.g., Yun et al., (2020ab) and Kratsios et al., (2022)). From the technical perspective, the authors consider a novel contextual mapping, which facilitates the analysis extending to the sparse attention settings. Moreover, the authors also use numerical experiments to verify the theoretical findings.", "summary_of_the_review": "This paper proves that Transformers are capable of memorizing $N$ seq2seq samples with length up to $n$ with $\\tilde{\\mathcal{O}}(d + n + \\sqrt{nN})$ parameters. The authors also use experiments to verify the theoretical findings.\n\n\nMinor Issue:\n\nFigure 2: From my point of view, figures 2(a)/(b) show the linear dependence even if in the beginning. The claim *..the square root relation between the training data size and the model size in the beginning...* is not well supported. I suggest the authors using separate figures to explicitly show square root dependence for the beginning stage.\n\nFigure 3: The description of the x-axis in Figure 3 is missing and I suggest the authors adding some detailed explanation of this experiment, as based on the current presentation, it is confusing to me how the entropy of parameters is computed.\n", "strength_and_weaknesses": "Strength:\nThis paper contains the technical in-depth results and sheds the light on explaining the efficiency of the transformer model.\n\nWeakness:\n1. The theoretical analysis is mainly an extension based on Vardi et al. (2022).\n2. the current analysis requires $\\mathcal{O}(n)$ transformer layers, which is a little bit too large for the common practice.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster