My takeaway series follow a Q&A format to explain AI concepts at three levels:
Anyone with general knowledge can understand them.
For anyone who wants to dive into the code implementation details of the concept.
For anyone who wants to understand the mathematics behind the technique.
Attention is a type of neural network layer that allows neural networks to focus on specific parts of the input sequence when making predictions.
Attention was first introduced in 2014 (Bahdanau, Cho, and Bengio 2014) for machine translation tasks.
The background was the need to improve the performance of sequence-to-sequence models, particularly for tasks like machine translation. Traditional RNNs struggled with long-range dependencies and often produced poor translations for long sentences. The attention mechanism allowed the model to focus on specific parts of the input sequence when generating each word in the output sequence, significantly improving translation quality.
Later, the self-attention mechanism was proposed in the Transformer architecture (Vaswani et al. 2017) in 2017, which further advanced the use of attention in deep learning. Please refer to AI Concept Takeaway: Transformer for more details.
We first discuss the attention mechanism in general, then the attention layer in neural networks.
Given a set of vectors, we want to know an other vector should pay attention to which vectors, and how much attention should be paid to each vector. This is what the attention mechanism does.
Input:
- Query: a vector \(\mathbf{q} \in \mathbb{R}^{d_k}\) that represents the element we want to compute attention for.
- Targets: a dictionary of vectors \(\{\mathbf{k}_1: \mathbf{v}_1, \cdots, \mathbf{k}_n: \mathbf{v}_n\}\):
- Key \(\mathbf{k}_i \in \mathbb{R}^{d_k}\) represents the identity of each element in the dictionary. They are what the query will be compared against. This can be represented as a matrix \(\mathbf{K} = (\mathbf{k}_1; \cdots; \mathbf{k}_n) \in \mathbb{R}^{n\times d_k}\). Note that the dimension \(d_k\) of the key vector \(\mathbf{k}_i\) is the same as that of the query vector \(\mathbf{q}\).
- Value \(\mathbf{v}_i \in \mathbb{R}^{d_v}\) represents the information content of each element in the dictionary. This can be represented as a matrix \(\mathbf{V} = (\mathbf{v}_1; \cdots; \mathbf{v}_n) \in \mathbb{R}^{n\times d_v}\). Note that the dimension \(d_v\) of the value vector \(\mathbf{v}_i\) can be different from that of the query and key vectors \(\mathbf{q}\) and \(\mathbf{k}_i\).
Output:
- Attention Scores: a vector of attention scores \(\mathbf{a} \in \mathbb{R}^n\) that indicates how much attention should be paid to each target element. Each score is a 0-1 value, and the sum of all scores is 1.
- Weighted Sum Value: a vector \(\mathbf{z} \in \mathbb{R}^{d_v}\) that is the weighted sum of the values, where the weights are given by the attention scores.
There can be multiple query vectors \(\mathbf{q}_1, \cdots, \mathbf{q}_m\) to compute attention at the same time, and the attention mechanism can be applied to each query vector independently. This is similar to the batch processing of data in neural networks. Therefore, the input query can be a matrix \(\mathbf{Q} = (\mathbf{q}_1; \cdots; \mathbf{q}_m) \in \mathbb{R}^{m \times d_k}\), the output attention scores can be a matrix \(\mathbf{A} \in \mathbb{R}^{m \times n}\), and the output weighted sum values can be a matrix \(\mathbf{Z} \in \mathbb{R}^{m \times d_v}\).
The attention scores are computed by the similarity between the query \(\mathbf{q}\) and the keys \(\mathbf{k}_i\). The similarity can be computed in various ways, such as dot product, cosine similarity, or learned similarity functions. The most common method is the dot product \(\mathbf{q} \cdot \mathbf{k}_i\).
The attention scores are each 0-1 value, and the sum of all scores is 1. This can be achieved by applying the softmax function to the similarity scores. But before applying the softmax function, the similarity scores are often scaled by the square root of the dimension of the key vectors \(\sqrt{d_k}\) to prevent the scores from being too large or too small, which can lead to numerical instability:
\[ \mathbf{a}_i = \text{softmax}\left(\frac{\mathbf{q} \cdot \mathbf{k}_i}{\sqrt{d_k}}\right), i = 1, \cdots, n \]
\[ \mathbf{z} = \sum_{i=1}^n \mathbf{a}_i \mathbf{v}_i \]
This is known as scaled dot-product attention.
The matrix form of the attention computation is:
\[ \mathbf{A} = \text{softmax}\left(\frac{\mathbf{Q} \mathbf{K}^T}{\sqrt{d_k}}\right), \mathbf{Z} = \mathbf{A} \mathbf{V} \]
It is often seen as a single attention function: \[ \text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q} \mathbf{K}^T}{\sqrt{d_k}}\right) \mathbf{V} \]
The attention scores are computed by the similarity between the query \(\mathbf{q}\) and the keys \(\mathbf{k}_i\). The similarity can be computed in various ways, such as dot product, cosine similarity, or learned similarity functions. The most common method is the dot product \(\mathbf{q} \cdot \mathbf{k}_i\).
The attention scores are each 0-1 value, and the sum of all scores is 1. This can be achieved by applying the softmax function to the similarity scores. But before applying the softmax function, the similarity scores are often scaled by the square root of the dimension of the key vectors \(\sqrt{d_k}\) to prevent the scores from being too large or too small, which can lead to numerical instability:
\[ \mathbf{a}_i = \text{softmax}\left(\frac{\mathbf{q} \cdot \mathbf{k}_i}{\sqrt{d_k}}\right), i = 1, \cdots, n \]
\[ \mathbf{z} = \sum_{i=1}^n \mathbf{a}_i \mathbf{v}_i \]
This is known as scaled dot-product attention.
The matrix form of the attention computation is:
\[ \mathbf{A} = \text{softmax}\left(\frac{\mathbf{Q} \mathbf{K}^T}{\sqrt{d_k}}\right) \]
\[ \mathbf{Z} = \mathbf{A} \mathbf{V} \]
It is often seen as a single attention function: \[ \text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q} \mathbf{K}^T}{\sqrt{d_k}}\right) \mathbf{V} \]
In neural networks, the attention mechanism is typically implemented as an attention layer. The input to the attention layer is a set of query, key, and value vectors, which are usually obtained by applying linear transformations to the input data \(\mathbf{X} \in \mathbb{R}^{n \times d_x}\). The output of the attention layer has the same size \(n\) but different dimension \(d\), and can be used as input to subsequent layers in the neural network. Please note that \(\mathbf{X}\) is usually a piece of sequential data, so \(n\) is the sequence length, not the batch size.
Attention layer is a parameterised layer. The weights of the linear transformations are learned during training, allowing the model to adaptively learn how to compute attention based on the input data. These are the parameters of the attention layer.
The earliest use of attention in neural networks was in sequence-to-sequence models for machine translation (Bahdanau, Cho, and Bengio 2014). In these years, RNN-based encoder-decoder architectures were widely used for sequence-to-sequence tasks. Here the query vector is (a linear transformation of) the decoder hidden state at the current time step, and the key and value vectors are (linear transformations of) the encoder hidden states at all time steps. The attention output represents a context vector that summarizes the relevant information from the input sequence based on the current decoder state’s attention. The prediction is then made based on this, which is connected to the output layer.
Later, the self-attention mechanism was proposed in the Transformer architecture (Vaswani et al. 2017), which further advanced the use of attention in deep learning. In self-attention, the query, key, and value vectors are all derived from the same input sequence \(\mathbf{X}\):
\[ \mathbf{Q} = \mathbf{X} \mathbf{W}_Q, \mathbf{K} = \mathbf{X} \mathbf{W}_K, \mathbf{V} = \mathbf{X} \mathbf{W}_V \]
\(\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V\) are learnable weights. The attention scores are now computed between all pairs of elements in the input sequence (where \(\mathbf{A}\) is a square matrix), allowing the model to capture dependencies between different parts of the sequence. Please refer to AI Concept Takeaway: Transformer for more details.