SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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} \to \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 known 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^{\top}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},\dots ,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} := \exp(\kappa(x_{i,j}, x_{j,i}, y_i^\top QK^\top 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^{\top} 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_{a}^{N}$  and  $V_{b}^{N}$ .
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_{a}^{N}$  if and only if  $i \leq \frac{N}{2}$ , and  $v_{i,j} \in V_{a}^{N}$  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:
(a) If  $w_{\ell}, w_{r} \in V_{a}^{N}$ , then let  $w_{p} \in V_{a}^{N}$ .
(b) If  $w_{\ell}, w_{r} \in V_{b}^{N}$ , then let  $w_{p} \in V_{b}^{N}$ .
(c) Otherwise, let  $w_{p}\in V_{a}^{N}$  if and only if root  $u_{i}\in V_{a}^{N}$
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}$ .
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.
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 \left( \frac{N}{5}, \frac{3N}{5} \right]$  and  $j \in \left( \frac{2N}{5}, \frac{4N}{5} \right]$ , then  $x_{i,j} = x_{j,i} = \delta_{i,j-N/5}$ .
- If  $i \in \left(\frac{3N}{5}, \frac{4N}{5}\right]$  and  $j \in \left(\frac{4N}{5}, N\right]$ , then  $x_{i,j} = x_{j,i} = b_{j - 4N/5, i - 3N/5}$ .
- If  $i \in \left(\frac{4N}{5}, N\right]$  and  $j \in \left(0, \frac{N}{5}\right]$ , then  $x_{i,j} = x_{j,i} = \delta_{i,j+4N/5}$ .
- Otherwise,  $x_{i,j} = 0$ .
This ensures that  $X$  has a 5-cycle if and only if 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.
Consider any transformer architecture  $f \in \mathcal{T}G_{N,m,1,p}^{D,H}$  that computes Cycle5. By Lemma 25, there exists a protocol on the CONGEST graph  $G^{N}$  that computes Cycle5 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 Cycle5, 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 \left( \frac{N}{2}, \frac{3N}{4} \right]$  and  $j \in \left( \frac{3N}{4}, N \right]$ , then  $x_{i,j} = b_{j - 3N/4, i - N/2}$ .
- If  $i \in \left(\frac{3N}{4}, N\right]$  and  $j \in \left(0, \frac{N}{4}\right]$ , 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.
This section describes the experimental setup behind Figure 2, and provides further experiments suggesting an implicit bias of transformers for qSA, 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 \coloneqq \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.
<sup>6</sup>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}]\}$ .
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.
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.
(a) Min loss,  $T = 0$ .
(b) Min loss,  $T = 1000$ .
(c) Min loss,  $T = 40000$ .
(d) Median loss,  $T = 0$ .
(e) Median loss,  $T = 1000$ .
(f) Median loss,  $T = 40000$ .
(g) Max loss,  $T = 0$ .
(h) Max loss,  $T = 1000$ .
(i) Max loss,  $T = 40000$ .
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.