Transformer Wrap-Up Lab: A Tiny Attention Block¶

Unit: Appendix: Transformer Wrap-Up Worksheet
Role: Optional synthesis
Textbook sections: Transformer data card; Linearity audit; Output layer and forgotten hidden directions; Reading a training objective
Core path: token matrix, query-key-value matrices, attention scores, causal mask, weighted values, logits, and one-token loss
Extension path: attention heatmap and final-layer-only training demo
This lab is not about training a real language model. It is a small numerical model showing how course ideas appear in a transformer-style attention block.

Goals:

  1. Represent a token sequence as a matrix.
  2. Compute query, key, and value matrices.
  3. Interpret dot-product attention scores.
  4. Visualize an attention matrix.
  5. Compute a weighted average of value vectors.
  6. Connect output scores and loss to optimization.

Computational tools used in this lab¶

Before starting, review the NumPy and SymPy Quick Reference for the Labs sections on these tools:

  • Indexing and slicing
  • Reductions, axes, and row-wise computations
  • Numerical checks and roundoff
  • Boolean masks, copying, and attention helpers
  • Plotting helpers

The goal is to interpret the mathematical computation, not to memorize every command.

Core path: helper functions¶

In [ ]:
import numpy as np

def softmax_rows(S):
    """Apply softmax to each row of a matrix."""
    S_shifted = S - np.max(S, axis=1, keepdims=True)
    E = np.exp(S_shifted)
    return E / np.sum(E, axis=1, keepdims=True)

def cross_entropy_from_probs(probs, target_index):
    """Negative log probability of the target index."""
    return -np.log(probs[target_index])

Expected meaning¶

Expected meaning: the shapes show how token vectors become query, key, value, score, attention, and output matrices.

Core path: tokens become vectors¶

Before running the next cell:

  1. How many token vectors are there?
  2. What is the dimension of each token vector?
  3. Which row of X represents "cat"?
In [ ]:
tokens = ["the", "cat", "sat", "."]

X = np.array([
    [1.0, 0.0, 0.5],
    [0.8, 0.7, 0.2],
    [0.1, 1.0, 0.4],
    [0.0, 0.0, 1.0],
])

print("tokens:", tokens)
print("X shape:", X.shape)
print(X)

Core path: similarities by dot product¶

  1. What does each entry of scores_cat compute?
  2. Which token has the largest dot product with "cat"?
  3. Is this exactly cosine similarity? Why or why not?
In [ ]:
scores_cat = X @ X[1]
print(dict(zip(tokens, np.round(scores_cat, 3))))

Core path: compute Q, K, and V¶

Predict the shapes of Q, K, and V before running. Which matrix multiplication rule determines these shapes?

In [ ]:
W_Q = np.array([
    [1.0, 0.0],
    [0.0, 1.0],
    [0.5, 0.5],
])

W_K = np.array([
    [0.5, 0.5],
    [1.0, 0.0],
    [0.0, 1.0],
])

W_V = np.array([
    [1.0, 0.0],
    [0.0, 1.0],
    [1.0, 1.0],
])

Q = X @ W_Q
K = X @ W_K
V = X @ W_V

print("Q shape:", Q.shape)
print("K shape:", K.shape)
print("V shape:", V.shape)

Core path: attention scores¶

  1. Why is S a square matrix?
  2. What does S[i, j] measure?
  3. Which row corresponds to the token "cat" attending to all tokens?
In [ ]:
d_k = Q.shape[1]
S = Q @ K.T / np.sqrt(d_k)

print("S shape:", S.shape)
print(np.round(S, 3))

Core path: causal mask and attention weights¶

  1. Why do the rows of A sum to 1?
  2. What does the mask prevent?
  3. Which entries are forced to be approximately 0 after softmax?
In [ ]:
mask = np.triu(np.ones_like(S), k=1).astype(bool)
S_masked = S.copy()
S_masked[mask] = -1e9

A = softmax_rows(S_masked)

print(np.round(A, 3))
print("row sums:", np.round(A.sum(axis=1), 3))

Expected meaning: mask value¶

Mask note: the value -1e9 is a numerical trick. After softmax, those entries receive nearly zero weight, so future-token positions are effectively ignored.

Extension: visualize the attention matrix¶

The heatmap is optional. It helps you see the causal mask, but the required mathematics is already in the score and weight matrices.

In [ ]:
import matplotlib.pyplot as plt

plt.figure(figsize=(5, 4))
plt.imshow(A)
plt.xticks(range(len(tokens)), tokens)
plt.yticks(range(len(tokens)), tokens)
plt.xlabel("attended-to token")
plt.ylabel("query token")
plt.title("Tiny causal attention matrix")
plt.colorbar(label="attention weight")
plt.show()

Core path: weighted values¶

  1. Why does A @ V compute weighted averages of rows of V?
  2. Which course concept explains weighted averages?
  3. Which course concept explains matrix multiplication as composition?
In [ ]:
H = A @ V
print("H shape:", H.shape)
print(np.round(H, 3))

Core path: output logits and one-token loss¶

  1. What is the shape of h_final?
  2. What is the shape of W_out?
  3. What does each output score represent?
  4. What happens to the loss if the model assigns higher probability to the correct token?
  5. Why is this an optimization problem?
In [ ]:
W_out = np.array([
    [1.0, 0.0, 0.5],
    [0.0, 1.0, 0.2],
])

# Use the final token representation to score three possible next tokens.
h_final = H[-1]
logits = h_final @ W_out

probs = np.exp(logits - logits.max())
probs = probs / probs.sum()

candidate_next_tokens = ["cat", "sat", "."]
target_index = 2  # suppose the correct next token is "."
loss = cross_entropy_from_probs(probs, target_index)
print(dict(zip(candidate_next_tokens, np.round(probs, 3))))
print("loss:", round(loss, 4))

Extension: train only the final output matrix¶

  1. Which parameter matrix changed?
  2. Which parts of the attention block were frozen?
  3. Where does the gradient appear?
  4. Why is this not training a real language model?
In [ ]:
def softmax(z):
    z = z - np.max(z)
    e = np.exp(z)
    return e / e.sum()

W_out_train = W_out.copy()
alpha = 0.5

for step in range(20):
    logits = h_final @ W_out_train
    probs = softmax(logits)
    loss = -np.log(probs[target_index])

    grad_logits = probs.copy()
    grad_logits[target_index] -= 1.0

    grad_W = np.outer(h_final, grad_logits)
    W_out_train -= alpha * grad_W

    if step in [0, 1, 2, 5, 10, 19]:
        print(step, "loss =", round(loss, 4), "probs =", np.round(probs, 3))

Interpretation check¶

Explain the role of each matrix: X, Q, K, V, S, A, and H. Then identify one step from Unit 1, one idea from Unit 3, and one idea from Unit 5.

Exam-style check¶

Given only the shapes X: 4 by 3, W_Q,W_K,W_V: 3 by 2, and A: 4 by 4, determine the shapes of Q, K, V, S, and H. Explain why A @ V is a weighted-average step.

Core path: validation assertions¶

In [ ]:
assert X.shape == (4, 3)
assert Q.shape == K.shape == V.shape == (4, 2)
assert S.shape == (4, 4)
assert A.shape == (4, 4)
assert H.shape == (4, 2)
assert np.allclose(A.sum(axis=1), 1.0)
assert np.allclose(A[mask], 0.0)
print("shape checks passed")