mirror geolip 0.8.8: model/relay.py (RelayPatchwork + RelayEMA functional form; additive, unimported by modeling path)

relay.py

@@ -0,0 +1,171 @@
+"""RelayPatchwork + RelayEMA — the relay organ and its causal-EMA memory taps.
+
+The organ is the certified residual patchwork (amoe-lora's arm geometry,
+brought here as a native model component):
+
+    slots_t = proj(x_t)                      (B, n, n_slots, D)  orthogonal, no bias
+    f_t     = m_hat(slots_t) flattened       the reconstructive signed read
+              m_hat(x) = sum_k sinh(u_k) a_hat_k / sum_j cosh(u_j)
+                       = signed(x) @ A_hat   (composed from AlephAddress)
+    y_t     = x_t + sigmoid(gate) * consume(f_t)
+
+consume = Linear(nD -> hidden), ReLU^2, LayerNorm, Linear(hidden -> d) with the
+output layer zero-initialized (weight AND bias — the strict C6 null path: a
+fresh organ is bit-inert, y == x exactly). Gate initialized -3. Reconstructive,
+never comparative: no softmax-over-choices, no argmax, no top-k.
+
+RelayEMA adds ONE mechanism: two fixed-decay causal EMAs of the organ's own
+read, fed to the head through input columns that are ZERO at birth —
+
+    F1_t = (1 - r1) F1_{t-1} + r1 f_t        r1 = 1/16
+    F2_t = (1 - r2) F2_{t-1} + r2 f_t        r2 = 1/64
+    y_t  = x_t + sigmoid(gate) * consume(cat(f_t, F1_t, F2_t))
+
+so at birth RelayEMA's forward EQUALS the memory-free organ. Training uses the
+chunked closed-form scan (ema_chunked, renormalized so q^-t never overflows);
+decode carries (F1, F2) state and updates one step at a time — exact, because
+the head is position-wise.
+
+The record (canon btx_e003, E-G2 2026-09-14, on frozen mini-beatrix-2s):
+the memory-free organ form, trained answer-only with the rule-order shortcut
+blocked in the data, internalized a 5-step rule chain — .527 as-is / .553
+rules-shuffled on out-of-vocabulary words (order-invariant; heuristic ceiling
+~.2), 1.000 in-vocabulary. The EMA taps added nothing measurable on that task
+(.527/.587). Packaged with the taps because they are free at birth and the
+accumulator is the certified form. Default geometry below is the validated
+wide arm (n_slots 32, K 64, D 4, hidden 256 at d 1024: 428k params/site).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .address import AlephAddress
+
+
+@dataclass
+class RelaySpec:
+    n_slots: int = 32
+    K: int = 64
+    D: int = 4
+    tau: float = 0.1
+    hidden: int = 256
+    rho1: float = 1.0 / 16.0
+    rho2: float = 1.0 / 64.0
+    gate_init: float = -3.0
+    zero_init_head: bool = True
+
+
+class SquaredReLU(nn.Module):
+    def forward(self, x):
+        return F.relu(x) ** 2
+
+
+def ema_chunked(f: torch.Tensor, rho: float, s0: torch.Tensor, chunk: int = 256):
+    """Causal EMA over dim 1, closed form per chunk (renormalized so q^-t
+    never overflows; autograd-safe). F_t = q^t (s0 + rho * sum_{s<=t} q^-s f_s),
+    q = 1 - rho; state carried between chunks. -> (F (B,T,C), F_last (B,C))."""
+    B, T, C = f.shape
+    q = 1.0 - rho
+    outs = []
+    s = s0
+    for a in range(0, T, chunk):
+        fb = f[:, a:a + chunk]
+        t = fb.shape[1]
+        idx = torch.arange(1, t + 1, device=f.device, dtype=f.dtype)
+        acc = torch.cumsum(fb * torch.pow(q, -idx).view(1, t, 1), dim=1)
+        Fb = torch.pow(q, idx).view(1, t, 1) * (s.unsqueeze(1) + rho * acc)
+        s = Fb[:, -1]
+        outs.append(Fb)
+    return torch.cat(outs, dim=1), s
+
+
+def _feats(m, x: torch.Tensor) -> torch.Tensor:
+    """The reconstructive read, flattened: m_hat = signed(slots) @ A_hat."""
+    B, n, _ = x.shape
+    slots = m.proj(x).view(B, n, m.n_slots, m.spec.D)
+    A = F.normalize(m.addr.codebook, dim=-1)
+    return (m.addr.signed(slots) @ A).reshape(B, n, -1)
+
+
+class RelayPatchwork(nn.Module):
+    """The memory-free relay organ: proj -> reconstructive aleph read ->
+    zero-born squared-ReLU patch head, residual write behind a sigmoid gate."""
+
+    def __init__(self, d: int, spec: RelaySpec | None = None):
+        super().__init__()
+        s = spec or RelaySpec()
+        self.spec = s
+        self.n_slots = s.n_slots
+        self.nD = s.n_slots * s.D
+        self.proj = nn.Linear(d, self.nD, bias=False)
+        nn.init.orthogonal_(self.proj.weight)
+        self.addr = AlephAddress(s.K, s.D, s.tau)
+        self.consume = nn.Sequential(
+            nn.Linear(self.nD, s.hidden), SquaredReLU(),
+            nn.LayerNorm(s.hidden), nn.Linear(s.hidden, d))
+        if s.zero_init_head:
+            nn.init.zeros_(self.consume[-1].weight)
+            nn.init.zeros_(self.consume[-1].bias)
+        self.gate = nn.Parameter(torch.tensor(float(s.gate_init)))
+
+    def feats(self, x: torch.Tensor) -> torch.Tensor:
+        return _feats(self, x)
+
+    def forward(self, x):
+        return x + torch.sigmoid(self.gate) * self.consume(self.feats(x))
+
+
+class RelayEMA(nn.Module):
+    """The organ plus zero-born EMA memory taps. Built by widening an organ's
+    head to 3nD input columns, the added columns zeroed — at birth the forward
+    equals the organ it came from (shared weights)."""
+
+    def __init__(self, d: int, spec: RelaySpec | None = None):
+        super().__init__()
+        self._widen(RelayPatchwork(d, spec))
+
+    @classmethod
+    def from_organ(cls, organ: RelayPatchwork) -> "RelayEMA":
+        """Upgrade a (possibly trained) organ in place: shares proj/addr/gate
+        and the head tail, widens the head's first layer with zeroed columns."""
+        self = cls.__new__(cls)
+        nn.Module.__init__(self)
+        self._widen(organ)
+        return self
+
+    def _widen(self, organ: RelayPatchwork):
+        self.spec = organ.spec
+        self.nD = organ.nD
+        self.proj, self.addr, self.gate = organ.proj, organ.addr, organ.gate
+        self.n_slots = organ.n_slots
+        wide = nn.Linear(3 * self.nD, organ.spec.hidden)
+        with torch.no_grad():
+            wide.weight[:, :self.nD] = organ.consume[0].weight
+            wide.weight[:, self.nD:] = 0.0
+            wide.bias.copy_(organ.consume[0].bias)
+        self.consume = nn.Sequential(
+            wide, organ.consume[1], organ.consume[2], organ.consume[3])
+
+    def feats(self, x: torch.Tensor) -> torch.Tensor:
+        return _feats(self, x)
+
+    def run(self, x, state=None):
+        """-> (y = x + write, (F1_last, F2_last)). state None starts both EMAs
+        at zero (fresh context). Decode: call on the single new position with
+        the carried state — exact, the head is position-wise."""
+        f = self.feats(x)
+        B = f.shape[0]
+        s1 = state[0] if state is not None else f.new_zeros(B, self.nD)
+        s2 = state[1] if state is not None else f.new_zeros(B, self.nD)
+        F1, s1 = ema_chunked(f, self.spec.rho1, s1)
+        F2, s2 = ema_chunked(f, self.spec.rho2, s2)
+        y = x + torch.sigmoid(self.gate) * self.consume(
+            torch.cat([f, F1, F2], dim=-1))
+        return y, (s1, s2)
+
+    def forward(self, x, state=None):
+        return self.run(x, state)[0]