SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {T} \left(p _ {\beta}\right) = W _ {\text {H e a d}} g ^ {(2)} \left(\frac {v + \beta \left(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)\right)}{\| v + \beta \left(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)\right) \| _ {2}} + e \left(\tau_ {2 j, d}\right)\right). \tag {13}
$$
The proof proceeds by showing a contradiction. Suppose  $u(\tau_{2j,d}, \tau_{2i,d'} - 1) + u(\tau_{2j,d}, \tau_{2i - 1,d'}) \neq 0$ . Based on the continuity of the projection function and the LayerNorm Layer, we can show that  $\lim_{\beta \to \infty} \mathcal{T}(p_{\beta})$  depend only on grammar depths  $d, d'$  and types  $2j, 2i - 1, 2i$ . However, these are not sufficient to determine the next-token probability from  $p_{\beta}$ , since the latter depends on the type of the last unmatched open bracket in  $p$ . This contradicts the assumption that the model can minimize the loss for any length  $N$ . Hence we must have
$$
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) = 0. \tag {14}
$$
Finally, as we assumed that  $\mathcal{P}_{\perp}W_V^{(2)}e(\tau_{t,d})\neq 0$  , we conclude that
$$
\left(e \left(\tau_ {2 i - 1, d ^ {\prime}}\right) - e \left(\tau_ {2 i, d ^ {\prime} - 1}\right)\right) ^ {\top} \left(W _ {K} ^ {(2)}\right) ^ {\top} W _ {Q} ^ {(2)} e \left(\tau_ {2 j + 1, d}\right) = \ln \left(\frac {\| \mathcal {P} _ {\perp} W _ {V} e \left(\tau_ {2 i , d ^ {\prime} - 1}\right) \| _ {2}}{\| \mathcal {P} _ {\perp} W _ {V} e \left(\tau_ {2 i - 1 , d ^ {\prime}}\right) \| _ {2}}\right),
$$
where the right hand side is independent of  $j, d$ , concluding the proof for necessity. The proof of sufficiency are given in Appendix C.1.
3The 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.
Note that the perfect balance condition is an orthogonal consideration to interpretability. For example, even the uniform attention satisfies the condition and can solve Dyck:
Corollary 1. There exists a 2-layer Transformer with uniform attention and no position embedding (but with causal mask and a starting token  $^5$ ) that generates the Dyck language of arbitrary length.
Since uniform attention patterns are hardly reflective of any structure of Dyck, Corollary 1 proves that attention patterns can be oblivious about the underlying task, violating the "faithfulness" criteria for an interpretation (Jain & Wallace, 2019). We will further show in Appendix B.1 that empirically, seemingly structured attention patterns may not accurately represent the inherent structure of the task.
Extension to approximate balance condition: 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. Details are deferred to Appendix C.4.
3.2 Interpretability Requires Inspecting More Than Any Single Weight Matrix
Another line of interpretability works involves inspecting the weight matrices of the model (Li et al., 2016; Dar et al., 2022; Eldan & Li, 2023). Some of the investigations are done locally, neglecting the interplay between different parts of the model. Our result in this section shows that from a representational perspective, isolating single weights can also be misleading for interpretability. For this section only, we will assume the linear head  $W_{\mathrm{Head}}$  is identity for simplicity. To consider the effect of pruning, we will also extend the parameterization of LayerNorm module (Equation (7)) as
$$
\mathrm {L N} _ {C _ {L N}} [ b ] (A) _ {:, j} = b \frac {\mathcal {P} _ {\perp} A _ {: , j}}{\operatorname* {m a x} \left\{\| \mathcal {P} _ {\perp} A _ {: , j} \| _ {2} , \epsilon \right\}} + (1 - b) A _ {:, j},
$$
which corresponds to a weighted residual branch; note that the original LayerNorm corresponds to  $\mathrm{LN}_C[1]$ . Let  $\hat{\theta}$  denote the set of parameters of this extended parameterization.
We define the nonstructural pruning as:
Definition 4 (Nonstructural pruning). Under the extended parameterization, a nonstructural pruning of a Transformer with parameter  $\hat{\theta}$  is a Transformer with the same architecture and parameter  $\hat{\theta}'$ , so that for any weight matrix  $W$  in  $\hat{\theta}$ , the corresponding matrix  $W'$  in  $\hat{\theta}'$  has  $W_{i,j}' \in \{W_{i,j}, 0\}, \forall i,j$ .
To measure the quality of the pruning, define the  $\epsilon$ -approximation:
Definition 5 ( $\epsilon$ -approximation). Given two metric spaces  $A, B$  with the same metric  $\|\cdot\|$ , a function  $f: A \to B$  is an  $\epsilon$ -approximation of function  $g$  with respect to that metric, if and only if,
$$
\forall x \in A, \| f (x) - g (x) \| \leq \epsilon \| x \|.
$$
The metric, unless otherwise specified, will be the 2-norm for vectors and the 1, 2-norm for matrices:
Definition 6. The 1, 2-norm of a matrix  $A$  is the max row norm, i.e.  $\| A \|_{1,2} = \max_{i \in [d']} \| A_{:,i} \|_2$ .
With these definitions, we are ready to state the main result of this section:
Theorem 2 (Indistinguishability From a Single Component). Consider any  $L$ -layer Transformer  $\mathcal{T}$  (Equation (10)) with embedding dimension  $m$ , attention dimension  $m_a$ , and projection function  $g$  as 2-layer ReLU MLP with width  $w$ . For any  $\delta \in (0,1)$  and  $N \in \mathbb{N}^+$ , consider a 4L-layer random Transformer  $\mathcal{T}_{large}$  with embedding dimension  $m_{\mathrm{large}} = O(m\log (Lm / \delta))$ , attention dimension  $m_{\mathrm{large},a} = O(m_aL\log \frac{m_amLN}{\epsilon\delta})$ , and projection function  $g_{\mathrm{large}}$  as 4-layer ReLU MLP with width  $w_{\mathrm{large}} = O(\max \{m,w\} L\log \frac{wmLN}{\epsilon\delta})$ .
4This is verified empirically: the uniform-attention models have attention weights fix to 0 and are to fit the distribution almost perfectly ( $>99\%$  accuracy).
Here the starting token is necessary because otherwise, the Transformer with uniform attention will have the same outputs for prefix  $p$  and prefix  $p \oplus p$ , in which  $\oplus$  denotes concatenation, i.e.  $p \oplus p$  means the same string  $p$  repeated twice.
This is as opposed to structural pruning, which prunes entire rows/columns of weight matrices.
Assume that  $\| W\| _2\leq 1$  for every weight matrix  $W$  in  $\mathcal{T}$ , and suppose the weights are randomly sampled as  $W_{i,j}\sim U(-1,1)$  for every  $W\in \mathcal{T}_{large}$ . Then, with probability  $1 - \delta$  over the randomness of  $\mathcal{T}_{large}$ , there exists a nonstructural pruning (Definition 4) of  $\mathcal{T}_{large}$ , denoted as  $\tilde{\mathcal{T}}_{large}$ , which  $\epsilon$ -approximates  $\mathcal{T}$  with respect to  $\| \cdot \|_{1,2}$  for any input  $X\in \mathbb{R}^{m\times N}$  satisfying  $\| X\|_{1,2}\leq 1$ .<sup>7</sup>
Proof sketch: connection to Lottery Tickets. Theorem 2 can be interpreted as a lottery ticket hypothesis (Frankle & Carbin, 2018; Malach et al., 2020) for randomly initialized Transformers, which can be of independent interest. The proof repeatedly uses an extension of Theorem 1 of Pensia et al. (2020), where it 1) first prunes the  $(2l - 1)$ -th and  $2l$ -th layers of  $\mathcal{T}_{\mathrm{large}}$  to approximate  $\mathcal{T}^{(l)}$  for each  $l \in [L]$  (Lemma 6), and 2) then prunes the remaining  $2L + 1$  to  $4L$ -th layers of  $\mathcal{T}_{\mathrm{large}}$  to approximate the identity function. The full proof is deferred to Appendix C.5.
Noting that the layers used to approximate the identity can appear at arbitrary depth in  $\mathcal{T}_{\text{large}}$ , a direct corollary of Theorem 2 is that one cannot distinguish between two functionally different Transformers by inspecting any single weight matrix only:
Corollary 2. Let  $\mathcal{T}_1, \mathcal{T}_2$  and  $\mathcal{T}_{\text{large}}$  follow the same definition and assumptions as  $\mathcal{T}$  and  $\mathcal{T}_{\text{large}}$  in Theorem 2. Pick any weight matrix  $W$  in  $\mathcal{T}_{\text{large}}$ , then with probability  $1 - \delta$  over the randomness of  $\mathcal{T}_{\text{large}}$ , there exist two Transformers  $\mathcal{T}_{\text{Large},1}, \mathcal{T}_{\text{Large},2}$  pruned from  $\mathcal{T}_{\text{large}}$ , such that  $\mathcal{T}_{\text{Large},i} \in \text{approximate } \mathcal{T}_i$ ,  $\forall i \in \{1,2\}$ , and  $\mathcal{T}_{\text{Large},1}, \mathcal{T}_{\text{Large},2}$  coincide on the pruned versions of  $W$ .
Hence, one should be cautious when using methods based solely on individual components to interpret the overall function of a Transformer.
Our theory in Section 3 proves the existence of abundant non-stack-like attention patterns, all of which suffice for (near-)optimal generalization on Dyck. However, could it be that stack-like solutions are more frequently discovered empirically, due to potential implicit biases in the architecture and the training procedure? In this section, we show there is no evidence for such implicit bias in standard training (Section 4.1). Additionally, we propose a regularization term based on the balance condition (Theorem 1), which leads to better length generalization (Section 4.2).
4.1 Different Attention Patterns Can Be Learned To Generate Dyck
We empirically verify our theoretical findings that Dyck solutions can give rise to a variety of attention patterns, by evaluating the accuracy of predicting the last bracket of a prefix (Equation 2) given the rest of the prefix. We only consider prefixes ending with a closing bracket, so that there exists a unique correct closing bracket which a correct parser should be able to determine. The experiments in this section are based on Transformers with 2 layers and 1 head, hidden dimension 50 and embedding dimension 50, trained using Adam. Additional results for three-layer Transformers are provided in Appendix D.3. The training data consists of valid  $\mathrm{Dyck}_{2,4}$  sequences of length less than 28 generated with  $q = 0.5$ . When tested in-distribution, all models are able to achieve  $\geq 97\%$  accuracy.
Variation in attention patterns First, as a response to (Q1), we observe that attention patterns of Transformers trained on Dyck are not always stack-like (Figure 1). In fact, the attention patterns differ even across different random initialization. Moreover, while Theorem 1 implies that position encoding is not necessary for a Transformer to generate Dyck, adding the position encoding does affect the attention patterns (Figures 1c and 1d).
Here the input and output dimension of  $\tilde{T}_{\mathrm{large}}$  is actually  $m_{\mathrm{large}}$  which is larger than  $m$ ; additional dimensions are padded with zeroes. The norm constraint can be easily extended to an arbitrary constant.
<sup>9</sup>We use the linear positional encoding following Yao et al. (2021): for the  $i_{th}$  position, the encoding is defined to be  $e_p(i) \coloneqq i / T_{\max}$  for some  $T_{\max}$ .
(a) Embedding 15, run 1
(b) Embedding 15, run 2
(c) Embedding 17
(d) Embedding 16
Figure 2: Second-layer attention patterns of two-layer Transformers with a minimal first layer: (a), (b) are based on embedding 15 with different learning rates, where the attention patterns show much variance as Theorem 1 predicts. (c), (d) are based on embedding 17 and 16. Different embedding functions lead to diverse attention patterns, most of which are not stack-like.
Figure 3: Relationship Between Balance Violation and Length Generalization. Accuracy from Transformers with minimal first layer with embedding 15, using both standard training and contrastive regularization (Equation (18)). Standard training leas to high balance violations which negatively correlate with length generalization performance. Contrastive regularization helps reduce the balance violation and improve the length generalization performance.
Specifically, for 2-layer Transformers with a minimal first layer, we experiment with three different types of embeddings  $e$ : let  $o_{t}$  denote the one-hot embedding where  $o_{t}[t] = 1$ ,
$$
\boldsymbol {e} \left(\tau_ {t, d}\right) = \boldsymbol {o} _ {(t - 1) D + d} \in \mathbb {R} ^ {2 k D}, \tag {15}
$$
$$
\boldsymbol {e} \left(\tau_ {t, d}\right) = \boldsymbol {o} _ {t} \oplus \boldsymbol {o} _ {d} \in \mathbb {R} ^ {2 k + D}, \tag {16}
$$
$$
\boldsymbol {e} \left(\tau_ {t, d}\right) = \boldsymbol {o} _ {t} \oplus \left[ \cos \left(\theta_ {d}\right), \sin \left(\theta_ {d}\right) \right] \in \mathbb {R} ^ {2 k + 2}, \theta_ {d} = \arctan \left(d / (D + 2 - d)\right), \tag {17}
$$
where  $\oplus$  denotes vector concatenation. Equation (15) is the standard one-hot embedding for  $\tau_{t,d}$ ; Equation (16) is the concatenation of one-hot embedding of types and depths. Finally, Equation (17) is the embedding constructed in Yao et al. (2021). As shown in Figure 2, the attention patterns learned by Transformers exhibit large variance between different choices of architectures and learning rates, and most learned attention patterns are not stack-like.
Quantifying the variation We now quantify the variation in attention by comparing across multiple random initializations. We define the attention variation between two attention patterns  $A_{1}, A_{2}$  as  $\mathrm{Variation}(A_{1}, A_{2}) = \|A_{1} - A_{2}\|_{F}^{2}$ , for  $A_{1}, A_{2} \in \mathbb{R}^{N \times N}$  over an length- $N$  input sequence. We report the average attention variation of each architecture based on 40 random initializations.
On the prefix  $\left[[\left[\right]\right]\left(\left(\left(\left(\right)\right)\right)\right)^{10}$ , we observe that for standard two layer training, the average attention variation is 2.20 with linear position embedding, and is 2.27 without position embedding. Both numbers are close to the random baseline value of  $2.85^{11}$ , showing that the attention head learned by different initializations indeed tend to be very different. We also experiment with Transformer with a minimal first layer and the embedding in Equation (15), where the average variation is reduced to 0.24. We hypothesize that the structural constraints in this setting provide sufficiently strong inductive bias that limit the variation.
This prefix contains brackets of all types and depths. Results with different prefixes are provided in Appendix D.3.
11The random baseline is calculated by generating purely random attention patterns (from the simplex, i.e. random square matrices s.t. each row sums up to 1) and calculate the average attention variation between them.
4.2 Guiding The Transformer To Learn Balanced Attention
In our experiments, we observe that although models learned via standard training that can generalize well in distribution, the length generalization performance is far from optimal. This implies that the models do not correctly identify the parsing algorithm for Dyck when learning from finite samples. A natural question is: can we guide Transformers towards correct algorithms, as evidenced by improved generalization performance on longer Dyck sequences?
In the following, we measure length generalization performance by the model accuracy on valid Dyck prefixes with length randomly sampled from 400 to 500, which corresponds to around 16 times the length of the training sequences. Inspired by results in Section 3, we propose a regularization term to encourage more balanced attentions, which leads to better length generalization.
Regularizing for balance violation improves length generalization accuracy We denote the balance violation of a Transformer as  $\beta \coloneqq \mathbb{E}_{d,d^{\prime},i,j}[S_{d,d^{\prime},i,j} / P_{d,j}]$  for  $S, P$  defined in Equations (31) and (33). Theorem 1 predicts that for models with a minimal first layer, perfect length generalization requires  $\beta$  to be zero. Inspired by this observation, we design a contrastive training objective to reduce the balance violation, which ideally would lead to improved length generalization. Specifically, let  $p_r$  denote a prefix of  $r$  nested pairs of brackets of for  $r \sim U([D])$ , and let  $\mathcal{T}(s \mid p_r \oplus s)$  denote the logits for  $s$  when  $\mathcal{T}$  takes as input the concatenation of  $p_r$  and  $s$ . We define the contrastive regularization term  $R_{\mathrm{contrastive}}(s)$  as the mean squared error between the logits of  $\mathcal{T}(s)$  and  $\mathcal{T}(s \mid p_r \oplus s)$ , taking expectation over  $r$  and  $p_r$ :
$$
\mathbb {E} _ {r \sim U ([ D ]), p _ {r}} \left[ \| \mathcal {T} (s \mid p _ {r} \oplus s) - \mathcal {T} (s) \| _ {F} ^ {2} \right]. \tag {18}
$$
Following the same intuition as in the proof of Theorem 1, if the model can perfectly length-generalize, then the contrastive loss will be zero. Models trained with contrastive loss show reduced balance violation as well as improved length generalization performance, as shown in Figure 3.
Why interpreting individual components sometimes leads to misconceptions? Through a case study of the Dyck grammar, we provide theoretical and empirical evidence that even in this simple and well-understood setup, Transformers can implement a rich set of non-interpretable solutions. This is reflected both by diverse attention patterns and by the absence of task-specific structures in local components. Our results directly imply similar conclusions for more complex Transformer models; see Appendix C.2 for technical details. Together, this work provides definite proof that myopic interpretability, i.e. methods based on examining individual components only, are not sufficient for understanding the functionality of a trained Transformer.
Our results do not preclude that interpretable attention patterns can emerge; however, they do suggest that interpretable patterns can be infrequent. We discuss the implications for multi-head, overparameterized Transformers trained on more complex data distributions in Appendix B. Moreover, our current results pertain to the existence of solutions; an interesting next step is to study how "inductive biases" given by the synergy of the optimization algorithm and the architecture affect the solutions found.
This work was in part supported by NSF awards IIS-2211907, CCF-2238523, and Amazon Research Award.
Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Computational Linguistics, 48(1):207-219, March 2022. doi: 10.1162/coli_a_00422. URL https://aclanthology.org/2022.cl-1.7.
Satwik Bhattachamishra, Kabir Ahuja, and Navin Goyal. On the Ability and Limitations of Transformers to Recognize Formal Languages. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 7096-7116, Online, November
Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 3770–3785, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.acl-long.292. URL https://aclanthology.org/2021.acl-long.292.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? In International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=ByxRMONtvr.
Marvin Zhang, Henrik Marklund, Abhishek Gupta, Sergey Levine, and Chelsea Finn. Adaptive risk minimization: A meta-learning approach for tackling group shift. arXiv preprint arXiv:2007.02931, 2020.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task, 2022. URL https://arxiv.org/abs/2206.04301.
Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word?, 2023.
Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks. arXiv preprint arXiv: 2306.17844, 2023.
1 Introduction 1
1.1 Related Work 3
2 Problem Setup 3
3 Theoretical Analyses 5
3.1 Interpretability Requires Inspecting More Than Attention Patterns 5
3.1.1 Perfect Balance Condition: Ideal Generalization of Unbounded Length 5
3.2 Interpretability Requires Inspecting More Than Any Single Weight Matrix 7
4 Experiments 8
4.1 Different Attention Patterns Can Be Learned To Generate Dyck 8
4.2Guiding The Transformer To Learn Balanced Attention 10
5 Conclusion 10
A Additional Related Work 18
B Are interpretable attention patterns useful? 20
B.1 Can interpretable attention patterns be misleading? 20
B.2 Are attention heads with interpretable patterns more important? 22
C Omitted Proofs in Section 3 23
C.1 Proof of Theorem 1 23
C.2 Implication of our results to larger models 26
C.3 Proof of Corollary 1 27
C.4 Approximate Balance Condition For Finite Length Training Data 27
C.4.1 Proof of Theorem 3 29
C.5 Proof of Theorem 2 31
C.5.1 Helper lemmas for Theorem 2 33
C.6 Technical Lemmas 39
C.7 Discussion on Architecture Choices 41
D Experiments 42
D.1 Training Details 42
D.2 Additional Results on Dyck Prefix 42
D.3 Extended Experiments 43
Interpreting Transformer solutions Prior empirical works show that Transformers trained on natural language data can produce representations that contain rich syntactic and semantic information, by designing a wide range of "probing" tasks (Raganato & Tiedemann, 2018; Liu et al., 2019; Hewitt & Manning, 2019; Clark et al., 2019; Tenney et al., 2019; Hewitt & Liang, 2019; Kovaleva et al., 2019; Lin et al., 2019; Wu et al., 2020; Belinkov, 2022) (or other approaches using the attention weights or parameters in neurons directly Vig & Belinkov, 2019; Htut et al., 2019; Sun & Marasovic, 2021; Eldan & Li, 2023). However, there is no canonical way to probe the model, partially due to the huge design space of probing tasks, and even a slight change in the setup may lead to very different (sometimes even seemingly contradictory) interpretations of the result (Hewitt & Liang, 2019). In this work, we tackle such ambiguity through a different perspective—by developing formal (theoretical) understanding of solutions learned by Transformers. Our results imply that it may be challenging to try to interpret Transformer solutions based on individual parameters (Li et al., 2016; Dar et al., 2022), or based on constructive proofs (unless the Transformer is specially trained to be aligned with a certain algorithm, as in Weiss et al., 2021).