Attention Mechanism: The Breakthrough Innovation

Sentence: "The cat that ate the mouse that lived in the barn ran away"

Problem: By the time model gets to "ran", information about "cat" 
has been compressed through many timesteps and starts to fade.

When predicting "ran":
- Attention can directly look at "cat" (high attention)
- Also look at "that", "ate", etc. (lower attention)
- Weights determine how much to focus on each word

No information loss!

Question: "Where is the Eiffel Tower?"
Context:  "The Eiffel Tower is located in Paris, France."

Attention weights:
  The      [0.05]  ▁
  Eiffel   [0.30]  ████████
  Tower    [0.25]  ███████
  is       [0.02]  ▁
  located  [0.05]  █
  in       [0.03]  ▁
  Paris    [0.25]  ███████
  France   [0.05]  █

Model focuses heavily on: "Eiffel", "Tower", "Paris"

"The cat sat"

# Simplified 4-dimensional embeddings
embeddings = {
    "The": [0.2, 0.1, 0.5, 0.3],
    "cat": [0.5, 0.8, 0.2, 0.1],
    "sat": [0.1, 0.3, 0.9, 0.4]
}
 
# Stack into matrix
X = [[0.2, 0.1, 0.5, 0.3],  # The
     [0.5, 0.8, 0.2, 0.1],  # cat
     [0.1, 0.3, 0.9, 0.4]]  # sat
# Shape: (3, 4) - 3 words, 4 dimensions

# Weight matrices (learned during training)
W_q = random_matrix(4, 4)  # Query weights
W_k = random_matrix(4, 4)  # Key weights
W_v = random_matrix(4, 4)  # Value weights
 
# Transform embeddings
Q = X @ W_q  # Query: What am I looking for?
K = X @ W_k  # Key: What do I contain?
V = X @ W_v  # Value: What do I output?
 
# All have shape (3, 4)

scores = Q @ K.T  # Matrix multiplication
# Shape: (3, 3)
 
# Example result:
scores = [[2.1, 1.5, 0.8],  # The attends to: The, cat, sat
          [1.5, 3.2, 1.9],  # cat attends to: The, cat, sat
          [0.8, 1.9, 2.7]]  # sat attends to: The, cat, sat

d_k = 4  # Dimension of keys
scaled_scores = scores / sqrt(d_k)

attention_weights = softmax(scaled_scores)
 
# Example result:
attention_weights = [
    [0.55, 0.32, 0.13],  # The: 55% to itself, 32% to cat, 13% to sat
    [0.15, 0.60, 0.25],  # cat: 15% to The, 60% to itself, 25% to sat
    [0.10, 0.30, 0.60]   # sat: 10% to The, 30% to cat, 60% to itself
]

output = attention_weights @ V
# Shape: (3, 4)
 
# For word "cat" (row 1):
output[1] = 0.15 * V[0] + 0.60 * V[1] + 0.25 * V[2]
#           ↑            ↑             ↑
#        from "The"  from "cat"   from "sat"

Attention(Q, K, V) = softmax(QK^T / √d_k) V

Where:
- Q = queries (what I'm looking for)
- K = keys (what I have)
- V = values (what I output)
- d_k = dimension of keys
- / √d_k = scaling factor
- softmax = converts to probabilities

# Database with (key, value) pairs
database = {
    "Paris": "Capital of France, home to Eiffel Tower",
    "London": "Capital of UK, home to Big Ben",
    "Tokyo": "Capital of Japan, largest city"
}
 
# Query
query = "Where is the Eiffel Tower?"
 
# Step 1: Match query to keys
scores = {
    "Paris": 0.85,   # High match!
    "London": 0.15,
    "Tokyo": 0.10
}
 
# Step 2: Softmax to get attention weights
attention = softmax(scores)  # [0.70, 0.17, 0.13]
 
# Step 3: Retrieve weighted combination of values
result = 0.70 * database["Paris"] + 
         0.17 * database["London"] + 
         0.13 * database["Tokyo"]
# Mostly Paris information!

# Input embeddings
X = word_embeddings  # Shape: (seq_len, d_model)
 
# Linear transformations (learned)
Q = X @ W_q  # "What to search for"
K = X @ W_k  # "How to identify relevant info"
V = X @ W_v  # "What info to pass forward"
 
# Attention computation
scores = Q @ K.T           # Similarity
weights = softmax(scores)  # Probability
output = weights @ V       # Weighted combination

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Args:
        Q: Queries (batch, seq_len, d_k)
        K: Keys (batch, seq_len, d_k)
        V: Values (batch, seq_len, d_v)
        mask: Optional mask (batch, seq_len, seq_len)
    
    Returns:
        output: (batch, seq_len, d_v)
        attention_weights: (batch, seq_len, seq_len)
    """
    d_k = Q.shape[-1]
    
    # Compute attention scores
    scores = Q @ K.transpose(-2, -1) / sqrt(d_k)
    
    # Apply mask (for padding or causality)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # Softmax to get attention weights
    attention_weights = softmax(scores, dim=-1)
    
    # Apply weights to values
    output = attention_weights @ V
    
    return output, attention_weights

# Sentence: "The cat <PAD> <PAD>"
mask = [[1, 1, 0, 0]]  # Only attend to real words

# When predicting word 2, can only see words 0, 1
mask = [[1, 0, 0, 0],
        [1, 1, 0, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 1]]
# Lower triangular matrix

"The cat sat on the mat"

Single head might learn: Subject-Verb relationships
- "cat" → "sat" (subject-verb)

But misses:
- Spatial relationships: "sat" → "on"
- Object relationships: "on" → "mat"

Head 1: Subject-Verb relationships
  "The" → [0.1, 0.1, 0.8, 0.0, 0.0, 0.0]  # Focuses on "sat"
  "cat" → [0.1, 0.7, 0.2, 0.0, 0.0, 0.0]  # Focuses on itself and "sat"

Head 2: Object-Preposition relationships
  "sat" → [0.0, 0.0, 0.1, 0.8, 0.1, 0.0]  # Focuses on "on"
  "on"  → [0.0, 0.0, 0.0, 0.1, 0.2, 0.7]  # Focuses on "mat"

Head 3: Positional relationships
  Each word attends to neighbors

... up to 8 or 12 heads

def multi_head_attention(X, num_heads=8):
    """
    Args:
        X: Input (batch, seq_len, d_model)
        num_heads: Number of attention heads
    
    Returns:
        output: (batch, seq_len, d_model)
    """
    d_model = X.shape[-1]
    d_k = d_model // num_heads  # Split dimensions across heads
    
    # Create Q, K, V for all heads at once
    Q = X @ W_q  # (batch, seq_len, d_model)
    K = X @ W_k
    V = X @ W_v
    
    # Reshape to separate heads
    # (batch, seq_len, num_heads, d_k)
    Q = Q.reshape(batch, seq_len, num_heads, d_k)
    K = K.reshape(batch, seq_len, num_heads, d_k)
    V = V.reshape(batch, seq_len, num_heads, d_k)
    
    # Transpose to (batch, num_heads, seq_len, d_k)
    Q = Q.transpose(1, 2)
    K = K.transpose(1, 2)
    V = V.transpose(1, 2)
    
    # Apply attention for each head in parallel
    attention_output = scaled_dot_product_attention(Q, K, V)
    # Shape: (batch, num_heads, seq_len, d_k)
    
    # Concatenate heads
    attention_output = attention_output.transpose(1, 2)
    # Shape: (batch, seq_len, num_heads, d_k)
    
    attention_output = attention_output.reshape(batch, seq_len, d_model)
    # Shape: (batch, seq_len, d_model)
    
    # Final linear transformation
    output = attention_output @ W_o
    
    return output

English: "I love machine learning"
French:  "J'adore l'apprentissage automatique"

When generating "apprentissage":
Attention weights to English words:
  I          [0.05]
  love       [0.10]
  machine    [0.35] ████████
  learning   [0.50] ████████████

Model attends heavily to "machine learning" when generating "apprentissage"

Question: "Who invented the transformer?"
Context:  "The transformer architecture was invented by Vaswani et al. 
           in 2017 at Google Brain. The paper 'Attention is All You Need' 
           introduced this revolutionary architecture."

Attention when answering:
  architecture [0.20] ████
  invented     [0.30] ████████
  by           [0.05] █
  Vaswani      [0.35] █████████
  et           [0.05] █
  al           [0.05] █

Answer: "Vaswani et al."

Review: "The movie was good but the ending was terrible"

When predicting sentiment:
Attention weights:
  The      [0.05]
  movie    [0.15] ███
  was      [0.02]
  good     [0.25] ██████
  but      [0.08] ██
  the      [0.02]
  ending   [0.18] ████
  was      [0.02]
  terrible [0.23] ██████

Model focuses on: "good" and "terrible"
Result: Mixed sentiment (conflicting signals)

Time: T₁ → T₂ → T₃ → T₄ → T₅
Must wait for each step to complete

All positions computed simultaneously
Time: T₁ (single forward pass for entire sequence)
100x faster training on GPUs

import matplotlib.pyplot as plt
import seaborn as sns
 
sns.heatmap(attention_weights, 
            xticklabels=words, 
            yticklabels=words)
plt.title("Attention Weights")
plt.show()