How transformers can compute distances along a curve locally.

[Alex Gibson]

Overview:

There is an interesting mechanism GPT-2 appears to use to measure distances between duplicate tokens. The mechanism reminds me a lot of Twisted pair cabling in communications.

The mechanism is fiddly to explain in context, so i’ve tried to abstract out most of the details, and give a clean toy version of the mechanism. I think some of the structures for this mechanism could be used in other contexts than computing distances.

Setup:

We have the set of points $\{x_i{\}}_{i=0}^n\in R^d\subset R^{d_{\text{model}}}$, with $x_i=f\left(i/n\right)$ where $f:[0,1]\to R^d$ is a smooth function. We take as input $x_i$, $x_j$ with $|i-j|\le k<<n$ . We want to construct a transformer which can estimate $|i-j|$ given $x_i,x_j$. For this transformer, we additionally assume that $d$ is relatively small compared to the embedding dimension $d_{\text{model}}$.

The mechanism:

We set up $M+1<n$ “sentry points” $\{s_i{\}}_{i=0}^M$ uniformly along $[0,1]$. And we define $g:[0,1]\to [0,M]\cap Z$ sending $t\in [0,1]$ to the index of the closest sentry point.

Then we have, $||x_i-x_j-(\frac{i}{n}-\frac{j}{n})f^{\prime }\left(s_{g\left(\frac{i}{n}\right)}\right)|{|}_2=||{\int }_{\frac{j}{n}}^{\frac{i}{n}}{f}^{\prime }\left(t\right)-f^{\prime }\left(s\right)dt|{|}_2\le L(\frac{1}{2M}+\frac{|i-j|}{n})\cdot \frac{|i-j|}{n}$

where $L$ is such that $||f^{\prime }\left(x\right)-f^{\prime }\left(y\right)|{|}_2\le L|x-y|$ for all $x,y\in [0,1]$.

So $|\frac{(x_i-x_j)\cdot f^{\prime }\left(s_{g\left(\frac{i}{n}\right)}\right)}{||f^{\prime }\left(s_{g\left(\frac{i}{n}\right)}\right)|{|}_2^2}-(i-j)|\le \frac{||{\int }_{\frac{j}{n}}^{\frac{i}{n}}{f}^{\prime }\left(t\right)-f^{\prime }\left(s\right)dt|{|}_2}{||f^{\prime }\left(s_{g\left(\frac{i}{n}\right)}\right)|{|}_2}\le \frac{L}{||f^{\prime }\left(s_{g\left(\frac{i}{n}\right)}\right)|{|}_2}(\frac{1}{2M+2}+\frac{|i-j|}{n})\cdot \frac{|i-j|}{n}$.

Therefore if we can approximate $\frac{(x_j-x_i)\cdot f^{\prime }\left(s_{g\left(\frac{i}{n}\right)}\right)}{||f^{\prime }\left(s_{g\left(\frac{i}{n}\right)}\right)|{|}_2^2}$, then we can approximate $j-i$.

Attention mechanism:

If we are given a two token input to the transformer $x_j$, $x_i$, then assuming that $d\le d_{\text{head}}$, two attention heads is sufficient to compute $x_j-x_i$ (have one head which outputs -$x_i$, and the other $x_j$). We write $x_j-x_i$ to an orthogonal subspace from $x_i$ so that the MLP can cleanly access $x_i$ later.

MLP mechanism:

The MLP mechanism consists of $2M+2$ neurons, with a pair of two neurons associated with each sentry point.

For each sentry point $s_m\in \{0,\frac{1}{M+1},\frac{2}{M+1},\dots ,1\}$, we define a neuron pair:

${\text{Neuron}}_m^{+}=\text{ReLU}\left(b_{m,1}f\right(s_m)\cdot x_i+b_{m,2}+b_{m,3}(x_j-x_i)\cdot f^{\prime }(s_m\left)\right)$

${\text{Neuron}}_m^{-}=\text{ReLU}\left(b_{m,1}f\right(s_m)\cdot x_i+b_{m,2}-b_{m,3}(x_j-x_i)\cdot f^{\prime }(s_m\left)\right)$

where $W_m$, $b_{m,l}$ are tuned so that $b_{m,1}f\left(s_m\right)\cdot x_i+b_{m,2}>ϵ$ when $m=g\left(\frac{i}{n}\right)$, and $b_{m,1}f\left(s_m\right)\cdot x_i+b_{m,2}<0$ otherwise. Additionally we set up $b_{m,3}(x_j-x_i)\cdot f^{\prime }\left(s_m\right)$ so that it has a magnitude of less than $ϵ$ when $|i-j|\le k$.

Example sentry neuron activations when $x_i=x_j$ with M = 5. Each different colour corresponds to the activation of a different sentry neuron. We can pick M+1 coprime to n so that the sentries don’t vanish at $\frac{i}{n}$ for any $i$, and so they are bounded above by some $ϵ>0$. A signal of magnitude $ϵ$ can be encoded in the difference between the activations of pairs of these sentry neurons.

Setting up these sentries relies on f not coming close to intersecting itself, so that the dot product with $f\left(s_m\right)$ is only high on a connected interval.

We wrote $x_j-x_i$ in an orthogonal subspace to $x_i$ so the sentry neurons can all be written in the standard form $\text{ReLU}(w^t{h}_0+b)$ where $h_0$ is the residual stream post attention.

We then output $\frac{C_m}{b_{m,3}}$(${\text{Neuron}}_m^{+}$-${\text{Neuron}}_m^{-}$)$v$ from the $m$th neuron pair.

Under this construction ${\text{Neuron}}_m^{+}$ and ${\text{Neuron}}_m^{-}$ always activate at the same time as each other, so the output of the $m$th neuron pair is $0$ if $m\ne g\left(\frac{i}{n}\right)$, and $C_m(x_j-x_i)\cdot f^{\prime }\left(s_m\right)v$ if $m=g\left(\frac{i}{n}\right)$.

Since only a single neuron pair activates at once, the complete output of this MLP layer is $C_{g\left(\frac{i}{n}\right)}(x_j-x_i)\cdot f^{\prime }\left(s_{g\left(\frac{i}{n}\right)}\right)v$.

Then setting $C_{g\left(\frac{i}{n}\right)}$ proportional to $\frac{1}{||f^{\prime }\left(s_{g\left(\frac{i}{n}\right)}\right)|{|}_2^2}$, we get an output proportional to $(j-i)v$

Extensions of mechanism:

The above mechanism is clean, and captures the key ideas. A nice thing about the mechanism is that the $f\left(s_m\right)\cdot x_i$ term can be noisy, but it doesn’t matter because the common noise will get cancelled out, similar to twisted pair encoding.

However, there can still be potential issues caused by noise at the transitions between sentries. It is also not robust to $f$ intersecting itself, and the number of sentries required grows with $n$.

There are cases where this kind of two neuron interference cancellation could come in useful outside of computing distance. For example, if you want to distinguish between British and Canadian English, you could have:

${\text{Neuron}}_{\text{canada}}=\text{ReLU}(\text{Commonwealth}+ϵ(\text{Canadian}\left)\right)$

${\text{Neuron}}_{\text{british}}=\text{ReLU}(\text{Commonwealth}+ϵ(\text{British}\left)\right)$

And then take the difference between the two. The interference that would usually make it difficult to distinguish between the two very similar dialects gets cancelled out.

Though this is probably just PCA??

[anaguma]

Is there evidence that GPT-2 uses this mechanism and not some other mechanism?