Title: Transformers are uninterpretable with myopic methods: a case study with bounded Dyck grammars
Venue: NeurIPS 2023
Accepted: yes
Keywords: Transformer

Self Attention

Dyck Language

Context Free Grammar

Formal Language

Theory

Interpretability
Abstract: Transformer interpretability aims to understand the algorithm implemented by a learned Transformer by examining various aspects of the model, such as the weight matrices or the attention patterns.
In this work, through a combination of theoretical results and carefully controlled experiments on synthetic data, we take a critical view
of methods that exclusively focus on individual parts of the model, rather than consider the network as a whole.
We consider a simple synthetic setup of learning a (bounded) Dyck language. Theoretically, we show that the set of models that (exactly or approximately) solve this task satisfy a structural characterization derived from ideas in formal languages (the pumping lemma).
We use this characterization to show that the set of optima is qualitatively rich; in particular, the attention pattern of a single layer can be "nearly randomized", while preserving the functionality of the network.
We also show via extensive experiments that these constructions are not merely a theoretical artifact: even with severe constraints to the architecture of the model, vastly different solutions can be reached via standard training. Thus, interpretability claims based on inspecting individual heads or weight matrices in the Transformer can be misleading.

Transformers are uninterpretable with myopic methods: a case study with bounded Dyck grammars

Abstract

Transformer interpretability aims to understand the algorithm implemented by a learned Transformer by examining various aspects of the model, such as the weight matrices or the attention patterns. In this work, through a combination of theoretical results and carefully controlled experiments on synthetic data, we take a critical view of methods that exclusively focus on individual parts of the model, rather than consider the network as a whole. We consider a simple synthetic setup of learning a (bounded) Dyck language. Theoretically, we show that the set of models that (exactly or approximately) solve this task satisfy a structural characterization derived from ideas in formal languages (the pumping lemma). We use this characterization to show that the set of optima is qualitatively rich; in particular, the attention pattern of a single layer can be 'nearly randomized', while preserving the functionality of the network. We also show via extensive experiments that these constructions are not merely a theoretical artifact: even with severe constraints to the architecture of the model, vastly different solutions can be reached via standard training. Thus, interpretability claims based on inspecting individual heads or weight matrices in the Transformer can be misleading.

1 Introduction

Transformer-based models power many leading approaches to natural language processing. With their growing deployment in various applications, it is increasingly essential to understand the inner working of these models. Towards addressing this, there have been great advancement in the field of interpretability presenting various types of evidence (Clark et al., 2019; Vig & Belinkov, 2019; Wiegreffe & Pinter, 2019; Nanda et al., 2023; Wang et al., 2023), some of which, however, can be misleading despite being highly intuitive (Jain & Wallace, 2019; Serrano & Smith, 2019; Rogers et al., 2020; Grimsley et al., 2020; Brunner et al., 2020; Meister et al., 2021).

In this work, we aim to understand the theoretical limitation of certain interpretability methods by characterizing the set of viable solutions. We focus on myopic interpretability methods, i.e. methods based on examining individual components only. We adopt a particular toy setup in which Transformers are trained to generate Dyck grammars , a classic type of formal language grammar consisting of balanced parentheses of multiple types. Dyck is a useful sandbox, as it captures properties like long-range dependency and hierarchical tree-like structure that commonly appear in natural and programming language syntax, and has been an object of interest in many theoretical studies (Hahn, 2020; Yao et al., 2021; Liu et al., 2022b, 2023). Dyck is canonically parsed using

(a) With Position

Embedding

(b) With Position

Embedding

(c) Without Position

Embedding

(d) Without Position

Embedding

S

S

S

S

S

S

）

S

（（

）

S

）

a stack-like data structure. Such stack-like patterns (Figure 1) have been observed in the attention heads (Ebrahimi et al., 2020), which was later bolstered by mathematical analysis in Yao et al. (2021).

From a representational perspective and via explicit constructions of Transformer weights, recent work (Liu et al., 2023; Li et al., 2023) show that Transformers are sufficiently expressive to admit very different solutions that perform equally well on the training distribution. Thus, the following questions naturally arise:

(Q1) Do Transformer solutions found empirically match the theoretical constructions given in these representational results (Figure 1)? In particular, are interpretable stack-like pattern in Ebrahimi et al. (2020) the norm or the exception in practice?

(Q2) More broadly, can we understand in a principled manner the fundamental obstructions to reliably 'reverse engineering' the algorithm implemented by a Transformer by looking at individual attention patterns?

(Q3) Among models that perform (near-)optimally on the training distribution, even if we cannot fully reverse engineer the algorithm implemented by the learned solutions, can we identify properties that characterize performance beyond the training distribution?

Our contributions. We first prove several theoretical results to provide evidence for why individual components (e.g. attention patterns or weights) of a Transformer should not be expected to be interpretable. In particular, we prove:

A perfect balance condition (Theorem 1) on the attention pattern that is sufficient and necessary for 2-layer Transformers with a minimal first layer (Assumption 1) to predict optimally on Dyck of any length. We then show that this condition permits abundant non-stack-like attention patterns that do not necessarily reflect any structure of the task, including uniform attentions (Corollary 1).

An approximate balance condition (Theorem 3), the near-optimal counterpart of the condition above, for predicting on bounded -length Dyck. Likewise, non-stack-like attention patterns exist.

Indistinguishability from a single component (Theorem 2), proved via a Lottery Ticket Hypothesis style argument that any Transformer can be approximated by pruning a larger random Transformer, implying that interpretations based exclusively on local components may be unreliable.

We further accompany these theoretical findings with an extensive set of empirical investigations.

Is standard training biased towards interpretable solutions? While both stack-like and non-stack like patterns can process Dyck theoretically, the inductive biases of the architecture or the optimization process may prefer one solution over the other in practice. In Section 4.1, based on a wide range of Dyck distributions and model architecture ablations, we find that Transformers that generalize near-perfectly in-distribution (and reasonably well out-of-distribution) do not typically produce stack-like attention patterns, showing that the results reported in prior work (Ebrahimi et al., 2020) should not be expected from standard training.

Do non-interpretable solutions perform well in practice? Our theory predicts that balanced (or even uniform) attentions suffice for good in- and out-of-distribution generalization. In Section 4.2, we

empirically verify that with standard training, the extent to which attentions are balanced is positively correlated with generalization performance. Moreover, we can guide Transformers to learn more balanced attention by regularizing for the balance condition, leading to better length generalization.

1.1 Related Work

There has been a flourishing line of work on interpretability in natural language processing. Multiple 'probing' tasks have been designed to extract syntactic or semantic information from the learned representations (Raganato & Tiedemann, 2018; Liu et al., 2019; Hewitt & Manning, 2019; Clark et al., 2019). However, the effectiveness of probing often intricately depend on the architecture choices and task design, and sometimes may even result in misleading conclusions (Jain & Wallace, 2019; Serrano & Smith, 2019; Rogers et al., 2020; Brunner et al., 2020; Prasanna et al., 2020; Meister et al., 2021). While these challenges do not completely invalidate existing approaches (Wiegreffe & Pinter, 2019), it does highlight the need for more rigorous understanding of interpretability.

Towards this, we choose to focus on the synthetic setup of Dyck whose solution space is easier to characterize than natural languages, allowing us to identify a set of feasible solutions. While similar representational results have been studied in prior work (Yao et al., 2021; Liu et al., 2023; Zhao et al., 2023), our work emphasizes that theoretical constructions do not resemble the solutions found in practice. Moreover, the multiplicity of valid constructions suggest that understanding Transformer solutions require analyzing the optimization process, which a number of prior work has made progress on (Jelassi et al., 2022; Li et al., 2023; Deng et al., 2023).

Finally, it is worth noting that the challenges highlighted in our work do not contradict the line of prior work that aim to improve mechanistic interpretability into a trained model or the training process (Elhage et al., 2021; Olsson et al., 2022; Nanda et al., 2023; Chughtai et al., 2023; Li et al., 2023), which aim to develop circuit-level understanding of a particular model or the training process.

We defer discussion on additional related work to Appendix A.

2 Problem Setup

Dyck languages A Dyck language (Schützenberger, 1963) is generated by a context-free grammar, where the valid strings consist of balanced brackets of different types (for example, ' [()] ' is valid but ' ([)] ' is not). Dyck k denote the Dyck language defined on k types of brackets. The alphabet of Dyck k is denoted as [2 k ] ≡ { 1 , 2 , · · · , 2 k } , where for each type t ∈ [ k ] , tokens 2 t -1 and 2 t are a pair of corresponding open and closed brackets. Dyck languages can be recognized by a push-down automaton. For a string w and i ≤ j ∈ Z + , we use w i : j to denote the substring of w between position i and position j (both ends included). For a valid prefix w 1: i , the grammar depth of w 1: i is defined as the depth of the stack after processing w 1: i :

We overload depth( w 1: i ) to also denote the grammar depth of the bracket at position i . For example, in each pair of matching brackets, the closing bracket is one depth smaller than the open bracket. We will use τ i,d to denote a token of type i ∈ [2 k ] placed at grammar depth d ∈ N .

We consider bounded-depth Dyck languages following Yao et al. (2021). Specifically, Dyck k,D is a subset of Dyck k such that the depth of any prefix of a word is bounded by D ,

While a bounded grammar depth might seem restrictive, it suffices to capture many practical settings. For example, the level of recursion occurring in natural languages is typically bounded by a small constant (Karlsson, 2007; Jin et al., 2018). We further define the lengthN prefix set of Dyck k,D as

Our theoretical setup uses the following data distribution D q,k,D,N :

Definition 1 (Dyck distribution) . The distribution D q,k,D,N , specified by q ∈ (0 , 1) , is defined over Dyck k,D,N such that ∀ w 1: N ∈ Dyck k,D,N ,

That is, q ∈ (0 , 1) denote the probability of seeing an open bracket at the next position, except for two corner cases: 1) the next bracket has to be open if the current grammar depth is 0 (1 after seeing the open bracket); 2) the next bracket has to be closed if the current grammar depth is D .

Training Objectives. Given a model f θ parameterized by θ , we train with a next-token prediction language modeling objective on a given D q,k,D,N . Precisely, given a loss function l ( · , · ) → R , f θ is trained to minimize the loss function min θ L ( θ ; D q,k,D,N ) with

in which z ( w i ) ∈ { 0 , 1 } 2 k denotes the one-hot embedding of token w i . We will omit the distribution D q,k,D,N when it is clear from the context. We will also consider a ℓ 2 -regularized version L reg ( θ ) = L ( θ ) + λ ∥ θ ∥ 2 2 2 with parameter λ > 0 .

For our theory, we will consider the mean squared error as the loss function: 1

In our experiments, we apply the cross entropy loss following common practice.

Transformer Architecture. We consider a general formulation of Transformer in this work: the l -th layer is parameterized by θ ( l ) := { W ( l ) Q , W ( l ) K , W ( l ) V , param(g ( l ) ) } ∈ Θ , where W ( l ) K , W ( l ) Q ∈ R m a × m , and W ( l ) V ∈ R m × m are the key, query, and value matrices of the attention module; param(g ( l ) ) are parameters of a feed-forward network g ( l ) , consisting of fully connected layers, (optionally) LayerNorms and residual links. Given X ∈ R m × N , the matrix of m -dimensional features on a lengthN sequence, the l -th layer of a Transformer computes the function

where σ is the column-wise softmax operation defined as σ ( A ) i,j = exp( A i,j ) ∑ N k =1 exp( A k,j ) , C is the causal mask matrix defined as C i,j = -inf · ✶ [ i > j ] where inf denotes infinity. We call σ ( C +( W ( l ) K X ) ⊤ ( W ( l ) Q X ) ) the Attention Pattern of the Transformer layer l . LN represents columnwise LayerNorm operation, whose j th output column is defined as

Here P ⊥ denotes the projection orthogonal to the 11 ⊤ subspace 2 and C LN is called the normalizing constant for LayerNorm.

We will further define the attention output at the l -th layer as

When C LN = 0 , we will also consider the unnormalized attention output as

where ˜ σ ( A ) i,j = exp( A i,j ) and it holds by definition that LN 0 (˜ a l ( X ; θ ( l ) )) = LN 0 ( a l ( X ; θ ( l ) )) .

An L -layer Transformer T L consists of a composition of L of the above layers, along with a word embedding matrix W E ∈ R m × 2 k and a linear decoding head with weight W Head ∈ R 2 k × w . When

inputting a sequence of tokens into Transformer, we will append a starting token t S that is distinct from any token in the language at the beginning of the sequence. Let Z ∈ R 2 k × ( N +1) denote the one-hot embedding of a lengthN sequence, then T L computes for Z as

3 Theoretical Analyses

Many prior works have looked for intuitive interpretations of Transformer solutions by studying the attention patterns of particular heads or some individual components of a Transformer (Clark et al., 2019; Vig & Belinkov, 2019; Dar et al., 2022). However, we show in this section why this methodology can be insufficient even for the simple setting of Dyck. Namely, for Transformers that generalize well on Dyck (both in-distribution and out-of-distribution), neither attention patterns nor individual local components are guaranteed to encode structures specific for parsing Dyck. We further argue that the converse is also insufficient: when a Transformer does produce interpretable attention patterns, there could be limitations of such interpretation as well, as discussed in Appendix B. Together, our results provide theoretical evidence that careful analyses (beyond heuristics) are required when interpreting the components of a learned Transformer.

3.1 Interpretability Requires Inspecting More Than Attention Patterns

This section focuses on Transformers with 2 layers, which are sufficient for processing Dyck (Yao et al., 2021). We will show that even under this simplified setting, attention patterns alone are not sufficient for interpretation. In fact, we will further restrict the set of 2-layer Transformers by requiring the first-layer outputs to only depend on information necessary for processing Dyck:

Assumption 1 (Minimal First Layer) . We consider 2-layer Transformers with a minimal first layer f 1 . That is, let Z ∈ R 2 k × ( N +1) denote the one-hot embeddings of any input sequence t S , t 1 , . . . , t N ∈ [2 k ] , then the ( j +1) th column of the output f 1 ( W E Z ) only depends on the type and depth of t j , ∀ j ∈ [ N ] .

Assumption 1 requires the first layer output to depend only on the bracket type and depth, disregarding any other information such as positions; one such example is given by Yao et al. (2021). The construction of a minimal first layer can vary, hence we directly parameterize its output instead:

Definition 2 (Minimal first layer embeddings) . Given a minimal first layer, e ( τ t,d ) ∈ R m denotes its output embedding of τ t,d for t ∈ [2 k ] , d ∈ [ D ] . e ( t S ) ∈ R m is the embedding of the starting token.

It is important to note that while the minimal first layer is a strong condition, it does not weaken our results: We will show that the function class allows for a rich set of solutions, none of which are necessarily interpretable. Relaxing to more complex classes will only expand the solution set, and hence our conclusion will remain valid. See Appendix C.2 for more technical details.

3.1.1 Perfect Balance Condition: Ideal Generalization of Unbounded Length

Some prior works have tried to understand the model by inspecting the attention patterns (Ebrahimi et al., 2020; Clark et al., 2019; Vig & Belinkov, 2019). However, we will show that the attention patterns alone are too flexible to be helpful, even for the restricted class of a 2-layer Transformer with a minimal first layer (Assumption 1) and even on a language as simple as Dyck. In particular, the Transformer only needs to satisfy what we call the balanced condition :

Definition 3 (Balance condition) . A 2-layer Transformer (Equation (10) ) with a minimal first layer (Assumption 1 and Definition 2) is said to satisfy the balance condition , if for any i, j 1 , j 2 ∈ [ k ] and d ′ , d 1 , d 2 ∈ [ D ] ,

The following result shows that under minor conditions the balance condition is both necessary and sufficient:

Theorem 1 (Perfect Balance) . Consider a two-layer Transformer T (Equation (10) ) with a minimal first layer (Assumption 1) and C LN = 0 (Equation (7) ). Let O denote the optimal prediction scenario, that is, when the first layer embeddings { e ( τ i,d ) } d ∈ [ D ] ,i ∈ [2 k ] (Definition 2) and second layer parameters θ (2) satisfy

where the objective L is defined in Equation (4) . Then,

Equation (11) a necessary condition of O , if W (2) V satisfies P ⊥ W (2) V e ( τ t,d ) = 0 , ∀ t ∈ [ k ] , d ∈ [ D ] .

̸

Equation (11) is a sufficient condition of O , if the set of 2 k + 1 encodings { e ( τ 2 i -1 ,d ) , e ( τ 2 i,d ) } i ∈ [ k ] ∪ { e ( t S ) } are linearly independent for any d ∈ [ D ] , and the projection function g (2) is a 6-layer MLP 3 with O ( k 2 D 2 ) width.

Remark : Recall from Equation (7) that P ⊥ projects to the subspace orthogonal to 11 ⊤ . The assumption in the necessary condition can be intuitively understood as requiring all tokens to have nonzero contributions to the prediction after the LayerNorm.

Recall that e ( τ 2 i -1 ,d ′ ) , e ( τ 2 i,d ′ -1 ) denote the first-layer outputs for a matching pair of brackets. Intuitively, Equation (11) says that since matching brackets should not affect future predictions, their embeddings should balance out each other. The balance condition Equation (11) is 'perfect' in the sense that the theory assumes the model can minimize the loss for any length N ; we will see an approximate version later in Theorem 3.

Proof of the necessity of the balance condition. The key idea is reminiscent of the pumping lemma for regular languages. For any prefix p ending with a closed bracket τ 2 j,d for d ≥ 1 and containing brackets of all depths in [ D ] , let p β be the prefix obtained by inserting β pairs of { τ 2 i -1 ,d ′ , τ 2 i,d ′ -1 } for arbitrary i ∈ [ k ] and d ′ ∈ [ D ] . Denote the projection of the unnormalized attention output by

We ignored the normalization in softmax above, since the attention output will be normalized directly by LayerNorm according to Equation (6).

By Equation (10), there exists a vector v ∈ R m such that for any β ∈ N , the next-token logits given by Transformer T are

̸

The proof proceeds by showing a contradiction. Suppose u ( τ 2 j,d , τ 2 i,d ′ -1 ) + u ( τ 2 j,d , τ 2 i -1 ,d ′ ) = 0 . Based on the continuity of the projection function and the LayerNorm Layer, we can show that lim β →∞ T ( p β ) depend only on grammar depths d, d ′ and types 2 j, 2 i -1 , 2 i . However, these are not sufficient to determine the next-token probability from p β , 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

̸

Finally, as we assumed that P ⊥ W (2) V e ( τ t,d ) = 0 , we conclude that

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.

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: 4

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 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 Head is identity for simplicity. To consider the effect of pruning, we will also extend the parameterization of LayerNorm module (Equation (7)) as

which corresponds to a weighted residual branch; note that the original LayerNorm corresponds to LN C [1] . Let ˆ θ denote the set of parameters of this extended parameterization.

We define the nonstructural pruning 6 as:

Definition 4 (Nonstructural pruning) . Under the extended parameterization, a nonstructural pruning of a Transformer with parameter ˆ θ is a Transformer with the same architecture and parameter ˆ θ ′ , so that for any weight matrix W in ˆ θ , the corresponding matrix W ′ in ˆ θ ′ has W ′ i,j ∈ { W i,j , 0 } , ∀ i, j .

To measure the quality of the pruning, define the ϵ -approximation:

Definition 5 ( ϵ -approximation) . Given two metric spaces A,B with the same metric ∥ · ∥ , a function f : A → B is an ϵ -approximation of function g with respect to that metric, if and only if,

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 ∈ [ 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 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 δ ∈ (0 , 1) and N ∈ N + , consider a 4 L -layer random Transformer T large with embedding dimension m large = O ( m log( Lm/δ )) , attention dimension m large ,a = O ( m a L log m a mLN ϵδ ) , and projection function g large as 4-layer ReLU MLP with width w large = O (max { m,w } L log wmLN ϵδ ) .

Assume that ∥ W ∥ 2 ≤ 1 for every weight matrix W in T , and suppose the weights are randomly sampled as W i,j ∼ U ( -1 , 1) for every W ∈ T large. Then, with probability 1 -δ over the randomness of T large, there exists a nonstructural pruning (Definition 4) of T large, denoted as ˜ T large, which ϵ -approximates T with respect to ∥ · ∥ 1 , 2 for any input X ∈ R m × N satisfying ∥ X ∥ 1 , 2 ≤ 1 . 7

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 (2 l -1) -th and 2 l -th layers of T large to approximate T ( l ) for each l ∈ [ L ] (Lemma 6), and 2) then prunes the remaining 2 L +1 to 4 L -th layers of T 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 T 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 T 1 , T 2 and T large follow the same definition and assumptions as T and T large in Theorem 2. Pick any weight matrix W in T large, then with probability 1 -δ over the randomness of T large , there exist two Transformers T Large , 1 , T Large , 2 pruned from T large, such that T Large ,i ϵ -approximate T i , ∀ i ∈ { 1 , 2 } , and T Large , 1 , T 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.

4 Experiments

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 Dyck 2 , 4 sequences of length less than 28 generated with q = 0 . 5 . When tested in-distribution, all models are able to achieve ≥ 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, 8 adding the position encoding 9 does affect the attention patterns (Figures 1c and 1d).

(a) Embedding 15, run 1

(b) Embedding 15, run 2

(c) Embedding 17

(d) Embedding 16

S

S

S

S

S

S

S

S

）

0.825

0.850

0.875

0.900

0.925

0.950

0.975

1.000

Length Generalization Acc

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

Balance Violation

Contrastive

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 ,

where ⊕ denotes vector concatenation. Equation (15) is the standard one-hot embedding for τ 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 Variation( A 1 , A 2 ) = ∥ A 1 -A 2 ∥ 2 F , for A 1 , A 2 ∈ R N × N over an lengthN input sequence. We report the average attention variation of each architecture based on 40 random initializations.

On the prefix [[[[]]]](((()))) 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.

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 β := E d,d ′ ,i,j [ S d,d ′ ,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 β 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 ∼ U ([ D ]) , and let T ( s | p r ⊕ s ) denote the logits for s when T takes as input the concatenation of p r and s . We define the contrastive regularization term R contrastive ( s ) as the mean squared error between the logits of T ( s ) and T ( s | p r ⊕ s ) , taking expectation over r and p r :

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.

5 Conclusion

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.

References

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 Bhattamishra, 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

Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of transformers and its implications in sequence modeling. In Proceedings of the 24th Conference on Computational Natural Language Learning , pp. 455-475, Online, November 2020b. Association for Computational Linguistics. doi: 10.18653/v1/2020.conll-1.37. URL https: //aclanthology.org/2020.conll-1.37 .

Tolga Bolukbasi, Adam Pearce, Ann Yuan, Andy Coenen, Emily Reif, Fernanda Viégas, and Martin Wattenberg. An interpretability illusion for bert. arXiv preprint arXiv: Arxiv-2104.07143 , 2021.

Gino Brunner, Yang Liu, Damian Pascual, Oliver Richter, Massimiliano Ciaramita, and Roger Wattenhofer. On identifiability in transformers. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id=BJg1f6EFDB .

Nick Cammarata, Shan Carter, Gabriel Goh, Chris Olah, Michael Petrov, Ludwig Schubert, Chelsea Voss, Ben Egan, and Swee Kiat Lim. Thread: Circuits. Distill , 2020. doi: 10.23915/distill.00024. https://distill.pub/2020/circuits.

Chughtai, Lawrence Chan, and Neel Nanda. A toy model of universality: Reverse engineering how networks learn group operations. International Conference on Machine Learning , 2023. doi: 10.48550/arXiv.2302.03025.

Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? an analysis of BERT's attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP , pp. 276-286, Florence, Italy, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W19-4828. URL https: //aclanthology.org/W19-4828 .

Guy Dar, Mor Geva, Ankit Gupta, and Jonathan Berant. Analyzing transformers in embedding space, 2022.

Yichuan Deng, Zhihang Li, and Zhao Song. Attention scheme inspired softmax regression, 2023.

Javid Ebrahimi, Dhruv Gelda, and Wei Zhang. How can self-attention networks recognize Dyck-n languages? In Findings of the Association for Computational Linguistics: EMNLP 2020 , pp. 4301-4306, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/ 2020.findings-emnlp.384. URL https://aclanthology.org/2020.findings-emnlp.384 .

Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pp. 5793-5831. PMLR, 17-23 Jul 2022. URL https://proceedings.mlr.press/v162/edelman22a.html .

Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english?, 2023.

Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread , 2021. https://Transformer-circuits.pub/2021/framework/index.html.

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412 , 2020.

Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. International Conference On Learning Representations , 2018.

Yeqi Gao, Sridhar Mahadevan, and Zhao Song. An over-parameterized exponential regression, 2023.

Gers and J. Schmidhuber. Lstm recurrent networks learn simple context-free and context-sensitive languages. IEEE transactions on neural networks , 12 6:1333-40, 2001.

Christopher Grimsley, Elijah Mayfield, and Julia R.S. Bursten. Why attention is not explanation: Surgical intervention and causal reasoning about neural models. In Proceedings of the Twelfth Language Resources and Evaluation Conference , pp. 1780-1790, Marseille, France, May 2020. European Language Resources Association. ISBN 979-10-95546-34-4. URL https://aclanthology.org/2020.lrec-1.220 .

Jonathan Haab, Nicolas Deutschmann, and María Rodríguez Martínez. Is attention interpretation? a quantitative assessment on sets. In Machine Learning and Principles and Practice of Knowledge Discovery in Databases , pp. 303-321, Cham, 2023. Springer Nature Switzerland. ISBN 978-3031-23618-1.

Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Trans. Assoc. Comput. Linguistics , 8:156-171, 2020. doi: 10.1162/tacl\_a\_00306. URL https://doi.org/ 10.1162/tacl_a_00306 .

John Hewitt and Percy Liang. Designing and interpreting probes with control tasks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pp. 27332743, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19-1275. URL https://aclanthology.org/D19-1275 .

John Hewitt and Christopher D Manning. A structural probe for finding syntax in word representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pp. 4129-4138, 2019.

John Hewitt, Michael Hahn, Surya Ganguli, Percy Liang, and Christopher D Manning. Rnns can generate bounded hierarchical languages with optimal memory. arXiv preprint arXiv:2010.07515 , 2020.

Phu Mon Htut, Jason Phang, Shikha Bordia, and Samuel R. Bowman. Do attention heads in bert track syntactic dependencies?, 2019.

Sarthak Jain and Byron C. Wallace. Attention is not Explanation. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pp. 3543-3556, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-1357. URL https://aclanthology.org/N19-1357 .

Samy Jelassi, Michael Eli Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id= eMW9AkXaREI .

Lifeng Jin, Finale Doshi-Velez, Timothy Miller, William Schuler, and Lane Schwartz. Unsupervised grammar induction with depth-bounded pcfg. Transactions of the Association for Computational Linguistics , 6:211-224, 2018.

Fred Karlsson. Constraints on multiple center-embedding of clauses. Journal of Linguistics , 43(2): 365-392, 2007.

Goro Kobayashi, Tatsuki Kuribayashi, Sho Yokoi, and Kentaro Inui. Attention is not only a weight: Analyzing transformers with vector norms. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pp. 7057-7075, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.emnlp-main.574. URL https: //aclanthology.org/2020.emnlp-main.574 .

Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the dark secrets of BERT. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language

Processing and the 9th International Joint Conference on Natural Language Processing (EMNLPIJCNLP) , pp. 4365-4374, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19-1445. URL https://aclanthology.org/D19-1445 .

Jiwei Li, Xinlei Chen, Eduard Hovy, and Dan Jurafsky. Visualizing and understanding neural models in NLP. In Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pp. 681-691, San Diego, California, June 2016. Association for Computational Linguistics. doi: 10.18653/v1/ N16-1082. URL https://aclanthology.org/N16-1082 .

Xian Li and Hongyu Gong. Robust optimization for multilingual translation with imbalanced data. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems , volume 34, pp. 25086-25099. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper/2021/file/ d324a0cc02881779dcda44a675fdcaaa-Paper.pdf .

Yuchen Li and Andrej Risteski. The limitations of limited context for constituency parsing. 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. 2675-2687, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.acl-long.208. URL https://aclanthology.org/2021.acl-long.208 .

Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pp. 19689-19729. PMLR, 23-29 Jul 2023. URL https://proceedings.mlr.press/v202/ li23p.html .

Yongjie Lin, Yi Chern Tan, and Robert Frank. Open sesame: Getting inside BERT's linguistic knowledge. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP , pp. 241-253, Florence, Italy, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W19-4825. URL https://aclanthology.org/W19-4825 .

Bingbin Liu, Daniel Hsu, Pradeep Kumar Ravikumar, and Andrej Risteski. Masked prediction: A parameter identifiability view. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems , 2022a. URL https://openreview.net/forum?id=Hbvlb4D1aFC .

Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=De4FYqjFueZ .

Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. arXiv preprint arXiv:2210.14199 , 2022b.

Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pp. 5747-5763, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.emnlp-main.463. URL https://aclanthology.org/2020.emnlp-main.463 .

Nelson F. Liu, Matt Gardner, Yonatan Belinkov, Matthew E. Peters, and Noah A. Smith. Linguistic knowledge and transferability of contextual representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pp. 1073-1094, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-1112. URL https: //aclanthology.org/N19-1112 .

Eran Malach, Gilad Yehudai, Shai Shalev-Schwartz, and Ohad Shamir. Proving the lottery ticket hypothesis: Pruning is all you need. In International Conference on Machine Learning , pp. 6682-6691. PMLR, 2020.

Clara Meister, Stefan Lazov, Isabelle Augenstein, and Ryan Cotterell. Is sparse attention more interpretable? In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 2: Short Papers) , pp. 122-129, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.acl-short.17. URL https://aclanthology.org/2021.acl-short.17 .

William Merrill. Sequential neural networks as automata. In Proceedings of the Workshop on Deep Learning and Formal Languages: Building Bridges , pp. 1-13, Florence, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W19-3901. URL https://www.aclweb.org/ anthology/W19-3901 .

Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/file/ 2c601ad9d2ff9bc8b282670cdd54f69f-Paper.pdf .

Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=9XFSbDPmdW .

Toan Q. Nguyen and Julian Salazar. Transformers without tears: Improving the normalization of self-attention. In Proceedings of the 16th International Conference on Spoken Language Translation , Hong Kong, November 2-3 2019. Association for Computational Linguistics. URL https://aclanthology.org/2019.iwslt-1.17 .

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread , 2022. https://Transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html.

Ankit Pensia, Shashank Rajput, Alliot Nagle, Harit Vishwakarma, and Dimitris Papailiopoulos. Optimal lottery tickets via subset sum: Logarithmic over-parameterization is sufficient. Advances in neural information processing systems , 33:2599-2610, 2020.

Jorge Perez, Pablo Barcelo, and Javier Marinkovic. Attention is turing-complete. Journal of Machine Learning Research , 22(75):1-35, 2021. URL http://jmlr.org/papers/v22/20-302.html .

Sai Prasanna, Anna Rogers, and Anna Rumshisky. When BERT Plays the Lottery, All Tickets Are Winning. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pp. 3208-3229, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.emnlp-main.259. URL https://aclanthology.org/2020. emnlp-main.259 .

Alessandro Raganato and Jörg Tiedemann. An analysis of encoder representations in transformerbased machine translation. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP , pp. 287-297, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi: 10.18653/v1/W18-5431. URL https://aclanthology.org/W18-5431 .

Anna Rogers, Olga Kovaleva, and Anna Rumshisky. A primer in bertology: What we know about how bert works. Transactions of the Association for Computational Linguistics , 8:842-866, 2020.

M.P. Schützenberger. On context-free languages and push-down automata. Information and Control , 6(3):246-264, 1963. ISSN 0019-9958. doi: https://doi.org/10.1016/S0019-9958(63)90306-1. URL https://www.sciencedirect.com/science/article/pii/S0019995863903061 .

Sofia Serrano and Noah A. Smith. Is attention interpretable? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , pp. 2931-2951, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/P19-1282. URL https: //aclanthology.org/P19-1282 .

Hava T. Siegelmann and Eduardo D. Sontag. On the computational power of neural nets. In Proceedings of the Fifth Annual Workshop on Computational Learning Theory , COLT '92, pp. 440-449, New York, NY, USA, 1992. Association for Computing Machinery. ISBN 089791497X. doi: 10.1145/130385.130432. URL https://doi.org/10.1145/130385.130432 .

Kaiser Sun and Ana Marasovi´ c. Effective attention sheds light on interpretability. In Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021 , pp. 4126-4135, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.findings-acl.361. URL https://aclanthology.org/2021.findings-acl.361 .

Mirac Suzgun, Yonatan Belinkov, Stuart Shieber, and Sebastian Gehrmann. LSTM networks can perform dynamic counting. In Proceedings of the Workshop on Deep Learning and Formal Languages: Building Bridges , pp. 44-54, Florence, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W19-3905. URL https://www.aclweb.org/anthology/W19-3905 .

Ian Tenney, Dipanjan Das, and Ellie Pavlick. Bert rediscovers the classical nlp pipeline. arXiv preprint arXiv:1905.05950 , 2019.

Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in a transformer language model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP , pp. 63-76, Florence, Italy, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W19-4808. URL https://aclanthology.org/W19-4808 .

Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , pp. 5797-5808, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/P19-1580. URL https://aclanthology.org/P19-1580 .

Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. In International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id= NpsVSN6o4ul .

Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers, 2021. URL https://arxiv.org/abs/2107. 13163 .

Gail Weiss, Yoav Goldberg, and Eran Yahav. On the practical computational power of finite precision rnns for language recognition. arXiv preprint arXiv:1805.04908 , 2018.

Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pp. 11080-11090. PMLR, 18-24 Jul 2021. URL https://proceedings.mlr.press/v139/weiss21a.html .

Sarah Wiegreffe and Yuval Pinter. Attention is not not explanation. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pp. 11-20, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19-1002. URL https://aclanthology.org/D19-1002 .

Zhiyong Wu, Yun Chen, Ben Kao, and Qun Liu. Perturbed masking: Parameter-free probing for analyzing and interpreting BERT. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics , pp. 4166-4176, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.acl-main.383. URL https://aclanthology.org/2020. acl-main.383 .

Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yan Liu. On layer normalization in the transformer architecture. In Proceedings of the 37th International Conference on Machine Learning , ICML'20. JMLR.org, 2020.

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=ByxRM0Ntvr .

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.

Contents

A Additional Related Work

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 & Marasovi´ c, 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).

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, Wiegreffe & 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 (Bhattamishra 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 & Marasovi´ c, 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 T (Equation (10) ). For any W ( l ) K , W ( l ) Q ∈ R m a × m ( l ∈ [ L ]) , there exist W Head ∈ R 2 k × w and b Head ∈ R 2 k such that T ( Z ) = 0 , ∀Z .

While its proof is trivial (simply setting W Head = 0 and b 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 .

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).

sentenceolayeroheado

START (b (b b)b) (a a)(a (b b)a) (a (aa) a) (b (a a) (b b)b) (a (a a) a)

(bSTART

b

b

b

e

b

b

START

(b

a)

a)

a)

(b

(a

a)

a)

(a) layer 1 of 4

(b) layer 3 of 4

START (b (b b) b) (a a) (a (b b) a) (a (a a) a) (b (a a) (b b) b) (a (a a) a)

sentence0 layer0 heado

START (b (b b) b) (a a) (a (b b) a) (a (a a) a) (b (a a) (b b) b) (a (a a) a)

sentence0 layer2 heado

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

C.1 Proof of Theorem 1

The key step is already shown in Section 3. We will restate the proof rigorously here.

Theorem 1 (Perfect Balance) . Consider a two-layer Transformer T (Equation (10) ) with a minimal first layer (Assumption 1) and C LN = 0 (Equation (7) ). Let O denote the optimal prediction scenario, that is, when the first layer embeddings { e ( τ i,d ) } d ∈ [ D ] ,i ∈ [2 k ] (Definition 2) and second layer parameters θ (2) satisfy

where the objective L is defined in Equation (4) . Then,

Equation (11) a necessary condition of O , if W (2) V satisfies P ⊥ W (2) V e ( τ t,d ) = 0 , ∀ t ∈ [ k ] , d ∈ [ D ] .

̸

Equation (11) is a sufficient condition of O , if the set of 2 k + 1 encodings { e ( τ 2 i -1 ,d ) , e ( τ 2 i,d ) } i ∈ [ k ] ∪ { e ( t S ) } are linearly independent for any d ∈ [ D ] , and the projection function g (2) is a 6-layer MLP 12 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 ( τ t,d ) as m .

For any i ∈ [ k ] , d ′ ∈ [ D ] , by Equation (11), we can assume that there exists a i,d ′ ∈ R such that for any j ∈ [ k ] , d ∈ [ D ] , it holds that,

We will first define the possible index sets of τ t,d as I = { (2 t, d ) | t ∈ [ k ] , 0 ≤ d ≤ D -1 } ∪ { (2 t -1 , d ) | t ∈ [ k ] , 1 ≤ d ≤ D } , and we will define the rank of ( t, d ) as

Then it is clear that r ( t, d ) is a one-to-one mapping from I to [2 kD ] . We will then define the collection of all e ( τ t,d ) as E , satisfying that E : ,r ( t,d ) = e ( τ t,d ) , E : , 2 kD +1 = e ( t S ) .

̸

Because e ( τ t,d ) are linearly independent, for any ( i, d ) = ( j, d ′ ) ∈ I , it holds that e ( τ i,d ) -e ( τ j,d ′ ) = 0 . Then based on Lemma 16, there exists a set of orthonormal vectors { b i } i ∈ [ m -2] , such that for any ( i, d ) , ( j, d ′ ) ∈ I , it holds that

̸

̸

We will further construct matrix O as 13

for t ∈ [ k ] , d ∈ [ D ] .

We can then choose W (2) V ∈ R m × m such that

Such W (2) V 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 . 14

For any valid Dyck prefix p of length n ending with a closed bracket τ 2 j,d satisfying d ≥ 1 , suppose the list of unmatched open brackets in p is [ τ 2 j 1 -1 , 1 , τ 2 j 2 -1 , 2 , . . . , τ 2 j d -1 ,d ] . Then, the remaining tokens in p are pairs of matching brackets. Denote them by τ 2 t k -1 ,d k , τ 2 t k ,d k -1 for k ∈ [ K ] . Then the input of the second layer of Transformer X , up to a permutation is

We will focus on the last column of the unnormalized attention output

in which the last line is by definition of u ( · , · ) in Equation (12).

For any indices s, j s , j, d , we can simplify the expression for u ( τ 2 j s -1 ,s , τ 2 j,d ) by observing that

Likewise by Equation (12), Equation (24), Equation (23), Equation (22)

By Equation (26) and Equation (27),

in which the last line is because the terms inside [ · · · ] cancel each other, because by Equation (19)

Plugging Equation (28) and Equation (26) into Equation (25),

Therefore, ˜ a 2 ( X ; θ (2) ) : ,n +1 lies in the span of { b j s D + s } s ∈ [ d ] . We will from now on assume ⟨ LN(˜ a 2 ( X ; θ (2) ) : ,n ) , b j s D + s ⟩ > M for all possible choices of p ending with a closed bracket with grammar depth at least 1 for some constant M ∈ (0 , 1) . Here M exists because

for all possible combination of j k , k ∈ [ 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 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 P ( p ) being the groundtruth next-token probability vector given q 15 , it holds that, g (2) ( LN(˜ a 2 ( X ; θ (2) ) : ,n +1 ) + X : ,n +1 ) = P ( q ) .

We will assume the last token of q is τ t,d . Suppose that b m -1 , b m is an orthonormal basis of the normal space of span { b 1 , .., b m -2 } , then we can first observe that for U = b m b ⊤ m +b m -1 b ⊤ m -1 , it holds that

is unique for every t, d . Then based on Lemma 15, there exists a 2-layer MLP with width 4 kD that maps U (LN(˜ a 2 ( X ; θ (2) ) : ,n +1 ) + X : ,n +1 ) to ( t, d ) . This implies that there exists a 2-layer MLP with width 4 kD that maps LN((˜ a 2 ( X ; θ (2) ) : ,n ) + X : ,n to ( t, d ) .

Further, let matrix U ′ = ∑ Dk j =1 o j b ⊤ j 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 ≥ 1 , based on Equation (29) and the definition of M ,

Then based on Lemma 18, there exists 2-layer MLP with width kD that operates on ( U ′ (LN(˜ a 2 ( X ; θ (2) ) : ,n +1 ) + X : ,n +1 ) t ′ D + d ′ ) t ′ ∈ [ k ] for a fixed d ′ and output the nonzero index in it, if such index exists. Hence, we can choose the weight of the first and second layer of g (2) , such that the output of the second layer is ( t, d ) ⊕ x , where 2 x d ′ -1 is the type of the unmatched open brackets with grammar depth d ′ if t is an even number, d ≥ d ′ ≥ 1 .

Now based on Lemma 17, we can choose the third and fourth layer of g (2) to perform indexing and let the output of the fourth layer be ( t, d, y ) , where y = x d when d ≥ 1 . 16 Notice that this triplet contains all the necessary information to infer P ( q ) because it uniquely determines the type of last unmatched open bracket,

If t is odd (i.e. the last bracket is open), and then the type of last unmatched open bracket is t .

If t is even and d = 0 , then all the brackets is matched.

If t is even and d ≥ 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,

the output of the constructed 4 layers also has a bounded norm. Hence, we can assume there exists constant M ′ > 1 , such that y ≤ M ′ . Now we will discuss by the value of t ,

t is odd, then one can neglect the third dimension and the correct probability is determined by d and can be represented by a width2 D network based on Lemma 15.

t is even. When d = 0 , one can construct a width1 network mapping any y to the correct probability distribution as it is unique. When d ≥ 1 , one can construct a width2 K network mapping x d ∈ [ K ] to the correct probability distribution based on Lemma 15. Then by Lemma 19, one can construct a width4 KD network that maps ( d, y ) to the corresponding probability distribution.

Putting together and using Lemma 19 again, one can construct a width8 K 2 D 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

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.

C.3 Proof of Corollary 1

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 ∈ [2 k ] is encoded as o t and the starting token is encoded as o 2 k +1 . Then it holds that

Then we can choose W (1) V such that for x ∈ R 2 k +1 ,

Hence it holds That

It is then easy to check LN ( [ W (1) V Zσ ( C · ( W (1) K Z ) ⊤ ( W (1) Q Z ) ))] : ,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 LN ( [ W (1) V Zσ ( C · ( W (1) K Z ) ⊤ ( W (1) Q Z ) ))] : ,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 (2) K = W (2) Q = 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 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 ′ ,i,j is defined as

which measures how much a matching pair of brackets ( τ 2 i,d ′ -1 , τ 2 i -1 ,d ′ ) changes the input to the LayerNorm upon seeing the last token τ 2 j,d . Note that under the perfect balance condition, S d,d ′ ,i,j [ θ (2) ] = 0 .

The second quantity P d,j [ θ (2) ] is defined via an intermediate quantity Q (2 j, d, ˜ t ) : for any i ∈ [ k ] , d ∈ [ D ] and a length-( d -1) prefix ˜ t ∈ [2 k ] d -1 , Q ( i, d, ˜ t ) is defined as

where ˜ t d ′ denotes the d ′ th entry of ˜ t . Intuitively, Q ( i, d, ˜ t ) denotes the unnormalized second-layer attention output at the last position, given the input sequence ˜ t ⊕ τ 2 i -1 ,d τ 2 i,d -1 , 18

For results in this subsection, it suffices to consider prefixes consisting only of open brackets. Let t := arg min ˜ t ∈{ 2 i -1 } d -1 i ∈ [ k ] ∥ Q (2 j, d, ˜ t ) ∥ 2 , and let t ′ denote the prefix that minimizes ∥ Q (2 j, d, ˜ t ) ∥ 2 subject to the constraint that t ′ differs from t at the last (i.e. ( d -1) th ) position, i.e.

̸

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

In the following theorem, P d,j will be used as a reference to upper bound S d,d ′ ,i,j [ θ (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 T (Equation (10) ) with a minimal first layer (Assumption 1) and a γ -Lipschitz g (2) for γ > 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 ¯ θ (2) N satisfies L ( T [ ¯ θ (2) N ] , D q,k,D,N ) ≤ q -N ϵ , for any positive integer N > 8 D and sufficiently small ϵ > 0 . Then, there exists a constant C γ,ϵ,D , such that ∀ 0 ≤ d ′ ≤ D, 1 ≤ d ≤ D,i, j ∈ [ k ] , it holds that

Intuitively, Theorem 3 states that when the loss L ( θ ) is sufficiently small, S d,d ′ ,i,j [ θ (2) ] must be small relative to P d,j [ ¯ θ (2) N ] . Inequality 34 can be interpreted as a relaxation of Equation (14), which is equivalent to S d,d ′ ,i,j [ θ (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 L reg λ = L θ ( x ) + λ ∥ θ ∥ 2 2 2 , which is the squared loss with λ weight decay. Suppose g (1) of the Transformer is a 2 -layer fully connected network and 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 θ λ,N minimizes L reg λ , then it holds ∀ 0 ≤ d ′ ≤ D, 1 ≤ d ≤ D,i, j ∈ [ k ] that,

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 t := arg min ˜ t ∈ [ k ] d -1 ∥ Q (2 j, d, ˜ t ) ∥ 2 , and let t ′ denote the prefix that minimizes ∥ Q (2 j, d, ˜ t ) ∥ 2 subject to the constraint that t ′ must differ from t in the last (i.e. ( d -1) th ) position, i.e.

̸

The motivation for such choices of t , t ′ 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 [ ¯ θ (2) ] = ∥ Q (2 j, d, t ′ ) ∥ .

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 ′ ∈ [ D ] , i, j ∈ [ k ] , consider a common prefix p = τ 2 i -1 . . . τ 2 i -1 ︸ ︷︷ ︸ d ′ open brackets τ 2 i -1 τ 2 i . . . τ 2 i -1 τ 2 i ︸ ︷︷ ︸ ( ⌊ N 2 ⌋-d ′ -d -1) pairs τ 2 i . . . τ 2 i ︸ ︷︷ ︸ d ′ closed brackets , where τ i denotes a token with type i whose depth is implicit from the context. Set p 1 , p 2 as

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 T with second layer parameters ¯ θ (2) N , its outputs on p 1 , p 2 satisfy

where TV ( p 1 , p 2 ) denotes the TV distance in the next-word distributions from p 1 and p 2 , and o ϵ (1) means the term will go to zero as ϵ goes to zero. The TV distance is lower bounded by the construction of p 1 , p 2 , where t , t ′ differ at the last open bracket. The error ϵ is upper bounded because of the assumption on ¯ θ (2) N , i.e. L ( T [ ¯ θ (2) N ] , D q,k,D,N ) ≤ q -N ϵ with sufficiently small ϵ .

Define by A p the contribution of p to the attention output (before LayerNorm) of the last position of p 1 , p 2 :

The attention outputs (before LayerNorm) of p 1 , p 2 , denoted by A ( p 1 ) and A ( p 2 ) , satisfy that

Note that for any prefix p ′ ,

where g (2) is γ -Lipschitz. Hence the Lipschitz constant with respect to and we have

We show that A p should not be too much larger in norm than Q (2 j, d, t ) or Q (2 j, d, t ′ ) . First, let's state a helper lemma about the contrapositive:

Lemma 1. For any ϵ > 0 , there exists a constant R ϵ , such that for any a, b ∈ R d and any r ∈ R d such that ∥ r ∥ 2 ≥ R ϵ · max {∥ a ∥ 2 , ∥ b ∥ 2 } , it holds that

Proof. Denote r 0 := max {∥ a ∥ 2 , ∥ b ∥ 2 } . Then R ϵ := 4 r 0 ϵ +1 suffices:

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 [ ¯ θ (2) N ] be defined as in Equation (33) and let R ϵ be the constant in Lemma 1, we need to bound ∥P ⊥ A p ∥ by

As Equation (45) holds for p with any d, d ′ , if one choose d ′ = 1 , this shows

Further, it holds that for any 1 < d ′ ≤ d -1 ,

The triangle inequality then yields,

Because N ≥ 8 D , we have that ⌊ N -2 d ′ -2 d -2 2 ⌋ ≥ N 8 , hence it holds that

Combined with Equation (46), one can conclude that,

The proof is then completed.

Proof of Corollary 3. This proof is in fact a direct combination of Theorems 1 and 3. By Theorem 1 we know there exists a weight θ (2) ∗ that can reach zero loss for arbitrarily length N . Then it holds that ∥ θ λ,N ∥ 2 ≤ ∥ θ (2) ∗ ∥ as θ λ,N minimizes the regularized loss. Noticing that bounded weight implies bounded Lipschitzness of g (2) , the rest follows as Theorem 3.

C.5 Proof of Theorem 2

We now show the limitation of interpretability from a single component, using a Lottery-Ticket-style argument by pruning from large random Transformers.

Theorem 2 (Indistinguishability From a Single Component) . Consider any L -layer Transformer 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 δ ∈ (0 , 1) and N ∈ N + , consider a 4 L -layer random Transformer T large with embedding dimension m large = O ( m log( Lm/δ )) , attention dimension m large ,a = O ( m a L log m a mLN ϵδ ) , and projection function g large as 4-layer ReLU MLP with width w large = O (max { m,w } L log wmLN ϵδ ) .

Assume that ∥ W ∥ 2 ≤ 1 for every weight matrix W in T , and suppose the weights are randomly sampled as W i,j ∼ U ( -1 , 1) for every W ∈ T large. Then, with probability 1 -δ over the randomness of T large, there exists a nonstructural pruning (Definition 4) of T large, denoted as ˜ T large, which ϵ -approximates T with respect to ∥ · ∥ 1 , 2 for any input X ∈ R m × N satisfying ∥ X ∥ 1 , 2 ≤ 1 . 19

Proof. We will first introduce some notation. For vector x ∈ R a and y ∈ R b , we will use x ⊕ y to denote their concatenation. We will use 0 a to denote the all-zero vector with dimension a . We will also assume without loss of generality that w ≥ 2 m . 20

In the following, a random network refers to a network whose weights have entries sampled from a uniform distribution, i.e. W i,j ∼ U ( -1 , 1) for every weight W in the random network.

We will first recall Lemma 2 from Pensia et al. (2020) which shows that a pruned 2-layer random network can approximate a linear function.

Lemma 2 (Approximating a linear function; Theorem 1 of Pensia et al. (2020) restated) . Let W ∈ R m ′ × m , ∥ W ∥ 2 = O (1) , then for σ ∈ { ReLU , I} , where I represents the identity operator, for a random network g ( x ) = W 2 σ ( W 1 x ) with W 2 ∈ R m ′ × h , W 1 ∈ R h × m for hidden dimension

h = O ( m log( mm ′ min { ϵ,δ } )) , with probability 1 -δ , there exists boolean masking matrices M 1 , M 2 , such that for any x ∈ R w ,

where ⊙ denotes the Hadamard product.

We then derive two approximation results Lemmas 3 and 4 based on Lemma 2.

Lemma3. Under the setting of Theorem 2, with probability 1 -2 δ/ 3 , for any l ∈ [ L ] , l ′ ∈ [4 L -1] , let T ( l ) be the l -th layer of T , there exists a pruning of the ( l ′ -1) -th and the ( l ′ ) -th layer T ( l ′ -1) large , T l ′ large , named ˜ T ( l ′ -1) large , ˜ T l ′ large such that when defined on domain ∥ X ∥ 1 , 2 ≤ 2 L, X ∈ R m × N ,

˜ T ( l ′ -1) large is independent of the last m large -m rows of the input.

˜ T l ′ large ◦ ˜ T ( l ′ -1) large ( ¯ X ) is an ( C 1000 L 2 ) 4 L -3 ϵ -approximation of T ( l ) ( X ) with respect to 1 , 2 -norm.

Lemma 4. Under the setting of Theorem 2, for any matrix W ∈ R 4 m × 4 m , ∥ W ∥ 2 ≤ 1 , with probability 1 -δ/ 4 , for any l ′ ∈ [4 L ] , there exists a pruning of the l -th layer T ( l ′ ) large , named ˜ T ( l ′ ) large , such that when defined on domain X ∈ R m × N ,

˜ T ( l ′ ) large is independent of the last m large -4 m rows of the input.

a ( x ) = ˜ T ( l ′ ) large ( ¯ X ) is an ( C 1000 L 2 ) 4 L ϵ -approximation of ˆ g ( X ) = W X with respect to 1 , 2 -norm.

The proof of Lemmas 3 and 4 is deferred to Appendix C.5.1 We can now prove the theorem.

We will first show with induction that if we 1) prune the (2 l -1) -th and 2 l -th layers of T large to approximate T ( l ) for each l ∈ [ L ] , and 2) prune the 2 L +1 to 4 L -th layers of T large to approximate identity, then the pruned large transformer will be an ϵ -approximation of T for any input ∥ X ∥ 1 , 2 ≤ 1 .

We will perform induction on l : Let T (1: l ) define the composition of layer 1 to l , i.e. T (1: l ) ( X ) := T ( l ) ◦ T ( l -1) ◦ · · · ◦ T (1) ( X ) , and define ϵ l := ( C 1000 L 2 ) 4 L -3 -l ϵ . Suppose that T (1:2 l ) large is an ϵ l -approximation of T (1: l ) . Note that ∥T (1: l ) ( X ) ∥ 1 , 2 ≤ ( l +1) , since each attention output has a bounded norm of 1 and every weight matrix in projection function g has spectral norm smaller than 1 , hence the norm will at most increment 1 (due to residual connection) after each layer. We have that

Then according to Lemma 13, T ( l +1) is (1 + 200 L 2 /C ) -Lipschitz on the set of intermediate outputs { ( ˜ T (1:2 l ) large ( ¯ X ) ) 1: m | ∥ X ∥ 1 , 2 ≤ 1 } . We also have that T (1: l ) ( X ) is (1 + 200 L 2 /C ) l -Lipschitz. Now we can apply Lemma 5 to show that T (1:2 l +2) large can ϵ ′ -approximate T (1: l +1) with

The induction is then completed and we have the composition of ˜ T i large for i ∈ [2 L ] ϵ L -approximates the composition of T with ϵ L = ( C 1000 L 2 ) 3 L -3 ϵ . We will then perform another induction showing that the composition of ˜ T i large for i ∈ [2 L + l ] ϵ l + L -approximates T with ϵ l + L = ( C 1000 L 2 ) 3 L -3 -l ϵ . Suppose the statement holds for L -1 ≥ l ≥ 0 .

The induction step is similar, because we have T is (1 + 200 L 2 /C ) L Lipschitz, by Lemma 5, it holds that the composition of T i large for i ∈ [2 L + l +1] ϵ ′ -approximates T with,

This concludes the induction and prove the first claim of the theorem. For the second claim, notice that through similar induction steps, we can prune arbitrary layer of T large to approximate identity function and obtain the same approximation rate, this concludes the proof for the second claim.

C.5.1 Helper lemmas for Theorem 2

Error Analysis Our first lemma shows that the composition of ϵ -approximation can approximate the composition of the original function.

Lemma 4. Under the setting of Theorem 2, for any matrix W ∈ R 4 m × 4 m , ∥ W ∥ 2 ≤ 1 , with probability 1 -δ/ 4 , for any l ′ ∈ [4 L ] , there exists a pruning of the l -th layer T ( l ′ ) large , named ˜ T ( l ′ ) large , such that when defined on domain X ∈ R m × N ,

˜ T ( l ′ ) large is independent of the last m large -4 m rows of the input.

a ( x ) = ˜ T ( l ′ ) large ( ¯ X ) is an ( C 1000 L 2 ) 4 L ϵ -approximation of ˆ g ( X ) = W X with respect to 1 , 2 -norm.

Proof. One can prune the value matrix on layer l ′ to zero and the rest is a direct consequence of Lemmas 2 and 20.

Lemma 5. Given three metric spaces A,B,C equipped with same metric ∥ · ∥ . Suppose f 1 : A → B,f 2 : B → C are ϵ 1 , ϵ 2 -approximations of g 1 , g 2 with respect to ∥ · ∥ , where g 1 is a Lipschitz function with constant λ 1 with respect to ∥ · ∥ and ∥ g 2 ( x ) ∥ ≤ λ 2 x , then it holds that, f 1 ◦ f 2 is an ϵ ′ -approximation of g 1 ◦ g 2 , with ϵ ′ = ( λ 2 + ϵ 1 )( λ 1 + ϵ 2 ) -λ 1 λ 2

Proof. For any x ∈ R d 1 , it holds that,

This then suggests that,

Approximating ReLU MLP We will first show an extension of Lemma 2, illustrating that a pruned wide 4-layer ReLU MLP can approximate any 2-layer ReLU MLP.

Lemma 6. Consider any 2-layer ReLU MLP g : R 4 m → R 4 m parameterized by W 1 ∈ R 4 m × w , W 2 ∈ R w × 4 m , ∥ W 1 ∥ 2 , ∥ W 2 ∥ 2 ≤ 2 √ 2 , for any δ, ϵ ∈ (0 , 1) , consider a random 4-layer ReLU MLP f with input and output dimension 4 m and width w ′ = O ( w log( wm min { ϵ,δ } )) parameterized by W large ,i , with probability 1 -δ over the randomness of weight of f , there exists a nonstructural pruning of f named ˜ f , such that ˜ f is an ϵ -approximation of f with respect to 2 -norm.

Proof. Choose ϵ 0 = ϵ/ 8 . We only need to show there exists boolean matrices M 1 , M 2 , M 3 , M 4 , such that,

By Lemma 2, there exists boolean matrices M 1 ∈ R w ′ × 4 m and M ′ 2 ∈ R w × w ′ , such that for any x ∈ R 4 m ,

Hence we can choose M 2 = [ M ′ 2 0 ( w ′ -w ) × w ′ ] and have f 1 ( x ) = ReLU ( ( M 2 ⊙ W large , 2 )ReLU ( ( M 1 ⊙ W , 1 ) x )) is ϵ 0 -approximation of g 1 ( x ) = [ ReLU( W 1 x ) 0 w ′ -w ] .

Again by Lemma 2, there exists boolean matrices M ′ 3 ∈ R w ′ × w and M 4 ∈ R 4 m × w ′ , such that for any y ∈ R w ,

Hence we can choose M 3 = [ M ′ 3 , 0 w ′ × ( w ′ -w ) ] , and have f 2 ( x ) = ReLU ( ( M 4 ⊙ W large , 4 )ReLU ( ( M 3 ⊙ W large , 3 ) x )) is ϵ 0 -approximation of g 2 ( x ) = W 2 x .

It is also easy to check g 1 and g 2 are both 2 √ 2 -lipschitz and g 1 (0) = 0 . By Lemma 5, we conclude that ˜ f = f 1 ⊙ f 2 is ϵ ′ -approximation of g = g 1 ⊙ g 2 , with ϵ ′ = 4 √ 2 ϵ 0 + ϵ 2 0 ≤ ϵ .

This lemma then yields the following corollaries.

Corollary 1. Under the setting of Theorem 2, with probability 1 -δ/ 4 , for any l ∈ [ L ] , l ′ ∈ [4 L ] , there exists a pruning of the projection function g ( l ′ ) large , named ˜ g ( l ′ ) large , such that

1. ˜ g ( l ′ ) large is independent of the last m large -m dimension of the input.

Proof. One can construct such pruning by pruning the last m large -m rows of the weight of the last layer and the last m large -m columns of the weight of the first layer of g ( l ′ ) large to zero and then apply Lemma 6.

Approximating Attention Patterns We will now show that the attention pattern can be approximated by pruning random Transformer layers.

Lemma 7. For any δ, ϵ ∈ (0 , 1) , for any W ∈ R m , ∥ W ∥ 2 ≤ 1 , for two random matrix W 1 , W 2 ∈ R m ′ × m where m ′ = O ( m log( m min { ϵ,δ } )) , suppose X ∈ R m × N , then there exists nonstructural pruning of W 1 , W 2 , named ˜ W 1 , ˜ W 2 , such that

Here we adopt ∥∥ ∞ in vector sense, meaning the entry with largest absolute value.

Proof. Suppose without loss of generality, ∥ X ∥ : ,i ≤ 1 . According to Lemma 2, there exists nonstructural pruning of W 1 , W 2 , named ˜ W 1 , ˜ W 2 , such that for any x ∈ R m , ∥ x ∥ 2 ≤ 1 ,

This then suggests that,

This concludes the proof.

The next lemma shows how error propogates through the softmax operators.

Lemma 8. For any dimension d , suppose x, y ∈ R d satisfies ∥ x -y ∥ ∞ ≤ ϵ , then it holds that,

Proof. One can observe that,

This then suggests,

Hence,

This concludes the proof.

Approximating Attention Module We will need the following lemma showing there exists a pruning of the value matrix in T large such that it has eigenvalues with magnitude Θ(1) .

Lemma 9. For a matrix W ∈ R m large × m large , with probability at least 1 -δ 10 L , there exists a pruning of W , named W ′ , such that all the nonzero entries is contained in a d × d submatrix of W ′ that satisfies that (1) all its eigenvalues are within ( 1 2 , 1) , (2) the index of row specifying the submatrix and the index of column specifying the submatrix are disjoint.

Proof. As w large = Ω( m log( dL δ )) , hence we can split W 1: ⌈ m large / 2 ⌉ , ⌈ m large / 2 ⌉ +1: m large into ( m × ( m blocks, each with width at least O (log( ( m δ )) 21 . Within each block, with probability 1 -δ 10 Lm large , there exists at least one entry that has value at least 1 2 . We can then choose d disjoint entries in W that are all at least 1 2 , indexed with { ( a i , b i ) } i ∈ [ d ] where a i < a j and b i < b j for i < j . We can then prune all other entries to zero. Consider the submatrix defined by entries ( a, b ) for a ∈ { a i } i ∈ m and b ∈ { b i } i ∈ m . Then, this submatrix will be diagonal and contains eigenvalues within ( 1 2 , 1) . Further { a i } i ∈ m and { b i } i ∈ m must be disjoint because a i ≤ ⌈ m large / 2 ⌉ < b i . The proof is then completed.

We will also prove that LayerNorm with nonzero normalization constant is Lipschitz.

Lemma 10. For LayerNorm function defined as LN( x ) = P ⊥ x max {∥P ⊥ x ∥ 2 ,C } , x ∈ R m , for any x, y ∈ R m , it holds that,

Proof. We will proceed by a case analysis:

If ∥P ⊥ x ∥ 2 , ∥P ⊥ y ∥ 2 ≤ C , then ∥ ∥ ∥ LN( x ) -LN( y ) ∥ ∥ ∥ 2 = ∥P ⊥ x -P ⊥ y ∥ 2 C ≤ 1 C ∥ x -y ∥ 2 .

The cases exhaust all possibilities, thus the proof is completed.

Finally, we will need a lemma showing how error accumulates when we consider both attention patterns and the value matrices.

Lemma 11. For any dimension d and positive number N , for P, Q ∈ R d × d satisfying that ∥ P ∥ 2 , ∥ Q ∥ 2 ≤ 1 , for any x ∈ R d × N , if matrix A ∈ R N × N , B ∈ R d × N satisfy that,

Then it holds that,

Proof. For any i ∈ N , we will have

The rest follows from Lemma 10

A LayerNorm of larger dimension can be made to be functionally equivalent to a LayerNorm of a smaller dimension. Precisely:

Lemma 12. Given any dimension d < d ′ , it holds that for any x ∈ R d ,

Proof. The proof follows directly from definition.

We will now formally define attention module.

Definition 7 (Attention Module) . We will define attention module a ( X | W V , W K , W Q ) as

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 ≤ 1 with input domain ∥ X ∥ 1 , 2 ≤ 4 L , a ( X ) is 200 L 2 /C -lipschitz with respect to 1 , 2 -norm.

Proof. We have that

Choose ϵ to be a sufficiently small constant, such that, exp(32 Lϵ ) -1 ≤ 64 Lϵ . Consider X and ˜ X satisfying that ∥ X -˜ X ∥ 1 , 2 ≤ ϵ and ∥ X ∥ 1 , 2 ≤ 4 L, ∥ ˜ X ∥ 1 , 2 ≤ 4 L , we will have

By Lemma 8, this implies,

We also have

Lemma 11 then implies that

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 -δ/ 2 , for any l ∈ [ L ] , l ′ ∈ [4 L -1] , let a ( l ) be the attention module on the l -th layer of T , there exists a pruning of the ( l ′ -1) -th layer T ( l ′ -1) large , named ˜ T ( l ′ -1) large and the attention module on l ′ -th layer a l ′ large named ˜ a l ′ large , such that when defined on domain ∥ X ∥ 1 , 2 ≤ 2 L ,

˜ T ( l ′ -1) large is independent of the last m large -m rows of the input.

( ˜ a l ′ large ◦ ˜ T ( l ′ -1) large ([ x 0 ( m large -m ) × N ])) 1: m is an ( C 1000 L 2 ) 4 L -1 ϵ -approximation of a ( l ) ( x ) with respect to 1 , 2 -norm.

Proof. We will use the shorthand ϵ 0 = ( C 1000 L 2 ) 4 L ϵ and prune in the following order. It holds that for ϵ ≤ 1 , exp(8 L 2 ϵ 0 ) -1 ≤ 16 L 2 ϵ 0 .

We will prune W large , ( l ′ ) V according to Lemma 9 and name the pruned matrix ˜ W large , ( l ′ ) V . By Lemma 9, all the nonzero entries is contained in a d × d submatrix of W ′ that satisfies that all its eigenvalues are within ( 1 2 , 1) . We will assume WLOG the submatrix is the one specified by row 1 . . . d and column d +1 . . . 2 d and name the submatrix as W .

We will then prune T ( l ′ -1) large according to Lemma 4 to output ϵ 0 -approximation of X ∈ R m × N →   X W -1 P ⊥ W ( l ) V X 0 ( m large -2 m ) × N   . As W is defined as the submatrix pruned by W ( t +1) V , it holds that ˜ W large , ( l ′ ) V   X W -1 P ⊥ W ( l ) V X 0 ( m large -m ) × N   = [ P ⊥ W ( l ) V X 0 ( m large -m ) × N ] .

Finally we will prune W large , ( l ′ ) K , W large , ( l ′ ) Q according to Lemma 7 to approximate ( W ( l ) K ) ⊤ W ( l ) Q up to ϵ 0 error.

we can now calculate the approximation error. For any X ∈ R m × N , ∥ X ∥ 1 , 2 ≤ 2 L , suppose

Then by our constrution, it holds that ∀ i ∈ { 1 , 2 } , ∥ δ i ∥ 1 , 2 ≤ ϵ 0 ∥ X ∥ 1 , 2 . We would then have

By our construction, it holds that ∥ ˜ W large , ( l ′ ) V δ 2 ∥ 1 , 2 ≤ 2 ∥ δ 2 ∥ 1 , 2 ≤ 2 ϵ 0 ∥ X ∥ 1 , 2 .

Further, by the construction of ˜ W large , ( l ′ ) K , ˜ W large , ( l ′ ) Q , it holds that,

As for any i, j ∈ [ N

combined with Equation (49),

By Lemma 8, this implies

By Lemma 11, Equations (48) and (51) imply,

Now according to Lemmas 10 and 12, it holds that

This concludes the proof.

Approximating Transformer Layers We will finally show that two random Transformer layers can be pruned to approximate a given Transformer layer.

Lemma3. Under the setting of Theorem 2, with probability 1 -2 δ/ 3 , for any l ∈ [ L ] , l ′ ∈ [4 L -1] , let T ( l ) be the l -th layer of T , there exists a pruning of the ( l ′ -1) -th and the ( l ′ ) -th layer T ( l ′ -1) large , T l ′ large , named ˜ T ( l ′ -1) large , ˜ T l ′ large such that when defined on domain ∥ X ∥ 1 , 2 ≤ 2 L, X ∈ R m × N ,

˜ T ( l ′ -1) large is independent of the last m large -m rows of the input.

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 ∥ ∥ ∥ a ( l ) ( X ) + X ∥ ∥ ∥ 1 , 2 ≤ ( 2 C +1) ∥ X ∥ 1 , 2 and g ( l ) is 1 -lipschitz, according to Lemma 5, ( ˜ T l ′ large ⊙ ˜ T ( l ′ -1) large ([ x 0 ( m large -m ) × N ])) 1: m is an ϵ ′ -approximation of T ( l ) ( x ) , with

This concludes the proof.

C.6 Technical Lemmas

Lemma 15. Given any dimension d and number of samples n , for any sizen dataset { ( x i , y i ) } i ∈ [ n ] with x i ∈ R d and y i ∈ R , there exists a width2 n two-layer MLP f : R d → R with ReLU activation such that, f ( x i ) = y i for any i ∈ [ n ] .

̸

Proof. We will first choose direction w ∈ R d , ∥ w ∥ 2 = 1 and margin γ > 0 such that for any i = j in [ n ] , it holds that,

We will assume WLOG w ⊤ x i is increasing in i .

Then we will construct an auxilliary series z i for i ∈ [ n ] such that,

Finally consider the following two-layer MLP with ReLU activation,

we will show that f ( x i ) = y i for any i ∈ [ n ] . Notice that

Thus it holds,

̸

Lemma 16. Given any sets { x i } i ∈ m satisfying that x i ∈ R n and x i = 0 , there exists a set of orthonormal vectors { u j } j ∈ [ n -2] of R n such that (1) u ⊤ j 1 n = 0 for any j ∈ [ n -2] and (2) ∑ j ∈ [ n -2] u ⊤ j x i u j = x i for any i ∈ [ m ] .

̸

Proof. There exists a vector v ∈ R n such that v ⊤ x i = 0 for any i ∈ [ m ] . We can then construct an orthonormal basis { u j } j ∈ [ n -2] of R n as the basis of the normal space of span ( v, 1 n ) . Then the lemma holds.

Lemma 17. Given any dimension n and constant M , there exists a 2-layer width2 n ReLU network f : R n +1 → R such that for any x ∈ [0 , M ] n , y ∈ [ n ] , f ( x ⊕ y ) = x y .

Proof. The construction is as followed, we will choose f as

Then as we have

The proof is completed.

̸

Lemma 18. Given any dimension n and constant M > 0 , there exists a 2-layer width2 n ReLU network f : R n → R such that for any x ∈ R n satisfying there exists i ∈ [ n ] , x i > M and ∀ j = i, x j = 0 , it holds that f ( x ) = i .

Proof. The construction is as followed, we will choose f as

The proof is completed.

Lemma 19. Given any dimension n and natural numbers K,m,M , if there exists K different 2-layer widthm ReLU networks f : R n → R , then there exists a 2-layer width2 Km ReLU network

Proof. Suppose that

Then we can construct

where c k,i satisfies k .

The proof is then completed.

Lemma 20. Given any dimension n and W ∈ R n × n , ∥ W ∥ 2 ≤ 2 , there exists a 2-layer width2 n ReLU network f : R n → R such that for any x ∈ R n , it holds that f ( x ) = Wx and both weight matrices parameterizing f has spectral norm less than 2 √ 2 .

̸

Proof. The construction is straightforward, one can choose

C.7 Discussion on Architecture Choices

The reader may notice that Equation (6) is not the same as the standard GPT architecture,

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 (2) V Xσ ( C +( W (2) K X ) ⊤ ( W (2) Q X ) ) will converge to zero when input length converges to infinity. Hence, if not all e ( τ t,d ) where τ 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 Experiments

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 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.

(a) With Position

Embedding

(b) With Position

Embedding

(c) Without Position

Embedding

(d) Without Position

Embedding

S

S

S

S

（

S

S

S

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.

(a) With Position

(b) With Position

(c) Without Position

(d) Without Position

S(（((）)））[[[[]]]]

S(（(（）)））[[[[]]]]

S（（（(））））[[[[]]]]

S(((（）)））[[[[]]]]

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.

Attention pattern visualization for three-layer experiments. The first-layer attention is close to uniform, while the higher-layer attention shows no clear patterns.

(a) Embedding 15, run 1

(b) Embedding 15, run 2

(c) Embedding 17

(d) Embedding 16

S((((）)））[[[[]]]]

S((（()）））[[[[]]]]

S（（((））））[[[[]]]]

S((((）)）)[[[[]]]]

0.800

0.825

0.850

0.875

0.900

0.925

0.950

0.975

Length Generalization Acc

0.5

1.0

1.5

2.0

2.5

3.0

3.5

Balance Violation

Contrastive

(a)

seed

= 0

(b)

seed

= 1

(c)

seed

= 2

(d)

seed

= 3

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides mathematical evidence that typical methods for interpreting the inner workings of transformers (looking at attention patterns, looking at feedforward weights) cannot be relied upon. The main argument of the paper is this: There is a task (specifically, language modeling of the bounded-depth Dyck-k language) such that it is possible for a two-layer transformer language model (assuming that it is trained on infinite data) to solve it perfectly using only uniform attention in the second layer (Corollary 1). The first layer can be one of multiple constructions proposed in prior work. Since uniform attention can hardly be considered \"interpretable,\" this shows that a transformer can solve a task perfectly without an interpretable attention pattern, so we cannot reliably draw conclusions from attention patterns alone. Theorem 2 removes the assumption that the model is trained on infinite data. Theorem 3 shows (I think) that for any sufficiently large transformer, it's likely that there is a fully-connected layer in one of the feedforward sublayers such that, if you prune weights in the whole transformer two different ways, resulting in two transformers that have different behavior, that layer can be pruned exactly the same way in both networks, implying that that layer can be interpreted ambiguously. So, we can't rely on analyzing feedforward layers in isolation for interpretability. The authors give some examples of attention patterns in transformers trained on this task that don't have clear interpretations and run experiments validating some aspects of their theory. They also propose a regularization term that helps the transformer generalize to out-of-distribution examples on the Dyck language.\n\nSTRENGTHS:\nThis paper, in a mathematically rigorous way, answers a very important question: How much can we rely on naive ways of attempting to interpret the inner workings of transformers, i.e. looking at attention weights or feedforward weights? It turns out the answer is: not much. Although it might have been a bit obvious beforehand that looking at only one component of a transformer, and not the interplay between different parts of the network, is not necessarily reliable for interpretability, the authors prove this fact rigorously, which is an extremely valuable contribution. I think the bounded Dyck language was an excellent choice of case study: not so easy that the solution is trivial, but easy enough that the analysis is manageable. I think the overall logical argumentation of the paper is solid.\n\nA particularly cool result is that it's possible to train a transformer to solve the Dyck language while forcing the attention in the second layer to be uniform, which agrees with the theory. The experiments nicely back up the theoretical work.\n\nThe fact that Theorem 3 is a provable Lottery Ticket Hypothesis is also very interesting.\n\nOriginality: AFAIK this is an original and creative contribution.\n\nQuality: The methodology and mathematical rigor in this paper seem to be sound.\n\nClarity: Overall good, but I think this could be improved in places.\n\nSignificance: High. It has important implications for anyone working on interpretability and transformers.\n\nWEAKNESSES:\n**Edit:** I have read the rebuttal, and it addressed my biggest concerns.\n\nThe only reasons I have not immediately given this paper a higher score are that (1) I have some important questions in the Questions section I would like answered (2) I think the clarity of the paper could be improved.\n\nThe clarity issues have to do with the fact that the mathematical definitions and notation can get a bit overwhelming, and I would appreciate any attempt by the authors to simplify it and make it more readable.\n\nI think Section 3.2 requires a lot more clarification. I had a very difficult time understanding the logical argument, as there are several levels of nested quantifiers. I did not have time to check the proof in the appendix.\n\nQUESTIONS:\nMost important questions:\n\n1. Eq 5: In this equation, layer norm comes before the residual connection. Isn't post-norm supposed to go after the residual? Is this a mistake?\n1. Eq 5: I think there's a mistake with how the causal mask is included -- it should be added to the softmax inputs, not multiplied, and the entries should have values of 0 or -inf, not 1 or 0.\n\nImportant questions:\n\n1. Eq 5: Will these results hold true if you use pre-norm instead of post-norm?\n1. 162: It might be quite difficult for a transformer to eliminate positional encodings, because of the residual connections. Would this interfere with your construction?\n1. Thm 3: Doesn't non-structural pruning have the ability to encode information in $T_{large}'$, merely with the pattern of which weights are pruned? Am I correct in understanding that the last part of Theorem 3 is meant to account for this? And that the significance of Theorem 3 is not that $T$ can approximate $T_{large}$, but that $T_{large}$ could be approximating one of multiple transformers $T_1$ and $T_2$, which have different behaviors? Won't $T_1$ and $T_2$ necessarily have similar behavior because they are similar to $T_{Large,1}$ and $T_{Large,2}$, and $T_{Large,1}$, and $T_{Large,2}$ are similar to $T_{Large}$? I think the answer is no because $T_{Large,1}$ and $T_{Large,2}$ are not necessarily similar to $T_{Large}$, but I wanted to clarify.\n1. Sec 3.2: I think this whole section would benefit from a more concise, intuitive description of the logical argument inside Theorem 3, and how it implies the statement made on 247. I'm still quite confused as to how this argument works. If $T_{Large,1}$ and $T_{Large,2}$ coincide on $W$ after pruning, how do you know that the weights pruned in $W$ do not account for the difference in behavior between them? Doesn't this argument only apply to pruned transformers then?\n1. 227: What is the main message of Corollary 2? Is it that Approximate Balance still holds true in the presence of l2 regularization, which is a common technique, but to a lesser degree for large $\\lambda$? Is there a large enough $\\lambda$ that it will fail to hold true?\n1. Do Theorems 1-2 work for transformers with more than 2 layers?\n1. 237: What is \"width\"?\n1. 279: As mentioned before, don't positional encodings complicate the construction, possibly explaining why there's a change?\n1. Fig 3: There are points that make it hard to believe there is a strong negative correlation. For example, some blue points have lower accuracy than orange points that have much higher balance violation. It seems like contrastive regularization helps, but likely in a way that isn't fully explained by balance violation.\n1. You showed that the transformer can solve Dyck without an interpretable attention pattern in the second layer. How do we know it doesn't have an interpretable pattern in the first layer? Do your experiments show that as well?\n\nOther questions:\n\n1. 120: Why is there a difference in loss function between the theory and experiments?\n1. Eq 5: What about dropout?\n1. 184: Which pumping lemma? For CFGs or DFAs? Can you elaborate on how the proof of Thm 1 resembles that of the pumping lemma?\n1. 208: Why can the loss not be optimal?\n1. 277: How many random restarts for these experiments did you run?\n1. Fig 1: How do you pick the models shown in this figure?\n1. 299: It doesn't seem that close to me.\n\nMinor questions and comments from reading:\n\n1. 27: Is it generating \"grammars\" or \"strings\"?\n1. 31: It's not like a stack, it is a stack. It can be recognized by a one-state real-time DPDA.\n1. Finite-precision transformers cannot recognize Dyck up to arbitrary depth.\n1. 35: works -> papers\n1. 37: I like the list of research questions.\n1. 51: Can you elaborate more on the significance of the non-stack-like attention patterns, and what non-stack-like attention is?\n1. 62: This finding is really interesting.\n1. 69: Interesting!\n1. 71: Works -> Work\n1. 88: works -> work\n1. 95: These brackets are balanced; the problem is that they aren't nested.\n1. Eq 1: This could probably be made a bit clearer.\n1. Eq 2: This is a bit hard to understand. Can you express it as an equivalent PCFG? The fact that all Dyck prefixes are included complicates this. Why is it important to include all prefixes?\n1. Eq 3: What is $e_{w_i}$?\n1. Eq 5: $g^{(l)}$ is a full feedforward sublayer, not just a feedforward network, right? The phrasing of this section makes it unclear.\n1. 135: What about positional encodings?\n1. 138: Is the whole LayerNorm operation set to the identity, or just some of its entries?\n1. 196: It would be helpful to point out why the proof of sufficiency is not necessary for the main argument of the paper.\n1. 200: Could you please include an intuitive proof sketch of Corollary 1 in the main text?\n1. 220: Can you elaborate on how this constraint can be satisfied?\n1. 238: With respect to what variables is O(1) constant? How is the l2 norm defined on a matrix?\n1. 241: What is the significance of using $1 - \\delta$ instead of $\\delta$?\n1. 242: Make it clearer that \\epsilon-approximate refers to the equation below.\n1. 243: It might be helpful to point out that this holds true when X is one-hot vectors.\n1. 280: Cool!\n1. 291: Quantitative is misspelled.\n1. 306: Can you include these results?\n1. 329: Ungrammatical sentence.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the methods for interpretability of transformers that focus on individual parts of the model (instead of the entire network) using a synthetic setup of learning a Dyck language. Models that can solve this task need to satisfy the pumping lemma, and this is used to show that the attention pattern of a single layer can be nearly randomized without changing the network functionality. The consequential claim is that interpretability based on individual heads or weight matrices can be incorrect.\n\nSTRENGTHS:\n+ Theoretical characterization shows a perfect balance condition on the attention pattern for 2-layer transformer corresponding to Dyck language. This condition admits non-stack-like attention pattern. \n\n+ As a consequence of lottery ticket hypothesis, it is shown that any transformer can be approximated. by a random larger transformer and thus, interpretability based on local components may not be reliable.\n\nWEAKNESSES:\n* The paper appears to have a very narrow focus - examining the claim of stack-like patterns in the attention heads from Ebrahimi et al. Both stack-like and non-stack-like patterns can process Dyck languages. The presented analysis is not a \"proof\" or very convincing argument that the second is preferred by transformers. The presentation can be shifted to a critical examination rather than non-arguable claims. This is a minor presentation issue.\n\n* The title and initial discussion can emphasize the focus on bounded-depth Dyck language.\n\nQUESTIONS:\n* The fact that we are considering bounded-depth Dyck language makes the classical results from formal automata theory not directly applicable - if I know the maximum depth of Dyck language, does the pumping lemma still apply?\n\nLIMITATIONS:\nThere are no concerns with respect to broader societal impacts.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors discover interpretability conclusions based off attention weights can be unreliable with a case study on a Dyck language with a shallow transformer. Concretely, the authors show non-distinct attention patterns, even an uniform one can be theoretically learned to solve the Dyck language task which invalidates any interpretability conclusion based off solely on attention weights. The authors also try to prove this from formal theory as well as empirical results. Overall, the paper studies a timely topic and provides some insights, and has some open questions that may be addressed to make the paper more convincing.\n\nSTRENGTHS:\n- The paper overall is easy to follow.\n- The research topic about mechanistic interpretability based off attention weights seems like timely, and informs us about future works.\n- The formal proof on a shallow two-layer transformer with assumptions seems to valid the point that there are many attention patterns can be learned to solve the task, which makes the conclusions made based off a single set of attention weights less convincing.\n\nWEAKNESSES:\n-**Formal proof on a simple setup with assumption is good, but missing connections to larger transformer blocks**\n\nIt is great to have a starting point to start looking at the potential attention patterns can be learned to solve structural tasks like Dyck language with a lot of assumptions. But, it limits its connections to larger transformer blocks where there is no assumption about the first layer. After all, multiple stacked transformer blocks approximate high-degree of linearity than a single layer. Can also provide more explanations why their formal proof can be extended to larger transformer blocks?\n\n- **Different Attention Patterns Can Be Learned To Generate Dyck is not convincing enough**\n\nQualitatively to visualize attention patterns is fine as a start, but quantitatively only measuring attention weights difference among different initialization is bit lacking. It is well-known that, for structural generalization tests like this one, random seeds, learning rate and schedulers can large even affect model's factual behavior, let alone its hidden representations. The fact that attention weights are different is not surprising, and I feel like it is not enough to say different **effective pattern** can be learned to generate Dyck. For instance, if we can fix the attention weights to be uniform and enable gradients on other parameters, and show the model can still solve the task, it would be more convincing, correct me if i am wrong.\n\n- **Balanced attention increases generalizability is a weak evidence about the un-interpretability of attention weights?**\n\nIt is interesting to add regularizations to the attention weights can make the model generalizes better. But, by making the attention weights to distribute in a certain pattern can be viewed as a counter argument about the un-interpretability. Making them more balance is an interpretable mechanism to enforce; since it increases model performance, which means that interpretable attention distribution is somewhat better than randomly optimized ones. So, when performance is good, attention weights are better to be interpretable? I think some clarifications are needed here to connect to the main points of the paper.\n\n- **(Un)interpretability of Transformers is a very broad claim**\n\nI think this paper only studies the (Un)interpretability of Transformers' attention weights under a specific use case. So, it is may be better to position the paper to a narrower scope.\n\nQUESTIONS:\nSee my comments above.\n\nLIMITATIONS:\nSee my comments above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nRecently several efforts have been made to understand the exact mechanism that a trained Transformer uses to make its predictions. A number of those works rely on analysing the attention patterns or other components of Transformers. Broadly, this paper argues that in some cases it may not be possible to identify the mechanism used by Transformers by analysing individual components such as attention patterns. \n\nThey confine their analyses to Dyck languages with bounded depths -- which is the task of determining whether a sequence of brackets is well-balanced. For a two-layer Transformer which is expressive enough to recognize Dycks, they first show a condition for attention operation which is necessary and sufficient to construct a Transformer that can generate Dycks optimally (or almost optimally). The first layer in their construction is something called a minimal first layer which they define and then they show that if the attention in the second layer satisfies certain properties then one could construct the ReLU FFNs in such a way that the 2-layer Transformer can generate Dyck and obtain near-optimal mean squared error. The result essentially implies that there exists multiple Transformers with different attention patterns that can generate Dyck near-optimally including uniform attention and therefore attention patterns may not necessarily be informative of the mechanism used by the network.\n\nThey then show that Transformers with constant normed weight matrices can be approximated by pruning randomly initialized large Transformers (which have 4x layers and width). Lastly, they conduct experiments with Transformers trained on bounded Dyck languages. They compare the attention patterns of trained Transformers with different initializations and find that there is significant variance in the attention pattern on a string from the Dyck language.\n\nSTRENGTHS:\n(S1) **Interesting theoretical results.** In my opinion both Theorem 2 and 3 are interesting to some degree. The fact that Transformers with various types of attention patterns and FFNs can solve the same problem in almost optimal manner is important to note while working on interpretability. \n\nDyck languages are also very important from a formal language theory perspective given the Chomsky–Schützenberger representation theorem -- which states that every context-free language can be represented by a regular and a Dyck language. It may be useful to readers to point this somewhere like the introduction.\n\n(S2) **Well-motivated problem.** Recently, mechanistic interpretability as well as experiments with Dyck languages have been the subject of various works, and these results could be useful to researchers working in the area. Although there have been several works in the past regarding the unreliability of attention for interpretability, the results in this paper provide more concrete evidence for a more well-defined task which could be useful. \n\n(S3) **Length generalization experiment.** I also found the experiments with more balanced attention in Section 4.2 somewhat interesting. Although it is for a specific case, it is interesting to see that insights from sections are directly applicable to improve the length-generalization performance.\n\nWEAKNESSES:\n(W1) **Implications of theoretical results.** An important thing to note is that the results prove the existence of Transformers with various attention patterns (and other components) that can solve the same task in a near-optimal way. It does not imply that Transformers with different initializations and hyperparameters can converge to each one of them uniformly at random (or a similar distribution). \n\nTo be clear, the authors make no such claims as well. The reason I think it is important to note is that similar evidence and arguments regarding generalization have been incorrect before. The fact that deep learning models have a large capacity (VC dimension, Rademacher complexity) which indicates that networks with various different parameterizations can fit the training set perfectly. This could be seen as an argument for the fact that it is not necessary for a network to converge to parameters that will generalize well. However, in practice when neural networks are initialized in a certain way and trained with GD, then they often converge to one that generalizes well and further work was needed to theoretically understand that. Similarly, I think it is important to keep in mind that even if Transformers can solve a task in many different ways, it is still possible that it is not the case when trained with gradient-based methods in practice. I think it could be useful to clearly explain this in the paper so that the reader does not infer incorrect conclusions.\n\n\n(W2) **Robustness of Experiments.** I do not think W1 is a major weakness or is needed to be immediately addressed, but because of the reasons mentioned in W1, I believe the empirical results are quite important. However, the experiments do not seem sound and rigorous enough to draw conclusions in a robust manner. For quantitative experiments in Section 4.1, is the average attention variation computed for just one input sequence \"[[[[]]]](((())))\" from Dyck-(2,4)? Let me know if I have interpreted it incorrectly. However, if it is the case, then it would be useful to conduct more extensive experiments indicating attention variation across multiple sampled inputs and for different Dyck languages such as Dyck-(2, 8), Dyck-(3, 24), etc. The number of random initializations (40) also seems quite small to me. It could be useful to see how the average attention variation is affected by the width and depth of the model. It should not be very computationally expensive to compute average attention variation across a large number of inputs (since it requires only inference) and the training time for these tasks should not be huge since it involves only synthetic data.\n\n\nI am happy to hear the author's response and change my score depending on that.\n\nQUESTIONS:\n(Q1) MSE vs cross entropy: Is there any particular reason why the theoretical results are with mean-squared error and the experiments are with cross-entropy?\n(Q2) Dyck-1: Do the results hold for Dyck-1 as well? Previous works have shown how to interpret Transformers [1] and LSTMs [2] trained on Dyck-1. For languages like Dyck-1, a stack-like mechanism is not necessary and these works show that Transformers and LSTMs learn to do it with a counting mechanism. However, since bounded Dyck languages are regular, it is not necessary to emulate stack-like mechanism to recognize bounded Dycks as well. Maybe it could be useful to conduct experiments with Dyck-1 and see if the mechanism is robust. The task setup in [1][2] is different from the one considered in this paper. \n\n\n\n[1] On the Ability and Limitations of Transformers to Recognize Formal Languages.  \n[2] LSTM networks can perform dynamic counting.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}