Multi-Head Self-Attention

Multiple self-attention mechanisms are usually applied in parallel. Now, $H$ different sets of values, keys, and queries are computed:

$$\begin{array}{rl}{V}_h& =1{\beta }_{vh}+X{\Omega }_{vh}\\ Q_h& =1{\beta }_{qh}+X{\Omega }_{qh}\\ K_h& =1{\beta }_{kh}+X{\Omega }_{kh}\end{array}$$

The $h$-th self-attention mechanism or head can be written as:

$${\text{Sa}}_h[X]=\text{Softmax}\left[\frac{Q_h{K}_h^T}{\sqrt{D_k}}\right]\cdot V_h$$

where we have different parameters $\{{\beta }_{vh},{\Omega }_{vh}\},\{{\beta }_{qh},{\Omega }_{q,h}\},\{{\beta }_{kh},{\Omega }_{kh}\}$ for each head.

Typically, if the dimensions of the input $x_m$ is $D$ and there are $H$ heads, the values, queries, and keys will all be of size $D/H$, as this allows for an efficient implementation.

The outputs for these self-attention mechanisms are concatenated along the feature dimension, and another linear transform ${\Omega }_c$ is applied to combine them:

\text{MhSA}[X]
=
[\text{Sa}_1[X],\ \text{Sa}_2[X],\ \dots,\ \text{Sa}_H[X]]
\,\Omega_c

• Note that diagram is using data of shape $(D,N)$ whereas my equations use $(N,D)$. In the $(D,N)$ case, concatenating along the feature dimension means we concatenate vertically.

Multiple heads seem to be necessary to make self-attention to work well. It has been speculated that they make the self-attention network more robust to bad initializations.

dl Multi-head self-attention ? Run multiple self-attention heads in parallel, concatenate their outputs along the feature dimension, then apply a learned linear transformation to combine them.

\text{MhSA}[X]
=
[\text{Sa}_1[X],\ \text{Sa}_2[X],\ \dots,\ \text{Sa}_H[X]]
\,\Omega_c

+++

In multi-head self-attention, if the embedding dimension is $D$ and there are $H$ heads, what is the dimension of the values, queries, and keys for each head?::$D/H$