What is multi-head attention doing mathematically, and how is it different from self-attention?

Question

I'm trying to understand the difference between the concept of self-attention and multi-head attention. The latter is not actually too clear to me.

I understand that, in the case of self-attention, we start with a feature matrix $X \in \mathbb{R}^{n \times d}$, and then we use the same linear transformation $W$ to produce

\begin{align} Q &= XW \\ K &= XW \\ V &= XW \end{align}

and then we compute the following

$$X' = \text{softmax} \left(\frac{Q\cdot K^T}{\sqrt{d}} \right)V$$

where $X' \in \mathbb{R}^{n \times d}$ is a new version of the input matrix, where the pairwise interactions between the points will be encoded.

What is multi-head attention doing, from a mathematical point of view, and what's the difference? I know we are going to use three different linear transformations in this case (so no weight-sharing), but what exactly will be encoded using three different $W$? Maybe it's more the conceptual view that it's not too clear in this case.

Answer 1

First of all, I believe that in self-attention mechanism for Query, Key and Value vectors the different linear transformations are used, $$ Q = XW_Q,\,K = XW_K,\,V = XW_V; W_Q \neq W_K, W_K \neq W_V, W_Q \neq W_V $$ The self-attention itself is a way of using more general attention mechanism. You can check this post for examples of other mechanisms that can be used.

Multi-head attention is a way of grouping together a bunch of attention mechanism ( Usually they are all the same type ), which consists in just running multiple mechanism in parallel and aggregating the resulting set in some way. In here for example, the aggregation is done by concatenation and weighted sum of the outputs, e.g. \begin{equation*} \begin{aligned} &\text{MultiHead}(Q, K, V)= [\text{head}_1, .., \text{head}_n]W^O & [..] \text{ meaning concatenation}\\ &\text{where head}_i = \text{Attention}(QW_i^Q,KW_i^K,VW_i^V) & \text{each head}_i\text{ is unique attention mechanism} \\ &\text{Attention}(\textbf{Q},\textbf{K},\textbf{V}) = \text{softmax}(\frac{\textbf{Q}\textbf{K}}{\sqrt{d_k}})\textbf{V} & \text{this can be other attention mechanism}\\ \end{aligned} \end{equation*}