| """ |
| 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 |
|
|
| |
| |
| EVENT_KEYS = [ |
| "settling", |
| "churning", |
| "converged", |
| "reorganising", |
| "amplifying", |
| "attn_dominant", |
| "local_dominant", |
| "mlp_dominant", |
| "memory_active", |
| "late_shift", |
| "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 |
|
|
|
|
| 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) |
| |
| 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) |
|
|
|
|
| |
| |
| |
| @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) |
| med = dn.median(dim=1, keepdim=True).values |
| hn = torch.stack([h.flatten(1).norm(dim=1) for h in traj], 1) |
| 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() |
| |
| 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()) |
|
|
| |
| 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] |
|
|
| |
| 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)} |
|
|
