Title: Compact Proofs of Model Performance via Mechanistic Interpretability
Venue: NeurIPS 2024
Accepted: yes
Keywords: mechanistic interpretability

verification

proof

guarantees

interpretability
Abstract: We propose using mechanistic interpretability -- techniques for reverse engineering model weights into human-interpretable algorithms -- to derive and compactly prove formal guarantees on model performance.
We prototype this approach by formally proving accuracy lower bounds for a small transformer trained on Max-of-$K$, validating proof transferability across 151 random seeds and four values of $K$.
We create 102 different computer-assisted proof strategies and assess their length and tightness of bound on each of our models.
Using quantitative metrics, we find that shorter proofs seem to require and provide more mechanistic understanding.
Moreover, we find that more faithful mechanistic understanding leads to tighter performance bounds.
We confirm these connections by qualitatively examining a subset of our proofs.
Finally, we identify compounding structureless errors as a key challenge for using mechanistic interpretability to generate compact proofs on model performance.

Compact Proofs of Model Performance via Mechanistic Interpretability

Abstract

We propose using mechanistic interpretability - techniques for reverse engineering model weights into human-interpretable algorithms - to derive and compactly prove formal guarantees on model performance. We prototype this approach by formally proving accuracy lower bounds for a small transformer trained on Max-ofK , validating proof transferability across 151 random seeds and four values of K . We create 102 different computer-assisted proof strategies and assess their length and tightness of bound on each of our models. Using quantitative metrics, we find that shorter proofs seem to require and provide more mechanistic understanding. Moreover, we find that more faithful mechanistic understanding leads to tighter performance bounds. We confirm these connections by qualitatively examining a subset of our proofs. Finally, we identify compounding structureless errors as a key challenge for using mechanistic interpretability to generate compact proofs on model performance.

1 Introduction

One approach to ensuring the safety and reliability of powerful AI systems is via formally verified proofs of model performance [48, 11]. If we hope to deploy formal verification on increasingly large models [24, 27] with powerful emergent capabilities [56] across more diverse and broader domains [5, 46], we will need compact proofs of generalization bounds on specific models that certify global robustness. However, existing approaches tend to use proof strategies that suffer from bad asymptotic complexity, while verifying either generalization properties of training procedures or local robustness properties of specific models.

One key challenge to verification is that neural network architectures are highly expressive [51, 58], and models with similar training procedure and performance may still have learned significantly different weights [38, 9]. This expressivity makes it difficult to adequately compress explanations of global model behavior in ways that correspond closely enough to the model's actual mechanisms to be useful for efficient verification without being too lossy , especially when using only knowledge of the architecture or training procedure. We propose verifying model performance using understanding derived from mechanistic interpretability (Section 2) - that is, reverse engineering the specific implementation of the algorithm from the learned weights of particular models. Knowledge of the specific implementation allows us to construct less lossy simplifications of the model, and more efficiently reason about model performance over possible inputs.

In this work, we provide a case study of translating mechanistic interpretations into compact proofs. We train an attention-only transformer on a Max-ofK task with 151 random seeds (Section 3), and

Unembed

Unembed

Embed

Embed

OV

QK

OV Circuit

QKCircuit

t

0

t

1

t

2

t

3

t

0

t

1

t

2

t

3

t

0

t

1

t

2

t

3

t

0

t

1

t

2

t

3

True Model

Brute Force Proof

(

ℓ

0

, ℓ

1

, . . . , ℓ

63

)

(

ℓ

0

, ℓ

1

, . . . , ℓ

63

)

(

ℓ

0

, ℓ

1

, . . . , ℓ

63

)

(

ℓ

0

, ℓ

1

, . . . , ℓ

63

)

Input

Logits

?

. . .

. . .

. . .

. . .

FLOPs Required:

Unexplained Dimension:

Asymptotic Complexity:

1

.

41

×

10

14

99

.

92 %

1

.

07

×

10

9

O

(

d

vocab

n

ctx

)

3

.

51

×

10

7

95

.

31 %

1

.

28

×

10

4

O

(

d

vocab

3

·

n

ctx

)

4

.

68

×

10

6

28

.

41 %

4

.

42

×

10

3

O

(

d

vocab

·

d

model

2

·

n

ctx

)

QK Circuit decomposes

into large 'size' and small

'noise' components

Size component w/

singular value

7

.

4

×

10

3

Other components have

singular value

<

1

.

5

×

10

1

then reverse engineer the models using standard mechanistic interpretability techniques. We use our understanding to define a set of 102 different computer-assisted proof strategies with varying tightness of bound and with different asymptotic complexity and number of required FLOPs (Section 4). 4 We validate our technique against an additional 604 models for varying values of K (Appendix A.2.1).

We define a quantitative metric to assess the mechanistic understanding used in a proof strategy by the dimensionality of the function space that the proof strategy must consider, which we deem the unexplained dimensionality of the proof strategy (Sections 5.1, and A.5). Using this metric, we find a negative relationship between proof length and degree of understanding. We qualitatively examine proof strategies to confirm and explain this relationship, finding that more compact proofs both require and provide more mechanistic understanding. We also find suggestive evidence that the trade-off between proof length and tightness of bound is modulated by the faithfulness of the mechanistic understanding used to derive the proof (Section 5.2). 5

However, we also identify compounding structureless error terms as a key challenge for generating compact proofs on model behavior (Sections 5.3, and G.2.5). The implementation of algorithms inside of neural networks may contain components that defy mechanistic understanding and appear to us as 'noise'. When we don't know how noise composes across model components, establishing a bound requires pessimizing over the ways the composition could occur. Worst-case noise can quickly grow even when the empirical noise is small, leading to vacuous performance bounds.

2 Mechanistic interpretability for proofs

Generalization bounds on global performance In the style of prior mechanistic interpretability evaluation work [6], we target theorem templates that establish bounds on the expected global performance of the model. Let M : X → Y be a model (here assumed to be a neural network), D be a probability distribution over input-label pairs ( l, t ) ∈ L × X , notated as D| X when marginalized over labels, and f : L × Y → R be a scoring function for evaluating the performance of the model. Then, we seek to establish lower bounds b on the expected ¯ s as the form:

As f can be any metric, this is a fully general template for theorems that can capture any aspect of model performance for which we have a formal specification. However, in this work we restrict f to be the accuracy and D| X to be uniform, so our theorems lower bound the accuracy of the model. Our proof methodology generalizes straightforwardly to other input distributions (Appendix A.8), and only a little work is required to generalize from accuracy to log-loss (Appendix A.11).

Proof template The proofs of model performance in this work have two components: a computational component C : model weights → R and a non-computational component Q arguing that for any model M ′ , C ( M ′ ) ≤ E ( l, t ) ∼D f ( l, M ′ ( t )) , thus implying that C generates a valid lower bound for the performance of M . The whole proof is Q paired with a trace of running C that certifies its output on M . 6 Here, b = C ( M ) . As even the size of the model parameters is much larger than any reasonable Q , we approximate the length of a proof pair C, Q by the length of a trace of C ( M ) .

Proof compactness vs. tightness of bound Different proof strategies make different tradeoffs between compactness and tightness of bound. For example, consider two extreme proof strategies: We can 'prove' a vacuous bound using a null proof. On the other hand, in the brute-force proof, we simply run the model on the entirety of D to achieve b = ¯ s , albeit with a very long proof.

We quantify the length of C ( M ) using two metrics: the asymptotic time complexity of C as we scale the size of the model and the input t , as well as the empirical average number of floating point operations required to evaluate C ( M ′ ) over a given set of models {M i } . We measure tightness of bound of C ( M ) using the ratio of the bound to the true accuracy: b/ ¯ s .

Proof as pessimal ablation A standard way of assessing the faithfulness of mechanistic interpretability is by ablating the parts of the model that your interpretation does not explain [54, 6, 23]. In this framework, proofs can be thought of as performing a pessimal ablation over the unexplained parts of the model - we set the remaining components of the model (the 'noise' or error terms) to values over X that minimize the performance of the model. However, the number of ablations required for a complete argument might be quite high. Thus, we construct relaxations (Appendix A.4) over input sequences, such that performing pessimal ablations on a smaller number of relaxed input sequences is sufficient to lower bound the performance on D .

3 Experimental setting

We study our approach to generating compact proofs in a simple toy setting: Max-ofK .

Model Architecture We study one-layer, one-head, attention-only transformers with no biases but with learned positional embeddings, with vocabulary size d vocab , model and head dimension d = d model = d head , and context length n ctx := k . The model parameters consist of the n ctx × d model positional embedding P ; the d vocab × d model token embed E ; the d model × d model query, key, value, and output matrices of the attention head Q , K , V , and O ; as well as the d model × d vocab unembed matrix U . We assume (as is standard in language modeling) that d model < d vocab .

For an n ctx × d vocab one-hot encoding x = [ x 0 , x 1 , . . . , x n ctx -1 ] of an input sequence t = [ t 0 , t 1 , . . . , t n ctx -1 ] , we compute the logits of the model as follows:

Initial residual stream ( n ctx × d model )

Attention matrix ( n ctx × n ctx )

Final residual stream ( n ctx × d model )

Final seq. position logits ( d vocab )

where σ ∗ is the masked softmax function used in causal attention. Because we only look at outputs of the model above the final sequence position i = n ctx -1 , we also denote this position as the 'query position' and the value of the token in this position as t query , one-hot encoded as x query . The model's prediction is the token corresponding to the max-valued logit ℓ max .

Task Specifically, we study the setting with n ctx = k = 4 because it is the largest sequence length for which we can feasibly evaluate the brute-force proof. We set hidden dimension d model = 32

̸

and a vocabulary of size d vocab = 64 comprising integers between 0 and 63 inclusive. For an input sequence t , we denote the true maximum of the sequence by t max . Outputting the correct behavior is equivalent to outputting logits ℓ such that ∆ ℓ t ∗ := ℓ t ∗ -ℓ max < 0 for all t ∗ = t max . We trained 151 models on this task. Models achieved average accuracy 0 . 9992 ± 0 . 0015 over the entire distribution.

Path decomposition Following prior work [13], we expand the logits of the model and split the paths through the model into three components - the QK circuit, the OV circuit, and the direct path:

Intuitively, the QK circuit determines which tokens the model attends to from a particular query token and sequence position, while the OV circuit processes the tokens and sequence positions the model attends to. The direct path is simply the skip connection around the attention head.

We further divide the QK and OV circuits into token (position-independent) and position-dependent components. Let P avg = ∑ i P i /n ctx be the average position embeds across positions (of size d model ), and let ¯ P denote either 1 n ctx ⊗ P avg or 1 d vocab ⊗ P avg depending on context, the result of broadcasting P avg back into the shape of P or E (that is, n ctx × d model or d vocab × d model ). Similarly, let P q = 1 d vocab ⊗ P query be the result of broadcasting P query . Then for one-hot encoded x , we can rewrite the QK and OV circuits, as well as the direct path, as follows:

where ˆ P = P -¯ P and ¯ E = E + ¯ P and E q = E + P q (since h (0) = x ¯ E + ˆ P ).

3.1 Mechanistic interpretation of learned models

Using standard empirical mechanistic interpretability techniques, we interpret one of our learned models (our 'mainline' model) by independently examining the QK and OV circuits and the direct path. 7 We find that the model outputs the largest logit on the true max token t max by attending more to larger tokens via the QK circuit and copying the tokens it attends to via the OV circuit. We then quantitatively confirm that these interpretations hold for all 151 models by reporting the mean plus minus standard deviation for various summary statistics. Plots for this section are available in Appendix B.2.

QK circuit By qualitatively examining the positionindependent QK component EQKE , we find the amount of pre-softmax attention paid to a key token is approximately in-

Input

Logits

(

ℓ

0

, ℓ

1

, . . . , ℓ

63

)

Unembed

Embed

t

0

t

1

t

2

t

3

Direct path

'does nothing'

OV Circuit performs

low-rank copying

QK Circuit attends

to larger tokens more

dependent of the value of the query token t query , and increases monotonically based on the size of the key token. We confirm this hypothesis by performing a singular-value decomposition (SVD) of the EQKE matrices (Appendix G.2.3), and find that it contains a single large rank-one component with singular value around 7800 ± 380 , around 620 ± 130 times larger than the second largest component with singular value 13 ± 3 . The left (query-side) singular vector is approximately constant in all dimensions, with value 0 . 1243 ± 0 . 0003 ≈ 1 / 8 = 1 / √ d vocab . The right (key-side) singular vector of this component is monotonically increasing as we increase the size of the key token, with ( 1 / √ d -scaled) pre-softmax attention increasing by an average of 1 . 236 ± 0 . 056 when the key token increases by 1. 8

√

In comparison, each 1 / d -scaled entry of the position-dependent QK component EQKP has negligible size (average 0 . 31 ± 0 . 18 ), suggesting that EQKP is unimportant to the functioning of the model.

Weconfirm this by zero ablating EQKP , which changes the models' accuracies from 0 . 9992 ± 0 . 0015 to 0 . 9993 ± 0 . 0011 . Combined with our interpretation of EQKE , this implies that the attention pattern of the model depends only on the token values and not the ordering of the sequence.

OV circuit Then, by qualitatively examining the position-independent OV component EVOU , we see that it has large positive entries along the diagonal. In fact, the entry along the diagonal is the largest in the row for all rows corresponding to t > 6 . 6 ± 1 . 2 . Since each entry in the sequence is uniformly sampled and d vocab = 64 , this means that EVOU is a good approximation for the identity matrix for all but ≈ (7 / 64) 4 ≈ 1 . 2 × 10 -2 % of the sequences.

As with the position-dependent QK component, the position-dependent OV component PVOU also has negligible size and is unimportant to model performance. Taken together with the above results on EVOU , this suggests that the attention head copies the tokens it attends to.

Direct path As with the two position-dependent components, the entries in EU have small absolute magnitude 2 . 54 ± 0 . 20 , 9 and contribute negligibly to model performance.

4 Proofs of model performance

In this section we describe intuitions for three categories of proof that are developed around different mechanistic interpretations and methods for using the interpretations. The strategies result in proofs of different complexities with varying bound tightness (Table 1). We provide detailed theorem statements, proofs, algorithms, and explanations of proof search in Appendices C, D, E, F, and G.

Our theorem statements for Q will all be of the form

We leave implicit the traces of running C specific strategy on our specific models to give the overall theorem. We report the computational complexity or estimated FLOPs of running C specific strategy as approximations for our proof lengths.

4.1 The brute-force baseline

We start by considering the brute-force proof (Appendix D), which treats the model as a black box and evaluates it on all possible sequences. 10 However, this proof strategy has bad asymptotic complexity and is untenable for larger models and larger input distributions. So in subsequent sections, we use knowledge of the model drawn from the interpretation in Section 3.1 to derive more compact proofs.

4.2 A cubic proof

Next, we use the fact that the model is composed of the direct path and the QK and OV circuits (Section 3) to decrease the number of sequences that we need to consider, and the fact that only the position-independent components EQKE and EVOU contribute meaningfully to performance (Section 3.1) to pessimize over sequence ordering.

First, let a pure sequence ξ be a sequence with at most three distinct tokens: the max token t max , the final token t query ≤ t max , and optionally a third token t ′ < t max , and let Ξ pure be the set of all pure sequences in X . 11 For a given input sequence t , define the adjacent pure sequences Adj( t ) as the set of sequences that share the same max and query token, and only take on values in t :

Using the convexity of softmax and the fact that the model contains three paths, we can show that one-layer attention-only transformers satisfies a variant of the following convexity property: for a given t , if M ( ξ ) is correct for all ξ ∈ Adj( t ) , then M ( t ) is correct. That is, for these transformers, we can bound the accuracy on all sequences by evaluating M on only the O ( d vocab 3 ( n ctx -1)!)

pure sequences. This allows us to bound the accuracy of our actual M on all d vocab n ctx sequences, while evaluating it on O ( d vocab 3 ( n ctx -1)!) sequences.

We can reduce the number of sequences that we need to evaluate by pessimizing over the order of a sequence. For a given tuple of ( t max , t query , t ′ ) , there are ( n ctx -1)! pure sequences, corresponding to the permutations of the tuple. Pessimizing over the order of sequences reduces the number of sequences to consider for each ( t max , t query , t ′ ) tuple to the number of t ′ in the pure sequence, and the total number of sequences to O ( d vocab 3 n ctx ) . By precomputing the five component matrices EU , EQKE , EQKP , EVOU , PVOU and cleverly caching intermediate outputs, we can reduce the additional work of each sequence to the O ( n ctx ) required to compute the softmax over n ctx elements, resulting in asymptotic complexity O ( d vocab 3 n ctx 2 ) (Theorem 12, additional details in Appendix E).

4.3 Sub-cubic proofs

We now consider proofs that are more compact than O ( d vocab 3 ) . These require avoiding iteration over any set of size O ( d vocab 3 ) (e.g. the set of pure sequences) and performing operations that take O ( d vocab ) time on each of O ( d vocab 2 ) combinations. Unfortunately, some methods of avoiding these operations can lead to vacuous bounds (i.e. accuracy lower bounds near 0% ). In order to recover non-vacuous bounds, we introduce two tricks: the 'mean+diff trick' to better approximate the sum of two components with unequal variance, and the 'max row diff trick' to improve upon the low-rank approximations for EU and EQKE . We consider applying variants of these tricks at different locations in the naïve subcubic proof, leading to 100 distinct subcubic proof strategies. See Appendix G.2 for a formal description of these strategies.

4.3.1 Removing cubic-time computations

Reducing the number of cases by pessimizing over sufficiently small tokens Previously, we considered Θ( d vocab 3 n ctx ) pure sequences ξ , with ξ parameterized by ( t max , t query , t ′ , c ) . Recall from our mechanistic interpretation in Section 3.1 that the pre-softmax attention paid from t query to a key token t ′ is broadly invariant in t query and increases roughly linearly with the size of t ′ . This allows us to pessimize over the OV circuit over all 'sufficiently small' tokens.

More formally, suppose we are given some gap g ∈ N . For each pure sequence ξ with max token t max , query token t query , such that t query ≤ t max -g , and c copies of the third token type t ′ ≤ t max -g , we pessimally ablate the OV circuit over the set Ξ pure ( t max , t query , c ; g ) of pure sequences ξ ′ with the same max and query tokens and c copies of the third token type t ′ . If the model gets all sequences in Ξ pure ( t max , t query , c ; g ) correct, then we can conclude that it gets ξ correct, otherwise, we treat the model as having gotten ξ wrong. This means that it suffices to only consider the O ( d vocab 2 n ctx ) pessimal pure sequences of each of the O ( d vocab 2 n ctx ) sets of the form Ξ pure ( t max , t query , c ; g ) .

Decoupling and pessimizing computations that require O ( d vocab 3 ) computations Many parts of our cubic certificate require iterating through O ( d vocab 2 ) cases parameterized by ( t max , t query ) or ( t max , t ′ ) . For example, as part of the pessimization procedure over pure sequences, for each of the d vocab possible values of t max , we need to consider the relative effects on the d vocab -sized logits of attending to each of the O ( d vocab ) other tokens t ′ < t max , and for each t max and t query , we need to check that the contribution of the direct path on logits x query EU is not sufficiently large as to overwhelm the contribution from x max EVOU . We independently pessimize over each of these components over one of the d vocab -sized axes: for example, instead of computing x max EVOU + x query EU for each t max , t query pair, we first pessimally ablate the direct path along the query token (which takes O ( d vocab 2 ) time as it does not depend on the t max , and then consider the sum x max EVOU+max x ′ x ′ EU . Since this sum no longer depends on t query , we only need to perform it O ( d vocab ) times, for a total cost of O ( d vocab 2 ) .

Low rank approximations to EQKE and EU Recall from Section 3.1 that EQKE is approximately rank 1, where the sole direction of variation is the size of the key token. By computing only the low rank approximation to EQKE , we can more cheaply compute the most significant component of the behavior in the QK circuit. To bound the remaining error, we can use the fact that after pulling off the first principal component from each of the four matrices we multiply, very little structure remains.

We can find the rank 1/2 approximations by performing SVD on EQKE . We can efficiently compute the SVD in O ( d vocab d model 2 ) time by using the fact that EQKE can be written as the product of a d vocab × d model matrix and a d model × d vocab matrix. This allows us to avoid performing the O ( d vocab 2 d model ) -cost matrix multiplications to explicitly compute EQKE .

Similarly, we can more efficiently check that the direct path EU contributes negligibly to the model outputs, by using SVD to decompose EU into a sum of rank 1 products (which we can evaluate exactly) and a high-rank error term that we can cheaply bound.

4.3.2 Additional subcubic proof strategies

Tighter bounds for sums of variables with unequal variance via the 'mean+diff trick' Suppose we want to lower bound the minimum of the sum of two functions over three variables h ( x, y, z ) = f ( x, y ) + g ( y, z ) , while only iterating over two variables at a time. The naïve way is to minimize f ( x, y ) and g ( x, y ) independently:

Here, the error comes from setting the y s in f and g to different values. But in cases where g ( y, z ) varies significantly with y and only slightly with z , rewriting g as a sum of a component that is independent of z (only varying along y ), and one that depends on z , yields a better lower bound:

This estimate will have error at most ε , while the naïve estimator can have arbitrarily large error. We refer to this rewrite as the 'mean+diff trick'. 12 From the mechanistic interpretation in Section 3.1, we know that some of the components barely vary among one or more axes. So we can apply the mean+diff trick to get tighter lower bounds.

Avoiding matrix multiplications using the 'max row-diff trick' Using properties of linear algebra, we derive a cheap approximation to the max row-diff for the product of matrices AB in terms of the product of the max row-diff of B and the absolute value of A , which we deem the 'max row-diff'

2

29

2

39

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

brute force (acc: 0.9992 ± 0.0015)

cubic (rel acc: 0.9539 ± 0.0080)

subcubic (rel acc: 0.700 ± 0.036)

attention-

d

vocab

d

model

2

(rel acc: 0.675 ± 0.035)

direct-quadratic (rel acc: 0.633 ± 0.062)

attention-

d

vocab

d

model

2

, direct-quadratic (rel acc: 0.610 ± 0.060)

attention-quadratic (rel acc: 0.316 ± 0.037)

attention-quadratic, direct-quadratic (rel acc: 0.283 ± 0.037)

brute-force linear baseline

trick. We apply this trick to get a better cheap bound on the error terms of low-rank approximations, without having to multiply out the full matrices. See Appendices G.2.2, and F for more details.

5 Results

We run each of 151 transformers on the various proof strategies of different asymptotic complexity, and analyze these proofs to empirically examine the relationship between proof length, bound tightness, and degree of understanding. For each proof on each transformer, we approximate the length of the proof by estimating the number of FLOPs used, and plot this against the ratio of certified bound the true accuracy b/ ¯ s (Equation 2) in Figure 3. There exists a clear trade-off between bound tightness and compactness of the proof - more compact proofs yield looser bounds, and tighter bounds are associated with more expensive proofs.

5.1 Compact proofs both require and provide mechanistic understanding

Quantifying mechanistic understanding using unexplained dimensionality We first quantify the amount of mechanistic understanding used in a proof by measuring its unexplained dimensionality - the number of free parameters required to fully describe model behavior, assuming the structural assumptions of the proof are correct. More detailed interpretations leave fewer free parameters needing to be filled in via empirical observation (Appendix A.5). In Figure 5, we plot these axes and find a suggestive correlation: proofs based on less mechanistic understanding are longer.

More mechanistic understanding allows for more compact proofs In addition to the constructions in Section 4, the parts of proofs we were unable to compact seem to correspond to components that we do not mechanistically understand. For example, we could not cheaply bound the behavior of EVOU without multiplying out the matrices, and this seems in part because we do have a mechanistic understanding of how EVOU implements low-rank copying.

Compact proofs seem to provide understanding By examining compact proofs, we can extract understanding about the model. For example, the fact that replacing each row of EU with its average across rows has little effect on the bound implies that EU does not vary much based on t query .

5.2 Proof length vs. bound tightness trade-off is modulated by faithfulness of interpretation

Compact proofs are less faithful to model internals To derive more compact proofs, we use our mechanistic understanding to simplify the model computation in ways that diverge from the original model internals. For example, in some subcubic proofs (Section 4.3), we approximate EQKE with a rank-1 approximation corresponding to the 'size direction'. However, while other components are small, they're nonzero; this approximation harms model internals.

Less faithful interpretations lead to worse bounds on performance To confirm that faithfulness of understanding affects the tightness of bound independent of proof length, we plot the normalized accuracy bound of subcubic proofs that perform a rank-1 approximation to EQKE , versus the ratio of the first two singular components. A larger ratio between the components implies that the rank-1 approximation is more faithful. In Figure 4, we see a positive correlation between the two axes: when the interpretation is more faithful, the bounds are tighter, even at a fixed proof length.

400

600

800

1000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

EPQKE Singular Ratio:

σ

1

/σ

2

mean+max-diff

mean+max-diff-subproduct

max-diff-subproduct

max-diff-exact

svd

mean-recursive+max-diff-subproduct-recursive

mean+max-diff-subproduct-recursive

max-diff-subproduct-recursive

2

23

2

27

2

31

2

35

2

39

2

43

2

47

2

14

2

17

2

20

2

23

2

26

2

29

Unexplained Dimension (Estimated)

brute force

cubic

subcubic

attention-

d

vocab

d

model

2

direct-quadratic

attention-

d

vocab

d

model

2

, direct-quadratic

attention-quadratic

attention-quadratic, direct-quadratic

5.3 Compounding structureless noise is a big challenge for compacting global-behavior proofs

Pessimal error terms compound in the absence of known structure The rank-1 approximation of EQKE has small error. However, when making rank-1 approximations of each of the constituent matrices E,Q,K , pessimizing over the worst way to composing the individual small error terms leads to a bound on the error term of EQKE that is orders of magnitude larger than the actual error term. Because we don't understand how the matrices compose in a way that doesn't cause errors to compound (without just multiplying out the matrices), this approximation leads to a trivial bound on performance (Appendix G.2.5). We speculate that in many cases, there is no short human-interpretable description for why random noise or approximation errors do not compound across layers of neural networks (e.g., see the error correction results on randomly initialized neural networks from Hänni et al. [21]), and thus that compounding structureless errors may be an issue in practice.

6 Related Work

Generalization Bounds Prior work in the PAC-Bayes framework [58, 36, 12] proves generalization bounds over learning procedures, which are similar to the global performance bounds we consider in this work. These proofs tend to provide statistical guarantees [25, 26] about the outputs of a known stochastic training procedure, while we seek to bound the performance of particular trained models.

Formally verifying neural networks Most prior work formally verifies neural networks either via model checking [28, 7] or by relaxing the problem setting and taking an automated theorem proving approach [17, 50, 18, 35, 43] to verify local robustness properties. These proof strategies tend to be derived by examining only the network architecture. We take an approach more akin to interactive theorem proving [22] and verify global performance properties by reverse-engineering the weights.

Mechanistic Interpretability Finally, mechanistic interpretability is the subfield of the broader field of understanding model internals [45], which is too large to faithfully summarize. Our work takes most direct inspiration from efforts to deeply understand how either toy models [38, 9, 53, 2] or small pretrained text transformers [54, 20] implement algorithmic tasks, generally by performing ablations and SVD. In contrast, we formally prove that a transformer implements an algorithm.

Nichani et al. [39] proves that, in a significantly simplified 2-layer, 1-head attention-only transformer model and for the task of in-context bigram statistics, gradient descent will create induction heads [40]. Our results concern transformers with fixed weights. In concurrent work, Michaud et al. [34] use techniques inspired by mechanistic interpretability to perform automated program synthesis on 2-dimensional RNNs, while our work works with significantly larger transformer models.

7 Conclusion and Future Work

Summary In this work, we used a Max-ofK setting to prototype the use of mechanistic interpretability to derive compact proofs of model behavior. Using varying amounts of understanding, we derived more efficient proof computations lower bounding model accuracy. We found preliminary evidence that mechanistic understanding can compactify proofs. Moreover, we observed that the tightness of the lower bound offered by various proof strategies can be used to grade the faithfulness our mechanistic interpretation. Finally, we identified compounding structureless errors as a key obstacle to deriving compact proofs of model behavior.

Limitations and future work We study one-layer attention-only transformers on a toy algorithmic task. Future work should explore the viability of deriving proofs via interpretability using larger models featuring MLPs or layernorm on more complex domains. In addition, we were unable to significantly compact the part of the proof involving the OV circuit, which future work can explore. The proofs we explored in this work also did not lead to qualitatively novel insights; future work may be able to derive such insights with improved techniques. Finally, future work can address the problem of compounding structureless errors, perhaps by relaxing from worst-case pessimal ablations to typical-case heuristic guarantees [8].

Acknowledgments and Disclosure of Funding

We are immensely grateful to Paul Christiano for providing the initial support for this project and for his invaluable research advice, encouragement, and feedback throughout its duration.

Additionally, we are thankful for clarifying discussions and feedback from Jacob Hilton, Matthew Coudron, Adrià Garriga-Alonso, Aryan Bhatt, Leo Gao, Jenny Nitishinskaya, Somsubhro Bagchi, Gabriel Wu, Erik Jenner, Ryan Greenblatt, Ronak Mehta, Louis Jaburi, and many others. Louis Jaburi in particular contributed the text of the final proof of Theorem 11 in Appendix E.

We are indebted to various organizations for their support: Alignment Research Center for funding this project and making it possible at all; Mentorship for Alignment Research Students (MARS) program of the Cambridge AI Safety Hub (CAISH) for setting up the collaboration between a subset of authors, and providing funding for compute and in-person research sprints; Constellation and FAR Labs for hosting a subset of the authors and providing an excellent research environment, including as part of the Visiting Fellows Program and Astra Fellowship.

Author Contributions

Jason Gross led the project, including managing the team and conceptualizing the proofs approach. He ran the Max-of-4 experiments, devised the proof strategies, and wrote up the formal proofs. He worked on various case studies and developed general methodology for computing complexity and length bounds for proofs. He also developed the particular convex relaxations presented in the paper.

Rajashree Agrawal was invaluable in steering the direction of the project, including contributing to the preliminary experiment on Max-of-2 and developing the pessimal ablation approach. She worked on framing the results, and contributed text to the paper.

Thomas Kwa and Euan Ong extended the preliminary experiments to larger values of k and contributed substantially to the cubic proof. Chun Hei Yip , Alex Gibson , and Soufiane Noubir worked on case studies other than the Max-ofK task and informed discussion on proof complexity.

Lawrence Chan spearheaded the writing of the paper, including turning informal claims into formal theorem statements, creating figures, and writing the core text. He also developed the unexplained dimensionality metric for clarifying the takeaway of the paper.

References

Behzad Akbarpour, Amr Abdel-Hamid, Sofiène Tahar, and John Harrison. Verifying a synthesized implementation of IEEE-754 floating-point exponential function using HOL. The Computer Journal , 53:465-488, May 2010. doi: 10.1093/comjnl/bxp023.

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models, 2022.

Andrew Appel and Ariel Kellison. VCFloat2: Floating-point error analysis in Coq. In Proceedings of the 13th ACM SIGPLAN International Conference on Certified Programs and Proofs , CPP 2024, pages 14-29, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400704888. doi: 10.1145/3636501.3636953.

Sylvie Boldo and Guillaume Melquiond. Flocq: A unified library for proving floating-point algorithms in Coq. In 2011 IEEE 20th Symposium on Computer Arithmetic , pages 243-252, July 2011. doi: 10.1109/ARITH.2011.40.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL https://arxiv.org/abs/2005.14165 .

Lawrence Chan, Adrià Garriga-Alonso, Nicholas Goldwosky-Dill, Ryan Greenblatt, Jenny Nitishinskaya, Ansh Radhakrishnan, Buck Shlegeris, and Nate Thomas. Causal scrubbing, a method for rigorously testing interpretability hypotheses. AI Alignment Forum , 2022. URL https://www.alignmentforum.org/posts/JvZhhzycHu2Yd57RN/causalscrubbing-a-method-for-rigorously-testing .

Chih-Hong Cheng, Georg Nührenberg, and Harald Ruess. Maximum resilience of artificial neural networks. In Automated Technology for Verification and Analysis: 15th International Symposium, ATVA 2017, Pune, India, October 3-6, 2017, Proceedings 15 , pages 251-268. Springer, 2017.

Paul Christiano, Eric Neyman, and Mark Xu. Formalizing the presumption of independence. arXiv preprint arXiv:2211.06738 , 2022. doi: 10.48550/arxiv.2211.06738.

Bilal Chughtai, Lawrence Chan, and Neel Nanda. A toy model of universality: Reverse engineering how networks learn group operations, 2023.

Edmund M. Clarke, William Klieber, Miloš Nováˇ cek, and Paolo Zuliani. Model Checking and the State Explosion Problem , pages 1-30. Springer Berlin Heidelberg, Berlin, Heidelberg, 2012. ISBN 978-3-642-35746-6. doi: 10.1007/978-3-642-35746-6_1. URL https://doi.org/10. 1007/978-3-642-35746-6_1 .

David Dalrymple, Joar Skalse, Yoshua Bengio, Stuart Russell, Max Tegmark, Sanjit Seshia, Steve Omohundro, Christian Szegedy, Ben Goldhaber, Nora Ammann, et al. Towards guaranteed safe AI: A framework for ensuring robust and reliable AI systems. arXiv preprint arXiv:2405.06624 , 2024.

Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. Proceedings of the Thirty-Third Conference on Uncertainty in Artificial Intelligence, UAI 2016, August 11-15, 2017, Sydney, NSW, Australia, 2017.

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. URL https://transformer-circuits.pub/2021/framework/index.html .

Martín H. Escardó. Synthetic topology of data types and classical spaces. Electronic Notes in Theoretical Computer Science , 87:21-156, November 2004.

Martín H. Escardó. Infinite sets that admit fast exhaustive search. In Proceedings of the 22nd Annual IEEE Symposium on Logic in Computer Science (LICS 2007) , Wrocław, Poland, July 2007.

Martín H. Escardó. Seemingly impossible functional programs, 2007. URL https://math. andrej.com/2007/09/28/seemingly-impossible-functional-programs/ . Accessed: 2024-05-15.

T. Gehr, M. Mirman, D. Drachsler-Cohen, P. Tsankov, S. Chaudhuri, and M. Vechev. AI2: Safety and robustness certification of neural networks with abstract interpretation. In 2018 IEEE Symposium on Security and Privacy (SP) , pages 3-18, Los Alamitos, CA, USA, May 2018. IEEE Computer Society. doi: 10.1109/SP.2018.00058.

Sven Gowal, Krishnamurthy Dvijotham, Robert Stanforth, Rudy Bunel, Chongli Qin, Jonathan Uesato, Relja Arandjelovic, Timothy Mann, and Pushmeet Kohli. On the effectiveness of interval bound propagation for training verifiably robust models. arXiv preprint arXiv:1810.12715 , 2018.

Jason S. Gross. Performance Engineering of Proof-Based Software Systems at Scale . PhD thesis, Massachusetts Institute of Technology, February 2021. URL https://dspace.mit. edu/handle/1721.1/130763 .

Michael Hanna, Ollie Liu, and Alexandre Variengien. How does GPT-2 compute greater-than. Interpreting mathematical abilities in a pre-trained language model , 2:11, 2023.

Kaarel Hänni, Jake Mendel, Dmitry Vaintrob, and Lawrence Chan. Mathematical models of computation in superposition. In ICML 2024 Workshop on Mechanistic Interpretability , 2024. URL https://openreview.net/forum?id=OcVJP8kClR .

John Harrison, Josef Urban, and Freek Wiedijk. History of interactive theorem proving. In Handbook of the History of Logic , volume 9, pages 135-214. Elsevier, 2014.

Stefan Heimersheim and Neel Nanda. How to use and interpret activation patching. arXiv preprint arXiv:2404.15255 , 2024.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556 , 2022.

Xiaowei Huang, Daniel Kroening, Wenjie Ruan, James Sharp, Youcheng Sun, Emese Thamo, Min Wu, and Xinping Yi. A survey of safety and trustworthiness of deep neural networks: Verification, testing, adversarial attack and defence, and interpretability. arXiv preprint arXiv:1812.08342 , 2018.

Xiaowei Huang, Wenjie Ruan, Wei Huang, Gaojie Jin, Yi Dong, Changshun Wu, Saddek Bensalem, Ronghui Mu, Yi Qi, Xingyu Zhao, Kaiwen Cai, Yanghao Zhang, Sihao Wu, Peipei Xu, Dengyu Wu, Andre Freitas, and Mustafa A. Mustafa. A survey of safety and trustworthiness of large language models through the lens of verification and validation, 2023.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020.

Guy Katz, Clark Barrett, David L. Dill, Kyle Julian, and Mykel J. Kochenderfer. Reluplex: An efficient SMT solver for verifying deep neural networks. In Computer Aided Verification: 29th International Conference, CAV 2017, Heidelberg, Germany, July 24-28, 2017, Proceedings, Part I 30 , pages 97-117. Springer, 2017.

A. E. Kellison, A. W. Appel, M. Tekriwal, and D. Bindel. LAProof: A library of formal proofs of accuracy and correctness for linear algebra programs. In 2023 IEEE 30th Symposium on Computer Arithmetic (ARITH) , pages 36-43, Los Alamitos, CA, USA, September 2023. IEEE Computer Society. doi: 10.1109/ARITH58626.2023.00021.

Gerwin Klein, Kevin Elphinstone, Gernot Heiser, June Andronick, David Cock, Philip Derrin, Dhammika Elkaduwe, Kai Engelhardt, Rafal Kolanski, Michael Norrish, Thomas Sewell, Harvey Tuch, and Simon Winwood. seL4: Formal verification of an OS kernel. In Proceedings of the ACM SIGOPS 22nd Symposium on Operating Systems Principles , SOSP '09, pages 207220, New York, NY, USA, 2009. Association for Computing Machinery. ISBN 9781605587523. doi: 10.1145/1629575.1629596. URL https://doi.org/10.1145/1629575.1629596 .

Xavier Leroy. A formally verified compiler back-end. Journal of Automated Reasoning , 43: 363-446, 2009.

Chuan Li. OpenAI's GPT-3 language model: A technical overview, June 2020. URL https: //lambdalabs.com/blog/demystifying-gpt-3 . Lambda Labs Blog, accessed October 30, 2024.

Wes McKinney. Data Structures for Statistical Computing in Python. In Stéfan van der Walt and Jarrod Millman, editors, Proceedings of the 9th Python in Science Conference , pages 56-61, 2010. doi: 10.25080/Majora-92bf1922-00a.

Eric J. Michaud, Isaac Liao, Vedang Lad, Ziming Liu, Anish Mudide, Chloe Loughridge, Zifan Carl Guo, Tara Rezaei Kheirkhah, Mateja Vukeli´ c, and Max Tegmark. Opening the AI black box: program synthesis via mechanistic interpretability, 2024.

Matthew Mirman, Timon Gehr, and Martin Vechev. Differentiable abstract interpretation for provably robust neural networks. In International Conference on Machine Learning , pages 3578-3586. PMLR, 2018.

Vaishnavh Nagarajan and J. Zico Kolter. Uniform convergence may be unable to explain generalization in deep learning. Advances in Neural Information Processing Systems , 32, 2019.

Neel Nanda and Joseph Bloom. TransformerLens. https://github.com/ TransformerLensOrg/TransformerLens , 2022.

Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint , 2023. doi: 10.48550/ arXiv.2301.05217.

Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. arXiv preprint , 2024. doi: 10.48550/arXiv.2402.14735.

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. URL https://transformer-circuits.pub/2022/in-contextlearning-and-induction-heads/index.html .

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-performancedeep learning library. Advances in neural information processing systems , 32, 2019.

Plotly Technologies Inc. Collaborative data science, 2015. URL https://plot.ly .

Aditi Raghunathan, Jacob Steinhardt, and Percy S. Liang. Semidefinite relaxations for certifying robustness to adversarial examples. Advances in neural information processing systems , 31, 2018.

Tahina Ramananandro, Paul Mountcastle, Benoıˆt Meister, and Richard Lethin. A unified Coq framework for verifying C programs with floating-point computations. In Proceedings of the 5th ACM SIGPLAN Conference on Certified Programs and Proofs , CPP 2016, pages 15-26, New York, NY, USA, 2016. Association for Computing Machinery. ISBN 9781450341271. doi: 10.1145/2854065.2854066.

Tilman Räuker, Anson Ho, Stephen Casper, and Dylan Hadfield-Menell. Toward transparent AI: A survey on interpreting the inner structures of deep neural networks. In First IEEE Conference on Secure and Trustworthy Machine Learning , 2022. doi: 10.48550/arxiv.2207.13243.

Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent, 2022. URL https://arxiv.org/abs/2205.06175 .

Alex Rogozhnikov. Einops: Clear and reliable tensor manipulations with Einstein-like notation. In International Conference on Learning Representations , 2022. URL https://openreview. net/forum?id=oapKSVM2bcj .

Sanjit A. Seshia, Dorsa Sadigh, and S. Shankar Sastry. Toward verified artificial intelligence making AI more trustworthy with a formal methods-based approach to AI system verification and validation.

Alex K. Simpson. Lazy functional algorithms for exact real functionals. In Luboš Brim, Jozef Gruska, and Jiˇ rí Zlatuška, editors, Mathematical Foundations of Computer Science 1998 , pages 456-464, Berlin, Heidelberg, 1998. Springer Berlin Heidelberg. ISBN 978-3-540-68532-6.

Gagandeep Singh, Timon Gehr, Markus Püschel, and Martin Vechev. An abstract domain for certifying neural networks. Proc. ACM Program. Lang. , 3(POPL), January 2019. doi: 10.1145/3290354.

Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus. Intriguing properties of neural networks. arXiv preprint arXiv:1312.6199 , 2013.

Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C. J. Carey, ˙ Ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental algorithms for scientific computing in Python. Nature Methods , 17:261-272, 2020. doi: 10.1038/s41592-019-0686-2.

Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174. PMLR, 2023.

Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. arXiv preprint , 2022. doi: 10.48550/arXiv.2211.00593.

Michael L. Waskom. seaborn: statistical data visualization. Journal of Open Source Software , 6 (60):3021, 2021. doi: 10.21105/joss.03021.

Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682 , 2022.

Eric Wong and Zico Kolter. Provable defenses against adversarial examples via the convex outer adversarial polytope. In International conference on machine learning , pages 5286-5295. PMLR, 2018.

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM , 64(3): 107-115, 2021.

A Subtleties of our approach

In this section, we address some subtleties and frequently asked questions about our approach.

A.1 Why study this simple task?

Formal reasoning is computationally expensive; very few large software projects have ever been verified [31, 30], none of them comparable to large transformer models [10, 19]. Separately, there is a high fixed cost to taking on any verification project, regardless of computational efficiency of the verification itself. Thus, we picked the simplest setting to study the question of interest: Is it even possible to formally reason more efficiently than by brute force about model behavior?

A.2 Scalability

In this section, we address concerns about the scalability of our approach.

A.2.1 Larger input spaces

We demonstrate that our proof strategies can be reused on larger input spaces while scaling better than the brute force approach does.

We applied our proof strategies to models trained for Max-of-5, Max-of-10, and Max-of-20. While running the brute force proof on Max-of-20 would require approximately 2 148 FLOPs, which is about 2 70 × the cost of training GPT-3 [32], our cubic proof achieves bounds of (94 . 1 ± 1 . 1) % (Max-of-5), (91 . 4 ± 2 . 1) % (Max-of-10), and (88 . 4 ± 4 . 0) % (Max-of-20) in under two minutes. See Tables 2, 3, 4, and 5 for more detailed numbers, and Figures 6, and 7 for visualizations. These results demonstrate that proof strategies can be reused on larger input spaces while scaling better than the brute force approach does.

A.2.2 Different tasks

In this paper, we worked on highly optimizing our relaxation to make our bounds as tight as possible when incorporating as little understanding as possible. This is not necessary for deriving proofs. Our general formalization of mechanistic interpretability is replicable: (1) theorem statements are exact expressions for the difference between the actual behavior of the model and the purported behavior, and (2) proofs are computations that bound the expression. Furthermore, our convexity theorems and proofs are applicable much more generally generally to element retrieval tasks.

A.2.3 More complicated architectures

We worked on a simple model studied in A Mathematical Framework for Transformer Circuits [13]. In follow-up work, we will extend this approach to proving bounds on 1L transformers with ReLU MLP trained on modular addition.

A.2.4 Larger models

It is an open question whether or not the mechanistic interpretability approach to proofs can scale to larger models. However, a large part of this question lies in the feasibility of deriving a high degree of faithful mechanistic understanding from large models - that is, whether mechanistic interpretability itself will scale. This is widely recognized in the field, and scaling interpretability approaches while getting both a high degree of mechanistic understanding and assurances that said understanding is faithful to the model is an active area of research. Broadly, we see the compact proofs approach as a metric on the quality of mechanistic understanding - we are not purporting to have a general solution to the problem of scaling interpretability, but instead claim that the challenges in proofs are in fact challenges in understanding networks.

2

31

2

43

2

55

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

brute force (estimated) (acc: 0.9990 ± 0.0018)

cubic (rel acc: 0.942 ± 0.010)

subcubic (rel acc: 0.702 ± 0.033)

direct-quadratic (rel acc: 0.649 ± 0.055)

attention-

d

vocab

d

model

2

(rel acc: 0.682 ± 0.033)

attention-

d

vocab

d

model

2

, direct-quadratic (rel acc: 0.629 ± 0.053)

attention-quadratic (rel acc: 0.354 ± 0.034)

attention-quadratic, direct-quadratic (rel acc: 0.336 ± 0.033)

(a)

n

ctx

= 5

,

d

vocab

= 64

2

42

2

65

2

88

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

brute force (estimated) (acc: 0.9988 ± 0.0013)

cubic (rel acc: 0.915 ± 0.021)

subcubic (rel acc: 0.673 ± 0.030)

direct-quadratic (rel acc: 0.640 ± 0.032)

attention-

d

vocab

d

model

2

(rel acc: 0.658 ± 0.028)

attention-

d

vocab

d

model

2

, direct-quadratic (rel acc: 0.626 ± 0.031)

attention-quadratic (rel acc: 0.393 ± 0.030)

(b)

n

ctx

= 10

,

d

vocab

= 64

2

47

2

74

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

brute force (estimated) (acc: 0.9972 ± 0.0031)

cubic (rel acc: 0.885 ± 0.012)

direct-quadratic (rel acc: 0.609 ± 0.030)

subcubic (rel acc: 0.613 ± 0.037)

attention-

d

vocab

d

model

2

, direct-quadratic (rel acc: 0.596 ± 0.030)

attention-

d

vocab

d

model

2

(rel acc: 0.595 ± 0.035)

attention-quadratic, direct-quadratic (rel acc: 0.385 ± 0.025)

(c)

n

ctx

= 10

and

d

vocab

= 128

2

63

2

109

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

brute force (estimated) (acc: 0.995 ± 0.015)

cubic (rel acc: 0.888 ± 0.036)

subcubic (rel acc: 0.563 ± 0.041)

attention-

d

vocab

d

model

2

(rel acc: 0.549 ± 0.041)

direct-quadratic (rel acc: 0.540 ± 0.040)

attention-

d

vocab

d

model

2

, direct-quadratic (rel acc: 0.528 ± 0.039)

attention-quadratic, direct-quadratic (rel acc: 0.322 ± 0.034)

attention-quadratic (rel acc: 0.323 ± 0.035)

(d)

n

ctx

= 20

,

d

vocab

= 64

400

600

800

1000

1200

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

EPQKE Singular Ratio:

σ

1

/σ

2

max-diff-exact

max-diff-subproduct

mean+max-diff-subproduct

mean+max-diff

svd

mean+max-diff-subproduct-recursive

max-diff-subproduct-recursive

mean-recursive+max-diff-subproduct-recursive

400

600

800

1000

1200

1400

1600

1800

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

EPQKE Singular Ratio:

σ

1

/σ

2

max-diff-exact

max-diff-subproduct

mean+max-diff-subproduct

mean+max-diff

svd

max-diff-subproduct-recursive

mean+max-diff-subproduct-recursive

mean-recursive+max-diff-subproduct-recursive

800

1000

1200

1400

1600

1800

2000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

EPQKE Singular Ratio:

σ

1

/σ

2

max-diff-exact

max-diff-subproduct

mean+max-diff-subproduct

mean+max-diff

svd

max-diff-subproduct-recursive

mean-recursive+max-diff-subproduct-recursive

mean+max-diff-subproduct-recursive

400

600

800

1000

1200

1400

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

EPQKE Singular Ratio:

σ

1

/σ

2

max-diff-exact

max-diff-subproduct

mean+max-diff-subproduct

mean+max-diff

svd

mean-recursive+max-diff-subproduct-recursive

mean+max-diff-subproduct-recursive

max-diff-subproduct-recursive

A.3 Why is more mechanistic understanding correlated with worse bounds?

Figure 3 exhibits Simpson's Paradox: although more faithful mechanistic understanding is correlated with better bounds within each class of proof (and moreover the most extensive mechanistic understanding results in the greatest improvement in bound tightness over baseline), when we aggregate across all proof strategies, we find that more mechanistic understanding is correlated with worse bounds.

This relationship is summarized in Figure 8.

From the compression perspective, more mechanistic understanding is about having more compression. Unless the model is losselessly compressible, we should expect that more compression will inherently be more lossy, no matter how good our compression scheme is. Correspondingly, using more understanding to get more compression will often result in a weaker bound, no matter how good our understanding is.

Conversely, we can think of the quality of proofs (the combination of tightness of bound, and length of proof) as a metric for how good our mechanistic understanding is. From this perspective, the fact that mechanistic-interpretability-derived bounds are bad suggests gaps in our mechanistic understanding. As the field matures and we develop tools that enable more faithful and complete understanding of model behavior, we expect that the quality of bounds we derive from mechanistic understanding will improve.

A.4 Convex relaxation

In this work, we construct convex relaxations to perform the pessimal ablations for our proofs.

In what sense are we using 'convexity'? The intuition is that we are attempting to optimize a function f over its domain X by incrementally making local changes to the input sequence, such as replacing one token by another, or by changing the order of tokens. The reason that convex optimization problems are easy to solve is that all local extrema are global extrema. This is not the case for our optimization problem, so we find a relaxation of f and its domain such that all local extrema are in fact global extrema.

Furthermore, most convex optimizers perform optimization at runtime by repeatedly stepping towards extrema. In this work, we 'optimize by hand', performing the optimization in the proof of our general theorems. The computation of the bound then only needs to instantiate the precomputed possible extrema with the actual values of the model's parameters to determine the the extrema actually are.

We now give a formal description of what we mean by 'convex relaxation'.

For a set of inputs X i , we define a set of 'relaxed inputs' X relaxed i with an injection T i : X i ↪ → P ( X relaxed i ) mapping input to the model to the set of corresponding relaxed inputs. On the relaxed input, we define a function h i : X relaxed i → R such that for all t ∈ X i and all labels l for which ( l, t ) is supported by (has non-zero probability in) D , we can find t relaxed ∈ T i ( t ) with f ( l, M ( t )) ≥ h i ( t relaxed ) . We proceed by finding a small subset of 'boundary' examples B i ⊂ X relaxed i , proving that if h i ( t relaxed ) ≥ b i for all t relaxed ∈ B i then h i ( t relaxed ) ≥ b i for all t relaxed ∈ X relaxed i .

Then, the computational component C of the proof validates that that h i ( t relaxed ) ≥ b i for some b i for all t relaxed ∈ X relaxed i . This allows us to conclude that f ( l, M ( t )) ≥ b i for all t ∈ X i .

A.5 Computing unexplained dimensionality

We claim in Figure 5 that we can use unexplained dimensionality as a metric for understanding. Here we describe how we compute the unexplained dimensionality of a proof strategy.

As in Figure 1, for any given proof, we can separate our treatment of transformer components into 'black-box' (e.g., matrix multiplication) and 'white-box' components (e.g., specifying that the QK circuit is approximately rank one; pessimizing over non-max tokens). Considering the performance score as a large white-box component which may reference black-boxes internally, we define the unexplained dimensionality of a single black-box computation as the log-cardinality of it function space (so, e.g, 2 · 64 for a function 64 → R 2 , whose cardinality is ( R 2 ) 64 , where 64 denotes the finite

Performance

Lower Bound

Brute Force Proof

True Performance

Trivial Proof

Impractical baseline of

guaranteeing

performance via

inference

(a) The baseline of using inference to generate proofs.

Performance

Lower Bound

Understanding

partially recovers

bound tightness

Decreased length

of proof

leads to

looser bounds

True Performance

(b) Shorter proofs by default have a worse performance bound b . Faithful understanding allows us to recover significant - but not complete - bound tightness with minimal proof-length overhead.

set on 64 elements). The unexplained dimensionality of the entire proof is the sum of the unexplained dimensions of all black-box components.

Intuitively speaking, unexplained dimensionality tries to capture the degrees of freedom that we have to check via brute enumeration over black-box computations. Proofs with less unexplained dimensionality contain more mechanistic understanding, and vice versa.

A.6 Computing approximate FLOPs

In Figure 3 and Table 1 on page 6 and on page 8, we display approximate floating point operations. We instrument our code to execute on phantom tensors that track their shape and accumulate an approximate count of floating point operations. We compute matrix additions and multiplications in the obvious way. We take the instruction count of SVD to be the cost of verifying that the output of SVD is a valid decomposition: that we have a pair of orthonormal bases which when multiplied out give the original basis.

A.7 IEEE 754 vs. R

In Section 2 we defined C and Q and glossed over whether we were reasoning over reals or floats. Here we clarify this point that we've so far been sweeping under the rug.

Let F denote the set of the relevant flavor of IEEE 754 Floating Point numbers (generally 32-bit for our concrete models, but everything would hold just as well for 64-bit). Let F ∗ denote F restricted to finite numbers (that is, without NaNs and without ±∞ ).

We parameterize C , M , and D over the real field 13 they operate on, so that, e.g., C F : model weights → F . Then we have Q establishing that for any model M ′ , C R ( M ′ R ) ≤ E ( l, t ) ∼D R f R ( l, M ′ R ( t )) , and we have a trace demonstrating that C F ( M F ) = b .

Let i : F ∗ → R be any injection that maps each floating point number to some real number that it is 'closest to'. Supposing that b ∈ F ∗ and thus b ∈ R , we need two additional components of the proof. We need to find ε, ε ′ ∈ R + prove that

Then we can chain these proofs to prove that

Such ε -ball robustness proofs should be well within the scope of existing approaches to formal methods on neural nets, see, e.g., [44, 3, 4, 29, 1, 57]. We leave actually dealing with the gap between floating point numbers and real numbers to future work.

A.8 Non-uniform distributions

In Equation 1 in Section 2 we defined the expected model performance as the expectation of the distribution D :

We then immediately specialized to the case where the marginalization D| X of D over labels is uniform. As we'll see in Theorem 1 in Appendix D and Algorithm 3 in Appendix E, the bound computation is modularized between a function that bounds the performance f ( l, M ( t )) over a restricted collection of inputs, and a much simpler function that combines the bounds on individual cases into a bound on the expectation over the entire distribution. The per-input bound computation is CORRECTNESS in Algorithm 1 and RELAXED-CORRECTNESS-PESSIMIZING-OVER-POSITION in Algorithm 3; the expectation computation is BRUTE-FORCE in Algorithm 1 and CUBIC in Algorithm 3.

Since the expectation computation is modularized, it is straightforward to extend our approach to non-uniform distributions simply by adjusting the weighting of each region of inputs. However, if the distribution is too far off from the uniform training distribution, the bound we get may not be very good, as we may not be allocating adequate computation to the high-probability regions of the input space.

A.9 Adversarial robustness via flexibility in D

There is flexibility inherent in Equation 1. Normally, by out-of-distribution (OOD) or adversarial inputs, we suppose that there is a distribution D in that's used for training and (in-distribution) validation, and another distribution D ′ that is the deployment distribution or generated by an adversary. If we had knowledge of D ′ , we could compute the expected performance from inputs sampled from D ′ . Even if we don't have exact knowledge of D ′ , we can still define a very broad distribution D that covers possible D ′ s.

In this work, D is the distribution of all 64 4 possible valid input sequences. In addition, as our proofs partition D into subdistributions, and bound the performance on each subdistribution, we can bound the model's performance on any possible distribution over valid input sequences.

A.10 Infinite distributions

In the brute force proof in Section 4.1, we run the model on the entirety of D . This operation is straightforward when X is finite. Perhaps surprisingly, we can do this even if X is infinite as long as the PDF L × X → R of D is computable and the natural computational topology of X is compact [16, 15, 14], because integration of computable functions on computable reals is computable [49].

A.11 Using alternate loss functions

Building on the point from Appendix A.8, it is also relatively straightforward to extend our approach from bounding expected accuracy to bounding log-loss. We will see in Figure 12 that the accuracy and log-loss share a subterm ∆ ℓ i . Since we compute this subterm in all of our algorithms, we can

easily extend our approach to log-loss by combining ∆ ℓ i directly rather than merely checking that the value is negative as we currently do in RELAXED-CORRECTNESS-PESSIMIZING-OVER-POSITION in Algorithm 3. Although this is sufficient for the brute-force and cubic proofs, for the subcubic proof using Algorithm 6 in Appendix F, we would additionally have to compute a log-loss bound for the sequences where the largest non-max token is 'too close' to the max token, which we currently neglect by considering the model to get them wrong in the worst case.

A.12 Proving upper bounds

In this work, we focus on proving lower bounds on model performance. Most of our theorems, for example in Appendices E, and F, prove two-sided bounds. Most of the other theorems can be straightforwardly adapted to proving upper bounds by swapping uses of min and max . Therefore, we expect that proving upper bounds on model performance should be straightforward.

A.13 What proof system?

Length of proof depends on what proof system we use. We permit any proof system where proofchecking time is linear in the length of the proof. This excludes dependently typed proof systems such as Martin-Löf type theory, but such proof systems can easily be accommodated by considering a proof-checking-trace rather than the proof object itself. Alternatively, a more conventional proof system like ZF, ZFC, or the proof system underlying Isabelle/HOL should suffice.

B Experimental details

B.1 Training details

To train each model, we generate 384,000 random sequences of 4 integers picked uniformly at random, corresponding to less than 2.5% of the input distribution. We use AdamW with batch_size = 128 , lr = 0 . 001 , betas = (0 . 9 , 0 . 999) , weight_decay left at the default 0 . 01 . We train for 1 epoch (3000 steps). Over our 151 seeds, models trained with this procedure achieve (99 . 92 ± 0 . 15) % train accuracy and a loss of (4 ± 8) × 10 -3 . 14 When qualitatively examining a single model (for example in Section 3.1 or Appendix H.1), we use the model with config seed 123, model seed 15 613947648.

As our models as sufficiently small, we did not have to use any GPUs to accelerate training our inference. Each training run takes less than a single CPU-hour to complete. In total, the experiments in this paper took less than 1000 CPU-hours.

We use the following software packages in our work: Paszke et al. [41], Plotly Technologies Inc. [42], Nanda and Bloom [37], Rogozhnikov [47], Virtanen et al. [52], McKinney [33], Waskom [55]

B.2 Additional details supporting our mechanistic interpretation of the model

We provide heatmaps of the matrices corresponding to the five components described/defined in Section 3, for the mainline model.

0

10

20

30

40

50

60

0

20

40

60

key token

query token

-

200

-

100

0

100

200

(a)

EQKE =

E

q

QK

T

T

0

1

2

3

0

20

40

60

key position

query token

-

3

-

2

-

1

0

1

(b)

EQKP =

E

q

QK

T

ˆ P

T

0

10

20

30

40

50

60

0

20

40

60

output logit token

input token

-

30

-

20

-

10

0

10

20

30

(a)

EVOU =

V OU

0

10

20

30

40

50

60

0

1

2

3

output logit token

input position

-

0

.

6

-

0

.

4

-

0

.

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

(b)

PVOU =

ˆ P

V OU

0

10

20

30

40

50

60

0

20

40

60

output logit token

input token

-

2

-

1

0

1

2

(c)

Direct Path

=

E

q

U

B.3 Distribution of model mechanisms

We provide some analysis of the distribution of the mechanisms of the models trained on the same configuration. At a glance, there is not that much variation across models.

The statistics of interest are: (1) σ 1 /σ 2 , the ratio of the first two singular values of EQKE , a measure of the extent to which the attention score computation is low-rank; (2) ¯ s , the average score (accuracy) of the model across the entire input distribution; (3) b cubic / ¯ s , the percent-score-recovered accuracy bound achieved by the cubic proof from Section 4.2; (4) b subcubic / ¯ s , the percent-score-recovered accuracy bound achieved by the (per-model best) 16 subcubic proof from Section 4.3.

For each statistic of interest, Table 6 presents an eleven-number summary of the statistic. Plots, seeds, and statistic values are shown for models whose values are closest to each of the corresponding summary statistics. 17 Additionally, each group contains a boxplot of the summary:

the minimum, maximum; the first and third quartiles; the median and mean; percentiles 2 . 15 % , 97 . 85 % , 8 . 87 % , and 91 . 13 % ; these are displayed as:

top and bottom of the vertical whisker lines; top and bottom of the box; horizontal line inside the box, and the square; horizontal whisker lines and whisker crosshatches.

seed

statistic of interest

EQKE

EQKP

EVOU

PVOU

E

q

U

distribution

value

E

q

QK

T

T

E

q

QK

T

ˆ P

T

V OU

ˆ P

V OU

-

400

-

200

0

200

400

-

40

-

20

0

20

40

123

(selected seed)

24262

400

600

800

1

,

000

σ

1

/σ

2

1061

.

1

13654

626

.

5

4810

350

.

0

6204

333

.

4

12457

0

.

992

0

.

994

0

.

996

0

.

998

1

¯

s

0

.

9999

19451

0

.

9992

15662

0

.

9932

16197

0

.

9915

32103

0

.

92

0

.

93

0

.

94

0

.

95

0

.

96

0

.

97

b

cubic

/

¯

s

0

.

971

6082

0

.

954

6155

0

.

926

2306

0

.

922

20415

0

.

6

0

.

65

0

.

7

0

.

75

0

.

8

b

subcubic

/

¯

s

0

.

80

29725

0

.

70

20976

0

.

63

6155

0

.

61

C Mathematical definitions

We provide a detailed breakdown of the mathematical notation used in the appendix.

Let

n

N

<n

σ

v

v

(

)

v

d

vocab elements; we write

n

1

when we care about the ele- be the finite set on

:=

0

,

1

, . . . , n

{

-

}

ments of

n

be the softmax function

i

e

e

/

∑

be the size of the vocabulary

d

be the dimension of the attention head, in our case, equal to the hidden dimension of the model d model (assumption: d < d vocab )

n ctx be the context length, the number of tokens in the input sequence, equal to K in Max-ofK

P

be the n ctx × d model positional embedding

E

Q

,

V

U

t

x

x

query

t

query be the d vocab × d model token embed

be the d model × d model query, key, value, and output matrices of the attention head be the d model × d vocab unembed matrix

be the input token sequence be the

ctx vocab

n

d

×

sequence

[

:=

1

-

x

:=

x

x

:=

1

-

t

token t max be the true maximum token in the input sequence, max i t i

[

t

0

, t

1

, . . . , t

n

ctx

-

1

]

one-hot-encoded input token

, x

1

, . . . , x

0

ctx

n

t

n

ctx

-

1

-

1

n

ctx

1

-

]

be the query token be the one-hot encoded query

M

ℓ

D

D|

X

f

P

avg

¯ P

P

ˆ P

q

,

,

E

q

of type X → Y be the model 18 ; sometimes we write for the logits of the model M ( t )

be a probability distribution over input-label pairs ( t , l ) ∈ X × L

be D marginalized to a distribution over t ∈ X of type L × Y → R be a scoring function for evaluating the performance of the model be the average position embeds across positions (of size d model ), 1 n ctx ∑ i P i

be either 1 n ctx ⊗ P avg or 1 d vocab ⊗ P avg depending on context - that is the result of broadcasting P avg back into the shape of P or E (that is, n ctx × d model or d vocab × d model )

be

, the broadcasting of be

1

d

P

vocab

-

¯ P

⊗

,

E

P

query

+

¯ P

, and

E

+

P

q

P

query respectively

For any vector-valued function v of length d vocab , parameterized over the input sequence t , let

∆ v i := v i -v t max be the difference between the i th element of v and the element of v corresponding to the true maximum of the input sequence t max .

For our particular model, we will have

Figure 11: Preliminary model definitions

We define the two typical performance functions corresponding to accuracy and log-loss. Note the shared subterm ∆ ℓ i .

̸

We present the model definition in four different regroupings to define via underbrace labels various useful quantities:

:=

K

,

,

O

D Brute-force proof

Theorem 1. For BRUTE-FORCE ( d vocab , n ctx , M ) as defined in Algorithm 1,

Proof. In fact the two sides of the inequality are equal by definition. Hence the inequality follows by reflexivity of ≥ .

E Details of cubic proof

In this section, we prove formally the result used in Section 4.2, A cubic proof.

At its heart, the convexity of softmax 19 is an extension to a simple idea: a weighted average of scalar values is extremized by putting 100% of the weight on an extremal value.

Using this simple version of the theorem, however, gives a useless bound of 0% accuracy: if we pay no attention to the maximum of the sequence, of course we're going to get the wrong answer. Since in fact the space of possible weightings we may see in practice is much smaller (finite, in fact, with at most d vocab n ctx values), we may look for a more general version of this idea that gives us tighter bounds that still cover the space of possible weightings.

The weights are not linearly independently choosable (softmax is non-linear), so extremal values do not necessarily result from putting maximal attention on the worst token. It may be, when trying to find the worst case, that some positions are so dis-preferred that it makes more sense to choose a token that is 'less bad' for those positions, if it draws enough attention away from the correct token. See Lemma 3 for details.

We thus spend this section characterizing a relaxation of the constraints on weights:

that contains all actually possible weightings,

that is extremized at weights that still correspond to some notion of 'put the most weight on the extremal tokens', and

for which computing the extremal weightings is computationally efficient.

Before diving in, let's recall the proof that a weighted average of scalar values is extremized by putting 100% of the weight on extremal values:

Theorem 2 (Warmup: Extremizing weighted averages) . Fix a set of values v i ∈ R . The weighted average is bounded by the extremal values: for any w i such that ∑ i w i = 1 and 0 ≤ w i ≤ 1 ,

Proof. The proof is simple. We have

̸

and so the result follows.

E.1 Proof strategy

The model computes the true maximum t max when its outputs logits ℓ are such that ∆ ℓ t ∗ := ℓ t ∗ -ℓ t max < 0 for all t ∗ = t max . 20 As a result, it suffices to lower-bound the proportion of sequences where (an upper bound on) ∆ ℓ t ∗ is negative for all t ∗ = t max . In particular, we will upper-bound the contribution from incorrect tokens t in positions i to the difference ∆ ℓ i between incorrect ( t ∗ ) and correct ( t max ) output tokens ∆ ℓ i t ∗ = ℓ i t ∗ -ℓ i t max .

̸

We do this by arguing that the logit difference ∆ ℓ t ∗ satisfies a certain notion of convexity over the space of a relaxation of sequences (Theorem 6), and constructing a set of Θ( d vocab 3 n ctx ) 'extremal' relaxed sequences where the position and token embedding components of attention are pessimized independently.

We start by first rewriting the contribution of each token through the attention head to the logit difference into the contributions involving PVOU and EVOU :

We then upper bound ∆ ℓ PVOU ,i t ∗ ( t ) by noting that because the softmax attention is a weighted average of PVOU ,

Since ∑ i α ∗ i ( t ) = 1 , we have

We then construct a set Ξ pure of 'pure sequences' consisting of only three types of tokens in one of two orders, and show that for each input sequence t and readoff logit t ∗ , we bound the logit difference from the token embeddings ∆ ℓ EVOU ,i t ∗ ( t ) using a small subset X of Ξ pure :

We construct a set X relaxed of relaxed sequences, where each relaxed sequence t relaxed consists of a sequence and a position ( t , i ) , where ∆ ℓ t ∗ ( t , i ) is evaluated by separately considering the positional

contribution through attention (that is, the attention weighted PVOU ) and the token contribution (that is, the attention-weighted EVOU ) and direct contribution (the logit difference through the skip connection EU ). Note that i indicates the position that we pay 100% of the attention to for the PVOU contribution.

We argue that ∆ ℓ t ∗ ( t , i ) satisfies a certain notion of convexity over mixtures of sequences, such that we can evaluate it only on a set of Θ( d vocab 3 n ctx ) 'extremal' sequences in a way that takes O ( d vocab 3 n ctx ) total time to bound ∆ ℓ t ∗ ( t , i ) for every possible input sequence. We then use the extremal sequences that the model gets correct to lower bound the proportion of all sequences that the model will get correct. Specifically, we argue that Algorithm 3 provides a valid lower bound on the proportion of sequences the model gets correct.

E.2 Proof outline

We now proceed to the main results of this section.

Math fact: For each token t ∗ , the logit difference ∆ ℓ t ∗ for any sequence t can be decomposed into the direct contribution from the embeds ℓ EU , the attention-weighted position contribution ( PVOU ), and the attention-weighted token contribution ( EVOU ). Therefore, it suffices to upper bound each of the three components independently, since summing these upper bounds gives a valid upper bound on the logit difference.

We can compute the direct contribution ℓ EU exactly by first computing EU = E q U and then, for each max, subtracting the logit of the max token from each row of the matrix. No theorems needed. For each max token, we can bound the position contribution by its maximum over positions (Theorem 6).

In order to upper bound the token contribution, we argue that any mixed sequence will be upper bounded by the maximum of the corresponding pure sequences (Theorem 7). We then argue that for pure sequences, it suffices to consider orderings where same tokens appear contiguously (Theorem 4).

E.3 Formal proof

For this subsection, all theorems are parameterized over the following quantities.

Definition 1 (Common theorem parameters) . Fix a token value function (à la a row difference in EVOU ) v : N <d vocab → R and a token attention function (à la EQKE for a fixed query token) a : N <d vocab → R . Fix a position value function (à la a row difference in PVOU ) w : N <n ctx → R and a position attention function (à la EQKP for a fixed query token) b : N <n ctx → R .

In practice, we'll take, for fixed query token t query , fixed output token of interest t ∗ , and fixed maximum token t max ,

Definition 2 (of a sequence via sorted tokens and a position permutation) . We can define a sequence of tokens via sorted tokens and a position permutation by specifying a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N <d vocab paired with a permutation σ : N <n ctx → N <n ctx .

Definition 3 (sequence score) . Given a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N <d vocab and a permutation σ : N <n ctx → N <n ctx define the sequence score s t 0 ,...,t n ctx -1 ,σ as:

We will drop the token subscript, writing only s σ , when the token values are unambiguous by context.

The sequence score here will be computing ∆ ℓ EVOU t ∗ for some fixed t ∗ and t max . The way we've set up our definitions, high scores predict t ∗ (and are thus bad), negative scores predict t max (and are thus good), and more negative the scores, the stronger the prediction of t max .

Definition 4 (swap permutation) . Given a permutation σ : N <n ctx → N <n ctx of the n ctx positions and two indices 0 ≤ i, j < n ctx , define the swap permutation σ i ↔ j to be the permutation that is σ

except swapping i and j :

Define ∆ σ,i ↔ j to be the difference in sequence scores when you swap i and j :

Lemma3 (Characterization of swapping tokens) . Fix a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N . Fix σ : N → N be a permutation of the n ctx positions. Fix indices 0 ≤ i, j < n ctx .

Then there are two cases for sign (∆ σ,i ↔ j ) :

Intuitively, Lemma 3 says that, if the token contribution to attention is equal between tokens t i and t j , then the impact of swapping their positions σ ( i ) and σ ( j ) is entirely determined by how much attention is paid to the positions of i and j and the relative difference in their value. (Notably, by swapping these tokens, we don't affect the attention paid on other tokens, and so the effect of the change does not depend on the values of the other tokens.) Alternatively, if the attentions are not equal, then swapping the positions changes the allocation of attention to other tokens in the sequence, and so it may the case that this change in allocation in attention dominates the attention-weighted values of these two tokens.

Proof. First note that the theorem is trivial for i = j .

̸

For the rest of the proof, we take i = j .

The proof proceeds just by algebraic manipulation with no deep insight. We first list the facts we use, the proceed to computing sign (∆ σ,i ↔ j ) . We abbreviate σ i ↔ j as σ ′ for brevity.

sign (∆ σ,i ↔ j ) = sign ( s σ ′ -s σ

Now multiply through by the denominator, which is positive

Divide through by non-zero values when possible

Definition 5 ( σ fixes F ) . Fix a set of fixed indices F ⊆ N <n ctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . ( F is the set of positions for which we are not pessimizing over the value of the token in that position.) Fix a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N .

Given a permutation σ : N <n ctx → N n ctx , say that σ fixes F (relative to t 0 , . . . , t n ctx -1 ) if t i = t F ( σ ( i )) whenever σ ( i ) ∈ F .

Note that in this section, for the cubic proofs, we will in fact generally take F = { n ctx -1 } , so that we are fixing the final query token, though in Theorems 7, 8, and 9 F will also contain all positions with the maximum token t max . In Appendix F, we will take F = ∅ or F to be the set of positions of the maximum token. However, none of these theorems are specific to F being subsingleton, and we prove them in generality.

Definition 6 (position-sorting permutation) . Fix a set of fixed indices F ⊆ N <n ctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab .

Define the position-sorting permutation fixing indices in F σ s : N <n ctx → N <n ctx to be the permutation that sorts the indices not in F according to b : for 0 ≤ i, j < n ctx with i, j ̸∈ F , b i ≤ b j whenever σ s ( i ) < σ s ( j ) ; and σ s ( i ) = i for i ∈ F .

Definition 7 (contiguous on equal tokens) . Fix a set of fixed indices F ⊆ N <n ctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N .

̸

Say that the sequence represented by a permutation σ : N <n ctx → N <n ctx is contiguous on equal tokens if, for all 0 ≤ i, j, k < n ctx with t i = t j = t k and i, j, k ̸∈ σ -1 ( F ) , it is never the case that σ s ( σ ( i )) < σ s ( σ ( k )) < σ s ( σ ( j )) .

Theorem 4 (Pessimization over sequence ordering is possible and results in contiguous sequences) . Fix a set of fixed indices F ⊆ N <n ctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N .

Let σ min , σ max : N → N be permutations of the n ctx positions, fixing positions in F , satisfying the following property: For all σ : N → N a permutation fixing F , we have

(Such permutations are guaranteed to exist because the permutation group on n ctx elements is finite.)

Then σ max and σ min may be taken to be contiguous on equal tokens. That is, there exist σ max and σ min satisfying the property of Equation 10 which additionally satisfy the definition of Definition 7.

The basic idea is that we will assume that one of σ max and σ min cannot be contiguous on equal tokens and derive a contradiction. We will pick the extremal permutation that is closest to being contiguous, take a contiguity violation, and then show that either we can correct the contiguity violation without changing the score-thus violating the presumption that the permutation is closest

to being contiguous-or we will find one swap of indices that decreases the score and another swap of indices that increases the score, thus violating the presumption of extremality.

In slightly more detail, but still informally, we will consider the sign of the difference between scores of our purported extremal permutation and a permutation that has swapped some indices. The theorem follows from showing that there exists a triple of indices i , j , k such that the sign of the score difference from swapping i and j is different from the sign of the score difference from swapping j and k .

First, a definition and some helpful facts about it.

Definition 8 (contiguous on equally-attended positions) . Fix a set of fixed indices F ⊆ N <n ctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N .

Say that a permutation σ is contiguous on equally-attended positions if, for all 0 ≤ i < n ctx with i ̸∈ σ -1 ( F ) , the sorting order according σ s on the contiguous block of positions with contribution to the attention score equal to that of σ ( i ) , { σ ( j ) ∣ ∣ b σ ( j ) = b σ ( i ) and σ ( j ) ̸∈ F } , is the same as the sorting order according to the fraction of tokens equal to t j with b -values greater than b σ ( i ) , with ties broken by the value of t j . Equationally, this second sorting order is defined by the score

̸

Most importantly, any permutation that is contiguous on equally-attended positions has the property that for any indices 0 ≤ i, j, k < n ctx with i, j, k ̸∈ σ -1 ( F ) and t i = t j = t k and σ s ( σ ( i )) < σ s ( σ ( k )) < σ s ( σ ( j )) , we will have the strict inequality b σ ( i ) < b σ ( k ) < b σ ( j ) . Additionally, we may always sort equally-attended positions to make any permutation contiguous on equally-attended positions.

We will define an additional notion of contiguity-violations which we avoid up-front by arbitrarily swapping involved indices without changing the score s σ .

Definition 9 (needlessly non-contiguous) . Fix a set of fixed indices F ⊆ N <n ctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N .

̸

Say that a permutation σ is needlessly non-contiguous at i , j , k (for i, j, k ̸∈ σ -1 ( F ) ) if ∆ σ,i ↔ k = 0 or ∆ σ,j ↔ k = 0 , for 0 ≤ i, j, k < n ctx with i, j, k ̸∈ σ -1 ( F ) with t i = t j = t k and σ s ( σ ( i )) < σ s ( σ ( k )) < σ s ( σ ( j )) .

Say that a permutation σ is needlessly non-contiguous if it is needlessly non-contiguous at any i, j, k ̸∈ σ -1 ( F ) .

Lemma 5. Fix a set of fixed indices F ⊆ N <n ctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N .

Any needlessly non-contiguous sequence σ which fixes F can be made into a sequence σ ′ which still fixes F and is both simultaneously contiguous on equally-attended positions and not needlessly non-contiguous, and for which s σ = s σ ′ .

Proof. First, sort regions of equally-attended positions to make σ contiguous on equally-attended positions. If the resulting permutation is not needlessly non-contiguous, then we are done.

̸

Otherwise, we have ∆ σ,i ↔ k = 0 or ∆ σ,j ↔ k = 0 for some i , j , k , for 0 ≤ i, j, k < n ctx with i, j, n ̸∈ σ -1 ( F ) and t i = t j = t k and σ s ( σ ( i )) < σ s ( σ ( k )) < σ s ( σ ( j )) . Since the sequence is contiguous on equally-attended positions, we have the strict inequality b σ ( i ) < b σ ( k ) < b σ ( j ) .

By Lemma 3, we have two cases. Noting that t i = t j , we can write them as

̸

In the first case, we may fully freely interchange tokens equal to t i with tokens equal to t k without changing the score; in this case we may use the token value as a sorting tie-breaker and swap tokens until there are no more needlessly non-contiguous triples falling into case (1).

In the second case, since swapping tokens does not change s σ , the property will continue to hold for these tokens after the swap. We may then swap tokens, again using token value as a tie-breaker, until there are no more needlessly non-contiguous triples falling into case (2).

We can now finally make our argument for Theorem 4 more precise.

̸

Proof of Theorem 4. Choose σ max and σ min to be contiguous on equally-attended positions and not needlessly non-contiguous, and suppose that we have σ ∈ { σ max , σ min } such that for some 0 ≤ i, j, k < n ctx with i, j, k ̸∈ σ -1 ( F ) and t i = t j = t k , we have b σ ( i ) < b σ ( k ) < b σ ( j ) . We will derive a contradiction with the presumption that σ is extremal by showing that we can swap i and k to change the score in one direction and that we can swap j and k to change the score in the other direction.

Take σ ′ 0 to be σ but swapping i and k , and take σ ′ 1 to be σ but swapping j and k .

̸

Nowwewill consider the cases for the sign of the score difference ∆ 0 := s σ ′ 0 -s σ and ∆ 1 := s σ ′ 1 -s σ . By the presumption of not being needlessly non-contiguous, ∆ z = 0 for z ∈ { 0 , 1 } . If we can show that the sign of ∆ 0 is distinct from the sign of ∆ 1 , then we will have a contradiction with extremality because we will have either s σ ′ 0 < s σ < s σ ′ 1 or s σ ′ 1 < s σ < s σ ′ 0 . That is, we would be able to swap i ↔ k and j ↔ k to get a lower and higher score, making σ not extremal.

Noting that t i = t j ,

Noting that the product is non-zero by presumption, that right multiplicand is equal for ∆ 0 and ∆ 1 , and sign ( b σ ( i ) -b σ ( k ) ) = -1 and sign ( b σ ( j ) -b σ ( k ) ) = 1 , we have our desired contradiction.

Note that the proof of Theorem 4 does not go through if we include the position value function w in the score, because we may trade off the position value function against the token value function. We now show that we can independently pessimize over positional attention.

Definition 10 (full sequence score) . Given a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N <d vocab and a permutation σ : N <n ctx → N <n ctx define the full sequence score s ′ t 0 ,...,t n ctx -1 ,σ as:

We will drop the token subscript, writing only s ′ σ , when the token values are unambiguous by context.

The sequence score here will be computing ∆ ℓ EPVOU t ∗ := ∆ ℓ EVOU t ∗ +∆ ℓ PVOU t ∗ for some fixed t ∗ and t max . As with Definition 3, with the way we've set up our definitions, high scores predict t ∗ (and are thus bad), negative scores predict t max (and are thus good), and more negative the scores, the stronger the prediction of t max .

Definition 11 (relaxed sequence score) . Given a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N <d vocab and a permutation σ : N <n ctx → N <n ctx define the relaxed sequence scores r t 0 ,...,t n ctx -1 ,σ, min and r t 0 ,...,t n ctx -1 ,σ, max as:

We will drop the token subscript, writing only r σ, min or r σ, max , when the token values are unambiguous by context.

Theorem 6 (Independent pessimization over positional contributions is possible) . Fix non-decreasing sequences of tokens t 0 ≤ ·· · ≤ t n ctx -1 ∈ N and t ′ 0 ≤ ·· · ≤ t ′ n ctx -1 ∈ N and permutations σ, σ ′ : N <n ctx → N <n ctx . Let r σ, min and r σ, max denote r t 0 ,...,t n ctx -1 ,σ, min and r t 0 ,...,t n ctx -1 ,σ, max ; let s σ denote s t 0 ,...,t n ctx -1 ,σ ; and let s σ ′ and s ′ σ ′ denote s t ′ 0 ,...,t ′ n ctx -1 ,σ ′ and s ′ t ′ 0 ,...,t ′ n ctx -1 ,σ ′ .

Then we have

That is, the difference between the relaxed sequence score and the sequence score of any given sequence always bounds the difference between the full sequence score and the sequence score for any (related or unrelated) sequence.

Proof. This proof follows straightforwardly from the softmax weighting being an affine weighting. We have

Since e x is non-negative for all real x , we have

Thus we get as desired

Note that we could prove a more fine-grained theorem, that pessimizes over attention paid to positions only for sequences compatible with the chosen fixed tokens F and t F , but since the positional contribution is so small we do not bother.

Theorem 7 (For a fixed ordering, softmax is convex over token counts and only pure sequences need be considered) . Fix a set of fixed indices F ⊆ N <n ctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a set S ⊆ N <d vocab of valid other tokens in the sequence. (In our uses of this theorem, S will be the largest subset of N <t max for which we can guarantee that the model behaves correctly on all sequences compatible with F and t max and with tokens otherwise drawn from S .)

Define a comparison on non-negative integers less than d vocab :

Let t cmpmin and t cmpmax be the minimum and maximum elements of S according to cmp . 21

For a given choice of a non-decreasing sequence of tokens t 0 ≤ · · · ≤ t n ctx -1 ∈ N compatible with F and S and a given choice of permutation σ : N → N of the n ctx positions fixing F ( t i = t F ( σ ( i ))

for σ ( i ) ∈ F ; and t i ∈ S for σ ( i ) ̸∈ F ): let s σ, min (and s σ, max ) denote s t 0 ,...,t n ctx -1 ,σ when t i = t cmpmin for all σ ( i ) ̸∈ F (or t cmpmax , respectively).

Then for all such choices of sequence-permutation pairs,

This theorem follows by chaining two lemmas: that scores are extremized by considering pure sequences, and that the extremal pure sequences match the comparison function defined in the theorem statement.

Lemma 8 (Sequences scores are extremized on purer sequences) . Fix all the same quantities as in Theorem 7.

̸

For any indices 0 ≤ i < j < n ctx , token values x, y ∈ S , the score for a sequence with t i = x = y = t j is bounded on both sides by sequences with t i = t j = x and t i = t j = y .

Proof. Let s α,β be the sequence score with t i = α and t j = β , and define the score differences ∆ x := s x,x -s x,y and ∆ y := s y,y -s x,y . It suffices to show that sign(∆ x ∆ y ) ≤ 0 . To show this, we must only compute the sign of ∆ α for α ∈ { x, y } and show that whenever both ∆ x and ∆ y are non-zero, they have opposite signs.

We proceed by computation after defining some convenience variables for brevity:

̸

̸

Multiply through by positive denominators and simplify

Pulling out e b σ ( i ˜ α )

Note that swapping α and ˜ α negates the sign. Hence, we have sign(∆ x ) = -sign(∆ y ) and hence s x,x ≤ s x,y ≤ s y,y or s y,y ≤ s x,y ≤ s x,x as desired.

Lemma 9 (Pure sequences are sorted according to cmp in Theorem 7) . Fix all the same quantities as in Theorem 7.

Fix tokens x, y ∈ S . Let n := n ctx -| F | be the number of non-fixed tokens. Fix sequences with n copies of x and y respectively: fix t x, 0 ≤ · · · ≤ t x,n ctx -1 ∈ N and t y, 0 ≤ · · · ≤ t y,n ctx -1 ∈ N compatible with F and S and given choices of permutations σ x , σ y : N → N of the n ctx positions fixing F : t x,i = t F ( σ x ( i )) for σ x ( i ) ∈ F ; t y,i = t F ( σ y ( i )) for σ y ( i ) ∈ F ; t x,i = x for σ x ( i ) ̸∈ F ; and t y,i = y for σ y ( i ) ̸∈ F .

Then

Proof. The proof goes by straightforward computation.

Multiply through by non-negative denominators

Use f > 0

Corollary 10. Define the relation ≤ cmp by x ≤ cmp y if and only if cmp( x, y ) ∈ {-1 , 0 } . The relation ≤ cmp is always transitive.

Proof. Note that by Lemma 9, cmp is comparing two sequence scores. Since ≤ is transitive over the reals, the relation ≤ cmp is also transitive.

Finally, we combine the previous lemmas to complete our proof of Theorem 7:

Proof of Theorem 7. Extremal sequences with scores s σ, min and s σ, max are guaranteed to exist because there are only finitely many elements of S and therefore only finitely many sequences. By Lemma 8, the extremal sequences must be pure (have t i = t j whenever σ ( i ) , σ ( j ) ̸∈ F ). By Lemma 9, the extremal sequences must have tokens that are extremal according to cmp .

We now have all the tools necessary to prove the following theorem. We refer to Algorithm 3 and Algorithm 4 or the proof of Theorem 11 for a definition of the CUBIC algorithm.

Theorem 11.

Before we give the proof of this theorem, we introduce some helpful notation.

Definition 12. Fix an element ( r m , r q , c ) ∈ { 0 , ..., d vocab } 2 ×{ 0 , ... 3 } such that r m ≥ r q . We define X ( r m ,r q ,c ) to be the set of tokens t such that

1. The max token t max is equal to r m ,

Algorithm 2 Counting Correct Sequences in Cubic Time: Preliminaries

̸

1: function CORRECTNESS( M , input-sequence) 2: return MODEL-BEHAVIOR( M , input-sequence) == MAX(input-sequence) 3: end function 4: function MODEL-BEHAVIOR( M , input-sequence) Require: input-sequence is a tensor of shape ( n ctx , ) with values in N <d vocab 5: t max ← MAX ( input-sequence ) ▷ t max ← max-token 6: t ← input-sequence 7: skip-score t ∗ ← ∆ ℓ EU t ∗ ( t n ctx -1 , t max ) 8: attn-weights-unscaled i ← EQKE( t n ctx -1 , t i ) + EQKP( t n ctx -1 , i ) 9: attn-weights ← SOFTMAX ( attn-weights-unscaled / √ d ) 10: v t ← EVOU( t ) 11: w i ← PVOU( i ) 12: ∆ v t,t ∗ ← v t,t ∗ -v t,t max 13: ∆ w i,t ∗ ← w i,t ∗ -w i,t max 14: return max t ∗ = t max ( skip-score t ∗ + ∑ n ctx -1 i =0 (∆ v i,t ∗ +∆ w i,t ∗ ) · attn-weights i ) 15: end function 16: function CORRECTNESS-PESSIMIZING-OVER-POSITION-SLOW( M , input-sequence) 17: t ← input-sequence 18: return ALL(CORRECTNESS( M , perm + [ t -1 ]) for all perm ∈ PERMUTATIONS( t 0: -1 )) 19: end function

The query token t query is equal to r q ,

The cardinality of tokens that are not at the query position and not equal to t max is equal to c .

For clarity, we list all the possible cases. We always take t query ≤ t max and let S 3 act on sequences by permuting the first three factors (i.e. keeping the query position fixed).

Definition 13. Let t ∈ X be a sequence. We say t is pure, if it has at most three distinct tokens: the max token t max , the query token t query , and optionally a third token t ∗ < t max .

We denote by X pure the subset of pure tokens. For any subset Y ⊂ X , we set Y pure := Y ∩ X pure .

We now come to the proof of Theorem 11. We will show how to use the previous theorems to get explicit bounds and explain how CUBIC ( d vocab , n ctx , M ) computes these bounds.

Proof of Theorem 11. First of all, we note that the algorithm CUBIC = CUBIC ( d vocab , n ctx , M ) yields a lower bound for the accuracy on the set X ( t max ,t query ,c ) . We can therefore compute the bound on X = ∐ ( t max ,t query ,c ) X ( t max ,t query ,c ) by computing it for each such choice ( t max , t query , c ) and summing over them

So from now on we will fix one such subset X ( t max ,t query ,c ) .

We begin by defining a map

which sends a sequence to the subsequence of elements which are not at the query position and not equal to t max . Then Theorem 7 can be restated as follows 22 :

Algorithm 3 Counting Correct Sequences in Cubic Time, Part I. Lines are annotated with comments indicating the parameters for a cache to avoid duplicate computations.

1: function MODEL-BEHAVIOR-RELAXED( M , query-tok, max-tok, non-max-tok, n-copiesnonmax)

2: t ← query-tok, t ← max-tok, t ′ ← non-max-tok, c ← n-copies-nonmax

̸

25: end if

26: end function

27: function RELAXED-CORRECTNESS-PESSIMIZING-OVER-POSITION( M , t query , t max , t ′ , c )

28: ▷ runs the model on a relaxed variant of input sequences compatible with the arguments Ensure: return is False if CORRECTNESS-PESSIMIZING-OVER-POSITION-SLOW( M , t ) is False for any t with specified t query , c copies of t ′ in non-query positions, and the remainder of the tokens equal to t max

29: return MODEL-BEHAVIOR-RELAXED ( M , t query , t max , t ′ , c ) < 0

30:

end function

̸

̸

̸

̸

Now instead of computing the output of the model for every element f -1 ( S c ) pure , we use Theorem 4 (combined with Theorem 6) to run a relaxed version of this. In particular, we may assume that the pure sequence is contiguous on equal tokens. Here contiguous on equal tokens means that for the positional part of the attention (i.e. the EQKP part), we have either b t max < { b i , b j } or b t max > { b i , b j } , where i, j ∈ { 0 , ..., n ctx -1 } are indices of tokens not equal to t max .

For the algorithm CUBIC ( d vocab , n ctx , M ) we fix a t ∗ ∈ { 0 , ..., t max -1 } (unless c = 0 , in which case there is no such choice). We then run the relaxed accuracy computation RCPOP ( t query , t max , t ′ , c ) as described in Theorem 6. If RCPOP ( t query , t max , t ′ , c ) < 0 , we add t ′ to S . If we do, we add t ∗ to S . Therefore by construction of S we know that we get full accuracy on f -1 ( S c ) pure and therefore we get full accuracy on f -1 ( S c ) .

Now we count the cardinality of f -1 ( S c ) and add it to the count of correct sequences.

Theorem 12. The running time of Algorithm 3, after using caching to avoid duplicate computations, is O ( d vocab 3 n ctx 2 ) .

̸

Proof. The nested loops in CUBIC execute the innermost body O ( d vocab 2 n ctx ) times, and the summation on Line 13 costs O ( n ctx ) per iteration. What remains is to show that the call to RELAXED-CORRECTNESS-PESSIMIZING-OVER-POSITION ( M , t query , t max , t ′ , c ) costs O ( n ctx ) when c = 0 and at most O ( d vocab n ctx ) when c = 0 and t ′ = t max .

The matrix multiplications in EQKE , EQKP , EVOU , PVOU , and ℓ EU can be cached upfront, costing O (max( d vocab , d model , n ctx ) 2 d model ) ≤ O ( d vocab 3 ) since we assume d vocab > d model and d vocab > n ctx .

The sorting on Line 9 can also be cached upfront (per t query ), costing O ( d vocab n ctx log n ctx ) .

Note that each variable assignment in RELAXED-CORRECTNESS-PESSIMIZING-OVER-POSITION can be cached into a table parameterized over at most three variables which range over d vocab and over at most two variables that range over n ctx .

What remains is the return statements.

̸

̸

When c = 0 , we have on Line 24:

̸

̸

̸

̸

We can cache max t ∗ = t max skip-score t ∗ per t max and t query , costing O ( d vocab 3 n ctx ) . We can cache max t ∗ = t max (∆ w max ,t ∗ + ∆ v i,t ∗ ) per t max and t ′ costing O ( d vocab 3 ) , since each ∆ v i,t ∗ will be ∆ v t,t ∗ for some t ∈ { t query , t max , t ′ } . Finally, we can compute the summation in cost O ( n ctx ) per loop iteration, as required.

F Quadratic counting for a sub-cubic proof

In this section we fill in the details lacking from Section 4.3.

In Appendix E we proved an intricate version of convexity of softmax where, modulo pessimizing in unrealistic ways over the attention paid to positions for the computation done on positional encodings, all extremal relaxed sequences correspond to actual sequences.

When we only get a budget of O ( d vocab 2 n ctx ) extremal relaxed cases to consider, though, we must pessimize more, which gives us a simpler version of the convexity theorem and proof. Notably, when we restrict our sequences to have only two tokens (the max token t max and the non-max token t ′ ), most of the theorems from Appendix E.3 get significantly simpler.

Additionally, we must pessimize separately over the token value ( v ) and token attention ( b ) computations in order to allow efficient computation (Theorem 15).

F.1 Proof of baseline sub-cubic result

For this subsection, all theorems are parameterized over the following quantities.

Definition 14 (Common theorem parameters) . Fix a total number of tokens n ctx . Fix a token value function (à la a row-difference in EVOU ) v : N <d vocab → R and a token attention function (à la EQKE for a fixed query token) a : N <d vocab → R . Fix a position value function (à la a row-difference in PVOU ) w : N <n ctx → R and a position attention function (à la EQKP for a fixed query token) b : N <n ctx → R .

In practice, as in Appendix E.3, we'll take, for fixed query token t query ,

̸

Note that unlike in Appendix E.3, we pessimize independently over the query token and the non-max token, so the 'fixed' query token may not in fact appear in any key-side position in the relaxed sequence we consider.

Definition 15 (of a sequence via mapping from positions) . We can define a sequence of tokens via mapping from positions by specifying a subset of valid tokens S ⊆ N <d vocab paired with a function T : N <n ctx → S specifying which token is in each position.

Definition 16 (sequence score) . Given a subset of valid tokens S ⊆ N <d vocab and a function T : N <n ctx → S specifying which token is in each position, define the sequence score

Definition 17 (swapped mapping) . Given a subset of valid tokens S ⊆ N <d vocab and a function T : N <n ctx → S specifying which token is in each position and two indices 0 ≤ i, j < n ctx , define the swapped mapping T i ↔ j be the function that is T except swapping i and j :

Lemma 13 (Characterization of swapping tokens in a two-token sequence) . Fix two tokens t 0 < t 1 ∈ N and a function T : N <n ctx →{ t 0 , t 1 } specifying which token is in each position.

Define ∆ T,i ↔ j to be the difference in sequence scores when you swap i and j :

Then

Proof. Lemma 3 gives us the result directly when a T ( i ) = a T ( j ) . Otherwise, we get

Hence all that remains is to show that

Define ¯ v := 1 2 ( v T ( i ) + v T ( j ) ) and define ∆ v := 1 2 ( v T ( i ) -v T ( j ) ) so that v T ( i ) = ¯ v + ∆ v and v T ( j ) = ¯ v -∆ v . Assume WLOG that T ( i ) = 0 and T ( j ) = 1 so that v T ( p ) = ¯ v +( -1) T ( p ) ∆ v for all p .

Then we have

Define so that we get

Multiply through by the positive denominator and expand out so that we get

Theorem 14 (Pessimization over sequence ordering for two-token sequences is simple) . Let σ s : N → N denote a permutation of the n ctx positions that sorts them according to b : for 0 ≤ i, j < n ctx , b i ≤ b j whenever σ s ( i ) < σ s ( j ) . Fix two tokens t 0 < t 1 ∈ N .

Let n t 0 be the number of p ∈ [0 , n ctx ) with T ( p ) = t 0 and let n 1 be the number of p ∈ [0 , n ctx ) with T ( p ) = t t 1 . Note that n t 0 + n t 1 = n ctx .

Define t min := argmin t ∈{ t 0 ,t 1 } v t and define t max := argmax t ∈{ t 0 ,t 1 } v t .

Define T min , T max : N <n ctx →{ t 0 , t 1 } to be the assignment of tokens to positions that pays the least (respectively, most) attention to t max :

Proof. The extremality of s T min and s T max follows straightforwardly from Theorem 4.

All that remains is s T min ≤ s T max .

This follows from noting by Lemma 13 that swapping two tokens in T min increases the sequence score, while the reverse is true of s T max , thus showing that it must be s T min that is the minimum and s T max that is the maximum and not vice versa.

Definition 18 (full sequence score) . Given a subset of valid tokens S ⊆ N <d vocab and a function T : N <n ctx → S specifying which token is in each position define the full sequence score s ′ T :

Theorem 15 (Independent pessimization over positional contributions and token attention and token value is possible) . Fix two tokens t 0 < t 1 ∈ N . Let T min , T max : N <n ctx →{ t 0 , t 1 } and t max , t min be as in Theorem 14. Fix a set S of valid tokens with t 0 , t 1 ∈ S .

Define relaxed versions T ′ max , T ′ min : N <n ctx → S of T max and T min :

̸

̸

That is, T ′ max replaces t min with whatever token in S draws the least attention away from t max , while T ′ min replaces t min with whichever token in S draws the most attention away from t max .

Then we have that

Define relaxed extremal sequence scores r T max , r T min :

Then r T min ≤ s ′ T min and s ′ T max ≤ r T max .

Proof. (sketch) Essentially the same as the proof of Theorem 6.

Note that in practice, we take S to be the set of all tokens less than t max -g for some minimum gap g . This allows us to share computation across the various maximum tokens to reduce overall computational complexity.

Algorithm 5 Counting Correct Sequences in Subcubic Time, Preliminaries

̸

̸

̸

Algorithm 6 Counting Correct Sequences in Subcubic Time

̸

̸

35: ▷ runs the model on a relaxed variant of input sequences compatible with the arguments Ensure: CORRECTNESS-PESSIMIZING-OVER-GAP-SLOW is False = ⇒ result is False

Ensure: return is False if CORRECTNESS-PESSIMIZING-OVER-GAP-SLOW( M , t ) is False for any t with specified t max , t query , and c tokens not equal to t max

36:

37:

return

MODEL-BEHAVIOR-RELAXED-OVER-GAP

end function

(

M

, t query

, c, g, g

)

<

0

∗

̸

Theorem 16. For all G ,

Proof. (sketch) Apply preceding lemmas and theorems to Algorithm 6

Theorem 17. The running time of Algorithm 6, after using caching to avoid duplicate computations, is O ( d vocab 2 d model + d vocab 2 n ctx 2 ) .

Proof. (sketch) Sum the complexities indicated along the right side of Algorithm 3. The d vocab 2 d model term comes from the precomputing EVOU , EU , and EQKP . The d vocab 2 n ctx 2 term comes from the softmax over n ctx tokens for O ( d vocab 2 n ctx ) pessimized pure sequences. Confirming that none of the complexities on the right side exceeds O ( d vocab 2 d model + d vocab 2 n ctx 2 ) completes the proof.

G Subcubic proof strategies

In this section, we present a number of proof strategies that we use to reduce the computational cost of the proof, ultimately driving down the cost of EU and EQKE verification to O ( d vocab d model ) , while unfortunately leaving the cost of EVOU verification at O ( d vocab 2 d model ) .

The three main tricks we cover are the mean+diff trick (Appendix G.1), the max row-diff trick (Appendix G.2.2), and the rank one / rank two SVD decomposition of EQKE (Appendix G.2.3). While the mean+diff trick is useful for getting slightly better bounds, the SVD decomposition of EQKE is the place where we get to insert the most understanding (without which we'd have no hope of non-vacuous bounds below O ( d vocab 2 d model ) ), and the max row-diff trick is the workhorse that allows us to drive down the error term computations from cubic to quadratic without getting completely vacuous bounds.

G.1 The mean+diff trick

Suppose we have quantities f x,y and g y,z and we want to pessimize (WLOG, suppose minimize) the quantity f x,y + g y,z over x , y , and z in time less than O ( n x n y n z ) , say we allow O ( n x n y + n y n z + n x n z ) . Also suppose the variation of f over the y axis is much larger than the variation of f over the x-axis.

We can of course say

But we can do better!

Note that

Suppose that f x,y varies much less over x than it does over y , and much less than g y,z varies over either of y and z . This will make the following bound a good approximation, though the bound is sound even without this assumption. We can write

By averaging the variation over certain axes, we have

Theorem 18 (Mean+Diff) .

and the RHSs can be computed in time O ( n x n y + n y n z + n x n z ) for n x , n y , and n z the number of possible values of x , y , and z , respectively.

Example for how this helps with small variation:

Take any function k ( y ) and then take

Then we have

If ε 1 and ε 2 are small compared to min y k ( y ) -max y k ( y ) , then using E x f x,y gives a much better bound.

Note, though, that this could be a worse bound if the assumption of small variation does not hold.

Note also that this trick is not restricted to adding and subtracting E x f x,y . If f is a matrix indexed by x and y , we might also try taking SVD and using the first principal component instead. A basic application of the triangle inequality gives the following, more general, result:

Theorem 19 (Summarize+Diff) . For any h y which can be computed in time O ( n h ) ,

and the RHSs can be computed in time O ( n x n y + n y n z + n h ) for n x , n y , and n z the number of possible values of x , y , and z , respectively.

We see that if the variation of f in the x -axis is indeed much smaller than the variation in the y -axis, then letting

we get

so indeed this bound isn't too much worse and we are able to compute it in quadratic rather than cubic time.

G.2 Details of SVD of QK proof

As discussed in Section 4.3.1, to further reduce the computation cost of proof, we need to avoid computing the residual stream, EVOU, and EPQKE matrices fully. Using mechanistic insight or otherwise, we observe that these matrices (apart from EVOU) can be well-approximated by rank one matrices. This will remove the dominant computation cost of O ( d vocab 2 · d model ) .

G.2.1 Comments on relationship between mechanistic insight and proof size

Up to this point, we haven't really said much in our proofs about what the model is doing. All the mechanistic insight has been of the form 'the model varies more along this axis than this other axis' or 'the input data is distributed such that handling these inputs is more important than handling these other inputs' or, at best, 'the model computes the answer by attending to the maximum token of the sequence; everything else is noise'.

Here, finally, our proof-size constraints are tight enough that we will see something that we could plausibly call 'how the model pays attention to the maximum token more than anything else', i.e., (if we squint a bit) 'the model pays more attention to larger tokens in general.

G.2.2 The max row-diff trick

As stated above, we are breaking matrices into their rank one approximation and some error term. To bound the error, i.e. to bound expressions of the form ∏ i ( A i + E i ) -∏ i A i , where E i denote the matrix errors, we can use the following trick:

Lemma 20 (Max Row-Diff (vector-matrix version)) . For a row vector a and a matrix B ,

Moreover, for a collection of n row vectors A r , if the shape of B is m × p , the right hand side can be computed for all r in time O ( nm + mp ) .

Proof.

The asymptotic complexity of computing the result follows from caching the computation of max i,j ( B k,i -B k,j ) for each k independently of r , as the computation does not depend on A r .

Theorem 21 (Max Row-Diff) . For matrices A and B ,

Proof. By taking the max of Lemma 20 over rows r of A .

Lemma 20 can also be applied recursively for a product of more than two matrices.

Lemma 22 (Max Row-Diff (vector-matrix recursive version)) . For a row vector a and a sequence of n matrices B p of shapes r p × c p ,

Moreover, for a collection of q row vectors A α , the right hand side can be computed for all α in time O ( qr 0 + ∑ p r p c p ) .

Proof. We proceed by induction on n .

For n = 1 , the statement is identical to Lemma 20.

Suppose the theorem holds for all positive n = s ; we show the theorem holds for n = s +1 . We reassociate the matrix multiplication as

Using the induction hypothesis gives

The triangle inequality gives

and algebra gives

The asymptotic complexity of computing the right hand side also follows straightforwardly by induction.

Theorem 23 (Max Row-Diff (recursive)) . For a sequence of n +1 matrices A 0 , . . . , A n ,

Proof. By taking the max of Lemma 22 over rows r of A 0 .

Note that Theorem 21 is compatible with the mean+diff trick of Appendix G.1.

Theorem 24 (Combined Mean+Diff and Max Row-Diff) . For matrices A and B , and any columnwise summary vector H k of A (for example we may take H k := E r A r,k )

Proof.

Moreover, for a collection of q row vectors A α , the right hand side can be computed for all α in time O ( qr 0 + ∑ p r p c p ) .

Proof sketch. Apply the triangle inequality recursively, fusing the proofs of Lemmas 22, and 24.

Theorem 26 (Combined Mean+Diff and Recursive Max Row-Diff) . For a sequence of n +1 matrices A 0 , . . . , A n , and corresponding column-wise summary vectors h 0 , . . . , h n -1 of A 0 , . . . , A n -1 ,

Moreover, if the matrices A p have shapes r p × c p , the right hand side can be computed in time O ( ∑ p r p c p ) .

Proof. By taking the max of Theorem 25 over rows r of A 0 .

G.2.3 Exploring rank one approximation via SVD

Let us first look at

From Figure 9a, we see that there is not much variation along long query token direction. We can confirm this by performing a singular value decomposition (SVD) on EQKE , as seen in Figure 14.

0

10

20

30

40

50

60

0

10

20

30

40

50

60

Singular Index

Query-Side SVD

0

10

20

30

40

50

60

0

1000

2000

3000

4000

5000

6000

7000

Singular Index

Value

Singular Values

0

10

20

30

40

50

60

0

10

20

30

40

50

60

Singular Index

Key-Side SVD

-

20

-

15

-

10

-

5

0

5

10

15

20

1

The first singular value is just over 7440 ( 7800 ± 380 across all seeds), while the second singular value is just under 15 ( 13 . 1 ± 2 . 8 across all seeds). The ratio across all seeds is 620 ± 130 . There's really not much going on here beyond the first singular component. 23

Call the first singular component of EQKE the 'query direction' d q and the 'size direction' d k on the query-side and key-side, respectively.

There are two ways that we can decompose EQKE into a low-rank component that we can compute exactly, and a full-rank error term that we approximate bounds for.

G.2.4 The simple SVD decomposition of QK

In time O ( d vocab d model 2 ) we can perform SVD on each of the four component matrices E q , Q , K , ¯ E and perform low-rank SVD on the matrix product E q QK T ¯ E T .

We can then bound the difference between two elements in the same row of EQKE by computing exactly the difference between the two elements in the same row of the rank one approximation of EQKE , and adding to that a bound on the difference between the two elements in the same row of the error term.

That is, we can decompose E into a part parallel to d q and a part orthogonal to d q , say E q = E q + E ⊥ q , and similarly ¯ E = E k + E ⊥ k . Note that E q and E k are both rank one, and hence can be multiplied with other matrices of shape d model × a in time O ( d model a ) rather than time O ( d vocab d model a ) .

Hence we can define EQKE _ err 1 (subscript one for 'rank one') and decompose EQKE as

Define for any vector v

so that we get

̸

̸

Then we may use any method we please to pessimize ∆ i,j (EQKE _ err 1 ) t query quickly. For example, since for any matrix M we have σ 1 ( M ) = sup x ∥ Mx ∥ / ∥ x ∥ , considering vectors with one 1 , one

0

10

20

30

40

50

60

0

20

40

60

key token

query token

-

1

.

00

-

0

.

75

-

0

.

50

-

0

.

25

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-1 , and zero elsewhere, the maximum difference between elements in a row upper bounded by √ 2 σ 1 ( M ) :

G.2.5 The complicated SVD decomposition of QK

While the 'most mechanistic' interpretation would proceed with the analysis in terms of E q and E k , perhaps decomposing them further, we can get more bang for our buck by extracting out all the low-rank structure available E , Q , and K , so as to make our error bounds as tight as possible.

To this end, we perform SVD on E ⊥ q , E ⊥ k , Q , and K and peel off the first singular components so as to get the decomposition

Then EQKE , a product of these four matrices, can be expressed as a sum of 2 2 3 2 -1 = 35 rank one products and one high-rank error term. We can compute the sum of the rank one products in time O ( d vocab 2 ) and express EQKE as, say, EQKE 2 + E ⊥ q, 2 Q ⊥ ( E ⊥ k, 2 K ⊥ ) T . Call the second term EQKE _ err (Figure 15). We must now bound for each q and m the quantity max i ≤ m -G EQKE _ err[ q, i ] -EQKE _ err[ q, m ] .

How big is this?

Even if we relax to max i,j EQKE _ err[ q, i ] -EQKE _ err[ q, j ] , the maximum such value across all rows is under 1 . 85 ( 1 . 99 ± 0 . 68 across all seeds). And the rows don't have any particular structure to them; the maximum absolute element of the entire matrix is just barely over 1 ( 1 . 12 ± 0 . 40 across all seeds), so doubling that doesn't give too bad an estimate.

But we somehow need to compute this value without multiplying out the four matrices.

One option is to try to use singular value decomposition again. Since σ 1 ( M ) = sup x ∥ Mx ∥ / ∥ x ∥ , considering vectors with one 1 , one -1 , and zero elsewhere, the maximum difference between elements in a row upper bounded by √ 2 σ 1 ( M ) . The largest singular value of EQKE _ err (Figure 16) is just under 7 . 6 ( 8 . 4 ± 2 . 0 across all seeds), giving a row-diff bound of about 10 . 7 ( 11 . 8 ± 2 . 8 across all seeds), which is large but not unusably so.

If we perform SVD before multiplying out the matrices (Figure 17), however, their first singular values are about 4 , 1 . 4 , 1 . 4 , and 4 , giving a product of about 30 , which when multiplied by √ 2 is about 43 . (Across all seeds, these numbers are 3 . 79 ± 0 . 12 , 1 . 525 ± 0 . 067 , 1 . 513 ± 0 . 073 , and 3 . 78 ± 0 . 12 , giving a product of about 33 . 1 ± 2 . 9 , which when multiplied by √ 2 is about 46 . 8 ± 4 . 2 .) This works because σ 1 ( AB ) ≤ σ 1 ( A ) σ 1 ( B ) , but note that we can do factored SVD without needing to use this technique. This bound is still usable, but pretty big.

Can we use Frobenius? Note that using anything close to this method to drop below d vocab d model 2 might seem infeasible (it'll eventually turn out not to be). For example, the best bound we know on the largest singular value that can be verified even in the worst-case in strictly less time than

0

10

20

30

40

50

60

0

10

20

30

40

50

60

U

0

10

20

30

40

50

60

0

1

2

3

4

5

6

7

Singular Values for EQKE_err

0

10

20

30

40

50

60

0

10

20

30

40

50

60

V

0

5

10

15

20

25

30

0

10

20

30

40

50

60

U

0

5

10

15

20

25

30

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

3

.

5

4

.

0

Singular Values for

E

⊥

q,

2

0

5

10

15

20

25

30

0

5

10

15

20

25

30

V

0

5

10

15

20

25

30

0

5

10

15

20

25

30

U

0

5

10

15

20

25

30

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

1

.

4

Singular Values for

Q

⊥

0

5

10

15

20

25

30

0

5

10

15

20

25

30

V

0

5

10

15

20

25

30

U

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

1

.

4

Singular Values for

K

⊥

0

5

10

15

20

25

30

V

0

5

10

15

20

25

30

U

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

3

.

5

4

.

0

Singular Values for

E

⊥

k,

2

0

10

20

30

40

50

60

V

0

5

10

15

20

25

30

0

5

10

15

20

25

30

0

5

10

15

20

25

30

0

5

10

15

20

25

30

0

5

10

15

20

25

30

0

5

10

15

20

25

30

it takes to compute the full SVD is the Frobenius norm, which is defined as tr ( MM T ) , can be computed in d model d vocab time, and is equal to the square root of the sum of the squares of the singular values. While the Frobenius norm of EQKE _ err is only about 12 (giving a bound of about 17 on the row-diff), the Frobenius norms of the four multiplicand matrices are a bit over 10 , 4 , 4 , and 10 , giving a product of 1932 and a bound of 2732 (!). (Across all seeds, the Frobenius norm of EQKE _ err is about 13 . 1 ± 1 . 9 (giving a bound of about 18 . 6 ± 2 . 7 on the row-diff), the Frobenius norms of the four multiplicand matrices are a bit over 9 . 92 ± 0 . 19 , 4 . 43 ± 0 . 01 , 4 . 361 ± 0 . 095 , and 9 . 85 ± 0 . 19 , giving a product of 1888 ± 99 and a bound of 2670 ± 140 .) This is unusably large.

However, we can get a much better bound on the max row-diff of EQKE _ err without having to multiply out all four matrices. We can use an approach vaguely similar to the mean+diff trick, as follows.

If we want to compute the max row-diff of a product of matrices AB , we can compute by Theorem 21

or by combining this approximation with Theorem 18 via Theorem 24 we may compute

taking whichever bound is better.

The first gives us a bound of 7 . 94 on the maximum row-diff, which is better than we can get by doing SVD on the product of the matrices! We can get an even better bound by peeling off the first two singular values of all four matrices before multiplying them; this gives us a bound of 5 . 67 . Combining it with the avg+diff trick wouldn't give us much ( 8 . 05 and 5 . 66 respectively), as we've effectively already done this by peeling off the leading singular contributions; the mean of EQKE _ err over dimension zero has norm 0 . 025 ( 0 . 030 ± 0 . 012 across all seeds).

Although this error bound is no longer the leading asymptotic bottleneck, we can peek ahead to what we get if we want to be linear in parameter count. In this case, we can apply the recursive version of Equation 12 via Theorem 23, giving a bound of 97 . 06 on the maximum row-diff.

The mechanistic understanding we get here is roughly 'for any given basis vector of the residual stream, the difference between the overlap of any two input tokens with this direction is small once we factor out the first two singular components', and this is sufficient to drive a low error term overall if we factor out the leading singular components in other places. We don't mechanistically understand how to combine the E q QK T (without multiplying them out) in a way that allows getting a good bound, though, which corresponds to our inability to drop below d vocab d model 2 here.

If we use this trick on QK only, and use the mean+diff trick on final attention handling (without which we lose about 19 % ), we can achieve a bound of 0 . 7840 ( 0 . 661 ± 0 . 035 across all seeds).

If we use this trick on the skip connection ( EU ) only, we can achieve a bound of 0 . 6768 ( 0 . 632 ± 0 . 061 across all seed).

Using this trick on both EU and QK drops us down only to 0 . 6354 ( 0 . 601 ± 0 . 060 across all seeds).

If we use this trick on EU and use the recursive version of this trick on QK, we get a bound of 0 . 2927 ( 0 . 281 ± 0 . 036 across all seeds).

Unfortunately, it's not clear how this trick would apply to EVOU . A fancier convex hull checking algorithm seems required, and an analysis thereof is in progress.

G.3 The algorithm

We now put all of these tricks together into the subcubic algorithm Algorithm 7, which is the full version of Algorithm 6. The format we give here is parameterized over the summarization strategy (from Theorem 19 in Appendix G.1), the decomposition of EQKE , and the handling of EQKE _ err and EU .

Algorithm 7 Counting Correct Sequences in Subcubic Time

̸

̸

̸

̸

H Comparison of proof strategies

In this section, we compare the various proof strategies that we have developed in Appendix G. We do some traditional mechanistic interpretability analysis to justify that the choices that we made could be expected to lead to reasonably good bounds in Appendix H.1. We then compare the complexities and performance of various proof strategies in Appendix H.2 to line up with the legends of Figures 3, and 4. We close with a figure relating the various categories of proof strategies.

H.1 Justification of pessimization choices

In Sections 4.3, F, and G we make a number of choices about which axes of variation are more or less important to track at various points in the bound computation.

Here we do some more traditional mechanistic interpretability analysis to justify that the choices that we made could be expected to lead to reasonably good bounds.

H.1.1 Justifying the gap

We take advantage of the fact that attention is mostly monotonically increasing in input integers and that for most sequences, the attentional contribution of the particular query token matters much more than the particular non-max token in the sequence.

We justify this as follows.

We can look at the typical diff, when attending to the max token, between the largest non-max logit and the max logit. As shown in Figure 18a, the largest difference between an off-diagonal entry of EVOU and the diagonal of that row is typically at most -7 . 24 The typical worst contribution to the wrong logit from a non-max token (this is typical over non-max tokens, worst over choice of output token-logit index) is around 43 , as shown in Figure 18b.

The difference in attention between tokens is approximately linear in the gap between the tokens, as seen in Figure 19. The slope of the line, that is, the difference in pre-softmax attention scores divided by the gap between the key token and the max token, is approximately 1 . 2 .

Exponentiating, the post-softmax attention paid to the max is typically about 3 × larger than to the token one below the max; here the logit difference between the max and non-max token is significant, typically being around 13 ( 43 / 3 ) for the worst output logit. But by the time the gap is 3, this difference has dropped to about 1 . 1 , and by the time the gap is 4 it is around 0 . 3 .

-

10

-

5

0

5

10

0

1

2

3

4

5

6

7

·

10

6

logit - diag

count

×

# sequences with given max

̸

30

35

40

45

50

55

60

65

70

0

.

0

2

.

5

5

.

0

7

.

5

10

.

0

12

.

5

15

.

0

17

.

5

20

.

0

logit diff

count

-

60

-

40

-

20

0

20

40

60

-

75

-

50

-

25

0

25

50

75

token gap

attention difference

(a)

(EQKE

i

-

EQKE

j

)

/

d

vs.

i

-

j

-

2

0

2

4

6

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

·

10

9

attention difference / gap

count

×

# sequences with given max

(b) (EQKE i -EQKE j ) / ( √ d ( i -j )) , weighted by sequence count. µ ± σ = 1 . 22 ± 0 . 13

1

2

3

4

5

6

7

8

9

0

1

2

3

4

5

6

·

10

6

gap

count

×

# sequences

So for sequences where the largest non-max and the max are close together, the particular structure of the non-max EVOU matters a lot; but when the max is separated from the largest non-max by a modest gap, the structure of the non-max EVOU does not matter so much.

̸

The upshot is that to handle most sequences, we need only ask an oracle for the minimum gap g > 0 between the max token t max and largest non-max tokens t ′ = t max , such that the model outputs the correct answer for all sequences where the non-max, non-query tokens have value at most t max -g .

While computing this gap may be expensive (and indeed the naïve computation of the oracle takes longer than the brute-force proof-though it should be very easy to optimize), we don't have to pay the cost of computing the gap in the size of the proof, only the cost of storing the gap table ( O ( d vocab 2 n ctx ) ) and of verifying the gap. Empirically, gaps are typically 1-5, as seen in Figure 20.

If we rely on the gaps, this results in leaving behind about 6 . 9 % of sequences.

Picking up more sequences In this paragraph / bulleted list, we sketch out how we might go about picking up more sequences to get a tighter bound. This is not coded up, and is left as future work. We propose computing the following quantities:

First, we could build in time ( O ( d vocab 2 ) ) a table indexed on pairs ( t, t max ) of the maximum token and a non-maximum token: the table would store pessimal logit contributions from t to maximum output tokens ≤ the t max parameter. The table could be further split to pessimize separately for tokens within and outside of the gap window.

Compute a table of pre-softmax attention differences between tokens t and t +1 in time ( O ( d vocab 2 ) ).

Next sort the queries by overlap with the query direction.

√

-

0

.

6

-

0

.

4

-

0

.

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

0

25

50

75

100

125

150

175

200

matrix element value

count

E

⊥

q,

2

N

(

-

0

.

00

,

0

.

22)

-

0

.

4

-

0

.

2

0

.

0

0

.

2

0

.

4

0

20

40

60

80

100

120

matrix element value

count

Q

⊥

N

(

-

0

.

01

,

0

.

14)

-

0

.

4

-

0

.

2

0

.

0

0

.

2

0

.

4

0

20

40

60

80

100

120

matrix element value

count

K

⊥

N

(

-

0

.

00

,

0

.

14)

-

0

.

6

-

0

.

4

-

0

.

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

0

25

50

75

100

125

150

175

200

matrix element value

count

E

⊥

k,

2

N

(0

.

00

,

0

.

22)

Compute for each number of queries handled (where we assume we handle all queries with greater overlap than the current one) and for each maximum input token t max , how many of the query tokens t query fall strictly below the max t max (and whether or not the model succeeds when t max = t query ). This will tell us how many query tokens we can count for a given maximum token.

Compute a table indexed on pairs of # of queries handled and input tokens t which stores the smallest difference in more attention paid to t +1 than to t ( O ( d vocab 2 ) ).

Compute a table indexed on pairs t max , t storing an upper bound on amount more attention paid to non-maximum tokens than to t max by Oracle-permitted query tokens (the Oracle is indexed only on t max ) ( O ( d vocab 2 ) ).

For each # queries permitted: compute for each t max , t , c , if the non-maximum token t contributes little enough to incorrect logits that even with the worst skip connection the model still gets the correct answer.

H.1.2 Stopping after 1-2 principal components of QK

Did we miss out on any structure in the error term of EQKE ? The distribution of entries of the four matrices looks pretty close to normal as seen in Figure 21.

2

25

2

31

2

37

2

43

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

brute force (acc: 0.9992 ± 0.0015)

cubic (rel acc: 0.9539 ± 0.0080)

subcubic (rel acc: 0.821 ± 0.013)

attention-

d

vocab

d

model

2

(rel acc: 0.795 ± 0.014)

direct-quadratic (rel acc: 0.653 ± 0.060)

attention-

d

vocab

d

model

2

, direct-quadratic (rel acc: 0.628 ± 0.060)

attention-quadratic (rel acc: 0.390 ± 0.032)

attention-quadratic, direct-quadratic (rel acc: 0.286 ± 0.036)

2

25

2

31

2

37

2

43

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

brute force (acc: 0.9992 ± 0.0015)

cubic (rel acc: 0.9539 ± 0.0080)

subcubic (rel acc: 0.821 ± 0.013)

attention-

d

vocab

d

model

2

(rel acc: 0.795 ± 0.014)

direct-quadratic (rel acc: 0.653 ± 0.060)

attention-

d

vocab

d

model

2

, direct-quadratic (rel acc: 0.628 ± 0.060)

attention-quadratic (rel acc: 0.390 ± 0.032)

attention-quadratic, direct-quadratic (rel acc: 0.286 ± 0.036)

If we replace the entries of E ⊥ q, 2 , E ⊥ k, 2 , Q ⊥ , and K ⊥ with randomly sampled values, we get (sample size 100) that the maximum row-diff of the product of the matrices is approximately 1 . 31 ± 0 . 13 (sampling without replacement from the empirical distribution) or 1 . 31 ± 0 . 14 (sampling from the normal distribution). So in fact our max row-diff is unusually high (by about 4 σ ). 25

H.2 How various combinations of tricks perform

Recall Figures 3 and 4 on page 8 and on page 9 from Section 5, recapitulated here without captions for convenience as Figures 22, and 23.

We describe what each subcubic proof strategy in the legend means. Note that all subcubic proof strategies (that is, all proof strategies except for 'brute force' and 'cubic') use the quadratic counting algorithm of Appendix F.

H.2.1 Proof strategies grouped by complexity

In Figures 3, and 22, proof strategies are grouped by computational complexity.

The 102 proof strategies break down into 1 + 1 + 2 × 5 × 10 × 2 strategies.

The brute force and cubic proofs ( 1 + 1 ) were fully covered in Appendices D, and E.

There are 5 options for handling EU :

direct-quadratic refers to handling EU in time O ( d vocab d model ) with either the max row-diff trick (Appendix G.2.2) 26 or the max row-diff trick fused with mean+diff or some other summary statistic (Theorem 24) 27 .

When direct is not mentioned, this indicates that we handle EU in time O ( d vocab 2 d model ) by first multiplying out E q U and then either taking the maximum row-diff in each row 28 or by taking the maximum row-diff across all rows 29 . The latter is included purely for comparison's sake, and never gives a tighter bound than the former.

There are 10 options for handling the high-rank attention error term EQKE _ err :

400

600

800

1000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

EPQKE Singular Ratio:

σ

1

/σ

2

max-diff-exact

mean+max-diff-subproduct

max-diff-subproduct

mean+max-diff

svd

mean+max-diff-subproduct-recursive

max-diff-subproduct-recursive

mean-recursive+max-diff-subproduct-recursive

attention-quadratic refers to handling the high-rank attention error term EQKE _ err from Appendix G.2.5 in time O ( d vocab d model ) either with the recursive max row-diff trick (Theorem 23) 30 or with the recursive max row-diff trick fused with the mean+diff trick either just on the query side 31 or throughout 32 (Theorem 26).

attentiond vocab d model 2 indicates that we use one of the various O ( d vocab d model 2 ) strategies for handling EQKE _ err 1 from Appendix G.2.4 or EQKE _ err from Appendix G.2.5. These include using √ 2 σ 1 -computed via low-rank SVD-as the bound (Equation 11) 33 , considering all ways of multiplying out a subset of the matrices and taking the maximum row-diff of the resulting pair of matrices 34 (Theorem 21), or fusing the max row-diff trick with the mean+diff trick 35 (Theorem 24).

When attention is not mentioned , this indicates that we handle the attention error term in time O ( d vocab 2 d model ) , either by taking the per-row maximum row-diff 36 or by using the full rank EQKE matrix and taking the per-row maximum row diff 37 .

Finally, note that in combining the rank one attention computation with EVOU , PVOU , and EU , we may either use the mean+diff trick 38 (Appendix G.1) or not 39 ; this makes up the final factor of 2 .

H.2.2 Proof strategies grouped by attention handling

This section slightly reorganizes the information just covered in Appendix H.2.1, for convenience of legend correspondence. Here we group by the strategy used to handle the attention error term. Strategies that involve using the full rank EQKE matrix are elided. The dashed descriptors here correspond to underscore-joined descriptors in footnotes of Appendix H.2.1.

max-diff-exact ( O ( d vocab 2 d model ) ) corresponds to taking the full rank EQKE _ err 1 term and taking the maximum row-diff in each row.

mean+max-diff-subproduct ( O ( d vocab d model ) ) corresponds to fusing the max row-diff trick with the mean+diff trick (Theorem 24) and considering all ways of associating the multiplication of EQKE _ err .

max-diff-subproduct ( O ( d vocab d model ) ) corresponds to using the max row-diff trick (Theorem 21) and considering all ways of associating the multiplication of EQKE _ err .

max-diff ( O ( d vocab d model 2 ) ) corresponds to using the max row-diff trick (Theorem 21) on the factored SVD of EQKE _ err 1 .

mean+max-diff ( O ( d vocab d model 2 ) ) corresponds to fusing the max row-diff trick with the mean+diff trick (Theorem 24) and applying it on the factored SVD of EQKE _ err 1 .

svd ( O ( d vocab d model 2 ) ) corresponds to using √ 2 σ 1 -computed via low-rank SVD-as the bound (Equation 11).

mean+max-diff-subproduct-recursive ( O ( d vocab d model ) ) corresponds to handling the high-rank attention error term EQKE _ err from Appendix G.2.5 with the recursive max row-diff trick fused with the mean+diff trick on the query-side only (Theorem 26, taking all but the first summary vector to be zero).

max-diff-subproduct-recursive ( O ( d vocab d model ) ) corresponds to handling the high-rank attention error term EQKE _ err from Appendix G.2.5 with the recursive max row-diff trick (Theorem 23).

mean-recursive+max-diff-subproduct-recursive ( O ( d vocab d model ) ) corresponds to handling the high-rank attention error term EQKE _ err from Appendix G.2.5 with the recursive max row-diff trick recursively fused with the mean+diff trick (Theorem 26).

H.2.3 What understanding do we get from each proof strategy?

Throughout most of this paper, we talk about doing mechanistic interpretability and using understanding to allow more compact proofs to have tighter bounds. We can also look at the reverse problem: we can take a collection of proof strategies, check by brute force which strategies give the tightest bounds for each model, and ask what this implies about how that model works. We do this here.

In general, which proof methods perform best is an indication of where structure exists in the model. For example, in quadratic EU proofs, when max_diff performs worse than mean_query+max_diff and svd_query+max_diff , this indicates that E has a relatively strong behavioral component shared across query tokens that U is not that good at filtering out. Similarly, when, e.g., mean_recursive+max_diff_subproduct_recursive performs better than max_diff_subproduct_recursive , this indicates that even after removing the first one or two principle components from E q , Q , K , and ¯ E , there is still enough common structure that it is worth factoring out the mean behavior.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We present the key challenges in the field of formal verification for neural networks, and describe our proposed solution. Then we summarize our experimental setup and results.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: See Section 7.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide intuitions for our proof constructions in the main body of the paper. In the supplemental material, we lay out in full detail theorem statements and proofs. Along with this, we provide algorithms and plots that are useful for understanding how the proofs were constructed.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We directly link to a codebase with the specific implementation of our case study.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: In supplementary materials, we provide the full details of our model training set up. In the main body of the paper, we describe our experimental setting and provide reasoning for why we chose this experimental setup as a very simple case study of our theoretical work.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We run our experiment on models trained with 151 different hyperparameters. For most reported computations, we provide statistical significance information. We do not need to perform explicit t-tests or such since it is not relevant to our setup.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: See Appendix B.1.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We confirm that we have read the Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This paper seeks to advance the fields of mechanistic interpretability and formal verification of machine learning systems. While there are many indirect societal consequences of our work through the impacts on these fields, we feel that none are sufficiently consequential as to be highlighted here.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper does not involve any such assets.

Guidelines:

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We list all important Python packages used in the paper in Appendix B.1.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We introduce no new assets except for the codebase needed to reproduce our experiments, which does contain appropriate documentation.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper makes progress towards the formal verification of model performance by using insights from mechanistic interpretability to reduce proof length. It presents a case study in a simplified setup: a single-layer, single-head attention-only transformer applied to a synthetic task of finding the maximum of four integers. The authors train a large number of models on this task and derive formal guarantees for their performance using different proof strategies that make use of different insights derived by mechanistically interpreting the model. The first strategy with cubic complexity, uses the model's decomposition into direct and indirect paths, and the observation that the model makes minimal use of positional embeddings. They also consider different sub-cubic proof strategies that rely on different simplifications, such as the observation the EQKE matrix is approximately rank 1, corresponding to the size of the key token. These strategies are compared to a brute-force solution in terms of proof length (estimated using the number of FLOPs) and the ratio of the certified bound (lower bound b divided by the expected value s). The authors observe a trade-off between proof length and the certified bound, noting that establishing tighter bounds is more computationally expensive. They also find that compact proofs are less faithful to model internals. Finally, the authors identify compounding noise as a challenge for proofs of global behaviour.\n\nSTRENGTHS:\n- The paper demonstrates great ambition by working towards provable guarantees about global model behaviour. Developing techniques to provide such guarantees is of significant interest, especially for deploying these models in safety-critical environments. \n- The authors employ a variety of methods to analyse the model and cleverly use mechanistic insights to reduce proof length. To the best of my knowledge, this is the first paper that attempts to verify global model behaviour using insights from reverse-engineering the neural network. \n- The proposed measure to assess mechanistic understanding (degrees of freedom that have to be evaluated using brute-force) is interesting, as current evaluations of interpretability results often rely on proxy metrics.\n\nWEAKNESSES:\nThe results appear to be merely an interesting case study in a highly simplified setting. The proposed proof strategies rely on various unrealistic simplifications, including the use of attention-only transformers, and the learned algorithm for the task allows for further simplifications, such as ignoring positional embeddings. These simplifications raise doubts about the transferability of the proposed proof strategies to other settings. For example, the observation that compact proofs, which rely on various simplifying assumptions, are already less faithful in this synthetic setting suggests that obtaining compact proofs in more realistic settings would be incredibly challenging. The paper could have been significantly strengthened by either studying a more realistic setting or characterising fundamental methods to leverage mechanistic understanding for any task to improve upon brute-force search.\n\nMinor Issues:\n- The title on the paper (“Compact Proofs of Model Performance via Mechanistic Interpretability”) does not match the title in OpenReview (“Provable Guarantees for Model Performance via Mechanistic Interpretability”).\n- Typo in line 95: “we denote the true maximum of the by t_{max}.”\n- Line 162: one “the” too many: “… sequence x, [the] define the ….”\n- Line 189: “consider” was probably meant to be “considered.”\n\nQUESTIONS:\n1. You only compare sequence lengths of up to four because the brute-force solution is infeasible beyond that. Could your proposed proof strategies be used to make guarantees beyond this length? I believe that demonstrating that your approach allows us to provide guarantees for settings where brute-force solutions are (nearly) impossible to derive would greatly strengthen the results.\n\nLIMITATIONS:\nThe authors have properly discussed the limitations of their approach.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors apply techniques from mechanistic interpretability to ``compact-ify'' the process required to prove model performance on the max-of-$k$ task, specifically where $k=4$. Through this procedure, the authors claim that (1) there is a positive relationship between mechanisitic understanding and the length of the proof required; (2) compact proofs provide better mechanistic understanding; (3) absence of known structure/interpretations make it difficult to compactify proofs in a non-vacuous manner.\n\nSTRENGTHS:\n- The proposed approach demonstrates connections between two important yet distinct subfields of deep learning---interpretability and robustness. \n- For this specific max-of-$k$ task, the authors reverse-engineer the circuitry and algorithm implemented by transformers to perform this task. To the best of this reviewer's knowledge, this has not been done before in existing works. \n- The authors introduce neat tricks such as mean+diff, max row-diff, as well as a nice convexity approximation in the appendix to achieve compact, subcubic proofs. \n- The paper provides some insight into how theorists may utilize tools from mechanistic interpretability which so far has been a purely empirical endeavor.\n\nWEAKNESSES:\n- The task the authors choose seems very contrived even for mechanistic interpretability studies. I'm not sure why specifically the max-of-$k$ task was chosen or what its significance is in the formal verification literature and the authors do not provide this context.\n \n- The model the authors use is also very contrived (one-layer, one attention-head transformer). Even though the results are interesting, it is unclear to this reviewer how any of these methods would generalize as the task/model changes in both complexity and size, respectively. If in fact they are not meant to generalize, it is unclear to this reviewer how the insights of this paper fit into the larger conversation of robustness and interpretability. \n\n- The authors claim that Fig. 4 shows how less faithful interpretations lead to tighter bounds. It is extremely unclear to me that there is a positive correlation between the Singular ratio and accuracy bound. The authors should perform a statistical test here based on the $R^2$ to demonstrate that there is indeed a significant correlation. \n\n- The paper is difficult to read in the sense that there are undefined/inconsistent terms and notations. For example,  \n- Line 16, ''we will need compact proofs of global robustness.'' The notion of a ''compact proof'' is never explicitly defined. I am not sure if this is jargon from formal verification. Even if it is, it should have citations and some explanation for audience. \n- Line 55, ''we restrict $f$ to be the 0-1 loss so our theorems bound the accuracy of the model.'' In Eq. 1, we seek to lowerbound $f$. Based on my understanding of 0-1 loss, the model is penalized 1 for incorrect classification and 0 otherwise. Is this not the opposite of what we want?\n- Line 87, Line 106, inconsistent bold-facing. What is difference between $\\mathbf{E}$ and $E$. Moreover, what does $\\mathbf{E}_q$ mean?\n- Table 1, in the expressions of the complexity cost, specifically for rows ''Low-rank $QK$'', ``Low-rank $EU$'', and ''Low-rank $QK$ &$EU$'', are the terms (EU & OV), (QK & OV), (OV) annotations or a part of the mathematical expression?\n- Line 501, citation deadlink. \n- Fig. 6 ''Let $d = \\sqrt{d}$''?\n- Line 589 and beyond, notation $\\mathbb{N}^{<n_{ctx}}$ is nonstandard. I think I understand from context, but it would be nice to clearly state what you mean by this.\n- Algorithm 3, 5 is very difficult to parse, break this up into multiple subroutines?\n\nQUESTIONS:\n- Why did the authors choose to investigate a task like max-of-$k$ instead of modular addition (or IOI, colored-objects, etc.) where circuits for more general models have been discussed extensively in related literature?\n- In lines 16-18, the authors emphasize that ``we need compact proofs of global robustness.'' Neither of these terms are rigorously defined in the paper nor is there reference to related literature. Since there are many working definitions of robustness what exactly do you mean by this? Based on Eq. 1, it does not seem that these guarantees fend against out-of-distribution or adversarial examples?\n- In mechanistic interpretability, researchers generally choose a specific task, then reverse engineer the mechanisms of models that perform this task well to uncover the algorithms driving these models. The authors also follow this principle when reverse engineering the max-of-$k$ task in their toy transformer. Given that these mechanistic interpretability techniques require models with almost perfect performance a priori why is this not tautological with respect to their formal verification? \n- How would this method generalize to different models that implement different circuits or algorithms on the same task? It seems that the inductive biases from the model used to compactify the proofs would not hold if a different model was using a different algorithm to complete the task (also perfectly).\n- In line 73, why is pessimal ablation done by simply sampling from the data distribution instead of using carefully crafted counterfactuals that is common in mechanistic interpretability [1, 2, 3] \n- Throughout the paper, the authors denote the residual from the rank-1 approximation of the EQKE matrix as ''unstructured noise.'' Could the authors elaborate on why they came to this conclusion? I understand in lines 1071-1077 the residuals follow ''pretty close'' to a normal distribution, did the authors run statistical tests to rigorize this claim? Moreover, why does this imply that the noise is unstructured? \n- Why is the ratio of the largest to second-largest singular value a good metric for the faithfulness of the interpretation?\n\n[1] Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Circuit Component Reuse Across Tasks in Transformer Lan-\nguage Models. In The Twelfth International Conference on Learning Representations, September 2023. URL\nhttps://openreview.net/forum?id=fpoAYV6Wsk.\n\n[2] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking\nvia mechanistic interpretability. In The Eleventh International Conference on Learning Representations,\nSeptember 2022. URL https://openreview.net/forum?id=9XFSbDPmdW.\n\n[3] Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in\nthe Wild: a Circuit for Indirect Object Identification in GPT-2 Small. In The Eleventh International Conference\non Learning Representations, September 2022. URL https://openreview.net/forum?id=NpsVSN6o4ul.\n\nLIMITATIONS:\nI am unconvinced that the procedure the authors propose can be generalized to different tasks (even of the same complexity). I think the large amount manual quantitative/qualitative analysis/engineering required to derive the compact proofs for even this trivial task should be addressed as limitations. Other than this, the limitations section is well-written.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper aims to generate formal guarantees certifying model’s performance on max of K task. Using brute force gives the bound close to true performance of the model but the complexity cost i.e. compactness is exponential in the size of the vocabulary. The paper then aims to use understanding derived from mechanistic interpretability to reduce the complexity cost while ensuring the bound remains tight. For this, they first use the fact that model is composed of OV, QK circuits and a direct path. This aids in reducing the complexity to cubic in size of vocabulary. The bound still remains tight. Then the authors utilize few observations from their setup, for instance, low rank of EQKE matrices to further reduce the complexity to sub-cubic. However, there is a sharp drop seen in the tightness of the bound as these insights are used.\n\nSTRENGTHS:\n1) This work can motivate future works to try to generate more compact bounds on model’s performance. This is especially important in context of alignment, where these bounds could help in providing certificates for safe performance. Although this is a very preliminary study, it attempts to make a contribution in this direction.\n\n2) It is interesting to see how authors were able to achieve better trade-offs between complexity and tightness of bounds on model performance by using insights from mechanistic interpretability. This work is very original in this regard.\n\nWEAKNESSES:\n1) The tightness of bounds seem to degrade very quickly as additional insights from mechanistic interpretability are used. Thus it seems difficult if future works would be able to use insights from mechanistic interpretability to derive tight bounds on performance of the model.\n\n2) The author analyze a very simple task and techniques used by authors to decrease computational complexity are specifically tailored for this task. It therefore remains unclear how insights from this work can be applied to any other tasks. It would have been great if the authors would have attempted to investigate a few other tasks which are simple (fundamental) but still very different in nature, e.g., counting, sorting, etc.\n\n3) There isn't enough evidence to adequately support the noise hypothesis presented by the authors i.e. the noise in estimating rank-1 approximations of E, Q, K results in larger magnitude in the error term. Currently this seems based on speculation and needs more detailed investigation. Therefore it would be great if the authors could try to artificially inject noise (of varying magnitude) in these matrices and see its effect. Else, the authors should not claim this being the major cause for not so tight bounds obtained on using insights from mechanistic interpretability.\n\nQUESTIONS:\nI request the authors to kindly address the questions in the weaknesses section.\n\nLIMITATIONS:\nYes, the authors have addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper does a detailed and careful case study of the trade-offs and design space of formal verifications for meaningful lower bounds of the accuracy of a transformer model on a chosen task:\n- Specifically, a one-layer, one-head, no-MLP, no-layernorm transformer is studied (akin to the model studied in detail in \"A Mathematical Framework for Transformer Circuits\", https://transformer-circuits.pub/2021/framework/index.html). A crucial difference from \"A mathematical framework...\" is that the model here also includes positional encodings.\n- The task is to output the maximal element from a context of length 4, where tokens are synthetically generated to correspond to integers in a limited range, and the model is trained (successfully) solely on this task.\n- Then, several verification approaches are considered:\n\t- Brute force: this means running the model on every possible input and calculating the accuracy\n\t- \"Cubic\": this proof uses some ingenious ideas to avoid iterating over all 4-tuples of possible inputs, and instead only over a particular set of triples $\\Xi$. Specifically, they show that given an input $x=(t_1, t_2, t_3, t_4)$, then, *if* the model is 100% accurate on a certain subset of $\\Xi$ derived from $x$, the model will also be correct on $x$. \n\t- \"Sub-cubic\": even more ingenious/compact proof strategies are possible if one uses mechanistic knowledge of how the chosen model architecture tends to solve the given task. For instance, a major part of the algorithm seems to be: the queries and keys combine to put attention on the token representing the largest integer in the input, and the output and value matrix then combine to copy this to the last token position. By using this and other properties (e.g. that the important structure of the query/key/value/output matrices is low-rank), it becomes possible to shorten the proof further.\n- The main result is that approaches using more mechanistic knowledge of the\napproximate algorithm implemented by the transformer require less computation for verification, but also lead to less faithful/tight bounds on the accuracy.\n\nSTRENGTHS:\n- this paper really engages with the computations of a transformer on a very fundamental level, and presents some ingenious ways to combine high level and heuristic approaches from the mechanistic interpretability literature with a fully verified method of deriving correctness. \n- This is to my knowledge the first work to carry out such a complete analysis. One may hope that in the (not too distant) future, similar methods can be scaled up via LLM automation to apply to more interesting, non-toy reasoning tasks, and deliver ground-truth insights about the structure and weaknesses of LLM computations.\n- the paper is very clear and rigorous about what it achieves, and all claims are backed up by evidence and proof.\n\nWEAKNESSES:\n- the setting is quite toy, though this is probably very hard to overcome for such an approach with human labor.\n\nQUESTIONS:\n- this is very interesting work! You probably have a good intuition for the most promising/low-hanging future directions in which it can be taken - can you get into specifics about this in the conclusion?\n- how can this method be carried over to more realistic tasks? If I understand correctly, much of the theory in the appendices relies on a careful analysis of the structure of this particular transformer model **and** the particular (arithmetic) task in question.\n- you mention: \"models with similar training procedure and performance may still\nhave learned significantly different weights\" - I believe a standard reference is \"Breiman, Leo. Statistical modeling: The two cultures (with comments and a rejoinder by the author).\", cf. \"Rashomon\"\n\nLIMITATIONS:\n- it is quite unclear how such results, based on careful analysis of the structure of the particular task, will carry over to more realistic tasks. \n- similarly, it is unclear whether re-introducing transformer components, and especially MLP layers, won't obliterate the accuracy bounds derived here to be practically useless (e.g. below chance levels, which would be 25% here)", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}