NLP and transformers / Transformer architecture
Attention Mechanism in Python
Contents
- Introduction to Attention Mechanism
- Traditional Sequence Models
- Attention Mechanism Overview
- Attention Mechanism in Machine Translation
- Attention Mechanism in Image Captioning
- Attention Mechanism in Transformers
- Attention Mechanism in Sequence-to-Sequence Models
- Attention Mechanism in Transformer Language Models
- Attention Mechanism in Vision Transformers
- Attention Mechanism in Sequence Labeling Tasks
- Attention Mechanism in Graph Neural Networks
- Attention Mechanism in Reinforcement Learning
- Attention Mechanism in Multimodal Learning
Introduction to Attention Mechanism
The Attention Mechanism is a powerful technique used in various deep learning models, particularly in natural language processing (NLP) and computer vision tasks. It allows the model to focus on the most relevant parts of the input data, enabling more accurate and contextual representations. The Attention Mechanism has proven to be a game-changer in areas like machine translation, text summarization, and image captioning.
Traditional Sequence Models
Traditional sequence models, such as Recurrent Neural Networks (RNNs) and Long Short-Term Memory (LSTM) networks, process input data sequentially. While these models can capture long-range dependencies, they often struggle with very long sequences or prioritizing relevant information. The Attention Mechanism addresses these limitations by allowing the model to selectively focus on the most important parts of the input.
import torch
import torch.nn as nn
class SimpleRNN(nn.Module):
def __init__(self, input_size, hidden_size, num_layers):
super(SimpleRNN, self).__init__()
self.rnn = nn.RNN(input_size, hidden_size, num_layers)
def forward(self, input_seq, hidden_state):
output, hidden_state = self.rnn(input_seq, hidden_state)
return output, hidden_state
Attention Mechanism Overview
The Attention Mechanism is a technique that allows the model to learn where to focus its attention within the input data. It computes a set of attention weights, which represent the importance or relevance of each part of the input concerning the current context. These weights are then used to compute a weighted sum of the input representations, capturing the most relevant information.
import torch
import torch.nn as nn
class AttentionLayer(nn.Module):
def __init__(self, input_dim, output_dim):
super(AttentionLayer, self).__init__()
self.query = nn.Linear(input_dim, output_dim)
self.key = nn.Linear(input_dim, output_dim)
self.value = nn.Linear(input_dim, output_dim)
def forward(self, query, key, value):
scores = torch.matmul(self.query(query), self.key(key).transpose(-2, -1))
attention_weights = nn.functional.softmax(scores, dim=-1)
output = torch.matmul(attention_weights, self.value(value))
return output
Attention Mechanism in Machine Translation
One of the most successful applications of the Attention Mechanism is in machine translation tasks. Traditional sequence-to-sequence models, such as RNNs, often struggle to capture long-range dependencies and handle large vocabularies. The Attention Mechanism allows the model to focus on relevant parts of the source sentence when generating the target translation, improving translation quality and handling longer sequences more effectively.
import torch
import torch.nn as nn
class Encoder(nn.Module):
def __init__(self, input_dim, embed_dim, hidden_dim, num_layers):
super(Encoder, self).__init__()
self.embed = nn.Embedding(input_dim, embed_dim)
self.rnn = nn.LSTM(embed_dim, hidden_dim, num_layers)
def forward(self, input_seq):
embedded = self.embed(input_seq)
outputs, (hidden, cell) = self.rnn(embedded)
return outputs, hidden, cell
class Decoder(nn.Module):
def __init__(self, output_dim, embed_dim, hidden_dim, num_layers, attention):
super(Decoder, self).__init__()
self.embed = nn.Embedding(output_dim, embed_dim)
self.rnn = nn.LSTMCell(embed_dim + hidden_dim, hidden_dim)
self.attention = attention
self.out = nn.Linear(hidden_dim, output_dim)
def forward(self, input, hidden, cell, encoder_outputs):
embedded = self.embed(input)
attended = self.attention(hidden, encoder_outputs)
rnn_input = torch.cat([embedded, attended], dim=2)
hidden, cell = self.rnn(rnn_input, (hidden, cell))
output = self.out(hidden)
return output, hidden, cell
Attention Mechanism in Image Captioning
The Attention Mechanism has also been successfully applied to image captioning tasks, where the model generates a natural language description of an input image. In this context, the Attention Mechanism allows the model to focus on different regions of the image when generating each word in the caption, capturing the most relevant visual information for that particular part of the description.
import torch
import torch.nn as nn
class Encoder(nn.Module):
def __init__(self, cnn):
super(Encoder, self).__init__()
self.cnn = cnn
def forward(self, images):
features = self.cnn(images)
return features
class Decoder(nn.Module):
def __init__(self, embed_dim, hidden_dim, vocab_size, attention):
super(Decoder, self).__init__()
self.embed = nn.Embedding(vocab_size, embed_dim)
self.rnn = nn.LSTMCell(embed_dim + hidden_dim, hidden_dim)
self.attention = attention
self.out = nn.Linear(hidden_dim, vocab_size)
def forward(self, input, hidden, cell, encoder_features):
embedded = self.embed(input)
attended = self.attention(hidden, encoder_features)
rnn_input = torch.cat([embedded, attended], dim=1)
hidden, cell = self.rnn(rnn_input, (hidden, cell))
output = self.out(hidden)
return output, hidden, cell
Attention Mechanism in Transformers
The Transformer architecture, introduced in the seminal paper "Attention Is All You Need," relies entirely on the Attention Mechanism and has revolutionized many NLP tasks. Transformers use self-attention mechanisms to capture dependencies within the input sequence and enable parallel processing, leading to improved performance and faster training times compared to recurrent models.
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, embed_dim, num_heads):
super(MultiHeadAttention, self).__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.query = nn.Linear(embed_dim, embed_dim)
self.key = nn.Linear(embed_dim, embed_dim)
self.value = nn.Linear(embed_dim, embed_dim)
self.output = nn.Linear(embed_dim, embed_dim)
def forward(self, query, key, value):
batch_size = query.size(0)
query = self.query(query).view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
key = self.key(key).view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
value = self.value(value).view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dim)
attention_weights = nn.functional.softmax(scores, dim=-1)
attended_values = torch.matmul(attention_weights, value)
attended_values = attended_values.permute(0, 2, 1, 3).contiguous().view(batch_size, -1, self.embed_dim)
output = self.output(attended_values)
return output
Attention Mechanism in Sequence-to-Sequence Models
Sequence-to-sequence models, such as those used in machine translation and text summarization, have greatly benefited from the Attention Mechanism. In these models, the encoder processes the input sequence and produces a set of hidden states, while the decoder generates the output sequence one token at a time. The Attention Mechanism allows the decoder to selectively focus on relevant parts of the input sequence when generating each output token, improving the quality and coherence of the generated output.
import torch
import torch.nn as nn
class Encoder(nn.Module):
def __init__(self, input_dim, embed_dim, hidden_dim, num_layers):
super(Encoder, self).__init__()
self.embed = nn.Embedding(input_dim, embed_dim)
self.rnn = nn.LSTM(embed_dim, hidden_dim, num_layers, bidirectional=True)
def forward(self, input_seq, input_lengths):
embedded = self.embed(input_seq)
packed = nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)
outputs, (hidden, cell) = self.rnn(packed)
outputs, _ = nn.utils.rnn.pad_packed_sequence(outputs)
return outputs, hidden, cell
class Decoder(nn.Module):
def __init__(self, output_dim, embed_dim, hidden_dim, num_layers, attention):
super(Decoder, self).__init__()
self.embed = nn.Embedding(output_dim, embed_dim)
self.rnn = nn.LSTMCell(embed_dim + hidden_dim, hidden_dim)
self.attention = attention
self.out = nn.Linear(hidden_dim, output_dim)
def forward(self, input, hidden, cell, encoder_outputs):
embedded = self.embed(input)
attended = self.attention(hidden, encoder_outputs)
rnn_input = torch.cat([embedded, attended], dim=1)
hidden, cell = self.rnn(rnn_input, (hidden, cell))
output = self.out(hidden)
return output, hidden, cell
Attention Mechanism in Transformer Language Models
Transformer-based language models, such as BERT, GPT, and XLNet, have achieved remarkable success in various NLP tasks. These models heavily rely on the self-attention mechanism, which allows them to capture long-range dependencies and contextual information within the input text. The self-attention mechanism computes attention weights for each token in the input sequence concerning all other tokens, enabling the model to effectively model complex relationships and dependencies.
import torch
import torch.nn as nn
class TransformerEncoderBlock(nn.Module):
def __init__(self, embed_dim, num_heads, ffn_dim, dropout):
super(TransformerEncoderBlock, self).__init__()
self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
self.ffn = nn.Sequential(
nn.Linear(embed_dim, ffn_dim),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(ffn_dim, embed_dim),
nn.Dropout(dropout)
)
def forward(self, x, mask=None):
attended, _ = self.attention(x, x, x, key_padding_mask=mask)
x = x + attended
x = x + self.ffn(x)
return x
Attention Mechanism in Vision Transformers
The Attention Mechanism has also been successfully applied to computer vision tasks, particularly with the introduction of Vision Transformers (ViTs). ViTs treat images as sequences of patches and use self-attention mechanisms to capture long-range dependencies and spatial relationships within the image. This approach has achieved state-of-the-art performance on various computer vision tasks, such as image classification, object detection, and semantic segmentation.
import torch
import torch.nn as nn
class PatchEmbedding(nn.Module):
def __init__(self, image_size, patch_size, embed_dim):
super(PatchEmbedding, self).__init__()
self.image_size = image_size
self.patch_size = patch_size
self.num_patches = (image_size // patch_size) ** 2
self.proj = nn.Conv2d(3, embed_dim, kernel_size=patch_size, stride=patch_size)
def forward(self, images):
x = self.proj(images)
x = x.flatten(2).transpose(1, 2)
return x
class VisionTransformer(nn.Module):
def __init__(self, embed_dim, num_heads, ffn_dim, num_layers, num_classes):
super(VisionTransformer, self).__init__()
self.patch_embedding = PatchEmbedding(image_size=224, patch_size=16, embed_dim=embed_dim)
self.transformer_blocks = nn.ModuleList([
TransformerEncoderBlock(embed_dim, num_heads, ffn_dim, dropout=0.1)
for _ in range(num_layers)
])
self.classifier = nn.Linear(embed_dim, num_classes)
def forward(self, images):
x = self.patch_embedding(images)
for block in self.transformer_blocks:
x = block(x)
x = x.mean(dim=1)
x = self.classifier(x)
return x
Attention Mechanism in Sequence Labeling Tasks
The Attention Mechanism has also been applied to sequence labeling tasks, such as named entity recognition (NER) and part-of-speech (POS) tagging. In these tasks, the model needs to assign a label to each token in the input sequence. The Attention Mechanism allows the model to focus on relevant parts of the input sequence when making predictions for each token, improving accuracy and capturing long-range dependencies.
import torch
import torch.nn as nn
class BiLSTMAttentionNER(nn.Module):
def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers, num_labels):
super(BiLSTMAttentionNER, self).__init__()
self.embed = nn.Embedding(vocab_size, embed_dim)
self.bilstm = nn.LSTM(embed_dim, hidden_dim, num_layers, bidirectional=True)
self.attention = nn.Linear(hidden_dim * 2, 1)
self.classifier = nn.Linear(hidden_dim * 2, num_labels)
def forward(self, input_seq):
embedded = self.embed(input_seq)
outputs, _ = self.bilstm(embedded)
attention_scores = self.attention(outputs).squeeze(-1)
attention_weights = nn.functional.softmax(attention_scores, dim=1).unsqueeze(-1)
attended_outputs = (outputs * attention_weights).sum(dim=1)
logits = self.classifier(attended_outputs)
return logits
Attention Mechanism in Graph Neural Networks
Graph Neural Networks (GNNs) have emerged as powerful models for processing graph-structured data. The Attention Mechanism has been incorporated into GNNs to enable them to selectively focus on the most relevant nodes and edges when making predictions. This approach, known as Graph Attention Networks (GATs), has shown improved performance on various graph-related tasks, such as node classification, link prediction, and graph classification.
import torch
import torch.nn as nn
import torch.nn.functional as F
class GraphAttentionLayer(nn.Module):
def __init__(self, in_features, out_features, dropout, alpha, concat=True):
super(GraphAttentionLayer, self).__init__()
self.dropout = dropout
self.in_features = in_features
self.out_features = out_features
self.alpha = alpha
self.concat = concat
self.W = nn.Parameter(torch.empty(size=(in_features, out_features)))
nn.init.xavier_uniform_(self.W.data, gain=1.414)
self.a = nn.Parameter(torch.empty(size=(2*out_features, 1)))
nn.init.xavier_uniform_(self.a.data, gain=1.414)
self.leakyrelu = nn.LeakyReLU(self.alpha)
def forward(self, h, adj):
Wh = torch.mm(h, self.W) # h.shape: (N, in_features), Wh.shape: (N, out_features)
e = self._prepare_attentional_mechanism_input(Wh)
zero_vec = -9e15*torch.ones_like(e)
attention = torch.where(adj > 0, e, zero_vec)
attention = F.softmax(attention, dim=1)
attention = F.dropout(attention, self.dropout, training=self.training)
h_prime = torch.matmul(attention, Wh)
if self.concat:
return F.elu(h_prime)
else:
return h_prime
def _prepare_attentional_mechanism_input(self, Wh):
Wh1 = torch.matmul(Wh, self.a[:self.out_features, :])
Wh2 = torch.matmul(Wh, self.a[self.out_features:, :])
e = Wh1 + Wh2.T
return self.leakyrelu(e)
class GAT(nn.Module):
def __init__(self, num_features, hidden_size, num_classes, dropout, alpha, nheads):
super(GAT, self).__init__()
self.dropout = dropout
self.attentions = [GraphAttentionLayer(num_features, hidden_size, dropout=dropout, alpha=alpha, concat=True) for _ in range(nheads)]
for i, attention in enumerate(self.attentions):
self.add_module('attention_{}'.format(i), attention)
self.out_attention = GraphAttentionLayer(hidden_size * nheads, num_classes, dropout=dropout, alpha=alpha, concat=False)
def forward(self, x, adj):
x = F.dropout(x, self.dropout, training=self.training)
x = torch.cat([att(x, adj) for att in self.attentions], dim=1)
x = F.dropout(x, self.dropout, training=self.training)
x = F.elu(self.out_attention(x, adj))
return F.log_softmax(x, dim=1)
In this implementation, the GraphAttentionLayer computes the attention coefficients using a self-attention mechanism and applies them to the input features. The GAT module combines multiple GraphAttentionLayers in parallel and performs the final node classification or prediction task. This code provides a complete implementation of Graph Attention Networks for processing graph-structured data.
Attention Mechanism in Reinforcement Learning
The Attention Mechanism has also found applications in Reinforcement Learning (RL), particularly in tasks involving complex and high-dimensional state spaces. By incorporating attention mechanisms, RL agents can selectively focus on the most relevant parts of the state representation, improving their ability to learn optimal policies and make better decisions. This approach has shown promising results in areas such as robotics, game playing, and control systems.
import torch
import torch.nn as nn
class AttentionQNetwork(nn.Module):
def __init__(self, state_dim, action_dim, hidden_dim):
super(AttentionQNetwork, self).__init__()
self.state_encoder = nn.Sequential(
nn.Linear(state_dim, hidden_dim),
nn.ReLU()
)
self.attention = nn.MultiheadAttention(hidden_dim, num_heads=4)
self.q_value = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, action_dim)
)
def forward(self, state):
encoded_state = self.state_encoder(state)
attended_state, _ = self.attention(encoded_state, encoded_state, encoded_state)
q_values = self.q_value(attended_state)
return q_values
Attention Mechanism in Multimodal Learning
Multimodal learning involves integrating and processing information from multiple modalities, such as text, images, and audio. The Attention Mechanism has played a crucial role in this area by enabling models to selectively attend to relevant information across different modalities. This approach has been successfully applied to tasks like visual question answering, image captioning, and multimodal sentiment analysis, where the model needs to effectively combine and reason over different types of data.
import torch
import torch.nn as nn
class MultimodalAttentionModel(nn.Module):
def __init__(self, text_dim, image_dim, hidden_dim, output_dim):
super(MultimodalAttentionModel, self).__init__()
self.text_encoder = nn.LSTM(text_dim, hidden_dim, batch_first=True)
self.image_encoder = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2, 2),
nn.Flatten(),
nn.Linear(64 * 112 * 112, image_dim)
)
self.attention = nn.MultiheadAttention(hidden_dim, num_heads=4)
self.classifier = nn.Linear(hidden_dim, output_dim)
def forward(self, text, images):
text_features, _ = self.text_encoder(text)
image_features = self.image_encoder(images)
combined_features = torch.cat([text_features, image_features.unsqueeze(1)], dim=1)
attended_features, _ = self.attention(combined_features, combined_features, combined_features)
output = self.classifier(attended_features[:, -1, :])
return output
Slide 14 (Additional Resources): Additional Resources on Attention Mechanism
If you want to dive deeper into the Attention Mechanism and related topics, here are some recommended resources from ArXiv.org:
- "Attention Is All You Need" (https://arxiv.org/abs/1706.03762) - The seminal paper that introduced the Transformer architecture and popularized the use of self-attention mechanisms in deep learning.
- "Self-Attention with Relative Position Representations" (https://arxiv.org/abs/1803.02155) - A paper that proposes a variation of self-attention that incorporates relative position information, improving performance on various NLP tasks.
- "Graph Attention Networks" (https://arxiv.org/abs/1710.10903) - The paper that introduces Graph Attention Networks (GATs), which apply attention mechanisms to graph-structured data.
- "Efficient Attention Mechanisms for Vision Transformers" (https://arxiv.org/abs/2112.13492) - A paper that explores efficient attention mechanisms for Vision Transformers, addressing the computational complexity of self-attention in computer vision tasks.
These resources provide in-depth theoretical and practical insights into the Attention Mechanism and its applications across various domains.