Reframe as research artifact: rich card + Apache-2.0 license + clean runnable code subset; remove internal design files
91d91c5 verified | """ | |
| Prizma-Seq — a Predictive-Coding Gated-DeltaNet sequence model (design spec PRIZMA_SEQ_SPEC.md). | |
| Mixer = a carried associative workspace state S_t in R^{d_h x d_h} per head, updated by a | |
| precision-gated targeted erase-and-write (the delta rule), which is exactly ONE gradient step on | |
| Prizma's per-token free energy F_t(S)=1/2||v_t - S k_t||^2. Read = S_{t-1} q_t (recognition-by- | |
| reconstruction, strictly pre-write/causal) + a small exact local window head. FFN is byte-identical | |
| to the Transformer baseline so ONLY the mixer differs. | |
| Honest design note: the write gate beta is INPUT-dependent (sigma(W_beta x_t)), which keeps the | |
| chunk-parallel training form valid. Surprise-proportionality is intrinsic: the delta write is | |
| u_t = beta_t * (v_t - S_{t-1} k_t) = beta_t * epsilon_t — it writes the prediction error itself | |
| (Prizma's dW ~ (Pi*eps) (x) r). An optional surprise-gated variant (two-pass) is exposed for B6. | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from .transformer import RMSNorm, SwiGLU, TFConfig | |
| from .delta import chunked_delta | |
| class PrizmaSeqConfig: | |
| vocab: int = 64 | |
| d_model: int = 64 | |
| n_layers: int = 2 | |
| n_heads: int = 2 | |
| chunk: int = 64 | |
| window: int = 16 | |
| d_ff: int = None | |
| max_len: int = 1024 | |
| rope: bool = False # delta keys/queries are POSITION-FREE (DeltaNet/Mamba standard): | |
| # RoPE makes the recall dot-product distance-dependent and | |
| # scrambles content-based recall. Position comes from the conv + | |
| # causal write-order. (RoPE on state keys empirically blocks MQAR.) | |
| beta_cap: float = 0.99 | |
| gated: bool = False # data-dependent forget gate alpha (Gated-DeltaNet); off for diagnostics | |
| learned_pos: bool = False # add a learned absolute pos-embedding (used for char-LM parity) | |
| short_conv: int = 4 # short causal depthwise conv before qkv (Mamba/Based/DeltaNet std); | |
| # lets a token's k/v encode its predecessor -> enables recall. 0=off. | |
| # --- ablation knobs (B6) --- | |
| precision_gate: str = "input" # 'input' (sigma(W_beta x)) | 'uniform' | 'random' (B6 controls) | |
| write_mode: str = "delta" # 'delta' (targeted erase-and-write) | 'additive' (linear-attn) | |
| use_workspace: bool = True # False -> no carried state (window head only) | |
| use_window: bool = True # False -> no local window head (state only) | |
| route_readout: bool = True # False -> read state with a FIXED (input-independent) query | |
| # (B6 noRouteReadout: kills content-based recall) | |
| # --- capacity lever: parameter-free quadratic key/query feature map (design rank #1) - # | |
| feat_map: str = "none" # 'none' | 'quad2' | 'quad2_lowrank' | 'rand_linear'. | |
| # 'quad2': expands delta keys/queries from d_h to | |
| # d_phi = d_h + feat_n2 via FIXED random quadratic monomials | |
| # -> rectangular state S in R^{d_h x d_phi}, raising | |
| # associative-recall key-rank toward D=128 at ZERO trainable | |
| # params (buffers, not Parameters). O(1) inference intact. | |
| # 'quad2_lowrank': leaner variant — projects x (d_h) -> r dims | |
| # via a FIXED seeded matrix P in R^{d_h x r}, then takes ALL | |
| # r*(r+1)/2 quadratic monomials in the r-dim projected space, | |
| # giving d_phi = d_h + r*(r+1)/2. Default r=feat_rank (or 14 | |
| # when feat_rank=0), which yields d_phi=137 (~54% of quad2's | |
| # 256 at d_h=32) while preserving recall capacity (crosstalk | |
| # gap from quad2 < 0.01). Still ZERO trainable params — | |
| # P and monomial indices are buffers seeded at 1234. | |
| feat_n2: int = 96 # number of quadratic monomials (d_phi = d_h + feat_n2); used by | |
| # 'quad2' and 'rand_linear' only (ignored by 'quad2_lowrank'). | |
| feat_rank: int = 0 # low-rank projection dim r for 'quad2_lowrank'. When 0, defaults | |
| # to 14 -> d_phi = d_h + 14*15//2 = d_h + 105 (≈137 for d_h=32). | |
| # Effective d_phi = d_h + r*(r+1)//2 (documented in __post_init__). | |
| out_gate: bool = False # per-token output gate g=sigma(W_g x); o = o * g before W_o (RWKV-7/GLA) | |
| state_norm: bool = False # per-head RMSNorm on the delta-state read o_delta before merge | |
| banded_window: bool = False # O(T*w) banded sliding-window kernel (exact-equal to _window, default off) | |
| decoupled_gate: bool = False # GDN-2: decouple erase gate beta_e (key-side) from write gate beta_w | |
| n_delta: int = 1 # DeltaProduct: number of sequential delta sub-steps per token (k=1 = today) | |
| # --- surprise-gated write (Lever A) --- | |
| surprise_gate: bool = False # if True, scale write by g_t = 1+tanh(||eps_t||) (default OFF = identical) | |
| surprise_mode: str = 'norm' # 'norm' | 'random' | 'constant' — controls for R9 ablation | |
| def __post_init__(self): | |
| if self.d_ff is None: | |
| self.d_ff = int(round(8 / 3 * self.d_model / 8) * 8) | |
| assert self.d_model % self.n_heads == 0 | |
| self.d_h = self.d_model // self.n_heads | |
| assert self.d_h % 2 == 0, "d_h must be even for RoPE" | |
| assert self.feat_map in ("none", "quad2", "quad2_lowrank", "rand_linear") | |
| assert self.surprise_mode in ('norm', 'random', 'constant'), \ | |
| f"surprise_mode must be 'norm', 'random', or 'constant', got {self.surprise_mode!r}" | |
| # d_phi = delta key/query dim after the optional feature map (= d_h when 'none'). | |
| # 'rand_linear' = a FIXED random linear map d_h->d_phi (a CONTROL: it stays in a d_h-rank | |
| # subspace so it must give NO capacity gain, proving the quad2 MONOMIALS are what help). | |
| # 'quad2_lowrank': effective r = feat_rank if feat_rank > 0 else 14 (default); | |
| # d_phi = d_h + r*(r+1)//2 (all upper-triangular monomial pairs in the projected space). | |
| if self.feat_map == "quad2_lowrank": | |
| _r = self.feat_rank if self.feat_rank > 0 else 14 | |
| self._feat_rank_eff = _r | |
| self.d_phi = self.d_h + _r * (_r + 1) // 2 | |
| elif self.feat_map in ("quad2", "rand_linear"): | |
| self._feat_rank_eff = 0 | |
| self.d_phi = self.d_h + self.feat_n2 | |
| else: | |
| self._feat_rank_eff = 0 | |
| self.d_phi = self.d_h | |
| # ----------------------------------- RoPE ------------------------------------------------- # | |
| def _rope_cache(T, d_h, device, dtype, offset=0): | |
| inv = 1.0 / (10000 ** (torch.arange(0, d_h, 2, device=device, dtype=torch.float32) / d_h)) | |
| pos = torch.arange(offset, offset + T, device=device, dtype=torch.float32) | |
| ang = torch.outer(pos, inv) # [T, d_h/2] | |
| return torch.cos(ang).to(dtype), torch.sin(ang).to(dtype) | |
| def _apply_rope(x, cos, sin): | |
| # x: [B,H,T,d_h]; cos/sin: [T,d_h/2] | |
| x1, x2 = x[..., 0::2], x[..., 1::2] | |
| rx1 = x1 * cos - x2 * sin | |
| rx2 = x1 * sin + x2 * cos | |
| out = torch.empty_like(x) | |
| out[..., 0::2] = rx1 | |
| out[..., 1::2] = rx2 | |
| return out | |
| def _l2(x, eps=1e-6): | |
| return x / (x.norm(dim=-1, keepdim=True) + eps) | |
| # --------------------------------- the block ---------------------------------------------- # | |
| class PrizmaSeqBlock(nn.Module): | |
| def __init__(self, cfg: PrizmaSeqConfig): | |
| super().__init__() | |
| self.cfg = cfg | |
| d, H, dh = cfg.d_model, cfg.n_heads, cfg.d_h | |
| self.H, self.dh = H, dh | |
| self.norm1 = RMSNorm(d) | |
| self.kc = cfg.short_conv | |
| if self.kc > 0: | |
| self.conv = nn.Conv1d(d, d, self.kc, groups=d, bias=True) # depthwise causal short conv | |
| self.W_qkv = nn.Linear(d, 3 * d, bias=False) | |
| self.W_beta = nn.Linear(d, H, bias=True) | |
| self.beta_logit = nn.Parameter(torch.zeros(H)) # uniform (input-independent) gate, B6 | |
| self.q_fixed = nn.Parameter(torch.randn(H, dh) * 0.02) # noRouteReadout fixed query, B6 | |
| self.W_alpha = nn.Linear(d, H, bias=True) if cfg.gated else None | |
| self.W_o = nn.Linear(d, d, bias=False) | |
| self.W_g = nn.Linear(d, d, bias=True) if cfg.out_gate else None | |
| self.state_rms = RMSNorm(dh) if cfg.state_norm else None | |
| self.W_e = nn.Linear(d, H, bias=True) if cfg.decoupled_gate else None # erase gate beta_e | |
| # DeltaProduct n_delta>=2: (n_delta-1) additional kv projections + beta heads. | |
| # Each extra sub-step j>=1 gets its own W_kv_j and W_beta_j (independent params). | |
| # Total param cost: (n_delta-1) * (2*d*d + H) trainable weights per block. | |
| # NOTE: the plan explicitly allows repeated projections for simplicity (documented). | |
| self.W_kv_extra = nn.ModuleList([ | |
| nn.Linear(d, 2 * d, bias=False) for _ in range(cfg.n_delta - 1) | |
| ]) if cfg.n_delta >= 2 else None | |
| self.W_beta_extra = nn.ModuleList([ | |
| nn.Linear(d, H, bias=True) for _ in range(cfg.n_delta - 1) | |
| ]) if cfg.n_delta >= 2 else None | |
| self.norm2 = RMSNorm(d) | |
| self.mlp = SwiGLU(TFConfig(d_model=d, d_ff=cfg.d_ff)) | |
| self.win_scale = dh ** -0.5 | |
| self.d_phi = cfg.d_phi | |
| if cfg.feat_map == "quad2": | |
| # FIXED random quadratic monomials phi(x) = [x ; x[I]*x[J]] -> d_phi. Seeded => a fixed | |
| # architectural choice (NOT tuned per task; disclosed). Registered as BUFFERS, so | |
| # param_count is unchanged (byte-identical param-match to the Transformer preserved). | |
| g = torch.Generator().manual_seed(1234) | |
| self.register_buffer("feat_I", torch.randint(0, dh, (cfg.feat_n2,), generator=g)) | |
| self.register_buffer("feat_J", torch.randint(0, dh, (cfg.feat_n2,), generator=g)) | |
| elif cfg.feat_map == "quad2_lowrank": | |
| # Leaner variant (Task 1.D): project x (d_h) -> r dims via FIXED seeded P in R^{d_h x r}, | |
| # then take ALL r*(r+1)/2 upper-triangular monomial pairs in the r-dim space. | |
| # d_phi = d_h + r*(r+1)//2. At r=14 (default): d_phi = 32+105=137 for d_h=32. | |
| # P is a buffer (ZERO trainable params), seeded same discipline as quad2 (seed 1234). | |
| r = cfg._feat_rank_eff | |
| g = torch.Generator().manual_seed(1234) | |
| self.register_buffer("feat_P", torch.randn(dh, r, generator=g) * (dh ** -0.5)) | |
| n_pairs = r * (r + 1) // 2 | |
| I_lr = torch.tensor([i for i in range(r) for j in range(i, r)], dtype=torch.long) | |
| J_lr = torch.tensor([j for i in range(r) for j in range(i, r)], dtype=torch.long) | |
| self.register_buffer("feat_I_lr", I_lr) | |
| self.register_buffer("feat_J_lr", J_lr) | |
| elif cfg.feat_map == "rand_linear": | |
| # CONTROL: fixed random linear map d_h -> d_phi. rank <= d_h, so it CANNOT | |
| # raise recall key-rank -> expected NO gain over 'none'. Buffer => param_count unchanged. | |
| g = torch.Generator().manual_seed(1234) | |
| self.register_buffer("W_rand", torch.randn(dh, self.d_phi, generator=g) * (dh ** -0.5)) | |
| def _apply_conv(self, x): | |
| """Causal depthwise short conv + SiLU on the normed input (Mamba/Based-style). x:[B,T,d].""" | |
| if self.kc == 0: | |
| return x | |
| xc = F.pad(x.transpose(1, 2), (self.kc - 1, 0)) # left-pad -> causal | |
| return F.silu(self.conv(xc).transpose(1, 2)) # [B,T,d] | |
| def _phi(self, x): | |
| """Quadratic key/query feature map for the DELTA path only. x:[...,d_h] (already L2-normed) | |
| -> [...,d_phi]. Identity when feat_map='none'; else _l2([x ; x[I]*x[J]]) over FIXED random | |
| monomial indices, which escapes the d_h subspace and cuts associative-recall crosstalk | |
| (D=128: ~0.142 -> ~0.117 for quad2/d_phi=256, key_crosstalk metric). The final _l2 preserves | |
| ||k||=1, the invariant the delta kernel relies on. The local-window head keeps linear L2 keys. | |
| 'quad2_lowrank': project x -> z=x@P (r-dim), take all r*(r+1)/2 monomials z[I]*z[J], | |
| concat [x; monomials] -> d_phi = d_h + r*(r+1)/2 (leaner than quad2 at ~half d_phi).""" | |
| if self.cfg.feat_map == "none": | |
| return x | |
| if self.cfg.feat_map == "rand_linear": | |
| return _l2(x @ self.W_rand) # control: rank <= d_h -> no capacity gain | |
| if self.cfg.feat_map == "quad2_lowrank": | |
| z = x @ self.feat_P # [..., r] (fixed seeded projection) | |
| two = z[..., self.feat_I_lr] * z[..., self.feat_J_lr] # [..., n_pairs] | |
| return _l2(torch.cat([x, two], dim=-1)) # [..., d_phi = d_h + n_pairs] | |
| # quad2: fixed random monomials from the d_h-dim input | |
| two = x[..., self.feat_I] * x[..., self.feat_J] | |
| return _l2(torch.cat([x, two], dim=-1)) | |
| def _encode(self, x, cos, sin): | |
| B, T, _ = x.shape | |
| qkv = self.W_qkv(x).view(B, T, self.H, 3, self.dh) | |
| q, k, v = qkv.unbind(3) # each [B,T,H,dh] | |
| q = q.transpose(1, 2); k = k.transpose(1, 2); v = v.transpose(1, 2) # [B,H,T,dh] | |
| if self.cfg.rope: | |
| q = _apply_rope(q, cos, sin) | |
| k = _apply_rope(k, cos, sin) | |
| q, k = _l2(q), _l2(k) | |
| B, T = x.shape[0], x.shape[1] | |
| if not self.cfg.route_readout: # B6: read with a FIXED query -> no content-based recall | |
| q = _l2(self.q_fixed[None, :, None, :].expand(B, self.H, T, self.dh).contiguous()) | |
| if self.cfg.precision_gate == "uniform": | |
| beta = torch.sigmoid(self.beta_logit)[None, :, None].expand(B, self.H, T) * self.cfg.beta_cap | |
| elif self.cfg.precision_gate == "random": # input-independent random write gate (B6 control) | |
| beta = torch.rand(B, self.H, T, device=x.device, dtype=x.dtype) * self.cfg.beta_cap | |
| else: | |
| beta = torch.sigmoid(self.W_beta(x)).transpose(1, 2) * self.cfg.beta_cap # [B,H,T] | |
| if self.W_alpha is not None: | |
| alpha = torch.sigmoid(self.W_alpha(x)).transpose(1, 2) # [B,H,T] | |
| alpha = 0.5 + 0.5 * alpha # keep decay in [0.5,1] (stability) | |
| else: | |
| alpha = None | |
| if self.W_e is not None: | |
| beta_e = torch.sigmoid(self.W_e(x)).transpose(1, 2) * self.cfg.beta_cap # [B,H,T] | |
| else: | |
| beta_e = None # => chunked_delta will use beta_e = beta (byte-identical) | |
| return q, k, v, beta, alpha, beta_e | |
| def _window(self, q, k, v): | |
| """Exact causal attention restricted to the last `w` tokens (incl. self). Fused SDPA + band | |
| mask (fast on MPS); scaling 1/sqrt(d_h) matches the streaming step() path.""" | |
| T = q.shape[2] | |
| w = self.cfg.window | |
| idx = torch.arange(T, device=q.device) | |
| band = (idx[None, :] <= idx[:, None]) & (idx[None, :] > idx[:, None] - w) # [T,T] True=allow | |
| mask = torch.zeros(T, T, device=q.device, dtype=q.dtype).masked_fill(~band, float("-inf")) | |
| return F.scaled_dot_product_attention(q, k, v, attn_mask=mask) # [B,H,T,dh] | |
| def _window_banded(self, q, k, v): | |
| """Sliding-window causal attention in O(T*w) via fixed-size chunks of size w. Each query in | |
| chunk c attends keys in chunks {c-1, c} masked to [i-w+1, i]. Numerically equals _window.""" | |
| B, H, T, dh = q.shape | |
| w = self.cfg.window | |
| outs = [] | |
| for c0 in range(0, T, w): | |
| c1 = min(c0 + w, T) | |
| qc = q[:, :, c0:c1] # [B,H,Cq,dh] | |
| k0 = max(0, c0 - w) | |
| kc = k[:, :, k0:c1]; vc = v[:, :, k0:c1] # span <= 2w | |
| qi = torch.arange(c0, c1, device=q.device)[:, None] | |
| ki = torch.arange(k0, c1, device=q.device)[None, :] | |
| band = (ki <= qi) & (ki > qi - w) | |
| mask = torch.zeros(c1 - c0, c1 - k0, device=q.device, dtype=q.dtype).masked_fill(~band, float("-inf")) | |
| outs.append(F.scaled_dot_product_attention(qc, kc, vc, attn_mask=mask)) | |
| return torch.cat(outs, dim=2) | |
| def forward(self, h): | |
| B, T, d = h.shape | |
| x = self._apply_conv(self.norm1(h)) | |
| cos, sin = _rope_cache(T, self.dh, h.device, h.dtype) if self.cfg.rope else (None, None) | |
| q, k, v, beta, alpha, beta_e = self._encode(x, cos, sin) | |
| o = torch.zeros(B, self.H, T, self.dh, device=h.device, dtype=h.dtype) | |
| if self.cfg.use_workspace: | |
| # DeltaProduct: for n_delta>=2, build multi-sub-step k,v,beta tensors | |
| if self.cfg.n_delta >= 2: | |
| # sub-step 0 uses the main projection; sub-steps 1..n_delta-1 use W_kv_extra | |
| ks = [k] # each [B,H,T,d_h] | |
| vs = [v] | |
| bs = [beta] | |
| for i, (wkv, wbeta) in enumerate(zip(self.W_kv_extra, self.W_beta_extra)): | |
| kv_i = wkv(x).view(B, T, self.H, 2, self.dh) # [B,T,H,2,dh] | |
| k_i, v_i = kv_i.unbind(3) # each [B,T,H,dh] | |
| k_i = k_i.transpose(1, 2); v_i = v_i.transpose(1, 2) # [B,H,T,dh] | |
| k_i = _l2(k_i) # unit-norm | |
| b_i = torch.sigmoid(wbeta(x)).transpose(1, 2) * self.cfg.beta_cap # [B,H,T] | |
| ks.append(k_i); vs.append(v_i); bs.append(b_i) | |
| # Stack on a new sub-step axis: [B,H,T,n_delta,d_h] | |
| k_nd = torch.stack(ks, dim=3) | |
| v_nd = torch.stack(vs, dim=3) | |
| b_nd = torch.stack(bs, dim=3) | |
| # phi applied per sub-step; for simplicity phi(k_sub_0) uses main k (already phi'd below) | |
| # For k>=2 we apply phi to sub-step 0 key; extra sub-steps use linear keys (d_h, not d_phi) | |
| # NOTE: n_delta>=2 uses d_h-dimensional state (no feat_map for sub-steps 1+, linear keys) | |
| # Sub-step 0 key goes through phi; but the state dim must be consistent: we use d_h for all | |
| # sub-steps when n_delta>=2 (phi expansion is not combined with n_delta in this impl). | |
| o_delta, _ = chunked_delta(q, k_nd, v_nd, b_nd, alpha, | |
| chunk=self.cfg.chunk, write_mode=self.cfg.write_mode, | |
| beta_e=None, n_delta=self.cfg.n_delta) # [B,H,T,d_h] | |
| else: | |
| # n_delta=1: standard path (byte-identical when surprise_gate=False, with phi and beta_e) | |
| o_delta, _ = chunked_delta(self._phi(q), self._phi(k), v, beta, alpha, | |
| chunk=self.cfg.chunk, write_mode=self.cfg.write_mode, | |
| beta_e=beta_e, | |
| surprise=self.cfg.surprise_gate, | |
| surprise_mode=self.cfg.surprise_mode) # [B,H,T,d_h] | |
| # delta state keyed by phi(q),phi(k) (dim d_phi); values stay d_h -> state [B,H,d_h,d_phi] | |
| if self.state_rms is not None: | |
| o_delta = self.state_rms(o_delta) # per-head RMSNorm over d_h | |
| o = o + o_delta | |
| if self.cfg.use_window: | |
| win_fn = self._window_banded if self.cfg.banded_window else self._window | |
| o = o + win_fn(q, k, v) # window head keeps the LINEAR L2 keys (dim d_h) | |
| o = o.transpose(1, 2).reshape(B, T, d) # merge heads | |
| if self.W_g is not None: | |
| o = o * torch.sigmoid(self.W_g(self.norm1(h))) # gate on block input (pre-conv normed) | |
| h = h + self.W_o(o) | |
| h = h + self.mlp(self.norm2(h)) | |
| return h | |
| # ---- O(1)-per-step inference path (for B5 latency / true streaming) ---- # | |
| def step(self, h_t, state): | |
| """h_t:[B,1,d]; state=(S, ring_k, ring_v, conv_ring, pos). Returns o_t, new_state. O(1).""" | |
| B = h_t.shape[0] | |
| S, rk, rv, cring, pos = state | |
| xin = self.norm1(h_t) # [B,1,d] | |
| if self.kc > 0: | |
| buf = torch.cat([cring, xin], dim=1) # [B,kc,d] | |
| w = self.conv.weight.squeeze(1) # [d,kc] | |
| xc = (buf.transpose(1, 2) * w).sum(-1) + self.conv.bias # [B,d] | |
| x = F.silu(xc)[:, None, :] # [B,1,d] | |
| cring = buf[:, 1:, :] # keep last kc-1 | |
| else: | |
| x = xin | |
| cos, sin = _rope_cache(1, self.dh, h_t.device, h_t.dtype, offset=pos) if self.cfg.rope else (None, None) | |
| q, k, v, beta, alpha, beta_e = self._encode(x, cos, sin) # [B,H,1,dh], beta [B,H,1] | |
| q1, k1, v1 = q[:, :, 0], k[:, :, 0], v[:, :, 0] # [B,H,dh] (linear L2; window keys) | |
| b1 = beta[:, :, 0] # [B,H] write gate beta_w | |
| a1 = alpha[:, :, 0] if alpha is not None else torch.ones_like(b1) | |
| # pre-write read (always uses state from end of previous token) | |
| if self.cfg.n_delta >= 2: | |
| o_delta = torch.einsum("bhij,bhj->bhi", S, q1) # [B,H,d_h] (d_h state for n_delta>=2) | |
| else: | |
| q1p = self._phi(q)[:, :, 0] # [B,H,d_phi] (delta state keys) | |
| o_delta = torch.einsum("bhij,bhj->bhi", S, q1p) # pre-write read S_{t-1} phi(q) | |
| if self.cfg.n_delta >= 2: | |
| # DeltaProduct: apply n_delta sub-steps sequentially (mirrors the chunked/reference form) | |
| # Sub-step 0: main kv + alpha decay | |
| k1_step = k1; v1_step = v1 | |
| Sk = torch.einsum("bhij,bhj->bhi", S, k1_step) | |
| u = b1[..., None] * v1_step - b1[..., None] * (a1[..., None] * Sk) | |
| S = a1[..., None, None] * S + torch.einsum("bhi,bhj->bhij", u, k1_step) | |
| # Sub-steps 1..(n_delta-1): extra projections, no additional alpha decay | |
| x1 = x[:, 0, :] # [B,d] (squeeze T=1 dim) | |
| for wkv, wbeta in zip(self.W_kv_extra, self.W_beta_extra): | |
| kv_j = wkv(x1).view(B, self.H, 2, self.dh) # [B,H,2,dh] | |
| k_j, v_j = kv_j[:, :, 0], kv_j[:, :, 1] # [B,H,dh] | |
| k_j = _l2(k_j) | |
| b_j = torch.sigmoid(wbeta(x1)).view(B, self.H) * self.cfg.beta_cap # [B,H] | |
| Sk_j = torch.einsum("bhij,bhj->bhi", S, k_j) | |
| u_j = b_j[..., None] * (v_j - Sk_j) | |
| S = S + torch.einsum("bhi,bhj->bhij", u_j, k_j) | |
| else: | |
| k1p = self._phi(k)[:, :, 0] # [B,H,d_phi] (delta state keys) | |
| be1 = beta_e[:, :, 0] if beta_e is not None else b1 # [B,H] erase gate beta_e | |
| Sk = torch.einsum("bhij,bhj->bhi", S, k1p) # [B,H,d_h] | |
| # Prediction error (free-energy gradient at S_{t-1}): eps = v - alpha*S*k | |
| eps1 = v1 - a1[..., None] * Sk # [B,H,d_h] | |
| if self.cfg.surprise_gate: | |
| # g_t = 1 + tanh(||eps_t||); same formula as _delta_reference / _surprise_gate 'norm' | |
| # (only 'norm' mode is used in step() — the 'random'/'constant' modes are ablation-only | |
| # and are not wired through step() since they are not meaningful at inference time). | |
| # For step==forward parity with surprise_mode='norm', this is exact. | |
| g1 = (1.0 + torch.tanh(eps1.norm(dim=-1)))[..., None] # [B,H,1] | |
| # Apply g to full write vector: u = g * (beta_w*v - beta_e*alpha*Sk) | |
| u = g1 * (b1[..., None] * v1 - be1[..., None] * (a1[..., None] * Sk)) | |
| else: | |
| # decoupled: u = beta_w * v - beta_e * (alpha * S k) | |
| u = b1[..., None] * v1 - be1[..., None] * (a1[..., None] * Sk) # [B,H,d_h] | |
| S = a1[..., None, None] * S + torch.einsum("bhi,bhj->bhij", u, k1p) # [B,H,d_h,d_phi] | |
| if self.state_rms is not None: | |
| o_delta = self.state_rms(o_delta) # per-head RMSNorm [B,H,d_h], mirrors forward | |
| # window ring | |
| rk = torch.cat([rk, k1[:, :, None]], dim=2)[:, :, -self.cfg.window:] | |
| rv = torch.cat([rv, v1[:, :, None]], dim=2)[:, :, -self.cfg.window:] | |
| sc = torch.einsum("bhd,bhwd->bhw", q1, rk) * self.win_scale | |
| aw = torch.softmax(sc, dim=-1) | |
| o_win = torch.einsum("bhw,bhwd->bhd", aw, rv) | |
| o = (o_delta + o_win).reshape(B, 1, -1) | |
| if self.W_g is not None: | |
| o = o * torch.sigmoid(self.W_g(self.norm1(h_t))) # gate on block input, mirrors forward | |
| h = h_t + self.W_o(o) | |
| h = h + self.mlp(self.norm2(h)) | |
| return h, (S, rk, rv, cring, pos + 1) | |
| class PrizmaSeqLM(nn.Module): | |
| def __init__(self, cfg: PrizmaSeqConfig): | |
| super().__init__() | |
| self.cfg = cfg | |
| self.tok = nn.Embedding(cfg.vocab, cfg.d_model) | |
| self.pos = nn.Embedding(cfg.max_len, cfg.d_model) if cfg.learned_pos else None | |
| self.blocks = nn.ModuleList([PrizmaSeqBlock(cfg) for _ in range(cfg.n_layers)]) | |
| self.nf = RMSNorm(cfg.d_model) | |
| self.head = nn.Linear(cfg.d_model, cfg.vocab, bias=False) | |
| self.head.weight = self.tok.weight | |
| self.apply(self._init) | |
| def _init(self, m): | |
| if isinstance(m, nn.Linear): | |
| nn.init.normal_(m.weight, std=0.02) | |
| if m.bias is not None: | |
| nn.init.zeros_(m.bias) | |
| elif isinstance(m, nn.Embedding): | |
| nn.init.normal_(m.weight, std=0.02) | |
| def forward(self, idx): | |
| B, T = idx.shape | |
| h = self.tok(idx) | |
| if self.pos is not None: | |
| h = h + self.pos(torch.arange(T, device=idx.device))[None] | |
| for blk in self.blocks: | |
| h = blk(h) | |
| return self.head(self.nf(h)) | |
| def init_state(self, batch, device): | |
| st = [] | |
| kc1 = max(self.cfg.short_conv - 1, 0) | |
| # For n_delta>=2, state uses d_h-dim keys (no phi expansion); for n_delta=1 uses d_phi | |
| state_k_dim = self.cfg.d_h if self.cfg.n_delta >= 2 else self.cfg.d_phi | |
| for _ in self.blocks: | |
| S = torch.zeros(batch, self.cfg.n_heads, self.cfg.d_h, state_k_dim, device=device) | |
| rk = torch.zeros(batch, self.cfg.n_heads, 0, self.cfg.d_h, device=device) | |
| rv = torch.zeros(batch, self.cfg.n_heads, 0, self.cfg.d_h, device=device) | |
| cring = torch.zeros(batch, kc1, self.cfg.d_model, device=device) | |
| st.append((S, rk, rv, cring, 0)) | |
| return st | |
| def step(self, tok, state): | |
| """tok:[B,1] -> (logits[B,1,V], new_state). O(1) in sequence length.""" | |
| h = self.tok(tok) | |
| if self.pos is not None: | |
| p = state[0][4] if state else 0 | |
| h = h + self.pos(torch.tensor([p], device=tok.device))[None] | |
| new = [] | |
| for blk, st in zip(self.blocks, state): | |
| h, st2 = blk.step(h, st) | |
| new.append(st2) | |
| return self.head(self.nf(h)), new | |
| def prizma_seq_factory(d_model=64, n_layers=2, n_heads=2, **kw): | |
| def f(vocab, max_len): | |
| return PrizmaSeqLM(PrizmaSeqConfig(vocab=vocab, d_model=d_model, n_layers=n_layers, | |
| n_heads=n_heads, max_len=max_len + 8, **kw)) | |
| return f | |
| if __name__ == "__main__": | |
| from .common import param_count, get_device | |
| dev = get_device() | |
| # O(1) GUARD (protocol guardrail): for ALL feat_map settings the streaming step() must equal | |
| # the parallel forward() to <1e-4, AND param_count must be identical (feature map = 0 params). | |
| ref_params = None | |
| for feat in ("none", "quad2", "quad2_lowrank"): | |
| cfg = PrizmaSeqConfig(vocab=64, d_model=64, n_layers=2, n_heads=2, feat_map=feat) | |
| m = PrizmaSeqLM(cfg).to(dev) | |
| x = torch.randint(0, 64, (2, 48), device=dev) | |
| y = m(x) | |
| m.train(False) | |
| st = m.init_state(2, dev) | |
| outs = [] | |
| for t in range(x.shape[1]): | |
| lg, st = m.step(x[:, t:t + 1], st) | |
| outs.append(lg) | |
| yo = torch.cat(outs, dim=1) | |
| d = (y - yo).abs().max().item() | |
| p = param_count(m) | |
| if ref_params is None: | |
| ref_params = p | |
| param_ok = (p == ref_params) | |
| print(f"[feat_map={feat:<12} d_phi={cfg.d_phi:<4}] forward {tuple(y.shape)} " | |
| f"params {p} {'(MATCH)' if param_ok else '(MISMATCH!)'} " | |
| f"step-vs-forward max|d|={d:.2e} {'OK' if d < 1e-4 else 'MISMATCH'}") | |
| # SURPRISE GATE O(1) guard: surprise_gate=True step()==forward() < 1e-4 | |
| cfg_s = PrizmaSeqConfig(vocab=64, d_model=64, n_layers=2, n_heads=2, | |
| feat_map='none', surprise_gate=True, surprise_mode='norm') | |
| m_s = PrizmaSeqLM(cfg_s).to(dev) | |
| m_s.train(False) | |
| torch.manual_seed(7) | |
| x_s = torch.randint(0, 64, (2, 48), device=dev) | |
| y_s = m_s(x_s) | |
| st_s = m_s.init_state(2, dev) | |
| outs_s = [] | |
| for t in range(x_s.shape[1]): | |
| lg_s, st_s = m_s.step(x_s[:, t:t + 1], st_s) | |
| outs_s.append(lg_s) | |
| yo_s = torch.cat(outs_s, dim=1) | |
| d_s = (y_s - yo_s).abs().max().item() | |
| print(f"[surprise_gate=True norm] step-vs-forward max|d|={d_s:.2e} {'OK' if d_s < 1e-4 else 'MISMATCH'}") | |