1"""Global latent-recurrence adapter (Huginn/Ouro/ELR-*inspired*, simplified): the whole top representation is2refined r times by ONE shared block operating at workspace width w, with a zero-initialised residual back to d.34Same frozen backbone, same forced mixed-budget training recipe and same budget semantics as PRISME; the only5difference is that every "action" rewrites the entire latent (global, dense) instead of 1-2 of 16 slots, and6there is no slot selection to learn. Disclosed simplifications: no prelude/coda retraining, no truncated7backprop through time, adapter-level recurrence on a frozen backbone rather than from-scratch recurrent pretraining.8Parameter count is matched to the PRISME workspace (ffn width chosen in the experiment config)."""9from __future__ import annotations1011import torch12import torch.nn as nn1314from prisme.attention import RMSNorm, SwiGLU15from prisme.backbone import Backbone16from prisme.config import PrismeConfig17import torch.nn.functional as F181920class RecurrentBlock(nn.Module):21 def __init__(self, w: int, heads: int, ffn: int, eps: float):22 super().__init__()23 self.h, self.hd = heads, w // heads24 self.n1 = RMSNorm(w, eps); self.q = nn.Linear(w, w, bias=False); self.k = nn.Linear(w, w, bias=False)25 self.v = nn.Linear(w, w, bias=False); self.o = nn.Linear(w, w, bias=False)26 self.n2 = RMSNorm(w, eps); self.ffn = SwiGLU(w, ffn)2728 def forward(self, x, ctx_mask):29 B, T, w = x.shape30 xn = self.n1(x)31 q = self.q(xn).view(B, T, self.h, self.hd).transpose(1, 2)32 k = self.k(xn).view(B, T, self.h, self.hd).transpose(1, 2)33 v = self.v(xn).view(B, T, self.h, self.hd).transpose(1, 2)34 a = F.scaled_dot_product_attention(q, k, v, attn_mask=ctx_mask)35 x = x + self.o(a.transpose(1, 2).reshape(B, T, w))36 return x + self.ffn(self.n2(x))373839class GlobalRecurrenceModel(nn.Module):40 """Interface-compatible with PrismeModel.forward(tokens, actions): any nonterminal action = one global41 refinement step; STOP = no step. n_edits semantics identical."""4243 def __init__(self, cfg: PrismeConfig, ffn_w: int | None = None, heads: int | None = None):44 super().__init__()45 self.cfg = cfg46 self.backbone = Backbone(cfg)47 w = cfg.w48 self.bank_norm = RMSNorm(cfg.d_model, cfg.norm_eps)49 self.down = nn.Linear(cfg.d_model, w, bias=False)50 self.block = RecurrentBlock(w, heads or cfg.n_editor_heads, ffn_w or cfg.ffn_w, cfg.norm_eps)51 self.ctx_proj = nn.Linear(2, w, bias=True)52 self.up = nn.Linear(w, cfg.d_model, bias=False)53 nn.init.zeros_(self.ctx_proj.weight); nn.init.zeros_(self.ctx_proj.bias)54 if cfg.readout_zero_init:55 nn.init.zeros_(self.up.weight)56 else:57 nn.init.normal_(self.up.weight, std=cfg.readout_init_std)58 self.controller = nn.Identity() # placeholder so parameter-group helpers work5960 def forward(self, tokens, actions=None, return_state=False):61 cfg = self.cfg62 H = self.backbone(tokens)63 if actions is None or actions.shape[-1] == 0:64 return {"logits": self.backbone.head(H), "H": H, "n_edits": None}65 B, T, R = actions.shape66 x0 = self.down(self.bank_norm(H))67 x = x068 mask = torch.ones(T, T, dtype=torch.bool, device=H.device).tril().view(1, 1, T, T)69 stopped = torch.zeros(B, T, dtype=torch.bool, device=H.device)70 n_edits = torch.zeros(B, T, dtype=torch.long, device=H.device)71 for r in range(R):72 a = torch.where(stopped, torch.full_like(actions[:, :, r], cfg.stop_action), actions[:, :, r])73 act = (a != cfg.stop_action)74 stopped |= ~act75 ctx = torch.stack([torch.full((B, T), r / cfg.r_max, device=H.device), torch.full((B, T), (R - r - 1) / cfg.r_max, device=H.device)], -1)76 xn = self.block(x + self.ctx_proj(ctx), mask)77 x = torch.where(act.unsqueeze(-1), xn, x) # positions that stopped keep their state78 n_edits += act.long()79 delta = self.up(x - x0) * self.cfg.readout_scale80 edited = n_edits > 081 delta = torch.where(edited.unsqueeze(-1), delta, torch.zeros_like(delta))82 return {"logits": self.backbone.head(H + delta), "H": H, "n_edits": n_edits, "delta_out": delta}83