SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Clayton Sanford, Daniel Hsu
{clayton, djhsu}@cs.columbia.edu
Matus Telgarsky
Courant Institute
matus.telgarsky@nyu.edu
Attention layers, as commonly used in transformers, form the backbone of modern deep learning, yet there is no mathematical description of their benefits and deficiencies as compared with other architectures. In this work we establish both positive and negative results on the representation power of attention layers, with a focus on intrinsic complexity parameters such as width, depth, and embedding dimension. On the positive side, we present a sparse averaging task, where recurrent networks and feedforward networks all have complexity scaling polynomially in the input size, whereas transformers scale merely logarithmically in the input size; furthermore, we use the same construction to show the necessity and role of a large embedding dimension in a transformer. On the negative side, we present a triple detection task, where attention layers in turn have complexity scaling linearly in the input size; as this scenario seems rare in practice, we also present natural variants that can be efficiently solved by attention layers. The proof techniques emphasize the value of communication complexity in the analysis of transformers and related models, and the role of sparse averaging as a prototypical attention task, which even finds use in the analysis of triple detection.
In recent years, transformer networks [Vaswani et al., 2017] have been established as a fundamental neural architecture powering state-of-the-art results in many applications, including language modeling [OpenAI, 2023], computer vision [Dosovitskiy et al., 2021], and protein folding [Jumper et al., 2021]. The key building block of transformer models is the self-attention unit, a primitive that represents interactions among input elements as inner-products between low-dimensional embeddings of these elements.
The success of transformer models is linked to their ability to scale their training and generalization performance to larger datasets and sequence lengths. Their representational capacity, however, underlies this scaling power, and is tied to the inductive biases of their learning algorithms. Empirically, transformer models trained with gradient-based learning algorithms exhibit biases towards certain algorithmic primitives [Edelman et al., 2022, Liu et al., 2022] and learn representations that may encode domain-specific information in the self-attention units [Clark et al., 2019, Hewitt and Manning, 2019, Rogers et al., 2020, Chen et al., 2022]. These examples indicate that transformer architectures not only provide computational benefits, but also have representational capabilities that are particularly well-matched to practical tasks.
In this paper, we investigate these inductive biases by identifying "natural" computational tasks for which transformers are well-suited, especially compared to other neural network architectures, as well as tasks that highlight the limitations of transformers. The tasks—sparse averaging, pair-matching,
and triples-matching—represent primitive operations that aggregate structural information encoded in embeddings. We use these tasks to elucidate the relationship between the embedding dimension  $m$  of a self-attention unit and its expressivity, and to showcase the fundamental representational limitations of self-attention layers.
In our model, the primary computational bottleneck faced by a transformer in computing a "sequence-to-sequence" function  $f\colon \mathcal{X}^N\to \mathcal{Y}^N$  is the constrained processing of pairs of input elements  $\{x_i,x_j\} \in \binom{\mathcal{X}}{2}$ ; we allow transformers unbounded computational power when processing the individual elements  $x_{i}\in \mathcal{X}$ . This is motivated by modern scaling regimes where the context length  $N$  has rapidly increased, the self-attention embedding dimension  $m$  remains much smaller than  $N$ , and the parameterization of multi-layer perceptrons (MLPs) that operate on individual elements is much larger than  $m$ . Indeed, the largest GPT-3 model [Brown et al., 2020] features a context length  $N = 2048$ , an embedding dimension  $m = 128$ , and MLPs with a 12288-dimensional parameterization; the context length of GPT-4 is as large as  $N = 32000$ . As such, we are interested in the capabilities of transformers with  $N^{o(1)}$  total "size", as opposed to  $N^{\Omega (1)}$ . The nature of the bottleneck in our model makes the tools of communication complexity indispensable for formalizing computational limits.
Sparse averaging separations among atomic self-attention units. The  $q$ -sparse averaging task  $q\mathrm{SA}$  aims to capture the essential approximation-theoretic properties of self-attention units. In  $q\mathrm{SA}$ , the  $i$ th input  $x_{i}$  is a pair  $(y_{i},z_{i})$ , where  $z_{i} \in \mathbb{R}^{d^{\prime}}$  is the data part of  $x_{i}$ , simply a vector in  $\mathbb{R}^{d^{\prime}}$ , whereas  $y_{i} \in \binom{[N]}{q}$  is the indexing part, which specifies  $q$  locations in the input sequence; the  $i$ th output element in  $q\mathrm{SA}$  is obtained by averaging the  $q$  data parts  $z_{j}$  given by  $j \in y_{i}$ , meaning
$$
q \operatorname {S A} \left(\left(y _ {1}, z _ {1}\right), \dots , \left(y _ {N}, z _ {N}\right)\right) = \left(\frac {1}{q} \sum_ {j \in y _ {1}} z _ {j}, \dots , \frac {1}{q} \sum_ {j \in y _ {N}} z _ {j}\right).
$$
(See also Definition 4.) As summarized in the following informal theorem, our analysis of  $q\mathrm{SA}$  in Section 3 and Appendix A illustrates the ability of the self-attention primitive to associate arbitrary subsets of input elements (as opposed to just "local" subsets, as specified by some sequential/topological structure), measures the expressive power accrued by increasing the embedding dimension  $m$  of a self-attention unit, and indicates the representational limitations of "traditional" neural architectures on basic computational tasks.
Informal Theorem 1. The task qSA for  $q \in \mathbb{Z}_{+}$  satisfies the following properties (see Definition 4 for a formal definition and approximation metric).
1. There exists a unit of self-attention  $f$  with an  $m$ -dimensional embedding that approximates qSA if and only if  $m \gtrsim q$  (Theorems 2 and 4).
2. Any fully-connected neural network whose output approximates qSA requires its first hidden layer to have width at least  $\Omega (Nd)$  (Theorem 10).
3. Any recurrent neural network whose iterates approximate qSA requires a hidden state of at least  $\Omega(N)$  bits (Theorem 11).
We consider the  $q$  SA implementation in Item 1 efficient since the dimension of the model parameters grows with  $\mathrm{poly}(q,d,\log N)$ , whereas the latter two are inefficient since their parameter (or state) dimension grows as  $\mathrm{poly}(N)$ . The proofs of the positive results employ embeddings for each index  $j$  and each subset  $y_{i}$  that have large inner products if and only if  $j\in y_{i}$ . The negative results involve communication complexity reductions and geometric arguments. These arguments naturally introduce a dependence on bits of precision, which we suppress above within the notation " $\gtrsim$ "; we note that these bounded-precision results are arguably more relevant to modern networks, which uses as few as 4 or even 2 bits of numerical precision.
Contrast between pairwise and triple-wise matching with self-attention layers. We frame standard transformer architectures as being able to efficiently represent functions that are decomposable
<sup>1</sup>Note, however, that attention units are permutation equivariant, so the order of elements in the input "sequence"  $X \in \mathcal{X}^N$  is irrelevant. In practice, positional encodings are used when the sequence order is relevant.
into sparse pairwise interactions between inputs. To do so, we introduce two sequential tasks and prove a collection of constructions and hardness results that characterize the abilities of transformers to solve these tasks.
Given an input sequence  $X = (x_{1},\ldots ,x_{N})\in [M]^{N}$  (for some  $M = \mathrm{poly}(N)$ ), we formalize the problems of similar pair detection (Match2) and similar triple detection (Match3) as
$$
\operatorname {M a t c h} 2 (X) _ {i \in [ N ]} = \mathbb {1} \left\{\exists j \text {s . t .} x _ {i} + x _ {j} = 0 (\mathrm {m o d} M) \right\}, \tag {1}
$$
$$
\operatorname {M a t c h} 3 (X) _ {i \in [ N ]} = \mathbb {1} \left\{\exists j _ {1}, j _ {2} \text {s . t .} x _ {i} + x _ {j _ {1}} + x _ {j _ {2}} = 0 (\mathrm {m o d} M) \right\}. \tag {2}
$$
For both tasks, note that the output is an  $N$ -dimensional vector whose  $i$ th element is 1 if and only if the sequence  $X$  includes a pair or triple containing  $x_{i}$ . In this sense, the problems differ from 2SUM and 3SUM, which are not sequence-to-sequence tasks.
We believe these two tasks are intrinsically "pairwise" and "triple-wise", respectively; moreover, since we also believe self-attention performs a fundamentally "pairwise" operation, we will use Match2 and Match3 to show a sharp gap in the representation power of self-attention.
Informal Theorem 2.
1. A single unit of standard self-attention with input and output MLPs and an  $O(d)$ -dimensional embedding can compute Match2 (Theorem 6).
2. A single layer of standard multi-headed self-attention cannot compute Match3 unless its number of heads  $H$  or embedding dimension  $m$  grows polynomially in  $N$  (Theorem 7).
3. A standard transformer model can efficiently compute a modified version of Match3 that makes assumptions about embedding structure or locality (Theorems 8 and 9).
4. Under a generalized notion of "third-order tensor self-attention" introduced in Appendix C.3, Match3 is efficiently computable with a single unit of third-order attention (Theorem 18).
While the above result demonstrates the limitations of multi-headed self-attention and illustrates the importance of learning embeddings with contextual clues, we believe that a stronger result exists. Specifically, we conjecture that even multi-layer transformers are unable to efficiently compute Match3 without hints or augmentation.
Informal Conjecture 1. Every multi-layer transformer that computes Match3 must have width, depth, embedding dimension, or bit complexity at least  $N^{\Omega(1)}$ .
In Appendices C.5 and C.6, we give a heuristic information-theoretic argument to support this conjecture, prove a matching upper-bound, and finally prove analogous results for graph-augmented transformers with respect to the problem of cycle detection in directed and undirected graphs.
Several computational and learning-theoretic aspects of transformers, distinct from but related to the specific aims of the present paper, have been mathematically studied in previous works.
Universality and Turing-completeness. To demonstrate the power of transformers, universal approximation results for transformers [Yun et al., 2020, Wei et al., 2022]—analogous to results for feedforward networks [Hornik et al., 1989]—establish the capability for sufficiently large networks to accurately approximate general classes of functions. Note, however, that the precise minimal dependence of the required size (e.g., number of attention units, depth of the network) as a function of the input size  $N$  does not directly follow from such results, and it is complicated by the interleaving of other neural network elements between attention layers. (Approximate) Turing-completeness of transformers demonstrates their power in a different manner, and such results have been established, first assuming infinite precision weights [Pérez et al., 2019] and later also with finite-precision [Wei et al., 2022]. Such results are more closely aligned with our aims, because Turing machines represent a uniform model of computation on inputs of arbitrary size. Wei et al. [2022] showed that Turing machines that run for  $T$  steps can be approximated by "encoder-decoder" transformers of depth  $\log(T)$  and size polynomial in  $\log(T)$  and the number of states of the Turing machine (but the decoder runs for  $T$  steps).
Formal language recognition. The ubiquity of transformers in natural language understanding has motivated the theoretical study of their ability to recognize formal languages. On the positive side, Bhattamishra et al. [2020] constructed transformers that recognize counter languages, and Yao et al. [2021] showed that transformers of bounded size and depth can recognize Dyck languages that have bounded stack depth. Liu et al. [2022] showed that the computations of finite-state automata on sequences of length  $N$  can be performed by transformers of depth  $\log (N)$  and size polynomial in the number of states. On the negative side, Hahn [2020] showed limitations of modeling distributions over formal languages (including Dyck) with fixed-size transformers (though this result does not imply quantitative lower bounds on the size of the transformer). Hahn [2020], as well as Hao et al. [2022], also establish the inability of "hard attention" Transformers to recognize various formal languages and circuit classes by leveraging depth reduction techniques from circuit complexity [Furst et al., 1984].
Learnability. The sample complexity of learning with low-weight transformers can be obtained using techniques from statistical learning theory and, in turn, establish learnability of certain boolean concept classes (e.g., sparse parity) [Edelman et al., 2022, Bhattachamishra et al., 2022] using transformer-based hypothesis classes. Our  $q$ SA function is inspired by these classes, and we establish concrete size lower bounds for approximation (and hence also learnability) by transformers. We note that our constructions use bounded-size weights, and hence, in principle, the aforementioned sample complexity results can be combined with our results to analyze empirical risk minimization for learning transformers. Prior work of Likhosherstov et al. [2021] also shows how sparse attention patterns can be achieved by self-attention units (via random projection arguments); however, when specialized to  $q$ SA, their construction is suboptimal in terms of the sparsity level  $q$ .
Related models. Graph neural networks (GNNs), like transformers, process very large inputs (graphs) using neural networks that act only on small collections of the input parts (vertex neighborhoods). Many classes of GNNs are universal approximators for classes of invariant and equivariant functions [Maron et al., 2019, Keriven and Peyre, 2019]. At the same time, they are restricted by the distinguishing power of certain graph isomorphism tests [Xu et al., 2018, Morris et al., 2019, Chen et al., 2019], and lower bounds have been established on the network size to approximate such tests [Aamand et al., 2022]. Loukas [2019] established a connection between GNNs and the LOCAL [Angluin, 1980] and CONGEST [Peleg, 2000] models for distributed computation, and hence directly translates lower bounds for CONGEST—notably cycle detection problems—into size lower bounds for GNNs. Our lower bounds for cycle detection using transformers also leverage a connection to the CONGEST model. However, transformers do not have the same limitations as GNNs, since the computational substrate of a transformer does not depend on the input graph in the way it is with GNNs. Thus, we cannot directly import lower bounds for CONGEST to obtain lower bounds for transformers.
Transformers are also related to other families of invariant and equivariant networks. Our focus on Match2 and Match3 (and related problems) was inspired by the separation results of Zweig and Bruna [2022] between models for processing sets: Deep Sets [Qi et al., 2017, Zaheer et al., 2017], which are "singleton symmetric", and the more expressive Relational Pooling networks [Santoro et al., 2017], which are only "pairwise symmetric".
1.3 Conclusion and future work
Our primary contributions are to present a multi-faceted story about transformer approximation: firstly, qSA separates transformer models approximation-theoretically from RNNs and MLPs, and moreover the attention embedding dimension both necessary and sufficient for qSA scale directly with  $q$ , meaning qSA also functions to characterize representation power amongst different transformers. Secondly, while single units of self-attention can solve the Match2 task, even wide layers of self-attention with high-dimensional embeddings cannot solve Match3, and we believe that deeper models cannot as well. This question of deeper models is stated as a formal conjecture and addressed heuristically in Appendix C.6, using both information- and communication-theoretic proof techniques, both of which we feel are significant steps towards a complete proof.
While our investigation is purely approximation-theoretic, we also include in Appendix D a preliminary empirical study, showing that attention can learn  $q\mathrm{SA}$  with vastly fewer samples than recurrent
networks and MLPs; we feel this further emphasizes the fundamental value of  $q\mathrm{SA}$ , and constitutes an exciting direction for future work.
Beyond the explicit open question in Informal Conjecture 1, we anticipate that future research could connect the separation results proved in this work to formal linguistic theory and empirical work on attention matrix interpretation. This work examines Match2 and Match3 because we believe that the former could represent a key primitive for language processing tasks such as co-referencing, while the latter represents a natural extension of the former that likely is not necessary for language modeling. Rather, it may be possible that language modeling performs triple-wise modeling for tasks such as the identification of subject, verb, and object components by relying on pairwise matching constructions and "clues" learned within an embedding, such as those encoded in the toy problems Match3Bigram and Match3Local. That is, transformers serve as a useful foundational model for language modeling because of their abilities to integrate contextual clues and pairwise communication, and while they are not extensible to "purely triple-wise problems," most practical sequential problems have some efficient decomposition to pairwise structures that can be easily exploited by these architectures. Future work by linguists, theoretical computer scientists, and empirical NLP practitioners could assess how foundational our primitives are and study whether there are any practical triple-wise problems that transformer models fail to solve.
Let  $\mathbb{B}^d = \{x\in \mathbb{R}^d:\| x\| _2\leq 1\}$  denote the unit ball in  $\mathbb{R}^d$  , and let  $[n] = \{1,2,\dots ,n\}$  denote the first  $n$  positive integers. The expression  $\mathbb{1}\{P\}$  equals 1 if predicate  $P$  is true and 0 otherwise. The row-wise softmax operator applied to matrix  $A\in \mathbb{R}^{N\times M}$  returns
$$
\operatorname {s o f t m a x} (A) _ {i, j} = \frac {\exp (A _ {i , j})}{\sum_ {j ^ {\prime} = 1} ^ {M} \exp (A _ {i , j ^ {\prime}})}.
$$
2.1 Attention units and transformer architectures
We first introduce the concept of self-attention, which is used as the building block of all transformer architectures included in this paper.
Definition 1. For input dimension  $d$ , output dimension  $d'$ , embedding dimension  $m$ , precision  $p$ , and matrices  $Q, K \in \mathbb{R}^{d \times m}$  and  $V \in \mathbb{R}^{d \times d'}$  (encoded using  $p$ -bit fixed-point numbers), a self-attention unit is a function  $f_{Q,K,V}: \mathbb{R}^{N \times d} \to \mathbb{R}^{N \times d'}$  with
$$
f _ {Q, K, V} (X) = \operatorname {s o f t m a x} (X Q K ^ {\top} X ^ {\top}) X V.
$$
Let  $\mathcal{A}_{d,m,d',p} = \{f_{Q,K,V}:Q,K,V\}$  denote all such self-attention units.
Self-attention units can be computed in parallel to create multi-headed attention.
Definition 2. For head-count  $H$  and self-attention units  $f_{1},\ldots ,f_{H}\in \mathcal{A}_{d,m,d^{\prime},p}$ , a multi-headed attention layer is a function  $L_{f_1,\dots,f_H}:\mathbb{R}^{N\times d}\to \mathbb{R}^{N\times m}$  with  $L_{f_1,\dots,f_H}(X) = \sum_{h = 1}^{H}f_h(X)$ . Let  $\mathcal{A}_{d,m,d^{\prime},p}^{H}$  contain all such  $L_{f_1,\dots,f_H}$ .
Transformer models are composed of two components: multi-headed attention layers (as above) and element-wise multi-layer perceptrons. Due to universal approximation results, we model multi-layer perceptrons as arbitrary functions mapping fixed-precision vectors to themselves.
Definition 3. A multi-layer perceptron (MLP) layer is represented by some  $\phi : \mathbb{R}^d \to \mathbb{R}^{d'}$ , whose real-valued inputs and outputs can be represented using  $p$ -bit fixed-precision numbers. We apply  $\phi$  to each element (i.e., row) of an input  $X \in \mathbb{R}^{N \times d}$ , abusing notation to let  $\phi(X) = (\phi(x_1), \ldots, \phi(x_N)) \in \mathbb{R}^{N \times d'}$ . Let  $\Phi_{d,d',p}$  denote all such MLPs.
We concatenate the notation of each class of functions to denote function composition. For example, for output dimension  $d'$ , we use  $\mathcal{A}_{d,m,d',p}' := \mathcal{A}_{m,m,d',p} \Phi_{d,m,p}$  and  $\mathcal{A}_{d,m,d',p}^{H'} := \mathcal{A}_{m,m,d',p}^{H} \Phi_{d,m,p}$  to represent single-headed and multi-headed attention units with an input MLP respectively. (The capabilities and limitations of these models are studied in Section 3.) For depth  $D$ , we let
$$
\mathcal {T} _ {d, m, d ^ {\prime}, p} ^ {D, H} = \Phi_ {m, d ^ {\prime}, p} (\mathcal {A} _ {m, m, m, p} ^ {H \prime}) ^ {D - 1} \mathcal {A} _ {d, m, m, p} ^ {H \prime}
$$
represent a full transformer model comprising  $D$  layers of  $H$ -headed self-attention with interspersed MLPs.
While two key features of transformer architectures—the residual connection and the positional embedding—are conspicuously missing from this formalism, the two can be implemented easily under the framework. We can include a positional embedding by encoding the index as a coordinate of the input, i.e.  $x_{i,1} = i$ . Then, the subsequent MLP transformation  $\phi(X)$  can incorporate  $i$  suitably into the embedding. A residual connection can be included additively as input to a multi-layer perceptron layer (as is standard) by implementing an "approximate identity" attention head  $f$  with  $Q, K$  and  $V = I_m$  set to ensure that  $f(X) \approx X$ .<sup>2</sup>
We periodically consider transformers implemented with real-valued arithmetic with infinite bit complexity; in those cases, we omit the bit complexity  $p$  from the notation.
Finally, we assume for the proof of Theorem 3 that the model is permitted to append a single  $<\mathrm{END}>$  token at the end of a sequence. That is, we say that a model  $f \in \mathcal{T}_{d,m,d',p}^{D,H}$  represents a target  $h: \mathbb{R}^{N \times d} \to \mathbb{R}^{N \times d'}$  if  $f(X')_{1:N} = g(X)$  when  $X' = (x_1, \ldots, x_N, x')$  for constant-valued  $x' \in \mathbb{R}^d$ .
3 Sparse averaging with attention units
We present the sparse averaging task to highlight the ability of transformer architectures to simulate a wide range of meaningful interactions between input elements. This task demonstrates how the embedding dimension of a self-attention unit modulates the expressive capabilities of the architecture, while showcasing the inabilities of fully-connected and recurrent neural networks to capture similar interactions (see Appendix A).
Definition 4. For sparsity  $q$ , problem dimension  $d'$ , and input dimension  $d = d' + q + 1$ , consider an input  $X = (x_{1}, \ldots, x_{N}) \in \mathbb{R}^{N \times d}$  with  $x_{i} = (z_{i}; y_{i}; i)$  for  $z_{i} \in \mathbb{B}^{d'}$  and  $y_{i} \in \binom{[N]}{q}$ . Let the  $q$ -sparse average be
$$
q \operatorname {S A} (X) = \left(\frac {1}{q} \sum_ {j = 1} ^ {q} z _ {y _ {i, j}}\right) _ {i \in [ N ]}.
$$
For accuracy  $\epsilon > 0$ , a function  $f: \mathbb{R}^{N \times d} \to \mathbb{R}^{N \times d'}$ $\epsilon$ -approximates qSA if for all  $X$ ,
$$
\max  _ {i \in [ N ]} \| f (X) _ {i} - q \mathrm {S A} (X) _ {i} \| _ {2} \leq \epsilon .
$$