Lab Guide

Format: Jupyter notebook (lab.ipynb) · Time estimate: 45–60 minutes · Prerequisites: Python 3.9+, numpy, matplotlib

Setup

This lab requires only NumPy and Matplotlib — no deep learning frameworks needed. Install dependencies if you have not already:

pip install numpy matplotlib

All exercises implement transformer components from scratch using matrix operations, so you can see exactly what happens at each step.

Exercise 1: Scaled Dot-Product Attention

Goal

Implement the core attention equation from scratch and verify it against known inputs.

Steps

Create a function scaled_dot_product_attention(Q, K, V, mask=None) that:
- Computes attention scores: scores = Q @ K.T / sqrt(d_k)
- Optionally applies a causal mask (set masked positions to -1e9 before softmax)
- Applies softmax row-wise to get attention weights
- Returns both the output (weights @ V) and the attention weights matrix

Implement a helper softmax(x) function using the numerically stable trick:

def softmax(x):
    """Numerically stable softmax along last axis."""
    e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return e_x / np.sum(e_x, axis=-1, keepdims=True)

Test with a small example — 4 tokens, d_k = 8:

np.random.seed(42)
n_tokens, d_k = 4, 8

Q = np.random.randn(n_tokens, d_k)
K = np.random.randn(n_tokens, d_k)
V = np.random.randn(n_tokens, d_k)

output, weights = scaled_dot_product_attention(Q, K, V)
print("Attention weights shape:", weights.shape)  # (4, 4)
print("Output shape:", output.shape)              # (4, 8)
print("Weights row sums:", weights.sum(axis=-1))  # [1, 1, 1, 1]

Add a causal mask and verify that future positions get zero attention weight:

# Create causal mask: True where attention is BLOCKED
causal_mask = np.triu(np.ones((n_tokens, n_tokens), dtype=bool), k=1)

output_masked, weights_masked = scaled_dot_product_attention(Q, K, V, mask=causal_mask)
print("Masked weights:\n", np.round(weights_masked, 3))
# Upper triangle should be all zeros

Expected Output

Attention weights shape: (4, 4)
Output shape: (4, 8)
Weights row sums: [1. 1. 1. 1.]

Masked weights (causal):
 [[1.    0.    0.    0.   ]
  [0.63  0.37  0.    0.   ]
  [0.28  0.41  0.31  0.   ]
  [0.19  0.33  0.22  0.26 ]]

(Exact values will vary with random seed; the key check is that the upper triangle is zero and each row sums to 1.)

Understanding Check

Why do we divide by sqrt(d_k)? Try removing the scaling and observe how the attention weights change. What happens to the softmax distribution?
What does the first row of the causal-masked weights represent? Why is it always [1, 0, 0, 0]?

Exercise 2: Visualise Attention Patterns

Goal

Create heatmap visualisations of attention weight matrices to build intuition about what attention "looks like".

Steps

Use a sequence of real words to make the visualisation interpretable:

tokens = ["The", "cat", "sat", "on", "the", "mat"]
n = len(tokens)
d_k = 16

np.random.seed(123)
Q = np.random.randn(n, d_k)
K = np.random.randn(n, d_k)
V = np.random.randn(n, d_k)

Compute both unmasked and causal-masked attention weights.

Plot both attention matrices side by side as heatmaps:

import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 2, figsize=(12, 5))

for ax, weights, title in [
    (axes[0], w_full, "Bidirectional Attention (Encoder)"),
    (axes[1], w_causal, "Causal Attention (Decoder)")
]:
    im = ax.imshow(weights, cmap="Blues", vmin=0, vmax=1)
    ax.set_xticks(range(n))
    ax.set_yticks(range(n))
    ax.set_xticklabels(tokens, rotation=45)
    ax.set_yticklabels(tokens)
    ax.set_xlabel("Key (attended to)")
    ax.set_ylabel("Query (attending)")
    ax.set_title(title)
    # Annotate cells with weight values
    for i in range(n):
        for j in range(n):
            ax.text(j, i, f"{weights[i,j]:.2f}", ha="center", va="center",
                    fontsize=7, color="black" if weights[i,j] > 0.3 else "gray")

plt.colorbar(im, ax=axes, shrink=0.8, label="Attention Weight")
plt.tight_layout()
plt.savefig("attention_heatmaps.png", dpi=150, bbox_inches="tight")
plt.show()

Create a multi-head visualisation: generate 4 random heads and plot their attention patterns in a 2x2 grid. Observe how each head attends to different positions.

Key Observations to Note

The causal (decoder) heatmap has a clear triangular shape — no attention to future tokens.
The first token in causal mode always has weight 1.0 on itself (it has nothing else to attend to).
Different random heads produce different patterns, illustrating why multi-head attention is powerful — each head can specialise.

Exercise 3: Positional Encodings

Goal

Implement sinusoidal positional encodings and visualise how position information is represented.

Steps

Implement the sinusoidal encoding function:

def sinusoidal_encoding(max_len, d_model):
    """Generate sinusoidal positional encodings.

    PE(pos, 2i)   = sin(pos / 10000^(2i / d_model))
    PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))
    """
    pe = np.zeros((max_len, d_model))
    position = np.arange(max_len)[:, np.newaxis]       # [max_len, 1]
    div_term = 10000 ** (np.arange(0, d_model, 2) / d_model)  # [d_model/2]

    pe[:, 0::2] = np.sin(position / div_term)
    pe[:, 1::2] = np.cos(position / div_term)
    return pe

Generate encodings for 128 positions with d_model = 64 and plot as a heatmap:

pe = sinusoidal_encoding(128, 64)

plt.figure(figsize=(10, 4))
plt.imshow(pe, cmap="RdBu", aspect="auto")
plt.xlabel("Encoding dimension")
plt.ylabel("Position")
plt.title("Sinusoidal Positional Encodings")
plt.colorbar(label="Value")
plt.tight_layout()
plt.savefig("sinusoidal_pe.png", dpi=150, bbox_inches="tight")
plt.show()

Compute and plot the dot-product similarity between positions:

similarity = pe @ pe.T  # [128 x 128]

plt.figure(figsize=(6, 5))
plt.imshow(similarity, cmap="viridis", aspect="auto")
plt.xlabel("Position j")
plt.ylabel("Position i")
plt.title("Position Similarity (dot product of PE vectors)")
plt.colorbar()
plt.tight_layout()
plt.savefig("pe_similarity.png", dpi=150, bbox_inches="tight")
plt.show()

Implement a simplified RoPE encoding and compare the similarity patterns:

def rope_encoding(max_len, d_model):
    """Simplified RoPE: rotate pairs of dimensions by position-dependent angle."""
    pe = np.zeros((max_len, d_model))
    position = np.arange(max_len)[:, np.newaxis]
    freq = 1.0 / (10000 ** (np.arange(0, d_model, 2) / d_model))

    angles = position * freq  # [max_len, d_model/2]
    pe[:, 0::2] = np.cos(angles)
    pe[:, 1::2] = np.sin(angles)
    return pe

rope = rope_encoding(128, 64)
rope_sim = rope @ rope.T

# Plot both similarity matrices side by side
fig, axes = plt.subplots(1, 2, figsize=(12, 5))
for ax, sim, title in [
    (axes[0], similarity, "Sinusoidal PE Similarity"),
    (axes[1], rope_sim, "RoPE Similarity")
]:
    ax.imshow(sim, cmap="viridis", aspect="auto")
    ax.set_xlabel("Position j")
    ax.set_ylabel("Position i")
    ax.set_title(title)
plt.tight_layout()
plt.savefig("pe_comparison.png", dpi=150, bbox_inches="tight")
plt.show()

Key Observations to Note

Low-frequency dimensions change slowly across positions (smooth gradients), while high-frequency dimensions oscillate rapidly.
The position similarity matrix shows that nearby positions are more similar than distant ones, but the relationship is smooth, not discrete.
Both sinusoidal and RoPE produce similar banded patterns in the similarity matrix, but RoPE's pattern depends purely on relative distance (the bands are more uniform along diagonals).

Bonus Challenge

Modify your attention function to use RoPE: instead of adding positional encodings to the embeddings, apply the rotation to Q and K before computing the dot product. Compare the attention patterns with and without RoPE on the same input.

Finishing Up

Once you have completed all three exercises:

Run all cells from top to bottom to verify everything executes without errors (Kernel → Restart & Run All).
Check that your saved images (attention_heatmaps.png, sinusoidal_pe.png, pe_similarity.png, pe_comparison.png) are generated correctly.
Compare your implementation with lab_answer.ipynb.

Self-Check

✓ Exercise 1: Your attention weights sum to 1.0 along each row.
✓ Exercise 1: The causal mask produces zero attention weights in the upper triangle.
✓ Exercise 1: You can explain why removing the sqrt(d_k) scaling causes the softmax to become "peaky".
✓ Exercise 2: Your heatmaps clearly show the difference between bidirectional and causal attention.
✓ Exercise 2: Your multi-head visualisation shows that different heads learn different patterns.
✓ Exercise 3: Your sinusoidal encoding heatmap shows the characteristic wave pattern with varying frequencies.
✓ Exercise 3: Your similarity matrix shows that nearby positions have higher similarity.