4. Transformers & Attention

Activity: Building Attention and Transformers from Scratch

This activity solidifies your understanding of attention mechanisms and Transformer architecture by implementing them from the ground up using only NumPy and Python (no PyTorch or TensorFlow for the core logic).

Exercise 1 — Scaled Dot-Product Attention

Implement the full scaled dot-product attention function:

Instructions

1. Implement softmax(x) — numerically stable version (subtract max before exponentiating)
2. Implement scaled_dot_product_attention(Q, K, V, mask=None):
3. Compute raw scores: scores = Q @ K.T / sqrt(d_k)
4. Apply mask if provided (set masked positions to -inf before softmax)
5. Apply softmax row-wise
6. Return weighted sum of Values: output = attn_weights @ V
7. Test with the following inputs:

import numpy as np

d_k = 4
Q = np.array([[1.0, 0.0, 1.0, 0.0],   # token 1 query
              [0.0, 1.0, 0.0, 1.0]])   # token 2 query
K = np.array([[1.0, 0.0, 1.0, 0.0],   # token 1 key
              [0.0, 1.0, 0.0, 1.0],   # token 2 key
              [1.0, 1.0, 0.0, 0.0]])  # token 3 key
V = np.array([[1.0, 0.0],
              [0.0, 1.0],
              [0.5, 0.5]])

1. Plot the attention weight matrix as a heatmap (use matplotlib). What pattern do you observe? Does token 1 attend more to token 1 or token 3? Why?

2. Apply a causal mask (lower-triangular) and re-run. Show how the attention weights change and explain why this mask is necessary for autoregressive generation.

Expected output

Report:

• The attention weight matrix (2×3) before and after masking
• The output matrix (2×2)
• Heatmap visualizations
• A brief explanation of why Q1 attends more to K1 than to K2

Exercise 2 — Multi-Head Attention from Scratch

Extend your implementation to Multi-Head Attention with \(h=2\) heads.

Architecture

where each head uses its own projection matrices.

Instructions

1. Implement the class MultiHeadAttention with:
2. __init__(d_model, num_heads) — initialize random weight matrices \(W_Q^i, W_K^i, W_V^i \in \mathbb{R}^{d_{model} \times d_k}\) and \(W^O \in \mathbb{R}^{d_{model} \times d_{model}}\) for each head \(i\)
3. forward(Q, K, V) — project, apply per-head attention, concatenate, project output
4. Use fixed random seed np.random.seed(42) for reproducibility
5. Test with a sequence of 5 tokens, each with d_model=8, num_heads=2 (so d_k=4 per head)
6. Verify the output shape is (5, 8) — same as input

Questions to answer in your report

• Why does using \(h=2\) heads with \(d_k = d_{model}/h\) keep the total computation similar to \(h=1\)?
• If head 1 learns to attend to nearby tokens and head 2 to distant tokens, how does the concatenated output benefit from both?

Exercise 3 — Single-Layer Transformer Block

Combine your attention with a Feed-Forward Network to implement a Transformer Encoder Block:

Instructions

1. Implement layer_norm(x) — normalize per row (subtract mean, divide by std + ε), with learnable γ=1, β=0
2. Implement ffn(x, W1, b1, W2, b2) — two linear layers with ReLU: \(\text{FFN}(x) = \text{ReLU}(xW_1 + b_1)W_2 + b_2\), where \(d_{ff} = 4 \times d_{model}\)
3. Stack it all: build transformer_encoder_block(x, mha, W1, b1, W2, b2) using your implementations above
4. Test with 5 tokens, d_model=8

Visualization

Plot the token representations before and after the block as a heatmap (tokens × dimensions). Do the representations become richer after the block? Compute and report the cosine similarity matrix between tokens before and after the block.

Exercise 4 — Positional Encoding

Implement sinusoidal positional encoding and visualize it.

Instructions

1. Implement positional_encoding(max_len, d_model) → returns matrix of shape (max_len, d_model)
2. Plot two visualizations:
3. Heatmap: rows=positions (0–99), cols=dimensions, color=PE value
4. Line plot: PE values for positions 0, 10, 50 across all dimensions

Questions

• Which dimensions encode high-frequency oscillations and which encode low-frequency?
• Why does adding PE to the token embedding allow the Transformer to distinguish position 1 from position 50?
• What happens if you add the same positional encoding to shuffled tokens?

Evaluation Criteria

Submission format: GitHub Pages (using the course template). No other format accepted.

AI Collaboration: Allowed, but every student must be able to explain all code and analysis. Oral exams may be required.