SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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).
B Supplementary results for Section 3
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' = \frac{m - d}{2}$ , let  $\phi(X)K = (u_1, \ldots, u_N) \in \mathbb{R}^{N \times m'}$  be embedded key vectors, where  $u_1, \ldots, u_N \in \{\pm 1/\sqrt{m'}\}^{m'}$  are the columns of a  $m' \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'$ . 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'}$  with  $\| w_y \|_2 \leq 2\sqrt{q}$  satisfying
$$
\langle u _ {i ^ {\prime}}, w _ {y} \rangle = 1 \quad \text {f o r a l l} i ^ {\prime} \in y,
$$
$$
| \langle u _ {i ^ {\prime}}, w _ {y} \rangle | \leq \frac {1}{2} \quad \text {f o r a l l} 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'}$  for every  $w_{y}$  with  $\| w_{y}\|_{2} \leq 2\sqrt{q}$  satisfying  $\| \widetilde{w_{y}} - w_{y}\|_{2} \leq \frac{\epsilon}{4\alpha}$ . As an immediate consequence,  $|\langle u_{i'}, \widetilde{w_{y}} \rangle - \langle u_{i'}, w_{y} \rangle| \leq \frac{\epsilon}{4\alpha}$  for all  $i'$  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 \to \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' \in y_i$ , then
$$
\begin{array}{l} \operatorname {s o f t m a x} (\phi (X) Q K ^ {\intercal} \phi (X) ^ {\intercal}) _ {i, i ^ {\prime}} = \frac {\exp (\alpha \langle u _ {i} , \widetilde {w _ {i ^ {\prime}}} \rangle)}{\sum_ {i ^ {\prime \prime} \in y _ {i}} \exp (\alpha \langle u _ {i} , \widetilde {w _ {i ^ {\prime \prime}}} \rangle) + \sum_ {i ^ {\prime \prime} \not \in y _ {i}} \exp (\alpha \langle u _ {i} , \widetilde {w _ {i ^ {\prime \prime}}} \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}}{q e ^ {\alpha} + N e ^ {\alpha / 2}} \cdot \exp \left(- \frac {\epsilon}{2}\right) \\ \geq \left(\frac {1}{q} - \frac {N e ^ {\alpha / 2}}{q e ^ {\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). \\ \end{array}
$$
An analogous argument shows that
$$
\operatorname {s o f t m a x} (\phi (X) Q K ^ {\intercal} \phi (X) ^ {\intercal}) _ {i, i ^ {\prime}} \leq \frac {1}{q} \left(1 + \frac {\epsilon}{2}\right).
$$
Likewise, if  $i' \notin y_i$ , then
$$
\operatorname {s o f t m a x} (\phi (X) Q K ^ {\intercal} \phi (X) ^ {\intercal}) _ {i, i ^ {\prime}} \leq \frac {\exp \left(\frac {\alpha}{2} + \frac {\epsilon}{4}\right)}{q \exp \left(\alpha - \frac {\epsilon}{4}\right)} \leq \exp \left(- \frac {\alpha}{2} + \frac {\epsilon}{2}\right) \leq \frac {\epsilon}{2 N}.
$$
We thus conclude that that we meet the desired degree of approximation for such  $m$ :
$$
\begin{array}{l} \left\| f (X) _ {i} - q \mathrm {S A} (X) _ {i} \right\| _ {2} = \left\| \sum_ {i ^ {\prime} \in y _ {i}} \left(\frac {1}{q} - \operatorname {s o f t m a x} (\phi (X) Q K ^ {\intercal} \phi (X) ^ {\intercal}) _ {i, i ^ {\prime}}\right) z _ {i ^ {\prime}} \right\| _ {2} \\ + \left\| \sum_ {i ^ {\prime} \notin y _ {i}} \left(\operatorname {s o f t m a x} (\phi (X) Q K ^ {\intercal} \phi (X) ^ {\intercal}) _ {i, i ^ {\prime}}\right) z _ {i ^ {\prime}} \right\| _ {2} \\ \leq q \cdot \frac {\epsilon}{2 q} + (N - q) \cdot \frac {\epsilon}{2 N} \leq \epsilon . \\ \end{array}
$$
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
$$
\| U v \| _ {2} ^ {2} \in [ (1 - \delta) \| v \| _ {2} ^ {2}, (1 + \delta) \| v \| _ {2} ^ {2} ] \quad \text {a n d} \quad | \langle U v, U v ^ {\prime} \rangle | \leq \delta \| v \| _ {2} \| v ^ {\prime} \| _ {2}
$$
for all vectors  $v, v' \in \mathbb{R}^N$  with  $|\operatorname{supp}(v)| \leq q$ ,  $|\operatorname{supp}(v')| \leq 2q$ , and  $\operatorname{supp}(v) \cap \operatorname{supp}(v') = \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  $\mathrm{supp}(v)\leq q$ , there exists  $w\in \mathbb{R}^m$  satisfying
$$
\begin{array}{l} \| w \| _ {2} \leq \sqrt {q} / (1 - 2 \delta), \\ \langle u _ {i}, w \rangle = 1 \quad \text {i f} v _ {i} = 1, \\ | \langle u _ {i}, w \rangle | \leq \delta / (1 - 2 \delta) \quad i f v _ {i} = 0. \\ \end{array}
$$
Theorem 3 (Infinite-precision). For fixed  $N$ ,  $m \geq \Omega(d' + q)$  and  $\epsilon > 0$ , there exists some  $f \in \mathcal{A}_{d,m,d'}'$  that  $\epsilon$ -approximates qSA.
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' \leq q$  vertices forms a  $(q' - 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'}$ , let  $\mathrm{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, \dots, t^{m'}) \in \mathbb{R}^{m'}$ . Then, for all distinct  $t_1, \dots, t_N \in \mathbb{R}$ , the cyclic polytope  $\mathrm{Conv}(\theta(t_1), \dots, \theta(t_N))$  is  $\frac{m'}{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'}$  such that  $\mathrm{Conv}(u_1, \ldots, u_N)$  is  $q$ -neighborly. Then, for any  $\epsilon > 0$ , there exists some  $f \in \mathcal{A}_{d,m,d'}'$  with fixed key vectors  $\phi(X)K = (u_1, \ldots, u_N)$  that  $\epsilon$ -approximates qSA.
Proof. The construction employs a similar look-up table  $\mathrm{MLP}\phi$  to the one used in the proof of Theorem 2. We let the key and value embeddings be
$$
\phi (X) K = \left(\left(u _ {1}, 1\right), \dots , \left(u _ {N}, 1\right)\right) \in \mathbb {R} ^ {N \times \left(m ^ {\prime} + 1\right)}, \text {a n d} \phi (X) V = \left(z _ {1}, \dots , z _ {N}\right) \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_y' \in \mathbb{R}^{m'}$  and  $b_y \in \mathbb{R}$  such that
$$
w _ {y} ^ {\prime \intercal} u _ {i} + b _ {y} \left\{ \begin{array}{l l} = 1 & \text {i f} i \in y, \\ <   1 & \text {o t h e r w i s e .} \end{array} \right.
$$
For  $\alpha > 0$ , let  $\phi(x_i)^\top Q = \alpha w_y = \alpha (w_y', 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]}\| f(X)_i - q\mathrm{SA}(X)_i\| _2\leq \epsilon$
Theorem 4. For any sufficiently large  $q$ , any  $N \geq 2q + 1$ , and any  $d' \geq 1$ , there exists a universal constant  $c$  such that if  $mp \leq cq$ , then no  $f \in \mathcal{T}_{d,m,d',p}^{1,1}$  exists that  $\frac{1}{2q}$ -approximates qSA.
Proof. We first embed every instance of 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}_{d,m,d',p}^{1,1}$  that  $\frac{1}{2q}$ -approximates  $q\mathrm{SA}$  and implies the existence of an  $O(mp)$ -bit communication protocol that computes DISJ. An application of Fact 5 concludes the proof.
Consider an instance of 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$  SA. For each  $j \in [2q]$ , let  $y_{2q+1} = \{2i + a_i - 1 : i \in [q]\}$ . Additionally, let
$$
z _ {j} = \left\{ \begin{array}{l l} e _ {1} & \text {i f j i s o d d a n d b _ {(j - 1) / 2} = 1}, \\ - e _ {1} & \text {o t h e r w i s e}. \end{array} \right.
$$
All other inputs are set arbitrarily. Then,
$$
\begin{array}{l} q \operatorname {S A} (X) _ {2 q + 1} = \frac {1}{q} | \{j \in [ 2 q ]: j \in y _ {2 q + 1}, j \text {i s o d d , a n d} a _ {(j - 1) / 2} = 1 \} | e _ {1} \\ - \frac {1}{q} | \left\{j \in [ 2 q ]: j \in y _ {2 q + 1} \text {a n d} (j \text {i s e v e n o r} a _ {(j - 1) / 2} = 0) \right\} | e _ {1} \\ = \frac {\left| \left\{i \in [ q ] : a _ {i} b _ {i} = 1 \right\} \right| - \left| \left\{i \in [ q ] : a _ {i} b _ {i} = 0 \right\} \right|}{q} e _ {1}. \\ \end{array}
$$
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'}$  such that
$$
f (X) _ {2 q + 1} = \psi \left(\frac {\sum_ {i = 1} ^ {N} \exp \left(Q (x _ {2 q + 1}) ^ {\intercal} K (x _ {i})\right) V (x _ {i})}{\sum_ {i = 1} ^ {N} \exp \left(Q (x _ {2 q + 1}) ^ {\intercal} 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}^{\top} e_{1} \geq -1 + \frac{1}{q}$ .
The protocol computes  $\mathrm{DISJ}(a,b)$  because  $f$  is a  $\frac{1}{2q}$ -approximation of qSA. Because any such protocol requires sharing  $\Omega(q)$  bits of information, we conclude that  $mp \leq cq$  for some  $c$ .
B.4 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)^\top Q = w(y_i,i)$ ,  $\phi (x_i)^\top K = u(i)$ , and  $\phi (x_i)^\top V = z_i$ . Then, for any  $q < N$  and  $m\leq q(1 - C\log_Nq)$  for some universal  $C$ , there do not exist  $w:\mathbb{R}^d\times [N]\to \mathbb{R}^m$  and  $u:[N]\to \mathbb{R}^m$  such that the resulting self-attention unit  $\frac{1}{2q}$ -approximates qSA.
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 {s o f t m a x} (w (X) U ^ {\intercal}) Z) _ {i} - \frac {1}{q} \sum_ {i ^ {\prime} \in y _ {i}} z _ {i ^ {\prime}} \right\| _ {2} \geq \frac {1}{2 q},
$$
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)^{\intercal} u_{i'} \geq \tau$  if and only if  $i' \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)^{\intercal} u_{i_2} > w(y_1, 1)^{\intercal} u_{i_1}$ . Given the value embeddings  $z_i = e_1, z_{i_2} = e_2$  and  $z_i = e_3$  for all  $i \notin \{i_1, i_2\}$ , we have
$$
\begin{array}{l} \left\|\left(\operatorname {s o f t m a x} (w (X) U ^ {\intercal}) Z\right) _ {1} - \frac {1}{q} \sum_ {i ^ {\prime} \in y _ {1}} z _ {i ^ {\prime}} \right\| _ {2} ^ {2} \geq \left(\operatorname {s o f t m a x} (w (X) U ^ {\intercal}) Z\right) _ {1, i _ {1}} - \frac {1}{q}\left. \right) ^ {2} + \left(\operatorname {s o f t m a x} (w (X) U ^ {\intercal}) Z\right) _ {1, i _ {2}}) ^ {2} \\ \geq \max  \left(\left(\operatorname {s o f t m a x} (w (X) U ^ {\intercal}) Z) _ {1, i _ {1}} - \frac {1}{q}\right) ^ {2}, \operatorname {s o f t m a x} (w (X) U ^ {\intercal}) Z) _ {1, i _ {1}} ^ {2}\right) \\ \geq \frac {1}{4 q ^ {2}}. \\ \end{array}
$$
Fact 17. If  $m' < q(1 - \log_N Cq)$ , then the columns of any  $U = (u_1, \ldots, u_N) \in \mathbb{R}^{N \times m'}$  can be partitioned into sets  $U_1$  and  $U_2$  with  $|U_1| = q$  that are not linearly separable. Hence,  $\operatorname{Conv}(u_1, \ldots, u_N)$  is not  $q$ -neighboringly.
Proof. By the Sauer-Shelah Lemma [Sauer, 1972, Shelah, 1972, Vapnik and Chervonenkis, 1968] and the fact that the VC dimension of  $m'$ -dimensional linear thresholds is  $m' + 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}}{(C q) ^ {q}} \leq \binom {N} {q},
$$
for a sufficiently large choice of  $C$  given universal constant  $C'$ . If the fact were to be false, then at least  $\binom{N}{q} \geq \left(\frac{N}{q}\right)^q$  such partitions must exist, which contradicts the above bound.
C Supplementary results for Section 4
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' = (x_1, \ldots, x_N, x')$  with  $x_{i,0} = i$  and  $x' = \vec{0}$  to be the input to the target attention model. We define input MLP  $\phi: \mathbb{R} \to \mathbb{R}^3$  and parameterizations  $Q, K, V \in \mathbb{R}^{3 \times 3}$  such that
$$
Q ^ {\intercal} \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 ^ {\intercal} \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^{\top}\phi (x_{i}) = \vec{1},Q^{\top}\phi (x^{\prime}) = \vec{0},K^{\top}\phi (x^{\prime}) = e_{3},$  and  $V^{\top}\phi (x^{\prime}) = \vec{0}$  . By elementary trigonometric identities, the following is true about the corresponding inner products:
$$
\left(Q ^ {\intercal} \phi (x _ {i})\right) ^ {\intercal} K ^ {\intercal} \phi (x _ {j}) = c \cos \left(\frac {2 \pi (x _ {i} + x _ {j})}{M}\right)
$$
$$
\left(Q ^ {\intercal} \phi (x _ {i})\right) ^ {\intercal} K ^ {\intercal} \phi (x ^ {\prime}) = c d.
$$
As a result,  $(Q^{\top}\phi (x_{i}))^{\top}K^{\top}\phi (x_{j}) = cd$  if and only if  $x_{i} + x_{j} = \mathbb{Q}\bmod M$ . Otherwise,  $(Q^{\top}\phi (x_{i}))^{\top}K^{\top}\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 (\mathrm {m o d} M) \} |
$$
represent the total number of matches the input belongs to. If we take  $c = M^2 \log(6N)$ , then
$$
(\operatorname {s o f t m a x} (\phi (X) Q K ^ {\intercal} \phi (X) ^ {\intercal})) _ {i, j} \in \left\{ \begin{array}{l l} [ 0, \frac {1}{6 N} ] & \text {i f} x _ {i} + x _ {j} \neq 0 (\text {m o d} M) \text {a n d} i, j \in [ N ]; \\ [ \frac {1}{\beta_ {i} + 1} \pm \frac {1}{6 N} ] & \text {i f} x _ {i} + x _ {j} = 0 (\text {m o d} M) \text {a n d} i, j \in [ N ]; \\ [ \frac {1}{\beta_ {i} + 1} \pm \frac {1}{6 N} ] & \text {i f} i \in [ N ], j = N + 1. \end{array} \right.
$$
We conclude that for any  $i\in [N]$
$$
(\operatorname {s o f t m a x} (\phi (X) Q K ^ {\intercal} \phi (X) ^ {\intercal}) V \phi (X)) _ {i} \left\{ \begin{array}{l l} \leq \frac {1}{6} \cdot \vec {1} & \text {i f} \not \exists j \text {s . t .} x _ {i} + x _ {j} = 0 (\text {m o d} M) \\ \geq \left(\frac {\beta_ {i}}{\beta_ {i} + 1} - \frac {1}{6}\right) \cdot \vec {1} & \text {i f} \exists j \text {s . t .} x _ {i} + x _ {j} = 0 (\text {m o d} M), \end{array} \right.
$$
where  $\leq$  is a partial ordering with  $v\leq v^{\prime}$  if  $v_{i}\leq v_{i}^{\prime}$  for all  $i$ . Since the latter case holds only when  $\beta_{i}\geq 1$ , the final step of the proof is design an output MLP  $\psi$  such that  $\psi (z) = 1$  if  $z\geq \frac{1}{3}$  and  $\psi (z) = 0$  if  $z\leq \frac{1}{6}$ , which can be crafted using two ReLU gates.
C.2 Proof of Theorem 7
Proof. The proof relies on a reduction to Fact 5 that embeds inputs to the set-disjointness problem of cardinality  $n = \frac{N - 1}{2}$  into a subset of instances passed to Match3. For the sake of simplicity, we assume in the construction that  $N$  is odd; if it were not, we could replace it with  $N - 1$  and set the final element such that it never belongs to a triple.
We consider the following family of inputs to Match3:
$$
x _ {i} \in \left\{ \begin{array}{l l} \{0 \} & \text {i f} i = 1, \\ \{1, i \} & \text {i f} i \in \{2, \dots , \frac {N + 1}{2} \}, \\ \{1, (M - i + \frac {N - 1}{2}) \} & \text {i f} i \in \{\frac {N + 3}{2}, \dots , N \}. \end{array} \right. \tag {3}
$$
Note that  $\mathrm{Match3}(X)_1 = 1$  if and only if there exists  $i\in \{2,\dots ,\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 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) = \operatorname{Match3}(X)$  for all  $X \in [M]^N$  for some  $f \in \mathcal{T}_{1,m,1,p}^{1,H}$ . We show that  $f$  simulates an  $O(mpH)$ -bit communication protocol for testing 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}) ^ {\intercal} K _ {h} (x _ {i})) V _ {h} (x _ {i})}{\sum_ {i = 1} ^ {N} \exp (Q _ {h} (x _ {1}) ^ {\intercal} 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 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 + 3}{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 \left(Q _ {h} \left(\phi \left(x _ {1}\right)\right) ^ {\top} K _ {h} \left(\phi \left(x _ {i}\right)\right)\right)\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 \left(L _ {h, a}\right) + \sum_ {i = \frac {N + 3}{2}} ^ {N} \exp \left(Q _ {h} \left(\phi \left(x _ {1}\right)\right) ^ {\top} K _ {h} \left(\phi \left(x _ {i}\right)\right)\right)\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 \left(Q _ {h} \left(\phi \left(x _ {1}\right)\right) ^ {\top} K _ {h} \left(\phi \left(x _ {i}\right)\right)\right) V _ {h} \left(\phi \left(x _ {i}\right)\right)}{\exp \left(L _ {h}\right)} \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