A Deep Dive into Attention Mechanisms for Image Caption Generation

Introduction

Image captioning is a challenging task that requires understanding and describing the salient objects and events depicted in an image using natural language. In recent years, deep learning approaches have achieved remarkable progress on this multimodal task by leveraging encoder-decoder architectures inspired by neural machine translation. The key idea is to first encode the input image into a compact vector representation using a convolutional neural network (CNN), and then decode it into a caption by generating one word at a time using a recurrent neural network (RNN).

However, compressing an entire image into a single fixed-length vector poses some limitations, especially for complex scenes with many objects. The model may struggle to capture all the important details needed for generating an informative caption. This is where attention mechanisms come to the rescue! Instead of relying on a single vector, attention allows the model to focus on different parts of the image at each step of the caption generation process, adaptively selecting the most relevant visual features. Attention has been one of the major breakthroughs in image captioning, consistently boosting performance on standard benchmarks.

In this blog post, we‘ll take a deep dive into attention mechanisms for image captioning. We‘ll start by explaining the intuition behind attention and how it improves upon the basic encoder-decoder framework. Then we‘ll go through the technical details of two popular variants: Bahdanau (additive) attention and Luong (dot-product) attention. Finally, we‘ll walk through a Python implementation of an attention-based image captioning model in TensorFlow and Keras. Let‘s get started!

Encoder-Decoder Framework

Before diving into attention, let‘s briefly review the typical encoder-decoder architecture used for image captioning without attention. The model consists of two main components:

1. Encoder (CNN): A pre-trained CNN (e.g. VGG16, ResNet) is used to extract a feature map from the input image. The feature map is then flattened or pooled into a single vector which serves as the initial hidden state of the decoder RNN. This can be seen as a "summary" of the entire image.

2. Decoder (RNN): An RNN language model (usually an LSTM) takes the image vector and generates the caption one word at a time in a sequential manner. At each time step, it outputs a probability distribution over the vocabulary, from which the word with the highest probability is sampled. The RNN maintains a hidden state that gets updated based on the previously generated words.

Mathematically, let I denote the input image, (y_1, …, y_T) denote the ground-truth caption of length T, and (ŷ_1, …, ŷ_T) denote the generated caption.

The CNN encoder produces the image vector v:

v = CNN(I)

At each time step t, the LSTM decoder takes the image vector v and the previous word y_(t-1) as input, and outputs the conditional probability of the next word y_t:

ht = LSTM(v, y(t-1), h_(t-1))
p(y_t | y1, …, y(t-1), I) = softmax(W_p h_t)

The model is trained end-to-end by minimizing the negative log-likelihood of the correct caption:

L = – ∑_t log p(y_t | y1, …, y(t-1), I)

While this basic encoder-decoder architecture can generate plausible captions, the bottleneck of compressing the entire image into a single vector limits its ability to model fine-grained details. The model may miss out on important visual cues present in certain regions of the image. This is especially problematic for long captions that need to mention multiple objects or describe complex scenes.

Attention Mechanism

Attention provides a way to overcome the limitations of the basic encoder-decoder by allowing the model to dynamically focus on different parts of the image as it generates each word in the caption. Instead of relying on a single vector, attention computes a weighted combination of all the vectors in the feature map, with the weights indicating the relevance of each region to the current time step.

The attention mechanism can be seen as a way for the decoder RNN to "query" the relevant information from the image feature map at each step. This allows for a more fine-grained and adaptive use of visual information, leading to better captions that capture the details of the image.

There are two main variants of attention commonly used in image captioning:

1. Bahdanau attention (additive)
2. Luong attention (dot-product)

Let‘s go through the technical details of each.

Bahdanau Attention

In Bahdanau attention (also known as additive attention), the alignment scores between the RNN hidden state and the image feature vectors are computed using a multi-layer perceptron (MLP).

At each time step t, let h_(t-1) denote the previous hidden state of the decoder RNN, and V = {v_1, …, v_N} denote the set of N feature vectors extracted from the image.

The alignment scores are computed as:

e_ti = w^T tanh(W_a[vi; h(t-1)] + b_a)

where w, W_a and b_a are learnable parameters of the MLP. The alignment scores are then normalized using a softmax to obtain the attention weights α_ti:

α_ti = softmax(e_ti) = exp(e_ti) / ∑_j exp(e_tj)

The context vector c_t is computed as a weighted sum of the image features:

c_t = ∑_i α_ti v_i

Finally, the context vector is concatenated with the RNN hidden state and fed through another MLP to obtain the output probability distribution:

p(y_t | y1, …, y(t-1), I) = softmax(W_p[c_t; h_t] + b_p)

The attention weights α_ti can be visualized as a heatmap over the image, indicating which regions the model is focusing on at each time step.

Luong Attention

Luong attention (also known as dot-product attention) is a simplified variant that computes the alignment scores using a dot product between the RNN hidden state and the image features.

The alignment scores are computed as:

e_ti = h_t^T v_i

The rest of the attention mechanism remains the same as in Bahdanau attention, with the scores being normalized by a softmax to obtain α_ti, and the context vector c_t computed as a weighted sum.

Luong attention is more computationally efficient than Bahdanau attention since it doesn‘t require an additional MLP. However, Bahdanau attention is more flexible and can potentially model more complex relationships between the RNN state and image features.

Implementing Attention in Python/TensorFlow

Now that we understand the technical details of attention, let‘s see how to implement it in Python using TensorFlow and Keras. We‘ll go through the key components of an attention-based image captioning model.

Extracting Image Features

First, we use a pre-trained CNN to extract a feature map from the input image. Popular choices include VGG16, Inception, ResNet, etc. We extract the features from the penultimate convolutional layer before the final pooling.

image_model = tf.keras.applications.InceptionV3(include_top=False, weights=‘imagenet‘)
new_input = image_model.input
hidden_layer = image_model.layers[-1].output
image_features_extract_model = tf.keras.Model(new_input, hidden_layer)

This gives us a feature map of shape (8, 8, 2048) for an input image of size (299, 299, 3). We then flatten the feature map into a set of N = 64 feature vectors V = {v_1, …, v_N}, each of dimension D = 2048.

Attention Layers

Next, we define the attention layers that compute the alignment scores, attention weights, and context vector at each time step. We can use the built-in tf.keras.layers.Attention layer for this.

class Attention(tf.keras.Model):
  def __init__(self, units):
    super(Attention, self).__init__()
    self.W1 = Dense(units) 
    self.W2 = Dense(units) 
    self.V = Dense(1)

  def call(self, features, hidden):
    # features shape: (batch_size, 64, 2048)  
    # hidden shape: (batch_size, hidden_size)

    # Expand hidden to shape: (batch_size, 1, hidden_size)  
    hidden_exp = tf.expand_dims(hidden, 1)  

    # Compute alignment scores  
    score = tanh(self.W1(features) + self.W2(hidden_exp))
    score = self.V(score) 

    # Normalize with softmax
    attention_weights = softmax(score, axis=1) 

    # Compute context vector 
    context_vector = attention_weights * features
    context_vector = tf.reduce_sum(context_vector, axis=1)

    return context_vector, attention_weights

This implements Bahdanau attention. For Luong attention, we can simply replace the MLP with a dot product:

score = tf.matmul(hidden_exp, features, transpose_b=True) 

Decoder RNN

Finally, we define the decoder RNN that takes the context vector and generates the caption one word at a time. We use an LSTM with an embedding layer to encode the input words. The initial hidden state of the LSTM is set to the mean-pooled image features.

class Decoder(tf.keras.Model):
  def __init__(self, embed_dim, units, vocab_size):
    super(Decoder, self).__init__()
    self.units = units
    self.embedding = Embedding(vocab_size, embed_dim) 
    self.lstm = LSTM(self.units, return_state=True) 
    self.fc1 = Dense(self.units)
    self.fc2 = Dense(vocab_size)  
    self.attention = Attention(self.units)

  def call(self, x, features, hidden):
    # x shape: (batch_size, 1)
    # features shape: (batch_size, 64, 2048)
    # hidden shape: (batch_size, hidden_size)

    x = self.embedding(x) # shape: (batch_size, 1, embed_dim)
    context_vector, attention_weights = self.attention(features, hidden)
    lstm_input = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)
    output, state, _ = self.lstm(lstm_input) 
    output = self.fc1(output)
    output = tf.reshape(output, (-1, output.shape[2])) 
    x = self.fc2(output)
    return x, state, attention_weights

  def reset_state(self, batch_size):
    return tf.zeros((batch_size, self.units))

The decoder RNN is trained using teacher forcing, where the ground-truth caption words are fed as input during training. For inference, we generate the caption by sampling the word with the highest probability at each time step and feeding it back as input for the next time step.

Latest Developments and Future Directions

Attention mechanisms have revolutionized image captioning and have become an indispensable component of state-of-the-art models. Recent advancements include:

• Transformer-based models such as MeshedMemory Transformer and Visual Transformer, which use self-attention to capture long-range dependencies between image regions and text tokens.

• Unsupervised pre-training of vision-language models on large-scale image-text datasets, which learn rich visual-semantic representations that can be fine-tuned for image captioning.

• Novel attention variants such as Adaptive Attention that learns an additional gate to decide whether to attend to the image or to the visual sentinel at each time step.

There are still many exciting research directions to further improve attention-based image captioning, such as:

• Incorporating object detection and attributes to generate more fine-grained and grounded captions
• Designing better attention mechanisms to capture spatial and semantic relationships between image regions
• Exploring multimodal attention to integrate information from multiple sources (e.g. image, video, audio, knowledge bases)
• Generating diverse and controllable captions with specific styles, sentiments, or language complexity

Conclusion

Attention mechanisms have become a core component of modern image captioning models, enabling them to generate more accurate, detailed, and grounded captions by selectively focusing on relevant image regions. In this post, we dove into the technical details of attention, including the two popular variants: Bahdanau and Luong attention. We also walked through a Python implementation using TensorFlow and Keras.

Attention is a powerful and versatile concept that has been successfully applied to various vision-language tasks beyond image captioning, such as visual question answering, visual dialog, and vision-language navigation. As research in this area continues to advance, we can expect to see more breakthroughs in the coming years that will enable machines to understand and communicate about visual data in more intelligent and human-like ways.

I hope this post gave you a comprehensive overview of attention mechanisms for image captioning. Feel free to check out the references below for more details and state-of-the-art models. Happy reading and coding!

References

• Show, Attend and Tell: Neural Image Caption Generation with Visual Attention (Xu et al., 2015)
• Knowing When to Look: Adaptive Attention via a Visual Sentinel for Image Captioning (Lu et al., 2017)
• Bottom-Up and Top-Down Attention for Image Captioning and Visual Question Answering (Anderson et al., 2018)
• M2 Transformer: Multi-modal Multi-scale Transformers for Image Captioning (Zheng et al., 2022)