I was trying to understand the tensor product formulation in transformer circuits and I had basically forgotten all I ever knew about tensor products, if I ever knew anything. This very brief post is aimed at me from Wednesday 22nd when I didn't understand why that formulation of attention was true. It basically just gives a bit more background and includes a few more steps. I hope it will be helpful to someone else, too.

Tensor product

For understanding this, it is necessary to understand tensor products. Given two finite-dimensional vector spaces $V,W$ we can construct the tensor product space $V\otimes W$ as the span^[1] of all matrices $v\otimes w$, where $v\in V,w\in W,$ with the property $(v\otimes w{)}_{ij}=v_i{w}_j$ ^[2]. We can equivalently define it as a vector space with basis elements $e_i^V\otimes e_j^W$, where we used the basis elements of $V$ and $W$ respectively.

But not only can we define tensor products between vectors but also between linear maps that map from one vector space to the other (i.e. matrices!):

Given two linear maps (matrices) $A:V\to X,B:W\to Y$ we can define $A\otimes B:V\otimes W\to X\otimes Y$, where each map simply operates on its own vector space, not interacting with the other:

$$(A\otimes B)(v\otimes w)=A\left(v\right)\otimes B\left(w\right)$$

For more information on the tensor product, I recommend this intuitive explanation and the Wikipedia entry.

How does this connect to the attention-only transformer?

In the "attention-only" formulation of the transformer we can write the "residual" of a fixed head as $AX{W}_V{W}_O$, with the values weight matrix $W_V$, the attention matrix $A$, the output weight matrix $W_O$, and the current embeddings at each position $X$

Let $E$ be the embedding dimension, $L$ the total context length and $D$ the dimension of the values, then we have that

• $X$ is an $L\times E$ matrix,
• $A$ is a $L\times L$ matrix,
• $W_V$ is a $E\times D$, and
• $W_O$ is a $D\times E$ matrix

Let's identify the participating vector spaces:

$A$ maps from the "position" space back to the "position" space, which we will call $P$ (and which is isomorphic to $R^L$). Similarly, we have the "embedding" space $E\cong R^E$ and the "value" space $V\cong R^D$.

It might become clear now that we can identify $X$ with an element from $P\otimes E$, i.e. that we can write $X=X_{ij}(e_i^P\otimes e_j^E)$.

From that lense, we can see that right-multiplying $X$ with $W_V$ is equivalent to multiplying with $\mathrm{Id}\otimes W_V$, which maps an element from $P\otimes E$ to an element from $P\otimes V$, by applying $W_V$ to the $E$-part of the tensor ^[3]:

$$(\mathrm{Id}\otimes W_V)\left(X\right)=(\mathrm{Id}\otimes W_V)\sum _{ij}{X}_{ij}{e}_i^P\otimes e_j^E=\sum _{ij}{X}_{ij}{e}_i^P\otimes W_V\left(e_j^E\right)=\sum _{ij}{X}_{ij}{e}_i^P\otimes \sum _k{W}_{jk}{e}_k^V=\sum _{ik}\sum _j\left(X_{ij}{W}_{jk}\right)e_i^P\otimes e_k^V=\sum _{ik}(X{W}_V{)}_{ik}{e}_i^P\otimes e_k^V=X{W}_V$$

Identical arguments hold for $W_O$ and $A$, so that we get the formulation from the paper:

$$AX{W}_O{W}_V=(A\otimes W_O{W}_V)\cdot X$$

Note that there is nothing special about this in terms of what these matrices represent. So it seems that a takeaway message is that whenever you have a matrix product of the form $ABC$ you can re-write it as $(A\otimes C)\cdot B$ (Sorry to everyone who thought that was blatantly obvious from the get-go ;P).^[4]