Attention Mechanism

import torch
import torch.nn.functional as F
import math
def scaled_dot_product_attention(query, key, value, mask=None):
    """
    Scaled Dot-Product Attention
    Args:
        query: (batch_size, seq_len, d_k)
        key: (batch_size, seq_len, d_k)
        value: (batch_size, seq_len, d_v)
        mask: Optional mask (batch_size, seq_len, seq_len)
    Returns:
        output: (batch_size, seq_len, d_v)
        attention_weights: (batch_size, seq_len, seq_len)
    """
    # Get dimension of key
    d_k = key.size(-1)
    # Step 1: Compute attention scores (Q @ K^T)
    # (batch, seq_len, d_k) @ (batch, d_k, seq_len) -> (batch, seq_len, seq_len)
    scores = torch.matmul(query, key.transpose(-2, -1))
    # Step 2: Scale by sqrt(d_k) to prevent saturation
    scores = scores / math.sqrt(d_k)
    # Step 3: Apply mask if provided (for padding or causal attention)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    # Step 4: Apply softmax to get attention weights
    attention_weights = F.softmax(scores, dim=-1)
    # Step 5: Apply attention weights to values
    output = torch.matmul(attention_weights, value)
    return output, attention_weights
# Example usage
batch_size, seq_len, d_model = 2, 5, 64
Q = torch.randn(batch_size, seq_len, d_model)
K = torch.randn(batch_size, seq_len, d_model)
V = torch.randn(batch_size, seq_len, d_model)
output, weights = scaled_dot_product_attention(Q, K, V)
print(f"Output shape: {output.shape}")  # (2, 5, 64)
print(f"Attention weights shape: {weights.shape}")  # (2, 5, 5)
print(f"Weights sum to 1: {weights.sum(dim=-1)}")  # All ones

import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        """
        Multi-Head Attention module
        Args:
            d_model: Total dimension of model (e.g., 512)
            num_heads: Number of attention heads (e.g., 8)
        """
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # Dimension per head
        # Linear projections for Q, K, V
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        # Output projection
        self.W_o = nn.Linear(d_model, d_model)
    def split_heads(self, x):
        """Split into multiple heads: (batch, seq_len, d_model) -> (batch, num_heads, seq_len, d_k)"""
        batch_size, seq_len, d_model = x.size()
        return x.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        # 1. Linear projections
        Q = self.W_q(query)  # (batch, seq_len, d_model)
        K = self.W_k(key)
        V = self.W_v(value)
        # 2. Split into multiple heads
        Q = self.split_heads(Q)  # (batch, num_heads, seq_len, d_k)
        K = self.split_heads(K)
        V = self.split_heads(V)
        # 3. Apply scaled dot-product attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attention_weights = F.softmax(scores, dim=-1)
        attention_output = torch.matmul(attention_weights, V)
        # 4. Concatenate heads
        attention_output = attention_output.transpose(1, 2).contiguous()
        attention_output = attention_output.view(batch_size, -1, self.d_model)
        # 5. Final linear projection
        output = self.W_o(attention_output)
        return output, attention_weights
# Example usage
d_model, num_heads = 512, 8
mha = MultiHeadAttention(d_model, num_heads)
x = torch.randn(2, 10, d_model)  # (batch, seq_len, d_model)
output, weights = mha(x, x, x)  # Self-attention
print(f"Output shape: {output.shape}")  # (2, 10, 512)
print(f"Attention weights shape: {weights.shape}")  # (2, 8, 10, 10)

import torch
import torch.nn as nn
class SelfAttentionClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_heads, num_classes):
        super().__init__()
        # Embedding layer
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        # Multi-head self-attention
        self.attention = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        # Classification head
        self.fc = nn.Linear(embed_dim, num_classes)
    def forward(self, x):
        # x: (batch, seq_len) - token indices
        # Embed tokens
        embedded = self.embedding(x)  # (batch, seq_len, embed_dim)
        # Apply self-attention
        # Each token attends to all other tokens
        attn_output, attn_weights = self.attention(embedded, embedded, embedded)
        # Pool by taking mean across sequence
        pooled = attn_output.mean(dim=1)  # (batch, embed_dim)
        # Classify
        logits = self.fc(pooled)  # (batch, num_classes)
        return logits, attn_weights
# Example: Sentiment classification
vocab_size = 10000
embed_dim = 128
num_heads = 4
num_classes = 2  # Positive/Negative
model = SelfAttentionClassifier(vocab_size, embed_dim, num_heads, num_classes)
# Sample batch of tokenized sentences
sentences = torch.randint(0, vocab_size, (8, 20))  # 8 sentences, 20 tokens each
logits, attention_weights = model(sentences)
predictions = torch.argmax(logits, dim=-1)
print(f"Predictions: {predictions}")
print(f"Attention weights shape: {attention_weights.shape}")  # (8, 20, 20)
# Attention weights show which words attend to which
# For example, "not" might attend strongly to "good" in "not good"

import torch
import matplotlib.pyplot as plt
import seaborn as sns
def visualize_attention(attention_weights, tokens, layer=0, head=0):
    """
    Visualize attention weights as a heatmap
    Args:
        attention_weights: (num_layers, num_heads, seq_len, seq_len) or (num_heads, seq_len, seq_len)
        tokens: List of tokens/words
        layer: Which layer to visualize
        head: Which attention head to visualize
    """
    # Extract attention for specific layer and head
    if attention_weights.ndim == 4:
        attn = attention_weights[layer, head].cpu().numpy()
    else:
        attn = attention_weights[head].cpu().numpy()
    # Create heatmap
    plt.figure(figsize=(10, 8))
    sns.heatmap(attn, xticklabels=tokens, yticklabels=tokens,
                cmap='viridis', cbar=True, square=True)
    plt.xlabel('Key (What we attend to)')
    plt.ylabel('Query (Current position)')
    plt.title(f'Attention Weights - Layer {layer}, Head {head}')
    plt.tight_layout()
    plt.show()
# Example with BERT
from transformers import BertTokenizer, BertModel
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased', output_attentions=True)
sentence = "The animal didn't cross the street because it was too tired"
inputs = tokenizer(sentence, return_tensors='pt')
with torch.no_grad():
    outputs = model(**inputs)
# outputs.attentions is a tuple of (num_layers) tensors
# Each tensor is (batch, num_heads, seq_len, seq_len)
attention = outputs.attentions  # 12 layers for BERT-base
tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
print(f"Tokens: {tokens}")
# Visualize layer 11 (last layer), head 0
visualize_attention(attention[11][0], tokens, layer=0, head=0)
# Find what "it" attends to
it_index = tokens.index('it')
attn_scores = attention[11][0, 0, it_index, :].cpu().numpy()
print("\nWhat 'it' attends to:")
for token, score in zip(tokens, attn_scores):
    if score > 0.1:  # Only show significant attention
        print(f"{token}: {score:.3f}")
# Output might show "it" attends strongly to "animal" (resolving pronoun reference)

from transformers import pipeline, AutoTokenizer, AutoModel
import torch
# 1. Text Summarization with Attention-based Model (BART)
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
long_text = """
Attention mechanisms have revolutionized natural language processing.
Unlike traditional RNNs that process sequences sequentially, attention allows
models to focus on relevant parts of the input simultaneously. The Transformer
architecture, introduced in 2017, uses only attention mechanisms and has become
the foundation for models like BERT, GPT, and T5. These models achieve
state-of-the-art results on various NLP tasks including translation,
question answering, and text generation.
"""
summary = summarizer(long_text, max_length=50, min_length=25, do_sample=False)
print("Summary:", summary[0]['summary_text'])
# 2. Machine Translation with Cross-Attention (MarianMT)
translator = pipeline("translation_en_to_fr", model="Helsinki-NLP/opus-mt-en-fr")
translation = translator("Attention is all you need")
print("Translation:", translation[0]['translation_text'])
# 3. Question Answering with BERT (uses self-attention)
qa_pipeline = pipeline("question-answering", model="bert-large-uncased-whole-word-masking-finetuned-squad")
context = """
The Transformer architecture was introduced in the paper 'Attention Is All You Need'
by Vaswani et al. in 2017. It relies entirely on attention mechanisms, dispensing
with recurrence and convolutions entirely.
"""
question = "When was the Transformer introduced?"
answer = qa_pipeline(question=question, context=context)
print(f"Question: {question}")
print(f"Answer: {answer['answer']} (confidence: {answer['score']:.2f})")
# 4. Extract attention weights from custom model
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased", output_attentions=True)
text = "Attention mechanisms are powerful"
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
# Get attention from all layers
attentions = outputs.attentions  # Tuple of 12 tensors for BERT-base
print(f"Number of layers: {len(attentions)}")
print(f"Attention shape per layer: {attentions[0].shape}")  # (batch, heads, seq, seq)
# Average attention across all heads and layers
avg_attention = torch.stack(attentions).mean(dim=(0, 1))  # Average over layers and heads
print(f"Average attention shape: {avg_attention.shape}")