From 11ba8a2ddf75667ceb4ff953bd2f631461ab08d5 Mon Sep 17 00:00:00 2001 From: Thomas Hart 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 | 45 +++---- src/transformer.py | 250 ++++++++++++++++++++++++++++++++++++++ tests/test_transformer.py | 178 +++++++++++++++++++++++++++ 3 files changed, 451 insertions(+), 22 deletions(-) create mode 100644 src/transformer.py create mode 100644 tests/test_transformer.py diff --git a/README.md b/README.md index c9e83e1..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,11 +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` with the 1/√d_k temperature from "Attention Is All You Need" -- Self-attention as Q, K, V linear projections of the same sequence (single head) -- Causal (autoregressive) masking via lower-triangular boolean masks and `-inf` logits -- Numerically stable softmax (max-subtraction) so large logits stay finite -- Attention weight matrices as row-stochastic distributions over keys +- 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 @@ -54,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. -- **Scaled dot-product self-attention, single head**: `src/attention.py` implements the Vaswani core formula in plain numpy. `scaled_dot_product_attention(Q, K, V)` builds the score matrix `Q Kᵀ / √d_k`, applies an optional boolean mask (False positions become `-inf` so their softmax weight is zero), row-normalizes with a max-stable softmax, and returns both the weighted values and the attention weights. `SelfAttentionHead` learns three projections `W_q`, `W_k`, `W_v` (and an optional `W_o`) from the same input sequence, so every token can attend over the sequence. `causal_mask` builds the lower-triangular mask used in decoder-style autoregressive models. Tests cover shape contracts, uniform-key averaging, near one-hot retrieval when a query matches one key, the entropy effect of the √d_k scale, causal future-blocking, empty and single-token sequences, and a copy-style retrieval check with identity weights. +- **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 @@ -154,26 +156,25 @@ print(forward_variance_profile(10, 64, "relu", "xavier")[-1]) print(forward_variance_profile(6, 64, "linear", "naive")[-1]) ``` -Run single-head self-attention on a short sequence: +Run a Pre-LN transformer block on a short sequence: ```python import numpy as np -from src.attention import SelfAttentionHead, causal_mask, scaled_dot_product_attention +from src.transformer import TransformerBlock, causal_mask, MultiHeadAttention -# bare formula: Q, K, V already projected rng = np.random.default_rng(0) -q = rng.normal(size=(4, 8)) -k = rng.normal(size=(4, 8)) -v = rng.normal(size=(4, 8)) -out, weights = scaled_dot_product_attention(q, k, v) -print(out.shape, weights.shape) # (4, 8), (4, 4); each row of weights sums to 1 - -# self-attention head: same sequence projects to Q, K, V -head = SelfAttentionHead(d_model=16, d_k=8, seed=0) -x = rng.normal(size=(2, 6, 16)) # batch=2, T=6 -result = head.forward(x, mask=causal_mask(6)) -print(result.values.shape) # (2, 6, 16) -print(result.weights[0].sum(-1)) # ~[1, 1, 1, 1, 1, 1] +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: 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 7c1f2b039ad36ec8b1920c6761327cbdfc1947c8 Mon Sep 17 00:00:00 2001 From: Thomas Hart 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