Question:
You are a highly experienced, conscientious, and fair academic reviewer, please help me review this paper. The review should be organized into nine sections:
1. Summary: A summary of the paper in 100-150 words.
2. Strengths/Weaknesses/Questions: The Strengths/Weaknesses/Questions of paper, which should be listed in bullet points, with each point supported by specific examples from the article where possible.
3. Soundness/Contribution/Presentation: Rate the paper's Soundness/Contribution/Presentation, and match this score to the corresponding description from the list below and provide the result. The possible scores and their descriptions are:
    1 poor
    2 fair
    3 good
    4 excellent
4. Rating: Give this paper an appropriate rating, match this rating to the corresponding description from the list below and provide the result. The possible Ratings and their descriptions are:
    1 strong reject
    2 reject, significant issues present
    3 reject, not good enough
    4 possibly reject, but has redeeming facets
    5 marginally below the acceptance threshold
    6 marginally above the acceptance threshold
    7 accept, but needs minor improvements
    8 accept, good paper
    9 strong accept, excellent work
    10 strong accept, should be highlighted at the conference
5. Paper Decision: It must include the Decision itself(Accept or Reject) and the reasons for this decision, based on the criteria of originality, methodological soundness, significance of results, and clarity and logic of presentation.

Here is the template for a review format, you must follow this format to output your review result:
**Summary:**
Summary content

**Strengths:**
- Strength 1
- Strength 2
- ...

**Weaknesses:**
- Weakness 1
- Weakness 2
- ...

**Questions:**
- Question 1
- Question 2
- ...

**Soundness:**
Soundness result

**Presentation:**
Presentation result

**Contribution:**
Contribution result

**Rating:**
Rating result

**Paper Decision:**
- Decision: Accept/Reject
- Reasons: reasons content

Please ensure your feedback is objective and constructive. The paper is as follows:

# Representational Strengths and Limitations of Transformers

Clayton Sanford, Daniel Hsu

Department of Computer Science

Columbia University

New York, NY 10027

{clayton,djhsu}@cs.columbia.edu

&Matus Telgarsky

Courant Institute

New York University

New York, NY 10012

matus.telgarsky@nyu.edu

###### Abstract

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.

## 1 Introduction

In recent years, transformer networks [24] have been established as a fundamental neural architecture powering state-of-the-art results in many applications, including language modeling [13], computer vision [15], and protein folding [16]. 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 [1, 17] and learn representations that may encode domain-specific information in the self-attention units [18, 19, 20, 16]. 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"1 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{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.

Footnote 1: 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.

### Our contributions

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 and \(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\mathrm{SA}\left((y_{1},z_{1}),\ldots,(y_{N},z_{N})\right)=\left(\frac{1}{q} \sum_{j\in y_{1}}z_{j},\ldots,\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 \(q\mathrm{SA}\) 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_ \(q\mathrm{SA}\) _if and only if_ \(m\gtrsim q\) _(Theorems 2 and 4)._
2. _Any fully-connected neural network whose output approximates_ \(q\mathrm{SA}\) _requires its first hidden layer to have width at least_ \(\Omega(Nd)\) _(Theorem_ 10_)._
3. _Any recurrent neural network whose iterates approximate_ \(q\mathrm{SA}\) _requires a hidden state of at least_ \(\Omega(N)\) _bits (Theorem_ 11_)._

We consider the \(q\mathrm{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 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_ (\(\mathrm{Match2}\)) and _similar triple detection_ (\(\mathrm{Match3}\)) as

\[\mathrm{Match2}(X)_{i\in[N]} =1\left\{\exists j\;\mathrm{s.t.}\;x_{i}+x_{j}=0\left(\mathrm{mod} \;M\right)\right\}, \tag{1}\] \[\mathrm{Match3}(X)_{i\in[N]} =1\left\{\exists j_{1},j_{2}\;\mathrm{s.t.}\;x_{i}+x_{j_{1}}+x_{ j_{2}}=0\left(\mathrm{mod}\;M\right)\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 \(\mathrm{Match2}\) and \(\mathrm{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_ \(\mathrm{Match2}\) _(Theorem_ 6_)._
2. _A single layer of standard multi-headed self-attention cannot compute_ \(\mathrm{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_ \(\mathrm{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_,_ \(\mathrm{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 \(\mathrm{Match3}\) without hints or augmentation.

**Informal Conjecture 1**.: _Every multi-layer transformer that computes \(\mathrm{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.

### Related work

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 (Perez 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, Bhatamishra 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; Bhatamishra et al., 2022) using transformer-based hypothesis classes. Our \(q\mathrm{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 Likhoshenstov 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\mathrm{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 \(\mathrm{Match2}\) and \(\mathrm{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".

### Conclusion and future work

Our primary contributions are to present a multi-faceted story about transformer approximation: firstly, \(q\mathrm{SA}\) separates transformer models approximation-theoretically from RNNs and MLPs, and moreover the attention embedding dimension both necessary and sufficient for \(q\mathrm{SA}\) scale directly with \(q\), meaning \(q\mathrm{SA}\) also functions to characterize representation power amongst different transformers. Secondly, while single units of self-attention can solve the \(\mathrm{Match2}\) task, even wide layers of self-attention with high-dimensional embeddings cannot solve \(\mathrm{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 \(\mathrm{Match2}\) and \(\mathrm{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 \(\mathrm{Match3Bigram}\) and \(\mathrm{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.

## 2 Preliminaries

Let \(\mathbb{B}^{d}=\{x\in\mathbb{R}^{d}:\left\lVert x\right\rVert_{2}\leq 1\}\) denote the unit ball in \(\mathbb{R}^{d}\), and let \([n]=\{1,2,\ldots,n\}\) denote the first \(n\) positive integers. The expression \(\mathbb{1}\left\{P\right\}\) 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

\[\mathrm{softmax}(A)_{i,j}=\frac{\exp(A_{i,j})}{\sum_{j^{\prime}=1}^{M}\exp(A_{ i,j^{\prime}})}.\]

### 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^{\prime}\), embedding dimension \(m\), precision \(p\), and matrices \(Q,K\in\mathbb{R}^{d\times m}\) and \(V\in\mathbb{R}^{d\times d^{\prime}}\) (encoded using \(p\)-bit fixed-point numbers), a _self-attention unit_ is a function \(f_{Q,K,V}:\mathbb{R}^{N\times d}\rightarrow\mathbb{R}^{N\times d}\) with

\[f_{Q,K,V}(X)=\mathrm{softmax}(XQK^{\intercal}X^{\intercal})XV.\]

Let \(\mathcal{A}_{d,m,d^{\prime},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}\), \(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\)\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\(\)\nexists\)\(\)\(\)\(\nexists\)\(\)\(\)\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(\)\(\)\nexists\)\(\nexists\)\(\)\(\nexists\)\(\)\(\nexists\)\(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\).2

Footnote 2: A simple construction involves letting \(XQ=XK\) with \(\operatorname{iid}\) Gaussian columns fixed for every index \(i\). Then, the diagonals of \(XQK^{\top}X^{\top}\) are far larger than all other entries and its softmax is approximately \(I_{N}\).

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 implement a single <END> token at the end of a sequence. That is, we say that a model \(f\in\mathcal{T}^{D,D}_{d,m,d^{\prime},p}\) represents a target \(h:\mathbb{R}^{N\times d}\rightarrow\mathbb{R}^{N\times d^{\prime}}\) if \(f(X^{\prime})_{1:N}=g(X)\) when \(X^{\prime}=(x_{1},\ldots,x_{N},x^{\prime})\) for constant-valued \(x^{\prime}\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^{\prime}\), and input dimension \(d=d^{\prime}+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^{\prime}}\) and \(y_{i}\in\binom{[N]}{q}\).3 Let the \(q\)-_sparse average_ be

Footnote 3: We may encode a \(q\) element subset of \([N]\) as a vector in \([N]^{q}\) constrained to have distinct components.

\[q\mathrm{SA}(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}\rightarrow\mathbb{R}^{N\times d^{\prime}}\)\(\epsilon\)-_approximates_\(q\mathrm{SA}\) if for all \(X\),

\[\max_{i\in[N]}\left\lVert f(X)_{i}-q\mathrm{SA}(X)_{i}\right\rVert_{2}\leq\epsilon.\]

Figure 0(a) visualizes the sparse averaging task as a bipartite graph between subsets \(y_{i}\) and elements \(z_{i}\) with corresponding averages. Theorems 2 and 4 jointly show that the minimum embedding dimension \(m\) of single self-attention units \(\mathcal{A}^{\prime}_{d,m,d^{\prime},p}\) that \(O(\frac{1}{q})\)-approximate \(q\mathrm{SA}\) scales linearly with \(q\). We believe that the sparse averaging problem is thus a canonical problem establishing the representational capabilities and inductive biases of self-attention units.

### Self-attention can approximate \(q\mathrm{SA}\) when \(m\gtrsim q\)

Our principle positive result shows that the sparse averaging task \(q\mathrm{SA}\) can be approximately solved using fixed-precision arithmetic self-attention units with embedding dimension \(m\) growing with \(q\log N\).

**Theorem 2** (Fixed-precision).: _For any \(N\), any \(m\geq\Omega(d^{\prime}+q\log N)\), any \(\epsilon\in(0,1)\), and \(p=\Omega(\log(\frac{q}{\epsilon}\log N))\), there exists some \(f\in\mathcal{A}^{\prime}_{d,m,d^{\prime},p}\) that \(\epsilon\)-approximates \(q\mathrm{SA}\)._

While the full proof appears in Appendix B.1, we briefly sketch the argument here. Because the output of a self-attention unit is a convex combination of rows of the value matrix \(\phi(X)V\in\mathbb{R}^{N\times d^{\prime}}\), a natural way to approximate \(q\mathrm{SA}\) with a unit of self-attention is to let each value be the corresponding vector in the average (i.e. \(V^{\mathrm{ T}}\phi(x_{i})=z_{i}\)) and choose the key and query functions in order to ensure that the attention matrix satisfies

\[\mathrm{softmax}(\phi(X)QK^{\mathrm{ T}}\phi(X)^{\mathrm{ T}})_{i,j}\approx\begin{cases}\frac{1}{q}&\text{if $j\in y_{i}$,}\\ 0&\text{otherwise.}\end{cases}\]

To do so, let each key \(K^{\mathrm{ T}}\phi(x_{i})\) represent a fixed vertex on a convex polytope, which depends only on index \(i\) and is constructed from random binary vectors. We select each query \(Q^{\mathrm{ T}}\phi(x_{i})\) to ensure that \(\phi(x_{i})^{\mathrm{ T}}QK^{\mathrm{ T}}\phi(x_{j})\) is a fixed large value if \(j\in y_{i}\) and a slightly smaller value otherwise. We obtain the precise query, key, and value embeddings by employing tools from dual certificate analysis from the theory of compressed sensing.

We visualize this construction in Figure 0(b) and 0(c) for \(q=3\) and \(d^{\prime}=4\), which presents the associated attention and value matrices necessary for the construction, and plots a polytope of keys (red dots) with each face corresponding to each subset \(y_{i}\) (green dots). The construction is empirically relevant; Figure 2 shows that a unit of self-attention trained on data generated by the \(q\mathrm{SA}\) task recovers a similar attention matrix to the one stipulated in our construction and visualized in Figure 0(b).

The logarithmic dependence of the embedding dimension \(m\) on the sequence length \(N\) can be eliminated by considering self-attention units with real-valued arithmetic with infinite bit complexity.

**Theorem 3** (Infinite-precision).: _For fixed \(N\), \(m\geq\Omega(d^{\prime}+q)\) and \(\epsilon>0\), there exists some \(f\in\mathcal{A}^{\prime}_{d,m,d^{\prime}}\) that \(\epsilon\)-approximates \(q\mathrm{SA}\)._

Figure 1: A visualization of the \(q\mathrm{SA}\) function outputs given a sequence of inputs \((z_{i};y_{i};i)_{i\in[N]}\) as a bipartite graph between subsets \(y_{i}\) and vectors \(z_{i}\) (a), and of the attention matrix (b) and underlying embeddings (c) that produce the self-attention construction in Theorem 2.

Figure 2: Attention matrix \(\mathrm{softmax}(\phi(X)QK^{\mathrm{ T}}\phi(X)^{\mathrm{ T}})\in\mathbb{R}^{20\times 20}\) for a fixed example after \(T\) epochs of training a self-attention unit to solve \(q\mathrm{SA}\) for \(q=3\). Each row \(i\) corresponds to subset \(y_{i}\), and each cell \(j\in y_{i}\) is outlined in red. See Appendix D for experimental details.

The proof of Theorem 3 employs a similar polytope-based construction in Appendix B.2, relying on a cyclic polytope rather than one drawn from discrete boolean vectors. Theorem 16 proves the near-optimality of _that_ bound by employing a geometric argument to show that a variant of \(q\mathrm{SA}\) can only be approximated by a restricted family of self-attention units with a sufficiently high-dimensional embedding.

### Self-attention cannot approximate \(q\mathrm{SA}\) when \(m\lesssim q\)

We show that the construction used to prove Theorem 2 is nearly optimal.

**Theorem 4**.: _For any sufficiently large \(q\), any \(N\geq 2q+1\), and any \(d^{\prime}\geq 1\), there exists a universal constant \(c\) such that if \(mp\leq cq\), then no \(f\in\mathcal{T}^{1,1}_{d,m,d^{\prime},p}\) exists that \(\frac{1}{2q}\)-approximates \(q\mathrm{SA}\)._

(By choosing \(p=O(\log(q\log N))\), Theorem 2 is shown to be optimal up to logarithmic factors of \(q\) and doubly-logarithmic factors of \(N\).)

The proof of Theorem 4 employs a standard communication complexity argument based on a reduction from the following _set disjointness_ problem in the two-party communication model, in which each party possesses a subset of an \(n\) element domain (encoded as \(n\)-bit strings), and they wish to jointly determine whether their subsets are disjoint. We note that communication complexity is commonly-used technique for proving lower bounds on the representational power of circuits and feedforward neural networks (see, e.g., Karchmer and Wigderson, 1988; Ben-David et al., 2002; Martens et al., 2013; Vardi et al., 2021).

**Fact 5** (Set disjointness communication lower bound (Yao, 1979)).: _Suppose Alice and Bob are given inputs \(a,b\in\{0,1\}^{n}\), respectively, with the goal of jointly computing \(\mathrm{DISJ}(a,b)=\max_{i}a_{i}b_{i}\) by alternately sending a single bit message to the other party over a sequence of communication rounds. Any deterministic protocol for computing \(\mathrm{DISJ}(a,b)\) requires at least \(n\) rounds of communication._

Our proof designs a communication protocol that Alice and Bob use to jointly compute \(\mathrm{DISJ}(a,b)\) when \(n=q\) in \(O(mp)\) rounds of communication, under the assumption that such an \(f\) exists that closely approximates \(q\mathrm{SA}\).

* Alice encodes her input \(a\) in a single subset by letting \(y_{2q+1}=\{2i+a_{i}-1:i\in[q]\}\).
* Bob uses his input \(b\) to assign \(z_{2i-1}\) to \(2b_{i}-1\) and \(z_{2i}=-1\) for all \(i\in[q]\).
* All other input components are set to constant values known by both parties.

Alice sends her \(mp\)-bit query embedding \(Q^{\mathrm{r}}\phi(x_{2q+1})\) bit-by-bit to Bob, who approximately computes \(q\mathrm{SA}\) by determining the outcome of \(f\). The crux of the reduction shows that \(q\mathrm{SA}(X)_{2q+1}=-1\) if and only if \(a_{i}b_{i}=0\) for all \(i\in[q]\), which allows Bob to determine \(\mathrm{DISJ}(a,b)\).

We visualize the protocol in Figure 3 and give the proof in Appendix B.3. The proofs of Theorems 7, 11, 21, and 23 employ similar communication complexity reductions to \(\mathrm{DISJ}\).

Figure 3: The \(mp\)-bit communication protocol used to reduce the hardness of computing \(q\mathrm{SA}\) with a single unit of self-attention to the hardness of solving the \(\mathrm{DISJ}\) communication problem for the proof of Theorem 4 for \(q=4\).

Standard transformer models can only efficiently represent intrinsically pairwise functions

In this section, we argue that the standard transformer architecture is unable to efficiently represent functions that do not decompose into a small number of pairwise-symmetric functions. We do this by contrasting the (in)approximability of intrinsically pairwise and triple-wise functions, respectively \(\mathrm{Match2}\) and \(\mathrm{Match3}\) (defined in (1) and (2)), and their variants.

### Efficient computation of \(\mathrm{Match2}\) with standard self-attention

We first show that \(\mathrm{Match2}\) can be efficiently approximated by a single standard (pairwise) self-attention unit.

**Theorem 6**.: _For any input size \(N\), input range \(M=N^{O(1)}\), and fixed-precision bit complexity \(p=O(\log M)\), there exists a transformer architecture \(f\in\mathcal{T}^{1,1}_{1,m,1,p}\) with a single self-attention unit with embedding dimension \(m=3\) such that for all \(X\in[M]^{N}\), \(f(X)=\mathrm{Match2}(X)\)._

The proof, given in Appendix C.1 uses both a "blank token" and a trigonometric positional embedding, which ensures that

\[\phi(x_{i})^{\intercal}QK^{\intercal}\phi(x_{j})=c\sum_{k=1}^{d}\cos\left( \frac{2\pi(x_{i,k}+x_{j,k})}{M}\right)\]

for some sufficiently large constant \(c\). This embedding ensures that a cell of the attention matrix \(\mathrm{softmax}(\phi(X)QK^{\intercal}\phi(X)^{\intercal})_{i,j}\) is extremely close to zero, unless \(x_{i}=-x_{j}\pmod{M}\).

### Hardness of computing \(\mathrm{Match3}\) with a multi-headed self-attention layer

Although \(\mathrm{Match2}\) can be efficiently represented using a single unit of standard self-attention, representing \(\mathrm{Match3}\) using an entire layer of multi-headed attention units is impossible unless either the number of heads \(H\), the embedding dimension \(m\), or the precision \(p\) grows as \(N^{\Omega(1)}\).

**Theorem 7**.: _There is universal constant \(c>0\) such that for sufficiently large \(N\), and any \(M\geq N+1\), if \(mpH\leq cN/\log\log N\), then there is no \(f\in\mathcal{T}^{1,H}_{1,m,1,p}\) satisfying \(f(X)=\mathrm{Match3}(X)\) for all \(X\in[M]^{N}\)._

We give the proof in Appendix C.2. Like that of Theorem 4, the proof relies on a reduction from set disjointness in two-party communication. The proof of the lower bound applies a domain-restricted variant of \(\mathrm{Match3}\), which actually makes the problem substantially simpler to solve. In Remark 1, we show how this variant of \(\mathrm{Match3}\) introduces a _depth separation_ between the representational powers of single-layer and two-layer transformer models.

As mentioned in the introduction, we also conjecture that multiple layers of multi-headed attention are subject to the same impossibility (Conjecture 19). The impossibility is specific to standard (pairwise) attention; in Appendix C.4, we show that \(\mathrm{Match3}\)_can_ be efficiently computed with a single unit of _third-order_ self-attention.

### More efficient constructions for simplified \(\mathrm{Match3}\) computations

While the previous sections suggests that no efficient construction exists to compute \(\mathrm{Match3}\) with standard transformer models, practical examples of triple detection bound. For example, a transformer-based language model will likely succeed in linking a subject/verb/object triple because all three tokens likely inhabit the same local region and because the model could agglomerate the triple by first identifying a pair and then adding the third. Here, we introduce two variants on the \(\mathrm{Match3}\) problem that have additional structure to serve as hints. The first variant specifies triple sums comprising the input element and a neighboring pair elsewhere in the sequence: for each \(i\in[N]\),

\[\mathrm{Match3Bigram}(X)_{i}=1\left\{\exists j\;\mathrm{s.t.}\;x_{i}+x_{j}+x_{ j+1}=0\left(\mathrm{mod}\;M\right)\right\}.\]

The second focuses on localized sums, where are all components of a triple must be within a fixed range of constant width \(K\ll N\): for each \(i\in[N]\),

\[\mathrm{Match3Local}(X)_{i}=1\left\{\exists j_{1},j_{2}\;\mathrm{s.t.}\;x_{i }+x_{j_{1}}+x_{j_{2}}=0\left(\mathrm{mod}\;M\right),\left|i-j_{1}\right|, \left|i-j_{2}\right|\leq K\right\}.\]

We show that the two can be efficiently represented using compact standard transformer models.

**Theorem 8**.: _For any \(N\), \(M=N^{O(1)}\), and \(p=O(\log M)\), there exists a transformer architecture \(f\in\mathcal{T}^{1,1}_{1,m,1,p}\) with embedding dimension \(m=3\) and depth \(D=2\) such that for all \(X\in[M]^{N\times d}\), \(f(X)=\operatorname{Match3Big}(X)\)._

Informally, the first layer of the construction uses a sinusoidal positional encoding to compute each bigram sum \(x_{j}+x_{j+1}\) in the \(j\)th element of the sequence. The second layer applies the \(\operatorname{Match2}\) construction provided by Theorem 6 to determine whether there exists a \(j\) for each \(i\) such that \(x_{i}+x_{j}+x_{j+1}=0\pmod{M}\).

**Theorem 9**.: _For any \(d\), \(N\), \(M=N^{O(1)}\), \(p=O(\log M)\), and \(K\leq N\), there exists a transformer architecture \(f\in\mathcal{T}^{1,1}_{1,m,1,p}\) with embedding dimension \(m=O(K\log N)\) and bit-complexity \(p=O(\log(K\log N))\) such that for all \(X\in[M]^{N\times d}\), \(f(X)=\operatorname{Match3Local}(X)\)._

Proof.: We implement the localized construction by using Theorem 2 to construct a specific sparse simultaneous average of the inputs with \(q:=2K+1\) and \(d^{\prime}:=2K+1\). To do so, we use the input MLP to convert \(x_{i}\) to the embedding \((z_{i};y_{i};i)\), for zero-padded input

\[z_{i}=x_{i}e_{\bar{i}}\in\mathbb{R}^{2K+1}\]

for \(\bar{i}=i\pmod{2K+1}\) and subset

\[y_{i}=\{i-K,i-K+1,\ldots,i+K\}\in\binom{[N]}{2K+1}.\]

This construction ensures that the \(i\)th element of self-attention output computes (a rotation of) \((x_{i-K},x_{i-K+1},\ldots,x_{i+K})\). An output MLP can then verify whether any matching triples involving \(x_{i}\) exist among those vectors.

## Acknowledgments and Disclosure of Funding

We are grateful for many discussions with and feedback from Navid Ardeshir, Peter Bartlett, Alberto Bietti, Yuval Efron, Christos Papadimitriou, Shivam Nadimpalli, Rocco Servedio, Yusu Wang, and Cyril Zhang. This work was supported in part by NSF grants CCF-1740833 and IIS-1563785, a JP Morgan Faculty Award, and an NSF Graduate Research Fellowship.

## References

* Aamand et al. [2022] Anders Aamand, Justin Chen, Piotr Indyk, Shyam Narayanan, Ronitt Rubinfeld, Nicholas Schiefer, Sandeep Silwal, and Tal Wagner. Exponentially improving the complexity of simulating the Weisfeiler-Lehman test with graph neural networks. In _Advances in Neural Information Processing Systems 35_, 2022.
* Angluin [1980] Dana Angluin. Local and global properties in networks of processors. In _Proceedings of the Twelfth Annual ACM Symposium on Theory of Computing_, 1980.
* Ben-David et al. [2002] Shai Ben-David, Nadav Eiron, and Hans Ulrich Simon. Limitations of learning via embeddings in euclidean half spaces. _Journal of Machine Learning Research_, 3(Nov):441-461, 2002.
* Bhattacharya et al. [2020] Satwik Bhattacharya, Kabir Ahuja, and Navin 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_, 2020.
* Bhattacharya et al. [2022] Satwik Bhattacharya, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions. _arXiv preprint arXiv:2211.12316_, 2022.
* Brown et al. [2020] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. _arXiv preprint arXiv:2005.14165_, 2020.
* Chen et al. [2020]Emmanuel J Candes and Terence Tao. Decoding by linear programming. _IEEE transactions on information theory_, 51(12):4203-4215, 2005.
* Chen et al. (2022) Nuo Chen, Qiushi Sun, Renyu Zhu, Xiang Li, Xuesong Lu, and Ming Gao. Cat-probing: A metric-based approach to interpret how pre-trained models for programming language attend code structure. _arXiv preprint arXiv:2210.04633_, 2022.
* Chen et al. (2019) Zhengdao Chen, Soledad Villar, Lei Chen, and Joan Bruna. On the equivalence between graph isomorphism testing and function approximation with GNNs. In _Advances in Neural Information Processing Systems 32_, 2019.
* Clark et al. (2019) Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. What does bert look at? an analysis of bert's attention. _arXiv preprint arXiv:1906.04341_, 2019.
* Daniely (2017) Amit Daniely. Depth separation for neural networks. In Satyen Kale and Ohad Shamir, editors, _Proceedings of the 2017 Conference on Learning Theory_, volume 65 of _Proceedings of Machine Learning Research_, pages 690-696. PMLR, 07-10 Jul 2017. URL [https://proceedings.mlr.press/v65/daniely17a.html](https://proceedings.mlr.press/v65/daniely17a.html).
* Dosovitskiy et al. (2021) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. _arXiv preprint arXiv:2010.11929_, 2021.
* Edelman et al. (2022) Benjamin L. Edelman, Surbhi Goel, Sham M. Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In _International Conference on Machine Learning_, 2022.
* Eldan and Shamir (2016) Ronen Eldan and Ohad Shamir. The power of depth for feedforward neural networks. In Vitaly Feldman, Alexander Rakhlin, and Ohad Shamir, editors, _29th Annual Conference on Learning Theory_, volume 49 of _Proceedings of Machine Learning Research_, pages 907-940, Columbia University, New York, New York, USA, 23-26 Jun 2016. PMLR. URL [https://proceedings.mlr.press/v49/eldan16.html](https://proceedings.mlr.press/v49/eldan16.html).
* Furst et al. (1984) Merrick Furst, James B Saxe, and Michael Sipser. Parity, circuits, and the polynomial-time hierarchy. _Mathematical systems theory_, 17(1):13-27, 1984.
* Gale (1963) David Gale. Neighborly and cyclic polytopes. In _Proc. Sympos. Pure Math_, volume 7, pages 225-232, 1963.
* Hahn (2020) Michael Hahn. Theoretical limitations of self-attention in neural sequence models. _Trans. Assoc. Comput. Linguistics_, 8:156-171, 2020. doi: 10.1162/tacl_{a}{_0}{0}{3}}{0}6. URL [https://doi.org/10.1162/tacl_a_00306](https://doi.org/10.1162/tacl_a_00306).
* Hao et al. (2022) Yiding Hao, Dana Angluin, and Robert Frank. Formal language recognition by hard attention transformers: Perspectives from circuit complexity. _Trans. Assoc. Comput. Linguistics_, 10:800-810, 2022. URL [https://transacl.org/ojs/index.php/tacl/article/view/3765](https://transacl.org/ojs/index.php/tacl/article/view/3765).
* Hewitt and Manning (2019) John Hewitt and Christopher D Manning. A structural probe for finding syntax in word representations. In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_, 2019.
* Hornik et al. (1989) Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. _Neural Netw._, 2(5):359-366, July 1989.
* Jumper et al. (2021) John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. _Nature_, 596(7873):583-589, 2021.
* Karchmer and Wigderson (1988) Mauricio Karchmer and Avi Wigderson. Monotone circuits for connectivity require super-logarithmic depth. In _Proceedings of the Twentieth Annual ACM Symposium on Theory of Computing_, 1988.
* Keriven and Peyre (2019) Nicolas Keriven and Gabriel Peyre. Universal invariant and equivariant graph neural networks. In _Advances in Neural Information Processing Systems 32_, 2019.
* Karchmer et al. (2019)Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. _arXiv preprint arXiv:2106.03764_, 2021.
* Liu et al. (2022) Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. _CoRR_, abs/2210.10749, 2022. doi: 10.48550/arXiv.2210.10749.
* Loukas (2019) Andreas Loukas. What graph neural networks cannot learn: depth vs width. _arXiv preprint arXiv:1907.03199_, 2019.
* Maron et al. (2019) Haggai Maron, Ethan Fetaya, Nimrod Segol, and Yaron Lipman. On the universality of invariant networks. In _International Conference on Machine Learning_, 2019.
* Martens et al. (2013) James Martens, Arkadev Chattopadhya, Toni Pitassi, and Richard Zemel. On the representational efficiency of restricted boltzmann machines. In _Advances in Neural Information Processing Systems 26_, 2013.
* Mendelson et al. (2007) Shahar Mendelson, Alain Pajor, and Nicole Tomczak-Jaegermann. Reconstruction and subgaussian operators in asymptotic geometric analysis. _Geometric and Functional Analysis_, 17(4):1248-1282, 2007.
* Morris et al. (2019) Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In _AAAI Conference on Artificial Intelligence_, 2019.
* OpenAI (2023) OpenAI. Gpt-4 technical report, 2023.
* Peleg (2000) David Peleg. _Distributed computing: a locality-sensitive approach_. SIAM, 2000.
* Perez et al. (2019) Jorge Perez, Javier Marinkovic, and Pablo Barcelo. On the turing completeness of modern neural network architectures. _arXiv preprint arXiv:1901.03429_, 2019.
* Qi et al. (2017) Charles R Qi, Hao Su, Kaichun Mo, and Leonidas J Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_, 2017.
* Rogers et al. (2020) Anna Rogers, Olga Kovaleva, and Anna Rumshisky. A primer in bertology: What we know about how bert works. _Transactions of the Association for Computational Linguistics_, 8:842-866, Dec 2020. ISSN 2307-387X. doi: 10.1162/tacl_a_00349. URL [http://dx.doi.org/10.1162/tacl_a_00349](http://dx.doi.org/10.1162/tacl_a_00349).
* Santoro et al. (2017) Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. A simple neural network module for relational reasoning. In _Advances in Neural Information Processing Systems 30_, 2017.
* Sauer (1972) Norbert Sauer. On the density of families of sets. _Journal of Combinatorial Theory, Series A_, 13(1):145-147, 1972.
* Shelah (1972) Saharon Shelah. A combinatorial problem; stability and order for models and theories in infinitary languages. _Pacific Journal of Mathematics_, 41(1):247-261, 1972.
* Telgarsky (2016) Matus Telgarsky. Benefits of depth in neural networks. In Vitaly Feldman, Alexander Rakhlin, and Ohad Shamir, editors, _29th Annual Conference on Learning Theory_, volume 49 of _Proceedings of Machine Learning Research_, pages 1517-1539, Columbia University, New York, New York, USA, 23-26 Jun 2016. PMLR. URL [https://proceedings.mlr.press/v49/telgarsky16.html](https://proceedings.mlr.press/v49/telgarsky16.html).
* Vapnik and Chervonenkis (1968) Vladimir Naumovich Vapnik and Aleksei Yakovlevich Chervonenkis. The uniform convergence of frequencies of the appearance of events to their probabilities. _Doklady Akademii Nauk_, 181(4):781-783, 1968.
* Vardi et al. (2021) Gal Vardi, Daniel Reichman, Toniann Pitassi, and Ohad Shamir. Size and depth separation in approximating benign functions with neural networks. In _Conference on Learning Theory_, 2021.
* Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In _Advances in Neural Information Processing Systems 30_, 2017.
* Wang et al. (2018)Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. _Advances in Neural Information Processing Systems_, 35:12071-12083, 2022.
* Xu et al. (2018) Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? _arXiv preprint arXiv:1810.00826_, 2018.
* Yao (1979) Andrew Chi-Chih Yao. Some complexity questions related to distributive computing (preliminary report). In _Proceedings of the Eleventh Annual ACM Symposium on Theory of Computing_, 1979.
* Yao et al. (2021) Shunyu Yao, Binghui Peng, Christos H. Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing_, 2021.
* Yun et al. (2020) Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? In _International Conference on Learning Representations_, 2020.
* Zaheer et al. (2017) Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. In _Advances in Neural Information Processing Systems 30_, 2017.
* Ziegler (2006) Gunter M Ziegler. Lectures on polytopes. _Graduate Texts in Mathematics_, 152, 2006.
* Zweig and Bruna (2022) Aaron Zweig and Joan Bruna. Exponential separations in symmetric neural networks. _CoRR_, abs/2206.01266, 2022. doi: 10.48550/arXiv.2206.01266.

Fully-connected neural networks and recurrent neural networks cannot efficiently approximate \(q\mathrm{SA}\)

### Only wide fully-connected neural networks can approximate \(q\mathrm{SA}\)

In this section, we show that any fully-connected neural network that approximates \(q\mathrm{SA}:\mathbb{R}^{Nd}\rightarrow\mathbb{R}^{Nd^{\prime}}\) must have width \(m=\Omega(N)\).4 We consider networks of the form \(f(x)=g(Wx)\) for some weight matrix \(W\in\mathbb{R}^{m\times Nd}\) (the first layer weights) and arbitrary function \(g:\mathbb{R}^{m}\rightarrow\mathbb{R}^{Nd^{\prime}}\) (computed by subsequent layers of a neural network).

Footnote 4: We regard inputs as \(Nd\)-dimensional vectors rather than \(N\times d\) matrices.

**Theorem 10**.: _Suppose \(q\leq\frac{N}{2}\). Any fully-connected neural network \(f\) defined as above that \(\frac{1}{2q}\) approximates \(q\mathrm{SA}\) satisfies \(m\geq\mathrm{rank}(W)\geq\frac{Nd^{\prime}}{2}\)._

Proof.: For simplicity, we arrange the input as

\[x=(1;\ldots;N;y_{1};\ldots;y_{N};z_{1};\ldots;z_{N})\]

and \(W=[\tilde{W};V_{1};\ldots;V_{N}]\) with \(z_{1},\ldots,z_{N}\in\mathbb{B}^{d^{\prime}}\), \(\tilde{W}\in\mathbb{R}^{m\times N(d-d^{\prime})}\), and \(V_{1},\ldots,V_{N}\in\mathbb{R}^{m\times d^{\prime}}\). If \(\mathrm{rank}(W)\leq\frac{Nd^{\prime}}{2}-1\), then so too is \(\mathrm{rank}([V_{q};\ldots;V_{N}])\leq\frac{Nd^{\prime}}{2}-1\), and \([V_{q};\ldots;V_{N}]\) has a nontrivial null space containing a nonzero vector \(u=(u_{q};\ldots;u_{N})\in\mathbb{R}^{(N-q)d^{\prime}}\). Let

\[\xi=\frac{1}{\max_{j\in\{q,\ldots,N\}}\left\|u_{j}\right\|_{2}}(u_{q};\ldots;u_ {N}),\]

\(z=(\vec{0};\ldots;\vec{0};\xi_{q};\ldots;\xi_{N})\), and \(z^{\prime}=(\vec{0};\ldots;\vec{0};-\xi_{q};\ldots;-\xi_{N})\). Then,

1. \(z_{j},z^{\prime}_{j}\in\mathbb{B}^{d^{\prime}}\) for all \(j\in[N]\);
2. \(V_{j}z_{j}=V_{j}z^{\prime}_{j}=0\) for all \(j\in[N]\); and
3. \(\|z_{j^{*}}-z^{\prime}_{j^{*}}\|_{2}=2\) for some \(j^{*}\in\{q,\ldots,N\}\).

Therefore, for any \(y_{1},\ldots,y_{N}\in\binom{[N]}{q}\), respective \(x=(1;\ldots;N;y_{1};\ldots;y_{N};z_{1};\ldots;z_{N})\) and \(x^{\prime}=(1;\ldots;N;y_{1};\ldots;y_{N};z^{\prime}_{1};\ldots;z^{\prime}_{N})\) satisfy \(f(x)=f(x^{\prime})\). Consider \(y\) with \(y_{j}=(1,\ldots,q-1,j)\) for each \(j\in\{q,\ldots,N\}\). Then,

\[q\mathrm{SA}(x)_{j}=\frac{1}{q}\xi_{j}\text{ and }q\mathrm{SA}(x^{\prime})_{j}= -\frac{1}{q}\xi_{j}.\]

Hence, \(\left\|q\mathrm{SA}(x)_{j^{*}}-q\mathrm{SA}(x^{\prime})_{j^{*}}\right\|_{2} \geq\frac{2}{q}\). Because \(f(x)=f(x^{\prime})\),

\[\max\left(\left\|f(x)-q\mathrm{SA}(x)_{j^{*}}\right\|_{2},\left\|f(x^{\prime}) -q\mathrm{SA}(x^{\prime})_{j^{*}}\right\|_{2}\right)\geq\frac{1}{q},\]

so \(f\) can approximate \(q\mathrm{SA}\) to accuracy no better than \(\frac{1}{q}\).

### Only high-memory recurrent neural networks can approximate \(q\mathrm{SA}\)

In this section, we show that any memory-bounded algorithm that approximates \(q\mathrm{SA}:\mathbb{R}^{N\times d}\rightarrow\mathbb{R}^{N\times d^{\prime}}\) must use a large "hidden state" (memory) as it processes the input elements. This lower bound applies to various recurrent neural network (RNN) architectures.

A memory-bounded algorithm with an \(m\)-bit memory processes input \(X\in\mathbb{R}^{N\times d}\) sequentially as follows. There is an initial memory state \(h_{0}\in\{0,1\}^{m}\). For \(i=1,2,\ldots,N\), the algorithm computes the \(i\)-th output \(f(X)_{i}\in\mathbb{R}^{d^{\prime}}\) and the updated memory state \(h_{i}\) as a function of the input \(x_{i}\in\mathbb{R}^{d}\) and previous memory state \(h_{i-1}\):

\[(f(X)_{i},h_{i})=g_{i}(x_{i},h_{i-1}),\]where \(g_{i}\colon\mathbb{R}^{d}\times\{0,1\}^{m}\to\mathbb{R}^{d^{\prime}}\times\{0,1 \}^{m}\) is permitted to be an arbitrary function, and \(f\colon\mathbb{R}^{N\times d}\to\mathbb{R}^{N\times d^{\prime}}\) is the function computed by the algorithm.

Our lower bound applies to algorithms that only need to solve the subclass of "causal" instances of \(q\mathrm{SA}\) in which the input \(X=((z_{i},y_{i},i))_{i\in[N]}\in\mathbb{R}^{N\times d}\) is promised to satisfy \(y_{i}=\emptyset\) for all \(i\leq N/2+1\), and \(y_{i}\subseteq\{1,\ldots,N/2+1\}\) for all \(i>N/2+1\).

**Theorem 11**.: _For any \(\varepsilon\in(0,1)\), any memory-bounded algorithm that \(\varepsilon\)-approximates \(q\mathrm{SA}\) (for \(q=1\) and \(d^{\prime}=1\)) on the subclass of "causal" instances must have memory \(m\geq(N-1)/2\)._

Proof.: Consider an \(m\)-bit memory-bounded algorithm computing a function \(f\colon\mathbb{R}^{N\times d}\to\mathbb{R}^{N}\) that \(\varepsilon\)-approximates \(q\mathrm{SA}\) (for \(q=1\) and \(d^{\prime}=1\)). We construct, from this algorithm, a communication protocol for DISJ (with \(N=2n+1\)) that uses \(m\) bits of communication.

Let \(a,b\in\{0,1\}^{n}\) be the input for DISJ provided to Alice and Bob, respectively. The protocol is as follows.

1. Alice constructs inputs \(x_{i}=(z_{i},\emptyset,i)\) for \(i=1,\ldots,n+1\), where for each \(i=1,\ldots,n\), \[z_{i}=\begin{cases}+1&\text{if }a_{i}=0,\\ -1&\text{if }a_{i}=1,\end{cases}\] and \[z_{n+1}=+1.\] Bob constructs inputs \(x_{n+1+i}=(0,y_{n+1+i},n+1+i)\) for \(i=1,\ldots,n\), where \[y_{n+1+i}=\begin{cases}\{n+1\}&\text{if }b_{i}=0,\\ \{i\}&\text{if }b_{i}=1.\end{cases}\] Observe that, for this input \(X=(x_{1},\ldots,x_{2n+1})\), we have \[q\mathrm{SA}(X)_{n+1+i}=\begin{cases}+1&\text{if }a_{i}b_{i}=0,\\ -1&\text{if }a_{i}b_{i}=1.\end{cases}\]
2. Alice simulates the memory-bounded algorithm on the first \(n+1\) inputs \(x_{1},\ldots,x_{n+1}\), and sends Bob the \(m\)-bit memory state \(h_{n+1}\). This requires \(m\) bits of communication.
3. Starting with \(h_{n+1}\), Bob continues the simulation of the memory-bounded algorithm on these \(n\) additional inputs \(x_{n+2},\ldots,x_{2n+1}\).
4. If any output \(f(X)_{n+1+i}\) for \(i=1,\ldots,n\) satisfies \[f(X)_{n+1+i}<0,\] then Bob outputs \(1\) (not disjoint); otherwise Bob outputs \(0\) (disjoint).

The approximation guarantee of \(f\) implies that \(\mathrm{sign}(f(X)_{n+1+i})=q\mathrm{SA}(X)_{n+1+i}\) for all \(i=1,\ldots,n\), so Bob outputs \(1\) if and only if \(a\) and \(b\) are not disjoint. Because this protocol for DISJ uses \(m\) bits of communication, by Fact 5, it must be that \(m\geq n=(N-1)/2\).

We note that the proof of Theorem 11 can be simplified by reducing from the INDEX problem, which has a 1-way communication lower bound of \(n\) bits. This suffices for "single pass" algorithms, such as standard RNNs. However, the advantage of the above argument (and reducing from DISJ) is that it easily extends to algorithms that make multiple passes over the input. Such algorithms are able to capture bidirectional recurrent neural net and related models. A straightforward modification of the protocol in the proof of Theorem 11 shows that \(\Omega(N)\) memory is required for any algorithm that makes \(O(1)\) passes over the input (and computes the outputs in a final pass).

Supplementary results for Section 3

### Proof of Theorem 2

**Theorem 2** (Fixed-precision).: _For any \(N\), any \(m\geq\Omega(d^{\prime}+q\log N)\), any \(\epsilon\in(0,1)\), and \(p=\Omega(\log(\frac{q}{\epsilon}\log N))\), there exists some \(f\in\mathcal{A}^{\prime}_{d,m,d^{\prime},p}\) that \(\epsilon\)-approximates \(q\)SA._

Proof.: Before explaining how they are produced by the input MLP, we introduce the corresponding key, value, and query inputs. The values will simply be \(\phi(X)V=(z_{1},\ldots,z_{N})\). For some \(m^{\prime}=\frac{m-d}{2}\), let \(\phi(X)K=(u_{1},\ldots,u_{N})\in\mathbb{R}^{N\times m^{\prime}}\) be embedded key vectors, where \(u_{1},\ldots,u_{N}\in\{\pm 1/\sqrt{m^{\prime}}\}^{m^{\prime}}\) are the columns of a \(m^{\prime}\times N\) matrix satisfying the \((q,1/4)\)-restricted isometry and orthogonality property (Definition 5), as guaranteed to exist by Lemma 12 and the assumption on \(m^{\prime}\). Let \(\alpha:=\lceil 2\log(4N/\epsilon)\rceil\). By Lemma 13, for each \(y\in\binom{[N]}{q}\), there exists \(w_{y}\in\mathbb{R}^{m^{\prime}}\) with \(\|w_{y}\|_{2}\leq 2\sqrt{q}\) satisfying

\[\langle u_{i^{\prime}},w_{y}\rangle =1\quad\text{ for all }i^{\prime}\in y,\] \[|\langle u_{i^{\prime}},w_{y}\rangle| \leq\frac{1}{2}\quad\text{ for all }i^{\prime}\notin y.\]

Given the bounded precision of the model, we are not free to represent the vectors \(w_{y}\) exactly. Under \(p\)-bit precision for \(p\) sufficiently large, we there exists a vector of \(p\)-bit floating point numbers \(\widetilde{w_{y}}\in\mathbb{R}^{m^{\prime}}\) for every \(w_{y}\) with \(\left\|w_{y}\right\|_{2}\leq 2\sqrt{q}\) satisfying \(\left\|\widetilde{w_{y}}-w_{y}\right\|_{2}\leq\frac{\epsilon}{4\alpha}\). As an immediate consequence, \(|\langle u_{i^{\prime}},\widetilde{w_{y}}\rangle-\langle u_{i^{\prime}},w_{y} \rangle|\leq\frac{\epsilon}{4\alpha}\) for all \(i^{\prime}\) and \(y\) (by Cauchy-Schwarz). The remainder of the proof demonstrates that the necessary properties of the argument hold even with this approximation.

We now describe how to structure the neural network. We define an MLP \(\phi:\mathbb{R}^{d}\rightarrow\mathbb{R}^{m}\) as \(\phi(x_{i})=\phi(z_{i};y_{i};i)=(z_{i};\alpha\widetilde{w_{y_{i}}};u_{i})\), which works simply by using a look-up table on the values of \(u_{i}\) and \(\widetilde{w_{y_{i}}}\) from keys \(i\) and \(y_{i}\) respectively. Then, we define \(Q,K,V\) as sparse boolean-valued matrices that simply copy their respective elements from \(\phi(X)\).

We analyze the output of the softmax. If \(i^{\prime}\in y_{i}\), then

\[\operatorname{softmax}(\phi(X)QK^{\mathsf{T}}\phi(X)^{\mathsf{T} })_{i,i^{\prime}} =\frac{\exp(\alpha\left\langle u_{i},\widetilde{w_{i^{\prime}}} \right\rangle)}{\sum_{i^{\prime\prime}\in y_{i}}\exp(\alpha\left\langle u_{i},\widetilde{w_{i^{\prime\prime}}}\right\rangle)+\sum_{i^{\prime\prime}\not\in y _{i}}\exp(\alpha\left\langle u_{i},\widetilde{w_{i^{\prime\prime}}}\right\rangle)}\] \[\geq\frac{\exp(\alpha-\frac{\epsilon}{4})}{q\exp(\alpha+\frac{ \epsilon}{4})+N\exp(\frac{\alpha}{2}+\frac{\epsilon}{4})}=\frac{e^{\alpha}}{ qe^{\alpha}+Ne^{\alpha/2}}\cdot\exp\left(-\frac{\epsilon}{2}\right)\] \[\geq\left(\frac{1}{q}-\frac{Ne^{\alpha/2}}{qe^{\alpha}}\right) \left(1-\frac{\epsilon}{2}\right)\geq\frac{\left(1-\frac{\epsilon}{4}\right) \left(1-\frac{\epsilon}{4}\right)}{q}\geq\frac{1}{q}\left(1-\frac{\epsilon}{2 }\right).\]

An analogous argument shows that

\[\operatorname{softmax}(\phi(X)QK^{\mathsf{T}}\phi(X)^{\mathsf{T} })_{i,i^{\prime}}\leq\frac{1}{q}\left(1+\frac{\epsilon}{2}\right).\]

Likewise, if \(i^{\prime}\not\in y_{i}\), then

\[\operatorname{softmax}(\phi(X)QK^{\mathsf{T}}\phi(X)^{\mathsf{T} })_{i,i^{\prime}}\leq\frac{\exp(\frac{\alpha}{2}+\frac{\epsilon}{4})}{q\exp( \alpha-\frac{\epsilon}{4})}\leq\exp\left(-\frac{\alpha}{2}+\frac{\epsilon}{2} \right)\leq\frac{\epsilon}{2N}.\]

We thus conclude that that we meet the desired degree of approximation for such \(m\):

\[\left\|f(X)_{i}-q\mathrm{SA}(X)_{i}\right\|_{2} =\left\|\sum_{i^{\prime}\in y_{i}}\left(\frac{1}{q}- \operatorname{softmax}(\phi(X)QK^{\mathsf{T}}\phi(X)^{\mathsf{T}})_{i,i^{ \prime}}\right)z_{i^{\prime}}\right\|_{2}\] \[\quad+\left\|\sum_{i^{\prime}\not\in y_{i}}\left(\operatorname{ softmax}(\phi(X)QK^{\mathsf{T}}\phi(X)^{\mathsf{T}})_{i,i^{\prime}}\right)z_{i^{\prime}} \right\|_{2}\] \[\leq q\cdot\frac{\epsilon}{2q}+(N-q)\cdot\frac{\epsilon}{2N}\leq\epsilon.\qed\]

#### b.1.1 Restricted isometry and orthogonality property

The proof relies on the restricted isometry and orthogonality property from the compressed sensing literature. For \(v\in\mathbb{R}^{N}\), let \(\operatorname{supp}(v)=\{i\in[N]:v_{i}\neq 0\}\).

**Definition 5**.: We say a matrix \(U\in\mathbb{R}^{m\times N}\) satisfies the _\((q,\delta)\)-restricted isometry and orthogonality property_ if

\[\|Uv\|_{2}^{2}\in[(1-\delta)\|v\|_{2}^{2},(1+\delta)\|v\|_{2}^{2}]\quad\text{ and}\quad|\langle Uv,Uv^{\prime}\rangle|\leq\delta\|v\|_{2}\|v^{\prime}\|_{2}\]

for all vectors \(v,v^{\prime}\in\mathbb{R}^{N}\) with \(|\operatorname{supp}(v)|\leq q\), \(|\operatorname{supp}(v^{\prime})|\leq 2q\), and \(\operatorname{supp}(v)\cap\operatorname{supp}(v^{\prime})=\emptyset\).

The first result shows the existence of a sign-valued matrix \(U\) that satisfies the desired distance-preserving property.

**Lemma 12** (Consequence of Theorem 2.3 of Mendelson et al. [2007] and Lemma 1.2 of Candes and Tao [2005]).: _There is an absolute constant \(C>0\) such that the following holds. Fix \(\delta\in(0,1/2)\) and \(q\in\mathbb{N}\). Let \(U\) denote a random \(m\times N\) matrix of independent Rademacher random variables scaled by \(1/\sqrt{m}\). If \(m\geq C(q\log N)/\delta^{2}\), then with positive probability, \(U\) satisfies the \((q,\delta)\)-restricted isometry and orthogonality property._

Sparse subsets of the columns of such a \(U\) can then be linearly separated from all other columns.

**Lemma 13** (Consequence of Lemma 2.2 in Candes and Tao [2005]).: _Fix \(\delta\in(0,1/2)\) and \(q\in\mathbb{N}\). Let matrix \(U=[u_{1},\ldots,u_{N}]\in\mathbb{R}^{m\times N}\) satisfy the \((q,\delta)\)-restricted isometry and orthogonality property. For every vector \(v\in\{0,1\}^{N}\) with \(\operatorname{supp}(v)\leq q\), there exists \(w\in\mathbb{R}^{m}\) satisfying_

\[\|w\|_{2} \leq\sqrt{q}/(1-2\delta),\] \[\langle u_{i},w\rangle =1 \text{if }v_{i}=1,\] \[|\langle u_{i},w\rangle| \leq\delta/(1-2\delta) \text{if }v_{i}=0.\]

### Proof of Theorem 3

**Theorem 3** (Infinite-precision).: _For fixed \(N\), \(m\geq\Omega(d^{\prime}+q)\) and \(\epsilon>0\), there exists some \(f\in\mathcal{A}^{\prime}_{d,m,d^{\prime}}\) that \(\epsilon\)-approximates \(q\mathrm{SA}\)._

The proof relies on the properties of _neighborly polytopes_, which we define.

**Definition 6** (Ziegler [2006]).: A polytope \(P\) is _\(q\)-neighborly_ if every subset of \(q^{\prime}\leq q\) vertices forms a \((q^{\prime}-1)\)-face.

We give a \(q\)-neighborly polytope below that we use for the construction. For vectors \(v_{1},\ldots,v_{N}\in\mathbb{R}^{m^{\prime}}\), let \(\operatorname{Conv}(v_{1},\ldots,v_{N})=\{\sum_{i=1}^{N}\alpha_{i}v_{i}:\alpha \in[0,1]^{N},\sum_{i}\alpha_{i}=1\}\) denote their convex hull.

**Fact 14** (Theorem 1 of Gale [1963]).: _For \(t\in\mathbb{R}\), let \(\theta(t)=(t,\ldots,t^{m^{\prime}})\in\mathbb{R}^{m^{\prime}}\). Then, for all distinct \(t_{1},\ldots,t_{N}\in\mathbb{R}\), the cyclic polytope \(\operatorname{Conv}(\theta(t_{1}),\ldots,\theta(t_{N}))\) is \(\frac{m^{\prime}}{2}\)-neighborly._

The proof of Theorem 3 is immediate from the aforementioned fact and the following lemma.

**Lemma 15**.: _Suppose there exists \(u_{1},\ldots,u_{N}\in\mathbb{R}^{m^{\prime}}\) such that \(\operatorname{Conv}(u_{1},\ldots,u_{N})\) is \(q\)-neighborly. Then, for any \(\epsilon>0\), there exists some \(f\in\mathcal{A}^{\prime}_{d,m,d^{\prime}}\) with fixed key vectors \(\phi(X)K=(u_{1},\ldots,u_{N})\) that \(\epsilon\)-approximates \(q\mathrm{SA}\)._

Proof.: The construction employs a similar look-up table MLP \(\phi\) to the one used in the proof of Theorem 2. We let the key and value embeddings be

\[\phi(X)K=((u_{1},1),\ldots,(u_{N},1))\in\mathbb{R}^{N\times(m^{\prime}+1)}\text {, and }\phi(X)V=(z_{1},\ldots,z_{N})\in\mathbb{R}^{N\times d}.\]

To set the query vectors, observe that for any face \(F\) of a polytope \(P\), there exists a hyperplane \(H_{F}\) such that \(F\subset H_{F}\) and \(P\setminus F\) lies entirely on one side of \(H_{F}\). Thus, for every \(y\in\binom{[N]}{q}\), there exists \(w^{\prime}_{y}\in\mathbb{R}^{m^{\prime}}\) and \(b_{y}\in\mathbb{R}\) such that

\[w^{\prime\prime}_{y}u_{i}+b_{y}\begin{cases}=1&\text{if }i\in y\text{,}\\ <1&\text{otherwise.}\end{cases}\]For \(\alpha>0\), let \(\phi(x_{i})^{\mathsf{T}}Q=\alpha w_{y}=\alpha(w_{y}^{\prime},b_{y})\).

We construct the MLP to satisfy \(\phi(x_{i})=(z_{k};w_{y_{i}};u_{i};1)\in\mathbb{R}^{m}\) for \(m=2m^{\prime}+2\) and set parameter weights accordingly. Following the softmax analysis of Theorem 3, a sufficiently large choice of \(\alpha\) ensures that \(\max_{i\in[N]}\left\lVert f(X)_{i}-q\mathrm{SA}(X)_{i}\right\rVert_{2}\leq\epsilon\).

### Proof of Theorem 4

**Theorem 4**.: _For any sufficiently large \(q\), any \(N\geq 2q+1\), and any \(d^{\prime}\geq 1\), there exists a universal constant \(c\) such that if \(mp\leq cq\), then no \(f\in\mathcal{T}^{1,1}_{d,m,d^{\prime},p}\) exists that \(\frac{1}{2q}\)-approximates \(q\mathrm{SA}\)._

Proof.: We first embed every instance of \(\mathrm{DISJ}\) with \(n=q\) into an instance of \(q\mathrm{SA}\) and prove that they correspond. We assume the existence of the a transformer \(f\in\mathcal{T}^{1,1}_{d,m,d^{\prime},p}\) that \(\frac{1}{2q}\)-approximates \(q\mathrm{SA}\) and implies the existence of an \(O(mp)\)-bit communication protocol that computes \(\mathrm{DISJ}\). An application of Fact 5 concludes the proof.

Consider an instance of \(\mathrm{DISJ}\) with \(a\in\{0,1\}^{q}\) and \(b\in\{0,1\}^{q}\) known by Alice and Bob respectively. We design an instance \(X=(z_{i};y_{i};i)_{i\in[N]}\) of \(q\mathrm{SA}\). For each \(j\in[2q]\), let \(y_{2q+1}=\{2i+a_{i}-1:i\in[q]\}\). Additionally, let

\[z_{j}=\begin{cases}e_{1}&\text{if $j$ is odd and $b_{(j-1)/2}=1$},\\ -e_{1}&\text{otherwise}.\end{cases}\]

All other inputs are set arbitrarily. Then,

\[q\mathrm{SA}(X)_{2q+1} =\frac{1}{q}\left\lvert\left\{j\in[2q]:\;j\in y_{2q+1},\;j\text{ is odd, and }a_{(j-1)/2}=1\right\}\right\rvert e_{1}\] \[\quad-\frac{1}{q}\left\lvert\left\{j\in[2q]:\;j\in y_{2q+1}\text { and }(j\text{ is even or }a_{(j-1)/2}=0)\right\}\right\rvert e_{1}\] \[=\frac{|\{i\in[q]:a_{i}b_{i}=1\}|-|\{i\in[q]:a_{i}b_{i}=0\}|}{q}e_ {1}.\]

Hence, \(q\mathrm{SA}(X)_{2q+1}=-e_{1}\) if and only if \(\mathrm{DISJ}(a,b)=0\).

It remains to show that this implies the existence of an efficient communication protocol that computes \(\mathrm{DISJ}(a,b)\). By the existence of \(f\), there exist \(Q,K,V:\mathbb{R}^{d}\to\mathbb{R}^{m}\) and \(\psi:\mathbb{R}^{m}\to\mathbb{R}^{d^{\prime}}\) such that

\[f(X)_{2q+1}=\psi\left(\frac{\sum_{i=1}^{N}\exp\left(Q(x_{2q+1})^{\mathsf{T}}K( x_{i})\right)V(x_{i})}{\sum_{i=1}^{N}\exp\left(Q(x_{2q+1})^{\mathsf{T}}K(x_{i}) \right)}\right).\]

The protocol is as follows:

1. From \(a\), Alice determines \(y_{2q+1}\) and then computes \(Q(x_{2q+1})\in\mathbb{R}^{m}\), which she sends to Bob. This transmission uses \(O(mp)\) bits.
2. Bob determines \(z_{1},\ldots,z_{2q}\) from \(b\). Using those and the information from Alice, he computes \(f(X)_{2q+1}\). He returns \(1\) if and only if \(f(X)_{2q+1}^{\mathsf{T}}e_{1}\geq-1+\frac{1}{q}\).

The protocol computes \(\mathrm{DISJ}(a,b)\) because \(f\) is a \(\frac{1}{2q}\)-approximation of \(q\mathrm{SA}\). Because any such protocol requires sharing \(\Omega(q)\) bits of information, we conclude that \(mp\leq cq\) for some \(c\).

### Optimality of Theorem 3 under restricted architectures

While the near-optimality of the bounded-precision self-attention construction in Theorem 2 is assured by the communication complexity argument of Theorem 4, it is not immediately apparent whether Theorem 3 is similarly optimal among infinite-precision self-attention models. Theorem 16 proves that this is indeed the case for a restricted family of architectures that resembles _cross-attention_ rather than self-attention.

**Theorem 16**.: _For input \(x_{1},\ldots,x_{N}\) satisfying \(x_{i}=(z_{i};y_{i};i)\), suppose \(\phi(x_{i})^{\mathsf{T}}Q=w(y_{i},i)\), \(\phi(x_{i})^{\mathsf{T}}K=u(i)\), and \(\phi(x_{i})^{\mathsf{T}}V=z_{i}\). Then, for any \(q<N\) and \(m\leq q(1-C\log_{N}q)\) for some universal \(C\), there do not exist \(w:\mathbb{R}^{d}\times[N]\rightarrow\mathbb{R}^{m}\) and \(u:[N]\rightarrow\mathbb{R}^{m}\) such that the resulting self-attention unit \(\frac{1}{2q}\)-approximates \(q\mathrm{SA}\)._

The architectural assumptions of this statement are strong. For each element \(x_{i}=(z_{i};y_{i};i)\), its value embedding must reproduce its target \(z_{i}\); its key embedding depends exclusively on the index \(i\); and its query embedding only on the indices \(y_{i}\) and \(i\). Indeed this attention unit more closely resembles _cross-attention_ rather than self-attention, in which the problem is formulated as two sequences \(((z_{1},1),\ldots,(z_{N},N))\) and \((y_{1};1),\ldots,(y_{N};N)\) that are passed to the key and value inputs and the query inputs respectively. We leave open the problem of generalizing this result to include all infinite-precision cross-attention or self-attention architectures, but we note that the constructions in Theorems 2 and 3 can be implemented under such architectural assumptions.

The proof relies on a geometric argument about how the convex hull of fixed key embeddings \(U=(u(1),\ldots,u(N))\) lacks neighborliness and hence cannot separate every size-\(q\) subsets of values embeddings \(z_{1},\ldots,z_{N}\) from the other values.

Proof.: It suffices to show that for any fixed key embedding \(U\), there exists some \(y_{i}\) and setting of \(z_{1},\ldots,z_{N}\) such that

\[\left\|(\mathrm{softmax}(w(X)U^{\mathsf{T}})Z)_{i}-\frac{1}{q}\sum_{i^{\prime} \in y_{i}}z_{i^{\prime}}\right\|_{2}\geq\frac{1}{2q},\]

where \(w(X)=(w(y_{1},1),\ldots,w(y_{N},N))\in\mathbb{R}^{N\times m}\) and \(U=(u(1),\ldots,u(N))\in\mathbb{R}^{N\times m}\).

By Fact 17, for some \(y_{1}\in\binom{[N]}{q}\), there are no \(w\) and \(\tau\in\mathbb{R}\) satisfying \(w(y_{1},1)^{\mathsf{T}}u_{i^{\prime}}\geq\tau\) if and only if \(i^{\prime}\in y_{1}\). Hence, for any fixed \(w\), there exists \(i_{1}\in y_{1}\) and \(i_{2}\in[N]\setminus y_{1}\) such that \(w(y_{1},1)^{\mathsf{T}}u_{i_{2}}>w(y_{1},1)^{\mathsf{T}}u_{i_{1}}\). Given the value embeddings \(z_{i_{1}}=e_{1},z_{i_{2}}=e_{2}\) and \(z_{i}=e_{3}\) for all \(i\not\in\{i_{1},i_{2}\}\), we have

\[\left\|(\mathrm{softmax}(w(X)U^{\mathsf{T}})Z)_{1}-\frac{1}{q}\sum _{i^{\prime}\in y_{1}}z_{i^{\prime}}\right\|_{2}^{2} \geq\left(\mathrm{softmax}(w(X)U^{\mathsf{T}})Z)_{1,i_{1}}-\frac{1 }{q}\right)^{2}+(\mathrm{softmax}(w(X)U^{\mathsf{T}})Z)_{1,i_{2}})^{2}\] \[\geq\max\left(\left(\mathrm{softmax}(w(X)U^{\mathsf{T}})Z)_{1,i_ {1}}-\frac{1}{q}\right)^{2},\mathrm{softmax}(w(X)U^{\mathsf{T}})Z)_{1,i_{1}}^{2}\right)\] \[\geq\frac{1}{4q^{2}}.\qed\]

**Fact 17**.: _If \(m^{\prime}<q(1-\log_{N}Cq)\), then the columns of any \(U=(u_{1},\ldots,u_{N})\in\mathbb{R}^{N\times m^{\prime}}\) can be partitioned into sets \(U_{1}\) and \(U_{2}\) with \(|U_{1}|=q\) that are not linearly separable. Hence, \(\mathrm{Conv}(u_{1},\ldots,u_{N})\) is not \(q\)-neighborly._

Proof.: By the Sauer-Shelah Lemma [22, 23, 24] and the fact that the VC dimension of \(m^{\prime}\)-dimensional linear thresholds is \(m^{\prime}+1\), the maximum number of partitions of the columns of \(U\) that can be linearly separated is at most

\[\sum_{k=0}^{m^{\prime}+1}\binom{N}{i}\leq C^{\prime}N^{m^{\prime}+1}<C^{\prime }\cdot\frac{N^{q}}{(Cq)^{q}}\leq\binom{N}{q},\]

for a sufficiently large choice of \(C\) given universal constant \(C^{\prime}\). If the fact were to be false, then at least \(\binom{N}{q}\geq(\frac{N}{q})^{q}\) such partitions must exist, which contradicts the above bound.

Supplementary results for Section 4

### Proof of Theorem 6

**Theorem 6**.: _For any input size \(N\), input range \(M=N^{O(1)}\), and fixed-precision bit complexity \(p=O(\log M)\), there exists a transformer architecture \(f\in\mathcal{T}^{1,1}_{1}\) with a single self-attention unit with embedding dimension \(m=3\) such that for all \(X\in[M]^{N}\), \(f(X)=\mathrm{Match2}(X)\)._

Proof.: As discussed in Section 2.1, we allow a single blank token to be appended to the end of the sequence \(X\) and assume the existence of a positional encoding. That is, we consider input \(X^{\prime}=(x_{1},\ldots,x_{N},x^{\prime})\) with \(x_{i,0}=i\) and \(x^{\prime}=\vec{0}\) to be the input to the target attention model. We define input \(\text{MLP}\;\phi:\mathbb{R}\to\mathbb{R}^{3}\) and parameterizations \(Q,K,V\in\mathbb{R}^{3\times 3}\) such that

\[Q^{\mathrm{ T}}\phi(x_{i})=c\left(\cos\left(\frac{2\pi x_{i}}{M}\right),\sin \left(\frac{2\pi x_{i}}{M}\right),1\right),\] \[K^{\mathrm{ T}}\phi(x_{i})=\left(\cos\left(\frac{2\pi x_{i}}{M}\right),-\sin \left(\frac{2\pi x_{i}}{M}\right),0\right),\]

\(V^{\mathrm{ T}}\phi(x_{i})=\vec{1}\), \(Q^{\mathrm{ T}}\phi(x^{\prime})=\vec{0}\), \(K^{\mathrm{ T}}\phi(x^{\prime})=e_{3}\), and \(V^{\mathrm{ T}}\phi(x^{\prime})=\vec{0}\). By elementary trigonometric identities, the following is true about the corresponding inner products:

\[(Q^{\mathrm{ T}}\phi(x_{i}))^{\mathrm{ T}}K^{\mathrm{ T}}\phi(x_{j})=c\cos\left(\frac{2\pi(x_{i}+x_{j})}{M}\right)\] \[(Q^{\mathrm{ T}}\phi(x_{i}))^{\mathrm{ T}}K^{\mathrm{ T}}\phi(x^{\prime})=cd.\]

As a result, \((Q^{\mathrm{ T}}\phi(x_{i}))^{\mathrm{ T}}K^{\mathrm{ T}}\phi(x_{j})=cd\) if and only if \(x_{i}+x_{j}=(\!\!\!\pmod{M}\). Otherwise, \((Q^{\mathrm{ T}}\phi(x_{i}))^{\mathrm{ T}}K^{\mathrm{ T}}\phi(x_{j})\leq c(1-\frac{1}{M^{2}})\). (Here, the \(O(\log M)\)-bit fixed-precision arithmetic is sufficient to numerically distinguish the two cases.) For each \(i\in[N]\) let

\[\beta_{i}=|\{j\in[N]:x_{i}+x_{j}=0\,(\!\!\!\pmod{M})\}|\]

represent the total number of matches the input belongs to. If we take \(c=M^{2}\log(6N)\), then

\[(\mathrm{softmax}(\phi(X)QK^{\mathrm{ T}}\phi(X)^{\mathrm{ T}}))_{i,j}\in\begin{cases}[0,\frac{1}{6N}]&\text{if }x_{i}+x_{j}\neq 0\,(\!\!\!\pmod{M})\text{ and }\,i,j\in[N];\\ \lfloor\frac{1}{\beta_{i}+1}\pm\frac{1}{6N}\rfloor&\text{if }x_{i}+x_{j}=0\,(\!\!\!\pmod{M})\text{ and }\,i,j\in[N];\\ \lfloor\frac{1}{\beta_{i}+1}\pm\frac{1}{6N}\rfloor&\text{if }i\in[N],j=N+1.\end{cases}\]

We conclude that for any \(i\in[N]\),

\[(\mathrm{softmax}(\phi(X)QK^{\mathrm{ T}}\phi(X)^{\mathrm{ T}})V\phi(X))_{i}\begin{cases}\leq\frac{1}{6}\cdot\vec{1}&\text{if }

Note that \(\mathrm{Match3}(X)_{1}=1\) if and only if there exists \(i\in\{2,\ldots,\frac{N+1}{2}\}\) such that \(x_{i}=i\) and \(x_{i+\frac{N-1}{2}}=(M-i)\). Given input \((a,b)\in\{0,1\}^{n}\times\{0,1\}^{n}\) to \(\mathrm{DISJ}\), let \(x_{i+1}=1\) if and only if \(a_{i}=0\), and let \(x_{i+\frac{N+1}{2}}=1\) if and only if \(b_{i}=0\). Then, \(\mathrm{Match3}(X)_{1}=1\) iff \(\mathrm{DISJ}(a,b)=1\).

Suppose \(f(X)=\mathrm{Match3}(X)\) for all \(X\in[M]^{N}\) for some \(f\in\mathcal{T}^{1,H}_{1,m,1,p}\). We show that \(f\) simulates an \(O(mpH)\)-bit communication protocol for testing \(\mathrm{DISJ}\). By definition of the standard self-attention unit with multi-layer perceptrons, note that \(f(X)_{1}=\psi(\sum_{h=1}^{H}f_{h}(\phi(X)))\) for \(\phi:\mathbb{R}\to\mathbb{R}^{m}\), \(\psi:\mathbb{R}^{m}\to\{0,1\}\), and

\[f_{h}(X)=\frac{\sum_{i=1}^{N}\exp(Q_{h}(x_{1})^{\mbox{\tiny\sf T}}K_{h}(x_{i} ))V_{h}(x_{i})}{\sum_{i=1}^{N}\exp(Q_{h}(x_{1})^{\mbox{\tiny\sf T}}K_{h}(x_{i} ))},\]

for \(Q_{h},K_{h},V_{h}:\mathbb{R}^{m\times m}\).

If we assume that this construction exists and is known explicitly by both Alice and Bob, we design a communication protocol for Alice and Bob to solve \(\mathrm{DISJ}\) by sharing \(O(mpH)\) bits with one another. Let Alice possess \(a\in\{0,1\}^{n}\) and Bob \(b\in\{0,1\}^{n}\), with \(n=\frac{N-1}{2}\).

1. Alice and Bob compute \((x_{2},\ldots,x_{\frac{N+1}{2}})\) and \((x_{\frac{N+2}{2}},\ldots,x_{N})\) from \(a\) and \(b\) respectively.
2. Alice computes an \(O(p\log\log N)\)-bit approximation of the logarithm of the first half of the softmax normalization term for each attention head and sends the result to Bob. That is, she sends Bob \[L_{h,a}=\log\left(\sum_{i=1}^{\frac{N+1}{2}}\exp(Q_{h}(\phi(x_{1}))^{\mbox{ \tiny\sf T}}K_{h}(\phi(x_{i})))\right)\] for each \(h\in[H]\). This requires transmitting \(O(pH\log\log N)\) bits.
3. Bob finishes the computation of normalization terms \[L_{h}=\log\left(\exp(L_{h,a})+\sum_{i=\frac{N+3}{2}}^{N}\exp(Q_{h}(\phi(x_{1} ))^{\mbox{\tiny\sf T}}K_{h}(\phi(x_{i})))\right)\] for each \(h\) and sends the result back to Alice (up to \(O(p\log\log N)\)-bits of precision). This again requires transmitting \(O(pH\log\log N)\) bits.
4. Alice computes the partial convex combination of the first \(\frac{N+1}{2}\) value vectors stipulated by the attention matrix \[S_{h,a}=\frac{\sum_{i=1}^{\frac{N+1}{2}}\exp(Q_{h}(\phi(x_{1}))^{\mbox{\tiny \sf T}}K_{h}(\phi(x_{i})))V_{h}(\phi(x_{i}))}{\exp(L_{h})}\in\mathbb{R}^{m}\] for each \(h\) and sends the partial combinations to Bob. This requires transmitting \(O(mpH\log\log N)\) bits (using the same precision as above).
5. Bob finishes the computation of the convex combinations \[f_{h}(X)=S_{h,a}+\frac{\sum_{i=\frac{N+3}{2}}^{N}\exp(Q_{h}(\phi(x_{1}))^{ \mbox{\tiny\sf T}}K_{h}(\phi(x_{i})))V_{h}(\phi(x_{i}))}{\exp(L_{h})}\in \mathbb{R}^{m}.\] Bob concludes the protocol by computing and outputting \(f(X)_{1}\), using his knowledge of each \(f_{h}(X)\) and of \(\psi\).

By the equivalences previously established, Bob returns 1 if and only if \(\mathrm{DISJ}(a,b)=1\). Because the protocol requires \(O(mpH\log\log N)\) bits of communication, we can only avoid contradicting Fact 5 if \(mpH\geq\Omega(n/\log\log N)=\Omega(N/\log\log N)\).

**Remark 1**.: _The domain restrictions to \(\mathrm{Match3}\) stipulated in Equation (3) make the \(\mathrm{Match3}\) problem substantially easier to solve than the full-domain case. Indeed, under the domain restrictions,_

\[\mathrm{Match3}(X)_{1}=\max_{i\in\{2,\ldots,\frac{N+1}{2}\}}\mathrm{Match2}(X) _{i},\]which is computable by a two-layer single-headed transformer network with constant embedding dimension. The first layer computes each \(\mathrm{Match2}(X)_{i}\) with the construction in the proof of Theorem 6, and the second computes the maximum of the previous outputs by using those outputs as key vectors. While Informal Conjecture 1 suggests that two layers are insufficient to compute the full-domain version of \(\mathrm{Match3}\), this restricted variant introduces a concise depth separation (see Eldan and Shamir [2016], Telgarsky [2016], Daniely [2017]) between one- and two-layer transformer models.

### Higher-order tensor attention

We introduce a novel category of higher-order tensor-based transformer models in order to show that problems like \(\mathrm{Match3}\) that are hard to compute with standard transformer models can be made solvable. An \(s\)-order transformer is designed to efficiently compute dense \(s\)-wise interactions among input elements in an analogous manner to how standard transformers compute pairwise interactions. (We think of a standard transformer as second-order.) Before defining the new type of attention, we introduce notation to express the needed tensor products.

For vectors \(v^{1}\in\mathbb{R}^{N_{1}}\) and \(v^{2}\in\mathbb{R}^{N_{2}}\), let \(v^{1}\otimes v^{2}\in\mathbb{R}^{N_{1}N_{2}}\) denote their _Kronecker product_ by \((v^{1}\otimes v^{2})_{(i_{1}-1)N_{2}+i_{2}}=v^{1}_{i_{1}}v^{2}_{i_{2}}\). The _column-wise Kronecker product_ of matrices \(A^{1}\in\mathbb{R}^{N_{1}\times m}\) and \(A^{2}\in\mathbb{R}^{N_{2}\times m}\) is

\[A^{1}\star A^{2}=[A^{1}_{1}\mid\dots\mid A^{1}_{m}]\star[A^{2}_{1}\mid\dots \mid A^{2}_{m}]=[A^{1}_{1}\otimes A^{2}_{1}\mid\dots\mid A^{1}_{m}\otimes A^{2 }_{m}]\in\mathbb{R}^{N_{1}N_{2}\times m}.\]

The following generalizes the definition of self-attention.

**Definition 7**.: For order \(s\geq 2\), input dimension \(d\), output dimension \(d^{\prime}\), embedding dimension \(m\), bit complexity \(p\), and matrices \(Q,K^{1},\dots,K^{s-1}\in\mathbb{R}^{d\times m}\) and \(V^{1},\dots,V^{s-1}\in\mathbb{R}^{d\times d^{\prime}}\) (encoded with \(p\)-bit fixed-point numbers), an \(s\)-_order self-attention unit_ is a function \(f_{Q,K,V}:\mathbb{R}^{N\times d}\to\mathbb{R}^{N\times d^{\prime}}\) with

\[f_{Q,K,V}(X)=\mathrm{softmax}(\underbrace{XQ}_{\in\mathbb{R}^{N\times m}} \underbrace{((XK^{1})\star\dots\star(XK^{s-1}))^{\gamma}}_{\in\mathbb{R}^{m \times N^{s-1}}})(\underbrace{((XV^{1})\star\dots\star(XV^{s-1}))}_{\in \mathbb{R}^{N^{s-1}\times d^{\prime}}}.\]

The input to the row-wise softmax is an \(N\times N^{s-1}\) matrix. Let \(\mathcal{A}_{d,m,d^{\prime},p}^{\otimes s}\) denote the set containing all such attention units.

Note that \(\mathcal{A}_{d,m,d^{\prime},p}^{\otimes 2}=\mathcal{A}_{d,m,d^{\prime},p}\). Because \(s\)-order self-attention units have the same domain and codomain as standard self-attention, multiple units can be analogous combined to construct multi-headed attention units and full transformer models. We define \(\mathcal{A}_{d,m,d^{\prime},p}^{\otimes s}\) and \(\mathcal{T}_{d,m,d^{\prime},p}^{D,H,\otimes s}\) accordingly.

The purpose of the \(s\)-order transformer model as a theoretical construct is to posit how strictly generalizing the architecture in order to permit higher order outer products transfers the expressive powers of standard transformer architectures to more sophisticated interactions among elements of the input sequence \(X\). The model is not defined to be immediately practical, due to its steep computational cost of evaluation.

However, the trade-offs involved in using such architectures resemble those already made by using transformer models instead of fully-connected networks. Transformers are already computationally wasteful relative to the number of the parameters, and these models likely succeed only because extremely efficient factorized parameterization exist. Likewise, third-order transformers could indeed be practical if even more factorization proves useful, since the computational costs may prove mild if the embedding dimension \(m\), number of heads \(H\), and depth \(D\) necessary to succeed on a task exceed the sequence length \(N\) for standard second-order transformers.

### Efficient representation of \(\mathrm{Match3}\) with third-order self-attention

**Theorem 18** (\(\mathrm{Match3}\) construction with third-order self-attention).: _For any sequence length \(N\), input range \(M=N^{O(1)}\), and fixed-precision bit complexity \(p=O(\log M)\), there exists a third-order transformer architecture \(f\in\mathcal{T}_{1,m,1,p}^{1,\otimes 3}\) with a single self-attention unit with embedding dimension \(m=5\) such that for all \(X\in[M]^{N}\), \(f(X)=\mathrm{Match3}(X)\)._Proof of Theorem 18.: The proof is almost identical to that of Theorem 6, except that we instead use a different key and query transforms to express a different trigonometric function:

\[Q\phi(x_{i}) =c\left(\cos\left(\frac{2\pi x_{i}}{M}\right),-\cos\left(\frac{2\pi x _{i}}{M}\right),\sin\left(\frac{2\pi x_{i}}{M}\right),\sin\left(\frac{2\pi x_{i }}{M}\right),1\right),\] \[K^{1}\phi(x_{i}) =\left(\cos\left(\frac{2\pi x_{i}}{M}\right),\sin\left(\frac{2\pi x _{i}}{M}\right),-\cos\left(\frac{2\pi x_{i}}{M}\right),\sin\left(\frac{2\pi x_ {i}}{M}\right),0\right),\] \[K^{2}\phi(x_{i}) =\left(\cos\left(\frac{2\pi x_{i}}{M}\right),\sin\left(\frac{2\pi x _{i}}{M}\right),\sin\left(\frac{2\pi x_{i}}{M}\right),-\cos\left(\frac{2\pi x_ {i}}{M}\right),0\right).\]

Together, these ensure that the resulting tensor products reduce to a trigonometric expression that is maximized when \(x_{i}+x_{j_{1}}+x_{j_{2}}=0\pmod{M}\). That is,

\[(\phi(X)Q((\phi(X)K^{1})\star(\phi(X)K^{2}))^{\mathsf{T}})_{i,(j_{1}-1)+j_{2}} =c\cos\left(\frac{2\pi(x_{i}+x_{j_{1}}+x_{j_{2}})}{M}\right).\]

We similarly let \(V^{1}\phi(x_{i})=V^{2}\phi(x_{i})=\vec{1}\) and \(V^{1}\phi(x^{\prime})=V^{2}\phi(x^{\prime})=\vec{0}\). The remaining choice of \(c\) and the output MLP, and the analysis of the softmax proceeds identically to the previous proof.

### Heuristic argument for Informal Conjecture 1

**Conjecture 19** (Formal version of Informal Conjecture 1).: _For sufficiently large \(N\) and any \(d\geq 1\), for all \(M\geq N+1\) and \(mpHD\leq N^{\Omega(1)}\), there is no \(f\in\mathcal{T}^{D,H}_{1,m,1,p}\) satisfying \(f(X)=\mathrm{Match3}(X)\) for all \(X\in[M]^{N}\)._

We believe that the conjecture holds due to a heuristic information-theoretic argument. Define the distribution \(\mathcal{D}\) over inputs \(X\in\mathbb{R}^{N}\) that will be used to show that the model cannot compute \(\mathrm{Match3}\) for \(M=N^{4}\) with high probability. We draw \(\mathbf{X}\) from \(\mathcal{D}\) as follows:

1. With probability \(\frac{1}{2}\), draw each \(\mathbf{x}_{i}\) iid from \(\mathrm{Unif}([M])\).
2. With probability \(\frac{1}{2}\), draw \(j_{1},j_{2},j_{3}\) iid from \(\mathrm{Unif}(\binom{[N]}{3})\). For all \(i\neq j_{3}\), draw each \(\mathbf{x}_{i}\) iid from \(\mathrm{Unif}([M])\). Let \(\mathbf{x}_{j_{3}}=-\mathbf{x}_{j_{1}}-\mathbf{x}_{j_{2}}\pmod{M}\).

Note that under event \(E_{1}\), a three matching elements exist with probability at most \(\frac{1}{N}\), and

\[\Pr\left[\mathrm{Match3}(\mathbf{X})=\vec{0}\mid E_{1}\right]\geq 1-\frac{1}{N}.\]

Under event \(E_{2}\), a triple of matching elements is always planted, so \(\mathrm{Match3}(\mathbf{X})\neq\vec{0}\). It would suffice to prove that--unless a transformer is sufficiently large--it is impossible to determine whether \(\mathrm{Match3}(\mathbf{X})=\vec{0}\) with probability at least \(0.9\).

Under \(\mathcal{D}\), any subset of \(\{\mathbf{x}_{1},\ldots,\mathbf{x}_{N}\}\) consists of iid integers drawn uniformly from \([M]\), unless all of \(\mathbf{x}_{j_{1}},\mathbf{x}_{j_{2}},\mathbf{x}_{j_{3}}\) appear in the subset. Consider a transformer architecture with \(p\)-bit precision, \(m\)-dimensional embeddings, \(H\) heads per layer, and \(D\) layers. We argue informally that a single-element output of a self-attention unit can take into account information about \(mp\) more inputs \(\mathbf{x}_{1},\ldots,\mathbf{x}_{N}\) than that it had in the previous layer. By induction, after \(D\) layers of \(H\)-headed self-attention with interleaved MLPs, each element is a function of at most \(mpHD\) inputs. Until an element exists that is a function of at least two of the three of \(\mathbf{x}_{j_{1}},\mathbf{x}_{j_{2}},\mathbf{x}_{j_{3}}\), we assume that the elements "known" by each output are chosen independently of the indices \(j_{1},j_{2},j_{3}\). (Given two elements of the triple, the third element can be identified with a single self-attention unit.) Hence, we argue that it suffices to show that the probability any two elements of the triple \(j_{1},j_{2},j_{3}\) occurring within any of the \(N\) sets of \(mpHD\) inputs is vanishingly small for sufficiently large transformer parameters. The probability of single collection having any of two of the three inputs is at most

\[\frac{3\binom{mpHD}{2}}{\binom{N}{2}}\leq 3\left(\frac{empHD}{N}\right)^{2}.\]

Thus, the probability that any collection has all three inputs is no more than \(3(empHD)^{2}/N\). If \(mpHD=O(\sqrt{N})\), then the randomly chosen triple will not jointly appear as the outcome of a single element of a self-attention unit with probability at least \(0.9\), and the transformer will be unexpected to successfully distinguish between the two cases.

Should the conjecture hold, it would represented a tight lower bound on the size of the smallest standard transformer architecture necessary to compute \(\mathrm{Match3}\).

**Theorem 20** (Tightness of Conjecture 19).: _For any sequence length \(N\), if the input range satisfies \(M=N^{O(1)}\) and the transformer size parameters satisfy \(p\geq\log(M)\), \(H=1\), \(m\geq 4\), and \(mD\geq CN^{2}\) for some universal constant \(C\), then there exists a transformer architecture \(f\in\mathcal{T}^{D,H}_{1,m,1,p}\) such that \(f(X)=\mathrm{Match3}(X)\)._

Proof.: We construct an architecture that collects a group of candidate pairs in each layer of single-headed self-attention and verifies whether there exists a triple incorporating each pair that satisfies the summation property. Then, all candidate triples are disposed of, and the subsequent layer collects a new family of candidates.

To do so, we first let \(\ell:=\left\lfloor\frac{m}{2}\right\rfloor-1\geq 1\) represent the total number of pairs shared in each layer of attention. We let \(P=\binom{|N|}{2}\) represent a collection of all pairs of indices and partition it into \(D\) subsets \(P_{1},\ldots,P_{D}\), each containing \(\ell\) distinct pairs. (Since \(|P|=\frac{N(N+1)}{2}\), any \(D\) satisfying the theorem's preconditions is sufficiently large for this to be a proper partition.) Our construction ensures that there exist \(x_{i}+x_{j_{1}}+x_{j_{2}}=0\pmod{M}\) for \((j_{1},j_{2})\in P_{k}\), then the \(k\)th layer of self attention will verify its existence and mark \(x_{i}\) as belonging to the match. Throughout the network, we maintain that the first two dimensions of any embedding of the \(i\)th element correspond to \(x_{i}\in[M]\) and a bit indicating whether a match has been found yet containing \(x_{i}\).

Consider the first layer of self-attention, and let \(P_{1}=\{(i_{1},j_{1}),\ldots,(i_{\ell},j_{\ell})\}\). We set the input MLP \(\phi_{1}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{m}\) and respective matrices \(Q^{1},K^{1}\in\mathbb{R}^{m\times m}\) such that

\[Q^{1}\phi_{1}(x_{i})=ce_{1}\text{ and }K^{1}\phi_{1}(x_{i})=\begin{cases}e_{1}& \text{if }i\in P_{1}\\ \vec{0}&\text{otherwise,}\end{cases}\]

for sufficiently large \(c\). We additionally let

\[V^{1}\phi_{1}(x_{i})=\begin{cases}(2\ell+1)\cdot(x_{i};0;\vec{0})&i\not\in P_{ 1},\\ (2\ell+1)\cdot(x_{i};0;x_{i}e_{2-1})&i=i_{\iota},\\ (2\ell+1)\cdot(x_{i};0;x_{i}e_{2\iota})&i=j_{\iota}.\end{cases}\]

By making use of a residual connection, we ensure that the \(i\)th outcome of the self-attention is \((x_{i},0,x_{i_{1}},x_{j_{1}},\ldots,x_{i_{\ell}},x_{j_{\ell}})\). We encode an MLP to compute

\[(x_{i},0,x_{i_{1}},x_{j_{1}},\ldots,x_{i_{\ell}},x_{j_{\ell}})\mapsto\left(x_ {i},\mathbbm{1}\left\{\exists\iota\in[\ell]\operatorname{s.t.}x_{i}+x_{i_{ \iota}}+x_{j_{\iota}}=\vec{0}\pmod{M}\right\};\vec{0}\right).\]

We repeat this construction \(D\) times, with the only modifications being the replacement of \(P_{1}\) and the fact that the second dimension of the embedding remains 1 after being set to that value. After \(D\) layers, the final MLP outputs the value of the second dimension, which will be 1 if and only if the respective \(x_{i}\) belongs to a three-way match.

### Sharper separations for embedded subgraph detection problems

In pursuit of proving separations analogous to the one between Theorem 18 and Conjecture 19, we draw techniques for proving lower bounds for graph problems in the Congest model of distributed computation with restricted bandwidth [20].5The problems we consider take, as input, the adjacency matrix \(X\in\{0,1\}^{N\times N}\) of an \(N\)-vertex graph \(G=(\mathcal{V},\mathcal{E})\) with \(\mathcal{V}=[N]\), so \(x_{i,j}=\mathbbm{1}\left\{(i,j)\in\mathcal{E}\right\}\). We may regard each row of \(X\) as a high-dimensional (\(d=N\)) embedding of the \(i\)-th vertex containing information about which (outgoing) edges are incident to the \(i\)-th vertex. We consider the following problems:

\[\mathrm{DirectedCycle3}(X) =(\mathbbm{1}\left\{\exists j_{1},j_{2}\in[N]\;\mathrm{s.t.}\;x_{ i,j_{1}}x_{j_{1},j_{2}}x_{j_{2},i}=1\right\})_{i\in[N]}\,;\] \[\mathrm{Cycle5}(X) =(\mathbbm{1}\left\{\exists j_{1},j_{2},j_{3},j_{4}\in[N]\; \mathrm{s.t.}\;x_{i,j_{1}}x_{j_{1},j_{2}}x_{j_{2},j_{3}}x_{j_{3},j_{4}}x_{j_{4 },i}=1\right\})_{i\in[N]}\,,\] \[\text{with }\mathrm{dom}(\mathrm{Cycle5})=\{X:\,X=X^{\uptau}\}.\]

The former treats \(X\) as a directed graph (where \(X\) need not be symmetric) and asks whether each input belongs to a directed 3-cycle. The latter insists that \(X\) be an undirected graph by enforcing symmetry and determines membership in (undirected) 5-cycles.

However, solving these problems with any transformer model of constant order trivially requires having the product of the precision \(p\), embedding dimension \(m\), heads per layer \(H\), and depth \(D\) grow polynomially with \(N\), since each attention unit is limited to considering at most \(pm\) bits of information from each input. Such a lower bound is not interesting for dense graphs, where every vertex may have \(\Omega(N)\) incident edges; the bottleneck is not due to any feature of standard attention units (and would persist with higher-order attention).

To circumvent this issue, we consider an augmented self-attention unit, which permits each element of the self-attention tensor to depend on both its respective inner product and on the presence of edges among corresponding inputs.

**Definition 8**.: For order \(s\geq 2\), input dimension \(d\), output dimension \(d^{\prime}\), embedding dimension \(m\), bit complexity \(p\), matrices \(Q,K^{1},\ldots,K^{s-1}\in\mathbb{R}^{d\times m}\) and \(V^{1},\ldots,V^{s-1}\in\mathbb{R}^{d\times d^{\prime}}\) (encoded with \(p\)-bit fixed-point numbers), and cell-wise attention tensor function \(\kappa:\{0,1\}^{s(s-1)}\times\mathbb{R}\rightarrow\mathbb{R}\), an _s-order graph self-attention unit_ is a function \(f_{Q,K,V}:\mathbb{R}^{N\times d}\rightarrow\mathbb{R}^{N\times d^{\prime}}\) with

\[f_{Q,K,V}(X)=\mathrm{softmax}(\kappa(X,XQ((XK^{1})\star\cdots\star(XK^{s-1}))^ {\uptau}))((XV^{1})\star\cdots\star(XV^{s-1})).\]

For attention tensor \(A\in\mathbb{R}^{N^{\otimes s}}\), we abuse notation by writing \(\kappa(X,A)\) as short-hand for the particular cell-wise application of a fixed function, incorporating information about all relevant edges:

\[\kappa(X,A)_{i_{1},\ldots,i_{s}}=\kappa(x_{i_{1},i_{2}},x_{i_{1},i_{3}}, \ldots,x_{i_{s},i_{s-1}},x_{i_{s},i_{s-2}},A_{i_{1},\ldots,i_{s}}).\]

Let \(\mathcal{AG}_{d,m,d^{\prime},p}^{\otimes s}\) and \(\mathcal{TG}_{d,m,d^{\prime},p}^{D,H,\otimes s}\) denote all such attention units and all such transformers respectively.

Now, we provide four results that exhibit separations between orders of graph self-attention.

**Theorem 21** (Hardness of representing \(\mathrm{Cycle5}\) with standard graph transformer).: _For sufficiently large \(N\), any \(f\in\mathcal{TG}_{N,m,1,p}^{D,H}\) satisfying \(f(X)=\mathrm{Cycle5}(X)\) for all \(X\in\{0,1\}^{N\times N}\) with \(X=X^{\uptau}\) requires \(mpHD=\Omega(N/\log^{2}N)\)._

**Theorem 22** (Efficient construction of \(\mathrm{Cycle5}\) with fifth-order graph transformer).: _For sequence length \(N\) and bit-complexity \(p=O(\log N)\), there exists a fourth-order graph transformer architecture \(f\in\mathcal{TG}_{N,1,1,p}^{1,1,\otimes 3}\) with a single graph self-attention unit such that for all \(X\in\{0,1\}^{N\times N}\) with \(X=X^{\uptau}\), \(f(X)=\mathrm{Cycle5}(X)\)._

**Theorem 23** (Hardness of representing \(\mathrm{DirectedCycle3}\) with standard graph transformer).: _For sufficiently large \(N\), any \(f\in\mathcal{TG}_{N,m,1,p}^{D,H}\) satisfying \(f(X)=\mathrm{DirectedCycle3}(X)\) for all \(X\in\{0,1\}^{N\times N}\) requires \(mpHD=\Omega(N/\log^{2}N)\)._

**Theorem 24** (Efficient construction of \(\mathrm{DirectedCycle3}\) with fourth-order graph transformer).: _For sequence length \(N\) and bit-complexity \(p=O(\log N)\), there exists a third-order graph transformer architecture \(f\in\mathcal{TG}_{N,1,1,p}^{1,1,\otimes 3}\) with a single graph self-attention unit such that for all \(X\in\{0,1\}^{N\times N}\), \(f(X)=\mathrm{DirectedCycle3}(X)\)._

The proofs of Theorems 22 and 24 are immediate from the construction. Because each cell of the self-attention tensor has explicit access the the existence of all relevant edges, \(\kappa\) can be configured to ensure that cell's value is large if and only if the requisite edges for the desired structure all exist. Taking a softmax with a blank element (like in Theorem 6) ensures that the outcome of the self-attention unit for a given element distinguishes between whether or not it belongs to a 5-cycle or a directed 3-cycle. The output MLP ensure that the proper output is returned.

We prove Theorems 21 and 23 by introducing a particular Congest communication graph that can be used to simulate any model in \(\mathcal{T}\mathcal{G}^{D,H}_{d,m,d^{\prime},p}\) (and hence, also any model in \(\mathcal{T}^{D,H}_{d,m,d^{\prime},p}\)) in \(O(mHD\log N)\) rounds of communication. Then, we show for each problem that we can encode each instance of the set disjointness communication problem as an instance of \(\mathrm{Cycle5}\) (or \(\mathrm{DirectedCycle3}\)) and derive a contradiction from the communication graph.

#### c.6.1 A Congest communication graph that generalizes standard graph transformer computation

The key principle of our analysis is that the predominant limitation of a transformer model is in its communication bandwidth and _not_ its computational abilities. We model transformers as having element-wise multi-layer perceptron units with unbounded computational ability (but bounded precision inputs and outputs) and self-attention units, which compute linear combinations of inputs in a carefully regimented way that limits the ability of individual elements to share information with one another. Here, we introduce a specific Congest graph for each sequence length \(N\) and show that every transformer has a communication protocol that simulates its computation in this graph.

For fixed \(N\), we design an undirected Congest graph \(G^{N}=(V^{N},E^{N})\) with \(O(N^{2})\) nodes, each having degree at most 3. (Note that this graph is _not_ the same as the graph provided as input \(X\) to a transformer; this graph is consistent across all transformers taking input of sequence size \(N\).) Let \(u_{1},\ldots,u_{N}\) be nodes in \(V^{N}\) corresponding to each input. For every pair \(i,j\in[N]\), let \(v_{i,j}\) be a node as well. For each \(i\in[N]\), let \(B_{i}=(V_{i},E_{i})\) be a balanced binary trees having root \(u_{i}\) and leaves \(v_{i,1},\ldots,v_{i,N},v_{1,i},\ldots,v_{N,i}\). Hence, each \(B_{i}\) has \(O(N)\) vertices of degree 3 and is of depth \(O(\log N)\). Let \(V^{N}=V_{1}\cup\cdots\cup V_{N}\) and \(E^{N}=E_{1}\cup\cdots\cup E_{N}\). Noting that \(E_{1},\ldots,E_{N}\) are disjoint and that \(V_{1},\ldots,V_{N}\) are disjoint, except for leaves \(v_{i,j}\), we ascertain that \(G^{N}\) contains \(O(N^{2})\) vertices of degree at most 3 and has diameter \(O(\log N)\). We visualize the graph \(G^{N}\) with a highlighted tree \(B_{1}\) in Figure 4.

**Lemma 25**.: _For any transformer \(f\in\mathcal{T}\mathcal{G}^{D,H}_{d,m,d^{\prime},p}\) and any \(X\in\mathbb{R}^{N\times d}\) with \(p\)-bit fixed-precision numbers, there exists a Congest communication protocol on the graph \(G^{N}\) that shares \(p\) bits of information between adjacent vertices per round satisfying the following characteristics:_

Figure 4: The Congest graph \(G^{N}\) visualized for \(N=6\) with root nodes \(\{u_{i}\}_{i\in[N]}\) in blue, leaf nodes \(\{v_{i,j}\}_{i,j\in[N]}\) in green, and the nodes \(V_{1}\) of the binary tree \(B_{1}\) shaded red and edges \(E_{1}\) colored red.

* _Before any communication begins, each node_ \(u_{i}\) _is provided with_ \(x_{i}\) _and each node_ \(v_{i,j}\) _is provided with_ \(x_{i,j}\) _and_ \(x_{j,i}\)_._
* _After_ \(T=O(HD(m+\log N))\) _rounds of communication, each node_ \(u_{i}\) _outputs_ \(f(X)_{i}\)_._

Proof.: It suffices to give a protocol that computes the outcome of a single-headed unit of graph self-attention with parameters \(Q,K,V\in\mathbb{R}^{m\times m}\) and \(\kappa:\{-1,1\}^{2}\times\mathbb{R}\rightarrow\mathbb{R}\) and transmits its \(i\)th output back to \(u_{i}\) in \(O(m\log N)\) rounds of \(p\)-bit communication. The remainder of the argument involves computing the outcomes of all element-wise MLPs within respective vertices \(u_{1},\ldots,u_{N}\) (since we assume each node to have unbounded computational power in the Congest model) and to repeat variants of the protocol \(HD\) times for every individual self-attention unit. Because the protocol is designed for a particular transformer architecture \(f\), we can assume that every node in the Congest graph has knows every parameter of \(f\).

We give the protocol in stages. We assume inductively that every input to \(f\), \(y_{1},\ldots,y_{N}\in\mathbb{R}^{m}\), is known by its respective vertex \(u_{1},\ldots,u_{N}\).

1. Every vertex \(u_{i}\) computes \(Q^{\intercal}y_{i}\in\mathbb{R}^{m}\) and propagates it to every vertex \(v_{i,1},\ldots,v_{i,N}\). This can be done in \(O(m+\log N)\) rounds by transferring one \(p\)-bit fixed-precision number per round from an element of the binary tree \(B_{i}\) to each of its children per round. Because the respective edges \(E_{1},\ldots,E_{N}\) are disjoint, this operation can be carried out in parallel.
2. Each \(u_{i}\) computes \(K^{\intercal}y_{i},V^{\intercal}y_{i}\in\mathbb{R}^{m}\) and propagates them to \(v_{1,i},\ldots,v_{N,i}\) in \(O(m+\log N)\) rounds.
3. Each \(v_{i,j}\), using their knowledge of \(x_{i,j}\) and \(x_{j,i}\), computes \(\alpha_{i,j}\coloneqq\exp(\kappa(x_{i,j},x_{j,i},y_{i}^{\intercal}QK^{\intercal }y_{j}))\). This takes zero rounds.
4. Each \(u_{i}\) computes \(\sum_{j=1}^{N}\alpha_{i,j}\) by propagating each \(\alpha_{i,j}\) in \(v_{i,j}\) up \(B_{i}\) to \(u_{i}\), iteratively summing terms passed up. This takes \(O(\log N)\) rounds.
5. Similarly, \(u_{i}\) computes \(\sum_{j=1}^{N}\alpha_{i,j}V^{\intercal}y_{j}\) in \(O(m\log N)\) rounds. Then, it computes \[\frac{\sum_{j=1}^{N}\alpha_{i,j}V^{\intercal}y_{j}}{\sum_{j=1}^{N}\alpha_{i,j }},\] which is the target output of the self-attention unit.

Because all steps are achievable in parallel with \(O(m+\log N)\) rounds, the claim follows.

#### c.6.2 Reduction from set disjointness

Before proving Theorems 21 and 23 by embedding an instance of a transformer model into an instance of each subgraph identification problem, we first introduce a partition of the vertices \(V^{N}\) of the Congest graph into those possessed by Alice and Bob for use in a two-party communication protocol. We call those two sets \(V^{N}_{a}\) and \(V^{N}_{b}\).

Note that the previous section made no assumptions about the organization of edges in the binary tree. We thus add an additional condition: that each binary tree \(B_{i}\) can be oriented to respect the left-to-right ordering \(v_{i,1},v_{1,i},\ldots,v_{i,N},v_{N,i}\). Let \(u_{i}\in V^{N}_{a}\) if and only if \(i\leq\frac{N}{2}\), and \(v_{i,j}\in V^{N}_{a}\) if and only if \(\min(i,j)\leq\frac{N}{2}\). We label are remaining nodes in \(B_{i}\) by labeling a parent node \(w_{p}\) as a function of its child nodes \(w_{\ell}\) and \(w_{r}\) using the following rules:

1. If \(w_{\ell},w_{r}\in V^{N}_{a}\), then let \(w_{p}\in V^{N}_{a}\).
2. If \(w_{\ell},w_{r}\in V^{N}_{b}\), then let \(w_{p}\in V^{N}_{b}\).
3. Otherwise, let \(w_{p}\in V^{N}_{a}\) if and only if root \(u_{i}\in V^{N}_{a}\).

This partition, which we visualize in Figure 5, bounds the number of bits Alice and Bob can exchange by simulating a protocol on Congest graph \(G^{N}\).

**Lemma 26**.: _Suppose Alice and Bob simulate an \(R\)-round \(p\)-bit protocol on Congest communication graph \(G^{N}\) where Alice has access to all vertices \(V_{a}^{N}\) and Bob \(V_{b}^{N}\). No other communication is permitted besides sharing bits as permitted by the Congest protocol between neighboring vertices. Then, Alice and Bob exchange at most \(O(pRN\log N)\) bits._

Proof.: It suffices to show that the partition \(V_{a}^{N},V_{b}^{N}\) induces a cut of size at most \(O(N\log N)\); this ensures that each can send no more than \(O(pN\log N)\) bits per round.

Per the rules defined above, an edge in \((w_{p},w_{\ell})\) and \((w_{p},w_{r})\) is cut if and only if they are described by case (c). Within each tree \(B_{i}\) under the orientation described above, an inductive argument shows that in every layer, all elements in \(V_{a}^{N}\) are to the left of all elements in \(V_{b}^{N}\). Thus, there exists at most one parent of that layer that belongs to case (c), and thus, no more than one cut edge per layer. Because each tree has \(O(\log N)\) layers and because there are \(N\) trees, the partition cuts at most \(O(N\log N)\) edges.

It remains to embed an instance of DISJ in \(V_{a}^{N},V_{b}^{N}\) for each problem such that its output corresponds identically with that of DISJ.

Proof of Theorem 21.: Assume for the sake of simplicity that \(N\) is divisible by 5. Let \(a,b\in\{0,1\}^{n}\) for \(n=\frac{N^{2}}{25}\) be an input to DISJ, and let Alice and Bob possess \(a\) and \(b\) respectively. We index those vectors as \(a=(a_{1,1},a_{1,2},\ldots,a_{N/5,N/5-1},a_{N/5,N/5})\) and \(b=(b_{1,1},\ldots,b_{N/5,N/5})\) for ease of analysis. We design input matrix \(X\in\{0,1\}^{N\times N}\) as follows:

* If \(i\in(0,\frac{N}{5}]\) and \(j\in(\frac{N}{5},\frac{2N}{5}]\), then \(x_{i,j}=x_{j,i}=a_{i,j-N/5}\).
* If \(i\in(\frac{N}{5},\frac{3N}{5}]\) and \(j\in(\frac{2N}{5},\frac{4N}{5}]\), then \(x_{i,j}=x_{j,i}=\delta_{i,j-N/5}\).
* If \(i\in(\frac{3N}{5},\frac{4N}{5}]\) and \(j\in(\frac{4N}{5},N]\), then \(x_{i,j}=x_{j,i}=b_{j-4N/5,i-3N/5}\).
* If \(i\in(\frac{4N}{5},N]\) and \(j\in(0,\frac{N}{5}]\), then \(x_{i,j}=x_{j,i}=\delta_{i,j+4N/5}\).
* Otherwise, \(x_{i,j}=0\).

Figure 5: The Congest graph \(G^{N}\) with vertices partitioned into sets \(V_{a}^{N}\) (violet) and \(V_{b}^{N}\) (orange) for \(N=6\). The six edges cut by the partition are colored red.

This ensures that \(X\) has a 5-cycle if and only there exist \(i,j\in(0,\frac{N}{5}]\) such that \(a_{i,j}b_{i,j}=1\)6. In addition, note that under the protocol in Lemma 25, Alice's and Bob's inputs \(a\) and \(b\) are known exclusively by nodes belonging to \(V_{a}^{N}\) and \(V_{b}^{N}\) respectively.

Footnote 6: We consider 5-cycles rather than 4-cycles because a spurious 4-cycle could exist among edges \(\{x_{i,j}:i\in(0,\frac{N}{5}],\;j\in(\frac{N}{5},\frac{2N}{5}]\}\).

Consider any transformer architecture \(f\in\mathcal{TG}_{N,m,1,p}^{D,H}\) that computes \(\mathrm{Cycle}5\). By Lemma 25, there exists a protocol on the Congest graph \(G^{N}\) that computes \(\mathrm{Cycle}5\) after \(O(HD(m+\log N))\) rounds of communication of \(p\)-bits each. If Alice and Bob simulate this protocol, and output 1 if and only if at least one of their outputs indicates the existence of a \(\mathrm{Cycle}5\), then they successfully decide DISJ. By Lemma 26, this communication algorithm solves DISJ after exchanging \(O(mpHDN\log^{2}N)\) bits of communication. However, Fact 5 implies that no communication algorithm can do so without exchanging \(\Omega(n)=\Omega(N^{2})\) bits, which concludes the proof.

Proof of Theorem 23.: The proof is identical to its predecessor, but uses a different embedding of an instance \(a,b\in\{0,1\}^{n}\) to DISJ. Let \(n=\frac{N^{2}}{16}\). Then:

* If \(i\in(0,\frac{N}{4}]\) and \(j\in(\frac{N}{2},\frac{3N}{4}]\), then \(x_{i,j}=a_{i,j-N/2}\).
* If \(i\in(\frac{N}{2},\frac{3N}{4}]\) and \(j\in(\frac{3N}{4},N]\), then \(x_{i,j}=b_{j-3N/4,i-N/2}\).
* If \(i\in(\frac{3N}{4},N]\) and \(j\in(0,\frac{N}{4}]\), then \(x_{i,j}=\delta_{i,j+3N/4}\).
* Otherwise, \(x_{i,j}=0\).

This construction ensures that a directed 3-cycle exists if and only if a corresponding pair of elements in \(a\) and \(b\) are both 1.

## Appendix D Experiment details

This section describes the experimental setup behind Figure 2, and provides further experiments suggesting an _implicit bias_ of transformers for \(q\mathrm{SA}\), in particular when compared with MLPs and RNNs.

Experimental setup.Experiments used synthetic data, generated for \(q\mathrm{SA}\) with \(n=1000\) training and testing examples, a sequence length \(N=20\), \(q=3\), with the individual inputs described in more detail as follows.

* The positional encoding of element \(i\) is a random vector sampled uniformly from the sphere in \(\mathbb{R}^{d_{0}}\) with \(d_{0}:=\lceil 1+2\ln(N)\rceil\), a quantity which agrees with the theory but was not tuned.
* A sequence element then consists of the data portion \(z\in\mathbb{R}^{d_{1}}\) where \(d_{1}=4\), also sampled from the unit sphere, then the positional encoding of this sequence element, and then \(q\) further positional encodings identifying elements to average to produce the output; this differs from (and is more tractable than) the presentation in Section 3, where the positional encoding is provided as an integer and the MLP layer input to our attention layers is expected to choose a sufficient positional encoding.

As such, the total dimension of a sequence element is \(d_{1}+(q+1)d_{0}=32\). The architectures are detailed as follows.

* The attention is identical to the description in the paper body, with the additional detail of the width and embedding dimension \(m\) being fixed to \(100\).
* Figure 6 also contains an MLP, which first flattens the input, then has a single hidden ReLU layer of width 256, before a final linear layer and an output reshaping to match the desired output sequence shapes.
* Figure 6 also contains an LSTM, which is a standard pytorch LSTM with \(2\) layers and a hidden state size \(800\), which is \(200\) times larger than the target output dimension \(4\).

Experiments fit the regression loss using Adam and a minibatch size of 32, with default precision, and take a few minutes to run on an NVIDIA TITAN XP, and would be much faster on standard modern hardware.

Further discussion of Figure 2 and Figure 7.In Figure 2 and Figure 7, we plot (post-softmax) alignment matrices after \(T\in\{0,1000,40000\}\) iterations of Adam. The alignment matrices in Figure 2 are taken from the training example whose loss is the median loss across all examples. Figure 7 is similar, but additionally shows the examples of minimal and maximal loss.

Further discussion of Figure 6.Figure 6 plots training and testing error curves for the same attention architecture as in Figure 2, but with further MLP and LSTM architectures as described above. but also an MLP trained on flattened (vectorized) error bars reflect \(5\) separate training runs from random initialization. A few variations of these architectures were attempted, however curves did not qualitatively change, and in particular, only the attention layer achieves good generalization across all attempts.

Figure 6: Test and train error curves of fitting various architectures to \(q\mathrm{SA}\), where the horizontal axis denotes thousands of training iterations, and the vertical axis denotes the regression objective; see Section D for further details.

Figure 7: Alignment plots as in Figure 2, but using examples with minimum, median, and maximum loss, whereas Figure 2 only uses the example with median loss.

Response:
## Review

### Summary

This paper investigates the representational power of attention layers in transformer networks, focusing on their advantages and limitations compared to other neural network architectures. It introduces three computational tasks: sparse averaging (qSA), pair matching, and triple matching, demonstrating that transformers excel in the qSA task with logarithmic scaling in input size, whereas their performance diminishes in the triple matching task, which requires polynomial scaling. The authors utilize communication complexity principles to derive these results, providing both theoretical insights and implications for future research in deep learning, particularly in natural language processing. The paper makes a compelling first step toward understanding transformer limitations and their potential for practical applications, despite some concerns about the relevance of the tasks to real-world scenarios.

### Strengths

- The paper offers original contributions to the understanding of attention layers in transformer networks, providing a mathematical analysis of their representation power.
- It rigorously compares transformers with fully connected networks and recurrent networks, establishing both strengths and limitations.
- The use of communication complexity techniques enhances the depth and reliability of the analysis.
- The mathematical formulations and proofs are of high quality, showcasing the authors' expertise.
- The paper addresses a notable gap in the literature regarding the theoretical understanding of transformers, with implications for model design in NLP.

### Weaknesses

- The fixed precision result in Theorem 2 is still dependent on growing precision with data size, leaving open the question of constant precision.
- The relevance of the computational tasks (qSA, Match2, Match3) to practical applications is not clearly established, which may limit its impact.
- Some sections, particularly the introduction and notation, could be clearer to enhance reader comprehension.
- The paper lacks empirical evaluation of the proposed tasks, which could validate theoretical findings and strengthen conclusions.
- The complexity of mathematical formulations may pose challenges for readers without a strong background in the subject.

### Questions

- How do the tasks presented connect to real-world applications in NLP or other domains?
- What are the specific bottlenecks in developing transformers that can approximate q-SA with fixed constant precision?
- Is the application of communication complexity in this context novel?
- Can the authors provide insights on how the theoretical results relate to practical performance in deep learning tasks?

### Soundness

**Score:** 3

**Description:** 3 = good: The paper is technically sound with solid theoretical foundations, though some concerns about precision and practical connections remain.

### Presentation

**Score:** 3

**Description:** 3 = good: The paper is generally well-written, but some sections could benefit from clearer explanations and more intuitive presentations of complex concepts.

### Contribution

**Score:** 3

**Description:** 3 = good: The contributions are significant in advancing theoretical understanding of transformers, although further empirical validation would enhance their relevance.

### Rating

**Score:** 6

**Description:** 6 = Weak Accept: The paper is technically solid with moderate-to-high impact potential, but it requires improvements in clarity and empirical evaluation.

### Paper Decision

**Decision:** Accept (poster)

**Reasons:** The paper presents a technically sound exploration of the representational limits of transformers, making original contributions to the field. While it has areas for improvement, especially concerning practical relevance and empirical validation, the theoretical insights are valuable and can influence future research directions. The overall assessment suggests a strong foundation that justifies acceptance.