Transformer Architecture

# Conceptual Comparison: RNN vs Transformer
"""
SENTENCE: "The cat sat on the mat"
RNN/LSTM Processing (Sequential):
Step 1: Process "The" → hidden state h1
Step 2: Process "cat" with h1 → hidden state h2
Step 3: Process "sat" with h2 → hidden state h3
Step 4: Process "on" with h3 → hidden state h4
Step 5: Process "the" with h4 → hidden state h5
Step 6: Process "mat" with h5 → hidden state h6
❌ Problems:
- Must wait for each step (no parallelization)
- By step 6, information from "The" is diluted
- Slow for long sequences
- Training time: O(n) sequential steps
---
TRANSFORMER Processing (Parallel):
ALL AT ONCE:
"The" attends to ["The", "cat", "sat", "on", "the", "mat"]
"cat" attends to ["The", "cat", "sat", "on", "the", "mat"]
"sat" attends to ["The", "cat", "sat", "on", "the", "mat"]
"on"  attends to ["The", "cat", "sat", "on", "the", "mat"]
"the" attends to ["The", "cat", "sat", "on", "the", "mat"]
"mat" attends to ["The", "cat", "sat", "on", "the", "mat"]
✅ Advantages:
- All words processed simultaneously (full parallelization)
- Each word directly attends to all others (no information loss)
- Fast for any sequence length
- Training time: O(1) parallel (but O(n²) complexity per step)
"""
# Simple Transformer usage with Hugging Face
from transformers import BertTokenizer, BertModel
import torch
# Load pre-trained BERT (a Transformer encoder)
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
# Example sentence
text = "The bank by the river is beautiful."
inputs = tokenizer(text, return_tensors='pt')
# Process with Transformer (all tokens processed in parallel!)
with torch.no_grad():
    outputs = model(**inputs)
# Get contextualized embeddings for each word
embeddings = outputs.last_hidden_state
print(f"Input shape: {inputs['input_ids'].shape}")
print(f"Output embeddings shape: {embeddings.shape}")
print(f"\nEach word has a 768-dimensional embedding based on FULL context!")
# The word "bank" now has different embeddings in:
# "The bank by the river" (shoreline) vs "I went to the bank" (financial)
# because Transformer looks at ALL words via self-attention!

# Implementing Self-Attention from Scratch
import torch
import torch.nn.functional as F
import math
def scaled_dot_product_attention(Q, K, V):
    """
    Self-Attention implementation
    Args:
        Q: Query matrix [batch, seq_len, d_k]
        K: Key matrix [batch, seq_len, d_k]
        V: Value matrix [batch, seq_len, d_v]
    Returns:
        output: Context vectors [batch, seq_len, d_v]
        attention_weights: Attention scores [batch, seq_len, seq_len]
    """
    # Step 1: Calculate attention scores (Q · K^T)
    # How much should each word attend to every other word?
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1))  # [batch, seq_len, seq_len]
    print(f"Attention scores shape: {scores.shape}")
    print(f"Scores[0,0,:] (how word 0 attends to all words): {scores[0,0,:].tolist()}\n")
    # Step 2: Scale by √d_k (prevents large values)
    scores = scores / math.sqrt(d_k)
    # Step 3: Apply softmax (convert to probabilities)
    attention_weights = F.softmax(scores, dim=-1)
    print(f"Attention weights (after softmax):")
    print(f"Weights[0,0,:] (probabilities sum to 1): {attention_weights[0,0,:].tolist()}")
    print(f"Sum: {attention_weights[0,0,:].sum().item()}\n")
    # Step 4: Weighted sum of Values
    output = torch.matmul(attention_weights, V)
    return output, attention_weights
# Example usage
batch_size = 1
seq_len = 5  # "The cat sat on mat"
d_model = 64  # embedding dimension
# Simulate input embeddings
X = torch.randn(batch_size, seq_len, d_model)
# Create Q, K, V through linear transformations (learned during training)
W_Q = torch.randn(d_model, d_model)
W_K = torch.randn(d_model, d_model)
W_V = torch.randn(d_model, d_model)
Q = torch.matmul(X, W_Q)  # Query: what am I looking for?
K = torch.matmul(X, W_K)  # Key: what do I offer?
V = torch.matmul(X, W_V)  # Value: what info do I pass?
print("INPUT:")
print(f"Sequence length: {seq_len} words")
print(f"Embedding dimension: {d_model}\n")
# Apply self-attention
output, attention_weights = scaled_dot_product_attention(Q, K, V)
print("OUTPUT:")
print(f"Contextualized embeddings shape: {output.shape}")
print("\n✅ Each word now has context from ALL other words!")
# Visualize attention pattern
print("\nATTENTION PATTERN (who attends to whom):")
print("     Word0  Word1  Word2  Word3  Word4")
for i in range(seq_len):
    print(f"Word{i}: ", end="")
    for j in range(seq_len):
        weight = attention_weights[0, i, j].item()
        print(f"{weight:.2f}  ", end="")
    print()
# Output shows attention weights:
# High weight = strong attention (important relationship)
# Low weight = weak attention (less relevant)
# Each row sums to 1.0