From 0a9d5a78a0dcf29a8dffe698d5f6d31bfd7a071d Mon Sep 17 00:00:00 2001
From: Thomas Hart <thomashartdev@gmail.com>
Date: Thu, 6 Aug 2026 20:59:59 +0000
Subject: [PATCH 1/2] feat: Add minimal transformer block with multi-head
 attention

Pre-LN encoder block: multi-head self-attention, position-wise GELU FFN,
residual connections, and LayerNorm. Numpy only, with shape and mask tests.
---
 README.md                 |  31 ++++-
 src/transformer.py        | 250 ++++++++++++++++++++++++++++++++++++++
 tests/test_transformer.py | 178 +++++++++++++++++++++++++++
 3 files changed, 458 insertions(+), 1 deletion(-)
 create mode 100644 src/transformer.py
 create mode 100644 tests/test_transformer.py

diff --git a/README.md b/README.md
index 90dfc41..dd2e99a 100644
--- a/README.md
+++ b/README.md
@@ -4,7 +4,7 @@ Core machine-learning and deep-learning building blocks implemented from scratch
 
 ## What this demonstrates
 
-Every deep learning framework is, at its core, a graph that records operations and replays them backward to compute gradients. This repo builds that machinery by hand and then uses it to train models, so the chain rule, backpropagation, and gradient descent are all visible in a few hundred lines you can read end to end. Each concept lands as its own tested module with the intuition and the math written out, working up to a from-scratch attention block that learns a toy task.
+Every deep learning framework is, at its core, a graph that records operations and replays them backward to compute gradients. This repo builds that machinery by hand and then uses it to train models, so the chain rule, backpropagation, and gradient descent are all visible in a few hundred lines you can read end to end. Each concept lands as its own tested module with the intuition and the math written out, working up to a from-scratch multi-head transformer block: the same residual + LayerNorm + attention + FFN layout that stacks into modern language models.
 
 ## Concepts demonstrated
 
@@ -40,6 +40,13 @@ Every deep learning framework is, at its core, a graph that records operations a
 - Heavy-ball / Polyak momentum: velocity accumulation that carries steps through elongated valleys
 - Adam (adaptive moment estimation): bias-corrected first and second moments, per-parameter step sizes
 - Fair convergence comparison: same MLP init, same minibatches, only the update rule changes
+- Scaled dot-product attention: `softmax(Q Kᵀ / √d_k) V`, including why the scale matters
+- Multi-head self-attention: h parallel heads on `d_model / h` slices, concat, then output projection
+- Causal (autoregressive) masking via `-inf` logits so future positions get zero softmax mass
+- Position-wise feed-forward network (two-layer MLP with GELU) shared across tokens
+- Residual connections as a gradient highway around attention and the FFN
+- Layer normalization on the last feature axis (per-token mean/variance, learnable γ/β)
+- Pre-LN transformer block layout: `x + Sublayer(LayerNorm(x))`, the stable stacking order used in GPT-style models
 
 ## What's implemented
 
@@ -49,6 +56,7 @@ Every deep learning framework is, at its core, a graph that records operations a
 - **Multilayer perceptron with hand-derived backprop**: `src/mlp.py` stacks affine layers with `tanh` or `relu` and a softmax head, and trains a multiclass classifier by minibatch SGD. The backward pass is written out by hand as one recursion on the per-layer delta rather than delegated to an autodiff engine: the output delta is the `softmax - onehot` residual, each hidden delta is `(delta_next @ W_nextᵀ) ⊙ act'(z)`, and the parameter gradients are `dW = a_prevᵀ @ delta` and `db = Σ delta`. Weights use He init for `relu` and Xavier for `tanh` so the signal variance holds across depth. The gradients are verified against central finite differences to a tight tolerance, and the model learns XOR and a three-arm spiral, targets a single hyperplane provably cannot separate. Ships with `make_xor` and `make_spiral` toy generators.
 - **Activation functions + weight initialization (Xavier/He) and why they matter**: `src/activations.py` is the dedicated treatment of the nonlinearity and the initial scale. Each activation (`linear`, `tanh`, `sigmoid`, `relu`, `leaky_relu`) exposes `forward` and a local `backward(z, grad_out)` that multiplies by `act'(z)`, so a hand-written backprop step can drop it in. Xavier/Glorot draws `N(0, 2/(fan_in+fan_out))` (or the matching uniform bound) to keep both forward and backward variance stable for symmetric activations; He/Kaiming draws `N(0, 2/fan_in)` so a ReLU stack does not quietly die after a few layers. `forward_variance_profile` stacks affine+activation layers from unit-variance noise and returns the per-layer activation variance: naive `N(0,1)` weights explode, He keeps a ReLU stack `O(1)`, and Xavier on the same ReLU stack fades, which is the usual silent failure mode when the scheme and the nonlinearity disagree.
 - **SGD, Momentum, Adam from scratch, convergence compared on the same net**: `src/optimizers.py` implements the three standard first-order update rules as numpy-only classes that own their state (velocity for momentum, bias-corrected moments for Adam) and mutate a flat list of parameter arrays in place. The MLP training loop calls `optimizer.step(params, grads)` after the hand-written backprop pass, so swapping the rule never touches the gradient math. `compare_optimizers` retrains the same architecture on the same data with the same seed for each factory, which keeps init and minibatch order fixed and isolates the update rule; on XOR, all three cut loss, and Adam typically pulls ahead of plain SGD early because its per-parameter rates absorb the uneven scale of the gradient.
+- **Minimal transformer block (multi-head attention + FFN + residual + LayerNorm)**: `src/transformer.py` is one Pre-LN encoder block. Multi-head attention projects the sequence to Q, K, and V, splits into `n_heads` slices of size `d_k = d_model / h`, runs scaled dot-product attention per head (with optional causal mask), concatenates, and applies `W_o`. The FFN is a position-wise `d_model → 4·d_model → d_model` MLP with GELU. Each sublayer sits inside a residual branch after LayerNorm (`x + Sub(LN(x))`), so a stack of blocks stays finite without Post-LN warmup tricks. Attention weights are returned for inspection; parameters flatten to a list ready for the same optimizers used by the MLP.
 
 ## Usage
 
@@ -148,6 +156,27 @@ print(forward_variance_profile(10, 64, "relu", "xavier")[-1])
 print(forward_variance_profile(6, 64, "linear", "naive")[-1])
 ```
 
+Run a Pre-LN transformer block on a short sequence:
+
+```python
+import numpy as np
+from src.transformer import TransformerBlock, causal_mask, MultiHeadAttention
+
+rng = np.random.default_rng(0)
+x = rng.normal(size=(2, 8, 32))  # batch=2, T=8, d_model=32
+
+block = TransformerBlock(d_model=32, n_heads=4, d_ff=128, seed=0)
+out = block.forward(x, mask=causal_mask(8))
+print(out.values.shape)          # (2, 8, 32)
+print(out.attn_weights.shape)    # (2, 4, 8, 8); rows sum to 1 per head
+print(out.attn_weights[0, 0].sum(-1))
+
+# multi-head alone, no residual or FFN
+mha = MultiHeadAttention(d_model=32, n_heads=4, seed=1)
+y, weights = mha.forward(x)
+print(y.shape, weights.shape)
+```
+
 Differentiate an arbitrary scalar expression:
 
 ```python
diff --git a/src/transformer.py b/src/transformer.py
new file mode 100644
index 0000000..f7c6d88
--- /dev/null
+++ b/src/transformer.py
@@ -0,0 +1,250 @@
+"""Minimal transformer encoder block: multi-head attention, FFN, residual, LayerNorm.
+
+Pre-LN layout (stable for stacking):
+
+    y = x + MultiHeadAttn(LayerNorm(x))
+    z = y + FFN(LayerNorm(y))
+
+Multi-head attention runs h scaled-dot-product heads on d_k = d_model/h slices,
+concatenates, and projects with W_o. The FFN is a shared two-layer MLP per token.
+Residuals keep a gradient highway; LayerNorm holds the residual-stream scale.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import NDArray
+
+Array = NDArray[np.float64]
+
+
+def softmax(x: Array, axis: int = -1) -> Array:
+    x = np.asarray(x, dtype=np.float64)
+    if x.size == 0:
+        return np.zeros_like(x, dtype=np.float64)
+    shifted = x - np.max(x, axis=axis, keepdims=True)
+    exp = np.exp(shifted)
+    return exp / np.sum(exp, axis=axis, keepdims=True)
+
+
+def causal_mask(seq_len: int) -> Array:
+    if seq_len < 0:
+        raise ValueError("seq_len must be non-negative")
+    return np.tril(np.ones((seq_len, seq_len), dtype=bool))
+
+
+def layer_norm(
+    x: Array,
+    gamma: Array,
+    beta: Array,
+    *,
+    eps: float = 1e-5,
+) -> Array:
+    """Per-token normalize last axis, then affine (γ, β shaped (d_model,))."""
+    x = np.asarray(x, dtype=np.float64)
+    if x.shape[-1] != gamma.shape[0] or gamma.shape != beta.shape:
+        raise ValueError("gamma/beta must match last dim of x")
+    mean = x.mean(axis=-1, keepdims=True)
+    var = x.var(axis=-1, keepdims=True)
+    return gamma * (x - mean) / np.sqrt(var + eps) + beta
+
+
+def gelu(x: Array) -> Array:
+    """GELU via the tanh approximation (Hendrycks & Gimpel)."""
+    x = np.asarray(x, dtype=np.float64)
+    return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3)))
+
+
+def scaled_dot_product_attention(
+    q: Array,
+    k: Array,
+    v: Array,
+    *,
+    mask: Array | None = None,
+) -> tuple[Array, Array]:
+    """Attention(Q,K,V) = softmax(Q Kᵀ / √d_k) V. Returns (out, weights)."""
+    q = np.asarray(q, dtype=np.float64)
+    k = np.asarray(k, dtype=np.float64)
+    v = np.asarray(v, dtype=np.float64)
+    if q.ndim < 2 or k.ndim < 2 or v.ndim < 2:
+        raise ValueError("q, k, v need at least (seq, feature)")
+    if q.shape[-1] != k.shape[-1]:
+        raise ValueError(f"d_k mismatch: {q.shape[-1]} vs {k.shape[-1]}")
+    if k.shape[-2] != v.shape[-2]:
+        raise ValueError(f"T_k mismatch: {k.shape[-2]} vs {v.shape[-2]}")
+    d_k = q.shape[-1]
+    if d_k == 0:
+        raise ValueError("d_k must be positive")
+    scores = np.matmul(q, np.swapaxes(k, -1, -2)) / np.sqrt(float(d_k))
+    if mask is not None:
+        mask_arr = np.asarray(mask)
+        if mask_arr.dtype != np.bool_:
+            mask_arr = mask_arr.astype(bool)
+        scores = np.where(mask_arr, scores, np.float64(-np.inf))
+    weights = np.nan_to_num(softmax(scores, axis=-1), nan=0.0, posinf=0.0, neginf=0.0)
+    return np.matmul(weights, v), weights
+
+
+def _xavier(rng: np.random.Generator, fan_in: int, fan_out: int) -> Array:
+    std = np.sqrt(2.0 / (fan_in + fan_out))
+    return rng.normal(0.0, std, size=(fan_in, fan_out)).astype(np.float64)
+
+
+@dataclass
+class BlockOutput:
+    values: Array
+    attn_weights: Array
+
+
+class MultiHeadAttention:
+    """h parallel heads on d_k = d_model/h, then concat and W_o."""
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        *,
+        rng: np.random.Generator | None = None,
+        seed: int | None = None,
+    ) -> None:
+        if d_model <= 0 or n_heads <= 0:
+            raise ValueError("d_model and n_heads must be positive")
+        if d_model % n_heads != 0:
+            raise ValueError(
+                f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
+            )
+        if rng is None:
+            rng = np.random.default_rng(seed)
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_k = d_model // n_heads
+        self.W_q = _xavier(rng, d_model, d_model)
+        self.W_k = _xavier(rng, d_model, d_model)
+        self.W_v = _xavier(rng, d_model, d_model)
+        self.W_o = _xavier(rng, d_model, d_model)
+
+    def parameters(self) -> list[Array]:
+        return [self.W_q, self.W_k, self.W_v, self.W_o]
+
+    def _split_heads(self, x: Array) -> Array:
+        *lead, t, _ = x.shape
+        return np.moveaxis(x.reshape(*lead, t, self.n_heads, self.d_k), -2, -3)
+
+    def _merge_heads(self, x: Array) -> Array:
+        x = np.moveaxis(x, -3, -2)
+        *lead, t, _, _ = x.shape
+        return x.reshape(*lead, t, self.d_model)
+
+    def forward(
+        self, x: Array, *, mask: Array | None = None
+    ) -> tuple[Array, Array]:
+        x = np.asarray(x, dtype=np.float64)
+        if x.ndim < 2:
+            raise ValueError("x must have shape (..., T, d_model)")
+        if x.shape[-1] != self.d_model:
+            raise ValueError(
+                f"last dim of x is {x.shape[-1]}, expected d_model={self.d_model}"
+            )
+        q = self._split_heads(x @ self.W_q)
+        k = self._split_heads(x @ self.W_k)
+        v = self._split_heads(x @ self.W_v)
+        out, weights = scaled_dot_product_attention(q, k, v, mask=mask)
+        return self._merge_heads(out) @ self.W_o, weights
+
+
+class FeedForward:
+    """Position-wise MLP: d_model → d_ff → d_model with GELU between."""
+
+    def __init__(
+        self,
+        d_model: int,
+        d_ff: int | None = None,
+        *,
+        rng: np.random.Generator | None = None,
+        seed: int | None = None,
+    ) -> None:
+        if d_model <= 0:
+            raise ValueError("d_model must be positive")
+        d_ff = 4 * d_model if d_ff is None else d_ff
+        if d_ff <= 0:
+            raise ValueError("d_ff must be positive")
+        if rng is None:
+            rng = np.random.default_rng(seed)
+        self.d_model = d_model
+        self.d_ff = d_ff
+        self.W1 = _xavier(rng, d_model, d_ff)
+        self.b1 = np.zeros(d_ff, dtype=np.float64)
+        self.W2 = _xavier(rng, d_ff, d_model)
+        self.b2 = np.zeros(d_model, dtype=np.float64)
+
+    def parameters(self) -> list[Array]:
+        return [self.W1, self.b1, self.W2, self.b2]
+
+    def forward(self, x: Array) -> Array:
+        x = np.asarray(x, dtype=np.float64)
+        if x.shape[-1] != self.d_model:
+            raise ValueError(
+                f"last dim of x is {x.shape[-1]}, expected d_model={self.d_model}"
+            )
+        return gelu(x @ self.W1 + self.b1) @ self.W2 + self.b2
+
+
+class TransformerBlock:
+    """Pre-LN encoder block: residual MHA then residual FFN.
+
+    `parameters()` flattens submodule + LayerNorm tensors for an external
+    optimizer. Attention weights have shape (..., n_heads, T, T).
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        d_ff: int | None = None,
+        *,
+        eps: float = 1e-5,
+        rng: np.random.Generator | None = None,
+        seed: int | None = None,
+    ) -> None:
+        if d_model <= 0:
+            raise ValueError("d_model must be positive")
+        if eps <= 0.0:
+            raise ValueError("eps must be positive")
+        if rng is None:
+            rng = np.random.default_rng(seed)
+        self.d_model = d_model
+        self.eps = float(eps)
+        self.attn = MultiHeadAttention(d_model, n_heads, rng=rng)
+        self.ffn = FeedForward(d_model, d_ff, rng=rng)
+        self.ln1_g = np.ones(d_model, dtype=np.float64)
+        self.ln1_b = np.zeros(d_model, dtype=np.float64)
+        self.ln2_g = np.ones(d_model, dtype=np.float64)
+        self.ln2_b = np.zeros(d_model, dtype=np.float64)
+
+    def parameters(self) -> list[Array]:
+        return [
+            *self.attn.parameters(),
+            *self.ffn.parameters(),
+            self.ln1_g,
+            self.ln1_b,
+            self.ln2_g,
+            self.ln2_b,
+        ]
+
+    def forward(self, x: Array, *, mask: Array | None = None) -> BlockOutput:
+        x = np.asarray(x, dtype=np.float64)
+        if x.ndim < 2:
+            raise ValueError("x must have shape (..., T, d_model)")
+        if x.shape[-1] != self.d_model:
+            raise ValueError(
+                f"last dim of x is {x.shape[-1]}, expected d_model={self.d_model}"
+            )
+        h, weights = self.attn.forward(
+            layer_norm(x, self.ln1_g, self.ln1_b, eps=self.eps),
+            mask=mask,
+        )
+        x = x + h
+        x = x + self.ffn.forward(layer_norm(x, self.ln2_g, self.ln2_b, eps=self.eps))
+        return BlockOutput(values=x, attn_weights=weights)
diff --git a/tests/test_transformer.py b/tests/test_transformer.py
new file mode 100644
index 0000000..783cea3
--- /dev/null
+++ b/tests/test_transformer.py
@@ -0,0 +1,178 @@
+import numpy as np
+import pytest
+
+from src.transformer import (
+    FeedForward,
+    MultiHeadAttention,
+    TransformerBlock,
+    causal_mask,
+    gelu,
+    layer_norm,
+    scaled_dot_product_attention,
+)
+
+
+def test_layer_norm_zero_mean_unit_var():
+    rng = np.random.default_rng(0)
+    x = rng.normal(loc=5.0, scale=3.0, size=(4, 6, 16))
+    gamma = np.ones(16)
+    beta = np.zeros(16)
+    y = layer_norm(x, gamma, beta)
+    assert np.allclose(y.mean(axis=-1), 0.0, atol=1e-6)
+    assert np.allclose(y.var(axis=-1), 1.0, atol=1e-4)
+
+
+def test_layer_norm_affine_shift_and_scale():
+    x = np.array([[1.0, 2.0, 3.0, 4.0]])
+    gamma = np.full(4, 2.0)
+    beta = np.full(4, -1.0)
+    y = layer_norm(x, gamma, beta)
+    # after standardizing, affine is 2 * x_hat - 1
+    x_hat = (x - x.mean()) / np.sqrt(x.var() + 1e-5)
+    assert np.allclose(y, 2.0 * x_hat - 1.0)
+
+
+def test_gelu_known_values():
+    assert gelu(np.array([0.0]))[0] == pytest.approx(0.0)
+    # tanh approx of Φ: positive inputs stay positive, negatives shrink toward 0
+    assert gelu(np.array([1.0]))[0] > 0.8
+    assert gelu(np.array([-1.0]))[0] < 0.0
+    assert gelu(np.array([-1.0]))[0] > -0.2
+
+
+def test_attention_shapes_and_row_stochastic():
+    rng = np.random.default_rng(1)
+    q = rng.normal(size=(2, 5, 8))
+    out, w = scaled_dot_product_attention(q, q, q)
+    assert out.shape == (2, 5, 8)
+    assert w.shape == (2, 5, 5)
+    assert np.allclose(w.sum(axis=-1), 1.0)
+
+
+def test_causal_mask_blocks_future():
+    rng = np.random.default_rng(2)
+    t, d = 6, 4
+    q = rng.normal(size=(t, d))
+    _, w = scaled_dot_product_attention(q, q, q, mask=causal_mask(t))
+    assert np.allclose(np.triu(w, k=1), 0.0)
+    assert np.allclose(w.sum(axis=-1), 1.0)
+
+
+def test_empty_sequence_attention():
+    q = np.zeros((0, 4))
+    out, w = scaled_dot_product_attention(q, q, q)
+    assert out.shape == (0, 4)
+    assert w.shape == (0, 0)
+
+
+def test_single_token_attention_is_identity_weight():
+    q = np.array([[1.0, 0.0]])
+    v = np.array([[3.0, 4.0]])
+    out, w = scaled_dot_product_attention(q, q, v)
+    assert w[0, 0] == pytest.approx(1.0)
+    assert np.allclose(out, v)
+
+
+def test_multihead_shapes():
+    mha = MultiHeadAttention(d_model=32, n_heads=4, seed=0)
+    x = np.random.default_rng(3).normal(size=(2, 7, 32))
+    out, w = mha.forward(x)
+    assert out.shape == (2, 7, 32)
+    assert w.shape == (2, 4, 7, 7)
+    assert np.allclose(w.sum(axis=-1), 1.0)
+    assert len(mha.parameters()) == 4
+
+
+def test_multihead_requires_divisible_d_model():
+    with pytest.raises(ValueError, match="divisible"):
+        MultiHeadAttention(d_model=30, n_heads=4)
+
+
+def test_multihead_rejects_bad_feature_dim():
+    mha = MultiHeadAttention(d_model=16, n_heads=2, seed=0)
+    with pytest.raises(ValueError, match="d_model"):
+        mha.forward(np.zeros((3, 15)))
+
+
+def test_multihead_causal():
+    mha = MultiHeadAttention(d_model=16, n_heads=2, seed=1)
+    x = np.random.default_rng(4).normal(size=(5, 16))
+    _, w = mha.forward(x, mask=causal_mask(5))
+    assert np.allclose(np.triu(w, k=1), 0.0)
+
+
+def test_ffn_expands_and_contracts():
+    ffn = FeedForward(d_model=16, d_ff=64, seed=0)
+    x = np.random.default_rng(5).normal(size=(3, 8, 16))
+    y = ffn.forward(x)
+    assert y.shape == x.shape
+    assert ffn.W1.shape == (16, 64)
+    assert ffn.W2.shape == (64, 16)
+    assert len(ffn.parameters()) == 4
+
+
+def test_block_preserves_sequence_shape():
+    block = TransformerBlock(d_model=32, n_heads=4, d_ff=64, seed=0)
+    x = np.random.default_rng(6).normal(size=(2, 9, 32))
+    result = block.forward(x)
+    assert result.values.shape == (2, 9, 32)
+    assert result.attn_weights.shape == (2, 4, 9, 9)
+    assert np.all(np.isfinite(result.values))
+
+
+def test_block_residual_keeps_input_when_sublayers_near_zero():
+    # zero the branch weights so the residual stream is the only path
+    block = TransformerBlock(d_model=8, n_heads=2, d_ff=16, seed=0)
+    for p in block.attn.parameters():
+        p.fill(0.0)
+    for p in block.ffn.parameters():
+        p.fill(0.0)
+    x = np.random.default_rng(7).normal(size=(4, 8))
+    y = block.forward(x).values
+    assert np.allclose(y, x)
+
+
+def test_block_causal_mask():
+    block = TransformerBlock(d_model=16, n_heads=2, seed=2)
+    x = np.random.default_rng(8).normal(size=(6, 16))
+    result = block.forward(x, mask=causal_mask(6))
+    assert np.allclose(np.triu(result.attn_weights, k=1), 0.0)
+
+
+def test_block_empty_sequence():
+    block = TransformerBlock(d_model=8, n_heads=2, seed=0)
+    x = np.zeros((0, 8))
+    result = block.forward(x)
+    assert result.values.shape == (0, 8)
+    assert result.attn_weights.shape == (2, 0, 0)
+
+
+def test_block_single_token():
+    block = TransformerBlock(d_model=8, n_heads=2, seed=3)
+    x = np.random.default_rng(9).normal(size=(1, 8))
+    result = block.forward(x)
+    assert result.values.shape == (1, 8)
+    assert result.attn_weights.shape == (2, 1, 1)
+    assert result.attn_weights[0, 0, 0] == pytest.approx(1.0)
+
+
+def test_block_rejects_bad_shapes_and_ctor():
+    block = TransformerBlock(d_model=16, n_heads=2, seed=0)
+    assert len(block.parameters()) == 12  # 4 attn + 4 ffn + 4 LN
+    with pytest.raises(ValueError, match="d_model"):
+        block.forward(np.zeros((3, 10)))
+    with pytest.raises(ValueError):
+        TransformerBlock(d_model=0, n_heads=1)
+    with pytest.raises(ValueError):
+        TransformerBlock(d_model=8, n_heads=3)
+    with pytest.raises(ValueError):
+        TransformerBlock(d_model=8, n_heads=2, eps=0.0)
+
+
+def test_stacked_blocks_stay_finite():
+    rng = np.random.default_rng(10)
+    x = rng.normal(size=(2, 5, 32))
+    for i in range(4):
+        x = TransformerBlock(d_model=32, n_heads=4, seed=100 + i).forward(x).values
+    assert np.all(np.isfinite(x))
+    assert x.shape == (2, 5, 32)

From 66afdd0f19ed7ca1cb9d0017e6d45f752324f3a9 Mon Sep 17 00:00:00 2001
From: Thomas Hart <thomashartdev@gmail.com>
Date: Thu, 6 Aug 2026 21:07:02 +0000
Subject: [PATCH 2/2] fix: honest forward-only docs, golden attention, masked
 softmax

Document transformer as forward-only (no optimizer/grad claim).
Pin scaled_dot_product_attention to softmax(QK^T/sqrt(d_k))V with a
scale-sensitive golden test. Softmax all-(-inf) rows return zeros
without RuntimeWarning/nan; fully-masked attention covered in tests.
---
 README.md                 |  2 +-
 src/transformer.py        | 31 +++++++++++++++-------
 tests/test_transformer.py | 55 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index dd2e99a..67070ce 100644
--- a/README.md
+++ b/README.md
@@ -56,7 +56,7 @@ Every deep learning framework is, at its core, a graph that records operations a
 - **Multilayer perceptron with hand-derived backprop**: `src/mlp.py` stacks affine layers with `tanh` or `relu` and a softmax head, and trains a multiclass classifier by minibatch SGD. The backward pass is written out by hand as one recursion on the per-layer delta rather than delegated to an autodiff engine: the output delta is the `softmax - onehot` residual, each hidden delta is `(delta_next @ W_nextᵀ) ⊙ act'(z)`, and the parameter gradients are `dW = a_prevᵀ @ delta` and `db = Σ delta`. Weights use He init for `relu` and Xavier for `tanh` so the signal variance holds across depth. The gradients are verified against central finite differences to a tight tolerance, and the model learns XOR and a three-arm spiral, targets a single hyperplane provably cannot separate. Ships with `make_xor` and `make_spiral` toy generators.
 - **Activation functions + weight initialization (Xavier/He) and why they matter**: `src/activations.py` is the dedicated treatment of the nonlinearity and the initial scale. Each activation (`linear`, `tanh`, `sigmoid`, `relu`, `leaky_relu`) exposes `forward` and a local `backward(z, grad_out)` that multiplies by `act'(z)`, so a hand-written backprop step can drop it in. Xavier/Glorot draws `N(0, 2/(fan_in+fan_out))` (or the matching uniform bound) to keep both forward and backward variance stable for symmetric activations; He/Kaiming draws `N(0, 2/fan_in)` so a ReLU stack does not quietly die after a few layers. `forward_variance_profile` stacks affine+activation layers from unit-variance noise and returns the per-layer activation variance: naive `N(0,1)` weights explode, He keeps a ReLU stack `O(1)`, and Xavier on the same ReLU stack fades, which is the usual silent failure mode when the scheme and the nonlinearity disagree.
 - **SGD, Momentum, Adam from scratch, convergence compared on the same net**: `src/optimizers.py` implements the three standard first-order update rules as numpy-only classes that own their state (velocity for momentum, bias-corrected moments for Adam) and mutate a flat list of parameter arrays in place. The MLP training loop calls `optimizer.step(params, grads)` after the hand-written backprop pass, so swapping the rule never touches the gradient math. `compare_optimizers` retrains the same architecture on the same data with the same seed for each factory, which keeps init and minibatch order fixed and isolates the update rule; on XOR, all three cut loss, and Adam typically pulls ahead of plain SGD early because its per-parameter rates absorb the uneven scale of the gradient.
-- **Minimal transformer block (multi-head attention + FFN + residual + LayerNorm)**: `src/transformer.py` is one Pre-LN encoder block. Multi-head attention projects the sequence to Q, K, and V, splits into `n_heads` slices of size `d_k = d_model / h`, runs scaled dot-product attention per head (with optional causal mask), concatenates, and applies `W_o`. The FFN is a position-wise `d_model → 4·d_model → d_model` MLP with GELU. Each sublayer sits inside a residual branch after LayerNorm (`x + Sub(LN(x))`), so a stack of blocks stays finite without Post-LN warmup tricks. Attention weights are returned for inspection; parameters flatten to a list ready for the same optimizers used by the MLP.
+- **Minimal transformer block (multi-head attention + FFN + residual + LayerNorm)**: `src/transformer.py` is one Pre-LN self-attention block (forward-only; no backward/grad path yet). Multi-head attention projects the sequence to Q, K, and V, splits into `n_heads` slices of size `d_k = d_model / h`, runs scaled dot-product attention per head (with optional causal mask), concatenates, and applies `W_o`. The FFN is a position-wise `d_model → 4·d_model → d_model` MLP with GELU. Each sublayer sits inside a residual branch after LayerNorm (`x + Sub(LN(x))`), so a stack of blocks stays finite without Post-LN warmup tricks. Attention weights and `parameters()` (weight tensors for inspection) are returned; training with the MLP optimizers needs a real backward pass, not yet implemented. No positional encodings in this module: inputs are raw token embeddings, so position must be mixed in by the caller if needed.
 
 ## Usage
 
diff --git a/src/transformer.py b/src/transformer.py
index f7c6d88..92f6834 100644
--- a/src/transformer.py
+++ b/src/transformer.py
@@ -1,13 +1,15 @@
-"""Minimal transformer encoder block: multi-head attention, FFN, residual, LayerNorm.
+"""Minimal Pre-LN self-attention block: multi-head attention, FFN, residual, LayerNorm.
 
-Pre-LN layout (stable for stacking):
+Forward-only: no backward/grad path. Pre-LN layout (stable for stacking):
 
     y = x + MultiHeadAttn(LayerNorm(x))
     z = y + FFN(LayerNorm(y))
 
 Multi-head attention runs h scaled-dot-product heads on d_k = d_model/h slices,
-concatenates, and projects with W_o. The FFN is a shared two-layer MLP per token.
-Residuals keep a gradient highway; LayerNorm holds the residual-stream scale.
+concatenates, and projects with W_o. Optional causal mask supports decoder-style
+use; omit the mask for bidirectional encoder-style attention. The FFN is a
+shared two-layer MLP per token. Residuals keep a residual stream; LayerNorm
+holds its scale.
 """
 
 from __future__ import annotations
@@ -21,12 +23,20 @@
 
 
 def softmax(x: Array, axis: int = -1) -> Array:
+    """Numerically stable softmax. All-(-inf) rows return zeros (no mass)."""
     x = np.asarray(x, dtype=np.float64)
     if x.size == 0:
         return np.zeros_like(x, dtype=np.float64)
-    shifted = x - np.max(x, axis=axis, keepdims=True)
-    exp = np.exp(shifted)
-    return exp / np.sum(exp, axis=axis, keepdims=True)
+    max_x = np.max(x, axis=axis, keepdims=True)
+    # All-masked rows: max is -inf; emit zeros without nan/RuntimeWarning.
+    valid = np.isfinite(max_x)
+    with np.errstate(invalid="ignore"):
+        shifted = np.where(valid, x - max_x, 0.0)
+    exp = np.where(valid, np.exp(shifted), 0.0)
+    denom = np.sum(exp, axis=axis, keepdims=True)
+    out = np.zeros_like(exp)
+    np.divide(exp, denom, out=out, where=denom > 0)
+    return out
 
 
 def causal_mask(seq_len: int) -> Array:
@@ -192,10 +202,11 @@ def forward(self, x: Array) -> Array:
 
 
 class TransformerBlock:
-    """Pre-LN encoder block: residual MHA then residual FFN.
+    """Pre-LN self-attention block: residual MHA then residual FFN.
 
-    `parameters()` flattens submodule + LayerNorm tensors for an external
-    optimizer. Attention weights have shape (..., n_heads, T, T).
+    Forward-only: `parameters()` returns weight tensors for inspection, not a
+    trainable optimizer interface (no backward/grad path). Attention weights
+    have shape (..., n_heads, T, T). Optional causal mask for decoder-style use.
     """
 
     def __init__(
diff --git a/tests/test_transformer.py b/tests/test_transformer.py
index 783cea3..e3ae013 100644
--- a/tests/test_transformer.py
+++ b/tests/test_transformer.py
@@ -1,3 +1,5 @@
+import warnings
+
 import numpy as np
 import pytest
 
@@ -9,6 +11,7 @@
     gelu,
     layer_norm,
     scaled_dot_product_attention,
+    softmax,
 )
 
 
@@ -49,6 +52,58 @@ def test_attention_shapes_and_row_stochastic():
     assert np.allclose(w.sum(axis=-1), 1.0)
 
 
+def test_scaled_dot_product_attention_golden():
+    """Core formula: softmax(Q Kᵀ / √d_k) V, with scale load-bearing."""
+    q = np.array([[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]], dtype=np.float64)
+    k = np.array([[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]], dtype=np.float64)
+    v = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float64)
+    d_k = float(q.shape[-1])
+    scores = (q @ k.T) / np.sqrt(d_k)
+    shifted = scores - scores.max(axis=-1, keepdims=True)
+    exp = np.exp(shifted)
+    expected_w = exp / exp.sum(axis=-1, keepdims=True)
+    expected_out = expected_w @ v
+
+    out, w = scaled_dot_product_attention(q, k, v)
+    assert np.allclose(w, expected_w, atol=1e-12)
+    assert np.allclose(out, expected_out, atol=1e-12)
+
+    unscaled = q @ k.T
+    u_shift = unscaled - unscaled.max(axis=-1, keepdims=True)
+    u_exp = np.exp(u_shift)
+    unscaled_w = u_exp / u_exp.sum(axis=-1, keepdims=True)
+    assert not np.allclose(w, unscaled_w, atol=1e-8)
+    assert not np.allclose(out, unscaled_w @ v, atol=1e-8)
+
+
+def test_softmax_all_neg_inf_returns_zeros_without_warning():
+    x = np.full((2, 3), -np.inf)
+    with warnings.catch_warnings(record=True) as rec:
+        warnings.simplefilter("always")
+        y = softmax(x)
+    assert y.shape == (2, 3)
+    assert np.allclose(y, 0.0)
+    assert np.all(np.isfinite(y))
+    assert not any(issubclass(w.category, RuntimeWarning) for w in rec)
+
+
+def test_fully_masked_attention_row_is_zero_finite():
+    q = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float64)
+    k = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float64)
+    v = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float64)
+    # Second query attends to nothing (all keys masked).
+    mask = np.array([[True, True], [False, False]])
+    with warnings.catch_warnings(record=True) as rec:
+        warnings.simplefilter("always")
+        out, w = scaled_dot_product_attention(q, k, v, mask=mask)
+    assert np.allclose(w[1], 0.0)
+    assert np.allclose(out[1], 0.0)
+    assert np.all(np.isfinite(w))
+    assert np.all(np.isfinite(out))
+    assert w[0].sum() == pytest.approx(1.0)
+    assert not any(issubclass(warn.category, RuntimeWarning) for warn in rec)
+
+
 def test_causal_mask_blocks_future():
     rng = np.random.default_rng(2)
     t, d = 6, 4