Compact Proofs of Model Performance via Mechanistic Interpretability
Authors: Jason Gross
Rajashree Agrawal
Thomas Kwa
Euan Ong
Chun Hei Yip
Alex Gibson
Soufiane Noubir
Lawrence Chan
Venue: NeurIPS 2024 poster (2024)
Keywords: mechanistic interpretability
verification
proof
guarantees
interpretability
Primary area: interpretability_and_explainability
TL;DR: We prototype using mechanistic interpretability to derive and formally verify guarantees on model performance in a toy setting.

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.

In this work, we provide a case study of translating mechanistic interpretations into compact proofs. Max-of-K task with 151 random seeds (Section 3), and Compact Proofs of Model Performance via Mechanistic Interpretability
Jason Gross"
Rajashree Agrawal
Thomas Kwat
Euan Ong
Chun Hei Yip
Alex Gibson
Soufiane Noubir
Lawrence Chan
Abstract
We propose using mechanistic interpretability- techniques for reverse engincering model weights into inmaaninterpretable algorithms- to derive and compactly prove formal guarantees 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 each of models. Using quantitative metries, find that shorter proofs seem to require and provide more mechanistic understanding. Moreover, 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.
Introduction
One approach to ensuring the safety and reliability of powerful AI systems is via formally verified proofs of model performance [48, 11]. 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 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 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 training procedure. We propose verifying model performance using understanding derived from mechanistic interpretabillity (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.
We train attention-only transformer
Conresponding author. Please direct correspondence to jgross@mit.cco edu.
These authors contributed equally to this work.
$ 'These authors contributed equally to this work.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Tree Modal
Brude Furce
Presf
Sabcatic Prodd
Loger
oo fa)
Ma fo fea)
P (fa)
QK Ciesai decorsposs inko lage "Sax" and seull "asin" coppurans
Uteerteed
Unremell
?
Siee compacei singal value
EPV
QX
Explel
Other corrponezee hake srggaar valee
Tapst
6g 1
60 11 6
6o
FLAIPs Recsared:
1.43 10th
5.33 10
4,88 11
Accurncy Lower Boaad:
95.31 95
28.4115
Lneesplaized Dissmuion: Augaptoii Complexay.
LO0 HP Ohoaaa (11)
1.28 y 10
4.42 10
Officcal Teed
We construct proofs using different degrees of mechanistic interpretation. (Left) The models we consider in this paper are one-layer attention-only transformers, and so contain three "paths": the OV circuit, the QK circuit, and the direct path. (Right) For the brute-fonce proof (Section 4.3.1), we treat the model as a black box and thus need to check all possible combinations of inputs. For the cubic proof (Section 4.3.1), we decompose the model into its three corresponding paths, but still check the correctness of each path via brute force. Finally, some subcubie proof' (Section 4.3), we use all parts of the mechamistic interpretation presented Scction 3. (Bottom) For each of the three categories of proof, we report the number of FLOPs used in computing the certificate (lower=better, Appendix A.6), lower bound on model accuracy (higher=better). effective dimension of the unexplained parts of the model (lower-hetter, Appendix A.5), and asymptotic complexity of the proof strategy as we scale the inputs and model (lowersbetter). Significantly more compact proofs have vacuous accuracy bounds by default. Using more mechanistic understanding allows us to recover some, but not all, of the accuracy bounds on these more compact proofs, as our understanding is not fully faithful to the model internals.
Figure 1:
then reverse engineer the models using standard mechanistic interpretability techniques. We 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). We validate our technique against an additional 604 models for varying values of K (Appendix 1.2.
We define quantitative metric to assess the mechanistic understanding used in 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, find a negative relationship between proof length and degree of understanding. We qualitatively examine proof strategies 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).
However, also identify compounding structureless error terms as 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 '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.
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 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)
X, notated as D)x when marginalized over labels, and f L x Y + R be scoring function for evaluating the performance of the model. Then, we seek to establish lower bounds 6 on the expected 5 as the form:
= B(,.)~D M(t)) 2 b.
'Our 102 proof strategies be broken up 1 1+1 + 10 5 x 2: two standalone strategies, and class of strategies parameterized on three axes cardinality 10, 5, and (Appendix H).
$ "Code: https: //gi /gtbtttttt.....// com JammJJaoonrros/////o guarantees- based- weweedd-eehanitti...m iechannssie--nttrrretttli.. interpretabilityy

---

As f can be any metric, this is 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 to be uniform, so our theorems lower bound the accuracy of the model. Our proof methodology generalizes straightforvardly to other input distributions (Appendix A.8), and only 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 + I and a non-computational component Q arguing that for any model M', C(M') I Ef Mt))).. thus implying that C generates 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. Here, b = C(M). As even the size of the model parameters much larger than any reasonable Q, we approximate the length of proof pair C, Q by the length of 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 = 3, albeit with a very long proof.
We quantify the length of ((M) using two metrics: the asymptotic time complexity of C' as we scale the size of the model and the imput t. as well as the empirical average number of floating point operations required to evaluate C(M') over given set of models [M, We measure rightmess of bound of C(M) using the ratio of the bound to the true accuracy: b/5.
Proof as pessimal ablation A standard way of assessing the faithfulness of mechanistic interpretabil- ity is by ablating the parts of the model that your interpretation does not explain [54, 6, 23]. 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 complete argument might be quite high. Thus, we construct relaxations (Appendix A.4) over input sequences, such that performing pessimal ablations 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-of-K.
Model Architecture We study one-layer, one-head, attention-only transformers with no biases but with learned positional embeddings, with vocabulary size dyneab, model and head dimension d = dmodel dhead, and context length Netx k. The model parameters consist of the fltx dmmdee positional embedding P; the dvocab X dmodel token embed E; the dmodel X Amodel query, key, value, and output matrices of the attention head Q, K, V, and O; as well as the dmodel dvocab unembed matrix U. We assume (as is standard in language modeling) that dmodel < dyocab-
For an neex x dvveab one-hot encoding x [xo, T1....... ..ox21...PPPax11 of an input sequence t = [to, f1,.... &nar - 1), we compute the logits of the model as follows:
(0)) xE + P
Initial residual stream (Metx X duuode))
a= 10 QKT /Va
Attention matrix (Netx X nexx))
(11) da()) - -0 v000
Final residual stream (netx X dmodel)
M(t) = 8 11 U
Final seq, position logits (dvocat)
where o² is the masked softmax function used in causal attention. Because we only look at outputs of the model above the final sequence position i = netx "query position" and the value of the token in this position as tquery. one-hot encoded as Tquery- The model's prediction is the token corresponding the max-valued logit &max-
also denote this position as the
Task Specifically, we study the setting with Netx = k = because is the largest sequence length for which we can feasibly evaluate the brute-force proof. We set hidden dimension dmodel = 32
'Other components of the proof to account for the difference between floating point numbers and reals are described in Appendix A.7. Note that all proofs explicitly given in this paper are of Q only; we do not include any traces of running C.
3

---

and vocabulary of size dvocab = 64 comprising integers between 0 and 63 inclusive. For an input sequence t, we denote the true maximum of the sequence by maxx- Outputting the correct behavior is equivalent to outputting logits e such that A := l- max < 0 for all 1* * /maxx. 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:
M(t)) o² ("Qqquery E + P querry QK (xE+ (Va) (xE+P)VOU (aquery E + Pquery) )
QK circuit
OV circuit
direct path
(2)
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.
and position-dependent components. Let Pavg = >; Ps/netx be the average position embeds across positions (of size dmodel), and let P denote either nee @ Pavg or Idyocas @ Pavg depending on context, the result of broadcasting Pavg back into the shape of P or E (that is, Netx dmmdel dvocab dmodde). Similarly, let Pg ddwca 8 Pquery be the result of broadcasting P Pquory- Then for one-hot encoded x, we can rewrite the QK and OV circuits, well as the direct path, as follows:
We further divide the QK and OV circuits into token pooiiion-inddeennde
QK circuit
Equery
E,QKTE' XJ
EOKTT
EQKE
EQKP
OV circuit
EVOU- + PVOU
Direct Path Equery Egy
EVOU
PVOU
EU
where P P P P and E E aad E.g E Pg (since (0)) xxE + P).
3.1
Mechanistic interpretation of learned models
Using standard empirical mechanistic interpretability tech- niques, we interpret one of our learned models (our "mainline" model) by independently examining the QK and OV circuits and the direct path.' We find that the model outputs the largest logit on the true max token maxx 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.
Logis (4.4
Uneutbed
OK Cinari stund to loger tolons mood
OV Chouu perfocer, JoA rank copping
Diin path nottingttigg
Ertbed
The models our setting implement Max-of-K by attending ex- ponentially more to larger tokens and copying the attended-to tokens (Sec- tion 3.1).
Figure 2:
QK circuit By qualitatively examining the position- independent QK component EQKE, we find the amount of pre-softmax attention paid to a key token is approximately in- dependent of the value of the query token fquery. and increases monotonically based on the size of the key token. We confirm this hypothesis by performing a singular-valle value decomposition (SVD) of the EQKE matrices (Appendix G.2.3), and find that it contains 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/4 = 1/Varocat- The right (key-side) singular vector of this component is monotonically - increasing as we increase the size of the key token, with (1/ dd-scaled) pre-softmax attention increasing by an average of 1.236 = 0.056 when the key token increases by 1.a
In comparison, each 1/V d/Valecdeed entry the ble size (average 0.31 + 0.18), suggesting that EQKP is unimportant to the functioning of the model.
QK component EQKP has negligi-
7 our trained models hehave similarly; see Appendix B.3. This implies that the ratio of attention paid to token and 1 approximately exp(1.236) 3.442.

---

We confirm this by zero ablating EQKP, which changes the models' accuracies from 0.9992 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, 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 > 6.6 + 1.2, Since each entry in the sequence is uniformly sampled and dyocab = 64, this means that EVOU is a good approximation for the identity matrix for all but (7/64)* = 1.2 10 2 % of the sequences.
As with the position-dependee 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.
QK component, the resiition-dppence
OV component PVOU also
Direct path As with the two position-dependen components, the entries in EU have small absolute magnitude 2.54 0.20," and contribute negligibly to model performance.
Proofs of model performance
In this section we describe intuitions for three categories of proof that 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
VM', Cspecific strutegy (M') S B (1max, Mt).).
We leave implicit the traces of running C Cspecific struteyy on our specific models to give the overall theorem. We report the computational complexity or estimated FLOPs of running Cspecific strutegy 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 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 need to consider, and the fact that only the position-indeeeendet components EQKE and EVOU contribute meaningfully to performance (Section 3) to pessimize over sequence ordering.
First, let a pure sequence & be a sequence with at most three distinct tokens: the max token fmax the final token fquery tmax. and optionally third token t' < fmaxx and let E be the set of all pure sequences in X. 11 For 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:
Adj(t) = {s ppure max & = tmaxs Equery
Equery: Vi < nexx: & E 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 variant of the following convexity property: for given t, if M({) is correct for all & E Ad)(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(dyocab 3 "neexx - 1)!)
FFor comparison, the average off dif-oiaaonal element of EVOU is 21.68 + 0.83 below the corresponding diagonal element.
10 Appenndi A. discusses how to compute the "brute-force" accuracy of model on an infinite distribution. In Section 4.3, we will consider smaller set of "pure sequences".
5

---

Table 1:
We report the proof complexity. accuracy bound, and estimated flops required (Equation 2), well as unexplained dimensionality (Section 5). We round the FLOP and unexplained dimension counts to the closest power of 2, and report the mean/standard deviation of the bound averaged across all 151 models. As we include more aspects of the mechanistic interpretation (reflected by lower number of unexplained dimensions), we get more compact proofs (in terms of both asymptotic complexity and FI.OPs), alheit with worse hounds, For space reasons, k Retx, d dmodel, and dyocab.
Description of Proof
Complexity Cost
Bound
Est. FLOPs
Unexplained Dimensions
Brute force
0.9992 0.0015
247
2
Cubic
Olr*k²)
0.9531 +: 0.0087
2
2
Sub-cubic
O(1 12 o2 .d.
0.702 + 0.033
221
213
w/o mean+diff
0.349 + 0.080
2a1
2
Low-rank QK
011 + nd + v²d) QK
0.675 + 0.035
2
2
SVD only
BUROV
0.284 1: 0.072
2
2
Low-rank EU
0 vd vdd)
(1633 + 0.062
2
2
SVD only
EU
QKAOV
(3.38.10.06) x x10
2a1
2
Low-rank QK&EU
O61 + nd
+ ed
+
0.610 + 0.060
221
2
SVD only
QK
BU
Ov
(3.38 13.38100009 00.00.0000000 x00
2
2
Quadratic QK
On k3 vd (
0.316 + 0.037
2
2
EUROV
Quadratic QK&EU
O + ud +
0.283 + 0.036
221
213
QK&ED
OY
pure sequences. This allows to bound the accuracy of actual M on all dvocab sequences,
while evaluating it (dvocabb ( 1)!) sequences.
We can reduce the number of sequences that we need to evaluate by pessimizing over the order of sequence. For given tuple of (fmax: !query, t'), there are (Netx - 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 /qurry, (') tuple to the number of l' in the pure sequence, and the total number of sequences (ddvccab 3 netx)) 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 the ((eex) required to compute the softmax over netx elements, resulting asymptotic complexity Oddvccab 3 3 nnxx2) (Theorem 12, additional details in Appendix E).
4.3 Sub-cubic proofs
We now consider proofs that are more compact than O(dvocab 3). These require avoiding iteration over any set of size ((dvocab 3) (e.g. the set of pure sequences) and performing operations that take (ddvecab) time on each of Oddvecab 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 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 (dvorab Retx) pure sequences ξ, with parameterized by (Imax, 'quurr, e). Recall from our mechanistic interpretation in Section 3.1 that the pre-softmax attention paid from lquery to a key token L' broadly invariant in lquery and increases roughly linearly with the size of L'. This allows us pessimize over the OV circuit over all "sufficiently small" tokens.

---

More formally, suppose we are given some gap 9 E N. For each pure sequence with max token mmax. query token Equryy such that Equery I max 9. and e copies of the third token type t' I fma 9. we pessimally ablate the OV circuit over the set E (Imax, /query, C; g) of pure sequences ç with the same max and query tokens and e copies of the third token type 1'. If the model gets all sequences in Epure (fmax: 'query+ , (; 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(dvocab Paiiaabnnneu nexx) pessimal pure sequences of each of the O(dvocab 2 netx) sets of the form Epuro Epuro (Imax Qquery C; g).
Decoupling and pessimizing computations that require (ddvocah 9) computations Many parts of our cubic certificate require iterating through O(dvocab 3) cases parameterized by (fmaxx Equery or (fmax.t') For example, as part of the pessimization procedure over pure sequences, for of the dyocab possible values of fmax. we need to consider the relative effects the dwucab -sized logits of attending to each of the (ddvocab) other tokens t' < fmax. and for each fmax and query, we need to check that the contribution of the direct path on logits EqueryyUU is not sufficiently large as to overwhelm the contribution from FmaxEVOO We independently pessimize over each of these components over one of the d axes: for example, instead of computing Tmax EVOU- Equery-EU EU for each (max. lquery pair, we first pessimally ablate the direct path along the query token (which takes Oddrueab? ?) time as it does not depend on the fmax. and then consider the sum Tmax EVOU + max,' E'E Since this sum no longer depends on 'query, we only need to perform (dvoeah)) times, for total cost of (dvoraah)) ")-
Low rank approximations to EQKE and EU Recall from Section 3.1 that EQKE is approximately rank I, where the sole direction of variation 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 dddoocaddmodel ?) time by using the fact that EQKE can be written the product of a dyocab X dmodel matrix and a dmacel X dyocab matrix. This allows us to avoid performing the (dvocab "dmodet)-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 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, 2) f(z, y) + g(y. 2). while only iterating over two variables at a time. The naive way to minimize fxx y) and g(x, v) independently:
min (I, 2 min f(x,)) v) + min g(y,2) z.y.2
z.y
9.2
Here, the error comes from setting the ys in J and y to different values. But in cases where g(v, =) varies significantly with y and only slightly with z, rewriting g as a sum of a component that is independent of (only varying along y), and one that depends on z, yields a better lower bound:
min ((, y, 2) 2 min (f 3) - E,,(y, z²)) + min(g(y, 2) - E,)(y,z))) 2,312
3,2
This estimate will have error at most E, while the naive estimator can have arbitrarily large error. We refer to this rewrite as the "mfan+diff trick". 12 From the mechanistic interpretation in Section 1.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- trick" Using properties of linear algebra, we derive 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"
In fact, this is the motivation behind the standard rewrites of QK and OV into position- independent and Dusition-dependen components (Section 3).
7

---

Nomalizal Accuney Bound
beuse force 0,9992 0,0015)
cubic
0.0080)
0.6
subcubuc trol acc: 0,700 0,036)
0.4
0.035)
diret-qustratic 0.613 0.063)
0.2
atcation-dlousbllaal R, ddiect-qudrrti fsel 0.610
atteation quadratic noc: 0.316 0.157)
altestiongiiiiaaii. dirns--uaaltati tral se: 0. 283 (.637)
Vediy (opgrosimate)
beune Socce lincur baseline
Figure For the proofs in Section 4, we plot the number of FLOPs used to compute the certificate, as well as the normalized accuracy oower-bound (b/8). s). The brute-force proof (Section 4.1) computes the exact performance uses orders of magnitude more compute than other approaches. The cubic proof (Section 4.3) uses small amount of mechanistic understanding and less compute, while still retaining good accuracy lower bounds. Finally, subcubic proofs (Section 4.3) require the entirety of the mechanistic interpretation of the model to attain non-vacuous bounds; this understanding allows us to further reduces compute costs, hut we still achieve worse bounds. See Appendix H.2.1 for a detailed description of the various proof strategies.
trick. We apply this trick get a better cheap bound the terms of low-rank approximations, without having to multiply out the full matrices. See Appendices G.2.2, and F for more details.
Results
We run each of 151 transformers on the various proof strategies of different asymptotic complexity, and analyze these proof' 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/3 (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 fully describe model behavior, assuming the structural assumptions of the proof are correct. More detailed interpretations leave fewer free parameters needing 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, 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.
could not cheaply bound the behavior of
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 qquery-
5.2 Proof length vs. bound tightness trade-off 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 ways that diverge from the original model internals. For example, in some subcubic proofs (Section 4.3), we approximate EQKE with rank-1 I approximation corresponding to the "size direction".. However, while other components are small, they're re nonzero; this approximation harms model internals.
Less faithful interpretations lead 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 1 approximation 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 fixed proof length.

---

Normalzzed Accurac Bound
1.0
mean-max
0.8
max-diff
mean-max moca-mas--iff--upoooodd dilf-subproduct
0.6
max diff subproduct
max-diff-exad
0.4
conn--oooisioennnn--- www.aaa-ii----pprrrtttt--e-e-eee subproduct-occusiii
0.2
mean-max diff subproduct reccuuuutctrruvvvee max diff- subproduct recursiiv
400
600
800
1000
EPQKR Singulor Ratio: fg//
Figure 4: We plot the normalized accuracy bound versus the ratio of first and second singular values of EQKE, for various types subcubic proofs that depend closer EQKE is to rank-1, the tighter the accuracy bound. This suggests that more faithful interpretations lead to tighter bounds even holding proof length fixed. See Appendix H.2.2 for a detailed description of proof strategies.
rank-1 approximation EQKE, each class proof,
Unexpfained Dimension (Esimated)
220
cubic
26
subcubic
843
attention- Aamii
direct- quact-quadratti
200
attetion-dleusslfaaae direct-quadratic
24T
attention quadratic
24
attention quadratic, autnntio-quadratis.
dinect-quadratic
225
247
2
24b
2000
2
247
F1.0Ps Verify Proof (approximate)
Figure 5: We plot, for each proof, the approximate number of flops required to evaluate the proof, versus the unexplained dimensionality (Section 5.1). More mechanistic understanding leaves fewer dimensions unexplained. We observe that more compact proofs to leave fewer unexplained dimensions, which is indicative of the relationship of mechanistic understanding and compact proofs. See Appendix H.2.1 for more detail. proof's
5.3 Compounding structureless noise is a big challenge for compacting
Pessimal error terms compound in the absence of known structure The rank-1 approximation of EQKE has small error. However, when making rank-1 I 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 't cause errors to compound (without just multiplying out the matrices), this approximation leads to trivial bound on performance (Appendix G.2.5). We speculate that many cases, there no short iamaaninterpreetbble 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.
Related Work
Generalization Bounds Prior work in the PAC-Bayes framework [58, 36, 12] proves generalization bounds over learning procedures, which are similar the global performance bounds consider in this work. These proofs tend to provide statistical guarantees [25, 26) about the outputs of a known stochastic training procedure, while seek 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 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 approach more akin to interactive theorem proving [22] and verify global performance properties by reverse-eng enggenerenginniiingg the weights.
Mechanistic Interpretability Finally, mechanistic interpretability is the subfield of the broader field of understanding model internals [45], which 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 transformer implements algorithm.

---

Nichani et al. [39] proves that, in a significantly simplified 2-layer, 1-head aitention-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-of-K setting to prototype the use of mechanistic interpretabil- ity to derive compact proofs of model behavior. Using varying amounts of understanding, derived more efficient proof computations lower bounding model accuracy. We found preliminary evidence that mechanistic understanding 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.
compactify proofs. Moreover, observed that the tightness
Limitations and future work We study one-layer attention-only transformers on toy algorithmic task. Future work should explore the viability of deriving proofs via interpretability using larger models featuring MLPs or layernorm significantly compact the part of the proof involving the OV circuit, which future work can explore. The proofs explored in this work also did not lead 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].
more complex domains. In addition, we were unable to
Acknowledgments and Disclosure of Funding
We are immensely grateful 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-Alos 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 subset of authors, and providing funding for compute and in- in-person research sprints; Constellation and FAR Labs for hosting subset of the authors and providing 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 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-of-K 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.
10

---

References
[1] Behzad Akharpour, Amr Abdel-Hamid, Sofiène Tahar, and John Harrison. Verifying syn- thesized implementation of IEEE-754 floating-point exponential function using HOL. The Computer Journal, 53:465-488, May 2010. doi: 10. 1011093cojjn//xp022.
[2] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma. and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models, 2022.
[3] Andrew Appel and Ariel Kellison. VCFloat2: Floating-poin error analysis in Coq. In Proceed- ings of the 13th ACM SIGPLAN International Conference Certified Programs and Proofs, CPP 2024, pages 14-29, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400704888. doi: 10.1 10.1145/3636501.33 3663366355013336655.
[4] 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//RTHH.2111
[5] 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- 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, llya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. org/ aa/aabas//005/ 14165.
URL https::///aaatt
[6] Lawrence Chan, Adrià Gloniga-Alonso, Nicholas Goldwosky-Dill. Ryan Greenblatt, Jenny Nitishinskaya, Ansh Radhakrishnan, Buck Shlegeris, and Nate Thomas. Causal scrub- bing, method for rigorously testing interpretability hypotheses. AI Aligument Forum, 2022. URL https: http:///wwwwaiiggneefffrorrrm wwwppottt/////o/tt././/t..t.sstt http://////////hhacaaaaaaaaa. htth///ahyyyyy22///////aaaaa causal-
[7] 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.
[8] Paul Christiano, Eric Neyman, and Mark Xu. Formalizing the presumption of independence. arXiv preprint arXiv:221 1.06738,, 2022. doi: 10.48550/arxiv.2211 1/..://xxx..22......... .06738.
[9] Bilal Chughtai, Lawrence Chan, and Neel Nanda. A toy model of universality: Reverse engineering how networks learn group operations, 2023.
[10] Edmund M. Clarke, William Klieber, Miloš Nováčck. 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:
URL https///:://oo org/10 10.
[11] David Dalrymple, Joar Skalse. Yoshua Bengio, Stuart Russell, Max Tegmark, Sanjit Seshia, Steve Omohundro, Christian Szegedy, Ben Goldhaber, Nora Ammann, et al. Towards guar- anteed safe AI: A framework for ensuring robust and reliable Al systems. arXiv preprint 266664240, 2024.
[12] 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.
[13] 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-Dodde Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. URL https: / http//ranfffot ormer circuits. wwt.ciiss..ubp//////// httpt///ppp///////eeeeeeeo/ fta.//orr////////// wtarrkiiiind...eee html.

---

[14] Martín H. Escardó. Synthetic topology of data types and classical spaces. Electronic Notes in Theoretical Computer Science, 87:21-156, November 2004.
[15] 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), Wroctaw, Poland, July 2007.
[16) Martín H. Escardó. Seemingly impossible functional programs, 2007. URI. htw.:::tttaatatttta /mmth.. andrej http:/ccmmmmmm////// /000007/0///0 /09//0/9///////// /etp//////eeenggy////e/ -ontt//ammmmiii--ttotooooooalala proggamm//. Accessed: 2024-05-15.
[17] T. Gehr, M. Mirman, D. Drachsler-Cohen, P. Tsankov, S. Chaudhuri, and M. Vechev, A12: 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..1099/P2218.0005
[18] 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.12777 2018.
[19] Jason S. Gross. Performance Engineering of Proof- Proof-Based Based Software Systems at Scale. Phl) thesis, Massachusetts Institute of Technology, February 2021. URL https: mit. Sttuhhandee/721..11 1713.77
[20] 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.
[21] Kaarel Hänni, Jake Mendel, Dmitry Vaintrob, and Lawrence Chan. Mathematical models of computation in superposition. In ICML 2024 Workshop Mechanistic Interpretabilityy 2024. URL
[22] 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.
[23] Stefan Heimersheim and Neel Nanda. How to use and interpret activation patching. arXiv preprint 1rX1::2404.15255, 2024.
[24] Jordan Hoffmann. Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
[25] 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: 2018
[26] 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.
[27] 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.088 000661, 2020.
[28] 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.
[29] 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. IEFE Computer Society. doi:
12

---

[30] 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. sel Formal verification of an OS kernel. In Proceedings of the ACM SIGOPS 22nd Symposium on Operating Systems Principles, SOSP '09, pages 207- 220, New York, NY, USA, 2009. Association for Computing Machinery. ISBN 978 9781605587523. doi: 10.1145/1629555.16229 1629599. URL https: 1tt../..10/..1112222222
1629596.
[31] Xavier Leroy. A formally verified compiler back-end. Journal of Automated Reasoning, 43: 363 446, 2009.
[32] Chuan Li. OpenAI's GPT-3 language model: technical overview, June 2020. URL https:
/ 1/lambdalabs cow www/blog///////cccccccccooowoooo /lllo//dmystifyinng Settiiinnggggg----- Lambda Labs Blog, accessed October 30, 2024.
[33] 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.250800Maaorra 1000-00.
[34] Eric J. Michaud, Isaac Liao, Vedang Lad, Ziming Liu, Anish Mudide, Chloe Loughridge, Zifan Carl Guo, Tara Rezaei Kheirkhah, Mateja Vukelić, and Max Tegmark. Opening the AI black box: program synthesis via mechanistic interpretability. 2024.
[35] 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
[36] Vaishnavh Nagarajan and J. Zico Kolter. Uniform convergence may be unable explain generalization in deep learning, Advances in Neural Information Processing Systems, 32, 2019.
[37] Neel Nanda and Joseph Bloom. TransformerL.ens.
https://github. com/
TransformerreensSrg/Traaa wwwpieannnrrrrraafsrrooee www.rreesssse..sss.. 2022.
[38] 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.
[39] Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. arXiv preprint, 2024. doi: 104485500/2X1.22422.47755
[40] 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 Seiffeld--oddde 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: learning- and www.iniiinnnneeieeeeoo
circuits. pub/2022/ contetext-- html.
[41] Adam Paszke. Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, al. PyTorch: An imperative style, highhperformmacccdcpp learning library. Advances in neural information processing systems, 32, 2019.
[42] Plotly Technologies Inc. Collaborative data science, 2015. URL https://plot.1u/
[43] Aditi Raghunathan, Jacob Steinhardt, and Percy S. Liang. Semidefinite relaxations for certifying robustness adversarial examples. Advances in neural information processing systems, 31, 2018.
[44] Tahina Ramananandro, Paul Mountcastle, Benor't Meister, and Richard Lethin. A unified Coq framework for verifying C programs with floating-poin 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 978 450341271. doi: 10.1145/2854065.285555 28544000.....66..
13

---

[45] Tilman Räuker, Anson Ho, Stephen Casper, and Dylan Madfield-Menelle Toward transparent AI: survey on interpreting the inner structures of deep neural networks. First IEEE Conference on Secure and Trustworthy Machine Learning, 2022. doi: 10.485550arxi..2227.13333 13243..
[46] 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/ wtw./ 06175.
[47] Alex Rogozhnikov. Einops:
Clear and reliable tensor manipulations with Einstein-like notation. In International Conference on Learning Representations, 2022. URL httph:htt//pemeetee /pperreeier.errrrr.. net/forum? htrf/formmmddddde id=oapkSWW2bbb
[48] Sanjit A. Seshia, Dorsa Sadigh, and S. Shankar Sastry. Toward verified artificial intelligence making Al more trustworthy with formal methods- based approach to Al system verification and validation.
[49] Alex K. Simpson. Lazy functional algorithms for exact real functionals. Luboš Brim, Jozef Gruska, and Jifí Zlatuška, editors, Mathematical Foundations of Computer Science 1998, pages 456 464, Berlin, Heidelberg, 1998. Springer Berlin Heidelberg. ISBN 978-3-540-68532-6.
[50] Gagandeep Singh, Timon Gehr, Markus Püschel, and Martin Vechev. An abstract domain for certifying neural networks. Proc. ACM Program. Lang., 3(PPLL), January 2019. doi: 10.1 10.1345/3290344
[51] Christian Szegedy, Wojciech Zaremba, llya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfel- low, and Rob Fergus. Intriguing properties neural networks. arXiv preprint arXiv: 13126699, 2013.
[52] 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. Carey, Îlhan 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. 11110////199--0991...
[53] 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.
[54] 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//rxiv.2211 1.00593.
[55] Michael L. Waskom. seabor:: statistical data visualization. Journal of Open Source Software, (60):3021, 2021. doi: 0022110//oss.03021.
[56] 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.0766 2022.
[57] 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.
[58] 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.
14

---

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 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 study the question of interest: 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 FLOPs, which is about 2 x the cost of training GPT-3 [32], our cubic proof achieves bounds (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 for more detailed numbers, and Figures 6, and 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, worked on highly optimizing our relaxation to make 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 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, convexity theorems and proofs are applicable much more generally generally to element retrieval tasks.
exact
A.2.3 More complicated architectures
We worked 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, large part of this question lies in the feasibility of deriving 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 metric on the quality of mechanistic understanding we are not purporting to have general solution to the problem of scaling interpretability, but instead claim that the challenges in proofs are in fact challenges in understanding networks.
15

---

Table 2:
Version of Table from Section 4.1 with Netx 5, dvocab = 64. We report the proof complexity, accuracy hound, and estimated flops required (Equation 2), well unexplained dimensionality (Section 5). Unlike Table I, which computes the brute force bound exactly, we instead use importance sampling to estimate the bound; estimated FLOPs are reported for what the full brute force proof would take. We round the FLOP and unexplained dimension counts to the closest power of 2, and report the mean/standard deviation of the bound averaged across all 151 models. For space reasons, use k = Netx: d == dmodel, and dvocab.
Description of Proof
Complexity Cost
Bound
Est. FLOPs
Unexplained Dimensions
Brute force
kd)
0.9990 0.0018
2
2
Cubic
O(v
0.941 = 0.011
226
2
Sub-cubic
Olv k² +
0.705 + 0.031
2
2
w/o mean+diff'
0.405 = 0.073
2
2
Low-rank QK
O + vd? + v²d) QK EUROY
0.682 + 0.033
2
2
SVD only
0.335 = 0.066
2
2
Low-rank EU
O vdl. (d
0.649 1 0.055
221
2
SVD only
HI
QKAOV
(4.8 0.1) 10 10-8
2
2
Low-rank QK&EU
O + ed
+ ed
0.628 + 0.053
2
213
SVD only
QK
(4.8 + 0.1) x 10
2a
2
Quadratic QK
O vd
+ dd
0.354 + 0.034
221
212
QK
FUROW
Quadratic QK&EU
OluRk² ud
+
0.335 + 0.033
2
213
QK&EU
OV
Table 3: Version of Table I from Section 4.1 with Beix 10, dvocab 64. We report the proof complexity. accuracy bound, and estimated flops required (Equation 2), well unexplained dimensionality (Section 5). Unlike Table 1, which computes the brute force bound exactly, we instead use importance sampling to estimate the bound; estimated FLOPs are reported for what the full brute force proof would take. We round the FLOP and unexplained dimension counts to the closest power of 2, and report the mean/standard deviation of the bound averaged across all 151 models. For space reasons, we use k == Metx: d == dmodel, and dvocab-
Description of Proof
Complexity Cost
Bound
Est. FLOPs
Unexplained Dimensions
Brute force
O kd)
0.9988 0.0013
2
2
Cubic
O(v*k²)
0.914 = 0.021
2a
2
Sub-cubic
Olv2 12 +.
0.674 = 0.028
223
213
w/o mean+diff
0.539 = 0.061
22
2
Low-rank QK
Olv²k3 + vd? + v²d) EUROV
0.657 = 0.028
233
212
SVD only
0.469 = 0.059
223
2
Low-rank EU
Olu2k3 vd + v'd)
0.639 = 0.032
233
213
SVD only
EU
QKAOV
(0+ 100) 10
222
2
Low-rank QK&EU
Olu2k² ed
+ vd
0.625 = 0.031
23
213
SVD only
(2.9 0.1) 10
223
2
Quadratic QK
k² vd + vnd)
0.392 = 0.030
2
212
QK EUSOV
Quadratic QK&EU
O + ud
+ v²dd)
0.390 = 0.028
2
213
QKREU
16

---

Table 4: Version of Table from Section 4.1 with netx 10 and dvocab 128. We report the proof complexit, accuracy hound, and estimated flops required (Equation 2), well unexplained dimensionality (Section 5). Unlike Table I, which computes the brute force bound exactly, we instead use importance sampling to estimate the bound; estimated FLOPs are reported for what the full brute force proof would take. We round the FLOP and unexplained dimension counts to the closest power of 2, and report the mean/standard deviation of the bound averaged across all 151 models. For space reasons, use k == Netxi d == dmodel, and dvocab.
Description of Proof
Complexity Cost
Bound
Est. FLOPs
Unexplained Dimensions
Brute force
O
0.9972 0.0031
206
2
Cubic
O(v
0.882 = 0.012
2a1
2
Sub-cubic
Olv k² +
0.622 + 0.031
2
245
w/o mean+diff'
0.390 = 0.070
224
2
Low-rank QK
O + vd? + v²d)
0.594 + 0.035
2
24
SVD only
QK EUROV
0.320 = 0.053
2
24
Low-rank EU
O
vdl. (d
0.607 + 0.031
224
2
SVD only
QKAOV
(5.4 0.2) 10
2
2
Low-rank QK&EU
O + ed
+ ed
0
0.595 + 0.030
2
214
SVD only
QK
(5.4 = 0.2) 10 20
2
2
Quadratic QK
O vd
+ dd
0.350 + 0.029
224
214
QK
FUROW
Quadratic QK&EU
OluRk² ud
+
0.384 + 0.025
2
214
QK&EU
OV
Table 5: Version of Table I from Section 4.1 with Beix 20, dvocab 64. We report the proof complexity. accuracy bound, and estimated flops required (Equation 2), well unexplained dimensionality (Section 5). Unlike Table 1, which computes the brute force bound exactly, we instead use importance sampling to estimate the bound; estimated FLOPs are reported for what the full brute force proof would take. We round the FLOP and unexplained dimension counts to the closest power of 2, and report the mean/standard deviation of the bound averaged across all 151 models. For space reasons, we use k == Metx: d == dmodel, and dvocab-
Description of Proof
Complexity Cost
Bound
Est. FLOPs
Unexplained Dimensions
Brute force
O kd)
0.995 + 0.015
2
2
Cubic
O(v*k²)
0.884 = 0.040
2
2
Sub-cubic
Olu2 12 +
0.561 = 0.043
224
213
w/o mean+diff
0.486 = 0.060
224
2
Low-rank QK
Olv²k3 + vd? + v²d) EUROV
0.547 = 0.043
2
212
SVD only
0.431 = 0.060
224
21
Low-rank EU
Olu2k3 vd + v'd)
0.538 = 0.043
22
213
SVD only
EU
QKAOV
(1.0 6.0) 10-4
2
2
Low-rank QK&EU
Olu2k² ed
+ vd
0.526 = 0.041
2
213
SVD only
(1.0 + 5.0) x 10-4
22
2
Quadratic QK
k² vd + vnd)
0.322 = 0.035
2
212
QK EUSOV
Quadratic QK&EU
O + ud
+ v²dd)
0.321 + 0.035
22
213
QKREU
17

---

bruse (escimated) 0,9990 0.0018)
Nemalizal Actnaay
0.010)
subcubic frel nec: 0.702 (LA33)
0.4
divect-quaalttic collaae:: 0.649 0.055)
sitenion-derniifleae (rel avc: 0.682 (0033)
0.2
0.0
atteniion dntemiion-uuddzae (rel sec: 0.354 (.004)
211
20
alteniionquacarrs direci-padrntic (rel soc: 0,336 36 0.033)
FLOPs Verily Proof (sppeoximate)
Plttx 5, dvocab = 64
Nernlltel Acannyy Bound
1.0
brate fonoz (estimated) cace: 0,9088 (.0013)
0.8
cubic (rel acc: 0.95 (LE21)
0.6
0.4
direct quadiatic (rel acc: 0.540 0.002)
attemionn foenssae
0.2
attemiinnoffeeaaaeaes direct-qunaliiie stic ace: 0.636 0.031)
attenion-quudrtts
0,000)
gd FLOPs Verify Proof (appeotimene)
(b) Tetx 10, dyneab 64
Nomalizal Acomcy Bound
1.0
(estimated)
cubic (rel acc: 1 0.012)
direct quaett--uuuiiicccc (rel acc: 11,09 (.000)
sulcubic
"esamiin-feccasuuussss disect-quadiite (eel acc: 0.5% (L08I)
attenton dwmaadiaerd (rel acc: 0.595 0.005)
0.0
antmtton-quaae quaiiinnn-uaaattte,... dircct-quaadd sacct--uaddittie (vel noc: 0,385 0.025)
2
FLOPs to Verify Proof (uppexximate)
(c) Retx 10 and dvocab = 128
Nomslbed Accnncy Boand
brute force (estimated) (acc: 0,955 (.015)
cubic (rel
0.036)
0.6
subcubic noc: 0,563 0.041)
0.4
attentionnff...d
direct quadratic (rel scc: 0.540 0.040)
0.2
attemion-ff....d.... dinect-quadraiicc
0.528
attention attemton-qpucctiii. direct quadmticc (rel aoc: 0,322 0.004)
y11
200
satemioonquudrrii (rel soc: 0.333 e 0.005)
Verify Proof (apgentinnat))
(d) Netx 20, dyocab 64
Figure 6: Version of Figure 3 from page 8 with varying fetx and dyocab- The brute-force proof (Section 4.1) computes the exact performance uses orders of magnitude more compute than other approaches; unlike in Figure 3, here importance sampling estimate the bound.
18

---

Normalizdd Accuncy Hound
1.0
maax-diff-exxc
0.8
mat-dill-subpu subproduct
mean-max subproduct
max-diff
mean-max- dilf
0.4
mean max-diff- diff- subproduct- recurssive
nux--diff-subproddee recuuussee
mccan waaa--rouriineteeeee woo.a----------oooooooooooooooo-
600
800
1000
1200
EPOKE Singular Ratio: fs/
(a)
Textx 5, dyocab = 64. The "svd" proof strategy best-fit line has equation b/² = 0.000 15(01/02) 0.48,
R
0.37.
Normalized Accumey Bound
1.0
maax-diff-exxc
max-ditl-subppe subdoouut
mean-max subproductt
max diff
svd
macn- has.a--oousssis--aa------------- oawp:///oooiioooooaoooooosossoor www.-l----s-oooooooooooreeeeeee
EPQKE Singular Ratio: f
(b)
Textx
10, dvocab 64. The "svd" proof strategy best-fit line has equation b/ 0.000077((11/2))+033)
H
0.23.
Normalized Accuracy Bound
max diff--xactt
max diff subproduct
max-difff-uubpee -iff-uuuhuuhh max-diff
0.4
svd
max sexx-dil--suuppe sutf--suppouuc--ooe recuusiee--
recursive +max diff' subproduct recursivv
mean max diff subproduct recursive
0.0
EPQKE Singular Ratio: its fir
(c) Retx
10 and dvocab 128. The "svd" proof strategy best fit line has equation b/$ - 0.000
0,40, R²
0.42.
Normalized Accurccy Hound
max-difl-cxas
max diff- subproduct
diff subproduuu
mean-mma-diff
max-dilf
0.4
svd
recunsive max subproduct-reecunitt
0.2
mean-max- dift suhppooduut--eeuusse
max diff wuwsssuuupctteeecceeeeee recursive
0.0
EPQKE Singular Ratio:
/
(d) Netx 20, dvocab 64. The "svd" proof strategy best-fit fit line has equation b/2 0.000 0.000098(01//0))+.... R² 0.22.
Figure 7: Version of Figure 4 from page with varying Bexx and dyocab- Note that the "svd" proof strategy has a clear upward trend, especially on early points.
19

---

A.3 Why is more mechanistic understanding correlated with worse bounds?
Figure 3 exhibits Simpson's 's Paradox: although more faithful mechanistic understanding is correlated with better bounds within each class of proof (and moreover the most extensive mechanistic under- standing 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 compres- sion. Unless the model is losselessly compressible, 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 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 intiiitiiiinterrrtabbittt
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, expect that the quality of bounds 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 we using "convexity"? The intuition is that are attempting to optimize 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 instantiate the precomputed possible extrema with the actual values of the model's parameters to determine the the extrema actually are.
We now give formal description of what we mean by "convex relaxation".
For set of inputs X,, we define a set of "relaxed inputs" X.elaxod with an injection T: : X; > P ((Xxollxxed) mapping input the model to the set of corresponding relaxed inputs. On the relaxed input, define function h; X.relxeed - R such that for all t X; and all labels I for which (1, t) is supported by (has non-zero probability in) D, we can find t trelaxed E T:(() with M(t) 2 (t (elaaed)) ). We proceed by finding a small subset of "boundary" examples B; C C Xxellxedd proving that if rellaxe) > by for all pelaxed & B, then ht 2 b, for all trelaxed E Xrelaxed
Then, the computational component C' of the proof validates that that h; h, (telaxdd) 2 by for some b; for all prelaxed xeelaxed This allows us to conclude that M(t) 2 b. for all t X..
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 proof strategy.
As in Figure 1, for any given proof, we can separate our treatment of transformer components into "black-box" box" (e.g., matrix multiplication) and "white-box" components (e.g., specifying that the QK circuit approximately rank one; pessimizing over non-max tokens). Considering the performance score 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 function 64 + R², whose cardinality is (R2))4 where 64 denotes the finite
20

---

Brute Force Proof
True Performance
Performance
Impractical baseline guaranteeing performance via inference
Lower Bound
Trivial Proof
FLOPs Verify
(a) The baseline of using inference to generate proofs.
True Performance
Understanding partially recovers bound tightness
Decreased length of proof leads to looser bounds
FLOPs to Verify Proof
(b) Shorter proofs by default have worse performance bound b. Faithful under- standing allows us to recover significant but not complete bound tightness with minimal proof-length overhead.
Figure The thcoretical relationship between proof length and bound tightness.
set on 64 elements). The unexplained dimensionality of the entire proof the sum of the unexplained dimensions of all black-box components.
Intuitively speaking, unexplained dimensionality tries to capture the degrees of freedom that 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
Figure 3 and Table on page and on page we display approximate floating point operations. We instrument our code execute phantom tensors that track their shape and accumulate approximate count floating point operations. We compute matrix additions and multiplications in the obvious way. We take the instruction count of SVD be the cost of verifying that the output of SVD is valid decomposition: that we have pair of orthonormal bases which when multiplied out give the original basis.
A.7 IEEE 754 vs. IR
In Section we defined and Q and glossed over whether were reasoning over reals or floats. Here we clarify this point that we' ve so far been sweeping under the rug.
Let IF 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 denote restricted to finite numbers (that is, without NaNs and without +co).
We parameterize C, M,, and over the real "3 they operate on, so that, e.g., CF model weights + F. Then we have Q establishing that for any model M', CR(M)) < E(,,))Da f (1, (M((t))), and we have a trace demonstrating that Cr(M?) =bb
Teehhically the floating point numbers are not field. We gloss over this point here, since they define all of the tield operations, if those operations do not satisfy the field axioms.
21

---

Let : : IF" - I 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 E. we need two additional components the proof. We need to find e, E' E IR+ prove that
(CR(MR) - (CC < and
((B.))o..f(t,Ma(()) - -(LLL---..L.. <e'
Then we can chain these proofs to prove that (Bu.).... -
Such e-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, 175.. We leave actually dealing with the gap between floating point numbers and real numbers to future work.
A.8 Non-uniform distributions
In Equation in Section we defined the expected model performance as the expectation of the distribution D:
= Ba M(t)) bb
We then immediately specialized to the case where the marginalization Dix of 'D over labels is uniform. As we'll see in Theorem 1 in Appendix and Algorithm 3 in Appendix E, the bound computation is modularized between function that bounds the performance f(t, over restricted collection of inputs, and a much simpler function that combines the bounds on individual cases into bound the expectation over the entire distribution. The per-input bound computation in
PESSIMIZING-OVER CORRECTNESS in Algorithm and RELAXED- RELAXED-CORRECTNES Algorithm 3; the expectation computation is FORTE-FORCE in Algorithm I and CUBIC in Algorithm 3.
Since the expectation computation is modularized, it 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 a distribution Din that's used for training and (in-distribution) validation, and another distribution D' that is the deployment distribution or generated by an adversary. we had knowledge
D', we could compute the expected performance from inputs sampled from D'. Even if we don't it 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 possible valid input sequences. In addition, as our proofs partition 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 I, x X - IR of 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 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 lossloss share subterm A6, Since compute this subterm in all of our algorithms, we can
22

---

easily extend our approach to log-loss by combining A6; directly rather than merely checking that the value negative as we currently do in CORREEEEE in Algorithm 3. Although this is sufficient for the brute- force and cubic proofs, for the subcubic proof using Algorithm Appendix F, we would additionally have to compute 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.
POSITION POTIPIIII
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 prove two-sided bounds. Most of the other theorems can be straightforwardy 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 proof- checking 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 proof-checking trace rather than the proof object itself. Alternatively, more conventional proof system like ZF, ZFC, or the proof system underlying lsabelle/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 Adam with batch_sizz size = 128, Ir = 0.001, betas (0.9, 0.999), weight decay left at the default 0.01. We train for epoch (3000 steps). Over our 151 seeds, models trained with this procedure achieve (99.92 + 0.15) % train accuracy and loss of (4 + 8) x 10 3. When qualitatively examining single model (for example in Section 3.1 or Appendix H.1), we use the model with config seed 123, model seed 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 the following software packages in work: Paszke 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.
NNumbeer reported as mean across training runs + std dev across training runs mean accuracy and loss. The model seed is deterministically pseudorandomly derived from the seed 123.
23

---

(a) EQKE E. QKQ
(b) EQKP = E,QKTPT
components EQKE and EQKP. It computes the pre-softmax attention score for the model. The positional contribution the attention score, as shown Figure (b). minimal. In Figure (a), the gradient from left to right along the key axis indicates that the single attention head pays more attention to larger tokens. The uniformity along the query axis suggests that this behavior is largely independent of the query token. Further, the light and dark bands imply that some queries better than others at focusing
Figure 9: The QK circuit can be decomposed into the position-independlee and
larger tokens.
(a) EVOU- = EVOU
(b) PVOU = PVOU
(c) Direct Path = E.U
Figure 10: The OV circuit is sum of EVOU and PVOU. In Figure (a) see that EVOU "copies" with the exception of input tokens < 5 (6.6 1.2 across all models) - by virtue of the fact that ahove 5, the diagonal larger than all the other elements the same row. We see that the range on Figure (b) is much smaller than Figure (a), indicating that positional contribution to the copying is minimal. In Figure (c) we see that direct path values matter a bit more than PVOU, being only z 20x smaller than the typical EVQU difference. They don't 't matter that much, though, being small. Additionally, the vertical banding indicates that the primary effect of this hargely-query-indeee geurr-iindeeeneen bias towards larger numbers, reflecting the fact that the input distribution biased towards larger numbers being the maximum. The weak diagonal pattern indicates a slight bias towards upweighting the query token itself as (possible) maximum token.
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 /22 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) 8, the average score (accuracy) accuracy bound achieved by the cubic proof from Section 4.2; (4) Dsubcubic/ 5, the percent- reoeeeeeee accuracy bound achieved by the (per-model best) subcubic proof from Section 4.3.
of the model across the entire input distribution; (3) heabie/ 8, the percent-scorr--ecoveee
For each statistic of interest, Table 6 presents 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." Additionally, each group contains 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 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.
16-Per-model best" here means that for each model seed, select the variant of the subcubic proof with the highest bound. "f a single model is the closest to two statistics, for example when the mean and median are very similar, the model is shown only once.
24

---

Table 6: Plots of various models. The statistics of interest /12 the ratio of the first singular values EQKE, measure of extent to which the attention score computation low-rank; (2) 8, average score (accuracy) of the model across the entire input distribution: (3) Dcubie/S. s. the sereettscore-recoveee bound achieved by the cubic proof from Section 4.2; and (4) buncuuie s. the percent-score score-recovered accuracy bound achieved by the sper-model best) subcubic proof from Section 4.3. The y axes for EQKE and EQKP, the query token; for EVOU, the input token; for E,UU, the input query token; for PVOU, the input position. The z axes are: for EQKE, the key token; for EQKP, the key position; for EVOU, PVOU, and E,U, the output logit token, All token axes range from at the top (or left) to dyocab - = 63 at the bottom (or right). All position axes range from 0 at the top (or left) to Tetx = 3 at the bottom (or right).
accuracy
statistic of interest distribution
EQKE EKKTET
EQKP EOKTPT
EVOU EVOU
PVOU PVOU
seed
EgU
value
123
(selected seed)
24262
01/02
1061.1
13654
626.5
4810
350.0
6204
333.4
12457
0.9999
19451
0.9992
C565
15662
0.9932
16197
0.9915
32103
bubiic
0.971
6082
0.954
6155
0.926
2306
0.922
20415
subcubic /
0.80
29725
0.70
20976
0.63
6155
0.61
25

---

C
Mathematical definitions
We provide a detailed breakdown of the mathematical notation used in the appendix.
Let z
M
of type X - Y be the model sometimes we write
be the finite set on elements; we write = 10.1....,,? 1} when we care about the ele-
Nen
for the logits of the model M(t)
ments of be the softmax function e²/ e²/E,
D
be a probability distribution over input-label pairs (t. E X X L
olv)
o (v)
be the casually-masked softmax function,
D x
be DD marginalized to a distribution over t E X of type L x Y - R be a scoring function for evaluating the performance of the model be the average position embeds across positions (of size dmodel). - P.
o"(v); : /Ejeee be the size of the vocabulary
f
dvocab
d
be the dimension of the attention head, in our case, equal to the hidden dimension of the model dmodel (assumption:
Pavg
dyocab)
P
be either 1m @ Pavg or Iduecat @ Pavg depending on context that is the result of broadcasting Pavg back into the shape of P or E (that is, netx x &nmodel or dyocab X dmodel)
be the context length, the number of tokens in the
input sequence, equal to K in Max-of- K be the fletx x mmoddl positional embedding be the d'vocab x dmodel token embed
P
E
P P. E. E.,
be 1 dunech @ P &query, the broadcasting of P Pquery be P P-PP P. E P. and F + P. respectively
Q, K,
be the dmddde X modde query, key, value, and output matrices of the attention head
V.O
U
he the dmodel dyocah unembed matrix
For any vector-valued function v of length dyocab. parameter- ized over the input sequence t. Au == 1/=1t -/=utmm be the difference between the ith element of v and the element of v corresponding to the true maximum of the input sequence fmax.
t.
be the input token sequence (to, 41 mnax-1] be the netx X dvocab one-hot-encoded input token sequence [zo, , _Tnex-11
x
Equery
= --1 == Tnax-1 be the query token
'query
= 11 == m be the one-hot encoded query
For our particular model, we will have
token
dvocab := 64 d dmodel == 32
Reex := 4
tmax
be the true maximum token in the input sequence, max; t
x (N (eddvooa Mete 2 644 L: Ned 2 64 DIx = U(0, 1,.., dyocab - nots, the uniform distribution
Figure 11: Preliminary model definitions
BaDYYD (F(,
b lower bound
(theorem statement)
(3)
8 (average model score)
We define the two typical performance functions corresponding to accuracy and log-loss. Note the shared subterm A4,
faccumay (tmax, E) := 1 [argmax &; +max] = 1 >
max E Llmax infhlmma
(4)
AC,
plog-loss (fmaxx E) = (o()) = log(E; exp(e; (ltmee)
(5)
De.
We present the model definition in four different regroupings to define via underbrace labels various useful quantities:
M(t) = E(t) = 0² (Iqueey E + P querry) ) QK" (xE + +P)
/Va)
(xE + P)VOU + (Iquery E + P Pquery) U
(6)
circuit
OV circuil
direct path
= 00 Equeery
E,OKTE XT
E,KKTPT /Va)
x EVOU + PVOU
+ Th IQquery E,U
(7)
BQKE
EQKP
EVOU
PVOU
EU
EPVOU(E)
(t)
=a'(t) XÉVOU + +t) (at())-PvO PVOu + Equery E,U
(8)
gRVOU (t)
(P))
GEU (Squary)
Di=0
-1
(a (t) EN EVOU + (a (tt())),.VOU
Equery E.U
9
(EVOU.(()
(PVOU..(E)
gRU (diquory) )
Figure 12: Definitions of the model behavior
26

---

D Brute-force proof
Theorem 1. For BRUTE- netx: M) as defined in Algorithm I,
(0.1. 1)xx
argmax(M(()) -axx
(RUTE-FORCE(Awwcah Metx M)
Proof. In fact the two sides of the inequality are equal by definition. Hence the inequality follows by reflexivity of 2.
Algorithm 1 Counting Correct Sequences By Brute Force
1: function CORRECTNESS( input-sequence)
2: return MODEL-BEULAVOR((MM input-sequence) == MAX(inutt-eequenc))
3: end function
4: function BRUTE-FORRCH(4<occab. Netxx M)
5: return Zvorsh Fesx SUMICORRECTNESSS/ tokens) for tokens (RANGE dyocab)) Resx)
6: end function
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 is extension 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 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 dvocab Nets values), we may look for 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 token that is "less bad" for those positions, if it draws enough attention away from the correct token. See Lemma for details.
We thus spend this section characterizing a relaxation of the constraints on weights:
1. that contains all actually possible weightings,
2. that is extremized weights that still correspond to some notion of "put the most weight on the extremal tokens", and
3. for which computing the extremal weightings is computationally efficient.
Before diving in, let's recall the proof that a weighted average of scalar values extremized by putting 100% of the weight on extremal values:
Theorem 2 (Warmup: Extremizing weighted averages). Fix set of values v; R. The weighted average is bounded by the extremal values: for any w; such that E, w; and 0 S w, < 1,
min Us Σ way, I maxu;
Note that while in the main body, MMt)) referred the pre- softmax output logits, in the appendix abuse notation and occasionally use it to refer to maximum token indicated by the logits where appropriate. 'eeee Appendix A.4 for the reason that we call this "convexity". Note that our use of "convexity" is purely descriptive in this section; all theorems written out explicitly.
27

---

Proof. The proof is simple. We have
Σ wii; min = Σ mi min vy) 2
and
max Σ Wip = Σ wi vi m)) 2 0
so the result follows.
E.1 Proof strategy
The model computes the true maximum max when its outputs logits l are such that AG,- -.-6tmas < O for all 1* A mmxx." 20 As a result, it suffices to lower-bound the proportion of sequences where (an upper bound on) Arr is negative for all L" * 'max- In particular, will upper-bound bound the contribution from incorrect tokens in positions to the difference A between incorrect (t*) and correct (tmax) output tokens A E. Li f/
We do this by arguing that the logit difference A satisfies a certain notion of convexity over the space of relaxation of sequences (Theorem 6), and constructing set of (d nnxx) "extremal" relaxed sequences where the position and token embedding components of attention are pessimized independently.
We start by first rewriting the contribution each token through the attention head to the logit difference into the contributions involving PVOU and EVOU:
AC. (t) PVOU, t (t) AVOUU
2 (t)
We then upper bound s by noting that because the softmax attention is weighted average PVOU,
APPOOU (.((t) (t)
- -PVO. (t) 1
(V) - PU)U t
= ai(()) (PVOU - PVOU,,aas (s)
I (t) max (PVOU PVOU 1,)
Since Σ; (t) 1, we have
netre
A AVVOU. 1- (t) < max (PVOU,,- PVOU, PVOU,,ax i,max 1y)
We then construct a set pure of "prre sequences" consisting of only three types of tokens in of two orders, and show that for each input sequence and readoff logit t" we bound the logit difference from the token embeddings A (t) using small subset X' of pprre
Retx
Phet
AEVOO (-(t)) < max EEX
AVOU, P- (E)
We construct set rrelaxed relaxed sequences, where each relaxed sequence trelaxed consists of sequence and position (t, where AlT- (t, i) is evaluated by separately considering the positional
We use logit difference Alr- because: (a) it shared linear function of the various paths through the model, which can therefore be analyzed separately; (c) it leaves open both the options of pessimizing over output logit before or after combining contributions of various paths through the model.
computation for plog-koss (b)
28

---

contribution through attention (that is, the attention weighted PVOU) and the token contribution (that is, the eeeetion-weighted EVOU) and direct contribution (the logit difference through the skip connection F.U). Note that indicates the position that we pay 100% of the attention to for the PVOU contribution.
We argue that (t, i) satisfies a certain notion of convexity over mixtures of sequences, such that we can evaluate it only on set of (dvocab BBexx) "extremal" sequences in way that takes Oddvocab Nexx) total time to bound 2- (t, i) for every possible input sequence. We then use the extremal sequences that the model gets correct lower bound the proportion of all sequences that the model will get correct. Specifically, we argue that Algorithm 3 provides 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 A6. for any sequence t can be decomposed into the direct contribution from the embeds (EU, the attention-weighted position contribution (PVOU), and the attention ertention-weighted the three components independently. since summing these upper bounds gives a valid upper bound on the logit difference.
weighted token contribution (EVOU). Therefore, it suffices to upper bound each of
We can compute the direct contribution &EU exactly by first computing EU E,, 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 token value function (a la row difference in EVOUJ v Neddooat + R and a token attention function (à la EQKE for a fixed query token) N <dheess + R. Fix position value function (à la row difference in PVOU) w N Anexx IR and a position attention functioo (à la EQKP for a fixed query token) b: Nenas + R.
In practice, we'll 'lll take, for fixed query token Equery. fixed output token of interest t", and fixed maximum token maxx
2 EVOUU - EVOUt, t,fmax
= EQKE, Kauerry ef Vd
PVOU - PVOU: i i,ana
b, = EQKP Lquery /dd
Definition (of sequence via sorted tokens and position permutation). We can define sequence of tokens via sorted tokens and position permutation by specifying non-decreasing sequence of tokens ln < ... < lmax-1 E Neddooar paired with permulation NN N Nenas'
Definition 3 (sequence score), Given a non-decreasing sequence of tokens to S ... S Unax-1 E Nedwoas and a permulation NN N Nenax define the sequence score Sta....ffers aS:
Σ eee
Σ OSiC
tbett
Sto tta
USi<neex
We will drop the token subscript. writing only So. when the token values unambiguous by context.
The sequence score here will be computing EVOU, (- for some fixed t* and tmax- The way we'r 'e set up our definitions, high scores predict {* (and thus bad), negative scores predict fmax (and thus good), and more negative the scores, the stronger the prediction fmaxx-
Definition 4 (swap permutation). Given a permutation a Nenas Nenan of the Dexx positions and two indices 0 & i.j <. Bexx define the swap permutation Dieej to be the permutation that is
29

---

except swapping and j:
(i())
j
Dieny (k)
o(j)
i otherwise
o(k)
Define Dojsesj to be the difference in sequence scores when you swap i and j:
Do,iti == Sas-;
Lemma 3 (Characterization swapping tokens). Fix a non-decreasing sequence of tokens to S. S.... t &ness E N. Fix o N N be permutation of the nexx positions. Fix indices 0 S i,j Netx
Then there are two eases for sign (Ao,ttz)):
1. af, = 0t, then sign (Do,ttz)) -sign (br(t) bo(s)) sign (ve, - vez).
2. Otherwise, sigu (Ao,ior)) = sign (as, at.) sign (b)()) bo()) sign 8a ve, 04 Mey
Intuitively, Lemma 3 says that, if the token contribution to attention is equal between tokens t; and ty. then the impact of swapping their positions (2) and (ji) is entirely determined by how much attention is paid to the positions of and j and the relative difference in their value. (Notably, by swapping these tokens, 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 -eention-weighter values of these tokens.
don't t affect the attention paid other tokens, and the effect of the
Proof. First note that the theorem is trivial for = j. For the rest of the proof, 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 (Ao,ttz)) We abbreviate Jzetj as of for brevity.
sign (c bots)) sign (bo(1) bo())
sign (Do,iets) = sign (Sor - 9/)
Eospenerr Ut etp Eospen cnty tba
sign
So
Now multiply through by the denominator, which is positive
ve - bol()) So Σ caoptborcss
sign
OSpenetx
OEp<nexx
Ui, E0g -e(()) he(r) Me,e00 (E)0(1) - -bo0(0)) - o2, (ebots) eb0()))
sign
Ospe
tbot + S0 (clocs) - -bo(1)) + soe (c cbo(c))
OSpenex
Σ
Uy e ctat - e0, (c - cbocs)) - De, en,e (clo0s) e
sign
Σ Cn + S (c )o55)) + S (clos) - closo)
DEpeReex
30

---

sign ((V, e - V1, e (clni) - -000)) ) + ( - - (c -c002))
sign (e (o0i)) sign (1r,, v2 + 80 (e e²",))
sign (ba(1) - bo())) sign (sg (E - - (u,, 0i, e e²
Divide through by non-zero values when possible
sign (ba(1) ba(i))
sign (Ut, - vt1z)
if oz = at otherwise
sign (e"4, e²T))) sign 87 "4, -une 22
6944-""
sign (bo(t) (()) sign (vi, - u1/))
if at. Otg otherwise
e - e 6 2.
sign (as, - at) sign (bo(s) - bo())) sign S0
Definition 5 (a fixes F). Fix a set of fixed indices F C Nenax and an assignment of token values to each of the fixed positions tp F - N Eedwear. (F is the set of positions for which we are not pessimizing over the value of the token that position.) Fix a non-decreasing sequence of tokens
to S... S..Ska Lnax-1 E N.
Given a permutation NNenex + Nuax, say that fixes F (relative to beb...,,beaa--- if t((a(()) whenever o E F.
Note that in this section, for the cubic proofs, we will in fact generally take = {Mety - }}, 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 /max Appendix F, we will take F or F to the set of positions of the maximum token. However, none of these theorems specific to F being subsingleton, and prove them in generality.
Definition poosition-sorting permutation). Fix a set fixed indices C Menor and an assignment of token values to each of the fixed positions iy F - Medvocab
Define the position-sorting permutation fixing indices in F 0g - Nenax N Nenax be the permutation that sorts the indices not in F according to b: for O S i.j < Beex with i, & F. b; S bj whenever (i) os (j); and o. (i) = i for E F.
equal tokens). Fix set fixed indices F C Ne and assignment of token values to each of the fixed positions tF : F - Nedwmarr Fix a non-decreasing sequence of tokens to S .... S. n N.
Definition 7 (contiguous
Say that the sequence represented by a permutation a Nenan - Nenas is contiguous on equal tokens if. for all ( < i,j,k < nexx with I; = 11 # by and i, k $ (F),, it is never the case that a. (((((i)) (k)) < (o(r))))
Theorem (Pessimization over sequence ordering is possible and results in contiguous sequences). Fix a set of fixed indices F C Nenax and assignment of token values to each of the fixed positions
tp : Nedvoerr- Fix a non-decreasing sequence of tokens to S ... t M.
Let minn max : N + N be permutations of the netx positions, fixing positions in F, satisfying the following property: For all a N 1 a permutation fixing F, we have
Samin
Sa S Samax
(10)
(Such permutations guaranteed to exist because the permutation group on Teexx elements is finite.)
Then Tmax and amin may be taken be contiguous on equal tokens. That is, there exist maax and Juin 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 Fmax and Fmin cannot be contiguous on equal tokens and derive 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 thus violating the presumption that the permutation is closest
31

---

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. & such that the sign of the score difference from swapping and is different from the sign of the score difference from swapping and k.
First, definition and some helpful facts about it.
Definition 8 (contiguous on equally-attended positions). Fix set of fixed indices F C Nenax and an assignment of token values to each of the fixed positions tF FF - Medhaaas. Fix non-derreasing sequence of tokens to S S fna--1 N.
Say that a permutation a is contiguous on equally-attended positions if. for all S < Betx with i (F), the sorting order according d, on the contiguous block of positions with contribution to the attention score equal to that of (i), {o(ā) bo(s) ba(e) and ojj) & F). is the same as the sorting order uccording to the fraction of tokens equal to I with b-values greater than bo(s), with ties broken by the value of t.. Equationally, this second sorting order is defined by the score
[K | tt and bo(k) ba(i) and (k) $ F]]+
tj dvvccab
/IEE) [ and o (k) & F}
Most importantly, any permutation that contiguous on equally-attendel positions has the property that for any indices 0 S j. k < nax with i, k & (F) and t; tt * t and (a)) < o,(a (k))) < az(a(j)), we will have the strict inequality bo(s) < by(x) ba(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-violationss which we avoid up-front by arbitrarily swapping involved indices without changing the score Sa.
Definition (needlessly non-contiguous). Fix set of fixed indices F S Nener and an assignment of token values to each of the fixed positions tF FF - Nidivocar' Fix non-decreasing sequence of tokens to S ... S fmax-11 E N.
Say that permutaation a is needlessly non-contiguous at i, k (for i.j.k & o (F)) D = 0 or Aajesk = 0, for 0 < i,j, k < Tetx with i, k & o (F) with L; = 11 # la and (o()(i)) < (,(a(k)) (a(()))
Say that a permutation a needlessly non-contiguous if it is needlessly non-contiguous at any i, j. & a (F)).
Lemma 5. Fix a set of fixed indices F C Nenez and an assignment of token values to each of the fixed positions tp : FF Nedwocar Fix non-decreasing sequence of tokens to S ... S 1 E N.
Any needlessly non-contiguous sequence which fixes F can be made into a sequence o' which still fixes F and is both simultaneously contiguous on equally-attedeed positions and not needlessly non-contiguous, and for which s, Sg¹.
Proof. First, sort regions of equally-attendee positions make contiguous on equally positions. If the resulting permutation is not needlessly non-contiguous, then we are done.
Otherwise, we have A 0 Ao.j+ 0 for some i, j. k, for 0 < i,j,k < netx with i, & (FF) and t, = # t and (I)( < a, (o(t)) < (o(()))). Since the sequence is contiguous on equally
positions, we have the strict inequality bo(1) < bo(k)) < bag)-
By Lemma 3, we have two cases. Noting that 4 = lj. we can write them as
1. Utg Ut, and dr: arz
ve, ct - clk
2. at; # ar and So
e
32

---

In the first case, we may fully freely interchange tokens equal to is with tokens equal to ty without changing the score; in this case we may use the token value as sorting ti-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 Sg, 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 or E {Tmaxx miin such that for some 0 S j, k < Petx with i, k $ (F)) and 4, l o lg. we have bo(() < bo(k) < be())- We will derive a contradiction with the presumption that a is extremal by showing that we can swap i and k to change the score in one direction and that we can swap and & change the score in the other direction.
Take % to be but swapping i and k; and take of to be but swapping j and k.
Now we will consider the cases for the sign of the score difference Ap = Sog -8 8. and D == = - S.- By the presumption of not being needlessly con-contiguous, A * for 2 E (0, 1}. If we can show that the sign of Ao is distinct from the sign of A., then we will have a contradiction with extremality because we will have either 506 < 8. < fof or S0; < 8. < S06- That is, we would be able to swap i <+ k and] ++ k to get lower and higher score, making not extremal.
Noting that t = tg.
sign (Utz Ut,)
if at, at. otherwise if ar, = atx otherwise
sign (Do) sign (bo(s) bo(k))
"e, ve,e" - Veg e'ta e -
sign (at. Orp) sign Sa
sign (Vix vt,)
sign (A1) = sign (bo() - (o))))
n"ee un n"te o"
sign (ae, az,,) sign 80
Noting that the product is non-zero by presumption, that right multiplicand is equal for Ao and A and sign (ba(s) bo())) - and sign (bo()) - bb(k)) = 1, we have our desired contradiction. O
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 non-decreasing sequence of tokens to S... S tn -1 N Noduo and a permutation a N + N Nenax define the full sequence score Stg...stmesx as:
E(vr, + (o(()) ttathoo
thato
Σ Deaithooo
t0
OSi<tiet
OSienaa
We will drop the token subscript, writing only s',, when the token values are unambiguous by context.
The sequence score here will be computing APEVVO t" and fmax. As with Definition 3, with the way we we've set our definitions, high scores predict t* (and are thus bad), negative scores predict mmax (and are thus good), and more negative the scores, the stronger the prediction of Lmaxx.
AEEOO :- + APPOU -- for some fixed
Definition 11 (relaxed sequence score). Given non-decreasing sequence of tokens to S I E Nedwoest and permutation o : Nenax - Nenas define the relaxed sequence scores Tto m, and Ito.....naa 1,,,, m,x, as!
Tta 1.mminn Stg...dfaa 1.0 min Wi
Osi<nax
To -,0 == Sta 1.0
max W/s
We will drop the token subscript, writing only Ta,min or T Ta,muxs when the token values are unambigu- ous by context.
33

---

Theorem 6 (Independent pessimization over positional contributions is possible), Fix non non-decreasing sequences of tokens to
-eets-1 N and to S S new-- E N and permutations N eners Let mom and max denote Ita m.ommin and Ita maxi
o, Nenna
let Se denote maa 1 and let Sat and s'g, denote 81% Sig. 1:00 and 8/6 Vg M -
Then we have
min Wi = Ta,min 8g S 8'g 80 S Ta,max 8g 0S1-Rei
max Wi.
0Si<nela
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
To,min
5o min (()) 0Si<naxx
min w 0Siches
Ta,max
max Wo(()
max DSiR
0Sicnes
Σ
Woreee a 44
Σ E tot
Osicnas
Osicnat
og (igtbor(1)
ΣEn "oo(tt)
By tb)())
OSichetx
Σ Ensjeneu
e
Since e" is non-negative for all real x, we have
Eosienax Eosiemal
Ensienous e4 antthhrr
min (
S S' So S max ()) 0si<nax
my tbargs
Eostenar en thrtr
Thus we get as desired
min S S'g OEicnesx
< max th: 0Si<netx
Note that could prove more fine-grained theorem, that pessimizes attention paid to positions only for sequences compatible with the chosen fixed tokens and tF, but since the positional contribution is small we do not bother.
Theorem 7 (For fixed ordering, softmax is convex over token counts and only pure sequences need be considered). Fix set of fixed indices F S N Nenes and an assignment of token values to each of the fixed positions ip = N Redunar Fix a set S C Neddoas of valid other tokens in the sequence. (In our uses of this theorem, will be the largest subset of Netoax for which we can guarantee that the model behaves correctly all sequences compatible with F and tmax and with tokens otherwise drawn from S.)
Define a comparison on non-negative integers less than dyocabi
c:= Σ
(
d d:=
+b)
f : Σe c
iEF
iEF
Osieneix 1
cmp(z, y) == sign (d(e"=uz - c'vNy cle== ent) + fea-tay (uzea bye
Let lemp min and l'mp max be the minimum and maximun elements of S according to cmp.." 21
For given choice of a non-decreasing sequence of tokens to S ... S tnax-11 N compatible with F and and given choice of permutation : N + NN of the netx positions fixing F (t: = tF(a(i))
We will prove that cmp is transitive in the process of proving this theorem.
34

---

for o(i) E F; and t, E S for o(i) & F): let So,iin (and max) denote -o 1,0 when t, = Sempmini for all (i) & F (or temp max- respectively).
Then for all such choices of requence-permuttin pairs,
so,min S Slp...n/nax-115 S Sa,max.
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
For any indices S < < Retx token values 2, y E S, the score for a sequence with t, = x + y is bounded on both sides by sequences with tj and tę
Proof. Let Sa.3 be the sequence
with 4 a and L 3, and define the differences
Ae == $2.x 82,y and Ay == Sv.y Sw.y- It suffices to show that sign(A, S 0. To show this, we must only compute the sign of Aa for a E {x,y) and show that whenever both Ag and Ay are non-zero, they have opposite signs.
We proceed by computation after defining
convenience variables for brevity:
C:= Σ Vta tbock)
tbock)
D c (a()))
USk<naxx 1ii
OSk<natx Kif
if a
i if a
if
ā
La
ia
y
if
j if
if y
Un e + vn catttt
t uyc aa- + ebutdors + D tbocii + vae vaeaatburn + C thotholiaa + e th(ag) +
sign (Aa) = sign
aatboll + tbocss +
eataatbolnn + Eae tateatbolia +
sign
thatholial + e thara) +
Multiply through by positive denominators and simplify
sign (c (easte (aatto(aa) ca((ta)taa) + D (ooa) Dac (acaabbolis))
(ol(ta) + (locio)) an catba(ea) v8 (e (clov)) e aa aattalue)too Va
Pulling out (e(ta)
sign (ctato (c'otta) + cloww) (va va) + C( =o)) (c"ava -
Note that swapping and à negates the sign. Henee, have sign (A,) sign ( and hence 8n, I 8x, S 8y.y or 8y. S 8r,y I 8z.a as desired.
L.
Lemma 9 (Pure sequences are sorted according to emp in Theorem 7). Fix all the same quantities as in Theorem 7.
Fix tokens I, E S. Let n := Bexx |F be the number of non-fixed tokens. Fix sequences with copies of x and respectively: fix t,,0 S ..... S tx E N and ty.0 S... S ty,max 1 E N compatible with F and S and given choices of permutations az, og. .N N of the netx positions fixing F: tz,: (i))((i) for (i) F: ty.s tt(aa(()) for (1) (i) E F; tz,i = x for (1))) & F: and tyy for o ( F.
Then
sign( (So, do oonn (Say (y cmp(x,))
35

---

EQ == B,
:= -
EVOU(4) I,EVOU
PVOU()) P,VOU
e (t-1) = t-
E (t-1,
1,-1,,ma maxt, f4)= == eFC (t-1)) - BEU (t-1(t))ma t;
Figure 13: Recapitulaiion of some relevant definitions from Figure 12, parameterized by the arguments they actually depend on.
Proof. The proof goes by straightforward computation.
sign((Boa.tt.ommtteeeee (90 fyaainntnnee
U f
Uye +
= ssign
e f
e00 f+d
Multiply through by non-negative denominators
sign (Vv,e c) (E"V d) (Uye"u (c)
= sign (-efe" + cfe" + dfu - dfuyemv + u f² yyenrta))
Use f >0
sign (-ce az 06" dvze - diiee" + Sugerrayy - Cesev++Sueee
sign (o(e-y-ee (2zear - vyeav) (Ugeartay uyeastav)
emp(æ,y)
Corollary 10. Define the relation Semp by Samp y if and only if cmp(x, v) E {-1,0]. The relation Semp always transitive.
Proof. Note that by Lemma 9, cmp is comparing two sequence scores. Since < is transitive over the reals, the relation Semp is also transitive.
Finally, we combine the previous lemmas to complete our proof of Theorem
Proof of Theorem? 7. Extremal sequences with scores 8 mim and $a,max 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 Lemma 9, the extremal sequences must have tokens that are extremal according to cmp.
t, whenever (()),(() F). By
L
We now have all the tools necessary to prove the following theorem. We refer Algorithm 3 and Algorithm 4 or the proof of Theorem 11 for definition of the CUBIC algorithm. Theorem 11.
Et-U(0 L... 1aaaab--"""" 1)tt
argmax( = maxt. CUBIC(dyocab netx: M)
Before we give the proof of this theorem, we introduce some helpful notation. )2 (voccd) {0....3) such that I'm 2 TTg. We define
Definition 12. Fix an element (rm, Tq. c) E (O..., dvocab)
X ()
e) to be the set of tokens such that
1. The max token tmax is equal Te.
36

---

Algorithm 2 Counting Correct Sequences in Cubic Time: Preliminaries
1: function CORRECTNESS(M, input-sequence)
2: return MODEL-BEHAVVOR((M, input-sequence) MAX(inut--eequence)
3: end function
4: function MODEL--BEHAVOOR(MM input-sequence)
Require: input-sequence is a tensor of shape (Netx,) with values in Nedrocab
5: 1max <- MAx(inuttssquuen
mmax f max-token
6:
input-sequence
7:
skip-score, k A
(- (tnre (max)
8: attn-weights-unscaled, EOKE( 4,44 EQKP((nav-1.
9:
attn attn-weights +- SOFTMAX SOFTMAX (oorrMAA(ttb--wiggtt--ssslla
10: U/ <- ())
11: W/ PVOU(1)
12:
At'r f- vt,-" - fmax
Aw,, Wi,z" Wi,fmax
14:
return maxe" ffema (skip-score,- + Me (Avr. + Aw,rr) attn-weights;)
15: end function
16: function
POO-IIIGOOOOOOPSSSSTTIOOOLOOOLO input-sequence)
17:
t input-sequence
18: return ALL(CORRECTNESS(M perm + [t-1]) for all perm E PERMUTATIONS(
19: end function
2. The query token lquery is equal to Tap
3. The cardinality of tokens that are not at the query position and equal tmax is equal to c.
For clarity, we list all the possible cases. We always take lquery S mmax and let S3 sequences by permuting the first three factors (i.e. keeping the query position fixed).
1. If O, then X (Memax Equery 0) { [I Imax mmax lquery]>.
2. If 1, then X (taanx 1) S3- (11 limnses lmnxx lquery] t < [mmax),
3. If 2, then X (Mmmax)quers 2) S3 (1-(11, 12, Imaax lquery] < (max).
4. If 3, then X ((tmaxx 3) Sz (1. (1,44 12, 13, mmax) t (max)..
Definition 13. Let t E X be a sequence. We say is pure, if it has at most three distinct tokens: the max token tmax. the query token Equery, and optionally a third token t² tmax-
We denote by Xpure the subset of pure tokens. For any subset C X, we set Ypure := Y n X
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(dyocab Retx: computes these bounds.
Proof of Theorem 11. First of all, we note that the algorithm CUBIC = CUBIC(dvocab. Retx M)) yields lower bound for the accuracy on the set fimmer) We can therefore compute the bound on X [ X (fmaxx by computing for each such choice (Imaxs lquery: e) and
summing over them
E0,U dvorab Pete
arg)) maxt
Σ
CUBIC(X (fmeee
So from now we will fix one such subset X ((Cmaxe We begin by defining map
X (Umae: {0, dyvoca)) J
which sends a sequence to the subsequence of elements which are not at the query position and not equal mmax. Then Theorem 7 can be restated as 22
22 In fact, the theorem yields stronger result, but will only need the following formulation.
37

---

Algorithm 3 Counting Correct Sequences in Cubic Time, Part Lines are annotated with comments indicating the parameters for a cache to avoid duplicate computations.
function MODEL-BEHAVIOR- REI.NOOR,RRR.EEDMMM query-tok, max-tok, non-max-tok, n-copies-
nonmax)
2:
Equery +- query-tok, max f max-tok, t' f- non-max-tok, c + n-copics-nonmax
Require: 0 S fquery S mmax < dyocab. 0 s t' max < dvocab. 0 S e < Netxx
Require: if n-copies-nonmax
then non-max-tok = max-tok
Require: if query-tok max-tok then n-copies-nonmat
Retx -1
Ensure: return 2
t) for all with specified fquery, c copies of in non- query positions, and the remainder of the tokens equal to max
3:
skip-score,- <- A A@EU, (- ((Qqquory bmaxx
D Cache by mmax lquory- t Cache by i Cache by max: t"
4:
Wy (VOU(() for 0 S i Netx
5:
A.max, < maxosicnax (Wi,r Wiftmax
6:
U/ < EVOU(!). AUr + U2,3" Vt.tmax
for {tquery mmxxs t'} Cache by tmax
t"
7:
a, < EQKE quurry t)/va for E {qquery Imaxs {'}
Cache by lquery
8:
D 1 f- EQKP query, neex - 1)/vd
> Cache by Equery
9:
bn.-1 + SORT (Lquery: 1))/va
Cache by lquery i
b1.-11 f REVERSE bo.--)
11:
attn attn-weights winn-wigghs--usaaleddd unscaled, -1 f Agquary + brees= -1
D Cache by Equery Cache by fmax C, i,
12:
attn-weighss-unccaldd, f- 0x + bo.i for 1 < Natx C
query
13:
alth-weights-unscalldd, + Atmax + b1, for S < 1letx
Cache by fmax. C, i,
query
14: attn-weights-uuscalldd. f Opp + bo,i
for Betx - c 1 < < Retx - P. Cache by t', c, i,
query
15:
attn hiin-weights-unsee unscaled,, + ov + b1ss for Betx c 1 i < Retx b Cache by t', c, i,
Equery
16:
attn- weights + SOFTMAX(atmn weiteetttttt-wwhht--ttt weights-wussaaleeddd unscaled,) attn-weights, f (eTTMAX(att--weihts--umm
> Cache by tmax t'. C, i, tquery b Cache by Imax 1', C, i, lquery
18:
if = 0) then In this case, attn-weights,, attn-weights
1, so we drop the first subscript
19:
return max; ttms (skip-score, + Awmax."
+ Av.. 1.2 attn-weighht
Ut Umaxx. else
2i aeighweights,)
20:
21:
c C- Avi f- Ammmat for S < N
22:
AUi Ave. for Netx
<i << Nexx -1
23:
AU &- Au Up query. Retx-
Mltx
max aaax (AWmax, + Avi.)) attn-weights
24:
return maxx+ scip-score, +max
1==
E
neta-1
max, Hma (A (AWmaxxtt + Ariee) aeighteights
25:
end if
26: end function
27: function RELAXED- RELAXED-CORRECT
POTTERR-OOOIOOOM,, fquerys fmax f', c) >> runss the model on relaxed variant of input sequences compatible with the arguments SEEZING-VERRPPO/TION--SS POw.TION-SLOWWWCCCC..... for any t with specified (query, e copies of t' in non-query positions, and the remainder of the tokens equal to max return REEELL-EEAAIORREEEEE LELAXED(M, Equery tmax, t² < 0
28:
Ensure: return is False if
False
29:
end function
38

---

Algorithm 4 Counting Correct Sequences in Cubic Time, Part
1: function CUBIC(dvncab. Netxi M))
2: count
> of correct sequences
3: for Emax RANGG(dvocsb) do
> max f max-token
4: for 0 S Equery I max do Cunxx t Actx if query max else Netx 2 D maximum copies of nonmax Cmax do
Equery + query-token
5:
6:
for 0 S
D number of valid choices for the non-max token RCPOP(X) < PARAAADDDDCRREECTTSSSSSPSEESEEEE TOA.H/IINNGOOEER-PPPIIIOOIOOOOO POSITION X)
7:
8:
c then
9:
t-count + if HERCCPP (Iquery, Imax, mmax, 0) else 0
10:
else
t-count Eimer -1 if (RCCOO((4qaae) Imaxs L', else 0
11:
12:
end if
13:
count +- count ("ave]) (t-count)
taking 0 0 conventionally
14:
end for end for
15:
16:
end for
17:
return count
deacab Netx
18:
end function
Let C .o..., d Then full accuracy f (S = X (fmaxe Aquaryy
Of ( implies full
accuracy on f f
Now instead of computing the output of the model for every element f (Se) we use Theorem (combined with Theorem 6) 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), have either b be_.x < {bi, by} or by bymax > (b,, (bs,by], where
relaxed version of this. In particular,
may assume that the pure
j E (0,..., Matx - }} are indices of tokens not equal to fmax-
For the algorithm CUBIC(dvocal fletx .M) we fix a t" E (0, fmax 1} (unless O, in which case there is no such choice). We then run the relaxed accuracy computation RCPOP(tquerss fmax f',c) as described in Theorem 6. If RCPOP((query) Lmaxx t', c) < O, we add e' to S. If we do, we add t to S. Therefore by construction of S we know that we get full accuracy on (Se) and therefore we get full accuracy on fS-(
Now we count the cardinality of (5*) and add it to the count of correct sequences.
Theorem 12. The rumning time of Algorithm 3, after using caching avoid duplicate computations, Odvocab 3 "eexx" 2).
Proof. The nested loops in CUBIC execute the innermost body O(dvocab ( times, and the summation Line 13 costs O(Rex)) per iteration. What remains to show that the call (M,TTION(,,tqueryys) tquery: fmax: f., costs ((netx)
to when e and at most (etocab when = 0 and t' fmaxx-
The matrix multiplications in EQKE, EQKP, EVOU, PVOU, and EEU can be cached upfront, costing O (mmax(d dmodel, ) 2 I Odlvocab 3) since we assume dvocah > dmoddl and dvocab Netx-
The sorting on Line can also be cached upfront (per lquery), costing Oddoeahh log Netx). POS-IIIIIOO
Note that each variable assignment in RELAXED- be cached into a table parameterized over at most three variables which range over dvocab and over at most two variables that range over nltx-
What remains is the return statements.
When c = O. we have on Line 19: return max." stase (skip-score,. + AWmax,
Att i attn- weights + Attams." Dmar= weightwights,). This is O(dvocabMetx) as desired.
39

---

When * O, we have on Line 24:
Mas-1 == max, +laxx (AWmax."* + Avi,t)) attn- weighwegh
return max skip-score,, + max 6 flmaa
2i=0 max,' fftmar (AWmax, + Aui,-) a
We can cache max; skip-score, per lmax and lquery. costing (deneab mexx)) We can cache m ffaax (AWmax.t' + Ati,))) per Imax and t' costing (dvocab "), since each A will be Avtr for some E q'query fmax: t'}. Finally, we can compute the summation in cost ( per loop iteration, as required.
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 convexity of softmax where, modulo pessimizing in unrealistic ways over the attention paid positions for the computation done positional encodings, all extremal relaxed sequences correspond to actual sequences.
When we only get budget of (dvocab 2 netxx 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 have only two tokens (the max token 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) computa- tions in order 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 (Common theorem parameters). Fix total number of tokens Retx- Fix a token value function (à la row-differene in EVOU) Nedgosab + IR and token attention function (à la EQKE for a fixed query token) a Nedvocat - Fix position value function (å la a row-difference in PVOU) Nenax -> R and a position attention function (å la EQKP for fixed query token) b b: Nen.x - IR
In practice, as in Appendix E.3, we'll take, for fixed query token lquery>
1/ EVOU,,, - EVOU t,bmace
= EQKE, Lo Vanerr.. t Vd
N0: PVOU PVOU E,maax
b; = EQQPP, Lquerry
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 sequence of tokens via mapping from positions by specifying subset of valid tokens S C Neldrocs paired with function T: Nenax + S specifying which token is in each position.
Definition 16 (sequence score). Given a subset of valid tokens S C Ned and function T N Nenax -> S specifying which token is in each position, define the sequence score
Σ TrO
8T
OSistiax
oSi<nav
Definition 17 (swapped mapping). Given a subset of valid tokens C Nedwoas and a function T Nenax -> S specifying which token is in each position and two indices 0 S i,j < Betx define the swapped mapping Tiesj be the function that except swapping and j:
T(i)
ik=
Teme(() (k)
T())
fk=
T(k)
otherwise
40

---

Lemma 13 (Characterization of swapping tokens in two-token sequence). Fix two tokens to < E N and function Nenex - (to, b1} specifying which token in each position.
Define Arjs to be the difference in sequence scores when you swap and j:
ATjio = ST.03 - ST
Then
sign (Ar,ses)) = sign (b; - by) sign (VT()) vTGi)
Proof. Lemma gives us the result directly when ar(s) aT(J)" Otherwise, we get sign (Ar,st)) = sign (aT(s) - aT())) sign (b; by) sign ST UT(G)e VTTT)
e e
Hence all that remains is show that
sign (ST (e"T(o - - VT(G) TT + oTG))eTro)) -iign (UT(3) vT(i))
Define 0 (VT()) + vTO)) and define Av ((TT() vTO) that vT()) 1 Av and
(TU) = Av. Assume WLOG that T(t) = and T(j) = so that UT(M) = + (-1) (p)) Au for all p.
Then we have
sign (sT (e (TT()) VT( + (T())
sign (sT (e eTT())) - (C -(TT)) Ave - Aveerun)
"T(
sign
Os DSpemaxx
(c - cTT6)) (c_T(a) (aT())) (C(TT))) + c
Σ
C thp
OSp<neix
Σ
(++ Ao) earint
sign
DSp
(E c - 8 (C [ Ao (c + coTUs))
tbp
OSpenetx
e (
6 (T() Σ c
OSpenetx T())=T(1)
OSp<neex
T(p)-T()
sign (A)) sign
(c eTTUI - er(() TT
Σ +dp
O<p<netx
exr(a)
bp e (Ti))
bp
OSpERess T(m==T()
T(p)=T()
sign(Vr(e) PTG)) sign
- eTT - 0T( eTT (6070)-000))
tbe
Define
P = Σ d
p=== = Σ
OSpCRax T()]=((3)
OSPcnas T()==())
so that get
sign
(ST (e (TT() "TT() (Ti) + (TO))
41

---

( P. c P3
=iign vT()) sign
e P. + e (e ( -e)(()) caT())
Multiply through by the positive denominator and expand out so that we get (ign(Ur(() - (TG)) sign (-2 (-2e " (eT()) P. tero)) Pş)
sign(VT()) *T()) sign (e P, er(a P3) sign(Vr(() *T())
Theorem 14 (Pessimization over sequence ordering for two-token sequences is simple). Let 0g N N denote permutation of the nexx positions that sorts them according b: for 0 S i, Nexx.
S by whenever oi) < ds( Fix two tokens to < t1 E N.
Let ret be the number of E [0, netx) with T(p) to and let nı be the number of E [0, Retx) with T(p) t, Note that nio nı, neexx
Define min = argmin, (fo.11) Ut and define mmax == argmaX (fa.t1] Ut-
Define Tiin: max Nenei + {to, }1} to be the assignment of tokens to positions that pays the least (respectively, most) attention to Imaxi
Fmax
if0 I (z( < Ratmxx
T Tmin (i) =
tmin
if f S (z netx
tiii
if0 S σ (i) < Aatin
T Timax (i)
tmax
if ft S og (1) Neexx
Then we have that
miin
S 8T S 8Taax
Proof. The extremality of 8Tiin and Tiaw follows straightforwardly from Theorem 4.
All that remains is STmi S STmav'
This follows from noting by Lemma 13 that swapping two tokens in Tmin increases the sequence score, while the reverse true of 8T.... thus showing that it must be ST..... that is the minimum and 8Tmax that is the maximum and not vice versa.
Definition 18 (full sequence score). Given a subset of valid tokens S C Nedrosat and a function T Nznax + S specifying which token is in each position define the full sequence score s
1/1 = Σ (UrG) + we +b)
Σ +b
USiSna
Theorem 15 (Independent pessimization over positional contributions and token attention and token value is possible). Fix two tokens to < ti E N. Let Tmin Tmux Nenax + {to, f1 and f fiin be as in Theorem 14. Fix a set of valid tokens with to, t1 S.
Define relaxed versions Th T :Nenav
+
S of Tmax and Tmin'
Timx(()
if (i) bmax otherwise
T T
argmin jES 0j
Tmin
Tm)x = +max
Thi (i)
argmax jES aj
otherwise
jfhnxx
That is, Tt Thax replaces Liin with whatever token in S draws the least attention away from bmaxs while T replaces trin with whichever token in S draws the most attention away from tmax-
42

---

Define relaxed extremal sequence scores T'mmxe T'min
Tiin
min Wi Ogicn
Σ (T)
axy tbs
Σ
any (1) tb;
0sianetx
USta
Tmmax
max +
Σ OSicnay
TTimm (ixr)))
Σ Ocienes
Timas (6)
USicnar
Then Timin T'mm and
Tmax
PTmmxe
Proof. (sketch) Essentially the same as the proof of Theorem 6.
Note that in practice, we take to be the of all tokens less than tmax 9 for some minimum gap 9. 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
function w dyocab. Metx Imaxs lquuryy C
9)
2:
t input-sequence
3: return False if & (Nedvoca)) Metx
D the sequence is not made of valid tokens wrong query token D. wrong max token wrong count of non-max toks check gap non-max toks
4: return False if 1_1 f 'query
5: return False if max; 44 4mmx
6: return False if H E Nenax + mmax
7: return ALL(t; Emax or tmax t, for 0 S Retx)
8: end function
9: function CORRECTNESS PEREIIINNSSS GOPPSSSLLM dvocab. Retx fmax Qquery + g)
10: return ALLCCOORECTNESSSMM, t) for all t s.t. INNUT-FEQUENCE--CMMRATTIE (tT
dyocab; Reix, maxx fquery, C, g))
11: end function
12: function SUBCUBIC(dracabr Bexxs M, G)
13:
count +0
D of correct sequences
Gimaattuwcryy MIN(4max. MAX(1, Gluaaxfquaaysd ))
Clip G to valid range D Cache running minima
15:
Gimserc f- min Gimastte for Emax RRANGE(dvooab do for 0 S tquery & tmax do
> mma + token
17:
Equery + query-token
Omin 0 if qqerry max else 1
minimum copies of nonmax
max netx - otherwise
19:
Cuax
+
> maximum copics of nonmax
20:
for Cain < C I Cmax do
D valid choices for the number of non-max tokens
21:
9 Gamarrauurry.e
22:
g" + G A
23:
q-gap < Emax
RCPOG(X) <- Equery
24:
25:
if (q-gap- 0 or q-gap 2 9) and CCPOG(dyocate Netx /maxx querry C, 9, 9
then
d + Equery tmax else c- count +- count + (Meey- (Imax g)
of non-max non-query tokens
27:
taking 0 = conventionally
28:
end if end for end for
29:
30:
31:
end for
32:
return count deocab
33:
end function
43

---

Algorithm Counting Correct Sequences Subcubic Time
1: function BOODLL-BBLINIIOR-RAAXXDD-----OO -GPMM fmax lquerys C, 9, 9²) Ensure: CORRECTNESS-PESST CORICTTNSSSPISSSIZZNNGVVOOOCCC.
is False = result is False
Require: 0 S g S S mmax
Require: if
then Equery = mmax
2: skip-score max," ( (fquery)(- min ( (tqueryy) )
Cache by Equery
3: Vt f EVOU())
4: W; PVOU(i)
5: A"mmxxt' & max; Wi,t* - Wiffmax
Cache by mmax t" Cache by fmax D Cache by
6:
Almax.max max, !
7:
Att f max.- 01,1- min.- Ut,
8:
AUmax maXostStmxx Avt.
Cache by tmax g" Cache by Imax Cache by 4max
9:
Anfrax t- Vtmaax tmax
Autmax t- maxT tl Aui
11:
c = O then
f.- f- {.U (tmaxx) Vmmaer A Amaax,1
13:
return maxt Hlaax (br- (
end if
15:
b D.mexx + EQKP( (t query, Netx 1)/Vd
D. Cache by lquery Cache by Equery
16:
bu.- SORT SORT(EQKP Equerr: 1))/Vd
17:
b1.--1 <- REVERSE (bo.-1)
18:
at < EQKE queery: t)/vd
Cache by fquery, t
19:
amin,t + minoge"st Oy!
b Cache by lquery. L, compute in amortized O Odvoeahh" 2)
20:
amax,1 + maX St Oe
Cache by Equery t, compute in amortized O (dvocaa" 2) Cache by lquery, mmaxx Cache by fquery. tmax C
21:
A0max 0tmax fmin,
22:
AOmin t Otmax amax,tmax - 1] if max F fquurr else {0..... Feexx
23:
idx-set < {0. .0, netx
netx
24:
attn-weights-unscaedd, +- bo,5 + (Aamin if i E idx-set else 0)
25:
atnn-weighss-unscaldd,; +- b1.5 + (Aamax if E idx-set else 0) Cache by lquery- Amaxx i
26:
attn attg-weights +
SOFTMAX(attn wwweaaattttttittsn-sssaalae
>> Cache by fquery. fmax i, C
atinhweights, SOFTMAX attn- wwightt-
Cache by Equery, maxx i,
28:
attn-maxo iEidx-set attn- wutn-weights,,
29:
attn-max, + iEidx-set attn witt-weights,,
30:
attn-max attn-maxo if Autmax < Aumax else attn-max; > Recall that Autmor is negative when the model outputs the correct answer
31:
32: return scip-score AWmaax,tax + attn-max A Aufmax + (1 attn- (aanmmmx) Aemax
33: end function
34: function RELAXED-CORRECTNEN PEWSIMIIING-OEERREEEE RELAEEBBRRRRRCCSSSSSSSSAAAAAACCE
dyocab. Retx tmaxs fquery
9- g.)
35:
the model relaxed variant of input sequences compatible with the arguments is False = result is False
Ensure:
Ensure: return False if t with specified mmax Equery, and e tokens not equal to tmax fmaxi lquery: C, G.
is False for any
return BADDEL-BEAVIIO-- WEW.L//IIRRRELLLE-VEEEEEEEEEEEEE
37: end function
44

---

Theorem 16. For all G.
(0. doorb 10
Argmax(M(())) max 1i
SUBCUBIC(dvoaab. Retx: M)
Proof. (sketch) Apply preceding lemmas and theorems Algorithm
Theorem 17. The running time of Algorithm 6, after using caching to avoid duplicate computations, (dloocabddmmdell dyvcca
2 "eexx" 2).
Proof. (sketch) Sum the complexities indicated along the right side of Algorithm 3. The dvncab mmodel term comes from the precomputing EVOU, EU, and EQKP. The 2 term comes from the softmax over Autx tokens for (dvoch Nexx) pessimized pure sequences. Con- firming that none of the complexities the right side exceeds (ddvocab dmodel deocab netx ?) completes the proof.
G Subeubic proof strategies
In this section, we present number of proof strategies that to reduce the computational cost of the proof, ultimately driving down the cost of EU and EQKE verification to ([dyocab@modell) while unfortunately leaving the cost of EVOU verification Oddooctt "(ddocat""dmodaa ddmddel).
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 FQKE (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 hope of non-vacuous bounds below (dmodel)))mmde).... and the max row-diff trick is the workhorse that allows to drive down the error term computations from cubic to quadratic without getting completely vacuous bounds.
get to insert the most understanding (without which we'r ''d have
G.1 The mean+diff' trick
Suppose we have quantities fz.y and Py.: and we want to pessimize (WLOG, suppose minimize) the quantity fzww + Iy.= over y, and in time less than Onn.,ynn)) say we allow Om + ny n. Also suppose the variation over the y axis is much larger than the variation of over the x-axis.
We of course say
min fay + min Oy.: I fz.y 9y.= 2,9
2.2
But we can do better!
Note that
Ja.y Eyy (Jr.yy F.
Suppose that fz,y varies much less over than it does over y, and much less than 9y,z varies either of y and z. This will make the following bound good approximation, though the bound is sound even without this assumption. We write
Jay 9y.z
min[Ja.y + 9y.z]
2.0.2
min Ex,fa.y + Iy.z fa. - Exfaaul 2.8.2
min E + my.=1 0.31.
minlfz.v E fz.y)
2,0.2
min
[E + My.el
minlfa.. - Ex fawl 2.8
9.2
By averaging the variation over certain have
45

---

Theorem 18 (Mean+Diff).
min fz.y 9y. 2 min fa. + gy.zl + min(fay E. Enfa
7:9.2
2.1
max Jan + 9n.= S [Ex + gy.= + maxx(aa - Enssy 2.y.2
9.2
and the RHSs can be computed in time ny + nyns + n. for for nz ny, and By the number of
possible values of z. y, and respectively.
Example for how this helps with small variation:
Take any function k(y) and then take
fay: = k(y) + Ei ()
99. == kky) + E2(8.
Then we have
min + -y.=]
minle | (2, y) + E2(U, =)]
v.y.2
min fa,y 2.y
min Jy.z v.2
min k(y) min k(y) min 8n (x, + min 02(1, =) via
y
V
2.9
min k(y) max k(y) + min a1 (x. + min (y 2)
9.a
min[fz.y - Ex.f.y + 2:31
min My. + Ex Exyyy N.z
mine, (2, v) min(s + E, [I)
2.
2.2
If E and E2 are small compared to min, k(y) max, k(y), then using E,, gives 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 Ex If f is matrix indexed by x and y, we might also try taking SVD and using the first principal component instead. basic application of the triangle inequality gives the following, more general, result:
Theorem 19 DSummarize+Diff), For any hy which can be computed in time O(ns),
min Sz. + 9y.x min[hy, [by + 9y,el + minlfa. hyl
2.y
fz. 99.x x,2,=
max(hy + 9y.x] + max hy] 9.2
x.y
and the RHSs can be computed in time O(nnny + ny nn.) for nw Ra, and ny the number of possible values of z, y, and 2, respectively.
We that if the variation of f in the x-axis indeed much smaller than the variation in the y-axis, then letting
Sa by Ex,9
we get
min Ja,y + 9y.2 0.1.
minnnn (hy + gy.2] min[fa.y lyg 2.0
2,31
min I2 + Hy.e2 min(h,, (hy + sy.z] 8.312
min
(5x.g)
9.a
< 2 max [Ew.yl
so indeed this bound isn't '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 Oddocabb -mooel).
46

---

G.2.1 Comments on relationship between mechanistic insight and proof size
Up to this point, we haven't it really said much in our proofs about what the model is doing. All the mechanistic insight has been 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 sequence; everything else is noise".
Here, finally, our proof-size constraints are tight enough that will see something that we could plausibly call "how the model pays attention to the maximum token more than anything else", i.c., (if we squint bit) "the model pays more attention 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, E,) II, A, where Eg denote the matrix errors, we can use the following trick: Lemma 20 (Max Row-Diff (vector-matrix version)). For a now vector a and matrix B,
max ((aB); (aB),) Σ la max (Ba, Bx.a)
is
Moreover, for a collection of n row vectors Ar, if the shape of B is m X p, the right hand side can be computed for all in time (nm + mp).
Proof.
max(aB); (aB)} 4.5
= mmax i.
ax (Bx,i - BB.s)
max ax (Bu,s - Bk.s)
maxi. (Br.i BR.j)
if 2 0
Σ
Ok
mini, (Bu, BR.))
if ak
k
maxxi (Ba, - BR.s)
if aa 2 0
Σ ak
maxis (Br,i - Bu.s))
if 0hk < 0
k
Σ
lax max (BR,i - Bx.s)
The asymptotic complexity of computing the result follows from caching the computation of max,, (Bx.; Bx,;)) for each is independently of r, as the computation does not depend on A,, 0
Theorem 21 (Max Row-Diff). For matrices A and B,
max ((AB),,, (AB)r,) S max Σ |r.
max (BR, Bx.))
1.
Proof. By taking the max of Lemma 20 over rows 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 matrices B, of shapes Tp Gp-
(
( [1 Bp
I[
Bp
Σ
Joko
(B
max ((Bn) -( (.)
(.j
Moreover, for a collection of row vectors Ans the right hand side can be computed for all a in time Olgro E, Cp)))
47

---

Proof. We proceed by induction on n.
For 1, the statement is identical to Lemma 20.
Suppose the theorem holds for all positive n 8; we show the theorem holds for + We reassociate the matrix multiplication
max f
Bp
Bp
max tff
(aBi)
Bp
I Bp
p=2
Using the induction hypothesis gives
ako ko
Σ (B2)
(By B Σ
max ((Bst))k44i (But1) in
k
The triangle inequality gives
akk (Bi)ko
Σ 1(22)a
(Bs))
max ((B (Bs+ (8s+1)
kg
kg
Kett
and algebra gives
|ako
Σ (Bo)aanket
max ((B41))as41 - (Botl)kers)
kp
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 + matrices Ao An
max r.k
[I
Ap
II Ap
S
max
Σ
[AAo)rtal ...
max ((An) (An)
r.a
Au
kn
Proof. By taking the max of Lemma 22 over rows r of Ag.
Note that Theorem 21 is compatible with the mean+diff trick of Appendix G.I.
Theorem 24 (Combined Mean+Diff and Max Row-Diff), For matrices A and B, and any column- wise summary vector H₃ of A (for example we may take H E, Ar,k)
max ((AB) - max rAJ
Hy (Br, - BR.s) +max Σ
AAr
H& max (Bu, - Bx.5)
k
Proof.
max ((AB),, (AB)r.)
maxx Mj.
Σ
Ar, (Bu.i - BBJ)
max
Σ
(Hy + (Ar. H2)) (Br, Bk.))
r.k.J
Σ Hi (BR.S BR.j) + max Σ (Ar,k H&) (Bu, Bk.s)
max
48

---

max
Hy (Bu. Brid))
max Σ
max f.
(Ar,
(4) Bk.r))
max 1
HI (BR,i BRs.))
max Σ
A
HA max (BR,s Bk.s)
13
Theorem 25 (Combined Mean+Diff and Vector-Matris Recursive Max Row-Diff). For a row vector vector of summaries h corresponding to a (for example, if is a row of matrix, h might be the average of the rows), a sequence of n matrices Bp of shapes Tp x Cp. and a corresponding sequence of column-wise summary vectors hp of Bp (for example we may take (ht)): == Ex, (Bp)r.)).
max W
II Bp p
Bp
max 1
Σ Ap
hiko
Σ
(B (Bn-1)Rn-1a 1)knn ((Bn) (Bn) (Bu)kn.))
An
((mar
+
hkaa hokal... 1..
Ehha-i)an (BBukknni - (Bn (n)nkn)))
1.j
ko
Am
(hn- 1nn
max ((Bn)knd (Bo))ion)
i.3
Moreover, for a collection of q row vectors Ags the right hand side can be computed for all a in time Olgro Ep TpCp)-
Proof sketch. Apply the triangle inequality recursively, fusing the proof's of Lemmas 22, and 24. 0
Theorem 26 (Combined Mean+Diff and Recursive Max Row-Diff), For a sequence of n matrices An, and corresponding column-wise summary vectors ho. .... ha-1 of Au. .... An-1
Au.
[I Ap p
II Ap
max r,i.J
rJ
Σ (ho)ko ko
Σ (An-1) ((An)mnd (An (dnnsn))))
< max r,K,]
Am
(mngg
Σ kg
(Ao)ks (ho))al...
max 1.
Σ
(hn- (hn-1)kn ((An)kard (An (An))mn))
kn
(An- - (Im--1)k, max ((An)sms - (An) in
An
Moreover, if the matrices Ap have shapes Tp X Cp. the right hand side can be computed in time OEE, TpCp).
Proof. By taking the max of Theorem 25 over rows of Ag.
G.2.3 Exploring rank one approximation via SVD
Let us first look at
EQKE E.QKTET.
From Figure 9a, we see that there not much variation along long query token direction. We can confirm this by performing singular value decomposition (SVD) on EQKE, in Figure 14.
49

---

Query Sale SVD
Singalar Values
Key Stde SVD
20
10.00
15
00.00
230
20
10
5000
Key Tokes
40.0
Quey
5
90
2000
50-
10
16
1000
Singsler Indles
Singslor Indes
Singslar Irales
Figure 14: SVD of EQKE for seed 123, with principal component vectors scaled by the square root of the corresponding singular value. This scaling allows to visually that there is not much going beyond the first singular component. Numerically: first singular value is just over 7440, while the second singular value is just under 15.
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.
Call the first singular component of EQKE the "query direction" dg and the "size direction" dy 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 full-rank error term that we approximate bounds for.
G.2.4 The simple SVD decomposition of QK
In time dldde) ?) perform SVD each of the four component matrices Eg. Q, K, E and perform low-rank SVD the matrix product E,QKK" E E.
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 bound on the difference between the two elements the same row of the error term.
That is, can decompose E into part parallel to dy and part orthogonal to dg. say E, F.g Fig. and similarly E = Es + Et. Note that E, and Eg are both rank one, and hence can be multiplied with other matrices of shape dmodel X in time (ddmodela) rather than time mlddoc) Hence we can define EQKE_er err, (subscript one for "rank one") and decompose EQKE as
EQKE E2 - EQKE_erp
Define for any vector
Aiju Us U-
so that get
(E (Ex) ) + min A. (EQKE CIT Duueryy ig
query
(E,(E,KK (EE)) Equary + max A., (EQKE erT ) )quary
Then we may use any method we please to pessimize A. (EQKE, err, Mypuryy quickly. For example, since for any matrix M we have 01 (M) = sup, || || | considering vectors with one 1, one
"e might be tempted to keep analyzing the SVD, and notice that the query direction is mostly uniform, while the key direction is monotonic (nearly linear, even). But the proof complexity doesn't demand this level of analysis, yet, and can't expect that any automated compact proof discovery system will give it us.
50

---

Mix
Figure 15: The error term EQKE_ern for seed 123.
-1, and zero elsewhere, the maximum difference between elements row upper bounded by (M))(M):
B.,KE _err Equery
S (EQKE_err)
(11)
G.2.5
The complicated SVD decomposition of QK
While the "most mechanistic" interpretation would proceed with the analysis in terms of Fgg and Eg, 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 to make our error bounds as tight possible.
To this end, we perform SVD on Eg Et. Q, and K and peel off the first singular components as to get the decomposition
E Fg + Fq. + F222
E Eş Ex.2 Ev.2
Q = Qu Q'
K = Ko + K'
Then EQKE, a product of these four matrices, can be expressed as a sum of 2 = 35 rank one products and one high-rank emrrr term. We can compute the sum of the rank one products in time (ddvccab") 2) and express EQKE as, say, EQKE, E (E ( LT Call the second term EQKE err (Figure 15). We must now bound for each and m the quantity maxism- EQKE errlg, i] EQKE err(q, m).
How big this?
Even if we relax to max.j EQKE LicK__errgg cer(c, s] BQKE_err(q Jl. the maximum such value across all 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), 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 (M) sup, ||Ma|]| ||=|]| considering vectors with one one 1, and zero elsewhere, the maximum difference between elements row upper bounded by /2a, (M). The largest singular value of EQKE_ (Figure 16) is just under 7.6 (8.4 + 2.0 across all seeds), giving row-diff bound of about 10.7 (11.8 + 2.8 across all seeds), which is large but not unusably so.
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 v2 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 about 33.1 2.9, which when multiplied by /2 is about 46.8 + 4.2.) This works because (AB) < a (B), but note that to use this technique. This bound is still usable, but pretty big.
do factored SVD without needing
Can we use Frobenius? Note that using anything close to this method to drop below dooea might seem infeasible (it'll eventually turn out not 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
51

---

Stagda: Valtee DOKE_as
20
33
20
Figure 16: SVD of EQKE__rr for seed 123.
Figure 17: SVD of the four component matrices of EQKE_err seed 123. Matrices look like noise.
it takes to compute the full SVD is the Frobenius norm, which is defined as tr(MM²), can be computed vooce time, and equal to the square root of the sum of the squares of the singular values. While the Frobenius norm of EQKE_er is only about 12 (giving 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 product of 1932 and a bound of 2732(!). (Across all seeds, the Frobenius 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 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.
of
However, we can get much better bound on the max row-diff of EQKE_ _err without having to multiply out all four matrices. can use an approach vaguely similar to the mean+diff trick, as follows.
If we want compute the row-diff of product of matrices AB, can compute by Theorem 21
max ((AB),,, (AB)r.) max Σ |Ar max (BR, BB.)
(12)
r.K.j
1.j
or by combining this approximation with Theorem 18 via Theorem 24 may compute
max ((AB)r.) (AB)r.))
< max Σ E, (Bx. Bus))
+ max Σ IA,. E,Ar, max (Bu, Br.5)
taking whichever bound is better.
The first gives bound of 7.94 on the maximum row-diff, which is better than 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 bound of 5.67. Combining 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 over dimension zero has norm 0.025 (0.030 + 0.012 across all seeds).
Although this emoor bound is no longer the leading asymptotic bottleneck, we can peek ahead to what we get we want to linear in parameter count. In this case, we can apply the recursive version of Equation 12 via Theorem 23, giving bound of 97. .66 on the maximum row-diff.
52

---

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 imput tokens with this direction small once we factor out the first two singular components", and this is sufficient to drive low error term overall we factor out the leading singular components in other places. We don't mechanistically understand how to combine the E (without multiplying them out) in a way that allows getting a good bound, though, which corresponds to our inability to drop below dvocabdmoder here.
If 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 bound of 0.6768 (0.632+-0.00 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 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. I),, the decomposition of EQKE, and the handling of EQKE_err and EU.
53

---

Algorithm Counting Correct Sequences Subcubic Time
1: function MODEL-BEHHAVORR-RELL.
GAPMM, 4max lquery: G, 9, 9") is False = result is False
Ensure:
Require: 0 S 9 S S max
Require: if then Equery mmax
2:
JUMMARIZEBU.Apprry (F (fquery) )--) 3: skip-score max,- {BU (Qquery) min {EU
scip-score,-
Cache by t*
( (fquery)e
Cache by lquery
4: UF - EV)U(()
5:
wy PVOU(i)
6:
A & max; Wi,e" Wifmmax
Cache by mmax t Cache by mma D Cache by t
7:
Atmax, f- max A
8:
Avt f max+- U1,1" min Ut,t" Av
9:
Ammax f maXosIS!mas
Cache by tmax 9" Cache by bmnax Cache by 4max
A f- itaxxxi" to Utmane Amax
11:
A Autmar t- m&xt of thaa Aereuus
max
12:
if cC- O then
13:
EG- &- Q (tmax)e + UUmmar, Ummaxi + Wmaax.!" return max, #lmax (Le- E, (lmax)
14:
15:
end
16:
De 11 EQKP( (1 query. Retx
Cache by lquery Cache by fquery: i
17:
Do.-1 +- SORT EQKP (qquery. 1))
18:
b1.--1 REVERSE (o,--1)
EQKE F.QKF__RR err DECOMPOSE(RQKF))
Require: EQKE (Lquory, 1) -EOKE" (Iquery, mmax) EQKE BerrEE_rr___aae
quary < (BKKE(Aquere ()
EQKE(tquery; mmax) S EQKE (fquery: 1) - EQKE (fquery- max) + E-KE_ERT err, e
at + EQKE( (tquery,t)
Cache by Iquery,
21:
amin, + minoge" St. Og"
Cache by lquery L, compute in amortized (dvocab 2)
22:
amax,1 maxog O4"
Cache by tquery, t, compute in amortized O Oddycabb 2)
23:
A0max 04 0axx Amint
EQKE__rr,, eerreeuery
Cache by lquery, bmaxx C
24:
AGuin < Ommax Omax,tmas y EQKE_eer, Equery
Cache by fquery. tmaxx C
25:
idx-set < (O..... Netx c - 1} if mmax * 'qqeery else (O...., Tlexx C 2, Pletx - 1}
attn autn-weights-und unsgalls- bo,: (Aamin if idx-set else 0)
atn--weihhs-unnsaaedd; b1.; + (Aamax if idx-set else 0) Cache by Equery fmax i,
attn- atin-weights SOFTMAX(attn, www.-aa-nntt-wehhtttttttt unnhchlls-ussslld,,y /Vd)
Cache by tquery- fmax i,
29:
attn-weights + SOFTMAX attn wwwi-www- unscaled, /
> Cache by fquery, fmax i, c
30:
attn-maxg iEidx-set weightss
31:
attn-max + iEidx-set attn- weights,,
atin-max attn-maxo if A Autfman Aumax else atin-max;
33:
attn-max + SUMMARIZE Equery (attn-max)
D> Cache by tmaxe
34:
attn-max' + attn-max attn-max
35:
summary,- Awmax,t- + skip-score,- attn-max A Amae
(1 attn-max A Aumax
Cache by tmax, t*
return skip-score attn-max's
Aptmos + (- Amax + max summary,, max
37: end function
54

---

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 line up with the legends of Figures 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 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 the fact that attention 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, approximately 1.2.
Exponentiating. the sost-softmax attention paid to the max is typically about 3x larger than to the token below the max; here the logit difference between the 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 it is around 0.3.
"*"Typically" here means about 96 % of the time.
1000
20.0
17.5
CoUnt x 9 siqueccce walh
10.0
7.5
5.0
2.5
30
35
10
45
50
55
60
65
70
logit - diag
logit diff
(a) The attention computation weighted by the number of sequences with the particular max. The computation
(b) Histogram of the maximum difference between
two logits contributed by single row of EVOU. The computation is, for each max, EVOU,; min, EVOU.,, p + a: 43.4 + 9.0; range: + 20
is max; jz; EVOU,, EVOU, μ a: 9.99 +
2.1; range: (-0 + 120) x 10-1
Figure 18: Plots of the difference in logit for the attention computation, EVOU EVOU for seed 123.
55

---

-10"
75
atteniion differeccc
25-
0.6
X 3F sequeices
0.4
-50
0.2
-75-
0.0
tokon 1pap
astention differenace
gep
(a) (EQKE,
(b) (EQKE; - E (/(v(( j)), weighted by sequence count. μr
1.22 0.13
Figure 19: Plots of attention difference vs. token gap, for EQKE = E, "OK"E" E for seed 123, The difference in attention between tokens is approximately linear the gap between the tokens.
10a
counE soquecces
6
Figure 20: Histogram of the minimum gap between the max token and the largest non-max token, for the seed 123.
So for sequences where the largest of the non-max EVOU matters lot; but when the max is separated from the largest non-max by modest gap, the structure of the non-max EVOU does not matter much.
and the max close together, the particular structure
The upshot that to handle most sequences, we need only ask an oracle for the minimum gap g > between the max token faax and largest non-max tokens / * mmax such that the model outputs the correct answer for all sequences where the non-max, non-query tokens have value at most fmax
While computing this gap may be expensive (and indeed the naïve computation of the oracle takes should be very easy to optimize), we don't 't have to
longer than the brute-force proof proof-though pay the cost of computing the gap in the size of the proof, only the cost of storing the gap table (Olvocab (Retx)) ) and of verifying the gap. Empirically, gaps are typically 1-5, as seen Figure 20.
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, could build in time (O(dvocab 2)) table indexed pairs (t, fmax) of the maximum token and non-maximum token: the table would store pessimal logit contributions from maximum output tokens < the (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 and t + in time (O(vvocab 2)).
Next sort the queries by overlap with the query direction.
56

---

200
120
[00000 0.22)
M(-000.0.14
150
100
640
40
50
-044 -0.
0.6
0.S
zuGZEN elessoon value
matii elemneat value
K
NI 0.00, 0.14)
200
N/00 00222
100
175
125
100
40
30
20
25
motris clement value
matris elemenn valae
Figure 21: The distribution of entries of the four residual matrices (after removing two principal components from E, and E and one principal component from Q and K). Distributions look pretty close to normal. Plots are for the seed 123.
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 /maxs how many of the query tokens 'query fall strictly below the max max (and whether or not the model succeeds when max tquery). This will tell how many query tokens we can count for
given maximum token.
Compute table indexed on pairs of # of queries handled and input tokens which stores the smallest difference in more attention paid to t 1 than t (Odvocab ?))
Compute table indexed on pairs maxx t storing upper bound on amount more attention paid mn-maxxmum tokens than to 'max by Oracle-permitted query tokens (the Oracle is indexed only on mmax (dvoocab)))
For each # queries permitted: compute for each maxx t, if the mon-maximum token t contributes little enough 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 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.
57

---

Necoslzeed Accutaty Barnl
10
bruie force (are: 0,9992 0.00155
0.8
cubue trel soc: 0.9339 0.0000))
0.6
suhaabic 0.331
0.4
attettoon-hoeaeelatt (rei acc:
disos-quuatticc ddirect-quaelt Sinect-quddaticc (nd ace: 0.628 0.060)
0.655
0.2
attention dowattfaadaa" aleettion
quidratic
(0032)
2
allestionquudaatic diocctquaalaticc ace: 0.086)
FLOPs Vediy Proo (approsimane)
Nemalied Accuracy Beand
104
0.8
cubic (rel boc: 0.9539 0.0080)
0.6
suhcabic (ne ace: 0.811 0.013)
atteitiotodoeeat//nnt " acc: a 0.014))
divet-quddtiic (rel sec: 0.653 0.0600
0.2
atcationoo dinest-qudratic
0.060)
0.0
attentionquudraaic trall acc: 0 390 (.032)
223
2
237
245
attentionquudeaii: direct-qualratic (nal aoc: 0.286 0.086)
FL.PPs to Verify Proof (ongroximaee)
Figure 22: Recreations of Figure 3 for case of viewing of the legend. Top is striet recreation; hottom includes points not the Pareto frontier.
If we replace the entries of E E 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 4cr).
H.2 How various combinations of tricks perform
Recall Figures 3 and on page 8 and on page from Section 5, recapitulated here without captions for convenience as Figures 22, and 23.
We describe what each subeubic 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 10 strategies.
The brute force and cubic proofs (1 + 1) were fully covered in Appendices D, and
There are 5 options for handling EU:
direct-quadratic refers to handling EU in time dmode) with either the row-diff trick (Appendix G.2 2.2.2) or the max row-diff trick fused with mean+diiff or some other summary statistic (Theorem 24)27
When direct is not mentioned, this indicates that we handle EU in time (dvocahh dmmode)) by first multiplying out E and then either taking the maximum row-diff in each row or by taking the maximum row-diff across all 29 The latter is included purely for comparison's 's sake, and never gives tighter bound than the former.
There are 10 options for handling the high-rank attention emror term EQKE_enn
TThis shows up in the bias towards having larger values (both positive and negative) in the lower-right corner of plot, indicating that errors are larger for larger query and key values. We hypothesize that this due to the distribution of data: larger values are more likely to have more space between the maximum and next-most- maximum next-most-maximm token, so bit of noise matters less for larger maxes than for smaller ones.
Thii strategy labeled "max "max__iff" diff"" in the Python source code.
Thhhese strategies are labeled tmean_query tuuerrtmxxddiff" and ssvd_query+maxx diff" in the Python source
code
3mmxx_ "'max_diff___ "iff_fxacct" _xxatt"
29., globbbll max diff exact"
58

---

Normalized Accuracy Bound
1.0
max max-diff- exact
0.8
diff- subroodctt
max-diff-enn ciax-diff-uhprrdeeet
0.6
mean mean-max-die max-diff
0.4
mean-max diff- subpoouuct-occusiie
0.2
max diff subproduct recunsive recursive max diff subproduct recursive
0.0
400
600
800
1000
EPQKRR Singulor Ratio: /i//
Figure Recreation of Figure for ease viewing the legend.
attention-quadratic refers to handling the high-rank attention error term EQKE from Ap- pendix G.2.5 in time modee)) either with the recursive max row-diff trick (Theorem 23).00 or with the recursive max row-diff trick fused with the mean+diff trick either just on the query side or throughout (Thcorem 26).
attention-drocaa indicates that we use one of the various dodve strategies for handling EQKE__rrs from Appendix G.2.4 or EQKE_ _rrr from Appendix G.2.5. These include using v201 1oopuumpteed via low-rank SVD as the bound (Equation 11)a5, considering all ways of multiplying out subset of the matrices and taking the maximum row-diff of the resulting pair of matrices* (Theorem 21), or fusing the max row-diff trick with the mean+diff trick (Theorem 24).
When attention is not mentioned, this indicates that we handle the attention error term in time (ddvccab 2 mmodet) either by taking the per-row maximum row-diff by using the full rank EQKE matrix and taking the per-row maximum row diff??.
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 29; this makes up the final factor
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(dvocal 2d ddmode)) corresponds to taking the full rank EQKE_ term and taking the maximum row-diff each row.
seaan--a--diffsuuprree (O(dvoea 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- max-diff-ssuprrdduc and considering all ways associating the multiplication of EQKE err.
corresponds to using the max row-diff trick (Theorem 21)
max-diff (ddvccabdmmotee)) 2)) corresponds to using the max row-diff trick (Theorem 21) on the factored SVD of EQKE_erT
mean-max-diff (ddvocabdmmode 2)) corresponds to fusing the max row-diff trick with the mean+diff' trick (Theorem 24) and applying it the factored SVD of EQKE_err,
"0max_diff diff ctcff__sbbroduuutreeeetee
iitraannrecurrsiveeea....... htt.://wwwweee......ooooooooreee 33=gvd"
$4"max_ diff"" for EQKE_ere "max_difff_ wwdddiffssuprrduucttt for EQKE_en
$5*mean=max_diff for EQKE_erT, "mean+max_diff dishauupoccuu for EQKE_er
"""ma__dff_eee diff_exafcec exact"
37=exact, DoxxacttEQKK+maa_.._. bitiaaxxdff_eexxtttt_
38-mean "ueeenn_uerr+diff""
39., trrop_average_uueryypr.. suwery-per_outut_tre
59

---

svd dmodee corresponds to using 2011-computd via low-rank SVD as the bound (Equation 11).
modd))) corresponds to handling the high-rank attention error term EQKF_ 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 be zero).
mean+max-diff hwtp:eaaaaddffuupddunrrrerriiiie
saa.dddfffubpprduuc---rrrss (oodd)) 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-ddde dis-ddff-upoouut---eueeeee high-rank attention enoor term EQKE_ _er from Appendix G.2.5 with the recursive max row-diff trick recursively fused with the mean+diff trick (Theorem 26).
(model) corresponds to handling the
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 understand- ing to allow more compact proofs have tighter bounds. We can also look at the reverse problem: we can take 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 quadratic EU proofs, when max_diff performs worse than
the model. For example, mean_ queryyuxrtttxaaa__m tirffffaaiifffffff diff and wvd_query+tn twma_miiffi.aaaaaaaaaaaaaaaaimie behavioral component shared across query tokens that U is that good filtering out. Sim- ilarly, when, e.g., rean_recursive wiifffffaaaffaa.eaaeeeeeeeaeeeee orw.oussssieeiaa.aaaauouurrrrr.e
this indicates that E has relatively strong
performs better than max_dife faalfffsubrrduutt..m seprrdductrreurriiiie ve, this indicates that even after removing the first one or two principle components from E,g. Q, K, and E, there is still enough common structure that it is worth factoring out the mean behavior.
60

---

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]
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.
Justification:
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: See Section
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
We provide intuitions for proof constructions in the main body of the paper. 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.
Justification:
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We directly link to codebase with the specific implementation of our study.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, 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. 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
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
experiment models trained with 151 different hyperparameters. For most reported computations, we provide statistical significance information. We do not need to perform explicit t-tests such since it is not relevant our setup.
We
8. Experiments Compute Resources
61

---

For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Appendix B.1.
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https:: /ttt::://nnurrps... http://www......ppppl/l///aaaaaa hc/ppub.i.//hiiicssueeeeeleeeee
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 deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration 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
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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.
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 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 research with human subjects.
62

---

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 equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
63


Reviews:
[{"review_id": "czVCx0YCDI", "reviewer": "Reviewer_KZFC", "summary": "This 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.", "questions": "- 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\"", "limitations": "- 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)", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- 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.", "weaknesses": "- the setting is quite toy, though this is probably very hard to overcome for such an approach with human labor.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your kind review!\n\n> how can this method be carried over to more realistic tasks?\n\nIt would be relatively straightforward to carry over the methodology to more realistic tasks if the following conditions were met:\n\n1. A sufficiently detailed mechanistic understanding of the components of the model not being brute-force evaluated.\n2. A model that is interpretable enough that the (not-understood) remainder of the model can be treated as independent error terms without blowing up the bound.\n\nIf these conditions were met, then the interpretation would broadly proceed as follows:\nEvery expression $x$ in the model with interpretation $y$ can be re-expressed as $x = y + (x - y)$.\nThe error terms $(x - y)$ must then be given worst-case bounds.\n\nAs far as we can tell, most realistic models do not meet these conditions.\nConsequently, in the project, we devote significant effort to getting reasonable bounds with minimal understanding (and it is surprising that, e.g., we can get a non-vacuous bound on the model that, modulo our lack of understanding of low-rank approximations to the identity matrix, is linear in the parameter count of the model!).\nIn follow-up work, we are making progress on applying the approach to models with MLPs, to models with with multiple attention heads, to models with layer norm, and to two-layer models. We see three big obstacles. We believe all the obstacles are important for the field of mech interp to address (with the possible exception of the third one):\n\n1. Unstructured approximation error terms (“noise”) accumulate quickly, necessitating either models with significantly less noise (more interpretable models) or more sophisticated error term analysis.\nIn follow-up work, we have preliminary results suggesting that fine-tuning models can make them sufficiently more interpretable as to bypass this problem.\n2. Path decomposition is exponential in the number of layers, necessitating adequate mechanistic understanding of how to simultaneously argue that multiple paths contribute negligibly.\nExisting approaches to solving instances of this problem in finding mechanistic interpretations -- such as presuming independence of irrelevant paths and evaluating them all simultaneously, or using gradient-based search methods that leave the path decomposition implicit -- unfortunately do not yield valid proof techniques.\nThis problem is especially pressing for us, as it interacts with the central problem of compounding error approximations (1).\nHowever, we are currently investigating the possibility of matching prefixes of paths to allow mechanistic understanding to be used to simultaneously argue that multiple paths contribute negligibly.\n3. Worst-case error bounds are both significantly looser by default and significantly more expensive to compute than average-case error bounds.\nWe would be excited to see a formalization and analysis of compact guarantees of average-case performance!  (Unfortunately, we are not aware of any adequate formalism of average-case performance that is independent of training method.)\n\n> 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\nThat is correct.\nHowever, we would like to highlight that it should be straightforward (if labor-intensive) to generalize our work to other toy models where the model computes the answer by paying attention to a single token or position, i.e. general element retrieval, example tasks include: first element, last element, binding adjective to key phrase, etc.\nThis is because our theorems in the appendices prove in full generality that any computation performed by a single attention head, when only one of the position and the token value matters, is extremized by pure sequences.\nSo the worst-case behavior of the computation at a single token can always be found by setting all other tokens to be identical, regardless of what the task is.\nFurthermore, generalizing our convexity arguments to multiple attention heads *whose error terms can be treated independently without degrading the bound too much* is also straightforward.\n\nGeneralizing the theory to handle coupled attention heads or more complex uses of attention will require additional theoretical work.\n\nWe address the high-level version of this concern in the top-level rebuttal where we talk about the general applicability of our approach to projects on larger input distributions, different tasks, more complex architectures, and larger models.\n\n\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)\n\nIn follow-up work, we conduct analysis of a modular addition model with an MLP where we achieve a bound of roughly 80%--90% on a sub-brute-force proof.\nHowever, it is indeed the case that integrating all the mechanistic understanding we have of that obliterates the accuracy bound.\nUltimately, we hope that fine-tuning will allow bypassing this issue.\n\nSeparately, we want to note that 25% would be *average case* chance levels for *argmax* (predicting the position of the maximum).\nAverage-case chance levels on max are actually $1/d_{\\mathrm{vocab}} = 1/64 = 1.56\\%$ for this task, since we are predicting tokens not positions.\nFurthermore, it's not entirely clear that the baseline for proofs should be average-case bounds rather than worst-case bounds.\nIndeed we typically see bound performance drop off to floating point error, not average-case chance, when error terms accumulate too quickly.\n\n> you mention: [...] I believe a standard reference is [...]\n\nThank you!\n\nPlease don't hesitate to ask any further questions!\nWe'd love the chance to highlight strengths of our submission that might encourage you to further increase your score!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rM6tQ3UcFD", "reviewer": "Reviewer_exsQ", "summary": "The 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.", "questions": "I request the authors to kindly address the questions in the weaknesses section.", "limitations": "Yes, the authors have addressed the limitations.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1) 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.", "weaknesses": "1) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review!\n\n> The tightness of bounds seem to degrade very quickly as additional insights from mechanistic interpretability are used.\nThus 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\nWe agree that the degrading tightness of bounds is an issue.\nHowever, note that we evaluate proofs on the additional criteria of compactness of proof length, and using information about model internals is how we obtain shorter proofs than brute force.\nThus, some degree of mechanistic interpretation is necessary for finding compact proofs in future work.\nWe address this question in more detail in the top-level rebuttal.\n\nOne claim that we will add more clearly to the paper is that proof length can be used as an evaluation metric on the quality of mechanistic interpretation (see Section 5).\nThe interpretations in our different proofs all appear to be sound and use the same high-level abstractions.\nHowever, by pushing them to rigorous formal proofs we get to see the various degrees of freedom that remain in the different interpretations.\nTherefore, we expect that future work will need mechanistic interpretations to get short proofs, and can use the gains in proof reduction as a way of distinguishing what interpretation techniques to use.\n\n\n> The author analyze a very simple task and techniques used by authors to decrease computational complexity are specifically tailored for this task.\nIt therefore remains unclear how insights from this work can be applied to any other tasks.\nIt 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\nVerification projects are quite intensive in human-labor, and our goal with this project was to prototype verification for transformer models, and focus on studying the core challenges to its viability.\nThus, we limited ourselves in the number of different tasks we worked on, and instead focused on deriving multiple proofs that varied the two core desiderata: compactness and bound tightness.\n\nHowever, we want to highlight that it should be straightforward (if labor-intensive) to generalize our work to other toy models where the model computes the answer by paying attention to a single token or position, i.e. general element retrieval, example tasks include: first element, last element, binding adjective to key phrase, etc.\nThis is because our theorems in the appendices prove in full generality that any computation performed by a single attention head, when only one of the position and the token value matters, is extremized by pure sequences.\nThis means that the worst-case behavior of the computation at a single token can always be found by setting all other tokens to be identical, regardless of what the task is.\nFurthermore, generalizing our convexity arguments to multiple attention heads *whose error terms can be treated independently without degrading the bound too much* is also straightforward.\n\nWe address the high-level version of this question in the top-level rebuttal where we talk about the general applicability of our approach to projects on larger input distributions, different tasks, more complex architectures, and larger models.\n\n\n> 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.\nCurrently this seems based on speculation and needs more detailed investigation.\nTherefore it would be great if the authors could try to artificially inject noise (of varying magnitude) in these matrices and see its effect.\nElse, the authors should not claim this being the major cause for not so tight bounds obtained on using insights from mechanistic interpretability.\n\nThank you for pointing out our confusing use of terminology! By “noise” we meant error terms in our approximations of the model that we use to generate more compact proofs, not error terms that the model learned.\nSo, we are not referring to model weights that we can ablate to improve performance bounds of our proofs.\nIn the improved version of Table 1 (included in the accompanying pdf, and which will replace Table 1 in the paper), consider the difference in bound between “Low-rank QK” (0.806) and “Quadratic QK” (0.407).\nThe only change between these results is in the strategy for bounding the error that comes from approximating QK as low rank.\nWe provide a more detailed clarification in the top-level rebuttal.\n\nPlease don't hesitate to ask any further questions! We would love the chance clarify our approach, and provide further details that might encourage you to increase your score.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "r3fG8vfzcO", "reviewer": "Reviewer_vjrV", "summary": "The 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.", "questions": "- 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.", "limitations": "I 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the extremely detailed review!\nWe also appreciate the list of typographic errors, which we’ve corrected in a revised draft of the paper.\n\nWe’d like to start by defining what we mean by “compact proof of global robustness” by clarifying our definition on lines 48-69 of Section 2.\n\n* By “compact proof”, we’re referring to a proof that is short.\nSpecifically, in our paper, we measure the length of proofs by considering the length of a trace of the computational component of the proof (as described in lines 56-61).\nWe agree that our use of the concept “compactness” could be made more explicit in the paper; we’ll add additional clarification.\nAs for why we want compact proofs: proofs that are too long are both infeasible to find and also infeasible to verify – hence our interest in using mechanistic understanding to generate shorter/more compact proofs.\n* By “global robustness”, we mean that the model performs correctly in expectation over an input distribution $\\mathcal D$ (lines 48-52).\nThis is in contrast to much of the prior work, which focuses on establishing correctness of model behavior for points close to a particular input $x$ (which we refer to as local robustness).\nWe agree that the use of the word “global robustness” in the introduction is confusing, because we use “global performance” to refer to the desired form of our verification target in the title, on line 49 in Section 2, and on line 294 and 301 in Section 6.\nWe’ll edit the paper to be more clear and consistent.\n\n> Based on Eq. 1, it does not seem that these guarantees fend against out-of-distribution or adversarial examples?\n\nIn addition, we’d like to clarify the flexibility inherent in equation 1.\nNormally, by OOD or adversarial inputs, we suppose that there is a distribution $\\mathcal D_{\\textrm{in}}$ that’s used for training and (in-distribution) validation, and another distribution $\\mathcal D’$ that is the deployment distribution or generated by an adversary.\nIf we had knowledge of $\\mathcal D’$, we could compute the expected performance from inputs sampled from $\\mathcal D’$.\nEven if we don’t have exact knowledge of $\\mathcal D’$, we can still define a very broad distribution $\\mathcal D$ that covers possible $\\mathcal D’$s.\n\nIn our work, $\\mathcal D$ is the distribution of all $64^4$ possible valid input sequences.\nIn addition, as our proofs partition $\\mathcal 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.\n\nWe agree that we don’t adequately discuss the flexibility inherent in our definition; we’ll add clarification to sec 2 as a well as a section to the appendix discussing this in more detail.\n\n----\n\nWe believe that much of the motivation for the work (and indeed the answers to your questions) follow from the goal of producing compact proofs.\n\n> Why did the authors choose to investigate a task like max-of-$k$ instead of modular addition (or IOI, colored-objects, etc.) \\[...\\] ?\n\nWe picked a very simple task because formal verification is quite challenging, even for Max-of-$k$, and believe that Max-of-$k$ captures some of the core difficulties that we might expect to see.\nBut to address some of your specific examples:\n\n1. **Modular addition:** In Nanda et al 2023, the authors don’t provide a mechanistic understanding of *how* the MLP layer performs multiplication – instead, they check this via brute force.\nThis lack of understanding prevents us from compactifying our proof of the model’s correctness.\n\n   In follow-up work, we extend our methodology to the modular addition task.\n   However, compactifying the proof required us to reverse engineer and provide a mechanistic description of how the MLP computes its output – which is, as far as we know, not previously known in the relevant literature.\n\n1. **IOI/colored-objects/other small circuits in LMs:** One of the concerns with using these small circuits is that they tend to only be valid on a very restricted subdistribution.\nFor example, the original IOI paper considers a small set of sentences with names ordered ABBA, and the circuit metrics become invalid on as small a distributional shift as ordering the names ABAB[1].\nIf we chose to study these circuits, we’d run into a dilemma: either we restrict our distribution to be small enough to check feasibly with brute force (as the authors tend to do), or we’d have to perform additional interp that advances SOTA.\nAs our focus is on the formal verification side, and not on the novel interpretability side, we chose instead to pick a task where the mech interp is relatively straightforward.\n\n> 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\nWe agree that on models where we can evaluate the behavior on *all* possible input-output pairs, we can indeed guarantee performance via the brute-force proof (lines 151-4).\nWe study such a case because it allows us to have “ground truth” on the performance of the model.\n\nIn general, however, mech interp is often studied with samples from a distribution – for example, in IOI, the mean ablation is computed with respect to a small handful of other sequences, and in modular addition the authors use sampling to evaluate the 5-digit addition, repeated subsequence, and skip trigram tasks.\nIn cases where the bad behavior is heavy tailed, sampling may not provide a viable estimate of performance – our hope is that understanding derived from studying the model with samples will allow us to construct formal guarantees that hold in the worst case.\nWe prototype this problem by providing formal guarantees on a max-of-20 transformer, where brute forcing the input distribution is infeasible and our understanding is derived via samples.\n\n[1]: Miller et al. Transformer Circuit Faithfulness Metrics are not Robust.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dWcu6xugAI", "reviewer": "Reviewer_ZhNY", "summary": "This 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.", "questions": "1. 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.", "limitations": "The authors have properly discussed the limitations of their approach.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- 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.", "weaknesses": "The 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.”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback!\nWe have also rectified the typographical issues.\n\n\n> You only compare sequence lengths of up to four because the brute-force solution is infeasible beyond that.\nCould 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\nThanks for the suggestion, we agree that this is a serious concern.\nWe have generated new results on Max-of-5, Max-of-10, and Max-of-20 (holding parameters other than n\\_ctx fixed).\nWe ran the cubic proof on each of these models, giving bounds of 94% (Max-of-5), 93% (Max-of-10), and 83% (Max-of-20) in under two minutes.\nFor comparison, the brute-force approach on Max-of-20 would take roughly $64^{20} \\cdot 20 \\cdot 64 \\cdot 32 \\approx 10^{40}$ FLOPs and it took about $10^{23}$ FLOPs to train GPT-3.\nWe will add a section on the scalability of proof strategies to larger models, which includes these experiments as validation that our strategies result in lower complexity proofs.\n\n\n> The results appear to be merely an interesting case study in a highly simplified setting.\nThe 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.\nThese simplifications raise doubts about the transferability of the proposed proof strategies to other settings.\nFor 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.\n\nOur goal in the project was to prototype verification of transformer models, lay out desiderata for verification (compactness of proof length for feasible cost, and tightness of bound for guarantees to be materially useful), and study what the key challenges to achieving the desiderata (faithfulness of interpretation, compounding error terms).\nObstacles that arise in prototyping are likely to be even more challenging in complex, real-world settings.\nThis approach is standard in the field of mechanistic interpretability: starting in simplified settings and pushing for extensive exploration [1]. Simplifications of the model architecture in our work are comparable to state of the art analysis of toy models [2]. And the challenges of decomposing even small models adequately are discussed in various papers in the field [3] [4].\n\nHowever, we agree that we could have taken the alternative approach of targeting the verification of a large model, which would have been an impressive result.\n\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\nWe appreciate the highlights about our contributions! We want to add that to the best of our knowledge, we are the first to (1) verify the performance of a transformer-architecture model, (2) guarantee global performance bounds on specific models, and (3) incorporate mechanistic information to obtain shorter proofs. While we work in a simplified setting, we are excited about follow-up work building upon these milestones. \n\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\nWe appreciate this highlight as well. Whereas current work relies upon human judgment, with ambitious projects pushing towards procedural arguments, by pushing mechanistic interpretability to formal proofs, we were able to find that degrees of freedom remain remain even when we have decent informal interpretations. \n\nPlease don't hesitate to ask any further questions. We would love the chance to clarify our approach and provide details that might encourage you to increase your score!\n\n[1]:  Elhage, et al., \"Toy Models of Superposition\", Transformer Circuits Thread, 2022.\n\n[2]:  Elhage, et al., \"A Mathematical Framework for Transformer Circuits\", Transformer Circuits Thread, 2021.\n\n[3]:  Stander, D., Yu, Q., Fan, H. and Biderman, S., 2023. Grokking Group Multiplication with Cosets. *arXiv preprint arXiv:2312.06581*.\n\n[4]:  Chan, et al., \"Causal Scrubbing: a method for rigorously testing interpretability hypotheses\", AI Alignment Forum, 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Thank you to the reviewers for their comments and feedback!
We’re happy that reviewers agree that our approach to verification is novel and that our approach to mech interp is rigorous.

We address the five common questions that reviewers raised.

**(1) Motivation for picking a toy setting (vjrV, KZFC, and ZhNY)**

Formal reasoning is computationally expensive; very few large software projects have ever been verified [1][2], none of them comparable to large transformer models [3][4].
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 reason more efficiently than by brute force about model behavior?
We did not adequately explain this motivation in the submission, and will add a section.

**(2) Applicability of results to other settings (exsQ, vjrV, KZFC, and ZhNY)**

1. **Larger input spaces:**
We applied our proof strategies to models trained for Max-of-10 and Max-of-20, and have attached our results.
While running the brute force proof on Max-of-20 would require $64^{20} \cdot 20 \cdot 64 \cdot 32 \approx 10^{40}$ FLOPs, which is about $10^{17} \times$ the cost of training GPT-3, our cubic proof achieves bounds of 94% (Max-of-5), 93% (Max-of-10) and 83% (Max-of-20) in under two minutes, demonstrating that proof strategies can be reused on larger input spaces (and indeed, scales better than the brute force or sampling approach do).
2. **Different tasks:**
In this work, we worked on highly optimized relaxations 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 mech interp 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.
3. **More complicated architectures:** We worked on a simple model studied in A Mathematical Framework for Transformer Circuits.
In follow-up work, we extend this approach to proving bounds on 1L transformers with ReLU MLP trained on modular addition.
4. **Larger models:** We agree that it’s an open question whether or not the mech interp 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 (i.e. whether mech interp itself will scale).
This is widely recognized in the field of mech interp, and scaling interp 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 in mech interp in general.

**(3) Why use mech interp at all, especially given that it “results” in worse bounds? (exsQ and ZhNY)**

First, we agree that ideally, we could maintain the tightness of the bound even as we decrease the length of the proof.
However, there’s some fundamental reasons to expect there to be a trade off: for example, from the compression perspective, any proof strategy that does not run the true model on all inputs needs to use a compressed approximation to the model, the input distribution, or both.
So unless the model weights are losslessly compressible, shorter proofs using lossy approximations will likely cause weaker bounds.
Consequently, the results of proofs derived via mech interp should be compared to other proofs of comparable length.
Indeed, we find better mech interp can give us tighter bounds at same proof length (figure 3), which provides evidence that better mech interp can be used to improve proofs.

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 mech interp derived bounds are worse suggests gaps in our mechanistic understanding.
As mech interp matures as a field 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.

**(4) Less faithful interpretations produce better bounds (exsQ and vjrV)**

We have corrected the plot attachment, which was a genuine mistake on our part.
The plot is more visually clear in the point it makes.

**(5) Validity of the noise hypothesis (exsQ and vjrV)**

We want to clarify that we use “noise” to refer to error terms in our approximation, which is confusing terminology that we will change.
In our work, we use approximations to model components in order to shorten proofs.
While we were able to approximate individual matrices with small error terms, when we compose the approximations, the estimate for the error term on the matrix composition is much larger than the empirical error term.
This is because lower-bounding (aka worst-case bounds) requires that we pessimize over the error terms.
Thus, we are not referring to noise in models (which we might ablate to improve our bounds or introduce to check our hypothesis).

Additionally, reviewers point out issues in our presentation.
We have incorporated feedback from this review process, and have made significant progress since the submission of the manuscript, including standardizing language and adding definitions, simplifying notation, and improving copyediting.
We address remaining issues in individual rebuttals.

[1]: Leroy. “A Formally Verified Compiler Back-end”

[2]: Klein et al. “seL4: Formal Verification of an OS Kernel”

[3]: Clarke. “State Explosion Problem in Model Checking”

[4]: Gross. “Performance Engineering of Proof-Based Software Systems at Scale”

Meta-review: The paper proposes strategies for verifying the performance of a simplified transformer model (one layer, one head, attention only) on a toy task (max of $k$ integers in a bounded range). The authors use clever tricks to reduce the complexity from brute force $d_{vocab}^{context}$ by making use of the precise computations of the transformer, and further use mechanistic interpretability techniques on the learned models to get worse lower bounds at faster computation. 

The main weakness of the paper is that the techniques are tailored to the simplified problem at hand and it seems to be very challenging to extend them efficiently to more complex models where we lack any mechanistic understanding (or can't really estimate the faithfulness of these). This seems like a severe limitation to me, though I agree with the other reviewers that the study is interesting, pretty rigorous (experimentally and theoretically), and could spark interesting future work. So I recommend accepting the paper. I encourage the authors to add a section about the applicability to other settings that they discussed in the rebuttal.
Decision: Accept (poster)