📂Machine Learning

Multi-Head Attention

Introduction¹

Suppose we are given three matrices $\mathbf{Q} \in M^{d_{k} \times n}$, $\mathbf{K} \in M^{d_{k} \times m}$, $\mathbf{V} \in M^{d_{v} \times m}$, called the query, key, and value respectively, together with a score function $f : M^{d_{k} \times n} \times M^{d_{k} \times m} \to M^{m \times n}$ of the query and the key. Attention refers to the function defined as follows.

$$ \begin{align*} \operatorname{Attention}: M^{d_{k} \times n} \times M^{d_{k} \times m} \times M^{d_{v} \times m} &\to M^{d_{v} \times n} \\ (\mathbf{Q},\mathbf{K},\mathbf{V}) &\mapsto \mathbf{V} \operatorname{Softmax}\left(f(\mathbf{Q},\mathbf{K})\right) \tag{1} \end{align*} $$

Here $\operatorname{Softmax}$ denotes the function that, given a matrix $\mathbf{X} = \begin{bmatrix} \mathbf{x}_{1} & \cdots & \mathbf{x}_{N}\end{bmatrix}$, applies the softmax $\operatorname{softmax}$ to each column vector so that the entries of each column sum to $1$.

$$ \operatorname{Softmax}(\mathbf{X}) := \begin{bmatrix} \underset{\vert}{\overset{\vert}{\operatorname{softmax}(\mathbf{x}_{1})}} & \cdots & \underset{\vert}{\overset{\vert}{\operatorname{softmax}(\mathbf{x}_{N})}} \end{bmatrix} $$

The value of the attention computed in this way is the matrix obtained by taking a weighted average of the values, with the compatibility between the query and the key serving as the weights (for a detailed explanation, refer to the attention document).

$$ \operatorname{Attention} (\mathbf{Q},\mathbf{K},\mathbf{V}) \in M^{d_{v} \times n} $$

Let us pause here and think about image data and convolutional neural networks (CNN). An input image of size $(1, N, N)$ is processed through a CNN into an array of size $(C, M, M)$. Here $C$ is called the dimension of the channels, and each channel holds abstract information about the input data. If this abstract information were color, a color photograph would have the three RGB channels and be represented as $(3, N, N)$, while a black-and-white photograph has only one channel and is represented as $(1, N, N)$. Of course, the information held in the channels of the $(C, M, M)$ array obtained by passing through a CNN is not something a human can intuitively grasp like this, but in any case one can understand it as the information of the input data having been decomposed into $C$ kinds. In the same way, by using several attentions one could represent the information of $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$ in several channels. In attention, such a channel is called a head.

However, looking at $(1)$, the $\operatorname{Attention}$ function itself contains no learnable parameters. It consists only of a softmax and matrix multiplications (of course, if $f$ is built as a neural network then learnable parameters are included, but let us assume that is not the case here). For this reason, before sending the query, key, and value into $\operatorname{Attention}$, each of them is mapped in advance into several spaces by learnable linear transformation matrices.

Suppose the dimension of the query/key/value vectors is $d_{\text{model}}$. In other words, $\mathbf{Q} \in M^{d_{\text{model}} \times n}$, $\mathbf{K} \in M^{d_{\text{model}} \times m}$, $\mathbf{V} \in M^{d_{\text{model}} \times m}$. For each $h = 1, \cdots, H$, suppose linear transformations that compress the query/key into $d_{k}$ dimensions and the value into $d_{v}$ dimensions are given as follows.

$$ \mathbf{W}_{h}^{Q} \in M^{d_{k} \times d_{\text{model}}}, \qquad \mathbf{W}_{h}^{K} \in M^{d_{k} \times d_{\text{model}}}, \qquad \mathbf{W}_{h}^{V} \in M^{d_{v} \times d_{\text{model}}} $$

Then $(\mathbf{Q}, \mathbf{K}, \mathbf{V})$ is mapped into a different space for each head, and thus decomposed into $H$ kinds of information, that is, into $H$ heads.

$$ \operatorname{head}_{h} := \operatorname{Attention} (\mathbf{W}_{h}^{Q}\mathbf{Q}, \mathbf{W}_{h}^{K}\mathbf{K}, \mathbf{W}_{h}^{V}\mathbf{V}) \in M^{d_{v} \times n} \tag{2} $$

Note that whereas $\mathbf{Q} \in M^{d_{k} \times n}$ in $(1)$, in $(2)$ we have $\mathbf{Q} \in M^{d_{\text{model}} \times n}$ and $\mathbf{W}_{h}^{Q}\mathbf{Q} \in M^{d_{k} \times n}$.

$$ \begin{align*} \text{in } (1) &: \quad \mathbf{Q} \in M^{d_{k} \times n} \\ \text{in } (2) &: \quad \mathbf{Q} \in M^{d_{\text{model}} \times n}, \quad \mathbf{W}_{h}^{Q}\mathbf{Q} \in M^{d_{k} \times n} \end{align*} $$

Collecting all of these, we can express them in the following form.

$$ \begin{bmatrix} \operatorname{head}_{1} \\ \operatorname{head}_{2} \\ \vdots \\ \operatorname{head}_{H} \end{bmatrix} = \begin{bmatrix} \operatorname{Attention}(\mathbf{W}_{1}^{Q}\mathbf{Q}, \mathbf{W}_{1}^{K}\mathbf{K}, \mathbf{W}_{1}^{V}\mathbf{V}) \\ \operatorname{Attention}(\mathbf{W}_{2}^{Q}\mathbf{Q}, \mathbf{W}_{2}^{K}\mathbf{K}, \mathbf{W}_{2}^{V}\mathbf{V}) \\ \vdots \\ \operatorname{Attention}(\mathbf{W}_{H}^{Q}\mathbf{Q}, \mathbf{W}_{H}^{K}\mathbf{K}, \mathbf{W}_{H}^{V}\mathbf{V}) \end{bmatrix} \in M^{H d_{v} \times n} $$

And by multiplying by the following block matrix $\mathbf{W}^{O} \in M^{d_{\text{model}} \times H d_{v}}$, we obtain a linear combination of the heads.

$$ \mathbf{W}^{O} = \begin{bmatrix} \mathbf{W}_{1}^{O} & \mathbf{W}_{2}^{O} & \cdots & \mathbf{W}_{H}^{O} \end{bmatrix}, \quad \mathbf{W}_{h}^{O} \in M^{d_{\text{model}} \times d_{v}} $$

$$ \begin{bmatrix} \mathbf{W}_{1}^{O} & \mathbf{W}_{2}^{O} & \cdots & \mathbf{W}_{H}^{O} \end{bmatrix} \begin{bmatrix} \operatorname{head}_{1} \\ \operatorname{head}_{2} \\ \vdots \\ \operatorname{head}_{H} \end{bmatrix} = \sum_{h=1}^{H} \mathbf{W}_{h}^{O} \operatorname{head}_{h} \in M^{d_{\text{model}} \times n} $$

This matrix is the output of multi-head attention. Now let us define the function $\operatorname{MultiHead}$ as below.

Definition

Let the matrices of vectors of dimension $d_{\text{model}}$, namely $\mathbf{Q} \in M^{d_{\text{model}} \times n}$, $\mathbf{K} \in M^{d_{\text{model}} \times m}$, $\mathbf{V} \in M^{d_{\text{model}} \times m}$, be the query, key, and value matrices respectively. For each $h = 1, \cdots, H$, let $\mathbf{W}_{h}^{Q} \in M^{d_{k} \times d_{\text{model}}}$, $\mathbf{W}_{h}^{K} \in M^{d_{k} \times d_{\text{model}}}$, $\mathbf{W}_{h}^{V} \in M^{d_{v} \times d_{\text{model}}}$ be the linear transformations that map the column vectors of these matrices into $d_{k}$, $d_{k}$, $d_{v}$ dimensions, respectively. Define $\operatorname{head}_{h}$ as below.

$$ \operatorname{head}_{h} := \operatorname{Attention} (\mathbf{W}_{h}^{Q}\mathbf{Q}, \mathbf{W}_{h}^{K}\mathbf{K}, \mathbf{W}_{h}^{V}\mathbf{V}) \in M^{d_{v} \times n} $$

The function multi-head attention $\operatorname{MultiHead} : M^{d_{\text{model}} \times n} \times M^{d_{\text{model}} \times m} \times M^{d_{\text{model}} \times m} \to M^{d_{\text{model}} \times n}$ is defined as below.

$$ \begin{align*} \operatorname{MultiHead} (\mathbf{Q}, \mathbf{K}, \mathbf{V}) &:= \mathbf{W}^{O} \begin{bmatrix} \operatorname{head}_{1} \\ \operatorname{head}_{2} \\ \vdots \\ \operatorname{head}_{H} \end{bmatrix} \\ &= \sum_{h=1}^{H} \mathbf{W}_{h}^{O} \operatorname{head}_{h} \\ &= \sum_{h=1}^{H} \mathbf{W}_{h}^{O} \operatorname{Attention} (\mathbf{W}_{h}^{Q}\mathbf{Q}, \mathbf{W}_{h}^{K}\mathbf{K}, \mathbf{W}_{h}^{V}\mathbf{V}) \\ &= \sum_{h=1}^{H} \mathbf{W}_{h}^{O} \left( \mathbf{W}_{h}^{V}\mathbf{V} \right) \operatorname{Softmax} \left( f(\mathbf{W}_{h}^{Q}\mathbf{Q}, \mathbf{W}_{h}^{K}\mathbf{K}) \right) \end{align*} $$

Here $\mathbf{W}^{O} = \begin{bmatrix} \mathbf{W}_{1}^{O} & \cdots & \mathbf{W}_{H}^{O} \end{bmatrix} \in M^{d_{\text{model}} \times H d_{v}}$ is a learnable weight matrix, and the size of each block is $\mathbf{W}_{h}^{O} \in M^{d_{\text{model}} \times d_{v}}$.

In Mathematical Papers^{2 3}

The definition mainly used in mathematical papers is far simpler than the above. Considering the case of self-attention, the query, key, and value are obtained from a single piece of data $\mathbf{X}$ as follows.

$$ \mathbf{Q} = \mathbf{W}^{QX}\mathbf{X},\qquad \mathbf{K} = \mathbf{W}^{KX}\mathbf{X},\qquad \mathbf{V} = \mathbf{W}^{VX}\mathbf{X} $$

Then the input of the attention is expressed in the messy form $\mathbf{W}^{X}\mathbf{W}^{QX}\mathbf{X}$, but since a composition of linear transformations is also a linear transformation, it is essentially the same as expressing it by $\mathbf{W}^{X}\mathbf{X}$. Moreover, mathematically the attention function is nothing but the special case $H=1$ of multi-head attention, so attention is defined as the following function, which includes a residual layer and takes the score function to be the inner product.

$$ \operatorname{Attention}(\mathbf{X}) = \mathbf{X} + \sum_{h=1}^{H} \left( \mathbf{W}_{h}^{V}\mathbf{X} \right) \operatorname{Softmax} \left[ (\mathbf{W}_{h}^{K}\mathbf{X})^{\mathsf{T}} (\mathbf{W}_{h}^{Q}\mathbf{X}) \right] $$

Explanation

In the transformer paper, everything is described in terms of row vectors, so the dimensions of the matrices are reversed and the notation is as below.

$$ \operatorname{head}_{h} = \operatorname{Attention} (\mathbf{Q}\mathbf{W}_{h}^{Q}, \mathbf{K}\mathbf{W}_{h}^{K}, \mathbf{V}\mathbf{W}_{h}^{V}) \in M^{n \times d_{v}} $$

$$ \operatorname{MultiHead} (\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \begin{bmatrix} \operatorname{head}_{1} & \operatorname{head}_{2} & \cdots & \operatorname{head}_{H} \end{bmatrix} \begin{bmatrix} \mathbf{W}_{1}^{O} \\ \mathbf{W}_{2}^{O} \\ \vdots \\ \mathbf{W}_{H}^{O} \end{bmatrix} $$