Attention in computer vision | by Javier Fernandez | Towards Data Science

Ever since the introduction of Transformer in the work “Attention is all you need”, there has been a transition in the field of NLP towards replacing Recurrent Neural Networks (RNN) with attention-based networks. In current literature, there are already many great articles that describe this method. Here are two of the best ones I found during my review: The Annotated Transformer and Transformers explained Visually.
However, after researching how to implement attention in computer vision (Best found articles: Understanding Attention Modules, CBAM, Papers with Code-Attention, Self-Attention, Self-Attention and Conv), I noticed that only a few of them clearly describe the attention mechanisms and include clean code along with the theory. Hence, the goal of this article is to describe in detail the two most important attention modules in computer vision and apply them to a practical case using PyTorch. The structure of the article is as follows:
Introduction to attention moduleAttention method in computer visionImplementation and results of attention-based networksConclusion
Introduction to attention module
Attention method in computer vision
Implementation and results of attention-based networks
Conclusion
In the context of machine learning, attention is a technique that mimics cognitive attention, defined as the ability to choose and concentrate on relevant stimuli. In other words, attention is a method that tries to enhance the important parts while fading out the non-relevant information.
Despite this mechanism can be divided into several families (Attention? Attention!), we focus on self-attention since it is the most popular type of attention for computer vision tasks. This one refers to the mechanism of relating different positions of a single sequence to compute a representation of the same sequence.
To better understand this concept, let’s think of the following sentence: Bank of a river. Do we agree that the word Bank losses its contextual information if we don’t see the word River? Well, that’s actually the main idea behind self-attention. It tries to give contextual information to each word since individual meanings of the words do not represent their meanings in the sentence.
As explained in An Intuitive Explanation of Self-attention, if we consider the example given above, self-attention works by comparing every word in the sentence to every other word and reweighing the word embeddings of each word to include contextual relevance. The input to the output module is the embedding for each word without contextual information while the output is a similar embedding with contextual information.
Here there is listed a continuously updating list of attention modules. From the listed ones, we focus on the two most popular ones for computer vision tasks: Multi-Head Attention and Convolutional Block Attention Module (CBAM).
2.1. Multi-Head Attention
Multi-Head Attention is a module for attention mechanism that runs an attention module several times in parallel. Hence, to understand its logic it is first needed to understand the Attention module. The two most commonly used attention functions are Additive Attention and Dot-Product Attention, being the latter the one of interest for this work.
The basic structure of the Attention module is that there are two lists of vectors x1 and x2, one which is attended and the other one which attends. The vector x2 generates a ‘query’ while the vector x1 creates a ‘key’ and a ‘value’. The idea behind the attention function is to map the query and the set key-value pairs to an output. “The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key” [Attention is all you need]. The output is computed as follows:
As mentioned in this discussion, the key/value/query concepts come from retrieval systems. For example, when typing a query on Youtube to search for some video, the search engine will map your query against a set of keys (video title, description, etc.) linked with candidate videos in the database. Then, it will present you with the best-matched videos (values).
Let’s run this Dot-Product Attention before moving towards the Multi-Head attention, which is an extension of this module. Below is the implementation in PyTorch. The input is [128, 32, 1, 256], where 128 corresponds to the batch, 32 to the sequence length, 1 to the number of heads (we will increase it for multiple attention heads), and 256 are the number of features.
The output is:
attn_output: [128, 32, 1, 256], attn_weights: [128, 1, 32, 32]attn_output: [128, 32, 1, 256], attn_weights: [128, 1, 32, 16]
Some takeaways from this basic implementation:
The output will have the same shape as the query input size.
The attention weights for each data have to be a matrix where the number of rows corresponds to the sequence length of the query and the number of columns to the sequence length of the key.
There is no learnable parameter in the Dot-Product Attention.
So, coming back to Multi-Head Attention, this one runs this explained Attention module several times in parallel. The independent attention outputs are then concatenated and linearly transformed into the expected dimension. Here is the implementation:
The output is:
attn_output: [128, 32, 256], attn_weights: [128, 8, 32, 32]attn_output: [128, 32, 256], attn_weights: [128, 8, 32, 32]
From the code, observe that:
The input to the linear layer for, for example, the query is [128, 32, 256]. However, and as is mentioned in this post, the Linear layer accepts tensors of arbitrary shape, where only the last dimension must match with the in_features argument you specified in the constructor. The output will have exactly the same shape as the input, only the last dimension will change to whatever you specified as out_features in the constructor. For our case, the input shape is a set of 128 * 32 = 4096 and 256 features. Therefore, we are applying the dense net to each of the elements of the sequence length and each data of the batch.
Also, we added the residual connection and layer normalization as it is implemented in the transformer neural network. But, those should be excluded if you just want to implement the Multi-Head Attention module.
So, at this point you might be wondering, why do we implement Multi-Head Attention instead of just a simple Attention module? According to the paper Attention is all you need, “multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.” In other words, dividing the features into heads allows each attention module to focus only on a set of features, giving greater power to encode multiple relationships and nuances for each word.
If at this point you still want to dig more into this type of attention, I encourage you to read this article, which explains in detail all this module with great illustrations.
Before finishing, I just want to mention that we have used this attention module as if we were working with sequences but this article is about images. If you understood to this point everything explained, the only difference between sequences and images is the input vector. What corresponds to the sequence length, for images will be the pixels. Therefore, if the input is [batch=128, no_channels=256, height=24, width=24] , one possible implementation could be:
The output is:
attn_output: [128, 256, 24, 24], attn_weights: [128, 8, 576, 576]
2.2. Convolutional Block Attention Module (CBAM)
In 2018, S. Woo et al. (2018) published a new attention module named Convolutional Block Attention Module (CBAM) that, as well as convolutions operations do, emphasizes meaningful features along the channel and spatial axes. Compared to Multi-Head attention, this type of attention was intentionally made for feed-forward convolutional neural networks and can be applied at every convolutional block in deep networks.
CBAM contains two sequential sub-modules called the Channel Attention Module (CAM) and the Spatial Attention Module (SAM). These two concepts are probably the two most important concepts when speaking about convolutions. While channel refers to the number of features or channels for each pixel, spatial refers to the feature maps of dimension (h x w).
Spatial Attention Module (SAM): This module is comprised of a three-fold sequential operation. The first part of it is called the Channel Pool and it consists of applying Max Pooling and Average Pooling across the channels to the input (c × h × w) to generate an output with shape (2 × h × w). This is the input to a convolution layer that outputs a 1-channel feature map (1 × h × w). After passing this output through a BatchNorm and an optional ReLU, the data goes to a Sigmoid Activation layer.
Channel Attention Module (CAM): This module first decomposes the input tensor into 2 subsequent vectors of dimensionality (c × 1 × 1) generated by Global Average Pooling (GAP) and Global Max Pooling (GMP). Thereafter, the output goes through a fully connected layer followed by a ReLu activation layer.
For further information about CBAM, I recommend reading this great post which has great explanatory images.
Here is the implementation:
The output is:
attn_output: [128, 256, 24, 24]
After the theoretical part presented above, this section focuses on the implementation of both attention layers on a practical case.
Specifically, we have selected the STL dataset and have included a white patch to some images as displayed below. The task is to create a neural network that classifies both types of images.
Then, we created three classes. The first one is just a CNN while the second contains the Multi-Head attention layer and the third includes the CBAM module.
Here is the code to run the training.
And these are the outputs:
CNN:
Min train error: 0.0011167450276843738Min test error: 0.05411996720208516
CNN + Multi-Head attention: The performance improves when adding the attention layer but the attention plots do not highlight the part of the image with the white patch.
Min train error: 9.811600781858942e-06Min test error: 0.04209221125441423
Since there was some overfitting and the attention layer did not do what it was supposed to, I reimplemented this layer using convolutional layers. If anyone has any suggestions, please leave them in the comments.
CNN + 1DConv-based Multi-Head attention: This time, the stability and performance improve notably. Also, it is possible to observe how the output of the attention layer highlights the white patch for the images that contain it.
Min train error: 0.00025470180017873645Min test error: 0.014278276459193759
CNN + CBAM attention: This one presents the best result. It is clearly possible to observe the white patch in the output of the attention layer and the training is really stable, achieving the lowest validation loss of all the models.
Min train error: 2.786791462858673e-05Min test error: 0.028047989653949175
In summary, this article presents Multi-Head Attention and CBAM modules, the two most popular attention modules in computer vision. Also, it includes an implementation in PyTorch where we classify images from the CIFAR dataset that contains a white patch (added manually).
For future work, I consider it interesting to include positional encoding along with the attention. Here I leave a link for anyone interested: