""" Internal-state transcription (Phase 9). A READ-ONLY decoder over the recurrent trajectory. It observes h0..hn and the state changes Δh_t, and emits a description of the *computational events* the model went through -- not a human-style chess explanation. Design rules enforced here: * The decoder receives DETACHED states. `assert_no_gradient_leak()` proves the chess pathway cannot be reshaped by the transcription loss. * It is trained (when trained at all) on labels derived from the model's OWN measurable internal events, never on human commentary. So a claim like "the transcript reflects internal computation" is testable, and tested in experiments/exp_faithfulness.py. * Two variants exist for the required control experiment: TrajectoryDecoder - sees h0..hn and Δh_t FinalStateDecoder - sees only hn (baseline) If trajectory access adds nothing, the two score equally. """ from __future__ import annotations from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F # The vocabulary of *computational events* the decoder predicts. These are # measurable properties of the trajectory, not chess concepts. EVENT_KEYS = [ "settling", # |Δh| decreasing "churning", # |Δh| roughly constant and large "converged", # |Δh| ~ 0 "reorganising", # direction of Δh flips vs previous step "amplifying", # |h| growing "attn_dominant", # router mass on global attention "local_dominant", # router mass on local mixing "mlp_dominant", "memory_active", # router mass on memory op "late_shift", # top-1 move changed in the second half of the trajectory "high_confidence", "low_confidence", ] @dataclass class ThoughtConfig: d_model: int = 64 d_hidden: int = 48 n_events: int = len(EVENT_KEYS) use_delta: bool = True final_only: bool = False # the control baseline class ThoughtDecoder(nn.Module): """Small GRU over the (detached) trajectory -> per-step event logits. Parameters are reported SEPARATELY from the chess model: the decoder is an observation instrument and must not be counted inside the chess budget. """ def __init__(self, cfg: ThoughtConfig): super().__init__() self.cfg = cfg d = cfg.d_model in_dim = d * (2 if (cfg.use_delta and not cfg.final_only) else 1) # pooled square states (mean+max) -> compact step descriptor self.inp = nn.Linear(in_dim * 2, cfg.d_hidden) self.gru = nn.GRUCell(cfg.d_hidden, cfg.d_hidden) self.out = nn.Linear(cfg.d_hidden, cfg.n_events) def param_count(self) -> int: return sum(p.numel() for p in self.parameters()) def _describe(self, h, h_prev): pooled = torch.cat([h.mean(1), h.amax(1)], -1) if self.cfg.use_delta and not self.cfg.final_only: d = h - (h_prev if h_prev is not None else torch.zeros_like(h)) pooled = torch.cat([pooled, d.mean(1), d.amax(1)], -1) return pooled def forward(self, trajectory): """trajectory: list of [B,64,d] DETACHED tensors. Returns [B,T,n_events].""" assert all(not t.requires_grad for t in trajectory), \ "trajectory must be detached before reaching the thought decoder" if self.cfg.final_only: trajectory = [trajectory[-1]] B = trajectory[0].shape[0] hid = torch.zeros(B, self.cfg.d_hidden, device=trajectory[0].device) outs = [] prev = None for h in trajectory: x = F.gelu(self.inp(self._describe(h, prev))) hid = self.gru(x, hid) outs.append(self.out(hid)) prev = h return torch.stack(outs, 1) # --------------------------------------------------------------------------- # Ground-truth events, measured from the model's own trajectory # --------------------------------------------------------------------------- @torch.no_grad() def extract_events(out: dict, ops: list[str]) -> torch.Tensor: """Derive per-step binary event labels from a traced forward pass. IMPORTANT (fixed after measurement): the first version used absolute thresholds and 7/12 events never fired, so a decoder scored 99% merely by predicting constants (majority baseline 94.5%). Events are now defined RELATIVE to each example's own trajectory (median splits), which guarantees a non-degenerate label distribution and makes the trajectory-vs-final-state comparison meaningful. """ traj = out["trajectory"] coll = out.get("collect", {}) or {} B = traj[0].shape[0] T = len(traj) ev = torch.zeros(B, T, len(EVENT_KEYS)) idx = {k: i for i, k in enumerate(EVENT_KEYS)} deltas = [None] + [traj[i] - traj[i - 1] for i in range(1, T)] dn = torch.stack([d.flatten(1).norm(dim=1) for d in deltas[1:]], 1) \ if T > 1 else torch.zeros(B, 1) # [B,T-1] med = dn.median(dim=1, keepdim=True).values hn = torch.stack([h.flatten(1).norm(dim=1) for h in traj], 1) # [B,T] hmed = hn.median(dim=1, keepdim=True).values for t in range(T): if t >= 1: d_t = dn[:, t - 1] ev[:, t, idx["churning"]] = (d_t > med.squeeze(1)).float() ev[:, t, idx["converged"]] = (d_t <= med.squeeze(1)).float() if t >= 2: ev[:, t, idx["settling"]] = (d_t < dn[:, t - 2]).float() cos = F.cosine_similarity(deltas[t].flatten(1), deltas[t - 1].flatten(1), dim=1) ev[:, t, idx["reorganising"]] = (cos < 0.9).float() ev[:, t, idx["amplifying"]] = (hn[:, t] > hmed.squeeze(1)).float() i = t - 1 if 0 <= i < len(coll.get("router", [])): w = coll["router"][i] top = w.argmax(-1) for j, op in enumerate(ops): key = {"attn": "attn_dominant", "local": "local_dominant", "mlp": "mlp_dominant", "mem": "memory_active"}.get(op) if key: ev[:, t, idx[key]] = (top == j).float() # memory_active also fires on non-trivial memory routing mass if "mem" in ops: ev[:, t, idx["memory_active"]] = torch.maximum( ev[:, t, idx["memory_active"]], (w[:, ops.index("mem")] > 1.0 / max(len(ops), 1)).float()) # confidence is measured against the batch median, so both classes occur if "logits" in out: p = torch.softmax(out["logits"], -1) conf = p.max(-1).values cmed = conf.median() ev[:, :, idx["high_confidence"]] = (conf > cmed).float()[:, None] ev[:, :, idx["low_confidence"]] = (conf <= cmed).float()[:, None] # late_shift: did the argmax move change between the first and second half? if "refine_logits" in out and out["refine_logits"] and \ out["refine_logits"][0] is not None and len(out["refine_logits"]) > 1: first = out["refine_logits"][0].argmax(-1) last = out["refine_logits"][-1].argmax(-1) ev[:, :, idx["late_shift"]] = (first != last).float()[:, None] return ev def render_transcript(event_probs, threshold=0.5, steps=None) -> str: """Human-readable transcription of ONE example: [T,n_events] -> text.""" lines = [] T = event_probs.shape[0] sel = steps or range(T) for t in sel: if t >= T: continue active = [(EVENT_KEYS[i], float(event_probs[t, i])) for i in range(len(EVENT_KEYS)) if event_probs[t, i] > threshold] active.sort(key=lambda x: -x[1]) if not active: desc = "no salient computational event" else: desc = ", ".join(f"{k} ({p:.2f})" for k, p in active[:4]) lines.append(f"Step {t}: {desc}") return "\n".join(lines) def assert_no_gradient_leak(model, decoder, batch, steps=4): """PROOF that the thought pathway cannot reshape the chess representation. Runs a transcription loss backward and checks every chess parameter has no gradient. Used as a test, not just documentation. """ model.zero_grad(set_to_none=True) decoder.zero_grad(set_to_none=True) out = model(**batch, steps=steps, trace=True) traj = [t.detach() for t in out["trajectory"]] logits = decoder(traj) tgt = extract_events(out, model.core.ops if hasattr(model, "core") else []) loss = F.binary_cross_entropy_with_logits(logits, tgt[:, :logits.shape[1]]) loss.backward() leaked = [n for n, p in model.named_parameters() if p.grad is not None and float(p.grad.abs().sum()) > 0] got_grad = [n for n, p in decoder.named_parameters() if p.grad is not None and float(p.grad.abs().sum()) > 0] return {"leaked_chess_params": leaked, "decoder_params_with_grad": len(got_grad), "loss": float(loss)}