SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Interpreting attention patterns Prior works (Jain & Wallace, 2019; Serrano & Smith, 2019; Rogers et al., 2020; Grimsley et al., 2020; Brunner et al., 2020; Prasanna et al., 2020; Meister et al., 2021; Bolukbasi et al., 2021; Haab et al., 2023, inter alia) present negative results on deriving explanations from attention weights using approaches by Vig & Belinkov (2019); Kobayashi et al. (2020, inter alia). However, Wegreffe & Pinter (2019) argues to the contrary by pointing out flaws in the experimental design and arguments of some of the prior works; they also call for theoretical analysis on the issue. Hence, a takeaway from these prior works is that expositions on explainability based on attention requires clearly defining the notion of explainability adopted (often task-specific). In our work, we restrict our main theoretical analysis to the fully defined data distribution of Dyck language (Definition 1), and define "interpretable attention pattern" as the stack-like pattern proposed in prior theoretical (Yao et al., 2021) and empirical (Ebrahimi et al., 2020) works. These concrete settings and definitions allow us to mathematically state our results and provide theoretical reasons.
Theoretical understanding of representability Methodologically, our work joins a long line of prior works that characterize the solution of neural networks via the lens of simple synthetic data, from class results on RNN representability (Siegelmann & Sontag, 1992; Gers & Schmidhuber, 2001; Weiss et al., 2018; Suzgun et al., 2019; Merrill, 2019; Hewitt et al., 2020), to the more recent Transformer results on parity (Hahn, 2020), Dyck (Yao et al., 2021), topic model (Li et al., 2023), and formal grammars in general (Bhattamishra et al., 2020a; Li & Risteski, 2021; Zhang et al., 2022; Liu et al., 2023; Zhao et al., 2023). Our work complements prior works by showing that although representational results can be obtained via intuitive "constructive proofs" that assign values to the weight matrices, the model does not typically converge to those intuitive solutions in practice. Similar messages are conveyed in Liu et al. (2023), which presents different types of constructions using different numbers of layers. In contrast, we show that there exist multiple different constructions even when the number of layers is kept the same.
There are also theoretical results on Transformers in terms of Turing completeness (Bhattachamishra et al., 2020b; Perez et al., 2021), universal approximatability (Yun et al., 2020), and statistical sample complexity (Wei et al., 2021; Edelman et al., 2022), which are orthogonal to our work.
Transformer optimization Given multiple global optima, understanding Transformer solutions requires analyzing the training dynamics. Recent works theoretically analyze the learning process of Transformers on simple data distributions, e.g. when the attention weights only depend on the position information (Jelassi et al., 2022), or only depend on the content (Li et al., 2023). Our work studies a syntax-motivated setting in which both content and position are critical. We also highlight that Transformer solutions are very sensitive to detailed changes, such as positional encoding, layer norm, sharpness regularization (Foret et al., 2020), or pre-training task (Liu et al., 2022a). On a related topic but towards different goals, a series of prior works aim to improve the training process of Transformers with algorithmic insights (Nguyen & Salazar, 2019; Xiong et al., 2020; Liu et al., 2020; Zhang et al., 2020; Li & Gong, 2021, inter alia). An end-to-end theoretical characterization of the training dynamics remains an open problem; recent works that propose useful techniques towards this goal include Gao et al., 2023; Deng et al., 2023.
Mechanistic interpretability Finally, it is worth noting that the challenges highlighted in our work do not contradict the line of prior works that aim to improve mechanistic interpretability into a trained model or the training process (Cammarata et al., 2020; Elhage et al., 2021; Olsson et al., 2022; Nanda et al., 2023; Chughtai et al., 2023; Li et al., 2023; Wang et al., 2023; Zhong et al., 2023): although we prove that components (e.g. attention scores) of trained Transformers do not generally admit intuitive interpretations based on the data distribution, it is still possible to develop circuit-level understanding about a particular model, or measures that closely track the training process, following these prior works.
B Are interpretable attention patterns useful?
Our results Section 3 and Section 4.1 demonstrate that Transformers are sufficiently expressive that a (near-)optimal loss on Dyck languages can be achieved by a variety of attention patterns, many of which may not be interpretable.
However, multiple prior works have shown that for multi-layer multi-head Transformers trained on natural language datasets, it is often possible to locate attention heads that produce interpretable attention patterns (Vig & Belinkov, 2019; Htut et al., 2019; Sun & Marasovic, 2021). Hence, it is also illustrative to consider the "converse question" of (Q1): when some attention heads do learn to produce attention patterns that suggest intuitive interpretations, what benefits can they bring?
We discuss this through two perspectives:
- Reliability of interpretation: Is the Transformer necessarily implementing a solution consistent with such interpretation based on the attention patterns? (Section B.1)
- Usefulness for task performance: Are those interpretable attention heads more important for the task than other uninterpretable attention heads? (Section B.2)
We present preliminary analysis on these questions, and motivate future works on the interpretability of attention patterns using rigorous theoretical analysis and carefully designed experiments.
B.1 Can interpretable attention patterns be misleading?
We show through a simple argument that interpretations based on attention patterns can sometimes be misleading, as we formalize in the following proposition:
Proposition 1. Consider an  $L$ -layer Transformer  $\mathcal{T}$  (Equation (10)). For any  $W_{K}^{(l)}, W_{Q}^{(l)} \in \mathbb{R}^{m_{a} \times m}$  ( $l \in [L]$ ), there exist  $W_{\mathrm{Head}} \in \mathbb{R}^{2k \times w}$  and  $b_{\mathrm{Head}} \in \mathbb{R}^{2k}$  such that  $\mathcal{T}(\mathcal{Z}) = 0, \forall \mathcal{Z}$ .
While its proof is trivial (simply setting  $W_{\mathrm{Head}} = 0$  and  $b_{\mathrm{Head}} = 0$  suffices), Proposition 1 implies that the solution represented by the Transformer could possibly be independent of the attention patterns in all the layers (1 through  $l$ ). Hence, it could be misleading to interpret Transformer solutions solely based on these attention patterns.
Empirically, Transformers trained on Dyck indeed sometimes produce misleading attention patterns.
We present one representative example in Figure 4, and Figure 5, in which all interpretable attention patterns are misleading.
We also present additional results in Figure 6, in which some interpretable attention patterns are misleading, and some are not.
Figure 4: Even interpretable attention patterns can be misleading: For a 4-layer Transformer trained on Dyck with the copying task (with  $>96\%$  validation accuracy), i.e. the output should be exactly the same as the input, the attention patterns in some layers seem interpretable: (layer 2) attending to bracket type a) or (b; (layer 3) attending to closing bracketss; (layer 4) neve attending to bracket type a); However, none of them are informative of the copying task. This is possible because Transformers can use the residual connections (or weights MLPs or the value matrices) to solve copying, bypassing the need of using attention.
Similar message has been conveyed in prior works Bolukbasi et al. (2021), and future works may aim to achieve the faithfulness, completeness, and minimality conditions in Wang et al. (2023).
Figure 5: Even interpretable attention patterns can be misleading: For a 1-layer Transformer trained on Dyck with the copying task (with  $>90\%$  validation accuracy), i.e. the output should be exactly the same as the input, the attention pattern seems to be attending to closing brackets only, but that is not informative of the copying task.
(a) layer 1 of 4
(b) layer 3 of 4
Figure 6: Even interpretable attention patterns can be misleading: For a 4-layer Transformer trained on Dyck with the copying task (with  $>96\%$  validation accuracy), i.e. the output should be exactly the same as the input, both types of attention patterns are common: (a) attending to closing bracketss, which is uninformative of the copying task; (b) attending to the current position, which solves the copying task.
B.2 Are attention heads with interpretable patterns more important?
Kovaleva et al. (2019) observes that, when the "importance" of an attention head is defined as the performance drop the model suffers when the head is disabled, then for most tasks they test, the most important attention head in each layer does not tend to be interpretable.
However, experiments by Voita et al. (2019) led to a seemingly contradictory observation: when attention heads are systematically pruned by finetuning the Transformer with a relaxation of  $L_{0}$ -penalty (i.e. encouraging the number of remaining attention heads to be small), most remaining attention heads that survive the pruning can be associated with certain functionalities such as positional, syntactic, or attending to rare tokens.
These works seem to bring mixed conclusions to our question: are interpretable attention heads more important for a task than uninterpretable ones? We interpret these results by conjecturing that the definition of "importance" (reflected in their experimental design) plays a crucial role:
- When the importance of an attention head is defined treating all other attention heads as fixed, motivating experiments that prune/disable certain heads while keeping other heads unchanged (Michel et al., 2019; Kovaleva et al., 2019), the conclusion may be mostly pessimistic: mostly no strong connection between interpretability and importance.
- On the other hand, when the importance of an attention head is defined allowing all other attention heads to adapt to its change, motivating experiments that jointly optimize all attention heads while penalizing the number of heads (Voita et al., 2019), the conclusion may be more optimistic: the heads obtained as a result of this optimization tend to be interpretable.
We think the following trade-offs apply:
- On one hand, the latter setting is more practical, since Transformers are typically not trained to explicitly ensure that the model performs well when a single attention head is individually disabled; rather, it would be more intuitive to think of a group of attention heads as jointly representing some transformation, so when one head is disabled, other heads should be fine-tuned to adapt to the change.
- On the other hand, when all other heads change too much during such fine-tuning, the resulting set of attention heads no longer admit an unambiguous one-to-one map with the original set of (unpruned) attention heads. As a result, the interpretability and importance obtained from the set of pruned heads do not necessarily imply those properties of the original heads.
A comprehensive study of this question involves multi-head extensions of our theoretical results (Section 3), and carefully-designed experiments that take the above-mentioned trade-offs into consideration. We think these directions are interesting future work.
C Omitted Proofs in Section 3
The key step is already shown in Section 3. We will restate the proof rigorously here.
$$
\theta := \left\{\boldsymbol {e} (\tau_ {i, d}) \right\} _ {d \in [ D ], i \in [ 2 k ]}, \theta^ {(2)} \} = \arg \min  _ {\tilde {\theta}} \mathcal {L} (\tilde {\theta}; \mathcal {D} _ {q, k, D, N}), \forall N,
$$
- Equation (11) a necessary condition of  $\mathcal{O}$ , if  $W_V^{(2)}$  satisfies  $\mathcal{P}_{\perp}W_V^{(2)}e(\tau_{t,d}) \neq 0, \forall t \in [k], d \in [D]$ .
- Equation (11) is a sufficient condition of  $\mathcal{O}$ , if the set of  $2k + 1$  encodings  $\{\pmb{e}(\tau_{2i - 1,d}),\pmb{e}(\tau_{2i,d})\}_{i\in [k]}\cup \{\pmb{e}(t_S)\}$  are linearly independent for any  $d\in [D]$ , and the projection function  $\mathrm{g}^{(2)}$  is a 6-layer MLP<sup>12</sup> with  $O(k^2 D^2)$  width.
Proof. We prove the sufficiency of the balanced condition below; the proof for the necessity has been given in Section 3.1.
We will denote the dimension of  $e(\tau_{t,d})$  as  $m$ .
For any  $i\in [k],d'\in [D]$ , by Equation (11), we can assume that there exists  $a_{i,d'}\in \mathbb{R}$  such that for any  $j\in [k],d\in [D]$ , it holds that,
$$
a _ {i, d ^ {\prime}} \triangleq \left(\boldsymbol {e} \left(\tau_ {2 i - 1, d ^ {\prime}}\right) - \boldsymbol {e} \left(\tau_ {2 i, d ^ {\prime} - 1}\right)\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} \boldsymbol {e} \left(\tau_ {2 j, d}\right). \tag {19}
$$
We will first define the possible index sets of  $\tau_{t,d}$  as  $\mathcal{I} = \{(2t,d)\mid t\in [k],0\leq d\leq D - 1\} \cup \{(2t - 1,d)\mid t\in [k],1\leq d\leq D\}$ , and we will define the rank of  $(t,d)$  as
$$
r (t, d) \triangleq \# \left\{\left(t _ {1}, d _ {1}\right) \mid t _ {1} <   t \text {o r} t _ {1} = t, d _ {1} \leq d, \left(t _ {1}, d _ {1}\right) \in \mathcal {I} \right\} \tag {20}
$$
Then it is clear that  $r(t,d)$  is a one-to-one mapping from  $\mathcal{I}$  to  $[2kD]$ . We will then define the collection of all  $e(\tau_{t,d})$  as  $E$ , satisfying that  $E_{:,r(t,d)} = e(\tau_{t,d}), E_{:,2kD + 1} = e(t_S)$ .
Because  $e(\tau_{t,d})$  are linearly independent, for any  $(i,d) \neq (j,d') \in \mathcal{I}$ , it holds that  $e(\tau_{i,d}) - e(\tau_{j,d'}) \neq 0$ . Then based on Lemma 16, there exists a set of orthonormal vectors  $\{\mathrm{b}_i\}_{i \in [m-2]}$ , such that for any  $(i,d), (j,d') \in \mathcal{I}$ , it holds that
$$
\sum_ {i = 1} ^ {m - 2} \mathrm {b} _ {i} \mathrm {b} _ {i} ^ {\top} \left(\boldsymbol {e} \left(\tau_ {i, d}\right) - \boldsymbol {e} \left(\tau_ {j, d ^ {\prime}}\right)\right) \neq \left(\boldsymbol {e} \left(\tau_ {i, d}\right) - \boldsymbol {e} \left(\tau_ {j, d ^ {\prime}}\right) \right. \tag {21}
$$
$$
\mathrm {b} _ {i} ^ {\top} 1 ^ {m} = 0 \tag {22}
$$
We will further construct matrix  $\pmb{O}$  as
$$
\boldsymbol {O} _ {:, r (2 t, d - 1)} = - \exp \left(a _ {t, d}\right) b _ {t D + d},
$$
$$
\boldsymbol {O} _ {:, r (2 t - 1, d)} = \mathrm {b} _ {t D + d}. \tag {23}
$$
$$
\boldsymbol {O} _ {: 2 k D + 1} = 0.
$$
12The 6 layers are by our construction. We will first use 4 layers to convert the input of the projection function to a triplet indicating the type and depth of the last token and the type of the last unmatched bracket when the last token is a closed bracket. We will then use another 2 layers to predict the next token probability based on the triplet. This construction may be improved.
Recall the definition of  $r$  in Equation (20). Comparing  $O_{:,r(2t,d - 1)}$  and  $O_{:,r(2t - 1,d)}$ : the idea is that a pair of matched brackets are represented by the same direction (i.e. the direction along  $b_{tD + d}$ ), just with different norms.
for  $t\in [k],d\in [D]$
We can then choose  $W_V^{(2)}\in \mathbb{R}^{m\times m}$  such that
$$
W _ {V} ^ {(2)} \boldsymbol {E} = \boldsymbol {O} \tag {24}
$$
Such  $W_V^{(2)}$  is guaranteed to exist, because  $E$  is of full column rank by the linear independence assumption.
Now based on this construction, we will show that the last column of unnormalized attention output (Equation (9)) depends only on the sequence of unmatched brackets when the last token is a closed bracket with depth  $d$  greater than or equal to 1.
For any valid Dyck prefix  $p$  of length  $n$  ending with a closed bracket  $\tau_{2j,d}$  satisfying  $d \geq 1$ , suppose the list of unmatched open brackets in  $p$  is  $[\tau_{2j_1 - 1,1},\tau_{2j_2 - 1,2},\dots ,\tau_{2j_d - 1,d}]$ . Then, the remaining tokens in  $p$  are pairs of matching brackets. Denote them by  $\tau_{2t_k - 1,d_k},\tau_{2t_k,d_k - 1}$  for  $k \in [K]$ . Then the input of the second layer of Transformer  $X$ , up to a permutation is
$$
X P = [ \pmb {e} (\tau_ {2 t _ {1} - 1, d _ {1}}), \pmb {e} (\tau_ {2 t _ {1}, d _ {1} - 1}), \dots , \pmb {e} (\tau_ {2 t _ {K} - 1, d _ {K}}), \pmb {e} (\tau_ {2 t _ {K}, d _ {K} - 1}), \pmb {e} (\tau_ {2 j _ {1} - 1, 1}), \dots \pmb {e} (\tau_ {2 j _ {d} - 1, d}), \pmb {e} (t _ {\mathcal {S}}) ].
$$
We will focus on the last column of the unnormalized attention output
$$
\begin{array}{l} \tilde {a} _ {2} (X; \theta^ {(2)}) _ {:, n + 1} = \mathcal {P} _ {\perp} \left[ W _ {V} ^ {(2)} X \cdot \tilde {\sigma} \left(\mathcal {C} \cdot \left(W _ {K} ^ {(2)} X\right) ^ {\top} \left(W _ {Q} ^ {(2)} X\right)\right) \right] _ {:, n + 1} \\ = \sum_ {s = 1} ^ {n + 1} \mathcal {P} _ {\perp} \left(W _ {V} ^ {(2)} X\right) _ {:, s} \left[ \tilde {\sigma} \left(\mathcal {C} \cdot \left(W _ {K} ^ {(2)} X\right) ^ {\top} \left(W _ {Q} ^ {(2)} X\right)\right) \right] _ {s, n + 1} \\ = \sum_ {s = 1} ^ {n + 1} \mathcal {P} _ {\perp} \left(W _ {V} ^ {(2)} X\right) _ {:, s} \exp \left(\left(\left(W _ {K} ^ {(2)} X\right) ^ {\top} \left(W _ {Q} ^ {(2)} X\right)\right) _ {s, n + 1}\right) \\ = \sum_ {s = 1} ^ {n + 1} \mathcal {P} _ {\perp} \left(W _ {V} ^ {(2)} X\right) _ {:, s} \exp \left(\left(W _ {K} ^ {(2)} X\right) _ {:, s} ^ {\top} \left(W _ {Q} ^ {(2)} X\right) _ {:, n + 1}\right) \\ = \sum_ {k = 1} ^ {K} \left(u \left(\tau_ {2 t _ {k}, d _ {k} - 1}, \tau_ {2 j, d}\right) + u \left(\tau_ {2 t _ {k} - 1, d _ {k}}, \tau_ {2 j, d}\right)\right) + \sum_ {s = 1} ^ {d} u \left(\tau_ {2 j _ {s} - 1, s}, \tau_ {2 j, d}\right) \tag {25} \\ \end{array}
$$
in which the last line is by definition of  $u(\cdot, \cdot)$  in Equation (12).
For any indices  $s, j_s, j, d$ , we can simplify the expression for  $u(\tau_{2j_s - 1,s}, \tau_{2j,d})$  by observing that
$$
\begin{array}{l} u \left(\tau_ {2 j _ {s} - 1, s}, \tau_ {2 j, d}\right) = \mathcal {P} _ {\perp} \exp \left(e \left(\tau_ {2 j _ {s} - 1, s}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} e \left(\tau_ {2 j, d}\right)\right) W _ {V} ^ {(2)} e \left(\tau_ {2 j _ {s} - 1, s}\right) \text {b y E q 1 2} \\ = \mathcal {P} _ {\perp} \exp \left(e \left(\tau_ {2 j _ {s} - 1, s}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} e \left(\tau_ {2 j, d}\right)\right) O _ {:, r (2 j _ {s} - 1, s)} \quad \text {b y} \\ = \mathcal {P} _ {\perp} \exp \left(e \left(\tau_ {2 j _ {s} - 1, s}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} e \left(\tau_ {2 j, d}\right)\right) b _ {j _ {s} D + s} \quad \text {b y} (23) \\ = \exp \left(e \left(\tau_ {2 j _ {s} - 1, s}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} e \left(\tau_ {2 j, d}\right)\right) b _ {j _ {s} D + s} \quad \text {b y} (26) \\ \end{array}
$$
Likewise by Equation (12), Equation (24), Equation (23), Equation (22)
$$
u \left(\tau_ {2 j _ {s}, s - 1}, \tau_ {2 j, d}\right) = - \exp \left(e \left(\tau_ {2 j _ {s}, s - 1}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} e \left(\tau_ {2 j, d}\right)\right) \exp \left(a _ {j _ {s}, s}\right) b _ {j _ {s} D + s} \tag {27}
$$
<sup>14</sup>When depth  $d = 0$ , all brackets are matched, the groundtruth next-token distribution is the prior distribution over the open brackets. Because in Equation (11)  $d_1, d_2 \geq 1$ , we handle the depth  $d = 0$  case separately in Case 2 "t is even,  $d = 0$ " towards the end of this proof. In the following, we focus on cases with depth  $d \geq 1$ .
By Equation (26) and Equation (27),
$$
\begin{array}{l} u \left(\tau_ {2 t _ {k}, d _ {k} - 1}, \tau_ {2 j, d}\right) + u \left(\tau_ {2 t _ {k} - 1, d _ {k}}, \tau_ {2 j, d}\right) \\ = \exp \left(\boldsymbol {e} \left(\tau_ {2 t _ {k} - 1, d _ {k}}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} \boldsymbol {e} \left(\tau_ {2 j, d}\right)\right) b _ {t _ {k} D + d _ {k}} \\ - \exp \left(\boldsymbol {e} \left(\tau_ {2 t _ {k}, d _ {k} - 1}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} \boldsymbol {e} \left(\tau_ {2 j, d}\right)\right) \exp \left(a _ {t _ {k}, d _ {k}}\right) \mathrm {b} _ {t _ {k} D + d _ {k}} \\ = \left[ \exp \left(e \left(\tau_ {2 t _ {k} - 1, d _ {k}}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} e \left(\tau_ {2 j, d}\right)\right) \right. \\ \left. - \exp \left(\boldsymbol {e} \left(\tau_ {2 t _ {k}, d _ {k} - 1}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} \boldsymbol {e} \left(\tau_ {2 j, d}\right) + a _ {t _ {k}, d _ {k}}\right) \right] \mathrm {b} _ {t _ {k} D + d _ {k}} \\ = 0 \tag {28} \\ \end{array}
$$
in which the last line is because the terms inside  $\left[\dots\right]$  cancel each other, because by Equation (19)
$$
\boldsymbol {e} \left(\tau_ {2 t _ {k} - 1, d _ {k}}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} \boldsymbol {e} \left(\tau_ {2 j, d}\right) = \boldsymbol {e} \left(\tau_ {2 t _ {k}, d _ {k} - 1}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} \boldsymbol {e} \left(\tau_ {2 j, d}\right) + a _ {t _ {k}, d _ {k}}
$$
Plugging Equation (28) and Equation (26) into Equation (25),
$$
\begin{array}{l} \tilde {a} _ {2} (X; \theta^ {(2)}) _ {:, n + 1} = \sum_ {s = 1} ^ {d} u (\tau_ {2 j _ {s} - 1, s}, \tau_ {2 j, d}) \\ = \sum_ {s = 1} ^ {d} \exp \left(e \left(\tau_ {2 j _ {s} - 1, s}\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} e \left(\tau_ {2 j, d}\right)\right) b _ {j _ {s} D + s} \tag {29} \\ \end{array}
$$
Therefore,  $\tilde{a}_2(X;\theta^{(2)})_{:,n + 1}$  lies in the span of  $\{\mathrm{b}_{j_sD + s}\}_{s\in [d]}$ . We will from now on assume  $\langle \mathrm{LN}(\tilde{a}_2(X;\theta^{(2)})_{:,n}),\mathrm{b}_{j_sD + s}\rangle >M$  for all possible choices of  $p$  ending with a closed bracket with grammar depth at least 1 for some constant  $M\in (0,1)$ . Here  $M$  exists because
$$
\langle \mathrm {L N} (\tilde {a} _ {2} (X; \theta^ {(2)}) _ {:, n}), \mathrm {b} _ {j _ {s} D + s} \rangle = \frac {\exp \left(\boldsymbol {e} \left(\tau_ {2 j _ {s} - 1 , s}\right) ^ {\top} (W _ {K} ^ {(2)}) ^ {\top} W _ {Q} ^ {(2)} \boldsymbol {e} \left(\tau_ {2 j , d}\right)\right)}{\sqrt {\sum_ {s ^ {\prime} = 1} ^ {d} \exp \left(2 \boldsymbol {e} \left(\tau_ {2 j _ {s} ^ {\prime} - 1 , s}\right) ^ {\top} (W _ {K} ^ {(2)}) ^ {\top} W _ {Q} ^ {(2)} \boldsymbol {e} \left(\tau_ {2 j , d}\right)\right)}} > 0,
$$
for all possible combination of  $j_{k}, k \in [d]$  and  $s$ , and there are only finite number of such combinations.
Constructing the projection function  $g^{(2)}$  We will finally show there exists a 6-layer MLP  $\mathrm{g}^{(2)}$  with width  $O(D^2 k^2)$ , such that for any dyck prefix  $q$  with  $n$  being the length of  $q$ ,  $X$  being the input of the second layer given  $q$  and  $\mathbb{P}(p)$  being the groundtruth next-token probability vector given  $q^{15}$ , it holds that,  $\mathrm{g}^{(2)}\left(\mathrm{LN}(\tilde{a}_2(X;\theta^{(2)})_{:,n + 1}) + X_{:,n + 1}\right) = \mathbb{P}(q)$ .
We will assume the last token of  $q$  is  $\tau_{t,d}$ . Suppose that  $\mathrm{b}_{m-1}, \mathrm{b}_m$  is an orthonormal basis of the normal space of  $\operatorname{span}\{\mathrm{b}_1,..,\mathrm{b}_{m-2}\}$ , then we can first observe that for  $U = \mathrm{b}_m\mathrm{b}_m^\top + \mathrm{b}_{m-1}\mathrm{b}_{m-1}^\top$ , it holds that
$$
U (\mathrm {L N} (\tilde {a} _ {2} (X; \theta^ {(2)}) _ {:, n + 1}) + X _ {:, n + 1}) = U e (\tau_ {t, d}).
$$
is unique for every  $t, d$ . Then based on Lemma 15, there exists a 2-layer MLP with width  $4kD$  that maps  $U(\mathrm{LN}(\tilde{a}_2(X; \theta^{(2)})_{:,n+1}) + X_{:,n+1})$  to  $(t, d)$ . This implies that there exists a 2-layer MLP with width  $4kD$  that maps  $\mathrm{LN}((\tilde{a}_2(X; \theta^{(2)})_{:,n}) + X_{:,n}$  to  $(t, d)$ .