Objective: Implement scaled dot-product attention as a reusable function, verify its behavior on constructed examples where you know the "correct" answer in advance, then extend it to multi-head attention. All in runnable, in-browser Python.
Task 1: Implement Scaled Dot-Product Attention
import numpy as np
def softmax(z):
z = z - np.max(z, axis=-1, keepdims=True)
e = np.exp(z)
return e / np.sum(e, axis=-1, keepdims=True)
def attention(Q, K, V):
d_k = K.shape[-1]
scores = (Q @ K.T) / np.sqrt(d_k)
weights = softmax(scores)
output = weights @ V
return output, weights
np.random.seed(0)
Q = np.random.randn(2, 4) # 2 queries
K = np.random.randn(5, 4) # 5 keys
V = np.random.randn(5, 4) # 5 values, same count as keys
output, weights = attention(Q, K, V)
print("Output shape: ", output.shape, " (one output row per query)")
print("Weights shape:", weights.shape, " (one row per query, one column per key)")
print("Each row of weights sums to:", np.round(weights.sum(axis=1), 6))
Task 2: Verify Against a Constructed Example
Build a query that should match exactly one key far more strongly than the others, and confirm the output is dominated by that key's corresponding value.
import numpy as np
def softmax(z):
z = z - np.max(z, axis=-1, keepdims=True)
e = np.exp(z)
return e / np.sum(e, axis=-1, keepdims=True)
def attention(Q, K, V):
d_k = K.shape[-1]
scores = (Q @ K.T) / np.sqrt(d_k)
weights = softmax(scores)
return weights @ V, weights
# 4 keys, each a one-hot-like vector, and a query that matches key index 2 exactly
K = np.array([
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 5.0, 0.0], # scaled up so this key dominates the dot product
[0.0, 0.0, 0.0, 1.0],
])
V = np.array([
[10.0, 0.0],
[0.0, 10.0],
[99.0, 99.0], # the value we expect to dominate the output
[0.0, 0.0],
])
Q = np.array([[0.0, 0.0, 1.0, 0.0]]) # matches the direction of key index 2
output, weights = attention(Q, K, V)
print("Attention weights:", np.round(weights[0], 4))
print("Output: ", np.round(output[0], 3))
print("(Should be close to the value at index 2: [99, 99])")
Task 3: Causal Masking (Looking Only at the Past)
When a decoder generates text one token at a time, it shouldn't be allowed to attend to future positions it hasn't generated yet. A causal mask enforces this by setting the attention score for any "future" position to negative infinity before the softmax, so its weight becomes exactly zero.
import numpy as np
def softmax(z):
z = z - np.max(z, axis=-1, keepdims=True)
e = np.exp(z)
return e / np.sum(e, axis=-1, keepdims=True)
def causal_attention(Q, K, V):
d_k = K.shape[-1]
scores = (Q @ K.T) / np.sqrt(d_k)
n = scores.shape[0]
mask = np.triu(np.ones((n, n)), k=1).astype(bool) # True above the diagonal
scores = np.where(mask, -np.inf, scores)
weights = softmax(scores)
return weights @ V, weights
np.random.seed(1)
n_positions, d_k = 5, 4
Q = K = V = np.random.randn(n_positions, d_k) # self-attention: Q, K, V from the same sequence
_, weights = causal_attention(Q, K, V)
print("Causal attention weight matrix (rows=query position, cols=key position):")
print(np.round(weights, 3))
Every row should have zeros above the diagonal: position 0 can only attend to itself, position 1 can attend to positions 0-1, and so on, exactly the "no looking at the future" constraint a decoder needs.
Task 4: Multi-Head Attention
Multiple attention heads run the same mechanism in parallel with different learned projections, then concatenate the results. Here we simulate that with random projection matrices to see the mechanics, without worrying about training them.
import numpy as np
def softmax(z):
z = z - np.max(z, axis=-1, keepdims=True)
e = np.exp(z)
return e / np.sum(e, axis=-1, keepdims=True)
def attention(Q, K, V):
d_k = K.shape[-1]
scores = (Q @ K.T) / np.sqrt(d_k)
return softmax(scores) @ V
def multi_head_attention(X, n_heads, d_model):
d_head = d_model // n_heads
np.random.seed(0)
head_outputs = []
for h in range(n_heads):
W_q = np.random.randn(d_model, d_head) * 0.3
W_k = np.random.randn(d_model, d_head) * 0.3
W_v = np.random.randn(d_model, d_head) * 0.3
Q, K, V = X @ W_q, X @ W_k, X @ W_v
head_outputs.append(attention(Q, K, V))
return np.concatenate(head_outputs, axis=-1) # back to shape (n_positions, d_model)
X = np.random.randn(6, 8) # 6 positions, d_model=8
output = multi_head_attention(X, n_heads=2, d_model=8)
print("Input shape: ", X.shape)
print("Output shape:", output.shape, "(concatenated heads, back to d_model)")
Each head computes attention independently, with its own learned query/key/value projections, before the results are concatenated back to the original dimensionality. In a real transformer, each head's projections are learned during training, letting different heads specialize in different kinds of relationships, exactly the structural piece this lab simulates with random weights.