Compact Proofs of Model Performance via Mechanistic Interpretability

Authors: jason  gross
rajashree  agrawal
thomas  kwa
euan  ong
chun  hei
alex  gibson
soufiane  noubir
lawrence  chan
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 23:19:36

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.

Main paper: 
 introduction: 
 One approach to ensuring the safety and reliability of powerful AI systems is via formally verified proofs of model performance [48, 11] . If we hope to deploy formal verification on increasingly large models [24, 27] with powerful emergent capabilities [56] across more diverse and broader domains [5, 46] , we will need compact proofs of generalization bounds on specific models that certify global robustness. However, existing approaches tend to use proof strategies that suffer from bad asymptotic complexity, while verifying either generalization properties of training procedures or local robustness properties of specific models.One key challenge to verification is that neural network architectures are highly expressive [51, 58] , and models with similar training procedure and performance may still have learned significantly different weights [38, 9] . This expressivity makes it difficult to adequately compress explanations of global model behavior in ways that correspond closely enough to the model's actual mechanisms to be useful for efficient verification without being too lossy, especially when using only knowledge of the architecture or training procedure. We propose verifying model performance using understanding derived from mechanistic interpretability (Section 2) -that is, reverse engineering the specific implementation of the algorithm from the learned weights of particular models. Knowledge of the specific implementation allows us to construct less lossy simplifications of the model, and more efficiently reason about model performance over possible inputs.In this work, we provide a case study of translating mechanistic interpretations into compact proofs. We train an attention-only transformer on a Max-of-K task with 151 random seeds (Section 3), and 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-force 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, in some subcubic proofs (Section 4.3), we use all parts of the mechanistic interpretation presented in Section 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=better, Appendix A.5), and asymptotic complexity of the proof strategy as we scale the inputs and model (lower=better). 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.then reverse engineer the models using standard mechanistic interpretability techniques. We use our understanding to define a set of 102 different computer-assisted proof strategies with varying tightness of bound and with different asymptotic complexity and number of required FLOPs (Section 4). 4 We validate our technique against an additional 604 models for varying values of K (Appendix A.2.1).We define a quantitative metric to assess the mechanistic understanding used in a proof strategy by the dimensionality of the function space that the proof strategy must consider, which we deem the unexplained dimensionality of the proof strategy (Sections 5.1, and A.5). Using this metric, we find a negative relationship between proof length and degree of understanding. We qualitatively examine proof strategies to confirm and explain this relationship, finding that more compact proofs both require and provide more mechanistic understanding. We also find suggestive evidence that the trade-off between proof length and tightness of bound is modulated by the faithfulness of the mechanistic understanding used to derive the proof (Section 5.2). 5However, we also identify compounding structureless error terms as a key challenge for generating compact proofs on model behavior (Sections 5.3, and G.2.5). The implementation of algorithms inside of neural networks may contain components that defy mechanistic understanding and appear to us as "noise". When we don't know how noise composes across model components, establishing a bound requires pessimizing over the ways the composition could occur. Worst-case noise can quickly grow even when the empirical noise is small, leading to vacuous performance bounds. 
mechanistic interpretability for proofs: 
 Generalization bounds on global performance In the style of prior mechanistic interpretability evaluation work [6] , we target theorem templates that establish bounds on the expected global performance of the model. Let M : X → Y be a model (here assumed to be a neural network), D be a probability distribution over input-label pairs (l, t) ∈ L × X, notated as D| X when marginalized over labels, and f : L × Y → R be a scoring function for evaluating the performance of the model. Then, we seek to establish lower bounds b on the expected s as the form:EQUATIONAs f can be any metric, this is a fully general template for theorems that can capture any aspect of model performance for which we have a formal specification. However, in this work we restrict f to be the accuracy and D| X to be uniform, so our theorems lower bound the accuracy of the model. Our proof methodology generalizes straightforwardly to other input distributions (Appendix A.8), and only a little work is required to generalize from accuracy to log-loss (Appendix A.11) .Proof template The proofs of model performance in this work have two components: a computational component C : model weights → R and a non-computational component Q arguing that for any model M ′ , C(M ′ ) ≤ E (l,t)∼D f (l, M ′ (t)), thus implying that C generates a valid lower bound for the performance of M. The whole proof is Q paired with a trace of running C that certifies its output on M. 6 Here, b = C(M). As even the size of the model parameters is much larger than any reasonable Q, we approximate the length of a proof pair C, Q by the length of a trace of C(M).Proof compactness vs. tightness of bound Different proof strategies make different tradeoffs between compactness and tightness of bound. For example, consider two extreme proof strategies:We can "prove" a vacuous bound using a null proof. On the other hand, in the brute-force proof, we simply run the model on the entirety of D to achieve b = s, albeit with a very long proof.We quantify the length of C(M) using two metrics: the asymptotic time complexity of C as we scale the size of the model and the input t, as well as the empirical average number of floating point operations required to evaluate C(M ′ ) over a given set of models {M i }. We measure tightness of bound of C(M) using the ratio of the bound to the true accuracy: b/s.Proof as pessimal ablation A standard way of assessing the faithfulness of mechanistic interpretability is by ablating the parts of the model that your interpretation does not explain [54, 6, 23] . In this framework, proofs can be thought of as performing a pessimal ablation over the unexplained parts of the model -we set the remaining components of the model (the "noise" or error terms) to values over X that minimize the performance of the model. However, the number of ablations required for a complete argument might be quite high. Thus, we construct relaxations (Appendix A.4) over input sequences, such that performing pessimal ablations on a smaller number of relaxed input sequences is sufficient to lower bound the performance on D. 
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 For an n ctx × d vocab one-hot encoding x = [x 0 , x 1 , . . . , x nctx-1 ] of an input sequence t = [t 0 , t 1 , . . . , t nctx-1 ], we compute the logits of the model as follows:h (0) = xE + P Initial residual stream (n ctx × d model ) α = h (0) QK T h (0) T / √ d Attention matrix (n ctx × n ctx ) h (1) = σ * (α) • h (0) V O + h (0) Final residual stream (n ctx × d model ) M(t) = ℓ = h (1)nctx-1 U Final seq. position logits (d vocab )where σ * is the masked softmax function used in causal attention. Because we only look at outputs of the model above the final sequence position i = n ctx -1, we also denote this position as the "query position" and the value of the token in this position as t query , one-hot encoded as x query . The model's prediction is the token corresponding to the max-valued logit ℓ max .Task Specifically, we study the setting with n ctx = k = 4 because it is the largest sequence length for which we can feasibly evaluate the brute-force proof. We set hidden dimension d model = 32and a vocabulary of size d vocab = 64 comprising integers between 0 and 63 inclusive. For an input sequence t, we denote the true maximum of the sequence by t max . Outputting the correct behavior is equivalent to outputting logits ℓ such that ∆ℓ t * := ℓ t * -ℓ max < 0 for all t * ̸ = t max . We trained 151 models on this task. Models achieved average accuracy 0.9992 ± 0.0015 over the entire distribution.Path decomposition Following prior work [13] , we expand the logits of the model and split the paths through the model into three components -the QK circuit, the OV circuit, and the direct path: M(t) =Intuitively, the QK circuit determines which tokens the model attends to from a particular query token and sequence position, while the OV circuit processes the tokens and sequence positions the model attends to. The direct path is simply the skip connection around the attention head.We further divide the QK and OV circuits into token (position-independent) and position-dependent components. Let P avg = i P i /n ctx be the average position embeds across positions (of size d model ), and let P denote either 1 nctx ⊗ P avg or 1 d vocab ⊗ P avg depending on context, the result of broadcasting P avg back into the shape of P or E (that is, n ctx × d model or d vocab × d model ). Similarly, let P q = 1 d vocab ⊗ P query be the result of broadcasting P query . Then for one-hot encoded x, we can rewrite the QK and OV circuits, as well as the direct path, as follows:QK circuit = x query E q QK T ĒT EQKE x T + E q QK T PT EQKP OV circuit = x ĒV OU EVOU + PV OU PVOU Direct Path = x query E q UEU where P = P -P and Ē = E + P and E q = E + P q (since h (0) = x Ē + P). Using standard empirical mechanistic interpretability techniques, we interpret one of our learned models (our "mainline" model) by independently examining the QK and OV circuits and the direct path. 7 We find that the model outputs the largest logit on the true max token t max by attending more to larger tokens via the QK circuit and copying the tokens it attends to via the OV circuit. We then quantitatively confirm that these interpretations hold for all 151 models by reporting the mean plus minus standard deviation for various summary statistics. Plots for this section are available in Appendix B.2.QK circuit By qualitatively examining the positionindependent QK component EQKE, we find the amount of pre-softmax attention paid to a key token is approximately independent of the value of the query token t query , and increases monotonically based on the size of the key token. We confirm this hypothesis by performing a singular-value decomposition (SVD) of the EQKE matrices (Appendix G.2.3), and find that it contains a single large rank-one component with singular value around 7800 ± 380, around 620 ± 130 times larger than the second largest component with singular value 13 ± 3. The left (query-side) singular vector is approximately constant in all dimensions, with value 0.1243 ± 0.0003≈ 1 ⁄8 = 1 / √ d vocab .The right (key-side) singular vector of this component is monotonically increasing as we increase the size of the key token, with (1/ √ d-scaled) pre-softmax attention increasing by an average of 1.236 ± 0.056 when the key token increases by 1. 8 In comparison, each 1/ √ d-scaled entry of the position-dependent QK component EQKP has negligible size (average 0.31 ± 0.18), suggesting that EQKP is unimportant to the functioning of the model.We confirm this by zero ablating EQKP, which changes the models' accuracies from 0.9992±0.0015 to 0.9993 ± 0.0011. Combined with our interpretation of EQKE, this implies that the attention pattern of the model depends only on the token values and not the ordering of the sequence.OV circuit Then, by qualitatively examining the position-independent OV component EVOU, we see that it has large positive entries along the diagonal. In fact, the entry along the diagonal is the largest in the row for all rows corresponding to t > 6.6 ± 1.2. Since each entry in the sequence is uniformly sampled and d vocab = 64, this means that EVOU is a good approximation for the identity matrix for all but ≈ (7/64) 4 ≈ 1.2 × 10 -2 % of the sequences.As with the position-dependent QK component, the position-dependent OV component PVOU also has negligible size and is unimportant to model performance. Taken together with the above results on EVOU, this suggests that the attention head copies the tokens it attends to.As with the two position-dependent components, the entries in EU have small absolute magnitude 2.54 ± 0.20, 9 and contribute negligibly to model performance. 
proofs of model performance: 
 In this section we describe intuitions for three categories of proof that are developed around different mechanistic interpretations and methods for using the interpretations. The strategies result in proofs of different complexities with varying bound tightness (Table 1 ). We provide detailed theorem statements, proofs, algorithms, and explanations of proof search in Appendices C, D, E, F, and G.Our theorem statements for Q will all be of the form∀M ′ , C specific strategy (M ′ ) ≤ E t∼D| X f (t max , M ′ (t)).We leave implicit the traces of running C specific strategy on our specific models to give the overall theorem. We report the computational complexity or estimated FLOPs of running C specific strategy as approximations for our proof lengths.We start by considering the brute-force proof (Appendix D), which treats the model as a black box and evaluates it on all possible sequences. 10 However, this proof strategy has bad asymptotic complexity and is untenable for larger models and larger input distributions. So in subsequent sections, we use knowledge of the model drawn from the interpretation in Section 3.1 to derive more compact proofs.Next, we use the fact that the model is composed of the direct path and the QK and OV circuits (Section 3) to decrease the number of sequences that we need to consider, and the fact that only the position-independent components EQKE and EVOU contribute meaningfully to performance (Section 3.1) to pessimize over sequence ordering.First, let a pure sequence ξ be a sequence with at most three distinct tokens: the max token t max , the final token t query ≤ t max , and optionally a third token t ′ < t max , and let Ξ pure be the set of all pure sequences in X. 11 For a given input sequence t, define the adjacent pure sequences Adj(t) as the set of sequences that share the same max and query token, and only take on values in t:Adj(t) = ξ ∈ Ξ pure max i ξ i = t max , ξ query = t query , ∀i < n ctx , ξ i ∈ tUsing the convexity of softmax and the fact that the model contains three paths, we can show that one-layer attention-only transformers satisfies a variant of the following convexity property: for a given t, if M(ξ) is correct for all ξ ∈ Adj(t), then M(t) is correct. That is, for these transformers, we can bound the accuracy on all sequences by evaluating M on only the O(d vocab 3 (n ctx -1)!) pure sequences. This allows us to bound the accuracy of our actual M on all d vocab nctx sequences, while evaluating it on O(d vocab 3 (n ctx -1)!) sequences.We can reduce the number of sequences that we need to evaluate by pessimizing over the order of a sequence. For a given tuple of (t max , t query , t ′ ), there are (n ctx -1)! pure sequences, corresponding to the permutations of the tuple. Pessimizing over the order of sequences reduces the number of sequences to consider for each (t max , t query , t ′ ) tuple to the number of t ′ in the pure sequence, and the total number of sequences to O(d vocab 3 n ctx ). By precomputing the five component matrices EU, EQKE, EQKP, EVOU, PVOU and cleverly caching intermediate outputs, we can reduce the additional work of each sequence to the O(n ctx ) required to compute the softmax over n ctx elements, resulting in asymptotic complexity O(d vocab 3 n ctx 2 ) (Theorem 12, additional details in Appendix E).We now consider proofs that are more compact than O(d vocab 3 ). These require avoiding iteration over any set of size O(d vocab 3 ) (e.g. the set of pure sequences) and performing operations that take O(d vocab ) time on each of O(d vocab 2 ) combinations. Unfortunately, some methods of avoiding these operations can lead to vacuous bounds (i.e. accuracy lower bounds near 0%). In order to recover non-vacuous bounds, we introduce two tricks: the "mean+diff trick" to better approximate the sum of two components with unequal variance, and the "max row diff trick" to improve upon the low-rank approximations for EU and EQKE. We consider applying variants of these tricks at different locations in the naïve subcubic proof, leading to 100 distinct subcubic proof strategies. See Appendix G.2 for a formal description of these strategies.Reducing the number of cases by pessimizing over sufficiently small tokens Previously, we considered Θ(d vocab 3 n ctx ) pure sequences ξ, with ξ parameterized by (t max , t query , t ′ , c). Recall from our mechanistic interpretation in Section 3.1 that the pre-softmax attention paid from t query to a key token t ′ is broadly invariant in t query and increases roughly linearly with the size of t ′ . This allows us to pessimize over the OV circuit over all "sufficiently small" tokens.More formally, suppose we are given some gap g ∈ N. For each pure sequence ξ with max token t max , query token t query , such that t query ≤ t max -g, and c copies of the third token type t ′ ≤ t max -g, we pessimally ablate the OV circuit over the set Ξ pure (t max , t query , c; g) of pure sequences ξ ′ with the same max and query tokens and c copies of the third token type t ′ . If the model gets all sequences in Ξ pure (t max , t query , c; g) correct, then we can conclude that it gets ξ correct, otherwise, we treat the model as having gotten ξ wrong. ) cases parameterized by (t max , t query ) or (t max , t ′ ). For example, as part of the pessimization procedure over pure sequences, for each of the d vocab possible values of t max , we need to consider the relative effects on the d vocab -sized logits of attending to each of the O(d vocab ) other tokens t ′ < t max , and for each t max and t query , we need to check that the contribution of the direct path on logits x query EU is not sufficiently large as to overwhelm the contribution from x max EVOU. We independently pessimize over each of these components over one of the d vocab -sized axes: for example, instead of computing x max EVOU + x query EU for each t max , t query pair, we first pessimally ablate the direct path along the query token (which takes O(d vocab2 ) time as it does not depend on the t max , and then consider the sum Similarly, we can more efficiently check that the direct path EU contributes negligibly to the model outputs, by using SVD to decompose EU into a sum of rank 1 products (which we can evaluate exactly) and a high-rank error term that we can cheaply bound.x max EVOU + max x ′ x ′ EU.Tighter bounds for sums of variables with unequal variance via the "mean+diff trick" Suppose we want to lower bound the minimum of the sum of two functions over three variables h(x, y, z) = f (x, y) + g(y, z), while only iterating over two variables at a time. The naïve way is to minimize f (x, y) and g(x, y) independently:min x,y,z h(x, y, z) ≥ min x,y f (x, y) + min y,z g(y, z)Here, the error comes from setting the ys in f and g to different values. But in cases where g(y, z) varies significantly with y and only slightly with z, rewriting g as a sum of a component that is independent of z (only varying along y), and one that depends on z, yields a better lower bound:min x,y,z h(x, y, z) ≥ min x,y (f (x, y) + E ′ z g(y, z ′ )) + min y,z (g(y, z) -E ′ z g(y, z ′ ))This estimate will have error at most ε, while the naïve estimator can have arbitrarily large error. We refer to this rewrite as the "mean+diff trick". 12 From the mechanistic interpretation in Section 3.1, we know that some of the components barely vary among one or more axes. So we can apply the mean+diff trick to get tighter lower bounds.Avoiding matrix multiplications using the "max row-diff trick" Using properties of linear algebra, we derive a cheap approximation to the max row-diff for the product of matrices AB in terms of the product of the max row-diff of B and the absolute value of A, which we deem the "max row-diff" FLOPs to Verify Proof (approximate) Normalized Accuracy Bound brute force (acc: 0.9992 ± 0.0015) cubic (rel acc: 0.9539 ± 0.0080) subcubic (rel acc: 0.821 ± 0.013) attention-dvocabdmodel 2 (rel acc: 0.795 ± 0.014) direct-quadratic (rel acc: 0.653 ± 0.060) attention-dvocabdmodel 2 , direct-quadratic (rel acc: 0.628 ± 0.060) attention-quadratic (rel acc: 0.390 ± 0.032) attention-quadratic, direct-quadratic (rel acc: 0.286 ± 0.036) Figure 3 : For each of the proofs in Section 4, we plot the number of FLOPs used to compute the certificate, as well as the normalized accuracy lower-bound (b/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 a 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, but we still achieve worse bounds. See Appendix H.2.1 for a detailed description of the various proof strategies.trick. We apply this trick to get a better cheap bound on the error terms of low-rank approximations, without having to multiply out the full matrices. See Appendices G.2.2, and F for more details. 
results: 
 We run each of 151 transformers on the various proof strategies of different asymptotic complexity, and analyze these proofs to empirically examine the relationship between proof length, bound tightness, and degree of understanding. For each proof on each transformer, we approximate the length of the proof by estimating the number of FLOPs used, and plot this against the ratio of certified bound the true accuracy b/s (Equation 2) in Figure 3 . There exists a clear trade-off between bound tightness and compactness of the proof -more compact proofs yield looser bounds, and tighter bounds are associated with more expensive proofs.Quantifying mechanistic understanding using unexplained dimensionality We first quantify the amount of mechanistic understanding used in a proof by measuring its unexplained dimensionality -the number of free parameters required to fully describe model behavior, assuming the structural assumptions of the proof are correct. More detailed interpretations leave fewer free parameters needing to be filled in via empirical observation (Appendix A.5). In Figure 5 , we plot these axes and find a suggestive correlation: proofs based on less mechanistic understanding are longer.More mechanistic understanding allows for more compact proofs In addition to the constructions in Section 4, the parts of proofs we were unable to compact seem to correspond to components that we do not mechanistically understand. For example, we could not cheaply bound the behavior of EVOU without multiplying out the matrices, and this seems in part because we do have a mechanistic understanding of how EVOU implements low-rank copying.Compact proofs seem to provide understanding By examining compact proofs, we can extract understanding about the model. For example, the fact that replacing each row of EU with its average across rows has little effect on the bound implies that EU does not vary much based on t query .Compact proofs are less faithful to model internals To derive more compact proofs, we use our mechanistic understanding to simplify the model computation in ways that diverge from the original model internals. For example, in some subcubic proofs (Section 4.3), we approximate EQKE with a rank-1 approximation corresponding to the "size direction". However, while other components are small, they're nonzero; this approximation harms model internals.To confirm that faithfulness of understanding affects the tightness of bound independent of proof length, we plot the normalized accuracy bound of subcubic proofs that perform a rank-1 approximation to EQKE, versus the ratio of the first two singular components. A larger ratio between the components implies that the rank-1 approximation is more faithful. In Figure 4 , we see a positive correlation between the two axes: when the interpretation is more faithful, the bounds are tighter, even at a fixed proof length. EPQKE Singular Ratio: σ1/σ2 Normalized Accuracy Bound max-diff-exact mean+max-diff-subproduct max-diff-subproduct max-diff mean+max-diff svd max-diff-subproduct-recursive mean-recursive+max-diff-subproduct-recursive mean+max-diff-subproduct-recursiveFigure 4 : We plot the normalized accuracy bound versus the ratio of first and second singular values of EQKE, for various types of subcubic proofs that depend on a rank-1 approximation EQKE. For each class of proof, the 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. 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 seem 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.Pessimal error terms compound in the absence of known structure The rank-1 approximation of EQKE has small error. However, when making rank-1 approximations of each of the constituent matrices E, Q, K, pessimizing over the worst way to composing the individual small error terms leads to a bound on the error term of EQKE that is orders of magnitude larger than the actual error term. Because we don't understand how the matrices compose in a way that doesn't cause errors to compound (without just multiplying out the matrices), this approximation leads to a trivial bound on performance (Appendix G.2.5). We speculate that in many cases, there is no short human-interpretable description for why random noise or approximation errors do not compound across layers of neural networks (e.g., see the error correction results on randomly initialized neural networks from Hänni et al. [21] ), and thus that compounding structureless errors may be an issue in practice. 
related work: 
 Generalization Bounds Prior work in the PAC-Bayes framework [58, 36, 12] proves generalization bounds over learning procedures, which are similar to the global performance bounds we consider in this work. These proofs tend to provide statistical guarantees [25, 26] about the outputs of a known stochastic training procedure, while we seek to bound the performance of particular trained models.Formally verifying neural networks Most prior work formally verifies neural networks either via model checking [28, 7] or by relaxing the problem setting and taking an automated theorem proving approach [17, 50, 18, 35, 43] to verify local robustness properties. These proof strategies tend to be derived by examining only the network architecture. We take an approach more akin to interactive theorem proving [22] and verify global performance properties by reverse-engineering the weights.Mechanistic Interpretability Finally, mechanistic interpretability is the subfield of the broader field of understanding model internals [45] , which is too large to faithfully summarize. Our work takes most direct inspiration from efforts to deeply understand how either toy models [38, 9, 53, 2] or small pretrained text transformers [54, 20] implement algorithmic tasks, generally by performing ablations and SVD. In contrast, we formally prove that a transformer implements an algorithm.Nichani et al. [39] proves that, in a significantly simplified 2-layer, 1-head attention-only transformer model and for the task of in-context bigram statistics, gradient descent will create induction heads [40] .Our results concern transformers with fixed weights. In concurrent work, Michaud et al. [34] use techniques inspired by mechanistic interpretability to perform automated program synthesis on 2-dimensional RNNs, while our work works with significantly larger transformer models. 
conclusion and future work: 
 Summary In this work, we used a Max-of-K setting to prototype the use of mechanistic interpretability to derive compact proofs of model behavior. Using varying amounts of understanding, we derived more efficient proof computations lower bounding model accuracy. We found preliminary evidence that mechanistic understanding can compactify proofs. Moreover, we observed that the tightness of the lower bound offered by various proof strategies can be used to grade the faithfulness our mechanistic interpretation. Finally, we identified compounding structureless errors as a key obstacle to deriving compact proofs of model behavior.Limitations and future work We study one-layer attention-only transformers on a toy algorithmic task. Future work should explore the viability of deriving proofs via interpretability using larger models featuring MLPs or layernorm on more complex domains. In addition, we were unable to significantly compact the part of the proof involving the OV circuit, which future work can explore.The proofs we explored in this work also did not lead to qualitatively novel insights; future work may be able to derive such insights with improved techniques. Finally, future work can address the problem of compounding structureless errors, perhaps by relaxing from worst-case pessimal ablations to typical-case heuristic guarantees [8] . 
 
 Appendix: In this section, we address some subtleties and frequently asked questions about our approach.A.1 Why study this simple task?Formal reasoning is computationally expensive; very few large software projects have ever been verified [31, 30] , none of them comparable to large transformer models [10, 19] . Separately, there is a high fixed cost to taking on any verification project, regardless of computational efficiency of the verification itself. Thus, we picked the simplest setting to study the question of interest: Is it even possible to formally reason more efficiently than by brute force about model behavior?In this section, we address concerns about the scalability of our approach.We demonstrate that our proof strategies can be reused on larger input spaces while scaling better than the brute force approach does.We applied our proof strategies to models trained for Max-of-5, Max-of-10, and Max-of-20. While running the brute force proof on Max-of-20 would require approximately 2 148 FLOPs, which is about 2 70 × the cost of training GPT-3 [32] , our cubic proof achieves bounds of (94.1 ± 1.1) % (Max-of-5), (91.4 ± 2.1) % (Max-of-10), and (88.4 ± 4.0) % (Max-of-20) in under two minutes. See Tables 2, 3 , 4, and 5 for more detailed numbers, and Figures 6, and 7 for visualizations. These results demonstrate that proof strategies can be reused on larger input spaces while scaling better than the brute force approach does.In this paper, we worked on highly optimizing our relaxation to make our bounds as tight as possible when incorporating as little understanding as possible. This is not necessary for deriving proofs. Our general formalization of mechanistic interpretability is replicable: (1) theorem statements are exact expressions for the difference between the actual behavior of the model and the purported behavior, and (2) proofs are computations that bound the expression. Furthermore, our convexity theorems and proofs are applicable much more generally generally to element retrieval tasks.We worked on a simple model studied in A Mathematical Framework for Transformer Circuits [13] .In follow-up work, we will extend this approach to proving bounds on 1L transformers with ReLU MLP trained on modular addition.It is an open question whether or not the mechanistic interpretability approach to proofs can scale to larger models. However, a large part of this question lies in the feasibility of deriving a high degree of faithful mechanistic understanding from large models -that is, whether mechanistic interpretability itself will scale. This is widely recognized in the field, and scaling interpretability approaches while getting both a high degree of mechanistic understanding and assurances that said understanding is faithful to the model is an active area of research. Broadly, we see the compact proofs approach as a metric on the quality of mechanistic understanding -we are not purporting to have a general solution to the problem of scaling interpretability, but instead claim that the challenges in proofs are in fact challenges in understanding networks. O(v 2 • k 2 + v 2 • d) 0.806 ± 0.014 2 22 w/o mean+diff 0.545 ± 0.066 2 22 Low-rank QK O(v 2 k 2 + vd 2 QK + v 2 d EU&OV ) 0.784 ± 0.014 2 22 SVD only 0.461 ± 0.063 2 22 Low-rank EU O(v 2 k 2 + vd EU + v 2 d QK&OV ) 0.666 ± 0.053 2 21 SVD only (4.8 ± 0.1) × 10 -8 2 21 Low-rank QK&EU O(v 2 k 2 + vd 2 QK + vd EU + v 2 d OV ) 0.645 ± 0.052 2 22 SVD only (4.8 ± 0.1) × 10 -8 2 22 Quadratic QK O(v 2 k 2 + vd QK + v 2 d EU&OV ) 0.423 ± 0.029 2 21 Quadratic QK&EU O(v 2 k 2 + vd QK&EU + v 2 d OV ) 0.339 ± 0.032 2 21Table 3 : Version of Table 1 from Section 4.1 with nctx = 10, d vocab = 64. We report the proof complexity, accuracy bound, and estimated flops required (Equation 2), as well as 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 := nctx, d := d model , and v := d vocab .Complexity Cost Bound Est. FLOPsBrute forceO(v k+1 kd) 0.9988 ± 0.0013 2 86 Cubic O(v 3 k 2 ) 0.914 ± 0.021 2 28 Sub-cubic O(v 2 • k 2 + v 2 • d) 0.733 ± 0.019 2 23 w/o mean+diff 0.643 ± 0.043 2 23 Low-rank QK O(v 2 k 2 + vd 2 QK + v 2 d EU&OV ) 0.717 ± 0.019 2 23 SVD only 0.562 ± 0.045 2 23 Low-rank EU O(v 2 k 2 + vd EU + v 2 d QK&OV ) 0.656 ± 0.032 2 23 SVD only (2.9 ± 0.1) × 10 -17 2 22 Low-rank QK&EU O(v 2 k 2 + vd 2 QK + vd EU + v 2 d OV ) 0.641 ± 0.032 2 23 SVD only (2.9 ± 0.1) × 10 -17 2 23 Quadratic QK O(v 2 k 2 + vd QK + v 2 d EU&OV ) 0.437 ± 0.028 2 22 Quadratic QK&EU O(v 2 k 2 + vd QK&EU + v 2 d OV ) 0.398 ± 0.028 2 22Table 4 : Version of Table 1 from Section 4.1 with nctx = 10 and d vocab = 128. We report the proof complexity, accuracy bound, and estimated flops required (Equation 2), as well as 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 := nctx, d := d model , and v := d vocab .Complexity Cost Bound Est. FLOPsBrute forceO(v k+1 kd) 0.9972 ± 0.0031 2 96 Cubic O(v 3 k 2 ) 0.882 ± 0.012 2 31 Sub-cubic O(v 2 • k 2 + v 2 • d) 0.758 ± 0.015 2 24 w/o mean+diff 0.573 ± 0.059 2 24 Low-rank QK O(v 2 k 2 + vd 2 QK + v 2 d EU&OV ) 0.740 ± 0.015 2 24 SVD only 0.476 ± 0.050 2 25 Low-rank EU O(v 2 k 2 + vd EU + v 2 d QK&OV ) 0.624 ± 0.030 2 24 SVD only (5.4 ± 0.2) × 10 -20 2 24 Low-rank QK&EU O(v 2 k 2 + vd 2 QK + vd EU + v 2 d OV ) 0.611 ± 0.029 2 24 SVD only (5.4 ± 0.2) × 10 -20 2 25 Quadratic QK O(v 2 k 2 + vd QK + v 2 d EU&OV ) 0.450 ± 0.025 2 24 Quadratic QK&EU O(v 2 k 2 + vd QK&EU + v 2 d OV ) 0.391 ± 0.024 2 24Table 5 : Version of Table 1 from Section 4.1 with nctx = 20, d vocab = 64. We report the proof complexity, accuracy bound, and estimated flops required (Equation 2), as well as 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 := nctx, d := d model , and v := d vocab .Complexity Cost Bound Est. FLOPsBrute force FLOPs to Verify Proof (approximate) Normalized Accuracy Bound brute force (estimated) (acc: 0.9990 ± 0.0018) cubic (rel acc: 0.942 ± 0.010) subcubic (rel acc: 0.806 ± 0.013) attention-dvocabdmodel 2 (rel acc: 0.785 ± 0.014) direct-quadratic (rel acc: 0.666 ± 0.053) attention-dvocabdmodel 2 , direct-quadratic (rel acc: 0.645 ± 0.052) attention-quadratic (rel acc: 0.423 ± 0.029) attention-quadratic, direct-quadratic (rel acc: 0.340 ± 0.032) FLOPs to Verify Proof (approximate) Normalized Accuracy Bound brute force (estimated) (acc: 0.9988 ± 0.0013) cubic (rel acc: 0.915 ± 0.021) subcubic (rel acc: 0.734 ± 0.019) attention-dvocabdmodel 2 (rel acc: 0.718 ± 0.019) direct-quadratic (rel acc: 0.657 ± 0.032) attention-dvocabdmodel 2 , direct-quadratic (rel acc: 0.642 ± 0.032) attention-quadratic (rel acc: 0.437 ± 0.028) FLOPs to Verify Proof (approximate) Normalized Accuracy Bound brute force (estimated) (acc: 0.995 ± 0.015) cubic (rel acc: 0.888 ± 0.036) subcubic (rel acc: 0.602 ± 0.034) attention-dvocabdmodel 2 (rel acc: 0.586 ± 0.036) direct-quadratic (rel acc: 0.554 ± 0.039) attention-dvocabdmodel 2 , direct-quadratic (rel acc: 0.542 ± 0.038) attention-quadratic (rel acc: 0.348 ± 0.035) attention-quadratic, direct-quadratic (rel acc: 0.328 ± 0.034) A.3 Why is more mechanistic understanding correlated with worse bounds? Figure 3 exhibits Simpson's Paradox: although more faithful mechanistic understanding is correlated with better bounds within each class of proof (and moreover the most extensive mechanistic understanding results in the greatest improvement in bound tightness over baseline), when we aggregate across all proof strategies, we find that more mechanistic understanding is correlated with worse bounds.O(v k+1 kd) 0.995 ± 0.015 2 148 Cubic O(v 3 k 2 ) 0.884 ± 0.040 2 29 Sub-cubic O(v 2 • k 2 + v 2 • d) 0.599 ± 0.036 2 24 w/o mean+diff 0.552 ± 0.049 2 24 Low-rank QK O(v 2 k 2 + vd 2 QK + v 2 d EU&OV ) 0.584 ± 0.039 2 24 SVD only 0.488 ± 0.048 2 24 Low-rank EU O(v 2 k 2 + vd EU + v 2 d QK&OV ) 0.552 ± 0.042 2 24 SVD only (1.0 ± 6.0) × 10 -4 2 24 Low-rank QK&EU O(v 2 k 2 + vd 2 QK + vd EU + v 2 d OV ) 0.539 ± 0.040 2 24 SVD only (0 ± 40) × 10 -5 2 24 Quadratic QK O(v 2 k 2 + vd QK + v 2 d EU&OV ) 0.347 ± 0.035 2 24 Quadratic QK&EU O(v 2 k 2 + vd QK&EU + v 2 d OV ) 0.327 ± 0.035 2 24(a) nctx = 5, d vocab = 64(b) nctx = 10, d vocab = 64(d) nctx = 20, d vocab = 64EPQKE Singular Ratio: σ1/σ2 Normalized Accuracy Bound max-diff-exact mean+max-diff-subproduct max-diff-subproduct mean+max-diff max-diff svd max-diff-subproduct-recursive mean+max-diff-subproduct-recursive mean-recursive+max-diff-subproduct-recursiveEPQKE Singular Ratio: σ1/σ2 Normalized Accuracy Bound max-diff-exact max-diff-subproduct mean+max-diff-subproduct mean+max-diff max-diff svd mean+max-diff-subproduct-recursive max-diff-subproduct-recursive mean-recursive+max-diff-subproduct-recursiveEPQKE Singular Ratio: σ1/σ2 Normalized Accuracy Bound max-diff-exact max-diff-subproduct mean+max-diff-subproduct mean+max-diff max-diff svd mean+max-diff-subproduct-recursive max-diff-subproduct-recursive mean-recursive+max-diff-subproduct-recursiveThis relationship is summarized in Figure 8 .From the compression perspective, more mechanistic understanding is about having more compression. Unless the model is losselessly compressible, we should expect that more compression will inherently be more lossy, no matter how good our compression scheme is. Correspondingly, using more understanding to get more compression will often result in a weaker bound, no matter how good our understanding is.Conversely, we can think of the quality of proofs (the combination of tightness of bound, and length of proof) as a metric for how good our mechanistic understanding is. From this perspective, the fact that mechanistic-interpretability-derived bounds are bad suggests gaps in our mechanistic understanding. As the field matures and we develop tools that enable more faithful and complete understanding of model behavior, we expect that the quality of bounds we derive from mechanistic understanding will improve.In this work, we construct convex relaxations to perform the pessimal ablations for our proofs.In what sense are we using "convexity"? The intuition is that we are attempting to optimize a function f over its domain X by incrementally making local changes to the input sequence, such as replacing one token by another, or by changing the order of tokens. The reason that convex optimization problems are easy to solve is that all local extrema are global extrema. This is not the case for our optimization problem, so we find a relaxation of f and its domain such that all local extrema are in fact global extrema. Furthermore, most convex optimizers perform optimization at runtime by repeatedly stepping towards extrema. In this work, we "optimize by hand", performing the optimization in the proof of our general theorems. The computation of the bound then only needs to instantiate the precomputed possible extrema with the actual values of the model's parameters to determine the the extrema actually are.We now give a formal description of what we mean by "convex relaxation".For a set of inputs X i , we define a set of "relaxed inputs" X relaxed i with an injection T i : X i → P(X relaxed i ) mapping input to the model to the set of corresponding relaxed inputs. On the relaxed input, we define a function h i : X relaxed i → R such that for all t ∈ X i and all labels l for which (l, t) is supported by (has non-zero probability in) D, we can find t relaxed ∈ T i (t) with f (l, M(t)) ≥ h i (t relaxed ). We proceed by finding a small subset of "boundary" examplesB i ⊂ X relaxed i , proving that if h i (t relaxed ) ≥ b i for all t relaxed ∈ B i then h i (t relaxed ) ≥ b i for all t relaxed ∈ X relaxed i .Then, the computational component C of the proof validates that that h i (t relaxed ) ≥ b i for some b i for all t relaxed ∈ X relaxed i . This allows us to conclude that f (l, M(t)) ≥ b i for all t ∈ X i .We claim in Figure 5 that we can use unexplained dimensionality as a metric for understanding. Here we describe how we compute the unexplained dimensionality of a proof strategy.As in Figure 1 , for any given proof, we can separate our treatment of transformer components into "black-box" (e.g., matrix multiplication) and "white-box" components (e.g., specifying that the QK circuit is approximately rank one; pessimizing over non-max tokens). Considering the performance score as a large white-box component which may reference black-boxes internally, we define the unexplained dimensionality of a single black-box computation as the log-cardinality of it function space (so, e.g, 2 • 64 for a function 64 → R 2 , whose cardinality is (R 2 ) 64 , where 64 denotes the finite set on 64 elements). The unexplained dimensionality of the entire proof is the sum of the unexplained dimensions of all black-box components.Intuitively speaking, unexplained dimensionality tries to capture the degrees of freedom that we have to check via brute enumeration over black-box computations. Proofs with less unexplained dimensionality contain more mechanistic understanding, and vice versa.In Figure 3 and Table 1 on page 6 and on page 8, we display approximate floating point operations. We instrument our code to execute on phantom tensors that track their shape and accumulate an approximate count of floating point operations. We compute matrix additions and multiplications in the obvious way. We take the instruction count of SVD to be the cost of verifying that the output of SVD is a valid decomposition: that we have a pair of orthonormal bases which when multiplied out give the original basis.A.7 IEEE 754 vs. RIn Section 2 we defined C and Q and glossed over whether we were reasoning over reals or floats.Here we clarify this point that we've so far been sweeping under the rug.Let F denote the set of the relevant flavor of IEEE 754 Floating Point numbers (generally 32-bit for our concrete models, but everything would hold just as well for 64-bit). Let F * denote F restricted to finite numbers (that is, without NaNs and without ±∞).We parameterize C, M, and D over the real field 13 they operate on, so that, e.g., C F : model weights → F . Then we have Q establishing that for any modelM ′ , C R (M ′ R ) ≤ E (l,t)∼D R f R (l, M ′ R (t)), and we have a trace demonstrating that C F (M F ) = b.Let i : F * → R be any injection that maps each floating point number to some real number that it is "closest to". Supposing that b ∈ F * and thus b ∈ R, we need two additional components of the proof. We need to find ε, ε′ ∈ R + prove that |C R (M R ) -i (C F (M F ))| < ε and E (l,t)∼D R f R (l, M R (t)) -i E (l,t)∼D F f F (l, M F (t)) < ε ′Then we can chain these proofs to prove thati E (l,t)∼D F f F (l, M F (t)) ≥ b -ε -ε ′Such ε-ball robustness proofs should be well within the scope of existing approaches to formal methods on neural nets, see, e.g., [44, 3, 4, 29, 1, 57] . We leave actually dealing with the gap between floating point numbers and real numbers to future work.In Equation 1in Section 2 we defined the expected model performance as the expectation of the distribution D:s := E (l,t)∼D [f (l, M(t))] ≥ b.We then immediately specialized to the case where the marginalization D| X of D over labels is uniform. As we'll see in Theorem 1 in Appendix D and Algorithm 3 in Appendix E, the bound computation is modularized between a function that bounds the performance f (l, M(t)) over a restricted collection of inputs, and a much simpler function that combines the bounds on individual cases into a bound on the expectation over the entire distribution. The per-input bound computation is CORRECTNESS in Algorithm 1 and RELAXED-CORRECTNESS-PESSIMIZING-OVER-POSITION in Algorithm 3; the expectation computation is BRUTE-FORCE in Algorithm 1 and CUBIC in Algorithm 3.Since the expectation computation is modularized, it is straightforward to extend our approach to non-uniform distributions simply by adjusting the weighting of each region of inputs. However, if the distribution is too far off from the uniform training distribution, the bound we get may not be very good, as we may not be allocating adequate computation to the high-probability regions of the input space.There is flexibility inherent in Equation 1. Normally, by out-of-distribution (OOD) or adversarial inputs, we suppose that there is a distribution D in that's used for training and (in-distribution) validation, and another distribution D ′ that is the deployment distribution or generated by an adversary.If we had knowledge of D ′ , we could compute the expected performance from inputs sampled from D ′ . Even if we don't have exact knowledge of D ′ , we can still define a very broad distribution D that covers possible D ′ s.In this work, D is the distribution of all 64 4 possible valid input sequences. In addition, as our proofs partition D into subdistributions, and bound the performance on each subdistribution, we can bound the model's performance on any possible distribution over valid input sequences.In the brute force proof in Section 4.1, we run the model on the entirety of D. This operation is straightforward when X is finite. Perhaps surprisingly, we can do this even if X is infinite as long as the PDF L × X → R of D is computable and the natural computational topology of X is compact [16, 15, 14] , because integration of computable functions on computable reals is computable [49] .Building on the point from Appendix A.8, it is also relatively straightforward to extend our approach from bounding expected accuracy to bounding log-loss. We will see in Figure 12 that the accuracy and log-loss share a subterm ∆ℓ i . Since we compute this subterm in all of our algorithms, we can easily extend our approach to log-loss by combining ∆ℓ i directly rather than merely checking that the value is negative as we currently do in RELAXED-CORRECTNESS-PESSIMIZING-OVER-POSITION in Algorithm 3. Although this is sufficient for the brute-force and cubic proofs, for the subcubic proof using Algorithm 6 in Appendix F, we would additionally have to compute a log-loss bound for the sequences where the largest non-max token is "too close" to the max token, which we currently neglect by considering the model to get them wrong in the worst case.In this work, we focus on proving lower bounds on model performance. Most of our theorems, for example in Appendices E, and F, prove two-sided bounds. Most of the other theorems can be straightforwardly adapted to proving upper bounds by swapping uses of min and max. Therefore, we expect that proving upper bounds on model performance should be straightforward.A.13 What proof system?Length of proof depends on what proof system we use. We permit any proof system where proofchecking time is linear in the length of the proof. This excludes dependently typed proof systems such as Martin-Löf type theory, but such proof systems can easily be accommodated by considering a proof-checking-trace rather than the proof object itself. Alternatively, a more conventional proof system like ZF, ZFC, or the proof system underlying Isabelle/HOL should suffice.To train each model, we generate 384,000 random sequences of 4 integers picked uniformly at random, corresponding to less than 2.5% of the input distribution. We use AdamW with batch_size = 128, lr = 0.001, betas = (0.9, 0.999), weight_decay left at the default 0.01. We train for 1 epoch (3000 steps). Over our 151 seeds, models trained with this procedure achieve (99.92 ± 0.15) % train accuracy and a loss of (4 ± 8) × 10 -3 . 14 When qualitatively examining a single model (for example in Section 3.1 or Appendix H.1), we use the model with config seed 123, model seed 15 613947648.As our models as sufficiently small, we did not have to use any GPUs to accelerate training our inference. Each training run takes less than a single CPU-hour to complete. In total, the experiments in this paper took less than 1000 CPU-hours.We use the following software packages in our work: Paszke et al. [41] , Plotly Technologies Inc.[42], Nanda and Bloom [37] , Rogozhnikov [47] , Virtanen et al. [52] , McKinney [33] , Waskom [55] B.2 Additional details supporting our mechanistic interpretation of the modelWe provide heatmaps of the matrices corresponding to the five components described/defined in Section 3, for the mainline model. 14 Numbers reported as mean across training runs ± std dev across training runs of mean accuracy and loss. 15 The model seed is deterministically pseudorandomly derived from the seed 123. They don't matter that much, though, being so small. Additionally, the vertical banding indicates that the primary effect of this is a largely-query-independent bias towards larger numbers, reflecting the fact that the input distribution is biased towards larger numbers being the maximum. The weak diagonal pattern indicates a slight bias towards upweighting the query token itself as a (possible) maximum token.We provide some analysis of the distribution of the mechanisms of the models trained on the same configuration. At a glance, there is not that much variation across models.The statistics of interest are: (1) σ 1 /σ 2 , the ratio of the first two singular values of EQKE, a measure of the extent to which the attention score computation is low-rank; (2) s, the average score (accuracy) of the model across the entire input distribution; (3) b cubic /s, the percent-score-recovered accuracy bound achieved by the cubic proof from Section 4.2; (4) b subcubic /s, the percent-score-recovered accuracy bound achieved by the (per-model best) 16 subcubic proof from Section 4.3.For each statistic of interest, Table 6 presents an eleven-number summary of the statistic. Plots, seeds, and statistic values are shown for models whose values are closest to each of the corresponding summary statistics. 17 Additionally, each group contains a boxplot of the summary:• the minimum, maximum; the first and third quartiles; the median and mean; percentiles 2.15 %, 97.85 %, 8.87 %, and 91.13 %; these are displayed as:• top and bottom of the vertical whisker lines; top and bottom of the box; horizontal line inside the box, and the square; horizontal whisker lines and whisker crosshatches. We provide a detailed breakdown of the mathematical notation used in the appendix. be 1 d vocab ⊗ P query , the broadcasting of P query P, Ē, E q be P -P, E + P, and E + P q respectively For any vector-valued function v of length d vocab , parameterized over the input sequence t, let ∆v i := v i -v tmax be the difference between the i th element of v and the element of v corresponding to the true maximum of the input sequence t max . For our particular model, we will have d vocab := 64 d = d model := 32 n ctx := 4 X := (N <d vocab ) nctx ∼ = 64 4 L := N <d vocab ∼ = 64 D| X := U (0, 1, . . . , d vocab -1) nctx , the uniform distributionWe define the two typical performance functions corresponding to accuracy and log-loss. Note the shared subterm ∆ℓ i .EQUATIONEQUATIONWe present the model definition in four different regroupings to define via underbrace labels various useful quantities: EQUATIONEQUATIONIn this section, we prove formally the result used in Section 4.2, A cubic proof.At its heart, the convexity of softmax 19 is an extension to a simple idea: a weighted average of scalar values is extremized by putting 100% of the weight on an extremal value.Using this simple version of the theorem, however, gives a useless bound of 0% accuracy: if we pay no attention to the maximum of the sequence, of course we're going to get the wrong answer. Since in fact the space of possible weightings we may see in practice is much smaller (finite, in fact, with at most d vocab nctx values), we may look for a more general version of this idea that gives us tighter bounds that still cover the space of possible weightings.The weights are not linearly independently choosable (softmax is non-linear), so extremal values do not necessarily result from putting maximal attention on the worst token. It may be, when trying to find the worst case, that some positions are so dis-preferred that it makes more sense to choose a token that is "less bad" for those positions, if it draws enough attention away from the correct token. See Lemma 3 for details.We thus spend this section characterizing a relaxation of the constraints on weights:1. that contains all actually possible weightings, 2. that is extremized at 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 is extremized by putting 100% of the weight on extremal values:Theorem 2 (Warmup: Extremizing weighted averages). Fix a set of values v i ∈ R. The weighted average is bounded by the extremal values: for any w i such that i w i = 1 and 0 ≤ w i ≤ 1,min i v i ≤ i w i v i ≤ max i v iProof. The proof is simple. We havei w i v i -min i v i = i w i (v i -min j v j ) ≥ 0 and max i v i - i w i v i = i w i (max j v j -v i ) ≥ 0so the result follows.The model computes the true maximum t max when its outputs logits ℓ are such that ∆ℓ t * := ℓ t * -ℓ tmax < 0 for all t * ̸ = t max . 20 As a result, it suffices to lower-bound the proportion of sequences where (an upper bound on) ∆ℓ t * is negative for all t * ̸ = t max . In particular, we will upper-bound the contribution from incorrect tokens t in positions i to the difference ∆ℓ i between incorrect (t * ) and correct (t max ) output tokens ∆ℓ i t * = ℓ i t * -ℓ i tmax . We do this by arguing that the logit difference ∆ℓ t * satisfies a certain notion of convexity over the space of a relaxation of sequences (Theorem 6), and constructing a set of Θ(d vocab 3 n ctx ) "extremal" relaxed sequences where the position and token embedding components of attention are pessimized independently.We start by first rewriting the contribution of each token through the attention head to the logit difference into the contributions involving PVOU and EVOU:∆ℓ t t * (t) = ∆ℓ PVOU,i t * (t) + ∆ℓ EVOU,i t * (t)We then upper bound ∆ℓ PVOU,i t * (t) by noting that because the softmax attention is a weighted average of PVOU,∆ℓ PVOU,i t * (t) = ℓ PVOU,i (t) t * -ℓ PVOU,i (t) maxj tj = α * i (t)PVOU i,t * -α * i (t)PVOU i,maxj tj = α * i (t) PVOU i,t * -PVOU i,maxj tj ≤ α * i (t) max i PVOU i,t * -PVOU i,maxj tj Since i α * i (t) = 1, we have nctx-1 i=0 ∆ℓ PVOU,i t * (t) ≤ max i PVOU i,t * -PVOU i,maxj tjWe then construct a set Ξ pure of "pure sequences" consisting of only three types of tokens in one of two orders, and show that for each input sequence t and readoff logit t * , we bound the logit difference from the token embeddings ∆ℓ EVOU,i t * (t) using a small subset X of Ξ pure :nctx-1 i=0 ∆ℓ EVOU,i t * (t) ≤ max ξ∈X nctx-1 i=0 ∆ℓ EVOU,i t * (ξ)We construct a set X relaxed of relaxed sequences, where each relaxed sequence t relaxed consists of a sequence and a position (t, i), where ∆ℓ t * (t, i) is evaluated by separately considering the positional contribution through attention (that is, the attention weighted PVOU) and the token contribution (that is, the attention-weighted EVOU) and direct contribution (the logit difference through the skip connection EU). Note that i indicates the position that we pay 100% of the attention to for the PVOU contribution.We argue that ∆ℓ t * (t, i) satisfies a certain notion of convexity over mixtures of sequences, such that we can evaluate it only on a set of Θ(d vocab 3 n ctx ) "extremal" sequences in a way that takes O(d vocab 3 n ctx ) total time to bound ∆ℓ t * (t, i) for every possible input sequence. We then use the extremal sequences that the model gets correct to lower bound the proportion of all sequences that the model will get correct. Specifically, we argue that Algorithm 3 provides a valid lower bound on the proportion of sequences the model gets correct.We now proceed to the main results of this section.Math fact: For each token t * , the logit difference ∆ℓ t * for any sequence t can be decomposed into the direct contribution from the embeds ℓ EU , the attention-weighted position contribution (PVOU), and the attention-weighted token contribution (EVOU). Therefore, it suffices to upper bound each of the three components independently, since summing these upper bounds gives a valid upper bound on the logit difference.We can compute the direct contribution ℓ EU exactly by first computing EU = E q U and then, for each max, subtracting the logit of the max token from each row of the matrix. No theorems needed. For each max token, we can bound the position contribution by its maximum over positions (Theorem 6).In order to upper bound the token contribution, we argue that any mixed sequence will be upper bounded by the maximum of the corresponding pure sequences (Theorem 7). We then argue that for pure sequences, it suffices to consider orderings where same tokens appear contiguously (Theorem 4).For this subsection, all theorems are parameterized over the following quantities. In practice, we'll take, for fixed query token t query , fixed output token of interest t * , and fixed maximum token t max , v ti e at i +b σ(i) 0≤i<nctx e at i +b σ(i)v t = EVOU t,We will drop the token subscript, writing only s σ , when the token values are unambiguous by context.The sequence score here will be computing ∆ℓ EVOU t * for some fixed t * and t max . The way we've set up our definitions, high scores predict t * (and are thus bad), negative scores predict t max (and are thus good), and more negative the scores, the stronger the prediction of t max . Definition 4 (swap permutation). Given a permutation σ : N <nctx → N <nctx of the n ctx positions and two indices 0 ≤ i, j < n ctx , define the swap permutation σ i↔j to be the permutation that is σ except swapping i and j:σ i↔j (k) =    σ(i) if k = j σ(j) if k = i σ(k) otherwiseDefine ∆ σ,i↔j to be the difference in sequence scores when you swap i and j: ∆ σ,i↔j := s σi↔j -s σ Lemma 3 (Characterization of swapping tokens). Fix a non-decreasing sequence of tokens t 0 ≤ • • • ≤ t nctx-1 ∈ N. Fix σ : N → N be a permutation of the n ctx positions. Fix indices 0 ≤ i, j < n ctx .Then there are two cases for sign (∆ σ,i↔j ):1. If a ti = a tj then sign (∆ σ,i↔j ) = -sign b σ(i) -b σ(j) sign v ti -v tj . 2. Otherwise, sign (∆ σ,i↔j ) = sign a ti -a tj sign b σ(i) -b σ(j) sign s σ - vt i ea t i -vt j e a t j e a t i -e a t j .Intuitively, Lemma 3 says that, if the token contribution to attention is equal between tokens t i and t j , then the impact of swapping their positions σ(i) and σ(j) is entirely determined by how much attention is paid to the positions of i and j and the relative difference in their value. (Notably, by swapping these tokens, we don't affect the attention paid on other tokens, and so the effect of the change does not depend on the values of the other tokens.) Alternatively, if the attentions are not equal, then swapping the positions changes the allocation of attention to other tokens in the sequence, and so it may the case that this change in allocation in attention dominates the attention-weighted values of these two tokens.Proof. First note that the theorem is trivial for i = j.For the rest of the proof, we take i ̸ = j.The proof proceeds just by algebraic manipulation with no deep insight. We first list the facts we use, the proceed to computing sign (∆ σ,i↔j ). We abbreviate σ i↔j as σ ′ for brevity.sign e b σ(i) -e b σ(j) = sign b σ(i) -b σ(j) sign (∆ σ,i↔j ) = sign (s σ ′ -s σ ) = sign 0≤p<nctx v tp e at p +b σ ′ (p) 0≤p<nctx e at p +b σ ′ (p)-s σ Now multiply through by the denominator, which is positive= sign   0≤p<nctx v tp e at p +b σ ′ (p) -s σ 0≤p<nctx e at p +b σ ′ (p)   = sign   0≤p<nctxv tp e at p +b σ(p) -v ti e at i e b σ(i) -e b σ ′ (i) -v tj e at j e b σ(j) -e b σ ′ (j)-s σ 0≤p<nctx e at p +b σ(p) + s σ e at i e b σ(i) -e b σ ′ (i) + s σ e at j e b σ(j) -e b σ ′ (j)  = sign   $ $ $ $ $ $ $ $ $ $ 0≤p<nctxv tp e at p +b σ(p) -v ti e at i e b σ(i) -e b σ(j) -v tj e at j e b σ(j) -e b σ(i)- $ $ $ $ $ $ $ $ $ $ 0≤p<nctxv tp e at p +b σ(p) + s σ e at i e b σ(i) -e b σ(j) + s σ e at j e b σ(j) -e b σ(i)   = sign v tj e at j -v ti e at i e b σ(i) -e b σ(j) + s σ (e at i -e at j ) e b σ(i) -e b σ(j)= sign e b σ(i) -e b σ(j) sign v tj e at j -v ti e at i + s σ (e at i -e at j ) = sign b σ(i) -b σ(j) sign s σ (e at i -e at j ) -v ti e at i -v tj e at j Divide through by non-zero values when possible= sign b σ(i) -b σ(j) •    sign v ti -v tj if a ti = a tjsign (e at i -e at j ) sign s σ -vt i e a t i -vt j e a t j e a t i -ea t j otherwise =    -sign b σ(i) -b σ(j) sign v ti -v tj if a ti = a tj sign a ti -a tj sign b σ(i) -b σ(j) sign s σ - vt i e a t i -vt j e a t je a t i -e a t j otherwise Definition 5 (σ fixes F ). Fix a set of fixed indices F ⊆ N <nctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . (F is the set of positions for which we are not pessimizing over the value of the token in that position.) Fix a non-decreasing sequence of tokenst 0 ≤ • • • ≤ t nctx-1 ∈ N.Given a permutation σ : N <nctx → N nctx , say that σ fixes F (relative to t 0 , . . . , t nctx-1 ) ift i = t F (σ(i)) whenever σ(i) ∈ F .Note that in this section, for the cubic proofs, we will in fact generally take F = {n ctx -1}, so that we are fixing the final query token, though in Theorems 7, 8, and 9 F will also contain all positions with the maximum token t max . In Appendix F, we will take F = ∅ or F to be the set of positions of the maximum token. However, none of these theorems are specific to F being subsingleton, and we prove them in generality. Definition 6 (position-sorting permutation). Fix a set of fixed indices F ⊆ N <nctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab .Define the position-sorting permutation fixing indices in F σ s : N <nctx → N <nctx to be the permutation that sorts the indices not in F according to b: for 0 ≤ i, j < n ctx with i, j ̸ ∈ F , b i ≤ b j whenever σ s (i) < σ s (j); and σ s (i) = i for i ∈ F . Definition 7 (contiguous on equal tokens). Fix a set of fixed indices F ⊆ N <nctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a non-decreasing sequence of tokenst 0 ≤ • • • ≤ t nctx-1 ∈ N.Say that the sequence represented by a permutation σ : N <nctx → N <nctx is contiguous on equal tokens if, for all 0 ≤ i, j, k < n ctx with t i = t j ̸ = t k and i, j, k ̸ ∈ σ -1 (F ), it is never the case thatσ s (σ(i)) < σ s (σ(k)) < σ s (σ(j)).Theorem 4 (Pessimization over sequence ordering is possible and results in contiguous sequences). Fix a set of fixed indices F ⊆ N <nctx and an assignment of token values to each of the fixed positionst F : F → N <d vocab . Fix a non-decreasing sequence of tokens t 0 ≤ • • • ≤ t nctx-1 ∈ N.Let σ min , σ max : N → N be permutations of the n ctx positions, fixing positions in F , satisfying the following property: For all σ : N → N a permutation fixing F , we haveEQUATION(Such permutations are guaranteed to exist because the permutation group on n ctx elements is finite.)Then σ max and σ min may be taken to be contiguous on equal tokens. That is, there exist σ max and σ min satisfying the property of Equation 10which additionally satisfy the definition of Definition 7.The basic idea is that we will assume that one of σ max and σ min cannot be contiguous on equal tokens and derive a contradiction. We will pick the extremal permutation that is closest to being contiguous, take a contiguity violation, and then show that either we can correct the contiguity violation without changing the score-thus violating the presumption that the permutation is closest to being contiguous-or we will find one swap of indices that decreases the score and another swap of indices that increases the score, thus violating the presumption of extremality.In slightly more detail, but still informally, we will consider the sign of the difference between scores of our purported extremal permutation and a permutation that has swapped some indices. The theorem follows from showing that there exists a triple of indices i, j, k such that the sign of the score difference from swapping i and j is different from the sign of the score difference from swapping j and k.First, a definition and some helpful facts about it.Definition 8 (contiguous on equally-attended positions). Fix a set of fixed indices F ⊆ N <nctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a non-decreasing sequence of tokenst 0 ≤ • • • ≤ t nctx-1 ∈ N.Say that a permutation σ is contiguous on equally-attended positions if, for all 0 ≤ i < n ctx with i ̸ ∈ σ -1 (F ), the sorting order according σ s on the contiguous block of positions with contribution to the attention score equal to that of σ(i), σ(j) b σ(j) = b σ(i) and σ(j) ̸ ∈ F , is the same as the sorting order according to the fraction of tokens equal to t j with b-values greater than b σ(i) , with ties broken by the value of t j . Equationally, this second sorting order is defined by the scorek t k = t j and b σ(k) > b σ(i) and σ(k) ̸ ∈ F + t j d vocab |{k | t k = t j and σ(k) ̸ ∈ F }| .Most importantly, any permutation that is contiguous on equally-attended positions has the property that for any indices 0 ≤ i, j, k < n ctx with i, j, k ̸ ∈ σ -1 (F ) andt i = t j ̸ = t k and σ s (σ(i)) < σ s (σ(k)) < σ s (σ(j)), we will have the strict inequality b σ(i) < b σ(k) < b σ(j). Additionally, we may always sort equally-attended positions to make any permutation contiguous on equally-attended positions.We will define an additional notion of contiguity-violations which we avoid up-front by arbitrarily swapping involved indices without changing the score s σ .Definition 9 (needlessly non-contiguous). Fix a set of fixed indices F ⊆ N <nctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a non-decreasing sequence of tokenst 0 ≤ • • • ≤ t nctx-1 ∈ N.Say that a permutation σ is needlessly non-contiguous at i, j, k (for i, j, k̸ ∈ σ -1 (F )) if ∆ σ,i↔k = 0 or ∆ σ,j↔k = 0, for 0 ≤ i, j, k < n ctx with i, j, k ̸ ∈ σ -1 (F ) with t i = t j ̸ = t k and σ s (σ(i)) < σ s (σ(k)) < σ s (σ(j)).Say that a permutation σ is needlessly non-contiguous if it is needlessly non-contiguous at any i, j, k ̸ ∈ σ -1 (F ).Lemma 5. Fix a set of fixed indices F ⊆ N <nctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a non-decreasing sequence of tokenst 0 ≤ • • • ≤ t nctx-1 ∈ N.Any needlessly non-contiguous sequence σ which fixes F can be made into a sequence σ ′ which still fixes F and is both simultaneously contiguous on equally-attended positions and not needlessly non-contiguous, and for which s σ = s σ ′ .Proof. First, sort regions of equally-attended positions to make σ contiguous on equally-attended positions. If the resulting permutation is not needlessly non-contiguous, then we are done.Otherwise, we have ∆ σ,i↔k = 0 or ∆ σ,j↔k = 0 for some i, j, k, for 0 ≤ i, j, k < n ctx with i, j, n ̸ ∈ σ -1 (F ) and t i = t j ̸ = t k and σ s (σ(i)) < σ s (σ(k)) < σ s (σ(j)). Since the sequence is contiguous on equally-attended positions, we have the strict inequality b σ(i) < b σ(k) < b σ(j) .By Lemma 3, we have two cases. Noting that t i = t j , we can write them as1. v t k = v ti and a ti = a t k 2. a ti ̸ = a t k and s σ = vt i e a t i -vt k e a t k e a t i -e a t kIn the first case, we may fully freely interchange tokens equal to t i with tokens equal to t k without changing the score; in this case we may use the token value as a sorting tie-breaker and swap tokens until there are no more needlessly non-contiguous triples falling into case (1).In the second case, since swapping tokens does not change s σ , the property will continue to hold for these tokens after the swap. We may then swap tokens, again using token value as a tie-breaker, until there are no more needlessly non-contiguous triples falling into case (2) .We can now finally make our argument for Theorem 4 more precise.Proof of Theorem 4. Choose σ max and σ min to be contiguous on equally-attended positions and not needlessly non-contiguous, and suppose that we have σ ∈ {σ max , σ min } such that for some 0 ≤ i, j, k < n ctx with i, j, k ̸ ∈ σ -1 (F ) andt i = t j ̸ = t k , we have b σ(i) < b σ(k) < b σ(j). We will derive a contradiction with the presumption that σ is extremal by showing that we can swap i and k to change the score in one direction and that we can swap j and k to change the score in the other direction.Take σ ′ 0 to be σ but swapping i and k, and take σ ′ 1 to be σ but swapping j and k. Now we will consider the cases for the sign of the score difference ∆ 0 := s σ ′ 0 -s σ and ∆ 1 := s σ ′ 1 -s σ . By the presumption of not being needlessly non-contiguous, ∆ z ̸ = 0 for z ∈ {0, 1}. If we can show that the sign of ∆ 0 is distinct from the sign of ∆ 1 , then we will have a contradiction with extremality because we will have eithers σ ′ 0 < s σ < s σ ′ 1 or s σ ′ 1 < s σ < s σ ′ 0 .That is, we would be able to swap i ↔ k and j ↔ k to get a lower and higher score, making σ not extremal. Note that the proof of Theorem 4 does not go through if we include the position value function w in the score, because we may trade off the position value function against the token value function. We now show that we can independently pessimize over positional attention. Definition 10 (full sequence score). Given a non-decreasing sequence of tokens t 0 ≤ • • • ≤ t nctx-1 ∈ N <d vocab and a permutation σ : N <nctx → N <nctx define the full sequence score s ′ t0,...,tn ctx -1 ,σ as:Noting that t i = t j , sign (∆ 0 ) = sign b σ(i) -b σ(k) sign (v t k -v ti ) if a ti = a t k sign (a ti -a t k ) sign s σ - vt i e a t i -vt k e a t k e a t i -e a t k otherwise sign (∆ 1 ) = sign b σ(j) -b σ(k) sign (v t k -v ti ) if a ti = a t k sign (a ti -a t k ) sign s σ - vt i e a t i -vt k e a t ks ′ t0,...,tn ctx -1 ,σ := 0≤i<nctx (v ti + w σ(i) )e at i +b σ(i) 0≤i<nctx e at i +b σ(i)We will drop the token subscript, writing only s ′ σ , when the token values are unambiguous by context.The sequence score here will be computing ∆ℓ EPVOU t * := ∆ℓ EVOU t * + ∆ℓ PVOU t * for some fixed t * and t max . As with Definition 3, with the way we've set up our definitions, high scores predict t * (and are thus bad), negative scores predict t max (and are thus good), and more negative the scores, the stronger the prediction of t max . Definition 11 (relaxed sequence score). Given a non-decreasing sequence of tokens t 0 ≤ • • • ≤ t nctx-1 ∈ N <d vocab and a permutation σ : N <nctx → N <nctx define the relaxed sequence scores r t0,...,tn ctx -1 ,σ,min and r t0,...,tn ctx -1,σ,max as: r t0,...,tn ctx -1 ,σ,min := s t0,...,tn ctx -1,σ + min 0≤i<nctx w i r t0,...,tn ctx -1,σ,max := s t0,...,tn ctx -1,σ + max 0≤i<nctx w iWe will drop the token subscript, writing only r σ,min or r σ,max , when the token values are unambiguous by context.0 ≤ • • • ≤ t nctx-1 ∈ N and t ′ 0 ≤ • • • ≤ t ′ nctx-1∈ N and permutations σ, σ ′ : N <nctx → N <nctx . Let r σ,min and r σ,max denote r t0,...,tn ctx -1 ,σ,min and r t0,...,tn ctx -1 ,σ,max ; let s σ denote s t0,...,tn ctx -1,σ ; and let s σ ′ and s ′ σ ′ denote s t ′ 0 ,...,t ′ n ctx -1 ,σ ′ and s ′ t ′ 0 ,...,t ′ n ctx -1 ,σ ′ . Then we havemin 0≤i<nctx w i = r σ,min -s σ ≤ s ′ σ ′ -s σ ′ ≤ r σ,max -s σ = max 0≤i<nctx w i .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 haver σ,min -s σ = min 0≤i<nctx w σ(i) = min 0≤i<nctx w i r σ,max -s σ = max 0≤i<nctx w σ(i) = max 0≤i<nctx w i s ′ σ ′ -s σ ′ = 0≤i<nctx w σ ′ (i) e a t ′ i +b σ ′ (i) 0≤i<nctx e a t ′ i +b σ ′ (i) = 0≤i<nctx w σ ′ (i) e a t ′ i +b σ ′ (i) 0≤j<nctx e a t ′ j +b σ ′ (j)Since e x is non-negative for all real x, we havemin 0≤j<nctx w σ ′ (j) 0≤i<nctx e a t ′ i +b σ ′ (i) 0≤j<nctx e a t ′ j +b σ ′ (j) ≤ s ′ σ ′ -s σ ′ ≤ max 0≤j<nctx w σ ′ (j) 0≤i<nctx e a t ′ i +b σ ′ (i) 0≤j<nctx e a t ′ j +b σ ′ (j).Thus we get as desired min0≤i<nctx w i ≤ s ′ σ ′ -s σ ′ ≤ max 0≤i<nctx w i .Note that we could prove a more fine-grained theorem, that pessimizes over attention paid to positions only for sequences compatible with the chosen fixed tokens F and t F , but since the positional contribution is so small we do not bother. Theorem 7 (For a fixed ordering, softmax is convex over token counts and only pure sequences need be considered). Fix a set of fixed indices F ⊆ N <nctx and an assignment of token values to each of the fixed positions t F : F → N <d vocab . Fix a set S ⊆ N <d vocab of valid other tokens in the sequence. (In our uses of this theorem, S will be the largest subset of N <tmax for which we can guarantee that the model behaves correctly on all sequences compatible with F and t max and with tokens otherwise drawn from S.)Define a comparison on non-negative integers less than d vocab :c := i∈F v t F (i) e a t F (i) +bi d := i∈F e a t F (i) +bi f := 0≤i<nctx i̸ ∈F e bicmp(x, y) := sign d(e ax v x -e ay v y ) -c(e ax -e ay ) + f e ax+ay v x e ax+ay -v y e ax+ay Let t cmp min and t cmp max be the minimum and maximum elements of S according to cmp. 21For a given choice of a non-decreasing sequence of tokens t 0 ≤ • • • ≤ t nctx-1 ∈ N compatible with F and S and a given choice of permutation σ : N → N of the n ctx positions fixingF (t i = t F (σ(i))for σ(i) ∈ F ; and t i ∈ S for σ(i) ̸ ∈ F ): let s σ,min (and s σ,max ) denote s t0,...,tn ctx -1,σ when t i = t cmp min for all σ(i) ̸ ∈ F (or t cmp max , respectively).Then for all such choices of sequence-permutation pairs, s σ,min ≤ s t0,...,tn ctx -1,σ ≤ s σ,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 7.For any indices 0 ≤ i < j < n ctx , token values x, y ∈ S, the score for a sequence with t i = x ̸ = y = t j is bounded on both sides by sequences with t i = t j = x and t i = t j = y.Proof. Let s α,β be the sequence score with t i = α and t j = β, and define the score differences ∆ x := s x,x -s x,y and ∆ y := s y,y -s x,y . It suffices to show that sign(∆ x ∆ y ) ≤ 0. To show this, we must only compute the sign of ∆ α for α ∈ {x, y} and show that whenever both ∆ x and ∆ y are non-zero, they have opposite signs.We proceed by computation after defining some convenience variables for brevity:C := 0≤k<nctx k̸ =i,j v t k e at k +b σ(k) D := 0≤k<nctx k̸ =i,j e at k +b σ(k) α := x if α = y y if α = x i α := i if α = x j if α = y i α := i if α = x j if α = y sign (∆ α ) = signv α e aα+b σ(i) + v α e aα+b σ(j) + C e aα+b σ(i) + e aα+b σ(j) + D -v x e ax+b σ(i) + v y e ay+b σ(j) + C e ax+b σ(i) + e ay+b σ(j) + D = sign v α e aα+b σ(iα ) + v α e aα+b σ(i α ) + C e aα+b σ(iα ) + e aα+b σ(i α) + D -v α e aα+b σ(iα ) + v αe a α +b σ(i α ) + C e aα+b σ(iα ) + e a α +b σ(i α) + DMultiply through by positive denominators and simplify= sign C e a α +b σ(i α) -e b σ(i α ) +aα + D v α e b σ(i α ) +aα -v αe a α+b σ(i α ) + v α e b σ(i α ) + e b σ(iα ) e a α+b σ(i α ) +aα -v α e b σ(i α ) + e b σ(iα) e a α +b σ(i α ) +aαPulling out e b σ(i α )= sign e a α +aα e b σ(i α) + e b σ(iα) (v α -v α) + C (e a α -e aα ) + D (e aα v α -e a α v α)Note that swapping α and α negates the sign. Hence, we have sign(∆ x ) = -sign(∆ y ) and hence s x,x ≤ s x,y ≤ s y,y or s y,y ≤ s x,y ≤ s x,x as desired.Lemma 9 (Pure sequences are sorted according to cmp in Theorem 7) . Fix all the same quantities as in Theorem 7.Fix tokens x, y ∈ S. Let n := n ctx -|F | be the number of non-fixed tokens. Fix sequences with n copies of x and y respectively: fixt x,0 ≤ • • • ≤ t x,nctx-1 ∈ N and t y,0 ≤ • • • ≤ t y,nctx-1 ∈ Ncompatible with F and S and given choices of permutations σ x , σ y : N → N of the n ctx positions fixingF : t x,i = t F (σ x (i)) for σ x (i) ∈ F ; t y,i = t F (σ y (i)) for σ y (i) ∈ F ; t x,i = x for σ x (i) ̸ ∈ F ; and t y,i = y for σ y (i) ̸ ∈ F .Then sign((s σx,tx,0,...,tx,n ctx -1 ) -(s σy,ty,0,...,ty,n ctx -1 )) = cmp(x, y) Proof. Note that by Lemma 9, cmp is comparing two sequence scores. Since ≤ is transitive over the reals, the relation ≤ cmp is also transitive.EQKE(t -1 , t i ) := t -1 E q QK T ĒT t i T / √ d EQKP(t -1 , i) := t -1 E q QK T PT i / √ d EVOU(t i ) := t i ĒV OU PVOU(i) := Pi V OU ℓ EU (t -1 ) := t -1 E q U ∆ℓ EU t * (t -1 , max i t i ) := ℓ EU (t -1 ) t * -ℓ EU (t -1 ) maxi tiFinally, we combine the previous lemmas to complete our proof of Theorem 7:Proof of Theorem 7. Extremal sequences with scores s σ,min and s σ,max are guaranteed to exist because there are only finitely many elements of S and therefore only finitely many sequences. By Lemma 8, the extremal sequences must be pure (have t i = t j whenever σ(i), σ(j) ̸ ∈ F ). By Lemma 9, the extremal sequences must have tokens that are extremal according to cmp.We now have all the tools necessary to prove the following theorem. We refer to Algorithm 3 and Algorithm 4 or the proof of Theorem 11 for a definition of the CUBIC algorithm. Theorem 11.E t∼U (0,1,...,d vocab -1) n ctx argmax i (M(t)) i = max i t i ≥ CUBIC(d vocab , n ctx , M)Before we give the proof of this theorem, we introduce some helpful notation. Definition 12. Fix an element (r m , r q , c) ∈ {0, ..., d vocab } 2 × {0, ...3} such that r m ≥ r q . We define X (rm,rq,c) to be the set of tokens t such that return MODEL-BEHAVIOR(M, input-sequence) == MAX(input-sequence) 3: end function 4: function MODEL-BEHAVIOR(M, input-sequence) Require: input-sequence is a tensor of shape (n ctx , ) with values in N <d vocab 5:t max ← MAX(input-sequence) ▷ t max ← max-token 6:t ← input-sequence return ALL(CORRECTNESS(M, perm + [t -1 ]) for all perm ∈ PERMUTATIONS(t 0:-1 )) 19: end function 2. The query token t query is equal to r q , 3. The cardinality of tokens that are not at the query position and not equal to t max is equal to c.For clarity, we list all the possible cases. We always take t query ≤ t max and let S 3 act on sequences by permuting the first three factors (i.e. keeping the query position fixed).1. If c = 0, then X (tmax,tquery,0) = {[t max , t max , t max , t query ]},2. If c = 1, then X (tmax,tquery,1) = S 3 .{[t 1 , t max , t max , t query ] | t 1 < t max }, 3. If c = 2, then X (tmax,tquery,2) = S 3 .{[t 1 , t 2 , t max , t query ] | t i < t max }, 4. If c = 3, then X (tmax,tmax,3) = S 3 .{[t 1 , t 2 , t 3 , t max ] | t i < t max }. Definition 13.Let t ∈ X be a sequence. We say t is pure, if it has at most three distinct tokens: the max token t max , the query token t query , and optionally a third token t * < t max .We denote by X pure the subset of pure tokens. For any subset Y ⊂ X, we set Y pure := Y ∩ X pure .We now come to the proof of Theorem 11. We will show how to use the previous theorems to get explicit bounds and explain how CUBIC(d vocab , n ctx , M) computes these bounds.Proof of Theorem 11. First of all, we note that the algorithm CUBIC = CUBIC(d vocab , n ctx , M) yields a lower bound for the accuracy on the set X (tmax,tquery,c) . We can therefore compute the bound on X = (tmax,tquery,c) X (tmax,tquery,c) by computing it for each such choice (t max , t query , c) and summing over themE t∼U (0,1,...,d vocab -1) n ctx argmax i (M(t)) i = max i t i ≥ (tmax,tquery,c)CUBIC(X (tmax,tquery,c) ).So from now on we will fix one such subset X (tmax,tquery,c) .We begin by defining a map f : X (tmax,tquery,c) → {0, ..., d vocab } c which sends a sequence to the subsequence of elements which are not at the query position and not equal to t max . Then Theorem 7 can be restated as follows 22 :Algorithm 3 Counting Correct Sequences in Cubic Time, Part I. Lines are annotated with comments indicating the parameters for a cache to avoid duplicate computations.1: function MODEL-BEHAVIOR-RELAXED(M, query-tok, max-tok, non-max-tok, n-copiesnonmax)2: t query ← query-tok, t max ← max-tok, t ′ ← non-max-tok, c ← n-copies-nonmax Require: 0 ≤ t query ≤ t max < d vocab , 0 ≤ t ′ ≤ t max < d vocab , 0 ≤ c < n ctx Require: if n-copies-nonmax = 0 then non-max-tok = max-tok Require: if query-tok ̸ = max-tok then n-copies-nonmax < n ctx -1 Ensure: return ≥ MODEL-BEHAVIOR(M, tEQUATIONattn-weights-unscaled :,nctx-1 ← a tquery + b nctx-1 ▷ Cache by t query 12:attn-weights-unscaled if c = 0 then ▷ In this case, attn-weights 0,i = attn-weights 1,i , so we drop the first subscript F Quadratic counting for a sub-cubic proof0,i ← a tmax + b 0,i for 0 ≤ i < n ctx -c -1 ▷ Cache by t max , c, i, t query 13: attn-weights-unscaled 1,i ← a tmax + b 1,i for 0 ≤ i < n ctx -c -1 ▷ Cache by t max , c, i, t query 14: attn-weights-unscaled 0,i ← a t ′ + b 0,i for n ctx -c -1 ≤ i < n ctx -1 ▷ Cache by t ′ , c, i, t query 15: attn-weights-unscaled 1,i ← a t ′ + b 1,i for n ctx -c -1 ≤ i < n ctx -1 ▷ Cache by t ′ ,∆v i,t * ← ∆v tmax,t * for 0 ≤ i < n ctx -c -1 22: ∆v i,t * ← ∆v t ′ ,t * for n ctx -c -1 ≤ i < n ctx -1 23: ∆v nctx-In this section we fill in the details lacking from Section 4.3.In Appendix E we proved an intricate version of convexity of softmax where, modulo pessimizing in unrealistic ways over the attention paid to positions for the computation done on positional encodings, all extremal relaxed sequences correspond to actual sequences.When we only get a budget of O(d vocab 2 n ctx ) extremal relaxed cases to consider, though, we must pessimize more, which gives us a simpler version of the convexity theorem and proof. Notably, when we restrict our sequences to have only two tokens (the max token t max and the non-max token t ′ ), most of the theorems from Appendix E.3 get significantly simpler.Additionally, we must pessimize separately over the token value (v) and token attention (b) computations in order to allow efficient computation (Theorem 15).For this subsection, all theorems are parameterized over the following quantities. In practice, as in Appendix E.3, we'll take, for fixed query token t query ,v t = EVOU t,t * -EVOU t,tmax a t = EQKE tquery,t / √ d w i = PVOU i,t * -PVOU i,tmax b i = EQKP tquery,i / √ dNote that unlike in Appendix E.3, we pessimize independently over the query token and the non-max token, so the "fixed" query token may not in fact appear in any key-side position in the relaxed sequence we consider. Definition 15 (of a sequence via mapping from positions). We can define a sequence of tokens via mapping from positions by specifying a subset of valid tokens S ⊆ N <d vocab paired with a function T : N <nctx → S specifying which token is in each position. Definition 16 (sequence score). Given a subset of valid tokens S ⊆ N <d vocab and a function T : N <nctx → S specifying which token is in each position, define the sequence scores T := 0≤i<nctx v T (i) e a T (i) +bi 0≤i<nctx e a T (i) +biDefinition 17 (swapped mapping). Given a subset of valid tokens S ⊆ N <d vocab and a function T : N <nctx → S specifying which token is in each position and two indices 0 ≤ i, j < n ctx , define the swapped mapping T i↔j be the function that is T except swapping i and j:T i↔j (k) =    T (i) if k = j T (j) if k = i T (k) otherwiseLemma 13 (Characterization of swapping tokens in a two-token sequence). Fix two tokens t 0 < t 1 ∈ N and a function T : N <nctx → {t 0 , t 1 } specifying which token is in each position.Define ∆ T,i↔j to be the difference in sequence scores when you swap i and j:∆ T,i↔j := s Ti↔j -s T Then sign (∆ T,i↔j ) = -sign (b i -b j ) sign v T (i) -v T (j)Proof. Lemma 3 gives us the result directly when a T (i) = a T (j) . Otherwise, we getsign (∆ T,i↔j ) = sign a T (i) -a T (j) sign (b i -b j ) sign s T - v T (i) e a T (i) -v T (j) e a T (j) e a T (i) -e a T (j)Hence all that remains is to show thatsign s T (e a T (i) -e a T (j) ) -v T (i) e a T (i) + v T (j) e a T (j) = -sign v T (i) -v T (j) Define v := 1 2 (v T (i) + v T (j)) and define ∆v := 1 2 (v T (i) -v T (j) ) so that v T (i) = v + ∆v and v T (j) = v -∆v. Assume WLOG that T (i) = 0 and T (j) = 1 so that v T (p) = v + (-1) T (p) ∆v for all p.Then we have sign s T (e a T (i) -e a T (j) ) -v T (i) e a T (i) + v T (j) e a T (j) = sign (s T (e a T (i) -e a T (j) ) -v (e a T (i) -e a T (j) ) -∆ve a T (i) -∆ve a T (j) ) = sign T (p) ∆v e a T (p) +bp 0≤p<nctx e a T (p) +bp (e a T (i) -e a T (j) ) -@ @ @ @ @ @ @ @ v (e a T (i) -e a T (j) ) -∆v (e a T (i) + e a T (j) )    0≤p<nctx v T (p) e a T (p) +bp 0≤p<nctx e a T (p) +bp (e a T (i) -e a T (j) ) -v (e a T (i) -e a T (j) ) -∆v (e a T (i) + e a T (j) )     = sign      0≤p<nctx ¡ v + (-1)     = sign(∆v) sign        e a T (i) 0≤p<nctx T (p)=T (i)e bp -e a T (j) 0≤p<nctx T (p)=T (j)e bp 0≤p<nctx e a T (p) +bp (e a T (i) -e a T (j) ) -e a T (i) -e a T (j)       = sign(v T (i) -v T (j) ) sign        e a T (i) 0≤p<nctx T (p)=T (i) e bp -e a T (j) 0≤p<nctx T (p)=T (j)e bp 0≤p<nctx e a T (p) +bp (e a T (i) -e a T (j) ) -e a T (i) -e a T (j)       Define P i := 0≤p<nctx T (p)=T (i) e bp P j := 0≤p<nctx T (p)=T (j)e bp so that we get sign s T (e a T (i) -e a T (j) ) -v T (i) e a T (i) + v T (j) e a T (j)= sign(v T (i) -v T (j) ) sign e a T (i) P i -e a T (j) P j e a T (i) P i + e a T (j) P j (e a T (i) -e a T (j) ) -e a T (i) -e a T (j)Multiply through by the positive denominator and expand out so that we get = sign(v T (i) -v T (j) ) sign -2e a T (i) +a T (j) P i -2e a T (i) +a T (j) P j = -sign(v T (i) -v T (j) ) sign e a T (i) +a T (j) P i + e a T (i) +a T (j) P j = -sign(v T (i) -v T (j) )Theorem 14 (Pessimization over sequence ordering for two-token sequences is simple). Let σ s : N → N denote a permutation of the n ctx positions that sorts them according to b: for 0 ≤ i, j < n ctx , b i ≤ b j whenever σ s (i) < σ s (j). Fix two tokens t 0 < t 1 ∈ N.Let n t0 be the number of p ∈ [0, n ctx ) with T (p) = t 0 and let n 1 be the number of p ∈ [0, n ctx ) with T (p) = t t1 . Note that n t0 + n t1 = n ctx .Define t min := argmin t∈{t0,t1} v t and define t max := argmax t∈{t0,t1} v t .Define T min , T max : N <nctx → {t 0 , t 1 } to be the assignment of tokens to positions that pays the least (respectively, most) attention to t max :T min (i) := t max if 0 ≤ σ s (i) < n tmax t min if n tmax ≤ σ s (i) < n ctx T max (i) := t min if 0 ≤ σ s (i) < n tmin t max if n tmin ≤ σ s (i) < n ctxThen we have that s Tmin ≤ s T ≤ s TmaxProof. The extremality of s Tmin and s Tmax follows straightforwardly from Theorem 4.All that remains is s Tmin ≤ s Tmax .This follows from noting by Lemma 13 that swapping two tokens in T min increases the sequence score, while the reverse is true of s Tmax , thus showing that it must be s Tmin that is the minimum and s Tmax that is the maximum and not vice versa.Definition 18 (full sequence score). Given a subset of valid tokens S ⊆ N <d vocab and a function T : N <nctx → S specifying which token is in each position define the full sequence score s ′ T :s ′ T := 0≤i<nctx (v T (i) + w i )e a T (i) +bi 0≤i<nctx e a T (i) +biTheorem 15 (Independent pessimization over positional contributions and token attention and token value is possible). Fix two tokens t 0 < t 1 ∈ N. Let T min , T max : N <nctx → {t 0 , t 1 } and t max , t min be as in Theorem 14. Fix a set S of valid tokens with t 0 , t 1 ∈ S.Define relaxed versions T ′ max , T ′ min : N <nctx → S of T max and T min :T ′ max (i) := T max (i) if T max (i) = t max argmin j∈S j̸ =tmax a j otherwise T ′ min (i) := T min (i) if T max (i) = t max argmax j∈S j̸ =tmax a j otherwiseThat is, T ′ max replaces t min with whatever token in S draws the least attention away from t max , while T ′ min replaces t min with whichever token in S draws the most attention away from t max .Define relaxed extremal sequence scores r Tmax , r Tmin :r Tmin := min Proof. (sketch) Essentially the same as the proof of Theorem 6.Note that in practice, we take S to be the set of all tokens less than t max -g for some minimum gap g. This allows us to share computation across the various maximum tokens to reduce overall computational complexity.Algorithm 5 Counting Correct Sequences in Subcubic Time, Preliminaries c max ←    0 if t max = 0 n ctx -1 if t max = t query n ctx -2 else ▷ maximum copies of nonmax 20:for c min ≤ c ≤ c max do ▷ valid choices for the number of non-max tokens 21:g ← G tmax,tquery,c 22:EQUATIONq-gap ← t max -t query 24:RCPOG(⃗ χ) ← RELAXED-CORRECTNESS-PESSIMIZING-OVER-GAP(M, ⃗ χ)25:if (q-gap = 0 or q-gap ≥ g) and RCPOG(d vocab , n ctx , t max , t query , c, g, g * ) then 26: attn-weights-unscaled 0,i ← b 0,i + (∆a min if i ∈ idx-set else 0) c ′ ← c if t query = t max else c -1 ▷ # off x,y + g y,z ≥ min y,z [E x f x,y + g y,z ] + min x,y [f x,y -E x f x,y ] max x,y,z f x,y + g y,z ≤ max y,z [E x f x,y + g y,z ] + max x,y [f x,y -E x f x,y ]and the RHSs can be computed in time O(n x n y + n y n z + n x n z ) for n x , n y , and n z the number of possible values of x, y, and z, respectively.Example for how this helps with small variation: Take any function k(y) and then takef x,y := k(y) + ε 1 (x, y) g y,z := -k(y) + ε 2 (y, z) Then we have min x,y,z [f x,y + g y,z ] = min x,y,z [ε 1 (x, y) + ε 2 (y, z)] min x,y f x,y + min y,z g y,z = min y k(y) + min y -k(y) + min x,y ε 1 (x, y) + min y,z ε 2 (y, z) = min y k(y) -max y k(y) + min x,y ε 1 (x, y) + min y,z ε 2 (y, z) min x,y [f x,y -E x f x,y ] + min y,z [g y,z + E x f x,y ] = min x,y ε 1 (x, y) + min y,z [ε 2 (y, z) + E x ε 1 (x, y)]If ε 1 and ε 2 are small compared to min y k(y) -max y k(y), then using E x f x,y gives a much better bound.Note, though, that this could be a worse bound if the assumption of small variation does not hold.Note also that this trick is not restricted to adding and subtracting E x f x,y . If f is a matrix indexed by x and y, we might also try taking SVD and using the first principal component instead. A basic application of the triangle inequality gives the following, more general, result: Theorem 19 (Summarize+Diff). For any h y which can be computed in time O(n h ), min x,y,zf x,y + g y,z ≥ min y,z [h y + g y,z ] + min x,y [f x,y -h y ] max x,y,z f x,y + g y,z ≤ max y,z [h y + g y,z ] + max x,y [f x,y -h y ]and the RHSs can be computed in time O(n x n y + n y n z + n h ) for n x , n y , and n z the number of possible values of x, y, and z, respectively.We see that if the variation of f in the x-axis is indeed much smaller than the variation in the y-axis, then letting f x,y = h y + ε x,y we getmin x,y,z f x,y + g y,z -min y,z [h y + g y,z ] -min x,y [f x,y -h y ] ≤ min x,y,z [f x,y + g y,z ] -min y,z [h y + g y,z ] + min x,y [ε x,y ] ≤ 2 maxx,y |ε x,y | so indeed this bound isn't too much worse and we are able to compute it in quadratic rather than cubic time.As discussed in Section 4.3.1, to further reduce the computation cost of proof, we need to avoid computing the residual stream, EVOU, and EPQKE matrices fully. Using mechanistic insight or otherwise, we observe that these matrices (apart from EVOU) can be well-approximated by rank one matrices. This will remove the dominant computation cost of O(d vocab 2 • d model ).Up to this point, we haven't really said much in our proofs about what the model is doing. All the mechanistic insight has been of the form "the model varies more along this axis than this other axis" or "the input data is distributed such that handling these inputs is more important than handling these other inputs" or, at best, "the model computes the answer by attending to the maximum token of the sequence; everything else is noise".Here, finally, our proof-size constraints are tight enough that we will see something that we could plausibly call "how the model pays attention to the maximum token more than anything else", i.e., (if we squint a bit) "the model pays more attention to larger tokens in general.As stated above, we are breaking matrices into their rank one approximation and some error term. To bound the error, i.e. to bound expressions of the form i (A i + E i )i A i , where E i denote the matrix errors, we can use the following trick: Lemma 20 (Max Row-Diff (vector-matrix version)). For a row vector a and a matrix B,max i,j ((aB) i -(aB) j ) ≤ k |a k | max i,j (B k,i -B k,j )Moreover, for a collection of n row vectors A r , if the shape of B is m × p, the right hand side can be computed for all r in time O(nm + mp).EQUATIONi -(aB) j = max i,j k a k (B k,i -B k,j ) ≤ k max i,j a k (B k,i -B k,j ) = k a k max i,j (B k,i -B k,j ) if a k ≥ 0 min i,j (B k,i -B k,j ) if a k < 0 = k a k max i,j (B k,i -B k,j ) if a k ≥ 0 -max i,j (B k,i -B k,j ) if a k < 0 = k |a k | max i,j (B k,i -B k,j )The asymptotic complexity of computing the result follows from caching the computation of max i,j (B k,i -B k,j ) for each k independently of r, as the computation does not depend on A r .Theorem 21 (Max Row-Diff). For matrices A and B, max r,i,j((AB) r,i -(AB) r,j ) ≤ max r k |A r,k | max i,j (B k,i -B k,j )Proof. By taking the max of Lemma 20 over rows r of A.Lemma 20 can also be applied recursively for a product of more than two matrices. Lemma 22 (Max Row-Diff (vector-matrix recursive version)). For a row vector a and a sequence of n matrices B p of shapes r p × c p ,max i,j   a p B p i -a p B p j   ≤ k0 |a k0 | • • • kn (B n-1 ) kn-1,kn max i,j ((B n ) kn,i -(B n ) kn,j )Moreover, for a collection of q row vectors A α , the right hand side can be computed for all α in time O(qr 0 + p r p c p ).Proof. We proceed by induction on n.For n = 1, the statement is identical to Lemma 20.Suppose the theorem holds for all positive n = s; we show the theorem holds for n = s + 1. We reassociate the matrix multiplication as a k0 (B 1 ) k0,k1 k2 |(B 2 ) k1,k2 | • • • ks+1 (B s ) ks,ks+1 max i,j (B s+1 ) ks+1,i -(B s+1 ) ks+1,jThe triangle inequality gives≤ k1 k0 |a k0 (B 1 ) k0,k1 | k2 |(B 2 ) k1,k2 | • • • ks+1 (B s ) ks,ks+1 max i,j (B s+1 ) ks+1,i -(B s+1 ) ks+1,jand algebra gives= k0 |a k0 | k1 |(B 1 ) k0,k1 | k2 |(B 2 ) k1,k2 | • • • ks+1 (B s ) ks,ks+1 max i,j (B s+1 ) ks+1,i -(B s+1 ) ks+1,jThe asymptotic complexity of computing the right hand side also follows straightforwardly by induction.Theorem 23 (Max Row-Diff (recursive)). For a sequence of n + 1 matrices A 0 , . . . , A n , max r,i,j  p A p r,i - p A p r,j   ≤ max r k0 |(A 0 ) r,k0 | • • • kn (A n-1 ) kn-1,kn max i,j ((A n ) kn,i -(A n ) kn,j )Proof. By taking the max of Lemma 22 over rows r of A 0 .Note that Theorem 21 is compatible with the mean+diff trick of Appendix G.1.Theorem 24 (Combined Mean+Diff and Max Row-Diff). For matrices A and B, and any columnwise summary vector H k of A (for example we may takeH k := E r A r,k ) max r,i,j ((AB) r,i -(AB) r,j ) ≤ max i,j k H k (B k,i -B k,j ) +max r k |A r,k -H k | max i,j (B k,i -B k,j ) Proof. max r,i,j ((AB) r,i -(AB) r,j ) = max r,i,j k A r,k (B k,i -B k,j ) = max r,i,j k (H k + (A r,k -H k )) (B k,i -B k,j ) = max i,j k H k (B k,i -B k,j ) + max r k (A r,k -H k ) (B k,i -B k,j ) ≤ max i,j k H k (B k,i -B k,j ) + max r k max i,j (A r,k -H k ) (B k,i -B k,j ) ≤ max i,j k H k (B k,i -B k,j ) + max r k |A r,k -H k | max i,j (B k,i -B k,j )Theorem 25 (Combined Mean+Diff and Vector-Matrix Recursive Max Row-Diff). For a row vector a, a vector of summaries h corresponding to a (for example, if a is a row of a matrix, h might be the average of the rows), a sequence of n matrices B p of shapes r p × c p , and a corresponding sequence of column-wise summary vectors h p of B p (for example we may take(h p ) k := E r (B p ) r,k ), max i,j   a p B p i -a p B p j   ≤ max i,j k0 h k0 • • • kn (B n-1 ) kn-1,kn ((B n ) kn,i -(B n ) kn,j ) + k0 |a k0 -h k0 | • • • max i,j kn (h n-1 ) kn ((B n ) kn,i -(B n ) kn,j ) + kn (B n-1 ) kn-1,kn -(h n-1 ) kn max i,j ((B n ) kn,i -(B n ) kn,j )Moreover, for a collection of q row vectors A α , the right hand side can be computed for all α in time O(qr 0 + p r p c p ).Proof sketch. Apply the triangle inequality recursively, fusing the proofs of Lemmas 22, and 24.Theorem 26 (Combined Mean+Diff and Recursive Max Row-Diff). For a sequence of n + 1 matrices A 0 , . . . , A n , and corresponding column-wise summary vectors h 0 , . . . , h n-1 of A 0 , . . . ,A n-1 , max r,i,j   p A p r,i - p A p r,j   ≤ max r,i,j k0 (h 0 ) k0 • • • kn (A n-1 ) kn-1,kn ((A n ) kn,i -(A n ) kn,j ) + k0 |(A 0 ) k0 -(h 0 ) k0 | • • • max i,j kn (h n-1 ) kn ((A n ) kn,i -(A n ) kn,j ) + kn (A n-1 ) kn-1,kn -(h n-1 ) kn max i,j ((A n ) kn,i -(A n ) kn,j )Moreover, if the matrices A p have shapes r p × c p , the right hand side can be computed in time O( p r p c p ).Proof. By taking the max of Theorem 25 over rows r of A 0 .Let us first look at EQKE := E q QK T ĒT .From Figure 9a , we see that there is not much variation along long query token direction. We can confirm this by performing a singular value decomposition (SVD) on EQKE, as seen in Figure 14 . However, we can get a much better bound on the max row-diff of EQKE_err without having to multiply out all four matrices. We can use an approach vaguely similar to the mean+diff trick, as follows.If we want to compute the max row-diff of a product of matrices AB, we can compute by Theorem 21 max r,i,j ((AB) r,i -(AB) r,j ) ≤ maxEQUATIONor by combining this approximation with Theorem 18 via Theorem 24 we may compute max r,i,j ((AB) r,i -(AB) r,j )≤ max i,j k E r A r,k (B k,i -B k,j ) + max r k |A r,k -E r A r,k | max i,j (B k,i -B k,j )taking whichever bound is better.The first gives us a bound of 7.94 on the maximum row-diff, which is better than we can get by doing SVD on the product of the matrices! We can get an even better bound by peeling off the first two singular values of all four matrices before multiplying them; this gives us a bound of 5.67. Combining it with the avg+diff trick wouldn't give us much (8.05 and 5.66 respectively), as we've effectively already done this by peeling off the leading singular contributions; the mean of EQKE_err over dimension zero has norm 0.025 (0.030 ± 0.012 across all seeds).Although this error bound is no longer the leading asymptotic bottleneck, we can peek ahead to what we get if we want to be linear in parameter count. In this case, we can apply the recursive version of Equation 12 via Theorem 23, giving a bound of 97.06 on the maximum row-diff. EQKE (1) , EQKE_err ← DECOMPOSE(EQKE) Require: EQKE (1) (t query , t) -EQKE (1) (t query , t max ) -EQKE_err tquery ≤ EQKE(t query , t) -EQKE(t query , t max ) ≤ EQKE (1) (t query , t) -EQKE (1) attn-weights-unscaled 0,i ← b 0,i + (∆a min if i ∈ idx-set else 0) • Compute for each number of queries handled (where we assume we handle all queries with greater overlap than the current one) and for each maximum input token t max , how many of the query tokens t query fall strictly below the max t max (and whether or not the model succeeds when t max = t query ). This will tell us how many query tokens we can count for a given maximum token. • Compute a table indexed on pairs of # of queries handled and input tokens t which stores the smallest difference in more attention paid to t + 1 than to t (O(d vocab 2 )).• Compute a table indexed on pairs t max , t storing an upper bound on amount more attention paid to non-maximum tokens than to t max by Oracle-permitted query tokens (the Oracle is indexed only on t max ) (O(d vocab 2 )).• For each # queries permitted: compute for each t max , t, c, if the non-maximum token t contributes little enough to incorrect logits that even with the worst skip connection the model still gets the correct answer.Did we miss out on any structure in the error term of EQKE? The distribution of entries of the four matrices looks pretty close to normal as seen in Figure 21 . attention-quadratic refers to handling the high-rank attention error term EQKE_err from Appendix G.2.5 in time O(d vocab d model ) either with the recursive max row-diff trick (Theorem 23) 30 or with the recursive max row-diff trick fused with the mean+diff trick either just on the query side 31 or throughout 32 11) 33 , considering all ways of multiplying out a subset of the matrices and taking the maximum row-diff of the resulting pair of matrices 34 (Theorem 21), or fusing the max row-diff trick with the mean+diff trick 35 (Theorem 24).When attention is not mentioned, this indicates that we handle the attention error term in time O(d vocab 2 d model ), either by taking the per-row maximum row-diff 36 or by using the full rank EQKE matrix and taking the per-row maximum row diff 37 .Finally, note that in combining the rank one attention computation with EVOU, PVOU, and EU, we may either use the mean+diff trick 38 (Appendix G.1) or not 39 ; this makes up the final factor of 2.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. )) corresponds to fusing the max row-diff trick with the mean+diff trick (Theorem 24) and applying it on the factored SVD of EQKE_err 1 . 30 "max_diff_subproduct_recursive" 31 "mean+max_diff_subproduct_recursive" 32 "mean_recursive+max_diff_subproduct_recursive" 33 "svd" 34 "max_diff" for EQKE_err 1 , "max_diff_subproduct" for EQKE_err 35 "mean+max_diff" for EQKE_err 1 , "mean+max_diff_subproduct" for EQKE_err 36 "max_diff_exact" 37 "exact_EQKE+max_diff_exact" 38 "mean_query+diff" 39 "drop_average_query_per_output_logit_reasoning"Our 102 proof strategies are can be broken up as 1 + 1 + 10 ×× 2: two standalone strategies, and a class of strategies parameterized on three axes of cardinality 10, 5, and 2 (Appendix H).5 Code: https://github.com/JasonGross/guarantees-based-mechanistic-interpretability/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.All of our trained models behave similarly; see Appendix B.3.This implies that the ratio of attention paid to token t and t -1 is approximately exp(1.236) ≈ 3.442.For comparison, the average off-diagonal element of EVOU is 21.68 ± 0.83 below the corresponding diagonal element.Appendix A.10 discusses how to compute the "brute-force" accuracy of a model on an infinite distribution.In Section 4.3, we will consider a smaller set of "pure sequences".In fact, this is the motivation behind the standard rewrites of QK and OV into position-independent and position-dependent components (Section 3).Technically the floating point numbers are not a field. We gloss over this point here, since they define all of the field operations, even if those operations do not satisfy the field axioms."Per-model best" here means that for each model seed, we select the variant of the subcubic proof with the highest bound.If 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.Note that while in the main body, M(t) referred to the pre-softmax output logits, in the appendix we abuse notation and occasionally use it to refer to maximum token indicated by the logits where appropriate.See 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 are written out explicitly.We use the logit difference ∆ℓt * because: (a) it is shared in the computation of f 0-1 and f log-loss ; (b) it is a 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.We will prove that cmp is transitive in the process of proving this theorem.In fact, the theorem yields a stronger result, but we will only need the following formulation.We 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 so we can't expect that any automated compact proof discovery system will give it to us.This shows up in the bias towards having larger values (both positive and negative) in the lower-right corner of the plot, indicating that errors are larger for larger query and key values. We hypothesize that this is due to the distribution of data: larger values are more likely to have more space between the maximum and next-most-maximum token, so a bit of noise matters less for larger maxes than for smaller ones.This strategy is labeled "max_diff" in the Python source code.These strategies are labeled "mean_query+max_diff" and "svd_query+max_diff" in the Python source code"max_diff_exact""global_max_diff_exact" 


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

Reviews: review: paper_summary: paper_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:
- 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.
- 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.
- Then, several verification approaches are considered:
	- Brute force: this means running the model on every possible input and calculating the accuracy
	- "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$. 
	- "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.
- The main result is that approaches using more mechanistic knowledge of the
approximate algorithm implemented by the transformer require less computation for verification, but also lead to less faithful/tight bounds on the accuracy.
strength_weakness: strength_weakness: - the setting is quite toy, though this is probably very hard to overcome for such an approach with human labor.
questions: 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?
- 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.
- you mention: "models with similar training procedure and performance may still
have 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: 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. 
- 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)
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

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

3) 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.
questions: questions: I request the authors to kindly address the questions in the weaknesses section.
limitations: limitations: Yes, the authors have addressed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
 
- 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. 

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

- The paper is difficult to read in the sense that there are undefined/inconsistent terms and notations. For example,  
- 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. 
- 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?
- Line 87, Line 106, inconsistent bold-facing. What is difference between $\mathbf{E}$ and $E$. Moreover, what does $\mathbf{E}_q$ mean?
- 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?
- Line 501, citation deadlink. 
- Fig. 6 ''Let $d = \sqrt{d}$''?
- 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.
- Algorithm 3, 5 is very difficult to parse, break this up into multiple subroutines?
questions: 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?
- 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?
- 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? 
- 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).
- 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] 
- 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? 
- Why is the ratio of the largest to second-largest singular value a good metric for the faithfulness of the interpretation?

[1] Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Circuit Component Reuse Across Tasks in Transformer Lan-
guage Models. In The Twelfth International Conference on Learning Representations, September 2023. URL
https://openreview.net/forum?id=fpoAYV6Wsk.

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

[3] Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in
the Wild: a Circuit for Indirect Object Identification in GPT-2 Small. In The Eleventh International Conference
on Learning Representations, September 2022. URL https://openreview.net/forum?id=NpsVSN6o4ul.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

Minor Issues:
- 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”).
- Typo in line 95: “we denote the true maximum of the by t_{max}.”
- Line 162: one “the” too many: “… sequence x, [the] define the ….”
- Line 189: “consider” was probably meant to be “considered.”
questions: 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: limitations: The authors have properly discussed the limitations of their approach.

Comments: title: Thank you for the detailed clarifications
Rebuttal by Authors (cont.)
comment: I thank the authors for their very thorough followup to my questions. Also, thanks for pointing out my mistake about the chance-level performance on the max task! 

I am looking forward to future work in this direction, and I remain strongly in favor of accepting this paper to the conference.
Thank you to the authors for their time and detailed response. My concerns have been addressed, and I think the methods proposed in this paper are a solid contribution. However, I believe that the presentation of the manuscript needs significant improvement including the many typographical errors, consistency of notation, explanation of the role of the data distribution, etc. 

I have raised my scores accordingly.
I thank the authors for their rebuttal. Their response helps me clarify my questions, therefore I will increase my score.
> Why would you need to restrict the distribution to be small enough to check with brute-force? Isn't the contribution of the paper the ability to achieve tight lower bounds without this type of brute-force verification?

We want to clarify that the paper's contribution is demonstrating how mechanistic understanding (aka compression of model internals) permits us to avoid brute force enumeration. Without such understanding, obtaining a more compact proof is not feasible.

Here are examples illustrating how existing circuits are developed on small distributions, and why extending these results to larger distributions would require additional mechanistic interpretation:
1. Names: The interpretation uses a small fixed dataset of names, lacking a broader interpretation to show that all tokens in specified positions would be parsed as names. For intuition: While we might expect a model to handle "John → Jon", it might struggle with "Mary → marry". Properly enumerating the set of names would require either proving that all elements in the dataset are equivalent, or resorting to brute force evaluation (which is the limiting behavior of their random sampling approach).
2. Patterns: The interpretation is restricted to the ABBA case. This means either limiting ourselves to the ABBA case similarly, or conducting additional interpretation to show how answers on the ABBA case apply to other cases like ABAB.

If we had taken on a project performing such interpretations, that would have been very cool! However, such further interpretation would be a separate contribution from the main objectives of our study.

Please let us know if you have any further questions or if this fails to clarify some question that you’ve asked.
Thank you for your time, effort, and engagement.

> Why is it the case that the role of every component in the network needs to be known for its proof to be compactified?

To be more precise: the role of every component must be known for the proof to be *asymptotically* compactified.

When measuring asymptotic behavior, it is standard to consider the leading order term(s) only. If the number of layers and attention heads is fixed, the number of parameters in a component is a constant fraction of (that is, linear in) the number of parameters of the overall network. Hence if we have to use the brute-force proof on any component, we will have an asymptotic cost equal to that for performing a brute-force analysis of the entire network.

As we are not able to compress the proof of a component's behavior in the absence of mechanistic understanding of the role and functioning of that component, the leading contribution to the asymptotic size will always be the components we don't understand at all.

Moreover, we want to clarify that compressing components using mechanistic understanding is how we obtain compact proofs with non-vacuous bounds *at all*, even if approximation errors from the compression can eventually tank the hoped-for bound. Without integrating mechanistic understanding, we are not able to compactify the brute-force proof.

> Where are the results for these experiments?

Sorry, “attached” was left-over from an earlier draft of the response, and we meant to say “included”: our results were included in the sentence “the cubic proof achieves bounds of 94% (Max-of-5), 93% (Max-of-10) and 83% (Max-of-20) in under two minutes”.

However, we can provide some additional results for Max-of-10: The equivalent graph to Fig. 3 for Max-of-10 has the following points:
- importance sampling (acc: 0.9988 ± 0.0013), $1.6 \cdot 2^{85}$ FLOPs for brute force
- cubic (rel acc: 0.915 ± 0.021), $1.4 \cdot 2^{27}$ FLOPs
- subcubic (rel acc: 0.734 ± 0.019) $(1.49 ± 0.10) \cdot 2^{22}$ FLOPs
- attention-$d_{\\mathrm{vocab}}d_{\\mathrm{model}}^2$ (rel acc: 0.718 ± 0.019), $(1.609 ± 0.090) \cdot 2^{22}$ FLOPs
- direct-quadratic (rel acc: 0.657 ± 0.032), $(1.392 ± 0.013) \cdot 2^{22}$ FLOPs
- attention-quadratic (rel acc: 0.437 ± 0.028), $(1.384 ± 0.013) \cdot 2^{22}$ FLOPs
- attention-quadratic, direct-quadratic (rel acc: 0.398 ± 0.028) $(1.318 ± 0.014) \cdot
2^{22}$ FLOPs

The equivalent graph to Fig. 4 has ratios $\sigma_1 / \sigma_2$ starting at around 400, and the best-fit line for the SVD proof approach has $bound ≈ 0.56 + 8.3\cdot 10^{-5} (σ₁/σ₂)$, with $r^2 \approx 0.5$.

> Also, as $k$ or $d$ changes how are the authors verifying that the underlying circuit implemented by the model is the same, as this could affect the assumptions of all subsequent theorems/algorithms?

We note that: 
1. We are working with proofs as computations (lines 56--59), i.e. the theorem statement doesn't exist independently of the proof, it doesn’t presuppose a lower bound $b$ that we then prove. 
2. Whatever bound $b$ we are able to compute is the theorem statement we have proven. Thus, the theorem template is valid regardless of what circuit is implemented by the model. 

The faithfulness of the interpretation used in the proof mediates the bound (Section 5.2). In this way, the tightness of bound becomes a metric for how faithful the interpretation is. Thus, the recovered bounds can be seen as a measurement of how well the interpretation for Max-of-4 generalized to the new models trained to compute Max-of-5, Max-of-10, and Max-of-20. We are clarifying this further in our writeup, and this is valuable feedback on what concepts are worth highlighting from verification. 

Note that if the EVOU circuit did not perform copying, or if the model did not compute its answer by paying more attention to the correct token, the bound would be vacuous. We want to highlight that in the Fig. 4 equivalent for Max-of-10, the smallest ratio $\sigma_1/\sigma_2$ that we saw was over 400, which is evidence that the interpretation that EQKE is approximately low-rank remains valid. We are making this point more clear in a new section on applicability of results to other models.
Thank you for your detailed clarifications. I have read both the authors' global response and the individual responses. 

The authors adequately addressed my concerns relating to (1) the motivation of studying the max-of-$k$ task, (2) the typographical errors, and (3) OOD/adversarial distributions. However, I still have some further questions:

> 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. This lack of understanding prevents us from compactifying our proof of the model’s correctness.

I do not really understand why this would prevent you from compactifying your proofs. Based on my understanding, understanding the underlying algorithm implemented by the model through mechanistic interpretability should only give more insight into the specific relaxations that can be made. In the case of modular arithmetic (specifically see [1]), as some algorithms implemented by the models are permutation invariant with respect to the input, one could potentially ignore those during the brute-force proof: shortening your proofs. 

Why is it the case that the role of every component in the network needs to be known for its proof to be compactified?

> For 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]. If 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)

Why would you need to restrict the distribution to be small enough to check with brute-force? Isn't the contribution of the paper the ability to achieve tight lower bounds without this type of brute-force verification?

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

This is also mentioned in the global response. Where are the results for these experiments? Also, as $k$ or $d$ changes how are the authors verifying that the underlying circuit implemented by the model is the same, as this could affect the assumptions of all subsequent theorems/algorithms?

Looking forward to your clarifications. 

[1] Zhiqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The Clock and the Pizza: Two Stories in
Mechanistic Explanation of Neural Networks. In Advances in Neural Information Processing Systems, November
2023. URL https://arxiv.org/pdf/2306.17844.
Thank you!  Please let us know if there are any other questions we can clarify.
Thank you for clarifying the motivation of your study and conducting the experiments on larger input spaces. I appreciate the effort and believe that these results significantly strengthen the results of the paper. After reading the author rebuttals and other reviews, I have reconsidered my initial assessment. I now believe this paper would be an interesting contribution to the conference and will increase my score accordingly.
> How would this method generalize to different models that implement different circuits or algorithms on the same task?

We agree that guarantees require taking into account the specific model – this is part of the mechanistic understanding that interp should provide.
That being said, we believe that this specificity can be useful: for example, if a proof strategy derived for one model does not apply to a second model, this provides evidence that the two models are using different algorithms.

> In line 73, why is pessimal ablation done \[...with samples\]?

We apologize for the typo.
In line 73, “with values drawn from $\mathcal{D}$” should instead read “to values over $X$”.
That is, instead of sampling, which might not find the actual worst-case, we use the actual global worst-case values as the counterfactuals for ablation.
The counterfactuals commonly constructed in mechanistic interpretability are attempts to estimate the typical (/average-case) behavior of a component, which is insufficient for getting a worst-case guarantee on the overall behavior of the model.
We’ll edit the paper to fix the typo and further clarify this point.

> \[T\]he authors denote the residual from the rank-1 approximation of the EQKE matrix as ''unstructured noise.'' \[...\]

We apologize for our lack of clarity.
First, by “unstructured”, we mean that insofar as structure exists, it doesn’t help the model’s performance.
Second, “If we sample elements randomly” should read “If we replace the entries of $E_{q,2}^\perp$, $E_{k,2}^\perp$, $Q^\perp$, and $K^\perp$ with randomly sampled values”.
In order to validate our claim that the noise contains no helpful structure, we performed this replacement 100 times and compared the max-row-diff of these samples with the diff in the original model (lines 1074-7), and found that the diff of the original model is *worse* than the mean value by about 4σ.
(This is analogous to the standard resampling or mean ablations that mech interp work uses to check for lack of structure.)

> Why is the ratio of the largest to second-largest singular value a good metric for the faithfulness of the interpretation?

Our interpretation claims that $EQKE$ is approximately rank 1 (the size direction).
If the first singular value is much larger than all other singular values, this suggests that the matrix is well-approximated by a rank-1 approximation.
We’ll clarify this in the text.

Responses to more minor points:

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

We had this terminology wrong, and indeed were intending to speak of 1 minus the 0-1 loss.
We’ll correct this, and thank you for pointing it out.

> Line 87, Line 106, inconsistent bold-facing.
What is difference between $\textbf{E}$ and $E$.
Moreover, what does $\mathbf{E}_q$ mean?

The boldfacing is significant and is defined on line 107:

$$\hat{\mathbf{P}} = P - \overline{\mathbf{P}} \quad \text{and} \quad \mathbf{x}\overline{\mathbf{E}} = \mathbf{x}E + \overline{\mathbf{P}} \quad \text{and} \quad \mathbf{x}\mathbf{E}_q = \mathbf{x}E + \mathbf{P}_q \quad (\text{since } h^{(0)} = \mathbf{x}\overline{\mathbf{E}} + \hat{\mathbf{P}}).$$

That is, we use the boldfaced versions of $\mathbf{E}$ to denote $E$ summed together with positional embeddings.
We’ll try to clarify this more in the text.

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

We deeply apologize for not catching the bug in our code that led to an incorrect version of this figure being included in the original submission.
We’ve included a corrected version of this figure in the accompanying pdf.
From this figure, it is apparent that for the smallest values of $\sigma_1/\sigma_2$ (the three values below 350), this ratio is strongly correlated with proof bound.
For the “svd” proof strategy, where $\sqrt{2}\sigma_1(\mathrm{EQKE\_err})$ is used to bound the error term, a linear regression gives an $R^2$ of about 0.71.

Please don't hesitate to ask any further questions! We appreciate the thoughtful engagement, and would love the chance to clarify results that might encourage you to increase your score.