lstate.py

#!/usr/bin/env python3
"""
lstate.py -- the NO-COPY INTELLIGENCE TRACK (beside the compressor): a LEARNED, tiny, diagonal-linear
RECURRENT STATE channel, trained online by RTRL, that is too small/smooth to replay spans -> any held-out
gain is ABSTRACTION (computation), not copy. This replaces §69's FIXED reservoir (which only overfit).

State (per byte t):  h_j(t) = a_j h_j(t-1) + sum_k W_jk x_k(t),   a_j = sigmoid(alpha_j),
  x(t) = the 8 SIGNED bits of the just-finished byte.  h (m=12 floats) is fed as continuous features to
  the mixer; m is tiny so the state PHYSICALLY cannot encode a verbatim 50-byte span (anti-copy by size).

Learned ONLINE by RTRL (diagonal recurrence -> O(m) eligibility traces, NO BPTT). Per byte:
  G_hj = sum over the 8 bit-predictions of (p-y)*v_j           # dL/dh_j via mixer weight v_j on feature j
  alpha_j -= lr * G_hj * e^alpha_j ;  W_jk -= lr * G_hj * e^W_jk
  e^alpha_j <- a_j(1-a_j) h_j(t-1) + a_j e^alpha_j ;  e^W_jk <- x_k + a_j e^W_jk ;  h_j <- a_j h_j + W_j.x

GATE (genmem.py protocol, reused): MATCH (copy) OFF, 13-byte-DECONTAMINATED held-out bits/byte,
state ON/OFF/SCRAMBLED, copy gain reported separately, AND a FROZEN-after-train read (train-attributable,
rules out in-context copy). Target: state beats the match-OFF baseline -> bridge to intelligence.
Run: python lstate.py
"""
import sys, math, random, time
from genmem import clean_mask, order_n_bpb, stretch, squash

ORDERS = [0, 1, 2, 3, 4, 5, 6]
M = 12        # state dimension (tiny -> cannot replay spans)
XDIM = 8      # input = the 8 signed bits of the just-finished byte


def sigmoid(z):
    if z > 30: return 1 - 1e-9
    if z < -30: return 1e-9
    return 1.0 / (1.0 + math.exp(-z))


class Model:
    def __init__(self, use_match=True, use_state=True, order_cap=6, lr=0.004, lr_rec=0.02, scramble=False,
                 fixed_long=False, seed=0):
        self.use_match, self.use_state = use_match, use_state
        self.lr, self.lr_rec, self.scramble, self.fixed_long = lr, lr_rec, scramble, fixed_long
        self.orders = [o for o in ORDERS if o <= order_cap]; self.NM = len(self.orders)
        self.NS = M if use_state else 0
        self.NIN = self.NM + (1 if use_match else 0) + self.NS
        self.w = [0.0] * (self.NIN + 1); self.wg = [0.0] * (self.NIN + 1)
        self.tab = [dict() for _ in self.orders]
        self.MINLEN, self.GATE = 16, 18
        self.mtab = {}; self.hist = bytearray(); self.mptr = -1; self.mlen = 0
        self.htail = 0; self.cur = 0; self.phase = 0
        self.sbase = self.NM + (1 if use_match else 0)        # index of first state feature in w/sts
        if use_state:
            rng = random.Random(seed)
            if fixed_long:
                # FORCE long-range: decays fixed at a in [0.90 .. 0.999] (timescales ~10 .. ~1000 bytes),
                # NOT learned -> the state cannot collapse to the short range the orders already cover.
                avals = [0.90 + (0.999 - 0.90) * (j / (M - 1)) for j in range(M)]
            else:
                avals = [0.5 + 0.49 * (j / (M - 1)) for j in range(M)]     # 0.5 .. 0.99, learned
            self.alpha = [math.log(a / (1 - a)) for a in avals]
            self.W = [[(rng.random() - 0.5) * 0.1 for _ in range(XDIM)] for _ in range(M)]
            self.h = [0.0] * M; self.ea = [0.0] * M; self.eW = [[0.0] * XDIM for _ in range(M)]
            self.Ghj = [0.0] * M
        if scramble: self._sr = random.Random(1234)

    def _octx(self, k):
        b = self.orders[k]; mask = (1 << (8 * b)) - 1 if b else 0
        return (self.phase, self.cur, (self.htail & mask) if b else 0)

    def predict(self):
        sts = [0.0] * (self.NIN + 1); i = 0
        for k in range(self.NM):
            c = self.tab[k].get(self._octx(k)); n0, n1 = (c[0], c[1]) if c else (0, 0)
            sts[i] = stretch((n1 + 0.2) / (n0 + n1 + 0.4)); i += 1
        if self.use_match:
            st = 0.0
            if self.mlen >= self.GATE and 0 <= self.mptr < len(self.hist):
                pb = self.hist[self.mptr]
                if self.phase == 0 or (pb >> (8 - self.phase)) == self.cur:
                    bit = (pb >> (7 - self.phase)) & 1; st = (1.6 + 0.35 * min(self.mlen, 28)) * (1 if bit else -1)
            sts[i] = st; i += 1
        if self.use_state:
            if self.scramble:
                for j in range(M): sts[i] = self._sr.uniform(-1, 1); i += 1
            else:
                for j in range(M): sts[i] = self.h[j]; i += 1
        sts[self.NIN] = 1.0
        d = sum(self.w[j] * sts[j] for j in range(self.NIN + 1))
        return squash(d), sts

    def step(self, y, learn):
        p, sts = self.predict()
        cost = -math.log2(p if y == 1 else 1 - p)
        if self.use_state and not self.scramble:
            for j in range(M):
                self.Ghj[j] += (p - y) * self.w[self.sbase + j]     # accumulate dL/dh_j over the byte
        if learn:
            err = y - p
            for j in range(self.NIN + 1):
                g = err * sts[j]; self.wg[j] = 0.999 * self.wg[j] + 0.001 * g * g
                self.w[j] += self.lr * g / (math.sqrt(self.wg[j]) + 1e-4)
            for k in range(self.NM):
                key = self._octx(k); c = self.tab[k].get(key)
                if c is None: c = [0, 0]; self.tab[k][key] = c
                c[y] += 1
        self.cur = (self.cur << 1) | y; self.phase += 1
        if self.phase == 8:
            b = self.cur & 0xFF; self.hist.append(b)
            if self.use_match: self._match_after(b)
            if self.use_state:
                x = [2.0 * ((b >> k) & 1) - 1.0 for k in range(8)]
                if self.scramble:
                    x = [self._sr.uniform(-1, 1) for _ in range(8)]
                if learn and not self.scramble:
                    for j in range(M):
                        gj = self.Ghj[j]
                        if not self.fixed_long:                  # learn decays only when not forced-long
                            self.alpha[j] -= self.lr_rec * gj * self.ea[j]
                        Wj, eWj = self.W[j], self.eW[j]
                        for k in range(XDIM):
                            Wj[k] -= self.lr_rec * gj * eWj[k]
                for j in range(M):
                    a = sigmoid(self.alpha[j]); hpr = self.h[j]
                    Wj, eWj = self.W[j], self.eW[j]; wr = 0.0
                    for k in range(XDIM):
                        wr += Wj[k] * x[k]
                    # proper EMA: h = a*hpr + (1-a)*wr  -> bounded state AND bounded RTRL traces
                    self.ea[j] = a * (1 - a) * (hpr - wr) + a * self.ea[j]
                    for k in range(XDIM):
                        eWj[k] = (1 - a) * x[k] + a * eWj[k]
                    self.h[j] = a * hpr + (1 - a) * wr
                self.Ghj = [0.0] * M
            self.htail = ((self.htail << 8) | b) & ((1 << 48) - 1); self.cur = 0; self.phase = 0
        return cost

    def _match_after(self, b):
        n = len(self.hist)
        if self.mlen > 0 and self.mptr < n - 1:
            if self.hist[self.mptr] == self.hist[n - 1]: self.mptr += 1; self.mlen = min(self.mlen + 1, 65535)
            else: self.mlen = 0; self.mptr = -1
        if n >= self.MINLEN:
            key = bytes(self.hist[n - self.MINLEN:n]); prev = self.mtab.get(key, -1); self.mtab[key] = n
            if self.mlen == 0 and 0 <= prev < n: self.mptr = prev; self.mlen = self.MINLEN

    def run(self, raw, learn, mask=None):
        tot = 0.0; nb = 0; bp = 0
        for b in raw:
            scored = (mask is None) or (bp < len(mask) and mask[bp])
            bc = 0.0
            for j in range(7, -1, -1):
                bc += self.step((b >> j) & 1, learn)
            if scored: tot += bc; nb += 1
            bp += 1
        return tot / max(1, nb)


def warm_at(train, test, mask, use_match, use_state, scramble):
    m = Model(use_match=use_match, use_state=use_state, scramble=scramble)
    m.run(train, learn=True)
    return m.run(test, learn=True, mask=mask)


def main():
    try:
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    except Exception:
        pass
    full_train = open("data/wt103_train.txt", "rb").read()
    full_test = open("data/wt103_test.txt", "rb").read()
    SIZES = [150, 450, 1200, 2700]                          # KB of TRAIN; test = half, capped
    print("=" * 100)
    print(f"NO-COPY INTELLIGENCE TRACK -- learned RTRL state (m={M}); DATA-SCALING the copy-ablated gate")
    print("  question: does the learned compact state cross BELOW the no-state baseline as data grows?")
    print("  (decontaminated wt103, copy/match OFF; state must beat SCRAMBLED = real signal, and baseline = win)")
    print("=" * 100)
    t0 = time.time()
    print(f"\n{'train':>8}{'baseline':>11}{'state':>10}{'scrambled':>11}{'  state-baseline':>17}{'  state-vs-noise':>17}")
    print("-" * 100)
    rows = []
    for kb in SIZES:
        tr = full_train[:kb * 1024]
        te = full_test[:min(400 * 1024, kb * 1024 // 2)]      # cap test slice for speed
        mask = clean_mask(tr, te, 13)
        L_base = warm_at(tr, te, mask, False, False, False)
        L_state = warm_at(tr, te, mask, False, True, False)
        L_scr = warm_at(tr, te, mask, False, True, True)
        gap = L_state - L_base          # < 0 => state BEATS baseline (the win)
        margin = L_scr - L_state        # > 0 => state beats the noise floor (real learned signal)
        rows.append((kb, L_base, L_state, L_scr, gap, margin))
        print(f"{kb:>6}KB{L_base:>11.4f}{L_state:>10.4f}{L_scr:>11.4f}{gap:>+17.4f}{margin:>+17.4f}", flush=True)
    print("-" * 100)
    gaps = [r[4] for r in rows]
    margins = [r[5] for r in rows]
    print(f"state-vs-baseline gap by data size: {['%+.4f' % g for g in gaps]}")
    print(f"state-vs-noise margin (real signal): {['%+.4f' % m for m in margins]}")
    crossed = any(g < -0.001 for g in gaps)
    shrinking = gaps[-1] < gaps[0] - 0.001
    real_signal = all(m > 0.003 for m in margins)
    print("\nVERDICT:", end=" ")
    if crossed and real_signal:
        print("PASS -- the learned compact state crosses BELOW the no-state baseline at scale, beating both")
        print("  the scrambled noise floor and the no-state baseline on copy-ablated decontaminated held-out.")
        print("  A no-copy (m=%d) memory adds GENERALIZATION. Next: port the state channel into strong.rs." % M)
    elif real_signal and shrinking:
        print("PROMISING -- the state has REAL learned signal (beats the noise floor at every size) and its")
        print("  gap to baseline SHRINKS with data, but has not crossed yet on this CPU-scale slice. The honest")
        print("  read: the learned no-copy memory works directionally; crossing needs more data/Rust speed or a")
        print("  richer input projection -- not a new idea. (Contrast: §69's FIXED reservoir matched noise.)")
    else:
        print("NEGATIVE -- no robust signal beyond the noise floor / no favorable trend. Honest; iterate or stop.")
    print(f"\n[{time.time()-t0:.0f}s]")
    print("=" * 100)


if __name__ == "__main__":
    main()