# Understanding the tensor product formulation in Transformer Circuits

I was trying to understand the [tensor product formulation](https://transformer-circuits.pub/2021/framework/index.html#architecture-attn-as-movement) 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 $\mathcal V, \mathcal W$ we can construct the tensor product space $\mathcal V \otimes \mathcal W$ as the [_span_](https://en.wikipedia.org/wiki/Linear_span)[^1] of all matrices $\mathbf v \otimes \mathbf w$, where $\mathbf v\in\mathcal V, \mathbf w\in \mathcal W,$ with the property $(\mathbf v \otimes \mathbf w)_{ij} = v_iw_j$ [^2]. We can equivalently define it as a vector space with basis elements $\mathbf e^{\mathcal V}_i \otimes \mathbf e^{\mathcal W}_j$, where we used the basis elements of $\mathcal V$ and $\mathcal 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) $\mathbf A:\mathcal V \to \mathcal X, \mathbf B:\mathcal W\to \mathcal Y$ we can define $\mathbf A \otimes \mathbf B : \mathcal V \otimes \mathcal W \to \mathcal X \otimes \mathcal Y$, where each map simply operates on its own vector space, not interacting with the other: 

$$(\mathbf A \otimes \mathbf B)(\mathbf v \otimes \mathbf w) = \mathbf A(\mathbf v) \otimes \mathbf B(\mathbf w)$$

For more information on the tensor product, I recommend this [intuitive explanation](https://www.math3ma.com/blog/the-tensor-product-demystified) and the [Wikipedia entry](https://en.wikipedia.org/wiki/Tensor_product).

## 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 $\mathbf A \mathbf X \mathbf W_V \mathbf W_O$, with the values weight matrix $\mathbf W_V$, the attention matrix $\mathbf A$, the output weight matrix $\mathbf W_O$, and the current embeddings at each position $\mathbf X$

Let $E$ be the embedding dimension, $L$ the total context length and $D$ the dimension of the values, then we have that 
- $\mathbf X$ is an $L \times E$ matrix,
- $\mathbf A$ is a $L\times L$ matrix, 
- $\mathbf W_V$ is a $E \times D$, and 
- $\mathbf W_O$ is a $D \times E$ matrix

### Let's identify the participating vector spaces: 

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

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

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

$$(\operatorname{Id}\otimes \mathbf W_V)(\mathbf X)\\ = (\operatorname{Id}\otimes \mathbf W_V)\sum_{ij} X_{ij} \mathbf e^{\mathcal P}_i \otimes \mathbf e^{\mathcal E}_j\\ = \sum_{ij}X_{ij} \mathbf e^{\mathcal P}_i \otimes \mathbf W_V(\mathbf e^{\mathcal E}_j)\\ = \sum_{ij}X_{ij} \mathbf e^{\mathcal P}_i \otimes \sum_k W_{jk}\mathbf e^{\mathcal V}_k\\ = \sum_{ik} \sum_j (X_{ij}W_{jk}) \mathbf e^{\mathcal P}_i \otimes \mathbf e^{\mathcal V}_k\\ = \sum_{ik} (\mathbf X\mathbf W_V)_{ik} \mathbf e^{\mathcal P}_i \otimes \mathbf e^{\mathcal V}_k\\ = \mathbf X\mathbf W_V$$


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

$$\mathbf A \mathbf X \mathbf W_O \mathbf W_V = (\mathbf A \otimes \mathbf W_O \mathbf W_V ) \cdot \mathbf 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 $\mathbf{ABC}$ you can re-write it as $(\mathbf A \otimes \mathbf C)\cdot \mathbf B$ (Sorry to everyone who thought that was blatantly obvious from the get-go ;P).[^4]

[^1]: A previous edition of this post said that it was the _space_ of all such matrices which is inaccurate. The _span_ of a set of vectors/matrices is the space of all linear combinations of elements from that set.

[^2]: I'm limiting myself to finite-dim spaces because that's what is relevant to the transformer circuits paper. The actual formal definition is more general/stricter but imo doesn't add much to understanding the application in this paper

[^3]: Note that the 'linear map' that we use here is basically right multiplying with $\mathbf W_V$, so that it maps $$\mathbf e^{\mathcal E}_k \mapsto \mathbf W_V^T\mathbf e^{\mathcal E}_k$$

[^4]: I should note that this is also what is mentioned in the paper's [introduction](https://transformer-circuits.pub/2021/framework/index.html#notation-tensor-product) on tensor products, but it didn't click with me, whereas going through the above steps did.