Deep Learning and Types of Attention

In an earlier post on reinforcement learning, a memory problem was described: an agent often needs some form of memory to act well when it cannot see the full state, and a recurrent network is one common fix. Sequence to sequence models in deep learning ran into a closely related problem years earlier, and the fix that emerged for it, attention, turned out to matter well beyond the original translation task it was built for. Attention is now the core building block behind transformers, and by extension behind most of the large language models in use today.

This post covers why attention was introduced, the general mechanism behind it, and the main types that have emerged since.

The bottleneck problem in sequence models

Before attention, the standard approach to a task like machine translation was an encoder decoder recurrent network. An encoder RNN reads an input sequence one token at a time and produces a final hidden state $h_T$. This single vector is expected to summarize the entire input, and a decoder RNN then generates the output sequence conditioned only on that one vector.

$$h_T = f(x_1, x_2, \ldots, x_T)$$

The problem is compression. A fixed length vector $h_T$ has to carry everything relevant about a sentence that could be five words or fifty words long, and empirically, translation quality dropped off sharply as sentence length grew, exactly as this bottleneck would predict. Bahdanau, Cho, and Bengio identified this as the central limitation of the encoder decoder architecture and proposed attention as the fix in 2014.

What attention actually computes

The general idea is to stop asking the encoder to compress everything into one vector, and instead let the decoder look back at every encoder hidden state at each output step, weighted by how relevant that state is to the token being generated right now.

This is usually described with three objects: a query, a set of keys, and a set of values. In the original translation setting, the query is the decoder's current state, the keys and values are the encoder's hidden states at every input position (often the same vectors serve as both). A score function compares the query against each key to produce an alignment score, the scores are normalized with a softmax into weights that sum to one, and the output is a weighted sum of the values.

$$e_{ij} = score(s_{i-1}, h_j)$$

$$\alpha_{ij} = \frac{\exp(e_{ij})}{\sum_{k} \exp(e_{ik})}$$

$$c_i = \sum_{j} \alpha_{ij} h_j$$

Here $c_i$ is the context vector used to generate the $i$th output token, and it is a different weighted combination of the encoder states for every output step, rather than the single fixed vector from before. The score function is where the different types of attention below mostly differ.

Types of attention

Additive attention

Additive attention, introduced in the original Bahdanau paper, learns the score function as a small feedforward network with a single hidden layer.

$$score(s, h) = v^{T} \tanh(W_1 s + W_2 h)$$

The vector $v$ and matrices $W_1$, $W_2$ are learned parameters. This form works well when the query and key vectors have different dimensions, since $W_1$ and $W_2$ can project them into a shared space before comparing, but it is more expensive to compute than the multiplicative forms below since it involves an extra nonlinearity and a learned projection for every query key pair.

Multiplicative attention

Luong, Pham, and Manning proposed a family of simpler and cheaper score functions in 2015, all based on a dot product rather than a learned feedforward layer.

$$score(s, h) = s^{T} h \quad \text{(dot)}$$

$$score(s, h) = s^{T} W h \quad \text{(general)}$$

The dot variant needs no learned parameters at all and is extremely cheap, but requires the query and key to already live in the same dimensional space. The general variant adds a single learned matrix $W$ to relax that requirement while staying much cheaper than the additive form.

The same paper also introduced a distinction that is easy to miss but useful in practice: global attention looks at every encoder position for every output step, exactly as described above, while local attention first predicts a single position in the input that is most relevant and then attends only to a small window around it. Local attention trades a small amount of accuracy for a large reduction in compute on long sequences, since the softmax and weighted sum only need to cover a window instead of the whole input.

Self attention

Every attention mechanism above connects a decoder to an encoder. Self attention instead lets a sequence attend to itself: the queries, keys, and values are all derived from the same set of vectors, so every position in a layer can directly attend to every other position in that same layer, in a single step, regardless of how far apart they are.

This is the mechanism behind the transformer architecture from Vaswani and colleagues in 2017, and it uses a scaled dot product score.

$$Attention(Q, K, V) = softmax\left(\frac{QK^{T}}{\sqrt{d_k}}\right)V$$

Here $Q$, $K$, and $V$ are matrices packing every query, key, and value vector in the sequence, and $d_k$ is the dimension of the key vectors. The scaling term $\sqrt{d_k}$ matters more than it looks: for large $d_k$, the dot products $QK^T$ grow large in magnitude, which pushes the softmax into regions with extremely small gradients. Dividing by $\sqrt{d_k}$ keeps the scores in a range where the softmax stays well behaved and gradients keep flowing during training.

Because self attention connects every position to every other position directly rather than through a chain of recurrent steps, it removed the sequential dependency that made recurrent networks slow to train on long sequences, which is a large part of why transformers scale so well.

Multi head attention

A single self attention computation forces every relationship between two positions through one shared set of query, key, and value projections. Multi head attention instead runs several smaller self attention computations in parallel, each with its own learned projections, then concatenates the results and passes them through one more learned projection.

$$MultiHead(Q, K, V) = Concat(head_1, \ldots, head_h) W^{O}$$

Each head is free to specialize: empirically, different heads in a trained transformer end up attending to different kinds of structure, such as nearby positions, syntactic relationships, or a specific earlier token that a task depends on. A single head cannot represent all of these relationships well at once, so splitting the representation across several smaller heads gives the model more expressive room without a large increase in total compute.

Cross attention

Cross attention is really the original encoder decoder attention mechanism, kept alive inside the transformer architecture and given a name to distinguish it from self attention. The query comes from one sequence, typically the decoder, while the keys and values come from a different sequence, typically the encoder output.

The same pattern shows up well beyond text translation. A vision language model can use cross attention with text token queries against image patch keys and values, letting a caption generator or a multimodal language model decide which part of an image is relevant to the word it is about to produce. The mechanism does not change, only which two sequences are on either side of it.

Soft versus hard attention

Every form of attention described so far is soft attention: the weights $\alpha_{ij}$ form a smooth probability distribution and the output is a weighted average, which makes the whole operation differentiable and trainable end to end with backpropagation.

Hard attention instead samples a single position to attend to, rather than averaging over all of them. Xu and colleagues used exactly this distinction in their 2015 image captioning paper, where a soft attention variant computed a weighted average over image regions and a hard attention variant stochastically selected one region per generated word. The stochastic choice is not differentiable, so the hard attention variant had to be trained with a policy gradient method, the same REINFORCE algorithm mentioned in the reinforcement learning post as a model free, on policy technique. In practice soft attention is far more common today, since it trains more easily, but hard attention is a useful reminder that attention is fundamentally a selection mechanism, and a differentiable weighted average is only the most convenient way to approximate that selection.

Channel and spatial attention in vision

The types above all attend over positions in a sequence or an image. Convolutional vision models added a different axis: attention over channels. A squeeze and excitation block, introduced by Hu, Shen, and Sun in 2018, learns a weight for each feature channel in a convolutional layer based on a global summary of that channel across the whole image, then rescales the channel by that weight before passing it to the next layer. This lets the network emphasize feature channels that matter for the current image and suppress ones that do not, which is a form of attention orthogonal to the positional attention described earlier in this post, and the two ideas are often combined in a single vision model.

Why this mattered

Attention started as a fix for one bottleneck in machine translation, but the self attention and multi head attention forms removed the need for recurrence entirely. That made it possible to process an entire sequence in parallel during training instead of one token at a time, which is a large part of why transformer based models were able to scale to the size of modern large language models. The types covered here are not mutually exclusive either: a typical transformer today uses scaled dot product self attention within an encoder or decoder, multi head attention to let different heads specialize, and cross attention wherever one sequence needs to condition on another.

References

Bahdanau, D., Cho, K. and Bengio, Y. 2014. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473.

Luong, M., Pham, H. and Manning, C.D. 2015. Effective approaches to attention based neural machine translation. Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, 1412 to 1421.

Vaswani, A. et al. 2017. Attention is all you need. Advances in Neural Information Processing Systems, 30.

Xu, K. et al. 2015. Show, attend and tell: Neural image caption generation with visual attention. Proceedings of the 32nd International Conference on Machine Learning, 2048 to 2057.

Hu, J., Shen, L. and Sun, G. 2018. Squeeze and excitation networks. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 7132 to 7141.

Back to posts