Separations in the Representational Capabilities of Transformers and Recurrent Architectures
Authors: Satwik Bhattamishra
Michael Hahn
Phil Blunsom
Varun Kanade
Venue: NeurIPS 2024 poster (2024)
Keywords: expressivity
Transformers
RNNs
deep learning theory
communication complexity
Primary area: deep_learning_architectures
TL;DR: We show that small-sized Transformers can theoretically express tasks like index lookup, nearest neighbor, and string matching, whereas RNNs require larger sizes; Also show limitations of one-layer Transformers on recognizing Dyck languages

Abstract: Transformer architectures have been widely adopted in foundation models. Due to their high inference costs, there is renewed interest in exploring the potential of efficient recurrent architectures (RNNs). In this paper, we analyze the differences in the representational capabilities of Transformers and RNNs across several tasks of practical relevance, including index lookup, nearest neighbor, recognizing bounded Dyck languages, and string equality.  For the tasks considered, our results show separations based on the size of the model required for different architectures. For example, we show that a one-layer Transformer of logarithmic width can perform index lookup, whereas an RNN requires a hidden state of linear size.  Conversely, while constant-size RNNs can recognize bounded Dyck languages, we show that one-layer Transformers require a linear size for this task. Furthermore, we show that two-layer Transformers of logarithmic size can perform decision tasks such as string equality or disjointness, whereas both one-layer Transformers and recurrent models require linear size for these tasks. We also show that a log-size two-layer Transformer can implement the nearest neighbor algorithm in its forward pass; on the other hand recurrent models require linear size. Our constructions are based on the existence of $N$ nearly orthogonal vectors in $O(\log N)$ dimensional space and our lower bounds are based on reductions from communication complexity problems.  We supplement our theoretical results with experiments that highlight the differences in the performance of these architectures on practical-size sequences.

Separations in the Representational Capabilities of Transformers and Recurrent Architectures
Satwik Bhattamishra' Michael Hahn
Phil Blunsom 1,3
Varun Kanade
UUniversity of Oxford
Saarland University
3 Cohhre
Abstract
Transformer architectures have been widely adopted in foundation models. Due their high inference costs, there is renewed interest in exploring the potential of efficient recurrent architectures (RNNs). In this paper. we analyze the differences the representational capabilities of Transformers and RNNs across several tasks of practical relevance, including index lookup, nearest neighbor, recognizing bounded Dyck languages, and string equality. For the tasks considered, our results show separations based on the size of the model required for different architectures. For example, we show that one-layer Transformer of logarithmic width can perform index lookup, whereas an RNN requires hidden state of linear size. Conversely, while constant-size RNNs can recognize bounded Dyck languages, we show that one-layer Transformers require linear size for this task. Furthermore. we show that two-layer Transformers of logarithmic size can perform decision tasks such string equality or disjointness, whereas both one-layer Transformers and recurrent models require linear size for these tasks. We also show that log-size two-layer Transformer can implement the nearest neighbor algorithm in its forward pass; on the other hand recurrent models require linear size. Our constructions on the existence N nearly orthogonal vectors in Olog N) dimensional space and our lower bounds based reductions from communication complexity problems. We supplement our theoretical results with experiments that highlight the differences in the performance of these architectures on practical-size sequences.
based
Introduction
Transformers [59] are the go-to architecture for building LLMs 11.. but recently, there has been significant interest in reviving recurrent architectures to build LLMs for practical tasks 421 45) to circumvent the quadratic complexity of inference for Transformers. While Transformers process all tokens in parallel, maintaining N vectors, and are in some sense stateless, recurrent models primarily information in fixed-size hidden state that they can update during the course of the computation. This raises a fundamental question that we explore: Are there tasks that are computationally easier for one architecture represent but significantly more difficult for the other one?
Ever since Transformers supplanted LSMMs, there has been significant interest understanding the differences in the computational capabilities of the models both theoretically and empirically. On the one hand, even as Transformers have been widely adopted for building LLMs, several studies found that they struggled at modeling various formal languages, particularly those that required modular counting or state-tracking [5] [6 38). At the same time, despite substantial effort, it has proven hard match the performance of Transformer-bases LLMs scale with recurrent architectures [15 45 21: in particular, it has been observed that LLMs based on recurrent models struggle on associative recall or extraction-related tasks [5] More recently, Arora al. and Bhattamishra
*Corresponding Authors: Satwik.bmishre varun.aande@ecsoo..cc
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

et al. [ demonstrated that Transformers are better than attention- free models at synthetic tasks based on associative recall and more general forms such as implementing nearest neighbors. Based on these observations, it is natural to wonder if such tasks are theoretically easier for Transformers to represent in comparison to recurrent models.
Our Contributions. We show differences in the representational capabilities of Transformers and recurrent models across several natural tasks of real-world relevance. In this paper, we show strong separation results: when a task is easy we show that it can be expressed by one architecture with size poly-logarithmic in the input length N: on the other hand we show the other type of architecture requires the size to be linear in N. We describe the tasks studied and our key results below.
(i) Index Lookup. Given a sequence of symbols $1...... S followed by a position p E [N. a model has to output the symbol in the y-th position 8p (Figure Tap. Our first result shows that one-layer Transformers with size poly-logarithmic in N express this task (Theorem whereas any recurrent model must have width 2 N) to perform this task (Theorem 3
(ii) Bounded Dycks. Dyck languages with bounded depth require model to recognize whether string of parentheses well-balanced (Figure [b). These formal languages have received great deal of attention in the study of neural sequence models because they capture hierarchical dependencies that occur in natural languages [7 25 23 6. They are also central to formal language theory as all context- free languages can be expressed in terms of Dyck languages [14]. In contrast to the results for index lookup, we show that one-layer Transformers must have a size that grows linearly in input length to represent this task (Theorem 5), whereas prior works found that constant-size recurrent models can express this task 25
= I[x = y] where y are two strings of length N; it models the natural task of matching two documents. We show that log-sized two-layer Transformers can represent the function EQ (Theorem 6). whereas both one-layer Transformers (Theorem (2) and recurrent models (Theorem [T] must have a size that grows linearly with the sequence length. These results extend to broader class of Boolean functions.
(iii) String Equality. This task is formalized by the Equality function y)(x,y)
(iv) Nearest Neighbor and Associative Recall. In this task, a model provided with a sequence of inputs and labels X Vi..... X 1, Mk followed by query imput xg. The model has to determine the label 2/ of the query input by applying the nearest neighbor algorithm in its forward pass (Figure Ic). This task subsumes various associative recall tasks considered in earlier works (cf. Appendix G and was introduced to understand in-context learning. Bhattamishra et al. [8] found empirical differences in the performance of Transformers and attention-free architectures, though theoretical understanding is lacking. We show that a two-layer Transformer with size logarithmic in the input length can perform this task (Theorem 7), whereas recurrent models require size linear in the length (Theorem
We also empirically investigated the performance of Transformers and standard recurrent models (26), including recently proposed state-space models [22 21 on these tasks. The observed behavior is along the lines indicated by our theoretical results.
Our Techniques. For constructing Transformers that perform the tasks, one of the key requirements is the ability to precisely attend to specific positions. In order to do this with low-dimensional embeddings, our constructions make use of nearly orthogonal vectors obtained using the Johnson- Lindenstrauss lemma (29]; furthermore these can be generated efficiently in logarithmic space which allows the size of the models to be poly-logarithmic in the input length (cf. Appendix [B.3).
For our lower bounds, we appeal to results from communication complexity. We show how to obtain protocols with communication complexity bounded by the size of the models. Together with established (and new) communication complexity lower bounds, we obtain lower bounds on the model size. For our lower bound for one-layer Transformers, we derive lower bound on the communication complexity for bounded Dyck languages (Lemma
1.1 Related Work
Expressivity of Sequence Models. The expressive power of neural sequence models has been
an active area of research, targeting both RNNs and Transformers [e.g. 23 parameters of infinite precision, Transformers are universal function approximators
52
40
55
37 With
67 and Turing
complete [47 77) In bounded precision settings, they relate to Boolean circuits [39 24 and logical

---

formalisms [[2]. Importantly, such studies usually consider asymptotic expressivity of a single model in the limit of unboundedly long imputs. Our study provides a more fine-grained and realistic picture by accounting for the size of the network, and its scaling with the input. The closest to our work is Sanford et al. [ST), who first used communication complexity to prove lower bounds for Transformers and RNNs for abstract tasks such as sparse averaging task and pair/triple detection tasks. Our work extends that line of work to show separations on natural tasks of practical and real-world relevance.
Formal languages and Algorithmic tasks. A strand of work has sought to understand sequence
models via empirical analysis on formal languages and algorithmic tasks [ 57 68 Numerous works
have examined the ability of recurrent models [38] [44 666
6 25 62 and Transformers [17 65 62 to
model Dyck languages. More recently, significant effort [19 60 has been devoted to investigating how Transformers can learn to implement learning algorithms in their forward pass to understand the in-context learning phenomenon. Bhattamishra al. g empirically observed that attention free architectures struggle to implement the nearest neighbor algorithm in comparison to Transformers. Our result takes a step toward understanding this phenomenon by showing that nearest neighbors can be implemented by small-sized Transformers but not by recurrent architectures.
2 Definitions
We consider two types of models: Transformers [59] and recurrent models. We use recurrent models to refer more generally to nonlinear RNNs such as LSTMs 26) state-space models [2 21 well as variants of linear Transformer [1 56 which can process inputs in a recurrent manner 31
For some finite alphabet E, a sequence model is given sequence in ΣN and depending on the task outputs either (0, }} or sequence of outputs. Each 8; E from a sequence 81 88 E Σ mapped to vector in Rd via an embedding map ø E + R For Transformers, the input embedding function further takes the position as input, along with 9;, Each layer for Transformers and recurrent models maps inputs from IR + R xd A model M has fixed precision p if all the parameters of the model as well as the values in the intermediate vectors can be implemented with p-bit precision numbers (cf. Appendix[B. [B.2..
Transformers. Each layer of a Transformer has an attention block followed by an MLP block. The attention block takes as input X R.xxd and applies the operation (x)
softmax(X Q W X JXWT where W Q, WK, Wv e Rmxdd For simplicity, we will use Q(x,) (and likewise K (xx)) and V(x:)) to denote W xxx The width of the Transformer is max(m, d). where x d is the shape of the projection matrices WQ. WK. For any matrix A E RN N the
exp( (n)))
softmax operator is applied row-wise as follows softmax AA)
Multi-head EMI (xp((A..) attention with H heads is defined as M-Att (X) [Attı(X), Att (X) where each (X) has its own set of parameters. The matrix Wo E R xd projects the concatenated vector to vector of dimension d. For an input X E Nd the output of layer of Transformer will be (M-Attt(X)
ig where sh ndd + Rd is a feedforward network. We use TF". p. to denote the class of all Transformers operating over p-bit precision numbers with width m, H heads, and at most L layers.
Recurrent Models. A general recurrent neural network (RNN) takes as input the sequence XI,...,XX, Xx where x; E Rd and produces an output sequence Vi.-.., yN; in this paper we will mostly consider the case when y/ E (0, 13. An RNN with a hidden state of size m over p-bit numbers can be defined as follows. The hidden state is an mp-bit memory h; E {0, 10,1)mpp and for some ho E {0, 10,p, the RNN computes h, = [(e) (X,, h:-1) ) and y/ = fay (h,) for = ,..., N. and g(t) and f(a) are arbitrary functions. Since the transition function is allowed to be arbitrary, this definition captures the general family of multi-layer recurrent architectures including LSTMs, state-space models, and linear Transformers, each of which differs in the way the transition function is defined. For a recurrent model, we say that the representation size of hidden state is mp, and its width is m, ie. the hidden state consists of m units each of which uses p-bits. Throughout the paper, we will use recurrent models or RNNs to refer to the general family of recurrent architectures mentioned above.
By the representation size of a model, we will refer to the total number of bits required to represent the model including all the parameters and embeddings. For Transformers, this is (mddHH). For a recurrent model, the representation size is at least mp.
3

---

S
Nearni Niighte
X. Y1 N2 Y2
X X3
Prediction
(a) Index Lookup Task
(b) Dyck-2 with depth 2
(c) Nearest Neighbor Task
Figure 1: Illustration of a few key tasks considered in our work.
3 Index Lookup Task
Task Description. We introduce simple task called Index Lookup (IdxL). In this task model receives a sequence of tokens 81,.... 8N (possibly with reperitions) followed by an index p where p E [N] and the goal of the model is to output the token 8p. Here the symbols 8; belong finite alphahet Σ.
This simple and natural task helps illustrate the key tools and building blocks we use to obtain other more general results in this paper: On the one hand, we show how one-layer Transformer with width Ollo N) can perform this task; on the other hand, we use communication complexity arguments to show that any type of recurrent or state-space model performing this task needs a hidden state with representation size (NN.
Our first result shows that, for any length N E N. there is a 1-layer Transformer with width Olog N) that performs the Index Lookup task for all input sequences of length at most N. Naïvely one could construct such a transformer by using one-hot encodings as positional embeddings, as they are orthogonal and would allow to attend to the desired index. However, this would require the embedding dimension, and hence the width of the model, to be ((N). Key to our constructions of a width (log N) Transformer, both here and in other sections, is result (Lemma in the Appendix) which states that, in & = Olog N N/?²) dimensional space we can find N nearly orthogonal vectors. We use such vectors in construction of the Transformers to allow it to attend almost exactly over desired positions.
Theorem For all N E N, there is a 1-layer Transformer with width m = O(lg N) and precision p Ollog N) which performs the index lookup task all input sequences of lengths up to N.
Proof Sketch. For an input sequence (s),..., .......,... , the Transformer uses the embeddings of the position token p and the positional embeddings of the first N inputs to attend over 8p. so that the feedforward network can extract the label from the output of the attention block. Our key idea is to use the N almost orthogonal vectors provided by Lemma 2 both as positional embeddings and also as way to embed the numbers (1...., N}, any of which can be used the index p. Formally, let ((1), ((N) be N vectors of dimension = O(log N) such that (T (i), T()) I 1/4 for i + j and (T TG)) 2 3/4 for =
Formal details of the construction are in Appendix C.2 we provide a sketch. The embedding of each input token is of size log | + 2k where log | + k entries are used for the token embeddings and the last i entries are used for the positional embeddings. The query and key matrices are designed so
that the query vector (p)) n)](,)], the key vectors K K(x,) = [T()) and K K(p) [0.]. The value
vectors simply contain the token embeddings V(x,) [o(=:)]. where p Σ - {0. 1) is some binary encoding of Σ. With such query and key vectors, the dot products in attention, (Q(p), K(x,)) are 2 3n/4 /4 if p. and < 4 otherwise. The dot product of the query vector with itself will be (Q(p), (p)) = We choose n > 0 to scale the dot products to amplify the difference between the high and low dot products. Thus we have,
exp(Q(p) K(x;)) softmax(Q(p)) K(X)) exp(Q(p) K(x,)) )) Ejj, exp(((p) K(x,))
exp( 20
exp
31 N exp(
which is at least for some 7 (log N); the total attention weight over the remaining tokens at most < 1. Recall that the value vectors contain the binary encodings of the input symbols, (xx)) = [p(s)).. The attention- average value vector aligns closely with p(sp) as 3/4

---

weight is on it. It is then straightforward to design a ReLU-FFN that can act as threshold function to retrieve p(s,) from it, which leads to the desired output.
We use results from communication complexity to show that RNNs require essentially ((N) width to solve this and several other problems. In communication complexity, there are two parties, typically called Alice and Bob, each of whom has part of the (discrete) input, and their goal is to compute function of the combined input using as little communication as possible. Our first key insight here is that the output of an RNN can be computed with a bounded amount of communication when Alice has prefix of the input and Bob has the remaining part. The resulting protocol will be one-way (Alice Bob) and one-round. We first state more general result and then discuss the implications for the Index Lookup problem.
Theorem 2. If an RNN with a hidden state of representation size mp computes any function f: EN -> {0,1), then for any K < N, if Alice has access 31..-.,8K and Bob has access to 8K+1...... SN, then there exists a one-way communication protocol with mp bits from Alice to Bob, by which Bob can compute the output of the function (s1....8N)
Proof. Assume that both Alice and Bob have access to the RNN that represents the function f. Alice can provide the sequence $1...., K to the recurrent model and iteratively update the hidden state from the initial state ho to obtain the Kth hidden state lg. Alice can then send the hidden state to Bob which requires mp bits. Bob can then update the hidden state using 5k+1..... sy to obtain hy, from which he can obtain the output of the RNN. Note that Alice and Bob can compute the output using one- way communication of mp bits.
Problems similar to Index Lookup are well-studied in communication complexity; specifically, the INDEX problem (See Appendix B.1] has a one-way communication complexity of ()(N) (Fact 3). We deduce lower bound on the size of the hidden state of RNNs by showing that any RNN that can represent the Index Lookup task can also compute the INDEX problem and since that implies the existence of one-way communication protocol with vpp bits (Theorem 2). it follows that the width of the hidden state must be (p((p)) (cf. Appendix C.I).
Theorem 3. Any recurrent model with hidden state of width m using p p-bits of precision that computes the Index Lookup task for all sequences of length N must have m 2 N/p.
The above results theoretically formalize intuitive differences between the way Trans- formers and recurrent models process sequences. Since Transformers have access to N input vectors during their computation, small-sized attention block attend over the desired input vector to make the correct prediction. On the other hand, any recurrent model-even with arbitrary positional embeddings- must store all the required information in its hidden state, which lower bounds the size of such models to compute the right output. These intuitions are made rigorous by showing (i) how soft-attention can do lookup using the almost orthogonal vectors, and (ii) small-width RNNs yield a short one-way communication protocol. These lower bounds also apply to causal forms of linear attention architectures where softmax is removed and attention weights become dot products [3I].
Discussion.
At first glance, might seem unfair to compare Transformers and RNNs with the same number of parameters: Transformers have access to N input vectors, whereas RNNs have fixed-size hidden state. But note that, in practice, empirical research on language models typically compares models of the same size c.g., 7B Transformer Transformers of a particular size can express something that recurrent models cannot.
7B space model. Hence, it is natural to ask if
4 Lower Bounds for RNNs and 1-layer Transformers
Whereas Section 3 established a case where one-layer Transformers can be more powerful than RNNs, we next exhibit an example of the opposite phenomenon. Here, the key tool will again be a communication complexity argument, but this time applies to one-layer Transformers: We establish a communication protocol by which Alice and Bob can compute the output of a one-layer Transformer by exchanging a number of bits that is bounded by the representation size of the Transformer and an overhead that is logarithmic in the input length. The key property here is that this protocol works not just when Alice and Bob have access to prefix and suffix of string, but instead works for an arbitrary partitioning of the input string (proof is in Appendix D):
5

---

Theorem 4. Consider a one-layer Transformer f E TF m.p.H operating over inputs of length N. Consider any disjoint subsets SA J Sg = [I....,N] N. SA n SB = 0. Assume Alice has access to s; for E S.. and Bob has access to 8, for i E SB- Then Alice and Bob can communicate 3m(p+ log H)H bits to compute the output [[S)...8N).
The proof idea is that Alice and Bob first compute their parts of the numerator and denominator of the softmax and exchange these compute the overall attention output. A naive implementation of this idea runs into the issue that the exponentiation of logits may exceed the bounds of p-bit precision; we circumvent this by first communicating the maximum logit and subtracting it from each logit, keeping the exponentials bounded without altering the resulting attention weights. Theorem 4 is slightly more general and formal version of a result in Sanford et al. [5] Theorem. 7].
4.1 Separation on Bounded Hierarchical Languages
We now use the communication protocol for one-layer Transformers to establish a separation between these and RNNs on bounded Dyck languages. Dyck languages are of central importance in formal language theory, as any context-free language can be expressed in terms of Dyck languages [14] Due to the boundedness of human memory [11, natural language tends to have more bounded levels of embedding (30 10]. This has motivated the study of bounded-depth Dyck languages as plausible
simple models of the hierarchical structure underlying language (25 65 6 62]
Task. Formally, Dyck- (n, k) (cf. Appendix E is the language of well-matched strings over n types of parenthesis pairs (1, )1, (2. )2,.................... (n. )n. where any prefix has at most & opening parentheses not yet closed. For instance, the string '( ()" has a maximum depth 2 corresponding to the prefix '( "['- Dyck-(n, k) can be recognized with access to bounded stack that never holds more than lc elements. In fact, each Dyck- (n, k) is a regular language and is accepted by a finite automaton.
We show that there is a linear communication complexity lower bound for Dyck-(n, k), already at n = k ===2 2. However, unlike the communication bound we used in Theorem 3 Alice and Bob now have access not to two halves of the input: rather, Alice and Bob have access to the even and odd positions in the string, respectively. Intuitively, in such situation, Alice needs to know almost all of the bits available to Bob in order to decide whether given string is well-bracketed-ane vice versa. More formally, they need to exchange at least - bits to decide membership in Dyck-(2, 2):
Lemma 1. Suppose Alice and Bob have the symbols in the odd and even indices of a string & E ΣN respectively. To each compute whether Dyck-(2, 2). they must exchange at least N bits.
The proof of Lemma iss in Appendix B and is based on fooling sets which is standard technique to prove lower bounds on communication complexity. Combining Lemmas 4 and 1 bound on the width of Transformer for computing Dyck-(2, 2):
entails a lower
Theorem 5. Consider a one-layer Transformer f TF! TF',.p.11 deciding membership in Dyck-(2, 2).
N-1 Then mll N).
This result establishes a second separation between one-layer Transformers and RNNs, but now in the other direction: Bounded-depth Dyck languages are regular, and previous works have shown that constant width RNNs can recognize them with standard activation functions such as Sigmoid (25) and ReLU [6]. We further note that two-layer Transformers of sublinear size can model bounded-depth Dyck languages (65)
4.2 Lower Bounds on Boolean Functions
There are some notable differences between the types of communication complexity lower bounds for one-layer Transformers (Thcorem 9 and for RNNs (Theorem (2) RNNs computing f yield one-way protocol for contiguous partitions of the input; thus showing a one way communication lower bound for such partitions is sufficient to obtain lower bounds on the size of RNNs. Transformers computing f yield a multi-way protocol that work for arbitrary partitions of the input; thus showing lower bound for any partition is sufficient to establish lower bounds on the size of the Transformer. For (2, 2). contiguous partitions are not a hard case, and in fact, communicating < 2 open brackets from the first half is sufficient. This is why the lower bound of Lemma doos not apply RNNs.
Lower Bounds. Despite the differences discussed above, there are several Boolean functions, for which we can establish that when widths are bounded, neither one-layer Transformers nor RNNs

---

can compute them. The Equality function EQ 10 + {0.1] is a Boolean function defined as EQ(x)
(XN/2+1..... XN)]. A related problem is Disjointness: given two vectors x, y E (0. 10,1) the function y) máx x,Y; I[x*s y > 0]. Both the functions Equality and Disjointness are known to have communication complexity (2N) (see Appendix [B.I) and Theorems 4 and 2 imply that both one-layer Transformers and RNNs must have width {TN) to represent them. In the next section, we show that these lower bounds do not apply to two-layer Transformers. Additionally, it is worth noting that the functions EQ and DISJ can also be expressed in the form of 2-CNFs with (NV) terms. Hence, more general consequence of the limitations of RNNs (and one-layer Transformers) is that with width (NV), they cannot compute certain functions in the class of uniform AC It is interesting since the class of uniform AC is considered one of the simplest classes of Boolean circuits and even the expressive power of hard-attention Transformers has been shown to be within this class (24). In Appendix [F.] we provide an alternate proof of the lower bound for RNNs computing Equality based on their relation to DFAs.
5 Representational Capabilities of 2-layer Transformers
In Section 4 we showed that single-layer Transformers and recurrent models must have size linear in the input length express natural Boolean functions such as EQ. In this section, we show that two-layer transformers overcome these limitations by efficiently expressing such Boolean functions and more general forms of associative recall tasks, such as simulating the nearest neighbor algorithm.
5.1 Representing Boolean Functions
We start by showing that two-layer Transformers of poly-logarithmic size can express the Equality function (proof is in Appendix (.2). The input domain need not necessarily be the Boolean vectors {0. 13,; rather, the construction works for sequences over any finite alphabet E.
Theorem 6. For any N N, there exists 2-layer Transformer f E TF² T..pp. where width m Ollog N) and precision p = O(log N) such that fxx) = EQ(x) for all x E {0.1)
The construction is based on tools developed in Section 3
The broad idea is as follows. In the N/2, an attention head attends to position N/2 and copies the input N/22 A feedforward network then checks whether the retrieved value is equal to z;. The second layer simply uses uniform attention over the outputs of the previous layer to check if there is a mismatch at any position. Importantly, we show that the above strategy can be implemented with a representation size O(log )²)
first layer, at each position
Generalizing this result, we find that two-layer Transformers with logarithmic width can express more general class of Boolean functions: thresholds of k-sparse features, class including func- tions such as Equality and Disjointness. Since such functions cannot be expressed by one-layer Transformers and recurrent models with width o(N), these results imply a separation on Equality and Disjointness: these functions can be expressed by small-sized two-layer Transformers whereas one-layer Transformers and recurrent models must grow linearly with input length represent them.
5.2 Implementing the Nearest Neighbors Algorithm
The goal of the nearest neighbor task (NsTNB) is to analyze whether sequence modeling architecture can implement the well- kell-known nearest neighbor algorithm to make predictions. Our description follows closely to the one used by Bhattamishra et al. [8] for their experiments.
Nearest Neighbors. In the NSTNB task, a model is provided with a sequence of vectors and labels (x1, ........ X&-1; 1, where N/2 h < N, the input unit vectors x; E IRd and labels E 13. For each Xx where k > N/2, the output is the label corresponding to the nearest neighbor in (xi,..., 1), that is, if = arg max,e[k-1] x x; or j = arg min; 1] ||Xe - -x,llz then the output for x is the label y/- Since we are working with unit vectors, maximizing the inner product is equivalent to minimizing the &z distance. If the second half of the sequence N , Xx is permutation of the first half xi,..., Xy then the task reduces to the Multi-Query Associative Recall (MQAR) task 2 (cf. Appendix
'The class A contains polynomial size AND/OR circuits with unbounded fan-in and constant depth.
7

---

Assumptions. We will make two assumptions about the problem. The first assumption is that all input vectors are of unit norm, i... |اا|| 1 and the second is the existence of a margin between the dot product with the nearest neighbor and the dot product with other input vectors, i.e. there exists 7 2 N for some universal constant e, such that for any N/2 k S N, j arg m]j- X X5, then x X, x x; + for any i j
The following is one of our main results which states that two-layer Transformers of logarithmic size can implement the nearest neighbor algorithm in their forward pass and as a corollary can also perform associative recall tasks like MQAR (Proofs in Appendix G.2).
Theorem For any N N. there exists a 2-layer Transformer fNN TF T...2 with width m = Olog N) and precision p- (log N) such that fNN computes the nearest-neighhor task all sequences of length at most N satisfying the assumptions above.
The broad idea of the construction is to identify the nearest neighbor input x,. and retrieve the position of the corresponding label yy- in the first layer. The second layer then uses this positional information to retrieve the desired label. There are few challenges implementing this strategy which we address in our construction. First, note that for input vectors X1, Xi,..., Xg. naively using them with dot-product attention will result in the query input x; having maximum dot product and hence maximum attention weight over itself. Second, the dot product with some label vectors y:s could be higher than the dot product with the nearest neighbor Xz-. Third, the positional information must be retrieved using soft-attention in a way that it can be used in the next layer to obtain the desired label. Our intuitive, though somewhat involved, construction deals with these issues to ensure that a two- two-layer Transformer with Ollog 2) total size implements the nearest neighbor algorithm.
Theorem 8. Any recurrent model with a hidden state of width m with p-bits of precision that can perform the nearest neighbor task for all inputs of length N must have m 2 N/2p..
The lower bound for recurrent models follows via a reduction from the Disjointness problem.
Discussion. Prior works S 2 have empirically demonstrated that Transformer-based LLMs can exhibit mechanisms such as nearest neighbors and MQAR. Further, on synthetic setups, they have observed that recurrent models struggle to perform these tasks compared to Transformers. Our results take step towards understanding the differences in the performance between the two architectures.
6 Empirical Analysis
While we focus on the differences in the representational capabilities of Transformers and recurrent models, it is natural to examine if differences of a similar nature arise in their empirical performance. One thing to keep in mind is that positive results regarding expressiveness presented in earlier sections do not imply that models can learn such tasks. With regard to negative results, they do imply that when the sequence length is much larger than the size of the hidden state or width of the model, then the model will be incapable of representing the task and consequently fail to learn the task. However, even for one-layer recurrent models with hidden states of size 128 with 64 bits of precision, our lower bound applies at lengths over 8k.
In this section, investigate the performance of Transformers and recurrent models tasks such as Index Lookup and recognizing bounded Dyck languages on sequences of small lengths (< 1000). Our experiments
designed to answer the following questions: (1) Are one-layer Transformers better than larger recurrent models on the Index Lookup task? (2) Are recurrent models and two- layer Transformers better than one-layer Transformers recognizing the Dyck-(2, 2) language? Importantly, as our results concern the scaling of the model size with the input length, we are specifically interested in the behavior of different models across input lengths.
We also explore the performance of models on string equality in Appendix H.2 Tasks like NSTNB and MQAR have already been analyzed empirically in prior works [
we do not include them.
Setup and Training details. We train the models with cross-entropy loss using the Adam optimizer [32. The models are trained for up to 250k steps where at each step we sample a fresh batch of 64 training examples- - resulting in R 16 million examples over 250k steps. The models are evaluated on 5000 examples for each task. For each model, we tune the various hyperparameters, notably across learning rates E (1e-2, Se-3. 1e-6) to find the best-per model. The details of the data generation method, hyperparameters, and implementation details can be found in Appendix H
8

---

Index
Lookup Task
Valiéution Curves lungth 200
Validation Carves of length 400
LSTMGG 2364
100
28:2
2.8
TT-1, 641
LSTH43, 256
040 Acconne 1961
Archtctuue
Aurivee 36, 2514
0.5
000 2040
L1rfF46, 251
0.6
0.4
TE-D 641
Mamba- 2361
L00.0
06.2
283
Mombe 2581
1943 454
108:0
1000
101.2
115-4
144
02
330
Inpul Langth
Tarre ons
Ilerations
Figure 2: Performance of models on the Index Lookup and bounded Dyck task. Labels such as TF-(I, 64) denote Transformers with layer and 64 widths. See Section for more details.
Index Lookup Task. We compare the performance of one-layer Transformers with five different recurrent models - LSTMs (26), state space models such as DSS [22] and Mamba [21] linear Transformers [31, and its variant RetNet [56). We explore the performance across various lengths E {20, 50, 100, 200, 400). We evaluate relatively small-sized Transformers with widths 64 against recurrent models with up to layers and widths or hidden state size of 256. The size of the alphabet in the experiments is | 64. Figure 2 (left) depicts the performance of all models across various lengths and Figure 2 (middle) depicts the validation curves during training on examples of length 200. As depicted by the figures, while one-layer Transformers with width 64 achieve near-perfect accuracy within a few thousand steps, the performance of relatively larger recurrent or state-space models degrades on lengths over 100 and they fail to learn even with 10x training iterations. We explore the influence of width on the performance of Mamba in Appendix H..I
Bounded Dycks. For 2 with depth at most 2, our separation results apply to one-layer Transformers and nonlinear recurrent models such as LSTMs but not to linear RNNs such as state- space models and linear Transformers. Hence, we are primarily interested in the difference in performance between one-layer Transformers and LSTMs In our experiments, we compare the performance of one-layer Transformers with relatively smaller recurrent models such as LSTMs and two-layer Transformers. We also include Mamba for reference. We consider LSTMs and Mamba with hidden state sizes of 64 and similarly, two-layer Transformers with width 64. We evaluate one-layer Transformer with a width of 256 across lengths E {20..... 400} most of which are smaller than the width of the model. We observe that one-layer Transformers achieve near-perfect accuracy up to lengths 100 but struggle on higher lengths. In contrast, small-sized recurrent models as well as two-layer Transformers can achieve near-perfect accuracy for lengths up to 400. Figure 2 (right) depicts the validation curve of the models on examples of length 400.
7 Discussion and Final Remarks
Based on prior theoretical results, it is known that, while recurrent models can express any regular language 34 Transformers with logarithmic precision can only express languages in the class of uniform constant depth threshold circuits (TC') 44]. These results indicate conjectures- Transformers are unable to represent certain state- trackeng tasks that recurrent models can represent. With such results, it might appear that Transformers are less expressive than recurrent models-potentiallyy at odds with the persistent practical success of Transforrmer-based LLMs. Our findings, however, show that when the model size is constrained relative to the sequence length, a varicty of tasks relevant to practice can be represented by small-sized Transformers but not by recurrent models. Our results suggest that the attention mechanism does lead to expressiveness that cannot be replicated by recurrent architectures even with arbitrary transition functions.
standard
Limitations. A general limitation of this line of work is that positive expressivity results do not imply that the problems under consideration are learnable. Additionally, while lower bounds for an architecture imply difficulty in learning, when using double precision these results only apply to very long sequences practice. Our results (and probably techniques) do not imply any limitations on two-layer Transformers; this is left as an open question. We note that communication complexity- based techniques akin to Theorem cannot exist for two-layer Transformers (cf. Appendix (F.4). Hence, we believe that other tools will be needed to prove lower bounds for two-layer Transformers.

---

Acknowledgements
We thank Dibyayoti Dhananjay Jena, Arkil Patel, and Charles London for insightful and helpful discussions on this work. We also thank Clayton Sanford, Will Merrill, and the anonymous reviewers for their valuable feedback and constructive suggestions.
References
[1] D. Achlioptas. Database-friendly random projections: Johnson-lindenstraus with binary coins. Journal of computer and System Sciences, 66(4):677 1-687, 2003.
[2] S. Arora, S. Eyuboglu, A. Timalsina, I. Johnson, M. Poli, J. Zou, A. Rudra, and C. Re. Zoology: Measuring and improving recall in efficient language models. In The Twelfth International Conference on Learning Representations, 2024. URL https://openrevieee net/forum? id=LY3ukUANko
[3] J. Ba, G. E. Hinton, Mnih, J. Z. Leibo. and C. Ionescu. Using fast weights to attend to the recent past, Advances in neural information processing systems, 29, 2016.
[4] Y. Bai, Chen, H. Wang, C. Xiong, and Mei. Transformers as statisticians: Provable in- context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
[5] S. Bhattamishra, Ahuja, and N. Goyal. On the Ability and Limitations of Transformers to Recognize Formal Languages. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 7096-7116, Online, Nov. 2020. Association for Computational Linguistics. doi: 10. 8653/v 12000eemmmmmmmmin..5.... URL https://
ac aclanthology. org/2020. emn emnlemmmmmmm www...maa.....mmm............... 576
[6] S. Bhattamishra, K. Ahuja, and N. Goyal. On the practical ability of recurrent neural networks to recognize hierarchical languages. In Proceedings of the 28th International Conference on Computational Linguistics, pages 1481-1494, Barcelona, Spain (Online). Dec. 2020. Interna- tional Committee on Computational Linguistics. doi: 11116663/112222.cooinngg main. 129. URL https://accaathologgg org/ 2020. cowing main. 129
[7] S. Bhatttamisha,, A. Patel, and N. Goyal. On the computational power of transformers and its implications in sequence modeling. In R. Fernández and T. Linzen, editors, Proceedings of the 24th Conference on Computational Natural Language Learning, pages 455-475. Online, Nov. 2020. Association for Computational Linguistics. doi: 10.18653/112020.comlll https://aclantholloy org/ 2020/202..co0 wo1.0...o11-..3.7.7..
URL
[8] S. Bhattamishra, A. Patel, P. Blunsom, and V. Kanade. Understanding in-context learning in transformers and LLMs by learning learn discrete functions. In The Twelfth International Conference on Learning Representations, 2024. URL htttps://opeerevieee wwwprreviiiieweeffrreeeeee id=ekeyCgeRfC
[9] Blanchard, D. J. Higham, and N. J. Higham. Accurately computing the log-sum-exp and softmax functions. IMA Journal of Numerical Analysis, (1(4):2311-2330 2021.
[10] D. Blasi, R. Cotterell, L. Wolf-Sonkin, S. Stoll, B. Bickel, and M. Baroni. On the distribution of deep clausal embeddings: A large cross-linguistic study. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 3938-3943 2019.
[11] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, Agarwal, A. Merbert-Voss, G. Knueger, Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, Clark, C. Berner, S. McCandlish, A. Radford, Sutskever, and D. Amodei. Lan- guage models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, pages 1877- 1901. Curran Associates, Inc., 2020. URL http:httts///occediiiiii iew.eeerrdingn.ee.e...........pe htwpggg.errrip...ppppp//// 1e00//ile11557c0eeee 4ww.a//0e00000444410888811112222
pdff
[12] D. Chiang. P. Cholak, and Pillay. Tighter bounds on the expressivity transformer encoders. In International Conference Machine Learning, pages 5544-5562, PMLR, 2023,
[13] N. Chomsky, Syntactic structures. Mouton, The Hague, 1957.
10

---

[14] N. Chomsky and M. P. Schützenberger. The algebraic theory of context-free languages. In Studies in Logic and the Foundations of Mathematics, volume 35, pages 118-161, Elsevier, 1963.
[15] S. De, S. L. Smith, A. Fernando, A. Botev, G. Cristian-Muranu A. Gu, R. Haroun, L. Berrada, Y. Chen, S. Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402. 19427, 2024.
[16] G. Deletang. A. Ruoss, J. Grau-Moya, T. Genewein, L.. K. Wenliang, E. Catt, C. Cundy, M. Hutter, S. Legg. J. Veness, and P. A. Ortega. Neural networks and the chomsky hierarchy. In The Eleventh International Conference Learning Representations, 2023. URL https
/ //opnrevieww //penrevree.nnt//ffrrmmm dibbbuikkeddb
[17] J. Ebrahimi, D. Gelda, and W. Zhang. How can self-attention networks recognize Dyck-n languages? In T. Cohn, Y. He, and Y. Liu, editors, Findings of the Association for Computational Linguistics: EMNLP 2020, pages 4301-4306, Online. Nov. 2020. Association for Computational Linguisties. doi: 10.18653/v1/20200
http:///.0.....insseeeeeeo Andings-emnlp.335 URL /tLLatt////aalaabollly org/
2020. findings-----ii emn emnlp..m.m 384
[18] M. B. Everaert, M. A. Huybregts, N. Chomsky, R. C. Berwick, and J. Bolhuis. Structures, not strings: linguisties as part of the cognitive sciences. Trends in cognitive sciences, 19(12):
729-743, 2015.
[19] S. Garg. D. Tsipras, P. S. Liang, and G. Valiant. What can transformers learn in-context? case study of simple function classes. Advances in Neural Information Processing Systems, 35: 30583-30598, 2022.
[20] A. Gers and E. Schmidhuber. LSTM recurrent networks learn simple context-free and context-sensitive languages. IEEE Transactions on Neural Networks, 12(6):1333-1340, 2001.
[21] A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint 1XXi:22312.00552 2023.
[22] A. Gupta, A. Gu, and Berant. Diagonal state spaces are as effective as structured state spaces. Advances in Neural Information Processing Systems, 33:22982-229991 2022.
[23] M. Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Lingusstics 8:156-171 171, 2020.
[24] Y. Hao, D. Angluin, and R. Frank. Formal language recognition by hard attention transformers: Perspectives from circuit complexity. Transactions of the Association for Computational Linguistics, 10:800-8 8100008000000000 2022.
[25] J. Hewitt, M. Hahn, S. Ganguli. P. Liang, and C. D. Manning. RNNs can generate bounded hierarchical languages with optimal memory. In B. Webber, T. Cohn, Y. He, and Y. Liu, editors, Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 1978-2010, Online, Nov. 2020, Association for Computational Linguistics. doi: 10.18653/, 1tt333/11222.eom/ 11::00.mmllp-mait 156.. URL https://aclaatthollsst org/ 2020. emnlp-main. 156. [26] S. Hochreiter and Schmidhuber. Long short-term memory. Neural computation, 9()):
1735-1780, 1997.
[27] T. S. Jayram, R. Kumar, and D. Sivakumar. The one-way communication complexity of hamming distance. Theory of Computing, 4(1):129-133, 2008
[28] S. Jelassi, D. Brandfonbrener, S. M. Kakade, and E. Malach. Repeat after me: Transformers are better than state space models at copying. arXiv preprint arXiv:2402.01032 2024.
[29] W. B. Johnson and J. Lindenstrauss. Extensions of lipschitz mappings into hilbert space. Contemporary mathematics, 26: 18:1099 1984. URL http://appp.eemnniccchhh...a.a
org/CorpusID::112 117819116
[30] Karlsson. Constraints on multiple center-embedding of clauses. Journal of Linguistics, 43(2): 365-392, 2007.
[31] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are rnss: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5165. PMLR, 2020.

---

[32] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[33] J. F. Kolen and S. C. Kremer. field guide dynamical recurrent networks. John Wiley & Sons, 2001.
[34] S. A. Korsky and R. C. Berwick. On the computational power of mnn. arXiv preprint arXiv:1906.0688 1906.06406 2019.
[35] E. Kushilevitz and N. Nisan. Communication Complexity. Cambridge University Press, 1996.
[36] D. Lindner, J. Kramár, S. Farquhar, M. Rahtz, T. McGrath, and V. Mikulik. Tracr: Compiled transformers a laboratory for interpretability. Advances in Neural Information Processing Systems, 36, 2024.
[37] B. Liu, T. Ash, S. Goel, A. Krishnamurthy. and Zhang. Transformers learn shorteuts to automata. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreet
[38] B. Liu, J. Ash, S. Goel, A. Krishnamurthy, and C. Zhang. Exposing attention glitches with flip-flop language modeling. Advances in Neural Information Processing Systems, 36, 2024.
[39] W. Merrill and A. Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 11:531-544 2023.
[40] W. Merrill and A. Sabharwal. A logic for expressing log- cog-precision transformers. Advances in Neural Information Processing Systems, 36, 2024.
[41] G. A. Miller and N. Chomsky. Finitary models of language users. In R. D. Luce, R. R. Bush, and E. Galanter, editors, Handbook of Mathematical Psychology, pages 1963.
John Wiley,
[42] A. Orvieto, S. L. Smith, A. Gu, A. Fernando, C. Gulcehre, R. Pascanu, and S. De. Resurrecting recurrent neural networks for long sequences. International Conference on Machine Leurning, pages 16670-26698 PMLR, 2023.
[43] R. Pascanu, T. Mikolov, and Y. Bengio. On the difficulty of training recurrent neural networks.
In International conference on machine learning, pages 1310-1318. Pmlr, 2013.
[44] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[45] B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, H. Cao, X. Cheng, M. Chung, M. Grella, K. K. GVV, et al. Rwkv: Reinventing rnns for the transformer era. arXiv preprint arXiv:2305 13048, 2023.
[46] B. Peng, S. Narayanan, and C. Papadimitriou. On limitations of the transformer architecture, 2024.
[47] J. Pérez, J. Marinković, and P. Barceló. On the turing completeness of modern neural network architectures. In International Conference on Learning Representations, 2019. URL https //openreview net/forumiid--HyBBBaa
[48] A. Rao and A. Yehudayoff. Communication Complexity: and Applications. Cambridge University Press, 2020.
[49] J.-F. Raymond, P. Tesson, and D. Thérien. An algebraic approach to communication complexity. In Automata, Languages and Programming: 25th International Colloquium, ICAI.P'98 Aalborg, Denmark, July 13-17, 1998 Proceedings 25, pages 29-40. Springer, 1998.
[50] P. Rodriguez. Simple recurrent networks learn context-free and context-sensitive languages by counting. Neural computation, (3(9):2093-22 2818 2001.
[51] C. Sanford, D. Hsu, and M. Telgarsky. Representational strengths and limitations of transformers.
In seirth- Conference on Neural Information Processing Systems, 2023. URL https
//openreview. wwwparrwwwieenee/oouummmmeeeeeee
[52] C. Sanford, D. Hsu, and M. Telgarsky. Transformers, parallel computation, and logarithmic
depth. arXiv preprint 24X.224022.09668. 2024.
[53] D. Sivakumar. Algorithmic derandomization via complexity theory. In Proceedings of the foury- annual ACM symposium on Theory of computing, pages 619-626, 2002.
12

---

[54] N. Skachkova, T. A. Trost, and D. Klakow. Closing brackets with recurrent neural networks. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP,, pages 232-239, 2018.
[55] L. Strobl, W. Merrill, G. Weiss, D. Chiang, and D. Angluin. What formal languages can transformers express? a survey. Transactions of the Association for Computational Linguistics, 12:543-561, 2024.
[56] Y. Sun, L. Dong, S. Huang, S. Ma, Y. Xia, J. Xue, J. Wang, and F. Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307 2011::000021 2023.
[57] M. Suzgun, Y. Belinkov, S. Shieher, and S. Gehrmann. LSTM networks can perform dynamic counting. In Proceedings of the Workshop on Deep Learning and Formal Languages: Building Bridges, pages 44-54, Florence, Aug. 2019. Association for Computational Linguistics. doi: 10.18653/v1/W
URL (tthhtt::: www. www.waclwebb...... orggaanhhologg//19---- 3905
[58] M. Suzgun, S. Gehrmann, Y. Belinkov, and S. M. Shieber. Memory-augmeented recurrent neural networks can learn generalized dyck languages. arXiv preprint arXiv: 1,33399, 2019.
[59] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, t.. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-50000 0003-600,222 2017.
[60] J. Von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174, PMLR, 2023.
[61] G. Weiss, Y. Goldberg, and Yahav. Thinking like transformers. In International Conference on Machine Learning, pages 11080-11900, PMLR, 2021.
[62] K. Wen, Y. Li, B. Liu, and A. Risteski. Transformers are uninterpretable with myopic methods: a case study with bounded dyck grammars. Advances in Neural Information Processing Systems, 36, 2024.
[63] T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, al. Transformers: State-of-the taae--of-theaart art natural language processing. In Proceed- ings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pages 38-45, 2020.
[64] A. C.-C. Yao. Some complexity questions related to distributive computing (preliminary report). Proceedings of the eleventh annual ACM symposium on Theory of computing, pages 209-213, 1979.
[65] S. Yao, B. Peng, C. Papadimitriou, and K. Narasimhan. self-attention networks can process bounded hierarchical languages. C. Zong, F. Xia, W. Li, and R. Navigli, editors, Proceedings of the 59th Anmual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Lnguage Processing (Volume Long Pupers), pages 3770-3785, Online, Aug. 2021. Association for Computational Linguistics. doi: 10. 18653/v1/2021.acll acl--........ URL https:tt::/// //acl htc///llantttt awl.nthologgyy..oooooooooooooooe org/ 2021 wwl..acl-....n........ammmooall. 292
[66] X. Yu, N. T. Vu, and J. Kuhn. Learning the dyck language with attention-based seq2seq models. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pages 138-146, 2019.
[67] C. Yun, S. Bhojanapalli, A. S. Rawat, S. Reddi, and S. Kumar. Are transformers universal approximators of Representations, 2020. https:::::///
functions? In International Conference on Learning
mettfformmidd-a ByRRott
[68] H. Zhou, A. Bradley, E. Littwin, N. Razin, O. Saremi, J. Susskind, S. Bengio, and P. Nakkiran. What algorithms can transformers learn? a study in length generalization. arXiv preprint arXiv: 160. 2023.
13

---

Contents
1
Introduction
Related Work
2
Definitions
3
Index Lookup Task
4
Lower Bounds for RNNs and 1-layer Transformers
Separation on Bounded Hierarchical Languages
4.2
Lower Bounds on Boolean Functions
5
Representational Capabilities of 2-layer Transformers
5.1
Representing Boolean Functions
5.2
Implementing the Nearest Neighbors Algorithm
Empirical Analysis
Discussion and Final Remarks
Clarifications
14
Preliminaries
15
B.1
Communication Complexity
15
B.2
Finite Precision Implementation
16
B.
Technical Tools
16
C
Index Lookup Task
17
C.I
Recurrent models must be wide to perform Index Lookup
17
1-layer Transformer with small width can perform Index Lookup
18
D
Lower Bounds for 1-layer Transformers
19
Dyck with Bounded Depths
21
Transformers and Boolean functions
23
F.
Lower Bounds
24
F2
Transformer Construction for Equality
25
F3
Representing more general class Boolean functions
27
F4
Difliculty of Deriving Communication--an based Lower Bounds for 2-layer Transtormers
28
Nearest Neighbors and Associative Recall
29
G.1 Lower Bounds for Recurrent Modelsi Transformer Construction for Nearest Neighbor
30
G.2
30
Empirical Analysis: Additional Details and Experiments
34
H.I Additional Experiments and Data Generation
34
H.2 String Equality Task].
A
Clarifications
(1) What do the results presented in the paper imply about the learnability of the tasks considered?
The lower bounds on recurrent models and one-layer Transformers have negative implications for learnability but the positive results do not have any nontrivial implications. The lower bounds on size of recurrent models and one-layer Transformers imply that unless the width of the models grows linearly with respect to the input length, they cannot represent and consequently cannot learn such tasks. Note, however, that even though the input length need not be unboundedly long for our lower bounds to apply, they still need to be sufficiently large N 2 and do not apply the lengths
14

---

considered in our experiments. The results indicating that Transformers can express tasks like Index lookup, string equality or nearest neighbors do not imply that they can learn those tasks in practice. The experiments in Sections and H as well as from prior works g seem to indicate that small-sized Transformers perform reasonably well on these tasks.
(2) For the upper bounds on the total size of Transformers to express functions, does it include the imput and positional embeddings or just the Transformer model with attention and feedforward block?
Yes, when we say that Transformers with total size (()(N) or poly-logarithmic in N can express a certain task, it includes the input and positional embeddings. The total representation size indicates the total number of bits required to represent all the parameters of the model. Our results imply that a Transformer with total size Ollog N) can represent tasks such as nearest neighbors. Equality, etc, whereas the size of any recurrent model must be ((N). For our constructions of Transformers, we ensure that the positional embeddings can be generated in log space (discussed in Appendix (B) and need not be stored as a N x d matrix for the required computations. Lastly, it is worth noting that all our lower bounds for recurrent models apply even if they have arbitrary positional embeddings stored in N x d matrix, and hence the size of recurrent models is (2(N) excluding the embeddings.
(3) For the constructions with Transformers, why can't 't those results follow from some frameworks like RASP?
is programming language aiming to abstract computations that transformers can compute. While one might able to construct RASP programs for some of the tasks we have considered, such as Index Lookup, such constructions would not entail results similar to ours, because RASP substantially abstracts from implementational aspects. For instance, RASP allows MLPs to compute arbitrary functions, and attention is not computed from dot products of keys and queries. is not clear if general-purpose translations from RASP to realistic transformers would be able to recover our efficient bounds, e.g., logarithmic size.
RASP 61 36
(4) Do the negative results experiments imply that those architectures cannot learn those tasks?
The differences in the performance of models apply more to the rate of learning than a binary form of success/failure. For instance, in the Index Lookup task, one can see that in Figure 2 one-layer Transformers achieve near-perfect accuracy in a few thousand steps whereas recurrent models fail to achieve high accuracy even after being trained for 25x steps. It can still be true that if the models trained for much longer or are much bigger, they might achieve high accuracy. See Figure which depicts the performance of Mamba across various sizes and lengths. For tasks like Index lookup, we observe that Transformers achieve near-perfect accuracy across several learning rates whereas recurrent architectures fail all the learning rates we tried. A similar thing is true for other tasks such as Dyck-(2, (2, 2) where one-layer Transformers seem to learn at a much slower rate compared to LSTMs and two-layer Transformers even for the lengths where they do achieve near cear-perfect accuracy.
(5) Do the lower bounds for Transformers and RNNs apply with additional components such as layer norm, residual connections, positional embeddings, etc?
Yes, for our lower bounds for lne-layer Transformers and Recurrent models, they still apply if the models have additional components including arbitrary positional embeddings, layer norms, etc. We only assume that the attention mechanism operates over finite precision and the hidden state of RNN is in finite precision. The results do not make any assumptions about the computational limits of the remaining components. They apply even when the output of the attention block or the hidden state of an RNN is processed by an arbitrary function.
B Preliminaries
B.1
Communication Complexity
Our lower bounds for recurrent models and 1 I-layer Transformers are based on communication- complexity bounds. We assume some familiarity with communication complexity (see Rao and Yehudayoff [48), Kushilevitz and Nisan [35] for an introduction). We will primarily focus on the two-party setting where the communication complexity of a function indicates the number of bits two parties must exchange in order to compute the output of a function. If Alice and Bob want to compute a function f =00, 1 + (0. where Alice has the bits in the indices / C [N] and Bob has the bits in indices [N] I, the communication complexity of f over that partition is the minimum number
15

---

of bits they must exchange to compute the output for all inputs r E {0,1]% Alice and Bob are allowed unbounded computational power. If Alice and Bob must exchange least & bits compute function f(x:) Vrr E {0, 1 1 over any partition of the input then we say that the communication complexity C(J) 2 k. If they use a communication protocol where only one party is allowed to send bits to the other party, then it is called one-way communication, and the communication complexity of the function in that setting is referred to as one-way communication complexity.
For our results, we will use the following three well-known facts about the communication complexity of the Disjointness, Equality, and the INDEX problem.
Disjointness. The disjointness function takes two sets B, C [N] as input and returns 0 the two sets are disjoint and returns 1 otherwise. This can also be seen (0. 1) x (0. 1 > {0, } over two Boolean inputs such that DIsI(a, b) max agh4 = T[a" b > 0j. If Alice has the input vector complexity of DISJ indicates the minimum number of bits that Alice and Bob will have to exchange to
function DISJ
{0, 1) and Bob has the vector b {0, 1 then the communication
determine the output of the DISJ function. The following is a well-known fact about the disjointness problem,
Fact (Disjointness (641) If Alice and Bob have two inputs a, E {0, 10,1)% then any deterministic communication protocol used by them to compute DisI(a, b) = max a,b; must exchange at least N-bits. Moreover, the randomized communication complexity of the DISJ is (2
Fact 2. (Equality 148 Ch. 1]) If Alice and Bob have two inputs b E {0, 1],, then any deterministic communication protocol used by them to compute EQ(a, b) = Iaa = b] must exchange at least N -bits.
Fact 3. (INDEX (27)) If Alice and Bob have two inpuls any deterministic communication protocol used by Alice to send bits to Bob must require N bits for Bob compute INDEX(a, b) 0b- Moreover, the one-way randomized communication complexity of the INDEX is (2NN
{0, 1) and b [N] respectively, then
One thing to note about the Equality problem is that although the deterministic communication complexity of the EQ problem is Q(N), the randomized communication complexity O(og N). For the Disjointness and INDEX problems, the randomized communication complexity is (?(N) as well. In other words, even if the two parties are allowed to compute the output of the function correctly with high probability (say² 2/3), even then the number of bits they must exchange is ((N).
B.2 Finite Precision Implementation
In this work, we are interested in the expressiveness of finite precision models. For our constructions with Transformers, we will work with p-bit numbers where p = O(log N) where N is the maximum length of the input string.
In particular, for some sufficiently large constant Kg > O, we will work with numbers between -NK, to N with a step size of A = N... If Qp is the set of all such numbers then Op {-NK: -v +
A.... 0, A, 2A, ...., NK).. Hence, the size of the set is IQ, 2 p= = 2K, log N = O(log N). For any real number 2 E IR, in the finite precision implementation, a is rounded down to the nearest & such that &N E Σ. If 2 > N.. then z is rounded down to N,. In our constructions with Transformers, all the parameters and intermediate values follow the above implementation.
B.3 Technical Tools
For our constructions of Transformers, we will use the following result about high dimensional vectors. The statement essentially says that at a high dimension & the number of vectors that are almost orthogonal to each other is exponential in ie even though the number of orthogonal vectors can be at most k.
Lemma 2. For any N, there exists N k-dimensional vectors Ti...., TN where = O log N)
and each entry of the vectors is in {- VRa 3 such that
2 1-7 if==
(Tz, T
<>
otherwise.
16

---

It is quite straightforward to see why this is true. It follows from simple application of the probabilistic method. Suppose, one samples N k-dimensional vectors X _...... XN independently at random such that each entry of each vector is drawn uniformly from vi VI ). Hence, each vector
X; 6 - JE and the expectation of the dot product of any two vectors E[(X;; X;)l = 0 for all i + j. Since the dot products of any two vectors (Xi, Xj) is bounded random variable, can apply Hoeffding's inequality to get
kn
PIKX,, X,MI 2 & eexp
4
Taking union bound over at most 2 pairs of dot products and setting 2N exp (-4-2)
< 1/2, we
get that for k _ log 2N,, the probability of all dot products (X,; X3) being less than y the same time is at least 1/2. Since the probability of the event is nonzero, follows that the statement in the lemma is true.
In construction, will use such vectors positional embedding vectors. While Lemma implies the existence of vectors, storing N such vectors will require O(N) space. We would like to generate i-th vector Ti when necessary in polynomial time using log space without storing all of them together. To achieve that, we use a derandomization of the Johnson-Lindenstraan (JL) Lemma 29 that can output the i-th vector in log-space and polynomial time 5
We use a slightly modified version of the JL lemma which preserves inner products over unit vectors.
Lemma 3. (Inner product preservation (1)) Let e E (0, 1/2) and let Q S be a set of N unit norm vectors of dimension d. For k
os log N tthree exists a linear map T(x) = vE Ax where each ²²
entry of TR.d R is in [-1,1] such that for all u, E Q
(u, v) - (T)),, T(w))
The above result has interesting implications which we will use in our constructions. Let X1...., Xy E EN be N unit norm vectors that form the basis of R which implies (x,, x,) = if
==j and is 0 otherwise. Then there exists a map T such that the vectors T(x))...... T(xN) R
have dot products (xx,)).(xx,)
c if = j and 0 + otherwise.
We will use JL transformations of the standard basis vectors CI...., ON where T(1),... T(N) will
refer k Olog N) dimensional vectors such that their inner product is z with themselves and is zs O. Intuitively, we can use such vectors to get Transformer to attend to unique positions.
Corollary 8.1. For any N, there exists N k-dimensional vectors T(1)...., T(N) where = Ollog N) and each entry of the vectors is in { Vk V such that
2 3/4 ifi=
(T(I,T()
< 1/4 otherwise.
C Index Lookup Task
Index Lookup. The index lookup task (IdxL) is multi-class classification task where model receives a sequence of tokens I1...., followed by a position token where p E [N] and the goal of the model is to output the token rp at position p. Here the symbols z; belong to vocabulary Σ, finite set of symbols. The sequence I = (#), .., N) can have repetitions. More precisely, we say model computes the function IdxL EeN [N] E if for all inputs (x,p), the model outputs (x.L(r.p).
C.1 Recurrent models must be wide to perform Index Lookup
Theorem Any recurrent model with hidden state of width m using p-bits of precision that computes the Index Lookup task all sequences of length N must have m 2 N N/p.
17

---

Proof. The proof follows naturally via a reduction from the INDEX problem in communication complexity. Assume the vocabulary size for the IdxL is at least 2. pick any two symbols from the vocabulary, and map them to 0 and 1. Suppose Alice has a sequence a E {0. 1) and Bob has an index i [N]. Both of them have access to a recurrent model as described in Section 2 which can perform the IdxL task perfectly. Alice can then provide the sequence to the recurrent model using any two symbols in the vocabulary 2 and iteratively update the hidden state to obtain the N hidden state N. Alice can then send the hidden state to Bob which requires mp bits. Bob can then provide the position token based on the index and compute the output to figure out whether a; is 0 or
Note that Alice and Bob can compute the output using one-way communication of mp bits and hence based on Fact 3 it must be the case that mp 2 N.
Associative Recall. A similar argument can be used to show that any recurrent model that can correctly perform the single query associative recall task [ must have width or hidden state of size at least ((N/p)) In the associative recall task, model is presented with sequence of symbols and labels æ V/,
N, y/N followed by a query symbol Ti which one of the symbols presented earlier (#1, (#1,-... .NN )- The goal of a model is to output the label y/1 corresponding to the symbol Ty = H2 where i = I....,,.. In this task, the symbols (#1....,IN must be distinct.
Proposition 9. Any recurrent model with a hidden state of size m over p-bits of precision that computes the associative recall task for all sequences of length N must have m 2 N/p.
A similar reduction from the INDEX problem can be used to show a lower bound for recurrent models. Both Alice and Bob have the description of the recurrent model and both of them can have a predetermined protocol for the sequence of symbols Zi....; NN. Alice can use the recurrent model to compute the hidden state h2 by providing it with inputs æ1, 01,... N; UN where the bits in a00, 1}* are provided labels. Alice can send the hidden state using mp bits and Bob can provide the symbol x; corresponding to the query index i and compute the output. Hence, the size of the hidden state of the RNN must be at most N/p.
C.2 1-layer Transformer with small width can perform Index Lookup
While any form of recurrent model must have a width or hidden state of size (2(N) to perform the index lookup task, we show that l-layer Transformer with width Ollog N) can perform the index lookup task.
Theorem For all N E EN, there is a 1-layer Transformer with width m = O(log N) and precision p= Ollg N) which performs the index lookup task all input sequences of lengths up to N.
Proof. For an input sequence (s1,...,8N, P). the Transformer uses the embeddings of the position token p and the positional embeddings of the first N inputs to attend over sp, so that the feedforward network can extract the label from the output of the attention block. Our key idea is use the N almost orthogonal vectors provided by Lemma [ both as positional embeddings and also as way to embed the numbers {l...., ], any of which can be used as the index p. Formally, let T(1. T(N) be N vectors of dimension k = O(log N) such that (T(1), TG)) S 1/4 for i 4 j and (T(1), T()) Z 3/4 for j.
Recall that, in the index lookup task, the alphabet V = -Y=UP [N] consists of symbols s, from a set E and the index tokens [N] {1,...,N] N). The embeddings of each input token will be of size log [X]] + 2k where log [Σ] + k indices are reserved for the token embeddings and the last l: indices used for the positional embeddings. Suppose we use a binary encoding for symbols in Σ, and
p: Σ + (0, 1] 1l |S| is such that p(s) is the binary encoding of E Σ. The token embedding of each symbol sj E E is p(s;) of length log Σ followed by l zeros. For inputs sequence $1...., SN, the embedding vectors will be of the form x; = [p(s:), 02. [())-
The embedding for any index token p E [N] contains log |Σ] zeros followed by the vector T(p). In other words, the embedding for the index token will be identical to the positional embeddings used in the first N tokens. The embedding vector corresponding to the index token p will be of the form
p logg E T(p), 0g].
18

---

The output of the l-layer Transformer is computed by applying attention over the input sequence with the query vector corresponding to the last token, followed by an application of the feedforward network over the output of the attention block. We can define the matrices We = 7/ O0 O WK O O, I Wv = [I, O], where > O is parameter that will be specified later, I is square identity matrix and O is zero-matrix of the appropriate shape. With these definitions we get that the query vector ((p) = [p[(p) contains the middle part of the input embedding, the key vectors K (x,) [T()] contain the last part of the input embedding and the value vectors
V(x)) - [o(s)] contain the first part of the input embedding.
With these query and key vectors, the dot products in attention satisfy:
2 3n/4 /
i P.
(Q(P), (xx))
S n/4
i
Additionally, the dot product of the query vector with itself will be (Q(p),((p))) 0. To retrieve the required token with the softmax operator, consider
exp((Q(p), (xxp))
exp((Q(p), (xpp)) + Σソ exp (((),
exp( 3n
exp N exp
which is 1/2 for any 7) > 2og(3.V). That is. for n) log(33.V) the attention weight over input token Tp with query token p will be greater than 3/4, 4, and hence the total attention weight over the remaining tokens will be less than 1/4.
Recall that the value vectors contain the binary encodings of the input symbols, Vxx)) = [o(s;)
Let Q1, "0l.....9 QN be the probabilities assigned the N tokens by the attention operator. Note that 4p 2 3/4. Let = = >. (:p():). Note that if the j-th bit of ((p)) is 1, then ž₁ 2 3/4 /4 and otherwise, 2, I 1/4. From there a ReLU-FFN can transform the vector the vector p(sp) which is the desired output.
Discussion. While we showed that a one-layer Transformer with O(log N) width can represent the index lookup task, it is unclear whether a one-layer Transformer with a small width can express the associative recall task. The construction above cannot be adapted in a straightforward manner to show that lne-l Transformers can represent the associative recall task. From the results in Section G on nearest neighbors, it follows that two-layer Transformers with logarithmic width can express the associative recall task. At the same time, the lower bound techniques for one-layer Transformers presented in Section do not directly apply to the associative recall task and hence is not straightforward to prove that Transformers must have two layers in order to perform the associative recall task.
The index lookup task serves a few purposes in our work. The result that it is in some sense easier for Transformers and difficult for recurrent models may not be very surprising based on the intuitive understanding of the architectures. Our results help theoretically formalize the intuitions that Transformers can use attention to arbitrarily retrieve tokens whereas recurrent models must compress the information their hidden states. Secondly, the task helps us introduce the techniques that will be used for the constructions and lower bounds to obtain more general results in the later sections. Lastly, it serves as a simple task that separates one-layer Transformer and recurrent models. As described above, it is not straightforward to show that the associative recall task can cannot be expressed by one-layer Transformers with small width with the index lookup task, we have that one-layer Transformers represent them efficiently.
DD Lower Bounds for 1-layer Transformers
As described in Section B.2 and in keeping with real-world implementations, the outputs of interme- diate computations are rounded p-bit precision. Further in keeping with real implementations, and to avoid overflow in exponentiation, softmax is implemented by first subtracting the maximum logit,
19

---

as
softmax(A)x, exp(Ar,) max, A.i))
ER=1 exp(Ar, max; A.i)
This is a popular approach to implementing softmax [9] it ensures that the exponentiated intermediate results are [0, 1]. avoiding possible overflow when exponentiating in finite precision.
Theorem 4. Consider a one-layer Transformer f E TF m.p.H operating over inputs of length N. Consider any disjoint subsets SA U SB = fI...., N), SA n SB = Assume Alice has access to 9; for S.. and Bob has access 9; for i E Sg. Then Alice and Bob can communicate 3mpp + log N)H bits to compute the output f(s1...
We note that conceptually related arguments were used in Sanford et al. [1 Theorem 7] and Peng al. [46 proof of Theorem 1]). Our approach here generalizes by stating this for arbitrary partitions over the input.
Proof. Without loss of generality, assume that N E SA. i.e. Alice has access to J-N-
In the first step, Alice sends z, to Bob using dp bits of communication. Then, Alice and Bob compute for each head the attention logits for each position within their respective sets:
AN.: == (Q(xw), K (xx:))
(1)
These numbers are rounded to p bits of precision.
In the second step of the protocol, Alice and Bob exchange 2p bits determine M max; AN,s and compute
ANa+ = AN,I - M
(2)
for their respective positions SA, Sg. Note that, AN, 0, exp (A.s) (0, 1], there no overflow issues arising from the p-bit representation. Then they can individually compute
ZA: =
Σ
(A.))
(Alice)
KESA
ZB == exp(AN.)
Σ
(Bob)
SESp
each with p bits of precision; both numbers are in [0. N). and at least one of them is 2 1. As all intermediate computations rounded to bits, exp(Ax,) ) is in [0, 1] and rounded to bits, Z Zp can be represented with + log bits.
In the third step of the protocol, they exchange 2(p log N) bits to exchange these. They then both have access
:: Σ exp((AN)) ZA + ZB
Here, we note that the definition of finite precision arithmetic Appendix B.2 makes addition of nonnegative numbers associative; hence, the outcome here is independent of the partition SA, Sp and agrees with the result when directly summing the exponentiated logits
The result in [1, N] as max; exp( (A,,)) attention scores 0ş:
This permits Alice and Bob to compute the
a = (aftmax(A)x.. = cxp(ANs)
(4)
Σn cxp(ANg)
each with p bits of precision.
'A similar protocol still works other bounded suunded-precision schemes where addition is not associative; as all exponentiated logits are in [0, 1]. one can use extended precision with + log N bits to perform the summation exactly and then round back to bits.
20

---

Then they both each compute:
UA softmax(A),,
=
(Alice)
UB voftmax(AA),, (x..)
(Bob)
rESn
with bits of precision. As each entry (tt((A),, x, has bits of precision, and E, Att(AA), = the sums can exactly represented 2p bits of precision.
In the fourth step of the protocol, they exchange these, which amounts to 4mp bits of communication. Then they compute
softmax(AAx, = UA + UB
(5)
and round it to p bits of precision. From this, they can obtain the result.
In total, the four steps took < 3m(p + log N) bits of communication. Performing this protocol separately for every head leads to S 3m(p + log N) bits of communication.
E Dyck with Bounded Depths
It is generally agreed that processing and comprehending natural language requires processing hierarchical structures [e.g. [3] [8]. The fundamental computational problem here consists of matching and relating material that matches hierarchically, even if it appears at great linear distance. This problem is formalized by the family of Dyck languages [[4]: languages of well-matched words over one more types of parentheses. These languages formalize problems that can be solved with access to a stack, where opening parentheses are pushed and closing parentheses are popped. Beyond fundamental model of hierarchical structure, they are of central importance in formal language theory, as any context-free language can be expressed in terms of Dyck languages [14. Perhaps due
to the boundedness of human memory [41]. natural language tends to have more bounded levels of
embedding [30 [0]. This has motivated the study of bounded-depth Dyck languages as plausible
simple models of the hierarchical structure underlying language, with substantial interest in the
abilities of neural architectures to model them 25 65 6
We will primarily focus Dyck-2 languages with depth at most k denoted Dyck-(2, k). The Dyck-2 2 language contains two types of brackets, such as round brackets 'C, 'Y' and square brackets '.' T. The D2c 2) language with depth most 2 contains well- bell-balanced parenthesis where for any prefix, the number of unbalanced parentheses can be at most 2. For instance, the string ((I)) has depth at most 2 whereas the string '([I]) )' has depth of 3. We say model recognizes language if it can correctly classify whether or not a string belongs to the language. We will primarily focus on strings with maximum length N and study how the size of a model depends on that.
Our main result in this section is that any 1-layer Transformer that can recognize Dyck-2 with bounded depths must have width that grows linearly with the input length N. To show that, we will first show that the communication complexity of the Dyck-2 2 language with depth at most 2 is at least N -- 1. The lower bound on the width of I-layer Transformers will follow from the lower bound on the communication complexity of Dyck-(2, 2).
Problem. Let Σ = {',", ",, 1'} be the vocabulary of language Dyck-(2, k). Let Σ" denote the set of all strings of length exactly n and Σ denote the set of all strings of lengths up to n. The communication problem between Alice and Bob is defined as follows. Assume the length N even for this problem. For a string x E ΣN, Alice has the symbols in the odd indices of the string and Bob has the symbols in the even indices of the string. They have to compute the function foyck: N/2 x N/2 > {0, 1} which outputs if the string x is in the language Dyck-(2, 2) and outputs otherwise.
For any function, the fooling set is defined in the following way, Definition 1. [Fooling set] A fooling set for any function EN/2 NN/2 + (0,1) is set S C N/2 x N/2 and value b {0, 1} such that,
21

---

For every (z. E S, (x, = b.
For every two distinct pairs (I1, M1), (I2, y2) E S, either (21))2) o b or f(22. # b.
The following is well-known fact in communication complexity.
Fact 4. For any function f. if there exists a fooling set of size [S], then the communication complexity C(J) log2 |
Lemma 1. Suppose Alice and Bob have the symbols in the odd and even indices of a string 8 E ΣN respectively. To each compute whether s Dyck-(2, 2). they must exchange least N bits.
Proof. The proof follows from the fact that there exists fooling set of size 2 2 for the language Dyck-(2, 2) with strings of length N. The fooling set S is constructed with all strings in s (zz y) y. such that fDyck(s) 8) = 1. Each string s in S satisfies:
(x, where E 2 N/2
odd y 2N12 and foyck (z. y) 1.
Here, x (2), 22.... N/2)) and y (v1, V2. ... UN/22 represent the sequences of symbols odd and even indices, respectively.
Constructing the fooling set. Note that, if x is the string of symbols in the odd indices of a string
s E (2, 2), then the string of symbols y in the even indices such that fDyak(1, y) = is unique. Suppose one is provided with the string
(F1,.... 2222), then can deterministically determine the symbols in the even indices y = (y1, ........ M2N/2) in the following way. Iterate through the symbols in æ, starting with T which must be an open bracket. After the open bracket 11, if there are one or more closing brackets 22.... Zg before encountering another open bracket 2K+1, then the symbols y1...., YK can be constructed deterministically using the following mapping.
' if = ')' for K, T if fzz+ = 'T' for j < K,
')' if xf = '(" for j = K,
T if 'I' for j K.
In other words, the symbols Vi,-..,Urre will be the open brackets corresponding to the closing brackets ,.,,,... KKK After the symbol 31, whenever you encounter another open bracket IK+1 where K >> O, then the symbol уи will be the closing bracket corresponding to the symbol æ Once you have matched the first open bracket symbol x1 with a closing bracket, you bound to encounter another open bracket. Follow the same process until the end of the string and one can obtain the string y, such that (x, y) E Dyck-(2, 2).
Hence, if x and y are the symbols in the odd and even indices of a string s E Dyck- (2, 2). then placing any other string y * y in the even indices leads to string which is not in Dyck-(2, 2). Our fooling set 5 contains all strings of length N in the language Dyck-(2, 2). Hence, by construction, we have that fDyck()) (5) fDyck v) for all (z, v) E S and
( w22 SDyck (22, vi) = 0 for all (I1,21) + (x2, E S.
Thus, such a set S is fooling set by Definition 1
Size of the fooling set. One can show that the total number of strings of length N in the language (2, 2) is exactly 2 with some elementary combinatorics. First, see that the number of Dyck-2 sequences of depth at most 1 and length N is exactly 2/1/2, This is because the string is made up of blocks of '()' and '[I', and hence there are two choices for every block. Then, consider a block of Dyck-2 2 string of length & (where & is even) which starts and ends with an open and closing bracket respectively. Morcover, the Dyck-2 2 string has a depth exactly 2, e.g. ' [ O J'. Note that, the total number of such strings of length le is exactly 2²/2 since there are two choices for the first bracket and there is D string of depth 1 and length k - 2 inside it. For simplicity, we will call such strings as belonging to the language Dyckb-(2, 2) which is subset of the language Dyck-(2, 2). In simple terms, strings of length m in the subset Dyckb-(2, 2) have an overall depth and have a depth string of length m between the first symbol (open bracket) and the last symbol (closing bracket); for instance 1O
22

---

The total number of Dyck-2 strings of depth at most 2 and length exactly N can be computed as follows. Partition the indices into & contiguous blocks where each partition is of an even length. Suppose the indices begin with 1, then the points of partition could be between 2 and 3, 4 and 5, and
so on. For instance, a valid partition of the indices [1, 2....... 8] into 3 blocks is [1. 2). [3..... 6). and
[7, 8]. The total number of partition points is N/2 1. For any such partition, if the length of a block is 2 then that block can only have Dyck-2 2 strings of depth I and if the block is of length 2 4, then consider all possibilities of Dyck-(2, (2, 2) strings starting and ending with open and closing brackets respectively or in other words, strings in Dyckb-(2, 2) described earlier. the number of partitions is N/2 1, then all possible strings have depth at most and if the number of partitions is O, then the entire string is in Dyckb-(2, 2).
See that, no matter how you partition the inputs, the number of possible strings at each block of length m is 2mm Further if P = {PI...., } denotes the set of lengths of each block in the partition, then
k
the total number of strings with such a partition is
II 2 = 2 p:/2 2N/2, In other words,
==1
regardless of how the indices are partitioned if each block of size > is required to be a string in Dyckb-(2, 2). then the total number of possible strings is 2N/22
The total number of Dyck-2 strings of depth at most 2 can be computed by considering partitions of all sizes A O..... N/2 1 and all possible partitions for given size. With this, we get the total number of valid strings in Dyck-(2, 2) of length N to be
N
N
2
2/12 2N
k
k=
Hence from fact 4 it follows that the communication complexity of Dyck-(2, 2) is at least N
While we have given a self-contained proof based on fooling sets, an alternative proof of Lemmal could proceed using varieties of finite monoids, by proving that the syntactic monoid of Dyck-(2, 2) not the variety DA, and then applying the result of Raymond et al. 49
Using Lemma I and Theorem it follows that any 1-layer Transformer that can recognize (2, 2) must have width that grows linearly with the input length.
Theorem 10. Consider one-layer transformer E F. deciding membership in Dyck-(2, 2).
Then dH 3(p+log NY
N-1
Suppose we allow the transition function a recurrent model to be any arbitrary function as defined in Section 2 In that case, it naturally follows that such RNN can represent any DFA with & states with a hidden state of size log k. Any bounded dyck language Dyck- (n, k) can be represented by DFA with O(2*) states which is independent of the input length N. In particular, the language (2, 2) considered here can be represented by a DFA with just 7 states. Hence, it follows that recurrent model of constant size can represent the D2ck-(2, 2) language for arbitrary lengths in a finite precision setting. Additionally, prior works have described constructions of such recurrent models with practical activation functions such as Sigmoid (25 and ReLU [6] Thus, the bounded Dyck language Dyck-(2, 2) provides us a separation between one-layer Transformer and recurrent models since constant-sized RNNs can represent them whereas one-layer Transformers must have at least linear width to represent them.
F Transformers and Boolean functions
Communication protocols and lower bounds. There are some notable differences between the types of communication complexity lower bounds we have for one-layer Transformers (Theorem 4 and for RNNs (Theorem 2). If an RNN can compute a function f then Theorem 2 implies that there exists a one-way communication protocol with mp bits over any contiguous partitions of the input. On the other hand, a one-layer Transformer can compute a function f, then Theorem implics the existence of communication protocol with O(mpH log N) bits over any partition of the
23

---

input - but not one-way. Hence. the types of lower bounds that apply to these two architectures are different. For any function/ f such as the INDEX problem, the one-way communication complexity is lower bounded by ( over some contiguous partitions then it implies that for RNNs, the width m = (N However since the two-way communication complexity of such problems might be lower, those lower bounds need not apply to one-layer Transformers. For instance, the INDEX problem can be solved with log N bits of communication if two-way communication is allowed or in other words, Bob is allowed to send bits to Alice. Conversely, to prove lower bounds for Transformers computing a certain function f. proving a communication complexity lower bound for f over any partition of inputs suffices. For a particular partitioning of inputs, we showed lower bound on the communication complexity of bounded Dyck languages. While that result implies a lower bound on one-layer Transformers, the partitioning is not contiguous and hence does not apply to recurrent models. For Dyc 2), contiguous partitions are not a hard case, and in fact, communicating S 2 open brackets from the first half is sufficient. This is why the lower bound of Lemma does not apply to RNNs. Despite these differences, for a large class of Boolean functions including functions such Equality and Disjointness, the lower bounds apply to both of these architectures.
Equality. For convenience, we discuss the complement of the Equality function INEQ = 1 -- (=)()) which is the inequality function. It does not influence any of our results or the constructions for Transformers but it helps make the intuitions for the construction clearer. The function INEQ {0, 1) - {0, 1) is a Boolean function defined as,
(NEQ(x) = I[( 1.001.......... (((=/2)..,,N//)) & (IN/2-11....NN))
This can also be represented as a 2-DNF with N/2 terms,
[NEQ(x)= = (z1 A -IN/2+1) V (-0) A xN/2-1) V... V (#N/2 A (-ax//
F1 Lower Bounds
We show that for any RNN that computes the INEQ function over {0. 1) with hidden state of size m and p bits of precision, it is necessary that mp = (NN). In other words, the size of the hidden state of the RNN grows linearly with the input length N. Similar to other results, this will be based on communication complexity but we also provide an alternate way to prove this lower bound based on the state complexity of DFAs that compute the INEQ function.
The result states that if any RNN with a hidden state of size m over p-bit precision can simulate a DFA with n states, then the representation size of the hidden state mp must be at least log n. There is one caveat related to this DFA-based result which does not apply to the communication complexity- lower bounds. For the DFA-based result, while the transition function is allowed to be an arbitrary function, it has to be of the form g(x,, h ).i.e., the functions cannot be different based on the timestep t. It is unclear if the result still applies when the function is allowed to be different based on the timestep. However, since the transition function is allowed to be any arbitrary function, it still captures all the recurrent and state-space architectures used in practice. Additionally, this fairly simple technique could be used to prove lower bounds for RNNs based on lower bounds in formal language theory.
Lemma 4. Let A be a DFA with n states. Let R be an RNN such that for all x E E", A(x) = R(x)) Then the size of the hidden state of the RNN: mp 2 log n.
Proof. Suppose an RNN R exists such that the dimension of its hidden state m < log n. It is straightforward to see that can construct another DFA A' from the RNN R such that the DFA A' accepts the same language as A but it has 2mp < n states.
Create state for each vector h {0, 1 and assign the state corresponding to the vector ho as the start state. For each state corresponding to vector h, make q a final state if f(h) = 1. For each h E {0, 1m and each symbol x E Σ, compute h' g(x, h). Add a transition between the state corresponding to h and h using the symbol x. Hence, we get the entire DFA with the transition map and the set of start and final states that accept the same language.
Since we are given that the automata A is the minimum-siae sized DFA that accepts the language, the automata A' having fewer states and accepting the same language is contradiction. Hence, such an RNN cannot exist and the dimension of the hidden state vector m of the RNN must be at least log
D
24

---

Theorem 11. Any recurvent model with a hidden state of size m over p-bits of precision that computes the (NEQ(=) for all E {0,1] must have mp N/2.
Proof. We show two ways to prove the above statement. The first proof is based on the communication complexity of the Equality problem.
Proof 1. Suppose Alice and Bob have the input vectors a, {0, (0,1)./2 respectively and have access to an RNN that computes the INF.Q function over {0. 1 Alice can first use the RNN starting with input d1 and the vector hg and iteratively update the hidden state until the input oN/2- Alice can then send the vector h2/2 to Bob who can provide b,,..., .1.....bn/2 as inputs and compute
y = INEQ(a b) = 1 - EQ b). Since sending the hidden state vector hxy2 requires mp bits, must be that mp 2 / due to Theorem 2
Proof 2. The second proof uses the following relation between recurrent models and a DFA.
Let AINEQ be the minimum-sizer DFA that computes the INEQ function over {0, 13.V When we say a DFA computes INEQ function over {0, 10 we mean that ArNEQ( (NEQ(=) for all E {0. 1J and ArN (2)) can be defined arbitrarily for x & {0, 1) to obtain the DFA with minimum number of states.
Note that, any DFA that agrees with INEQ over all x in {0, 11. must have at least 2 states. This follows from the fact that for any two distinct 21, E {0, 1/ there is a distinguishing suffix 8 E {0, 13/822 such that (=1 8) + INFQ( : 8). Hence, any DFA that agrees with INFQ on all inputs in {0, 11 must have at least 2/2 states even if is defined arbitrarily over other inputs & (0, 1 From Lemma 4 it follows that mp 2 N/2.
Theorem 12.
Any one-layer Transformer with width m and H heads operating over p-bits of precision that computes the INEQ(=) for all x E 13 must have mpH- = S2(N).
The statement immediately follows from Theorem and Fact 2
F22 Transformer Construction for Equality
We now show how a log-sized 2-layer Transformer operating over log-precision numbers can compute the INEQ function over all Boolean inputs.
Theorem 6. For any N N, there exists 2-layer Transformer f E TF where width Ollog N) and precision p = Ollo N) such that fxx) F.Q(x) for all {0, 13
We first describe the broad idea behind the construction. We will consider the input domain to be {-1, 1 and our goal is to compute LNEQ(z) for all x E {11,,11 This can also be formulated as,
INEQ(=) = (a1 e @ (12 (12 V V (x 2 @xx)
Let fineq denote the first layer of our construction fineq that computes INEQ and let p(1) f1l)(()) denote the ith output vector of the first layer. Our construction will be such that for i > N/2. the ith input the attention mechanism at the first layer will retrieve the (i - N/2)th input using tools described (4); will compute x; @ x;- for each > N/2 where 24 - if x; = - v and is otherwise. The second layer will then take an OR over all those values which will result in computing (NE)(z).
(z);
in Section B.3 With the MLP, the first layer (ll((a))
Proof. Let (a1, nN) denote our input vector. The input to the Transformer model will include the positions as well ã = ((x1, 1), ... (IN,N) N)). Let E Rd denote the embedding of the äth input ã; where d = 2 + 2k. Here. = (log NN the dimension of vectors T(1)..... T(N) described in Corollary 8.1 The input embeddings will contain four parts and will be of the following form
[r 0, T())
if < N/2, otherwise.
x,
(o T() N )
25

---

The query vectors Q(x,) x, W E will be the last part of the embeddings, i.e., Q(x,) = [ for S 2/2 and is (( 222 for N/2. Similarly, the key vectors K(x,) = (112 for all i. The value vector will be d-dimensional vector V(x;) = [0, z4. 02 0s.]- The query, key, and value transformations can implemented with block matrices containing zero and identity matrices similar to the one described in Section C2.2
By construction, the dot products Aij - (Q((s),,(Xx)) (T(i-N/2)) T()) 1/2 will be such that for i > 2/2, Ao, - (Q(x),
-
2 1/4 $ 1/4 otherwise.
if +
For i < N/2, the dot products A., will be greater than 1/4 if i and will be less than -1/4 for * j. If this was a Transformer with the hard-attention mechanism, then, note that Att(X); 0, rg. Og, 0g] for i S N/2 and (tt(X); = [0, ",, N/2; 02, 0&] for N/2. With residual connections or another attention-head, the output of the attention block will include the original input as well which will lead
{ [pi, D; (1), for i S N/2, [2 (()). [ for > N/2.
Att(X); +
Then simple ReLU FFN can compute the XOR of the first two values of the output vector from the attention block. Hence, by construction, the output vector from the first layer Jiey (xr); = [c; 2 T:-N/2, for i > N/2 and [o. for i - N/2.
Second layer computing OR. The second layer of the Transformer will compute the OR over the first coordinate of the input vector which is quite straightforward. Let the query vector Q(x; ()) = The value vectors will be of the form ((x(")) [(2, @ z;-N/2), 0,.,.,... 0] for i > N/2 and they will be zero vectors construction for i < N/2. Then, regardless of the keys, the dot products will be 0 for each position, and hence
(2) N Att(x)
Σ N/
(a+ O
If Att(X) (2) 2 $.. then the FFN can output 1 and it can output 0 otherwise.
Softmax attention. We now describe how to retrieve the N/- values with softmax attention in the first layer of the model. The approach is slightly different from the one used in Section C.2 and makes use of finite precision rounding.
Consider another construction that is identical to the construction descrihed above with the exception that Q(x,)) = 7)(x,) = (t), 1]. We show that for large enough r) and i >> N/2, the weight softmax(A);, will be close to 1 for i N/2 that it will be rounded 1. Similarly, it will be
rounded to 0 for j * i N/2. For i < N/2, the weight will be rounded to for = j and will be rounded to 0 otherwise.
Recall that the finite precision implementation parameterized by a large constant K. (cf. Ap- pendix B.2 For > N/2 and j = N/2, there exists such that,
((xn(Q(x),,
1
exp(n(Q(xx)) K (xx)))
2NK.
For such an n, the weight softmax( (f)t) will be rounded to 1.
12- cxp(n(Q(x.),
cxp(11)
Σ E (xp(n(Q(x;)). K ((xx))) exp(1n) + (N 1) exp(c) > 1 2/
kc expl 1 (n) 2 (2N 1)(epp( (1) (N 11) exp( 1 1n
26

---

exp( 2) ()) (N- 1)(2Nks -1) -1xxpp-2
2 2 (IN- ((V-1)(2Vk) - 1))
Thus, for 1) = log N + K2 log 2NN, the softmax attention weight at j = 2/2 will be rounded to 1. Similarly, one may verify that if 17 2 Kg log 2N,, then the weight softmax(A);, for j + i N/2 will be less than 1/ and hence will be rounded to O. Hence, for log N + log 2o, N, the softmax attention will behave like hard attention, and as described earlier the Transformer will compute the INEQ function. This completes the proof.
The scaling in the attention mechanism in the last part can also be implemented in almost exactly the same way as the construction for Theorem [ Implementing it that way then requires the ReLU FFN act as threshold function. The scaling described above makes use of the finite precision setting to amplify the dot products to the point that it acts as hard attention.
F.3 Representing more general class of Boolean functions
We now describe how the construction for Equality in Section F.2 can be extended to class of Boolean functions namely, thresholds of at most. N k-SPARSE features. By threshold functions, we mean functions of the form Th, {0, 1 + {0, 1} where for {0, 1 the function is defined as (2) = -=> 24 - > O). The function is parameterized by a constant b. For b = n 1, the function effectively is the AND over all bits. Similarly, for b n//2, the function outputs 1 if the majority the input bits are 1 and outputs 0 otherwise.
A S-SPARSE function is simply function whose output on any input x E {0, 1)"h depends on at most indices of the input. More formally, function f {0, 13"h + {0, 1] is k-SPARSE if there exist indices 11 < i2 <... < ig S n and a function g [0, - (0, 1), such that for every E {0, 13", f(z1, T2...., IT) = g g +Tz. ). A k-SPARSE function can be any Boolean function (AND, OR, XOR, etc) with the constraint that it can depend on most k O(1) bits of input.
We now define the class of threshold of at most N k-SPARSE features denoted as THRES Let gr be a k- k-SPARSE function depending on I C [N] indices where [ < k. A function f {0, 1)8 - {0, 1} is in the class T if it is of the form f(x) = (h, (z), g1, (r), 9IN (z)) where each
1 C [N] and [4, S k for all j E [N]. Further g1, (z) = 0 for all x E {0, 1 if the set L = 0.
The following result states that for any function h in the class of threshold of at most N k SPRRSE features, there exists a two-layer Transformer with logarithmic width that can express the function
Theorem 13. For any N E N and any function h THRESK,N, there exists a 2-layer Transformer fIr E TF2 TFR with width m = O(log N). H heads, and precision p O(log N) such that fTF(x)= = h(x) for all x E 10.
The result follows from a straightforward extension of the construction in Theorem 6 We will again remap the domain (-1, 1) For this problem, will also prepend the input 1)M with an additional beginning of sequence token [BOS]. Hence, the Transformer will receive N + 1 tokens Xo, , XN.
Proof.
For any function E THRES the first layer of the Transformer will compute the k-SPARSE features, 91, (z), 91, (z)....., 91 (z). The second layer will then compute the threshold function The over the S-SPARS features.
Computing Threshold. The second layer is quite trivial to construct. Let x (0) the output vectors of the first layer and inputs to the second layer. Suppose the ith output vector contains the ith k-SPARSE feature, x; = [gr, (r),.. for i = 1,..., N and x,=) = 0. If the query transformation is a null matrix, Q(x = o, and the value transformation is such that v(x(") [gn, (x)). it follows that the output of the attention block will be N4i ΣN 91, (z). The
..... x be
27

---

ReLU-based feedforward network can then be used to subtract with constant N+1 and implement a threshold function to compute the desired output.
Computing K-SPARSE features. In the first layer, we compute the S-SRES features by making use of the almost orthogonal positional vectors (0)), T(.)..... T(N) of size r = (log N). From the construction for equality, it should be clear that using one attention head, one can attend to any desired position and retrieve a single bit. Extending that, using k heads, we can retrieve & bits from different positions. To compute 9r. (z)), we will have input embeddings of size + (A: + 1)r. For the ith input zi, the first coordinate will contain the input x; E [[1,1]} and will contain O for the [BOS] token 55. The next indices will contain the positional vector T(i) corresponding to the position i. The remaining kr indices can be divided into & blocks each of which will contain a positional vector. If the set 1; contains indices 1," ...., 1. then the & blocks will contain vectors T(4)).... ...).......................... T(1,)) If the set 4 has less than k indices, then the last A 44] blocks will have the positional vectors (0).
See that if the input embeddings are designed in the output of the attention block at the ith token. The value vectors will be of i dimension such that (x,,), head can be K I₁ at the coordinate h and is 0 everywhere else. The key transformation for each (x,) = [T()) containing the vector corresponding to the token's 's position. The query transformation is distinct for each head, and for the jth head, the query vector contains jth block from the last & blocks containing the positional vectors for each index in the set I;. If the query and key transformations can be seen that the output of the attention block at the ith position will be a l dimensional vector containing (..,,..., 2 gt.) J. If the set 4 has a size less than k. then the output vector will be followed by zeros after | coordinates and will be a zero vector if 7; = Finally, the feedforward network can then compute the function g1, (Ir) at every position which can be done by network of constant size because & = O(1) does not depend on N. The second layer can then compute the threshold over the k- K-SPARSE features described earlier which leads to the desired output.
described above then one can obtain input bits 2T,
designed in such a way, then using the arguments described in Theorem @ it
Discussion.
The class of thresholds of S-SPARSE features contains certain functions of interest such as Disjointness and Equality as well as more general classes such as k-DNF's and k-CNF's with at most N terms or clauses. As described earlier, the (In) function can be also defined as INEQ(=) (z) x +1) V V ( ① xN). See that it contains 2 2-SPARSE features where the set of indices 4;- = {i, i N
2 for i = i,..., N/2 and the feature function g(a. b) @ b for a. E 10. Similarly. the Disjointness function can be represented as 2-CNF. The complement of the Disjointness function can be described more simply
(0) A Fy +1) V (x2 A z +2) V... V (x A AxN)
which is 2-DNF that outputs if the first and second half of the input x E {0, 13 are disjoint and outputs 1 otherwise. Thus, it follows from Theorem 13 that two-layer Transformers with logarithmic width can represent functions such as Disjointness as well as k-DNFs and k-CNFs with at most N terms or clauses.
F.4 Difficulty of Deriving Communication Lower Bounds for 2-layer Transformers
Our lower bounds both for RNNs and for one-layer transformers are based on communication complexity arguments. Here, we provide evidence that other techniques may be needed to establish lower bounds for two-layer transformers by showing that, in certain sense, no short communication protocol of the same kind as Theorem 4 can exist for Alice and Bob to obtain the output of any two-layer transformer. We start from the following lemma:
Lemma S. Assume N is even. For every partition SA SB of (L,..,,N where |SAl = |S there is a 2-layer transformer f E TF2,2 with width m = OOlog N) with precision p = O N) with f(x) E {0,1) for each x E {0, 13N, such that Alice and Bob, having access to A and x respectively, need to exchange 2 2 bits to compute f(x).
N
Proof. For any partitioning SA, Sn, consider the task of determining whether ISA and ES identical. That is, define (for x E (0, )
{1 if ISA # ISB
fs (z)
(6)
else
28

---

If Alice and Bob have access to 29 and TSa. respectively, they need to exchange 2 bits to compute (xr). Now by Theorem 6 there is a transformer fineg E TF m.pp that computes (mn.n////////////////////// where m ( N) and p= Ollo N). Now renumbering the positional encodings results in a transformer computing SSA.Sn.
Corollary 13.1. If there is any communication protocol by which Alice and Bob can compute the output for any nwo-layer Transformer, for any partition, with O(mpHg(n)) bits, then it must be the case that (n)) Ω
N (log N
Proof. Suppose there is communication protocol such that for any two-layer Transformer and over any partition, the protocol can compute the output of the Transformer with
N
o mpH
(log N)²
bits. But by Lemma 5 for each partition, there exists a two-layer Transformer with
mpH O00o( N)
(8)
such that Alice and Bob need to exchange 2 N/2 bits to compute its output. We obtain a contradiction.
G Nearest Neighbors and Associative Recall
Recall from Section 5.2 that in the NSTNB task, model is provided with sequence of vectors and labels (x1, 11......... X
1k-1, x1) Xx) where N/2 < k S N and the goal of the model is to predict the
label corresponding to the nearest neighbor of x for each k = + 1,..., N. We first show that any recurrent model that performs this task must have width or hidden state of size 2
Relation to Associative Recall. The nearest neighbor task is closely related to the single and multi- query associative recall (MQAR) task introduced Ba et al. and Arora al. [2]. In the associative recall task, a model receives a sequence $1, 1/1 V/.............................. -k 1, where the symbols s; belong to an alphabet EX|. Assume the symbols the $1,.... - to distinct and the labels y: E (0. for simplicity. The query input 8æ is repetition of one of the preceding inputs. The goal of the model is to predict the label of the query input S by finding the exact match from the context and producing the corresponding label. A model that can compute the nearest neighbor algorithm can perform these associative recall tasks by embedding the sequence of symbols 81, y1. 8k-1 yk 1, 8% as a sequence of vectors (x1, y1....... X;-1, 1,, Xx).
If a model receives a sequence (x1, 31...... ... Xg-1, -k-1, xx) distinct and the query vector X₁ is repetition then the task of applying nearest neighbor to predict the label for the query vector reduces the single query associative recall task. Implementing the nearest neighbor algorithm is equivalent to finding the exact match for the query vector X and producing the label y corresponding to that input. In the case, where the models are required to predict iteratively for all X3 for h: = N 1,..., N, the task becomes equivalent to MQAR. T
where the vectors X1 X2, X are
The MQAR task is simplified version is first provided with the prompt (x1, ..... XN/2, WN/2) and the subsequent N/2 input vectors are a permutation of the first N/2 vectors. In other words, for N/2 < h & N and a sequence (x1, y1,...,X4, -;-1; 1/-1, Xx), the model has find the exact match of x2 in the first N/2 vectors and output the corresponding label. It is straightforward to see that any model that can perform the nearest neighbor task also perform the MQAR task. Assume the size of the alphabet |Σ| (NN) The embedding of each symbol s; Σ can be distinct vector from the hypercube {- --va= va 1d where d [log ||| = Olog N). The embeddings can be thought of as the normalized binary encodings of each symbol. Thus, model receives sequence of vectors (X1, $31./1.-. 1/1................... ..... .......,,..11 X4 1/--1, Xx) corresponding to a sequence of symbols $1, 1/....., Sk V/-11 $3. Since the query vectors x; for k N/2 are repetitions, there is an exact match for each of them in X1, XN/2- Suppose the exact match for query X Xy", then their dot products X,)) and the dot product with
subset of the nearest neighbor task where the model
29

---

log N. Since the margin between the dot product with the nearest neighbor and other vectors satisfy & = Oloo N) and all the input embeddings have unit norms, the problem satisfies the assumptions of the nearest neighbor task.
every other vector xx) -
1
G.1 Lower Bounds for Recurrent Models
Theorem 8. Any recurrent model with a hidden state of width m with p-bits of precision that can perform the nearest neighbor task for all inputs of length N must have m 2 N/2p..
Proof. The proof is via reduction from the disjointness problem. We show that the lower bound true even for the restricted problem of multi-query associative recall (MQAR). Recall that in the MQAR task the sequence of query inputs x +1..... X is permutation of the sequence of labelled vectors X X/22-
If there exists recurrent model R that can solve the MQAR task, then we show that Alice and Bob can use follow communication protocol and compute the DISJ function.
The communication protocol is as follows. Alice and Bob have two Boolean vectors a, E {0, 1 1 respectively. Both of them know the description of the recurrent model R. Alice and Bob have decided a set of N/2 vectors Vi,..., V2/2 in gd-1 that they will use in the communication protocol. Choosing any N/2 distinct vectors from the hypereube satisfy the assumptions mentioned in Section 5.2
1 will suffice and also Va val
To reiterate, both Alice and Bob have the model parameters/deesriiption and have decided on a set of N/2 unit vectors as a part of their communication protocol. Alice then the recurrent model R and provides the sequence of pairs (V;, a;) in any arbitrary order. Alice then sends the hidden state vector hy to Bob, Bob then uses the model and iteratively provides V, vectors along with the generated output in any order. Bob knows that the output produced by the recurrent model for the vector v, corresponds to the value a, Hence, Bob obtains the entire vector a and computes DIsI(a, b).
Since Alice sent the hidden state vector, which requires mp bits, and they could compute DIsI b) over (0, 1] N/2 by exchanging mp bits, by Fact we have that mpp 2 N/2.
G.2 Transformer Construction for Nearest Neighbor
We now show how a log-sized 2-layer Transformer operating over rog-precision numbers can compute the NSTNB function over all inputs that satisfy the constraints described in Section 5.2 Recall the assumptions that
Norm. All vectors x; have unit norms.
Margin. For any N/2 < k s N, and let j* = arg max{c(k-1) x7,x,, then xx[x,. 2
xf{x +7 y for any + j"
Theorem For any N E N, there exists a 2-layer Transformer fNN E TF2 TF with width m = Olog N) and precision p = OOlog N) such that JINN computes the nearest-neighbor task all sequences of length at most N satisfying the assumptions above.
For N/2 < k N, the model will be provided the sequence (X1: Xx 9k-1, Xx) where X; E Rdd vectors contain the input points and yıs contain the labels. For clarity, we will refer to the embedding of all inputs z; E d where Z2:-1 = o(x;, 2i 1) is the embedding of input points for i = I,..., k. Similarly, the vectors Z2 = Φ(v, 21) will contain the embedding for the corresponding labels. Hence, technically the sequence of input vectors the Transformer model will be (z), Z2, 12)
Overview of Idea. The explicit construction is bit tedious but the key ideas straightforward to understand. The construction itself does not rely on input sequences of fixed lengths, unlike the one for the Equality problem.
Recall that for an input sequence (Z), Z2,... ... 22k 1), the odd indices (2i 1) correspond to the inputs x; and the even indices (2i) contain the label information y:- The final input 22k-1 contains
30

---

the query input X and the goal of the model is to find the nearest neighbor input x,- and output the label corresponding to that By-.
Intuitively, a two-layer Transformer can do that in the following way: the first layer can find the nearest neighbor Xj- and retrieve the position of the label y-- The second layer can then attend over that position and produce the desired label.
Challenges. There are few challenges to executing this strategy which our construction will address. the input embeddings were identical to the input vectors, then X would have maximum dot product with itself and not with its nearest neighbor X,. Second, if you remove that, even then the dot product with some label vector could be larger than the dot product with the nearest neighbor (which could be close -1). Lastly, suppose by design you have the maximum dot product with the desired input, you still need to retrieve the required information in a useful way since we are working with softmax attention which will also put weight over other inputs.
Key ideas. One can think of the vectors T(1), ((N) as some form of positional or address vectors of dimension Olog N). All vectors 2 corresponding inputs x; will have the address vectors of their next position T(2i). Along with that, all the vectors z; will also have their own address/position vectors T (t). If in the first layer, the model can attend over the vector --zz-11 with high attention weight, then it will be able to retrieve the required address vector T(2)*) which it can then use in the second layer to retrieve the required label.
The input embeddings are designed in such a way that the maximum dot product will be with the desired input vector Xj- or more specifically 22j*-11 The embeddings are such that the dot products between the query vector of input 1 and the key vectors of all other input vectors z; will be of the following form,
(Q(22-11), ((Z2i-1)) = (x xi)
-2, - 1), 1)) for i = 1,...,,
See that for i + k. (Q(22x-1), K (222-1)) -2 or much smaller based on the constant 01. Hence, attention weight will be smaller on the query input itself compared to other inputs.
R (x) whereas for i = k. (Q(22x-1) 1), K (222-1) ()} <
Secondly, for i 1....,k k-- 1, the dot products with the label vectors will be of the following form,
1), 7(2 (Q(Z2=-1), ((za)) = (xx. 0) ci (12 1)) - 02 Re -
for i = 1,...,ke k.
which will be small depending on the constant C2. Hence, the dot product will be maximum with input with the nearest neighbor Xz. with a margin 12(7).
To retrieve and use T(2)*) in the next layer, we do not need to hard-attend on -1 Since the address vectors are of the form T(i) E { { Vi' ) we only retrieve the corresponding sign vectors Ti(1) ViT(i) e 13 Since the attention weight will be maximum on -2j--1 with a margin, we can scale the query vectors to increase the weight to a sufficient value and then use ReLU as form of threshold obtain T, (2)"). The second layer is then straightforward and will retrieve the desired label.
Proof. We will now describe the explicit construction for Transformers compute nearest neighbors. Input Embeddings. The embedding vectors are defined in the following way,
(x,, 2i 1) = Z = [x, 0, 1, (2i - 1), T(2)), 0,] @uv,, 21) = Z22 10x Vi. 2, ((2)) 0x)
(9)
Here, the vectors T(i) E (- v 3 are JL transformations as described in Section B.3 and hence k = Olog N). The dimension of the embedding vectors d = 3k + d' + 2 = Olog N). The vectors T(I)..... T(2N) are such that,
7/1000 ifi=j. j, (TO.TU) = 0 otherwise.
31

---

Query, key and value vectors. The query and key transformations in the first layers are designed in the following way,
(( = [xx: 1, 107 - 1)]
(222-11) = [x,, (2i - 1)] for = ,....,k. k ((22;) [0.. 2, ((:))
for i 1, k-
The value vectors will only have the 5th part (2i)) in the last slot and all other values will be 0. Let T.(d) (i) E {- that is, (i) contains the signs of the vector (i)). The value vector will be of the form V (22 (Ozr, 0, 0, 0g, Os; 27, (2i) and V ((22))
[Ox, 0,O, O2, O2, (2)) for all i = 1..... N.. The goal of the attention in the first layer is to retrieve the vector (2j")
corresponding to the label of the nearest neighbour.
Inner products. The design of such query and key transformation ensures that for an input query Xk. the dot product is maximum with its nearest neighbour X,' and not with itself any of the embeddings of the labels yis.
The inner products A 1,21-1 (Q(z2=-)) (Z2;-1)) have the following form,
(Q(22x- K ((221-1)) = (xx: x;) 10 - 1), T(i 1))
(x) + 3 1/10
if k,
(xx, x:) + 10 7/10
if k.
21 if - if k.
(Q(
Similarly, the inner product with the label vectors Z2,8 is less than 0 for all as well,
(Q(z2x-1); K (22:))
6 + 7/10 -
To summarize, the query and key vectors are such that for the query input z2 1, the dot product with itself (Q)2)e-1). (()) I -5 and the dot product with all label vectors containing y; (Q(z K (22)) s -5. The remaining dot products are with the vectors of interest which include the nearest neighbour input point,
(Q ((zz) = (X.. x.) + 7/110 for i = 1,...,k- 1.
Suppose j arg [a x x; is the index for the input point which has the minimum L2 distance maximum inner product with the query vector Xt. Then we have that,
(Q (((22=1). K(225' -1) - (Q(z2t-1). (0( ((2z;-1)
2 7/5
(10)
for all i j³ This indicates the maximum inner product will be with its nearest neighbor and all other inner products will have margin of at least 7 = 7.
If we have Transformer with hard-attention mechanism, then it will attend only to the input which is the nearest neighbor of the query vector XĄ. However, with softmax attention, it is non-trivial to only attend over single input.
The embedding of all input vectors x; contains vector T(2i) (see Eq. 9 which serves as a key vector to retrieve the label following that input vector. Note that, we only need to retrieve the signs of 1 & We can then use retrieve the v Je
the vector (24) since the vectors are of the form { (-
label of the corresponding input in the next layer.
Retrieving with softmax. We show using softmax attention and ReLU FFN we can retrieve the required T(2**). In particular. we will obtain (2j*)) vkT(2j") E 13 As shown carlier, if j" = arg max k- x x4, then the maximum dot product (Q( (Q (Q(226-)) 1); (((22-1)) )} + is greater by margin = - The value vector corresponding to Zyr. i.e., V(2,-) = [0...., 0, (2,"))) has the key vector which will allow the next layer to retrieve the required label.
Z
32

---

The basic idea is that even if at least 9/10 of the attention weight is on V(zz-) and essentially 2T, (2j") then that suffices to preserve the signs using a ReLU FFN.
Let (a) be a function such that (a) = for a > 1, a()) = -1 for a < -1, and ( = for 11 S I 1. See that (a) can easily be implemented by ReLU FFN since it is essentially ReLU(x + 1) Re - 1) 1
Without loss of generality, let's say (2)")) = +2. If the attention weight on it is 9/10 and the remaining 1/10 weight distributed among the rest of the inputs, then the worst case that value will be 2 1.2 1. Hence, applying o-)) over the value will result in -1. The same goes for the case when 2T, (2)"), = -2 and for other indices of attention weight over V (Z2z5-1) can be greater than 1010 since it implies that with the ReLU FFN, the output for 1 in the first layer will contain T, (2)").
Thus, suffices to show that the
Consider another construction which is identical to the construction described above with the ex- ception that ((2,) r/Q(z.). We show that for large enough 7) the weight softmax( (A) will be greater than 10 and the remaining weights combined will be less than 1/10. Let ( and Z E NNd contain vectors (Z1,..., Z28-1). Then, for any r E N.
(0lmax(((0xx--)X((2))))> (2j))
exp(n(Q(22x-1), (2)jj- (xp((
exp(7(Q(zex-1) ((2g)-)))) 1o =
exp(n13) + (2k 1) exp(11(3 - -)) exp(n3) + 2N exp(n( - -)) 2 r- 1
exp(1)3)
exp(n3)
r
= exp(113) 2 (r 12 exp(n(3 - -T))
n 2 log(r 1)2N
Hence, for r 10, 11 greater than 9/10. Similarly, one can verify that for 7j = In log 18N, the attention weight for the rest of the inputs combined is at most 1/10.
5 log 18N then the attention weight on the 2j* - 1th input vector will be
Output of the first layer. Let 2 denote the output of the first layer on the ith input vector. By construction, with MLP and residual connection after the attention block, the output of the first layer is such that,
2- = [x,,0,1, 0,,, 1), TT(j")
(1)
(1)
2 x,, 0,
for i 1...., k
(1) 02., T(2i),, ((2)), for i 1, 1. 2
Second Layer. Since we have the address/key vector (22")) for the target label as the input in the first layer, it is straightforward to apply attention and (-)) with ReLU FFN to produce the desired label.
See that if the query vector Q( ( = ( T2j""), and the key vectors contain the 4th part of the
input vector, that is, K [zz (1) = [i))., then the dot product will be greater than 1 - 7-/100 with the desired input 2ÿ" and will be less than 7/100 for the rest of the inputs. The value vectors will be
assigned the second part of the input V(z
[O,, yi. ..... 0] and will V(zl2n) = [O..... 0] for all i 1,...,k k. Using the techniques used for the first layer, it is straightforward to that scaling of the query vector and applying (-)) to the output will produce the desired label yj*
33

---

Index Lookup Task
Index Lookup: Validation Curves
Number(233 31)
37.3
143
21.4
18.8
9.8
34
Nambs-22 5321
30
Validation Acourary (%6
Mamber2 1224)
Namba (2, 64)
Mamba <2, 32)
Model
0.6
Mambaa-k 2581
Nambe-2. 2580
96.7
28.7
07:1
521
30.4
11.5
0.4
Monber 1024)
00.5
29.4
97.5
823
34.1
22.0
20
75
100
202
400
50000
200000
150000
200260
250003
Input Length
Iterations
Figure 3: Performance of Mamba on the Index Lookup task across various lengths and widths. See Section H. for more details.
H Empirical Analysis: Additional Details and Experiments
In this section, we discuss the details of implementation and data generation for experiments described in Section Further, discuss some additional experiments the string equality task.
Implementation and hyperparameters. All of our implementations for experiments are based on PyTorch 44 For our experiments with Transformers, we use the Huggingface Transformers library [63] with the GPT-2 backbone well as our own custom implementation. We use PyTorch's 's standard or in-built implementation for experiments with LSTMs. For state-space models like Mamba [21] and Diagonal state-space models [22]. we use the official implementation provided by the authors for our experiments. For RetNet [56] and Linear Transformers, we use our own custom implementation.
For each task, we tune the models across several hyperparameters and report the results based on the best-performing model. We use grid search to tune the models for each task. For each architecture excluding LSTMs, we tuned across depths E {1, 2, 4, 6} and widths (64, 128, 256, 512). Since LSTMs with large depths are hard to train 43), we only consider LSTMs with depths up to 3. For Transformers and linear Transformer variants, we tune the heads across {4, 8). For all models, we tune the learning rate across {0.01, 0.005, 0.001, 0.0005, 0.0001, 0.00005, 0. 0.000001). J. We primarily use Transformers with absolute encodings for the result presented in the main paper. We train the models for up to 250k steps unless they achieve (almost) perfect validation accuracy earlier. After training, we evaluate the models on fresh set of 5000 examples.
Note, however, that the focus of our paper is slightly different and for some problems, we are interested in the behavior of small-sized (width or depth) models of one architecture and a relatively larger model for another architecture, For instance, in the index lookup task, we are concerned with how well one-layer Transformers with small widths fare against relatively larger recurrent models. Additionally, our focus is on the comparison of different architectures of fixed sizes across lengths and not primarily on how well models of the scale used practice perform these tasks on much larger lengths. Hence, we do not consider models with much larger depth or width.
Compute. All experiments 16GB memory and 16 NVIDIA GTX 1080 Ti GPUs each with 12GB memory. Each run for 500k steps could take between hour and 16 hours depending of the models. Some runs are much shorter if the model achieves high accuracy quite early (e.g. on lengths 20). While each individual run does
conducted using NVIDIA Tesla V100 GPUs each with
the length of the inputs and the size
take significant amount of time, tuning the models, particularly for the tasks where they fail to learn requires several runs and consequently a much longer duration. We estimate that all the runs across hyperparameterss, architectures, and input lengths took < 1200 GPU hours in total on the GPUs mentioned above.
H.1 Additional Experiments and Data Generation
Index Lookup Task. The input sequences are comprised of symbols Σ and a positional token in
[N]. In our experiments, the size of the set of symbols is |Σ] = 64 and N varies from 20 to 400. To create each example, we first sample length & uniformly between 10 and N (20 - -400) and then sample N symbols independently from Σ uniformly at random. Lastly, sample position between
34

---

1 and k uniformly at random and append it to the sequence to produce a labeled example for training or evaluation. During evaluation, the model is tested on sequences of length exactly N.
Bounded Dycks. For the Dyck-(2, 2) task, we generate the examples in such a way that with 0.5 probability, the generated string is well-balanced with depth most 2, and with 0.5 probability it does not belong to the language and has label O. The generated strings are of length exactly N in both cases. To generate positive examples we use the following strategy: we iterate over N steps and for each step, we check whether the current depth of the stack is < 2. If the depth of the stack is the sequence continues with the closing bracket corresponding to the open bracket in the stack. If the depth is < 2, the sequence is either continued by choosing one of the open brackets uniformly if the stack is empty or by choosing uniformly between open brackets and the closing bracket if the stack is non- non-empty. After N 2 steps, the generated string could be a well-balanced string of length N 2 in which case we add a depth 1 string of length 2 at the end which results in a well-balanced string of length N. Otherwise, the stack could be nonempty we add the remaining closing brackets which also leads to string of length N.
To generate negative examples, we first sample a positive example and then corrupt some of the symbols in the string. For strings of length N, we first sample number i 1,..., N/ 10 uniformly at random and then pick & different indices uniformly at random. For each of those positions, with probability 1/2, we swap the types of brackets, e.g. round '(" to square 'I', and with probability 1/2 we switch open brackets to closing brackets (or vice versa). There is a very small probability that after the corruption the resulting string will still be a valid well-balle cell-balanced string of depth at most in which case we redo the corruption again. The probability of the event is too low to affect the efficiency of the generation process.
Additional Experiment with Mamba. We explore how the size of the hidden state influences the performance of a Mamba model across various lengths on the Index Lookup task. We evaluate two-layer models of different widths {32, 64, 256, 512, 1024] across various lengths ranging from 20 to 400. We find a clear trend where the performance of Mamba models increases monotonically with the increase in the width of the model (see Figure 3p. While the performance does exactly scale linearly with width it still somewhat interesting that the trend exists. We did a similar experiment with LSTM but did not observe such a trend and for lengths above 100 the performance remained at chance level even when the width was increased. Note, however, that even with width of 1024, the performance of Mamba is still much worse than a one-layer Transformer with a width of 64.
H.2 String Equality Task
We explore a few different strategies to evaluate the performance of models on the string equality task. In these experiments, the goal of the models is to determine whether the first half of the input string and the second half of the string are equal. The first two experiments are in the standard classification setting but differ in the way the negative examples are created. The third experiment is in the next character prediction which is commonly used in prior works to test models on formal languages. For the equality task, we find that it is not straightforward to generate negative examples in way that the problem cannot be solved using shortcuts. Hence, we include the next character prediction setting which does not require the generation of negative examples.
Summary of results. The results on the String Equality task are relatively more nuanced than the results for index lookup and 2 On standard binary classification setup, one must make design choice regarding the generation of negative examples that could influence the difficulty of the task. We first conduct experiments in the classification setup by generating negative examples using two different strategies. On one task, we find that while recurrent models like LSTMs struggle beyond certain length, state-space models like DSS and Mamba are able to match Transformer's performance. In the second task, we find that all models are able to achieve near-perfect accuracy. We note that both classification tasks can be solved by using shortcuts based on the way the negative examples are created. Hence, we explore another strategy called the next character prediction setting inspired by prior works [57] (0) on empirical analysis on formal languages. We note that the task in that setting becomes almost identical to the copying task where model observes a sequence and then has to produce the same sequence. In that setting, we observe that at small width like 64, state-space models fail to perform the task accurately at certain lengths whereas Transformers with the same width succeed. For models with larger widths, we find that DSS and Transformers succeed at performing the tasks for the lengths considered in our experiments.
35

---

Eq-random Task
Eq-ncp Task
LiSTF-22 256)
202.0
33.5
33.5
51.2
50.7
LnTT-12, 641
25.9
35.6
0.0
0.
90
80
Archnoccurr
Architecture
RecNct-2, 256)
202.0
103.2
25.7
52.3
40.3
ACCLOO/
PA2M66122 641
20.8
20.0
02
0.
0.0
DSS-2, 236)
202.0
92.5
202.2
92.3
90.3
DS3-12, 64)
300.0
30.9
22.5
01
Namto 256)
1020
99.3
99.5
99.5
99.3
Mamba-12 12, 64)
35.9
19.3
05.6
0.0
0.0
-20
TF-02. 256)
TF-12
20
196
200
400
100
200
400
600
850
Input Length
Input Length
Figure Performance of architectures on the Equality task. See Section H.2 for more details.
Setup. In all our experiments with string equality, we have a vocabulary E which contains one token that is reserved as a separator token '[SEP]' and is placed between the first half and the second half of the input string. For each task, while creating an example we first pick the length l uniformly N N at random from IO IO + 2, TO + 4,...NN .. The choice of lengths is due to the requirement that the length of the strings must be even. During evaluation, we test the models strings of length exactly N.
(i) Eq-random. The first experiment is pretty straightforward and contains binary strings, i.c., {0, 1, [SEP]]. With probability 1/2, we create a positive example by first sampling a string of length k/2 uniformly followed by the separator token and then the same string again. By construction, the created example has a positive label. With probability 1/2, we create the second half of the string by sampling another string of the same length where each symbol is sampled uniformly at random. We assign the label based on whether or not it is equal to the first half but with high probability (1 2472 ) the example has negative label.
(ii) Eq-one In the second experiment, we have vocabulary with 1024 and we create the positive example in the same way as earlier in the Eq-one setting by sampling strings uniformly at random. To generate a negative example, we set the second half be equal to the first half and then change the symbol at only one of the positions. In other words, the second half matches the first half at all positions except one.
The training details and hyperparameters are almost identical to the experiments described earlier.
Results. On both the Eq-random task and Eq-one task we find that Transformers achieve near- rear-perfect accuracy on lengths up to 400. Recurrent models like LSTMs and linear Transformers struggle at lengths beyond 100 on the Eq-random task. On the other hand, we find that recently proposed state-space models like DSS and Mamba are able to match Transformers' performance on both tasks (See Figure left). On the Eq-one task, we find that for lengths up to 400, all models are able to achieve near- parfeerfect accuracy.
Remark. One thing to note is that both of the experimental setups described above can be solved using some form of shortcuts. In the first case, is straightforward to see that the first half and the second half will differ at about half the positions on average. A recurrent algorithm does not necessarily have to store the first N/2 elements to compute the output correctly with high probability. Even if it stores the first few elements, then with high probability it can determine the label of the string correctly. For the second case (Eq-one), even if it might seem difficult at first look, a recurrent model can solve it perfectly by just maintaining a dictionary of the counts of each symbol in the vocabulary. Since the first half and second half differ at exactly one position, the number of occurrences of at least one symbol will be different in the two halves.
To rule out such phenomena, we adopt the next character prediction setting
20 50
(iii) Eq-nep. In the next character setting (NCP), a model is required to predict the next set of valid continuations for every prefix of given string. It can be seen as a multi-label classification problem for every prefix of a given string. The prediction for a particular input string is considered correct the prediction for every prefix is correct.
In the context of the string equality task with length N, for the first N/2 symbols, all symbols are valid continuations, and hence the predictions for the first half of the string are trivial. After observing
36

---

the separator token [SEP], only one of the Σ symbols is allowed at every prefix until the end of the string. We note that the prediction problem becomes equivalent copying [28] a sequence of symbols. For our experiments the vocabulary size |Σ] = 1024. We explore sequences of higher lengths up to 800. In this setting, we find that when the model sizes are restricted, i.e., the width of the models is 64, state-space models such as DSS and Mamba struggle to perform better than chance-level accuracy for lengths over 400. contrast, Transformers are able to achieve near-perfect accuracy (See Figure 4). However, unlike the case of Index Lookup, we find that the DSS architecture in particular is able to solve the task for lengths up to 800 with larger widths the NCP setting.
Discussion. We discuss a few takeaways from our experiments. For the Index Lookup task. our results indicate that even small-sized one-layer Transformers can learn a lot more efficiently than recurrent models of much larger sizes. The experiments with bounded Dycks are primarily for one-layer Transformers and indicate that they learn at much slower rate than recurrent models like LSTMs and even two-layer Transformers. On the string equality task, the difference in performance between Transformers and recurrent models not as stark as the Index Lookup task, particularly with state-space models such as DSS. However, unlike Transformers, they seem struggle on long sequences the NCP setting when the widths of models are small.
37

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims are mentioned in the introduction and abstract reflect the main contributions the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitations of the work discussed in Section 7
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
38

---

Justification:
The assumptions for each result are mentioned in their respective Sections before the results are stated.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Yes, all the information regarding experiments are provided either Section
or H
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
39

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes)
Justification:
Yes we have included the code in the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
/ //nips. wwtppblcccgguieeeeeeeeooeeppeeee httpp/////g..dee/eaeaaaaaaaaae// SondeSuumissionPPee
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: The details are provided Sections
and
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
synthetic tasks where models either achieve near- perfect accuracy if they succeed or near chance-level accuracy if they fail. This large quantitative difference performance allows for clear comparison between models. As the test sets have 5,000 entries, Binomial confidence intervals for the accuracy would be very small and would not add further information to this comparison. Additionally, since our primary focus is on whether any single run of an architecture can successfully solve a task, we do compute eтror bars across multiple runs and use our compute to tune the learning rates and other hyperparameters extensively to rule out the failure of models due to lack of sufficient tuning.
Justification: Our experiments
Guidelines:
The answer NA means that the paper does not include experiments.
40

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The details are provided in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http://///eeeeee.cccc////////eea
Answer: [Yes]
Justification:
The research conforms with the NeurlPs Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper primarily studies the capabilities of neural sequence models theoret- ically. While the models themselves might have some societal impact, the results presented in the paper do not have any direct socictal impact.
41

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
We do not release any models or data.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: We do not use any existing assets apart from standard Python libraries which are cited appropriately.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
42

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thood cww datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide the code to reproduce our experiments with appropriate instructions and the relevant details regarding the experiments mentioned in the paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The work does involve any crowdsoureing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The work does not involve any crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
43

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
44


Reviews:
[{"review_id": "TqkE3DpbGJ", "reviewer": "Reviewer_LT9S", "summary": "In this paper, the authors study the representational separation results about two widely used classes of language models: transformers and RNNs. For a set of practically well-motivated tasks, they establish lower and upper bounds for attention based one-layer (and some for two) transformers and arbitrary RNNs using and building upon techniques from communication complexity. They also empirically validate their conclusions.", "questions": "1. I realized that the soft-max attention in line 121 is non-causal. Would your results change if it's causal attention?", "limitations": "Yes.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "First of all, I want to appreciate the authors for clearly elucidating the main ideas and the intuitions in a very accessible manner even for non-experts. Such style of writing is a rarity these days and the authors deserve the credit for this.\n\nOverall, it's a very important topic and problem of research interest to understand the differences between the two widely used models in the form of transformers and RNNs. They present a set of intereseting results for well-motivated tasks, and establish theoretical results for transformers needing to be of linear size in the input dimension, whereas RNNs could get away with logarithmic size, and vice versa.", "weaknesses": "While I appreciate the technical results of the paper, I am not fully sure about how meaningful these results are inorder to decipher the fundamental differences between RNNs and transformers. In particular, for the index lookup task, it makes sense that the only way RNNs can retrieve the symbol $s_p$ at an unknown location $p$ revealed at the end is only via storing all the past information in its hidden state and hence $m \\geq N/p$ is logical. However, if you feed the sequence in the order $p, s_1, \\ldots, s_N$, would the same result hold? I believe you could get away with $\\log N$ here too though I could be wrong. Shedding light on things like this could yield more insights about how these architectures are fundamentally different? Because for an index lookup task, if all we care about is retrieving a symbol at some position, we don't really care about which order you feed the input sequence right?\n\nAlso on a minor note, as the authors themselves acknowledged, these results might not have full bearing on the learnability settings. It would be interesting to see how these results hold for SGD learnability on the same tasks. Food for future thought.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments and time.\n\nResponses to the two points raised in your review below.\n\n> “In particular, for the index lookup task, it makes sense that the only way RNNs can retrieve the symbol $s_p$ at an unknown location revealed at the end is only via storing all the past information in its hidden state and hence $m\\geq N/p$ is logical. However, if you feed the sequence in the order $p, s_1,\\ldots, s_N$, would the same result hold? I believe you could get away with $\\log N$ here too though I could be wrong. Shedding light on things like this could yield more insights about how these architectures are fundamentally different? Because for an index lookup task, if all we care about is retrieving a symbol at some position, we don't really care about which order you feed the input sequence right?”\n\nWhile what you have mentioned is technically correct, that approach is specific to the particular version of the Index Lookup task and it does not overcome the fundamental issues with how RNNs process inputs. \n\nTo be clear, for the case where the models are provided sequences in the order $p, s_1,\\ldots, s_N$, RNNs can indeed get away with a hidden state of size $\\log N$. However, consider the following variant of the Index Lookup task (let’s call it Multi-Index Lookup) where the models are provided with $k =O(N)$ (e.g. N/2, N/4, etc)  indices for which they have to look up and produce the respective symbols. In particular, the input is $s_1, \\ldots, s_N, p_1, p_2, \\ldots, p_{k}$.  For each $p_i$ after the symbols $s_1, \\ldots, s_N$, the model is required to output the respective symbol at position $p_i$. Note that for one-layer Transformers, the construction for the Index Lookup task can be directly extended for this Multi-Index Lookup task implying that poly-log size Transformers can solve this. \n\nHowever, for RNNs, they must have a hidden state of size $\\Omega(N)$ to solve this correctly even if the positions are prepended to the symbols. Instead of a reduction from the $\\mathrm{Index}$ problem, we can show via a simple reduction from the Disjointness problem. In other words, if the input sequence is of the form $p_1, p_2, \\ldots, p_{k}, s_1, \\ldots, s_N, p_1, p_2, \\ldots, p_{k}$ and an RNN can produce the required outputs then two parties Alice and Bob can compute Disjointness using $mp$ bits which implies that $mp$ must be $\\Omega(N)$. \n\nThe argument is quite straightforward. Let’s say Alice and Bob have two bits strings $a$ and $b$ of size $k=O(N)$ and both have access to the same RNN. For simplicity, they have agreed that Alice will place her bits in the first $k$ indices sequentially. Alice can then take an RNN, and provide it with indices $1$ to $k$ as input followed by $k$ symbols corresponding to $a$. The last $N - k$ symbols can be arbitrary. Alice can then send the hidden state to Bob which requires $mp$ bits. Bob can provide the same set of indices $1$ to $k$ to the RNN and then record the output. Bob can then compute Disjointness and hence it follows that $mp= \\Omega(N)$. \n\n**Summary:** While it is true that for the Index Lookup task prepending the input sequence with the position/indices can make a difference for RNNs, the limitation/lower bound still applies for a natural extension of the task involving multiple indices. For us, the goal of introducing the Index Lookup was to come up with the simplest task to describe our tools. However, based on the point raised by you,  we see that it could create confusion, and will make sure to clarify this in the next version of the paper. \n\n\n\n> “I realized that the soft-max attention in line 121 is non-causal. Would your results change if it's causal attention?”\n\nUsing causal attention does not affect any of the constructions or lower bounds for Transformers in the main paper, i.e., Index Lookup, Equality, Disjointness, Nearest Neighbor, etc. It only affects the construction of Transformers presented in Section F.3 for a more abstract yet general class of Boolean functions.", "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": "F6ECIw7NuE", "reviewer": "Reviewer_FtyK", "summary": "This paper analyzes the differences in terms of representations between Transformers and recurrent architectures. They highlight multiple cases: a) a setting where 1-layer Transformer can represent the task with a log number of parameters but not RNNs (index lookups) b) a case where RNNs can represent the task with a log number of parameters but not a 1-layer Transformer (Dyck language) c) a case where both 1-layer Transformer and RNNs cannot represent the task (boolean tasks like string equality). Then finally show that a 2-layer Transformer can represent the string equality task and other associative recall tasks such as nearest neighbors. They end up with some experiments to validate their theorems: they show that 1-layer transformers are great at learning the index lookup task and that RNNs learn very quickly the bounded dyck languages.", "questions": "I mentioned my questions in the weaknesses section.", "limitations": "The authors mention the limitations of their work.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "I think that the paper is of a great value to the community. There have been many papers in the community that the study the cases where  Transformers are superior over RNNs [1] and where RNNs are superior to Transformers [2]. It is great that this paper gives a unified picture of the strengths and weaknesses of both models. Besides, their analysis is clean: the theorems and proof sketch are very easy to follow. For this reason, I advocate for acceptance for this paper.\n\n\n\n\n[1] Jelassi, S., Brandfonbrener, D., Kakade, S. M., & Malach, E. (2024). Repeat after me: Transformers are better than state space models at copying. arXiv preprint arXiv:2402.01032.\n[2] Liu, B., Ash, J., Goel, S., Krishnamurthy, A., & Zhang, C. (2024). Exposing attention glitches with flip-flop language modeling. Advances in Neural Information Processing Systems, 36.", "weaknesses": "There are not a lot of flaws for this paper in my opinion. I will just give some suggestions: \n\n- I believe that the authors could further improve the presentation of their results. In particular, it is not clear at the beginning why they choose the tasks they propose and it sounds a bit like a \"catalog\". I think the authors should clearly say that they consider a case where a 1-layer transformer with log(L) params fail but not RNN, another case where RNN with log(L) params fail but not a 1-layer etc. Maybe it may be worth adding a table where the column headers are the tasks and the row headers are the two models and each entry contains the complexity of each model at each task. \n\n- Besides, I liked a lot the discussion from lines 210 to 218 and I believe that this should be earlier in the paper (even in the introduction). I think this point is central to understand one difference between Transformers and RNNs. I think that the point raised by the authors is not totally novel since a similar behavior has been reported by [1] in the case of the copy task. \n\n- Regarding the superiority of RNNs over Transformers in the case of the bounded Dyck languages, do the authors have a similar discussion to add? If yes, could they add it? Is the explanation similar to the one advanced by [2]?\n\n- Lastly, I don't know if it is possible to give any intuition about the limitation of 1-layer Transformers and RNNs at representing boolean functions? \n\n- I am not sure to follow the experiments for the bounded Dyck languages. The authors say that the 1-layer Transformer achieve near perfect accuracy up to lengths 100 but the red curve in figure 2 right is around 60% for all the iterations. Did I miss something?\n\n\n[1] Jelassi, S., Brandfonbrener, D., Kakade, S. M., & Malach, E. (2024). Repeat after me: Transformers are better than state space models at copying. arXiv preprint arXiv:2402.01032.\n[2] Liu, B., Ash, J., Goel, S., Krishnamurthy, A., & Zhang, C. (2024). Exposing attention glitches with flip-flop language modeling. Advances in Neural Information Processing Systems, 36.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments and time.\n\nResponses to the individual comments below.\n\n> “I believe that the authors could further improve the presentation of their results. In particular, it is not clear at the beginning why they choose the tasks they propose and it sounds a bit like a \"catalog\". I think the authors should clearly say that they consider a case where a 1-layer transformer with log(L) params fail but not RNN, another case where RNN with log(L) params fail but not a 1-layer etc. Maybe it may be worth adding a table where the column headers are the tasks and the row headers are the two models and each entry contains the complexity of each model at each task.”\n\nThank you for the suggestion. We agree that categorizing with respect to the separations (log L vs L) could be more helpful. We considered adding a table/figure at the top of the second page. We weren’t sure about the best way to depict the results and due to lack of space, we decided to go without it. We will consider adding these in the next version.\n\n\n> “Besides, I liked a lot the discussion from lines 210 to 218 and I believe that this should be earlier in the paper (even in the introduction). I think this point is central to understand one difference between Transformers and RNNs. I think that the point raised by the authors is not totally novel since a similar behavior has been reported by [1] in the case of the copy task.”\n\nThanks for the suggestion and we will consider adding it to the introduction in the next version. Regarding your point about [1], we do not claim novelty about the intuition. We believe other researchers in the community may have a similar intuition behind the differences in the two architectures. In our work, that intuition is formalized using communication complexity for lower bounds and the JL vectors for upper bounds. In their work [1], they use a different approach to derive a lower bound for the copying task. However, it is worth noting that using communication complexity-based techniques as used in our work, it is also possible to show lower bounds for RNNs to perform the copying task. We will consider adding that somewhere in the appendix for readers who might find it interesting.\n\n> “Regarding the superiority of RNNs over Transformers in the case of the bounded Dyck languages, do the authors have a similar discussion to add? If yes, could they add it? Is the explanation similar to the one advanced by [2]?”\n\nWe have some discussion on this at the beginning of Section F, which you might find helpful. However, the discussion is intertwined with the proof technique, so it might be less straightforward than the former discussion. The general intuition is that bounded Dycks can be represented by constant-sized DFAs, implying that RNNs require very little (constant) memory while processing them sequentially. In contrast, a 1-layer Transformer must compute the same based on a convex combination of all input symbols (attention) and classify 0 or 1 based on this vector (using MLP). The weights of the convex combination are determined by dot products of query and key projections, and hence cannot be arbitrary.\n\nIn simpler terms, the proof intuitively shows that if a 1-layer Transformer can recognize bounded Dyck languages, then the query and value vectors must contain sufficient information, $\\Omega(N)$ in bits. Think of it this way: the width of the query/key vectors determines the flexibility of the attention weights, and the width of the value vectors or input embeddings determines the amount of information stored in the output of the attention block. The proof essentially shows that both must be sufficiently wide (contain enough bits of information) for any arbitrary MLP block to be able to classify correctly based on the attention block's output.\n\n> “Lastly, I don't know if it is possible to give any intuition about the limitation of 1-layer Transformers and RNNs at representing boolean functions?”\n\nFor RNNs, even for Boolean functions, the key intuition is the same as the one explained for the Index Lookup task. The goal of introducing that task was to highlight that intuition in a natural manner since it may be less obvious in the context of Boolean functions. For 1-layer Transformers, the intuition is the same as the one described above. \n\n> “I am not sure to follow the experiments for the bounded Dyck languages. The authors say that the 1-layer Transformer achieve near perfect accuracy up to lengths 100 but the red curve in figure 2 right is around 60% for all the iterations. Did I miss something?”\n\nFor bounded Dyck languages, we find that 1-layer Transformers struggle at lengths 200 and above, unlike RNNs and 2-layer Transformers. Figure 2 right depicts the validation curve for models at length 400 (and not 100). The goal of that line about lengths 100 is to clarify that this difference is performance occurs at lengths above 100 or so which is natural since the difficulty of the tasks increases with length. For the Index Lookup task, all models solve the task at length 20 and the difference is more stark at length 100 and above (see Figure 2 left).", "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": "Xcdxj6TnZX", "reviewer": "Reviewer_hDB9", "summary": "This works studies the differences between Transformers and recurrent models with respect to 4 tasks: index lookup, associative recall, string equality and Bounded Dyck languages. The authors prove that for index lookup and nearest neighbor recall, there exists a 1-layer Transformer that needs poly-logarithmic width in the number of inputs while recurrent architectures require at least linear width. On the other hand, 1-layer Transformers require at least linear width for checking string equality and Bounded-Dyck languages. The authors strengthen their claims with experiments on synthetic datasets.", "questions": "To reiterate the points raised above\n1. How do we extend these results to deeper networks?\n2. Are the results relevant for other task or other domains and is the relevance of these results very narrow?\n3. Is there a way to motivate the selection of the tasks considered in this work?", "limitations": "The authors sufficiently acknowledge limitations. In particular they discuss that existence proofs do not guarantee such a solution is always found using gradient descent. The authors also comment on the limitations of their results to 1-layer Transformers and how they do not extend to 2-layer Transformers.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "I found the following techniques were interesting: (1) the usage of Johnson-Lindenstrauss lemma to construct solutions with poly-logarithmic width in the number of inputs. (2) Using results from communication complexity to prove results for the lower-bound for the model width. While the ideas relating relating to communication complexity have been used in Sanford et al. (which is acknowledged by the authors), I nevertheless found the presentation of the lower-bound results to be interesting. I also found the writing to be refreshing; the authors are careful when interpreting the results and do not overstate the importance of their results while being precise in their claims.", "weaknesses": "Overall, I worry that the contributions of the work don't make progress towards the motivating questions in the introduction: why Transformers have supplanted  LSTMs (L29) in many problems. The analysis is specific to synthetic tasks with 1-layer recurrent models and Transformers (https://arxiv.org/abs/2105.11115, https://arxiv.org/abs/2208.01066). The theorems comment on the width of 1-layer Transformers but they do not seem to have any relevance to Transformers with multiple layers or when trained on other tasks. \n\nWith the goal of differentiating Transformers from LSTMs, I wonder why these particular tasks were chosen and what makes them important. It would help to motivate why these tasks are important and for example, explains why Transformers are better for language modelling. The different tasks are often motivated for specific contexts such as in-context learning (associative recall) or natural language itself (Dyck languages - https://arxiv.org/abs/2105.11115). Furthermore, the results also suggest that scaling up both Transformers and recurrent networks suffices for tackling any of the tasks.\n\nAnother concern is that the results are primarily for 1/2-layer Transformers. However, it is more relevant to study Transformers with multiple layers and it is unclear how to extend these results or if they have any implications for deeper Transformers. For example, we know that Dyck-languages can be represented by deeper Transformers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments and time.\n\n> “The analysis is specific to synthetic tasks with 1-layer recurrent models and Transformers …”\n\n> “How do we extend these results to deeper networks?”\n\nWe wish to clarify what may be a misunderstanding here – we should have been more explicit about this when we defined recurrent models on pg. 3. While our lower bounds for Transformers apply only to 1-layer models, our results for recurrent models apply to any depth. The way we have defined RNNs (Section 2), the transition function is allowed to be an arbitrary function which is very general and includes multi-layer versions of all architectures used in practice. When we state that an RNN needs a hidden state of size $\\Omega(N)$ to solve a task, what we mean is that for any RNN with L layers, the sum of the size of hidden states of all layers has to be $\\Omega(N)$ and thus the size of the model is $\\Omega(N)$. Hence, our results show some tasks that can be solved with 1 or 2-layer Transformers (e.g. Index Lookup, Equality, Nearest Neighbour, etc) with poly-log size whereas any **multi-layer RNN** must have at least **linear size** (exponentially larger) to solve them. We will emphasize this point more clearly in Section 2 to avoid this confusion.\n\nWe agree with the reviewer that understanding the limitations of multi-layer Transformers (TF) is very important. However, proving lower bounds even for 1-layer TF has been challenging—there have been some very recent results [4] and most questions about (multi-layer) TF with sub-linear width remain open. Our lower bound results for 1-layer TF are more general and more interestingly apply to tasks solvable by small-sized RNNs or small-sized 2-layer TF, thus showing a clear separation. Existing lower bounds for multi-layer TF are not solely information-theoretic and rely on unproven conjectures, e.g. from circuit complexity.  The situation mirrors somewhat that in computational complexity, where proving circuit lower bounds even for depth 2 or 3 has proved challenging for decades (https://shorturl.at/6lwUh). \n\nEven for 1-layer Transformers (TF) we believe there are many interesting unanswered questions, some of which are answered in our work. (Q1) Are there tasks that cannot be solved by small 1-layer TFs but are solvable by much smaller 2-layer TF?  (Q2) Since 1-layer TFs can solve tasks like Index lookup with poly-log size, while any multi-layer RNN needs linear size, are there tasks that small RNNs can solve but not 1-layer TFs?\n\n (1) We show an exponential depth separation for tasks like Equality, Dyck, and Disjointness, where 1-layer TFs need linear width but 2-layer TFs solve them with log width. Such depth separation questions are widely studied for feedforward networks (see [1, 2, 3]). As far as we know, our results are the first to show exponential depth separation between 1-layer and 2-layer TFs on multiple tasks. \n\n(2) Since 1-layer TFs of poly-log size can express tasks that multi-layer RNNs with sublinear size cannot, the lower bound on Dyck shows a task solvable by constant-sized RNNs but not by 1-layer TFs.\n\nRegarding upper bounds or constructions for 1 or 2-layer TF, we see it as a strength rather than a limitation. We show multiple tasks that 1 or 2-layer TF can represent, whereas any multi-layer RNNs must be much larger regardless of depth. \n\n[1] The Power of Depth for Feedforward Neural Networks.  \n[2] Depth Separation for Neural Networks.  \n[3] Representation benefits of deep feedforward networks.  \n[4] Representational strengths and limitations of transformers.  \n\n> “With the goal of differentiating .. why these particular tasks were chosen and what makes them important …  for language modelling.”\n\n> “Is there a way to motivate the selection of the tasks in this work?”\n\nWe discussed the task choices in the Introduction and further for Dycks in Section F. Briefly, our goal was to focus on tasks that appear natural or have strong justification in ML or linguistics research. For instance, in the MQAR task introduced in [1], they observed a perplexity gap in natural language modeling between Transformers (TF) and non-TF LMs, finding that non-TF language models struggled with texts that precisely mimic the MQAR task. Similarly, for the general version, i.e., the nearest neighbor task, multiple works [2, 3] found TF-based LLMs can mimic such behavior in practice.  Dyck languages have been extensively studied in formal language theory and ML research. Tasks like String Equality and Index Lookup are natural tasks that are likely to be relevant as primitives in search and learning tasks. For example, in code execution, given `arr = [6, 2, 8, 2, …, 3]` followed by `print(arr[i])`, the task is essentially Index Lookup. Models might also perform such lookups when provided with line numbers for debugging.\n\n\n[1]  arxiv.org/abs/2312.04927  \n[2] arxiv.org/abs/2310.03016  \n[3] arxiv.org/abs/2404.11018  \n\n\n> “The analysis is specific to synthetic tasks…”\n\n> “Are the results relevant for other task or other domains and is the relevance of these results very narrow?”\n\nIn the main paper, we discuss upper bounds for Transformers for specific tasks relevant to practice. In Section F, we show Transformers can represent a broader class of Boolean functions which include subclasses like k-CNFs/k-DNFs. Our lower bounds for multi-layer RNNs and 1-layer Transformers apply to any function with communication complexity ~ N, including most Boolean functions. Since the work aims to understand the problems theoretically, the problems must be mathematically well-defined and hence synthetic by definition. We're unsure what is meant by \"other domains,\" but as these models are applied to more domains beyond NLP, like formal mathematics or algorithm design, understanding their capabilities/limitations along the lines we have considered in our paper would be imperative.  We hope the techniques/framework presented in our work will be useful for such analyses.", "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": "taMtVj9xv8", "reviewer": "Reviewer_BJPt", "summary": "This paper demonstrates theoretical separations in the representational abilities of Transformer and Recurrent Architectures on selected synthetic tasks, including index lookup, nearest neighbor, recognizing bounded Dyck languages, and string equality. The class of recurrent architectures examined includes popular networks such as RNNs, State Space Models, and Mamba. The results reveal significant differences in the required model sizes for each architecture to effectively represent these tasks. Additionally, the authors present some experiments that highlight the optimization efficiencies of architectures belonging to both classes on the selected tasks on practical-size sequences.", "questions": "- Do you have any guess of which Boolean functions are conjectured to be hard to compute for 2-layers Tranformers of poly-logarithmic size?\n\n- It seems your separation results are worst-case over the input space. Can you argue whether all/some separations still hold for typical inputs, if for instance a uniform distribution over the input space is assumed?", "limitations": "Limitations and societal impact are well addressed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1) This is a theoretically solid paper. The proof techniques are novel and well explained, and they might be useful for further analysis of modern deep learning architectures.\n\n2) The architectures considered are of practical interest, and the analysis contributes to a deeper understanding of the power and limitations of these architectures.\n\n3) The paper extends previously known separation results to a set of diverse and more real-world applicable tasks.", "weaknesses": "1) The experimental part is fairly brief. The analysis of computational or statistical learning complexity is left to future work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments and time.\n\nResponses to the individual comments below.\n\n\n\n> “The experimental part is fairly brief. The analysis of computational or statistical learning complexity is left to future work.”\n\nWe respectfully disagree with this assessment. The central claims of the paper are theoretical in nature and we believe the empirical results are sufficiently extensive to support them and provide some context for the theoretical results.  \n\nOur work does include experiments for all primary tasks considered in the paper – apart from those related to associative recall, where extensive empirical evidence has already been provided in earlier works (see citations in the paper). Additionally, we consider five different recurrent architectures spanning traditional architectures such as LSTMs to more recently proposed ones such as Mamba, DSS, etc. Given the kind of statements made in the theoretical results, it seems natural to explore with experiments whether differences of a similar nature are exhibited in practice. We believe we have studied that fairly extensively. While it is true that experiments analyzing the statistical learning complexity of models are not present, we believe that it is far and somewhat detached from the central claims of the paper. \n\n\n> “Do you have any guess of which Boolean functions are conjectured to be hard to compute for 2-layers Tranformers of poly-logarithmic size?”\n\nWe do not have any strong conjectures at the moment, but there are a couple of directions worth thinking about. \n\n- In our paper, we show that 2-layer Transformers with poly-log size can represent k-DNFs/k-CNFs with at most N terms. We suspect that for k-DNFs/k-CNFs with cubic ($\\Omega(N^3)$) or more terms may be difficult for two-layer Transformers to represent with poly-log size. \n- Due to circuit complexity bounds, functions outside of the complexity class TC0 are impossible for transformers even of polynomial size, and thus in particular for 2-layer transformers of poly-logarithmic size. This is likely to include foundational problems such as the evaluation of Boolean formulas. However, it is unclear if such tasks can represented by small-sized RNNs/SSMs.\n\n\n> “It seems your separation results are worst-case over the input space. Can you argue whether all/some separations still hold for typical inputs, if for instance a uniform distribution over the input space is assumed?”\n\nThanks for raising this question. We have some thoughts on this. Currently, most of the results in our paper are of a deterministic nature in the sense that a model is said to express a function $f$ if it can produce the same output as $f$ for all inputs in the domain of $f$. This is more typical when analyzing expressivity in the context of classification, i.e., $f: X \\rightarrow \\{0, 1\\}$.\n\nUpon some thinking over the past week,  we believe that some of our lower bounds can indeed be extended to the average case on certain distributions, such as the uniform one. For instance, there is a well-studied Boolean function called Inner-product mod 2 (IP2 in short), for which we can show that any RNN or 1-layer Transformer must have linear size to achieve low error with respect to the uniform distribution. The proof follows from the arguments in Section C.2 in Vardi et. al [1]. We think that some of the lower bounds in our paper, such as Index Lookup and nearest neighbor/MQAR, can be extended to various distributions based on reductions from the IP2 problem. We will verify the proofs and include them in the next version of our paper.\n\n\n[1] Vardi et. al. Size and Depth Separation in Approximating Benign Functions with Neural Networks. COLT 2021.", "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:
We thank all the reviewers for their thoughtful feedback and their time. We are encouraged to see that they found our results interesting (Rev *LT9S, FtyK, BJPt*), well-motivated (Rev *LT9S,  BJPt*), and to be of value to the community (Rev *FtyK*). We are further pleased to see that they found our proof techniques to be interesting (Rev *LT9S, BJPt, hDB9*), clean (Rev *FtyK*), and well-explained (*all reviewers*). 

In this work, we show that various tasks can be solved by small-sized (poly-log size with 1 to 2 layers) Transformers, whereas any multi-layer RNN must be exponentially larger in comparison. Additionally, we show that 1-layer Transformers cannot solve certain tasks with sublinear width, whereas they can be solved by either RNNs or 2-layer Transformers of much smaller size.


We have addressed the weaknesses and specific questions from each reviewer in the individual responses. Below, we summarize the key aspects of our responses. Please refer to the individual responses for more details.


------------------------------------------------


Reviewer LT9S mentioned that the Index Lookup task for which we provide a linear lower bound for RNNs can be substantially easier if the inputs are provided in a different way. In the rebuttal, we explain why it does not solve the core issue and show that a natural extension of the task will still be hard for RNNs even if the inputs are provided in a different manner.

Reviewer FtyK provided some suggestions and had some questions, which we have answered in the individual response.

Reviewer hDB9 mentioned that our analysis is specific to 1-layer recurrent models and Transformers. We believe there may be a misunderstanding. We clarify that our lower bounds for recurrent networks apply to models of any depth, not just 1-layer RNNs. Further, we discuss the implications and relevance of the lower bounds for 1-layer Transformers and other questions raised by the reviewer in the individual response.

Reviewer BJPt stated that the experimental part of our paper is brief and does not include experiments to analyze statistical learning complexity. We argue that the experiments are supportive in nature and are sufficiently extensive to provide context for the theoretical results in the paper. We address the specific questions raised by the reviewer in the individual response.

Meta-review: The work gives several theoretical separation results between Transformers and recurrent models. For example, the authors show that Transformers can implement index lookup, string equality and nearest neighbors, but RNNs require linear size to solve these problems. On the other hand, they show a lower bound for single layer Transformer on recognizing Dyck languages, which can be solved with RNNs.

As all reviewers gave a positive evaluation of this work, and agree that it is a strong submission, I recommend that it is accepted.
Decision: Accept (poster)