Mech Interp Puzzle 1: Suspiciously Similar Embeddings in GPT-Neo

What background knowledge do you think this requires? If I know a bit about how ML and language models work in general, should I be able to reason this out from first principles (or from following a fairly obvious trail of "look up relevant terms and quickly get up to speed on the domain?"). Or does it require some amount of pre-existing ML taste?

Also, do you have a rough sense of how long it took for MATS scholars?

Note that we have built-in spoiler tags [LW · GW].

(Potential spoilers!)

There is some relevant literature which explores this phenomenon, also looking at the cosine similarity between words across layers of transformers. I think the most relevant is (Cai et. al 2021), where they also find this higher than expected cosine similarity between residual stream vectors in some layer for BERT, D-BERT, and GPT. (Note that they use some somewhat confusing terminology: they define inter-type cosine similarity to be cosine similarity between embeddings of different tokens in the same input; and intra-type cosine similarity to be cosine similarity of the same token in different inputs. Inter-type cosine similarity is the one that is most relevant here).

They find that the residual stream vectors for GPT-2 small tend to lie in two distinct clusters. Once you re-centre these clusters, the average cosine similarity between residual stream vectors falls to close to 0 throughout the layers of the model, as you might expect.

 

What comes to mind:

Qbrf guvf rzretr va genvavat be vf vg cerfrag ng vavgvnyvmngvba? Purpx juvpu cnegf lbh pna erzbir sebz gur nepuvgrpgher. YnlreAbez fbhaqf yvxr gur xvaq bs gjrnx gung zvtug punatr guvf? Cebonoyl sbetrg guvf chmmyr naq chmmyr bire cybgf pbagnvavat zber guna ~1 ovg bs vasb.

These puzzles are great, thanks for making them!

The equation king-man+woman=queen suggests that vector addition is not necessarily a fundamental operation in the space of embedded words, but instead we want a weaker operation that is kind of like addition but where we cannot pinpoint a zero vector.

In mathematics, we sometimes want to define a new axiomatization for an algebraic structure where we do not have any presence of a zero element. The best way to do this is to replace our binary fundamental operation with a ternary operation. For example, the notion of a median algebra is an adaptation of the notion of a distributive lattice where we do not have any notion of which way is up and which way is down but if we have two points and we specify which one we want to be on top and which one we want to be on the bottom, then the interval between those two points will be a distributive lattice. As another example, the notion of a heap is what you would get if you took the notion of a group but you removed the identity element.

A heap is an algebraic structure $(X,t)$ where $X$ is a set and $t$ is a ternary operation that satisfies the identities

$t\left(t\right(v,w,x),y,z)=t(v,w,t(x,y,z\left)\right)=t(v,t(y,x,w),z)$ and $t(x,x,y)=t(y,x,x)=y$. 

If $(G,\cdot )$ is a group, then we can define a ternary heap operation $t:G^3\to G$ by setting $t(x,y,z)=x{y}^{-1}z$. For example, t(king,man,woman)=queen.

Similarly, if $t:G^3\to G$ is a heap operation on a set $G$ and $a\in G$ ($a$ can be called a basepoint), then we can define a group operation ${\cdot }_a$ by setting $x{\cdot }_{a}y=t(x,a,y)$. The category of groups is isomorphic to the category of heaps with a basepoint $a$. 

We can adapt other classes of mathematical structures to the case where we remove a basepoint including the classes of rings, vector spaces, and even inner product spaces if we really wanted to, and an affine space is the thing we obtain when we take a vector space and remove the origin but keep everything else intact in such a way that we can recover everything from the affine space and the origin. In an affine space $V$ over a field $K$ (the field $K$ has an origin, but the space $V$ does not), we cannot define scalar multiplication since that requires a basepoint, but we can declare a point in $V$ to be the origin for the purposes of performing scalar multiplication. We therefore define a scalar multiplication operation $s:K\times V\times V\to V$ by setting $S(\alpha ,x,y)=\alpha \cdot (y-x)+x=\alpha y+(1-\alpha )x$ which is scalar multiplication where $x$ is declared to be the origin. From this scalar multiplication, we can easily define the mean of several elements and weighed means (except in some finite fields). So if we are given data $x_1,\dots ,x_n$ in an affine space $V$, we can still average $x_1,\dots ,x_n$ and then declare the average to be the origin of $V$ and then we have obtained a vector space instead of just an affine space.

Suppose that $K,X$ are finite dimensional vector spaces over a field $F$, and let $\iota :K\to X$ be a vector space isomorphism. Let $t$ denote the heap operations on these vector spaces. Let $f:X\to X$ be a permutation, and let $E:K\times X\to X$ be the function defined by $E(k,x)=\iota \left(k\right)+f\left(x\right)$. If $k\in K$, then define $E_k:X\to X$ by setting $E_k\left(x\right)=E(k,x)$. Observe that $E_k^{-1}\left(y\right)=f^{-1}(y-\iota (k\left)\right)$.

We can define (in a model theoretic sense) the heap operations on $K$ and $X$ from $E$, but we are not necessarily able to define the group operations on $K$ and $X$ from $E$.

If $j,k,l\in K$, then $E_j\left(E_k^{-1}\right(x\left)\right)=\iota \left(j\right)+f\left(E_k^{-1}\right(x\left)\right)=\iota \left(j\right)+f\left(f^{-1}\right(x-\iota \left(k\right)\left)\right)=\iota \left(j\right)+x-\iota \left(k\right)=\iota (j-k)+x$. 

Given $i,j,k\in K$, there is a unique $l\in K$ where $E_i\circ E_j^{-1}=E_l\circ E_k^{-1}$ since $E_i\circ E_j^{-1}=E_l\circ E_k^{-1}$ precisely when $i-j=l-k$ which happens precisely when $i-j+l=k$. We can therefore define the heap operation $t$ on $K$ from $E$. 

Given $x,y,z\in X$, there are $j,k\in K$ where $E_j\left(E_k^{-1}\right(y\left)\right)=\iota (j-k)+y=x$. We can therefore define $t:X^3\to X$ from $E$ by setting $t(x,y,z)=E_j\left(E_k^{-1}\right(z\left)\right)=\iota (j-k)+z=x-y+z.$ We can define the heap operations on $K,X$ from $E$, but we cannot necessarily define the zero element from $E$.

For example, $K,X$ could be vector spaces over the finite field with $2$ elements and $E$ could be the round function in a block cipher (we observe that permutation-substitution encryption functions like the AES encryption are simply recurrent neural networks over finite fields; AES uses $x\mapsto x^{-1},0\mapsto 0$ as its activation function). More generally $K,X$ could be vector spaces, and $f:X\to X$ could be an RNN layer.

Let us not forget that anyone who has studied a little bit of physics knows that there is a lot of physical intuition for affine spaces without origins. We cannot pinpoint a precise location and declare that location to be the center of the universe. We cannot say that an object is at rest; we can only say that an object is at rest relative to another object.

My quick take would be that

Answer: (decode with rot13)

Should this link instead go to rot13.com?

Good puzzle, thanks!

Thinking outloud a bit, I haven't followed mechanistic interpretability besides skimming some transformer circuits posts and half understanding them  - 

Jung'f gur qvssrerapr orgjrra arb naq tcg-2 fznyy?
- arb unf ybpny nggragvba naq hfrf n qvssrerag cbfvgvbany rapbqvat? 
- ibpno fvmr ybbxf gur fnzr

V xabj va uvtu qvzrafvbany fcnprf enaqbz irpgbef fubhyq or begubtbany. Fb jung'f qvssrerag urer?
- gur rzorqqvatf zhfg or nyy pyhfgrerq arneol fbzrubj
- vzntvar nyy ohg gur ynfg 10 qvzrafvbaf ner svkrq ng 1, gung jbhyq qb gur gevpx, ohg vg frrzf hayvxryl gur zbqry jbhyq yrnea guvf
- pbhyq nyfb or nyy gur rzorqqvatf unir n pbafgnag ba gurz, r.t. vafgrnq bs h naq i jr trg (p+h) naq (p+i)

Bx abj V jnag gb fubj gur zbqry pbhyq pbzchgr gur fnzr guvat jvgu na nqqrq pbafgnag va gur rzorqqvat. 
nggragvba ba (p+h) sebz (p+i) vf fbzrguvat yvxr (p+h)^G Jdx (p+i)
= p^G Jdx p + h^G Jdx p + p^G Jdx i + h^G Jdx i

Fb nyzbfg gur fnzr guvat, rkprcg bar fznyy rkgen grez qrcraqvat ba h. 
- vg'f cebonoyl svar gb vtaber guvf? 
- gur ZYCf pbhyq whfg unir n yrnearq ovnf juvpu fhogenpgf bhg p, fb gubfr ner gbgnyyl svar

tcg2 fznyy -- V thrff vg whfg unf n fznyyre p, znlor whfg ol enaqbz punapr bs ubj vg jnf vavgvnyvmrq? Ohg vs guvf vf n obahf dhrfgvba, gur nafjre fubhyq or zber vagrerfgvat. Fbzrguvat gb qb jvgu gur cbfvgvbany rapbqvat be gur nggragvba zrpunavfz? fueht. 

@Neel Nanda [LW · GW] but isn't it that there are more layers that produce the actual vocab tokens? it's not only layer 0 or the embedding layer right?

Thanks for posting this! Coincidentally, just yesterday I was wondering if there were any mech interp challenges like these, it seems to lend itself to this kind of thing. Had been considering trying to come up with a few myself.