feat: Add scaled dot-product self-attention, single head

README.md

@@ -40,6 +40,11 @@ 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
 
 ## What's implemented
 
@@ -49,6 +54,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.
 
 ## Usage
 
@@ -148,6 +154,28 @@ 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:
+
+```python
+import numpy as np
+from src.attention import SelfAttentionHead, causal_mask, scaled_dot_product_attention
+
+# 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]
+```
+
 Differentiate an arbitrary scalar expression:
 
 ```python

src/attention.py

@@ -0,0 +1,183 @@
+"""Scaled dot-product self-attention, single head.
+
+The attention mechanism (Vaswani et al., 2017) lets each position build a
+weighted average of value vectors, with weights from the similarity of query
+and key vectors. For one head the core is:
+
+    Attention(Q, K, V) = softmax(Q Kᵀ / √d_k) V
+
+The 1/√d_k scale keeps the logits from growing like √d_k when entries are
+unit-variance, which would otherwise push softmax into a near-one-hot regime
+and kill gradients. Self-attention means Q, K, and V are linear projections of
+the same sequence X, so every token can look at every other token (or a
+restricted subset under a mask).
+"""
+
+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 along `axis`."""
+    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 scaled_dot_product_attention(
+    q: Array,
+    k: Array,
+    v: Array,
+    *,
+    mask: Array | None = None,
+    scale: bool = True,
+) -> tuple[Array, Array]:
+    """Compute Attention(Q,K,V) and the attention weight matrix.
+
+    Shapes (trailing dims; leading batch dims are broadcast):
+      q: (..., T_q, d_k)
+      k: (..., T_k, d_k)
+      v: (..., T_k, d_v)
+      mask: broadcastable to (..., T_q, T_k); True/1 keeps the position,
+            False/0 sets its logit to -inf before softmax (causal / padding).
+
+    Returns:
+      output: (..., T_q, d_v)
+      weights: (..., T_q, T_k) rows sum to 1 over the key axis
+    """
+    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 must have at least 2 dims (seq, feature)")
+    if q.shape[-1] != k.shape[-1]:
+        raise ValueError(
+            f"q and k must share d_k, got {q.shape[-1]} vs {k.shape[-1]}"
+        )
+    if k.shape[-2] != v.shape[-2]:
+        raise ValueError(
+            f"k and v must share T_k, got {k.shape[-2]} vs {v.shape[-2]}"
+        )
+
+    d_k = q.shape[-1]
+    # (..., T_q, d_k) @ (..., d_k, T_k) -> (..., T_q, T_k)
+    scores = np.matmul(q, np.swapaxes(k, -1, -2))
+    if scale:
+        if d_k == 0:
+            raise ValueError("d_k must be positive when scale=True")
+        scores = scores / 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)
+        # False positions become -inf so softmax weight is exactly 0
+        scores = np.where(mask_arr, scores, np.float64(-np.inf))
+
+    weights = softmax(scores, axis=-1)
+    # all-masked rows yield 0/0 -> nan; treat as zero mass
+    weights = np.nan_to_num(weights, nan=0.0, posinf=0.0, neginf=0.0)
+    output = np.matmul(weights, v)
+    return output, weights
+
+
+def causal_mask(seq_len: int) -> Array:
+    """Lower-triangular boolean mask: position i may attend to j <= i."""
+    if seq_len < 0:
+        raise ValueError("seq_len must be non-negative")
+    return np.tril(np.ones((seq_len, seq_len), dtype=bool))
+
+
+@dataclass
+class AttentionOutput:
+    """Forward result: projected values plus the weight matrix for inspection."""
+
+    values: Array
+    weights: Array
+
+
+class SelfAttentionHead:
+    """Single-head self-attention: Q, K, V from the same sequence via W_q/k/v.
+
+    X is (..., T, d_model). Optional W_o maps the head output back to d_model.
+    Parameters are plain numpy arrays so training can plug into the same
+    optimizer loop used by the MLP, or stay frozen for pure attention demos.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        d_k: int | None = None,
+        d_v: int | None = None,
+        *,
+        use_output_proj: bool = True,
+        rng: np.random.Generator | None = None,
+        seed: int | None = None,
+    ) -> None:
+        if d_model <= 0:
+            raise ValueError("d_model must be positive")
+        d_k = d_model if d_k is None else d_k
+        d_v = d_k if d_v is None else d_v
+        if d_k <= 0 or d_v <= 0:
+            raise ValueError("d_k and d_v must be positive")
+
+        if rng is None:
+            rng = np.random.default_rng(seed)
+
+        # Xavier-style scale for the linear maps into the head
+        def _init(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)
+
+        self.d_model = d_model
+        self.d_k = d_k
+        self.d_v = d_v
+        self.use_output_proj = use_output_proj
+
+        self.W_q = _init(d_model, d_k)
+        self.W_k = _init(d_model, d_k)
+        self.W_v = _init(d_model, d_v)
+        self.W_o: Array | None
+        if use_output_proj:
+            self.W_o = _init(d_v, d_model)
+        else:
+            self.W_o = None
+
+    def parameters(self) -> list[Array]:
+        params = [self.W_q, self.W_k, self.W_v]
+        if self.W_o is not None:
+            params.append(self.W_o)
+        return params
+
+    def forward(
+        self,
+        x: Array,
+        *,
+        mask: Array | None = None,
+        scale: bool = True,
+    ) -> AttentionOutput:
+        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 = np.matmul(x, self.W_q)
+        k = np.matmul(x, self.W_k)
+        v = np.matmul(x, self.W_v)
+        out, weights = scaled_dot_product_attention(q, k, v, mask=mask, scale=scale)
+        if self.W_o is not None:
+            out = np.matmul(out, self.W_o)
+        return AttentionOutput(values=out, weights=weights)