SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lemma 13. Attention module is lipschitz with respect to 1, 2-norm for bounded input. Precisely, consider attention module (Definition 7) parameterized by  $\| W_V\|_2$ ,  $\| W_K\|_2$ ,  $\| W_Q\|_2 \leq 1$  with input domain  $\| X\|_{1,2} \leq 4L$ ,  $a(X)$  is  $200L^2 / C$ -lipschitz with respect to 1, 2-norm.
$$
a (\boldsymbol {X}) = \mathrm {L N} _ {C} \left(W _ {V} \boldsymbol {X} \sigma (\boldsymbol {X} ^ {\top} W _ {K} ^ {\top} W _ {Q} \boldsymbol {X})\right).
$$
Choose  $\epsilon$  to be a sufficiently small constant, such that,  $\exp (32L\epsilon) - 1\leq 64L\epsilon$ . Consider  $\pmb{X}$  and  $\tilde{\pmb{X}}$  satisfying that  $\| X - \tilde{X}\|_{1,2}\leq \epsilon$  and  $\| X\|_{1,2}\leq 4L$ ,  $\| \tilde{X}\|_{1,2}\leq 4L$ , we will have
$$
\begin{array}{l} \left| \left(\boldsymbol {X} ^ {\top} W _ {K} ^ {\top} W _ {Q} \boldsymbol {X} - (\tilde {\boldsymbol {X}}) ^ {\top} W _ {K} ^ {\top} W _ {Q} (\tilde {\boldsymbol {X}})\right) _ {i, j} \right| \\ = \left| \left(\boldsymbol {X} _ {:,: i} - \tilde {\boldsymbol {X}} _ {:,: i}\right) ^ {\top} W _ {K} ^ {\top} W _ {Q} \boldsymbol {X} _ {:, j} + \left(\tilde {\boldsymbol {X}} _ {:, i}\right) ^ {\top} W _ {K} ^ {\top} W _ {Q} \left(\boldsymbol {X} _ {:, j} - \tilde {\boldsymbol {X}} _ {:, j}\right) + \left(\boldsymbol {X} _ {:,: i} - \tilde {\boldsymbol {X}} _ {:,: i}\right) ^ {\top} W _ {K} ^ {\top} W _ {Q} \left(\boldsymbol {X} _ {:, j} - \tilde {\boldsymbol {X}} _ {:, j}\right) \right| \\ \leq 8 L \epsilon + \epsilon^ {2} \leq 1 6 L \epsilon . \\ \end{array}
$$
By Lemma 8, this implies,
$$
\left\| \sigma \left(\boldsymbol {X} ^ {\top} W _ {K} ^ {\top} W _ {Q} \boldsymbol {X}\right) - \sigma \left(\left(\tilde {\boldsymbol {X}}\right) ^ {\top} W _ {K} ^ {\top} W _ {Q} \left(\tilde {\boldsymbol {X}}\right)\right) \right\| _ {1, 1} \leq \exp (3 2 L \epsilon) - 1 \leq 6 4 L \epsilon .
$$
$$
\begin{array}{l} \left\| W _ {V} \left(\boldsymbol {X} - \tilde {\boldsymbol {X}}\right) \right\| _ {1, 2} \leq \epsilon . \\ \| W _ {V} \boldsymbol {X} \| _ {1, 2} \leq 4 L \\ \end{array}
$$
Lemma 11 then implies that
$$
\left\| a (\boldsymbol {X}) - a \left(\tilde {\boldsymbol {X}}\right) \right\| _ {1, 2} \leq 2 0 0 L ^ {2} \epsilon / C.
$$
This then concludes the proof.
We can now prove that a large Transformer Layer and an attention module of the larger Transformer can be pruned to approximate the attention module of a smaller Transformer Layer module.
Lemma 14. Under the setting of Theorem 2, with probability  $1 - \delta /2$ , for any  $l\in [L],l^{\prime}\in [4L - 1]$  let  $a^{(l)}$  be the attention module on the  $l$ -th layer of  $\mathcal{T}$ , there exists a pruning of the  $(l^{\prime} - 1)$ -th layer  $\mathcal{T}_{large}^{(l^{\prime} - 1)}$ , named  $\tilde{\mathcal{T}}_{large}^{(l^{\prime} - 1)}$  and the attention module on  $l^{\prime}$ -th layer  $a_{\mathrm{large}}^{l^{\prime}}$  named  $a_{\mathrm{large}}^{l^{\prime}}$ , such that when defined on domain  $\| X\|_{1,2}\leq 2L$ ,
1.  $\tilde{\mathcal{T}}_{\text{large}}^{(l' - 1)}$  is independent of the last  $m_{\text{large}} - m$  rows of the input.
2.  $\left(a_{\mathrm{large}}^{\tilde{l}^{\prime}}\circ \tilde{\mathcal{T}}_{\mathrm{large}}^{(l^{\prime} - 1)}\left(\left[ \begin{array}{c}x\\ 0^{(m_{\mathrm{large}} - m)\times N} \end{array} \right]\right)\right)_{1:m}$  is an  $\left(\frac{C}{1000L^2}\right)^{4L - 1}\epsilon$  -approximation of  $a^{(l)}(x)$  with respect to 1, 2-norm.
3.  $\left(\tilde{T}_{large}^{(l' - 1)}\left(\left[0^{(m_{large} - m)\times N}\right]\right)\right)_{1:m}$  is an  $\left(\frac{C}{1000L^2}\right)^{4L}\epsilon$ -approximation of  $\mathbf{X}$  with respect to 1, 2-norm.
Proof. We will use the shorthand  $\epsilon_0 = \left(\frac{C}{1000L^2}\right)^{4L}\epsilon$  and prune in the following order. It holds that for  $\epsilon \leq 1$ ,  $\exp (8L^2\epsilon_0) - 1\leq 16L^2\epsilon_0$ .
1. We will prune  $W_V^{\mathrm{large},(l')}$  according to Lemma 9 and name the pruned matrix  $W_V^{\mathrm{large},(l')}$ . By Lemma 9, all the nonzero entries is contained in a  $d \times d$  submatrix of  $W'$  that satisfies that all its eigenvalues are within  $\left(\frac{1}{2}, 1\right)$ . We will assume WLOG the submatrix is the one specified by row  $1 \ldots d$  and column  $d + 1 \ldots 2d$  and name the submatrix as  $W$ .
2. We will then prune  $\mathcal{T}_{\mathrm{large}}^{(l' - 1)}$  according to Lemma 4 to output  $\epsilon_0$ -approximation of  $\pmb{X} \in \mathbb{R}^{m \times N} \to \left[ \begin{array}{c} \pmb{X} \\ W^{-1}\mathcal{P}_{\perp}W_{V}^{(l)}\pmb{X} \\ 0^{(m_{\mathrm{large}} - 2m) \times N} \end{array} \right]$ . As  $W$  is defined as the submatrix pruned by  $W_{V}^{(t + 1)}$ , it holds that
$$
\tilde {W _ {V} ^ {\mathrm {l a r g e ,} (l ^ {\prime})}} \left[ \begin{array}{c} \boldsymbol {X} \\ W ^ {- 1} \mathcal {P} _ {\perp} W _ {V} ^ {(l)} \boldsymbol {X} \\ 0 ^ {(m _ {\mathrm {l a r g e}} - m) \times N} \end{array} \right] = \left[ \begin{array}{c} \mathcal {P} _ {\perp} W _ {V} ^ {(l)} \boldsymbol {X} \\ 0 ^ {(m _ {\mathrm {l a r g e}} - m) \times N} \end{array} \right].
$$
3. Finally we will prune  $W_{K}^{\mathrm{large},(l')}, W_{Q}^{\mathrm{large},(l')}$  according to Lemma 7 to approximate  $(W_{K}^{(l)})^{\top} W_{Q}^{(l)}$  up to  $\epsilon_{0}$  error.
we can now calculate the approximation error. For any  $\mathbf{X} \in \mathbb{R}^{m \times N}$ ,  $\| \mathbf{X} \|_{1,2} \leq 2L$ , suppose
$$
\mathcal {T} ^ {(\tilde {l ^ {\prime}} - 1)} (\boldsymbol {X}) = \left[ \begin{array}{c} \boldsymbol {X} + \delta_ {1} \\ W ^ {- 1} \mathcal {P} _ {\perp} W _ {V} ^ {(l)} \boldsymbol {X} + \delta_ {2} \\ 0 ^ {(m _ {\mathrm {l a r g e}} - 2 m) \times N} \end{array} \right]
$$
Then by our construction, it holds that  $\forall i\in \{1,2\} ,\| \delta_i\|_{1,2}\leq \epsilon_0\| X\|_{1,2}$
We would then have
$$
W _ {V} ^ {\tilde {\text {l a r g e}}, \left(l ^ {\prime}\right)} \mathcal {T} ^ {\left(l ^ {\prime} - 1\right)} (\boldsymbol {X}) = \left[ \begin{array}{c} \mathcal {P} _ {\perp} W _ {V} ^ {(l)} \boldsymbol {X} + W _ {V} ^ {\tilde {\text {l a r g e}}, \left(l ^ {\prime}\right)} \delta_ {2} \\ 0 ^ {\left(m _ {\text {l a r g e}} - m\right) \times N} \end{array} \right] \tag {48}
$$
By our construction, it holds that  $\| W_V^{\mathrm{large},(l')} \delta_2 \|_{1,2} \leq 2 \| \delta_2 \|_{1,2} \leq 2 \epsilon_0 \| X \|_{1,2}$ .
Further, by the construction of  $W_{K}^{\mathrm{large},(l')}$ ,  $W_{Q}^{\mathrm{large},(l')}$ , it holds that,
$$
\begin{array}{l} \left\| \left(W _ {K} ^ {\tilde {\text {l a r g e}}, (l ^ {\prime})} \mathcal {T} ^ {(l ^ {\prime} - 1)} (\boldsymbol {X})\right) ^ {\top} \left(W _ {Q} ^ {\tilde {\text {l a r g e}}, (l ^ {\prime})} \mathcal {T} ^ {(l ^ {\prime} - 1)} (\boldsymbol {X})\right) \right. \\ \left. - \left(W _ {K} ^ {(l)} \boldsymbol {X} + W _ {K} ^ {(l)} \delta_ {1}\right) ^ {\top} \left(W _ {Q} ^ {(l)} \boldsymbol {X} + W _ {Q} ^ {(l)} \delta_ {1}\right) \right\| _ {\infty} \leq \epsilon_ {0} \tag {49} \\ \end{array}
$$
As for any  $i,j\in [N]$
$$
\begin{array}{l} \left| \left(\left(W _ {K} ^ {(l)} \boldsymbol {X} + W _ {K} ^ {(l)} \delta_ {1}\right) ^ {\top} \left(W _ {Q} ^ {(l)} \boldsymbol {X} + W _ {Q} ^ {(l)} \delta_ {1}\right) - \left(W _ {K} ^ {(l)} \boldsymbol {X}\right) ^ {\top} W _ {Q} ^ {(l)} \boldsymbol {X}\right) _ {i, j} \right| \\ \leq \left| \left(W _ {K} ^ {(l)} \boldsymbol {X} _ {:, i}\right) ^ {\top} \left(W _ {Q} ^ {(l)} \delta_ {1}\right) _ {:, j} \right| + \left| \left(W _ {K} ^ {(l)} \delta_ {1}\right) _ {:, i} ^ {\top} \left(W _ {Q} ^ {(l)} \boldsymbol {X}\right) _ {:, j} \right| + \left| \left(W _ {K} ^ {(l)} \delta_ {1}\right) _ {:, i} ^ {\top} \left(W _ {Q} ^ {(l)} \delta_ {1}\right) _ {:, j} \right| \\ \leq \| \boldsymbol {X} \| _ {1, 2} ^ {2} (2 \epsilon_ {0} + \epsilon^ {2}) \leq 4 \| \boldsymbol {X} \| _ {1, 2} ^ {2} \epsilon_ {0}. \\ \end{array}
$$
combined with Equation (49),
$$
\left. \right.\left\|\left(W _ {K} ^ {\tilde {\operatorname* {l a r g e}}, (l ^ {\prime})} \mathcal {T} ^ {(\tilde {l} ^ {\prime} - 1)} (\boldsymbol {X})\right) ^ {\top} \left(W _ {Q} ^ {\tilde {\operatorname* {l a r g e}}, (l ^ {\prime})} \mathcal {T} ^ {(\tilde {l} ^ {\prime} - 1)} (\boldsymbol {X})\right) - \left(W _ {K} ^ {(l)} \boldsymbol {X}\right) ^ {\top} W _ {Q} ^ {(l)} \boldsymbol {X} \right\| _ {\infty} \leq \epsilon_ {0} (1 + 4 \| \boldsymbol {X} \| _ {1, 2} ^ {2}). \tag {50}
$$
By Lemma 8, this implies
$$
\begin{array}{l} \left\| \sigma \left(\left(W _ {K} ^ {\tilde {\text {l a r g e}}, (l ^ {\prime})} \mathcal {T} ^ {(l ^ {\prime} - 1)} (\boldsymbol {X})\right) ^ {\top} \left(W _ {Q} ^ {\tilde {\text {l a r g e}}, (l ^ {\prime})} \mathcal {T} ^ {(l ^ {\prime} - 1)} (\boldsymbol {X})\right)\right) \right. \\ \left. - \sigma \left(\left(W _ {K} ^ {(l)} \boldsymbol {X}\right) ^ {\top} W _ {Q} ^ {(l)} \boldsymbol {X}\right) \right\rVert_ {1, 1} \leq 4 \epsilon_ {0} (1 + 4 \| \boldsymbol {X} \| _ {1, 2} ^ {2}). \tag {51} \\ \end{array}
$$
By Lemma 11, Equations (48) and (51) imply,
$$
\begin{array}{l} \left\| W _ {V} ^ {\tilde {\text {l a r g e}}, (l ^ {\prime})} \mathcal {T} ^ {(\tilde {l ^ {\prime}} - 1)} (\boldsymbol {X}) \sigma \left(\left(W _ {K} ^ {\tilde {\text {l a r g e}}, (l ^ {\prime})} \mathcal {T} ^ {(\tilde {l ^ {\prime}} - 1)} (\boldsymbol {X})\right) ^ {\top} \left(W _ {Q} ^ {\tilde {\text {l a r g e}}, (l ^ {\prime})} \mathcal {T} ^ {(\tilde {l ^ {\prime}} - 1)} (\boldsymbol {X})\right)\right) \right. \\ -\left[ \begin{array}{c}\mathcal{P}_{\perp}W_{V}^{(l)}\boldsymbol {X}\sigma \left((W_{K}^{(l)}\boldsymbol {X})^{\top}W_{Q}^{(l)}\boldsymbol {X}\right)\\ 0^{(m_{\text{large}} - m)\times N} \end{array} \right]\bigg|_{1,2}\leq 8\epsilon_{0}(1 + 4\| \boldsymbol {X}\|_{1,2}^{2})\| \boldsymbol {X}\|_{1,2}\leq 80L^{2}\epsilon_{0}. \\ \end{array}
$$
Now according to Lemmas 10 and 12, it holds that
$$
\| a _ {\text {l a r g e}} ^ {l ^ {\prime}} \circ \tilde {\mathcal {T}} _ {\text {l a r g e}} ^ {(l ^ {\prime} - 1)} \left(\left[ \begin{array}{c} x \\ 0 ^ {(m _ {\text {l a r g e}} - m) \times N} \end{array} \right]\right) _ {1: m} - a ^ {(l)} (x) \| _ {1, 2} \leq 1 6 0 L ^ {2} \epsilon_ {0} / C.
$$
Approximating Transformer Layers We will finally show that two random Transformer layers can be pruned to approximate a given Transformer layer.
Lemma 3. Under the setting of Theorem 2, with probability  $1 - 2\delta /3$  , for any  $l\in [L],l^{\prime}\in [4L - 1]$  , let  $\mathcal{T}^{(l)}$  be the  $l$  -th layer of  $\mathcal{T}$  , there exists a pruning of the  $(l^{\prime} - 1)$  -th and the  $(l^{\prime})$  -th layer  $\mathcal{T}_{large}^{(l^{\prime} - 1)},\mathcal{T}_{large}^{\prime},$  named  $\tilde{\mathcal{T}}_{large}^{(l^{\prime} - 1)},\tilde{\mathcal{T}}_{large}^{\prime}$  such that when defined on domain  $\| X\|_{1,2}\leq 2L,X\in \mathbb{R}^{m\times N},$
1.  $\tilde{\mathcal{T}}_{\text{large}}^{(l' - 1)}$  is independent of the last  $m_{\text{large}} - m$  rows of the input.
2.  $\tilde{\mathcal{T}}_{large}^{l^{\prime}}\circ \tilde{\mathcal{T}}_{large}^{(l^{\prime} - 1)}(\bar{X})$  is an  $\left(\frac{C}{1000L^2}\right)^{4L - 3}\epsilon$ -approximation of  $\overline{\mathcal{T}^{(l)}(\bar{X})}$  with respect to 1, 2-norm.
Proof. We will prune the  $(l' - 1)$ -th layer and the attention module of the  $l'$ -th layer according to Lemma 14 to approximate  $a^{(l)}$  and the projection function of the  $l'$ -th layer according to Corollary 1. Notice that  $\left\| a^{(l)}(\mathbf{X}) + \mathbf{X}\right\|_{1,2} \leq \left(\frac{2}{C} + 1\right)\|\mathbf{X}\|_{1,2}$  and  $\mathrm{g}^{(l)}$  is 1-lipschitz, according to Lemma 5,  $\left(\tilde{\mathcal{T}}_{\text{large}}^{\prime} \odot \tilde{\mathcal{T}}_{\text{large}}^{(l' - 1)}\left(\left[0^{(m_{\text{large}} - m) \times N}\right]\right)\right)_{1:m}$  is an  $\epsilon'$ -approximation of  $\mathcal{T}^{(l)}(x)$ , with
$$
\epsilon^ {\prime} \leq \left(\frac {2}{C} + 1\right) \left(\frac {C}{1 0 0 0 L ^ {2}}\right) ^ {4 L} \epsilon + \left(\frac {C}{1 0 0 0 L ^ {2}}\right) ^ {4 L - 2} \epsilon + \left(\frac {C}{1 0 0 0 L ^ {2}}\right) ^ {8 L - 2} \epsilon^ {2} \leq \left(\frac {C}{1 0 0 0 L ^ {2}}\right) ^ {4 L - 3} \epsilon .
$$
C.6 Technical Lemmas
Lemma 15. Given any dimension  $d$  and number of samples  $n$ , for any size- $n$  dataset  $\{(\pmb{x}_i, y_i)\}_{i \in [n]}$  with  $\pmb{x}_i \in \mathbb{R}^d$  and  $y_i \in \mathbb{R}$ , there exists a width-2n two-layer MLP  $f: \mathbb{R}^d \to \mathbb{R}$  with ReLU activation such that,  $f(\pmb{x}_i) = y_i$  for any  $i \in [n]$ .
Proof. We will first choose direction  $w \in \mathbb{R}^d$ ,  $\| w \|_2 = 1$  and margin  $\gamma > 0$  such that for any  $i \neq j$  in  $[n]$ , it holds that,
$$
\left| \langle w, \boldsymbol {x} _ {i} - \boldsymbol {x} _ {j} \rangle \right| \geq 2 \gamma .
$$
We will assume WLOG  $w^\top \pmb{x}_i$  is increasing in  $i$ .
Then we will construct an auxiliary series  $z_{i}$  for  $i\in [n]$  such that,
$$
z _ {1} = y _ {1} / \gamma
$$
$$
z _ {i} = y _ {i} / \gamma - 2 \sum_ {j = 1} ^ {i - 1} z _ {j}, i \in \{2, \dots n \}.
$$
Finally consider the following two-layer MLP with ReLU activation,
$$
f (\boldsymbol {x}) = \sum_ {i = 1} ^ {n} z _ {i} \operatorname {R e L U} \left(\langle w, \boldsymbol {x} - \boldsymbol {x} _ {i} \rangle + \gamma\right) - z _ {i} \operatorname {R e L U} \left(\langle w, \boldsymbol {x} - \boldsymbol {x} _ {i} \rangle - \gamma\right),
$$
we will show that  $f(\pmb{x}_i) = y_i$  for any  $i \in [n]$ . Notice that
$$
z _ {j} \mathrm {R e L U} \left(\langle w, \boldsymbol {x} _ {i} - \boldsymbol {x} _ {j} \rangle + \gamma\right) - z _ {j} \mathrm {R e L U} \left(\langle w, \boldsymbol {x} _ {i} - \boldsymbol {x} _ {j} \rangle - \gamma\right) = \left\{ \begin{array}{l l} 0, & j > i, \\ \gamma z _ {i}, & j = i, \\ 2 \gamma z _ {j}, & j <   i. \end{array} \right.
$$
Thus it holds,
$$
\begin{array}{l} f (\boldsymbol {x} _ {i}) = \sum_ {j = 1} ^ {n} z _ {j} \operatorname {R e L U} \left(\langle w, \boldsymbol {x} _ {i} - \boldsymbol {x} _ {j} \rangle + \gamma\right) - z _ {j} \operatorname {R e L U} \left(\langle w, \boldsymbol {x} _ {i} - \boldsymbol {x} _ {j} \rangle - \gamma\right) \\ = \sum_ {j = 1} ^ {i - 1} 2 \gamma z _ {j} + \gamma z _ {i} = y _ {i}. \\ \end{array}
$$
Lemma 16. Given any sets  $\{x_i\}_{i\in m}$  satisfying that  $x_{i}\in \mathbb{R}^{n}$  and  $x_{i}\neq 0$  , there exists a set of orthonormal vectors  $\{u_j\}_{j\in [n - 2]}$  of  $\mathbb{R}^n$  such that (1)  $u_{j}^{\top}1^{n} = 0$  for any  $j\in [n - 2]$  and (2)  $\sum_{j\in [n - 2]}u_j^\top x_iu_j\neq x_i$  for any  $i\in [m]$
Proof. There exists a vector  $v \in \mathbb{R}^n$  such that  $v^\top x_i \neq 0$  for any  $i \in [m]$ . We can then construct an orthonormal basis  $\{u_j\}_{j \in [n-2]}$  of  $\mathbb{R}^n$  as the basis of the normal space of  $\operatorname{span}(v, 1^n)$ . Then the lemma holds.
Lemma 17. Given any dimension  $n$  and constant  $M$ , there exists a 2-layer width-2n ReLU network  $f: \mathbb{R}^{n + 1} \to \mathbb{R}$  such that for any  $x \in [0, M]^n$ ,  $y \in [n]$ ,  $f(x \oplus y) = x_y$ .
Proof. The construction is as followed, we will choose  $f$  as
$$
f (x \oplus y) = \sum_ {i = 1} ^ {n} \operatorname {R e L U} \left(x _ {i} + M (y - i)\right) - \sum_ {i = 1} ^ {n} \operatorname {R e L U} \left(x _ {i} + M (y - i - 1)\right) - M (y - 1).
$$
Then as we have
$$
\operatorname {R e L U} (x _ {i} + M (y - i)) - \sum_ {i = 1} ^ {n} \operatorname {R e L U} (x _ {i} + M (y - i - 1)) = \left\{ \begin{array}{l l} M, & i \leq y - 1; \\ x _ {i}, & i = y; \\ 0, & i \geq y + 1. \end{array} \right.
$$
Lemma 18. Given any dimension  $n$  and constant  $M > 0$ , there exists a 2-layer width-2n ReLU network  $f: \mathbb{R}^n \to \mathbb{R}$  such that for any  $x \in \mathbb{R}^n$  satisfying there exists  $i \in [n]$ ,  $x_i > M$  and  $\forall j \neq i, x_j = 0$ , it holds that  $f(x) = i$ .
$$
f (x) = \sum_ {i = 1} ^ {m} i (\operatorname {R e L U} (x _ {i}) - \operatorname {R e L U} (x _ {i} - M) + M) / M.
$$
Lemma 19. Given any dimension  $n$  and natural numbers  $K, m, M$ , if there exists  $K$  different 2-layer width-  $m$  ReLU networks  $f_{k} : \mathbb{R}^{n} \to \mathbb{R}$ , then there exists a 2-layer width-  $2Km$  ReLU network  $f : \mathbb{R}^{n + 1} \to \mathbb{R}$ , such that  $f(\left[ \begin{array}{c} k \\ x \end{array} \right]) = f_{k}(x)$  when  $x \in [0, M]^{n}$ .
Proof. Suppose that
$$
f _ {k} (x) = \sum_ {i = 1} ^ {m} a _ {k, i} \operatorname {R e L U} \left(w _ {k, i} ^ {\top} x + b _ {k, i}\right) + b _ {k}.
$$
Then we can construct
$$
\begin{array}{l} f(\left[ \begin{array}{c}y\\ x \end{array} \right]) = \sum_{k = 1}^{K}\sum_{i = 1}^{m}a_{k,i}\mathrm{ReLU}(w_{k,i}^{\top}x + b_{k,i} + M(y - k)) - a_{k,i}\mathrm{ReLU}(w_{k,i}^{\top}x + b_{k,i} + M(y - k - 1)) \\ + b _ {k} - c _ {k, i} \operatorname {R e L U} (y + 1 - k), \\ \end{array}
$$
where  $c_{k,i}$  satisfies
$$
\forall i, k ^ {\prime}, \sum_ {k = 1} ^ {k ^ {\prime}} c _ {k, i} \left(k ^ {\prime} + 1 - k\right) = M \sum_ {k = 1} ^ {k ^ {\prime} - 1} a _ {k, i}.
$$
Lemma 20. Given any dimension  $n$  and  $W \in \mathbb{R}^{n \times n}$ ,  $\| W \|_2 \leq 2$ , there exists a 2-layer width-2n ReLU network  $f: \mathbb{R}^n \to \mathbb{R}$  such that for any  $x \in \mathbb{R}^n$ , it holds that  $f(x) = Wx$  and both weight matrices parameterizing  $f$  has spectral norm less than  $2\sqrt{2}$ .
Proof. The construction is straightforward, one can choose
$$
f (x) = [ I _ {n}, - I _ {n} ] ^ {\top} \mathrm {R e L U} \left(\left[ \begin{array}{c} W x \\ - W x \end{array} \right]\right).
$$
C.7 Discussion on Architecture Choices
The reader may notice that Equation (6) is not the same as the standard GPT architecture,
$$
f _ {l} (X; \theta^ {(l)}) = \mathrm {g} ^ {(l)} \left(\mathrm {L N} \left(W _ {V} ^ {(l)} X \sigma \left(\mathcal {C} + \left(W _ {K} ^ {(l)} X\right) ^ {\top} \left(W _ {Q} ^ {(l)} X\right)\right) + X\right)\right). \tag {52}
$$
We will shortly discuss the impact of considering Equation (52) here.
With similar arguments to Theorem 3 and the necessity part of Theorem 1, one can prove that similar balance conditions should also hold for a transformer with a layer specified by Equation (52) and a minimal first layer that can nearly perfectly generate bounded Dyck languages.
However, the sufficiency part of Theorem 1 no longer holds, when the balance condition holds, the last column of the term  $W_V^{(2)}X\sigma \Big(\mathcal{C} + (W_K^{(2)}X)^\top (W_Q^{(2)}X)\Big)$  will converge to zero when input length converges to infinity. Hence, if not all  $e(\tau_{t,d})$  where  $\tau_{t,d}$  is a closed bracket aligns with  $1^m$ , then it is impossible for the model to perfectly generate Dyck for arbitrary length. Although it remains possible to refine a sharper condition for standard GPT architecture to perfectly generate Dyck Language, we find considering Equation (6) more elegant in theory. We also verify with experiments that our architecture with standard training can learn bounded Dyck language to more than  $97\%$  accuracy. Also, the learned attention patterns are also similarly not interpretable as standard architectures.
D.1 Training Details
For Figure 1, we train 2-layer standard GPT on Dyck $_{2,4}$  with sequence length no longer than 28. For  $(a)$ , we train with hidden dimension and network width 200 and learning rate 3e-4. For  $(b)$ ,  $(c)$ ,  $(d)$ , we train with hidden dimension and FFN width 50 and learning rate 3e-3.
For Figure 2, for  $(a)$ , we train 1-layer transformer without residual link, FFN and the final LayerNorm before the linear head. The hidden dimensions and FFN widths are fixed as 500. For  $(a)$ , we train the network with learning rate 1e-2 and for  $(b)$ ,  $(c)$ ,  $(d)$  we train the network with learning rate 3e-3.
D.2 Additional Results on Dyck Prefix
In the experiment presented in the main text, we perform experiments on complete Dyck sequences, which is a special case of Dyck prefixes. In this section, we present additional experiments on Dyck prefixes  $\mathrm{Dyck}_{2,4,28}$ .
Attention Patterns We first perform experiments on attention patterns. The qualitative results are shown in Figures 7 and 9. We can observe that the attention patterns are still diverse and do not commonly show stack-like patterns. We also calculate the attention variation  $^{22}$ , and find that the attention variation is 0.34, based on 30 models with a minimal first layer and different random seeds. In contrast, for models with a standard first layer and without position encodings, the attention variation is surprisingly high, reaching 14.51. The high value is caused by the large distance between attention patterns like Figure 7 (c) and (d); that is, between patterns that attend more to the current positions, and patterns that attend more heavily to the initial position. The difference is even increased when we consider longer sequence (Figure 8). Similarly, the variation is also high for models with linear position embedding, reaching 11.92. This shows that the attention patterns are still diverse and do not commonly show stack-like patterns.
(c) Without Position Embedding
Figure 7: Second-layer attention patterns of two-layer Transformers on Dyck Prefix: Models for (a),(b) are under the same setup but different random seeds; similarly for (c),(d). All models reach  $\geq 97\%$  accuracy (defined in Section 4.1). In the heatmap, darker color indicates larger value. As we can observe, the attention patterns still show much variance.
Balanced Violations We also test the relationship with the balance violation with length generalization on Dyck prefixes, similar to Figure 3. We observe that although the negative correlation is not presented as in the case of Dyck sequences, contrastive regularization still helps reduce the balance violation and significantly improve the length generalization performance. This shows that for Dyck prefixes, while the balance violation may not be predictive of the length generalization performance, it is still possible to reduce the balance violation and improve the length generalization performance. The results are shown in Figure 10.
22Recall from Section 4.1 that the attention variation between two attention patterns  $A_{1},A_{2}\in \mathbb{R}^{N\times N}$  is defined as  $\mathrm{Variation}(A_1,A_2) = \| A_1 - A_2\| _F^2$
(a) With Position Embedding Run 1
(b) With Position Embedding Run 2
(c) Without Position Embedding Run 1
Figure 8: Second-layer attention patterns of two-layer Transformers on Longer Dyck Prefix: Models for (a),(b) are under the same setup but different random seeds. All models reach  $\geq 97\%$  accuracy (defined in Section 4.1). In the heatmap, darker color indicates larger value.
(d) Without Position Embedding Run 2
D.3 Extended Experiments
We include more experiments on the attention variation of different Dyck languages and architectures. The results are summarized in Table 1.
#types k	Grammar depth m	#Layers l	Layer 1	Layer 2	Layer 3
2	4	2	0.047(0.006)	7.721(0.908)
2	4	3	0.070(0.013)	5.072(0.645)	24.063(1.166)
2	8	2	0.087(0.012)	7.583(0.961)
2	8	3	0.059(0.011)	5.560(0.714)	23.590(0.829)
3	4	2	0.182(0.024)	9.313(0.815)
3	4	3	0.225(0.032)	8.426(0.877)	25.749(0.897)
3	8	2	0.178(0.028)	7.000(0.884)
3	8	3	0.154(0.036)	6.280(0.711)	25.451(0.871)
Table 1: Extended attention variation. "Layer  $i$  " shows the mean (and standard deviation) of the attention variation on layer  $i$  ,calculated on 40 sentences. The embedding width and FFN width are fixed as 50 in the experiments. We train using sentences from Dyck  ${}_{k,m}$  of length less than 28 and test the variation on 40 randomly sampled sentences with length 19 (the sampled sentence is fixed across different architectures). The random attention variation baseline here is 3.33 . The numbers in this table are different from previous discussion, since the results here are from a slightly different architecture than the standard GPT-2 architecture: a residue link is appended after the LayerNorm to match our theory better. The models are trained to convergence and have in-distribution accuracy higher than 97%.
Attention pattern visualization for three-layer experiments. The first-layer attention is close to uniform, while the higher-layer attention shows no clear patterns.