In this work, 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 lower bounds on the accuracy of 151 small transformers trained on a Max-of-K task.
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 noise as a key challenge for using mechanistic interpretability to generate compact proofs on model performance.

§ INTRODUCTION

One approach to ensuring the safety and reliability of powerful AI systems is via formally verified proofs of model performance [40, Dalrymple et al., 2024].
If we hope to deploy formal verification on increasingly large models [Hoffmann et al., 2022, Kaplan et al., 2020] with powerful emergent capabilities [Wei et al., 2022], we will need compact proofs of global robustness and generalization properties.
define compact somewhere?

However, existing approaches tend to use proof strategies that suffer from bad asymptotic complexity, while verifying either local robustness properties or correctness or generalization properties specific to training procedures instead of resulting models of interest.

We construct proofs using different degrees of mechanistic interpretation.
(Left) The models we consider in this paper are one-layer attention-only transformers, and so contain three “paths”: the OV circuit, the QK circuit, and the direct path.
(Right) For the brute-force proof (<ref>), we treat the model as a black box and thus need to check all possible combinations of inputs.
For the cubic proof (<ref>), 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 (<ref>), we use all parts of the mechanistic interpretation presented in <ref>.
(Bottom) For each of the three categories of proof, we report the number of FLOPs used in computing the certificate (lower=better, <ref>), lower bound on model accuracy (higher=better), effective dimension of the unexplained parts of the model (lower=better, <ref>), 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.

maybe include normalized accuracy bound / bound tightness in the diagram

One key challenge is that neural network architectures are highly expressive [Szegedy et al., 2013, Zhang et al., 2021], and models with similar training procedure and performance may still have learned significantly different weights [Nanda et al., 2023, Chughtai et al., 2023].
Thus, using only knowledge of the architecture or training procedure, it can be challenging to develop efficient proof strategies for verifying model performance.
We propose verifying model performance using understanding derived from mechanistic interpretability (<ref>) – 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.

add a sentence about generalization bounds
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  random seeds (<ref>), and 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 floating-point operations (<ref>).[
Our 102 proof strategies are can be broken up as $1+1+10\times5\times2$: two standalone strategies, and a class of strategies parameterized on three axes of cardinality 10, 5, and 2 (<ref>).

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 (sec:short-proofs-understanding,sec:computing-unexplained-dimensionality-reduction).
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 (
<ref>).[Code for reproducing our results can be found at <https://github.com/JasonGross/guarantees-based-mechanistic-interpretability/>.
A cache of generated data can be found at <https://github.com/JasonGross/guarantees-based-mechanistic-interpretability-with-data/>.

However, we also identify compounding structureless noise as a key challenge for generating compact proofs on model behavior (sec:discussion:compounding-noise,sec:compounding-structureless-noise:details).
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 over components even in cases when the empirical noise is small, and lead to vacuous performance bounds (<ref>).

§ MECHANISTIC INTERPRETABILITY FOR PROOFS

Bounds on global performance
In the style of prior mechanistic interpretability evaluation work [Chan et al., 2022], we target theorem templates that establish bounds on the expected global performance of the model.
Let $\Model : \InputSequenceType \to \SequenceLogitsType$ be a model (here assumed to be a neural network),  be a probability distribution over input-label pairs $\LabelInputSequence \in \LabelInputSequenceType$, notated as  when marginalized over labels, and $\PerformanceFn : \SequenceLabelLogitsType \to \mathbb{R}$ be a scoring function for evaluating the performance of the model.
Then, we seek to establish lower bounds $\bound$ on the expected $\AverageScore$ as the form:
should this be s star or sbar star

\begin{align}\label{eq:bound-template}
\AverageScore &:= \E_{\DrawnFrom{\LabelInputSequence}{\LabelInputDistribution}}[\PerformanceFn(\SequenceLabelLogitsPair{\SequenceLabel}{\Model(\InputSequence)})]  \ge \bound .
\end{align}

As $\PerformanceFn$ 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 0-1 loss and  to be uniform, so our theorems lower bound the accuracy of the model.
Our proof methodology generalizes straightforwardly to other input distributions (<ref>), and only a little work is required to generalize from accuracy to log-loss (<ref>).

Proof template
The proofs of model performance in this work have two components: a computational component $\Computation : \text{model weights} \to \mathbb{R}$ and a non-computational component  arguing that for any model $\Model'$, $\Computation(\Model') \le \E_{\DrawnFrom{\LabelInputSequence}{\LabelInputDistribution}} \PerformanceFn(\SequenceLabelLogitsPair{\SequenceLabel}{\Model'(\InputSequence)})$, thus implying that  generates a valid lower bound for the performance of $\Model$.
The whole proof is  paired with a trace of running  that certifies its output on .[Other components of the proof to account for the difference between floating point numbers and reals are described in <ref>.
Note that all proofs explicitly given in this paper are of  only;
we do not include any traces of running .
Here, $\bound = \Computation({\Model})$.
As even the size of the model parameters is much larger than any reasonable ,

we approximate the length of a proof pair $\Computation, \NonComputationalComponent$ by the length of a trace of $\Computation(\Model)$.

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 $\LabelInputDistribution$ to achieve  $\bound = \AverageScore$, albeit with a very long proof.

We quantify the length of $\Computation({\Model})$ using two metrics: the asymptotic time complexity of $\Computation$ as we scale the size of the model and the input $\InputSequence$, as well as the empirical average number of floating point operations required to evaluate $\Computation(\Model')$ over a given set of models $\{\Model{}_i\}$.
We measure tightness of bound of $\Computation(\Model)$ using the ratio of the bound to the true accuracy: $\bound/\AverageScore$.

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 [Wang et al., 2022, Chan et al., 2022, Heimersheim and Nanda, 2024].
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”) with values drawn from $\LabelInputDistribution$ 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 (<ref>) over input sequences, such that performing pessimal ablations on a smaller number of relaxed input sequences is sufficient to lower bound the performance on $\LabelInputDistribution$.

§ EXPERIMENTAL SETTING

We study our approach to generating compact proofs in a simple toy setting: one-layer transformers trained to output the max of $k$ integers.

Model Architecture We study one-layer, one-head, attention-only transformers with no biases but with learned positional embeddings, with vocabulary size $\dvocab$, model and head dimension $d = \dmodel=\modeldim{head}$, and context length $\nctx := k$.
The model parameters consist of the $\nctx \times\dmodel{}$ positional embedding $\Wpos$; the $\dvocab \times \dmodel{}$ token embed $\WE$; the $\dmodel{}\times\dmodel$ query, key, value, and output matrices of the attention head $\WQ$, $\WK$, $\WV$, and $\WO$; as well as the $\dmodel{} \times \dvocab$ unembed matrix $\WU$.
We assume (as is standard in language modeling) that $\dmodel < \dvocab$.

For an $\nctx{} \times \dvocab$ one-hot encoding $\InputSequenceOneHot%
\ifthenelse{\boolean{@twocolumn}}{}{%
= [\InputSequenceOneHot[0], \InputSequenceOneHot[1], \ldots, \InputSequenceOneHot[\nctx-1]]
}$ of an input sequence $\InputSequence\ifbool{neurips}{ = [\Token[0], \Token[1], \ldots, \Token[\nctx-1]]}{}$, we compute the logits of the model as follows:

\begin{align*}
% \begin{aligned}
\h{0} &= \InputSequenceOneHot \WE + \Wpos && \text{Initial residual stream }
% \\ &&&
(\nctx \times \dmodel)\\
\PreSoftmaxAttention{} &=  h^{(0)}  \WQ \WK^T  {h^{(0)}}^T/\sqrt{\dhead} && \text{Attention matrix }
% \\ &&&
(\nctx \times \nctx)\\
\h{1}  &=   \sigma^* (\PreSoftmaxAttention) \cdot \h{0} \WV \WO + \h{0} && \text{Final residual stream }
% \\ &&&
(\nctx \times \dmodel)\\
\Model(\InputSequence)  &= \Logits =  \h{1}_{\nctx-1} \WU && \text{Final seq.\@ position logits }
% \\ &&&
% \end{aligned}
\end{align*}

where $\sigma^*$ is the masked softmax function used in causal attention.
Because we only look at outputs of the model above the final sequence position $i= \nctx-1$, we also denote this position as the “query position” and the value of the token in this position as $\QueryTok$, one-hot encoded as .
The model's prediction is the token corresponding to the max-valued logit $\Logits_{\max}$.

Specifically, we study the setting with $\nctx = k=4$ because it is the largest sequence length for which we can feasibly evaluate the brute-force proof.
We set hidden dimension $\dmodel{}=32$ and a vocabulary of size $\dvocab = 64$ comprising integers between 0 and 63 inclusive.

For an input sequence $\InputSequence = [\Token[0], \Token[1], \Token[2], \Token[3]]$, we denote the true maximum of the sequence by $\MaxTok$.
Outputting the correct behavior is equivalent to outputting logits $\Logits$ such that $\LogitsDiff[\NonMaxOutputTok] := \Logits_{\NonMaxOutputTok} - \Logits_{\max} < 0$ for all $\NonMaxOutputTok \neq \MaxTok$.
We trained  models on this task.
Models achieved an average accuracy of  over the entire data distribution.

Path decomposition Following prior work [Elhage et al., 2021], 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 :

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  is simply the skip connection around the attention head.
consider the computational description vs “is” description

We further divide the QK and OV circuits into token (position-independent) and position-dependent components.
Let $\avgWpos = \frac{1}{\nctx}\sum_{i} \Wpos_i$ be the average position embeds across positions (of size $\dmodel$), and let $\barWpos$ denote either $\mathbf{1}_{\nctx} \otimes \avgWpos$ or $\mathbf{1}_{\dvocab} \otimes \avgWpos$ depending on context, the result of broadcasting $\avgWpos$ back into the shape of $\Wpos$ or $\WE$ (that is, $\nctx \times \dmodel$ or $\dvocab \times \dmodel$).
Similarly, let $\qWpos = \mathbf{1}_{\dvocab} \otimes \Wpos_{\QueryIndex}$ be the result of broadcasting $\Wpos_{\QueryIndex}$.
Then for one-hot encoded , we can rewrite the QK and OV circuits, as well as the direct path, as follows:

\begin{align*}
\textrm{QK circuit} &= \QueryTokOneHot \Big(\underbrace{\qWE \WQ \WK^T \barWE^T}_{\EPQKE} \InputSequenceOneHot^T + \underbrace{\qWE \WQ \WK^T\hatWpos^T}_{\EPQKP} \Big)\\
\textrm{OV circuit} &= \InputSequenceOneHot \underbrace{\barWE \WV \WO \WU}_{\EVOU} +  \underbrace{\hatWpos\WV \WO \WU}_{\PVOU} \qquad \textrm{Direct Path} = \QueryTokOneHot \underbrace{\qEU}_{\EU}
\end{align*}

where $\hatWpos = \Wpos - \barWpos$ and $\barWE = \WE + \barWpos$ and $\qWE = \WE + \qWpos$  (since  $\h{0} = \InputSequenceOneHot \barWE + \hatWpos$).

§.§ Mechanistic interpretation of learned models

Using standard empirical mechanistic interpretability techniques, we interpret one of our learned models (our “mainline” model) by independently examining the QK and OV circuits and the .[
All of our trained models behave similarly; see <ref>.
We find that the model outputs the largest logit on the true max token $\MaxTok$ 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  models by reporting the mean plus minus standard deviation for various summary statistics.
Plots for this section are available in <ref>.

QK circuit By qualitatively examining the position-independent QK component $\EPQKE$, we find the amount of pre-softmax attention paid to a key token is approximately independent of the value of the query token $\QueryTok$, and increases monotonically based on the size of the key token.
We confirm this hypothesis by performing a singular-value decomposition (SVD) of the $\EPQKE$ matrices (<ref>), and find that it contains a single large rank-one component with singular value around , around  times larger than the second largest component with singular value [1].
The left (query-side) singular vector is approximately constant in all dimensions, with value [1] $\approx$ 18 $=$ $\sfrac{1}{\sqrt{\dvocab}}$.
The right (key-side) singular vector of this component is monotonically increasing as we increase the size of the key token, with (0pt[0pt][0pt]$1/\sqrt{\dhead}$-scaled) pre-softmax attention increasing by an average of  when the key token increases by 1.[
This implies that the ratio of attention paid to token $t$ and $t-1$ is approximately $\exp({\EQKEHistAttentionDifferenceOverGapDupBySeqCountMeanMean[3]}) \approx \EQKEHistAttentionDifferenceOverGapDupBySeqCountExpMeanMean[3]$.]

In comparison, each $1/\sqrt{\dhead}$-scaled entry of the position-dependent QK component $\EPQKP$ has negligible size (average ), suggesting that $\EPQKP$ is unimportant to the functioning of the model.
We confirm this by zero ablating $\EPQKP$, which changes the models' accuracies from  to .
Combined with our interpretation of $\EPQKE$, 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 > \EPVOUInputsWithCopyingFailureMaxMeanPMStd$.
Since each entry in the sequence is uniformly sampled and $\dvocab=64$, this means that $\EVOU$ is a good approximation for the identity matrix for all but $\approx(\EPVOUInputsWithCopyingFailureMaxMean[0]/64)^4 \approx \qty[round-mode=places,round-precision=1,exponent-mode=scientific]{\fpeval{100*(\EPVOUInputsWithCopyingFailureMaxMeanFloat / 64) ^ 4}}{\percent}$ of the sequences.

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

Direct path As with the two position-dependent components,  the entries in $\EU$

have small absolute magnitude ,[For comparison, the average off-diagonal element of  is  below the corresponding diagonal element.] 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 (<ref>).
We provide detailed theorem statements, proofs, algorithms, and explanations of proof search in sec:model-math,sec:brute-force,sec:cubic:full,sec:subcubic:counting:details,sec:subcubic:tricks.

Our theorem statements for  will all be of the form

\[
\forall \Model',
\Computation_{\text{\emph{specific strategy}}}(\Model') \le \E_{\DrawnFrom{\InputSequence}{\InputDistribution}} \PerformanceFn(\SequenceLabelLogitsPair{\MaxTok}{\Model'(\InputSequence)}).
\]

We leave implicit the traces of running $\Computation_{\text{\emph{specific strategy}}}$ on our specific models to give the overall theorem.
We report the computational complexity or estimated FLOPs of running $\Computation_{\text{\emph{specific strategy}}}$ as approximations for our proof lengths.

§.§ The brute-force baseline

We start by considering the brute-force proof (<ref>), which treats the model as a black box and evaluates it on all possible sequences.[<ref> discusses how to compute the “brute-force” accuracy of a model on an infinite distribution.]
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 <ref> to derive more compact proofs.

We report the proof complexity, normalized accuracy bound, and estimated flops required (<ref>), as well as unexplained dimensionality (<ref>).
We round the FLOP and unexplained dimension counts to the closest power of 2, and report the mean/standard deviation of the bound averaged across all 151 models.
As we include more aspects of the mechanistic interpretation (reflected by a lower number of unexplained dimensions), we get more compact proofs (in terms of both asymptotic complexity and FLOPs), albeit with worse bounds.
For space reasons, we use $\tmpnctx := \nctx$, $\tmpdmodel := \dmodel$, and $\tmpdvocab := \dvocab$.

Description of Proof     Complexity Cost     Bound
Est. FLOPs
Unexplained Dimensions

Brute force     $\mathcal{O}(\dvocab^{\nctx+1} \nctx \dmodel)$

Cubic     $\mathcal{O}(\dvocab^3\nctx^2)$

Sub-cubic     $\mathcal{O}(\dvocab^2 \cdot \nctx^2 + \dvocab^2 \cdot \dmodel)$

without mean+diff

Low-rank QK

$\mathcal{O}(\dvocab^2 \nctx^2 + \dvocab \dmodel^2 + \secondpart{\text{(EU\&OV) } \dvocab^2 \dmodel}$

SVD only

Low-rank EU     $\mathcal{O}(\dvocab^2  \nctx^2 + \dvocab \dmodel + \secondpart{\text{(QK\&OV) } \dvocab^2 \dmodel}$

SVD only         [1,exponent-mode = threshold]

Low-rank QK&EU     $\mathcal{O}(\dvocab^2 \nctx^2 + \dvocab \dmodel^2 + \dvocab\dmodel + \secondpart{\text{(OV) } \dvocab^2 \dmodel}$

SVD only             [1,exponent-mode = threshold]

Quadratic QK     $\mathcal{O}(\dvocab^2 \nctx^2 + \dvocab\dmodel + \secondpart{\text{(EU\&OV) }\dvocab^2\dmodel}$

Quadratic QK&EU     $\mathcal{O}(\dvocab^2 \nctx^2 + \dvocab\dmodel + \secondpart{\text{(OV) }\dvocab^2\dmodel}$

§.§ A cubic proof

Next, we use the fact that the model is composed of the direct path and the QK and OV circuits (<ref>) to decrease the number of sequences that we need to consider, and the fact that only the position-independent components $\EPQKE$ and $\EVOU$ contribute meaningfully to performance (<ref>) to pessimize over sequence ordering.

First, let a pure sequence $\PureSeq$ be a sequence with at most three distinct tokens: the max token $\MaxTok$, the query token $\QueryTok \leq \MaxTok$, and optionally a third token $\NonMaxNonQueryTok < \MaxTok$, and let $\PureSeqSet$ be the set of all pure sequences in $\InputSequenceType$.[
In <ref>, we will consider a smaller set of “pure sequences”.
For a given input sequence $\InputSequence$, define the adjacent pure sequences $\AdjSeqs{\InputSequence}$ as the set of sequences that share the same max and query token, and only take on values in $\InputSequence$:

\begin{align*}
\AdjSeqs{\InputSequence} = \setst{\PureSeq \in \PureSeqSet}{
% \begin{aligned}
% &
\max_i \PureSeq_i = \MaxTok, ~\PureSeq_{\mathrm{query}} = \QueryTok, %\\
% &
% \quad
\forall i < \nctx,~ \PureSeq_i \in \InputSequence
% \end{aligned}
\end{align*}

Using the convexity of softmax and the fact that the model contains three paths, we can show that one-layer attention-only transformers
satisfies a variant of the following convexity property: for a given $\InputSequence$, if $\Model(\PureSeq)$ is correct for all $\PureSeq \in \AdjSeqs{\InputSequence}$, then $\Model(\InputSequence)$ is correct.
That is, for these transformers, we can bound the accuracy on all sequences by evaluating $\Model$ on only the $O(\dvocab^3 (\nctx-1)!)$ pure sequences.

This allows us to bound the accuracy of our actual $\Model{}$ on all $\dvocab^{\nctx}$ sequences, while evaluating it on $O(\dvocab^3 (\nctx-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 $(\MaxTok, \QueryTok, \NonMaxNonQueryTok)$, there are $(\nctx-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 $(\MaxTok, \QueryTok, \NonMaxNonQueryTok)$ tuple to the number of $\NonMaxNonQueryTok$ in the pure sequence, and the total number of sequences to $O(\dvocab^3 \nctx)$.
By precomputing the five component matrices $\EU$, $\EPQKE$, $\EPQKP$, $\EVOU$, $\PVOU$ and cleverly caching intermediate outputs, we can reduce the additional work of each sequence to the $O(\nctx)$ required to compute the softmax over $\nctx$ elements, resulting in asymptotic complexity $O(\dvocab^3 \nctx^2)$ (<ref>, additional details in <ref>).

§.§ Sub-cubic proofs

We now consider proofs that are more compact than $O(\dvocab^3)$.
These require avoiding iteration over any set of size $O(\dvocab^3)$ (e.g. the set of pure sequences) and performing operations that take $O(\dvocab)$ time on each of $O(\dvocab^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 $\EPQKE$.
We consider applying variants of these tricks at different locations in the naïve subcubic proof, leading to 100 distinct subcubic proof strategies.
See <ref> for a formal description of these strategies.

§.§.§ Removing cubic-time computations

Reducing the number of cases by pessimizing over sufficiently small tokens Previously, we considered $\Theta(\dvocab^3 \nctx)$ pure sequences $\PureSeq$, with $\PureSeq$ parameterized by $(\MaxTok, \QueryTok, \NonMaxNonQueryTok, \NumCopiesNonMaxTok)$.
Recall from our mechanistic interpretation in <ref> that the pre-softmax attention paid from $\QueryTok$ to a key token $\NonMaxNonQueryTok$ is broadly invariant in $\QueryTok$ and increases roughly linearly with the size of $\NonMaxNonQueryTok$.
This allows us to pessimize over the OV circuit over all “sufficiently small” tokens.

More formally, suppose we are given some gap $\MinGapSingular \in \N$.
For each pure sequence $\PureSeq$ with max token $\MaxTok$, query token , such that $\QueryTok \leq  \MaxTok- \MinGapSingular$, and $\NumCopiesNonMaxTok$ copies of the third token type $\NonMaxNonQueryTok \leq  \MaxTok - g$, we pessimally ablate the OV circuit over the set $\PureSeqSet(\MaxTok, \QueryTok, \NumCopiesNonMaxTok; \MinGapSingular)$ of pure sequences $\PureSeq'$ with the same max and query tokens and $\NumCopiesNonMaxTok$ copies of the third token type $\NonMaxNonQueryTok$.
If the model gets all sequences in $\PureSeqSet(\MaxTok, \QueryTok, \NumCopiesNonMaxTok; \MinGapSingular)$ correct, then we can conclude that it gets $\PureSeq$ correct, otherwise, we treat the model as having gotten $\PureSeq$ wrong.
This means that it suffices to only consider the $O(\dvocab^2 \nctx)$ pessimal pure sequences of each of the $O(\dvocab^2 \nctx)$ sets of the form $\PureSeqSet(\MaxTok, \QueryTok, \NumCopiesNonMaxTok; \MinGapSingular)$.

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

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

We can find the rank 1/2 approximations by performing SVD on $\EPQKE$.
We can efficiently compute the SVD in $\mathcal{O}(\dvocab \dmodel^2)$ time by using the fact that $\EPQKE$ can be written as the product of a $\dvocab \times \dmodel$ matrix and a $\dmodel \times \dvocab$ matrix.
This allows us to avoid performing the $\mathcal{O}(\dvocab^2 \dmodel)$-cost matrix multiplications to explicitly compute $\EPQKE$.

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.

§.§.§ Additional subcubic proof strategies

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

\begin{align*}
\min_{x,y,z} h(x,y,z) &\geq \min_{x,y} f(x, y) + \min_{y,z} g(y, z)
\end{align*}

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

\begin{align*}
\min_{x,y,z} h(x,y,z) &\geq \min_{x,y} \left( f(x, y) + \E_z' g(y, z')\right)
% \\ & \phantom{\geq}
+ \min_{y, z} (g(y, z) - \E_z' g(y, z'))
\end{align*}

This estimate will have error at most $\varepsilon$, while the naïve estimator can have arbitrarily large error.
We refer to this rewrite as the “mean+diff trick”.[
In fact, this is the motivation behind the standard rewrites of QK and OV into position-independent and position-dependent components (<ref>).
From the mechanistic interpretation in <ref>, 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 $A B$ 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” 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 <ref> for neuripsmore details.

See <ref> for more variants neuripsand combinations of these strategies.

every axis/.append style=legend pos = outer north east

For each of the proofs in <ref>, we plot the number of FLOPs used to compute the certificate, as well as the normalized accuracy lower-bound ($\bound/\AverageScore$).
The brute-force proof (<ref>) computes the exact performance uses orders of magnitude more compute than other approaches.
The cubic proof (<ref>) uses a small amount of mechanistic understanding and less compute, while still retaining good accuracy lower bounds.
Finally, subcubic proofs (<ref>) 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 <ref> for a detailed description of the various proof strategies.

§ RESULTS

We run each of $\NumSeeds$ 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 $\bound/\AverageScore$ (<ref>) in <ref>.
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.

§.§ Compact proofs both require and provide mechanistic understanding

Quantifying mechanistic understanding using unexplained dimensionality We first quantify the amount of mechanistic understanding used in a proof by measuring its unexplained dimensionality – the number of free parameters required to fully describe model behavior, assuming the structural assumptions of the proof are correct.
More detailed mechanistic interpretations will leave fewer free parameters that need to be filled in via empirical observation.
(Details in <ref>.)
In <ref>, we plot the two axes and find a suggestive correlation – that is, proofs based on less mechanistic understanding are longer.

More mechanistic understanding allows for more compact proofs In addition to the constructions in <ref>, 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 $\QueryTok$.

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

We plot the normalized accuracy bound versus the ratio of first and second singular values of $\EPQKE$, for various types of subcubic proofs that depend on a rank-1 approximation $\EPQKE$.
For each class of proof, the closer $\EPQKE$ 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 <ref> for a detailed description of the various proof strategies.

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 (<ref>), we approximate $\EPQKE$ with a rank-1 approximation corresponding to the “size direction”.
However, while other components are small, they're nonzero; this approximation harms model internals.

Less faithful interpretations lead to worse bounds on performance To confirm that faithfulness of understanding affects the tightness of bound independent of proof length, we plot the normalized accuracy bound of subcubic proofs that perform a rank-1 approximation to $\EPQKE$, 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 <ref>, 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.

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

Pessimal error terms compound in the absence of known structure Approximating $\EPQKE$ with a rank-1 matrix has small error.
However, when approximating each of the constituent matrices $E, Q, K$ with rank-1 approximations, pessimizing over the worst way to composing the individual small error terms leads to a bound on the error term of $\EPQKE$ 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 noise to compound (without just multiplying out the matrices), this approximation leads to a trivial bound on performance (<ref>).
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 Vaintrob et al., 2023), and thus that compounding structureless noise may be an issue in practice.

We plot, for each proof, the approximate number of flops required to evaluate the proof, versus the unexplained dimensionality (<ref>).
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 <ref> for a detailed description of the various proof strategies.

§ RELATED WORK

Generalization Bounds
Prior work in the PAC-Bayes framework [Zhang et al., 2021, Nagarajan and Kolter, 2019, Dziugaite and Roy, 2017] 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 [Huang et al., 2018, Huang et al., 2023] 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 [Katz et al., 2017, Cheng et al., 2017] or by relaxing the problem setting and taking an automated theorem proving approach [Gehr et al., 2018, Singh et al., 2019, Gowal et al., 2018, Mirman et al., 2018, Raghunathan et al., 2018] 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 [Harrison et al., 2014] and verify global performance properties by reverse-engineering the neural network weights.

Mechanistic Interpretability Finally, mechanistic interpretability is the subfield of the broader field of understanding model internals [Räuker et al., 2022], which is too large to faithfully summarize.
Our work takes most direct inspiration from efforts to deeply understand how either toy models [Nanda et al., 2023, Chughtai et al., 2023, Von Oswald et al., 2023, Akyürek et al., 2022] or small pretrained text transformers [Wang et al., 2022, Hanna et al., 2023] implement algorithmic tasks, generally by performing ablations and SVD.
In contrast, we formally prove that a transformer implements an algorithm.

Nichani et al., 2024 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 [Olsson et al., 2022].
Our results concern transformers with fixed weights.
In concurrent work, Michaud et al., 2024 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 noise 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 noise, perhaps by relaxing from worst-case pessimal ablations to typical-case heuristic guarantees [Christiano et al., 2022].

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

Additionally, we are thankful for clarifying discussions and feedback from Jacob Hilton, Matthew Coudron, Adrià Garriga-Alonso, Aryan Bhatt, Leo Gao, Jenny Nitishinskaya, Somsubhro Bagchi, Gabriel Wu, Erik Jenner, Ryan Greenblatt, Ronak Mehta, Louis Jaburi, and many others.

We are indebted to various organizations for their support:

* Alignment Research Center for funding this project and making it possible at all

* Mentorship for Alignment Research Students (MARS) program of the Cambridge AI Safety Hub (CAISH) for setting up the collaboration between a subset of authors, and providing funding for compute and in-person research sprints

* Constellation and FAR Labs for hosting a subset of the authors and providing an excellent research environment, including as part of the Visiting Fellows Program and Astra Fellowship

§.§ Author Contributions

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

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

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

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

[Akbarpour et al., 2010]
Behzad Akbarpour, Amr Abdel-Hamid, Sofiène Tahar, and John Harrison.
Verifying a synthesized implementation of IEEE-754 floating-point
exponential function using HOL.
The Computer Journal, 53:0 465–488, May 2010.

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

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

[Boldo and Melquiond, 2011]
Sylvie Boldo and Guillaume Melquiond.
Flocq: A unified library for proving floating-point algorithms in
In 2011 IEEE 20th Symposium on Computer Arithmetic, pages
243–252, July 2011.

[Chan et al., 2022]
Lawrence Chan, Adrià Garriga-Alonso, Nicholas Goldwosky-Dill, Ryan
Greenblatt, Jenny Nitishinskaya, Ansh Radhakrishnan, Buck Shlegeris, and Nate
Causal scrubbing, a method for rigorously testing interpretability
AI Alignment Forum, 2022.

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

[Christiano et al., 2022]
Paul Christiano, Eric Neyman, and Mark Xu.
Formalizing the presumption of independence.
arXiv preprint arXiv:2211.06738, 2022.

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

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

[Dziugaite and Roy, 2017]
Gintare Karolina Dziugaite and Daniel M. Roy.
Computing nonvacuous generalization bounds for deep (stochastic)
neural networks with many more parameters than training data, 2017.

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

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

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

[Escardó, 2007]
Martín H. Escardó.
Seemingly impossible functional programs, 2007b.
Accessed: 2024-05-15.

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

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

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

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

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

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

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

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

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

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

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

[McKinney, 2010]
Wes McKinney.
Data Structures for Statistical Computing in Python.
In Stéfan van der Walt and Jarrod Millman, editors,
Proceedings of the 9th Python in Science Conference, pages
56–61, 2010.

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

[Mirman et al., 2018]
Matthew Mirman, Timon Gehr, and Martin Vechev.
Differentiable abstract interpretation for provably robust neural
In International Conference on Machine Learning, pages
3578–3586. PMLR, 2018.

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

[Nanda and Bloom, 2022]
Neel Nanda and Joseph Bloom.
<https://github.com/TransformerLensOrg/TransformerLens>, 2022.

[Nanda et al., 2023]
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt.
Progress measures for grokking via mechanistic interpretability.
arXiv preprint, 2023.

[Nichani et al., 2024]
Eshaan Nichani, Alex Damian, and Jason D. Lee.
How transformers learn causal structure with gradient descent.
arXiv preprint, 2024.

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

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

[Plotly Technologies Inc., 2015]
Plotly Technologies Inc.
Collaborative data science, 2015.
URL <https://plot.ly>.

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

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

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

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

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

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

[Singh et al., 2019]
Gagandeep Singh, Timon Gehr, Markus Püschel, and Martin Vechev.
An abstract domain for certifying neural networks.
Proc. ACM Program. Lang., 30 (POPL), January 2019.

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

[Vaintrob et al., 2023]
Dimtry Vaintrob, Karell Hanni, and Jake Mendel.
Toward a mathematical framework for computation in superposition,
Accessed: 2024-05-22.

[Virtanen et al., 2020]
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy,
David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan
Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod
Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric
Larson, C. J. Carey, İlhan Polat, Yu Feng, Eric W. Moore, Jake
VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen,
E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro,
Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors.
SciPy 1.0: Fundamental Algorithms for Scientific Computing in
Nature Methods, 17:0 261–272, 2020.

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

[Wang et al., 2022]
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob
Interpretability in the wild: a circuit for indirect object
identification in gpt-2 small.
arXiv preprint, 2022.

[Waskom, 2021]
Michael L. Waskom.
seaborn: statistical data visualization.
Journal of Open Source Software, 60 (60):0
3021, 2021.

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

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

[Zhang et al., 2021]
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.
Understanding deep learning (still) requires rethinking
Communications of the ACM, 640 (3):0 107–115,

§ SUBTLETIES OF OUR APPROACH

§.§ Convex relaxation

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

In what sense are we using “convexity”?
The intuition is that we are attempting to optimize a function $f$ over its domain  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 $\InputSequenceType_i$, we define a set of “relaxed inputs” $\InputSequenceRelaxedType_i$ with an injection $\RelaxationFn_i : \InputSequenceType_i \hookrightarrow \mathcal{P}(\InputSequenceRelaxedType_i)$ mapping input to the model to the set of corresponding relaxed inputs.
On the relaxed input, we define a function $\RelaxedModelPerformanceFn_i: \InputSequenceRelaxedType_i \to \R$ such that for all $\InputSequence \in \InputSequenceType_i$ and all labels  for which  is supported by (has non-zero probability in) , we can find $\InputSequenceRelaxed \in T_i(\InputSequence)$ with $\PerformanceFn(\SequenceLabelLogitsPair{\SequenceLabel}{\Model(\InputSequence)}) \ge \RelaxedModelPerformanceFn_i(\InputSequenceRelaxed)$.
We proceed by finding a small subset of “boundary” examples $\Boundary_i \subset \InputSequenceRelaxedType_i$, proving that if $\RelaxedModelPerformanceFn_i(\InputSequenceRelaxed) \ge \bound_i$ for all $\InputSequenceRelaxed \in \Boundary_i$ then $\RelaxedModelPerformanceFn_i(\InputSequenceRelaxed) \ge \bound_i$ for all $\InputSequenceRelaxed \in \InputSequenceRelaxedType_i$.

Then, the computational component  of the proof validates that that $\RelaxedModelPerformanceFn_i(\InputSequenceRelaxed) \ge \bound_i$ for some $\bound_i$ for all $\InputSequenceRelaxed \in \InputSequenceRelaxedType_i$.
This allows us to conclude that $\PerformanceFn(\SequenceLabelLogitsPair{\SequenceLabel}{\Model(\InputSequence)}) \geq \bound_i$ for all $\InputSequence \in \InputSequenceType_i$.

note somewhere that only compositions are interpretable, how we break things up in the middle is arbitrary

§.§ Computing unexplained dimensionality

We claim in <ref> 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 <ref>, 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\cdot 64$ for a function $\FiniteSet{64} \to \R^2$, whose cardinality is $(\R^2)^{\FiniteSet{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.

§.§ Computing approximate FLOPs

In fig:pareto-frontier,tab:proof-complexity-bound, 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.

§.§ IEEE 754 vs.

In <ref> we defined  and  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 $\mathbb{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 $\mathbb{F}^*$ denote $\mathbb{F}$ restricted to finite numbers (that is, without NaNs and without $\pm\infty$).

We parameterize , , and  over the real field[
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.
] they operate on, so that, e.g., $\Computation_F : \text{model weights} \to F$.
Then we have  establishing that for any model $\Model'$, $\Computation_\R(\Model'_\R) \le \E_{\DrawnFrom{\LabelInputSequence}{\LabelInputDistribution_\R}} \PerformanceFn_\R(\SequenceLabelLogitsPair{\SequenceLabel}{\Model'_\R(\InputSequence)})$, and we have a trace demonstrating that $\Computation_{\mathbb{F}}(\Model_{\mathbb{F}}) = \bound$.

Let $i : \mathbb{F}^* \to \R$ be any injection that maps each floating point number to some real number that it is “closest to”.
Supposing that $\bound \in \mathbb{F}^*$ and thus $\bound\in\R$, we need two additional components of the proof.
We need to find $\varepsilon, \varepsilon'\in \R^+$ prove that

\[
\left|\Computation_\R(\Model_\R) - i\left(\Computation_{\mathbb{F}}(\Model_{\mathbb{F}})\right)\right| < \varepsilon
\qquad
\text{and}
\qquad
\left|
\left(\E_{\DrawnFrom{\LabelInputSequence}{\LabelInputDistribution_\R}} \PerformanceFn_\R(\SequenceLabelLogitsPair{\SequenceLabel}{\Model_\R(\InputSequence)})\right)
- i\left(
\E_{\DrawnFrom{\LabelInputSequence}{\LabelInputDistribution_{\mathbb{F}}}} \PerformanceFn_{\mathbb{F}}(\SequenceLabelLogitsPair{\SequenceLabel}{\Model_{\mathbb{F}}(\InputSequence)})
\right)
\right|
< \varepsilon'
\]

Then we can chain these proofs to prove that

\[
i\left(\E_{\DrawnFrom{\LabelInputSequence}{\LabelInputDistribution_{\mathbb{F}}}} \PerformanceFn_{\mathbb{F}}(\SequenceLabelLogitsPair{\SequenceLabel}{\Model_{\mathbb{F}}(\InputSequence)})\right)
\ge
\bound - \varepsilon - \varepsilon'
\]

Such $\varepsilon$-ball robustness proofs should be well within the scope of existing approaches to formal methods on neural nets, see, e.g., [Ramananandro et al., 2016, Appel and Kellison, 2024, Boldo and Melquiond, 2011, Kellison et al., 2023, Akbarpour et al., 2010, Wong and Kolter, 2018].
We leave actually dealing with the gap between floating point numbers and real numbers to future work.

§.§ Non-uniform distributions

In <ref> in <ref> we defined the expected model performance as the expectation of the distribution :

\begin{align*}\label{eq:bound-template:again}
\AverageScore &:= \E_{\DrawnFrom{\LabelInputSequence}{\LabelInputDistribution}}[\PerformanceFn(\SequenceLabelLogitsPair{\SequenceLabel}{\Model(\InputSequence)})]  \ge \bound .
\end{align*}

We then immediately specialized to the case where the marginalization  of  over labels is uniform.
As we'll see in <ref> in <ref> and <ref> in <ref>, the bound computation is modularized between a function that bounds the performance $\PerformanceFn(\SequenceLabelLogitsPair{\SequenceLabel}{\Model(\InputSequence)})$ 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 <ref> and relaxed-correctness-pessimizing-over-position in <ref>;
the expectation computation is brute-force in <ref> and cubic in <ref>.

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.

§.§ Infinite distributions

In the brute force proof in <ref>, we run the model on the entirety of $\LabelInputDistribution$.
This operation is straightforward when  is finite.
Perhaps surprisingly, we can do this even if  is infinite as long as the PDF $\LabelInputSequenceType \to \mathbb{R}$ of  is computable and the natural computational topology of  is compact [Escardó, 2007, Escardó, 2007, Escardó, 2004], because integration of computable functions on computable reals is computable [Simpson, 1998].

§.§ Using alternate loss functions

Building on the point from <ref>, it is also relatively straightforward to extend our approach from bounding 0-1 loss (expected accuracy) to bounding log-loss.
We will see in <ref> that the 0-1 loss 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 <ref>.
Although this is sufficient for the brute-force and cubic proofs, for the subcubic proof using <ref> in <ref>, 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.

§.§ Proving upper bounds

In this work, we focus on proving lower bounds on model performance.
Most of our theorems, for example in sec:softmax-convexity,sec:subcubic:counting:details, 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.

§.§ What proof system?

Length of proof depends on what proof system we use.
We permit any proof system where proof-checking time is linear in the length of the proof.

§ EXPERIMENTAL DETAILS

§.§ Training details

To train each model, we generate 384,000 random sequences of 4 integers picked uniformly at random, corresponding to less than 2.5% of the input distribution.
We use AdamW with $\text{batch\_size} = 128$, $\text{lr}=0.001$, $\text{betas}=(0.9, 0.999)$, weight_decay left at the default $0.01$.
We train for 1 epoch (3000 steps).

Over our  seeds, models trained with this procedure achieve  train accuracy and a loss of [1].[
Numbers reported as mean across training runs $\pm$ std dev across training runs of mean accuracy and loss.
When qualitatively examining a single model (for example in <ref> or <ref>), we use the model with config seed , model seed[The model seed is deterministically pseudorandomly derived from the seed .] .

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., 2019, Plotly Technologies Inc., 2015, Nanda and Bloom, 2022, Rogozhnikov, 2022, Virtanen et al., 2020, McKinney, 2010, Waskom, 2021

§.§ Additional details supporting our mechanistic interpretation of the model

We provide heatmaps of the matrices corresponding to the five components described/defined in <ref>, for the mainline model.

b0.490.5fig:eqke1$\EPQKE = \qWE \WQ \WK^T \barWE^T$

b0.490.5fig:eqkp$\EPQKP = \qWE \WQ \WK^T\hatWpos^T$

The $\text{QK circuit}$ can be decomposed into the position-independent and position-dependent components $\EPQKE$ and $\EPQKP$.
It computes the pre-softmax attention score for the model.
The positional contribution to the attention score, as shown in Figure (b), is minimal.
In Figure (a), the gradient from left to right along the key axis indicates that the single attention head pays more attention to larger tokens.
The uniformity along the query axis suggests that this behavior is largely independent of the query token.
Further, the light and dark bands imply that some queries are better than others at focusing more on larger tokens.

b0.251fig:evou1$\EVOU = \barWE \WV \WO \WU$

b0.251fig:pvou1$\PVOU = \hatWpos\WV \WO \WU$

b0.251fig:eupu$\textrm{Direct Path} = \qEU$

The $\text{OV circuit}$ is a sum of $\EVOU$ and $\PVOU$.
In Figure (a) we see that $\EVOU$ “copies” — with the exception of input tokens $\le \EPVOUInputsWithCopyingFailureMaxFloat{}$ ( across all models) — by virtue of the fact that above , the diagonal is larger than all the other elements in the same row.
We see that the range on Figure (b) is much smaller than Figure (a), indicating that positional contribution to the copying is minimal.
In Figure (c) we see that direct path values matter a bit more than , being only $\approx \num[round-mode=figures,round-precision=1]{\fpeval{\EPVOUMaxRowDiffMeanMeanFloat / \EUPUAbsMaxMeanFloat}}\times$ smaller than the typical  difference.
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.

§.§ Distribution of model mechanisms

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

The statistics of interest are:
(1) $\sigma_1/\sigma_2$, the ratio of the first two singular values of , a measure of the extent to which the attention score computation is low-rank;
(2) , the average score (accuracy) of the model across the entire input distribution;
(3) $b_{\text{cubic}}/\AverageScore$, the percent-score-recovered accuracy bound achieved by the cubic proof from <ref>;
(4) $b_{\text{subcubic}}/\AverageScore$, the percent-score-recovered accuracy bound achieved by the (per-model best)[
“Per-model best” here means that for each model seed, we select the variant of the subcubic proof with the highest bound.
subcubic proof from <ref>.

For each statistic of interest, <ref> 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.[
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.
Additionally, each group contains a boxplot of the summary:

* the minimum, maximum; the first and third quartiles; the median and mean; percentiles [2], [2], [2], and [2]; 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.

Plots of various models.

The statistics of interest are: (1) $\sigma_1/\sigma_2$, the ratio of the first two singular values of , a measure of the extent to which the attention score computation is low-rank; (2) , the average score (accuracy) of the model across the entire input distribution; (3) $b_{\text{cubic}}/\AverageScore$, the percent-score-recovered accuracy bound achieved by the cubic proof from <ref>; and (4) $b_{\text{subcubic}}/\AverageScore$, the percent-score-recovered accuracy bound achieved by the (per-model best) subcubic proof from <ref>.
The $y$ axes are: for  and , the query token; for , the input token; for , the input query token; for , the input position.
The $x$ axes are: for , the key token; for , the key position; for , , and , the output logit token.
All token axes range from 0 at the top (or left) to $\dvocab-1 = 63$ at the bottom (or right).
All position axes range from 0 at the top (or left) to $\nctx-1 = 3$ at the bottom (or right).

every axis/.append style=

max width=,height=,raise=-0.5
every axis/.append style=
scale only axis=true

2*seed     2cstatistic of interest

distribution     value
$\qWE \WQ \WK^T \barWE^T$
$\qWE \WQ \WK^T\hatWpos^T$
$\barWE \WV \WO \WU$
$\hatWpos\WV \WO \WU$

2c[$\qWE \WQ \WK^T \barWE^T$$\qWE \WQ \WK^T\hatWpos^T$]QK2

3c[$\barWE \WV \WO \WU$$\hatWpos\WV \WO \WU$$\qEU$]OV3

2c(selected seed)

§ MATHEMATICAL DEFINITIONS

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

, ,

$\FiniteSet{n}$     be the finite set on $n$ elements; we write

$\mathbb{N}^{<n}$     $:=\{0,1,\ldots,n-1\}$ when we care about the elements of $\FiniteSet{n}$

$\softmax(\Vector{v})$     be the softmax function $e^{\Vector{v}} / \sum_i e^{\Vector{v}[i]}$

$\maskedsoftmax(\Vector{v})$     be the casually-masked softmax function, $\maskedsoftmax(\Vector{v})_{i} := e^{\Vector{v}[i]} / \sum_{j\le i} e^{\Vector{v}[j]}$

be the size of the vocabulary

be the dimension of the attention head, in our case, equal to the hidden dimension of the model  (assumption: $\dhead < \dvocab$)

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

be the $\nctx \times\dmodel$ positional embedding

be the $\dvocab \times \dmodel$ token embed

, , ,      -3-be the $\dmodel\times\dmodel$ query, key, value, and output matrices of the attention head

be the $\dmodel \times \dvocab$ unembed matrix

$\InputSequence$     be the input token sequence $[\InputSequence[0], \InputSequence[1], \ldots, \InputSequence[\nctx-1]]$

$\InputSequenceOneHot$     be the $\nctx \times \dvocab$ one-hot-encoded input token sequence $[\InputSequenceOneHot[0], \InputSequenceOneHot[1], \ldots, \InputSequenceOneHot[\nctx-1]]$

$\QueryTokOneHot$     $:= \QueryTokShortOneHot := \QueryTokExplicitOneHot$ be the query token

$\QueryTok$     $:= \QueryTokShort := \QueryTokExplicit$ be the one-hot encoded query token

$\MaxTok{}$     be the true maximum token in the input sequence, $\max_i \Token[i]$

$\Model$     of type $\InputSequenceType \to \SequenceLogitsType$ be the model; sometimes we write

for the logits of the model $\Model(\InputSequence)$

be a probability distribution over input-label pairs $(\InputSequence, \SequenceLabel) \in \InputSequenceType \times \SequenceLabelType$

be  marginalized to a distribution over $\InputSequence \in \InputSequenceType$

$\PerformanceFn$     of type $\SequenceLabelType \times \SequenceLogitsType \to \R$ be a scoring function for evaluating the performance of the model

$\avgWpos$     be the average position embeds across positions (of size $\dmodel$), $\frac{1}{\nctx}\sum_{i} \Wpos_i$

$\barWpos$     be either $\mathbf{1}_{\nctx} \otimes \avgWpos$ or $\mathbf{1}_{\dvocab} \otimes \avgWpos$ depending on context – that is the result of broadcasting $\avgWpos$ back into the shape of $\Wpos$ or $\WE$ (that is, $\nctx \times \dmodel$ or $\dvocab \times \dmodel$)

$\qWpos$     be $\mathbf{1}_{\dvocab} \otimes \Wpos_{\QueryIndex}$, the broadcasting of $\Wpos_{\QueryIndex}$

$\hatWpos$, $\barWE$, $\qWE$     be $\Wpos - \barWpos$, $\WE + \barWpos$, and $\WE + \qWpos$ respectively

For any vector-valued function $\Vector{v}$ of length , parameterized over the input sequence , let

$\DeltaVector{v}[i] := \Vector{v}[i] - \Vector{v}[\MaxTok]$ be the difference between the $i^\text{th}$ element of $\Vector{v}$ and the element of $\Vector{v}$ corresponding to the true maximum of the input sequence .

For our particular model, we will have

\begin{align*}
\dvocab{} & :=64
\dhead{} & =\dmodel:=32
\nctx{} & :=4
\end{align*}

\begin{align*}
\InputSequenceType & :=(\N^{<\dvocab})^{\nctx} \cong \FiniteSet{64}^4
\SequenceLabelType & :=\N^{<\dvocab} \cong \FiniteSet{64}
\end{align*}

\[
\InputDistribution :=\UniformZeroN{\dvocab}^{\nctx}\text{, the uniform distribution}% over sequences of length \nctx{} of integers between 0 and $\dvocab-1$}
\]

Preliminary model definitions

\begin{align}
\underbrace{\E_{\DrawnFrom{(\LabelInputSequence)}{\LabelInputDistribution}}[\PerformanceFn(\SequenceLabelLogitsPair{\SequenceLabel}{\Model(\InputSequence)})]}_{\AverageScore\text{ (average model score)}} & \ge \underbrace{\bound}_{\mathclap{\text{lower bound}}}\qquad\text{(theorem statement)} \label{eq:bound-template:appendix}
\end{align}

We define the two typical performance functions corresponding to 0-1 accuracy and log-loss.
Note the shared subterm [i].

\begin{align}
\PerformanceFn^{\text{0-1}}(\MaxTok, \Logits) &:= {\Indicator}[{\argmax_i \Vector{\Logits}[i] = \MaxTok}] = \raisebox{0pt}[0pt]{\ensuremath{\displaystyle{\Indicator}[0 > \max_{i\neq \MaxTok} \underbrace{\Logits_i - \Logits_{\MaxTok}}_{\LogitsDiff[i]} ]}} \label{eq:performance-fn-0-1} \\
\PerformanceFn^{\text{log-loss}}(\MaxTok, \Logits) &:= \raisebox{0pt}[0pt]{\ensuremath{\Vector{(\softmax(\Logits))}[\MaxTok] = \left. 1 \middle/ \sum_i \right. \mathop{\mathrm{exp}}(\underbrace{\Logits_i - \Logits_{\MaxTok}}_{\LogitsDiff[i]})}} \label{eq:performance-fn-log-loss}
\end{align}

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

\begin{align}
\Model(\InputSequence) = \Logits(\InputSequence) &= \sigma^* \Big(\underbrace{\left(\QueryTokOneHot \WE + \Wpos_{\QueryIndex}  \right) \WQ \WK^T \left(\InputSequenceOneHot \WE + \Wpos \right)^T}_{\textrm{QK circuit}} /\sqrt{\dhead} \Big) \cdot \underbrace{\left(\InputSequenceOneHot \WE + \Wpos  \right) \WV \WO \WU}_{\textrm{OV circuit}} + \underbrace{\left(\QueryTokOneHot \WE + \Wpos_{\QueryIndex}  \right)\WU}_{\DirPath} \label{eqn:model-path-decomposition:appendix} \\
&= \underbrace{
\sigma^* \Big(
% \underbrace{
% \underbrace{
\QueryTokOneHot\Big(
\underbrace{\qWE \WQ \WK^T \barWE^T}_{\EPQKE} \InputSequenceOneHot^T
+ \underbrace{\qWE \WQ \WK^T\hatWpos^T}_{\EPQKP}
\Big)
% }_{\textrm{QK circuit}}
/ \sqrt{\dhead}
% }_{\PreSoftmaxAttention(\InputSequence)}
\Big)
\cdot
\underbrace{
\Big(
% \underbrace{
\InputSequenceOneHot \underbrace{\barWE \WV \WO \WU}_{\EVOU} + \underbrace{\hatWpos\WV \WO \WU}_{\PVOU}
% }_{\textrm{OV circuit}}
\Big)
}_{\mathclap{\EPVOU(\InputSequence)%_i := \EPVOU(\InputSequence[i], i)
}} +
% \underbrace{
\QueryTokOneHot \underbrace{\qEU}_{\EU}
% }_{\DirPath}
\label{eqn:model-path-decomposition:deeper} \\
&= \underbrace{\PostSoftmaxAttention(\InputSequence) \cdot \InputSequenceOneHot \barWE \WV \WO \WU}_{\LogitsContributionAttnTok[](\InputSequence)}
+ \underbrace{\PostSoftmaxAttention(\InputSequence) \cdot \hatWpos \WV \WO \WU}_{\LogitsContributionAttnPos[](\InputSequence)}
+ \underbrace{\QueryTokOneHot \qEU}_{\LogitsContributionSkip(\QueryTokOneHot)} \label{eqn:model-path-decomposition:dot-expanded} \\
% \underbrace{
\Big(\textstyle\sum_{i=0}^{\nctx-1}
\underbrace{
\underbrace{
(\PostSoftmaxAttention(\InputSequence))_i \InputSequenceOneHot[i]
% \underbrace{
\barWE \WV \WO \WU
% }_{\Vector\EVOU}
\underbrace{
% \underbrace{
\Vector\hatWpos[i]\WV \WO \WU
% }_{\Vector\PVOU[i]}
\Big)
% }_{\LogitsContributionAttn[](\InputSequence)}
+ \underbrace{\QueryTokOneHot
% \underbrace{
\qEU
% }_{\EU}
}_{\LogitsContributionSkip(\QueryTokOneHot)} \label{eqn:model-path-decomposition:contribs}
\end{align}

Definitions of the model behavior

Note that while in the main body, $\Model(\InputSequence)$ 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.

§ BRUTE-FORCE PROOF

For $\Call{brute-force}{\dvocab, \nctx, \Model}$ as defined in <ref>,

\[
\E_{\DrawnFrom{\InputSequence}{\UniformZeroN{\dvocab}^{\nctx}}}\left[
\argmax_i \Vector{(\Model(\InputSequence))}[i] = \max_i \InputSequence[i]
\right]
\ge
\Call{brute-force}{\dvocab, \nctx, \Model}
\]

In fact the two sides of the inequality are equal by definition.
Hence the inequality follows by reflexivity of $\ge$.

Counting Correct Sequences By Brute Force
[1]
correctness, input-sequence
model-behavior, input-sequencemaxinput-sequence
brute-force, ,
$\frac{1}{\dvocab^{\nctx}}$sumcorrectness, tokens    tokens $\in$ $\left(\Call{range}{\dvocab}\right)^{\nctx}$

§ DETAILS OF CUBIC PROOF

In this section, we prove formally the result used in sec:softmax-convexity-use.

At its heart, the convexity of softmax[
See <ref> 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.
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 $\dvocab^{\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 <ref> for details.

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

* that contains all actually possible weightings,

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

* for which computing the extremal weightings is computationally efficient.

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

Fix a set of values $v_i \in \R$.
The weighted average is bounded by the extremal values:
for any $w_i$ such that $\sum_i w_i = 1$ and $0 \le w_i \le 1$,

\[
\min_i v_i \le \sum_i w_i v_i \le \max_i v_i
\]

The proof is simple.
We have

\[
\sum_i w_i v_i - \min_i v_i = \sum_i w_i (v_i - \min_j v_j) \ge 0
\]

\[
\max_i v_i - \sum_i w_i v_i = \sum_i w_i (\max_j v_j - v_i) \ge 0
\]

so the result follows.

§.§ Proof strategy

The model computes the true maximum $\MaxTok$ when its outputs logits $\Logits$ are such that $\LogitsDiff[\NonMaxOutputTok] := \Logits_{\NonMaxOutputTok} - \Logits_{\MaxTok} < 0$ for all $\NonMaxOutputTok \neq \MaxTok$.[
We use the logit difference [] because:
(a) it is shared in the computation of $\PerformanceFn^{\text{0-1}}$ and $\PerformanceFn^{\text{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.
As a result, it suffices to lower-bound the proportion of sequences where (an upper bound on) $\LogitsDiff[\NonMaxOutputTok]$ is negative for all $\NonMaxOutputTok \neq \MaxTok$.
In particular, we will upper-bound the contribution from incorrect tokens $\NonMaxInputTok$ in positions $\NonMaxInputTokIndex$ to the difference $\LogitsDiffContributionAttn[\NonMaxInputTokIndex]{}$ between incorrect ($\NonMaxOutputTok$) and correct ($\MaxTok$) output tokens $\LogitsDiffContributionAttn[\NonMaxInputTokIndex]{\NonMaxOutputTok} = \LogitsContributionAttn[\NonMaxInputTokIndex]_{\NonMaxOutputTok} - \LogitsContributionAttn[\NonMaxInputTokIndex]_{\MaxTok}$.

We do this by arguing that the logit difference $\LogitsDiff[\NonMaxOutputTok]$ satisfies a certain notion of convexity over the space of a relaxation of sequences (<ref>), and constructing a set of $\Theta(\dvocab^3 \nctx)$ “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{}$:

\begin{align*}
%   \LogitsContributionAttn[k] &
%    \EVOU(\InputSequence[k]) & := \embed[{\InputSequence[k]}]\WV\WO\WU \\
% \PVOU(k) & := \posembed[k]\WV\WO\WU \\
% \EPVOU(\InputSequence[k], k) & := \EVOU(\InputSequence[k]) + \PVOU(k) \\
% \LogitsContributionSkip(\InputSequence) & := (\embed[{\InputSequence[-1]}] + \posembed[-1])\WU \\
%   \LogitsContributionAttn[k](\InputSequence) & = \PostSoftmaxAttention[k](\InputSequence) \left(\PVOU(k) +\EVOU(\InputSequence[k])\right)\\
%     & = \PostSoftmaxAttention[k](\InputSequence)\PVOU(k) + \PostSoftmaxAttention[k](\InputSequence)\EVOU(\InputSequence[k])\\
% &= \LogitsContributionAttnPos[k](\InputSequence) + \LogitsContributionAttnTok[k](\InputSequence)\\
\LogitsDiffContributionAttn[\NonMaxInputTok]{\NonMaxOutputTok}(\InputSequence) &=   \LogitsDiffContributionAttnPos[\NonMaxInputTokIndex]{\NonMaxOutputTok}(\InputSequence) +  \LogitsDiffContributionAttnTok[\NonMaxInputTokIndex]{\NonMaxOutputTok}(\InputSequence)\\
% &\leq \max_k \left(\Vector{\PVOU(k)}[t] - \Vector{\PVOU(k)}[{\max_j \InputSequence[j]}]\right) + \LogitsDiffContributionAttnTok[k]{t}(\InputSequence)
% \LogitsContributionAttnTok[k](\InputSequence) & := \todojg{\PostSoftmaxAttention[k] = (\InputSequence)\EVOU(\InputSequence[k])} \\
% \LogitsContributionAttnPos[k](\InputSequence) & := \todojg{\PostSoftmaxAttention[k](\InputSequence)\PVOU(k)} \\
% \Logits(\InputSequence) & := \LogitsContributionSkip(\InputSequence) + \sum_{k=0}^3 \LogitsContributionAttn[k](\InputSequence) \\
\end{align*}

We then upper bound $\LogitsDiffContributionAttnPos[\NonMaxInputTokIndex]{\NonMaxOutputTok}(\InputSequence)$ by noting that because the softmax attention is a weighted average of ,

\begin{align*}
\LogitsDiffContributionAttnPos[\NonMaxInputTokIndex]{\NonMaxOutputTok}(\InputSequence)
& = \Vector{\LogitsContributionAttnPos[\NonMaxInputTokIndex](\InputSequence)}[\NonMaxOutputTok] - \Vector{\LogitsContributionAttnPos[\NonMaxInputTokIndex](\InputSequence)}[{\max_{\NonMaxInputTokIndexAlt} \InputSequence[\NonMaxInputTokIndexAlt]}] \\
& = \PostSoftmaxAttention[\NonMaxInputTokIndex](\InputSequence)\Vector{\PVOU}[\NonMaxInputTokIndex,\NonMaxOutputTok]
- \PostSoftmaxAttention[\NonMaxInputTokIndex](\InputSequence)\Vector{\PVOU}[\NonMaxInputTokIndex,{\max_{\NonMaxInputTokIndexAlt} \InputSequence[\NonMaxInputTokIndexAlt]}] \\
& = \PostSoftmaxAttention[\NonMaxInputTokIndex](\InputSequence)\left(\Vector{\PVOU}[\NonMaxInputTokIndex,\NonMaxOutputTok]
- \Vector{\PVOU}[\NonMaxInputTokIndex,{\max_{\NonMaxInputTokIndexAlt} \InputSequence[\NonMaxInputTokIndexAlt]}]\right) \\
& \le \PostSoftmaxAttention[\NonMaxInputTokIndex](\InputSequence)\max_{\NonMaxInputTokIndex}\left(\Vector{\PVOU}[\NonMaxInputTokIndex,\NonMaxOutputTok]
- \Vector{\PVOU}[\NonMaxInputTokIndex,{\max_{\NonMaxInputTokIndexAlt} \InputSequence[\NonMaxInputTokIndexAlt]}]\right)
\end{align*}

Since $\sum_{\NonMaxInputTokIndex} \PostSoftmaxAttention[\NonMaxInputTokIndex](\InputSequence) = 1$, we have

\begin{align*}
\sum_{\NonMaxInputTokIndex=0}^{\nctx-1} \LogitsDiffContributionAttnPos[\NonMaxInputTokIndex]{\NonMaxOutputTok}(\InputSequence) & \le \max_{\NonMaxInputTokIndex}\left(\Vector{\PVOU}[\NonMaxInputTokIndex,\NonMaxOutputTok]
- \Vector{\PVOU}[\NonMaxInputTokIndex,{\max_{\NonMaxInputTokIndexAlt} \InputSequence[\NonMaxInputTokIndexAlt]}]\right)
\end{align*}

We then construct a set $\PureSeqSet$ of “pure sequences” consisting of only three types of tokens in one of two orders, and show that for each input sequence $\InputSequence$ and readoff logit $\NonMaxOutputTok$, we bound the logit difference from the token embeddings $\LogitsDiffContributionAttnTok[\NonMaxInputTokIndex]{\NonMaxOutputTok}(\InputSequence)$ using a small subset $\mathcal{X}$ of $\PureSeqSet$:

\begin{align*}
\sum_{\NonMaxInputTokIndex=0}^{\nctx-1}\LogitsDiffContributionAttnTok[\NonMaxInputTokIndex]{\NonMaxOutputTok}(\InputSequence) &\leq \max_{\xi \in \mathcal{X}} \sum_{\NonMaxInputTokIndex=0}^{\nctx-1}\LogitsDiffContributionAttnTok[\NonMaxInputTokIndex]{\NonMaxOutputTok}(\xi)
\end{align*}

indicate i comes from tuples of max, nonmax, etc indices
fill out template explicitly, recall from body
link notation to algorithms, model behavior relaxed computes s or w/e
We construct a set  of relaxed sequences, where each relaxed sequence  consists of a sequence and a position $(\InputSequence, \NonMaxInputTokIndex)$, where $\LogitsDiff_{\NonMaxOutputTok}{(\InputSequence, \NonMaxInputTokIndex)}$ 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 $\NonMaxInputTokIndex$ indicates the position that we pay 100% of the attention to for the  contribution.

We argue that $\LogitsDiff_{\NonMaxOutputTok}{(\InputSequence, \NonMaxInputTokIndex)}$ satisfies a certain notion of convexity over mixtures of sequences, such that we can evaluate it only on a set of $\Theta(\dvocab^3 \nctx)$ “extremal” sequences in a way that takes $O(\dvocab^3 \nctx)$ total time to bound $\LogitsDiff_{\NonMaxOutputTok}{(\InputSequence, \NonMaxInputTokIndex)}$ 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 <ref> provides a valid lower bound on the proportion of sequences the model gets correct.

§.§ Proof outline

We now proceed to the main results of this section.

Math fact: For each token $\NonMaxOutputTok$, the logit difference $\LogitsDiff[\NonMaxOutputTok]$ for any sequence $\InputSequence$ can be decomposed into the direct contribution from the embeds $\LogitsContributionSkip$, the attention-weighted position contribution (), and the attention-weighted token contribution ().
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 $\LogitsContributionSkip$ exactly by first computing $\EU = \qEU$ 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 (<ref>).

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 (<ref>).
We then argue that for pure sequences, it suffices to consider orderings where same tokens appear contiguously (<ref>).

insert instantiation of relaxed template

§.§ Formal proof

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

Fix a token value function (à la a row difference in ) $v : \N^{<\dvocab} \to \R$ and a token attention function (à la  for a fixed query token) $a : \N^{<\dvocab} \to \R$.
Fix a position value function (à la a row difference in ) $w : \N^{<\nctx} \to \R$ and a position attention function (à la  for a fixed query token) $b : \N^{<\nctx} \to \R$.

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

\begin{align*}
v_{\NonMaxInputTok} & = \Vector{\EVOU}[\NonMaxInputTok, \NonMaxOutputTok] - \Vector{\EVOU}[\NonMaxInputTok, \MaxTok] &
a_{\NonMaxInputTok} & = \EPQKE_{\QueryTok,\NonMaxInputTok} / \sqrt{\dhead} \\
w_{\NonMaxInputTokIndex} & = \Vector{\PVOU}[\NonMaxInputTokIndex, \NonMaxOutputTok] - \Vector{\PVOU}[\NonMaxInputTokIndex, \MaxTok] &
b_{\NonMaxInputTokIndex} & = \EPQKP_{\QueryTok,\NonMaxInputTokIndex} / \sqrt{\dhead}
\end{align*}

We can define a sequence of tokens via sorted tokens and a position permutation by specifying a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N^{<\dvocab}$ paired with a permutation $\sigma : \N^{<\nctx} \to \N^{<\nctx}$.

Given a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N^{<\dvocab}$ and a permutation $\sigma : \N^{<\nctx} \to \N^{<\nctx}$ define the sequence score $s_{t_0,\ldots,t_{\nctx-1},\sigma}$ as:

\[
s_{t_0,\ldots,t_{\nctx-1},\sigma} := \left. \sum_{\mathclap{0 \le i < \nctx}} v_{t_i} e^{a_{t_i} + b_{\sigma(i)}} \middle/ \hphantom{0}\sum_{\mathclap{0 \le i < \nctx}} e^{a_{t_i} + b_{\sigma(i)}} \right.
\]

We will drop the token subscript, writing only $s_\sigma$, when the token values are unambiguous by context.

The sequence score here will be computing $\LogitsDiffContributionAttnTok[]{\NonMaxOutputTok}$ for some fixed $\NonMaxOutputTok$ and $\MaxTok$.
The way we've set up our definitions, high scores predict  (and are thus bad), negative scores predict  (and are thus good), and more negative the scores, the stronger the prediction of .

Given a permutation $\sigma : \N^{<\nctx} \to \N^{<\nctx}$ of the  positions and two indices $0 \le i,j < \nctx$, define the swap permutation $\sigma_{i\leftrightarrow j}$ to be the permutation that is $\sigma$ except swapping $i$ and $j$:

\[
\sigma_{i\leftrightarrow j}(k) = \begin{cases}
\sigma(i) & \text{if }k = j \\
\sigma(j) & \text{if }k = i \\
\sigma(k) & \text{otherwise}
\end{cases}
\]

Fix a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N$.
Fix $\sigma : \N \to \N$ be a permutation of the $\nctx$ positions.
Fix indices $0 \le i,j < \nctx$.

Define $\Delta_{\sigma,i\leftrightarrow j}$ to be the difference in sequence scores when you swap $i$ and $j$:

\[
\Delta_{\sigma,i\leftrightarrow j} := s_{\sigma_{i\leftrightarrow j}} - s_\sigma
\]

Then there are two cases for $\sign\left(\Delta_{\sigma,i\leftrightarrow j}\right)$:

* If $a_{t_i} = a_{t_j}$ then $\sign\left(\Delta_{\sigma,i\leftrightarrow j}\right) = -  \sign\left(b_{\sigma(i)} - b_{\sigma(j)}\right)\sign\left(v_{t_i} - v_{t_j}\right)$.

* Otherwise, $\sign\left(\Delta_{\sigma,i\leftrightarrow j}\right) = \sign\left(a_{t_i} - a_{t_j}\right)
\sign\left(b_{\sigma(i)} - b_{\sigma(j)}\right)
\sign\left(
- \frac{v_{t_i} e^{a_{t_i}} - v_{t_j} e^{a_{t_j}}}{e^{a_{t_i}} - e^{a_{t_j}}}
\right)$.

Intuitively, <ref> says that, if the token contribution to attention is equal between tokens $t_i$ and $t_j$, then the impact of swapping their positions $\sigma(i)$ and $\sigma(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.

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

For the rest of the proof, we take $i \neq j$.

The proof proceeds just by algebraic manipulation with no deep insight.
We first list the facts we use, the proceed to computing $\sign\left(\Delta_{\sigma,i\leftrightarrow j}\right)$.
We abbreviate $\sigma_{i\leftrightarrow j}$ as $\sigma'$ for brevity.

\begin{align*}
\sign\left(e^{b_{\sigma(i)}} - e^{b_{\sigma(j)}}\right) & = \sign\left(b_{\sigma(i)} - b_{\sigma(j)}\right)
\end{align*}

\begin{align*}
\sign\left(\Delta_{\sigma,i\leftrightarrow j}\right)
& = \sign\left(s_{\sigma'} - s_\sigma\right) \\
& = \sign\left(\frac{\sum_{0 \le p < \nctx} v_{t_p} e^{a_{t_p} + b_{\sigma'(p)}}}{\sum_{0 \le p < \nctx} e^{a_{t_p} + b_{\sigma'(p)}}} - s_\sigma\right) \\
\intertext{Now multiply through by the denominator, which is positive}
& = \sign\left(\sum_{0 \le p < \nctx} v_{t_p} e^{a_{t_p} + b_{\sigma'(p)}} - s_\sigma\sum_{0 \le p < \nctx} e^{a_{t_p} + b_{\sigma'(p)}}\right) \\
& = \sign\left(\sum_{0 \le p < \nctx} v_{t_p} e^{a_{t_p} + b_{\sigma(p)}} - v_{t_i} e^{a_{t_i}}\left(e^{b_{\sigma(i)}} - e^{b_{\sigma'(i)}}\right) - v_{t_j} e^{a_{t_j}}\left(e^{b_{\sigma(j)}} - e^{b_{\sigma'(j)}}\right) \right. \\
& \phantom{{}=\sign\left(\vphantom{\sum_{0 \le p < \nctx}}\right.}\left.
{} - s_\sigma\sum_{0 \le p < \nctx} e^{a_{t_p} + b_{\sigma(p)}}
{} + s_\sigma e^{a_{t_i}}\left(e^{b_{\sigma(i)}} - e^{b_{\sigma'(i)}}\right)
{} + s_\sigma e^{a_{t_j}}\left(e^{b_{\sigma(j)}} - e^{b_{\sigma'(j)}}\right)
\right) \\
& = \sign\left(\cancel{\sum_{0 \le p < \nctx} v_{t_p} e^{a_{t_p} + b_{\sigma(p)}}} - v_{t_i} e^{a_{t_i}}\left(e^{b_{\sigma(i)}} - e^{b_{\sigma(j)}}\right) - v_{t_j} e^{a_{t_j}}\left(e^{b_{\sigma(j)}} - e^{b_{\sigma(i)}}\right) \right. \\
& \phantom{{}=\sign\left(\vphantom{\sum_{0 \le p < \nctx}}\right.}\left.
{} - \cancel{\sum_{0 \le p < \nctx} v_{t_p} e^{a_{t_p} + b_{\sigma(p)}}}
{} + s_\sigma e^{a_{t_i}}\left(e^{b_{\sigma(i)}} - e^{b_{\sigma(j)}}\right)
{} + s_\sigma e^{a_{t_j}}\left(e^{b_{\sigma(j)}} - e^{b_{\sigma(i)}}\right)
\right) \\
& = \sign\left(\left(v_{t_j} e^{a_{t_j}}- v_{t_i} e^{a_{t_i}}\right)\left(e^{b_{\sigma(i)}} - e^{b_{\sigma(j)}}\right)
+ s_\sigma \left(e^{a_{t_i}} - e^{a_{t_j}}\right)\left(e^{b_{\sigma(i)}} - e^{b_{\sigma(j)}}\right)
\right) \\
& = \sign\left(e^{b_{\sigma(i)}} - e^{b_{\sigma(j)}}\right)
\sign\left(\left(v_{t_j} e^{a_{t_j}}- v_{t_i} e^{a_{t_i}}\right)
+ s_\sigma \left(e^{a_{t_i}} - e^{a_{t_j}}\right)
\right) \\
& = \sign\left(b_{\sigma(i)} - b_{\sigma(j)}\right)
\sign\left(s_\sigma \left(e^{a_{t_i}} - e^{a_{t_j}}\right)
- \left(v_{t_i} e^{a_{t_i}} - v_{t_j} e^{a_{t_j}}\right)
\right) \\
\intertext{Divide through by non-zero values when possible}
& = \sign\left(b_{\sigma(i)} - b_{\sigma(j)}\right) \\
& \phantom{{}={}}\cdot
\begin{cases}
\sign\left(v_{t_i} - v_{t_j}\right) & \text{if }a_{t_i} = a_{t_j} \\
\sign\left(e^{a_{t_i}} - e^{a_{t_j}}\right)
\sign\left(s_\sigma
- \frac{v_{t_i} e^{a_{t_i}} - v_{t_j} e^{a_{t_j}}}{e^{a_{t_i}} - e^{a_{t_j}}}
\right)
& \text{otherwise}
\end{cases} \\
& = \begin{cases}
-\sign\left(b_{\sigma(i)} - b_{\sigma(j)}\right)\sign\left(v_{t_i} - v_{t_j}\right) & \text{if }a_{t_i} = a_{t_j} \\
\sign\left(a_{t_i} - a_{t_j}\right)
\sign\left(b_{\sigma(i)} - b_{\sigma(j)}\right)
\sign\left(
- \frac{v_{t_i} e^{a_{t_i}} - v_{t_j} e^{a_{t_j}}}{e^{a_{t_i}} - e^{a_{t_j}}}
\right)
& \text{otherwise}
\end{cases}
\end{align*}

Fix a set of fixed indices $F \subseteq \N^{<\nctx}$ and an assignment of token values to each of the fixed positions $t_F : F \to \N^{<\dvocab}$.
($F$ is the set of positions for which we are not pessimizing over the value of the token in that position.)
Fix a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N$.

Given a permutation $\sigma : \N^{<\nctx} \to \N^{\nctx}$, say that $\sigma$ fixes $F$ (relative to $t_0, \ldots, t_{\nctx-1}$) if $t_i = t_F(\sigma(i))$ whenever $\sigma(i) \in F$.

Note that in this section, for the cubic proofs, we will in fact generally take $F = \{\nctx-1\}$, so that we are fixing the final query token, though in thm:softmax-convexity-no-position,lem:softmax-convexity:purification,lem:softmax-convexity:pure-sorting $F$ will also contain all positions with the maximum token .
In <ref>, we will take $F=\emptyset$ 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.

Fix a set of fixed indices $F \subseteq \N^{<\nctx}$ and an assignment of token values to each of the fixed positions $t_F : F \to \N^{<\dvocab}$.

Define the position-sorting permutation fixing indices in $F$ $\sigma_s : \N^{<\nctx} \to \N^{<\nctx}$ to be the permutation that sorts the indices not in $F$ according to $b$:
for $0 \le i, j < \nctx$ with $i,j\not\in F$, $b_{i} \le b_{j}$ whenever $\sigma_s(i) < \sigma_s(j)$;
and $\sigma_s(i) = i$ for $i \in F$.

Fix a set of fixed indices $F \subseteq \N^{<\nctx}$ and an assignment of token values to each of the fixed positions $t_F : F \to \N^{<\dvocab}$.
Fix a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N$.

Say that the sequence represented by a permutation $\sigma : \N^{<\nctx} \to \N^{<\nctx}$ is contiguous on equal tokens if, for all $0 \le i, j, k < \nctx$ with $t_i = t_j \neq t_k$ and $i,j,k\not\in \sigma^{-1}(F)$, it is never the case that $\sigma_s(\sigma(i)) < \sigma_s(\sigma(k)) < \sigma_s(\sigma(j))$.

Fix a set of fixed indices $F \subseteq \N^{<\nctx}$ and an assignment of token values to each of the fixed positions $t_F : F \to \N^{<\dvocab}$.
Fix a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N$.

Let $\sigma_{\min}, \sigma_{\max} : \N \to \N$ be permutations of the  positions, fixing positions in $F$, satisfying the following property:
For all $\sigma : \N \to \N$ a permutation fixing $F$, we have

\begin{equation}\label{eqn:softmax-convexity:position-pessimization-order:sigma-extreme-ordering}
s_{\sigma_{\min}} \le s_\sigma \le s_{\sigma_{\max}}
\end{equation}

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

Then $\sigma_{\max}$ and $\sigma_{\min}$ may be taken to be contiguous on equal tokens.
That is, there exist $\sigma_{\max}$ and $\sigma_{\min}$ satisfying the property of <ref> which additionally satisfy the definition of <ref>.

The basic idea is that we will assume that one of $\sigma_{\max}$ and $\sigma_{\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.

Fix a set of fixed indices $F \subseteq \N^{<\nctx}$ and an assignment of token values to each of the fixed positions $t_F : F \to \N^{<\dvocab}$.
Fix a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N$.

Say that a permutation $\sigma$ is contiguous on equally-attended positions if, for all $0 \le i < \nctx$ with $i \not\in \sigma^{-1}(F)$, the sorting order according $\sigma_s$ on the contiguous block of positions with contribution to the attention score equal to that of $\sigma(i)$, $\setst{\sigma(j)}{b_{\sigma(j)} = b_{\sigma(i)}\text{ and }\sigma(j)\not\in F}$, is the same as the sorting order according to the fraction of tokens equal to $t_j$ with $b$-values greater than $b_{\sigma(i)}$, with ties broken by the value of $t_j$.
Equationally, this second sorting order is defined by the score

\[
\left.\left(\card{\setst{k}{t_k = t_j\text{ and }b_{\sigma(k)} > b_{\sigma(i)}\text{ and }\sigma(k)\not\in F}} + \frac{t_j}{\dvocab}\right) \middle/ \card{\setst{k}{t_k = t_j\text{ and }\sigma(k)\not\in F}} \right..
\]

Most importantly, any permutation that is contiguous on equally-attended positions has the property that for any indices $0 \le i, j, k < \nctx$ with $i,j,k\not\in \sigma^{-1}(F)$ and $t_i = t_j \neq t_k$ and $\sigma_s(\sigma(i)) < \sigma_s(\sigma(k)) < \sigma_s(\sigma(j))$, we will have the strict inequality $b_{\sigma(i)} < b_{\sigma(k)} < b_{\sigma(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_\sigma$.

Fix a set of fixed indices $F \subseteq \N^{<\nctx}$ and an assignment of token values to each of the fixed positions $t_F : F \to \N^{<\dvocab}$.
Fix a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N$.

Say that a permutation $\sigma$ is needlessly non-contiguous at $i$, $j$, $k$ (for $i,j,k \not \in \sigma^{-1}(F)$) if $\Delta_{\sigma,i\leftrightarrow k} = 0$ or $\Delta_{\sigma,j\leftrightarrow k} = 0$, for $0 \le i, j, k < \nctx$ with $i,j,k \not\in \sigma^{-1}(F)$ with $t_i = t_j \neq t_k$ and $\sigma_s(\sigma(i)) < \sigma_s(\sigma(k)) < \sigma_s(\sigma(j))$.

Say that a permutation $\sigma$ is needlessly non-contiguous if it is needlessly non-contiguous at any $i, j, k \not\in \sigma^{-1}(F)$.

Fix a set of fixed indices $F \subseteq \N^{<\nctx}$ and an assignment of token values to each of the fixed positions $t_F : F \to \N^{<\dvocab}$.
Fix a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N$.

Any needlessly non-contiguous sequence $\sigma$ which fixes $F$ can be made into a sequence $\sigma'$ which still fixes $F$ and is both simultaneously contiguous on equally-attended positions and not needlessly non-contiguous, and for which $s_\sigma = s_{\sigma'}$.

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

Otherwise, we have $\Delta_{\sigma,i\leftrightarrow k} = 0$ or $\Delta_{\sigma,j\leftrightarrow k} = 0$ for some $i$, $j$, $k$, for $0 \le i, j, k < \nctx$ with $i,j,n\not\in \sigma^{-1}(F)$ and $t_i = t_j \neq t_k$ and $\sigma_s(\sigma(i)) < \sigma_s(\sigma(k)) < \sigma_s(\sigma(j))$.
Since the sequence is contiguous on equally-attended positions, we have the strict inequality $b_{\sigma(i)} < b_{\sigma(k)} < b_{\sigma(j)}$.

By <ref>, we have two cases.
Noting that $t_i = t_j$, we can write them as

* $v_{t_k} = v_{t_i}$ and $a_{t_i} = a_{t_k}$

* $a_{t_i} \neq a_{t_k}$ and $s_\sigma = \frac{v_{t_i} e^{a_{t_i}} - v_{t_k} e^{a_{t_k}}}{e^{a_{t_i}} - e^{a_{t_k}}}$

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

In the second case, since swapping tokens does not change $s_\sigma$, 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 <ref> more precise.

Choose $\sigma_{\max}$ and $\sigma_{\min}$ to be contiguous on equally-attended positions and not needlessly non-contiguous, and suppose that we have $\sigma \in \{\sigma_{\max}, \sigma_{\min}\}$ such that for some $0 \le i, j, k < \nctx$ with $i,j,k\not\in \sigma^{-1}(F)$ and $t_i = t_j \neq t_k$, we have $b_{\sigma(i)} < b_{\sigma(k)} < b_{\sigma(j)}$.
We will derive a contradiction with the presumption that $\sigma$ 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 $\sigma'_0$ to be $\sigma$ but swapping $i$ and $k$, and take $\sigma'_1$ to be $\sigma$ but swapping $j$ and $k$.

Now we will consider the cases for the sign of the score difference $\Delta_0 := s_{\sigma'_0} - s_\sigma$ and $\Delta_1 := s_{\sigma'_1} - s_\sigma$.
By the presumption of not being needlessly non-contiguous, $\Delta_z \neq 0$ for $z\in\{0,1\}$.
If we can show that the sign of $\Delta_0$ is distinct from the sign of $\Delta_1$, then we will have a contradiction with extremality because we will have either $s_{\sigma'_0} < s_\sigma < s_{\sigma'_1}$ or $s_{\sigma'_1} < s_\sigma < s_{\sigma'_0}$.
That is, we would be able to swap $i\leftrightarrow k$ and $j\leftrightarrow k$ to get a lower and higher score, making $\sigma$ not extremal.

Noting that $t_i = t_j$,

\begin{align*}
\sign\left(\Delta_0\right)
& = \sign\left(b_{\sigma(i)} - b_{\sigma(k)}\right)
\begin{cases}
\sign\left(v_{t_k} - v_{t_i}\right) & \text{if }a_{t_i} = a_{t_k} \\
\sign\left(a_{t_i} - a_{t_k}\right)
\sign\left(
- \frac{v_{t_i} e^{a_{t_i}} - v_{t_k} e^{a_{t_k}}}{e^{a_{t_i}} - e^{a_{t_k}}}
\right)
& \text{otherwise}
\end{cases} \\
\sign\left(\Delta_1\right)
& = \sign\left(b_{\sigma(j)} - b_{\sigma(k)}\right)
\begin{cases}
\sign\left(v_{t_k} - v_{t_i}\right) & \text{if }a_{t_i} = a_{t_k} \\
\sign\left(a_{t_i} - a_{t_k}\right)
\sign\left(
- \frac{v_{t_i} e^{a_{t_i}} - v_{t_k} e^{a_{t_k}}}{e^{a_{t_i}} - e^{a_{t_k}}}
\right)
& \text{otherwise}
\end{cases}
\end{align*}

Noting that the product is non-zero by presumption, that right multiplicand is equal for $\Delta_0$ and $\Delta_1$, and $\sign\left(b_{\sigma(i)} - b_{\sigma(k)}\right) = -1$ and $\sign\left(b_{\sigma(j)} - b_{\sigma(k)}\right) = 1$, we have our desired contradiction.

Note that the proof of <ref> 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.

Given a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N^{<\dvocab}$ and a permutation $\sigma : \N^{<\nctx} \to \N^{<\nctx}$ define the full sequence score $s'_{t_0,\ldots,t_{\nctx-1},\sigma}$ as:

\[
s_{t_0,\ldots,t_{\nctx-1},\sigma} := \left. \sum_{\mathclap{0 \le i < \nctx}} (v_{t_i} + w_{\sigma(i)}) e^{a_{t_i} + b_{\sigma(i)}} \middle/ \hphantom{0}\sum_{\mathclap{0 \le i < \nctx}} e^{a_{t_i} + b_{\sigma(i)}} \right.
\]

We will drop the token subscript, writing only $s'_\sigma$, when the token values are unambiguous by context.

The sequence score here will be computing $\LogitsDiffContributionAttn[]{\NonMaxOutputTok} := \LogitsDiffContributionAttnTok[]{\NonMaxOutputTok} + \LogitsDiffContributionAttnPos[]{\NonMaxOutputTok}$ for some fixed $\NonMaxOutputTok$ and $\MaxTok$.
As with <ref>, with the way we've set up our definitions, high scores predict  (and are thus bad), negative scores predict  (and are thus good), and more negative the scores, the stronger the prediction of .

Fix a set of fixed indices $F \subseteq \N^{<\nctx}$ and an assignment of token values to each of the fixed positions $t_F : F \to \N^{<\dvocab}$.
Fix a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N$.
Let $\sigma_{\min}, \sigma_{\max} : \N \to \N$ be as in <ref>.

Define relaxed extremal sequence scores $r_{\sigma_{\max}}$, $r_{\sigma_{\min}}$:

\begin{align*}
r_{\sigma_{\min}} & := s_{\sigma_{\min}} + \min_{0 \le i < \nctx} w_i \\
r_{\sigma_{\max}} & := s_{\sigma_{\max}} + \max_{0 \le i < \nctx} w_i
\end{align*}

Then $r_{\sigma_{\min}} \le s'_{\sigma_{\min}}$ and $s'_{\sigma_{\max}} \le r_{\sigma_{\max}}$.

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

\begin{align*}
s'_\sigma & = \left. \sum_{0 \le i < \nctx} (v_{t_i} + w_{\sigma(i)}) e^{a_{t_i} + b_{\sigma(i)}} \middle/ \sum_{0 \le i < \nctx} e^{a_{t_i} + b_{\sigma(i)}} \right. \\
& = \frac{\sum_i v_{t_i} e^{a_{t_i} + b_{\sigma(i)}}}{\sum_i e^{a_{t_i} + b_{\sigma(i)}}}
+ \frac{\sum_i w_{\sigma(i)} e^{a_{t_i} + b_{\sigma(i)}}}{\sum_i e^{a_{t_i} + b_{\sigma(i)}}} \\
& = s_\sigma + \frac{\sum_i w_{\sigma(i)} e^{a_{t_i} + b_{\sigma(i)}}}{\sum_i e^{a_{t_i} + b_{\sigma(i)}}} \\
& = s_\sigma + \sum_i w_{\sigma(i)} \frac{e^{a_{t_i} + b_{\sigma(i)}}}{\sum_j e^{a_{t_i} + b_{\sigma(i)}}}
\end{align*}

\begin{align*}
s_\sigma + \sum_i w_{\sigma(i)} \frac{e^{a_{t_i} + b_{\sigma(i)}}}{\sum_j e^{a_{t_i} + b_{\sigma(i)}}}
& = s'_\sigma
= s_\sigma + \sum_i w_{\sigma(i)} \frac{e^{a_{t_i} + b_{\sigma(i)}}}{\sum_j e^{a_{t_i} + b_{\sigma(i)}}} \\
s_\sigma + \min_k w_{\sigma(k)} \sum_i \frac{e^{a_{t_i} + b_{\sigma(i)}}}{\sum_j e^{a_{t_i} + b_{\sigma(i)}}}
& \le s'_\sigma
\le s_\sigma + \max_k w_{\sigma(k)} \sum_i \frac{e^{a_{t_i} + b_{\sigma(i)}}}{\sum_j e^{a_{t_i} + b_{\sigma(i)}}}
\end{align*}

\[
\sum_i \frac{e^{a_{t_i} + b_{\sigma(i)}}}{\sum_j e^{a_{t_i} + b_{\sigma(i)}}}
= \frac{\sum_i e^{a_{t_i} + b_{\sigma(i)}}}{\sum_j e^{a_{t_i} + b_{\sigma(i)}}}
= 1
\]

we get

\[
s_\sigma + \min_k w_{\sigma(k)}
\le s'_\sigma
\le s_\sigma + \max_k w_{\sigma(k)}
\]

and hence $r_{\sigma_{\min}} \le s'_{\sigma_{\min}}$ and $s'_{\sigma_{\max}} \le r_{\sigma_{\max}}$ as desired.

Fix a set of fixed indices $F \subseteq \N^{<\nctx}$ and an assignment of token values to each of the fixed positions $t_F : F \to \N^{<\dvocab}$.
Fix a set $S \subseteq \N^{< \dvocab}$ of valid other tokens in the sequence.
(In our uses of this theorem, $S$ will be the largest subset of $\N^{\le\MaxTok}$ for which we can guarantee that the model behaves correctly on all sequences compatible with $F$ and  and with tokens otherwise drawn from $S$.)

Define a comparison on non-negative integers less than :

\begin{align*}
c & := \sum_{i\in F} v_{t_F(i)} e^{a_{t_F(i)} + b_{i}}
d & := \sum_{i\in F} e^{a_{t_F(i)} + b_{i}}
f & := \sum_{\mathclap{\substack{0 \le i < \nctx \\ i \not\in F}}} e^{b_{i}}
\end{align*}

\[
\mathop{\mathrm{cmp}}(x, y) := \sign\left(d (e^{a_x} v_x - e^{a_y} v_y) - c(e^{a_x} - e^{a_y}) + f e^{a_x + a_y} \left(v_x e^{a_x+a_y}-v_y e^{a_x+a_y}\right)\right)
\]

Let $t_{\mathrm{cmp}\min}$ and $t_{\mathrm{cmp}\max}$ be the minimum and maximum elements of $S$ according to $\mathop{\mathrm{cmp}}$.[
We will prove that $\mathop{\mathrm{cmp}}$ is transitive in the process of proving this theorem.

For a given choice of a non-decreasing sequence of tokens $t_0 \le \cdots \le t_{\nctx-1} \in \N$ compatible with $F$ and $S$ and a given choice of permutation $\sigma : \N \to \N$ of the  positions fixing $F$ ($t_i = t_F(\sigma(i))$ for $\sigma(i) \in F$; and $t_i \in S$ for $\sigma(i) \not\in F$):
let $s_{\sigma,\min}$ (and $s_{\sigma,\max}$) denote $s_{t_0,\ldots,t_{\nctx-1},\sigma}$ when $t_i = t_{\mathrm{cmp}\min}$ for all $\sigma(i) \not\in F$ (or $t_{\mathrm{cmp}\max}$, respectively).

Then for all such choices of sequence-permutation pairs,

\[
s_{\sigma,\min} \le s_{t_0,\ldots,t_{\nctx-1},\sigma} \le s_{\sigma,\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.

Fix all the same quantities as in <ref>.

For any indices $0 \le i < j < \nctx$, token values $x, y \in S$, the score for a sequence with $t_i = x \neq y = t_j$ is bounded on both sides by sequences with $t_i = t_j = x$ and $t_i = t_j = y$.

Let $s_{\alpha,\beta}$ be the sequence score with $t_i = \alpha$ and $t_j = \beta$, and define the score differences $\Delta_x := s_{x,x} - s_{x,y}$ and $\Delta_y := s_{y,y} - s_{x,y}$.
It suffices to show that $\sign(\Delta_x\Delta_y) \le 0$.
To show this, we must only compute the sign of $\Delta_\alpha$ for $\alpha \in \{x,y\}$ and show that whenever both $\Delta_x$ and $\Delta_y$ are non-zero, they have opposite signs.

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

\begin{align*}
C & := \sum_{\mathclap{\substack{0 \le k < \nctx \\ k \neq i,j}}} v_{t_k} e^{a_{t_k} + b_{\sigma(k)}}
D & := \sum_{\mathclap{\substack{0 \le k < \nctx \\ k \neq i,j}}} e^{a_{t_k} + b_{\sigma(k)}}
\end{align*}

\begin{align*}
\tilde{\alpha} & := \begin{cases} x & \text{if }\alpha = y \\ y & \text{if }\alpha = x \end{cases}
i_\alpha & := \begin{cases} i & \text{if }\alpha = x \\ j & \text{if }\alpha = y \end{cases}
i_{\tilde{\alpha}} & := \begin{cases} i & \text{if }\tilde{\alpha} = x \\ j & \text{if }\tilde{\alpha} = y \end{cases}
\end{align*}

\begin{align*}
\sign\left(\Delta_\alpha\right)
& = \sign\left(
\frac{v_{\alpha} e^{a_{\alpha} + b_{\sigma(i)}} + v_{\alpha} e^{a_{\alpha} + b_{\sigma(j)}} + C}{e^{a_{\alpha} + b_{\sigma(i)}} + e^{a_{\alpha} + b_{\sigma(j)}} + D}
- \frac{v_{x} e^{a_{x} + b_{\sigma(i)}} + v_{y} e^{a_{y} + b_{\sigma(j)}} + C}{e^{a_{x} + b_{\sigma(i)}} + e^{a_{y} + b_{\sigma(j)}} + D}
\right) \\
& = \sign\left(
\frac{v_{\alpha} e^{a_{\alpha} + b_{\sigma(i_\alpha)}} + v_{\alpha} e^{a_{\alpha} + b_{\sigma(i_{\tilde{\alpha}})}} + C}{e^{a_{\alpha} + b_{\sigma(i_\alpha)}} + e^{a_{\alpha} + b_{\sigma(i_{\tilde{\alpha}})}} + D}
- \frac{v_{\alpha} e^{a_{\alpha} + b_{\sigma(i_\alpha)}} + v_{\tilde{\alpha}} e^{a_{\tilde{\alpha}} + b_{\sigma(i_{\tilde{\alpha}})}} + C}{e^{a_{\alpha} + b_{\sigma(i_\alpha)}} + e^{a_{\tilde{\alpha}} + b_{\sigma(i_{\tilde{\alpha}})}} + D}
\right) \\
\intertext{Multiply through by positive denominators and simplify}
& = \sign\left(C \left(e^{a_{\tilde{\alpha }}+b_{\sigma
\left(i_{\tilde{\alpha }}\right)}}-e^{b_{\sigma
\left(i_{\tilde{\alpha }}\right)}+a_{\alpha }}\right)+D
\left(v_{\alpha } e^{b_{\sigma \left(i_{\tilde{\alpha
}}\right)}+a_{\alpha }}-v_{\tilde{\alpha }}
e^{a_{\tilde{\alpha }}+b_{\sigma \left(i_{\tilde{\alpha
\right. \\
&\phantom{{}=\sign\left(\vphantom{e^{a_{\tilde{\alpha }}}}\right.}\left.
{}+v_{\alpha } \left(e^{b_{\sigma
\left(i_{\tilde{\alpha }}\right)}}+e^{b_{\sigma
\left(i_{\alpha }\right)}}\right) e^{a_{\tilde{\alpha
}}+b_{\sigma \left(i_{\tilde{\alpha }}\right)}+a_{\alpha
}}-v_{\tilde{\alpha }} \left(e^{b_{\sigma
\left(i_{\tilde{\alpha }}\right)}}+e^{b_{\sigma
\left(i_{\alpha }\right)}}\right) e^{a_{\tilde{\alpha
}}+b_{\sigma \left(i_{\tilde{\alpha }}\right)}+a_{\alpha
}}\right) \\
\intertext{Pulling out $e^{b_{\sigma(i_{\tilde{\alpha}})}}$}
& = \sign\left(
e^{a_{\tilde{\alpha }}+a_{\alpha }}
\left(e^{b_{\sigma \left(i_{\tilde{\alpha}}\right)}}+e^{b_{\sigma\left(i_{\alpha}\right)}}\right)
\left(v_{\alpha } - v_{\tilde{\alpha }}\right)
\left(e^{a_{\tilde{\alpha }}}-e^{a_{\alpha }}\right)+D
\left(e^{a_{\alpha }} v_{\alpha }-e^{a_{\tilde{\alpha }}}
v_{\tilde{\alpha }}\right)\right)
\end{align*}

Note that swapping $\alpha$ and $\tilde{\alpha}$ negates the sign.

Hence, we have $\sign(\Delta_x) = -\sign(\Delta_y)$ and hence $s_{x,x} \le s_{x,y} \le s_{y,y}$ or $s_{y,y} \le s_{x,y} \le s_{x,x}$ as desired.

Fix all the same quantities as in <ref>.

Fix tokens $x, y \in S$.
Let $n := \nctx-\card{F}$ be the number of non-fixed tokens.
Fix sequences with $n$ copies of $x$ and $y$ respectively: fix $t_{x,0} \le \cdots \le t_{x,\nctx-1} \in \N$ and $t_{y,0} \le \cdots \le t_{y,\nctx-1} \in \N$ compatible with $F$ and $S$ and given choices of permutations $\sigma_x, \sigma_y : \N \to \N$ of the  positions fixing $F$: $t_{x,i} = t_F(\sigma_x(i))$ for $\sigma_x(i) \in F$; $t_{y,i} = t_F(\sigma_y(i))$ for $\sigma_y(i) \in F$; $t_{x,i} = x$ for $\sigma_x(i) \not\in F$; and $t_{y,i} = y$ for $\sigma_y(i) \not\in F$.

\[
\sign(
- (s_{\sigma_y,t_{y,0},\ldots,t_{y,\nctx-1}})
= \mathop{\mathrm{cmp}}(x, y)
\]

The proof goes by straightforward computation.

\begin{align*}
& \sign(
- (s_{\sigma_y,t_{y,0},\ldots,t_{y,\nctx-1}})
) \\
& = \sign\left(
\frac{v_x e^{a_x} f + c}{e^{a_x} f + d}
- \frac{v_y e^{a_y} f + c}{e^{a_y} f + d}
\right) \\
\intertext{Multiply through by non-negative denominators}
& = \sign\left(
\left(v_x e^{a_x} f + c\right)\left(e^{a_y} f + d\right)
- \left(v_y e^{a_y} f + c\right)\left(e^{a_x} f + d\right)
\right) \\
& = \sign\left(-c f e^{a_x}+c f e^{a_y}+d f v_x e^{a_x}-d f v_y e^{a_y}+f^2 v_x
e^{a_x+a_y}-f^2 v_y e^{a_x+a_y}\right) \\
\intertext{Use $f > 0$}
& = \sign\left(-c e^{a_x}+c e^{a_y}+d v_x e^{a_x}-d v_y e^{a_y}+f v_x
e^{a_x+a_y}-f v_y e^{a_x+a_y}\right) \\
& = \sign\left(c \left(e^{a_y}-e^{a_x}\right)+d \left(v_x e^{a_x}-v_y
e^{a_y}\right)+f \left(v_x e^{a_x+a_y}-v_y e^{a_x+a_y}\right)\right) \\
& = \mathop{\mathrm{cmp}}(x, y)
\end{align*}

Define the relation $\leq_{\mathop{\mathrm{cmp}}}$ by $x \leq_{\mathop{\mathrm{cmp}}} y$ if and only if $\mathop{\mathrm{cmp}}(x,y) \in \{-1, 0\}$.
The relation $\leq_{\mathop{\mathrm{cmp}}}$ is always transitive.

Note that by <ref>, $\mathop{\mathrm{cmp}}$ is comparing two sequence scores.
Since $\leq$ is transitive over the reals, the relation $\leq_{\mathop{\mathrm{cmp}}}$ is also transitive.

Finally, we combine the previous lemmas to complete our proof of <ref>:

Extremal sequences with scores $s_{\sigma, \min}$ and $s_{\sigma, \max}$ are guaranteed to exist because there are only finitely many elements of $S$ and therefore only finitely many sequences.
By <ref>, the extremal sequences must be pure (have $t_i = t_j$ whenever $\sigma(i),\sigma(j) \not\in F$).
By <ref>, the extremal sequences must have tokens that are extremal according to $\mathop{\mathrm{cmp}}$.

\begin{align*}
\EPQKE(\InputSequence[-1], \InputSequence[\NonMaxInputTokIndex]) & := \InputSequence[-1]\qWE \WQ \WK^T \barWE^T\InputSequence[\NonMaxInputTokIndex]^T / \sqrt{\dhead} \\
\EPQKP(\InputSequence[-1], \NonMaxInputTokIndex) & := \InputSequence[-1]\qWE \WQ \WK^T \hatWpos_i^T / \sqrt{\dhead} \\
% \PreSoftmaxAttention(\InputSequence[-1], \InputSequence[\NonMaxInputTokIndex], i)
% & := \EPQKE(\InputSequence[-1], \InputSequence[\NonMaxInputTokIndex]) + \EPQKP(\InputSequence[-1], i) \\
% \PostSoftmaxAttention[\NonMaxInputTokIndex](\InputSequence) & := \frac{1}{\sum_{\NonMaxInputTokIndexAlt} e^{\alpha(\InputSequence[-1], \InputSequence[\NonMaxInputTokIndexAlt], i)}} e^{\alpha(\InputSequence[-1], \InputSequence[\NonMaxInputTokIndex], \NonMaxInputTokIndex)} \\
\EVOU(\InputSequence[\NonMaxInputTokIndex]) & := \InputSequence[\NonMaxInputTokIndex]\barWE\WV\WO\WU \\
\PVOU(\NonMaxInputTokIndex) & := \hatWpos_{\NonMaxInputTokIndex}\WV\WO\WU \\
% \EPVOU(\InputSequence[\NonMaxInputTokIndex], \NonMaxInputTokIndex) & := \EVOU(\InputSequence[\NonMaxInputTokIndex]) + \PVOU(\NonMaxInputTokIndex) \\
\LogitsContributionSkip({\InputSequence[-1]}) & := \InputSequence[-1]\qEU \\
% \LogitsContributionAttn[\NonMaxInputTokIndex](\InputSequence) & := \PostSoftmaxAttention[\NonMaxInputTokIndex](\InputSequence)\EPVOU(\InputSequence[\NonMaxInputTokIndex], \NonMaxInputTokIndex) \\
% \Logits(\InputSequence) & := \LogitsContributionSkip(\InputSequence) + \sum_{\NonMaxInputTokIndex=0}^3 \LogitsContributionAttn[\NonMaxInputTokIndex](\InputSequence) \\
% \Model(\InputSequence) & := \text{softmax}\left(\Logits(\InputSequence)\right) \\
% \LogitsDiff[t](\InputSequence) & := \Logits(\InputSequence)_t - \Logits(\InputSequence)_{\max_i \InputSequence[\NonMaxInputTokIndex]} \\
\LogitsDiffContributionSkip[\NonMaxOutputTok]({\InputSequence[-1]}, {\max_{\NonMaxInputTokIndex} \InputSequence[\NonMaxInputTokIndex]}) & := \LogitsContributionSkip({\InputSequence[-1]})_{\NonMaxOutputTok} - \LogitsContributionSkip({\InputSequence[-1]})_{\max_{\NonMaxInputTokIndex} \InputSequence[\NonMaxInputTokIndex]}% \\
% \LogitsDiffContributionAttn[\NonMaxInputTokIndex]{t}(\InputSequence) & := \LogitsContributionAttn[\NonMaxInputTokIndex](\InputSequence)_t - \LogitsContributionAttn[\NonMaxInputTokIndex](\InputSequence)_{\max_{\NonMaxInputTokIndexAlt} \InputSequence[\NonMaxInputTokIndexAlt]} \\
% \LogitsDiff(\InputSequence) & = \max_{i \neq \max_j \InputSequence[j]} \LogitsDiff[\NonMaxInputTokIndex](\InputSequence) \\
% \LogitsDiffContributionSkip(\InputSequence) & = \max_{\NonMaxInputTokIndex \neq \max_{\NonMaxInputTokIndexAlt} \InputSequence[\NonMaxInputTokIndexAlt]} \LogitsDiffContributionSkip[\NonMaxInputTokIndex](\InputSequence) \\
% \LogitsDiffContributionAttn[k]{}(\InputSequence) & = \max_{i \neq \max_j \InputSequence[j]} \LogitsDiffContributionAttn[k]{i}(\InputSequence)
\end{align*}

Recapitulation of some relevant definitions from <ref>, parameterized by the arguments they actually depend on.

Counting Correct Sequences in Cubic Time: Preliminaries
[1]
correctness, input-sequence
model-behavior, input-sequencemaxinput-sequence
model-behavior, input-sequence
input-sequence is a tensor of shape $(\nctx, )$ with values in $\N^{<\dvocab}$
$\MaxTok{} \gets \Call{max}{\algorithmiclongvariable{input-sequence}}$
$\MaxTok{} \gets \algorithmiclongvariable{max-token}$
$\InputSequence \gets \algorithmiclongvariable{input-sequence}$
$\Vector{\algorithmiclongvariable{skip-score}}[\NonMaxOutputTok] \gets \Call{\LogitsDiffContributionSkip[\NonMaxOutputTok]}{\InputSequence[\nctx-1], \MaxTok}$
$\algorithmiclongvariable{attn-weights-unscaled}_{\NonMaxInputTokIndex} \gets \Call{\EPQKE}{\InputSequence[\nctx-1],\InputSequence[\NonMaxInputTokIndex]} + \Call{\EPQKP}{\InputSequence[\nctx-1], \NonMaxInputTokIndex}$
$\algorithmiclongvariable{attn-weights} \gets \Call{softmax}{\algorithmiclongvariable{attn-weights-unscaled} / \sqrt{\dhead}}$
$\Vector{v}[\NonMaxInputTok] \gets \Call{\EVOU}{\NonMaxInputTok}$
$\Vector{w}[\NonMaxInputTokIndex] \gets \Call{\PVOU}{\NonMaxInputTokIndex}$
$\Vector{\Delta v}[\NonMaxInputTok, \NonMaxOutputTok] \gets \Vector{v}[\NonMaxInputTok, \NonMaxOutputTok] - \Vector{v}[\NonMaxInputTok, \MaxTok{}]$
$\Vector{\Delta w}[\NonMaxInputTokIndex, \NonMaxOutputTok] \gets \Vector{w}[\NonMaxInputTokIndex, \NonMaxOutputTok] - \Vector{w}[\NonMaxInputTokIndex, \MaxTok{}]$
$\max_{{\NonMaxOutputTok \neq \MaxTok{}}} (\Vector{\algorithmiclongvariable{skip-score}}[i] + \sum_{\NonMaxInputTokIndex=0}^{\nctx-1} (\Vector{\Delta v}[\NonMaxInputTokIndex, \NonMaxOutputTok] + \Vector{\Delta w}[\NonMaxInputTokIndex, \NonMaxOutputTok]) \cdot \Vector{\algorithmiclongvariable{attn-weights}}[\NonMaxInputTokIndex])$
correctness-pessimizing-over-position-slow, input-sequence
$\InputSequence \gets \algorithmiclongvariable{input-sequence}$
allcorrectness, perm + [[-1]]    perm $\in$ permutations[0:-1]

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

[1]
model-behavior-relaxed, query-tok, max-tok, non-max-tok, n-copies-nonmax
$\QueryTok \gets \algorithmiclongvariable{query-tok}$,
$\MaxTok \gets \algorithmiclongvariable{max-tok}$,
$\NonMaxNonQueryTok \gets \algorithmiclongvariable{non-max-tok}$,
$\NumCopiesNonMaxNonQueryTok \gets \algorithmiclongvariable{n-copies-nonmax}$
$0\le \QueryTok \le \MaxTok < \dvocab$,
$0\le \NonMaxNonQueryTok \le \MaxTok < \dvocab$,
$0\le \NumCopiesNonMaxNonQueryTok < \nctx$
n-copies-nonmax${}=0$  non-max-tok${}={}$max-tok
query-tok $\neq$ max-tok n-copies-nonmax $< \nctx-1$