Self-Attention and Cross-Attention
Introduction
Despite their names, self-attention and cross-attention are not definitions of the attention function itself. The attention function refers to the function computed as follows when a query matrix $\mathbf{Q}$, a key matrix $\mathbf{K}$, and a value matrix $\mathbf{V}$ are given.
$$ \operatorname{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V}) = \mathbf{V} \operatorname{Softmax}\left(f(\mathbf{Q},\mathbf{K})\right) \tag{1} $$
- $\mathbf{Q} \in M^{d_{k} \times n}$
- $\mathbf{K} \in M^{d_{k} \times m}$
- $\mathbf{V} \in M^{d_{v} \times m}$
Here $\operatorname{Softmax}$ denotes the function that, for a given 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 attention function varies depending on how $f$ in $(1)$ is defined concretely, but self-attention and cross-attention are not about how to define the attention function; they are about how to define the query/key/value matrices.
Definition
Suppose a data matrix $\mathbf{X} \in M^{d \times n}$ and learnable weight matrices $\mathbf{W}^{Q} \in M^{d_{k} \times d}$, $\mathbf{W}^{K} \in M^{d_{k} \times d}$, $\mathbf{W}^{V} \in M^{d_{v} \times d}$ are given. The method of obtaining the query, key, and value matrices all from the same data $\mathbf{X}$ is called self-attention.
$$ \text{Self-Attention:} \quad \mathbf{Q} = \mathbf{W}^{Q}\mathbf{X}, \quad \mathbf{K} = \mathbf{W}^{K}\mathbf{X}, \quad \mathbf{V} = \mathbf{W}^{V}\mathbf{X} $$
On the other hand, for two different data matrices $\mathbf{X} \in M^{d_{x} \times n}$, $\mathbf{Y} \in M^{d_{y} \times m}$, the method of obtaining the query matrix from $\mathbf{X}$ and the key and value matrices from $\mathbf{Y}$ is called cross-attention.
$$ \text{Cross-Attention:} \quad \mathbf{Q} = \mathbf{W}^{Q}\mathbf{X}, \quad \mathbf{K} = \mathbf{W}^{K}\mathbf{Y}, \quad \mathbf{V} = \mathbf{W}^{V}\mathbf{Y} $$
In this case the sizes of the weight matrices are $\mathbf{W}^{Q} \in M^{d_{k} \times d_{x}}$, $\mathbf{W}^{K} \in M^{d_{k} \times d_{y}}$, $\mathbf{W}^{V} \in M^{d_{v} \times d_{y}}$.
Explanation
Summarized in the form of the attention function, they are as follows.
$$ \begin{align*} \text{Self-Attention:} \quad &\operatorname{Attention}(\mathbf{W}^{Q}\mathbf{X}, \mathbf{W}^{K}\mathbf{X}, \mathbf{W}^{V}\mathbf{X}) \\ &= \mathbf{W}^{V}\mathbf{X} \operatorname{Softmax}\left( \left( \mathbf{W}^{K}\mathbf{X} \right)^{\mathsf{T}} \left( \mathbf{W}^{Q}\mathbf{X} \right) \right) \end{align*} $$
$$ \begin{align*} \text{Cross-Attention:} \quad &\operatorname{Attention}(\mathbf{W}^{Q}\mathbf{X}, \mathbf{W}^{K}\mathbf{Y}, \mathbf{W}^{V}\mathbf{Y}) \\ &= \mathbf{W}^{V}{\color{red}\mathbf{Y}} \operatorname{Softmax}\left( \left( \mathbf{W}^{K}{\color{red}\mathbf{Y}} \right)^{\mathsf{T}} \left( \mathbf{W}^{Q}\mathbf{X} \right) \right) \end{align*} $$
Self-attention is a sequence attending to itself. Since the query, key, and value all come from the same data $\mathbf{X}$, the number of queries equals the number of keys/values ($m = n$), and each column vector of $\mathbf{X}$—for example, each word of a sentence—measures a score against every other word in the same sentence and mixes the values according to those weights. As a result, the representation of each word is updated so as to reflect the context of the entire sentence it belongs to. For instance, which noun in the sentence a pronoun such as 'it' refers to is revealed by the weights of self-attention. The encoder of the Transformer is a representative user of self-attention, and it is also called intra-attention.
Cross-attention is one of two different sequences referring to the other. Machine translation is a representative example: the feature vector on the side of the translated text being produced by the decoder becomes the query $\mathbf{X}$, and the feature vector on the side of the source text produced by the encoder becomes the key and value $\mathbf{Y}$. Bahdanau attention, which first introduced attention into recurrent neural network-based seq2seq models, takes exactly this form, and the attention layer in the Transformer decoder that refers to the encoder’s output is also cross-attention. In this context, cross-attention is also called encoder-decoder attention.