16 Transformer — First Principles¶

Goal. Build Transformer components from scratch in NumPy: scaled dot-product attention, multi-head attention, sinusoidal positional encoding, and a full encoder block. Cross-check against PyTorch nn.MultiheadAttention, visualise attention patterns, and demonstrate failure cases.

Prerequisites. theory.md — attention formula, scaling derivation, multi-head structure, positional encoding, encoder/decoder architecture.

In [1]:

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

# PyTorch (optional, for comparison)
try:
    import torch
    import torch.nn as nn
    torch.manual_seed(SEED)
    HAS_TORCH = True
    print(f"PyTorch {torch.__version__} available")
except ImportError:
    HAS_TORCH = False
    print("PyTorch not available — library comparison will be skipped")

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 # PyTorch (optional, for comparison) try: import torch import torch.nn as nn torch.manual_seed(SEED) HAS_TORCH = True print(f"PyTorch {torch.__version__} available") except ImportError: HAS_TORCH = False print("PyTorch not available — library comparison will be skipped")

PyTorch not available — library comparison will be skipped

1. Problem Setup — WHY¶

RNNs process tokens sequentially: token $t$ waits for token $t-1$. This creates two problems: (1) no parallelism during training, and (2) long gradient paths between distant tokens.

Self-attention solves both: every token directly attends to every other token in one matrix multiply. Let's visualise what "attention" looks like on a simple sentence.

In [2]:

# Simulate attention weights for a 5-token sentence.
# In a trained model these would be learned; here we construct a plausible pattern.
tokens = ["The", "cat", "sat", "on", "mat"]
n_tokens = len(tokens)

# Fake but plausible attention: each token mostly attends to nearby + semantically related
raw_scores = rng.standard_normal((n_tokens, n_tokens))
# Make "cat" attend strongly to "sat" and "mat"
raw_scores[1, 2] = 3.0  # cat -> sat
raw_scores[1, 4] = 2.5  # cat -> mat
# Make "sat" attend to "cat" and "on"
raw_scores[2, 1] = 3.0  # sat -> cat
raw_scores[2, 3] = 2.0  # sat -> on

# Softmax row-wise
def softmax(x, axis=-1):
    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)

attn_weights = softmax(raw_scores)

fig, ax = plt.subplots(figsize=(6, 5))
im = ax.imshow(attn_weights, cmap="Blues", vmin=0, vmax=1)
ax.set_xticks(range(n_tokens))
ax.set_yticks(range(n_tokens))
ax.set_xticklabels(tokens)
ax.set_yticklabels(tokens)
ax.set_xlabel("Key (attending TO)")
ax.set_ylabel("Query (attending FROM)")
ax.set_title("Self-Attention Weights\nEach row sums to 1 (softmax)")
for i in range(n_tokens):
    for j in range(n_tokens):
        ax.text(j, i, f"{attn_weights[i, j]:.2f}",
                ha="center", va="center", fontsize=10,
                color="white" if attn_weights[i, j] > 0.4 else "black")
plt.colorbar(im, ax=ax, label="Attention weight")
plt.tight_layout()
plt.show()

print("\nRow sums (should all be 1.0):")
print(attn_weights.sum(axis=1).round(6))

# Simulate attention weights for a 5-token sentence. # In a trained model these would be learned; here we construct a plausible pattern. tokens = ["The", "cat", "sat", "on", "mat"] n_tokens = len(tokens) # Fake but plausible attention: each token mostly attends to nearby + semantically related raw_scores = rng.standard_normal((n_tokens, n_tokens)) # Make "cat" attend strongly to "sat" and "mat" raw_scores[1, 2] = 3.0 # cat -> sat raw_scores[1, 4] = 2.5 # cat -> mat # Make "sat" attend to "cat" and "on" raw_scores[2, 1] = 3.0 # sat -> cat raw_scores[2, 3] = 2.0 # sat -> on # Softmax row-wise def softmax(x, axis=-1): 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) attn_weights = softmax(raw_scores) fig, ax = plt.subplots(figsize=(6, 5)) im = ax.imshow(attn_weights, cmap="Blues", vmin=0, vmax=1) ax.set_xticks(range(n_tokens)) ax.set_yticks(range(n_tokens)) ax.set_xticklabels(tokens) ax.set_yticklabels(tokens) ax.set_xlabel("Key (attending TO)") ax.set_ylabel("Query (attending FROM)") ax.set_title("Self-Attention Weights\nEach row sums to 1 (softmax)") for i in range(n_tokens): for j in range(n_tokens): ax.text(j, i, f"{attn_weights[i, j]:.2f}", ha="center", va="center", fontsize=10, color="white" if attn_weights[i, j] > 0.4 else "black") plt.colorbar(im, ax=ax, label="Attention weight") plt.tight_layout() plt.show() print("\nRow sums (should all be 1.0):") print(attn_weights.sum(axis=1).round(6))

Row sums (should all be 1.0):
[1. 1. 1. 1. 1.]

Reading. Each row is a probability distribution (sums to 1). The attention matrix tells us how much each token "looks at" every other token. In a real Transformer, these weights are computed from learned projections — not hand-crafted.

Key advantage over RNNs: the path from "The" to "mat" is direct (one attention step), not through 4 sequential hidden states.

2. Mathematical Core — WHAT¶

Full derivations are in theory.md. Here we state the equations used in code.

Scaled dot-product attention:

$$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V$$

Why scale by $\sqrt{d_k}$? If $Q, K$ entries have unit variance, then $\text{Var}(q^\top k) = d_k$. Without scaling, large $d_k$ pushes softmax into saturation (one-hot outputs, vanishing gradients). Dividing by $\sqrt{d_k}$ restores unit variance.

Multi-head attention:

$$\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) \, W^O, \quad \text{head}_i = \text{Attention}(X W_i^Q, X W_i^K, X W_i^V)$$

Positional encoding (sinusoidal):

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

3. Solution Method — HOW¶

We build the Transformer bottom-up:

1. scaled_dot_product_attention — the core operation
2. MultiHeadAttention — splits into heads, applies attention, concatenates
3. positional_encoding — injects position information
4. TransformerEncoderBlock — self-attention + FFN + LayerNorm + residuals

Each component is verified independently before stacking.

4. Implementation — BUILD¶

4.1 Scaled Dot-Product Attention¶

In [3]:

def scaled_dot_product_attention(Q, K, V, mask=None):
    """Scaled dot-product attention (theory.md eq. 2.1).

    Args:
        Q: queries, shape (n_q, d_k)
        K: keys, shape (n_k, d_k)
        V: values, shape (n_k, d_v)
        mask: optional, shape (n_q, n_k). 0 = allowed, -inf = blocked.

    Returns:
        output: shape (n_q, d_v)
        attn_weights: shape (n_q, n_k)
    """
    d_k = Q.shape[-1]
    # Step 1-2: compute scaled scores
    scores = Q @ K.T / np.sqrt(d_k)  # (n_q, n_k)
    # Step 3: apply mask
    if mask is not None:
        scores = scores + mask
    # Step 4: row-wise softmax
    attn_weights = softmax(scores, axis=-1)  # (n_q, n_k)
    # Step 5: weighted sum of values
    output = attn_weights @ V  # (n_q, d_v)
    return output, attn_weights


# --- Sanity check ---
# 3 tokens, d_k = 4, d_v = 4
Q_test = rng.standard_normal((3, 4))
K_test = rng.standard_normal((3, 4))
V_test = rng.standard_normal((3, 4))

out, weights = scaled_dot_product_attention(Q_test, K_test, V_test)
print(f"Output shape:  {out.shape}  (expected (3, 4))")
print(f"Weights shape: {weights.shape}  (expected (3, 3))")
print(f"Row sums:      {weights.sum(axis=1).round(6)}  (should be 1.0)")
assert out.shape == (3, 4)
assert weights.shape == (3, 3)
assert np.allclose(weights.sum(axis=1), 1.0, atol=1e-10)
print("All checks passed.")

def scaled_dot_product_attention(Q, K, V, mask=None): """Scaled dot-product attention (theory.md eq. 2.1). Args: Q: queries, shape (n_q, d_k) K: keys, shape (n_k, d_k) V: values, shape (n_k, d_v) mask: optional, shape (n_q, n_k). 0 = allowed, -inf = blocked. Returns: output: shape (n_q, d_v) attn_weights: shape (n_q, n_k) """ d_k = Q.shape[-1] # Step 1-2: compute scaled scores scores = Q @ K.T / np.sqrt(d_k) # (n_q, n_k) # Step 3: apply mask if mask is not None: scores = scores + mask # Step 4: row-wise softmax attn_weights = softmax(scores, axis=-1) # (n_q, n_k) # Step 5: weighted sum of values output = attn_weights @ V # (n_q, d_v) return output, attn_weights # --- Sanity check --- # 3 tokens, d_k = 4, d_v = 4 Q_test = rng.standard_normal((3, 4)) K_test = rng.standard_normal((3, 4)) V_test = rng.standard_normal((3, 4)) out, weights = scaled_dot_product_attention(Q_test, K_test, V_test) print(f"Output shape: {out.shape} (expected (3, 4))") print(f"Weights shape: {weights.shape} (expected (3, 3))") print(f"Row sums: {weights.sum(axis=1).round(6)} (should be 1.0)") assert out.shape == (3, 4) assert weights.shape == (3, 3) assert np.allclose(weights.sum(axis=1), 1.0, atol=1e-10) print("All checks passed.")

Output shape:  (3, 4)  (expected (3, 4))
Weights shape: (3, 3)  (expected (3, 3))
Row sums:      [1. 1. 1.]  (should be 1.0)
All checks passed.

4.2 Verifying the scaling argument¶

Theory.md §2.3 derives that without scaling, dot-product variance grows as $d_k$. Let's verify numerically.

In [4]:

dims = [4, 16, 64, 128, 256, 512]
var_unscaled = []
var_scaled = []

for d_k in dims:
    # Generate many q, k pairs to estimate variance
    q_samples = rng.standard_normal((5000, d_k))
    k_samples = rng.standard_normal((5000, d_k))
    dots = np.sum(q_samples * k_samples, axis=1)  # q^T k for each pair
    var_unscaled.append(np.var(dots))
    var_scaled.append(np.var(dots / np.sqrt(d_k)))

fig, axes = plt.subplots(1, 2, figsize=(11, 4))
axes[0].plot(dims, var_unscaled, "o-", color="crimson", label="Var(q·k)")
axes[0].plot(dims, dims, "--", color="gray", label="y = d_k (theory)")
axes[0].set_xlabel("d_k")
axes[0].set_ylabel("Variance")
axes[0].set_title("Without scaling: Var grows linearly")
axes[0].legend()

axes[1].plot(dims, var_scaled, "o-", color="steelblue", label="Var(q·k / √d_k)")
axes[1].axhline(1.0, color="gray", ls="--", label="y = 1 (theory)")
axes[1].set_xlabel("d_k")
axes[1].set_ylabel("Variance")
axes[1].set_title("With scaling: Var ≈ 1 for all d_k")
axes[1].legend()

plt.tight_layout()
plt.show()

# Numerical check
for d_k, vu, vs in zip(dims, var_unscaled, var_scaled):
    print(f"d_k={d_k:4d}  Var(unscaled)={vu:7.2f}  Var(scaled)={vs:.4f}")
    assert np.isclose(vs, 1.0, atol=0.15), f"Scaled variance should be ~1, got {vs}"

dims = [4, 16, 64, 128, 256, 512] var_unscaled = [] var_scaled = [] for d_k in dims: # Generate many q, k pairs to estimate variance q_samples = rng.standard_normal((5000, d_k)) k_samples = rng.standard_normal((5000, d_k)) dots = np.sum(q_samples * k_samples, axis=1) # q^T k for each pair var_unscaled.append(np.var(dots)) var_scaled.append(np.var(dots / np.sqrt(d_k))) fig, axes = plt.subplots(1, 2, figsize=(11, 4)) axes[0].plot(dims, var_unscaled, "o-", color="crimson", label="Var(q·k)") axes[0].plot(dims, dims, "--", color="gray", label="y = d_k (theory)") axes[0].set_xlabel("d_k") axes[0].set_ylabel("Variance") axes[0].set_title("Without scaling: Var grows linearly") axes[0].legend() axes[1].plot(dims, var_scaled, "o-", color="steelblue", label="Var(q·k / √d_k)") axes[1].axhline(1.0, color="gray", ls="--", label="y = 1 (theory)") axes[1].set_xlabel("d_k") axes[1].set_ylabel("Variance") axes[1].set_title("With scaling: Var ≈ 1 for all d_k") axes[1].legend() plt.tight_layout() plt.show() # Numerical check for d_k, vu, vs in zip(dims, var_unscaled, var_scaled): print(f"d_k={d_k:4d} Var(unscaled)={vu:7.2f} Var(scaled)={vs:.4f}") assert np.isclose(vs, 1.0, atol=0.15), f"Scaled variance should be ~1, got {vs}"

d_k=   4  Var(unscaled)=   4.19  Var(scaled)=1.0480
d_k=  16  Var(unscaled)=  16.23  Var(scaled)=1.0145
d_k=  64  Var(unscaled)=  65.35  Var(scaled)=1.0211
d_k= 128  Var(unscaled)= 125.85  Var(scaled)=0.9832
d_k= 256  Var(unscaled)= 253.15  Var(scaled)=0.9888
d_k= 512  Var(unscaled)= 500.18  Var(scaled)=0.9769

4.3 Multi-Head Attention¶

In [5]:

class MultiHeadAttention:
    """Multi-head attention from scratch (theory.md eq. 3.1).

    Parameters are stored as NumPy arrays for transparency.
    """

    def __init__(self, d_model, n_heads, rng=None):
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # = d_v
        rng = rng or np.random.default_rng()

        # Xavier initialization: scale = sqrt(2 / (fan_in + fan_out))
        scale = np.sqrt(2.0 / (d_model + self.d_k))
        self.W_Q = rng.normal(0, scale, (d_model, d_model))  # all heads packed
        self.W_K = rng.normal(0, scale, (d_model, d_model))
        self.W_V = rng.normal(0, scale, (d_model, d_model))
        self.W_O = rng.normal(0, scale, (d_model, d_model))

    def _split_heads(self, X):
        """Reshape (n, d_model) -> (n_heads, n, d_k)."""
        n = X.shape[0]
        return X.reshape(n, self.n_heads, self.d_k).transpose(1, 0, 2)

    def _merge_heads(self, X):
        """Reshape (n_heads, n, d_k) -> (n, d_model)."""
        return X.transpose(1, 0, 2).reshape(-1, self.d_model)

    def forward(self, X, mask=None):
        """Self-attention forward pass.

        Args:
            X: input, shape (n, d_model)
            mask: optional, shape (n, n)

        Returns:
            output: shape (n, d_model)
            attn_weights: shape (n_heads, n, n)
        """
        # Project to Q, K, V
        Q = X @ self.W_Q  # (n, d_model)
        K = X @ self.W_K
        V = X @ self.W_V

        # Split into heads
        Q_h = self._split_heads(Q)  # (n_heads, n, d_k)
        K_h = self._split_heads(K)
        V_h = self._split_heads(V)

        # Attention per head
        head_outputs = []
        all_weights = []
        for i in range(self.n_heads):
            out_i, w_i = scaled_dot_product_attention(Q_h[i], K_h[i], V_h[i], mask)
            head_outputs.append(out_i)
            all_weights.append(w_i)

        # Concatenate heads and project
        concat = np.concatenate(head_outputs, axis=-1)  # (n, d_model)
        output = concat @ self.W_O  # (n, d_model)

        return output, np.array(all_weights)  # weights: (n_heads, n, n)


# --- Sanity check ---
d_model, n_heads, seq_len = 16, 4, 5
mha = MultiHeadAttention(d_model, n_heads, rng=np.random.default_rng(SEED))
X_mha = rng.standard_normal((seq_len, d_model))

out_mha, weights_mha = mha.forward(X_mha)
print(f"Input shape:   {X_mha.shape}")
print(f"Output shape:  {out_mha.shape}  (should match input)")
print(f"Weights shape: {weights_mha.shape}  (n_heads, n, n)")
assert out_mha.shape == X_mha.shape
assert weights_mha.shape == (n_heads, seq_len, seq_len)
assert np.allclose(weights_mha.sum(axis=-1), 1.0, atol=1e-10)
print("All multi-head attention checks passed.")

class MultiHeadAttention: """Multi-head attention from scratch (theory.md eq. 3.1). Parameters are stored as NumPy arrays for transparency. """ def __init__(self, d_model, n_heads, rng=None): assert d_model % n_heads == 0, "d_model must be divisible by n_heads" self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // n_heads # = d_v rng = rng or np.random.default_rng() # Xavier initialization: scale = sqrt(2 / (fan_in + fan_out)) scale = np.sqrt(2.0 / (d_model + self.d_k)) self.W_Q = rng.normal(0, scale, (d_model, d_model)) # all heads packed self.W_K = rng.normal(0, scale, (d_model, d_model)) self.W_V = rng.normal(0, scale, (d_model, d_model)) self.W_O = rng.normal(0, scale, (d_model, d_model)) def _split_heads(self, X): """Reshape (n, d_model) -> (n_heads, n, d_k).""" n = X.shape[0] return X.reshape(n, self.n_heads, self.d_k).transpose(1, 0, 2) def _merge_heads(self, X): """Reshape (n_heads, n, d_k) -> (n, d_model).""" return X.transpose(1, 0, 2).reshape(-1, self.d_model) def forward(self, X, mask=None): """Self-attention forward pass. Args: X: input, shape (n, d_model) mask: optional, shape (n, n) Returns: output: shape (n, d_model) attn_weights: shape (n_heads, n, n) """ # Project to Q, K, V Q = X @ self.W_Q # (n, d_model) K = X @ self.W_K V = X @ self.W_V # Split into heads Q_h = self._split_heads(Q) # (n_heads, n, d_k) K_h = self._split_heads(K) V_h = self._split_heads(V) # Attention per head head_outputs = [] all_weights = [] for i in range(self.n_heads): out_i, w_i = scaled_dot_product_attention(Q_h[i], K_h[i], V_h[i], mask) head_outputs.append(out_i) all_weights.append(w_i) # Concatenate heads and project concat = np.concatenate(head_outputs, axis=-1) # (n, d_model) output = concat @ self.W_O # (n, d_model) return output, np.array(all_weights) # weights: (n_heads, n, n) # --- Sanity check --- d_model, n_heads, seq_len = 16, 4, 5 mha = MultiHeadAttention(d_model, n_heads, rng=np.random.default_rng(SEED)) X_mha = rng.standard_normal((seq_len, d_model)) out_mha, weights_mha = mha.forward(X_mha) print(f"Input shape: {X_mha.shape}") print(f"Output shape: {out_mha.shape} (should match input)") print(f"Weights shape: {weights_mha.shape} (n_heads, n, n)") assert out_mha.shape == X_mha.shape assert weights_mha.shape == (n_heads, seq_len, seq_len) assert np.allclose(weights_mha.sum(axis=-1), 1.0, atol=1e-10) print("All multi-head attention checks passed.")

Input shape:   (5, 16)
Output shape:  (5, 16)  (should match input)
Weights shape: (4, 5, 5)  (n_heads, n, n)
All multi-head attention checks passed.

4.4 Positional Encoding¶

In [6]:

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

    Returns:
        PE: shape (max_len, d_model)
    """
    PE = np.zeros((max_len, d_model))
    pos = np.arange(max_len)[:, np.newaxis]  # (max_len, 1)
    # Frequency: 1 / 10000^(2i/d)
    i = np.arange(0, d_model, 2)  # even indices
    freq = 1.0 / (10000.0 ** (i / d_model))  # (d_model/2,)

    PE[:, 0::2] = np.sin(pos * freq)  # even dimensions: sin
    PE[:, 1::2] = np.cos(pos * freq)  # odd dimensions: cos
    return PE


# Visualise positional encoding
pe_vis = positional_encoding(max_len=100, d_model=64)

fig, axes = plt.subplots(1, 2, figsize=(13, 4.5))

# Heatmap of all positions and dimensions
im = axes[0].imshow(pe_vis.T, aspect="auto", cmap="RdBu_r", vmin=-1, vmax=1)
axes[0].set_xlabel("Position")
axes[0].set_ylabel("Dimension")
axes[0].set_title("Sinusoidal Positional Encoding\n(PE[pos, dim])")
plt.colorbar(im, ax=axes[0])

# A few individual dimensions
positions = np.arange(100)
for dim_idx in [0, 1, 10, 11, 30, 31]:
    label = f"dim {dim_idx} ({'sin' if dim_idx % 2 == 0 else 'cos'})"
    axes[1].plot(positions, pe_vis[:, dim_idx], label=label, alpha=0.8)
axes[1].set_xlabel("Position")
axes[1].set_ylabel("Encoding value")
axes[1].set_title("Individual PE dimensions\nLow dims = fast oscillation")
axes[1].legend(fontsize=7, ncol=2)

plt.tight_layout()
plt.show()

# Check properties
assert pe_vis.shape == (100, 64)
assert np.all(np.abs(pe_vis) <= 1.0 + 1e-10), "PE values should be in [-1, 1]"
# Each position should have a unique encoding
dists = np.linalg.norm(pe_vis[:, np.newaxis] - pe_vis[np.newaxis, :], axis=-1)
np.fill_diagonal(dists, np.inf)
assert np.all(dists > 0.01), "Positions should have distinct encodings"
print("Positional encoding checks passed.")

def positional_encoding(max_len, d_model): """Sinusoidal positional encoding (theory.md eq. 4.1). Returns: PE: shape (max_len, d_model) """ PE = np.zeros((max_len, d_model)) pos = np.arange(max_len)[:, np.newaxis] # (max_len, 1) # Frequency: 1 / 10000^(2i/d) i = np.arange(0, d_model, 2) # even indices freq = 1.0 / (10000.0 ** (i / d_model)) # (d_model/2,) PE[:, 0::2] = np.sin(pos * freq) # even dimensions: sin PE[:, 1::2] = np.cos(pos * freq) # odd dimensions: cos return PE # Visualise positional encoding pe_vis = positional_encoding(max_len=100, d_model=64) fig, axes = plt.subplots(1, 2, figsize=(13, 4.5)) # Heatmap of all positions and dimensions im = axes[0].imshow(pe_vis.T, aspect="auto", cmap="RdBu_r", vmin=-1, vmax=1) axes[0].set_xlabel("Position") axes[0].set_ylabel("Dimension") axes[0].set_title("Sinusoidal Positional Encoding\n(PE[pos, dim])") plt.colorbar(im, ax=axes[0]) # A few individual dimensions positions = np.arange(100) for dim_idx in [0, 1, 10, 11, 30, 31]: label = f"dim {dim_idx} ({'sin' if dim_idx % 2 == 0 else 'cos'})" axes[1].plot(positions, pe_vis[:, dim_idx], label=label, alpha=0.8) axes[1].set_xlabel("Position") axes[1].set_ylabel("Encoding value") axes[1].set_title("Individual PE dimensions\nLow dims = fast oscillation") axes[1].legend(fontsize=7, ncol=2) plt.tight_layout() plt.show() # Check properties assert pe_vis.shape == (100, 64) assert np.all(np.abs(pe_vis) <= 1.0 + 1e-10), "PE values should be in [-1, 1]" # Each position should have a unique encoding dists = np.linalg.norm(pe_vis[:, np.newaxis] - pe_vis[np.newaxis, :], axis=-1) np.fill_diagonal(dists, np.inf) assert np.all(dists > 0.01), "Positions should have distinct encodings" print("Positional encoding checks passed.")

Positional encoding checks passed.

Reading. Low-index dimensions oscillate quickly (high frequency), high-index dimensions oscillate slowly (low frequency). This creates a unique "fingerprint" for each position. Nearby positions have similar encodings (small distance), while distant positions differ — a useful inductive bias for the model.

4.5 Layer Normalisation and Encoder Block¶

In [7]:

def layer_norm(x, gamma, beta, eps=1e-5):
    """Layer normalisation across the last dimension (theory.md eq. 5.4).

    Args:
        x: input, shape (n, d)
        gamma: scale, shape (d,)
        beta: shift, shape (d,)
    Returns:
        normalised x, shape (n, d)
    """
    mu = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)
    x_norm = (x - mu) / np.sqrt(var + eps)
    return gamma * x_norm + beta


def relu(x):
    return np.maximum(0, x)


class TransformerEncoderBlock:
    """Single Transformer encoder block (theory.md §5).

    Components: MultiHeadAttention + FFN + LayerNorm + residual connections.
    Post-LN variant (original Transformer).
    """

    def __init__(self, d_model, n_heads, d_ff=None, rng=None):
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_ff = d_ff or 4 * d_model
        rng = rng or np.random.default_rng()

        # Multi-head attention
        self.mha = MultiHeadAttention(d_model, n_heads, rng=rng)

        # FFN weights (theory.md eq. 5.3)
        scale_ff = np.sqrt(2.0 / (d_model + self.d_ff))
        self.W1 = rng.normal(0, scale_ff, (d_model, self.d_ff))
        self.b1 = np.zeros(self.d_ff)
        self.W2 = rng.normal(0, scale_ff, (self.d_ff, d_model))
        self.b2 = np.zeros(d_model)

        # LayerNorm parameters (initialise to identity: gamma=1, beta=0)
        self.gamma1 = np.ones(d_model)
        self.beta1 = np.zeros(d_model)
        self.gamma2 = np.ones(d_model)
        self.beta2 = np.zeros(d_model)

    def forward(self, X, mask=None):
        """Forward pass.

        Args:
            X: input, shape (n, d_model)
            mask: optional, shape (n, n)

        Returns:
            output: shape (n, d_model)
            attn_weights: shape (n_heads, n, n)
        """
        # Sub-layer 1: Multi-head self-attention + residual + LayerNorm
        attn_out, attn_weights = self.mha.forward(X, mask)
        Z1 = layer_norm(X + attn_out, self.gamma1, self.beta1)  # (5.1)

        # Sub-layer 2: FFN + residual + LayerNorm
        ffn_out = relu(Z1 @ self.W1 + self.b1) @ self.W2 + self.b2  # (5.3)
        Z2 = layer_norm(Z1 + ffn_out, self.gamma2, self.beta2)  # (5.2)

        return Z2, attn_weights


# --- Sanity check ---
d_model_enc, n_heads_enc = 32, 4
enc_block = TransformerEncoderBlock(d_model_enc, n_heads_enc,
                                    rng=np.random.default_rng(SEED))
X_enc = rng.standard_normal((8, d_model_enc))

out_enc, w_enc = enc_block.forward(X_enc)
print(f"Input shape:   {X_enc.shape}")
print(f"Output shape:  {out_enc.shape}  (should match input)")
print(f"Weights shape: {w_enc.shape}")
assert out_enc.shape == X_enc.shape
# Check LayerNorm output: each row should have mean ≈ 0, std ≈ 1
row_means = out_enc.mean(axis=-1)
row_stds = out_enc.std(axis=-1)
print(f"Row means (should be ~0): {row_means.round(4)}")
print(f"Row stds  (should be ~1): {row_stds.round(4)}")
assert np.allclose(row_means, 0.0, atol=1e-4)
assert np.allclose(row_stds, 1.0, atol=0.05)
print("Encoder block checks passed.")

def layer_norm(x, gamma, beta, eps=1e-5): """Layer normalisation across the last dimension (theory.md eq. 5.4). Args: x: input, shape (n, d) gamma: scale, shape (d,) beta: shift, shape (d,) Returns: normalised x, shape (n, d) """ mu = x.mean(axis=-1, keepdims=True) var = x.var(axis=-1, keepdims=True) x_norm = (x - mu) / np.sqrt(var + eps) return gamma * x_norm + beta def relu(x): return np.maximum(0, x) class TransformerEncoderBlock: """Single Transformer encoder block (theory.md §5). Components: MultiHeadAttention + FFN + LayerNorm + residual connections. Post-LN variant (original Transformer). """ def __init__(self, d_model, n_heads, d_ff=None, rng=None): self.d_model = d_model self.n_heads = n_heads self.d_ff = d_ff or 4 * d_model rng = rng or np.random.default_rng() # Multi-head attention self.mha = MultiHeadAttention(d_model, n_heads, rng=rng) # FFN weights (theory.md eq. 5.3) scale_ff = np.sqrt(2.0 / (d_model + self.d_ff)) self.W1 = rng.normal(0, scale_ff, (d_model, self.d_ff)) self.b1 = np.zeros(self.d_ff) self.W2 = rng.normal(0, scale_ff, (self.d_ff, d_model)) self.b2 = np.zeros(d_model) # LayerNorm parameters (initialise to identity: gamma=1, beta=0) self.gamma1 = np.ones(d_model) self.beta1 = np.zeros(d_model) self.gamma2 = np.ones(d_model) self.beta2 = np.zeros(d_model) def forward(self, X, mask=None): """Forward pass. Args: X: input, shape (n, d_model) mask: optional, shape (n, n) Returns: output: shape (n, d_model) attn_weights: shape (n_heads, n, n) """ # Sub-layer 1: Multi-head self-attention + residual + LayerNorm attn_out, attn_weights = self.mha.forward(X, mask) Z1 = layer_norm(X + attn_out, self.gamma1, self.beta1) # (5.1) # Sub-layer 2: FFN + residual + LayerNorm ffn_out = relu(Z1 @ self.W1 + self.b1) @ self.W2 + self.b2 # (5.3) Z2 = layer_norm(Z1 + ffn_out, self.gamma2, self.beta2) # (5.2) return Z2, attn_weights # --- Sanity check --- d_model_enc, n_heads_enc = 32, 4 enc_block = TransformerEncoderBlock(d_model_enc, n_heads_enc, rng=np.random.default_rng(SEED)) X_enc = rng.standard_normal((8, d_model_enc)) out_enc, w_enc = enc_block.forward(X_enc) print(f"Input shape: {X_enc.shape}") print(f"Output shape: {out_enc.shape} (should match input)") print(f"Weights shape: {w_enc.shape}") assert out_enc.shape == X_enc.shape # Check LayerNorm output: each row should have mean ≈ 0, std ≈ 1 row_means = out_enc.mean(axis=-1) row_stds = out_enc.std(axis=-1) print(f"Row means (should be ~0): {row_means.round(4)}") print(f"Row stds (should be ~1): {row_stds.round(4)}") assert np.allclose(row_means, 0.0, atol=1e-4) assert np.allclose(row_stds, 1.0, atol=0.05) print("Encoder block checks passed.")

Input shape:   (8, 32)
Output shape:  (8, 32)  (should match input)
Weights shape: (4, 8, 8)
Row means (should be ~0): [ 0.  0.  0. -0. -0. -0. -0.  0.]
Row stds  (should be ~1): [1. 1. 1. 1. 1. 1. 1. 1.]
Encoder block checks passed.

4.6 Putting it together: Simple sequence task¶

Apply the encoder to a synthetic sequence of token embeddings and observe how attention mixes information across positions.

In [8]:

# Create a simple sequence: 6 token embeddings of dimension 16
d_model_task = 16
n_tokens_task = 6
token_labels = ["I", "love", "machine", "learning", "from", "scratch"]

# Simulate token embeddings (in a real model, these come from an embedding layer)
X_task = rng.standard_normal((n_tokens_task, d_model_task)) * 0.5

# Add positional encoding
PE_task = positional_encoding(n_tokens_task, d_model_task)
X_with_pos = X_task + PE_task

# Pass through encoder block
enc_task = TransformerEncoderBlock(d_model_task, n_heads=4,
                                   rng=np.random.default_rng(123))
out_task, weights_task = enc_task.forward(X_with_pos)

# Visualise attention patterns for each head
fig, axes = plt.subplots(1, 4, figsize=(16, 4))
for h_idx in range(4):
    im = axes[h_idx].imshow(weights_task[h_idx], cmap="Blues", vmin=0, vmax=0.5)
    axes[h_idx].set_xticks(range(n_tokens_task))
    axes[h_idx].set_yticks(range(n_tokens_task))
    axes[h_idx].set_xticklabels(token_labels, rotation=45, fontsize=8)
    axes[h_idx].set_yticklabels(token_labels, fontsize=8)
    axes[h_idx].set_title(f"Head {h_idx + 1}")
    if h_idx == 0:
        axes[h_idx].set_ylabel("Query")
    axes[h_idx].set_xlabel("Key")
plt.suptitle("Multi-Head Attention Patterns (each head learns a different pattern)",
             fontsize=12, y=1.02)
plt.tight_layout()
plt.show()

print(f"Input  norms per token: {np.linalg.norm(X_with_pos, axis=1).round(3)}")
print(f"Output norms per token: {np.linalg.norm(out_task, axis=1).round(3)}")

# Create a simple sequence: 6 token embeddings of dimension 16 d_model_task = 16 n_tokens_task = 6 token_labels = ["I", "love", "machine", "learning", "from", "scratch"] # Simulate token embeddings (in a real model, these come from an embedding layer) X_task = rng.standard_normal((n_tokens_task, d_model_task)) * 0.5 # Add positional encoding PE_task = positional_encoding(n_tokens_task, d_model_task) X_with_pos = X_task + PE_task # Pass through encoder block enc_task = TransformerEncoderBlock(d_model_task, n_heads=4, rng=np.random.default_rng(123)) out_task, weights_task = enc_task.forward(X_with_pos) # Visualise attention patterns for each head fig, axes = plt.subplots(1, 4, figsize=(16, 4)) for h_idx in range(4): im = axes[h_idx].imshow(weights_task[h_idx], cmap="Blues", vmin=0, vmax=0.5) axes[h_idx].set_xticks(range(n_tokens_task)) axes[h_idx].set_yticks(range(n_tokens_task)) axes[h_idx].set_xticklabels(token_labels, rotation=45, fontsize=8) axes[h_idx].set_yticklabels(token_labels, fontsize=8) axes[h_idx].set_title(f"Head {h_idx + 1}") if h_idx == 0: axes[h_idx].set_ylabel("Query") axes[h_idx].set_xlabel("Key") plt.suptitle("Multi-Head Attention Patterns (each head learns a different pattern)", fontsize=12, y=1.02) plt.tight_layout() plt.show() print(f"Input norms per token: {np.linalg.norm(X_with_pos, axis=1).round(3)}") print(f"Output norms per token: {np.linalg.norm(out_task, axis=1).round(3)}")

Input  norms per token: [4.154 3.747 3.657 3.795 3.201 3.482]
Output norms per token: [4. 4. 4. 4. 4. 4.]

Reading. Each head develops a different attention pattern. With random weights, the patterns are not yet meaningful — but they demonstrate how multi-head attention provides $h$ independent "views" of the sequence. In a trained model, different heads would specialise (e.g., one for adjacent tokens, another for long-range dependencies).

5. Library Comparison¶

Compare our scaled_dot_product_attention against PyTorch's nn.MultiheadAttention by feeding the same input and weights.

In [9]:

if HAS_TORCH:
    # Test with single-head attention for direct comparison
    d_test = 8
    n_test = 4

    # Create input
    X_np = rng.standard_normal((n_test, d_test)).astype(np.float32)

    # Create PyTorch MHA with 1 head (simplifies comparison)
    torch_mha = nn.MultiheadAttention(embed_dim=d_test, num_heads=1, batch_first=False,
                                       bias=False)

    # Extract weights from PyTorch
    with torch.no_grad():
        # PyTorch stores in_proj_weight as [W_Q; W_K; W_V] stacked vertically
        in_proj = torch_mha.in_proj_weight.numpy()  # (3*d, d)
        W_Q_pt = in_proj[:d_test, :]      # (d, d)
        W_K_pt = in_proj[d_test:2*d_test, :]
        W_V_pt = in_proj[2*d_test:, :]
        W_O_pt = torch_mha.out_proj.weight.numpy()  # (d, d)

    # Compute with our implementation
    Q_ours = X_np @ W_Q_pt.T  # PyTorch uses X @ W^T convention
    K_ours = X_np @ W_K_pt.T
    V_ours = X_np @ W_V_pt.T
    out_ours, w_ours = scaled_dot_product_attention(Q_ours, K_ours, V_ours)
    out_ours = out_ours @ W_O_pt.T  # output projection

    # Compute with PyTorch
    # PyTorch MHA expects (seq_len, batch, d) for batch_first=False
    X_torch = torch.from_numpy(X_np).unsqueeze(1)  # (n, 1, d)
    with torch.no_grad():
        out_pt, w_pt = torch_mha(X_torch, X_torch, X_torch)
    out_pt = out_pt.squeeze(1).numpy()
    w_pt = w_pt.squeeze(0).numpy()

    # Compare
    output_diff = np.max(np.abs(out_ours - out_pt))
    weight_diff = np.max(np.abs(w_ours - w_pt))
    print(f"Max output difference: {output_diff:.2e}")
    print(f"Max weight difference: {weight_diff:.2e}")
    assert np.allclose(out_ours, out_pt, atol=1e-5), \
        f"Outputs differ: max diff = {output_diff}"
    assert np.allclose(w_ours, w_pt, atol=1e-5), \
        f"Weights differ: max diff = {weight_diff}"
    print("\n✓ Our attention output matches PyTorch nn.MultiheadAttention")
else:
    print("Skipping PyTorch comparison (not installed).")

if HAS_TORCH: # Test with single-head attention for direct comparison d_test = 8 n_test = 4 # Create input X_np = rng.standard_normal((n_test, d_test)).astype(np.float32) # Create PyTorch MHA with 1 head (simplifies comparison) torch_mha = nn.MultiheadAttention(embed_dim=d_test, num_heads=1, batch_first=False, bias=False) # Extract weights from PyTorch with torch.no_grad(): # PyTorch stores in_proj_weight as [W_Q; W_K; W_V] stacked vertically in_proj = torch_mha.in_proj_weight.numpy() # (3*d, d) W_Q_pt = in_proj[:d_test, :] # (d, d) W_K_pt = in_proj[d_test:2*d_test, :] W_V_pt = in_proj[2*d_test:, :] W_O_pt = torch_mha.out_proj.weight.numpy() # (d, d) # Compute with our implementation Q_ours = X_np @ W_Q_pt.T # PyTorch uses X @ W^T convention K_ours = X_np @ W_K_pt.T V_ours = X_np @ W_V_pt.T out_ours, w_ours = scaled_dot_product_attention(Q_ours, K_ours, V_ours) out_ours = out_ours @ W_O_pt.T # output projection # Compute with PyTorch # PyTorch MHA expects (seq_len, batch, d) for batch_first=False X_torch = torch.from_numpy(X_np).unsqueeze(1) # (n, 1, d) with torch.no_grad(): out_pt, w_pt = torch_mha(X_torch, X_torch, X_torch) out_pt = out_pt.squeeze(1).numpy() w_pt = w_pt.squeeze(0).numpy() # Compare output_diff = np.max(np.abs(out_ours - out_pt)) weight_diff = np.max(np.abs(w_ours - w_pt)) print(f"Max output difference: {output_diff:.2e}") print(f"Max weight difference: {weight_diff:.2e}") assert np.allclose(out_ours, out_pt, atol=1e-5), \ f"Outputs differ: max diff = {output_diff}" assert np.allclose(w_ours, w_pt, atol=1e-5), \ f"Weights differ: max diff = {weight_diff}" print("\n✓ Our attention output matches PyTorch nn.MultiheadAttention") else: print("Skipping PyTorch comparison (not installed).")

Skipping PyTorch comparison (not installed).

Reading. Using the same weights and input, our from-scratch scaled_dot_product_attention produces the same output as PyTorch's nn.MultiheadAttention to within floating-point tolerance. Key difference: PyTorch uses $X W^\top$ (right-multiply by transposed weight), while our formulation uses $X W$ — the weight matrices are transposed.

6. Experiments and Failures — VERIFY¶

6.1 Causal masking¶

In the decoder, future tokens must be masked to preserve the autoregressive property.

In [10]:

def causal_mask(n):
    """Create a causal (lower-triangular) mask.

    Returns mask of shape (n, n) with 0 for allowed and -inf for blocked.
    """
    mask = np.full((n, n), -np.inf)
    mask[np.tril_indices(n)] = 0.0
    return mask


n_causal = 5
mask_c = causal_mask(n_causal)

# Apply to attention
Q_c = rng.standard_normal((n_causal, 4))
K_c = rng.standard_normal((n_causal, 4))
V_c = rng.standard_normal((n_causal, 4))

_, w_unmasked = scaled_dot_product_attention(Q_c, K_c, V_c)
_, w_masked = scaled_dot_product_attention(Q_c, K_c, V_c, mask=mask_c)

fig, axes = plt.subplots(1, 3, figsize=(14, 4))

axes[0].imshow(mask_c, cmap="RdBu_r")
axes[0].set_title("Causal mask\n(0 = allowed, -∞ = blocked)")
for i in range(n_causal):
    for j in range(n_causal):
        val = "0" if mask_c[i, j] == 0 else "-∞"
        axes[0].text(j, i, val, ha="center", va="center", fontsize=10)

for ax_idx, (w, title) in enumerate([(w_unmasked, "Without mask"),
                                       (w_masked, "With causal mask")]):
    im = axes[ax_idx + 1].imshow(w, cmap="Blues", vmin=0, vmax=0.8)
    axes[ax_idx + 1].set_title(f"Attention weights\n{title}")
    for i in range(n_causal):
        for j in range(n_causal):
            axes[ax_idx + 1].text(j, i, f"{w[i, j]:.2f}",
                                   ha="center", va="center", fontsize=9,
                                   color="white" if w[i, j] > 0.4 else "black")
    axes[ax_idx + 1].set_xlabel("Key position")
    axes[ax_idx + 1].set_ylabel("Query position")

plt.tight_layout()
plt.show()

# Verify: upper triangle of masked weights should be 0
upper_triangle = w_masked[np.triu_indices(n_causal, k=1)]
assert np.allclose(upper_triangle, 0.0, atol=1e-10), \
    "Future positions should have zero attention weight"
print("Causal mask verification passed: no leakage to future positions.")

def causal_mask(n): """Create a causal (lower-triangular) mask. Returns mask of shape (n, n) with 0 for allowed and -inf for blocked. """ mask = np.full((n, n), -np.inf) mask[np.tril_indices(n)] = 0.0 return mask n_causal = 5 mask_c = causal_mask(n_causal) # Apply to attention Q_c = rng.standard_normal((n_causal, 4)) K_c = rng.standard_normal((n_causal, 4)) V_c = rng.standard_normal((n_causal, 4)) _, w_unmasked = scaled_dot_product_attention(Q_c, K_c, V_c) _, w_masked = scaled_dot_product_attention(Q_c, K_c, V_c, mask=mask_c) fig, axes = plt.subplots(1, 3, figsize=(14, 4)) axes[0].imshow(mask_c, cmap="RdBu_r") axes[0].set_title("Causal mask\n(0 = allowed, -∞ = blocked)") for i in range(n_causal): for j in range(n_causal): val = "0" if mask_c[i, j] == 0 else "-∞" axes[0].text(j, i, val, ha="center", va="center", fontsize=10) for ax_idx, (w, title) in enumerate([(w_unmasked, "Without mask"), (w_masked, "With causal mask")]): im = axes[ax_idx + 1].imshow(w, cmap="Blues", vmin=0, vmax=0.8) axes[ax_idx + 1].set_title(f"Attention weights\n{title}") for i in range(n_causal): for j in range(n_causal): axes[ax_idx + 1].text(j, i, f"{w[i, j]:.2f}", ha="center", va="center", fontsize=9, color="white" if w[i, j] > 0.4 else "black") axes[ax_idx + 1].set_xlabel("Key position") axes[ax_idx + 1].set_ylabel("Query position") plt.tight_layout() plt.show() # Verify: upper triangle of masked weights should be 0 upper_triangle = w_masked[np.triu_indices(n_causal, k=1)] assert np.allclose(upper_triangle, 0.0, atol=1e-10), \ "Future positions should have zero attention weight" print("Causal mask verification passed: no leakage to future positions.")

Causal mask verification passed: no leakage to future positions.

Reading. With the causal mask, all attention weights above the diagonal are exactly zero. Token 0 attends only to itself, token 1 attends to tokens 0-1, etc. This is critical for autoregressive generation.

6.2 Failure case: Quadratic memory scaling¶

In [11]:

# Demonstrate O(n²) memory cost of attention
import sys

seq_lengths = [64, 128, 256, 512, 1024, 2048]
d_k_mem = 64
memory_bytes = []

for n in seq_lengths:
    # The attention matrix is n × n of float64
    attn_matrix = np.zeros((n, n), dtype=np.float64)
    mem = attn_matrix.nbytes
    memory_bytes.append(mem)

fig, ax = plt.subplots(figsize=(8, 4.5))
memory_mb = [m / 1e6 for m in memory_bytes]
ax.plot(seq_lengths, memory_mb, "o-", color="crimson", lw=2, label="Attention matrix")

# Theoretical O(n²) line
n_ref = seq_lengths[0]
m_ref = memory_mb[0]
theoretical = [m_ref * (n / n_ref) ** 2 for n in seq_lengths]
ax.plot(seq_lengths, theoretical, "--", color="gray", label="O(n²) reference")

ax.set_xlabel("Sequence length n")
ax.set_ylabel("Memory (MB)")
ax.set_title("Attention matrix memory: O(n²) scaling\n"
             "This is the fundamental bottleneck of standard Transformers")
ax.legend()
ax.set_yscale("log")
ax.set_xscale("log")
ax.grid(True, alpha=0.3)

for n, mb in zip(seq_lengths, memory_mb):
    ax.annotate(f"{mb:.2f} MB", (n, mb), textcoords="offset points",
                xytext=(10, 5), fontsize=8)

plt.tight_layout()
plt.show()

print("Memory for attention matrix (single head, float64):")
for n, mb in zip(seq_lengths, memory_mb):
    print(f"  n = {n:5d}  →  {mb:.4f} MB")
print(f"\n4× longer sequence → 16× more memory (quadratic).")

# Verify quadratic scaling
ratio_4x = memory_bytes[-1] / memory_bytes[-3]  # 2048 vs 512
assert np.isclose(ratio_4x, 16.0, atol=0.1), \
    f"Expected 16× memory for 4× sequence length, got {ratio_4x:.1f}×"
print(f"Verified: n=2048 uses {ratio_4x:.0f}× memory vs n=512 (expected 16×).")

# Demonstrate O(n²) memory cost of attention import sys seq_lengths = [64, 128, 256, 512, 1024, 2048] d_k_mem = 64 memory_bytes = [] for n in seq_lengths: # The attention matrix is n × n of float64 attn_matrix = np.zeros((n, n), dtype=np.float64) mem = attn_matrix.nbytes memory_bytes.append(mem) fig, ax = plt.subplots(figsize=(8, 4.5)) memory_mb = [m / 1e6 for m in memory_bytes] ax.plot(seq_lengths, memory_mb, "o-", color="crimson", lw=2, label="Attention matrix") # Theoretical O(n²) line n_ref = seq_lengths[0] m_ref = memory_mb[0] theoretical = [m_ref * (n / n_ref) ** 2 for n in seq_lengths] ax.plot(seq_lengths, theoretical, "--", color="gray", label="O(n²) reference") ax.set_xlabel("Sequence length n") ax.set_ylabel("Memory (MB)") ax.set_title("Attention matrix memory: O(n²) scaling\n" "This is the fundamental bottleneck of standard Transformers") ax.legend() ax.set_yscale("log") ax.set_xscale("log") ax.grid(True, alpha=0.3) for n, mb in zip(seq_lengths, memory_mb): ax.annotate(f"{mb:.2f} MB", (n, mb), textcoords="offset points", xytext=(10, 5), fontsize=8) plt.tight_layout() plt.show() print("Memory for attention matrix (single head, float64):") for n, mb in zip(seq_lengths, memory_mb): print(f" n = {n:5d} → {mb:.4f} MB") print(f"\n4× longer sequence → 16× more memory (quadratic).") # Verify quadratic scaling ratio_4x = memory_bytes[-1] / memory_bytes[-3] # 2048 vs 512 assert np.isclose(ratio_4x, 16.0, atol=0.1), \ f"Expected 16× memory for 4× sequence length, got {ratio_4x:.1f}×" print(f"Verified: n=2048 uses {ratio_4x:.0f}× memory vs n=512 (expected 16×).")

Memory for attention matrix (single head, float64):
  n =    64  →  0.0328 MB
  n =   128  →  0.1311 MB
  n =   256  →  0.5243 MB
  n =   512  →  2.0972 MB
  n =  1024  →  8.3886 MB
  n =  2048  →  33.5544 MB

4× longer sequence → 16× more memory (quadratic).
Verified: n=2048 uses 16× memory vs n=512 (expected 16×).

Reading. Doubling the sequence length quadruples the attention matrix memory. For sequences of length 10,000+ (e.g., long documents), this becomes the primary bottleneck. Techniques like FlashAttention, sparse attention, and linear attention address this by avoiding explicit materialisation of the full $n \times n$ matrix.

6.3 Attention without scaling: softmax saturation¶

In [12]:

# Show that without scaling, attention weights become near-one-hot for large d_k
d_k_values = [4, 32, 128, 512]
n_demo = 8

fig, axes = plt.subplots(1, len(d_k_values), figsize=(16, 4))

for idx, d_k in enumerate(d_k_values):
    Q_demo = rng.standard_normal((n_demo, d_k))
    K_demo = rng.standard_normal((n_demo, d_k))

    # Without scaling
    scores_raw = Q_demo @ K_demo.T
    w_raw = softmax(scores_raw)

    im = axes[idx].imshow(w_raw, cmap="Blues", vmin=0, vmax=1)
    axes[idx].set_title(f"d_k = {d_k}\nmax weight = {w_raw.max():.4f}")
    axes[idx].set_xlabel("Key")
    if idx == 0:
        axes[idx].set_ylabel("Query")

plt.suptitle("Attention WITHOUT scaling: higher d_k → more peaked (saturated)",
             fontsize=12, y=1.02)
plt.tight_layout()
plt.show()

# Compute entropy of attention distributions as a measure of peakedness
print("Entropy of attention distributions (higher = more uniform):")
for d_k in d_k_values:
    Q_e = rng.standard_normal((n_demo, d_k))
    K_e = rng.standard_normal((n_demo, d_k))

    # Without scaling
    w_no = softmax(Q_e @ K_e.T)
    entropy_no = -np.sum(w_no * np.log(w_no + 1e-12), axis=-1).mean()

    # With scaling
    w_sc = softmax(Q_e @ K_e.T / np.sqrt(d_k))
    entropy_sc = -np.sum(w_sc * np.log(w_sc + 1e-12), axis=-1).mean()

    print(f"  d_k={d_k:4d}  entropy(no scale)={entropy_no:.3f}  "
          f"entropy(scaled)={entropy_sc:.3f}")

# Show that without scaling, attention weights become near-one-hot for large d_k d_k_values = [4, 32, 128, 512] n_demo = 8 fig, axes = plt.subplots(1, len(d_k_values), figsize=(16, 4)) for idx, d_k in enumerate(d_k_values): Q_demo = rng.standard_normal((n_demo, d_k)) K_demo = rng.standard_normal((n_demo, d_k)) # Without scaling scores_raw = Q_demo @ K_demo.T w_raw = softmax(scores_raw) im = axes[idx].imshow(w_raw, cmap="Blues", vmin=0, vmax=1) axes[idx].set_title(f"d_k = {d_k}\nmax weight = {w_raw.max():.4f}") axes[idx].set_xlabel("Key") if idx == 0: axes[idx].set_ylabel("Query") plt.suptitle("Attention WITHOUT scaling: higher d_k → more peaked (saturated)", fontsize=12, y=1.02) plt.tight_layout() plt.show() # Compute entropy of attention distributions as a measure of peakedness print("Entropy of attention distributions (higher = more uniform):") for d_k in d_k_values: Q_e = rng.standard_normal((n_demo, d_k)) K_e = rng.standard_normal((n_demo, d_k)) # Without scaling w_no = softmax(Q_e @ K_e.T) entropy_no = -np.sum(w_no * np.log(w_no + 1e-12), axis=-1).mean() # With scaling w_sc = softmax(Q_e @ K_e.T / np.sqrt(d_k)) entropy_sc = -np.sum(w_sc * np.log(w_sc + 1e-12), axis=-1).mean() print(f" d_k={d_k:4d} entropy(no scale)={entropy_no:.3f} " f"entropy(scaled)={entropy_sc:.3f}")

Entropy of attention distributions (higher = more uniform):
  d_k=   4  entropy(no scale)=1.055  entropy(scaled)=1.669
  d_k=  32  entropy(no scale)=0.257  entropy(scaled)=1.687
  d_k= 128  entropy(no scale)=0.243  entropy(scaled)=1.778
  d_k= 512  entropy(no scale)=0.041  entropy(scaled)=1.592

7. Connections¶

From	To	Relationship
Neural Networks (13)	Transformer	FFN sub-layer is a 2-layer MLP; softmax output layer
RNN/LSTM (15)	Transformer	Sequential → parallel; $O(n)$ memory → $O(n^2)$
Softmax	Attention	Softmax converts scores to probability distribution over keys
Residual connections	Encoder block	Same as ResNet: additive skip for gradient flow

Key takeaways:

1. Attention = soft dictionary lookup. Query looks up keys, returns weighted average of values. The weights are data-dependent — unlike fixed convolution kernels.
2. Scaling is essential. Without $\sqrt{d_k}$, dot products grow with dimension, saturating softmax and killing gradients.
3. Multi-head = multiple views. Same parameter count as single-head, but richer representation capacity.
4. Positional encoding breaks permutation symmetry. Without it, the Transformer cannot distinguish "cat sat" from "sat cat."
5. $O(n^2)$ is the price of global attention. This is the main limitation that motivates efficient attention variants.