Memorization Capacity of Multi-Head Attention in Transformers
Authors: Sadegh Mahdavi
Renjie Liao
Christos Thrampoulidis
Venue: ICLR 2024 spotlight (2024)
Keywords: Learning Theory
Expressivity
Multi-Head Attention
Transformers
Primary area: learning theory
TL;DR: We propose a lower bound for the memorization capacity of multi-head attention in Transformers

Abstract: Transformers have become the go-to architecture for language and vision tasks, yet their theoretical properties, especially memorization capacity, remain elusive. This paper investigates the memorization abilities of multi-head attention mechanisms, examining how many example sequences they can memorize, as a function of the number of heads and sequence length. Motivated by experimental findings on vision transformers, we introduce novel assumptions about the linear independence of input data, distinct from the commonly used general-position assumption. Under these assumptions, we demonstrate that an attention layer with $H$ heads, dimension $d$, and context size $n < d,$ featuring $\Theta(Hd^2)$ parameters, can memorize $\Omega(Hn)$ examples. Our analysis sheds light on how different attention heads handle various example sequences, aided by the softmax operator’s saturation property. We validate our findings through experiments on synthetic data.

Published as a conference paper at ICLR 2024
MEMORIZATION CAPACITY OF MULTI-HEAD ATTEN- TION IN TRANSFORMERS
Sadegh Mahdavi'2 Renjie Li Christos Thrampoulidis'
'University of British Columbia
Vector Institute for Al
(smahdavi, r rjliao, ethrampo lece. wbc.bbc...oo ca
ABSTRACT
Transformers have become the go-to architecture for language and vision tasks, yet their theoretical properties, especially memorization capacity, remain clusive. This paper investigates the memorization abilities of multi-head attention mechanisms, examining how many example sequences they can memorize, as a function of the number of heads and sequence length. Motivated by experimental findings vision transformers, we introduce novel assumptions about the linear independence of input data, distinct from the commonly used general-position assumption. Under these assumptions, we demonstrate that an attention layer with H heads, dimension d. and context size n d, featuring O (H) parameters, can memorize Q Hn) examples. Our analysis sheds light on how different attention heads handle various example sequences, aided by the softmax operator's saturation property. We validate our findings through experiments on synthetic data.
INTRODUCTION
Transformers have emerged as the de-facto architecture for achieving state-of-the-art performance across different domains such as natural language processing and vision tasks, as demonstrated by recent breakthroughs, e.g. (Dehghani et al. 2023 Touvron et al. 2023 OpenAl 2023). These models are characterized by their immense scales, often comprising billions of parameters to achieve their high performances. Given their sizes and complexities, natural question arises: How effectively can they memorize the training data? This query important both from privacy perspective (Carlini et al. 20220), and as stepping stone towards quantifying the model's ability to generalize to new data ZZaagg et al.) 2017). Additionally, studying how the memorization of transformers emerges can offer insights into the distinct roles played by their different building blocks.
Concretely, the memorization capacity of a model fw refers to the minimal set of parameters W such that (x) = y holds for all input-output pairs (z. y) in the training dataset Historically, this question has been extensively analyzed for fully connected networks (FCNs), dating back to at least the 1980s when Baum (9888) proved the first set of results for fingle-hidden-laye Threshold networks with binary labels. The best known bound on the memorization capacity of single-hidden- layer ReLUU FCNs has only recently been established by Bubeck et al. (2020), following numerous recent studies, e.g. (Zhang et al.) 2017 Yun et al.] 2019). Recent studies have also extensively considered extensions to deeper models, e.g. (Vershynin) 2020; Park et al. 2021: Vardi et al.] 2022 Yun et al. 2019). However, of now, the best known results consist of order-wise bounds, involving unknown constants and logarithmic factors. The majority these works on memorization of FCNs established for continuous inputs assumed in "General Position," a relatively mild assumption, as it holds true under scenarios like small random perturbations of input examples.
In contrast FCNs, understanding the memorization capacity of transformers is still in its infancy. A common Transformer architecture consists of two main components: (i) a Multi-head Attention (MHA) module and (ii) a ReLU fully-connected network. The MHA module creates convex combi- nations of its input elements by computing softmax similarities between input representations. This
'The definition of memorization in this context is distinct from concepts such as "associative memory" Demir cigil et al. (007); : RRamsauer et al. (2020) and "neural memory" " [Sukhhaatar et al. (2015); Geva et (2020)
which involve associating an input pattern with given query and have been recently studied for transformers.

---

Published as a conference paper at ICLR 2024
operation, unique to transformers when compared to FCNs, is leveraged in practice for its ability to selectively attend to information from various representations and effectively process long sequences (Vaswani et al. 2017). Yet, its theoretical underpinnings remain largely unexplored. This paper aims to contribute towards closing this gap by investigating the memorization capacity of the MHA module. To isolate its role from the ReLU FCN, take step toward analyzing Attention-only models.
Concretely, we study a single-layer MHA module with H heads that takes as inputs n x x key/context matrix comprising representations of tokens, and d x query vector. We are particularly interested in understanding how the number of attention heads, and the context size of the input affect the memorization in this architecture. Our contributions are as follows:
We introduce a new set of input-data assumptions that are more relaxed compared to General Position assumptions. Our first assumption requires the set of all query vectors have Kruskal rank at least while our second assumption needs context vectors of each example to be linearly independent. These are motivated and validated by experiments on real-world models, such as ViT, following a single layer, whether it be Attention-only or complete Transformer, that is either randomly initialized or pretrained. In contrast, General-Position does not hold in either of these settings.
For this set of practically relevant assumptions, we prove that a single-layer MHA module with H heads, embedding dimension d, key/query dimension da. value dimension de, context size n < d, output dimension dout & d,, equipped with a total of O(Hd(da, + d.)) trainable parameters, is capable of memorizing Ω (H min(n, dx)) input examples. Specifically when dout d, d and n = ((d), this finding recovers the optimal order of real-valued numbers that can be memorized with the given number of parameters.
Our proof illuminates how different attention heads play distinct roles in memorizing different sets of examples, and how this role allocation is facilitated by the softmax. To show memory capacity increases linearly with H, our idea is to make each head individually "responsible" for memorizing min(n, dx) 1 examples. By leveraging the saturation property of softmax and introducing specific key-query weight adjustments, we design desired softmax logits for individual heads while ensuring minimal interference with other examples.
In addition to experiments validating our assumptions on the input sequences, we design and discuss synthetic experiments corroborating our findings: namely, the impact of head number H, context length n and key/query dimension da on memorization, as well as, the occurrence of the saturation property of the softmax operator.
2 RELATED WORK
Theoretical understanding of Transformers. Driven by the success of Transformers, there have been several attempts at understanding their underlying mechanisms, e.g., (Li et al.] [023b] Luo et al.
2022: von Oswald et al. 2022,) Dong et al. 2021 EEelman et al. 2022, Jelassi et al.] 2022 Li et al.
2023a; Sahiner et al. 2022). Closer to us, Yun et al. 2200aa b) study the expressivity of Transformers perspective, by constructing networks with an exponential number of layers and parameters in input dimensions (context size and dimension) to cover the entire input space. In contrast, our aim is to find a small network to memorize a finite number of samples. [Bhojanapalii et al. (2200) studies the expressivity of MHA on a single input sequence, yielding guidelines for setting the head dimension d, at least as high the context size preserve expreissivity. To the best of our knowledge, our work is the first to explicitly address memorization of MHA with more than one input data points. Only very recently, Kim et al. (2023) has extended previous results for ReLU networks (Vardi et al. 2022) to Transformers with ReLU-activari feed-forward layer and a single Attention head, showing that sublinear number of parameters suffice for memorization. Here, in order to better understand the distinct role of the Attention mechanism, we study the memorization of lne-layer MHA. Also, unlike those previous works, we focus on single prediction per input. These make our results not directly comparable to theirs. Also, our assumptions on the input data and proof techniques are different.
from a aniversal-approximaat
Memory capacity of FCNs. Contrary to Transformers, the memory capacity of FCNs with SeLLUThreshhoddSigmooi activation has been extensively analyzed. Baum (1988) provided the first memory bound, proving under General Position that single hidden-layee Threshold can memo- rize T binary labels with OTT) parameters. The best known result on memorization under ReLU activation and real-valued inputs was obtained only very recently by Bubeck et al. (2020). These
2

---

Published as a conference paper at ICLR 2024
bounds are tight for the respective class of assumptions since memorization for a dense class of input data (e.g., General Position) with piece- wise analytic definable functions requires number of parame- ters at least linear the number of examples (Sontag 1997). Building these foundational works,
recent work have extended these results to different input-data assumptions, smooth activations, and
multi-layer networks (Vershynin 2020 Yun et al.) 2019 Rajput et al.) 2021 Zhang et al. 202T
Bombari et al., 2022). Additionally, recent works yield bounds on ReLU networks with (vVT)
parameters by utilizing -deep networks (Vardi et al. 2022; Park et al., 2021). Following the tradition of early studies on FCNs that explored memorization in simplified one-hidden layer
networks, we delve into the memory capacity of a single-layer MHA mechanism. Our objective is to isolate the MHA from the ReLU MLP and analyze its unique properties. We anticipate that our discoveries will inspire additional research and extensions for more advanced Transformer models.
3 PROBLEM SETUP
Notation, [n] denotes the set of positive integers 21, . m).. letters, such as A \ v, denote matrices column-vectors. [A; B) [A, B denote row-wise \ column- wise concatenation. Tin]] = {21, 22.... Fn denotes the set of [r]] indexed vectors.
Boldface upper-case \ lower-case
We study a single-layer multi-head Atlention (MHA) mechanism. An MHA layer A with H heads consists of three sets of matrices Wh:: Won E Rdxd,, Wv, E Ruxd, for each head E [#]. one weight matrix Wo E IR. Hd.dd to combine the outputs of different heads, and fi- nally, read-out weight matrix WD E Rdxdoa to get the final model's output. We denote the entire set of parameters as W = ([Won, WKL. Wval Wo. WD}. An input example to the MHA model consists of (i) key matrix E = [e :e:e; ::::...... ;e]] E Rnxd containing n to- function of the input exam- ple and of the set of parameters is given according to the following computational mechanism:
kens, and, (ii) query vector e E Rd. The output of the model
an == EW KA Wox e
(oth R") (1)
Ph Wv T h %h
(ph R
(4)
On Softmax (or)
(0, [R") (2)
0: Wo (p::Pa....:.
(o E Rd)
(5)
2h = ETO. 0,,
(ZA E RRd)
ÿ :: WD To
(y E Id)
(6)
Eqs. 1 (4) hold for all heads E [H] o is the output of the MHA layer, and, ý is the final output of the model. In Transformer terminology. the key matrix E is also called the context. Moreover, when e E eix), the MHA layer is called Self-Attention, and when the query vector is not among key vectors, the MHA layer is particular form of Cross-Attention. Here, we impose restrictions on the specific type of MHA as long as certain assumptions on the input data are satisfied (see Sec. 4
We consider a training set TT consisting of T input examples. To each example t E [T], con- sisting of context tokens and one query token, corresponds label ( Thus, we have T : {(E(()) e "y(") Our goal is to find parameter set W such that the MHA model is able to memorize any label set given the key matrices and query vectors, ie.., y(t) = fw(E e()) for all E (T). The maximum number T of examples for which we are able to achieve this serves as a lower bound on the memorization capacity of the model.
To illustrate the role of an input example distinct components (key matrix, query vector, label), consider image classification using (say) Vision Transformer (ViT) (Dosovitskiy et al. 2021). Two types of tokens are given as input to this model: A single "[CLS]" token, and several "image patch" tokens. The image patches contain information about the input image, while the "[CLS]" token is a special token in that the model's class output is read from the prediction of this token. For each image E [7), the key matrix E()) is the matrix containing image patches and the "[CLS]" token, the query vector (()) corresponds to the "[CLS]" token only, and y is the label of the image (here, the label is a scalar; thus, dour 1.) Note that, the type of MHA is here Self-Attentio since the "[CLS]" token is inside E+)) as well, otherwise, the context would have contained only the image patches.
Often, a positional encoding is also added to the embedding of each token before feeding them the Transformer. Namely, for all E [T]: ( = !!)) + pOSp, and (t) x; + pos,, for all i E [n]] where ( and (t) are the raw embedding of input tokens, and pos; E IRd are fixed and unique positional encoding vectors for each position index. In the case of Self-AAtention. the query vector takes only one positional encoding since is already present among C(n]- We will show in Sec. 4.4
that the use of positional encoding motivates our Assumption 2
3

---

Published as a conference paper at ICLR 2024
4 MAIN RESULTS
Two types of assumptions are typical in prior works on memorization: (1) norm based (Vershynin
2020
Park et al., 2021 Kim et al., 2023). and, (2) linear-independence (Baum) 1988 Bubeck et al. assumptions. The former assumes well-separated and/or bounded-norm input data, while the latter only requires some notion of linear independence among the data points. While norm- norm-based assumptions are well-suited for discrete data such language, linear independence is usually better suited in continuous data such as images, or continuous distributions, e-g. Gaussian and Uniform.
2020
In this paper, we focus on the second type of assumptions. Thus, it is useful recall the notions of Kruskal Rank and General Position.
Definition (Kruskal Rank), Kruskal Rank of a set of vectors is the largest number i such that every subset of size i is linearly independent.
Definition (General Position). A set of vectors (2 B2,..... Im C Rdd is in General Position if it has the maximal possible Kruskal Rank of d.
In contrast to previous works on the memorization capacity of FCNs, we depart from the commonly used General Position assumptions and tailor our assumptions to transformers. Specifically, we make the following assumptions.
Assumption 1. The set of all query vectors (e()) RT has Kruskal Rank at least Assumption 2. For each example [T). the context matrix E(t) E Fuxd has rank n.
These assumptions find their motivation in experiments conducted on real data, which reveal scenarios where these assumptions hold, while the General Position assumption consistently fails (see Sec. 5.1).
We are now ready state main result below.
Theorem 1. Consider multi-head attention layer A with H heads, embedding dimensions d, dy 2 dow 2 L, and dn 2 1. Let = {EE!) ett) y y 1Tm be training set with context size n < d. Define == min(n, dr)- If Assumptions/ M and 2 hold, and S H (r 1) + 1, then, there exists a set of parameters W such that A can memorize T.
Theorem D shows that an MHA with dHdd, + Hdd + dd parameters can memorize at least H (min((, d,.) - 1) + 1 examples under Assumptions 2 and To gain further insights the implications of this theorem consider three concrete settings: (1) Suppose da - d and scalar model 1, then MHA with ((Hd²) can memorize Q(Hn) examples. (2) Suppose d, = n < d, d, = dout = 1 then a MHA with @(Hdn) can memorize S2 (H)) examples. (3) Suppose and d-dimensional model output dout d. then MHA with ((Hd²) can memorize ((H)) examples with d-dimensional label each. The third scenario is particularly interesting. To see this further, consider the typical setting where and d are of the same order (n = (dd)), then the established bound becomes optimal. This optimality is reflected in the fact that f2 (2(dd) = 2(HHd) real-valued outputs are memorized, and a trivial upper bound for the memorization an MHA is the number of its parameters, which is O(Hd))
output d, dout
dj. = d,
Additionally, Theorem 1 suggests there is no gain memorization capacity provided d, > n. This agrees with using da < d in practice. We also confirm this in synthetic experiments in Section 5.2 A similar conclusion is also drawn by /Bhojanapalli et al. (020)). where authors show by having dy < Attention loses expressivity in terms of convex combinations O₁ the Attention represent. However, their finding alone does not allow direct conclusions about the impact on memorization. Our theorem directly connects the off between and d, to memorization.
The assumption that n < is general practice. For instance, in prominent models such as Vision Transformer (Dosovitskiy al. 2021), BERT (Devlin et al. 2019), and GPT-3 (Brown et al. 2020). the n ratio is 197/768, 512/ 768, and 2048/12288, respectively. Furthermore, in Sections 4. holds in practice due to positional encoding.
and 5.1, we demonstrate that Cotuumptiontttyiiaaat and Assumption I is valid after applying a single Attention layer. Unlike previous works on FCNs that regard General Position as an assumption (Baum, 1988, Bubeck et al. 2020). our approach innovatively incorporates Assumption I which imposes a weaker constraint than General Position
2 dddimensional outputs is common in prior works, e.g., consider equeene-to-seqqeenc outputs.
Yun et al. 2020b a), with the distinction that they
4

---

Published as a conference paper at ICLR 2024
since < d. In fact, our experiments reveal that Assumption | is satisfied in real data, even though the more commonly used General Position assumption is often not met.
4.1 PROOF OF THEOREM
Our proof consists of two steps: (1) We focus on the intermediate representations originating from 2h in Equation (3) We demonstrate with a suitable set of attention weights {(WK Wa ) we can achieve linear independence across different data points, resulting in a rich representation. (2) We prove the existence of { & /! Wo, WD} for achieving memorization by taking advantage of the rich representation from the previous step.
Let us define matrix Z, which is a function input data and WK.. Won h E [H] as follows:
Z = z (1) (2))
:....: E TRxxd
(t) where 20) == [1
(0)
(7)
Each vector 20)) is a concatenation of H d-dimensional vectors each representing a convex com- bination of context vectors of the I-th example parameterized in terms of key/query vectors and WKE, Won. The matrix Z gathers all these dH dimensional representations, one for each example in matrix form. Then, based on Eq. (T to Eq. (6). the predicted labels g! are given by
g
Z diagg( WV Wv WoWD E R, V [7].
(8)
Step
Step
In the first step, using an inductive proof technique, we show that Z has a high rank, and use this information in the second step to solve a linear system of equations and achieve memorization. Step 1. The Rank of Z. Our key technical result proving the following lower bound on rank ( Proposition 1. Under the conditions of Theorem there exists parameters {(Wk)) WQn li such that rank(22 Z) min (r - 1) + 1,T}.
K
We provide proof sketch of this proposition here and defer the full proof to Appendix B The high-level idea is to inductively add heads and use each new head for a distinct set of r examples. which in turn translates an increase on the rank of Z. The new head exploits new distinct set of examples without affecting the rank of the previous induction step. To accomplish this, we leverage the saturation property of the softmax and the linear independence assumptions. For simplicity, only consider here the case d, 2 d >> n, for which r n.
In order to get Hnn 1) rank, each time head is added, the rank of Z should be increased by n 1. For each head h, we only focus on (n - 1) + examples, therefore, the subsequent step can be viewed adding new head with tunable weights, and set of 1 fresh data points. Namely, the induction step of Z can be described in view of the following block matrix:
Z
Z
Z'a
Z"
T-=T-nt1
Z =
TOW ww-ogeerations
Z
Z
0
Z A
H'd
H'd
Z'. E TTT'' xHd denotes the matrix derived from the induction hypothesis, Z E R( contains the embeddings of new examples corresponding to the H' 2 heads of the previous induction steps. Z E RTT xd contains the embeddings of the previous examples corresponding to the new head H H' + 1, and finally, Z E (==xx) contains the embeddings of the new examples corresponding to the new head H. The upper-triangular block matrix on the right-hand side is obtained from Z only by row operations (see the appendix for details); thus, (i) its rank is the same as the rank of Z, and, (ii) the matrix A := A A" is a function the block Z
In view of the upper-triangular block formulation above, our objective becomes clear.
Goal: Adjust the weights WKH. WQ, of the new head H, to increase the rank of Z + A while dealing with the challenge that such tuning can also impact Z subsequently affecting the rank of A A A(Z ).). We will demonstrate method for precisely tuning these weights to control the rank of Zí without altering the rank of Z
5

---

Published as a conference paper at ICLR 2024
Approach: We approach this in two sub-steps, providing an overview of each below. Recall H (W) == Softmax (E()) Welt) is the vector of softmax coefficients, associated with the t-th input and the new head H, parameterized by W == WKH Wa T Also note that Z" and Z are parameterized by 0/ (W). Consequently, we will use (W) and (WW to explicitly denote parameterization with specific weight matrix.
Sub-step I.A. Tune the new head's weights and memorize the new set of examples. Given an arbitrary, but fixed matrix M E R( whose precise value will be chosen later in the second substep, we demonstrate that it is possible to design weights W = (M) to achieve a target rank for Z specifically achieve rank Z (W")) +M) 2 n The construction based on an n-step induction, making use of Assumption that the rows of E(t) span an n-dimensional space.
Sub-step 1.B. Suppress the impact of the new head's weights on the previous examples. To select the matrix M, we establish the following procedure for fixing Z We first establish (see Lemma 3 the existence of a matrix W satisfying two conditions: (i) E g+ (+(
0 for all t T + 1, T. and (ii) Will o (1) e for all i + j E [m]] and t E [T']. Given W+, define 0
1 [arg | e el) for t E [T"] The crucial insight here is that owing to the saturating property of the softmax (see Lemma 4 and property (ii) above, for any weight matrix W (which may depend on W+,, we have 04 = lim 0/ (W + +W+) for all t E [T]. Consequently, for any
W, the matrix lim (W + +w+)= = E T of.... .... ....E ...E 0 OTTIT== =: Za,+ is independent of the choice of W.
The proof is finalized by integrating the two sub-steps. Specifically, we utilize Z from sub-step I.B and select M = A (Zo,.),, Next, we apply sub-step 1.A to identify the weights W W (W(WW) in manner that ensures rank Z,, (W*) + A (Z,,+)) 2 n - 1. Finally, we set the overall weights of the new head as W W + +W+ As the parameter e approaches infinity, the softmax coefficients for the first T" examples will be exclusively determined by W (not W*), as observed in Sub-step I.B. Furthermore, thanks to property (i) of Sub-step I.B, we have Z (W- + + = 2í5 (W*) for all values of c. Collectively, these steps complete the proof of the proposition.
In the Appendix B we provide all the details. This involves showing that a finite value of scaling c above suffices; thus, there exist finite weights with the desired property. Moreover, for the scenario d, I d, we demonstrate that the matrices W*, W can be chosen to have rank at most da.
Step 2. Solving system of equations. We can find the remaining weights memorize at least rank Z) examples. WLOG, assume that = rank Z), otherwise we can ignore some of the data points and keep only rank(2 Z) data points that make Z full-rank. We find weights (1W Vn M Wo, WD) to memorize all 7 examples as follows:
y (1)
Wv1
Ia xd
Wvz
Id.xd
ZWv= =
Or (d. dout)
= Wv, Wo: =
WD == Idxdaa
(T)
WvH
Ja
where Wr is a solution to the first expression (note that such Wv exists since Z has full row rank).
Substituting in Eq. (8. this gives (*)) = (t)) for all t E [T]. Therefore, we have proved that the memorization capacity of one layer is at least rank Z). But, since also proved lower bound of rank Z) 2 H(r - + 1, we conclude that the memorization capacity is at least H(r- (1r--) 1) + 1.
Remark 1. In the setting of Theorem [ if Assumption iis violated (i.e the Kruskal Rank of query vectors is less than n), define Q < n be the Kruskal Rank of query vectors (e(0) (!()) R Then, a weaker lower bound on the memorization still holds: T S H(min( dw,) - 1)+ +1
Jt=1-
4.2 TIGHTNNSSS OF THE LOWER BOUND ON THE RANK
Our proof of Theorem relies on lower bounding rank Z) in Proposition I One might naturally wonder how tight this bound is. The proposition below proves that the bound indeed tight.
Propusition 2. Under the conditions of Theorem [ if all the contexts are shared, i.e., E == E for all & E (T), then for any {(W kr Won) ) H 1 we have rank(Z) & H(n 1) +1.

---

Published as a conference paper at ICLR 2024
We defer the proof to Appendix E The proposition shows that the lower bound of Proposition | is tight when the context is shared across data points and d, 2 n. Furthermore, the setting of Propositionn[better highlights the dependency factor of r - 1 rather than To illustrate this, consider this setting for r = n = i.e., each context contains a single token and is shared across all examples.
(t) Then, all ( become scalars, whose values equal to one. Consequently, the vector 2 copies the first (and only) context token, which is shared across all examples. Hence, no matter how large H or d are. the output of MHA for all examples is identical, and the memory capacity of the attention is exactly T = Hrr 1) + 1 = (i.e. a single example).
4.3 COMPARISON WITH TWO-LAYER RELU NETWORKS
Given Assumptions D and 2 we compare the memorization capacity of MHA with one-hidden layer FCN networks of the same size. An FCN does not distinguish between key/query vectors. Hence, we flatten out the concatenations of (!)) and e(), resulting in input oPpN == [e] j...: ;.... ::.:.;;:: e e(0] E p(++i)d t E [T]. The proposition below upper bounds the memorization power of a ReLU FCN.
Proposition 3. Suppose a two-layer ReLU network with impul dimension (n + 1)d, hidden dimension m, and output dimension dow- Then, the memorization power of this network for the class of datasets
T = ={EE(, e((), . {)]]]= with context size n < d and satisfying Assumptions 7 ladd 2 is at most T S (m h//m/ouu + (m + 1) examples.
We defer the proof to Appendix F A ReLU FCN with @(ndm + mndgu) parameters can memorize at most O /dau + m) examples under Assumptions I and 2 This setting is particularly interesting with m = (()) and daut = 1, where a two-layer ReLU network with e (Hdd) can memorize at most (HH) examples. This upper bound memorization of a ReLU network together with the lower bound of nn) as established in Theorem 1 for MHA with the same order of parameters, shows that the latter is least as powerful, in terms of memorization, the former (when data satisfy Assumptions D and 2). This also aligns with practical Transformer implementations, where a similar order of parameters allocated to both Attention and ReLU FCN (e.g., (Touvron et al. 2023).
4.4 ANALYSIS OF THE VALIDITY OF ASSUMPTIONS
We confirm our assumptions are valid either on the embedding layer or after just one Attention layer. Detailed experimental verification of this full-rank assumption is provided in Section 5.1
First, we motivate Assumption 2 by considering the common practice of incorporating positional encoding into token representations before inputting data into Transformer. A widely used positional encoding method utilizes sinusoidal positional encodings (Vaswani et al.] 2017). Consequently, the positional encoding matrix is inherently full rank. In view of the positional encoding equations (see Sec. 3). this characteristic extends to the context matrix E(!) establishing the validity of Asm. 2
Second, we provide an explanation for why Assumption D generally becomes valid after just one Attention layer. Consider the case of image classification using ViT, where query vectors e( are all composed of the "[CL-S]" token combined with positional encoding at position zero. Here, Assumption does not initially hold. However, we demonstrate below that it becomes valid after a single Self-Attentin layer with skip connection. The key insight lies in the behavior of the "[CLS]" token: as traverses the attention layer, it effectively mixes information from the surrounding context, which is distinct for each image. Therefore, after single (even randomly initialized and fixed) Self-Attention layer, we establish that Assumptions 2 and [ are satisfied. This allows us to directly apply Theorem to the second layer of Attention to achieve memorization. This line of argument is conceptually similar to the approach employed by Vershynin (2020), who first designs layers with random weights to enrich the representation of data points, and then proves the existence of learned subsequent perception layer for memorizing the labels in a FCN.
Now, we provide the details on how one Attention-layer with skip connection validates Assumption I To begin with, we recall how a Self-Attention layer with a skip connection works. Consider a single example with context matrix E and query token e. Given a set of attention weights W, the Self-Attention layer with skip connection transforms the tokens according to the following equations:
e' := e + Jw Jw e), e/ e= e, + Sw(EE,) (E,e:) for all i E [m), E' == [e ; ; i.. :e,] ),
where e', and E' correspond to the output of query, and context vectors, respectively. Note that the query matrix is part of the context matrix due to Self-Attention. Let us assume the query vector is in the first row of E, iee el e, and consequently e' e
7

---

Published as a conference paper at ICLR 2024
Remark 2. Theorem D holds true for Attention layers with skip connection (see Appendix
We are now ready to state proposition on the mixing power of this type of Attention with skip connection. To do this, we introduce the following relaxed assumption.
Assumption 3. For training set T (( e(t), y ( for all t E [T]. Then the set S := { E(2), (()))).....)) has Kruskal Rank at least
define ε(0)
= :00)=c)
+ 11 LiEnn c
Proposition 4. Define Self-Attention layer Ao with single head, skip connection, dy = d. and weights (()), y
WK- = Wo= = 0, Wv = Wo = I. Suppose training set T = (EE
satisfies )()))) }
Assumptions 2 and 3 Let E*(0) e!(t) be the output of Ao for each example & E [T] and define new
dataset T (E e((t) ,y ) Then, T' satisfies both Assumptions
and
Assumption 3 assumes that the input data as a whole is sufficiently distinctive across different examples. In partieular, this is more relaxed assumption compared to Assumption [ For instance, in the case of images, it corresponds to the average of all image patches (plus a constant vector) being sufficiently distinctive across different images. Proposition 4Jstablishes that under the given assumptions (specifically, Asm. 3 in conjunction with Asm. 2 the application of a trivial Self- Attention layer results in the fulfillment of Assumption at the output of the Attention layer, all while preserving Assumption 2 This implies that the "[CLS]" token is capable of performing token mixing even with a basic Attention mechanism. Consequently, Theorem [ caan be directly applied a second layer of Attention, facilitating effective memorization of the examples.
EXPERIMENTS
5.1 TESTING ASSUMPTIONS ON VISION TRANSFORMER (VIT)
As mentioned in Section we show that our assumptions on the input tokens hold after applying a single Attention layer. To test this, we mainly focus on image classification and Vision Transformer
(ViT) (Dosovitskiy et al. 2021: Wolf et al., 2020). We consider the following models: 1) Embedding Layer: A randomly initialized embedding layer. This layer contains only positional encoding and a linear embedding layer (i.e., the input tokens E("), (!)).. 2) Random Attention: A randomly initialized Self-Attention layer with similar architecture to VIT layer, without the FCN of a Transformer. 3) Random ViT: A randomly initialized base VIT. We only look at the output of the first layer and discard the subsequent layers. 4) Trained VIT: Similar the random ViT, with the difference that we take the weights from the ViT pre-trained on ImageNet (Deng et al.) 2009).
We evaluate the mentioned models on 2000 images sampled from ImageNet. To empirically test Assumption 2 we verify whether the context vectors are all linearly independent for each example. On the other hand, testing Assumption T is computationally difficult since computing Kruskal Rank is NP-Hard. Therefore, it can only be approximately tested in polynomial time. To address this, we randomly sample n query vectors from the training set T and construct a matrix of size n x d. We then test if the rank of the matrix is equal n. After repeating this test 5000 times, we report that the assumption holds if the rank check holds for at least 99% of the test instances. The results of our empirical tests are summarized in Table Specifically, we find that only Assumption is violated in the Embedding layer, whereas both assumptions hold true for the remaining models. Notably, the commonly made assumption of General Position fails in all cases, and the Kruskal Rank typically much less than d, as shown by Figure We also empirically test Assumption 3 and find that it already holds for the Embedding layer. Following up on Section 4.4 this may explain why Assumption [ already holds after applying any type of Attention.
5.2 EMPIRICAL VALIDATION OF MEMORIZATION IMPROVEMENT FACTORS
Theorem Pheeoremprrvvie provides three conclusions: (1) When fixing dimension d. increasing the number of heads H improves memorization. (2) When further fixing the number of heads H. increasing the context size n improves memorization. (3) When fixing d, n, increasing d, only helps up to d, < n, and there is no memorization gain beyond that. By conducting synthetic experiments, we verify the conclusions drawn by the theorem. Throughout, we fix dimension d = 64. Moreover, follow Assumptions 2 and I we generate random query inputs, and random shared context inputs (see Appendix [.I for
*Our code is available at http::/gthh..cccoommmmmaaaaaaaae

---

Published as a conference paper at ICLR 2024
Souze od Ractin
0.8
Model
Gen. Pos.
Asm.
Asm. 2
0.6
Embedding
Rand Attention
X
0.2
Rand VIT
X
v V
Trained VIT
X
0.0
Kruskal
Testing Kruskal Rank of query tokens on the oulput of one layer Random Attention on ImageNet. The Kruskal Rank only slightly larger than n (Assumption D and much smaller than d (General Position).
Figure 1:
Table I: Testing assumptions on ImageNet and ViT: while only one assumption is violated in the bedding layer, both hold true after a single layer of Attention or Transformer. The General Position assumption is violated in all cases. Classiticttion Accuracy
Classification Accuracy
Classification Accuracy
T=:00k
1.0
0.8-
0.6
0.9
0.6
Tolkk
0.6
0.4
0.4
T=10k
T-14k
0.7
T-15k
0.2
T-10k
0.6
Tw10k
Tolk
0.0
6.0
20
60
10
20
30
50
50
20
40
Heeds 0O
Contest Sie
(a)
(b)
(c)
Figure 2: Testing memorization as a function number heads (a), context size (b), and head size (c) for classification under Assumptions and Examples generated synthetically with d = 64 and shared context (see Proposition]? 2). Memorization increases linearly with H, monotonically with n, and monotonically with d2 long dn S
experiments with non-shared context) according uniform distribution, as follows:
01, 01, En (4,1 (1) 1))4,
2.
for all i E [no]. / [T],
where xr!" x are the raw inputs. Each input is then given to the model consisting of an embedding layer, positional encoding, and Attention layer. To test the first conclusion, fix 32, d, and increase H. To test the second conclusion, we fix H 20, dz. d and increase context size In addition classification tasks, we also include regression tasks the appendix.
The results of our experiments are presented in Figure 2 for classification task with various dataset sizes T and uniform labels y E ({(00))) from 100 classes (see Figure 5 in the appendix for the regression task). We measure memorization by computing average accuracy across examples. Observe that memorization increases monotonically with n and linearly with H. To quantify the linearity with H, we measure the R² score of the best linear fit of the non-saturated part in Figure 2a and observe R² 0.98 all cases, confirming high degree linearity. Also observe in Figure 2c that there appears no benefit on memorization fraction beyond dn >
CONCLUSION AND DISCUSSION
In this paper, we proved a lower bound on the memorization capacity of attention layers under a rather mild linear independence assumption. Our proof technique relies on establishing tight bound on the rank of a matrix of intermediate representations in the attention layer. We discussed implications of our lower bound which we also confirmed by additional synthetic experiments. Our work opens up several promising avenues for future research. First, it would be valuable to extend our theoretical results beyond one layer of attention networks and to woweeeneett-sequeeee learning scenarios. Second, while we studied some specific cases of interest, establishing general upper bounds on the memorization capacity of MHA remains an important open question. Finally, refining our data assumptions to narrow the gap between practical memorization abilities and theoretical guarantees could shed more light on the memorization capacity of Attention networks (see Appendix mb. Ultimately, we hope that gaining better understanding of the memorization abilities of Transformers could facilitate more efficient implementations and better use of data for training more effective, generalizable, and preeey-preserving models.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was funded, in part, by NSERC DG Grants (No. RGPIN-2022-04636 and No. RGPIN- 2021-03677), the Vector Institute for AI, Canada CIFAR AI Chair, and Oracle Cloud cred- its. Resources used in preparing this research were provided, in part, by the Province of On- tario, the Government of Canada through CIFAR, and companies sponsoring the Vector Insti- tute (www. instttreett aaltt.tte. ai), Advanced Research Computing at the University of British Columbia, the Digital Research Alliance of Canada (a liancecan. ca), and the Oracle for Re- search program.
REFERENCES
Eric B Baum. On the capabilities of multilayer perceptrons. Journal of Complexity, 4(3):193-215, 1988. ISSN 0885-064X. doi: https://doi.org org/ 10. (800888800004(88)))0)))0
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.). Proceedings of the 37th International Conference on Machine Learning. volume 119 of Proceedings of Machine Learning Research, pp. 864-873. PMLR, 13-18 Jul 2020.
Simone Bombari, Mohammad Hossein Amani, and Marco Mondelli. Memorization and op- timization in deep neural networks with minimum ver-paraameterization. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.). Advances in Neu- ral Information Processing Systems, volume 35, pp. 7628-7640. Curran Associates, Inc., files/paper/2022
www.ipe. 2022. URL https::/ http:://s.....rrr.e.rocoopaaapaa File/323 mtle/323746f0ee http/////.eerbb00000002255oeeee.
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-Voss 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, lly Sutskever, and Dario Amodei. Language models are 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 httpt::://// pttpLtpp:e/ppceeenss...uus wwwpppppccpppppppppppppppppppeee files/paper) 2020 1et1111114555005eee/eee1e oar.://c........bbbacca.22222aa.
Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, and Dan Mikulincer. Network size and size of the weights in memorization with two-layers neural networks. In Newral Information Processing Systems, 2020. wttttttss://////
Nicholas Carlini, Florian Tramèr, Eric Wallace, Matthew Jagielski, Ariel Herbert-Voss Katherine Lee. Adam Roberts, Tom B. Brown, Dawn Xiaodong Song, Ulfar Erlingsson, Alina Oprea, and Colin Raffel. Extracting training data from large language models. In USENIX Security Symposium, 2020.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenat- ton, Lucas Beyer, Michael Tschannen, Anurag Arna,, Xiao Wang. Carlos Riquelme Ruiz, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd Van Steenkiste, Gamaleldin Fathy Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Collier, Alexey A. Gritsenko, Vighnesh Birodkar, Cristina Nader Vasconcelos. Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetic, Dustin Tran, Thomas Kipf, Mario Lucic, Xiaohua Zhai, Daniel Keysers, Jeremiah J. Harmsen, and Neil Houlsby. Scaling vision transformers to 22 bil- lion parameters. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 7480-7512, PMLR, 23 29 Jul 2023. URL
ing...ngggs.. mlr. .pressprssss...... / v202 / dwtp//ihhhaaiinn 23a. .ttml
Mete Demireigil, Judith Heusel, Matthias Löwe, Sven Upgang, and Franck Vermet. On a model of associative memory with huge storage capacity. Journal of Statistical Physics, 168:288-291 2999, 2017.
10

---

Published as a conference paper at ICLR 2024
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale scale hierarchical image database. In 2009 JEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Jacob Devlin, Ming- Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pp. 4171- 4186. Association for Computational Linguistics, 2019. doi: 10.18653/v1/n19. 1423. URL
Yihe Dong. Jean- -aatistt 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 Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 2793-2803. PMLR, 18-24 Jul 2021. URL https::/prrceedii:
m www.prsss/viii. vttp/vv39////// 0tless//v133//ogg22ll.. html
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL dettps://openee
net / forum?? ?id=
Yi YicbFaNTTy
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and vari- able creation in self-attention mechanisms. In Kamalika Chaudhuri. Stefanic Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings Machine Learning Research, pp.
5793-5831. PMLR, 17-23 Jul 2022. URL https::/ wwwinngs.m...po. mttp://ssmm......sss///1//2/ Iidelman22a. html
Wikimedia Foundation. Wikimedia downloads. URL httph:httpp//nttit www.iipps..iiimmiii.. utL.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transform- ers learn in-context? a case study of simple function classes. In S. Koyejo, S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.). Advances in Neural In- formation Processing Systems, volume 35, pp. 30583-30598. Curran Associates, Inc.. neurips http:/ppprrriilssppppppp//////// 2022. URL 1tle:c/22obbbllleeee6ee/aaeeeee2
pdf
Mor Geva, R. Schuster, Jonathan Berant, and Omer Levy. feed--orw layers are key- value memories. ArXiv, abs/2012.14413, 2020, URL stctpp:://ap..semantiie
org/Corpus ID: ID: 229923700
Roger Horn and Charles R. Johnson. Matrix Analysis. Cambridge University Press, 1990. ISBN 0521386322. URL.
Samy Jelassi, Michael Sander, and Yuan Yuan-Fang Li. Vision transformers provably learn spatial structure. ArXiv, abs/2210.09221, 2022.
Junghwan Kim, Michelle Kim, and Barzan Mozafari. Provable memorization capacity of transformers. In The Eleventh International Conference on Learning Representations, 2023. URL https
net wtt.:/eeoooommn11-00o55ooooo
Hongkang Li, M. Wang, Sijia Liu, and Pin- Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. ArXiv, abs/2302.06015, 2023a.
11

---

Published as a conference paper at ICLR 2024
Yingeong Li, M Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and implicit model selection in in-context learning. arXiv preprint arXiv:2301.07067, 2023b.
Shengjie Luo, Shanda Li, Shuxin Zheng, Tie- Yan Liu, Liwei Wang, and Di He. Your transformer may not be as powerful as you expect. In S. Koyejo, S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35. pp- 4301-4315. Curran Associates, Inc., wwtpurris.cccppppee www.://aaeeee///////00
2022. URL
OpenAl. Gpt-4 technical report. ArXiv, abw2303.087744, 2023.
URL https://api
semanticschen Corpp
Sejun Park, Jaeho Lee, Chulhee Yun, and Jinwoo Shin. Provable memorization via deep neural networks using sub-linear parameters. In Mikhail Belkin and Samory Kpotufe (eds.), Proceedings Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pp. 3627-3661. PMLR, 15-19 Aug 2021. URL htttps://proceeeee
press/v33/parrkla..
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Shashank Rajput, Kartik Sreenivasan, Dimitris Papailiopoulos, and Amin Karbasi. An exponential improvement on the memorization capacity of deep threshold networks. In M. Ranzato, A. Beygelz- imer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 12674- 12685. Curran Associates, Inc., 2021.
Hubert Ramsauer, Bernhard Schaft, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitner, Milena Pavlovi'c, Geir Kjetil Ferkingstad Sandve, Victor Greiff, David P. Kreil, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. ArXiv, abs/2008.0224 2020. URL https: api semant icscho. ar orr/ Pegg/Copussin 220:229222990 958978
Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Un- raveling attention via convex duality: Analysis and interpretations of vision transformers. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 19050-19088. PMLR, 17-23 Jul 2022. URL mlr. www.sspppeee/eeeeee.eee v 62/ wwtp:/aaaiieeaaaaamammm
https: htttp::/ppxccedl.. hpppoceeeiinns..ppeppeppp.....se
Eduardo D. Sontag. Shattering All Sets of 'k' Points in "General Position" Requires (k 1)/2 Parameters. Neural Computation, 9(2)3337334, 02 1997. ISSN 0899-7667. doi: 10. 1661/nnnccc. 1997.9.2.337. URL httph:tttt//ttt org.
62. neco. 1997. 9.2.337
Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. End-to-end memory networks. Advances in neural information processing systems, 28, 2015.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
Gal Vardi, Gilad Yehudai. and Ohad Shamir. On the optimal memorization power of reLU neural networks. In International Conference on Learning Representations, 2022. URL https:// openreview net. meet/forumm httee/orruumddMMPPPttttae d=MKTPtn
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, 2017.
Roman Vershynin. Memory capacity of neural networks with threshold and rectified linear unit activations. SIAM J. Math. Data Sci., 1033, 2020. URL hetps:
semant csccollall ott.aarrrr///orrrttr2225555555
12

---

Published as a conference paper at ICLR 2024
Johannes Oswald, Eyvind Niklasson, E. Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. ArXiv, abs/2212.006777, 2022.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger. Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface 's transformers: State-of-the-ae natural language processing, 2020.
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. Gamett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to se-queene functions? In International Confer- ence on Learning Representations, 2020a. URL https: /tRL:ttt:///peeeeeeeee /pppreeeeeeieeeeeeeeeeeeeee net forum?id' ByxRMONtVE
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 transformers. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1788-113794, Curran Associates, Inc., 2020b. URL http://s///proceeigg....eee (tt/0//f1ll///
cc/ paper. cappaper files/ pdf
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017. URL http::tttpttteeeiiiiiie
Jiawei Zhang, Yushun Zhang, Mingyi Hong, Ruoyu Sun, and Zhi-Quan Luo. When expressivity meets trainability: Fewer than n neurons can work. Advances in Neural Information Processing Systems, 34:9167-9180, 2021.
13

---

Published as a conference paper at ICLR 2024
Symbol
Description
A
A single-layer multi-head Attention (MHA) mechanism Training set
T
H
Number of heads in the MHA layer Training set size Context length
T
WK0, Wa, E ddd Wvn E Rdxd,
Weight matrices for each head h E [H]
Wo E Hddddd
Weight matrix to combine the outputs of different heads Final read-out weight of the model
WD lRd
W
The entire set of parameters of the MHA layer Key/Context matrix containing n key/context vectors
E E Raxd
E Rdd
Query vector
y E R
Ground truth label
Ot,, On
Rh. Ph
Intermediate variables in the MHA layer
Output of the MHA layer
ÿ E R
Final output of the model
Table 2: Important notations used in the paper.
A EXPERIMENT DETAILS
In the synthetic experiment, for regression tasks, we generate uniform labels y ~ 4(0, 1), and for classification tasks, we draw uniform labels from 100 classes y(t) ((100))). For classification, we measure memorization by computing the average accuracy across examples in dataset. For regression, we use the mean squared error (MSE) loss to compare memorization abilities. We train each task for at least 500 epochs and at least 50, 000 optimization steps (whichever is larger) with a batch size of 256, using an Adam optimizer with a learning rate of 0.001, and a scheduler with linear warmup and Cosine decay. For the experiments reported in Figure 2c and Figure 3 since the optimization is more challenging, we train for a total of 300, 000 optimization steps.
While our experimental results in Section 5.2 conclusively show that increasing the number of heads H or context size n improves the model's s memorization capability, we caution the reader that the experiments are done purely through gradient descent optimization, which has limitations when comes to studying memorization. First, it is not guaranteed that if memorizing solution exists, GD would be able find it. Second, optimization requires hyperparameter tuning and a large number of gradient updates to achieve memorization, especially in large-context sizes (Garg al.] (2022). We mitigate these issues by running small-scale experiments with sufficiently high number of optimization steps and adequate learning rate scheduling. Running the experiments reported in our paper takes approximately 20 GPU days on an NVIDIA V100 GPU. However, this procedure becomes highly resource-exhaustive and prohibitive in larger-scale experiments. For instance, we find out that having large d, (beyond n) can generally improve optimization speed, and lead to better memorization. However, with sufficient large steps, this gap closes, and gains in memorization vanish in the d, > n regime (as observed in Figure
regimes and
B PROOF OF PROPOSITION
For convenience let us define r = min(n, d..). WLOG let us assume that T 2 H(r 1. This because for the case of T < H(r 1) + we can add extra dummy data points that satisfy Assumptions ] and 2 and make T = H(r 1) 1, then prove Z has full row-rank, and finally remove the extra rows of (i.e. dummy data points) to prove the statement.
We prove the lemma by induction on H. We will demonstrate that by adding each head, we can increase the rank of Z by r The key to proving the induction is the following two claims:
Claim 1 (Fit examples in one head). Assume that examples (EE!) e satisfy Assumptions
and 2
Fix any head index h [H). For any arbitrary set of convex combination vectors
14

---

Published as a conference paper at ICLR 2024
0 VI. E [r], there exists an Attention head with key/query weights WKh: Won which achieves such coefficients for the set of data points.
Claim 2 (Fit examples in one head, leave the rest intact). Assume that T 2 examples {(E(), } satisfy Assumptions I Lnd 2 and define 'R. = (T' - + +2....... 1, 1,T")
the set of last examples' indices and R (T) RR as its complement for convenience. Fix any head index h E [H]. Then, there exist 6 t R such that for any arbitrary set of desired comvex
combination vectors 0!!) tt E RR and any 5 > 0, there exists keylquery weights WKh. Won, such that both k/)) !! Vt R. and 10
Claim I states that we can design an attention head that produces any arbitrary 0, coefficients for examples. Furthermore, Claim 2 reinforces Claim by demonstrating the possibility of tuning coefficients 03 of any precision 8. However, this comes at the cost of fitting one less example into single head. The proofs of both claims will be deferred until the end of the proposition. We are now ready to prove the proposition by induction H.
examples while keeping the rest of the examples approximately fixed with
Step 1. The base of induction H 1. First, we prove the base of induction for 1. Let us only consider the first r examples (i.e., T = r). Notice that the matrix Z has rows and d columns in this case. Moreover, the f-th row of Z has the form
(2) Z, E 0(() ei
(1) 0
and is convex combination of the context tokens for the example t. According to Assumption
2 the
(t) span of context vectors eix] has dimension Therefore, we can inductively find o 00(²)
..... 0 (r)
such that each row Z, is not in the span of the previous t rows. We continue this process for all & < r since the span of the previous rows has at most - dimensions, but the convex hull (t) has at least n linearly independent elements (r S n), so it is not covered by the previous t - rows. Therefore, there exists (0()) Y== such that the resulting Z matrix has rank at least Then, using the result of Claim there exists Attention weights Wx We, that produce the desired set of convex combinations
[n]
yi=1= (0
and subsequently the first r rows of the desired feature matrix Z.
Step 2. Induction Step. Next, assume that for H = H' there exists heads with Attention weights {(WK,, Kb, Won) )²r such that matrix Z has rank H' 1) + 1. We prove for H = H' + there exists Attention weights {(WKr: Wo H' such that the resulting matrix Z attains rank (H' + 1)( 1) 1. Let refer to the matrices derived from the induction hypothesis by prime D.. Now, for the case of H H' we the weights of the first H' Attention heads to be equal to that of H'. Namely,
superscript (i.e., Z'. and {WWkh Kh. W We')) }.)))>))
WK == Wk'. Won == We'
for all h E [H']. Then, the feature matrix Z can be written as the following block form:
Z
Z'
z"
H'd
where Z' E RTxHH' is the matrix derived from the hypothesis and 2" E RTxd is the block corresponding to the new head H' + 1, whose parameters we have not yet determined. Now, we know rank(Z? Zz) > 1I) 1. Define l: == rank Z'). WLOG, assume that the first l rows of the matrix Z' are linearly independent (otherwise, we can rearrange the rows of Z, L.e. the examples, to satisfy this property). Now, take the last 1 rows that are linearly dependent We defer the case for when less than r - such rows exist (i.e., + r - T) to the later stages of the proof.
the first k rows
15

---

Published as a conference paper at ICLR 2024
Let us partition the matrix Z into two blocks: the first T r 1 rows and the last r 1 rows. Namely,
Z Zb
Z
Z"
Z
(9)
Z
Z
Observe that the rows of Zj are all linear combinations of rows of Z So, we can zero out the Z block by adding a linear combination of rows of Zz' Such operations add a combination rows of Z" to the Z as well. Notice that the rank of Z stays the same since such operations do not change rank. So let us call op the resulting matrix:
Z'
Z
op =
0
Z + A
where
A
A( Z z=)== C(Z' , E (-1xx
contains linear combination of rows of Z (whose coefficients are in C(Z', Zz) and are function of Za Zj). Since
Lemma
rank (Z) rank (op
rank(Zū + A) + rank (Z.) rank Z + A) + (r) 1,
we have find suitable Attention weights such that rank Z A) 2 r to achieve a (H' -)) rank for Z and complete the proof.
Now, consider Claim for some fixed convex combination vectors {0) HER,, and any desired (0!!) (0H/+))EER- For any 8 0, the claim statement gives keylquery weights, and subsequently resulting intermediate matrix Zz. Let us define Z as the limiting matrix lim Zg. and similarly define its sub-blocks corresponding to this head by the same sign, i.e., Zz ZK A., and similarly Z",5- Zh,s As. Notice that A+ is only a function of Z Zi Z", As a result, different choices of (0l) H' heR produce the same A+ since {
the
HER- remains fixed under Claim 2 Therefore, we
first find ((() 0 such that rank (Z",, + A+) 2 r 1, then find small enough ō that properly approximates the limiting matrices. Similar to the base of induction (Step 1), we inductively build the individual rows of z"' but the difference that the original matrix Z is added fixed
matrix A+. Hence, each inductive step t, the following modified constraint must be satisfied:
E" A+], & span 1" + A+), 1 A+|2....
Z + ++I1,1)
which can be satisfied by a stronger constraint:
1%*+,, $ span
[A+]. [2* A+], [2] + A+1 4 + A+],-1)
Hence, since the input context tokens are in an n dimensional space, the set of feasible valid rows [2 lies in an n dimensional space as well, which allows us continue for n 1, (n r) rows
(and not n, due to presence of [A.))), find suitable + hER., and build a full-rank matrix Z with rank (Zl', A+) 2 - 1. Notice that no well-defined key/weights generate Z+ as the weights need be infinite (see proof of Claim 2. However, we only need small enough 8 such that
rank Z As) 2 rank(2,". + A+) 2 1--1
(10)
for the weights be well-defined. In Appendix B.2 we prove the existence of such 8. Therefore. there exists W WHH+ W W1++11 that produces rank Z + A) 2 r - 1. and the proof is complete.
The case of less than 1 linearly dependent rows of Z'. In this case, we have k - 1 T. Let us define m : T l < r 1. Since we have assumed 2 (H' + 1)(r - 1) + 1, this means that that Z' already has rank at least (H' + ()(r - -1) + - m. So, the head index H' + 1 needs to increase the rank of Z' by m. Since m S r - 1, all the above steps also apply to m and the rest of the proof procedure is the same.
16

---

Published as a conference paper at ICLR 2024
B.1 PROOF OF CLAIMS 1 AND 2
First, note each % is derived from applying Softmax nonlinearity on the logits / So, we need (t) to solve system of equations for h, for each / [T] Since our focus is on single head, we omit the subscript h for simplicity. Furthermore, we define W := Wx W Now, for data indices mm [T] where n, we are interested in solving the following system of
l/ equations:
W Wel = a.. (b:)
Vi [m].
Since each E(4) E Raxd is full-rank with rank n, it has a pseudo right inverse E(4) E RRdx such that E(4) E(t))* SS, we solve the following system of equations
We E(('' (ta)
Vi E [m²]
E We a (b) Vi [m]]
Now, let us define (.4) = E(t)' ax (1)) Then, we get the following linear system of equations to solve:
W (e elta),
e(mn]]= - ( - B
Due to Assumption the query tokens have Kruskal Rank at least and S r S so the query token matrix the left-hand-side above has a pseudo left inverse. Thus, W has at least one solution
with rank at most Assuming the solution is W E Idd there exists Wan: WQn E Rdxds such that WKA KA Wen W for the weights of the Attention head. So, we can prove Claim by setting and (t1, 12, .,) = rT.
Next, to prove Claim 2 we prove if <r then we can not only achieve any desired logits for the examples (T m 1,7 1, m 2 T). but also keep the (0(*) "T= approximately unchanged. So far, we have found W that achieves the set of desired logits. Now, there exists ran linear transformation W+ IRdxd such that
E W (tt) O for all T m <t ST
(11)
E e())
+ (E( w e!),, for all 1 SIST - m, i,j E (m)
(12)
This is due Lemma
So, we set W == W + cW' and choose a large enough c. This way, the desired logits for T m t < T stay the same since cW+ has effect them, while for the rest of logits, eW+ dominates W and saturates the Softmax operator. Lemma@ppove the existence of such c, so, we
can prove Claim 2 by finding WK Won, E dddx such that WK.. Wol W + cW+ + Note that this construction is feasible because rank (W* + +cW+ S m + S S di-
B.2 PROOF OF EXISTENCE OF 5 FOR EQ.
10
Let us call the (()) ; derived from Claim as (()) So, for any 5 > we have:
< 5. for all t E R
0 l for all f E R
Th = z%
To prove the rank inequality argument, we use corollary of the Weyl's singular value theorem (Lemma 5) and prove:
As = (Zi's + A + A+), Fmax (As) < Omial ( + A+),
(13)
which in turn leads to rank Z + As) > rank 2h,+ + A.). First, let us assume Z + A, is not zero matrix (otherwise we must have 1 = O, in which case any 5 > 0 would satisfy the
'The set of logits that produce convex combination is not necessarily unique. However, uniqueness is not requirement in this case, we only require the existence of some set of logits that produces convex combination.
17

---

Published as a conference paper at ICLR 2024
inequality). We now bound the spectral norm of the perturbation matrix As, Lee, maax mmax(A6).
O max (As)
|
||As A+ll2
(2 (g,0)
|| Z,) 2",5 C( Zz) Zh,+I.
||C(Z, 2₁) (Z",5 Z", )l1
[C(Z) 1) Z",' 20,112
- |[ 21 Z",
-",+
where C(Z,, Z3,) is a constant matrix containing coefficients derived from zeroing out Z using the rows of Zz To bound the Frobenius norm of Z", T" notice that t-th row of this matrix is:
| Z", 10,t
( (1+2))
E (
Ett) T
0
Elt)
1.2
emax 6.
where 6max is the maximum token Euclidian norm of all tokens in the dataset. Therefore:
VH,6 D",+ S TTT'mmxo
Hence, choosing
mmin 1 + A+)
85 VT e max |C 2;)112
gives the desired inequality Eq. 13
C EXTENSION OF THEOREM TO ATTENTION WITH SKIP CONNECTION
Theorem T holds for when have Attention with skip connection. Recall that skip connection keeps all the computational mechanisms of MHA the same except for Eq. where the updated equation is:
0 = Wo PP::Pii...:pp ee
(o E F
The only modification that we have to make, is to modify the labels. Let us consider the setting of Theorem with dataset T = (EE)) e(!), (*) } and Attention weights W. Now, if add the skip connection to the Attention layer with the same set of weights, the new g(()) will be:
g = Wp T + (()) Ii ((e!!) + +o( = (*) Iimuxee!!
Therefore, we get labels (t + -Iauxdel which is the true label shifted by the first dour dimensions of the query vector. So, we need to define new dataset
T' T'= T'=((E') (E(1), (")),
Idmuxd 1=1
and find Attention weights W' for ''' according to Theorem 1 Then, we can use the set of Attention weights W to get the correct label (t)) + Iduxde((1) Io [ for an Attention layer with skip connection and dataset TT
D ADDITIONAL LEMMAS
Lemma 1. For any real-valued block matrix X of the form
A
C
X
0
B
we have rank(X) 2 rank(A) + raak B).
18

---

Published as a conference paper at ICLR 2024
Lemma 2. Assume vectors e1, ........ En E RR are linearly independent. Then,
e/ e; - (el + E2 ......................... en,,
also linearly independent for all i E [n].
Proof. Due to the linear independent assumption, the only solution to
01 el + 0262 + ....................... + 0n = s.1. d1 02,02..... . On R
is 02 0n = 0. Now, let us examine the same system of equations for e/p Namely, e + byez +... b
s.t. .n,b,....,be bn E IR.
We prove zero is the only solution to this system to conclude the proof.
biel - b2e + en
=bje bye2 + ... bn (b1 + b2 ... - (e) + +22 + ee,)
-(b1 + b)ei + (b2 + b)ez +.... + (bn b)e, 0. According to the linear independence assumption, we must have b, +b for all i [n].
By summing all the terms up, we conclude that 6; for all i E [n] is the only solution and this concludes the proof.
Lemma Consider a set of full-rank matrices E E
..... ,.... E E Baxd where n
Moreover, consider a set of vectors el e(?) ...., (T)) Bd with Kruskal Rank at least n. Then, for any m S n there exists a rank-one matrix W E idddd such that
( =0 for all sts m,
(14)
we * [E"We("), for all m < 1.j
(15)
Proof. Let us define a basis for Rd as bi, b2, b by E Rdd such that b; = ()) for all E [m], and basis. Then, linear transformation
bm+1.... b ... . ba are an arbitrary extension of e|m] W uniquely determined by defining Cı, C2, C2 Rd and setting
Wb: G, for all i E [d].
First, to satisfy Eq. 14 we set C; == 0 for all E [m]. which gives Wb; = (e = = 0 for all i [m]. Let us decompose every vector (t) into the mentioned basis, i.e.,
(t)) a bi + a b2 +.... a.... ba for all + -stt Since [e(")T= have Kruskal Rank at least n, then for any t [T] at least one of (t)
am++1 mm ................................ od must be nonzero, otherwise, e!!) can be written as linear combination of the first m vectors. Moreover, there exists fixed um Um+22...,,W Ma E R such that
am+ Mm++ + 2m+2 ... " ud + 0 for all m + 1 < tt d, for instance, we can choose u, iid from Standard Gaussian Distribution and satisfy this property almost surely. Now, let us define
(1)
(t)
c c~ C~,(0,,)
(16)
C; = u.c for all m 1sis d.
(17)
Therefore,
(e =a Wb, + d2 Wb, + +... a Wb.
(18)
an (1)ci az (1)
C2 + + a Ca
(1)
(19)
amt1 Cm + m+22mm22 ... +0
(20)
-c Um+1 + (1)
2
a us
(1)
for all m 1
(21)
19

---

Published as a conference paper at ICLR 2024
with nonzero coefficient. Now, since is standard gaussian, then
(1) (e N (o. (o.(mmt ( 1m-1 + a 4m+22 ... + ad ud) I
(6)
Moreover, E(() is full-rank, none of its rows are equal to each other. As a result,
c (e-" (e"- We
follows Gaussian distribution for all i,j [n], 7 j. Since {0) has zero Lebesgue measure, the probability of having two equal elements is zero. So, Eq. [3] satisfied almost surely. Note that this construction produces rank-one W since the rank of [C1, C2...., Cal is at most one. Therefore, there exists rank-one matrix W that satisfies Eq. 15) and this completes the proof.
Lemma 4. Consider an input matrix E Raxd, imput vector Rd, and weight matrix W+ E IR [adx such that
[EW*, e], * EW el, for all i,j [n], +
(22)
For any W Baxd define W. : W cW+ where E R is scalar. Then,
0+ lim 0g = lim Softmax ([EW,e)), c-000
c-
is only a function of W Consequently. 6g converges to 0 in sense, i.e., for any > 0, there exists Ca such that
|1+ - 0coll <.
Proof. Let us define
at = EW+e, e, and a" EW
Now, ith element of 0 is given by:
exp (a + caa))
0++
lim 0-000
0cl = lim
c0000
Eje(n) exp (az + cac.)
exp (a;))
lim
c exp (at)
541 exp (aj + claaa - -a a.))
arg max ai a ........,..!!!
where the last equality due to azt being distinct. Therefore, [0.]; is one-hot vector determined by W+ Consequently, since each individual entry of Og converges 0+, we get convergence in L and this completes the proof.
Lemma 5. Let A, B E RRxm be two matrices satisfying
Tmax B A) a min,
where mmin, Tmax are the smallest and largest non-zero singular values. Then, we have rank 2 rank(A).
Proof. From Weyl's singular value theorem (Hom & Johnson 1990), know
[xx(A + A) 0x (A) < F (A). for any perturbation matrix A, and singular value index k [min(n, (n, m)). Setting A B - and using the lemma's ' assumption we have:
(Ox(B) ax (x(A)] Fmax (B B A) mmin (A) (A(
for all k such that (A(A) > 0. Hence, for any non-zero singular value k, B has a corresponding non-zero singular value and rank(B) 2 rank(A)
PROOF OF PROPOSITION 2
First, note that when the contexts are shared, the matrix Z has the following decomposition:
Z (In ET O
20

---

Published as a conference paper at ICLR 2024
where @ is the Kronecker product and @ E R Hn is defined similar to Z (t.e., the block in t-th row and h-th column is 0%² ). Notice that @ E RTx Hn, therefore, the rank of Z can be at most Hn. However, we prove here that the column rank of O
most H 1) due the structure of (t) being convex combination vectors. With a slight abuse of notation, let us write Q as: O
1 2,
+g
where each b E RTxx corresponds block belonging to head h E [H]. Observe that θ is a right stochastic matrix (ie., sum up one) for all h [H]. Thus, each row of 6,, O sums up to zero and the matrix has at least one zero eigenvalue. Therefore.
rank Z) S rank (0) rank ((O1, 22 O3.
@r]
rank ((01, O2 O1, 2 O1, ..... O 91))
< rank ))) + rank ( ②1) + +.... (r) (()
2n (n 1) (n 1) + ....... (n (1--1) 1) H( 1) +1.
Hence rank (2) is upper bounded by (n - 1) + and this concludes the proof.
F PROOF OF PROPOSITION
We give proof by contradiction, let assume T > (n + ()///out + m + 1. Consider ReLU network f as
fdi, ddid2,,aoo den (z; W1, b1, W2, W. ReLU (W, + b1) + b2-
where Wi Rddxxd, bi Rd2 W₂ Rduxdd, and b2 Rdm are the network parameters. First, we consider d1 n, d2 m. This network possesses a total of (n + 1)m + (m + 1)dou parameters. Since a trivial upper bound on the number of real-valuee labels network can memorize is at most as high as the number of its parameters, for any
(n 1)m (m 1)dour
(n + 1)m
+ m
dom
daui
and inputs (1), (2)
(T) E RR, there exists labels ()), (2) ..... (T) Rd such that no
fn.m.du can memorize this dataset. In other words, for any W1, W2, b1, b2, there exists index i (T) such that
fnm (x;; W1, b1, W2. b2) yr.
Now, let us consider any arbitrary (1), (2) (2))) ..... (T) R" in General Position (i.e., maximal Kruskal rank of n), and their corresponding set of labels y (2))) .......() y Ruuk for which no fa,m,dm can memorize the labels. Consider the following extended input dataset:
.....
CFCN [8 82: 8ni 0 (d=n)) E IR (n)+))d-n
(t) FFCN [CFCN; [R++1)dd for all / [T], (23)
where s; are standard basis vectors in Rdd Let us call TENN == [(ffeN: (()) )E(T)- First, we claim that no network f(++I)d,m can memorize TECN. To show this, create a mapping between any network of f(n+1)d,m.d.. to networks fn.m,du For any W' e 1mxx(++1)d by E RT, W e
Idd xm and b/2 E Rd define:
Wi Wit: (n 1)d n + 1: (n + l)d
(W1 R=*)))
b, : bj- Wil: 1: (n + 1)d - n²] OFCN
(b, E R
W2 = W²
(W₂ E RRdux Xm)
b2 bbz
(b2 E Bdd
where Wil: a b] is the sub-block of matrix W/ selected by columns a to b. With the above assignment, we get:
fydmm (Ftex: W b'. W bb) = fn.m.dw. (1) W1. b,, W2, bz). for all t E [T]. Therefore, since no ReLU network fn,m,d.. is able to memorize ((x(*) y hemm no ReLU network f(1+1)d,m.d.. is able to memorize TreN either.
21

---

Published as a conference paper at ICLR 2024
So far, have found dataset TENN that two-layer ReLU network with hidden neurons
memorize. Next, we prove that TEON can be derived from a valid dataset T -= (E (E()) ()))) }
which satisfies both Assumptions |
and
2 To observe this, consider the following inputs:
E [s] :811 i..,....8 8,
e*)) : [d)=n)); x
and labels (()) the same as TreN. Then, we get
(t) (e
(+) e!)).
FFCN
So, what is left is to prove that T satisfies Assumptions and 2 First, hetr has the same Kruskal rank as fxx(1) hee which is n, so Assumption is satisfied. Second, each E(!) full-rank since each row i E [m] is the i-th standard basis for all t [T].
Therefore, we have found dataset T with size T' > (n + 1)mm//o/ + + that no ReLU network with m hidden dimensions can memorize, and this completes the proof.
G PROOF OF PROPOSITION 4
First, we derive the output tokens Ex(e) ex(0) as a function of input tokens. We omit the head subscript since we have only one head. Since WK We O, we have (+) O, and consequently g 1. Therefore, the attention head only takes the average of each token, i.e..
n
= 1 (1) + E2 (t)
e
+e("
Now, since the Attention head also has a skip connection, the output of each token becomes:
e0 e + o (t) + e,
(24)
e!!) + (1) e (1) +e for all [n]
(25)
{(t) in Assumption 3 this means that the set {e y has Kruskal Rank at least n, and Assumption holds for T'. Next, since the E(!) are full-rank, adding the average of tokens to (t) preserves the rank (Lemma 2). so E'(() are also full-rank and Assumption is satisfied for T.
First, since we have e0
H SOFTMAX SATURATION AND OPTIMIZATION-BASED MEMORIZATION
In the proof of Theorem T we took advantage of the saturation properties of Softmax. We generate dataset of size T with a shared context (as in Section 5.2 with = 64, n = 32, and discrete labels g/(t) E Z1( [100]) drawn iid from uniform distribution. We test with two different head sizes H 4 and H 8. For each head size, we first empirically find the maximal number of data points the model can perfectly memorize using gradient- based optimization. We find for H = 4, this number is around T = 550, and for H = 8 around T = 1300. We train both models for 300 000 optimization steps. We call head saturated for example if the query token has one softmax coefficient 0.99 to some context token in that head, i.e., head à E [ is saturated for example t E [T] if o has at least entry > 0.99.
The results are depicted in Figure 3 Interestingly, a significant number of heads are saturated for the majority of examples, while almost all examples have at least one corresponding non-saturated head. This is in line with the intuition we provided in the proof of Theorem [] where our proof relied on designating a set of examples one head and saturating that head for the rest of the examples.
I UPPER BOUND ON THE RANK OF MATRIX Z WITH STRONGER ASSUMPTIONS
In Proposition we provided upper bound on the rank of Z, which matches the lower bound of Proposition One might wonder whether under stronger assumptions (and perhaps far from reality), full-rank Z with rank Hd is achievable or not. Here, we prove that even the "General Position" on
22

---

Published as a conference paper at ICLR 2024
Sakinted hash overeto
etorrker weight per heed ovampe
Nr Setendae
0.00
0u:0 Maximrm coetiisint
(a)
(c)
Sckeral bast pEr Ourredo
200
1,000
1000
700
Hc Setimted
0.50
0.30
Domppe
sudlunedd houls
Muximum ccelkaeni
(d)
(e)
Figure 3: Testing the saturation property of Softmax on synthetic data. On the top row, we have H 8, and on the bottom row H = (a) and (d) show whether softmax in each head saturated for the first 20 examples of the dataset, results are qualitatively similar for the rest of the examples. (b) and (e) show the histogram of the number saturated heads for each example across the dataset (H T in total). (c) and (f) show the histogram of the maximum softmax coefficient for each head and example across all the datasets (H x T in total). All the figures suggest that with perfect memorization, the majority of heads become saturated across the examples, with almost all examples having at least one non-saturated head.
all tokens assumption is not enough, and there exists a very simple assumption that limits the rank of Z to (d +
Proposition 5. Assume that we have a multi-head attention layer A with H heads, embedding dimensions dę d, and d, Let T {E(t),,e)) ())) )im be a training set with context size n d. If all context vectors satisfy 1 ee for all E [T], E [m] and some universal )
scalar C E IR, then any set of weights {(W Kh: Wo,) we have rank Z) < H(d- +1
Proposition gives an upper bound the rank of Z when individual elements of each context token sum up to a constant C. Note that all the context vectors can be in general position and still satisfy this condition. Moreover, this upper bound puts no assumption on the query vectors. Hence even with the strongest assumptions, obtaining full-rank matrix Z is out of reach. However, whether stronger assumptions can maintain dependence on instead of n (i.e., give H(d 1) + memorization lower bound instead of H(r 1) + 1) is still an open question for future work.
Proof of Proposition have:
First, note that since 1 e = C, and 0.()) are convex combinations we
1 A = 1 E(t)
g = C1 T h = C for all t E [T].
h
With some abuse notation, let us write the feature matrix Z E RTxHd as follows:
Z Zi. Zz
Zh),
where each Z E Txxd is sub-block of Z corresponding to the h-th head for all h [H]. Then, since each row of Zn, (i.e., (t)) sums up to C, with an argument similar to the proof Proposition get:
rank( Z)
= rank (ZZ,, Zz Zz,
ZH])
= rank (ZZ,, Z₂ Z₁, Z3 Z,,
ZH Zı))
I rank(22 Z1) rank Zz Z1) rank
Zz Z1)
I <d+ (d 1) + (d 1) + ... (d
= H(d- ((d-1) 1) +.
23

---

Published as a conference paper at ICLR 2024
Classification Accuracy
Regression Loss
1.0
T-130k
10
T=110k
10-2
T-70k
Memarizdd Fortton
0.8
T-90k
T=50k
0.6
T-S0K
10-7
T-10k
MSSE
10
10
10-12
0.0
10
20
50
60
20
30
40
60
(a)
(b)
Classification Accuracy
degression Loss
10-1
T-90k
T=70k
Mcmorize Fracton
0.8
T=50k
10
T-30k
0.6
10
MSE
0.4
T=60k
10-*
T-50k
0.2
T-40k
10
T=30k
T=k0k
10-11
0.0
Context Size in)
Context Size (n)
Figure 4: Testing memorization under general position assumption for both context and query vectors. While the linearity H still remains, the linearity no longer strictly holds, suggesting [ 1) + 1 to be better alternative under general position assumption.
I.1 SYNTHETIC EXPERIMENTS UNDER GENERAL POSITION ASSUMPTION
We also synthetic experiments under the general position assumption all input tokens. We consider similar setting to Section 5.2 but we also generate context vectors from a uniform distribution well, i.e.
(t) id u(0, 1)61 for E [T]]
This way, all tokens (i.e., context and query) inside, and across all examples would satisfy the general position assumption. Preliminary results in Figure 4 suggest that the lincarity H still remains, while the linearity in is weakened. We conjecture that bound with complexity O(Hd) would hold, which could be interesting for future works. Yet, as we showed in Section 5.] from a practical point of view, the General Position assumption might be too strong be realistic.
TESTING ASSUMPTIONS ON NATURAL LANGUAGE TASKS
further test Assumptions 2 and on BERT Devlin et al. 2019). and GPT2 Radford et al. (2019) the task of pre-training Wikipedia Foundation Following the setting introduced in Section 5.1 we examine the embedding/first layer of trained/untrained GPT2/BERT.
We follow setting similar to the pre-training of BERT. For each example, an input of the format "/CLS] Sentencel [SEP] Sentence2" is formed, where Sentence / and Sentence are two pieces of text from the Wikipedia corpus. Here Sentence is the sentence after Sentence with probability 0.5 and randomly chosen from the corpus with probability
BERT:
24

---

Published as a conference paper at ICLR 2024
Table Testing assumptions Wikipedia for BERT and GPT2. LO refers the embedding layer. and L1 refers to the output of the first layer.
GPT2
BERT
Model
Random
Trained
Random
Trained
L0
L1
LO
L1
L0
L1
L0
L1
General Position
x
X
X
Assumption
x
x
v
X
Assumption
V
x
v
V
0.5. The "[CLS]" token's output predicts whether Sentence 2 follows the first sentence not. We consider the "[CLS]] token the query token e(!), and the whole input as the context E(t) for all t E [T]. The positional encoding is sinusoidal positional encoding.
GPT2: This model trained autoregressively to predict the next token in each piece of text. We consider last-token prediction from the context. That is, we set the query e to be the last token an example, and the context ( to be whole input example. In contrast to BERT, The positional encoding in this model is learned instead of fixed.
The results of testing Assumptions are shown in Table 3 While linear independence assumptions are not well-suited for discrete data, we observe that most assumptions still hold. For Assumption 2. the sinusoidal positional encodings in BERT preserve the linear independence for all cases. In GPT2, we observe that the learned positional encodings pre-trained GPT2 have lower rank and are no longer linearly independent. For Assumption 1 the GPT2 model already satisfies it after one layer. In the case BERT, the n/d ratio 512/ 1024, and while the Kruskal Rank of query vectors below 512 (violation of Assumption [ observe that the Kruskal Rank is high enough (291 1 for the case of trained and 349 for the case random BERT). Hence, the result of Remark still yields acceptable lower bound.
25

---

Published as a conference paper at ICLR 2024
Regression Loss
Regression Loss
10a2
10-2
10-4
MSE Loss
10-1
MSE
T 30k
10-1
10-4
T-25k
T= 20k
T=20k
10-10
10k
T=10k
10
50
10
Heads (M)
Context Size [n)
(b)
Figure 5: Testing memorization regression task with similar setting Figure 2 except for real-valued labels. The same results and conclusions terms of monotonicity with H and n hold true for regression. The variance in regression plots is magnified in some cases due to the log scale of the plot.
26


Reviews:
[{"review_id": "vIgCKMFOZi", "reviewer": "Reviewer_6Bvj", "summary": "This paper explores the memorization capabilities of multi-head attention mechanisms, a core component of transformer architectures widely used for language and vision tasks. The study aims to understand how effectively these mechanisms can remember example sequences, depending on the number of attention heads and the lengths of the sequences. The research introduces new assumptions about the linear independence of input data, which differ from previous models, and under these conditions, it's shown that an attention layer with a certain number of heads and parameters can memorize a proportional number of examples. The paper further discusses how different attention heads process different sequences and how the softmax function contributes to this capability. These theoretical findings are supported by experimental results using synthetic data.", "questions": "- Could you explain the rationale behind the specific assumptions made regarding the linear independence of input data? How do these assumptions align with real-world data scenarios in transformer applications?\n- How do your findings contribute to the ongoing discussion in machine learning between model generalization and memorization, particularly in the context of overfitting?\n- How do your derivation connect to the Hopfield Network as one classical memory network that has recently been proven to be connected to the Transformer network [1]?\n\n[1] Ramsauer, Hubert, et al. \"Hopfield networks is all you need.\" arXiv preprint arXiv:2008.02217 (2020).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper makes theoretical contributions by exploring the memorization capacity of transformers, an area that is not yet fully understood. This contributes to a deeper understanding of transformer architectures.\n- The paper introduces new assumptions about the linear independence of input data, distinct from commonly used assumptions. This novel approach provides a fresh perspective on analyzing transformer models. \n- The findings are validated through experiments on synthetic data. This empirical approach strengthens the theoretical claims made in the paper.\n- The paper includes a detailed analysis of assumptions, theoretical proofs, and experimental validation, making it a comprehensive study.", "weaknesses": "- Limited Empirical Testing: While the paper includes synthetic experiments, real-world data experiments might be needed to fully understand the practical implications of the findings.\n- Focus on Single-Layer MHA Module: The study primarily focuses on a single-layer Multi-head Attention (MHA) module. Expanding the analysis to multi-layered architectures could provide more comprehensive insights.\n- Potential for Broader Impact Analysis: The paper could benefit from a more in-depth discussion on how these findings impact current transformer-based models in various applications, like natural language processing or computer vision.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CvgrIwNYeQ", "reviewer": "Reviewer_SUAq", "summary": "This paper proposes a lower bound on the memorization capacity of MHA layers under a rather mild linear independence assumption. Based on two input-data assumptions, the authors theoretically and experimentally prove that (1) When fixing dimension d, increasing the number of heads H improves memorization. (2) When further fixing the number of heads H, increasing the context size n improves memorization. (3) When fixing d, n, increasing dh only helps up to dh < n, there is no memorization gain beyond that.", "questions": "1.\tThe authors demonstrate that Assumption 2 typically holds in practice due to positional encoding. Does the assumption still hold when the positional encoding is learnable?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe paper is well-organized and the proof makes sense.\n2.\tThe two input-data assumptions are milder than the General Position assumptions. Although it is impossible to fully verify its generalizability, the author demonstrated the reasonableness of the assumptions through sampling testing, which interests me.\n3.\tThe conclusion “When fixing d, n, increasing dh only helps up to dh < n, and there is no memorization gain beyond that” is enlightening and I believe it can bring more valuable thinking and discussion to the community.", "weaknesses": "1.\tIt might be significantly different between the image patch tokens (ViT) and the language tokens. Can the author's experimental verification of those assumptions be verified on NLP tasks?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tbwVxQsfYA", "reviewer": "Reviewer_mRPx", "summary": "this paper investigates the memorization abilities of MHA of transformers with theorectical analysis. Based on two new proposed assumptions, the authors find a lower bound on the memorization capacity of attention layers. The paper tested the rationality of the assumptions and validate the theorectial findings with synthetic experiments.", "questions": "refer to the weakness.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The assumptions in this paper are more relaxed The authors  verified the rationality of the assumptions on real data.\n2. The exploration of memorization capacity of transformers is meaningful for more advanced go-to architecture, while the memorization abilities of attention modules is quite interesting. \n3. The paper is well-written.", "weaknesses": "1. One of my main concern is the illustration or definition of \"memorization\" in this paper. The inputs of attention include both the key matrix and the query vector. In a common understanding, attention plays a role to capture knowledge from the context according to the \"attention\" on other tokens for each token.  So what does attention memorize? I think the paper should make it clearer before or after the theorectical analysis, or even verify the memorized knowledge with some visualization. \n2. Analysis and comparison with existing works are insufficient, for example,  <<Transformer Feed-Forward Layers Are Key-Value Memories>>. How are the memories different between this paper and the reference? There are also some works about transformer   interpretability and model editing. Do the observations in these related works support the findings in this paper? I think the authors should give more analysis, rather than demonstrating the their propsed theorem only.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Y3YvltGQ8X", "reviewer": "Reviewer_oLtn", "summary": "The authors prove a lower bound on the memorization capacity of a multi-head attention layer under a new set of input-data assumptions that relax General Position assumptions previously used in similar work. The main result is that memorization increases linearly with the number of heads, monotonically with the context size, and monotonically with the head dimension as long as it is smaller than the context size. The result validates some of the common practical choices of hyperparameters. The authors verify their data assumptions empirically by checking for linear independence of context vectors and approximately testing the Kruskal Rank. They also present a synthetic experiment that confirms their theoretical findings.", "questions": "How good is the approximation of the Kruskal Rank computation?\n\nHow did you come up with this particular set of data assumptions? Was this dictated by the proof technique you wanted to apply?\n\nCan your proof technique be applied to prove a similar result beyond a single attention layer or is there some fundamental obstacle?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper presents an important result on the memorization capacity of a multi-head attention layer using an original set of data assumptions that are shown to hold in practice. The result is verified through a properly designed experiment. The paper is well written and structured, and the logic of the argument is easy to follow. The figures are of high quality and they help to get the message across. While I haven't checked the proof in detail, the main steps seem to be sound.", "weaknesses": "It would be great to provide code for your synthetic experiments and input-data assumption validation to increase reproducibility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)