$\Return \ge \Call{model-behavior}{\Model, \InputSequence}$  with specified ,  copies of  in non-query positions, and the remainder of the tokens equal to $\Vector{\algorithmiclongvariable{skip-score}}[\NonMaxOutputTok] \gets \Call{\LogitsDiffContributionSkip[\NonMaxOutputTok]}{\QueryTok, \MaxTok}$
Cache by , , $\NonMaxOutputTok$
$\Vector{w}[\NonMaxInputTokIndex] \gets \Call{\PVOU}{\NonMaxInputTokIndex}$    $0 \le \NonMaxInputTokIndex < \nctx$
Cache by $\NonMaxInputTokIndex$

$\Vector{\Delta w}[\max, \NonMaxOutputTok] \gets \max_{0 \le \NonMaxInputTokIndex < \nctx}(\Vector{w}[\NonMaxInputTokIndex, \NonMaxOutputTok] - \Vector{w}[\NonMaxInputTokIndex, \MaxTok{}])$
Cache by , $\NonMaxOutputTok$
$\Vector{v}[\NonMaxInputTok] \gets \Call{\EVOU}{\NonMaxInputTok}$
$\Vector{\Delta v}[\NonMaxInputTok, \NonMaxOutputTok] \gets \Vector{v}[\NonMaxInputTok, \NonMaxOutputTok] - \Vector{v}[\NonMaxInputTok, \MaxTok{}]$    $\NonMaxInputTok \in \{\QueryTok, \MaxTok, \NonMaxNonQueryTok\}$
Cache by , $\NonMaxInputTok$, $\NonMaxOutputTok$
$\Vector{a}[\NonMaxInputTok] \gets \Call{\EPQKE}{\QueryTok,\NonMaxInputTok} / \sqrt{\dhead}$    $\NonMaxInputTok \in \{\QueryTok, \MaxTok, \NonMaxNonQueryTok\}$
Cache by , $\NonMaxInputTok$
$\Vector{b}[\nctx-1] \gets \Call{\EPQKP}{\QueryTok, \nctx-1} / \sqrt{\dhead}$
Cache by $\Vector{b}[:{-1}] \gets \Call{sort}{\Call{\EPQKP}{\QueryTok, :{-1}}} / \sqrt{\dhead}$
Cache by , $\NonMaxInputTokIndex$
$\algorithmiclongvariable{attn-weights-unscaled}_{:,\nctx-1} \gets \Vector{a}[\QueryTok] + \Vector{b}[\nctx-1]$
Cache by $\algorithmiclongvariable{attn-weights-unscaled}_{0,\NonMaxInputTokIndex} \gets \Vector{a}[\MaxTok] + \Vector{b}[\NonMaxInputTokIndex - \NumCopiesNonMaxNonQueryTok]$  $\NumCopiesNonMaxNonQueryTok \le \NonMaxInputTokIndex < \nctx - 1$
Cache by , , $\NonMaxInputTokIndex$, $\algorithmiclongvariable{attn-weights-unscaled}_{1,\NonMaxInputTokIndex} \gets \Vector{a}[\MaxTok] + \Vector{b}[\NonMaxInputTokIndex]$  $\NumCopiesNonMaxNonQueryTok \le \NonMaxInputTokIndex < \nctx - 1$
Cache by , , $\NonMaxInputTokIndex$, $\algorithmiclongvariable{attn-weights-unscaled}_{0,\NonMaxInputTokIndex} \gets \Vector{a}[\NonMaxNonQueryTok] + \Vector{b}[\NumCopiesNonMaxNonQueryTok+\NonMaxInputTokIndex]$  $0 \le \NonMaxInputTokIndex < \nctx - 1 - \NumCopiesNonMaxNonQueryTok$
Cache by , , $\NonMaxInputTokIndex$, $\algorithmiclongvariable{attn-weights-unscaled}_{1,\NonMaxInputTokIndex} \gets \Vector{a}[\NonMaxNonQueryTok] + \Vector{b}[\NonMaxInputTokIndex]$  $0 \le \NonMaxInputTokIndex < \nctx - 1 - \NumCopiesNonMaxNonQueryTok$
Cache by , , $\NonMaxInputTokIndex$, $\algorithmiclongvariable{attn-weights}_{0} \gets \Call{softmax}{\algorithmiclongvariable{attn-weights-unscaled}_{0}}$
Cache by , , , $\NonMaxInputTokIndex$, $\algorithmiclongvariable{attn-weights}_{1} \gets \Call{softmax}{\algorithmiclongvariable{attn-weights-unscaled}_{1}}$
Cache by , , , $\NonMaxInputTokIndex$, $c=0$
In this case, $\algorithmiclongvariable{attn-weights}_{0,\NonMaxInputTokIndex} = \algorithmiclongvariable{attn-weights}_{1,\NonMaxInputTokIndex}$, so we drop the first subscript

alg:cubic:full:line:return-c-zero:line:value$\max_{{\NonMaxOutputTok \neq \MaxTok{}}} (\Vector{\algorithmiclongvariable{skip-score}}[\NonMaxOutputTok] + \Vector{\Delta w}[\max, \NonMaxOutputTok] + \Vector{\Delta v}[\QueryTokShort, \NonMaxOutputTok] \Vector{\algorithmiclongvariable{attn-weights}}[\QueryIndexShort] + \Vector{\Delta v}[\MaxTok, \NonMaxOutputTok] \sum_{\NonMaxInputTokIndex=0}^{\nctx-2} \Vector{\algorithmiclongvariable{attn-weights}}[\NonMaxInputTokIndex])$

$\Delta v_{\NonMaxInputTokIndex,\NonMaxOutputTok} \gets \Vector{\Delta v}[\MaxTok, \NonMaxOutputTok]$  $\NumCopiesNonMaxNonQueryTok \le \NonMaxInputTokIndex < \nctx - 1$
$\Delta v_{\NonMaxInputTokIndex,\NonMaxOutputTok} \gets \Vector{\Delta v}[\NonMaxNonQueryTok, \NonMaxOutputTok]$  $0 \le \NonMaxInputTokIndex < \NumCopiesNonMaxNonQueryTok$
$\Delta v_{\nctx-1,\NonMaxOutputTok} \gets \Vector{\Delta v}[\QueryTok, \nctx-1]$

relaxed-correctness-pessimizing-over-position, , , ,
runs the model on a relaxed variant of input sequences compatible with the arguments
is False correctness-pessimizing-over-position-slow,  is False any  with specified ,  copies of  in non-query positions, and the remainder of the tokens equal to $\Call{model-behavior-relaxed}{\Model, \QueryTok, \MaxTok, \NonMaxNonQueryTok, \NumCopiesNonMaxNonQueryTok} < 0$
cubic, ,
$\algorithmiclongvariable{count} \gets 0$
# of correct sequences
$\in$ range
$\MaxTok{} \gets \algorithmiclongvariable{max-token}$
$0 \le \QueryTok \le \MaxTok$
$\QueryTok{} \gets \algorithmiclongvariable{query-token}$
$\NumCopiesNonMaxNonQueryTok_{\max} \gets \nctx-1 ~\algorithmicif~ \QueryTok{} = \MaxTok{} ~\algorithmicelse~ \nctx-2$
maximum copies of nonmax
$0 \le \NumCopiesNonMaxNonQueryTok \le \NumCopiesNonMaxNonQueryTok_{\max}$
number of valid choices for the non-max token
$\Call{RCPOP}{\ensuremath{\vec{\chi}}} \gets \Call{relaxed-correctness-pessimizing-over-position}{\Model, \vec{\chi}}$
$\algorithmiclongvariable{t-count} \gets %\begin{cases}
1 %&
~\algorithmicif~ \Call{RCPOP}{\QueryTok, \MaxTok, \MaxTok, 0}%\\
% 0 & \algorithmicelse
~\algorithmicelse~ 0
% \end{cases}
$\algorithmiclongvariable{t-count} \gets \sum_{\NonMaxNonQueryTok=0}^{%\mathclap
% \begin{cases}
1 %&
~\algorithmicif~ \Call{RCPOP}{\QueryTok, \MaxTok, \NonMaxNonQueryTok, \NumCopiesNonMaxNonQueryTok}%\\
% 0 & \algorithmicelse
~\algorithmicelse~ 0
% \end{cases}
$\algorithmiclongvariable{count} \gets \algorithmiclongvariable{count} + \sum_{i=0}^{\NumCopiesNonMaxNonQueryTok}\binom{\nctx-1}{i} \cdot (\algorithmiclongvariable{t-count})^i$

taking $0^0=1$ conventionally

\[
\E_{\DrawnFrom{\InputSequence}{\UniformZeroN{\dvocab}}^{\nctx}}\left[
\argmax_i \Vector{(\Model(\InputSequence))}[i] = \max_{\NonMaxInputTokIndex} \InputSequence[\NonMaxInputTokIndex]
\right]
\ge
\Call{cubic}{\dvocab, \nctx, \Model}
\]

Apply the previous theorems and lemmas to <ref>.

The running time of <ref>, after using caching to avoid duplicate computations, is $\mathcal{O}(\dvocab^3 \nctx^2)$.

The nested loops in cubic execute the innermost body $\mathcal{O}(\dvocab^2 \nctx)$ times, and the summation on alg:cubic:line:add-to-count costs $\mathcal{O}(\nctx)$ per iteration.
What remains is to show that the call to $\Call{relaxed-correctness-pessimizing-over-position}{\Model, \QueryTok, \MaxTok, \NonMaxNonQueryTok, \NumCopiesNonMaxNonQueryTok}$ costs $\mathcal{O}(\nctx)$ when $\NumCopiesNonMaxNonQueryTok \ne 0$ and at most $\mathcal{O}(\dvocab \nctx)$ when $\NumCopiesNonMaxNonQueryTok = 0$ and $\NonMaxNonQueryTok = \MaxTok$.

The matrix multiplications in , , , , and can be cached upfront, costing $\mathcal{O}(\max(\dvocab,\dmodel,\nctx)^2\dmodel) \le \mathcal{O}(\dvocab^3)$ since we assume $\dvocab > \dmodel$ and $\dvocab > \nctx$.

The sorting on alg:cubic:full:sort-line can also be cached upfront (per ), costing $\mathcal{O}(\dvocab\nctx \log \nctx)$.

Note that each variable assignment in relaxed-correctness-pessimizing-over-position can be cached into a table parameterized over at most three variables which range over and over at most two variables that range over .

What remains is the statements.

When $\NumCopiesNonMaxNonQueryTok = 0$, we have on alg:cubic:full:line:return-c-zero:
This is $\mathcal{O}(\dvocab \nctx)$ as desired.

When $c \neq 0$, we have on alg:cubic:full:line:return-c-nonzero:

\[
\csname alg:cubic:full:line:return-c-nonzero:line:value\endcsname
\]

We can cache $\max_{{\NonMaxOutputTok \neq \MaxTok{}}} \Vector{\algorithmiclongvariable{skip-score}}[\NonMaxOutputTok]$ per and , costing $\mathcal{O}(\dvocab^3\nctx)$.
We can cache $\max_{{\NonMaxOutputTok \neq \MaxTok{}}}(\Vector{\Delta w}[\max, \NonMaxOutputTok] + \Vector{\Delta v}[\NonMaxInputTokIndex, \NonMaxOutputTok])$ per and $\NonMaxNonQueryTok$ costing $\mathcal{O}(\dvocab^3)$, since each $\Vector{\Delta v}[\NonMaxInputTokIndex, \NonMaxOutputTok]$ will be $\Vector{\Delta v}[\NonMaxInputTok, \NonMaxOutputTok]$ for some $\NonMaxInputTok\in\{\QueryTok, \MaxTok, \NonMaxNonQueryTok\}$.
Finally, we can compute the summation in cost $\mathcal{O}(\nctx)$ per loop iteration, as required.

§ QUADRATIC COUNTING FOR A SUB-CUBIC PROOF

In this section we fill in the details lacking from <ref>.

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

When we only get a budget of $\mathcal{O}(\dvocab^2\nctx)$ extremal relaxed cases to consider, though, we must pessimize more, which gives us a simpler version of the convexity theorem and proof.
Notably, when we restrict our sequences to have only two tokens (the max token  and the non-max token ), most of the theorems from <ref> get significantly simpler.

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

§.§ Proof of baseline sub-cubic result

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

Fix a total number of tokens .
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, as in <ref>, we'll take, for fixed query 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*}

Note that unlike in <ref>, we pessimize independently over the query token and the non-max token, so the “fixed” query token may not in fact appear in any key-side position in the relaxed sequence we consider.

We can define a sequence of tokens via mapping from positions by specifying a subset of valid tokens $S \subseteq \N^{<\dvocab}$ paired with a function $T : \N^{<\nctx} \to S$ specifying which token is in each position.

Given a subset of valid tokens $S \subseteq \N^{<\dvocab}$ and a function $T : \N^{<\nctx} \to S$ specifying which token is in each position, define the sequence score

\[
s_T := \left. \sum_{\mathclap{0 \le i < \nctx}} v_{T(i)} e^{a_{T(i)} + b_i} \middle/ \hphantom{0}\sum_{\mathclap{0 \le i < \nctx}} e^{a_{T(i)} + b_i} \right.
\]

Given a subset of valid tokens $S \subseteq \N^{<\dvocab}$ and a function $T : \N^{<\nctx} \to S$ specifying which token is in each position and two indices $0 \le i,j < \nctx$, define the swapped mapping $T_{i\leftrightarrow j}$ be the function that is $T$ except swapping $i$ and $j$:

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

Fix two tokens $t_0 < t_1 \in \N$ and a function $T : \N^{<\nctx} \to \{t_0, t_1\}$ specifying which token is in each position.

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

\[
\Delta_{T,i\leftrightarrow j} := s_{T_{i\leftrightarrow j}} - s_T
\]

\[
\sign\left(\Delta_{T,i\leftrightarrow j}\right)
= -\sign\left(b_i - b_j\right)\sign\left(v_{T(i)} - v_{T(j)}\right)
\]

<ref> gives us the result directly when $a_{T(i)} = a_{T(j)}$.
Otherwise, we get

\[
\sign\left(\Delta_{T,i\leftrightarrow j}\right) = \sign\left(a_{T(i)} - a_{T(j)}\right)
\sign\left(b_{i} - b_{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)
\]

Hence all that remains is to show that

\[
\sign\left(
s_T\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- v_{T(i)} e^{a_{T(i)}} + v_{T(j)} e^{a_{T(j)}}
\right)
-\sign\left(v_{T(i)} - v_{T(j)}\right)
\]

Define $\bar v := \frac{1}{2}(v_{T(i)} + v_{T(j)})$ and define $\Delta v := \frac{1}{2}(v_{T(i)} - v_{T(j)})$ so that $v_{T(i)} = \bar v + \Delta v$ and $v_{T(j)} = \bar v - \Delta v$.
Assume WLOG that $T(i) = 0$ and $T(j) = 1$ so that $v_{T(p)} = \bar v + (-1)^{T(p)}\Delta v$ for all $p$.

Then we have

\begin{align*}
& \sign\left(
s_T\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- v_{T(i)} e^{a_{T(i)}} + v_{T(j)} e^{a_{T(j)}}
\right) \\
& = \sign\left(
s_T\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- \bar v\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- \Delta v e^{a_{T(i)}} - \Delta v e^{a_{T(j)}}
\right) \\
& = \sign\left(
\frac{\displaystyle\sum_{{0 \le p < \nctx}} v_{T(p)} e^{a_{T(p)} + b_p}}{\displaystyle\sum_{\mathclap{0 \le p < \nctx}} e^{a_{T(p)} + b_p}}\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- \bar v\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- \Delta v \left(e^{a_{T(i)}} + e^{a_{T(j)}}\right)
\right) \\
& = \sign\left(
\frac{\displaystyle\sum_{{0 \le p < \nctx}} \left(\cancel{\bar v} + (-1)^{T(p)}\Delta v\right) e^{a_{T(p)} + b_p}}{\displaystyle\sum_{\mathclap{0 \le p < \nctx}} e^{a_{T(p)} + b_p}}\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- \cancel{\bar v\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)}
- \Delta v\left(e^{a_{T(i)}} + e^{a_{T(j)}}\right)
\right) \\
& = \sign(\Delta v)\sign\left(
\frac{\displaystyle
e^{a_{T(i)}} \sum_{\mathclap{\substack{0 \le p < \nctx \\ T(p) = T(i)}}} e^{b_p}
- e^{a_{T(j)}} \sum_{\mathclap{\substack{0 \le p < \nctx \\ T(p) = T(j)}}} e^{b_p}
}{\displaystyle\sum_{\mathclap{0 \le p < \nctx}} e^{a_{T(p)} + b_p}}\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- e^{a_{T(i)}} - e^{a_{T(j)}}
\right) \\
& = \sign(v_{T(i)} - v_{T(j)})\sign\left(
\frac{\displaystyle
e^{a_{T(i)}} \sum_{\mathclap{\substack{0 \le p < \nctx \\ T(p) = T(i)}}} e^{b_p}
- e^{a_{T(j)}} \sum_{\mathclap{\substack{0 \le p < \nctx \\ T(p) = T(j)}}} e^{b_p}
}{\displaystyle\sum_{\mathclap{0 \le p < \nctx}} e^{a_{T(p)} + b_p}}\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- e^{a_{T(i)}} - e^{a_{T(j)}}
\right) \\
\intertext{Define
\begin{align*}
P_i & := \sum_{\mathclap{\substack{0 \le p < \nctx \\ T(p) = T(i)}}} e^{b_p} &
P_j & := \sum_{\mathclap{\substack{0 \le p < \nctx \\ T(p) = T(j)}}} e^{b_p} \\
\end{align*}
so that we get}
& \sign\left(
s_T\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- v_{T(i)} e^{a_{T(i)}} + v_{T(j)} e^{a_{T(j)}}
\right) \\
& = \sign(v_{T(i)} - v_{T(j)})\sign\left(
\frac{e^{a_{T(i)}} P_i - e^{a_{T(j)}} P_j}{e^{a_{T(i)}}P_i + e^{a_{T(j)}} P_j}\left(e^{a_{T(i)}} - e^{a_{T(j)}}\right)
- e^{a_{T(i)}} - e^{a_{T(j)}}
\right) \\
\intertext{Multiply through by the positive denominator and expand out so that we get}
& = \sign(v_{T(i)} - v_{T(j)})\sign\left(-2
e^{a_{T(i)}+a_{T(j)}} P_i-2 e^{a_{T(i)}+a_{T(j)}} P_j\right) \\
& = -\sign(v_{T(i)} - v_{T(j)})\sign\left(
e^{a_{T(i)}+a_{T(j)}} P_i + e^{a_{T(i)}+a_{T(j)}} P_j\right) \\
& = -\sign(v_{T(i)} - v_{T(j)}) \\
\end{align*}

Let $\sigma_s : \N \to \N$ denote a permutation of the $\nctx$ positions that sorts them according to $b$: for $0 \le i, j < \nctx$, $b_{i} \le b_{j}$ whenever $\sigma_s(i) < \sigma_s(j)$.
Fix two tokens $t_0 < t_1 \in \N$.

Let $n_{t_0}$ be the number of $p \in [0, \nctx)$ with $T(p) = t_0$ and let $n_1$ be the number of $p \in [0, \nctx)$ with $T(p) = t_{t_1}$.
Note that $n_{t_0} + n_{t_1} = \nctx$.

Define $t_{\min} := \argmin_{t\in \{t_0,t_1\}} v_t$ and define $t_{\max} := \argmax_{t\in \{t_0,t_1\}} v_t$.

Define $T_{\min}, T_{\max} : \N^{<\nctx} \to \{t_0,t_1\}$ to be the assignment of tokens to positions that pays the least (respectively, most) attention to $t_{\max}$:

\begin{align*}
T_{\min}(i) & := \begin{cases}
t_{\max} & \text{if }0 \le \sigma_s(i) < n_{t_{\max}} \\
t_{\min} & \text{if }n_{t_{\max}} \le \sigma_s(i) < \nctx
\end{cases} \\
T_{\max}(i) & := \begin{cases}
t_{\min} & \text{if }0 \le \sigma_s(i) < n_{t_{\min}} \\
t_{\max} & \text{if }n_{t_{\min}} \le \sigma_s(i) < \nctx
\end{cases}
\end{align*}

Then we have that

\[
s_{T_{\min}} \le s_T \le s_{T_{\max}}
\]

The extremality of $s_{T_{\min}}$ and $s_{T_{\max}}$ follows straightforwardly from <ref>.

All that remains is $s_{T_{\min}} \le s_{T_{\max}}$.

This follows from noting by <ref> that swapping two tokens in $T_{\min}$ increases the sequence score, while the reverse is true of $s_{T_{\max}}$, thus showing that it must be $s_{T_{\min}}$ that is the minimum and $s_{T_{\max}}$ that is the maximum and not vice versa.

Given a subset of valid tokens $S \subseteq \N^{<\dvocab}$ and a function $T : \N^{<\nctx} \to S$ specifying which token is in each position define the full sequence score $s'_T$:

\[
s'_T := \left. \sum_{\mathclap{0 \le i < \nctx}} (v_{T(i)} + w_i) e^{a_{T(i)} + b_i} \middle/ \hphantom{0}\sum_{\mathclap{0 \le i < \nctx}} e^{a_{T(i)} + b_i} \right.
\]

Fix two tokens $t_0 < t_1 \in \N$.
Let $T_{\min}, T_{\max} : \N^{<\nctx} \to \{t_0,t_1\}$ and $t_{\max}$, $t_{\min}$ be as in <ref>.
Fix a set $S$ of valid tokens with $t_0, t_1 \in S$.

Define relaxed versions $T'_{\max}, T'_{\min} : \N^{<\nctx} \to S$ of $T_{\max}$ and $T_{\min}$:

\begin{align*}
T'_{\max}(i) & := \begin{cases}
T_{\max}(i) & \text{if }T_{\max}(i) = t_{\max} \\
\argmin_{{\substack{j \in S \\ {j \neq t_{\max}}}}} a_j & \text{otherwise}
\end{cases} \\
T'_{\min}(i) & := \begin{cases}
T_{\min}(i) & \text{if }T_{\max}(i) = t_{\max} \\
\argmax_{{\substack{j \in S \\ {j \neq t_{\max}}}}} a_j & \text{otherwise}
\end{cases}
\end{align*}

That is, $T'_{\max}$ replaces $t_{\min}$ with whatever token in $S$ draws the least attention away from $t_{\max}$, while $T'_{\min}$ replaces $t_{\min}$ with whichever token in $S$ draws the most attention away from $t_{\max}$.

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

\begin{align*}
r_{T_{\min}} & := \min_{\mathclap{0 \le i < \nctx}} w_i + \left( \phantom{00}\sum_{\mathclap{0 \le i < \nctx}} v_{T_{\min}(i)} e^{a_{T'_{\min}(i)} + b_i} \middle/ \hphantom{0}\sum_{\mathclap{0 \le i < \nctx}} e^{a_{T'_{\min}(i)} + b_i} \right) \\
r_{T_{\max}} & := \max_{\mathclap{0 \le i < \nctx}} w_i + \left( \phantom{00}\sum_{\mathclap{0 \le i < \nctx}} v_{T_{\min}(i)} e^{a_{T'_{\max}(i)} + b_i} \middle/ \hphantom{0}\sum_{\mathclap{0 \le i < \nctx}} e^{a_{T'_{\max}(i)} + b_i} \right)
\end{align*}

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

Essentially the same as the proof of <ref>.

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

Counting Correct Sequences in Subcubic Time, Preliminaries

[1]
input-sequence-compatible-withinput-sequence, , , , , ,
$\InputSequence\gets \algorithmiclongvariable{input-sequence}$
False  $\InputSequence \not\in (\N^{<\dvocab})^{\nctx}$
the sequence is not made of valid tokens
False  $\InputSequence[-1] \neq \QueryTok$
wrong query token
False  $\max_{\NonMaxInputTokIndex} \InputSequence[\NonMaxInputTokIndex] \neq \MaxTok$
wrong max token
False  $\card{\setst{\NonMaxInputTokIndex\in \N^{<\nctx}}{\InputSequence[\NonMaxInputTokIndex] \neq \MaxTok}}\neq \NumCopiesNonMaxTok$
wrong count of non-max toks
all$\InputSequence[\NonMaxInputTokIndex] = \MaxTok ~\algorithmicor~ \MaxTok - \InputSequence[i] \ge \MinGapSingular$  $0 \le \NonMaxInputTokIndex < \nctx$
check gap on non-max toks
correctness-pessimizing-over-gap-slow, , , , , ,
allcorrectness,      s.t. input-sequence-compatible-with, , , , , ,
subcubic, , ,
$\algorithmiclongvariable{count} \gets 0$
# of correct sequences
$\MinGap[\MaxTok, \QueryTok, \NumCopiesNonMaxTok] \gets \Call{min}{\MaxTok,\Call{max}{1, \MinGap[\MaxTok, \QueryTok, \NumCopiesNonMaxTok]}}$
Clip  to valid range
$\MinGap[\MaxTok, \NumCopiesNonMaxTok]^* \gets \min_{\Token \le \MaxTok} \MinGap[\MaxTok, \Token, \NumCopiesNonMaxTok]$
Cache running minima
$\in$ range
$\MaxTok{} \gets \algorithmiclongvariable{max-token}$
$0 \le \QueryTok \le \MaxTok$
$\QueryTok{} \gets \algorithmiclongvariable{query-token}$
$\NumCopiesNonMaxTok_{\max} \gets 0 ~\algorithmicif~ \QueryTok{} = \MaxTok{} ~\algorithmicelse~ 1$
minimum copies of nonmax
$\NumCopiesNonMaxTok_{\max} \gets \begin{cases}
0 & \algorithmicif~ \MaxTok=0 \\
\nctx - 1 & \algorithmicif~\MaxTok=\QueryTok \\
\nctx - 2 & \algorithmicelse
\end{cases}$
maximum copies of nonmax
$\NumCopiesNonMaxTok_{\min} \le \NumCopiesNonMaxTok \le \NumCopiesNonMaxTok_{\max}$
number of valid choices for the non-max token
$\MinGapSingular \gets \MinGap[\MaxTok, \QueryTok, \NumCopiesNonMaxTok]$
$\MinGapSingular^* \gets \MinGap[\MaxTok, \NumCopiesNonMaxTok]^*$
q-gap $\gets \MaxTok - \QueryTok$
$\Call{RCPOG}{\ensuremath{\vec{\chi}}} \gets \Call{relaxed-correctness-pessimizing-over-gap}{\Model, \vec{\chi}}$
$\left(\algorithmiclongvariable{q-gap}=0 ~\algorithmicor~ \algorithmiclongvariable{q-gap} \ge \MinGapSingular\right) ~\algorithmicand~ \Call{RCPOG}{\dvocab, \nctx, \MaxTok, \QueryTok, \NumCopiesNonMaxTok, \MinGapSingular, \MinGapSingular^*}$
$\NumCopiesNonMaxNonQueryTok' \gets \NumCopiesNonMaxTok ~\algorithmicif~ \QueryTok=\MaxTok ~\algorithmicelse~ \NumCopiesNonMaxTok - 1$
# of non-max non-query tokens
$\algorithmiclongvariable{count} \gets \algorithmiclongvariable{count} + \binom{\nctx-1}{\NumCopiesNonMaxNonQueryTok'}(\MaxTok - \MinGapSingular)^{\NumCopiesNonMaxNonQueryTok'}$
taking $0^0=1$ conventionally

Counting Correct Sequences in Subcubic Time

[1]

model-behavior-relaxed-over-gap, , , , , $\MinGapSingular^*$
correctness-pessimizing-over-gap-slow is False $\Longrightarrow$ result is False
$0 \le \MinGapSingular^* \le \MinGap \le \MaxTok$
$\NumCopiesNonMaxTok = 0$ $\QueryTok=\MaxTok$

$\algorithmiclongvariable{skip-score} \gets \max_{\NonMaxOutputTok} \Vector{\Call{\LogitsContributionSkip}{\QueryTok}}[\NonMaxOutputTok] - \min_{\NonMaxOutputTok} \Vector{\Call{\LogitsContributionSkip}{\QueryTok}}[\NonMaxOutputTok]$
Cache by

$\Vector{v}[\NonMaxInputTok] \gets \Call{\EVOU}{\NonMaxInputTok}$
$\Vector{w}[\NonMaxInputTokIndex] \gets \Call{\PVOU}{\NonMaxInputTokIndex}$
$\Delta w_{\max,\NonMaxOutputTok} \gets \max_{\NonMaxInputTokIndex} \Vector{w}[\NonMaxInputTokIndex,\NonMaxOutputTok] - \Vector{w}[p,\MaxTok]$
Cache by , $\NonMaxOutputTok$
$\Delta w_{\max,\max} \gets \max_{\NonMaxOutputTok} \Delta w_{\max,\NonMaxOutputTok}$
Cache by

$\Vector{\Delta v}[\NonMaxInputTok] \gets \max_i\Vector{v}[\NonMaxInputTok,\NonMaxOutputTok] - \min_{\NonMaxOutputTok} \Vector{v}[\NonMaxInputTok,\NonMaxOutputTok]$
Cache by $\Delta v_{\max} \gets \max_{0 \le \NonMaxInputTok \le \MaxTok - \MinGapSingular^*} \Vector{\Delta v}[\NonMaxInputTok]$
Cache by , $\Vector{\Delta v^{\MaxTok}}[\NonMaxOutputTok] \gets \Vector{v}[\MaxTok,\NonMaxOutputTok] - \Vector{v}[\MaxTok,\MaxTok]$
Cache by $\Delta v^{\MaxTok}_{\max} \gets \max_{\NonMaxOutputTok\neq \MaxTok} \Vector{\Delta v^{\MaxTok}}[\NonMaxOutputTok]$
Cache by $\NumCopiesNonMaxTok = 0$
$\Vector\Logits[\NonMaxOutputTok] \gets \Vector{\Call{\LogitsContributionSkip}{\MaxTok}}[\NonMaxOutputTok] + \Vector{v}[\MaxTok,\NonMaxOutputTok] + \Delta w_{\max,\NonMaxOutputTok}$
$\max_{\NonMaxOutputTok\neq \MaxTok}\left(\Vector\Logits[\NonMaxOutputTok] - \Vector\Logits[\MaxTok]\right)$
$\Vector{b}[{:},\nctx-1] \gets \Call{\EPQKP}{\QueryTok, \nctx-1} / \sqrt{\dhead}$
Cache by $\Vector{b}[0,:{-1}] \gets \Call{sort}{\Call{\EPQKP}{\QueryTok, :{-1}}} / \sqrt{\dhead}$
Cache by , $\Vector{b}[1,:{-1}] \gets \Call{reverse}{\Vector{b}[0,:{-1}]}$
$\Vector{a}[\NonMaxInputTok] \gets \Call{\EPQKE}{\QueryTok,\NonMaxInputTok} / \sqrt{\dhead}$
Cache by , $\Vector{a}[\min,\NonMaxInputTok] \gets \min_{0 \le \NonMaxInputTokAlt \le \NonMaxInputTok} \Vector{a}[\NonMaxInputTokAlt]$
Cache by , , compute in amortized $\mathcal{O}(\dvocab^2)$
$\Vector{a}[\max,\NonMaxInputTok] \gets \max_{0 \le \NonMaxInputTokAlt \le \NonMaxInputTok} \Vector{a}[\NonMaxInputTokAlt]$
Cache by , , compute in amortized $\mathcal{O}(\dvocab^2)$
$\Delta a_{\max} \gets \Vector{a}[\MaxTok] - \Vector{a}[\min,\MaxTok-\MinGapSingular]$
Cache by , , $\Delta a_{\min} \gets \Vector{a}[\MaxTok] - \Vector{a}[\max,\MaxTok-\MinGapSingular]$
Cache by , , $\algorithmiclongvariable{idx-set} \gets \{0,\ldots,\nctx - \NumCopiesNonMaxTok-1\}~\algorithmicif~\MaxTok\neq\QueryTok~\algorithmicelse~\{0,\ldots,\nctx - \NumCopiesNonMaxTok-2,\nctx-1\}$
$\algorithmiclongvariable{attn-weights-unscaled}_{0,\NonMaxInputTokIndex} \gets \Vector{b}[0,\NonMaxInputTokIndex] + (\Delta a_{\min}~\algorithmicif~\NonMaxInputTokIndex\in \algorithmiclongvariable{idx-set}~\algorithmicelse~0)$
$\algorithmiclongvariable{attn-weights-unscaled}_{1,\NonMaxInputTokIndex} \gets \Vector{b}[0,\NonMaxInputTokIndex] + (\Delta a_{\max}~\algorithmicif~\NonMaxInputTokIndex\in \algorithmiclongvariable{idx-set}~\algorithmicelse~0)$
Cache by , , , $\algorithmiclongvariable{attn-weights}_{0} \gets \Call{softmax}{\algorithmiclongvariable{attn-weights-unscaled}_{0}}$
Cache by , , , $\algorithmiclongvariable{attn-weights}_{1} \gets \Call{softmax}{\algorithmiclongvariable{attn-weights-unscaled}_{1}}$
Cache by , , , $\algorithmiclongvariable{attn-max}_0 \gets \sum_{\NonMaxInputTokIndex\in\algorithmiclongvariable{idx-set}}\algorithmiclongvariable{attn-weights}_{0,\NonMaxInputTokIndex}$
$\algorithmiclongvariable{attn-max}_1 \gets \sum_{\NonMaxInputTokIndex\in\algorithmiclongvariable{idx-set}}\algorithmiclongvariable{attn-weights}_{1,\NonMaxInputTokIndex}$
$\algorithmiclongvariable{attn-max} \gets \algorithmiclongvariable{attn-max}_0~\algorithmicif~\Delta v^{\MaxTok}_{\max} \ge \Delta v_{\max}~\algorithmicelse~\algorithmiclongvariable{attn-max}_1$

$\algorithmiclongvariable{skip-score} + \Delta w_{\max,\max} + \algorithmiclongvariable{attn-max} \cdot \Delta v^{\MaxTok}_{\max} + (1-\algorithmiclongvariable{attn-max})\Delta v_{\max}$

relaxed-correctness-pessimizing-over-gap, , , , , , , $\MinGapSingular^*$
runs the model on a relaxed variant of input sequences compatible with the arguments
correctness-pessimizing-over-gap-slow is False $\Longrightarrow$ result is False
is False correctness-pessimizing-over-gap-slow,  is False any  with specified , , and  tokens not equal to $\Call{model-behavior-relaxed-over-gap}{\Model, \MaxTok, \QueryTok, \NumCopiesNonMaxTok, \MinGapSingular, \MinGapSingular^*} < 0$

For all ,

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

(sketch) Apply preceding lemmas and theorems to <ref>

The running time of <ref>, after using caching to avoid duplicate computations, is $\mathcal{O}(\dvocab^2 \dmodel + \dvocab^2 \nctx^2)$.

(sketch) Sum the complexities indicated along the right side of <ref>.
The $\dvocab^2 \dmodel$ term comes from the precomputing , , and .
The $\dvocab^2 \nctx^2$ term comes from the softmax over $\nctx$ tokens for $O(\dvocab^2 \nctx)$ pessimized pure sequences.
Confirming that none of the complexities on the right side exceeds $\mathcal{O}(\dvocab^2 \dmodel + \dvocab^2 \nctx^2)$ completes the proof.

§ SUBCUBIC PROOF STRATEGIES

In this section, we present a number of proof strategies that we use to reduce the computational cost of the proof, ultimately driving down the cost of  and  verification to $\mathcal{O}(\dvocab \dmodel)$, while unfortunately leaving the cost of  verification at $\mathcal{O}(\dvocab^2 \dmodel)$.

The three main tricks we cover are the mean+diff trick (<ref>), the max row-diff trick (<ref>), and the rank one / rank two SVD decomposition of  (<ref>).
While the mean+diff trick is useful for getting slightly better bounds, the SVD decomposition of  is the place where we get to insert the most understanding (without which we'd have no hope of non-vacuous bounds below $\mathcal{O}(\dvocab^2\dmodel)$), and the max row-diff trick is the workhorse that allows us to drive down the noise computations from cubic to quadratic without getting completely vacuous bounds.

§.§ The mean+diff trick

Suppose we have quantities $f_{x,y}$ and $g_{y,z}$ and we want to pessimize (WLOG, suppose minimize) the quantity $f_{x,y} + g_{y,z}$ over $x$, $y$, and $z$ in time less than $\mathcal{O}(n_x n_y n_z)$, say we allow $\mathcal{O}(n_x n_y + n_y n_z + n_x n_z)$.
Also suppose the variation of $f$ over the $y$ axis is much larger than the variation of f over the x-axis.

We can of course say

$$\min_{x,y} f_{x,y} + \min_{y, z} g_{y,z} \le f_{x,y} + g_{y,z}$$

But we can do better!

Note that

$$f_{x,y} = \E_x f_{x,y} + (f_{x,y} - \E_x f_{x,y})$$

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

\begin{align*}
f_{x,y} + g_{y,z}
& \ge \min_{x,y,z} [f_{x,y} + g_{y,z}] \\
& = \min_{x,y,z} [\E_x f_{x,y} + g_{y,z} + f_{x,y} - \E_x f_{x,y}] \\
& \ge \min_{x,y,z} [\E_x f_{x,y} + g_{y,z}] + \min_{x,y,z}[f_{x,y} - \E_x f_{x,y}] \\
& = \min_{y,z} [\E_x f_{x,y} + g_{y,z}] + \min_{x,y}[f_{x,y} - \E_x f_{x,y}]
\end{align*}

By averaging the variation over certain axes, we have

\begin{align*}
\min_{x,y,z} f_{x,y} + g_{y,z} & \ge \min_{y,z} [\E_x f_{x,y} + g_{y,z}] + \min_{x,y}[f_{x,y} - \E_x f_{x,y}] \\
\max_{x,y,z} f_{x,y} + g_{y,z} & \le \max_{y,z} [\E_x f_{x,y} + g_{y,z}] + \max_{x,y}[f_{x,y} - \E_x f_{x,y}]
\end{align*}

and the RHSs can be computed in time $\mathcal{O}(n_x n_y + n_y n_z + n_x n_z)$ for $n_x$, $n_y$, and $n_z$ the number of possible values of $x$, $y$, and $z$, respectively.

Example for how this helps with small variation:

Take any function $k(y)$ and then take

\begin{align*}
f_{x,y} & := k(y) + \varepsilon_1(x, y) \\
g_{y,z} & := -k(y) + \varepsilon_2(y, z)
\end{align*}

Then we have

\begin{align*}
\min_{x,y,z}[f_{x,y} + g_{y,z}]
& = \min_{x,y,z}[\varepsilon_1(x, y) + \varepsilon_2(y, z)] \\
\min_{x,y}f_{x,y} + \min_{y,z}g_{y,z}
& = \min_{y}k(y) + \min_{y} -k(y) + \min_{x,y}\varepsilon_1(x, y) + \min_{y,z}\varepsilon_2(y, z) \\
& = \min_{y}k(y) - \max_{y} k(y) + \min_{x,y}\varepsilon_1(x, y) + \min_{y,z}\varepsilon_2(y, z) \\
\min_{x,y}[f_{x,y} - \E_x f_{x,y}] + \min_{y,z}[g_{y,z} + \E_x f_{x,y}]
& = \min_{x,y}\varepsilon_1(x, y) + \min_{y,z}[\varepsilon_2(y, z) + \E_x\varepsilon_1(x, y)]
\end{align*}

If $\varepsilon_1$ and $\varepsilon_2$ are small compared to $\min_y k(y) - \max_y k(y)$, then using $\E_x f_{x,y}$ gives a much better bound.

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

Note also that this trick is not restricted to adding and subtracting $\E_x f_{x,y}$.
If $f$ is a matrix indexed by $x$ and $y$, we might also try taking SVD and using the first principal component instead.

A basic application of the triangle inequality gives the following, more general, result:

For any $h_y$ which can be computed in time $\mathcal{O}(n_h)$,

\begin{align*}
\min_{x,y,z} f_{x,y} + g_{y,z} & \ge \min_{y,z} [h_y + g_{y,z}] + \min_{x,y}[f_{x,y} - h_y] \\
\max_{x,y,z} f_{x,y} + g_{y,z} & \le \max_{y,z} [h_y + g_{y,z}] + \max_{x,y}[f_{x,y} - h_y]
\end{align*}

and the RHSs can be computed in time $\mathcal{O}(n_x n_y + n_y n_z + n_h)$ for $n_x$, $n_y$, and $n_z$ the number of possible values of $x$, $y$, and $z$, respectively.

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

\[
f_{x,y} = h_y + \varepsilon_{x,y}
\]

we get

\begin{align*}
&\left|\min_{x,y,z} f_{x,y} + g_{y,z} - \min_{y,z} [h_y + g_{y,z}] - \min_{x,y}[f_{x,y} - h_y]\right| \\
&\leq \left|\min_{x,y,z} [f_{x,y} + g_{y,z}] - \min_{y,z} [h_y + g_{y,z}]\right| + \left|\min_{x,y}[\varepsilon_{x,y}]\right| \\
&\leq 2 \max_{x,y} \left|\varepsilon_{x,y}\right|
\end{align*}

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

§.§ Details of SVD of QK proof

As discussed in <ref>, to further reduce the computation cost of proof, we need to avoid computing the residual stream, EVOU, and EPQKE matrices fully.
Using mechanistic insight or otherwise, we observe that these matrices (apart from EVOU) can be well-approximated by rank one matrices.
This will remove the dominant computation cost of $\mathcal{O}(\dvocab^2 \cdot \dmodel)$.

§.§.§ Comments on relationship between mechanistic insight and proof size

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

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

§.§.§ The max row-diff trick

As stated above, we are breaking matrices into their rank one approximation and some noise.
To bound the noise, i.e. to bound expressions of the form $\prod_i (A_i + E_i) - \prod_i A_i$, where $E_i$ denote the matrix errors, we can use the following trick:

For a row vector a and a matrix $B$,

\[
\max_{i,j} \left((\Vector{a}B)_{i} - (\Vector{a}B)_{j}\right)
\le \sum_k \left|\Vector{a}[k]\right|\max_{i,j}  \left(B_{k,i} - B_{k,j}\right)
\]

Moreover, for a collection of $n$ row vectors $A_r$, if the shape of $B$ is $m \times p$, the right hand side can be computed for all $r$ in time $\mathcal{O}(n m + m p)$.

\begin{align*}
&\max_{i,j} (\Vector{a}B)_{i} - (\Vector{a}B)_{j} \\
&= \max_{i,j} \sum_k \Vector{a}[k] \left(B_{k,i} - B_{k,j}\right) \\
&\le \sum_k \max_{i,j} \Vector{a}[k] \left(B_{k,i} - B_{k,j}\right) \\
&= \sum_k \Vector{a}[k]
\begin{cases}
\max_{i,j}  \left(B_{k,i} - B_{k,j}\right) & \text{if }\Vector{a}[k] \ge 0 \\
\min_{i,j} \left(B_{k,i} - B_{k,j}\right) & \text{if }\Vector{a}[k] <0
\end{cases} \\
&= \sum_k \Vector{a}[k]
\begin{cases}
\max_{i,j}  \left(B_{k,i} - B_{k,j}\right) & \text{if }\Vector{a}[k] \ge 0 \\
-\max_{i,j} \left(B_{k,i} - B_{k,j}\right) & \text{if }\Vector{a}[k] <0
\end{cases} \\
&= \sum_k \left|\Vector{a}[k]\right|\max_{i,j} \left(B_{k,i} - B_{k,j}\right)
\end{align*}

The asymptotic complexity of computing the result follows from caching the computation of $\max_{i,j} \left(B_{k,i} - B_{k,j}\right)$ for each $k$ independently of $r$, as the computation does not depend on $A_r$.

For matrices $A$ and $B$,

\[
\max_{r,i,j} \left((AB)_{r,i} - (AB)_{r,j}\right)
\le
\max_r \sum_k \left|A_{r,k}\right|\max_{i,j}  \left(B_{k,i} - B_{k,j}\right)
\]

By taking the max of <ref> over rows $r$ of $A$.

<ref> can also be applied recursively for a product of more than two matrices.

For a row vector a and a sequence of $n$ matrices $B_p$ of shapes $r_p \times c_p$,

\[
\max_{i,j} \left(\left(\Vector{a}\prod_p B_p\right)_{i} - \left(\Vector{a}\prod_p B_p\right)_{j}\right)
\le \sum_{k_0} \left|\Vector{a}[k_0]\right| \cdots \sum_{k_n} \left|(B_{n-1})_{k_{n-1},k_n}\right| \max_{i,j} \left((B_{n})_{k_n,i} - (B_{n})_{k_n,j}\right)
\]

Moreover, for a collection of $q$ row vectors $A_\alpha$, the right hand side can be computed for all $\alpha$ in time $\mathcal{O}(q r_0 + \sum_p r_p c_p)$.

We proceed by induction on $n$.

For $n = 1$, the statement is identical to <ref>.

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

\begin{align*}
& \max_{i,j} \left(\left(\Vector{a}\prod_{p=1}^{s+1} B_p\right)_{i} - \left(\Vector{a}\prod_{p=1}^{s+1} B_p\right)_{j}\right) \\
& = \max_{i,j} \left((\Vector{a}B_1)\left(\left( \prod_{p=2}^{s+1} B_p\right)_{i} - \left(\prod_{p=2}^{s+1} B_p\right)_{j}\right)\right) \\
\intertext{Using the induction hypothesis gives}
& \le \sum_{k_1}\left|\sum_{k_0} \Vector{a}[k_0](B_1)_{k_0,k_1}\right| \sum_{k_2} \left|(B_2)_{k_1,k_2}\right| \cdots \sum_{k_{s+1}} \left|(B_{s})_{k_{s},k_{s+1}}\right| \max_{i,j} \left((B_{s+1})_{k_{s+1},i} - (B_{s+1})_{k_{s+1},j}\right) \\
\intertext{The triangle inequality gives}
& \le \sum_{k_1}\sum_{k_0} \left|\Vector{a}[k_0](B_1)_{k_0,k_1}\right| \sum_{k_2} \left|(B_2)_{k_1,k_2}\right| \cdots \sum_{k_{s+1}} \left|(B_{s})_{k_{s},k_{s+1}}\right| \max_{i,j} \left((B_{s+1})_{k_{s+1},i} - (B_{s+1})_{k_{s+1},j}\right) \\
\intertext{and algebra gives}
& = \sum_{k_0} \left|\Vector{a}[k_0]\right| \sum_{k_1} \left|(B_1)_{k_0,k_1}\right| \sum_{k_2} \left|(B_2)_{k_1,k_2}\right| \cdots \sum_{k_{s+1}} \left|(B_{s})_{k_{s},k_{s+1}}\right| \max_{i,j} \left((B_{s+1})_{k_{s+1},i} - (B_{s+1})_{k_{s+1},j}\right)
\end{align*}

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

For a sequence of $n+1$ matrices $A_0$, …, $A_n$,

\[
\max_{r,i,j} \left(\left(\prod_p A_p \right)_{r,i} - \left(\prod_p A_p\right)_{r,j}\right)
\le
\max_r \sum_{k_0}\left|(A_0)_{r,k_0}\right|\cdots \sum_{k_n} \left|(A_{n-1})_{k_{n-1},k_n}\right|\max_{i,j} \left((A_n)_{k_n,i} - (A_n)_{k_n,j}\right)
\]

By taking the max of <ref> over rows $r$ of $A_0$.

Note that <ref> is compatible with the mean+diff trick of <ref>.

For matrices $A$ and $B$, and any column-wise summary vector $\Vector{H}[k]$ of $A$ (for example we may take $H_k := \E_r A_{r,k}$)

\[
\max_{r,i,j} \left((AB)_{r,i} - (AB)_{r,j}\right)
\le
\left(\max_{i,j} \sum_k \Vector{H}[k]\left(B_{k,i} - B_{k,j}\right)\right) + \max_r \sum_k \left|A_{r,k} - \Vector{H}[k]\right| \max_{i,j}\left(B_{k,i} - B_{k,j}\right)
\]

\begin{align*}
&\max_{r,i,j} \left((AB)_{r,i} - (AB)_{r,j}\right) \\
&= \max_{r,i,j} \sum_k A_{r,k} \left(B_{k,i} - B_{k,j}\right) \\
&= \max_{r,i,j} \sum_k \left(\Vector{H}[k] + \left(A_{r,k} - \Vector{H}[k]\right)\right) \left(B_{k,i} - B_{k,j}\right) \\
&= \max_{i,j} \left(\sum_k \Vector{H}[k]\left(B_{k,i} - B_{k,j}\right) + \max_r \sum_k \left(A_{r,k} - \Vector{H}[k]\right) \left(B_{k,i} - B_{k,j}\right)\right) \\
&\le \left(\max_{i,j} \sum_k \Vector{H}[k]\left(B_{k,i} - B_{k,j}\right)\right) + \max_r \sum_k \max_{i,j}\left(A_{r,k} - \Vector{H}[k]\right) \left(B_{k,i} - B_{k,j}\right) \\
&\le \left(\max_{i,j} \sum_k \Vector{H}[k]\left(B_{k,i} - B_{k,j}\right)\right) + \max_r \sum_k \left|A_{r,k} - \Vector{H}[k]\right| \max_{i,j}\left(B_{k,i} - B_{k,j}\right)
\end{align*}

For a row vector a, a vector of summaries h corresponding to a (for example, if a is a row of a matrix, h might be the average of the rows), a sequence of $n$ matrices $B_p$ of shapes $r_p \times c_p$, and a corresponding sequence of column-wise summary vectors h_p of $B_p$ (for example we may take $\Vector{(h_p)}[k] := \E_r \Vector{(B_p)}[r,k]$),

\begin{align*}
& \max_{i,j} \left(\left(\Vector{a}\prod_p B_p\right)_{i} - \left(\Vector{a}\prod_p B_p\right)_{j}\right) \\
& \le \max_{i,j} \left(\sum_{k_0} \Vector{h}[k_0] \cdots \sum_{k_n} \Vector{(B_{n-1})}[k_{n-1},k_n] \left((B_{n})_{k_n,i} - (B_{n})_{k_n,j}\right)\right) \\
& \phantom{{}\le{}}
+ \sum_{k_0} \left|\Vector{a}[k_0] - \Vector{h}[k_0]\right| \cdots \left(
\left(\max_{i,j} \sum_{k_n} \Vector{(h_{n-1})}[k_n] \left((B_n)_{k_n,i} - (B_n)_{k_n,j}\right)\right)
\right. \\
& \phantom{{}\le{} + \sum_{k_0} \left|\Vector{a}[k_0] - \Vector{h}[k_0]\right| \cdots \left(\vphantom{\left(\max_{i,j} \sum_{k_n} \Vector{(h_{n-1})}[k_n] \left((B_n)_{k_n,i} - (B_n)_{k_n,j}\right)\right)}\right.}\left.
+ \sum_{k_n} \left|(B_{n-1})_{k_{n-1},k_n} - \Vector{(h_{n-1})}[k_n]\right| \max_{i,j} \left((B_n)_{k_n,i} - (B_n)_{k_n,j}\right)
\right)
\end{align*}

Moreover, for a collection of $q$ row vectors $A_\alpha$, the right hand side can be computed for all $\alpha$ in time $\mathcal{O}(q r_0 + \sum_p r_p c_p)$.

Apply the triangle inequality recursively, fusing the proofs of thm:max-row-diff:vector:recursive,thm:max-row-diff+avg+diff.

For a sequence of $n+1$ matrices $A_0$, …, $A_n$, and corresponding column-wise summary vectors $\Vector{h_0}$, …, $\Vector{h_{n-1}}$ of $A_0$, …, $A_{n-1}$,

\begin{align*}
& \max_{r,i,j} \left(\left(\prod_p A_p \right)_{r,i} - \left(\prod_p A_p\right)_{r,j}\right) \\
& \le \max_{r,i,j} \left(\sum_{k_0} \Vector{(h_0)}[k_0] \cdots \sum_{k_n} \Vector{(A_{n-1})}[k_{n-1},k_n] \left((A_{n})_{k_n,i} - (A_{n})_{k_n,j}\right)\right) \\
& \phantom{{}\le{}}
+ \sum_{k_0} \left|\Vector{(A_0)}[k_0] - \Vector{(h_0)}[k_0]\right| \cdots \left(
\left(\max_{i,j} \sum_{k_n} \Vector{(h_{n-1})}[k_n] \left((A_n)_{k_n,i} - (A_n)_{k_n,j}\right)\right)
\right. \\
& \phantom{{}\le{} + \sum_{k_0} \left|\Vector{(A_0)}[k_0] - \Vector{(h_0)}[k_0]\right| \cdots \left(\vphantom{\left(\max_{i,j} \sum_{k_n} \Vector{(h_{n-1})}[k_n] \left((A_n)_{k_n,i} - (A_n)_{k_n,j}\right)\right)}\right.}\left.
+ \sum_{k_n} \left|(A_{n-1})_{k_{n-1},k_n} - \Vector{(h_{n-1})}[k_n]\right| \max_{i,j} \left((A_n)_{k_n,i} - (A_n)_{k_n,j}\right)
\right)
\end{align*}

Moreover, if the matrices $A_p$ have shapes $r_p \times c_p$, the right hand side can be computed in time $\mathcal{O}(\sum_p r_p c_p)$.

By taking the max of <ref> over rows $r$ of $A_0$.

§.§.§ Exploring rank one approximation via SVD

Let us first look at

\[
\EPQKE := \qWE\WQ\WK^T\barWE^T.
\]

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

group/horizontal sep=1.5cm

SVD of for seed , with principal component vectors scaled by the square root of the corresponding singular value.
This scaling allows us to see visually that there is not much going on beyond the first singular component.
Numerically: the first singular value is just over [0,round-direction=down], while the second singular value is just under [0,round-direction=up].

The first singular value is just over [0,round-direction=down] (across all seeds), while the second singular value is just under [0,round-direction=up] (across all seeds).
The ratio across all seeds is .
There's really not much going on here beyond the first singular component.
We might be tempted to keep analyzing the SVD, and notice that the query direction is mostly uniform, while the key direction is monotonic (nearly linear, even).
But the proof complexity doesn't demand this level of analysis, yet, and so we can't expect that any automated compact proof discovery system will give it to us.

Call the first singular component of EQKE the “query direction”  and the “size direction”
on the query-side and key-side, respectively.

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

§.§.§ The simple SVD decomposition of QK

In time $\mathcal{O}(\dvocab\dmodel^2)$ we can perform SVD on each of the four component matrices , , , and perform low-rank SVD on the matrix product $\qWE\WQ\WK^T\barWE^T$.

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

That is, we can decompose  into a part parallel to  and a part orthogonal to , say $\qWE = \WEq + \WEqPerp$, and similarly $\barWE = \WEk + \WEkPerp$.
Note that  and  are both rank one, and hence can be multiplied with other matrices of shape $\dmodel \times a$ in time $\mathcal{O}(\dmodel a)$ rather than time $\mathcal{O}(\dvocab \dmodel a)$.

Hence we can define  (subscript one for “rank one”) and decompose  as

\[\EPQKE = \WEq\WQ\WK^T(\WEk)^T + \EQKEerrOne.\]

Define for any vector $v$

\begin{align*}
\Delta_{i,j}v & := \Vector{v}[i] - \Vector{v}[j]
\end{align*}

so that we get

\begin{gather*}
\Delta_{i,j}\Vector{(\WEq\WQ\WK^T(\WEk)^T)}[\QueryTok] + \min_{i\neq j} \Delta_{i,j}\Vector{(\EQKEerrOne)}[\QueryTok] \\
\le \Delta_{i,j}\Vector\EPQKE[\QueryTok] \le \\
\Delta_{i,j}\Vector{(\WEq\WQ\WK^T(\WEk)^T)}[\QueryTok] + \max_{i\neq j} \Delta_{i,j}\Vector{(\EQKEerrOne)}[\QueryTok]
\end{gather*}

Then we may use any method we please to pessimize $\Delta_{i,j}\Vector{(\EQKEerrOne)}[\QueryTok]$ quickly.
For example, since for any matrix $M$ we have $\sigma_1(M) = \sup_x \norm{Mx} / \norm{x}$, considering vectors with one $1$, one $-1$, and zero elsewhere, the maximum difference between elements in a row upper bounded by $\sqrt2 \sigma_1(M)$:

\begin{align}
\left|\Delta_{i,j}\Vector{{\EQKEerrOne}}[\QueryTok]\right|
& \le \sqrt{2}\sigma_1(\EQKEerrOne)\label{eqn:sqrt2-sigma1-bound}% \\
% \left|\Delta_{i,j}\Vector{{\EQKEerrOne}}[\QueryTok]\right|
% & \le \sqrt{2}\sqrt{\sum_i (\sigma_i(\EQKEerrOne))^2} \qquad\text{(Frobenius norm bound)}
\end{align}

§.§.§ The complicated SVD decomposition of QK

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

To this end, we perform SVD on , , , and  and peel off the first singular components so as to get the decomposition

\begin{align*}
\qWE & = \WEq + \WEqq + \WEqqPerp \\
\barWE & = \WEk + \WEkk + \WEkkPerp \\
\WQ & = \WQq + \WQqPerp \\
\WK & = \WKk + \WKkPerp
\end{align*}

Then , a product of these four matrices, can be expressed as a sum of $2^2 3^2 - 1 = 35$ rank one products and one high-rank error term.
We can compute the sum of the rank one products in time $\mathcal{O}(\dvocab^2)$ and express as, say, $\EQKETwo + \WEqqPerp \WQqPerp (\WEkkPerp \WKkPerp)^T$.
Call the second term (<ref>).
We must now bound for each $q$ and $m$ the quantity $\max_{i \le m - \MinGap} \EQKEerr[q, i] - \EQKEerr[q, m]$.

How big is this?

max width=.3

The error term  for seed .

Even if we relax to $\max_{i,j} \EQKEerr[q, i] - \EQKEerr[q, j]$, the maximum such value across all rows is under [2,round-direction=up] (across all seeds).
And the rows don't have any particular structure to them; the maximum absolute element of the entire matrix is just barely over [0,round-direction=down] (across all seeds), so doubling that doesn't give too bad an estimate.

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

One option is to try to use singular value decomposition again.
Since $\sigma_1(M) = \sup_x \norm{Mx} / \norm{x}$, considering vectors with one $1$, one $-1$, and zero elsewhere, the maximum difference between elements in a row upper bounded by $\sqrt2 \sigma_1(M)$.

The largest singular value of  (<ref>) is just under [1,round-direction=up] (across all seeds), giving a row-diff bound of about [1] (across all seeds), which is large but not unusably so.

max width=0.75EQKE-err-svd
SVD of  for seed .

If we perform SVD before multiplying out the matrices (<ref>), however, their first singular values are about [0], [1], [1], and [0], giving a product of about [0], which when multiplied by $\sqrt2$ is about [0].
(Across all seeds, these numbers are , , , and , giving a product of about , which when multiplied by $\sqrt2$ is about .)
This works because $\sigma_1(AB) \le \sigma_1(A)\sigma_1(B)$, but note that we can do factored SVD without needing to use this technique.
This bound is still usable, but pretty big.

SVD of the four component matrices of  for seed .
Matrices look like noise.

Can we use Frobenius?
Note that using anything close to this method to drop below $\dvocab \dmodel^2$ might seem infeasible (it'll eventually turn out not to be).
For example, the best bound we know on the largest singular value that can be verified even in the worst-case in strictly less time than it takes to compute the full SVD is the Frobenius norm, which is defined as $\text{tr}(M M^T)$, can be computed in $\dmodel \dvocab$ time, and is equal to the square root of the sum of the squares of the singular values.
While the Frobenius norm of  is only about [0] (giving a bound of about [0] on the row-diff), the Frobenius norms of the four multiplicand matrices are a bit over [0], [0], [0], and [0], giving a product of [0] and a bound of [0](!).
(Across all seeds, the Frobenius norm of  is about (giving a bound of about on the row-diff), the Frobenius norms of the four multiplicand matrices are a bit over , , , and , giving a product of and a bound of .)
This is unusably large.

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

If we want to compute the max row-diff of a product of matrices $AB$, we can compute by <ref>

\begin{equation}\label{eqn:max-row-diff}
\max_{r,i,j} \left((AB)_{r,i} - (AB)_{r,j}\right)
\le
\max_r \sum_k \left|A_{r,k}\right|\max_{i,j}  \left(B_{k,i} - B_{k,j}\right)
\end{equation}

or by combining this approximation with <ref> via <ref> we may compute

\begin{align*}
& \max_{r,i,j} \left((AB)_{r,i} - (AB)_{r,j}\right) \\
& \le
\left(\max_{i,j} \sum_k \E_rA_{r,k}\left(B_{k,i} - B_{k,j}\right)\right)
+ \max_r \sum_k \left|A_{r,k} - \E_rA_{r,k}\right| \max_{i,j}\left(B_{k,i} - B_{k,j}\right)
\end{align*}

taking whichever bound is better.

The first gives us a bound of [2] on the maximum row-diff, which is better than we can get by doing SVD on the product of the matrices!
We can get an even better bound by peeling off the first two singular values of all four matrices before multiplying them; this gives us a bound of [2].
Combining it with the avg+diff trick wouldn't give us much ([2] and [2] respectively), as we've effectively already done this by peeling off the leading singular contributions; the mean of  over dimension zero has norm [round-mode=figures,round-precision=2] (across all seeds).

Although this noise bound is no longer the leading asymptotic bottleneck, we can peek ahead to what we get if we want to be linear in parameter count.
In this case, we can apply the recursive version of <ref> via <ref>, giving a bound of [2] on the maximum row-diff.

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

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

If we use this trick on the skip connection () only, we can achieve a bound of (across all seed).

Using this trick on both EU and QK drops us down only to (across all seeds).

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

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

§ COMPARISON OF PROOF STRATEGIES

In this section, we compare the various proof strategies that we have developed in <ref>.
We do some traditional mechanistic interpretability analysis to justify that the choices that we made could be expected to lead to reasonably good bounds in <ref>.
We then compare the complexities and performance of various proof strategies in <ref> to line up with the legends of fig:pareto-frontier,fig:norm-accuracy-vs-singular-ratio.
We close with a figure relating the various categories of proof strategies.

§.§ Justification of pessimization choices

In sec:results:proof:sub-cubic,sec:subcubic:counting:details,sec:subcubic:tricks we make a number of choices about which axes of variation are more or less important to track at various points in the bound computation.

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

§.§.§ Justifying the gap

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

We justify this as follows.

b0.490.8fig:logit-diff-max-tokenThe attention computation weighted by the number of sequences with the particular max.
The computation is $\max_{j\text{ s.t.\@ }j \neq i} \Vector\EVOU[i,j] - \Vector\EVOU[i,i]$.
$\mu\pm\sigma$: ; range:

b0.490.8fig:EVOU-hist-max-row-diffHistogram of the maximum difference between two logits contributed by a single row of .
The computation is, for each $i$, $\max_h \Vector\EVOU[i,j] - \min_j\Vector\EVOU[i,j]$.
$\mu\pm\sigma$: ; range:

Plots of the difference in logit for the attention computation, $\EVOU := \barWE \WV \WO \WU$ for seed .

We can look at the typical diff, when attending to the max token, between the largest non-max logit and the max logit.
As shown in <ref>, the largest difference between an off-diagonal entry of  and the diagonal of that row is typically at most $\EVOUHistMinAboveDiagDupBySeqCountMostBelowValue$.[
“Typically” here means about [0] of the time.
The typical worst contribution to the wrong logit from a non-max token (this is typical over non-max tokens, worst over choice of output token-logit index) is around , as shown in <ref>.

The difference in attention between tokens is approximately linear in the gap between the tokens, as seen in <ref>.
The slope of the line, that is, the difference in pre-softmax attention scores divided by the gap between the key token and the max token, is approximately [1].

$(\EPQKE_i - \EPQKE_j)/\sqrt{\dhead}$ vs. $i - j$

table/search path=generated/max-of-4/adjusted-figures/,

$(\EPQKE_i - \EPQKE_j)/(\sqrt{\dhead}(i-j))$, weighted by sequence count.
$\mu \pm \sigma = \EQKEHistAttentionDifferenceOverGapDupBySeqCountMeanPMStd$

Plots of attention difference vs. token gap, for $\EPQKE := \qWE\WQ\WK^T\barWE^T$ for seed .
The difference in attention between tokens is approximately linear in the gap between the tokens.

Exponentiating, the post-softmax attention paid to the max is typically about $\EQKEHistAttentionDifferenceOverGapDupBySeqCountExpMean\times$ larger than to the token one below the max;
here the logit difference between the max and non-max token is significant, typically being around [round-mode=places,round-precision=0]/exp() ($\EVOUMeanMaxRowDiff / \EQKEHistAttentionDifferenceOverGapDupBySeqCountExpMean$) for the worst output logit.
But by the time the gap is 3, this difference has dropped to about [round-mode=places,round-precision=1]/exp(3*), and by the time the gap is 4 it is around [round-mode=places,round-precision=1]/exp(4*).

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

The upshot is that to handle most sequences, we need only ask an oracle for the minimum gap $\MinGapSingular{} > 0$ between the max token  and largest non-max tokens $\NonMaxNonQueryTok \neq \MaxTok$, such that the model outputs the correct answer for all sequences where the non-max, non-query tokens have value at most $\MaxTok-g$.

While computing this gap may be expensive (and indeed the naïve computation of the oracle takes longer than the brute-force proof—though it should be very easy to optimize), we don't have to pay the cost of computing the gap in the size of the proof, only the cost of storing the gap table ($\mathcal{O}(\dvocab^2 \nctx)$) and of verifying the gap.
Empirically, gaps are typically 1–, as seen in <ref>.

max width=.32generated/max-of-4/figuresSubcubicGapHistogram.tex
Histogram of the minimum gap between the max token and the largest non-max token, for the seed .

If we rely on the gaps, this results in leaving behind about [1] of sequences.

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

* First, we could build in time ($\mathcal{O}(\dvocab^2)$) a table indexed on pairs $(\NonMaxInputTok, \MaxTok)$ of the maximum token and a non-maximum token:
the table would store pessimal logit contributions from  to maximum output tokens $\le$ the  parameter.
The table could be further split to pessimize separately for tokens within and outside of the gap window.

* Compute a table of pre-softmax attention differences between tokens $\NonMaxInputTok$ and $\NonMaxInputTok+1$ in time ($\mathcal{O}(\dvocab^2)$).

* Next sort the queries by overlap with the query direction.

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

* Compute a table indexed on pairs of # of queries handled and input tokens  which stores the smallest difference in more attention paid to $\NonMaxInputTok+1$ than to $\NonMaxInputTok$ ($\mathcal{O}(\dvocab^2)$).

* Compute a table indexed on pairs ,  storing an upper bound on amount more attention paid to non-maximum tokens than to  by Oracle-permitted query tokens (the Oracle is indexed only on ) ($\mathcal{O}(\dvocab^2)$).

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

§.§.§ Stopping after 1–2 principal components of QK

Did we miss out on any structure in the noise of ?
The distribution of entries of the four matrices looks pretty close to normal as seen in <ref>.

The distribution of entries of the four residual matrices (after removing two principal components from and and one principal component from  and ).
Distributions look pretty close to normal.
Plots are for the seed .

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

table/search path=generated/max-of-4/adjusted-figures/,
every axis/.append style=legend pos = outer north east,

Recreations of <ref> for ease of viewing of the legend.
Top is a strict recreation; bottom includes points not on the Pareto frontier.

If we sample elements randomly, we get (sample size ) that the maximum row-diff of the product of the matrices is approximately [2] (sampling without replacement from the empirical distribution) or [2] (sampling from the normal distribution).
So in fact our max row-diff is unusually high (by about $\num[round-mode=places,round-precision=0]{\fpeval{(\EQKEErrMaxRowDiffFloat - \ResampleEQKEErrMeanFloat) / \ResampleEQKEErrStdDevFloat}}\sigma$).
This shows up in the bias towards having larger values (both positive and negative) in the lower-right corner of the plot, indicating that noise is larger for larger query and key values.
We hypothesize that this is due to the distribution of data: larger values are more likely to have more space between the maximum and next-most-maximum token, so a bit of noise matters less for larger maxes than for smaller ones.

§.§ How various combinations of tricks perform

Recall fig:pareto-frontier,fig:norm-accuracy-vs-singular-ratio from <ref>, recapitulated here without captions for convenience as fig:pareto-frontier:again,fig:norm-accuracy-vs-singular-ratio:again.

We describe what each subcubic proof strategy in the legend means.
Note that all subcubic proof strategies (that is, all proof strategies except for “brute force” and “cubic”) use the quadratic counting algorithm of <ref>.

Recreation of <ref> for ease of viewing of the legend.

§.§.§ Proof strategies grouped by complexity

In fig:pareto-frontier,fig:pareto-frontier:again, proof strategies are grouped by computational complexity.

The 102 proof strategies break down into $1+1+2\times 5\times10\times2$ strategies.

The brute force and cubic proofs ($1+1$) were fully covered in sec:brute-force,sec:cubic:full.

There are 5 options for handling :

direct-quadratic refers to handling  in time $\mathcal{O}(\dvocab\dmodel)$ with either the max row-diff trick (<ref>)[This strategy is labeled “” in the Python source code.] or the max row-diff trick fused with mean+diff or some other summary statistic (<ref>)[These strategies are labeled “” and “” in the Python source code].

direct is not mentioned, this indicates that we handle  in time $\mathcal{O}(\dvocab^2\dmodel)$ by first multiplying out  and then either taking the maximum row-diff in each row[“”] or by taking the maximum row-diff across all rows[“”].
The latter is included purely for comparison's sake, and never gives a tighter bound than the former.

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

attention-quadratic refers to handling the high-rank attention error term  from <ref> in time $\mathcal{O}(\dvocab\dmodel)$ either with the recursive max row-diff trick (<ref>)[“”] or with the recursive max row-diff trick fused with the mean+diff trick either just on the query side[“”] or throughout[“”] (<ref>).

attention-$\dvocab\dmodel^2$ indicates that we use one of the various $\mathcal{O}{(\dvocab\dmodel^2)}$ strategies for handling  from <ref> or  from <ref>.
These include using $\sqrt{2}\sigma_1$—computed via low-rank SVD—as the bound (<ref>)[“”], considering all ways of multiplying out a subset of the matrices and taking the maximum row-diff of the resulting pair of matrices[“” for , “” for ] (<ref>), or fusing the max row-diff trick with the mean+diff trick[“”  for , “” for ] (<ref>).

When attention is not mentioned, this indicates that we handle the attention error term in time $\mathcal{O}(\dvocab^2\dmodel)$, either by taking the per-row maximum row-diff[“”] or by using the full rank  matrix and taking the per-row maximum row diff[“”].

Finally, note that in combining the rank one attention computation with , , and , we may either use the mean+diff trick[“”] (<ref>) or not[“”]; this makes up the final factor of $2$.

§.§.§ Proof strategies grouped by attention handling

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

max-diff-exact ($\mathcal{O}(\dvocab^2\dmodel)$) corresponds to taking the full rank  term and taking the maximum row-diff in each row.

mean+max-diff-subproduct ($\mathcal{O}(\dvocab\dmodel)$) corresponds to fusing the max row-diff trick with the mean+diff trick (<ref>) and considering all ways of associating the multiplication of .

max-diff-subproduct ($\mathcal{O}(\dvocab\dmodel)$) corresponds to using the max row-diff trick (<ref>) and considering all ways of associating the multiplication of .

max-diff ($\mathcal{O}(\dvocab\dmodel^2)$) corresponds to using the max row-diff trick (<ref>) on the factored SVD of .

mean+max-diff ($\mathcal{O}(\dvocab\dmodel^2)$) corresponds to fusing the max row-diff trick with the mean+diff trick (<ref>) and applying it on the factored SVD of .

svd ($\mathcal{O}(\dvocab\dmodel^2)$) corresponds to using $\sqrt{2}\sigma_1$—computed via low-rank SVD—as the bound (<ref>).

mean+max-diff-subproduct-recursive ($\mathcal{O}(\dvocab\dmodel)$) corresponds to handling the high-rank attention error term  from <ref> with the recursive max row-diff trick fused with the mean+diff trick on the query-side only (<ref>, taking all but the first summary vector to be zero).

max-diff-subproduct-recursive ($\mathcal{O}(\dvocab\dmodel)$) corresponds to handling the high-rank attention error term  from <ref> with the recursive max row-diff trick (<ref>).

mean-recursive+max-diff-subproduct-recursive ($\mathcal{O}(\dvocab\dmodel)$) corresponds to handling the high-rank attention error term  from <ref> with the recursive max row-diff trick recursively fused with the mean+diff trick (<ref>).

§.§.§ What understanding do we get from each proof strategy?

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

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