feat: Add minimal transformer block with multi-head attention

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 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
 
@@ -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

src/transformer.py

@@ -0,0 +1,261 @@
+"""Minimal Pre-LN self-attention block: multi-head attention, FFN, residual, LayerNorm.
+
+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. 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
+
+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:
+    """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)
+    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:
+    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 self-attention block: residual MHA then residual FFN.
+
+    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__(
+        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)