Upload janus/janus4_temporal_diff.py with huggingface_hub

janus/janus4_temporal_diff.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+"""
+janus4_temporal_diff.py — Janus 4-way attention reference implementation.
+
+4th attention mechanism: Temporal Diff — attends to CHANGES between positions.
+Based on Variant 4 (dedicated wtd+wvd) + Opus analysis fixes:
+  - Removed distance decay (RoPE handles this in full nanochat)
+  - Gate init biased against delta (-1.0) — model discovers if/when to use it
+  - Dedicated projections (no weight sharing with QKV or RRPRAM)
+
+Architecture: QKV (semantic) + RRPRAM (positional) + Echo (self-resonance) + TemporalDiff (change detection)
+
+What TemporalDiff captures that others don't:
+  - QKV with RoPE: encodes distance between positions, not content of change
+  - RRPRAM: positional patterns, not transitions
+  - Echo: self-similarity, not change rate
+  - TemporalDiff: "where did representation change, and do changes correlate?"
+
+Pure Python, zero deps. For reference/testing. Production = C or PyTorch.
+
+By Arianna Method, 2026-03-25.
+"""
+
+import argparse
+import math
+import random
+
+VOCAB = 256
+MAX_T = 48
+DIM = 48
+HEADS = 4
+HD = DIM // HEADS
+
+
+def bpe_encode(text):
+    return list(text.encode('utf-8', errors='ignore'))
+
+
+def bpe_decode(ids):
+    return bytes([i % 256 for i in ids]).decode('utf-8', errors='ignore')
+
+
+def rand_mat(r, c, s=0.02):
+    return [[(random.random() * 2 - 1) * s for _ in range(c)] for _ in range(r)]
+
+
+def vec_mat(v, m):
+    out = [0.0] * len(m[0])
+    for i, vi in enumerate(v):
+        row = m[i]
+        for j in range(len(out)):
+            out[j] += vi * row[j]
+    return out
+
+
+def softmax(xs):
+    mx = max(xs)
+    ex = [math.exp(x - mx) for x in xs]
+    s = sum(ex) + 1e-9
+    return [x / s for x in ex]
+
+
+class Janus4:
+    def __init__(self):
+        self.tok = rand_mat(VOCAB, DIM)
+        self.pos = rand_mat(MAX_T, DIM)
+        # QKV — semantic attention
+        self.wq = rand_mat(DIM, DIM)
+        self.wk = rand_mat(DIM, DIM)
+        self.wv = rand_mat(DIM, DIM)
+        # RRPRAM — positional resonance
+        self.wr = rand_mat(DIM, MAX_T)
+        self.wvr = rand_mat(DIM, DIM)
+        # Echo — self-resonance (W^T * W)
+        self.wj = rand_mat(DIM, DIM)
+        # Temporal Diff — dedicated projections (NOT shared with QKV/RRPRAM)
+        self.wtd = rand_mat(DIM, DIM)   # delta key projection
+        self.wvd = rand_mat(DIM, DIM)   # delta value projection
+        # 4-way gate — delta starts suppressed (Opus recommendation)
+        self.gate = [0.0, 0.0, 0.0, -1.0]  # QKV, RRPRAM, Echo, TemporalDiff
+        # Output
+        self.out = rand_mat(DIM, VOCAB)
+        self.bias = [0.0] * VOCAB
+
+    def _dot(self, a, b):
+        return sum(x * y for x, y in zip(a, b))
+
+    def _head(self, v, h):
+        return v[h * HD:(h + 1) * HD]
+
+    def forward(self, ids):
+        T = len(ids)
+        x = [[self.tok[ids[t]][e] + self.pos[t][e] for e in range(DIM)] for t in range(T)]
+
+        # Precompute all projections
+        q = [vec_mat(x[t], self.wq) for t in range(T)]
+        k = [vec_mat(x[t], self.wk) for t in range(T)]
+        v = [vec_mat(x[t], self.wv) for t in range(T)]
+        rv = [vec_mat(x[t], self.wvr) for t in range(T)]
+        je = [vec_mat(x[t], self.wj) for t in range(T)]
+
+        # Temporal diff: delta of input
+        dx = [[0.0] * DIM for _ in range(T)]
+        for t in range(1, T):
+            for e in range(DIM):
+                dx[t][e] = x[t][e] - x[t - 1][e]
+
+        # Dedicated projections for delta (not shared!)
+        dk = [vec_mat(dx[t], self.wtd) for t in range(T)]  # delta keys
+        dv = [vec_mat(dx[t], self.wvd) for t in range(T)]  # delta values
+
+        g = softmax(self.gate)
+
+        cat = [[0.0] * DIM for _ in range(T)]
+        for h in range(HEADS):
+            # 1) QKV attention — semantic content matching
+            a1 = [[-1e9] * T for _ in range(T)]
+            for i in range(T):
+                qi = self._head(q[i], h)
+                for j in range(i + 1):
+                    a1[i][j] = self._dot(qi, self._head(k[j], h)) / math.sqrt(HD)
+                a1[i] = softmax(a1[i])
+            ho = [[0.0] * HD for _ in range(T)]
+            for i in range(T):
+                for j in range(T):
+                    vv = self._head(v[j], h)
+                    for d in range(HD):
+                        ho[i][d] += a1[i][j] * vv[d]
+
+            # 2) RRPRAM — positional resonance
+            a2 = [[-1e9] * T for _ in range(T)]
+            for i in range(T):
+                for j in range(i + 1):
+                    a2[i][j] = sum(x[i][e] * self.wr[e][j] for e in range(DIM)) / math.sqrt(HD)
+                a2[i] = softmax(a2[i])
+            ro = [[0.0] * HD for _ in range(T)]
+            for i in range(T):
+                for j in range(T):
+                    rvh = self._head(rv[j], h)
+                    for d in range(HD):
+                        ro[i][d] += a2[i][j] * rvh[d]
+
+            # 3) Echo — self-resonance (W^T * W)
+            a3 = [[-1e9] * T for _ in range(T)]
+            for i in range(T):
+                ei = self._head(je[i], h)
+                for j in range(i + 1):
+                    a3[i][j] = self._dot(ei, self._head(je[j], h)) / math.sqrt(HD)
+                a3[i] = softmax(a3[i])
+            jo = [[0.0] * HD for _ in range(T)]
+            for i in range(T):
+                for j in range(T):
+                    ej = self._head(je[j], h)
+                    for d in range(HD):
+                        jo[i][d] += a3[i][j] * ej[d]
+
+            # 4) Temporal Diff — change detection attention
+            # No distance decay (Opus fix: RoPE handles this in full implementation)
+            a4 = [[-1e9] * T for _ in range(T)]
+            for i in range(T):
+                dki = self._head(dk[i], h)
+                for j in range(i + 1):
+                    a4[i][j] = self._dot(dki, self._head(dk[j], h)) / math.sqrt(HD)
+                a4[i] = softmax(a4[i])
+            to = [[0.0] * HD for _ in range(T)]
+            for i in range(T):
+                for j in range(T):
+                    dvh = self._head(dv[j], h)
+                    for d in range(HD):
+                        to[i][d] += a4[i][j] * dvh[d]
+
+            # Gate blend — 4-way softmax
+            for t in range(T):
+                base = h * HD
+                for d in range(HD):
+                    cat[t][base + d] = (g[0] * ho[t][d] + g[1] * ro[t][d] +
+                                        g[2] * jo[t][d] + g[3] * to[t][d])
+
+        logits = [[0.0] * VOCAB for _ in range(T)]
+        for t in range(T):
+            for vi in range(VOCAB):
+                logits[t][vi] = sum(cat[t][e] * self.out[e][vi] for e in range(DIM)) + self.bias[vi]
+        return logits, cat
+
+    def train_step(self, tok, tgt, lr):
+        logits, cat = self.forward(tok)
+        loss = 0.0
+        grad = [[0.0] * VOCAB for _ in range(len(tok))]
+        for t in range(len(tok)):
+            p = softmax(logits[t])
+            loss -= math.log(max(1e-9, p[tgt[t]]))
+            for vi in range(VOCAB):
+                grad[t][vi] = p[vi]
+            grad[t][tgt[t]] -= 1.0
+        loss /= len(tok)
+        # Gradient on output layer only (reference impl)
+        for t in range(len(tok)):
+            for e in range(DIM):
+                ce = cat[t][e]
+                if ce == 0.0:
+                    continue
+                row = self.out[e]
+                for vi in range(VOCAB):
+                    row[vi] -= lr * ce * grad[t][vi] / len(tok)
+            for vi in range(VOCAB):
+                self.bias[vi] -= lr * grad[t][vi] / len(tok)
+        return loss
+
+
+def generate(model, prompt, n=60):
+    ids = bpe_encode(prompt)[-MAX_T:]
+    for _ in range(n):
+        logits, _ = model.forward(ids)
+        p = softmax(logits[-1])
+        ids.append(max(range(VOCAB), key=lambda i: p[i]))
+        ids = ids[-MAX_T:]
+    return bpe_decode(ids)
+
+
+def train(model, text, steps, lr):
+    ids = bpe_encode(text)
+    losses = []
+    for step in range(1, steps + 1):
+        off = random.randint(0, max(0, len(ids) - MAX_T - 2))
+        tok = ids[off:off + MAX_T]
+        tgt = ids[off + 1:off + MAX_T + 1]
+        losses.append(model.train_step(tok, tgt, lr))
+        if step % 10 == 0:
+            print(f"step {step:4d}/{steps} loss={losses[-1]:.4f}")
+    return losses
+
+
+if __name__ == '__main__':
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--train', type=str)
+    ap.add_argument('--steps', type=int, default=40)
+    ap.add_argument('--lr', type=float, default=0.05)
+    ap.add_argument('--generate', type=str)
+    args = ap.parse_args()
+
+    random.seed(42)
+    m = Janus4()
+    if args.train:
+        txt = open(args.train, 'r', encoding='utf-8', errors='ignore').read()
+        losses = train(m, txt, args.steps, args.lr)
+        print(f'loss_start={losses[0]:.4f} loss_end={losses[-1]:.4f}')
+    if args.generate:
+        print(generate(m, args.generate))