1"""PrismeModel: backbone + projected source bank + ephemeral workspace + sparse editor + controller + readout.23Contracts implemented here (see docs/architecture_spec.md):4 * explicit bypass: max_budget == 0 (or actions None) returns the backbone head without touching the workspace;5 * causality: the bank visible to position t is 1..t; nothing from t+1 onward enters S, the editor or the controller;6 * locality: only the slots of the chosen action are written (bitwise identity elsewhere);7 * the editor never modifies H, K, V or z (they are inputs only);8 * the deployed controller scores actions from target-free features and never executes candidate edits.9"""10from __future__ import annotations1112from dataclasses import dataclass13from typing import Optional1415import torch16import torch.nn as nn17import torch.nn.functional as F1819from .attention import RMSNorm, LayerKVCache20from .backbone import Backbone21from .config import PrismeConfig22from .controller import ValueController23from .editor import SparseEditor24from .readout import DifferentialReadout25from .workspace import action_table, action_group_size, init_workspace262728@dataclass29class EncodedContext:30 """Everything the workspace needs for a batch of query positions. All tensors are read-only for the editor."""31 H: torch.Tensor # [B,Tq,d] backbone residual stream at the query positions32 K: torch.Tensor # [B,Tb,w] source bank keys (causally readable)33 V: torch.Tensor # [B,Tb,w]34 z: torch.Tensor # [B,Tq,w]35 bank_mask: torch.Tensor # [B,Tq,Tb] bool36 S0: torch.Tensor # [B,Tq,M,w]37 ctrl_in: Optional[torch.Tensor] = None # controller input ([B,Tq,w] z or [B,Tq,d] normalised H)38 entropy: Optional[torch.Tensor] = None # [B,Tq] backbone predictive entropy at budget 0 (if cfg.ctrl_use_entropy)394041class PrismeModel(nn.Module):42 def __init__(self, cfg: PrismeConfig):43 super().__init__()44 cfg.validate()45 self.cfg = cfg46 self.backbone = Backbone(cfg)47 if cfg.use_workspace:48 self.bank_norm = RMSNorm(cfg.d_model, cfg.norm_eps)49 self.W_K = nn.Linear(cfg.d_model, cfg.w, bias=False)50 self.W_V = nn.Linear(cfg.d_model, cfg.w, bias=False)51 self.P = nn.Linear(cfg.d_model, cfg.w, bias=False)52 self.slot_embed = nn.Parameter(torch.randn(cfg.n_slots, cfg.w) * cfg.init_std)53 self.editor = SparseEditor(cfg)54 self.readout = DifferentialReadout(cfg)55 self.controller = ValueController(cfg)56 for lin in (self.W_K, self.W_V, self.P):57 nn.init.normal_(lin.weight, std=cfg.init_std)58 self.register_buffer("action_slots", action_table(cfg), persistent=False)59 self.register_buffer("action_k", action_group_size(cfg), persistent=False)6061 # ------------------------------------------------------------------ parameter groups62 def backbone_parameters(self):63 return list(self.backbone.parameters())6465 def workspace_parameters(self):66 names = ("bank_norm", "W_K", "W_V", "P", "slot_embed", "editor", "readout")67 return [p for n, p in self.named_parameters() if n.split(".")[0] in names]6869 def controller_parameters(self):70 return list(self.controller.parameters())7172 # ------------------------------------------------------------------ encoding73 def encode(self, tokens: torch.Tensor) -> EncodedContext:74 """Teacher-forced parallel encoding of all positions; bank mask is strictly causal (key j <= query t)."""75 H = self.backbone(tokens)76 return self.context_from_H(H, H)7778 def context_from_H(self, H_bank: torch.Tensor, H_query: torch.Tensor, query_pos: Optional[torch.Tensor] = None,79 bank_len: Optional[torch.Tensor] = None) -> EncodedContext:80 """H_bank: [B,Tb,d] history (incl. the query positions); H_query: [B,Tq,d].81 query_pos: [B,Tq] absolute positions of the queries inside the bank (default: arange).82 bank_len: [B] valid bank length per request (for padded batches)."""83 B, Tb, _ = H_bank.shape84 Tq = H_query.shape[1]85 hb = self.bank_norm(H_bank)86 K, V = self.W_K(hb), self.W_V(hb)87 z = self.P(self.bank_norm(H_query))88 if query_pos is None:89 query_pos = torch.arange(Tq, device=H_bank.device).view(1, Tq).expand(B, Tq)90 key_pos = torch.arange(Tb, device=H_bank.device)91 bank_mask = key_pos.view(1, 1, Tb) <= query_pos.unsqueeze(-1)92 if bank_len is not None:93 bank_mask = bank_mask & (key_pos.view(1, 1, Tb) < bank_len.view(B, 1, 1))94 S0 = init_workspace(self.slot_embed, z)95 hq = self.bank_norm(H_query)96 ctrl_in = hq if self.cfg.ctrl_input == "h" else z97 entropy = None98 if self.cfg.ctrl_use_entropy:99 p = torch.softmax(self.backbone.head(H_query).float(), -1)100 entropy = -(p * torch.log(p.clamp_min(1e-12))).sum(-1)101 return EncodedContext(H=H_query, K=K, V=V, z=z, bank_mask=bank_mask, S0=S0, ctrl_in=ctrl_in, entropy=entropy)102103 # ------------------------------------------------------------------ one action104 def apply_action(self, S: torch.Tensor, ctx: EncodedContext, action: torch.Tensor, step, remaining):105 B, Tq = action.shape106 if not torch.is_tensor(step):107 step = torch.full((B, Tq), int(step), device=S.device, dtype=torch.long)108 if not torch.is_tensor(remaining):109 remaining = torch.full((B, Tq), int(remaining), device=S.device, dtype=torch.long)110 return self.editor(S, ctx.K, ctx.V, ctx.bank_mask, action, step, remaining)111112 # ------------------------------------------------------------------ readout113 def logits_from_state(self, S: torch.Tensor, ctx: EncodedContext, edited: Optional[torch.Tensor] = None):114 """h_corrected = H + U(r(S) - r(S0)). If `edited` ([B,Tq] bool) is given, positions that were never115 edited use the backbone head directly (exact bypass)."""116 delta = self.readout(S, ctx.S0, ctx.z)117 if edited is not None:118 delta = torch.where(edited.unsqueeze(-1), delta, torch.zeros_like(delta))119 return self.backbone.head(ctx.H + delta), delta120121 # ------------------------------------------------------------------ forced-action forward (training / tests)122 def forward(self, tokens: torch.Tensor, actions: Optional[torch.Tensor] = None, return_state: bool = False):123 """tokens: [B,T]; actions: [B,T,R] long (STOP for unused steps) or None (-> exact bypass).124 Returns dict(logits=[B,T,V], ...)."""125 cfg = self.cfg126 if actions is None or not cfg.use_workspace or actions.shape[-1] == 0:127 H = self.backbone(tokens)128 return {"logits": self.backbone.head(H), "H": H, "n_edits": None}129 R = actions.shape[-1]130 assert R <= cfg.r_max, f"hard action cap exceeded: {R} > r_max={cfg.r_max}"131 ctx = self.encode(tokens)132 S = ctx.S0133 n_edits = torch.zeros(actions.shape[:2], dtype=torch.long, device=tokens.device)134 stopped = torch.zeros(actions.shape[:2], dtype=torch.bool, device=tokens.device)135 infos = []136 for r in range(R):137 a = actions[:, :, r]138 a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a) # STOP is absorbing139 nonstop = a != cfg.stop_action140 stopped = stopped | ~nonstop141 remaining = (R - r) - 1142 S, info = self.apply_action(S, ctx, a, r, remaining)143 n_edits = n_edits + nonstop.long()144 infos.append(info)145 edited = n_edits > 0146 logits, delta = self.logits_from_state(S, ctx, edited)147 out = {"logits": logits, "H": ctx.H, "n_edits": n_edits, "delta_out": delta}148 if return_state:149 out.update({"ctx": ctx, "S_final": S, "infos": infos})150 return out151152 # ------------------------------------------------------------------ controller-driven forward (teacher forced)153 @torch.no_grad()154 def forward_policy(self, tokens: torch.Tensor, max_budget: int, cost: torch.Tensor,155 lambda_cost: float | None = None, kappa_risk: float | None = None, return_state: bool = False):156 """Parallel teacher-forced rollout of the deployed controller at every position (evaluation only).157 cost: [A] analytic cost proxy per nonterminal action. Returns logits and the action trace [B,T,max_budget]."""158 cfg = self.cfg159 if max_budget == 0:160 H = self.backbone(tokens)161 return {"logits": self.backbone.head(H), "actions": None, "n_edits": torch.zeros(tokens.shape, dtype=torch.long, device=tokens.device)}162 max_budget = min(max_budget, cfg.r_max)163 ctx = self.encode(tokens)164 S = ctx.S0165 B, T = tokens.shape166 stopped = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)167 trace = torch.full((B, T, max_budget), cfg.stop_action, dtype=torch.long, device=tokens.device)168 n_edits = torch.zeros(B, T, dtype=torch.long, device=tokens.device)169 for r in range(max_budget):170 step = torch.full((B, T), r, device=tokens.device, dtype=torch.long)171 remaining = torch.full((B, T), max_budget - r - 1, device=tokens.device, dtype=torch.long)172 preds = self.controller(ctx.ctrl_in, S, step, remaining, ctx.entropy)173 a = self.controller.select(preds, cost, lambda_cost=lambda_cost, kappa_risk=kappa_risk)174 a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a)175 stopped = stopped | (a == cfg.stop_action)176 trace[:, :, r] = a177 S, _ = self.apply_action(S, ctx, a, step, remaining)178 n_edits = n_edits + (a != cfg.stop_action).long()179 if bool(stopped.all()):180 break181 logits, delta = self.logits_from_state(S, ctx, n_edits > 0)182 out = {"logits": logits, "actions": trace, "n_edits": n_edits}183 if return_state:184 out.update({"ctx": ctx, "S_final": S})185 return out186