16 Transformer — Exercises¶

Test your understanding of scaled dot-product attention, causal masking, positional encoding, and the failure modes of the Transformer.

Prerequisites. Read theory.md and work through first_principles.ipynb before attempting these.

In [1]:

Copied!





import random

import matplotlib.pyplot as plt
import numpy as np

%load_ext autoreload
%autoreload 2

SEED = 42
random.seed(SEED)
rng = np.random.default_rng(SEED)

plt.rcParams["figure.dpi"] = 90
import random

import matplotlib.pyplot as plt
import numpy as np

%load_ext autoreload
%autoreload 2

SEED = 42
random.seed(SEED)
rng = np.random.default_rng(SEED)

plt.rcParams["figure.dpi"] = 90

Exercise 1 — Hand Calculation: Scaled Dot-Product Attention on a Tiny Example¶

Given queries, keys, and values for a 2-token sequence ($n = 2$, $d_k = d_v = 2$):

$$Q = \begin{pmatrix} 1 & 0 \\ 0 & 2 \end{pmatrix}, \quad K = \begin{pmatrix} 1 & 1 \\ 0 & 1 \end{pmatrix}, \quad V = \begin{pmatrix} 10 & 0 \\ 0 & 10 \end{pmatrix}$$

Task: compute by hand, following theory.md eq. (2.1):

Raw scores $S = Q K^\top$, where $S_{ij} = q_i^\top k_j$.
Scaled scores $S' = S / \sqrt{d_k}$ with $d_k = 2$.
Attention weights $A = \text{softmax}(S')$ row-wise. Hint: for two logits, $\text{softmax}(z_1, z_2)_1 = \sigma(z_1 - z_2)$, the sigmoid of the difference.
Output $= A V$.

Expected results (4 dp):

$$S = \begin{pmatrix} 1 & 0 \\ 2 & 2 \end{pmatrix}, \quad S' = \begin{pmatrix} 0.7071 & 0 \\ 1.4142 & 1.4142 \end{pmatrix}, \quad A = \begin{pmatrix} 0.6698 & 0.3302 \\ 0.5 & 0.5 \end{pmatrix}, \quad AV = \begin{pmatrix} 6.6976 & 3.3024 \\ 5 & 5 \end{pmatrix}$$

Row 1 of $A$: $\sigma(0.7071 - 0) = 1/(1 + e^{-0.7071}) \approx 0.6698$. Row 2 has equal logits, so the weights are uniform $(0.5, 0.5)$ — token 2's output is the plain average of the two value vectors.

In [2]:

Copied!





def softmax(x, axis=-1):
    """Numerically stable row-wise softmax."""
    x_shifted = x - x.max(axis=axis, keepdims=True)
    exp_x = np.exp(x_shifted)
    return exp_x / exp_x.sum(axis=axis, keepdims=True)


def scaled_dot_product_attention(Q, K, V, mask=None):
    """Scaled dot-product attention (theory.md eq. 2.1)."""
    d_k = Q.shape[-1]
    scores = Q @ np.swapaxes(K, -1, -2) / np.sqrt(d_k)
    if mask is not None:
        scores = scores + mask
    attn_weights = softmax(scores, axis=-1)
    return attn_weights @ V, attn_weights


# Verify the hand calculation
Q1 = np.array([[1.0, 0.0], [0.0, 2.0]])
K1 = np.array([[1.0, 1.0], [0.0, 1.0]])
V1 = np.array([[10.0, 0.0], [0.0, 10.0]])

S = Q1 @ K1.T
S_scaled = S / np.sqrt(2)
out1, A1 = scaled_dot_product_attention(Q1, K1, V1)

print("S =\n", S)
print("S' =\n", S_scaled.round(4))
print("A =\n", A1.round(4))
print("AV =\n", out1.round(4))

# Deterministic checks against the hand values
assert np.allclose(S, [[1.0, 0.0], [2.0, 2.0]], atol=1e-12)
assert np.allclose(S_scaled, [[0.7071, 0.0], [1.4142, 1.4142]], atol=1e-4)
assert np.allclose(A1, [[0.6698, 0.3302], [0.5, 0.5]], atol=1e-4)
assert np.allclose(out1, [[6.6976, 3.3024], [5.0, 5.0]], atol=1e-3)
assert np.allclose(A1.sum(axis=1), 1.0, atol=1e-12), "rows must sum to 1"
print("\nAll hand-calculation checks passed.")
def softmax(x, axis=-1):
    """Numerically stable row-wise softmax."""
    x_shifted = x - x.max(axis=axis, keepdims=True)
    exp_x = np.exp(x_shifted)
    return exp_x / exp_x.sum(axis=axis, keepdims=True)


def scaled_dot_product_attention(Q, K, V, mask=None):
    """Scaled dot-product attention (theory.md eq. 2.1)."""
    d_k = Q.shape[-1]
    scores = Q @ np.swapaxes(K, -1, -2) / np.sqrt(d_k)
    if mask is not None:
        scores = scores + mask
    attn_weights = softmax(scores, axis=-1)
    return attn_weights @ V, attn_weights


# Verify the hand calculation
Q1 = np.array([[1.0, 0.0], [0.0, 2.0]])
K1 = np.array([[1.0, 1.0], [0.0, 1.0]])
V1 = np.array([[10.0, 0.0], [0.0, 10.0]])

S = Q1 @ K1.T
S_scaled = S / np.sqrt(2)
out1, A1 = scaled_dot_product_attention(Q1, K1, V1)

print("S =\n", S)
print("S' =\n", S_scaled.round(4))
print("A =\n", A1.round(4))
print("AV =\n", out1.round(4))

# Deterministic checks against the hand values
assert np.allclose(S, [[1.0, 0.0], [2.0, 2.0]], atol=1e-12)
assert np.allclose(S_scaled, [[0.7071, 0.0], [1.4142, 1.4142]], atol=1e-4)
assert np.allclose(A1, [[0.6698, 0.3302], [0.5, 0.5]], atol=1e-4)
assert np.allclose(out1, [[6.6976, 3.3024], [5.0, 5.0]], atol=1e-3)
assert np.allclose(A1.sum(axis=1), 1.0, atol=1e-12), "rows must sum to 1"
print("\nAll hand-calculation checks passed.")

S =
 [[1. 0.]
 [2. 2.]]
S' =
 [[0.7071 0.    ]
 [1.4142 1.4142]]
A =
 [[0.6698 0.3302]
 [0.5    0.5   ]]
AV =
 [[6.6976 3.3024]
 [5.     5.    ]]

All hand-calculation checks passed.

Exercise 2 — Coding: Causal Masking¶

Implement a function apply_causal_mask(scores) that takes a score matrix of shape (batch, seq_len, seq_len) and blocks all future positions (theory.md §6.1):

$$M_{ij} = \begin{cases} 0 & \text{if } j \le i \\ -\infty & \text{if } j > i \end{cases}$$

Requirements:

The strict upper triangle of the returned scores must be a very large negative number (e.g. -1e9), so those positions get weight $\approx 0$ after softmax.
Past and diagonal entries ($j \le i$) must be unchanged.
Do not mutate the input array.

Deterministic check (run after implementing): after softmax, every weight above the diagonal must be 0 within atol=1e-9, each row must still sum to 1, and query 0 must attend only to itself (weight exactly 1 on position 0).

In [3]:

Copied!





def apply_causal_mask(scores):
    """Block future positions in a (batch, seq_len, seq_len) score tensor.

    Args:
        scores: raw attention scores, shape (batch, seq_len, seq_len)

    Returns:
        masked scores of the same shape (input not mutated)
    """
    # TODO: implement the causal mask
    # Hint:
    # 1. Build a boolean upper-triangular mask with np.triu(..., k=1)
    # 2. Copy the input, set masked entries to -1e9, return the copy
    pass
def apply_causal_mask(scores):
    """Block future positions in a (batch, seq_len, seq_len) score tensor.

    Args:
        scores: raw attention scores, shape (batch, seq_len, seq_len)

    Returns:
        masked scores of the same shape (input not mutated)
    """
    # TODO: implement the causal mask
    # Hint:
    # 1. Build a boolean upper-triangular mask with np.triu(..., k=1)
    # 2. Copy the input, set masked entries to -1e9, return the copy
    pass

Solution 2¶

In [4]:

Copied!





def apply_causal_mask(scores):
    """Block future positions in a (batch, seq_len, seq_len) score tensor."""
    seq_len = scores.shape[-1]
    mask = np.triu(np.ones((seq_len, seq_len)), k=1).astype(bool)
    masked_scores = scores.copy()
    masked_scores[..., mask] = -1e9  # -inf also works; -1e9 avoids inf-arithmetic warnings
    return masked_scores


# Deterministic checks on random scores
scores2 = rng.standard_normal((2, 4, 4))
masked2 = apply_causal_mask(scores2)
weights2 = softmax(masked2, axis=-1)
print("Causal attention weights (batch 0):\n", weights2[0].round(4))

tri_upper = np.triu(np.ones((4, 4)), k=1).astype(bool)
tri_lower = ~tri_upper

# 1. no attention leaks to the future
assert np.allclose(weights2[..., tri_upper], 0.0, atol=1e-9), "future positions must get zero weight"
# 2. allowed entries are untouched, input not mutated
assert np.allclose(masked2[..., tri_lower], scores2[..., tri_lower], atol=1e-15)
assert not np.allclose(scores2[..., tri_upper], -1e9), "input array must not be mutated"
# 3. rows still normalise, and query 0 can only see itself
assert np.allclose(weights2.sum(axis=-1), 1.0, atol=1e-10)
assert np.allclose(weights2[:, 0, 0], 1.0, atol=1e-9)
print("\nAll causal-mask checks passed: no leakage to future positions.")
def apply_causal_mask(scores):
    """Block future positions in a (batch, seq_len, seq_len) score tensor."""
    seq_len = scores.shape[-1]
    mask = np.triu(np.ones((seq_len, seq_len)), k=1).astype(bool)
    masked_scores = scores.copy()
    masked_scores[..., mask] = -1e9  # -inf also works; -1e9 avoids inf-arithmetic warnings
    return masked_scores


# Deterministic checks on random scores
scores2 = rng.standard_normal((2, 4, 4))
masked2 = apply_causal_mask(scores2)
weights2 = softmax(masked2, axis=-1)
print("Causal attention weights (batch 0):\n", weights2[0].round(4))

tri_upper = np.triu(np.ones((4, 4)), k=1).astype(bool)
tri_lower = ~tri_upper

# 1. no attention leaks to the future
assert np.allclose(weights2[..., tri_upper], 0.0, atol=1e-9), "future positions must get zero weight"
# 2. allowed entries are untouched, input not mutated
assert np.allclose(masked2[..., tri_lower], scores2[..., tri_lower], atol=1e-15)
assert not np.allclose(scores2[..., tri_upper], -1e9), "input array must not be mutated"
# 3. rows still normalise, and query 0 can only see itself
assert np.allclose(weights2.sum(axis=-1), 1.0, atol=1e-10)
assert np.allclose(weights2[:, 0, 0], 1.0, atol=1e-9)
print("\nAll causal-mask checks passed: no leakage to future positions.")

Causal attention weights (batch 0):
 [[1.     0.     0.     0.    ]
 [0.3432 0.6568 0.     0.    ]
 [0.2575 0.1116 0.6309 0.    ]
 [0.173  0.5    0.2585 0.0686]]

All causal-mask checks passed: no leakage to future positions.

Exercise 3 — Coding: Sinusoidal Positional Encoding¶

Implement positional_encoding(max_len, d_model) returning the matrix $\text{PE} \in \mathbb{R}^{\text{max\_len} \times d}$ of theory.md eq. (4.1):

$$\text{PE}(\text{pos}, 2i) = \sin\!\left(\frac{\text{pos}}{10000^{2i/d}}\right), \qquad \text{PE}(\text{pos}, 2i+1) = \cos\!\left(\frac{\text{pos}}{10000^{2i/d}}\right)$$

Deterministic checks (for max_len=8, d_model=4, so the two frequency bands are $\omega_0 = 1$ on dims $(0,1)$ and $\omega_1 = 1/100$ on dims $(2,3)$):

Entry	Formula	Expected (6 dp)
$\text{PE}(0, :)$	$(\sin 0, \cos 0, \sin 0, \cos 0)$	$(0, 1, 0, 1)$
$\text{PE}(1, 0)$	$\sin(1)$	$0.841471$
$\text{PE}(1, 1)$	$\cos(1)$	$0.540302$
$\text{PE}(5, 2)$	$\sin(5/100)$	$0.049979$
$\text{PE}(5, 3)$	$\cos(5/100)$	$0.998750$

Bonus property (theory.md §4.2, property 3): for a fixed offset $k$, each $(\sin, \cos)$ pair rotates by the constant angle $k\omega$:

$$\begin{pmatrix} \sin(\omega(p+k)) \\ \cos(\omega(p+k)) \end{pmatrix} = \begin{pmatrix} \cos(k\omega) & \sin(k\omega) \\ -\sin(k\omega) & \cos(k\omega) \end{pmatrix} \begin{pmatrix} \sin(\omega p) \\ \cos(\omega p) \end{pmatrix}$$

Verify this numerically for $k = 3$ on the low-frequency band — it is what lets attention express relative positions with a linear map.

In [5]:

Copied!





def positional_encoding(max_len, d_model):
    """Sinusoidal positional encoding (theory.md eq. 4.1).

    Args:
        max_len: number of positions
        d_model: embedding dimension (assumed even)

    Returns:
        PE: shape (max_len, d_model)
    """
    # TODO: implement
    # Hint:
    # 1. pos = np.arange(max_len)[:, None], even indices i = np.arange(0, d_model, 2)
    # 2. freq = 1 / 10000 ** (i / d_model)
    # 3. PE[:, 0::2] = sin(pos * freq), PE[:, 1::2] = cos(pos * freq)
    pass
def positional_encoding(max_len, d_model):
    """Sinusoidal positional encoding (theory.md eq. 4.1).

    Args:
        max_len: number of positions
        d_model: embedding dimension (assumed even)

    Returns:
        PE: shape (max_len, d_model)
    """
    # TODO: implement
    # Hint:
    # 1. pos = np.arange(max_len)[:, None], even indices i = np.arange(0, d_model, 2)
    # 2. freq = 1 / 10000 ** (i / d_model)
    # 3. PE[:, 0::2] = sin(pos * freq), PE[:, 1::2] = cos(pos * freq)
    pass

Solution 3¶

The rotation property follows from the angle-addition identities $\sin(a+b) = \sin a \cos b + \cos a \sin b$ and $\cos(a+b) = \cos a \cos b - \sin a \sin b$ applied with $a = \omega p$, $b = \omega k$.

In [6]:

Copied!





def positional_encoding(max_len, d_model):
    """Sinusoidal positional encoding (theory.md eq. 4.1)."""
    PE = np.zeros((max_len, d_model))
    pos = np.arange(max_len)[:, np.newaxis]
    i = np.arange(0, d_model, 2)
    freq = 1.0 / (10000.0 ** (i / d_model))
    PE[:, 0::2] = np.sin(pos * freq)
    PE[:, 1::2] = np.cos(pos * freq)
    return PE


PE = positional_encoding(max_len=8, d_model=4)
print("PE (max_len=8, d_model=4):\n", PE.round(6))

# Spot checks from the expected-results table
assert PE.shape == (8, 4)
assert np.allclose(PE[0], [0.0, 1.0, 0.0, 1.0], atol=1e-12)
assert np.isclose(PE[1, 0], 0.841471, atol=1e-6)   # sin(1)
assert np.isclose(PE[1, 1], 0.540302, atol=1e-6)   # cos(1)
assert np.isclose(PE[5, 2], 0.049979, atol=1e-6)   # sin(5/100)
assert np.isclose(PE[5, 3], 0.998750, atol=1e-6)   # cos(5/100)
assert np.all(np.abs(PE) <= 1.0 + 1e-12), "PE entries must lie in [-1, 1]"

# Bonus: relative-shift property on the low-frequency band (dims 2:4, omega = 1/100)
k, omega = 3, 1.0 / 100.0
R = np.array([[np.cos(k * omega), np.sin(k * omega)],
              [-np.sin(k * omega), np.cos(k * omega)]])
for p in range(8 - k):
    assert np.allclose(PE[p + k, 2:4], R @ PE[p, 2:4], atol=1e-10), f"rotation fails at pos {p}"
print("\nAll positional-encoding checks passed (values + rotation property).")
def positional_encoding(max_len, d_model):
    """Sinusoidal positional encoding (theory.md eq. 4.1)."""
    PE = np.zeros((max_len, d_model))
    pos = np.arange(max_len)[:, np.newaxis]
    i = np.arange(0, d_model, 2)
    freq = 1.0 / (10000.0 ** (i / d_model))
    PE[:, 0::2] = np.sin(pos * freq)
    PE[:, 1::2] = np.cos(pos * freq)
    return PE


PE = positional_encoding(max_len=8, d_model=4)
print("PE (max_len=8, d_model=4):\n", PE.round(6))

# Spot checks from the expected-results table
assert PE.shape == (8, 4)
assert np.allclose(PE[0], [0.0, 1.0, 0.0, 1.0], atol=1e-12)
assert np.isclose(PE[1, 0], 0.841471, atol=1e-6)   # sin(1)
assert np.isclose(PE[1, 1], 0.540302, atol=1e-6)   # cos(1)
assert np.isclose(PE[5, 2], 0.049979, atol=1e-6)   # sin(5/100)
assert np.isclose(PE[5, 3], 0.998750, atol=1e-6)   # cos(5/100)
assert np.all(np.abs(PE) <= 1.0 + 1e-12), "PE entries must lie in [-1, 1]"

# Bonus: relative-shift property on the low-frequency band (dims 2:4, omega = 1/100)
k, omega = 3, 1.0 / 100.0
R = np.array([[np.cos(k * omega), np.sin(k * omega)],
              [-np.sin(k * omega), np.cos(k * omega)]])
for p in range(8 - k):
    assert np.allclose(PE[p + k, 2:4], R @ PE[p, 2:4], atol=1e-10), f"rotation fails at pos {p}"
print("\nAll positional-encoding checks passed (values + rotation property).")

PE (max_len=8, d_model=4):
 [[ 0.        1.        0.        1.      ]
 [ 0.841471  0.540302  0.01      0.99995 ]
 [ 0.909297 -0.416147  0.019999  0.9998  ]
 [ 0.14112  -0.989992  0.029996  0.99955 ]
 [-0.756802 -0.653644  0.039989  0.9992  ]
 [-0.958924  0.283662  0.049979  0.99875 ]
 [-0.279415  0.96017   0.059964  0.998201]
 [ 0.656987  0.753902  0.069943  0.997551]]

All positional-encoding checks passed (values + rotation property).

Exercise 4 — Conceptual: Why Scale by $\sqrt{d_k}$?¶

Assume query and key entries are i.i.d. with zero mean and unit variance (theory.md §2.3).

Questions:

What are $\mathbb{E}[q^\top k]$ and $\text{Var}(q^\top k)$ as functions of $d_k$? What is the typical magnitude of an unscaled score when $d_k = 512$?
What does a softmax over logits of that magnitude look like, and why does it make gradients vanish? (Hint: $\partial \text{softmax}/\partial z \approx 0$ when one output is $\approx 1$.)
The scaled score $q^\top k / \sqrt{d_k}$ can be read as a softmax with temperature $T = \sqrt{d_k}$. What do the limits $T \to 0$ and $T \to \infty$ correspond to for the attention weights?

Write your answers, then run the verification cell.

Solution 4¶

$\mathbb{E}[q^\top k] = 0$ and $\text{Var}(q^\top k) = d_k$ (a sum of $d_k$ independent unit-variance products — see theory.md §2.3 for the derivation). Typical magnitude is one standard deviation, $\sqrt{d_k} \approx 22.6$ for $d_k = 512$.
With logit gaps of order 20+, softmax is essentially one-hot (saturated): the largest weight is $\approx 1$ and the rest $\approx 0$. In that regime the softmax Jacobian entries $A_i(\delta_{ij} - A_j)$ are all $\approx 0$, so almost no gradient flows back through the attention scores and the projections stop learning.
$T \to 0$: argmax (hard, one-hot attention — each query picks a single key). $T \to \infty$: uniform weights (attention becomes plain averaging, ignoring content). $T = \sqrt{d_k}$ keeps the logits at unit variance for any $d_k$, holding attention in the informative middle regime.

In [7]:

Copied!





# Verify: Var(q.k) grows like d_k, and unscaled softmax saturates (entropy drops)
d_k_list = [4, 16, 64, 256]
n_tokens = 8
entropies_unscaled, entropies_scaled = [], []

print(f"{'d_k':>5} {'Var(q.k)':>10} {'H(unscaled)':>12} {'H(scaled)':>10}")
for d_k in d_k_list:
    q = rng.standard_normal((2000, d_k))
    k = rng.standard_normal((2000, d_k))
    dots = np.sum(q * k, axis=1)
    assert np.isclose(np.var(dots), d_k, rtol=0.2), f"Var should be ~{d_k}, got {np.var(dots):.1f}"

    Qd = rng.standard_normal((n_tokens, d_k))
    Kd = rng.standard_normal((n_tokens, d_k))
    w_un = softmax(Qd @ Kd.T)                    # no scaling
    w_sc = softmax(Qd @ Kd.T / np.sqrt(d_k))    # scaled
    H_un = -np.sum(w_un * np.log(w_un + 1e-12), axis=-1).mean()
    H_sc = -np.sum(w_sc * np.log(w_sc + 1e-12), axis=-1).mean()
    entropies_unscaled.append(H_un)
    entropies_scaled.append(H_sc)
    print(f"{d_k:>5} {np.var(dots):>10.2f} {H_un:>12.3f} {H_sc:>10.3f}")

# Saturation is systematic for moderate-to-large d_k (max entropy = ln 8 ~ 2.08)
for d_k, H_un, H_sc in zip(d_k_list, entropies_unscaled, entropies_scaled):
    if d_k >= 64:
        assert H_sc > H_un + 0.5, f"scaled entropy should clearly exceed unscaled at d_k={d_k}"

fig, ax = plt.subplots(figsize=(6.5, 4))
ax.plot(d_k_list, entropies_unscaled, "o-", color="crimson", label="unscaled $QK^T$")
ax.plot(d_k_list, entropies_scaled, "o-", color="steelblue", label="scaled $QK^T/\\sqrt{d_k}$")
ax.axhline(np.log(n_tokens), color="gray", ls="--", label="uniform (max entropy)")
ax.set_xscale("log", base=2)
ax.set_xlabel("$d_k$")
ax.set_ylabel("mean attention entropy (nats)")
ax.set_title("Without scaling, attention saturates as $d_k$ grows")
ax.legend()
plt.tight_layout()
plt.show()
print("Saturation checks passed.")
# Verify: Var(q.k) grows like d_k, and unscaled softmax saturates (entropy drops)
d_k_list = [4, 16, 64, 256]
n_tokens = 8
entropies_unscaled, entropies_scaled = [], []

print(f"{'d_k':>5} {'Var(q.k)':>10} {'H(unscaled)':>12} {'H(scaled)':>10}")
for d_k in d_k_list:
    q = rng.standard_normal((2000, d_k))
    k = rng.standard_normal((2000, d_k))
    dots = np.sum(q * k, axis=1)
    assert np.isclose(np.var(dots), d_k, rtol=0.2), f"Var should be ~{d_k}, got {np.var(dots):.1f}"

    Qd = rng.standard_normal((n_tokens, d_k))
    Kd = rng.standard_normal((n_tokens, d_k))
    w_un = softmax(Qd @ Kd.T)                    # no scaling
    w_sc = softmax(Qd @ Kd.T / np.sqrt(d_k))    # scaled
    H_un = -np.sum(w_un * np.log(w_un + 1e-12), axis=-1).mean()
    H_sc = -np.sum(w_sc * np.log(w_sc + 1e-12), axis=-1).mean()
    entropies_unscaled.append(H_un)
    entropies_scaled.append(H_sc)
    print(f"{d_k:>5} {np.var(dots):>10.2f} {H_un:>12.3f} {H_sc:>10.3f}")

# Saturation is systematic for moderate-to-large d_k (max entropy = ln 8 ~ 2.08)
for d_k, H_un, H_sc in zip(d_k_list, entropies_unscaled, entropies_scaled):
    if d_k >= 64:
        assert H_sc > H_un + 0.5, f"scaled entropy should clearly exceed unscaled at d_k={d_k}"

fig, ax = plt.subplots(figsize=(6.5, 4))
ax.plot(d_k_list, entropies_unscaled, "o-", color="crimson", label="unscaled $QK^T$")
ax.plot(d_k_list, entropies_scaled, "o-", color="steelblue", label="scaled $QK^T/\\sqrt{d_k}$")
ax.axhline(np.log(n_tokens), color="gray", ls="--", label="uniform (max entropy)")
ax.set_xscale("log", base=2)
ax.set_xlabel("$d_k$")
ax.set_ylabel("mean attention entropy (nats)")
ax.set_title("Without scaling, attention saturates as $d_k$ grows")
ax.legend()
plt.tight_layout()
plt.show()
print("Saturation checks passed.")

  d_k   Var(q.k)  H(unscaled)  H(scaled)
    4       4.00        1.394      1.808
   16      16.75        0.435      1.559
   64      66.00        0.350      1.728
  256     255.17        0.085      1.689

No description has been provided for this image

Saturation checks passed.

Exercise 5 — Failure Analysis: Permutation Equivariance and Missing Masks¶

(a) Attention treats the input as a set. Self-attention with no positional encoding is permutation equivariant: for any permutation matrix $P$, $\text{Attn}(PX) = P\,\text{Attn}(X)$ (theory.md §4.1). Explain why this holds from the attention formula, and what it implies for a Transformer processing "cat sat" vs "sat cat" without positional encoding. Then run the demo cell to confirm both that the equivariance holds exactly, and that adding $\text{PE}$ breaks it.

(b) Forgetting the causal mask. A decoder trained without the causal mask often reaches excellent training loss but generates gibberish at inference. Explain the mechanism (theory.md §8, failure 4): what information leaks during training, and why is it unavailable at generation time?

Solution 5¶

(a) Every operation in $\text{softmax}(QK^\top/\sqrt{d_k})V$ is built from row-wise projections ($Q = XW^Q$, etc.) and inner products between rows. Permuting the rows of $X$ permutes the rows of $Q$, $K$, $V$ identically, which permutes the rows and columns of the score matrix consistently — so the output rows are the same vectors, just reordered: $\text{Attn}(PX) = P\,\text{Attn}(X)$. Without positional encoding the model therefore assigns "cat sat" and "sat cat" identical token representations (up to reordering) — word order is invisible. Adding $\text{PE}$ breaks the symmetry because each row of $X + \text{PE}$ is now tied to its absolute position: $PX + \text{PE} \ne P(X + \text{PE})$.

(b) Without the mask, position $t$ attends to positions $t+1, \dots, n$ — during teacher forcing the ground-truth future tokens sit right there in the input, so the network learns to copy the answer from the future instead of predicting it. Training loss looks great. At inference the future does not exist yet (tokens are generated one at a time), so the shortcut input is gone and the model, which never learned a genuine predictive mapping, produces garbage. The failure is silent during training — which is what makes it dangerous.

In [8]:

Copied!





# Demo (a): permutation equivariance of self-attention, and how PE breaks it
d_model, n = 8, 6
W_Q = rng.standard_normal((d_model, d_model)) * 0.5
W_K = rng.standard_normal((d_model, d_model)) * 0.5
W_V = rng.standard_normal((d_model, d_model)) * 0.5


def self_attention(X):
    out, _ = scaled_dot_product_attention(X @ W_Q, X @ W_K, X @ W_V)
    return out


X5 = rng.standard_normal((n, d_model))
perm = np.array([3, 0, 5, 1, 4, 2])  # a fixed permutation of the 6 tokens

# Without PE: Attn(PX) == P Attn(X) exactly
out_plain = self_attention(X5)
out_permuted = self_attention(X5[perm])
diff_no_pe = np.max(np.abs(out_permuted - out_plain[perm]))
print(f"max |Attn(PX) - P Attn(X)| without PE: {diff_no_pe:.2e}")
assert np.allclose(out_permuted, out_plain[perm], atol=1e-12), "equivariance should hold exactly"

# With PE: the symmetry is broken
PE5 = positional_encoding(n, d_model)
out_pe = self_attention(X5 + PE5)
out_pe_permuted = self_attention(X5[perm] + PE5)
diff_pe = np.max(np.abs(out_pe_permuted - out_pe[perm]))
print(f"max |Attn(PX+PE) - P Attn(X+PE)| with PE: {diff_pe:.2e}")
assert not np.allclose(out_pe_permuted, out_pe[perm], atol=1e-6), "PE should break equivariance"

print("\nWithout PE the Transformer cannot tell 'cat sat' from 'sat cat';")
print("with PE, each position carries a unique fingerprint and order matters.")
# Demo (a): permutation equivariance of self-attention, and how PE breaks it
d_model, n = 8, 6
W_Q = rng.standard_normal((d_model, d_model)) * 0.5
W_K = rng.standard_normal((d_model, d_model)) * 0.5
W_V = rng.standard_normal((d_model, d_model)) * 0.5


def self_attention(X):
    out, _ = scaled_dot_product_attention(X @ W_Q, X @ W_K, X @ W_V)
    return out


X5 = rng.standard_normal((n, d_model))
perm = np.array([3, 0, 5, 1, 4, 2])  # a fixed permutation of the 6 tokens

# Without PE: Attn(PX) == P Attn(X) exactly
out_plain = self_attention(X5)
out_permuted = self_attention(X5[perm])
diff_no_pe = np.max(np.abs(out_permuted - out_plain[perm]))
print(f"max |Attn(PX) - P Attn(X)| without PE: {diff_no_pe:.2e}")
assert np.allclose(out_permuted, out_plain[perm], atol=1e-12), "equivariance should hold exactly"

# With PE: the symmetry is broken
PE5 = positional_encoding(n, d_model)
out_pe = self_attention(X5 + PE5)
out_pe_permuted = self_attention(X5[perm] + PE5)
diff_pe = np.max(np.abs(out_pe_permuted - out_pe[perm]))
print(f"max |Attn(PX+PE) - P Attn(X+PE)| with PE: {diff_pe:.2e}")
assert not np.allclose(out_pe_permuted, out_pe[perm], atol=1e-6), "PE should break equivariance"

print("\nWithout PE the Transformer cannot tell 'cat sat' from 'sat cat';")
print("with PE, each position carries a unique fingerprint and order matters.")

max |Attn(PX) - P Attn(X)| without PE: 4.44e-16
max |Attn(PX+PE) - P Attn(X+PE)| with PE: 2.30e+00

Without PE the Transformer cannot tell 'cat sat' from 'sat cat';
with PE, each position carries a unique fingerprint and order matters.