Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators?
Authors: Tokio Kajitsuka
Issei Sato
Venue: ICLR 2024 poster (2024)
Keywords: Transformer
Self-Attention
Memorization
Universal Approximation Theorem
Contextual Mapping
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: One-layer and single-head self-attention with low-rank weight matrices is expressive enough to be a contextual mapping.

Abstract: Existing analyses of the expressive capacity of Transformer models have required excessively deep layers for data memorization, leading to a discrepancy with the Transformers actually used in practice. 
This is primarily due to the interpretation of the softmax function as an approximation of the hardmax function.
By clarifying the connection between the softmax function and the Boltzmann operator, we prove that a single layer of self-attention with low-rank weight matrices possesses the capability to perfectly capture the context of an entire input sequence.
As a consequence, we show that one-layer and single-head Transformers have a memorization capacity for finite samples, and that Transformers consisting of one self-attention layer with two feed-forward neural networks are universal approximators for continuous functions on a compact domain.

Published as a conference paper at ICLR 2024
ARE TRANSFORMERS WITH ONE LAYER SELF- ATTENTION USING LOW-RANK WEIGHT MATRICES UNIVERSAL APPROXIMATORS?
Tokio Kajitsuka Issei Sato
Department of Computer Science
The University of Tokyo ikajitsuka-tokio
ecc u-t.k-.o......ccc.....o.........
ABSTRACT
Existing analyses of the expressive capacity of T'aansformer models have required excessively deep layers for data memorization, leading discrepancy with the Transformers actually used in practice. This is primarily due to the interpretation the softmax function as an approximation of the hardmax function. By clari- fying the connection between the softmax function and the Boltzmann operator, we prove that a single layer of self-attention with low-rank weight matrices pos- sesses the capability to perfectly capture the context of an entire input sequence. As a consequence, we show that one-layer and single-head Transformers have a memorization capacity for finite samples, and that Transformers consisting of one self-attention layer with two feed-forward neural networks are universal approxi- mators for continuous permutation equivariant functions on a compact domain.
INTRODUCTION
The Transformer model has been ubiquitously used in deep learning since its proposal by Vaswani al.) (2017). Its widespread application spans several domains, not only revolutionizing Natural Language Processing (NLP) through models like BERT (Devlin et al.) 2019 Liu et al. 2019) and GPT (Brown al., 2020 Radford et al.) a b) but also making significant advancements in image alternative to conventional models like convolutional neural networks
and graph processing
(CNNs) and graph neural networks (GNNs) (Dosovitskiy et al.) 2022; Ying et al. 2022).
One of the key reasons behind the success of the Transformer model is its ability represent wide range of functions. Various studies investigated this aspect, including the universal approximation
theorem for Transformer models and its memorization capacity Yun al. (2019); Kim et al. (202));
Mahdavi et al. 2023); Edelman et al. (2022): Gurevych et al. (2021 1 Takakura & Suzuki (0233)
Likhosherstov al. (2023).
The main challenge in proving universal approximation theorems for Transformer models lies in the fact that the Transformer needs account for the context of the entire input sequence. Unlike feed- forward neural networks where each input is processed independently, the self-attention mechanism in Transformer models must take into account the dependencies between all elements in each input sequence. In constructive proofs (Edelman et al. 2022 Yun et al. 2019 Kim et al.) 2023 Mahdavi al. 2023; Gurevych al.) 2021] Takakura Suzuki 2023), these dependencies are often aggre- gated into a token-wise quantity, which we call a "context id" here, by self-attention mechanism, and then feed-forward neural networks map each context id to the desired output.
The drawback of existing analyses is that they require excessively deep layers (Yun et al. 2019 Kim et al. 2023) or quite lot of attention heads (Gurevych et al. [022 Takakura Suzuki 2023 Likhosherstov et al., 2023) for data memorization, which leads to discrepancy with Transformers being deployed in practice. This discrepancy primarily arises from the interpretation of the softmax function approximation of the hardmax function. Consequently, compute the "context id" within the self-attention mechanism, the number of required self-attention parameters scales linearly with the length of input sequence.

---

Published as a conference paper at ICLR 2024
In this work, we address this gap by closely examining the softmax function itself. First, we show that it is impossible to output the "context id" using just one layer of self celf-attention with the hardmax function. At the same time, we demonstrate that just layer of single-head and softmax -ased self-attention with low-rank weight matrices possesses the capability to perfectly capture the context of entire input sequence. This result implies that the Transformer with one self-attention layer is a universal approximator of continuous permutation equivariant functions by using two feed- forward neural networks connected before and after the single-head self-attention mechanism with an arbitrary head size.
Our contributions are summarized follows.
1. We show that one layer self-attention with the hardmax function is not contextual map- ping; that is, one layer hardmax- -ased Transformer has no memorization capacity.
2. In contrast, we provide a framework for constructing a context mapping with one-layer and single-head self-attention using the softmax function.
3. We prove that one layer Transformer has a memorization capacity for finite samples, and Transformers with one-layer and single-head self-attention are universal approximators of continuous permutation equivariant functions.
1.1 RELATED WORKS
Universal approximation theorems. The history of the universal approximation theorem begins around 1990 (Cybenko 1989 Carroll & Dickinson 1989; Hornik 1991 Funahashi 1989). Recent studies on this topic include analyses of how network width and depth affect the expressive power (Lu et al. [017], and analyses for specific architectures (Lin & Jegelka 2018). There have also been analyses of the memorization capacity of models (Baum) 1988; Huang & Babri 1998). The main focus of the memorization capacity is mainly on the analysis of parameter efficiency for storing
finite samples (Huang 2003 Vershynin] |PrkkPPark et al. [202 /ardi et al.) [2022) Yun et al. 20188] Bubeck et al., 2020; Hardt & Ma 2016 Rajput et al. [202]] Zhang et al. 2016). Notably, Zhang
et (2016) demonstrated that neural network of the size used in practice
perfectly memorize
randomly labeled data set. Belkin et al. (2019); Nakkiran et al. (2019) pointed out that the minimum
number of parameters required memorize dataset is related to the double descent threshold.
Expressive capacity of Transformer. Ever since Vaswani et al. (017)) first proposed the Trans- former architecture, there have been various theoretical analyses its expressive capacity. Yuun et al. (2019) proved for the first time the universal approximation theorem for Transformer models, showing that a continuous function on compact domain can be approximated if the number of Transformer blocks is on the order of the power of n1, where n is the length of each input sequence. Later, Kim al. (2023) showed that 2n self-attention blocks are sufficient for the memorization of finite samples. Since the studies of Yun et al.] (2019) and Kim et al. (2023) are closely related to our paper, we discuss the details in more depth in Section 3.2 later. Their results were based on the assumption that the inputs are separated to some extent, which is an assumption we also make in this paper. Alternatively, under the assumption that input sequences are linearly independent, Mahdavi ct al. (2023) showed that one-layer /hhhad self-attention mechanism ples. Relatedly, Edelman et al. (2022) demonstrated that the bounded self-attention head is capable of expressing sparse Boolean function while obtaining an upper bound on the covering number of self-attention. Gurevych et al. (2021) analyzed the theoretical performance of Transformers as a hierarchical composition model. Later, Takakura & Suzuki (2023) extended their result by utiliz- ing a sinusoidal positional encoding and multiple heads, and showed that a one-layer Transformer with an embedding layer is universal approximator for thift--quivariant y-smooth functions, Tiang & Li (2023) recently used the Kolmogorov representation theorem to provide a non-constructive proof of the existence of two-layer Transformer that approximates an arbitrary continuous func- tion on a certain domain. There are variants of universal approximation theorems for Transformers, such as analyses of sparse Transformers (Yun et al., (020)) and constrained universal approximation theorems (Kratsios et al.) 2021). Likhosherstov et al. (2023) showed that, given parameters, there exists an input such that self-attention approximates an arbitrary sparse pattern. While Bhojana- palli et al. (2020) proved that Transformers with a small head size, which is typical for multi-head self-attention, cannot express certain positive Sllumn-stochasticc matrices, Aghajanyan et al. (202T) demonstrated empirically that pre-trained Transformers have very low intrinsic dimension, and
memorize O(Hn) sam-
2

---

Published as a conference paper at ICLR 2024
Reif et al. (2019) visualized context embeddings in BERT. Luo et al. (2022) showed the existence of functions that cannot be approximated by Transformers with relative positional encoding. There is also a series of papers analyzing Transformer's 's expressive capabilities from the perspective of formal
languages (Hahn 2020 Bhattamishra et al.) 2020 Yao et al. 2021 Hao et al. 2022 Merrill et al.) 2022 Chiang & Cholak 2022 Chiang et al. 2023) where softmax function in self-attention
mechanism is treated as an averaging or hardmax function. As for a negative result of Transformers' expressive capacity, Zaheer et al. (2020) demonstrated that there exist some problems which sparse Transformers must require a superlinear number of layers to solve.
2
PRELIMINARIES
2.1 NOTATION
We use bold lowercase letters to represent vectors and bold uppercase letters to represent matrices. For any vector E IK,, we denote by the i-the element of v. For any matrix A E Ruxb, we denote its i-th row by A..: its k-th column by A. and the element at its i-th row and k-thh column by A,k-- For any positive integer m E N4. [m]] represents the set {l...., m.). For any real numbers a < b, [a, represents the interval {u R S S b). (-00, a) represents {a R represents (I E IR hardmax function, respectively. Note that when there are multiple indices with maximum values, the hardmax function is defined such that the sum of the values these indices equals one. By abuse of notation, for any input matrix A, as [A] and Øh [A] are defined as column-wise softmax and column-we wise hardmax, respectively. We denote the ReLU activation function by dR. Unlike øs and TH. FR is always an element- wise operator, regardless of whether the input is a vector or a matrix. Let || || be the E2 norm and || |Ip (1 S p 00) be the l norm. We define the distance between two functions fi, f2 Rdxx- RRdx by
a). and (b, oo) bj. Let o (v) and H (r) for any input vector v be the softmax function and
de (2) = ( - (X) 1/P
In this paper, n denotes the length of an input sequence, the number of input sequences, C' the number of output classes, and d the embedding dimension. In addition, i, are basically used for the indices finite samples and k for the indices in each input sequence.
2.2 TRANSFORMER BLOCK
Transformer was first introduced in Vaswani al. (017)). Here follow the definitions adopted in Kim et al. ) : the Transformer block is composed of the self-attention mechanism and the feed-forward neural network, each accompanied by skip connection. Given an input sequence Z E dx composed of n tokens each with an embedding dimension of size d, dot-product self-attention mechanism with h heads outputs the following values:
F (2) = 2 Σ w (0) (W!")
2) os (w. (w/
(W! ((, 2)
(2)
where (!( w!(K) (!()) E Raxd and W! ( dxxx are the weight matrices, and s is the head size. Note that here, as with Yun et al. (2019) and Kim et al. (2023), we adopt the definition of the self-attention mechanism, which excludes layer normalization from the original definition of
Vaswani et al. 2017) for the sake of simplicity.
In contrast, given an input H Rdxn, the output of feed-forward neural network with a skip connection at index k E [[]]
((FF) (H) H., + 2 oR (L)) H + b0w] + (2) E Rd.
(3)
where q the hidden dimension, ( E 129x and (2)) kdxa are weight matrices, and ((1) IR and (2) are bias terms.
On the basis the above definition, the Transformer block represented as a composition of a self-attention mechanism and feed-forward neural network: for any input sequence Z E Rdxn.
3

---

Published as a conference paper at ICLR 2024
composed of tokens each with embedding dimension of size d, the Transformer block F Rdx" + Rax= outputs
= = F(PFF ($(A)) (2)
(4)
From the above definition, we see that the interaction of each token occurs only in the self-attention mechanism.
ATTENTION IS A CONTEXTUAL MAPPING
3.1 PROBLEM SETTING
Let (X(1), Y(")...., ( C Rdxn x [C]axn be an N input-output pairs of sequences, each of which consists of sequence ( of n tokens with embedding dimension d, and an output Y where r! corresponds to the label of the token x (8) at index k. In addition, we define the
i-th vocabulary set for E [[V] by Dlo) Ukein X k,k 12d, and the whole vocabulary set V is defined by V ULE[N] 20() Rd.
3.2 BACKGROUND
Yun et al. 2019) proved affirmatively of the most fundamental questions the expressive capacity of Transformer models, namely, whether the universal approximation theorem for Trans- former models holds. Their proof approach is to quantize the input domain and reduce the universal approximation theorem to the memorization analysis of finite samples, i.e., the construction of model that achieves zero loss for finite number of training data, which was also analyzed later by Kim et al. (2023). In the analysis of memorization capacity, assumptions are usually made on the inputs order to perform a meaningful analysis beyond the lower bound of [Sontag (1997). Here, as with the assumptions adopted by Yun et al. (019);; Kim et al. (2023), we assume that the input tokens are separated by a certain distance.
Definition 1 (Tokenwise Separatedness). Let N and Z ..... Z(N) R be input se- are called tokenwise (Tmin m'ax, 8) -eeparate if the following three
quences. Then, (1) ..... conditions hold.
Z(N) ().....,
1. For any E [N] and & E [m].
(()
> Tiin holds.
2. For any E [N] and [n]. 2 r mmax holds.
3. For any i,j E [N] and kt I E [n] with 2
2
å holds.
Note that we refer ( satisfy conditions 2 and 3.
(N) tokenwise (Fmax, e)-spparated instead if the sequences
The achievement of Yun al. (019) was only to prove the universal approximation theorem for Transformers, but also to clarify the difficulties in the analysis of this kind of expressive capacity of Transformers and elucidated an approach to establishing the proof. Namely, what makes Transform- ers' memorization different from that of feed-forward neural networks is that Transformers need to capture the context of each input sequence as whole, rather than simply associating each token with label.
Remarkably, Yun et al. (2019); Kim al. 2023) formulated this concept as contextual mapping, which assigns a unique id to pair of an input sequence and each of their tokens. We define it here using the notion of (r, 5) sepparatednsss.
Definition (Contextual Mapping). Let ( ..... (N) Rd be input sequences. Then, map q Rdxn + ddxn is called an (r, (r,5)-com)) -contexxtllll mapping if the following two conditions hold:
1. For any E [N] and & E [m]. ||o (x(())... < holds.
4

---

Published as a conference paper at ICLR 2024
2. For any i, [N] and k,l [n]
such that 200 DG)
x + x(
(x q (x)) ) > holds.
In particular, q( )) for i [M] is called context id of (1)
we have such contextual mapping, label sequence can be associated with unique id for each input sequence using the existing analysis of memorization in feed-forward neural networks.
Thus, the central question is: how to construct contextual mapping in Transformer models? The only place in Transformer models where interaction between tokens can be taken into account is the self-attention mechanism; therefore, the self-attention mechanism must be used to construct contextual mapping, Yun et al. (2019) first constructed contextual mapping by using [V]d self-attention layers and later Kim et al. (2023) improved it to 2n self-attention layers. However, this is still far from the practical implementation of Transformers, and it remains unclear whether reasonably-sized Transformer would possess such memorization capacity or if the universal approx- imation theorem would hold. This leads to the following question.
How many self-attention layers are both necessary and sufficient to construct a contextual mapping?
We first point out the reason for requiring a significant number of self-attention layers the con- struction of contextual mapping in the analyses of Yun et al.] (20T9); Kim al. (2023). Their approach entails interpreting the softmax function the self-attention mechanism as an approxi- mation of the hardmax function, which also hinders a detailed analysis of the specific properties of the softmax function. As evidence of this, we illustrate in Section 3.3 that using a single layer of self-attention with the hardmax function does not suffice construct contextual mapping.
Next, in Section 3.4 we demonstrate that contextual mapping can be constructed by using only one self-attention layer with the softmax function. This is somewhat surprising because this implies the probability of fully capturing the context of each input sequence only through the attention coefficients computed by the pairwise dot-product of the softmax function and its weighted average.
3.3 SELF-ATTENTION WITH HARDMAX
In previous studies analyzing the memorization capacity of Transformers Yun et al. 2019 Kim ct al. 2023). the softmax function is taken to be an approximation of the hardmax function. However, show here that the attention block with the hardmax function is not contextual mapping.
First we define the attention block with the hardmax function: for an input sequence Z E RR the attention with the hardmax function is calculated
F F(SA)(
(Z) Z
W
(w!ZZ)
W
2) oo (w; (
(5)
where ((()) (K)
w!Q) ( E xxd and W(()) Bdxx are the weight matrices
The following theorem holds for such model. The proof is in Appendix A.1
Theorem 1. 1-layer multi-head self-attention ((SA) with the hardmax function cannot be contex- tual mapping.
Since the self-attention mechanism is the only place in Transformer models where interaction be- tween tokens happens, this theorem indicates that one-layer Transformers with hardmax attention do not have a memorization capacity.
3.4 SELF-ATTENTION WITH SOFTMAX
In this subsection, show that softmax-based 1-layer attention block with low-rank weight ma- trices a contextual mapping for almost all input sequences. This result is consistent with recent
'To be precise, when the continuous input range is quantized into 1/6 pieces for some < 8 < 1, they demonstrated that there exists contextual mapping composed of 8 self-attention layers.
5

---

Published as a conference paper at ICLR 2024
empirical evidence that pre-traines Transformers are low-rank (Aghajanyan et al. 2021 Choroman- ski et al. 2020: Wang et al. 2020 Lialin et al. 2023). and theoretically supports that the low-rank self-attention mechanism is sufficient to fully comprehend the contextual information of an imput sequence. It is worth noting that our construction allows for an arbitrary head size. By considering the case of head size of 1, this particularly indicates that the self-attention mechanism has the ability to compress the information of an input sequence through scalar value.
Theorem 2. Let X X E IRdx be input sequences with no duplicate word token in each sequence, that is,
X X
(6)
for any E [N] and k. E [n]. Also assume that X X()..... "I,,..,,,,,,, X separated. Then, there exist weight matrices W E Idx* and W W(k), , E Rxxd such that the ranks of W W(K) and W are all 1, and 1-layer single head attention with softmax, i... F. with h = 1 is an (r, (contecantlll mapping for the input sequences X(1) pdx= with r and 6 defined by
are tokenwise (rmins Tmax c)-
X (N) E
c 4"
r = F Jmax +
(7)
2(log (2)22r 2rinn
exp (- ([V) + 1)a² (2 (2log n + +) dr max
(8)
max (1V) + (2 log 3)rd
4ermin
Here we provide a simple proof sketch. The full proof can be found Appendix A.2
Proof Overview. For simplicity, we here assume corresponding to each input sequence (1) for i [N], context id can be constructed from suit- able linear combination of the sequence id and the value of cach token. Since this linear combination can be calculated by the output projection matrix (O)) and skip connection, the problem is how to configure weight parameters W(V), W(K), ((Q) E Rixd so that each row of the values' softmax weighted average,
we have a unique id, i.e., sequence id,
(W x os (wrl)xx1)) (wrl)x((i)
E IR 1x
outputs the unique sequence id (1).
Actually, even weaker condition is sufficient for attention block to be contextual mapping: there is no need to have just one unique sequence id for each input sequence. fact, it is possible to construct contextual mapping, provided that for each token E V. input sequences in which the token appears can be identified by some v-dependent sequence ids. This condition can be ex- pressed in mathematical form as follows: what have show is to construct weight matrices (V), (K)), (
IR with some >> 0 such that
(w x(() as (wick) T (w(xxx
-(w(X as (WiRix)N)) (
>
(10)
holds for any distinct i.j [N] and any k E [n] such that x!0 x! and v(a) * vg)
For simplicity, we choose (VV) (kk @ wn such that the linear operator w Rd projects each token to a scalar value while approximately preserving the distance between each pair of tokens: for any pair of tokens U, 17, V.
cllu, - 2, < www T U. w Un S l| --al
(11)
In our actual proof, there exist vectors v, o' e a2d such that (V),, ( (()) may defined by (()) u²v" (KK). u'v" and (((@) un T for arbitrary vectors u, , u² E B satisfying certain constraints.

---

Published as a conference paper at ICLR 2024
Then, by using the assumption t x x for some token
with some constant 0 < < t E Rdd, we have
w T t| .
w T (()
w T x(10) T
(w Tt)
- -(w0 (w0] vo['()) (w x(xx) T (w Tu)
08
(a())
T
og a
61)
(ab))²
0g [a6i1
(12)
- (w t) R" There- fore, in order to prove that a self-attention block serves as a contextual mapping, we only have to focus on the separability of the function
where we denote (f)
(w X (1)
(w t) E R² and a
boltz R + R. a a [s]a].
(13)
which is known as the Boltzmann operator (Littman
1996
Asadi & Littman 2017).
The following lemma shows that the Boltzmann operator a mapping that projects input sequences
to scalar values while preserving some distance, and is central to our proof that the self-attention function is a contextual mapping.
Lemma 1. Let a ..... a (m) E R" be tokenwise ( separated vectors with no duplicate element in each vector and
8 >> log +
(14)
Then, the outputs of the Boltzmann operator are (r, "'''--eppaated that is,
boltz(a (0) r
(15)
(oltz(a boltz(a (1) >0'= =(log e-
(16)
hold for each i,j [m] with a + a!))
Taking into account the above arguments, this separability of the Boltzmann operator allows us to construct one self-attentio layer to be contextual mapping.
Remark 1 (Masked self-attention). In practice, attention matrices are often masked to avoid directing
attention to undesired tokens. This performed, for example, for autoregressive text generation or padding of inputs with different lengths. It is relatively straightforward to extend Theorem 2 to masked self-attention mechanisms. See Appendix [] more details.
APPLICATIONS OF CONTEXTUAL MAPPING
4.1 MEMORIZATION CAPACITY OF ONE-LAYER TRANSFORMER
As first application of Theorem we prove that 1-layer Transformer can completely memorize finite samples, each of which has no duplicate token. This result emphasizes that in contrast the proof of Kim et al. (2023), which requires 2n self-attention layers for Transformer memorization, actually sufficient. In addition, it is worth noting that the hardmax- based Transformers do not have memorization capacity, which is implied straightforwardly from Theorem I
one layer of self-attention
Corollary (Memorization capacity of one-layer Transformer). Let E > 0,Tuax > Tmin 0 and (1) y (X(N), y()) C Rdxn x [C]dxn be sequences of input-oulputppaiss such that x ..... xN) are tokenwise (Tmin, Tmax, seeparraee input sequences with no duplicate token in
each sentence and consistently labeled, that is, y( k.k =Y(() holds for any i, E [N] and k; I E [n]] such that ((() vii and X k
x()) .11
Then, there exist 4(s + d) + d(2nN + d) weight parameters such that for any E [N)
JF (x()) F(Fr) (I!5 (x())) (
(17)
holds.

---

Published as a conference paper at ICLR 2024
Remark 2 (Parameter efficiency). To achieve the memorization with one-layer Transformer, the one-hidden-layer feed-forward block has to map each context id to the corresponding label. Since the possible number of context ids is at most n.N in the worst case, the linear dependency n.V of the number of parameters in Corollary | is optimal up to logarithmic factors (Bartlett et al. 2019).
worth mentioning that this linear dependency can be relaxed to the milder requirement Ô (vnV
by allowing for deeper layers in the feed-forward block (Vardi et al. 2022 under the assumption that the size [V] of the vocabulary set is independent of n and N.
In addition, it is straightforward show that 1-layer Transformer with trainable positional encod- ings has memorization capacity for arbitrary input sequences possibly with duplicate tokens.
Corollary (Memorization capacity of one-layer Transformer with positional encodings). Let > )....
O, mmax > Tuin > 0 and (X ( (X(N), nutput-pairs such that X X(),.... ..... X (N) are tokenwise (Tmin; m'max e)-spaaaateed input sequences and are consistently labeled, that is, Y y holds for any i,j [N] such that X x(O)
(N)))) R x [C]dxm be sequences of input-
Then, there exist 4(s + d) + d(nnN + d) weight parameters and positional encodings E E Rddx such that for any E [N],
F x F + E F(FF) (F! (x ((0++E)) = =Y(1)
(18)
holds.
4.2 TRANSFORMERS WITH ONE SELF-ATTENTION LAYER ARE UNIVERSAL APPROXIMATORS
As a further application of Theorem 2 we here provide proof that Transformer with one self- attention layer is a universal approximator. More precisely, let FPE be the set of all permutation cquivariant continuous functions that take values on a compact domain in ædxn, and let T2 be the set of all two layer Transformers with one-layer and single-head self-attention, that is,
To = ( ((SA) o F((F) Raxd + Raxd]
(19)
where F(FF), F (FF) and ((SA) are feed-forward neural network layers and a single-head self- attention layer with the softmax function, respectively. Then the following proposition holds (see the definition (.)
Proposition 1 (Transformers with one layer self-attention are universal approximators), Let 1 S p 30. Then, for any f E FPF and e >> there exists a Transformer g E To with one-layer and single-head self-attention such that (p(J,g) <. holds.
To the best of our knowledge, this is the first universal approximation theorem for two-layer Trans- formers with a self-attention of realistic size. Takakura & Suzuki (2023) showed that a one-layer Transformer with an embedding layer is capable of approximating shift-equivariant 7-smooth func- tions. However, their construction requires a high number of self-attention heads and large head size to flatten an input sequence into outputs of self-attention. Diang & Li (2023) used the Kol- mogorov representation theorem to give a non-constructive proof of the universal approximation theorem of two-layer Transformers with positional encoding for continuous functions on a compact domain. Acknowledging the contributions of prior studies, we would like to emphasize the novelty of our results again, that is, Transformers using a single-head self-attention are universal approx- imators for continuous permutation equivariant functions to Theorem 2 Our result can be readily extended for continuous but not necessarily permutation equivariant functions significant from the perspective of geometric deep learning.
arbitrary compact domain, thanks
compact domain by using positional encoding, and at the same time is
5 EXPERIMENTS
As shown in Theorem 2 self-attention mechanism with rank 1 weight matrices already has enough expressive capacity to become a contextual mapping. In particular, its proof leads us to consider the following simplified form of self-attention mechanism: for any input sequence Z E gdxn
F R1) (Z) Z + (
(w 1)0as as (0 2)
(v. z)
Rdxn
(20)
8

---

Published as a conference paper at ICLR 2024
where U1, 222 Rd and W(O) E Rd weight matrices. This architecture corresponds to a common self-attention with the head size s = 1, and value and query matrices having the same weight vector v, . In this section, test whether Transformers with self-attention layers replaced by equation 20 which we call rank-1 Transformers, actually have the theoretically predicted expressive capacity by using real-world dataset.
We train rank-1 Transformers on a token classification task with the CoNLL-20033 (Toog Kim Sang De Meulder 2003) dataset. The batch size is 32 and the training are conducted over 400 epochs.
We train three different depths rank-1 transformers on the dataset and do not use layer normaliza- tion match the situation with theoretical analysis.
1.00
CoNLL-200 L-203 Accuracy
0.98
0.96
1-layer
3-layer
0.90
6-layer
0.88
Training Epochs
(a) CoNLL-2003
Figure Training accuracy of rank-1 Transformers for the CoNLL-2205 dataset. 1-layer (solid line), 3-layer (dashed line) and 6-layer (dotted line) rank-1 Transformers are trained over 400 epochs (X-axis). For the 1-layer rank-1 Transformer, we observed that it reached an accuracy of 0.9872 at 800 epochs after further training.
Figure shows training accuracies of 1-layer, 3-layer and 6-layer rank-1 Transformers on each task over 400 epochs. It can be seen that the 1-layer rank-1 Transformer is already able to memorise the CoNLL-2003 dataset almost perfectly. On the other hand, while the accuracy curve for the 1- layer rank-1 Transformer shows that the accuracy is still increasing steadily at 400 epochs, reaching 0.9872 at 800 epochs, its rate of increase is much slower than for the 3-layer and 6-layer Transform-
From this observation, we conjecture that while theoretically 1-layer Transformers already have a memorisation capacity for finite samples, the advantage of deepening layers lies in speeding up the learning of such tasks. Since our analysis is on the expressive capabilities of Transformers, we leave this hypothesis on the optimisation aspect of Transformers as future work.
CONCLUSIONS
We demonstrated that a contextual mapping can be implemented in one-layer and single-head self- attention with low-rank matrices, by clarifying the connection between a self-attention mechanism and the Boltzmann operator. This particularly indicates that one-layer Transformers have a mem- orization capacity for finite samples, and that Transformers with one-layer and single-head self- attention are universal approximators for continuous permutation equivariant functions on a com- pact domain. Our proof of the universal approximation theorem requires one feed-forward neural network layer before the self-attention layer to quantize continuous inputs. We leave it as future work to clarify whether the one-layer Transformers without such quantization layer are universal approximators or not. We also expect that our analysis of the softmax function will have an impact on the evaluation of Transformer's expressive capability from the perspective of formal languages.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported by JSPS KAKENHI Grant Number 20H04239 Japan. We would like to thank all the collaborators and anonymous reviewers for constructive discussions.
REFERENCES
Armen Aghajanyan, Sonal Gupta, and Luke Zettlemoyer. Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine- Tuning In Proceedings of the 59th Anmual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 7319-7328, Online, August 2021. Association for Computational Linguistics. doi: 10. 18653/811/22211111 acl-long.568. URL https:
//aclanthology wttg//2g11....o.oooo. acl....ccl--m. acl-long.
Kavosh Asadi and Michael I., Littman. An Alternative Softmax Operator for Reinforcement Learn- ing. In Proceedings of the 34th International Conference on Machine Learning, pp. 243-252 PMLR, July 2017. URL ISSN: 2640-3498.
mlr, wtt..m...p..////////aaa......... html
Peter L. Bartlett, Nick Harvey, Christopher Liaw, and Abbas Mehrabian. Nearly-tight vc-dimension and pseudodimension bounds for piecewise linear neural networks. Journal of Machine Learning
Research, 20(63):1-17, 2019. htRl:tttt://mmmrrorr///eesee
htm.
Eric B Baum. On the capabilities of multilayer perceptrons. Journal of Complexity, 4(3):193-
215, September 1988. ISSN 0885-064X. doi:
URL https:
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine-
learning practice and the classical varinvari trade-off. Proceedings of the National Academy
of Sciences, 1 116(32):15849 15854, 2019. doi: 10.1073/pnas. 19000/000..99000166011 URL https:
// www. pnas wttp:pa....ooooo//////ooooo wtt.aa/.............aa.a. 1073//011/ppsp.. 1e0.0993000000
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the Ability and Limitations of Trans- formers to Recognize Formal Languages. In Proceedings of the 2020 Conference on Empiri- cal Methods in Natural Language Processing (EMNLP), pp. 7096-7116, Online, January 2020. Association for Computational Linguistics. doi: 10.18653/v 10t3l////220.0 /20000eemnlpppp-e/e. http220.emmlpmmain.. main.576. URL
ltttss:/aclanthhllla www.:yy.ooooooooo wttp///0.c.....oooo www.eennnalaeeeeeeeee man.mmain.555. 576
Srinadh Bhojanapalli, Chulhee Yun, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Low-Rank Bottleneck in Multi-head Attention Models. In Proceeuliggs of the 37th Interna- tional Conference on Machine Learning, pp. 864-873, PMLR, November 2020. URL hetps: ISSN: 2640-3498.
//proceedings.o .mw..nnn..snrr.mlmmm....rr....m. mlr. srress/vl19/bbojaree wwtp:aaaiiiaaa.hc.emaa
Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, 2004. URL httpt:tp/////.... htt:://wbbsaaffd..eeeemm -byd///by/////// cvxbook http://wwvxxxxoocodboooooo
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 McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learn- ers. In Advances in Neural Information Processing Systems, volume 33, pp. 1807-1901. Curran Associates, Inc., 2020. URL http::tttts:prrrreeeeeee www.:niipsccsee.ee..eeeeeee 1457c006bffc4447771188 (ww.777488888ac1ccc. 4eeee1224442 Abstract. Abstrrtttt httpppp.ccppaeer/22///aaaaa
Sebastien Bubeck, Ronen Eldan, Yin Tat Lee, and Dan Mikulincer. Network size and size of the weights in memorization with two-layers neural networks. In Advances in Neu- ral Information Processing Systems, volume 33, pp. 4977 Curran Associates, Inc., 2020. URL https: htt.pppsss/pppeeeng...erio.oooo. hash hhtr////000oaaaaa444444742777070
10

---

Published as a conference paper at ICLR 2024
Carroll and Dickinson. Construction of neural nets using the radon transform. In International 1989 Joint Conference on Neural Networks, pp. 607-611 vol.1, 1989. doi: 10.1 N.1109/UCCN..199. 118639.
David Chiang and Peter Cholak. Overcoming a Theoretical Limitation of Self-Attention. In Pro- ceedings of the 60th Annual Meeting of the Association for Computatioonl Linguistics (Volume I: Long Papers), pp. 7654-7664, Dublin, Ireland, May 2022. Association for Computational Lin- guistics. doi: 10.18653/ 12188663/11202222lllloooggg
acl acl-long. 527
David Chiang, Peter Cholak, and Anand Pillay. Tighter Bounds on the Expressivity of Transformer Encoders, May 2023. URL [es].
10743 1rXiv:2301.10743
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy Colwell, and Adrian Weller. Rethinking Attention with Per- formers. October 2020. |tttts::
G. Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of Con- trol, Signals and Systems. 244)3303-314, December 1989. ISSN 1435-568X. doi: 10.1007/ BF02551274. htt.tt.tss:////..rgg//// 1ttp://11111000///1111110 100.00002///F/2230..0 http:///oo...........0000///5522
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, May 2019. URL http://arxiv.
04805 arXiv: arXiv:1810.04805 10,04805 [cs].
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition Scale. March 2022.
Benjamin L. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive Biases and Variable Creation in Self-Attention Mechanisms. In Proceedings of the 39th International Conference on Machine Learning, 5793 5733-5831. PMLR, June 2022. URL httpps::/ppoceeee.onotoooooo .mlr. ISSN: 2640-3498.
pxess/v162/edelnnneee html
Ken-Ichi Funahashi. On the approximate realization of continuous mappings by neural net- works. Neural Networks, 183):1831122 1989. ISSN 0893-6080. doi: https://doi.org!
10.101660892-600008999900033 atticlee/LL1009986800000009
www www.sciencceiirretccoooooooooe
Iryna Gurevych, Michael Kohler, and Gözde Gül Sahin. On the rate of convergence of a classifier based on a Transformer encoder, November 2021. URL http:///arxiv..rr//bb////
14574 arXiv:2111.1455 [cs, math, stat].
Michael Hahn. Theoretical Limitations of Self-Attention in Neural Sequence Models. Transactions of the Association for Computational Linguisties, 171, January 2020. ISSN 2307-387X. doi: 10. 166/tacl 00306. URL htttl::tt:///// www.ddoo.oggg/0.. 1162/ _00360
Yiding Hao, Dana Angluin, and Robert Frank. Formal Language Recognition by Hard Attention Transformers: Perspectives from Circuit Complexity. Transactions of the Association for Com- putational Linguistics, 10:800 810, July 2022. ISSN 2307-387X. doi: 10.1 10.166/acc
URL https: /wii///oiiiia hwt.:oo..ogo.ggooooooo /0. 100.0.11221...0 taal a 00490
Moritz Hardt and Tengyu Ma. Identity Matters in Deep Learning. November 2016. URL https:
/ //openrevien ew. www.teettoonmn www.mmmwmdddxxxxxxxxxxxx.
Kurt Homik. Approximation capabilities of multilayer feedforward networks. Neural Ner- works, 422):551-257 1991. ISSN 0893-6080. doi: http:///teaangg/ll0oo00000000000
90009-T. T.
08 0893608991 90009T

---

Published as a conference paper at ICLR 2024
Guang-Bin Huang. Learning capability and storage capacity of two-hidden feedforward net- works. IEEE Transactions on Neural Networks, 14(2):274-28) 28,,2 2003. doi: 10.1109/TNN.003. 809401.
Guang-Bin Huang and H.A. Babri. Upper bounds on the number of hidden neurons feedforward networks with arbitrary bounded nonlinear activation functions. IEEE Transactions Neural Networks, 91):2224-22) 1998. doi: 10.1109/72.635055.
Haotian Jiang and Qianxiao Li. Approximation theory of transformer networks for sequence mod- eling, May 2023. URL
18475 arXiv:2305. [es].
Junghwan Kim, Michelle Kim, and Barzan Mozafari. Provable Memorization Capacity of Trans- formers. February 2023. URL http:: httpp//ppppeeieeeeeeeeeeee hwwp:/rwwwwwweeeeeeeeeeeeeeeeeee net/ forum 1or
Anastasis Kratsios, Behnoosh Zamanlooy, Tianlin Liu, and Ivan Dokmanić. Universal Approx- imation Under Constraints is Possible with Transformers. October 2021. URL hetps:
//openzeview.
Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky. Stack More Lay- ers Differently: High High-Rank Training Through Low-Rank Updates, July 2023. URL http:
//axxiv.orq////oo...v.aooo. org/ attaaabb//3/3./ 2300 05695 arXiv:2307.05699 [cs].
Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the Expressive Flexibility of Self-Attention Matrices. Proceedings of the AAAI Conference Artificial Intelligence, 37 (7):8773-8781, June 2023. ISSN 2374-3468. doi: 10. 20.1609/aaaivv3777 wtw....665..55553 URL https:
7/ojs. aaai. org/ index. php/ http://AAA/aaat.eeeeeeeee/////e/
Number: 7.
Hongzhou Lin and Stefanie Jegelka. ResNet with one-neuron hidden layers is Universal Approximator, In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL paper/2018/hast httaaa/aeaaaaaa3ee66666622222000
files/ .ttml
Michael Lederman Littman. Algorithms for Sequential Decision-Making. PhD thesis, USA, 1996. AA19709069.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A Robustly Optimized BERT Pretraining Approach, July 2019. URL http::/ org wwtpaawwwwwoo.///////..1.1../
16X1 [es].
Zhou Lu, Hongming Pu, Feicheng Wang, Zhiqiang Hu, and Liwei Wang. The Expres- sive Power of Neural Networks: A View from the Width. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL
https::/prroceedings..eee newrppss..........o............. wtcppaper_f.lessppppeeeeee html
Shengjie Luo, Shanda Li, Shuxin Zheng, Tie-Ya Yan Liu, Liwei Wang, and Di He. Your Transformer May Not be as Powerful as You Expect. October 2022. URL htttss://openreiiiiiiii htttarrrvieeeeeeeeeeeeeeee
Forum? 1000mm2NoONN
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization Capacity of Multi- Head Attention in Transformers, June 2023. hRRLlttppp//arrx.oorr//////236 arXiv:2306.02010 [cs].
William Merrill, Ashish Sabharwal, and Noah A. Smith. Saturated Transformers are Constant-Depth Threshold Circuits. Transactions of the Association for Computational Linguistics, 10:843 856, August 2022. ISSN 2307-387X. doi: 10.1 162/talll a 00493. URL httptttst////....oo
11 1162/tacl. 2t00\tacl___00933
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and llya Sutskever. Deep Double Descent: Where Bigger Models and More Data Hurt. September 2019. URL https:
//openrevien ew. net/ffrum d=B 1= g5sA4twr
12

---

Published as a conference paper at ICLR 2024
Sejun Park, Jaeho Lee, Chulhee Yun, and Jinwoo Shin. Provable Memorization via Deep Neural Networks using Sub-linear Parameters. In Proceedings of Thirty Fourth Conference on Learning Theory, pp. 3627-3661. PMLR, July 2021. URL htttp::/ppoceessee pres.
mmlr
v134/park21 html
ISSN: 2640-3498.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving Language Under- standing by Generative Pre-Training. a.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language Models are Unsupervised Multitask Learners, b.
Shashank Rajput, Kartik Sreenivasan, Dimitris Papailiopoulos, and Amin Karbasi. An Ex- ponential Improvement on the Memorization Capacity of Deep Threshold Networks. In Advances in Neural Information Processing Systems, volume 34, pp. 12674-12685. Curran Associates, Inc., 2021. URL paper/2021/hant /ttp://eeff.a2222c22222222222aaa Abttoe233bbbttttt.ta
www.:innnnn.euuuu.ooer...... nwwrips ww/ppperppppp...eeeepee files/
Emily Reif, Ann Yuan, Martin Wattenberg, Fernanda B Viegas, Andy Coenen, Adam Pearce, and Been Kim. Visualizing and Measuring the Geometry of BERT. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL httppp///paperrs..pp..ccc/ooea filess 159c1ffe5be1e4LbBB 1SS1ffebcsa1scdd44a11566------a2
pappr/2019/haa// /ttpapaee_..lllaaapar///////////
Eduardo D. Sontag. Shattering All Sets of 'k' Points in "General Position" Requires (k 1)/2 Parameters. Neural Computation, 92):3337-348, February 1997. ISSN 0899-7667, 1530-888X.
doi: 10.1162/neco. (ttp2//ecco.999...222
URL httpt:::///// //di:::///rrraa mit wtt.:mm..tt......eemm wwtpe////..//atice/e/////////
2/337-348/605 6035
Shokichi Takakura and Taiji Suzuki. Approximation and Estimation Ability of Transformers for Functions with Infinite Dimensional Input. In Proceedings of the 40th International Conference Machine Learning, pp. 33416-33447. PMLR, July 2023. URL
mll www.:ssmm....prrmmmmm wttps////2ttttaaauuuaaaa
html ISSN: 2640-3498.
Erik F. Tjong Kim Sang and Fien De Meulder. Introduction to the CoNLL-2003 Shared Task: Language-Independen Named Entity Recognition. In Proceedings of the Seventh Conference
on Natural Language Learning IILT-NACL 2003, pp. 142-147, 2003. URL https://
ac aclanthology org/W03-0411
Gal Vardi, Gilad Yehudai, and Ohad Shamir. On the Optimal Memorization Power of ReLU Neural http://oppeneeeeeeeeeeeeeerrrree
Networks. January 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances Neural Information Processing Systems, volume 30. Curran Asso- ciates,
Inc., 2017. 3 aat.5eeedddeeeeoaoaaoaaaaaaaaaa. Abst html
URL https://papers: .ips. cc/ httpp/pppp////aaaaaeeee
Roman Vershynin. Memory capacity of neural networks with threshold and rectified linear unit activations. SIAM Journal on Mathematics of Data Science, 204); 2020. doi: 10. 1137/20M1314884, URL https::/t/o/oio.. org///............ 11t..0.1111//21111111
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self- Attention with Linear Complexity, June 2020. URL httt:::/axxiv..or/bbbbtt. 04768 arXiv:2006.04766 [cs, stat].
Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-Attention Net- works Can Process Bounded Hierarchical Languages. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing (Volume 1: Long Papers), pp. 3770-3785, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v 1/2021.acll...c. long. URL
https: / antho wowphhoooy...ooo ogy. wgw.oog.o....... acl looll.loon..... 292
13

---

Published as a conference paper at ICLR 2024
Chengxuan Ying, Tianle Cai. Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do Transformrrs Really Perform Badly for Graph Representation? January 2022. URL htt.a//ss..eeeaaeaeeaeeee/a//oea
htttps::/cpenreeieettt
net /
Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Small ReLU networks powerful memorizers: tight analysis of memorization capacity. In Advances in Neural Information Processing Systems. volume 32. Curran Associates, Inc., 2018. URL https:://proceedd.ngggg////////s
neurii neurips cc/ h22p://24A444/aaa222222211233312
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are Transformers universal approximators of sequeecceto-sequeeee functions? December 2019, URL
https:://ppereviiee wttp:/wwweeetforeeeeeeeeee htt.:eeeeee/ooooo//ddd.yt.tt////
Chulhee Yun, Yin- Ven Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and San- jiv Kumar. O(n) Connections are Expressive Enough: Universal Approximability of Sparse Trans- formers, In Advances in Neural Information Processing Systems, volume 33, pp. 13783-13794, Curran Associates, Inc., 2020. URL hatpps://proeeednngee neurips...e..pp..p...... wttppppecccpppp/////ppp
2020 htth://cc22558339999444422222333 http:///oceeecnnnn......ioi..app
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Al- berti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big Bird: Transformers for Longer Sequences. In Advances in Neu- ral Information Processing Systems, volume 33, pp. 17283-17297 Curran Asso-
ciates. Inc.. 2020. URL lattss:/papers..ot nips.. httpaas..c/papprr/2////sss//// c8512d142a2d842
1aa9a7736lab3 Abbrts--bstrtttt html
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Under- standing deep learning requires rethinking generalization. November 2016. URL https:
7/openrevient net/forumm.dd-yyygerr
NOTATION TABLE
Numbers and Arrays
scalar
A vector
A
A matrix
n
The length of an input sequence
N
The number of input sequences
C
The number of output classes Embedding dimension
d
x
i-th input sequence, consisting of tokens of embedding dimension
Sets
R
Set of real numbers
N,
Set positive integers
[m]
Set all integers from to m
[a, b]
Closed interval from b
200
vocabulary set
Indexing
Element of vector with indexing starting at
Ai,j
Element i.j of matrix A
A.
Column é of matrix
A.:
Row of matrix A
Functions
14

---

Published as a conference paper at ICLR 2024
|||||
l norm of æ l norm of æ
|olla
condition if the condition true, otherwise
f2)
- [x [x(
as
Softmax function
Hardmax function
on
ReLU activation function
FH
F((4)
Hardmax-based self-attention mechanism with a skip-connection
F(
Softmax-based self-attention mechanism with a skip-connection
((FF)
Feed-forward neural network with a skip-connection
boltz
Boltzmann opeartor
A PROOF OF MAIN RESULTS
First, we introduce the Boltzmann operator, which frequently appears in proofs.
Definition 3. (Boltzmann operator) The Boltzmann operator is defined by
boltz R R, a T as|a].
(21)
By abuse notation, we use the same notation boltz for various dimension m N4.
A.1 PROOF OF THEOREM
Proof. Let Rd be an arbitrary nonzero vector, and consider the situation that all input tokens can be written as
V = (a1 a22,, (130,040) C Ied
(22)
for some scalars a 02 head i can be expressed as
< 03 a4. Then, the attention matrix inside the hardmax function
(w(K)ua) va
(w!o)ol va
(w)(k)) T
(w!() a
(23)
with input coefficients a {ar, 02, 03,, ) C IR. In particular, when we focus on a certain index, at which the token is, c.g., azu, the above expression can further be written
w K)) va 1)
( (
a20)
(w!kko) (w!
02
(w,(k) w) (w(())))
02 a
(24)
ER
The right-hand side is vector multiplied by some scalar. So it is evident that the maximum value of the vector on the right-hand side achieved only at the indices where the values of the input sequence va.7 are on a4. This implies that self-attention with the hardmax function invariably gets distracted by the indices where (a) or 04 are present, thereby overlooking information from other tokens in the input sequence. As result, no matter how many heads there are, one-layer self-attention with the hardmax function cannot distinguish input sequences, e.g., (a) agU aqv) and (a1e, 0,30, a.e).
A.2 PROOF OF THBOREM 2
Proof of Theorem 2 Recall that softmax-bas self-attention function ((SA). IRdx - Rdxx with h = defined as
(Z) Z (()) (w(V))) as (w(K))) (w(Q)))
F( (SA)
(25)
15

---

Published as a conference paper at ICLR 2024
where ( Rax and W w(k), w( Rexd are weight matrices.
We construct a softmax- based self-attention function F F(5A) with the property that
wio) (w(x) os (I (wle) x x
< 6
(26)
holds for any input sequence ( with i [N] and index & E [n]. When this property fulfilled, it is easy to show that
( w (WIM)(()) os (WWIK)x(0)) (w (0)
((S)) (x())
X 11)
S
+
max
+
4
(27)
holds for any i [N] and k E [r]], and also
F3 ((SA)
(x()))
- -F( (x()
(0o) ( (mrro) (wixxx!))
2 x i
x
(10
W(o) (wiXx(l)) 0s (IIC) (I(isi
>
(28)
4 4 2
for any i,j E [N] and k E [n] such that x! * x.P So all that remains to prove is to construct a self-attention function F(SA) that has the properties described above and can also distinguish input tokens x(1 = x such that v(t) * vt))
Let 5 = 2 2og n + 3 and fix any vectors u, u'
E R with
Rd 8 8 el T'min
u T u'l (VV) + 1)a
(29)
Then, according to Lemma such that
3 with d = 2log + 3, we see that there exists unit vector v E Rd
(m (wll)os))
(wi()ees) ( (wil)w))) >0,
(30)
1
Td well
S |00 T vel S vell
(31)
(IV) + 1)²
for any Va. Ub E V with va * vg. where W(K) = To E axd and ((Q) = u'vT E Rax
Furthermore, we configure (()) E Rddx and W(V) E Rexd to be ((V = u²u" T for any nonzero vector u" E IR. such that
w (0
(32)
4
maxx
holds. This can be accomplished, e.g., ( u" "T for any vector n" Bd which satisfies || </(4rmxx I=')²)) In this case, the value of the self-attention without skip-connection is
16

---

Published as a conference paper at ICLR 2024
upper-bounded by
w
(
(m)l T (
Σ w(o) (wr()
with up = os (wrlr)x(1)) (wrlo)
(())) (wir) x
W (V((x())
(from
Σ
skr
1)
kE x
k'=1
WrO0) x
max M MEM
.,k
wl0)
x
max k'E(n]
4rmax
max ME(n]
x
..ke
(from equation and equation 32
(33)
(34)
which means that equation 27 and equation 28 are satisfied with the weight matrices defined above.
Now, we see that the weight matrices W (O) w W(K), ((() configured above can distinguish the most subtle pattern of input tokens, i.e. x (
X ( with v(t) * v()
Pick up any i,j [N] and k. I [m] such that x x! and yia) + yu) In addition, we define (1)) (3) by
(
a (w(k) ( (w( x ER"
(35)
( (wik)x()) (wiQ)
E R..
(36)
Then, equation and and equation 31 imply that ( and a are tokenwise (r, (ssparrae,, where defined by
TO 8r max
r= + 1)a
(37)
efmin
because for any W E [n]. we have
(((
(1)
(w(Q)x!!))
vv) (1)
|u u I
( xx
$ (1V) + 1)a nd
8
(from equation 29 and equation 31
(38)
8 crmin
and the same upper-bound also holds for a
Since (+t)_ o vo) and there exists duplicate token in x and X(i respectively, it follows from
Lemma [ that
boltz(a()) - b0)) >8' 8' = -log n) 2 2r
(39)
that is.
(alo)) os [a(0] - -(a(t) as [al0]]
> ~
(40)
17

---

Published as a conference paper at ICLR 2024
Since x (6)
x
by assumption, equation 40 are further expanded
8' (a(1)) so [a()) - (aaii) os [aio]
(x(2)
(W1lO) w(k) (x(1) (a(1)
[alo1)
(x vul uv
(x(10s (a))) - -x
[aco1)
(i) 1 u'l
v x (i)
05 [a00] a x(xx o5 |
S Tmax (1V + 1)*
Td
(v"
x (1))) - x()x a
(1)
(41)
8
ermin
where the last inequality follows from equation
29 and equation 31
Therefore, the gap between the outputs of th self-attention function for X and x are lower- bounded as follows:
F SA)
(x(1)
(SA)
(x
|W()
(Wi')) x
oS ( wlo) (wo()) x()) os
(: (: x
=II u"|
(
o x
os ( (o
(x(xi))) os [a0]
8'
Ser
(42)
4rmax
([V) + 1)4
vdδrmax
where 8 and 8' are defined respectively
5 2 2og + 3,
(43)
dd 5r max
8' (log n) e-2r with r= ((V) 1)*
(44)
er'in
By plugging ō and 6', the above inequality is simplified as
(SA)
(
I(Sa)
(x(i)
i.k
2(log 222 2mi
(- 1)a (2 (2log + 3)r max
exp
(45)
(IV) + 1) 3)rd
4crmin
A.3 PROOF OF COROLLARY
Proof. According to Theorem we can construct such self-attention to be contextual mapping, that
is, there exist weight matrices (O)) E Rdxe and W (()), w(Q) e Teed such that F (SA) with h 1 is a (r, d'-contextual mapping for the input sequences X ..... X (N) with and 6 defined by
Tmax +
(46)
o'iin log
2 max (DV) + (2 n - +
exp - -(( 1)2 (2 log n + 3)
(47)
4crmin
So what remains to do is associate each context id with the corresponding output label using feed-forward neural network (FFF), Construction of such network is typical memorization
18

---

Published as a conference paper at ICLR 2024
task of a one-hidden-layer feed-forward neural network. Here adopt the implementation from Zhang et al. (2016). In this case, since the possible number of context ids is upper- bounded by upper-bounded
n.,, the required parameters for the FF layer with output dimension d is at most x (2n.V + d) (Zhang et al. 2016). As for the self-attention layer, rank-1 weight matrices W E Rdxx and W(V), W(k) W(Q) E R all require + d parameters each. Thus, the number of parameters of the self-attention layer 4(s d). In conclusion, the total number of parameters for one-layer Transformers to memorize the dataset is at most 4(s d) + d(2nN d).
A.4 PROOF OF COROLLARY
Proof. First, we define the positional encoding matrix E E RRdx" as follows:
2r max
4rmax
2nTmax
E =
(48)
max
4rmax
2nTmax
Then,
( +E
E are tokenwise (Mmax, (2n - mixx <) sespaarrte and each sentence
has no duplicate token.
From Theorem 2 there exist weight matrices W Rdxe and WV WK WQ Rexd such that F((A) with h = 1 is a (r, conotettllll mapping for the input sequences X + E E...., ( + with and 5 defined by
E r = (2n + max +
(49)
7
2log 2
8 (2n + 1)2r max (n. + 1) log n + 3)rd
exp (-(2n+ 1)2 (n.V + 1)a (2 log n + 3) 3) max
(50)
4e
because the size of the vocabulary set of X (1) + E...., X (N) + E is at most n.N Hence, hereafter do the same thing as in the proof of Corollary T that is, implementing feed-forward neural network (FFF) which associates each context id with the corresponding label. The total number of parameters required to implement this construction can be evaluated in the same manner as in Corollary
A.5 PROOF OF PROPOSITION 1
Proof.
We show the propositioin by the same steps as in Yun et al. (2019). Namely,
1. First, given a permutation equivariant continuous function f E JPP defined on a compact set, it follows from typical analysis that f can be approximated by step function with arbi- trary precision in terms of p-norm. Therefore, to show universal approximation theorem, it is sufficient to show that such step function can be approximated by Transformer with one self-attention layer.
2. Second, we use first fred-forward neural network layer F( quantize the input do- main, reducing the problem to memorization of finite samples.
3. Then, by similar analysis as in Corollary it can be shown that combination of the self-attention layer (SS)) and FF(FF) can memorize the step function almost everywhere, in the sense that quantized input domains corresponding to sentences with duplicate tokens are negligibly small.
We hereafter provide rough proofs of the three steps outlined above, because there are multiple ways construct model that satisfies the above requirements, and we do pursue the efficiency of feed-forward neural networks in this paper.
First, without loss of generality, we ignore skip-connections in in.((FF) (FFF) and (FF)
19

---

Published as a conference paper at ICLR 2024
1. Since f is continuous function compact set, f has maximum and minimum values on the domain. By scaling with F(FF) and F(FF), ff assumed be normalized without loss of generality: for any Rdx=
[0. 1d
f(z) =0
(51)
and for any Z E [-1,1] x
f() 51.
(52)
Let D E N be the granularity of grid
Cp (1D,,/DD.... 1] C Rd
(53)
such that a piece-wise constant approximation
[(Z) = Σ (LL 1zen+(-1/D.0pa
(54)
LEGn
satisfies
(p( < c/3.
(55)
Such a D always exists because of uniform continuity of f.
2. We use F ((FF) to quantize the input domain into Gp.
For any small > O, the following 9-approximated step function can be constructed with mee-hddee-llayee ferd-for neural network: for any E R
<
RR [2/0]
OR 12/6 1]
2/OD
2 2
(56)
D
1/D
8 S
By shifting and stacking this step function, we have an approximated multiple-step function
D-1 on 1z/0 1//D] on (2/0 1/8D)
0
z<0
1/D 0<=<1/1
s qqantp(z)
(57)
1 1/D
and subtracting the last step function from it,
O |= 1/8DD oR [=/6 1-16D (on (2/6 1/0] on (=/0 - 1/8)
D
(58)
approximately quantize [0, 1] into (1/D,...., 1}, while projects (8) to
These operations can be realized by one-hidden-layer neural network, and it is straight- forward to approximate its extension quant to dimension x n, which we denote quant d'x [gdxn + Rddx
D
In addition that, also add penalty term, with which identify whether an input sequence is in [0, _]dxn or not. This is defined by
oR [z 1)/5] OR 1( 1)/6
oR a2l-~166 oR [=2/ 1]
2 250
z penalty(2)
0<251,
(59)
-1 1<z
20

---

Published as a conference paper at ICLR 2024
which can also be implemented by one-hidden-layer feed-forward neural network.
Combining these components together, the first feed-forward neural network layer F| (PF) approximates the following function:
A quant dxn (Z) + Σ
penalty (Zt,k)
(60)
t=1 k=1
Note that this function quantizes inputs in [0, 1]4x= with granularity D, while every element of the output is non non-positive for inputs outside [0, 1|=x". In particular, the norm of the output is upper-bounded by
max F ( 1. = dn vd
(61)
2 ZEROXA
for any k e [n]]
3. Let GD C CD be sub-grid
-p= GD Vk, [n], L, + L,
(62)
and consider memorization of GD with its labels given by f(L) for each L E Gp. Note that the label sets are consistent because is permutation equivariant function. Then, Theorem alloows us construct self-attention F(SA) be contextual mapping for such input sequences, because GD can be regarded tokenwise (1/D, Vd, 1/D)-separated input sequences, each of which has no duplicate token by definition. The idea is that when the granularity D of GD is sufficiently large, the number of cells with duplicate tokens, that is, |GD \ @D| negligible compared to the total number |Gp| of cells, and thus the memorization of GD suffices for universal approximation theorem.
From the way the self-attention ((SA) constructed, we have
(SA) (Z) |
Z.,
max WEM
(63)
AV
for any k [n] and Z _dxn This follows from the fact that ( in equation 33 may actually be replaced with any input sequence Z, because in equation 33 is unit vector. In particular, combining this upper- bound with equation 61 we have
(SA)
((FF) (FF)
(Z.,) - F((FF) (Z (Z
dn
(64)
4D
Thus, if we take large enough D, every element of the output for Z E Rd \ (0, 1 is upper-bounded by
((AA) F( (FF) (Z), <
(vt e [d], & E [n]),
(65)
4D
while the output for Z E [0, 1]dm is lower- by
F o Fi ((FF) (Z),R² > 4D
(vt E [d], &: k [n))
(66)
Therefore, what remains to show is construct feed-forward neural network F (PF) which associates the context id of each L E Go (3/4D, oc))).o.dddee. )xn to its corresponding label, while it outputs 0 for any input matrix E (-o0, 1/4D) dxn, This can be accomplished by usual hump-function Precisely, construct bump function of scale R > 0
F(L) dn
bumpn(()
ΣΣe on [R(Z,, L2,,) - 1] Rn [R k k=
L(,k)
(67)
on [R(Ze, L(,k) + 1))
(68)
21

---

Published as a conference paper at ICLR 2024
for each input sequence L Gn and add up these functions implement F22
For large enough R > 0, F2 (FF maps each input sequence L E Gn to its labels f(L)
and Z 1/00,
O. In addition, the value of F(FF) is always bounded: (FFF) I 1. Thus, by taking sufficiently small ò > 0, we have
0 S F2
dp (F( F FSSA) oFFFF) , F ((SA)
(((F)
(69)
Lastly, there are |Go Gp/= = O (GDl) input sequences with duplicate tokens, each
corresponding to a cell of area D Thus, by taking sufficiently large D, areas of G GD becomes negligible and
dp ((FF) o (SA) oF (
(70)
3
Combining equation 55 error the Transformer:
equation 69 and equation 70 together, we get the approximation
dp F2 FF F(
oF (F(F), .)) <e.
(71)
B TECHNICAL LEMMAS
B.1 PROJECTION OF INPUT TOKENS
We cite Lemma 13 in Park et al. (2021), with which we configure weight matrices of a self-attention mechanism.
Lemma 2 Park et al. 2202 ) Let d E N. Then, for any finite subset X C IRd, there exists unit vector v E 12ah such that
|
x'1 S 10 (x - -x))( S 1|e0 - x'll
(72)
Rd
holds for any x. x' E X.
The following lemma follows immediately from Lemma 2 We W(K) and (QQ) in the lemma as low-rank weight matrices
Lemma 3. Given (Tmin: Fmax, e)-sparraeed finite vocabulary V Rd with Tmin 0. Then, for any
8 > 0, there exists unit vector v E Rd such that for any vectors u, u E R with
Fd 8
u T u'| =[V]= + 1)*
(73)
8 eT el'min
we have
(wik))o)) T (w(l)o.) ( (w(k)os) (wrcore)) =6
(74)
nd ||all| S Tvol vel S well
(75)
([V+ + 1)a
for any Va. Vb. Uc E V with va # Ub. where (K) uv
T sxd and ( u'u Raxd
Proof. By applying Lemma 2 to VU (0), we know that there exists unit vector v e Rd such that for any Va, Uú V/ (0} such that vn + vb. we have
(IVI 1)² v. - vbll s (v, vb)l S l| - vpll
(76)
'It is easy to show that different unit vectors v. v' E Ied may used (K) and ( respectively, that is. (K))) uv" and ((() u'b" as long as v and o' satisfy equation 72
22

---

Published as a conference paper at ICLR 2024
In partieular, this means that for any E V
1
1ell S 10 TTell
vell
(77)
(IV) 1)2 Td
holds. Thus, pick up arbitrary vectors u, u' E I2* with
u uu| ((V) + 1)² rd 8
(78)
8 Crmin
and by setting W(K) = ux E Roxd and ( = u'v E Rexd we have
(k())s)) (wrle)s) (w((k)
(w(1) vz)
(W(Kk)) (wil)o))
(va- v6)
(va v6)
1 T u'
1 vvl
8
1)* Td
1
1/e,
vbll
(12)
nd ||o||
([V) + 1)
Td
8 eTmin
((V) + 1)
(79)
where the last inequality follows from (Fminy maxx: sepparattesssss of V.
B.2 PROPERTIES OF BOLTZMANN OPERATOR
For any vector (.,,..., a2) E R", let denote p (Pi....,Py ..... os(a]. In addition, we define the Boltzmann operator, partition function, and entropy as follows.
boltz(a) a.Pi.
(80)
i=1
((a) llog
Σ c
(81)
S(p) Σ ps log
(82)
1=1
The following lemma shows that the Boltzmann operator is monotonically decreasing when the maximum and the rest of the arguments are far enough apart.
Lemma 4 (Monotonicity). Let N+, E [n] and (ai...., an) R". Then, the derivative of the Boltzmann operator boltz(a) = a as(a with respect a, is
i
boltz(a) = px(1 logp, + S(p)).
(83)
da,
In particular, the Boltzmann operator monotonically decreasing with respect to a, whenever
< max 0 log jE(n]
(84)
holds.
Proof. Since
n
C(a) S(p) = log
Σ
a
Pk log Pk
1=1
k=
n
Pk log
Pk
j=l
Pk log
a
pbap boltz(a),
(85)
k=1
23

---

Published as a conference paper at ICLR 2024
have
a
a
(a)t =
L(a)
S(p)-
(86)
da,
8a:
oa;
Notice that
Opt
0 ea.
daj
daj
> km1 cak
a
et 2x kiml
1 e0 eac
day Lim c0
8.
Bnz
(Ek=1 eaa)
Beje4 Ek=1 c0 tlr c
py - P:)
(87)
(k )al)²
holds for any i [n]. Here dij is the Kronecker delta, that is,
if + j, dij 1 if
(88)
Thus,
OC(a)
Ear
pi
(89)
0o;
E
j=1
0S(p)
0S(p) 0p
Oo;
Öp3
Oa,
==1
Σ
Σ Ph log pk k=1
-iilfj - Pj)
Öpi
Σ
-log Pj 1) pi( -Pi;
n
= Ps
Σ
18 (log p, + 1) - Py logp; p. Pjl
)=1
= (log p. + +1+ - 1)
p: (log P: + S(p)).
(90)
Plugging equation 89 and equation 90 into equation 86 we have
θ L(a) θ
boltz(a)=
S(p)
Da;
Da:
0a,
= p: + po p. + S(p))
p1 log p: + S(p)).
(91)
In particular, the derivative of the Boltzmann operator at index i negative when
1 1 S(p) + logp; < 0 + S(p) + (a: L(a)) <
0; C(a) S(p) 1.
(92)
Since maxje(n) Oj < (pa)(, a boltz(a) 0 whenever Bo,
Boyd & Vandenberghe 2004)) and S(p) < log n, we have
0: < max 0g- log 1 }E(n]
(93)
holds.
Lemma 5 (Concavity). Let E N_, i E [m] and a (a1,..., an) R.. Then, the Boltzmann operator boltz(a) is concave with respect to that is
2²
boltz(a) <0
(94)
du
holds in domain where a satisfies
04 < max - log n - 3. jE(n]
(95)
24

---

Published as a conference paper at ICLR 2024
Proof. According to Lemma
have
boltz(a) - pi(1 + logp: + S(p)).
(96)
aa
Thus,
8
à
boltz(a)
[p.(1 + logp: +(p))
dag
0a,
Öp:
(1 + log +S(p)) + Pi
(1 + log p; + S(p))
Oa,
0a;
pp - pi)
p: - p:) (1 + log pi + S()) + pi
pe + S(p))
= Pr [1 2p;) ) (logp: S(p) - 1) + 1],
(97)
where we used equation 87 and equation 90
Hereafter, we show that the right-hand side of the above equality is negative under the assumption that equation 95tiioonn55ho holds. First, the fact that ma Oj ((a)(.722, Boyd & Vandenberghe 2004 and S()) s log implies
a, < max dj logn -3 [(a) -(()) 3. ]E(n]
(98)
follows from a, L(a) log Pt that
< max aj log 3 = log < -S(p) - 33
(99)
jE
Next, since the entropy S(p) is always non-negative, we have log Pt < -3 under equation using inequality -- S 1+z
95 By
for > 11, this implies
1
P 3<e3
(100)
1 1+3 4
Thus, as long equation 95 holds, have
(1 2p, (logp: + S(p) + + (1-2. 2.) 1) (1-2) -1 (101)
which in turn implies 8² boltz(a) <. do?
Lemma 6. Let n 2 and (a1,..., an-1, an), (b .b. bn) E R² be sequences such that the first - elements of and b match, that is, = b, for all i E [n 1}. In addition, if
max a; 8 an bn <E[n-1]
(102)
with d log n 3, the difference boltz(a) boltz(b) is lower-bounded by
cba
boltz(b) boltz(a) > (a, b, + a, bn -llog n
(103)
Proof. According to the monotonicity (Lemma and concavity (Lemma 5) of the Boltzmann op- crator, we have
i
Moltz((i......bupy + (an x) Ox
1, > boltz bo,.. bn-1, On
(104)
for any x an In particular, by setting bn and using equation 83 this implies
(-
boltz(b) boltz(a) > (a,, bw)) (- dx
à boltz bn x))
(an bw)) [-Pr((t (1 logog +S(p))]
(an bn) -Pr [ bn max b: + log iE/n]
(a, b,)) (0 + an - bn - -log n 1), where p (P1,....Pr) Pn)) RD is the softmax function of b, i.e., p os (b)
(105)
25

---

Published as a conference paper at ICLR 2024
B.3 PROOF OF LEMMA
Before moving on to the proof, we first illustrate the proof sketch of Lemma | using a simple example.
Since the Boltzmann operator is permutation invariant, without loss of generality we assume the elements of ()) and ( are sorted in descending order, e.g., ( (8, 6)5 and a = (4,3, 1). In this case, since the Boltzmann operator boltz can be regarded approximation of max, have
boltz(a!)) ()) 8 4 boltz(a ())
(106)
and so the Boltzmann operator can readily separate these two inputs.
The subtle case is where the initial tokens of (*) and ( are identical, but the rest of each sequences differs, like (*) (8, 6, 5) and a (8, 3. 1). However, closer observation reveals that the first coordinate and the second one of the input E IR' are well-separated, then boltz(a) is
monotonically decreasing for each coordinate = 2...., n. In the above example, this intuitively
implies
boltz (al)) 2 boltz (
with σ = =8,,333)
(107)
and
boltz (a()) S boltz (a (6)
with (+) (8, 6, -00),
(108)
or by abuse of notation, boltiz (a) S boltz (8, 6).
Thus, if we can show boltz (a()) < boltz ( then boltz (a! < boltz (al1) holds and we know that the Boltzmann operator can also separate this pattern inputs. Fortunately, this indeed the case if each element of the inputs is sufficiently separated, because
Se + 3e3
3e3
8e 2 3
boltz (2())
c ee + e3
88 + 2-
8 + (3 + log 2
boltz (8, + log (109)
e + logl
and we have boltz (8, 6) boltz (8.3+logg2 by the monotonicity of the Boltzmann operator.
Proof of Lemma We only have to show the case of m 2, and for notational convenience, a (!)) (resp. a (2)) hereafter is denoted by a (resp. b). Also, since the Boltzmann operator is permutation invariant, we assume without loss of generality aj > ... > an and 6 > ... > b, (Since there is no duplicate element in a, ag is strictly greater than a; for any l. The same holds for b).
First, since the Boltzmann operator can be regarded as weighting averaging, we have
[oolz(aa) max([a1), la,.l) <
(110)
For 8' sessaaaaatssssss, a + b, w.l.o.g. we assume that there exists ls E {O...., - 1) such that (a1...., ax) = (b.....,by) by) and a&+1 > bx+1-
(111)
Then, Lemma implies that have
[boltz(a) - boltz(b) > (log n)2e (ai bbe1)).
(112)
Since and b are tokenwise (r, 8 -)ssppaaat a1 the above inequality is further lower-bounded by
bx+1 2r holds. Thus, the right-hand side of
[boltz(a) boltz(b)) > (logn)²e (an-bbi))) > (log 2
(113)
O
Lemma 7. Let a, E Rn be tokenwise d-separated vectors in a decreasing order with no duplicate element each vector and 8 2 log that is,
i {>ラ a; - ag. b; - b, > 5, a + by = |a; - -b 8
(114)
(115)
26

---

Published as a conference paper at ICLR 2024
for any i.j [n).
In addition, suppose there exists E {O,.... n 1) such that
(ai...., 0x) (b1,..., .,, bu) and ax+1 bk+1. Then, the outputs of the Boltzmann operator are (log n (ar - bbtt -separated, that is, boltz (a) boltz(b)) (log n)²e e (or--be+1)
(116)
(117)
holds.
Proof. We show the lemma by dividing it into the following two cases:
1. k 121
Let a and b be
(a), .2..... ........000. 0k1 ak+1) R
(118)
(a1, 02 .. ak. b be41s.... bx+1) E R".
(119)
Then, by abuse of notation, Lemma implies that boltz (a) a Øs [a] boltz (a) and boltz (b) boltz (b)
(120)
and only have to evaluate the magnitude of the difference boltz (b) boltz (a). Let 7) and &a be
7
ale at
and E = e
(121)
Then, boltz (b) be decomposed into
Th + (n - k. k)bx+1e
boltz (b)
& + (n k)chx11
+ b4+1c244+)oo(nn
& + ebettttog(n-1
"kk + (bx+1 + log(n k)) (i log(n k)) Eh 1oog(1
Ek chul log(nn
EL + σb 11
boltz (ai,..., ax bb+1 + log(n - k)) log(n k) c +log(n-))
&a + ebu +lgg(n
(122)
Therefore, the difference boltz (b)
boltz (a) can be written as
boltz (b) bollz (a) boltz (a1, abe bb+1 + log(n - k)) boltz (a)
log(n (n - k) tlggnnk)
(123)
&x +log(n)))
According Lemma 0 the first two terms on the right-hand side are evaluated as
boltz (01,.... ak: bat1 log(n k)) boltz (a)
(ak+1 bet1 log(n (5 + ax+1 bet1 log(n k)) log(A + 1) 1)
E + +lo
(6 128 2 loog 1).
(124)
Ex + cbe
27

---

Published as a conference paper at ICLR 2024
Since 8 2log by assumption, the above inequality is further lower-bounded by
boltz (01,..., 0k, bet1 log(n k)) boltz (a)
chesl chesstloR(nnA)
(6 log 2 2 n 1)
&k + tlogi-))
e (b+1+log((-h)
> (logn + 3) log 5)
(125)
Ek ch (+log(n-R))
Plugging the above inequality into equation 123 we see
boltz (6) boltz (a) boltz (ai,.... ag. bk+1 + log(n k)) boltz (a)
log(n k) bu+4 tlog(m-k))
&s + e Hlgg(a
b (1+tog((nn)))
EA ebe4i (log + 3) + 5) - log(n k)]
ch++1+log(nnh)
&A + cb tlgg(n-)) 2(logn)².
(126)
Lastly, notice that the following inequality follows from ô-separatedness of and and the assumption that a has no duplicate token:
Ek bbts +log(n-K))
Σ
(:: a:+1 > bk+1 log(n k))
k+1
e (1-1)8
1=1
< 2c
log 2). ( 5 >
(127)
By using this inequality, equation 126 is Nower-bounder by
bbetitoo(nn-)
boltz (b) boltz (a)
EN + tlessttog(nnx) 2(log n))
e e Mloiitog(nn-e
2ea 2(log n)²
(logn)² (aa
(128)
which implies boltz(b) boltz(a) > (log 2 -but1),
2. k ==0 0
Since the Boltzmann operator can be regarded as weighted averaging, boltz(b) S b always holds. Thus, is enough evaluate how much greater boltz(a) than bj.
Let E R" be
(a), S,...,a 8).
(129)
Then, boltz(a) > boltz(a) follows from Lemma 4 and its value is
boltz (a) asea + (n 1)(a) 8)
t0l + (n a)
+ (n - () o)e-
1 1 (n 1) )
(n 1)6E
0j
1 (n 1)e-5
(130)
28

---

Published as a conference paper at ICLR 2024
Therefore, the difference boltz (a) boltz (b) is
boltz (a) boltz (b) 2 boltz (T) - b1
(n 1 1) -
d1 1 (n 1)e
(a1 -
(n 1)6e
1 (n 1)e-
5
1 (n 1)
log n,
(131)
where the last inequality follows from the assumption 8 > 2 log n. Note that the right-hand side is greater than (log e e (01-b)), because 01 by >> log n implies log n e --(a- -b1) 1.
C EXTENSION TO MASKED SELF-ATTENTION
Masked self-attention mechanisms are formulated as follows: the Softmax part in equation replaced by
(w (w/L))
(Q) z) +c E pdxn
(132)
with masking matrix C E R=* whose elements are either 0 or -00.
Our main result Theorem can be readily extended to the case where attention masking are con- ducted. The idea is that
boltz(a + c) (a c) as la c]
=a =a os la + cl c os
t la
=a c
(133)
holds for any masking vector R whose elements either 0 or -o0 Thus, in order to ensure that the masked attention is a contextual mapping, it sufficient to verify boltz(a + e) are well- separated. The caveat that the Boltzmann operator now has to distinguish inputs consisting of the same attention a, but different masks C₁, C2. However, the separability in this case can also be proved in the same way as the proof of Lemma !
29


Reviews:
[{"review_id": "Ohw9BCVEpD", "reviewer": "Reviewer_A9Eo", "summary": "The paper proves a new approximation results for Transformers with a single layer and attention head and dim-1 head size. The key result is that under suitable conditions on the input sequences, there exists a self-attention layer with softmax function and the settings as above such that the output tokens are bounded and separated from each other by a given distance. This allows a feedforward layer to be trained on top of it and associate to each output token a given label due to the memorization capacity of FFN layers. Hence, two theorems regarding the memorization capacity of such Transformers are proven. Further, exploiting this result, a novel universality result is proved for two-layer and single-head Transformers for permutation equivariant functions.", "questions": "I have a couple questions:\n- What is a sequence id, similarly what is a $v$-dependent sequence id?\n- Is there any implication regarding the value of $\\delta$ in the contextual mapping? In particular, how are the following approximation results affected if $\\delta$ decreases?\n- How does the memorization capacity depend on the vocabulary size?\n- As is mentioned above, can we value of the rank-$1$ projection be computed as implied by the theory and kept at that value in the experiments to see how well that performs?\n- Is there any difficulty with extending Prop. 1 using positional encoding to consider functions that are not permutation equivariant?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Very interesting paper! The discussion is surprisingly readable, although I had some issues with interpreting some keywords which are not properly introduced. The results are quite significant, since this is the first paper according to my knowledge that proves approximation results for Transformers with a single layer and attention head. It's interesting that a dimension-1 head can already produce a contextual mapping (bounded and separated output tokens), hence allowing to train an FFN on top to map each token to the required label. Detailed proofs are provided in the appendix, which I did not check completely, but from what I've checked they seem correct.", "weaknesses": "I guess one weakness (which does not deduct from the value of the theoretical results) is that the experiments section simply demonstrates that such a Transformer (single layer single head head dim-1 and tied weights between Q K and V) can already memorize a dataset. Maybe one thing I would be interested in is whether keeping the projection vector fixed and set to the value found using the technique in the proof of Theorem 2 using Lemma 3 can already perform well in practice? \n\nI also had some issues with a few keywords that were not introduced, for example, in the discussion of Theorem 2 I had no idea what a sequence id was, or how it is used to construct the context id (which if I understood correctly is the output of the attention map?).\n\nSome interesting follow-up questions were left unanswered, for example, how should the size of the FFN layer scale theoretically given the separatedness of the output tokens, or how having multiple attention heads and/or with higher rank can improve the performance?", "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": "UX0F8C3z4d", "reviewer": "Reviewer_EXsH", "summary": "The paper critically examines the expressive capacity of Transformer models, specifically addressing the discrepancy between theoretical analyses and practical implementations of Transformers. Existing analyses have often necessitated overly deep layers or numerous attention heads for data memorization, which doesn't align with the Transformers used in real-world applications. This misalignment is largely attributed to the interpretation of the softmax function as an approximation to the hardmax function.", "questions": "Please see the Weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-organized and clearly written, which is easy to follow.\n2. The problem studied in this paper is interesting and valuable.\n3. The theoretical work of this paper is sufficient, which improves the value of the paper.", "weaknesses": "This paper delves into the theoretical underpinnings of one-layer Transformers. However, there are areas that could benefit from further exploration:\n1. While the study provides insights into one-layer Transformers, it raises the question of whether these findings can be extended to two-layer or even deeper architectures. How scalable is the presented theory?\n2. The role of the number of attention heads in determining the memorization capacity remains unclear. If it does have an impact, are there any quantitative metrics provided to elucidate its influence?", "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": "WIB6vApXPQ", "reviewer": "Reviewer_mvi8", "summary": "This paper proves that Transformers with one self-attention layer using low-rank weight matrices, plus two FFNs, are universal appoximators of sequence-to-sequence functions. It makes a connection between the softmax function and the Boltzmann operator, and argues that softmax attention can be more useful for universal approximation than hardmax which is commonly used in existing papers.", "questions": "- What are the size of FFNs and the number of model parameters in the construction?\n\n- Can you clarify the discussion on [1] in the paper?\n\n- Can you clarify the proof issue mentioned in **weaknesses**?\n\n[1] Approximation theory of transformer networks for sequence modeling.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- I appreciate the connection between the softmax function and the Boltzmann operator.\n\n- The negative result on hardmax attention looks interesting to me.", "weaknesses": "- **Significance of the result seems limited to me.** Some claims in the paper are weak or wrong.\n   - The authors claim that existing results on the expressiveness of Transformers require excessively deep layers or a lot of attention heads. While this paper requires less parameters in the attention layer, its construction requires large FFNs. \n   - Moreover, there is no characterization on the total number of model parameters in the memorization capacity result. It's not surprising at all the Transformers can memorize. It is the upper/lower bound on the required number of parameters for memorization that makes the result interesting.\n   - This paper is _not_ the first universal approximation theorem for two-layer Transformers with a self-attention of realistic size, and Section 4.2 _overclaims_ the contribution because of misinterpretation of [1]. The universal approximation result (Theorem 4.1) in [1] _is_ constructive and the construction is even simpler than this paper. The authors also claim that [1] make a particular assumption on the domain of functions. But those assumptions are only for the approximation rate result (Theorem 4.2), which is orthogonal to this paper.\n\n- **Proof issues.** Equation 10 seems buggy. First, I think the authors should specify $\\mathbf X^{(i)}\\neq \\mathbf X^{(j)}$ or $\\mathcal V^{(i)}\\neq \\mathcal V^{(j)}$. Second, if it holds for some $i,j$ and $k,l$, then one can consider the same inequality for $j,i$ and $l,k$ and obtain a contradiction. So I don't see why the condition can be true.\n\n- **Related works.** The year of [2] is marked as 2023 in the paper, but [2] was actually published in ICLR 2020. The authors should correct this to avoid misleading. In addition, I recommend the authors to also discuss [3,4,5].\n\n- **Presentation issues.** There are many typos and minor errors in the paper, to name a few:\n\n   - Page 3, \"the Transformer block is represented as a combination of ...\" $\\to$ \"the Transformer block is represented as a composition of ...\"\n   - Page 5, \"this theorem indicates that one-layer Transformer does not have a memorization capacity\": Please emphasize one-layer Transformer _with hardmax attention_.\n   - Theorem 2, $\\mathbf{W}^V$ $\\to$ $\\mathbf{W}^{(V)}$.\n   - The appendix uses eq. $i$ and equation $i$ interchangeably. Please make it consistent.\n\n\n[1] Approximation theory of transformer networks for sequence modeling.\n\n[2] Are Transformers universal approximators of sequence-to-sequence functions?\n\n[3] $O(n)$ Connections are Expressive Enough: Universal Approximability of Sparse Transformers.\n\n[4] Universal Approximation Under Constraints is Possible with Transformers.\n\n[5] Your Transformer May Not be as Powerful as You Expect.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)