AbstractPhil/mini-beatrix-1

1"""RelayPatchwork + RelayEMA — a gated residual patch module over a signed2address read, and a variant with causal-EMA memory taps.3 4RelayPatchwork (ported from the amoe-lora adapter of the same name,5github.com/AbstractEyes/amoe-lora, as a native model component):6 7    slots_t = proj(x_t)                      (B, n, n_slots, D)  orthogonal, no bias8    f_t     = m_hat(slots_t) flattened       the reconstructive signed read9              m_hat(x) = sum_k sinh(u_k) a_hat_k / sum_j cosh(u_j)10                       = signed(x) @ A_hat   (composed from AlephAddress)11    y_t     = x_t + sigmoid(gate) * consume(f_t)12 13consume = Linear(nD -> hidden), ReLU^2, LayerNorm, Linear(hidden -> d) with the14output layer zero-initialized (weight AND bias), so a fresh module is exactly15inert: y == x. Gate initialized -3. Reconstructive, never comparative: no16softmax-over-choices, no argmax, no top-k.17 18RelayEMA adds one mechanism: two fixed-decay causal EMAs of the module's own19read, fed to the head through input columns that start at zero —20 21    F1_t = (1 - r1) F1_{t-1} + r1 f_t        r1 = 1/1622    F2_t = (1 - r2) F2_{t-1} + r2 f_t        r2 = 1/6423    y_t  = x_t + sigmoid(gate) * consume(cat(f_t, F1_t, F2_t))24 25so at birth RelayEMA's forward equals the memory-free RelayPatchwork. Training26uses the chunked closed-form scan (ema_chunked, renormalized so q^-t never27overflows); decode carries (F1, F2) state one step at a time — exact, because28the head is position-wise.29 30Defaults are the validated geometry (n_slots 32, K 64, D 4, hidden 256; ~428k31params at d = 1024). Validation record and trained weights:32huggingface.co/AbstractPhil/mini-beatrix-2s, arms/btx_e003.33"""34from __future__ import annotations35 36from dataclasses import dataclass37 38import torch39import torch.nn as nn40import torch.nn.functional as F41 42from .address import AlephAddress43 44 45@dataclass46class RelaySpec:47    n_slots: int = 3248    K: int = 6449    D: int = 450    tau: float = 0.151    hidden: int = 25652    rho1: float = 1.0 / 16.053    rho2: float = 1.0 / 64.054    gate_init: float = -3.055    zero_init_head: bool = True56 57 58class SquaredReLU(nn.Module):59    def forward(self, x):60        return F.relu(x) ** 261 62 63def ema_chunked(f: torch.Tensor, rho: float, s0: torch.Tensor, chunk: int = 256):64    """Causal EMA over dim 1, closed form per chunk (renormalized so q^-t65    never overflows; autograd-safe). F_t = q^t (s0 + rho * sum_{s<=t} q^-s f_s),66    q = 1 - rho; state carried between chunks. -> (F (B,T,C), F_last (B,C))."""67    B, T, C = f.shape68    q = 1.0 - rho69    outs = []70    s = s071    for a in range(0, T, chunk):72        fb = f[:, a:a + chunk]73        t = fb.shape[1]74        idx = torch.arange(1, t + 1, device=f.device, dtype=f.dtype)75        acc = torch.cumsum(fb * torch.pow(q, -idx).view(1, t, 1), dim=1)76        Fb = torch.pow(q, idx).view(1, t, 1) * (s.unsqueeze(1) + rho * acc)77        s = Fb[:, -1]78        outs.append(Fb)79    return torch.cat(outs, dim=1), s80 81 82def _feats(m, x: torch.Tensor) -> torch.Tensor:83    """The reconstructive read, flattened: m_hat = signed(slots) @ A_hat."""84    B, n, _ = x.shape85    slots = m.proj(x).view(B, n, m.n_slots, m.spec.D)86    A = F.normalize(m.addr.codebook, dim=-1)87    return (m.addr.signed(slots) @ A).reshape(B, n, -1)88 89 90class RelayPatchwork(nn.Module):91    """The memory-free form: proj -> reconstructive address read ->92    zero-initialized squared-ReLU head, residual write behind a sigmoid gate."""93 94    def __init__(self, d: int, spec: RelaySpec | None = None):95        super().__init__()96        s = spec or RelaySpec()97        self.spec = s98        self.n_slots = s.n_slots99        self.nD = s.n_slots * s.D100        self.proj = nn.Linear(d, self.nD, bias=False)101        nn.init.orthogonal_(self.proj.weight)102        self.addr = AlephAddress(s.K, s.D, s.tau)103        self.consume = nn.Sequential(104            nn.Linear(self.nD, s.hidden), SquaredReLU(),105            nn.LayerNorm(s.hidden), nn.Linear(s.hidden, d))106        if s.zero_init_head:107            nn.init.zeros_(self.consume[-1].weight)108            nn.init.zeros_(self.consume[-1].bias)109        self.gate = nn.Parameter(torch.tensor(float(s.gate_init)))110 111    def feats(self, x: torch.Tensor) -> torch.Tensor:112        return _feats(self, x)113 114    def forward(self, x):115        return x + torch.sigmoid(self.gate) * self.consume(self.feats(x))116 117 118class RelayEMA(nn.Module):119    """RelayPatchwork plus the EMA memory taps. Built by widening a patchwork's120    head to 3nD input columns with the added columns zeroed, so at birth the121    forward equals the patchwork it came from (shared weights)."""122 123    def __init__(self, d: int, spec: RelaySpec | None = None):124        super().__init__()125        self._widen(RelayPatchwork(d, spec))126 127    @classmethod128    def from_patchwork(cls, base: RelayPatchwork) -> "RelayEMA":129        """Upgrade a (possibly trained) RelayPatchwork in place: shares its130        proj/addr/gate and head tail, widens the head's first layer with131        zeroed columns."""132        self = cls.__new__(cls)133        nn.Module.__init__(self)134        self._widen(base)135        return self136 137    def _widen(self, base: RelayPatchwork):138        self.spec = base.spec139        self.nD = base.nD140        self.proj, self.addr, self.gate = base.proj, base.addr, base.gate141        self.n_slots = base.n_slots142        wide = nn.Linear(3 * self.nD, base.spec.hidden)143        with torch.no_grad():144            wide.weight[:, :self.nD] = base.consume[0].weight145            wide.weight[:, self.nD:] = 0.0146            wide.bias.copy_(base.consume[0].bias)147        self.consume = nn.Sequential(148            wide, base.consume[1], base.consume[2], base.consume[3])149 150    def feats(self, x: torch.Tensor) -> torch.Tensor:151        return _feats(self, x)152 153    def run(self, x, state=None):154        """-> (y = x + write, (F1_last, F2_last)). state None starts both EMAs155        at zero (fresh context). Decode: call on the single new position with156        the carried state — exact, the head is position-wise."""157        f = self.feats(x)158        B = f.shape[0]159        s1 = state[0] if state is not None else f.new_zeros(B, self.nD)160        s2 = state[1] if state is not None else f.new_zeros(B, self.nD)161        F1, s1 = ema_chunked(f, self.spec.rho1, s1)162        F2, s2 = ema_chunked(f, self.spec.rho2, s2)163        y = x + torch.sigmoid(self.gate) * self.consume(164            torch.cat([f, F1, F2], dim=-1))165        return y, (s1, s2)166 167    def forward(self, x, state=None):168        return self.run(x, state)[0]169