feat: Add skip-gram word embeddings with nearest-neighbor demo

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
+- Distributional hypothesis: word meaning from co-occurrence context
+- Skip-gram language modeling: predict context words given a center word inside a sliding window
+- Two-matrix word2vec parameterization (center embeddings W_in, context embeddings W_out)
+- Full-softmax training objective with the `softmax − onehot` gradient identity
+- Dense word vectors and cosine nearest-neighbor retrieval as a semantic similarity demo
 
 ## 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.
+- **Word embeddings (skip-gram) trained on a small corpus, nearest-neighbor demo**: `src/embeddings.py` learns dense vectors by predicting each window neighbor of a center word (skip-gram). Two matrices store center and context embeddings; the score for context o given center c is the dot product `W_out[o] · W_in[c]`, turned into a distribution with a full softmax over the vocabulary. Minibatch SGD minimizes the mean negative log-likelihood, using the closed-form `p − onehot` gradient on the logits (verified against finite differences). After training, `nearest(word, k)` ranks the rest of the vocab by cosine similarity on the center rows. A built-in toy corpus repeats short sentences about capitals, animals, and people so co-occurrence clusters show up in the neighbor lists without any external data.
 
 ## Usage
 
@@ -128,6 +134,18 @@ for name, h in curves.items():
     print(name, h[0], "->", h[-1])
 ```
 
+Train skip-gram embeddings and query nearest neighbors:
+
+```python
+from src.embeddings import fit, make_toy_corpus, cosine_similarity
+
+model, history = fit(make_toy_corpus(), dim=32, window=2, epochs=100, seed=0)
+print(history[0], "->", history[-1])          # NLL falls over training
+print(model.nearest("paris", k=4))            # other capital cities rise
+print(model.nearest("cat", k=4))              # dog / mouse / cats nearby
+print(cosine_similarity(model.embed("king"), model.embed("queen")))
+```
+
 Compare init schemes by watching activation variance with depth:
 
 ```python

src/embeddings.py

@@ -0,0 +1,256 @@
+"""Skip-gram word embeddings trained from scratch.
+
+The distributional hypothesis says a word's meaning is the company it keeps.
+Skip-gram turns that into a supervised task: given a center word, predict each
+nearby context word inside a sliding window. After training, each vocabulary
+word is a dense vector (a row of the input embedding matrix). Cosine nearest
+neighbors recover words that shared similar contexts.
+
+The model keeps two matrices, matching the original word2vec form: W_in (V x D)
+embeds the center word and W_out (V x D) scores every context candidate. The
+probability is a full softmax over the vocabulary, p(o|c) ∝ exp(W_out[o] · W_in[c]).
+That is exact and fine for the small corpora here; production systems replace the
+softmax with negative sampling so each step is O(k) instead of O(V).
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import NDArray
+
+Array = NDArray[np.float64]
+
+_TOKEN = re.compile(r"[a-z]+")
+
+
+def tokenize(text: str) -> list[str]:
+    """Lowercase alphabetic tokens; punctuation and digits are dropped."""
+    return _TOKEN.findall(text.lower())
+
+
+def build_vocab(
+    tokens: list[str], min_count: int = 1
+) -> tuple[dict[str, int], list[str]]:
+    if min_count < 1:
+        raise ValueError("min_count must be at least 1")
+    counts: dict[str, int] = {}
+    for t in tokens:
+        counts[t] = counts.get(t, 0) + 1
+    words = sorted(w for w, c in counts.items() if c >= min_count)
+    if not words:
+        raise ValueError("vocabulary is empty after min_count filter")
+    word_to_idx = {w: i for i, w in enumerate(words)}
+    return word_to_idx, words
+
+
+def skipgram_pairs(
+    token_ids: list[int], window: int
+) -> tuple[NDArray[np.int64], NDArray[np.int64]]:
+    """Emit (center, context) id pairs for every token inside the window.
+
+    The window is symmetric and clipped at sentence ends when callers pass one
+    sentence at a time; a flat corpus just uses sequence edges.
+    """
+    if window < 1:
+        raise ValueError("window must be at least 1")
+    centers: list[int] = []
+    contexts: list[int] = []
+    n = len(token_ids)
+    for i, c in enumerate(token_ids):
+        lo = max(0, i - window)
+        hi = min(n, i + window + 1)
+        for j in range(lo, hi):
+            if j == i:
+                continue
+            centers.append(c)
+            contexts.append(token_ids[j])
+    if not centers:
+        return (
+            np.zeros(0, dtype=np.int64),
+            np.zeros(0, dtype=np.int64),
+        )
+    return (
+        np.asarray(centers, dtype=np.int64),
+        np.asarray(contexts, dtype=np.int64),
+    )
+
+
+def _softmax(z: Array) -> Array:
+    z = z - z.max(axis=-1, keepdims=True)
+    ez = np.exp(z)
+    return ez / ez.sum(axis=-1, keepdims=True)
+
+
+def cosine_similarity(a: Array, b: Array) -> float:
+    a = np.asarray(a, dtype=np.float64).reshape(-1)
+    b = np.asarray(b, dtype=np.float64).reshape(-1)
+    if a.shape != b.shape:
+        raise ValueError(f"shape mismatch: {a.shape} vs {b.shape}")
+    na = float(np.linalg.norm(a))
+    nb = float(np.linalg.norm(b))
+    if na == 0.0 or nb == 0.0:
+        return 0.0
+    return float(np.dot(a, b) / (na * nb))
+
+
+@dataclass
+class SkipGramModel:
+    """Trained skip-gram: vocab maps plus center and context embedding tables."""
+
+    word_to_idx: dict[str, int]
+    idx_to_word: list[str]
+    W_in: Array
+    W_out: Array
+
+    @property
+    def dim(self) -> int:
+        return int(self.W_in.shape[1])
+
+    @property
+    def vocab_size(self) -> int:
+        return int(self.W_in.shape[0])
+
+    def embed(self, word: str) -> Array:
+        if word not in self.word_to_idx:
+            raise KeyError(f"unknown word: {word!r}")
+        return self.W_in[self.word_to_idx[word]].copy()
+
+    def nearest(self, word: str, k: int = 5) -> list[tuple[str, float]]:
+        """Top-k cosine neighbors of `word` among the rest of the vocabulary."""
+        if k < 1:
+            raise ValueError("k must be at least 1")
+        if word not in self.word_to_idx:
+            raise KeyError(f"unknown word: {word!r}")
+        query = self.W_in[self.word_to_idx[word]]
+        qn = float(np.linalg.norm(query))
+        if qn == 0.0:
+            return []
+        norms = np.linalg.norm(self.W_in, axis=1)
+        # zero vectors contribute nothing useful; skip them and the query itself
+        scores = (self.W_in @ query) / np.maximum(norms * qn, 1e-12)
+        scores[self.word_to_idx[word]] = -np.inf
+        scores[norms == 0.0] = -np.inf
+        k_eff = min(k, self.vocab_size - 1)
+        if k_eff < 1:
+            return []
+        # partial sort is enough; vocab stays small in this module
+        top = np.argpartition(-scores, k_eff - 1)[:k_eff]
+        top = top[np.argsort(-scores[top])]
+        return [(self.idx_to_word[int(i)], float(scores[i])) for i in top]
+
+
+def make_toy_corpus() -> str:
+    """A tiny multi-sentence corpus with a few clear co-occurrence clusters.
+
+    City names share the "capital of" pattern; animals share sit/chase patterns;
+    royalty and people share ruled/walked patterns. Skip-gram should pull those
+    groups closer than unrelated words after enough epochs.
+    """
+    sentences = [
+        "paris is the capital of france",
+        "berlin is the capital of germany",
+        "london is the capital of england",
+        "rome is the capital of italy",
+        "madrid is the capital of spain",
+        "the cat sat on the mat",
+        "the dog sat on the rug",
+        "the cat chased the mouse",
+        "the dog chased the cat",
+        "cats and dogs live in the house",
+        "the king ruled the kingdom with the queen",
+        "the queen ruled the kingdom with the king",
+        "a man walked with a woman in the park",
+        "a woman walked with a man in the park",
+        "the boy played with the girl in the yard",
+        "the girl played with the boy in the yard",
+    ]
+    # repeat so rare content words see enough center/context pairs
+    return " . ".join(sentences * 8)
+
+
+def fit(
+    text: str,
+    dim: int = 32,
+    window: int = 2,
+    lr: float = 0.05,
+    epochs: int = 80,
+    batch_size: int = 64,
+    min_count: int = 1,
+    seed: int = 0,
+) -> tuple[SkipGramModel, list[float]]:
+    """Train skip-gram on `text` with minibatch SGD; return model and loss curve.
+
+    Loss is mean negative log-likelihood of the context word under the full
+    softmax. Parameters update from the closed-form softmax gradient
+    `p - onehot(context)`.
+    """
+    if dim < 1:
+        raise ValueError("dim must be at least 1")
+    if lr <= 0.0:
+        raise ValueError("lr must be positive")
+    if epochs < 1:
+        raise ValueError("epochs must be at least 1")
+    if batch_size < 1:
+        raise ValueError("batch_size must be at least 1")
+
+    tokens = tokenize(text)
+    if len(tokens) < 2:
+        raise ValueError("need at least two tokens to form a skip-gram pair")
+
+    word_to_idx, idx_to_word = build_vocab(tokens, min_count=min_count)
+    ids = [word_to_idx[t] for t in tokens if t in word_to_idx]
+    centers, contexts = skipgram_pairs(ids, window)
+    if centers.size == 0:
+        raise ValueError("no skip-gram pairs produced; try a larger window or corpus")
+
+    v = len(idx_to_word)
+    rng = np.random.default_rng(seed)
+    # small Gaussian so early softmax is not saturated
+    scale = 0.1
+    W_in = rng.standard_normal((v, dim)) * scale
+    W_out = rng.standard_normal((v, dim)) * scale
+
+    n = centers.shape[0]
+    history: list[float] = []
+    for _ in range(epochs):
+        order = rng.permutation(n)
+        total_loss = 0.0
+        seen = 0
+        for start in range(0, n, batch_size):
+            idx = order[start : start + batch_size]
+            c = centers[idx]
+            o = contexts[idx]
+            m = c.shape[0]
+
+            h = W_in[c]  # (m, D)
+            scores = h @ W_out.T  # (m, V)
+            p = _softmax(scores)
+            # NLL of the true context id
+            rows = np.arange(m)
+            total_loss += float(-np.sum(np.log(np.clip(p[rows, o], 1e-12, 1.0))))
+            seen += m
+
+            # dL/dscores = p - onehot(o); average over the batch
+            ds = p
+            ds[rows, o] -= 1.0
+            ds /= m
+
+            dW_out = ds.T @ h  # (V, D)
+            dh = ds @ W_out  # (m, D)
+
+            W_out -= lr * dW_out
+            # scatter-add dh into the center rows that appeared in the batch
+            np.add.at(W_in, c, -lr * dh)
+
+        history.append(total_loss / max(seen, 1))
+
+    model = SkipGramModel(
+        word_to_idx=word_to_idx,
+        idx_to_word=idx_to_word,
+        W_in=W_in,
+        W_out=W_out,
+    )
+    return model, history