Upload family.py with huggingface_hub

family.py

@@ -0,0 +1,446 @@
+"""SpikeWhale / Byrne family traits, ported to Quazimoto-LM (attention backbone).
+
+These are the *transformer-native* family blocks (their origin is the transformer
+`modeling_byrne_embed.py`), so unlike the SNN port they operate directly on the
+sequence hidden state [B,T,d] -- no per-step adaptation needed. Each keeps the
+family's safe-at-init contract: a tanh/zero gate makes the block a no-op at start,
+while the content (`up`/`down`) weights are NON-zero so the gate still receives
+gradient (the double-zero saddle would deadlock it). DERF soft_clamp bounds any
+new instability surface, in line with the family's stability discipline.
+
+Included: HRMRefinementBlock (signature), MoESwiGLU, MTPHead, JEPAPredictorBlock.
+Engram / ProgSem are bio/SNN-specific and SpikingLinearAttention is the SNN's
+stand-in for the real attention this model already has, so they are omitted.
+"""
+from __future__ import annotations
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import instrument as _viz   # live-visualizer capture hooks (no-op unless a recorder is active)
+
+# sqrt(pi)/2: soft_clamp is the identity for small inputs and saturates smoothly to
+# +/-bound with a non-zero gradient everywhere (no dead-gradient zones).
+_ERF_K = math.sqrt(math.pi) / 2.0
+
+
+def soft_clamp(x, bound):
+    return bound * torch.erf(x * (_ERF_K / bound))
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+
+def sqrtsoftplus(x):
+    """Family expert-scoring function: sqrt(softplus(x))."""
+    return torch.sqrt(F.softplus(x) + 1e-8)
+
+
+class HRMRefinementBlock(nn.Module):
+    """Signature family block: iterative gated refinement that, in the canonical
+    HRM spirit, starts the reasoning from a RANDOM initial state z0 and iterates
+    toward a solution conditioned on the input (`anchor`).
+
+    Unlike the other family traits this block does NOT start as a no-op: the gate
+    is initialised OPEN (gate_init_open). With an input-anchored no-op start the
+    gate received ~zero gradient and never woke up; a random z0 forces the block
+    to actively reconcile the random state against the input, so the open gate
+    carries real signal from step 0. To keep the deep trunk intact we contribute
+    only the reasoning DELTA (h - z0) as a residual -- z0 itself is never dumped
+    into the trunk, and a closed gate (h == z0) degrades cleanly to a no-op."""
+
+    def __init__(self, hidden_size, refine_dim, steps, eps=1e-3, gate_init_open=0.1):
+        super().__init__()
+        self.steps = steps
+        self.norm = RMSNorm(hidden_size, eps)
+        self.down = nn.Linear(hidden_size * 2, refine_dim, bias=False)
+        self.up = nn.Linear(refine_dim, hidden_size, bias=False)
+        # random initial reasoning state (learnable), broadcast over batch/time
+        self.z0 = nn.Parameter(torch.empty(hidden_size))
+        nn.init.trunc_normal_(self.z0, std=1.0, a=-2.0, b=2.0)
+        # gates start OPEN so the random-state reasoning reaches the output at init
+        go = math.atanh(min(gate_init_open, 0.9)) if gate_init_open > 0 else 0.0
+        self.gate = nn.Parameter(torch.full((steps,), go))
+        nn.init.normal_(self.down.weight, std=0.02)
+        nn.init.normal_(self.up.weight, std=0.02)
+
+    def forward(self, x):                       # x: [B,T,d]
+        B, T, _ = x.shape
+        anchor = x
+        h = self.z0.expand(B, T, -1)            # random initial reasoning state
+        for t in range(self.steps):
+            inp = torch.cat([self.norm(h), anchor], dim=-1)
+            update = soft_clamp(self.up(F.silu(self.down(inp))), 10.0)
+            h = h + torch.tanh(self.gate[t]) * update
+        return x + (h - self.z0)                # add reasoning delta, keep trunk
+
+
+class ExpertFFN(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class MoESwiGLU(nn.Module):
+    """Shared + top-k routed SwiGLU experts, sqrtsoftplus scoring, norm_topk_prob,
+    Switch-style load-balance aux (via `last_aux_loss`). down-projections zero-init
+    => no-op at start."""
+
+    def __init__(self, hidden_size, intermediate_size, n_routed=4, n_shared=1,
+                 top_k=2, aux_loss_coef=0.01):
+        super().__init__()
+        self.top_k = min(top_k, n_routed)
+        self.n_routed = n_routed
+        self.n_shared = n_shared
+        self.aux_loss_coef = aux_loss_coef
+        self.router = nn.Linear(hidden_size, n_routed, bias=False)
+        self.experts = nn.ModuleList([ExpertFFN(hidden_size, intermediate_size)
+                                      for _ in range(n_routed)])
+        self.shared = (ExpertFFN(hidden_size, intermediate_size * n_shared)
+                       if n_shared > 0 else None)
+        for e in self.experts:
+            nn.init.zeros_(e.down_proj.weight)
+        if self.shared is not None:
+            nn.init.zeros_(self.shared.down_proj.weight)
+        self.last_aux_loss = None
+
+    def forward(self, x):                       # x: [B,T,d]
+        flat = x.reshape(-1, x.shape[-1])
+        out = torch.zeros_like(flat)
+        if self.shared is not None:
+            s = self.shared(flat)
+            out = out + (s / self.n_shared if self.n_shared > 1 else s)
+
+        logits = self.router(flat)
+        scores = sqrtsoftplus(logits)
+        topv, topi = scores.topk(self.top_k, dim=-1)
+        topv = topv / (topv.sum(-1, keepdim=True) + 1e-8)
+        for slot in range(self.top_k):
+            idx = topi[:, slot]
+            w = topv[:, slot].unsqueeze(-1)
+            for e_id, expert in enumerate(self.experts):
+                mask = idx == e_id
+                if mask.any():
+                    out[mask] = out[mask] + w[mask] * expert(flat[mask])
+
+        probs = F.softmax(logits, dim=-1)
+        expert_mask = torch.zeros_like(probs)
+        expert_mask.scatter_(1, topi, 1.0)
+        self.last_aux_loss = (self.n_routed * (expert_mask.mean(0) * probs.mean(0)).sum()
+                              * self.aux_loss_coef)
+        return out.view_as(x)
+
+
+class MTPHead(nn.Module):
+    """Multi-token-prediction head: zero-init d->d residual, reuses the tied readout."""
+
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, hidden):
+        return hidden + self.proj(hidden)
+
+
+class TokenCompressor(nn.Module):
+    """Frozen LSH-style projection (gradient never reaches it through the hash cast)."""
+    def __init__(self, hidden_size, compress_dim):
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, compress_dim, bias=False)
+        nn.init.normal_(self.proj.weight, std=0.02)
+        self.proj.weight.requires_grad_(False)
+
+    def forward(self, x):
+        return self.proj(x)
+
+
+class MultiHeadHashLookup(nn.Module):
+    """N-gram hash memory: for n=1..max_ngram, hash the n-token compressed window
+    into per-head tables and average. Ported from v2 EngramModule."""
+    def __init__(self, num_heads, table_size, compress_dim, out_dim, max_ngram=3):
+        super().__init__()
+        self.num_heads, self.table_size = num_heads, table_size
+        self.max_ngram, self.out_dim = max_ngram, out_dim
+        self.tables = nn.ModuleList([nn.Embedding(table_size, out_dim) for _ in range(num_heads)])
+        for t in self.tables:
+            nn.init.normal_(t.weight, std=0.01)
+        for n in range(1, max_ngram + 1):
+            for k in range(n):
+                proj = torch.randn(num_heads, compress_dim)
+                proj = proj / (proj.norm(dim=1, keepdim=True) + 1e-8)
+                self.register_buffer(f"hash_proj_n{n}_p{k}", proj, persistent=True)
+
+    def forward(self, compressed):
+        B, S, _ = compressed.shape
+        dev = compressed.device
+        out = torch.zeros(B, S, self.out_dim, device=dev, dtype=compressed.dtype)
+        norm = torch.zeros(S, device=dev)
+        for n in range(1, self.max_ngram + 1):
+            if S < n:
+                continue
+            valid, start = S - n + 1, n - 1
+            h = torch.zeros(B, valid, self.num_heads, device=dev)
+            for k in range(n):
+                proj = getattr(self, f"hash_proj_n{n}_p{k}")
+                h = h + torch.matmul(compressed[:, k:k + valid, :].float(), proj.t())
+            idx = h.abs().long() % self.table_size
+            for hi, table in enumerate(self.tables):
+                out[:, start:, :] = out[:, start:, :] + table(idx[:, :, hi])
+            norm[start:] += self.num_heads
+        return (out / norm.view(1, -1, 1).clamp(min=1)).to(compressed.dtype)
+
+
+class DERFContextGate(nn.Module):
+    def __init__(self, obs_size, init_bias=-4.0):
+        super().__init__()
+        self.proj = nn.Linear(obs_size * 2, obs_size)
+        self.alpha = nn.Parameter(torch.ones(obs_size))
+        self.bias = nn.Parameter(torch.full((obs_size,), init_bias))
+        self.gamma = nn.Parameter(torch.ones(obs_size))
+
+    def forward(self, retrieved, obs):
+        logits = self.proj(torch.cat([retrieved, obs], dim=-1))
+        gate = self.gamma * ((torch.erf(self.alpha * logits + self.bias) + 1.0) / 2.0)
+        return retrieved * gate
+
+
+class PhaseAttentionRing(nn.Module):
+    """Interstitial ATTENTION ring: attends causally over the sequence in
+    oscillator-PHASE space ([cos,sin] of the two neighbor oscillator rings) and
+    returns an injection current of width m = n_r + n_{r+1} for those neighbors.
+    Zero-init gate => no-op at start; soft_clamp bounds the injected drive."""
+
+    def __init__(self, m, n_heads=4, head_dim=16, bound=10.0):
+        super().__init__()
+        self.h, self.d, self.bound = n_heads, head_dim, bound
+        self.qkv = nn.Linear(2 * m, 3 * n_heads * head_dim, bias=False)
+        self.out = nn.Linear(n_heads * head_dim, m, bias=False)
+        self.gate = nn.Parameter(torch.zeros(1))
+        nn.init.normal_(self.qkv.weight, std=0.02)
+        nn.init.normal_(self.out.weight, std=0.02)
+
+    def forward(self, theta_slice):              # [B,T,m] phases of the neighbor rings
+        B, T, m = theta_slice.shape
+        feat = torch.cat([torch.cos(theta_slice), torch.sin(theta_slice)], dim=-1)
+        q, k, v = self.qkv(feat).split(self.h * self.d, dim=-1)
+        shp = lambda z: z.view(B, T, self.h, self.d).transpose(1, 2)
+        y = F.scaled_dot_product_attention(shp(q), shp(k), shp(v), is_causal=True)
+        y = y.transpose(1, 2).reshape(B, T, self.h * self.d)
+        return soft_clamp(self.out(y) * torch.tanh(self.gate), self.bound)
+
+
+class EngramRing(nn.Module):
+    """Interstitial ENGRAM ring: absorbs n-gram context from the hidden state via
+    hash memory + DERF gate, projected to an injection current of width m for the
+    two neighbor oscillator rings. Two no-op gates at init (DERF bias -4 + scale)."""
+
+    def __init__(self, hidden_size, m, compress_dim=32, num_heads=2,
+                 table_size=2048, max_ngram=3, bound=10.0):
+        super().__init__()
+        self.bound = bound
+        self.compressor = TokenCompressor(hidden_size, compress_dim)
+        self.lookup = MultiHeadHashLookup(num_heads, table_size, compress_dim, m, max_ngram)
+        self.to_obs = nn.Linear(hidden_size, m, bias=False)
+        self.gate = DERFContextGate(m, init_bias=-4.0)
+        self.scale = nn.Parameter(torch.zeros(1))   # extra no-op gate at init
+        nn.init.normal_(self.to_obs.weight, std=0.02)
+
+    def family_reinit(self):
+        """Re-apply the inits the model's global self.apply would clobber (table
+        std 0.01, frozen-random compressor, DERF bias -4)."""
+        for t in self.lookup.tables:
+            nn.init.normal_(t.weight, std=0.01)
+        nn.init.normal_(self.compressor.proj.weight, std=0.02)
+        self.compressor.proj.weight.requires_grad_(False)
+        nn.init.constant_(self.gate.bias, -4.0)
+
+    def forward(self, h):                        # h: [B,T,hidden]
+        retrieved = self.lookup(self.compressor(h.detach()))
+        gated = self.gate(retrieved, self.to_obs(h))
+        return soft_clamp(gated * torch.tanh(self.scale), self.bound)
+
+
+class RingController(nn.Module):
+    """Tiny per-ring manager that OPTIMIZES ITSELF by a predictive / free-energy rule.
+
+    Core = a fast-weight linear predictor `W` (a BUFFER, excluded from the global
+    optimizer) updated online by the delta rule  W += lr * (f - W@prev) outer prev,
+    which is exactly one gradient step on the squared prediction error -- the
+    controller learns to predict its ring's next state, minimizing surprise, with no
+    backprop. A small backprop-trained decoder maps the self-organized feature +
+    surprise into ring-control modulations; zero-init => exact no-op at start."""
+
+    def __init__(self, d_obs=4, feat=384, n_ctrl=4, local_lr=0.01):
+        super().__init__()
+        self.feat, self.local_lr = feat, local_lr
+        self.enc = nn.Linear(d_obs, feat)
+        self.dec = nn.Linear(feat + 1, n_ctrl)
+        nn.init.zeros_(self.dec.weight)
+        nn.init.zeros_(self.dec.bias)                 # control == 0 at init (no-op)
+        self.register_buffer("W", torch.zeros(feat, feat))   # self-organizing fast weights
+        self.register_buffer("prev_f", torch.zeros(feat))
+
+    def family_reinit(self):
+        nn.init.zeros_(self.dec.weight)
+        nn.init.zeros_(self.dec.bias)
+
+    def forward(self, obs):                            # obs: [d_obs] (detached ring stats)
+        f = torch.tanh(self.enc(obs))                  # [feat]
+        pred = self.W @ self.prev_f                    # predicted current feature
+        surprise = F.mse_loss(f.detach(), pred)
+        if self.training:
+            with torch.no_grad():                      # predictive self-organization (no global grad)
+                err = f.detach() - pred
+                self.W.add_(self.local_lr * torch.outer(err, self.prev_f)).clamp_(-3.0, 3.0)
+                self.prev_f.copy_(f.detach())
+        ctrl = self.dec(torch.cat([f, surprise.detach().reshape(1)]))   # [n_ctrl]
+        return ctrl, surprise.detach()
+
+
+class RingControllerBank(nn.Module):
+    """One RingController per oscillator ring (shared across all layers)."""
+
+    def __init__(self, n_rings, d_obs=4, feat=384, local_lr=0.01):
+        super().__init__()
+        self.controllers = nn.ModuleList(
+            [RingController(d_obs, feat, 4, local_lr) for _ in range(n_rings)])
+        self.last_surprise = None
+
+    def forward(self, obs):                            # obs: [R, d_obs] -> ctrl [R, 4]
+        ctrls, surps = [], []
+        for r, c in enumerate(self.controllers):
+            ct, sp = c(obs[r])
+            ctrls.append(ct)
+            surps.append(sp)
+        self.last_surprise = torch.stack(surps).mean()
+        return torch.stack(ctrls, dim=0)
+
+
+class RingSpecialists(nn.Module):
+    """A MoE-style bank of `n_spec` MINI MEMORY SPECIALISTS for ONE oscillator ring.
+
+    Each specialist owns two fast-weight stores (test-time-mutable BUFFERS, like
+    RingController.W -- excluded from the optimizer):
+      * store_in  -- a memory of the INPUT context that routes to it, and
+      * store_out -- the OUTPUT information it injects back into the ring.
+    Tokens are routed to the top-k specialists (a small MoE) by similarity to each
+    specialist's address = its learnable identity key + a read of its accumulated
+    input memory. The routed store_out is decoded into an injection current for the
+    ring, and BOTH stores are written online (train AND inference) by a gated EMA
+    rule -- so a generation accumulates an addressable context memory as it runs.
+
+    Slow (backprop) weights -- q_proj/in_enc/val_enc/out_dec/in_read/key/active --
+    learn to route, encode, retrieve and decode; the stores are the fast memory.
+    Family contract: zero-init `scale` => exact no-op at start, and an empty
+    store_out is zero anyway, so the block is doubly safe until it learns to write
+    and open the gate. `active` is a per-specialist usage gate biasing the router."""
+
+    def __init__(self, ring_size, hidden_size, n_spec=7, key_dim=32, slot_dim=64,
+                 top_k=2, write_lr=0.1, bound=10.0):
+        super().__init__()
+        self.n_spec = n_spec
+        self.top_k = min(top_k, n_spec)
+        self.write_lr, self.bound = write_lr, bound
+        self.write_enabled = True
+        self.q_proj  = nn.Linear(hidden_size, key_dim, bias=False)   # router query   (learned)
+        self.out_dec = nn.Linear(slot_dim, ring_size, bias=False)    # retrieved -> injection (learned)
+        self.in_read = nn.Linear(slot_dim, key_dim, bias=False)      # input-store -> addr  (learned)
+        # write-side encoders are FROZEN RANDOM projections (cf. EngramRing's frozen
+        # compressor): they only ever run inside the no-grad write, so backprop can't
+        # train them -- as fixed random features the stores hold a stable encoding the
+        # learned read/route/decode path can address.
+        self.in_enc  = nn.Linear(hidden_size, slot_dim, bias=False)  # input  -> input-store  (frozen)
+        self.val_enc = nn.Linear(hidden_size, slot_dim, bias=False)  # input  -> output-store (frozen)
+        self.key     = nn.Parameter(torch.randn(n_spec, key_dim) * 0.02)  # specialist identity
+        self.active  = nn.Parameter(torch.zeros(n_spec))             # per-specialist usage gate
+        self.scale   = nn.Parameter(torch.zeros(1))                  # no-op output gate at init
+        self.register_buffer("store_in",  torch.zeros(n_spec, slot_dim))
+        self.register_buffer("store_out", torch.zeros(n_spec, slot_dim))
+        for m in (self.q_proj, self.out_dec, self.in_read, self.in_enc, self.val_enc):
+            nn.init.normal_(m.weight, std=0.02)
+        self.in_enc.weight.requires_grad_(False)
+        self.val_enc.weight.requires_grad_(False)
+
+    def family_reinit(self):
+        """Re-apply inits the model's global self.apply clobbers (key, gates, and
+        the frozen write encoders)."""
+        nn.init.normal_(self.key, std=0.02)
+        nn.init.zeros_(self.active)
+        nn.init.zeros_(self.scale)
+        nn.init.normal_(self.in_enc.weight, std=0.02)
+        nn.init.normal_(self.val_enc.weight, std=0.02)
+        self.in_enc.weight.requires_grad_(False)
+        self.val_enc.weight.requires_grad_(False)
+
+    def reset_memory(self):
+        """Clear both stores -- call between independent prompts/sequences so
+        context memory does not bleed across them."""
+        self.store_in.zero_()
+        self.store_out.zero_()
+
+    def forward(self, h):                        # h: [B,T,hidden]
+        B, T, _ = h.shape
+        # snapshot the fast-weight stores: the graph must hold an immutable copy
+        # because we mutate the buffers in-place for the online write below.
+        store_in, store_out = self.store_in.clone(), self.store_out.clone()
+        q = self.q_proj(h)                                           # [B,T,key_dim]
+        addr = self.key + self.in_read(store_in)                     # [n_spec,key_dim]
+        logits = q @ addr.t()                                        # [B,T,n_spec]
+        logits = logits + F.logsigmoid(self.active)                  # usage gate biases routing
+        if self.top_k < self.n_spec:                                 # top-k MoE sparsity
+            tv = torch.topk(logits, self.top_k, dim=-1).values
+            logits = logits.masked_fill(logits < tv[..., [-1]], float("-inf"))
+        route = torch.softmax(logits, dim=-1)                        # [B,T,n_spec]
+
+        retrieved = route @ store_out                                # [B,T,slot_dim]
+        inject = self.out_dec(retrieved) * torch.tanh(self.scale)    # [B,T,ring_size]
+
+        rec = _viz.get_rec()
+        if rec is not None and rec.enabled:                          # last-token routing
+            rec.push_spec(route[0, -1].tolist())
+
+        # online write: blend this step's input/value into the routed specialists
+        if self.write_enabled and self.write_lr > 0:
+            with torch.no_grad():
+                w = route.reshape(-1, self.n_spec)                   # [BT,n_spec]
+                denom = w.sum(0).clamp(min=1e-3).unsqueeze(1)        # [n_spec,1]
+                in_info  = (w.t() @ self.in_enc(h).reshape(-1, self.in_enc.out_features)) / denom
+                val_info = (w.t() @ self.val_enc(h).reshape(-1, self.val_enc.out_features)) / denom
+                a = self.write_lr
+                self.store_in.mul_(1 - a).add_(a * in_info).clamp_(-self.bound, self.bound)
+                self.store_out.mul_(1 - a).add_(a * val_info).clamp_(-self.bound, self.bound)
+        return soft_clamp(inject, self.bound)
+
+
+class JEPAPredictorBlock(nn.Module):
+    """Representation-space k-ahead prediction with stop-grad target (JEPA asymmetry).
+    Zero-init gate => identity at init; `up` normal so the gate gets gradient."""
+
+    def __init__(self, dim, pred_dim, horizon, eps=1e-3):
+        super().__init__()
+        self.horizon = horizon
+        self.norm = RMSNorm(dim, eps)
+        self.down = nn.Linear(dim, pred_dim, bias=False)
+        self.up = nn.Linear(pred_dim, dim, bias=False)
+        self.gate = nn.Parameter(torch.zeros(horizon))
+        nn.init.normal_(self.down.weight, std=0.02)
+        nn.init.normal_(self.up.weight, std=0.02)
+
+    def forward(self, h, k):                    # h: [B,T,dim]
+        update = self.up(F.silu(self.down(self.norm(h))))
+        return h + torch.tanh(self.gate[k - 1]) * update