add kolm.py

kolm.py

@@ -0,0 +1,362 @@
+#!/usr/bin/env python3
+"""
+KOLM — Kuramoto Oscillator Language Model (OLM v4).
+
+v3 (olm.py) trains damped *amplitude* oscillators (coRNN). v4 swaps the core
+for Kuramoto oscillators, AKOrN-style (Miyato et al., ICLR 2025): each hidden
+unit is a unit vector x_i on the sphere S^{N-1}. Units interact through one
+trained scalar coupling J_ij per pair and synchronize; binding-by-synchrony
+means units that align are bound into one active concept, and the same units
+re-align differently for another concept.
+
+One char step (all trained: J, A, U, Wo, bo):
+
+    y_i = sum_j J_ij x_j + U[char]_i              coupling + input drive
+    P_i = y_i - (x_i . y_i) x_i                   project onto tangent space
+    x_i <- normalize(x_i + dt (A_i x_i + P_i))    A_i antisymmetric: rotation
+
+Unit norm is enforced by construction, so stability is geometric — no
+gamma/eps clipping. Same zero-pretraining setup as v3: learns one novel as
+'?keyword;sentence' lines; chat picks a keyword and generates conditioned
+on it. Pure NumPy, hand-derived BPTT, ~141k params at H=256 N=4.
+
+    python3 kolm.py --gradcheck
+    python3 kolm.py --train --save kolm256.npz
+    python3 kolm.py --load kolm256.npz
+"""
+
+import argparse
+import sys
+import time
+import numpy as np
+
+from olm import CHARS, C2I, V, F32, build_corpus, pick_keyword
+
+
+# ---------------------------------------------------------------- model
+def init_params(H: int, N: int, seed: int = 0, om_init: float = 0.1,
+                groups: int = 0):
+    rng = np.random.default_rng(seed)
+    p = {
+        "J":  rng.normal(0, 1.0 / np.sqrt(H), (H, H)).astype(F32),
+        "Om": rng.normal(0, om_init, (H, N, N)).astype(F32),  # A = Om - Om^T
+        "U":  rng.normal(0, 0.5, (V, H, N)).astype(F32),
+        "Wo": rng.normal(0, 1.0 / np.sqrt(H * N),
+                         (H * N + groups, V)).astype(F32),
+        "bo": np.zeros(V, F32),
+    }
+    if groups:
+        # coherence readout: r_g = ||sum_i M_gi x_i|| — a trained order
+        # parameter, large only when the units row g weights are in sync
+        p["M"] = rng.normal(0, 1.0 / np.sqrt(H), (groups, H)).astype(F32)
+    return p
+
+
+def init_state(B: int, H: int, N: int, dtype=F32):
+    x = np.zeros((B, H, N), dtype)
+    x[..., 0] = 1.0        # all units start at the same pole; A and U disperse them
+    return x
+
+
+def antisym(Om):
+    return Om - Om.transpose(0, 2, 1)
+
+
+def couple(J, x):
+    """y[b,i,:] = sum_j J[i,j] x[b,j,:] as one BLAS matmul."""
+    B, H, N = x.shape
+    return (J @ x.transpose(1, 0, 2).reshape(H, B * N)) \
+        .reshape(-1, B, N).transpose(1, 0, 2)
+
+
+def readout_feats(p, x):
+    """[flat phases, coherence r] and the population vectors m (or None)."""
+    B, H, N = x.shape
+    if "M" not in p:
+        return x.reshape(B, H * N), None, None
+    m = couple(p["M"], x)                                # (B,G,N)
+    r = np.sqrt((m * m).sum(-1) + 1e-8)                  # (B,G)
+    return np.concatenate([x.reshape(B, H * N), r], axis=1), m, r
+
+
+def forward(p, x, X, Y, dt):
+    """One BPTT window. X,Y: (B,T) int ids. Returns loss, caches, final state."""
+    B, T = X.shape
+    H, N = p["Om"].shape[0], p["Om"].shape[1]
+    A = antisym(p["Om"])
+    caches, loss = [], 0.0
+    for t in range(T):
+        x0 = x
+        y = couple(p["J"], x0) + p["U"][X[:, t]]
+        d = (x0 * y).sum(-1, keepdims=True)
+        P = y - d * x0
+        R = np.einsum("inm,bim->bin", A, x0)
+        xt = x0 + dt * (R + P)
+        nrm = np.sqrt((xt * xt).sum(-1, keepdims=True))
+        x = xt / nrm
+        f, m, r = readout_feats(p, x)
+        logits = f @ p["Wo"] + p["bo"]
+        logits -= logits.max(axis=1, keepdims=True)
+        e = np.exp(logits)
+        probs = e / e.sum(axis=1, keepdims=True)
+        loss -= np.log(probs[np.arange(B), Y[:, t]] + 1e-9).sum()
+        caches.append((x0, y, d, nrm, x, f, m, r, probs, X[:, t], Y[:, t]))
+    return loss / (B * T), caches, x
+
+
+def backward(p, caches, dt):
+    B = caches[0][0].shape[0]
+    T = len(caches)
+    H, N = p["Om"].shape[0], p["Om"].shape[1]
+    A = antisym(p["Om"])
+    g = {k: np.zeros_like(v) for k, v in p.items()}
+    gA = np.zeros_like(A)
+    gx = np.zeros_like(caches[0][0])
+    for t in reversed(range(T)):
+        x0, y, d, nrm, x1, f, m, r, probs, xs, ys = caches[t]
+        # output head
+        dlog = probs.copy()
+        dlog[np.arange(B), ys] -= 1.0
+        dlog /= (B * T)
+        g["Wo"] += f.T @ dlog
+        g["bo"] += dlog.sum(0)
+        gf = dlog @ p["Wo"].T
+        gx1 = gx + gf[:, :H * N].reshape(B, H, N)
+        if m is not None:
+            # r = ||m||, m_g = sum_i M_gi x1_i
+            gm = (gf[:, H * N:] / r)[:, :, None] * m
+            gm_r = gm.transpose(1, 0, 2).reshape(-1, B * N)
+            x1_r = x1.transpose(1, 0, 2).reshape(H, B * N)
+            g["M"] += gm_r @ x1_r.T
+            gx1 += (p["M"].T @ gm_r).reshape(H, B, N).transpose(1, 0, 2)
+        # x1 = xt / ||xt||  =>  gxt = (I - x1 x1^T) gx1 / ||xt||
+        gxt = (gx1 - (gx1 * x1).sum(-1, keepdims=True) * x1) / nrm
+        # xt = x0 + dt*(R + P)
+        gx0 = gxt.copy()
+        gR = dt * gxt
+        gP = dt * gxt
+        # R_i = A_i x0_i
+        gA += np.einsum("bin,bim->inm", gR, x0)
+        gx0 += np.einsum("inm,bin->bim", A, gR)
+        # P = y - (x0.y) x0
+        gPdot = (gP * x0).sum(-1, keepdims=True)
+        gy = gP - gPdot * x0
+        gx0 -= gPdot * y + d * gP
+        # y = J x0 + U[xs]
+        gy_r = gy.transpose(1, 0, 2).reshape(H, B * N)
+        x0_r = x0.transpose(1, 0, 2).reshape(H, B * N)
+        g["J"] += gy_r @ x0_r.T
+        gx0 += (p["J"].T @ gy_r).reshape(H, B, N).transpose(1, 0, 2)
+        np.add.at(g["U"], xs, gy)
+        gx = gx0
+    g["Om"] = gA - gA.transpose(0, 2, 1)
+    return g
+
+
+def gradcheck():
+    """Finite-difference check of the hand-derived BPTT gradients."""
+    np.random.seed(1)
+    H, N, B, T, dt = 5, 3, 2, 4, 0.5
+    p = {k: v.astype(np.float64)
+         for k, v in init_params(H, N, seed=3, groups=3).items()}
+    X = np.random.randint(0, V, (B, T))
+    Y = np.random.randint(0, V, (B, T))
+    x = np.random.randn(B, H, N)
+    x /= np.sqrt((x * x).sum(-1, keepdims=True))
+    loss, caches, _ = forward(p, x, X, Y, dt)
+    g = backward(p, caches, dt)
+    worst = 0.0
+    for name in p:
+        flat = p[name].reshape(-1)
+        for idx in np.random.choice(flat.size, min(6, flat.size), replace=False):
+            eps_ = 1e-5
+            old = flat[idx]
+            flat[idx] = old + eps_
+            lp, _, _ = forward(p, x, X, Y, dt)
+            flat[idx] = old - eps_
+            lm, _, _ = forward(p, x, X, Y, dt)
+            flat[idx] = old
+            num = (lp - lm) / (2 * eps_)
+            ana = g[name].reshape(-1)[idx]
+            rel = abs(num - ana) / max(1e-8, abs(num) + abs(ana))
+            worst = max(worst, rel)
+    print(f"gradcheck worst relative error: {worst:.2e} "
+          f"({'PASS' if worst < 1e-4 else 'FAIL'})")
+    return worst < 1e-4
+
+
+# ---------------------------------------------------------------- training
+def adam_step(p, g, m, v, t, lr, clip=1.0):
+    norm = np.sqrt(sum(float((gi ** 2).sum()) for gi in g.values()))
+    scale = min(1.0, clip / (norm + 1e-8))
+    b1, b2, e = 0.9, 0.999, 1e-8
+    for k in p:
+        gk = g[k] * scale
+        m[k] = b1 * m[k] + (1 - b1) * gk
+        v[k] = b2 * v[k] + (1 - b2) * gk * gk
+        mh = m[k] / (1 - b1 ** t)
+        vh = v[k] / (1 - b2 ** t)
+        p[k] -= (lr * mh / (np.sqrt(vh) + e)).astype(F32)
+    return norm
+
+
+def train(args):
+    corpus, freq, kwc = build_corpus(args.text)
+    ids = np.array([C2I[c] for c in corpus], dtype=np.int64)
+    n_val = 20_000
+    tr, va = ids[:-n_val], ids[-n_val:]
+    B, T, H, N, dt = args.batch, args.seq, args.hidden, args.ndim, args.dt
+
+    p = init_params(H, N, om_init=args.om_init, groups=args.groups)
+    n_params = sum(x.size for x in p.values())
+    print(f"KOLM v4 | hidden {H} x S^{N - 1} | params {n_params:,} "
+          f"({n_params * 4 / 1e6:.2f} MB) | corpus {len(ids):,} chars", flush=True)
+
+    L = len(tr) // B
+    streams = tr[: B * L].reshape(B, L)
+    m = {k: np.zeros_like(x) for k, x in p.items()}
+    v = {k: np.zeros_like(x) for k, x in p.items()}
+    step = 0
+    for ep in range(1, args.epochs + 1):
+        lr = args.lr * (0.5 ** max(0, ep - args.epochs + 6) if ep > args.epochs - 6 else 1.0)
+        x = init_state(B, H, N)
+        tot = nb = 0
+        t0 = time.time()
+        for s in range(0, L - T - 1, T):
+            X, Y = streams[:, s:s + T], streams[:, s + 1:s + T + 1]
+            loss, caches, x = forward(p, x, X, Y, dt)
+            g = backward(p, caches, dt)
+            step += 1
+            adam_step(p, g, m, v, step, lr)
+            tot += loss
+            nb += 1
+        vl, vacc = evaluate(p, va, dt)
+        print(f"epoch {ep:2d} | train loss {tot / nb:.3f} | "
+              f"val loss {vl:.3f} | val acc {vacc:.1%} | "
+              f"{time.time() - t0:.0f}s", flush=True)
+        if ep % 4 == 0 or ep == args.epochs:
+            print("  sample:", generate(p, dt, "?whale;", seed=ep)[:110], flush=True)
+    save(args.save, p, dt, freq, kwc)
+    print(f"saved to {args.save}", flush=True)
+
+
+def evaluate(p, ids, dt, B=50):
+    H, N = p["Om"].shape[0], p["Om"].shape[1]
+    L = len(ids) // B
+    st = ids[: B * L].reshape(B, L)
+    x = init_state(B, H, N)
+    loss, caches, _ = forward(p, x, st[:, :-1][:, :400], st[:, 1:][:, :400], dt)
+    hits = tot = 0
+    for cache in caches[50:]:   # skip washout
+        probs, ys = cache[-3], cache[-1]
+        hits += int((probs.argmax(1) == ys).sum())
+        tot += len(ys)
+    return loss, hits / tot
+
+
+# ---------------------------------------------------------------- generation
+def step_one(p, x, cid, dt, A):
+    x0 = x
+    y = couple(p["J"], x0) + p["U"][cid]
+    d = (x0 * y).sum(-1, keepdims=True)
+    R = np.einsum("inm,bim->bin", A, x0)
+    xt = x0 + dt * (R + y - d * x0)
+    x = xt / np.sqrt((xt * xt).sum(-1, keepdims=True))
+    f, _, _ = readout_feats(p, x)
+    return x, f @ p["Wo"] + p["bo"]
+
+
+def generate(p, dt, prime, max_len=300, temperature=0.8, top_k=8, seed=None):
+    rng = np.random.default_rng(seed)
+    H, N = p["Om"].shape[0], p["Om"].shape[1]
+    A = antisym(p["Om"])
+    x = init_state(1, H, N)
+    logits = None
+    for ch in prime:
+        x, logits = step_one(p, x, np.array([C2I.get(ch, 1)]), dt, A)
+    out = []
+    for _ in range(max_len):
+        lg = logits[0] / temperature
+        lg -= lg.max()
+        pr = np.exp(lg)
+        if top_k and top_k < V:
+            cut = np.partition(pr, -top_k)[-top_k]
+            pr = np.where(pr >= cut, pr, 0.0)
+        pr /= pr.sum()
+        nxt = int(rng.choice(V, p=pr))
+        if CHARS[nxt] == "\n":
+            break
+        out.append(CHARS[nxt])
+        x, logits = step_one(p, x, np.array([nxt]), dt, A)
+    return "".join(out)
+
+
+# ---------------------------------------------------------------- persistence
+def save(path, p, dt, freq, kwc):
+    words = np.array(list(freq.keys()))
+    counts = np.array(list(freq.values()))
+    kwords = np.array(list(kwc.keys()))
+    kcounts = np.array(list(kwc.values()))
+    np.savez_compressed(path, dt=dt, words=words, counts=counts,
+                        kwords=kwords, kcounts=kcounts,
+                        **{f"p_{k}": x for k, x in p.items()})
+
+
+def load(path):
+    zf = np.load(path)
+    p = {k[2:]: zf[k] for k in zf.files if k.startswith("p_")}
+    freq = dict(zip(zf["words"].tolist(), zf["counts"].tolist()))
+    kwc = dict(zip(zf["kwords"].tolist(), zf["kcounts"].tolist()))
+    return p, float(zf["dt"]), freq, kwc
+
+
+# ---------------------------------------------------------------- main
+def main():
+    ap = argparse.ArgumentParser(description="KOLM v4 — Kuramoto oscillators")
+    ap.add_argument("--text", default="mobydick.txt")
+    ap.add_argument("--hidden", type=int, default=256)
+    ap.add_argument("--ndim", type=int, default=4, help="oscillator dimension N")
+    ap.add_argument("--seq", type=int, default=128)
+    ap.add_argument("--batch", type=int, default=64)
+    ap.add_argument("--dt", type=float, default=0.25)
+    ap.add_argument("--om-init", type=float, default=0.1,
+                    help="init scale of Om; with dt sets the gradient horizon")
+    ap.add_argument("--groups", type=int, default=32,
+                    help="coherence readout groups G (0 disables)")
+    ap.add_argument("--lr", type=float, default=2e-3)
+    ap.add_argument("--epochs", type=int, default=18)
+    ap.add_argument("--train", action="store_true")
+    ap.add_argument("--save", default="kolm256.npz")
+    ap.add_argument("--load", default=None)
+    ap.add_argument("--ask", default=None, help="one-shot question")
+    ap.add_argument("--gradcheck", action="store_true")
+    args = ap.parse_args()
+
+    if args.gradcheck:
+        sys.exit(0 if gradcheck() else 1)
+    if args.train:
+        train(args)
+        if not args.ask:
+            return
+    p, dt, freq, kwc = load(args.load or args.save)
+
+    def answer(msg):
+        kw = pick_keyword(msg, freq, kwc)
+        return kw, generate(p, dt, f"?{kw};")
+
+    if args.ask:
+        kw, resp = answer(args.ask)
+        print(f"you : {args.ask}\nkolm ({kw}) : {resp}")
+        return
+    print("KOLM chat — it answers about the topic word of your message (ctrl-d quits)")
+    while True:
+        try:
+            msg = input("\nyou > ")
+        except EOFError:
+            break
+        kw, resp = answer(msg)
+        print(f"kolm ({kw}) > {resp}")
+
+
+if __name__ == "__main__":
+    main()