Using the probabilistic method to bound the performance of toy transformers

post by Alex Gibson · 2025-01-21T23:01:38.067Z · LW · GW · 0 comments

Contents

  Introduction:
  max-of-k transformer:
  Equation for logits:
  Hyperplane form:
  Chernoff bounds:
  Full bounds:
  When to use probabilistic bound:
None
No comments

Introduction:

Transformers are statistical artefacts. If a model achieves 99% accuracy, it might have learnt an algorithm that relies upon some property of the input data that only holds 99% of the time.

Let's say that our train input to the transformer is 100 independent unbiased coin tosses. Then the transformer might implicitly rely upon there being between 35 and 65 heads, because there is a ~ 0.17% chance that there would be a more extreme number of heads.

This potentially makes transformers vulnerable to adversarial inputs. But it also tells us how we should be approaching formal proofs. We should not necessarily expect transformers to learn algorithms that work on every input, but instead for them to learn algorithms which work on the "typical set" of inputs, like "the set of inputs with between 35 and 65 heads". 

Then we can bound transformer performance in two stages:

1.) Identify a "typical set" of inputs, where we can show the transformer gets high accuracy

2.) Bound the volume of this "typical set" using probabilistic techniques.

max-of-k transformer:

To demonstrate this approach, I derive bounds on the performance of a toy one-layer attention-only model, the max-of-k transformer. This transformer takes as input $n_{\text{ctx}}$ numbers between $1$ and $63$, and is trained to output the maximum number in the sequence. It is a single layer transformer with a single attention head, and has no layer norm or MLP layer.

Intuitively, the way the transformer solves this task is by using its QK circuit to attend more to tokens the bigger they are, and then using its OV circuit to copy the tokens it attends to. So it will attend the most to the biggest token, and then copy this token the most, so that it outputs the maximum token as its highest logit.

But there are adversarial inputs where toy transformers trained on this task fail. For instance, take the sequence $62,62,62,62,....,62,63$. A sequence this extreme is super unlikely to appear in the training data, so the model won't be prepared to handle it. Even though the QK circuit attends slightly more to an individual $63$ token than $62$, the huge number of $62$'s in a single input sequence can overwhelm the transformer and lead it to erroneously outputting $62$ rather than $63$.

We'd like a condition on the input sequence, which measures how "extreme" it is, which is simple enough that we can use probabilistic techniques to estimate how many "extreme sequences" there are. 

For max-of-k, we can get a pretty simple condition by working through the algebra.

Equation for logits:

Let's say we have an input sequence of tokens between $0$ and $63$, say $x_1,x_2,...,x_{n_{\text{ctx}}}$.

Then let $A_i$ be the exponentiated attention score at $x_i$.

The output from the single attention head is ${\sum }_{i=1}^{n_{\text{ctx}}}\frac{A_i}{{\sum }_{j=1}^{n_{\text{ctx}}}{A}_j}\left(EVO\right[x_i]+PVO[i\left]\right)$.

This accounts for the EVO circuit and the PVO circuit. But we also have the direct circuit due to E and P, which we can bring inside the sum, so that our output logit at token k is:

${\text{logit}}_k={\sum }_{i=1}^{n_{\text{ctx}}}\frac{A_i}{{\sum }_{j=1}^{n_{\text{ctx}}}{A}_j}\left(EVOU\right[x_i,k]+PVOU[i,k]+EU[x_i,k]+PU[i,k\left]\right)$

We are only interested in the relative sizes of logits, so we can cancel out the softmax denominator, and just look at

${\text{out}}_k={\sum }_{i=1}^{n_{\text{ctx}}}{A}_i\left(EVOU\right[x_i,k]+PVOU[i,k]+EU[x_i,k]+PU[i,k\left]\right)$.

  Write ${\text{eu}}_{x_i,k}$ for $\text{EVOU}[x_i,k]+\text{EU}[x_i,k]$ and ${\text{pu}}_{i,k}$ for $\text{PVOU}[i,k]+\text{PU}[i,k]$.

Hyperplane form:

If we restrict to sequences whose maximum is M, then we are interested in finding conditions on having ${\text{out}}_M\ge {\text{out}}_k$ for all $k$ between $0$ and $63$.

${\text{out}}_M\ge {\text{out}}_k\leftrightarrow {\sum }_{i=1}^{n_{\text{ctx}}}{A}_i({\text{eu}}_{x_i,M}-{\text{eu}}_{x_i,k}+{\text{pu}}_{i,M}-{\text{pu}}_{i,k})\ge 0$

Now $A_i=e^{\frac{EQKE[x_{n_{\text{ctx}}},x_i]+EQKP[x_{n_{\text{ctx}}},i]+PQKP[n_{\text{ctx}},i]+PQKE[n_{\text{ctx}},x_i]]}{\sqrt{d_{\text{head}}}}}=e_{x_{n_{\text{ctx}}},x_i}{p}_{x_{n_{\text{ctx}}},i}$

Now we fix $x_{n_{ctx}}$ as well as fixing M.

Then ${\text{out}}_M\ge {\text{out}}_k\leftrightarrow {\sum }_{i=1}^{n_{\text{ctx}}}{e}_{x_i}{p}_i({\text{eu}}_{x_i,M}+{\text{pu}}_{i,M})\ge {\sum }_{i=1}^{n_{\text{ctx}}}{e}_{x_i}{p}_i({\text{eu}}_{x_i,k}+{\text{pu}}_{i,k})$.

Now bound $|{\text{pu}}_{i,M}-{\text{pu}}_{i,k}|\le C$. Then ${\text{out}}_M<{\text{out}}_k$ implies:$0<{\sum }_{i=1}^{n_{\text{ctx}}}{e}_{x_i}{p}_i({\text{eu}}_{x_i,k}-{\text{eu}}_{x_i,M}+C)$

Finally this is in a form we can work with, because the right hand side is just the set of sequences which lie above a certain hyperplane in sequence space.

Chernoff bounds:

It's in general quite hard to determine precisely the size of discrete sets which lie above hyperplanes, but we can bound their size from above using Chernoff bounds.

Let $f_{x_i,k,M,x_{n_{\text{ctx}}}}=e_{x_i}({\text{eu}}_{x_i,k}-{\text{eu}}_{x_i,M}+C)$

Then for arbitrary $\lambda$, we have:

 $P({\text{out}}_M<{\text{out}}_k)\le P(0<{\sum }_{i=1}^{n_{\text{ctx}}}{p}_i{f}_{x_i,k,M,x_{n_{\text{ctx}}}})=P(1<e^{\lambda {\sum }_{i=1}^{n_{\text{ctx}}}{p}_i{f}_{x_i,k,M,x_{n_{\text{ctx}}}}})$

$P(1<e^{\lambda {\sum }_{i=1}^{n_{\text{ctx}}}{p}_i{f}_{x_i,k,M,x_{n_{\text{ctx}}}}})\le E\left[e^{\lambda {\sum }_{i=1}^{n_{\text{ctx}}}{p}_i{f}_{x_i,k,M,x_{n_{\text{ctx}}}}}\right]={\prod }_{i=1}^{n_{\text{ctx}}}E\left[e^{\lambda p_i{f}_{x_i,k,M,x_{n_{\text{ctx}}}}}\right]$.

Where the last equality comes from independence of $x_1$,$x_2$,..,$x_{n_{ctx}}$. Although technically because we fixed $M$ we don't strictly have independence, but we can fix this with some bookkeeping by partitioning sequences by the positions which have $M$'s in.

There is no closed form for optimal $\lambda$, but we can use gradient descent to obtain a good numerical value for $\lambda$.

We can sum across different values of $x_{n_{\text{ctx}}}$ to get a probability just in terms of $M$ and $k$.

Full bounds:

Now we have $P({\text{logit}}_M\ge {\text{logit}}_k\text{for all}k\in [0,63\left]\right)=1-P({\text{logit}}_M<{\text{logit}}_k\text{for some}k\in [0,63\left]\right)$

And $P({\text{logit}}_M<{\text{logit}}_k\text{for some}k\in [0,63\left]\right)\le {\sum }_{k=0}^{63}{B}_{k,M}$, where $B_{k,M}$ is our chernoff bound from before, and we have used a union bound.

We should expect the union bound to work well, because there will tend to only be one $k$ which has a greater logit than $M$, if there exists such a $k$. So we get a lower bound on $P({\text{logit}}_M\ge {\text{logit}}_k\text{for all}k\in [0,63\left]\right)$.

Then we have:

$P\left(\text{Correct max logit}\right)={\sum }_{M=0}^{63}P\left(M\text{max token}\right)P\left(\text{Correct max logit}|M\text{max token}\right)$,

and we can get a lower bound on this.

 

When to use probabilistic bound:

This bound works best for longer sequences, because for shorter sequences it's easier to manually enumerate the sequences which lie above the hyperplane. The probabilistic method easily adapts to longer sequences, where manually enumerating sequences quickly becomes unfeasible.

0 comments

Comments sorted by top scores.