SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Further, let matrix  $U' = \sum_{j=1}^{Dk} o_j b_j^\top$  where  $o_j$  is the  $Dk$  dimension one-hot vector with the  $j$ -th entries being 1. Then when  $t$  is an even number and  $d \geq 1$ , based on Equation (29) and the definition of  $M$ ,
$$
U ^ {\prime} (\mathrm {L N} (\tilde {a} _ {2} (X; \theta^ {(2)}) _ {:, n + 1}) + X _ {:, n + 1}) _ {t ^ {\prime} D + d ^ {\prime}} \left\{ \begin{array}{l l} = 0, & \tau_ {2 t ^ {\prime} - 1, d ^ {\prime}} \text {i s n o t a n u m m a t c h e d o p e n b r a c k e t s i n} p. \\ > M, & \tau_ {2 t ^ {\prime} - 1, d ^ {\prime}} \text {i s a n u m m a t c h e d o p e n b r a c k e t s i n} p. \end{array} \right.
$$
15That is  $\mathbb{P}(q)_t = \mathbb{P}$  (The next token of  $q$  has type  $t$ )
Then based on Lemma 18, there exists 2-layer MLP with width  $kD$  that operates on  $\left(U^{\prime}(\mathrm{LN}(\tilde{a}_{2}(X;\theta^{(2)})_{:,n + 1}) + X_{:,n + 1})_{t^{\prime}D + d^{\prime}}\right)_{t^{\prime}\in [k]}$  for a fixed  $d^{\prime}$  and output the nonzero index in it, if such index exists. Hence, we can choose the weight of the first and second layer of  $\mathfrak{g}^{(2)}$ , such that the output of the second layer is  $(t,d)\oplus x$ , where  $2x_{d^{\prime}} - 1$  is the type of the unmatched open brackets with grammar depth  $d^{\prime}$  if  $t$  is an even number,  $d\geq d^{\prime}\geq 1$ .
Now based on Lemma 17, we can choose the third and fourth layer of  $\mathbf{g}^{(2)}$  to perform indexing and let the output of the fourth layer be  $(t,d,y)$ , where  $y = x_{d}$  when  $d\geq 1$ . 16 Notice that this triplet contains all the necessary information to infer  $\mathbb{P}(q)$  because it uniquely determines the type of last unmatched open bracket,
1. If  $t$  is odd (i.e. the last bracket is open), and then the type of last unmatched open bracket is  $t$ .
2. If  $t$  is even and  $d = 0$ , then all the brackets is matched.
3. If  $t$  is even and  $d \geq 1$ , then the type of last unmatched bracket is  $y$ .
One may finally construct a 2-layer MLP  $f$  that maps  $(t, d, y)$  to the corresponding probability vector. As the input of  $g$  has bounded norm,
$$
\| \mathrm {L N} (\tilde {a} _ {2} (X; \theta^ {(2)}) _ {:, n + 1}) + X _ {:, n + 1} \| _ {2} \leq 1 + \max _ {t, d} \| e (\tau_ {t, d}) \|,
$$
the output of the constructed 4 layers also has a bounded norm. Hence, we can assume there exists constant  $M' > 1$ , such that  $y \leq M'$ . Now we will discuss by the value of  $t$
1.  $t$  is odd, then one can neglect the third dimension and the correct probability is determined by  $d$  and can be represented by a width-2D network based on Lemma 15.
2.  $t$  is even. When  $d = 0$ , one can construct a width-1 network mapping any  $y$  to the correct probability distribution as it is unique. When  $d \geq 1$ , one can construct a width-2K network mapping  $x_{d} \in [K]$  to the correct probability distribution based on Lemma 15. Then by Lemma 19, one can construct a width-4KD network that maps  $(d,y)$  to the corresponding probability distribution.
Putting together and using Lemma 19 again, one can construct a width- $8K^2D$  network that maps  $(t,d,y)$  to the correct next token probability prediction. The proof is then completed.
C.2 Implication of our results to larger models
Recall that the main conclusion of our paper is that interpretability based on a single Transformer component (e.g. an attention pattern or an MLP block) can be unreliable, since the set of optimal solutions can give rise to a large set of attention patterns and pruned MLP weights. Section 3 has demonstrated this with simple two-layer Transformers. The simplicity of this architecture choice is intentional, since our theory on two-layer Transformers directly implies similar conclusions for larger models, as we discuss in this section.
Intuitively, when moving to more complex architectures, the set of solutions can only grow and complicate interpretability further, hence our main conclusion still stands. For example, even though Theorem 1 and Theorem 3 are stated for 2-layer Transformers only, the constructed solutions can be trivially extended to multiple layers by e.g. letting the higher layers perform the identity function, or removing Assumption 1 and allowing the model to flexibly use or ignore positional information. More precisely:
- For Transformers with greater width, our Theorem 1 applies directly, since the construction does not depend on the width.
- For Transformers with greater depth, it suffices to show that additional layers can perform the identity function. To this end, one can utilize the residue link in the Transformer layer and choose the value matrix to be zero and the FFN (with or without residue connection) to be identity. This
<sup>16</sup>When  $d = 0$ ,  $y$  does not matter since there is no unmatched open brackets.
construction is implicitly assuming LayerNorm will map zero vector to zero vector, which is true for the common PyTorch implementation and for our paper. Also, it is worth noting that this holds for both the architecture we considered in the paper and the standard GPT-2 architecture.
Corollary 1. There exists a 2-layer Transformer with uniform attention and no position embedding (but with causal mask and a starting token  $^{17}$ ) that generates the Dyck language of arbitrary length.
Proof. We will first construct a uniform attention first layer that can generate the embedding in Equation (15). Suppose  $Z$  is the one-hot embeddings of a prefix  $p$  of length  $n$ , where each token of type  $t$  for  $t \in [2k]$  is encoded as  $o_{t}$  and the starting token is encoded as  $o_{2k+1}$ . Then it holds that
$$
\left. \left[ Z \sigma \left(\mathcal {C} \cdot \left(W _ {K} ^ {(1)} Z\right) ^ {\top} \left(W _ {Q} ^ {(1)} Z\right)\right)\right) \right] _ {:, n + 1} = \sum_ {i = 1} ^ {2 k} \# \{\text {t o k e n o f t y p e} t \text {i n} p \} \boldsymbol {o} _ {t} + \boldsymbol {o} _ {2 k + 1}. \tag {30}
$$
Then we can choose  $W_V^{(1)}$  such that for  $x \in \mathbb{R}^{2k+1}$ ,
$$
\begin{array}{l} \left(W _ {V} ^ {(1)} x\right) _ {1} = \sum_ {i = 1} ^ {k} x _ {2 i - 1} - x _ {2 i}, \\ \left(W _ {V} ^ {(1)} x\right) _ {2} = x _ {2 k + 1}, \\ \left(W _ {V} ^ {(1)} x\right) _ {i} = 0, \forall i \geq 3. \\ \end{array}
$$
Hence it holds That
$$
\left. \left[ W _ {V} ^ {(1)} Z \sigma \left(\mathcal {C} \cdot \left(W _ {K} ^ {(1)} Z\right) ^ {\top} \left(W _ {Q} ^ {(1)} Z\right)\right)\right) \right] _ {:, n + 1} = \# \left\{\text {d e p t h o f} p _ {n} \right\} \boldsymbol {o} _ {1} + \boldsymbol {o} _ {2}.
$$
It is then easy to check  $\mathrm{LN}\left(\left[W_V^{(1)}Z\sigma \left(\mathcal{C}\cdot (W_K^{(1)}Z)^\top (W_Q^{(1)}Z)\right)\right)\right]_{:,n + 1}) + Z_{:,n + 1}$  is uniquely determined by the type and depth of  $p_n$  without repetition. Then by Lemma 15, there exists a 2-layer ReLU MLP with width  $O(k^2 D^2)$  that can map  $\mathrm{LN}\left(\left[W_V^{(1)}Z\sigma \left(\mathcal{C}\cdot (W_K^{(1)}Z)^\top (W_Q^{(1)}Z)\right)\right)\right]_{:,n + 1}) + Z_{:,n + 1}$  to the embedding in Equation (15). It is then easy to see that the condition in Theorem 1 is satisfied as  $W_{K}^{(2)} = W_{Q}^{(2)} = 0$ . Hence the second layer can be constructed to let the Transformer to output the correct next token probability.
C.4 Approximate Balance Condition For Finite Length Training Data
Theorem 1 assumes the model reaches the optimal loss for Dyck prefixes of any length. However, in practice, due to finite samples and various sources of randomness, training often does not end exactly at a population optima. In this case, the condition in Theorem 1 is not precisely met. However, even for models that approximately meet those conditions, we will prove that when the second-layer projection function  $\mathrm{g}^{(2)}$  is Lipschitz, a similar condition as in Equation (14) is still necessary.
We will show this by bounding the amount of deviations from the perfect balance. The idea is that for two long prefixes that differ in only the last open bracket, correct next token prediction requires the Transformer outputs on these prefixes to be sufficiently different, hence the part irrelevant to the prediction (i.e. matched brackets) should not have a large contribution.
To formalize this intuition, let's define two quantities: 1)  $S_{d,d',i,j}$  which measures the effect from one matching pair, and 2)  $P_{d,j}$  which measures the effect on the last position from all tokens in a prefix.
Let  $u$  be defined as in Equation (12).  $S_{d,d^{\prime},i,j}$  is defined as
$$
S _ {d, d ^ {\prime}, i, j} \left[ \theta^ {(2)} \right] = u \left(\tau_ {2 j, d}, \tau_ {2 i, d ^ {\prime} - 1}\right) + u \left(\tau_ {2 j, d}, \tau_ {2 i - 1, d ^ {\prime}}\right), \tag {31}
$$
which measures how much a matching pair of brackets  $(\tau_{2i,d' - 1},\tau_{2i - 1,d'})$  changes the input to the LayerNorm upon seeing the last token  $\tau_{2j,d}$ . Note that under the perfect balance condition,  $S_{d,d',i,j}[\theta^{(2)}] = 0$ .
The second quantity  $P_{d,j}[\theta^{(2)}]$  is defined via an intermediate quantity  $Q(2j,d,\tilde{t})$ : for any  $i\in [k],d\in [D]$  and a length-  $(d - 1)$  prefix  $\tilde{t}\in [2k]^{d - 1}$ ,  $Q(i,d,\tilde{t})$  is defined as
$$
\begin{array}{l} Q (i, d, \tilde {\boldsymbol {t}}) := u \left(\tau_ {2 i, d - 1}, t _ {\mathcal {S}}\right) + \sum_ {1 \leq d ^ {\prime} <   d} u \left(\tau_ {2 i, d - 1}, \tau_ {\tilde {\boldsymbol {t}} _ {d ^ {\prime}}, d ^ {\prime}}\right) \tag {32} \\ + u \left(\tau_ {2 i, d - 1}, \tau_ {2 i - 1, d}\right) + u \left(\tau_ {2 i, d - 1}, \tau_ {2 i, d - 1}\right), \\ \end{array}
$$
where  $\tilde{t}_{d'}$  denotes the  $d_{th}'$  entry of  $\tilde{t}$ . Intuitively,  $Q(i,d,\tilde{t})$  denotes the unnormalized second-layer attention output at the last position, given the input sequence  $\tilde{t} \oplus \tau_{2i-1,d}\tau_{2i,d-1}$ ,<sup>18</sup>
For results in this subsection, it suffices to consider prefixes consisting only of open brackets. Let  $\pmb{t} \coloneqq \arg \min_{\tilde{\pmb{t}} \in \{2i - 1\}_{i \in [k]}^{d - 1}} \| Q(2j,d,\tilde{\pmb{t}}) \|_2$ , and let  $\pmb{t}'$  denote the prefix that minimizes  $\|Q(2j,d,\tilde{\pmb{t}})\|_2$  subject to the constraint that  $\pmb{t}'$  differs from  $\pmb{t}$  at the last (i.e.  $(d - 1)_{th}$ ) position, i.e.
$$
\boldsymbol{t}^{\prime} = \arg \min_{\substack{\tilde{\boldsymbol{t}}^{\prime}\in \{2i - 1\}_{i\in [k]}^{d - 1},\boldsymbol{t}^{\prime}_{d - 1}\neq \boldsymbol{t}_{d - 1}}}Q(2j,d,\tilde{\boldsymbol{t}}^{\prime}).
$$
Such choices of  $t$ ,  $t'$  guarantees that the two prefixes differ at the last open bracket and hence must have different next-word distributions. Finally, define
$$
P _ {d, j} \left[ \bar {\theta} ^ {(2)} \right] = \| Q (2 j, d, \boldsymbol {t} ^ {\prime}) \|. \tag {33}
$$
In the following theorem,  $P_{d,j}$  will be used as a reference to upper bound  $S_{d,d',i,j}[\theta^{(2)}]$ , meaning that the model should not be sensitive to the insertion of a matching pair of brackets.
Theorem 3 (Approximate Balance). Consider a 2-layer Transformer  $\mathcal{T}$  (Equation (10)) with a minimal first layer (Assumption 1) and a  $\gamma$ -Lipschitz  $\mathrm{g}^{(2)}$  for  $\gamma > 0$ , trained on sequences of length  $N$  with the mean squared loss (Equation (5)).
Suppose the loss is approximately optimal, namely, the set of second-layer weights  $\bar{\theta}_N^{(2)}$  satisfies  $\mathcal{L}(\mathcal{T}[\bar{\theta}_N^{(2)}],\mathcal{D}_{q,k,D,N})\leq q^{-N}\epsilon$ , for any positive integer  $N > 8D$  and sufficiently small  $\epsilon >0$ . Then, there exists a constant  $C_{\gamma ,\epsilon ,D}$ , such that  $\forall 0\leq d^{\prime}\leq D,1\leq d\leq D,i,j\in [k]$ , it holds that
$$
\left\| S _ {d, d ^ {\prime}, i, j} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \right\| \leq \frac {C _ {\gamma , \epsilon , D}}{N} P _ {d, j} \left[ \bar {\theta} _ {N} ^ {(2)} \right]. \tag {34}
$$
Intuitively, Theorem 3 states that when the loss  $\mathcal{L}(\theta)$  is sufficiently small,  $S_{d,d^{\prime},i,j}[\theta^{(2)}]$  must be small relative to  $P_{d,j}[\bar{\theta}_N^{(2)}]$ . Inequality 34 can be interpreted as a relaxation of Equation (14), which is equivalent to  $S_{d,d^{\prime},i,j}[\theta^{(2)}] = 0$ . The proof of Theorem 3 shares a similar intuition as Theorem 1 and is given in Appendix C.4.1.
A direct corollary of Theorem 3 additionally consider weight decay, in which case approximate balance condition still holds, as the regularization strength goes to 0:
Corollary 3. Consider the setting where a Transformer with a fixed minimal first layer is trained to minimize  $\mathcal{L}_{\lambda}^{reg} = \mathcal{L}_{\theta}(x) + \lambda \frac{\|\theta\|_2^2}{2}$ , which is the squared loss with  $\lambda$  weight decay. Suppose  $\mathrm{g}^{(1)}$  of the Transformer is a 2-layer fully connected network and  $\mathrm{g}^{(2)}$  of the Transformer is a 6-layer fully connected network. Then, there exists constant  $C > 0$ , such that if a set of parameters  $\theta_{\lambda,N}$  minimizes  $\mathcal{L}_{\lambda}^{reg}$ , then it holds  $\forall 0 \leq d' \leq D, 1 \leq d \leq D, i, j \in [k]$  that,
$$
\forall N, \exists \lambda_ {N}, s u c h t h a t \forall \lambda \in [ 0, \lambda_ {N} ], S _ {d, d ^ {\prime}, i, j} \left[ \theta_ {\lambda , N} ^ {(2)} \right] \leq \frac {C}{N} P _ {d, i} \left[ \theta_ {\lambda , N} ^ {(2)} \right].
$$
<sup>18</sup>We use  $s \oplus t$  to denote the concatenation of two strings  $s, t$ , same as in Equation (15)-(16), and use  $\tau_i \tau_j$  to denote the concatenation of two tokens  $\tau_i, \tau_j$ .
C.4.1 Proof of Theorem 3
Proof. The key idea is similar to the proof of necessity in Theorem 1. That is, we will construct two input sequences with different next-word distributions, and show that the approximate balance condition must hold so that inserting (a bounded number of) pairs of matching brackets does not collapse the two predicted distributions given by the Transformer.
Constructing the input sequences.
Let  $\pmb{t} \coloneqq \arg \min_{\tilde{\pmb{t}} \in [k]^{d-1}} \| Q(2j, d, \tilde{\pmb{t}}) \|_2$ , and let  $\pmb{t}'$  denote the prefix that minimizes  $\| Q(2j, d, \tilde{\pmb{t}}) \|_2$  subject to the constraint that  $\pmb{t}'$  must differ from  $\pmb{t}$  in the last (i.e.  $(d - 1)_{th}$ ) position, i.e.
$$
\boldsymbol{t}^{\prime} = \arg \min_{\tilde{\boldsymbol{t}}^{\prime}\in [k]^{d - 1},\boldsymbol{t}_{d - 1}^{\prime}\neq \boldsymbol{t}_{d - 1}}Q(2j,d,\tilde{\boldsymbol{t}}^{\prime}).
$$
The motivation for such choices of  $\pmb{t},\pmb{t}^{\prime}$  is that since they differ at least by the last position which is an open bracket, they must lead to different next-word distributions. Note also that  $P_{d,j}[\bar{\theta}^{(2)}] = \| Q(2j,d,\pmb{t}^{\prime})\|$ .
With the above definition of  $t, t'$ , consider two valid Dyck prefixes  $p_1$  and  $p_2$  with length no longer than  $N$ , defined as follows: for any  $d, d' \in [D]$ ,  $i, j \in [k]$ , consider a common prefix  $p = \underbrace{\tau_{2i-1} \ldots \tau_{2i-1}}_{d' \text{ open brackets}} \underbrace{\tau_{2i-1} \tau_{2i} \ldots \tau_{2i-1} \tau_{2i}}_{\left( \left\lfloor \frac{N}{2} \right\rfloor - d' - d-1 \right) \text{ pairs}} \underbrace{\tau_{2i} \ldots \tau_{2i}}_{d' \text{ closed brackets}}$ , where  $\tau_i$  denotes a token with type  $i$  whose
depth is implicit from the context. Set  $p_1, p_2$  as
$$
p _ {1} = p \oplus \boldsymbol {t} \oplus \tau_ {2 j - 1} \tau_ {2 j},
$$
$$
p _ {2} = p \oplus \boldsymbol {t} ^ {\prime} \oplus \tau_ {2 j - 1} \tau_ {2 j}.
$$
That is,  $p_1, p_2$  differ in the last unmatched open bracket. In the following, we will show that the approximate balance condition must hold for the predictions on  $p_1, p_2$  to be sufficiently different.
Bounding the difference in Transformer outputs. For a Transformer  $\mathcal{T}$  with second layer parameters  $\bar{\theta}_N^{(2)}$ , its outputs on  $p_1, p_2$  satisfy
$$
\left\| \mathcal {T} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \left(p _ {1}\right) - \mathcal {T} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \left(p _ {2}\right) \right\| _ {2} \tag {35}
$$
$$
\geq \| p _ {1} - p _ {2} \| _ {2} - \left\| \mathcal {T} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \left(p _ {1}\right) - p _ {1} \right\| _ {2} - \left\| \mathcal {T} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \left(p _ {2}\right) - p _ {2} \right\| _ {2} \tag {36}
$$
$$
\geq \frac {1}{\sqrt {2 k}} \| p _ {1} - p _ {2} \| _ {1} - \left(\| \mathcal {T} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \left(p _ {1}\right) - p _ {1} \| _ {1} + \| \mathcal {T} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \left(p _ {2}\right) - p _ {2} \| _ {1}\right) \tag {37}
$$
$$
\geq \frac {1}{\sqrt {2 k}} \mathrm {T V} \left(p _ {1}, p _ {2}\right) - o _ {\epsilon} (1) \quad (\text {s i n c e} \mathcal {L} \left(\mathcal {T} \left[ \bar {\theta} _ {N} ^ {(2)} \right], \mathcal {D} _ {q, k, D, N}\right) \leq q ^ {- N} \epsilon) \tag {38}
$$
$$
= \Omega (1), \tag {39}
$$
where  $\mathrm{TV}(p_1,p_2)$  denotes the TV distance in the next-word distributions from  $p_1$  and  $p_2$ , and  $o_{\epsilon}(1)$  means the term will go to zero as  $\epsilon$  goes to zero. The TV distance is lower bounded by the construction of  $p_1,p_2$ , where  $\pmb {t},\pmb{t}^{\prime}$  differ at the last open bracket. The error  $\epsilon$  is upper bounded because of the assumption on  $\bar{\theta}_N^{(2)}$ , i.e.  $\mathcal{L}(\mathcal{T}[\bar{\theta}_N^{(2)}],\mathcal{D}_{q,k,D,N})\leq q^{-N}\epsilon$  with sufficiently small  $\epsilon$ .
Define by  $A_{p}$  the contribution of  $p$  to the attention output (before LayerNorm) of the last position of  $p_1, p_2$ :
$$
\begin{array}{l} A _ {p} = \sum_ {1 \leq d ^ {\prime \prime} <   d ^ {\prime}} \left(u \left(\tau_ {2 j, d - 1}, \tau_ {2 i, d ^ {\prime \prime} - 1}\right) + u \left(\tau_ {2 j, d - 1}, \tau_ {2 i - 1, d ^ {\prime \prime}}\right)\right) \\ + \left\lfloor \frac {N - 2 d ^ {\prime} - 2 d}{2} \right\rfloor \left(u \left(\tau_ {2 j, d - 1}, \tau_ {2 i, d ^ {\prime}}\right) + u \left(\tau_ {2 j, d - 1}, \tau_ {2 i - 1, d ^ {\prime} + 1}\right)\right). \tag {40} \\ \end{array}
$$
The attention outputs (before LayerNorm) of  $p_1, p_2$ , denoted by  $A(p_1)$  and  $A(p_2)$ , satisfy that
$$
\begin{array}{l} \mathcal {P} _ {\perp} A \left(p _ {1}\right) = \mathcal {P} _ {\perp} \left(A _ {p} + Q (2 j, d, t)\right), \\ \mathcal {P} _ {\perp} A \left(p _ {2}\right) = \mathcal {P} _ {\perp} \left(A _ {p} + Q \left(2 j, d, \boldsymbol {t} ^ {\prime}\right)\right). \tag {41} \\ \end{array}
$$
Note that for any prefix  $p'$ ,
$$
\begin{array}{l} \mathcal {T} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \left(p ^ {\prime}\right) = \mathrm {g} ^ {(2)} \left(\mathrm {L N} _ {C _ {L N}} \left(\mathcal {P} _ {\perp} A \left(p ^ {\prime}\right)\right)\right) + e \left(\tau_ {2 i, d ^ {\prime}}\right) (42) \\ = \mathrm {g} ^ {(2)} \left(\frac {\mathcal {P} _ {\perp} A \left(p ^ {\prime}\right)}{\| \mathcal {P} _ {\perp} A \left(p ^ {\prime}\right) \|}\right) + e \left(\tau_ {2 i, d ^ {\prime}}\right), (43) \\ \end{array}
$$
where  $\mathrm{g}^{(2)}$  is  $\gamma$ -Lipschitz. Hence the Lipschitz constant with respect to and we have
$$
\left\| \frac {\mathcal {P} _ {\perp} A \left(p _ {1}\right)}{\| \mathcal {P} _ {\perp} A \left(p _ {1}\right) \| _ {2}} - \frac {\mathcal {P} _ {\perp} A \left(p _ {2}\right)}{\| \mathcal {P} _ {\perp} A \left(p _ {2}\right) \| _ {2}} \right\| _ {2} \geq \frac {\operatorname {T V} \left(p _ {1} , p _ {2}\right) - o _ {\epsilon} (1)}{\sqrt {2 k} \gamma} = \Omega_ {\frac {1}{\gamma}, \epsilon} (1). \tag {44}
$$
We show that  $A_{p}$  should not be too much larger in norm than  $Q(2j, d, t)$  or  $Q(2j, d, t')$ . First, let's state a helper lemma about the contrapositive:
Lemma 1. For any  $\epsilon >0$ , there exists a constant  $R_{\epsilon}$ , such that for any  $a,b\in \mathbb{R}^d$  and any  $r\in \mathbb{R}^d$  such that  $\| r\| _2\geq R_\epsilon \cdot \max \{\| a\| _2,\| b\| _2\}$ , it holds that
$$
\left| \left| \frac {a + r}{\| a + r \| _ {2}} - \frac {b + r}{\| b + r \| _ {2}} \right| \right| _ {2} \leq \epsilon .
$$
Proof. Denote  $r_0 \coloneqq \max \{\|a\|_2, \|b\|_2\}$ . Then  $R_\epsilon \coloneqq \frac{4r_0}{\epsilon} + 1$  suffices:
$$
\begin{array}{l} \left| \left| \frac {r + a}{\| r + a \| _ {2}} - \frac {r + b}{\| r + b \| _ {2}} \right| \right| \leq \| r \| \cdot \left| \frac {1}{\| r + a \|} - \frac {1}{\| r + b \|} \right| + \frac {\| a \|}{\| r + a \|} + \frac {\| b \|}{\| r + b \|} \\ \leq \| r \| \cdot \left(\frac {1}{\| r \| - r _ {0}} - \frac {1}{\| r \| + r _ {0}}\right) + \frac {2 r _ {0}}{\| r \| - r _ {0}} \\ = \frac {2 r _ {0}}{\| r \| - r _ {0}} \cdot \left(\frac {\| r \|}{\| r \| + r _ {0}} + 1\right) \leq \frac {4 r _ {0}}{\| r \| - r _ {0}} \leq \frac {4 r _ {0}}{R _ {\epsilon} - r _ {0}} \leq \epsilon . \\ \end{array}
$$
Lemma 1 implies that if  $A_{p}$  is too large, then the output on  $p_1, p_2$  (Equation (44)) won't be sufficiently different. Let  $P_{d,j}[\bar{\theta}_N^{(2)}]$  be defined as in Equation (33) and let  $R_{\epsilon}$  be the constant in Lemma 1, we need to bound  $\| \mathcal{P}_{\perp}A_{p}\|$  by
$$
\left\| \mathcal {P} _ {\perp} A _ {p} \right\| _ {2} \leq R _ {\epsilon} \left\| P _ {d, j} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \right\| _ {2}. \tag {45}
$$
As Equation (45) holds for  $p$  with any  $d, d'$ , if one choose  $d' = 1$ , this shows
$$
\left\| u \left(\tau_ {2 j, d - 1}, \tau_ {2 i, 1}\right) + u \left(\tau_ {2 j, d - 1}, \tau_ {2 i - 1, 2}\right) \right\| _ {2} \leq \frac {4 R _ {\epsilon} \left\| P _ {d , j} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \right\| _ {2}}{N}. \tag {46}
$$
Further, it holds that for any  $1 < d' \leq d - 1$ ,
$$
\begin{array}{l} \| \sum_ {1 \leq d ^ {\prime \prime} <   d ^ {\prime}} \left(u \left(\tau_ {2 j, d - 1}, \tau_ {2 i, d ^ {\prime \prime} - 1}\right) + u \left(\tau_ {2 j, d - 1}, \tau_ {2 i - 1, d ^ {\prime \prime}}\right)\right) \\ + \left\lfloor \frac {N - 2 d ^ {\prime} - 2 d}{2} \right\rfloor \left(u (\tau_ {2 j, d - 1}, \tau_ {2 i, d ^ {\prime}}) + u (\tau_ {2 j, d - 1}, \tau_ {2 i - 1, d ^ {\prime} + 1})\right) \| _ {2} \leq R _ {\epsilon} \| P _ {d, j} [ \bar {\theta} _ {N} ^ {(2)} ] \| _ {2}. \\ \end{array}
$$
$$
\begin{array}{l} \| \sum_ {1 \leq d ^ {\prime \prime} <   d ^ {\prime} + 1} \left(u \left(\tau_ {2 j, d - 1}, \tau_ {2 i, d ^ {\prime \prime} - 1}\right) + u \left(\tau_ {2 j, d - 1}, \tau_ {2 i - 1, d ^ {\prime \prime}}\right)\right) \\ + \left\lfloor \frac {N - 2 d ^ {\prime} - 2 d - 2}{2} \right\rfloor \left(u (\tau_ {2 j, d - 1}, \tau_ {2 i, d ^ {\prime} + 1}) + u (\tau_ {2 j, d - 1}, \tau_ {2 i - 1, d ^ {\prime} + 2})\right) \| _ {2} \leq R _ {\epsilon} \| P _ {d, j} [ \bar {\theta} _ {N} ^ {(2)} ] \| _ {2}. \\ \end{array}
$$
The triangle inequality then yields,
$$
\begin{array}{l} \left\lfloor \frac {N - 2 d ^ {\prime} - 2 d - 2}{2} \right\rfloor \| \left(u \left(\tau_ {2 j, d - 1}, \tau_ {2 i, d ^ {\prime} + 1}\right) + u \left(\tau_ {2 j, d - 1}, \tau_ {2 i - 1, d ^ {\prime} + 2}\right)\right) \\ - \left(u \left(\tau_ {2 j, d - 1}, \tau_ {2 i, d ^ {\prime}}\right) + u \left(\tau_ {2 j, d - 1}, \tau_ {2 i - 1, d ^ {\prime} + 1}\right)\right) \| _ {2} \leq 2 R _ {\epsilon} \| P _ {d, j} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \| _ {2}. \\ \end{array}
$$
Because  $N \geq 8D$ , we have that  $\left\lfloor \frac{N - 2d' - 2d - 2}{2} \right\rfloor \geq \frac{N}{8}$ , hence it holds that
$$
\begin{array}{l} \left\| \left(u \left(\tau_ {2 j, d - 1}, \tau_ {2 i, d ^ {\prime} + 1}\right) + u \left(\tau_ {2 j, d - 1}, \tau_ {2 i - 1, d ^ {\prime} + 2}\right)\right) \right. \\ - \left(u \left(\tau_ {2 j, d - 1}, \tau_ {2 i, d ^ {\prime}}\right) + u \left(\tau_ {2 j, d - 1}, \tau_ {2 i - 1, d ^ {\prime} + 1}\right)\right) \| _ {2} \leq \frac {1 6 R _ {\epsilon} \| P _ {d , j} \left[ \bar {\theta} _ {N} ^ {(2)} \right] \| _ {2}}{N}. \\ \end{array}
$$
Combined with Equation (46), one can conclude that,