Show, Attend and Tell: Mastering Image Captioning with Visual Attention
Show, Attend and Tell: Mastering Image Captioning with Visual Attention
Imagine a computer looking at a photograph of a dog catching a frisbee. Instead of simply labeling the image as "dog" or "outdoor," the system generates a descriptive sentence: "A brown dog is leaping in the air to catch a yellow frisbee."
To achieve this, the model cannot simply "summarize" the image into a single vector; it must look at specific parts of the image while generating each specific word. This is the core intuition behind the seminal paper "Show, Attend and Tell," which introduced a dynamic "gaze" to image captioning.
In this post, we will break down the architecture, the mathematics of attention, and provide a production-ready PyTorch implementation.
🧠 The Core Intuition: From Static to Dynamic
Traditional image captioning models treated the problem as a simple translation: Image $\rightarrow$ Fixed Vector $\rightarrow$ Sentence. The problem? A single vector is a bottleneck; it cannot capture the spatial nuances of a complex scene.
Show, Attend and Tell re-images this as a translation task where the "source" is a grid of visual features. Instead of a static summary, the model uses an Encoder-Decoder framework with Attention:
- The Encoder (The Eyes): A CNN that doesn't just say "what" is in the image, but "where" things are, producing a grid of spatial feature vectors.
- The Decoder (The Voice): An LSTM that generates words one by one.
- The Attention (The Gaze): A mechanism that tells the Decoder: "To generate the next word, focus your attention on these specific pixels in the image grid."
🏗️ Architectural Deep Dive
The High-Level Workflow
The Mathematics of "Looking"
The model implements two types of attention: Soft (differentiable, weighted average) and Hard (stochastic, picking one location). Most modern implementations favor Soft Attention for its stability during training.
1. Calculating Attention Weights
The model uses the previous hidden state of the LSTM ($\mathbf{h}{t-1}$) to determine the importance of each spatial region $\mathbf{a}i$: $$\alpha_i = f{\text{att}}(\mathbf{h}{t-1})$$
2. Generating the Context Vector
For Soft Attention, the context vector $\hat{\mathbf{z}}_t$ is the weighted sum of all spatial annotation vectors: $$\hat{\mathbf{z}}t = \sum{i=1}^{L} \alpha_i \mathbf{a}_i$$
3. Predicting the Word
The final probability of the next word $y_t$ is conditioned on the current LSTM state, the previous word, and the visual context: $$p(y_t | y_{<t}, \text{image}) = \text{softmax}(\mathbf{L}o \mathbf{E} y{t-1} + \mathbf{L}_h \mathbf{h}_t + \mathbf{L}_z \hat{\mathbf{z}}_t)$$
💻 Implementation in PyTorch
Below is the complete implementation. We use a pre-trained ResNet50 as the encoder, stripping the final classification layers to preserve the $7 \times 7$ spatial grid.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
class Attention(nn.Module):
"""
Soft Attention: Computes a context vector as a weighted sum of image features.
"""
def __init__(self, encoder_dim, decoder_dim, attention_dim):
super(Attention, self).__init__()
self.encoder_attn = nn.Linear(encoder_dim, attention_dim)
self.decoder_attn = nn.Linear(decoder_dim, attention_dim)
self.full_attn = nn.Linear(attention_dim, 1)
def forward(self, encoder_out, decoder_hidden):
# Project encoder features and decoder hidden state into a shared attention space
enc_proj = self.encoder_attn(encoder_out)
dec_proj = self.decoder_attn(decoder_hidden).unsqueeze(1)
# Additive attention score: tanh(W_e*enc + W_d*dec)
attn_scores = self.full_attn(torch.tanh(enc_proj + dec_proj))
attn_weights = F.softmax(attn_scores, dim=1)
# Context vector = weighted sum of spatial features
context = torch.sum(attn_weights * encoder_out, dim=1)
return context, attn_weights
class Decoder(nn.Module):
"""
LSTM Decoder that generates captions conditioned on visual attention.
"""
def __init__(self, encoder_dim, decoder_dim, embed_dim, vocab_size):
super(Decoder, self).__init__()
self.decoder_dim = decoder_dim
self.embedding = nn.Embedding(vocab_size, embed_dim)
self.attention = Attention(encoder_dim, decoder_dim, 256)
self.lstm = nn.LSTMCell(embed_dim + encoder_dim, decoder_dim)
self.fc = nn.Linear(decoder_dim + encoder_dim, vocab_size)
def forward(self, encoder_out, captions):
batch_size = encoder_out.size(0)
seq_len = captions.size(1)
h = torch.zeros(batch_size, self.decoder_dim).to(encoder_out.device)
c = torch.zeros(batch_size, self.decoder_dim).to(encoder_out.device)
inputs = captions[:, 0] # Start token
outputs = []
for t in range(1, seq_len):
context, _ = self.attention(encoder_out, h)
embedded = self.embedding(inputs).squeeze(1)
# Combine visual context and word embedding
lstm_input = torch.cat([embedded, context], dim=1)
h, c = self.lstm(lstm_input, (h, c))
# Predict word based on LSTM state AND visual context
out = self.fc(torch.cat([h, context], dim=1))
outputs.append(out)
inputs = captions[:, t] # Teacher forcing
return torch.stack(outputs, dim=1)
class ShowAttendTell(nn.Module):
def __init__(self, vocab_size):
super(ShowAttendTell, self).__init__()
# Encoder: ResNet50 without the final pooling/FC layers
resnet = models.resnet50(pretrained=False)
self.encoder = nn.Sequential(*list(resnet.children())[:-2])
self.decoder = Decoder(encoder_dim=2048, decoder_dim=512, embed_dim=256, vocab_size=vocab_size)
def forward(self, images, captions):
features = self.encoder(images) # (batch, 2048, 7, 7)
# Flatten spatial grid: (batch, 49, 2048)
features = features.view(features.size(0), features.size(1), -1).permute(0, 2, 1)
return self.decoder(features, captions)🚀 Summary & Key Takeaways
The "Show, Attend and Tell" architecture shifted the paradigm of image captioning from global image representation to local spatial awareness.
| Feature | Traditional Encoder-Decoder | Show, Attend and Tell |
|---|---|---|
| Image Representation | Single Global Vector | Grid of Spatial Vectors |
| Visual Focus | Static (Whole image) | Dynamic (Attention-based) |
| Context | Fixed for the whole sentence | Changes for every word generated |
| Interpretability | Black box | High (via Attention Maps) |
By allowing the model to "gaze" at specific regions, we not only improve the accuracy of the generated captions but also gain a window into the model's "thought process," making AI more transparent and powerful.