initial upload: NLA demo (playground/trace/steer)

README.md

@@ -1,13 +1,19 @@
 ---
-title: Srt Nla Demo
-emoji: 📉
-colorFrom: yellow
+title: SRT-NLA Demo
+emoji: 🧠
+colorFrom: indigo
 colorTo: pink
 sdk: gradio
-sdk_version: 6.14.0
-python_version: '3.12'
-app_file: app.py
+sdk_version: 4.44.0
+app_file: demo/app.py
 pinned: false
+license: apache-2.0
+hardware: zero-a10g
+short_description: Ask a frozen LM what it is thinking, in plain English.
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SRT-NLA Demo
+
+Three views of the trained Activation Verbalizer over Qwen-2.5-7B,
+Llama-3.2-3B, and Gemma-2-2B (all frozen). See `demo/README.md` for
+details and the [companion paper](https://github.com/space-bacon/SRT/blob/main/paper_nla.md).

demo/README.md

@@ -0,0 +1,96 @@
+# SRT-NLA Demo
+
+A single Gradio app exposing three views of the trained Activation
+Verbalizer (AV) sitting on top of three frozen backbones (Qwen-2.5-7B,
+Llama-3.2-3B, Gemma-2-2B):
+
+1. **Playground** - prompt -> hidden state -> ranked candidate
+   verbalisations with raw and centred round-trip fidelity.
+2. **Live thought trace** - watch the backbone generate token by token
+   while the AV periodically verbalises its running hidden state.
+3. **Steer by editing** - verbalise the prompt's hidden into plain
+   English, edit the text, re-encode it, and patch the difference into
+   layer L while the backbone generates a new continuation.
+
+All three tabs share a backbone selector at the top and a lazy in-memory
+cache of loaded backbones. The first call on a given backbone downloads
+the AV checkpoint and the centring pool from HuggingFace.
+
+## Run locally on a GPU box
+
+```bash
+cd /workspace/srt-adapter
+pip install -r demo/requirements.txt
+python demo/app.py
+```
+
+By default the app launches on `http://0.0.0.0:7860` and serves all
+three backbones lazily. Set `NLA_DEFAULT_BACKBONE` to one of
+`qwen2.5-7b`, `llama-3.2-3b`, `gemma-2-2b` to change the initial
+selection.
+
+Hardware (bf16):
+
+| Backbone        | VRAM  | First-load time |
+|-----------------|------:|----------------:|
+| Qwen-2.5-7B     | ~16GB | ~60s            |
+| Llama-3.2-3B    | ~8GB  | ~30s            |
+| Gemma-2-2B      | ~6GB  | ~20s            |
+
+## Deploy as an HF Space
+
+The app supports both standard GPU Spaces and **ZeroGPU**. When the
+`spaces` package is importable and `SPACES_ZERO_GPU` is set in the
+environment (HF Spaces does this automatically on a ZeroGPU tier), the
+three callbacks are wrapped with `@spaces.GPU` and request CUDA on
+demand. Otherwise they run on whatever device `torch.cuda.is_available()`
+reports.
+
+1. Create a new Space (Gradio template). For ZeroGPU pick the
+   *zero-a10g* tier (free); otherwise pick *A10G* or larger.
+2. Mirror this `demo/` directory and the `srt/` package into the Space
+   repo (the app imports `srt.nla.verbalizer`).
+3. Add `HF_TOKEN` as a secret if any of the backbones is gated for your
+   account (Llama-3.2-3B is gated).
+4. Set `app_file: demo/app.py` in the Space `README.md` frontmatter.
+
+Notes on ZeroGPU:
+
+- First call on a backbone loads weights from HuggingFace; the duration
+  budget on each call is 120s (180s for the steer/trace tabs). Qwen-7B
+  cold-start can saturate this — prefer Gemma-2-2B as the default and
+  only switch to Qwen on Pro hardware.
+- Models stay resident in the persistent control process between calls,
+  so subsequent invocations only pay the GPU-attach overhead.
+
+## How it works
+
+Each `NLAPipeline` holds:
+
+- the frozen backbone (`AutoModelForCausalLM`, `requires_grad=False`),
+- the AV (`srt.nla.verbalizer.ActivationVerbalizer`) loaded from the
+  matching `RiverRider/srt-nla-av-*` HF repo,
+- a centring pool (2K real layer-L last-token hiddens from the matching
+  `RiverRider/srt-nla-targets-*` HF dataset) used to compute centred
+  cosine fidelity that corrects for the backbone's anisotropic mean.
+
+The verbaliser is invoked through `av.generate(v, ...)`, which prepends
+`v` (projected) plus a learned static prefix to the backbone's
+input-embeddings stream and runs greedy or sampled generation. The
+backbone's weights are never touched.
+
+Steering is a single `register_forward_hook` on
+`backbone.model.layers[L]` that adds the layer-L residual difference
+between the re-encoded edited text and the original prompt; once the
+generation is finished the hook is removed.
+
+## Caveats
+
+- The AV is trained on layer-L last-token hiddens of natural-language
+  prompts. Out-of-distribution hiddens (e.g. mid-token, code, foreign
+  scripts) will produce lower fidelity verbalisations.
+- Best-of-K with the centred metric is the default ranking. Raw cosine
+  is also reported for comparison; raw is biased upward by the
+  backbone's anisotropic mean (see paper §3).
+- Steering with `|alpha| > 1` often pushes the model off-distribution
+  fast; small values (0.1 - 0.5) are usually more interesting.

demo/app.py

@@ -0,0 +1,275 @@
+"""Gradio app for the SRT-NLA demo.
+
+Three tabs over a shared backbone selector:
+
+  1. Playground          - prompt -> hidden -> ranked candidate verbalisations
+  2. Live Thought Trace  - generate token by token, show AV verbalisation
+                           every N tokens alongside each token
+  3. Steer by Editing    - verbalise the prompt's hidden state, edit the
+                           text, re-encode it, run the model again with
+                           the difference patched into layer L
+
+Run locally on a GPU box:
+    pip install -r demo/requirements.txt
+    python demo/app.py
+
+Or deploy as an HF Space with hardware = A10G or larger; the Qwen-7B
+backbone needs ~16 GB of VRAM in bf16, the smaller backbones less.
+"""
+from __future__ import annotations
+
+import logging
+import os
+
+import gradio as gr
+import torch
+
+from nla_pipeline import BACKBONES, NLAPipeline
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
+log = logging.getLogger("nla_demo_app")
+
+# Optional ZeroGPU support. On HF Spaces with `hardware: zero-*`, `spaces`
+# is preinstalled and exposes the @spaces.GPU decorator that grants the
+# decorated function ephemeral CUDA access. Off-Spaces it is a no-op.
+try:
+    import spaces  # type: ignore
+
+    _ON_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU")) or hasattr(spaces, "GPU")
+
+    def _gpu(duration: int = 120):
+        if _ON_ZEROGPU:
+            return spaces.GPU(duration=duration)
+        return lambda fn: fn
+except ImportError:  # pragma: no cover - local dev path
+    _ON_ZEROGPU = False
+
+    def _gpu(duration: int = 120):
+        return lambda fn: fn
+
+
+DEFAULT_BACKBONE = os.environ.get("NLA_DEFAULT_BACKBONE", "gemma-2-2b")
+# On ZeroGPU torch.cuda.is_available() is False at import time but becomes
+# True inside @spaces.GPU functions. Pin device to cuda when we know the
+# Space has a GPU slice.
+DEVICE = "cuda" if (torch.cuda.is_available() or _ON_ZEROGPU) else "cpu"
+log.info("device=%s zero_gpu=%s", DEVICE, _ON_ZEROGPU)
+
+# Lazy cache: load each backbone the first time it's selected.
+_PIPES: dict[str, NLAPipeline] = {}
+
+
+def get_pipe(key: str) -> NLAPipeline:
+    if key not in _PIPES:
+        log.info("First use of backbone %s; loading.", key)
+        _PIPES[key] = NLAPipeline(BACKBONES[key], device=DEVICE)
+    return _PIPES[key]
+
+
+# ---------------------------------------------------------------------------
+# Tab callbacks
+# ---------------------------------------------------------------------------
+
+@_gpu(duration=120)
+def cb_playground(backbone_key: str, prompt: str, K: int, temperature: float):
+    if not prompt.strip():
+        return "_(enter a prompt)_", ""
+    pipe = get_pipe(backbone_key)
+    v = pipe.extract_hidden(prompt, token_index=-1)
+    ranked = pipe.verbalize(v, K=int(K), temperature=float(temperature))
+    rows = [
+        f"| {i+1} | {cen:.3f} | {raw:.3f} | {txt[:160].replace(chr(10), ' ')} |"
+        for i, (txt, raw, cen) in enumerate(ranked)
+    ]
+    table = "| rank | centred fve | raw fve | verbalisation |\n|---|---|---|---|\n" + "\n".join(rows)
+    mu_str = f"{pipe.mu.norm().item():.2f}" if pipe.mu is not None else "n/a"
+    summary = (
+        f"**Backbone:** {pipe.spec.label}  "
+        f"**hidden norm** = {v.norm().item():.2f}  "
+        f"**centring pool norm** = {mu_str}"
+    )
+    return summary, table
+
+
+@_gpu(duration=180)
+def cb_thought_trace(backbone_key: str, prompt: str, max_new: int, every: int):
+    if not prompt.strip():
+        yield "_(enter a prompt)_"
+        return
+    pipe = get_pipe(backbone_key)
+    rows: list[str] = ["| step | token | verbalisation of running hidden |", "|---|---|---|"]
+    yielded = 0
+    for step, (tok, verb) in enumerate(
+        pipe.thought_trace(prompt, max_new_tokens=int(max_new), every=int(every), K=4)
+    ):
+        if verb is None:
+            rows.append(f"| {step} | `{tok}` | _(skipped, not at every-{every} step)_ |")
+        else:
+            rows.append(f"| {step} | `{tok}` | {verb[:140].replace(chr(10), ' ')} |")
+        yielded += 1
+        # Stream every few rows
+        if yielded % 4 == 0 or yielded == int(max_new):
+            yield "\n".join(rows)
+    yield "\n".join(rows)
+
+
+@_gpu(duration=180)
+def cb_steer(
+    backbone_key: str,
+    prompt: str,
+    replacement: str,
+    alpha: float,
+    max_new: int,
+):
+    if not prompt.strip():
+        return "_(enter a prompt)_", "", ""
+    pipe = get_pipe(backbone_key)
+    # 1. Original verbalisation (best of 4) of the prompt's hidden
+    v = pipe.extract_hidden(prompt, token_index=-1)
+    ranked = pipe.verbalize(v, K=4, temperature=1.0)
+    base_verb = ranked[0][0] if ranked else ""
+
+    # 2. Original (unsteered) greedy continuation
+    ids = pipe.tokenizer(prompt, return_tensors="pt").to(pipe.device)
+    with torch.no_grad():
+        gen = pipe.backbone.generate(
+            **ids,
+            max_new_tokens=int(max_new),
+            do_sample=False,
+            pad_token_id=pipe.tokenizer.pad_token_id,
+        )
+    base_cont = pipe.tokenizer.decode(
+        gen[0, ids.input_ids.shape[1]:], skip_special_tokens=True
+    )
+
+    # 3. Steered continuation: re-encode replacement, patch into layer L
+    if not replacement.strip():
+        replacement = base_verb
+    steered_cont = pipe.steer_with_text(
+        prompt, replacement, alpha=float(alpha), max_new_tokens=int(max_new)
+    )
+    return base_verb, base_cont, steered_cont
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+
+with gr.Blocks(title="SRT-NLA: ask the model what it is thinking") as app:
+    gr.Markdown(
+        """# SRT-NLA Demo
+
+Ask a frozen language model what is in its own internal hidden state, in
+plain English, in three different ways. The trained Activation
+Verbalizer (AV) is a small adapter (~5–13M params depending on backbone)
+sitting on top of a fully frozen backbone. The backbone weights are
+unchanged from the public release. Companion paper:
+[`paper_nla.md`](https://github.com/space-bacon/SRT/blob/main/paper_nla.md).
+"""
+    )
+    backbone_key = gr.Radio(
+        choices=[(spec.label, key) for key, spec in BACKBONES.items()],
+        value=DEFAULT_BACKBONE,
+        label="Frozen backbone",
+    )
+
+    with gr.Tab("1 - Playground"):
+        gr.Markdown(
+            "Enter any prompt. We extract the last-token hidden state at the "
+            "configured layer, sample K candidate verbalisations, then for each "
+            "candidate we feed it back through the same frozen model and measure "
+            "how well its own re-encoded hidden state matches the original (raw "
+            "and anisotropy-centred cosine fidelity)."
+        )
+        with gr.Row():
+            with gr.Column():
+                pl_prompt = gr.Textbox(
+                    label="Prompt",
+                    value="The Eiffel Tower stands on the Champ de Mars in Paris,",
+                    lines=3,
+                )
+                pl_K = gr.Slider(1, 32, value=8, step=1, label="K (samples)")
+                pl_temp = gr.Slider(0.1, 1.5, value=1.0, step=0.05, label="Temperature")
+                pl_btn = gr.Button("Verbalize", variant="primary")
+            with gr.Column():
+                pl_summary = gr.Markdown()
+                pl_table = gr.Markdown()
+        pl_btn.click(
+            cb_playground,
+            inputs=[backbone_key, pl_prompt, pl_K, pl_temp],
+            outputs=[pl_summary, pl_table],
+        )
+
+    with gr.Tab("2 - Live thought trace"):
+        gr.Markdown(
+            "Watch the model generate token by token while we periodically "
+            "verbalise its running last-token hidden state. The right column is "
+            "the model talking *about* what its current hidden state is about, "
+            "while the left column is what it is choosing to *say next*."
+        )
+        with gr.Row():
+            with gr.Column():
+                tt_prompt = gr.Textbox(
+                    label="Prompt",
+                    value="In one paragraph, explain why the sky is blue.",
+                    lines=3,
+                )
+                tt_max = gr.Slider(8, 64, value=24, step=4, label="Tokens to generate")
+                tt_every = gr.Slider(1, 8, value=4, step=1, label="Verbalise every N tokens")
+                tt_btn = gr.Button("Trace", variant="primary")
+            with gr.Column():
+                tt_out = gr.Markdown()
+        tt_btn.click(
+            cb_thought_trace,
+            inputs=[backbone_key, tt_prompt, tt_max, tt_every],
+            outputs=tt_out,
+        )
+
+    with gr.Tab("3 - Steer by editing"):
+        gr.Markdown(
+            "Activation steering with a sentence editor. We verbalise the "
+            "prompt's hidden state into plain English, you edit the text, we "
+            "re-encode the edited text into a new hidden state, and we add "
+            "`alpha * (v_new - v_orig)` to layer L's residual stream while the "
+            "model generates. Compare the unsteered and steered continuations "
+            "side by side."
+        )
+        with gr.Row():
+            with gr.Column():
+                st_prompt = gr.Textbox(
+                    label="Prompt",
+                    value="Tell me a short fact about cats.",
+                    lines=3,
+                )
+                st_replacement = gr.Textbox(
+                    label="Replacement description (leave blank to use the AV's own verbalisation)",
+                    placeholder="e.g. 'a fact about deep-sea fish'",
+                    lines=2,
+                )
+                st_alpha = gr.Slider(-2.0, 2.0, value=1.0, step=0.05, label="alpha")
+                st_max = gr.Slider(8, 128, value=64, step=8, label="Tokens to generate")
+                st_btn = gr.Button("Steer", variant="primary")
+            with gr.Column():
+                st_base_verb = gr.Textbox(label="AV's verbalisation of the original prompt", lines=2)
+                st_base_cont = gr.Textbox(label="Unsteered continuation", lines=4)
+                st_steered = gr.Textbox(label="Steered continuation", lines=4)
+        st_btn.click(
+            cb_steer,
+            inputs=[backbone_key, st_prompt, st_replacement, st_alpha, st_max],
+            outputs=[st_base_verb, st_base_cont, st_steered],
+        )
+
+    gr.Markdown(
+        """---
+**Notes.** The first call on each backbone downloads weights from
+HuggingFace; subsequent calls reuse the in-memory cache.  Centred fve
+uses a 2K-sample anisotropy pool published in the matching dataset
+repo.  Steering is implemented as a single forward hook on layer L and
+does not retrain anything.  Hardware: roughly 16 GB VRAM for Qwen-7B
+in bf16; less for Llama-3B and Gemma-2B.
+"""
+    )
+
+
+if __name__ == "__main__":
+    app.queue().launch(share=False)

demo/nla_pipeline.py

@@ -0,0 +1,370 @@
+"""Shared NLA pipeline used by demo/app.py.
+
+Wraps the trained Activation Verbalizer (AV) for any of three backbones
+behind a single class with three operations:
+
+    1. extract_hidden(prompt, token_index) -> torch.Tensor
+       Take a prompt string, run the frozen backbone forward, return the
+       hidden state at the configured extraction layer at the given token
+       position (default: last token).
+
+    2. verbalize(v, K, temperature) -> list[(text, raw_fve, centred_fve)]
+       Generate K candidate verbalisations of vector v, score each one
+       round-trip, return them sorted by centred fidelity.
+
+    3. steer_with_text(prompt, replacement_text, alpha) -> str
+       Re-encode `replacement_text` to a layer-L hidden vector v_new, then
+       run the frozen backbone forward on `prompt` with v_new added (scaled
+       by alpha) to the L20 last-token hidden state at every position. Return
+       the greedy continuation.
+
+Loaded weights come from the public HF release pack; nothing in this
+file is repo-specific past the import path of `srt.nla`.
+"""
+from __future__ import annotations
+
+import logging
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable
+
+import torch
+import torch.nn.functional as F
+
+# Make `srt` importable when running from repo root or from demo/.
+_REPO = Path(__file__).resolve().parent.parent
+if str(_REPO) not in sys.path:
+    sys.path.insert(0, str(_REPO))
+
+from srt.nla.config import NLAConfig  # noqa: E402
+from srt.nla.verbalizer import ActivationVerbalizer  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
+from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402
+
+log = logging.getLogger("nla_demo")
+
+
+# ---------------------------------------------------------------------------
+# Backbone registry
+# ---------------------------------------------------------------------------
+
+@dataclass(frozen=True)
+class BackboneSpec:
+    label: str                # display name in UI
+    backbone_id: str          # HF model ID for the frozen LM
+    extraction_layer: int     # layer to read hidden states from
+    av_repo: str              # HF repo with best_av.pt
+    av_filename: str = "best_av.pt"
+    targets_repo: str | None = None   # HF dataset for centring pool
+    targets_filename: str | None = None
+    num_prefix_tokens: int = 1
+    num_inject_slots: int = 1
+
+
+BACKBONES: dict[str, BackboneSpec] = {
+    "qwen2.5-7b": BackboneSpec(
+        label="Qwen 2.5-7B (L20)",
+        backbone_id="Qwen/Qwen2.5-7B",
+        extraction_layer=20,
+        av_repo="RiverRider/srt-nla-av-v1",
+        av_filename="best_av.pt",
+        targets_repo="RiverRider/srt-nla-targets-v1",
+        targets_filename="targets_q7b_L20_seq64_30k_seed1.pt",
+        num_prefix_tokens=16,
+    ),
+    "llama-3.2-3b": BackboneSpec(
+        label="Llama 3.2-3B (L20)",
+        backbone_id="meta-llama/Llama-3.2-3B",
+        extraction_layer=20,
+        av_repo="RiverRider/srt-nla-av-llama32-3b",
+        av_filename="best_av.pt",
+        targets_repo="RiverRider/srt-nla-targets-llama32-3b-v1",
+        targets_filename="targets_L20_seq64_30k_seed1.pt",
+    ),
+    "gemma-2-2b": BackboneSpec(
+        label="Gemma 2-2B (L19)",
+        backbone_id="google/gemma-2-2b",
+        extraction_layer=19,
+        av_repo="RiverRider/srt-nla-av-gemma2-2b-v1",
+        av_filename="best_av.pt",
+        targets_repo="RiverRider/srt-nla-targets-gemma2-2b-v1",
+        targets_filename="targets_L19_seq64_30k_seed1.pt",
+    ),
+}
+
+
+# ---------------------------------------------------------------------------
+# Pipeline
+# ---------------------------------------------------------------------------
+
+class NLAPipeline:
+    """One backbone + AV + centring pool, kept resident on a single device."""
+
+    def __init__(
+        self,
+        spec: BackboneSpec,
+        device: str | torch.device = "cuda",
+        dtype: torch.dtype = torch.bfloat16,
+        pool_size: int = 2000,
+    ) -> None:
+        self.spec = spec
+        self.device = torch.device(device)
+        log.info("Loading frozen backbone %s on %s", spec.backbone_id, self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained(spec.backbone_id)
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.backbone = AutoModelForCausalLM.from_pretrained(
+            spec.backbone_id, torch_dtype=dtype
+        ).to(self.device)
+        self.backbone.eval()
+        for p in self.backbone.parameters():
+            p.requires_grad_(False)
+
+        cfg = NLAConfig(
+            backbone_id=spec.backbone_id,
+            extraction_layer=spec.extraction_layer,
+            num_prefix_tokens=spec.num_prefix_tokens,
+            num_inject_slots=spec.num_inject_slots,
+            max_new_tokens=64,
+        )
+        self.av = ActivationVerbalizer(cfg, backbone=self.backbone, tokenizer=self.tokenizer)
+
+        log.info("Downloading AV checkpoint from %s", spec.av_repo)
+        ckpt_path = hf_hub_download(spec.av_repo, spec.av_filename)
+        sd = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        # Released checkpoints are wrapped as {'trainable': {...}, 'step': int,
+        # 'val_fve_nrm': float}; older checkpoints used 'state_dict' or were
+        # raw state dicts.
+        if isinstance(sd, dict):
+            for wrap_key in ("trainable", "state_dict", "av_state_dict"):
+                if wrap_key in sd and isinstance(sd[wrap_key], dict):
+                    sd = sd[wrap_key]
+                    break
+        # Filter to AV-only keys (proj, prefix_embeds, etc.). The frozen
+        # backbone weights live under `backbone.*` in `av.state_dict()` but
+        # are not in the saved AV checkpoint, so we ignore those when
+        # checking for missing keys.
+        own = self.av.state_dict()
+        filtered = {k: v for k, v in sd.items() if k in own and own[k].shape == v.shape}
+        adapter_only = {k for k in own if not k.startswith("backbone.")}
+        missing = adapter_only - set(filtered)
+        if missing:
+            log.warning("AV checkpoint missing adapter keys: %s", sorted(missing))
+        self.av.load_state_dict(filtered, strict=False)
+        self.av.to(self.device)
+        self.av.eval()
+
+        # Centring pool
+        self.mu: torch.Tensor | None = None
+        if spec.targets_repo and spec.targets_filename:
+            try:
+                log.info("Downloading centring pool %s", spec.targets_repo)
+                pool_path = hf_hub_download(
+                    spec.targets_repo, spec.targets_filename, repo_type="dataset"
+                )
+                pool_obj = torch.load(pool_path, map_location="cpu", weights_only=False)
+                # The released targets file is a dict
+                # {'sequences': [...str], 'activations': [...Tensor], 'meta': {...}}
+                # but earlier shards were saved as a flat tensor or as a dict
+                # with a 'hiddens'/'targets' key.  Normalise to (N, d).
+                if isinstance(pool_obj, dict):
+                    for key in ("activations", "hiddens", "targets"):
+                        if key in pool_obj:
+                            pool_obj = pool_obj[key]
+                            break
+                if isinstance(pool_obj, list):
+                    # Items are per-sequence (T_i, d) hidden tensors saved
+                    # flat as (T_i*d,). Reshape and keep the last-token
+                    # row of each, since the AV is trained on last-token
+                    # last-token hiddens.
+                    d = self.backbone.config.hidden_size
+                    last_rows = []
+                    for t in pool_obj[:pool_size]:
+                        flat = t.flatten()
+                        T = flat.numel() // d
+                        last_rows.append(flat[(T - 1) * d : T * d])
+                    pool_obj = torch.stack(last_rows)
+                else:
+                    pool_obj = pool_obj[:pool_size]
+                self.mu = pool_obj.float().mean(0).to(self.device)
+                log.info("Centring pool ||mu|| = %.2f", self.mu.norm().item())
+            except Exception as e:  # noqa: BLE001
+                log.warning("Could not load centring pool (%s); centred fve disabled", e)
+
+        self.dtype = dtype
+
+    # ------------------------------------------------------------------
+    # Hidden-state extraction
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def extract_hidden(
+        self,
+        prompt: str,
+        token_index: int = -1,
+    ) -> torch.Tensor:
+        """Run the frozen backbone on `prompt`; return the layer-L hidden
+        state at `token_index` (default: last input token), shape (d,)."""
+        ids = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        out = self.backbone(
+            input_ids=ids.input_ids,
+            attention_mask=ids.attention_mask,
+            output_hidden_states=True,
+            use_cache=False,
+        )
+        h = out.hidden_states[self.spec.extraction_layer][0]  # (T, d)
+        if token_index < 0:
+            token_index = h.size(0) + token_index
+        return h[token_index].float().detach()
+
+    @torch.no_grad()
+    def extract_all_hidden(self, prompt: str) -> tuple[list[str], torch.Tensor]:
+        """Return (per-token decoded strings, hiddens of shape (T, d))."""
+        ids = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        out = self.backbone(
+            input_ids=ids.input_ids,
+            attention_mask=ids.attention_mask,
+            output_hidden_states=True,
+            use_cache=False,
+        )
+        h = out.hidden_states[self.spec.extraction_layer][0].float().detach()
+        toks = [self.tokenizer.decode([t]) for t in ids.input_ids[0].tolist()]
+        return toks, h
+
+    # ------------------------------------------------------------------
+    # Verbalisation + scoring
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def verbalize(
+        self,
+        v: torch.Tensor,
+        K: int = 8,
+        temperature: float = 1.0,
+        max_new_tokens: int = 64,
+    ) -> list[tuple[str, float, float]]:
+        """Sample K candidate verbalisations, round-trip score each, return
+        list of (text, raw_fve, centred_fve) sorted by centred_fve desc.
+
+        If centring pool is unavailable, centred_fve == raw_fve."""
+        v = v.to(self.device).float()
+        v_rep = v.unsqueeze(0).expand(K, -1)
+        ids = self.av.generate(
+            v_rep,
+            max_new_tokens=max_new_tokens,
+            do_sample=temperature > 0,
+            temperature=max(temperature, 1e-6),
+            top_p=1.0,
+        )
+        texts = self.tokenizer.batch_decode(ids, skip_special_tokens=True)
+
+        # Round-trip: feed each text into the same backbone, extract layer-L
+        # last-token hidden, score against v.
+        scored: list[tuple[str, float, float]] = []
+        for txt in texts:
+            if not txt.strip():
+                scored.append((txt, 0.0, 0.0))
+                continue
+            h = self.extract_hidden(txt, token_index=-1)
+            raw = 0.5 * (1.0 + F.cosine_similarity(h.unsqueeze(0), v.unsqueeze(0)).item())
+            if self.mu is not None:
+                hc, vc = h - self.mu, v - self.mu
+                cen = 0.5 * (1.0 + F.cosine_similarity(hc.unsqueeze(0), vc.unsqueeze(0)).item())
+            else:
+                cen = raw
+            scored.append((txt, raw, cen))
+        scored.sort(key=lambda r: r[2], reverse=True)
+        return scored
+
+    # ------------------------------------------------------------------
+    # Steering by edited verbalisation
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def steer_with_text(
+        self,
+        prompt: str,
+        replacement_text: str,
+        alpha: float = 1.0,
+        max_new_tokens: int = 64,
+    ) -> str:
+        """Re-encode `replacement_text` to a layer-L hidden vector v_new,
+        then run greedy generation on `prompt` with a forward hook that
+        adds `alpha * (v_new - v_orig)` to layer-L's output at every
+        position. Returns the generated continuation only.
+        """
+        v_orig = self.extract_hidden(prompt, token_index=-1)
+        v_new = self.extract_hidden(replacement_text, token_index=-1)
+        delta = (v_new - v_orig) * alpha
+
+        # Find the residual stream block at extraction_layer.
+        # Both LLaMA-like and Gemma-like architectures expose this as
+        # backbone.model.layers[L].
+        try:
+            target_block = self.backbone.model.layers[self.spec.extraction_layer]
+        except AttributeError:
+            target_block = self.backbone.transformer.h[self.spec.extraction_layer]
+
+        delta_b = delta.to(self.dtype)
+
+        def _hook(_mod, _inp, out):
+            if isinstance(out, tuple):
+                h = out[0]
+                h = h + delta_b
+                return (h,) + out[1:]
+            return out + delta_b
+
+        handle = target_block.register_forward_hook(_hook)
+        try:
+            ids = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+            gen = self.backbone.generate(
+                **ids,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,
+                pad_token_id=self.tokenizer.pad_token_id,
+            )
+        finally:
+            handle.remove()
+
+        new_ids = gen[0, ids.input_ids.shape[1]:]
+        return self.tokenizer.decode(new_ids, skip_special_tokens=True)
+
+    # ------------------------------------------------------------------
+    # Live thought-trace
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def thought_trace(
+        self,
+        prompt: str,
+        max_new_tokens: int = 32,
+        every: int = 4,
+        K: int = 4,
+    ) -> Iterable[tuple[str, str | None]]:
+        """Greedy-generate `max_new_tokens` tokens after `prompt`. Every
+        `every` tokens, run the AV on the running last-token hidden state
+        and yield (token_text, verbalisation) pairs. Other steps yield
+        (token_text, None)."""
+        ids = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        cur = ids.input_ids
+        attn = ids.attention_mask
+        for step in range(max_new_tokens):
+            out = self.backbone(
+                input_ids=cur,
+                attention_mask=attn,
+                output_hidden_states=True,
+                use_cache=False,
+            )
+            next_id = out.logits[0, -1].argmax().unsqueeze(0).unsqueeze(0)
+            tok_text = self.tokenizer.decode(next_id[0])
+            cur = torch.cat([cur, next_id], dim=1)
+            attn = torch.cat([attn, torch.ones_like(next_id)], dim=1)
+            if step % every == 0:
+                v = out.hidden_states[self.spec.extraction_layer][0, -1].float().detach()
+                ranked = self.verbalize(v, K=K, temperature=1.0, max_new_tokens=24)
+                yield tok_text, ranked[0][0] if ranked else None
+            else:
+                yield tok_text, None
+            if next_id.item() == (self.tokenizer.eos_token_id or -1):
+                break

demo/requirements.txt

@@ -0,0 +1,9 @@
+gradio>=4.40
+torch>=2.2
+transformers>=4.44
+huggingface_hub>=0.24
+accelerate>=0.33
+sentencepiece
+safetensors
+numpy
+spaces>=0.30 ; sys_platform == "linux"

requirements.txt

@@ -0,0 +1,9 @@
+gradio>=4.40
+torch>=2.2
+transformers>=4.44
+huggingface_hub>=0.24
+accelerate>=0.33
+sentencepiece
+safetensors
+numpy
+spaces>=0.30 ; sys_platform == "linux"

srt/__init__.py

@@ -0,0 +1,28 @@
+"""Semiotic-Reflexive Transformer (SRT) — Adapter Architecture."""
+
+__version__ = "0.1.0"
+
+
+def _nla_factory():
+    """Lazy entry point for the SRT-NLA verbalizer/reconstructor pair.
+
+    Imported on demand so callers that only want the v3 adapter don't
+    pay the cost of pulling in ``srt.nla`` (which lazily loads
+    ``transformers`` even further).
+    """
+    from srt.nla import (  # noqa: F401  (re-export)
+        ActivationReconstructor,
+        ActivationVerbalizer,
+        NLAConfig,
+    )
+
+    return {
+        "ActivationReconstructor": ActivationReconstructor,
+        "ActivationVerbalizer": ActivationVerbalizer,
+        "NLAConfig": NLAConfig,
+    }
+
+
+HEAD_FACTORIES = {
+    "nla": _nla_factory,
+}

srt/adapter.py

@@ -0,0 +1,316 @@
+"""SRT Adapter — Semiotic awareness bolted onto any frozen causal LM.
+
+The adapter wraps a HuggingFace AutoModelForCausalLM and runs its layers
+manually, tapping hidden states at MAH hook points and injecting corrections
+at RRM injection points.  The backbone's native embeddings and LM head are
+used directly — no bridges, no tied embeddings, no CE degradation.
+
+    model = SRTAdapter(config)
+    out = model(input_ids, labels=labels)
+    # out.ce_loss  — from backbone's native LM head
+    # out.r_hat    — per-position reflexivity estimate
+    # out.regime   — subcritical vs supercritical classification
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoConfig
+
+from srt.config import SRTConfig
+from srt.modules.mah import MetapragmaticAttentionHead, MAHOutput
+from srt.modules.rrm import ReflexiveRecurrentModule
+from srt.modules.ben import BifurcationEstimationNetwork, BENOutput
+from srt.modules.community import CommunityDiscoveryHead, CommunityOutput
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SRTAdapterOutput:
+    """Full output from the SRT adapter."""
+
+    logits: torch.Tensor  # (B, T, V)
+    ce_loss: torch.Tensor | None = None  # scalar
+    divergences: list[torch.Tensor] = field(default_factory=list)  # [(B, T, d_div)]
+    injections: list[torch.Tensor] = field(default_factory=list)  # [(B, T, d_backbone)]
+    ben_output: BENOutput | None = None
+    community_output: CommunityOutput | None = None
+    meta_state: torch.Tensor | None = None  # (B, T, d_meta)
+    chain_residual_per_token: torch.Tensor | None = None  # (B, T) mean chain residual
+
+
+def _make_causal_mask(
+    seq_len: int, dtype: torch.dtype, device: torch.device
+) -> torch.Tensor:
+    """Create 4D additive causal attention mask."""
+    mask = torch.full(
+        (seq_len, seq_len), torch.finfo(dtype).min, dtype=dtype, device=device
+    )
+    mask = torch.triu(mask, diagonal=1)
+    return mask[None, None, :, :]  # (1, 1, T, T)
+
+
+class SRTAdapter(nn.Module):
+    """Semiotic-Reflexive Transformer adapter for any causal LM backbone."""
+
+    def __init__(self, config: SRTConfig) -> None:
+        super().__init__()
+        self.config = config
+
+        # ── Load and freeze backbone ─────────────────────────────────
+        dtype_map = {
+            "float32": torch.float32,
+            "float16": torch.float16,
+            "bfloat16": torch.bfloat16,
+        }
+        load_dtype = dtype_map.get(config.backbone_dtype, torch.bfloat16)
+
+        logger.info("Loading backbone: %s in %s", config.backbone_id, config.backbone_dtype)
+        self.backbone = AutoModelForCausalLM.from_pretrained(
+            config.backbone_id, torch_dtype=load_dtype
+        )
+        for p in self.backbone.parameters():
+            p.requires_grad = False
+        self.backbone.eval()
+
+        # Extract backbone parts (works for LLaMA, Qwen, Mistral, Phi, Gemma)
+        inner = self.backbone.model
+        self._embed_tokens = inner.embed_tokens
+        self._layers = inner.layers
+        self._final_norm = inner.norm
+        self._lm_head = self.backbone.lm_head
+        self._rotary_emb = getattr(inner, "rotary_emb", None)
+
+        d_backbone = self.backbone.config.hidden_size
+        num_layers = self.backbone.config.num_hidden_layers
+        self._d_backbone = d_backbone
+        self._num_layers = num_layers
+
+        # Resolve auto layer indices
+        config.resolve_layer_indices(num_layers)
+
+        logger.info(
+            "Backbone: d=%d, L=%d, MAH@%s, inject@%s, community@%d",
+            d_backbone,
+            num_layers,
+            config.mah_layer_indices,
+            config.rrm_inject_indices,
+            config.community_layer_idx,
+        )
+
+        # ── Community discovery (early layer) ────────────────────────
+        self.community_head = CommunityDiscoveryHead(config.community, d_backbone)
+
+        # ── MAH heads (one per hook layer) ───────────────────────────
+        self.mah_heads = nn.ModuleList([
+            MetapragmaticAttentionHead(
+                config.mah, d_backbone, d_community=config.community.d_community
+            )
+            for _ in config.mah_layer_indices
+        ])
+
+        # ── RRM ──────────────────────────────────────────────────────
+        self.rrm = ReflexiveRecurrentModule(
+            config.rrm, d_divergence=config.mah.d_divergence, d_backbone=d_backbone
+        )
+
+        # Chain predictor: predict next divergence from current (self-supervised)
+        self.chain_predictor = nn.Linear(
+            config.mah.d_divergence, config.mah.d_divergence, bias=False
+        )
+
+        # ── BEN ──────────────────────────────────────────────────────
+        self.ben = BifurcationEstimationNetwork(config.ben, d_meta=config.rrm.d_meta)
+
+        # Build lookup sets for fast layer-index checking
+        self._mah_set = set(config.mah_layer_indices)
+        self._inject_set = set(config.rrm_inject_indices)
+        self._mah_index_map = {idx: i for i, idx in enumerate(config.mah_layer_indices)}
+
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        frozen = sum(p.numel() for p in self.parameters() if not p.requires_grad)
+        logger.info(
+            "SRT Adapter: %s trainable, %s frozen (backbone)",
+            f"{trainable:,}",
+            f"{frozen:,}",
+        )
+
+        # Cast adapter modules to backbone dtype so bf16 hidden states flow
+        # through without dtype mismatch (backbone is frozen bf16, adapter
+        # modules default to float32)
+        for module in [
+            self.community_head, self.mah_heads, self.rrm,
+            self.chain_predictor, self.ben,
+        ]:
+            module.to(load_dtype)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        forced_community: torch.Tensor | None = None,
+    ) -> SRTAdapterOutput:
+        """Forward pass: backbone with semiotic taps and injections.
+
+        Args:
+            input_ids: (B, T) token ids.
+            attention_mask: (B, T) padding mask (1 = real, 0 = pad). Optional.
+            labels: (B, T) target token ids for CE loss. Optional.
+            forced_community: (B, d_community) override community vector. Optional.
+                When provided, uses this instead of CommunityDiscoveryHead output
+                for conditioning MAH heads. Discovery still runs for diagnostics.
+
+        Returns:
+            SRTAdapterOutput with logits, losses, and semiotic intermediates.
+        """
+        device = input_ids.device
+        B, T = input_ids.shape
+
+        # 1. Native backbone embeddings
+        h = self._embed_tokens(input_ids)
+
+        # 2. Prepare position embeddings
+        position_ids = torch.arange(T, device=device).unsqueeze(0).expand(B, -1)
+        position_embeddings = None
+        if self._rotary_emb is not None:
+            position_embeddings = self._rotary_emb(h, position_ids)
+
+        # 3. Causal mask for MAH attention
+        mah_causal_mask = _make_causal_mask(T, h.dtype, device)
+
+        # 4. Prepare 4D causal+padding mask for backbone layers
+        # Must combine causal mask (T, T) with padding mask (B, T) into (B, 1, T, T)
+        # so that SDPA doesn't drop is_causal=True behavior
+        causal_4d = _make_causal_mask(T, h.dtype, device)  # (1, 1, T, T)
+        backbone_mask = None
+        if attention_mask is not None:
+            # (B, T) → (B, 1, 1, T) padding mask
+            pad_mask = (1.0 - attention_mask[:, None, None, :].to(h.dtype)) * torch.finfo(
+                h.dtype
+            ).min
+            backbone_mask = causal_4d + pad_mask  # (B, 1, T, T)
+        else:
+            backbone_mask = causal_4d  # (1, 1, T, T) — causal only
+
+        # 5. Layer-by-layer forward with semiotic taps
+        divergences: list[torch.Tensor] = []
+        injections: list[torch.Tensor] = []
+        meta_state: torch.Tensor | None = None
+        community_out: CommunityOutput | None = None
+        community_vec: torch.Tensor | None = None
+        mah_idx = 0
+
+        for layer_i, layer in enumerate(self._layers):
+            # Run backbone layer
+            layer_kwargs: dict = {"position_ids": position_ids}
+            if position_embeddings is not None:
+                layer_kwargs["position_embeddings"] = position_embeddings
+            if backbone_mask is not None:
+                layer_kwargs["attention_mask"] = backbone_mask
+
+            layer_out = layer(h, **layer_kwargs)
+            h = layer_out[0]
+
+            # Community discovery at early layer
+            if layer_i == self.config.community_layer_idx and community_out is None:
+                community_out = self.community_head(h.detach(), attention_mask)
+                # Use forced_community override if provided, else discovered
+                community_vec = (
+                    forced_community if forced_community is not None
+                    else community_out.vector
+                )
+
+            # MAH hook: extract divergence
+            if layer_i in self._mah_set:
+                mah_head = self.mah_heads[self._mah_index_map[layer_i]]
+                mah_out = mah_head(h, community_vec=community_vec, causal_mask=mah_causal_mask)
+                divergences.append(mah_out.divergence)
+
+                # Update RRM meta-state
+                meta_state = self.rrm.step(mah_out.divergence, meta_state)
+
+                # RRM injection (if this is also an injection layer)
+                if layer_i in self._inject_set:
+                    inj = self.rrm.inject(meta_state, h)
+                    h = h + inj
+                    injections.append(inj)
+
+        # 6. Final norm + native LM head
+        h = self._final_norm(h)
+        logits = self._lm_head(h)
+
+        # 7. CE loss (shifted, standard next-token prediction)
+        ce_loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            ce_loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        # 8. BEN
+        ben_out = None
+        if meta_state is not None:
+            ben_out = self.ben(meta_state)
+
+        # Per-token chain residual: mean across consecutive divergence pairs of
+        # squared error (chain_predictor(div_i) - div_{i+1})^2 averaged over
+        # the divergence dim. Shape (B, T). Same quantity that chain_loss
+        # reduces to a scalar; surfaced here for inference/probing.
+        chain_res = None
+        if len(divergences) >= 2:
+            B_, T_, _ = divergences[0].shape
+            acc = torch.zeros(B_, T_, dtype=divergences[0].dtype,
+                              device=divergences[0].device)
+            for i in range(len(divergences) - 1):
+                pred = self.chain_predictor(divergences[i])
+                acc = acc + (pred - divergences[i + 1]).pow(2).mean(dim=-1)
+            chain_res = acc / (len(divergences) - 1)
+
+        return SRTAdapterOutput(
+            logits=logits,
+            ce_loss=ce_loss,
+            divergences=divergences,
+            injections=injections,
+            ben_output=ben_out,
+            community_output=community_out,
+            meta_state=meta_state,
+            chain_residual_per_token=chain_res,
+        )
+
+    # Adapter module prefixes for save/load (everything else is backbone)
+    _ADAPTER_PREFIXES = (
+        "community_head.", "mah_heads.", "rrm.", "chain_predictor.", "ben.",
+    )
+
+    def save_adapter(self, path: str) -> None:
+        """Save only the trainable adapter weights (not the backbone)."""
+        state = {
+            k: v for k, v in self.state_dict().items()
+            if k.startswith(self._ADAPTER_PREFIXES)
+        }
+        torch.save(state, path)
+        logger.info("Saved adapter weights (%d tensors) to %s", len(state), path)
+
+    def load_adapter(self, path: str) -> None:
+        """Load adapter weights (backbone loaded separately from HF)."""
+        state = torch.load(path, map_location="cpu", weights_only=True)
+        missing, unexpected = self.load_state_dict(state, strict=False)
+        # Expected: all non-adapter keys will be "missing" (loaded from HF)
+        adapter_missing = [k for k in missing if k.startswith(self._ADAPTER_PREFIXES)]
+        if adapter_missing:
+            logger.warning("Missing adapter keys: %s", adapter_missing)
+        logger.info("Loaded adapter weights from %s", path)
+
+    def trainable_parameters(self):
+        """Yield only the trainable (adapter) parameters."""
+        return (p for p in self.parameters() if p.requires_grad)

srt/config.py

@@ -0,0 +1,155 @@
+"""Configuration dataclasses for SRT Adapter."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class MAHConfig:
+    """Metapragmatic Attention Head configuration."""
+
+    d_sub: int = 512  # semiotic subspace dimension
+    d_divergence: int = 256  # divergence vector dimension
+    num_heads: int = 4  # attention heads
+    dropout: float = 0.1
+
+
+@dataclass
+class RRMConfig:
+    """Reflexive Recurrent Module configuration."""
+
+    d_meta: int = 512  # GRU meta-state dimension
+    inject_scale: float = 1.0  # FiLM correction scale (v3 used 0.1 with linear inject; v4 uses 1.0 with FiLM)
+
+
+@dataclass
+class BENConfig:
+    """Bifurcation Estimation Network configuration."""
+
+    d_hidden: int = 256  # MLP hidden dimension
+
+
+@dataclass
+class CommunityConfig:
+    """Unsupervised community discovery configuration."""
+
+    num_prototypes: int = 32  # number of soft community clusters
+    d_community: int = 64  # community embedding dimension
+    temperature: float = 1.0  # softmax temperature for assignment
+    # v8a: when False, skip the discrete prototype basis entirely; the
+    # encoder output IS the community vector. Motivated by the v7 PCA
+    # finding that prototype tensors barely move from random init across
+    # v5/v6/v7 (mean abs delta ~3e-5) — the encoder was already doing all
+    # the discriminative work and the prototype-mixing readout was
+    # discarding information at the soft-argmax. With use_prototypes=False
+    # the community channel becomes a continuous 64-D coordinate rather
+    # than a soft assignment over K anchors.
+    #
+    # Env override: set SRT_USE_PROTOTYPES=0 (or "false") to flip this off
+    # globally. Lets probe / eval scripts run against v8a checkpoints
+    # without per-script flag plumbing.
+    use_prototypes: bool = True
+
+    def __post_init__(self) -> None:
+        import os
+        v = os.environ.get("SRT_USE_PROTOTYPES")
+        if v is not None and v.lower() in ("0", "false", "no", "off"):
+            self.use_prototypes = False
+
+
+@dataclass
+class LossConfig:
+    """Loss weights."""
+
+    ce_weight: float = 1.0
+    chain_weight: float = 0.5  # divergence chain prediction
+    bif_weight: float = 1.0  # bifurcation (r_hat vs r_true)
+    regime_weight: float = 5.0  # regime classification
+    div_alive_weight: float = 0.1  # prevent divergence collapse
+    # v4: dropped to 0 because v3 ablation showed the inject-norm regularizer
+    # was driving the optimizer to satisfy ||inj||=1 with arbitrary directions
+    # rather than directions useful for downstream loss.  FiLM init handles
+    # gradient flow without needing a norm prior.
+    inject_reg_weight: float = 0.0
+    inject_target_norm: float = 1.0
+    community_entropy_weight: float = 0.01  # diverse community usage
+    # v4/v5: SupCon loss on community ENCODER output keyed by source-id
+    # hash.  Forces prototypes apart by giving same-source pairs positive
+    # gradient and different-source pairs negative gradient through the
+    # encoder.  v5 raised the weight 0.5 -> 2.0 because v4's signal at 0.5
+    # was overwhelmed and the loss flatlined at log(B-1)=2.71.
+    community_supcon_weight: float = 2.0
+    community_supcon_temperature: float = 0.1
+    # v6 additions:
+    #   - divergence SupCon on mean-pooled last-MAH divergence (analog of v5
+    #     community SupCon, applied to the metapragmatic channel)
+    #   - ListNet ranking loss on r̂ within each sequence (sharpens ordering;
+    #     pointwise smooth-L1 alone tolerates large rank errors at the tails)
+    #   - chain-residual auxiliary floor: keeps inference signal alive after
+    #     chain_loss has driven the per-position residual near zero
+    divergence_supcon_weight: float = 1.0
+    divergence_supcon_temperature: float = 0.1
+    listnet_weight: float = 0.5
+    listnet_temperature: float = 1.0
+    chain_residual_aux_weight: float = 0.05
+    chain_residual_aux_target: float = 0.5
+    # v9: supervised contrastive loss keyed by archetype_id, applied to the
+    # same `community_output.encoded` representation as community_supcon. The
+    # 33 archetypes (Lancaster, paired with the Lexicon of Synthetic
+    # Interiority) are an external taxonomy that has only been a held-out
+    # probe through v8b. v9 promotes them to a training signal alongside
+    # Reddit subreddit ids. Rows whose archetype_id == -1 (Reddit corpus) are
+    # masked out of this loss; rows from the archetype-generations corpus
+    # carry archetype_id ∈ [1, 33] and contribute positive pairs.
+    archetype_supcon_weight: float = 0.0
+    archetype_supcon_temperature: float = 0.1
+
+
+@dataclass
+class TrainingConfig:
+    """Training hyperparameters."""
+
+    lr: float = 3e-4
+    weight_decay: float = 0.01
+    epochs: int = 3
+    batch_size: int = 16
+    max_seq_len: int = 512
+    val_every: int = 1000
+    log_every: int = 100
+    patience: int = 5
+    warmup_steps: int = 500
+    grad_clip: float = 1.0
+
+
+@dataclass
+class SRTConfig:
+    """Top-level SRT Adapter configuration."""
+
+    backbone_id: str = "Qwen/Qwen2.5-7B"
+    backbone_dtype: str = "bfloat16"
+
+    # Layer hook indices — empty means auto-compute from backbone depth
+    mah_layer_indices: list[int] = field(default_factory=list)
+    rrm_inject_indices: list[int] = field(default_factory=list)
+    community_layer_idx: int = -1  # -1 = auto
+
+    num_mah_layers: int = 3
+
+    mah: MAHConfig = field(default_factory=MAHConfig)
+    rrm: RRMConfig = field(default_factory=RRMConfig)
+    ben: BENConfig = field(default_factory=BENConfig)
+    community: CommunityConfig = field(default_factory=CommunityConfig)
+    loss: LossConfig = field(default_factory=LossConfig)
+    training: TrainingConfig = field(default_factory=TrainingConfig)
+
+    def resolve_layer_indices(self, num_layers: int) -> None:
+        """Auto-compute layer indices from backbone depth if not set."""
+        if not self.mah_layer_indices:
+            step = num_layers // (self.num_mah_layers + 1)
+            self.mah_layer_indices = [step * (i + 1) for i in range(self.num_mah_layers)]
+        if not self.rrm_inject_indices:
+            # Inject at all MAH layers except the first (let meta-state build up)
+            self.rrm_inject_indices = self.mah_layer_indices[1:]
+        if self.community_layer_idx < 0:
+            self.community_layer_idx = max(1, num_layers // 7)

srt/data/__init__.py

@@ -0,0 +1 @@
+"""SRT data utilities."""

srt/data/dataset.py

@@ -0,0 +1,217 @@
+"""Dataset for SRT Adapter training.
+
+Loads JSONL files produced by the Reddit corpus pipeline and aligns word-level
+labels (r_true) to BPE token boundaries using the backbone's own tokenizer.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizerBase
+
+logger = logging.getLogger(__name__)
+
+
+class SRTAdapterDataset(Dataset):
+    """Load JSONL samples and tokenize with a BPE tokenizer.
+
+    Expected JSONL format (all fields except ``text`` are optional):
+        {"text": "...", "r_true": [0.1, -0.2, ...], "community": "some_label"}
+    """
+
+    def __init__(
+        self,
+        path: str | Path,
+        tokenizer: PreTrainedTokenizerBase,
+        max_seq_len: int = 512,
+        max_samples: int | None = None,
+    ) -> None:
+        self.tokenizer = tokenizer
+        self.max_seq_len = max_seq_len
+        self.samples: list[dict[str, Any]] = []
+
+        path = Path(path)
+        logger.info("Loading dataset from %s", path)
+        with open(path) as f:
+            for i, line in enumerate(f):
+                if max_samples is not None and i >= max_samples:
+                    break
+                row = json.loads(line)
+                if "text" in row and row["text"].strip():
+                    self.samples.append(row)
+
+        logger.info("Loaded %d samples from %s", len(self.samples), path)
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        row = self.samples[idx]
+        text = row["text"]
+
+        # Tokenize
+        enc = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_seq_len,
+            return_tensors="pt",
+            return_offsets_mapping=True,
+        )
+        input_ids = enc["input_ids"].squeeze(0)  # (T,)
+        attention_mask = enc["attention_mask"].squeeze(0)  # (T,)
+        offsets = enc["offset_mapping"].squeeze(0)  # (T, 2)
+        T = input_ids.size(0)
+
+        # Align word-level r_true → token-level
+        r_true = torch.zeros(T)
+        r_mask = torch.zeros(T, dtype=torch.bool)
+
+        if "r_true" in row and row["r_true"]:
+            word_r = row["r_true"]
+            words = text.split()
+            token_r, token_mask = _align_word_labels_to_bpe(
+                words, word_r, offsets, text
+            )
+            r_true[:len(token_r)] = token_r[:T]
+            r_mask[:len(token_mask)] = token_mask[:T]
+
+        # Community id: prefer the explicit integer in the data; fall back
+        # to a stable hash of a string-typed source field if the integer is
+        # missing.  v5 fix: training data has `community_id: int` populated
+        # for every row; v3-v5 was reading a non-existent `community` field
+        # and hashing the empty string, giving every sample the same id and
+        # collapsing the SupCon signal.
+        if "community_id" in row and row["community_id"] is not None:
+            community_id = int(row["community_id"])
+        else:
+            for fld in ("community_label", "source", "community"):
+                community_str = row.get(fld) or ""
+                if community_str:
+                    break
+            community_id = _stable_hash(community_str)
+
+        # v9: archetype_id is present on rows from the archetype-generations
+        # corpus (data/archetype_train.jsonl); rows from the Reddit corpus
+        # carry -1 here so the v9 archetype-supcon loss masks them out.
+        if "archetype_id" in row and row["archetype_id"] is not None:
+            archetype_id = int(row["archetype_id"])
+        else:
+            archetype_id = -1
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": input_ids.clone(),  # next-token prediction
+            "r_true": r_true,
+            "r_mask": r_mask,
+            "community_id": torch.tensor(community_id, dtype=torch.long),
+            "archetype_id": torch.tensor(archetype_id, dtype=torch.long),
+        }
+
+
+def _stable_hash(s: str, modulus: int = 100003) -> int:
+    """Stable non-cryptographic hash → int in [0, modulus). FNV-1a 32-bit.
+
+    We use FNV-1a rather than the built-in `hash()` because the latter is
+    randomized per-process under PYTHONHASHSEED=random, which would give
+    different ids each run.
+    """
+    h = 2166136261
+    for ch in s.encode("utf-8"):
+        h ^= ch
+        h = (h * 16777619) & 0xFFFFFFFF
+    return h % modulus
+
+
+def _align_word_labels_to_bpe(
+    words: list[str],
+    word_labels: list[float],
+    offsets: torch.Tensor,
+    text: str,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Align word-level labels to BPE token positions.
+
+    For each BPE token, find which word it belongs to (by character offset)
+    and assign that word's label.
+
+    Args:
+        words: list of whitespace-split words.
+        word_labels: per-word label values.
+        offsets: (T, 2) character offsets from tokenizer.
+        text: original text string.
+
+    Returns:
+        (token_labels, token_mask) both of shape (T,).
+    """
+    T = offsets.size(0)
+    token_labels = torch.zeros(T)
+    token_mask = torch.zeros(T, dtype=torch.bool)
+
+    if len(words) != len(word_labels):
+        # Mismatched lengths — skip alignment
+        return token_labels, token_mask
+
+    # Build word → character span mapping
+    word_spans: list[tuple[int, int]] = []
+    pos = 0
+    for word in words:
+        start = text.find(word, pos)
+        if start == -1:
+            break
+        end = start + len(word)
+        word_spans.append((start, end))
+        pos = end
+
+    if len(word_spans) != len(words):
+        return token_labels, token_mask
+
+    # For each token, find its word
+    for tok_idx in range(T):
+        tok_start, tok_end = offsets[tok_idx].tolist()
+        if tok_start == 0 and tok_end == 0:
+            continue  # special token
+
+        tok_mid = (tok_start + tok_end) / 2
+        for word_idx, (ws, we) in enumerate(word_spans):
+            if ws <= tok_mid < we:
+                token_labels[tok_idx] = word_labels[word_idx]
+                token_mask[tok_idx] = True
+                break
+
+    return token_labels, token_mask
+
+
+def make_collate_fn(pad_token_id: int = 0):
+    """Create a collate function with the correct pad token id."""
+
+    def collate_fn(batch: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
+        """Collate with dynamic padding to the longest sequence in the batch."""
+        max_len = max(item["input_ids"].size(0) for item in batch)
+
+        padded: dict[str, list[torch.Tensor]] = {key: [] for key in batch[0]}
+        for item in batch:
+            T = item["input_ids"].size(0)
+            pad_len = max_len - T
+            padded["input_ids"].append(F.pad(item["input_ids"], (0, pad_len), value=pad_token_id))
+            padded["attention_mask"].append(
+                F.pad(item["attention_mask"], (0, pad_len), value=0)
+            )
+            padded["labels"].append(F.pad(item["labels"], (0, pad_len), value=-100))
+            padded["r_true"].append(F.pad(item["r_true"], (0, pad_len), value=0.0))
+            padded["r_mask"].append(F.pad(item["r_mask"], (0, pad_len), value=False))
+            # community_id is a scalar — stack without padding
+            if "community_id" in item:
+                padded.setdefault("community_id", []).append(item["community_id"])
+            if "archetype_id" in item:
+                padded.setdefault("archetype_id", []).append(item["archetype_id"])
+
+        return {k: torch.stack(v) for k, v in padded.items()}
+
+    return collate_fn

srt/modules/__init__.py

@@ -0,0 +1,16 @@
+"""SRT Adapter semiotic modules."""
+
+from srt.modules.mah import MetapragmaticAttentionHead, MAHOutput
+from srt.modules.rrm import ReflexiveRecurrentModule
+from srt.modules.ben import BifurcationEstimationNetwork, BENOutput
+from srt.modules.community import CommunityDiscoveryHead, CommunityOutput
+
+__all__ = [
+    "MetapragmaticAttentionHead",
+    "MAHOutput",
+    "ReflexiveRecurrentModule",
+    "BifurcationEstimationNetwork",
+    "BENOutput",
+    "CommunityDiscoveryHead",
+    "CommunityOutput",
+]

srt/modules/ben.py

@@ -0,0 +1,71 @@
+"""Bifurcation Estimation Network (BEN).
+
+Estimates the reflexivity coefficient r̂ ∈ [-1, 1] at each position from
+the RRM's accumulated meta-state.  Also classifies semiotic regime:
+  - Subcritical (r < 0): sign has stable, conventional meaning
+  - Supercritical (r > 0): sign is contested, meaning is actively forking
+
+r̂ is the core output of SRT — it tells you WHERE and HOW MUCH meaning
+is under contestation in a given text.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+from srt.config import BENConfig
+
+
+@dataclass
+class BENOutput:
+    """Output from BEN."""
+
+    r_hat: torch.Tensor  # (B, T) reflexivity coefficient (unbounded; supervised on log-compressed r_true)
+    regime_logits: torch.Tensor  # (B, T, 2) subcritical/supercritical
+
+
+class BifurcationEstimationNetwork(nn.Module):
+    """Estimates bifurcation from RRM meta-state."""
+
+    def __init__(self, cfg: BENConfig, d_meta: int) -> None:
+        super().__init__()
+
+        # r̂ prediction: meta-state → unbounded scalar.
+        # v3 used nn.Tanh() here, which capped output at ±1.  The training target
+        # is sign(r) * log1p(|r|) and r_true reaches ~12.77 (compressed ~2.55), so
+        # the tanh ceiling truncated ~25% of supercritical tokens and capped the
+        # achievable Pearson.  The smooth_l1 loss on a log-compressed target is
+        # numerically well-behaved without an output activation; we keep the head
+        # unbounded and init the final linear with small weights so early outputs
+        # start near zero and the supervised gradient does the shaping.
+        self.r_head = nn.Sequential(
+            nn.Linear(d_meta, cfg.d_hidden),
+            nn.SiLU(),
+            nn.Linear(cfg.d_hidden, 1),
+        )
+        r_out: nn.Linear = self.r_head[-1]  # type: ignore[assignment]
+        nn.init.normal_(r_out.weight, std=0.02)
+        nn.init.zeros_(r_out.bias)
+
+        # Regime classification: subcritical (0) vs supercritical (1)
+        self.regime_head = nn.Sequential(
+            nn.Linear(d_meta, cfg.d_hidden),
+            nn.SiLU(),
+            nn.Linear(cfg.d_hidden, 2),
+        )
+
+    def forward(self, meta_state: torch.Tensor) -> BENOutput:
+        """Estimate bifurcation from accumulated meta-state.
+
+        Args:
+            meta_state: (B, T, d_meta) from RRM.
+
+        Returns:
+            BENOutput with r_hat and regime_logits.
+        """
+        r_hat = self.r_head(meta_state).squeeze(-1)  # (B, T)
+        regime_logits = self.regime_head(meta_state)  # (B, T, 2)
+        return BENOutput(r_hat=r_hat, regime_logits=regime_logits)

srt/modules/community.py

@@ -0,0 +1,116 @@
+"""Unsupervised Community Discovery Head.
+
+Discovers discourse communities from backbone hidden states without
+predefined labels.  A discourse community (in Peirce's framework) is a
+group of language users who share interpretive norms — they assign similar
+interpretants to the same representamens.
+
+The community head runs at an early backbone layer (before MAH hooks) and
+produces a soft assignment over K learned prototypes.  The resulting community
+vector conditions how MAH computes divergence, so the same sign can produce
+different divergence patterns in different community contexts.
+
+Training signal: the community prototypes are pulled apart by the semiotic
+losses — if assigning text to different communities helps the model predict
+divergence better, it will learn to separate them.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from srt.config import CommunityConfig
+
+
+@dataclass
+class CommunityOutput:
+    """Output from community discovery.
+
+    When the head runs in continuous-trajectory mode (cfg.use_prototypes=False,
+    v8a), `logits` and `weights` are None and `vector == encoded`.
+    """
+
+    logits: torch.Tensor | None  # (B, K) raw assignment scores, or None
+    weights: torch.Tensor | None  # (B, K) soft assignment probabilities, or None
+    vector: torch.Tensor  # (B, d_community) community embedding (mixture or encoded)
+    encoded: torch.Tensor  # (B, d_community) pre-prototype-mixing encoder output
+
+
+class CommunityDiscoveryHead(nn.Module):
+    """Soft clustering of hidden states into discourse communities.
+
+    With cfg.use_prototypes=True (default): pooled hidden state → encoder →
+    cosine similarity to K learned prototypes → soft assignment weights →
+    weighted mixture of prototypes as the community vector. This is the
+    v3–v7 architecture.
+
+    With cfg.use_prototypes=False (v8a): pooled hidden state → encoder →
+    the encoder output IS the community vector. No discrete basis. Motivated
+    by the v7 PCA finding that prototype tensors barely move from random
+    init; the encoder was already doing the discriminative work and the
+    soft-argmax over K anchors was throwing information away.
+    """
+
+    def __init__(self, cfg: CommunityConfig, d_backbone: int) -> None:
+        super().__init__()
+        self.temperature = cfg.temperature
+        self.use_prototypes = cfg.use_prototypes
+
+        # Encode pooled hidden states → community space
+        self.encoder = nn.Sequential(
+            nn.Linear(d_backbone, cfg.d_community),
+            nn.SiLU(),
+        )
+
+        # Learnable community prototypes (only when enabled)
+        if cfg.use_prototypes:
+            self.prototypes = nn.Embedding(cfg.num_prototypes, cfg.d_community)
+        else:
+            self.prototypes = None  # type: ignore[assignment]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> CommunityOutput:
+        """Discover community from hidden states.
+
+        Args:
+            hidden_states: (B, T, d_backbone) from an early backbone layer.
+            attention_mask: (B, T) padding mask (1 = real, 0 = pad). Optional.
+
+        Returns:
+            CommunityOutput. In prototype mode, logits/weights are populated
+            and vector is the prototype-weighted mixture. In trajectory mode
+            (use_prototypes=False), logits and weights are None and vector
+            equals encoded.
+        """
+        # Masked mean pool across positions → document-level representation
+        if attention_mask is not None:
+            mask = attention_mask.unsqueeze(-1).to(hidden_states.dtype)  # (B, T, 1)
+            pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+        else:
+            pooled = hidden_states.mean(dim=1)  # (B, d_backbone)
+        encoded = self.encoder(pooled)  # (B, d_community)
+
+        if not self.use_prototypes:
+            # v8a: continuous-trajectory mode — no discrete basis.
+            return CommunityOutput(
+                logits=None, weights=None, vector=encoded, encoded=encoded,
+            )
+
+        # Cosine similarity to prototypes
+        encoded_norm = F.normalize(encoded, dim=-1)
+        proto_norm = F.normalize(self.prototypes.weight, dim=-1)
+        logits = (encoded_norm @ proto_norm.T) / self.temperature  # (B, K)
+
+        weights = F.softmax(logits, dim=-1)  # (B, K)
+        vector = weights @ self.prototypes.weight  # (B, d_community)
+
+        return CommunityOutput(
+            logits=logits, weights=weights, vector=vector, encoded=encoded,
+        )

srt/modules/mah.py

@@ -0,0 +1,109 @@
+"""Metapragmatic Attention Head (MAH).
+
+Detects where meaning diverges across positions by computing the gap between
+direct (local) interpretation and contextual (global) interpretation of each
+token's hidden state.  This is Peirce's "unlimited semiosis" made computational:
+each sign (representamen) receives an interpretation (interpretant) that depends
+on the surrounding discourse context.  MAH quantifies where that context
+*changes* the interpretation — i.e., where meaning forks.
+
+The divergence vector d_t at position t captures:
+  d_t = f(interp_t) - g(attend(interp_{0..t}))
+where f is direct projection, g is the contextual output after causal attention.
+High ||d_t|| → the sign at position t means something different in context
+than it would in isolation.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from srt.config import MAHConfig
+
+
+@dataclass
+class MAHOutput:
+    """Output from a single MAH layer."""
+
+    divergence: torch.Tensor  # (B, T, d_divergence)
+    attention_weights: torch.Tensor | None = None  # (B, H, T, T)
+
+
+class MetapragmaticAttentionHead(nn.Module):
+    """Single MAH layer that reads hidden states and produces divergence vectors."""
+
+    def __init__(self, cfg: MAHConfig, d_backbone: int, d_community: int = 0) -> None:
+        super().__init__()
+        d_sub = cfg.d_sub
+
+        # Project backbone hidden states → interpretant subspace
+        self.interp_proj = nn.Linear(d_backbone, d_sub, bias=False)
+
+        # Optional community conditioning
+        self.comm_proj: nn.Module | None = None
+        if d_community > 0:
+            self.comm_proj = nn.Linear(d_community, d_sub, bias=False)
+
+        # Multi-head self-attention in interpretant subspace
+        self.num_heads = cfg.num_heads
+        self.head_dim = d_sub // cfg.num_heads
+        assert d_sub % cfg.num_heads == 0
+
+        self.q_proj = nn.Linear(d_sub, d_sub, bias=False)
+        self.k_proj = nn.Linear(d_sub, d_sub, bias=False)
+        self.v_proj = nn.Linear(d_sub, d_sub, bias=False)
+        self.out_proj = nn.Linear(d_sub, d_sub, bias=False)
+        self.attn_dropout = nn.Dropout(cfg.dropout)
+
+        # Divergence output projection
+        self.div_proj = nn.Linear(d_sub, cfg.d_divergence, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        community_vec: torch.Tensor | None = None,
+        causal_mask: torch.Tensor | None = None,
+    ) -> MAHOutput:
+        """Compute divergence from backbone hidden states.
+
+        Args:
+            hidden_states: (B, T, d_backbone) from a transformer layer.
+            community_vec: (B, d_community) soft community vector.
+            causal_mask: (1, 1, T, T) additive causal mask.
+
+        Returns:
+            MAHOutput with divergence vectors and optional attention weights.
+        """
+        B, T, _ = hidden_states.shape
+
+        # Project to interpretant subspace
+        interp = self.interp_proj(hidden_states)  # (B, T, d_sub)
+
+        # Community conditioning: shift interpretant space
+        if community_vec is not None and self.comm_proj is not None:
+            comm_bias = self.comm_proj(community_vec)  # (B, d_sub)
+            interp = interp + comm_bias.unsqueeze(1)
+
+        # Multi-head causal self-attention
+        q = self.q_proj(interp).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(interp).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(interp).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+
+        attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        if causal_mask is not None:
+            attn = attn + causal_mask
+        attn_weights = F.softmax(attn, dim=-1)
+        attn_weights = self.attn_dropout(attn_weights)
+
+        contextual = (attn_weights @ v).transpose(1, 2).reshape(B, T, -1)
+        contextual = self.out_proj(contextual)  # (B, T, d_sub)
+
+        # Divergence = gap between direct and contextual interpretation
+        divergence = self.div_proj(interp - contextual)  # (B, T, d_divergence)
+
+        return MAHOutput(divergence=divergence, attention_weights=attn_weights.detach())

srt/modules/rrm.py

@@ -0,0 +1,109 @@
+"""Reflexive Recurrent Module (RRM).
+
+Tracks per-position semiotic meta-state via a GRU that processes divergence
+observations from MAH layers.  At injection points, produces a FiLM-style
+modulation (gamma, beta) that multiplicatively + additively biases the
+backbone's hidden states:
+
+    h' = h * (1 + gamma(meta_state)) + beta(meta_state)
+
+The meta-state h_meta_t represents the model's accumulated awareness of
+semiotic divergence at position t.  Each MAH observation updates it:
+    h_meta_t^{l+1} = GRU(divergence_t^l, h_meta_t^l)
+
+v3 used a single low-rank linear inject (gate * proj * scale) with
+zero-initialized projection.  Ablation showed the inject-back arm contributed
+exactly nothing (every benchmark metric was identical to four decimal places
+with injection forced to zero).  The diagnosis was that the zero init plus
+the inject-norm regularizer (which rewarded ||inj|| \u2248 1 regardless of
+direction) drove the optimizer to satisfy the norm penalty with arbitrary
+directions that were then orthogonal to the gradient signal from the frozen
+backbone's CE.
+
+v4 fixes both: FiLM modulation has a non-zero gradient pathway from the first
+step (beta is initialized to zero so the forward is identity at init, but
+gamma has small Gaussian init so dL/d(gamma_proj) flows immediately when the
+downstream MAH layer's divergence is supervised by the bif/regime losses).
+The inject-norm regularizer is dropped at the loss layer (LossConfig
+inject_reg_weight = 0.0 by default in v4).
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from srt.config import RRMConfig
+
+
+class ReflexiveRecurrentModule(nn.Module):
+    """GRU-based reflexive meta-state tracker with FiLM-style injection."""
+
+    def __init__(self, cfg: RRMConfig, d_divergence: int, d_backbone: int) -> None:
+        super().__init__()
+        self.d_meta = cfg.d_meta
+        self.inject_scale = cfg.inject_scale
+        self.d_backbone = d_backbone
+
+        # Per-position GRU: processes divergence \u2192 meta-state
+        self.gru = nn.GRUCell(d_divergence, cfg.d_meta)
+
+        # FiLM projections: meta-state \u2192 (gamma, beta) in backbone-dim.
+        # gamma is multiplicative on (1 + gamma); beta is additive.
+        # gamma init: small Gaussian (std=0.02) so identity-at-init holds in
+        # expectation but gradient flows from the first step.
+        # beta init: zeros so identity-at-init is exact, then learns offsets.
+        self.gamma_proj = nn.Linear(cfg.d_meta, d_backbone, bias=True)
+        self.beta_proj = nn.Linear(cfg.d_meta, d_backbone, bias=True)
+        nn.init.normal_(self.gamma_proj.weight, std=0.02)
+        nn.init.zeros_(self.gamma_proj.bias)
+        nn.init.zeros_(self.beta_proj.weight)
+        nn.init.zeros_(self.beta_proj.bias)
+
+    def step(
+        self, divergence: torch.Tensor, meta_state: torch.Tensor | None
+    ) -> torch.Tensor:
+        """Update per-position meta-state with new divergence observation.
+
+        Args:
+            divergence: (B, T, d_divergence) from MAH.
+            meta_state: (B, T, d_meta) or None for initial state.
+
+        Returns:
+            Updated meta-state (B, T, d_meta).
+        """
+        B, T, d_div = divergence.shape
+        div_flat = divergence.reshape(B * T, d_div)
+
+        if meta_state is None:
+            meta_flat = torch.zeros(
+                B * T, self.d_meta, device=divergence.device, dtype=divergence.dtype
+            )
+        else:
+            meta_flat = meta_state.reshape(B * T, self.d_meta)
+
+        meta_flat = self.gru(div_flat, meta_flat)
+        return meta_flat.reshape(B, T, self.d_meta)
+
+    def inject(
+        self, meta_state: torch.Tensor, hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        """Produce FiLM modulation correction for backbone hidden states.
+
+        Returns the *correction*  (h' - h) = h * gamma + beta, NOT h'.  The
+        caller adds this to h to get h'.  This keeps the rest of the adapter
+        and the diagnostic logging (injection norm tracking) unchanged: the
+        \"injection\" tensor is still the additive correction applied to h.
+
+        Args:
+            meta_state: (B, T, d_meta) current reflexive awareness.
+            hidden_states: (B, T, d_backbone) current hidden states.
+
+        Returns:
+            Correction vector (B, T, d_backbone) to add to hidden_states.
+        """
+        gamma = self.gamma_proj(meta_state)  # (B, T, d_backbone)
+        beta = self.beta_proj(meta_state)  # (B, T, d_backbone)
+        # FiLM: h' = h * (1 + gamma) + beta  \u2192  correction = h * gamma + beta
+        correction = hidden_states * gamma + beta
+        return correction * self.inject_scale

srt/modules/thead.py

@@ -0,0 +1,193 @@
+"""T-Head: Takens delay-embedding readout (v10 prototype).
+
+Inspired by Haylett (2025, 2026)'s Geofinitism program and MARINA architecture,
+which apply Takens (1981)'s delay-embedding theorem to the token stream.
+
+Where MARINA *replaces* attention with delay-coordinate phase-space
+reconstruction, the T-Head is the conservative cross-program experiment: keep
+the frozen attention backbone of SRT-Adapter v9 and add a delay-embedding
+readout over hidden states as an *additional* channel, orthogonal to the MAH's
+pairwise interpretant divergence.
+
+The motivation is that the SRT readouts (BEN r_hat, MAH divergence, Community
+Head prototype) all operate on per-token state. The trajectory through hidden
+state across positions has additional structure (curvature, recurrence,
+local-Lyapunov sensitivity) that none of the v9 readouts surface. A
+delay-embedding readout extracts two such signals per position:
+
+  - **Local Lyapunov estimate.** Approximates the local divergence rate of
+    nearby trajectories by comparing the displacement of a delay-coordinate
+    point at time t to the displacement of its k nearest neighbors in the
+    embedded space at time t+1. Positive values indicate locally chaotic
+    (sensitive) regions; near-zero values indicate locally stable regions.
+
+  - **Recurrence density.** Counts how many other positions in the sequence
+    fall within an epsilon-ball of the current position's delay coordinate.
+    High recurrence density indicates positions on a low-dimensional
+    attractor (a topic the model returns to); low density indicates novel
+    excursions.
+
+Both signals are computed from the same delay-coordinate stack and add a
+single learned linear readout each.
+
+This module is NOT wired into v9 SRTAdapter. It is a standalone prototype
+intended for v10 experiments where it will be added as a parallel readout
+alongside MAH/BEN/Community.
+
+References:
+  Takens, F. (1981). Detecting strange attractors in turbulence.
+  Haylett, K. R. (2025). Finite Tractus.
+  Haylett, K. R. (2026). Geofinitism / MARINA.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class THeadConfig:
+    """T-Head delay-embedding readout configuration."""
+
+    embedding_dim: int = 4  # m in Takens: number of delay coordinates
+    delay: int = 2  # tau: lag between successive coordinates (in tokens)
+    knn_k: int = 5  # neighbors used for local-Lyapunov estimate
+    recurrence_eps: float = 1.0  # eps-ball radius for recurrence density
+    project_to: int = 32  # dim of the per-position projection of hidden state
+
+
+@dataclass
+class THeadOutput:
+    """Per-token outputs from T-Head."""
+
+    lyapunov: torch.Tensor  # (B, T) local Lyapunov estimate
+    recurrence: torch.Tensor  # (B, T) recurrence density (count of neighbors)
+    delay_coord: torch.Tensor  # (B, T, m * project_to) raw delay-coord stack
+
+
+class TakensHead(nn.Module):
+    """Delay-embedding readout over a sequence of hidden states.
+
+    Given a hidden-state stream H of shape (B, T, D), the T-Head:
+      1. Projects H -> Z of shape (B, T, project_to) via a learned linear.
+      2. Builds a delay-coordinate stack X of shape (B, T, m, project_to)
+         where X[b, t, j] = Z[b, t - j*tau] (zero-padded for t < j*tau).
+      3. Flattens to V of shape (B, T, m * project_to).
+      4. Computes per-position local-Lyapunov estimate and recurrence density
+         from V.
+
+    The two scalar outputs (lyapunov, recurrence) can be supervised against
+    surrogate signals (e.g. CE-loss derivative for Lyapunov; topical
+    self-similarity for recurrence) or used as auxiliary inputs to downstream
+    heads in v10.
+
+    Notes
+    -----
+    The neighbor search and recurrence count are O(T^2) in sequence length;
+    for long sequences a chunked or random-projection approximation will be
+    needed. For prototype validation on the v9 evaluation sequences (all
+    <= 512 tokens) the naive implementation is acceptable.
+    """
+
+    def __init__(self, cfg: THeadConfig, d_hidden: int) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.project = nn.Linear(d_hidden, cfg.project_to, bias=False)
+        nn.init.normal_(self.project.weight, std=0.02)
+
+        flat_dim = cfg.embedding_dim * cfg.project_to
+        # Optional learned linear over the delay-coordinate vector for
+        # downstream wiring; v10 will decide whether to keep this raw or
+        # project it further.
+        self.readout = nn.Linear(flat_dim, flat_dim, bias=False)
+        nn.init.eye_(self.readout.weight)
+
+    def _build_delay_stack(self, z: torch.Tensor) -> torch.Tensor:
+        """Build (B, T, m, project_to) delay-coordinate stack."""
+        b, t, d = z.shape
+        m = self.cfg.embedding_dim
+        tau = self.cfg.delay
+        out = z.new_zeros(b, t, m, d)
+        for j in range(m):
+            shift = j * tau
+            if shift == 0:
+                out[:, :, j, :] = z
+            elif shift < t:
+                out[:, shift:, j, :] = z[:, : t - shift, :]
+            # else: leave zeros (positions earlier than t = j*tau)
+        return out
+
+    def _local_lyapunov(self, v: torch.Tensor) -> torch.Tensor:
+        """Estimate local Lyapunov exponent at each position.
+
+        For each position t, find the k nearest neighbors in the embedded
+        space and compare the mean inter-neighbor distance at t to the mean
+        inter-neighbor distance at t+1. The log of the ratio approximates the
+        local divergence rate (positive = chaotic).
+        """
+        b, t, d = v.shape
+        k = min(self.cfg.knn_k, t - 1)
+        if t < 2 or k <= 0:
+            return v.new_zeros(b, t)
+
+        # (B, T, T) pairwise distances
+        dist = torch.cdist(v, v)  # symmetric, zero diagonal
+        # exclude self by setting diagonal to +inf
+        eye_mask = torch.eye(t, device=v.device, dtype=torch.bool)
+        dist = dist.masked_fill(eye_mask.unsqueeze(0), float("inf"))
+
+        # nearest-neighbor indices (B, T, k)
+        nn_idx = dist.topk(k, dim=-1, largest=False).indices
+
+        # next-step indices (clamp at T-1)
+        nxt_idx = (nn_idx + 1).clamp(max=t - 1)
+
+        # gather neighbor positions at t and t+1
+        # v_nbr_t: (B, T, k, D); v_nbr_t1: (B, T, k, D)
+        v_nbr_t = torch.gather(
+            v.unsqueeze(1).expand(-1, t, -1, -1), 2,
+            nn_idx.unsqueeze(-1).expand(-1, -1, -1, d),
+        )
+        v_nbr_t1 = torch.gather(
+            v.unsqueeze(1).expand(-1, t, -1, -1), 2,
+            nxt_idx.unsqueeze(-1).expand(-1, -1, -1, d),
+        )
+
+        # current-position vector at t and t+1
+        v_t = v.unsqueeze(2)  # (B, T, 1, D)
+        nxt_self = torch.arange(t, device=v.device).clamp(max=t - 1) + 1
+        nxt_self = nxt_self.clamp(max=t - 1)
+        v_t1 = v[:, nxt_self, :].unsqueeze(2)  # (B, T, 1, D)
+
+        d_t = (v_nbr_t - v_t).norm(dim=-1).mean(dim=-1).clamp_min(1e-8)
+        d_t1 = (v_nbr_t1 - v_t1).norm(dim=-1).mean(dim=-1).clamp_min(1e-8)
+        return (d_t1 / d_t).log()
+
+    def _recurrence_density(self, v: torch.Tensor) -> torch.Tensor:
+        """Count neighbors within eps-ball at each position.
+
+        Returns a (B, T) float tensor with the number of other positions
+        within `recurrence_eps` of each position's delay coordinate.
+        """
+        b, t, _ = v.shape
+        if t < 2:
+            return v.new_zeros(b, t)
+        dist = torch.cdist(v, v)
+        eye_mask = torch.eye(t, device=v.device, dtype=torch.bool)
+        dist = dist.masked_fill(eye_mask.unsqueeze(0), float("inf"))
+        return (dist < self.cfg.recurrence_eps).float().sum(dim=-1)
+
+    def forward(self, hidden: torch.Tensor) -> THeadOutput:
+        z = self.project(hidden)  # (B, T, project_to)
+        stack = self._build_delay_stack(z)  # (B, T, m, project_to)
+        b, t, m, d = stack.shape
+        v = self.readout(stack.reshape(b, t, m * d))  # (B, T, m*project_to)
+        lyap = self._local_lyapunov(v)
+        rec = self._recurrence_density(v)
+        return THeadOutput(lyapunov=lyap, recurrence=rec, delay_coord=v)
+
+
+__all__ = ["THeadConfig", "THeadOutput", "TakensHead"]

srt/nla/__init__.py

@@ -0,0 +1,54 @@
+"""SRT-NLA — Natural Language Autoencoder over the SRT backbone.
+
+Two halves:
+
+- ``ActivationReconstructor`` (AR): run the frozen backbone forward and read
+  hidden states at layer ``L``. Zero new parameters.
+- ``ActivationVerbalizer`` (AV): inject a target activation vector as a
+  prefix input-embedding into the frozen backbone and decode text whose
+  re-encoded activation at layer ``L`` matches the target.
+
+The AV side carries the only learned parameters in N0–N2 (a small adapter
+on top of the frozen backbone). After N2 the SRT adapter (12.7M) becomes
+the AV trunk so divergence/community channels condition the generation.
+
+See ``docs/SRT_NLA_PLAN.md`` for the phasing.
+"""
+
+from __future__ import annotations
+
+from srt.nla.config import NLAConfig
+from srt.nla.loss import (
+    cosine_similarity,
+    fraction_variance_explained,
+    mse_nrm,
+    nla_reward,
+)
+from srt.nla.metrics import (
+    Anchors,
+    anisotropy_mu,
+    fve_nrm,
+    fve_nrm_centered,
+    rho_norm,
+)
+from srt.nla.reconstructor import ActivationReconstructor
+from srt.nla.sidecar import NLAMeta, load_meta, save_meta
+from srt.nla.verbalizer import ActivationVerbalizer
+
+__all__ = [
+    "ActivationReconstructor",
+    "ActivationVerbalizer",
+    "Anchors",
+    "NLAConfig",
+    "NLAMeta",
+    "anisotropy_mu",
+    "cosine_similarity",
+    "fraction_variance_explained",
+    "fve_nrm",
+    "fve_nrm_centered",
+    "load_meta",
+    "mse_nrm",
+    "nla_reward",
+    "rho_norm",
+    "save_meta",
+]

srt/nla/config.py

@@ -0,0 +1,76 @@
+"""Configuration for SRT-NLA."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, fields
+from pathlib import Path
+
+
+@dataclass
+class NLAConfig:
+    """Natural Language Autoencoder configuration.
+
+    Defaults target Qwen2.5-7B (28 hidden layers, hidden_size=3584). The
+    NLA reference paper extracts at L=20 (~2/3 depth); we keep the same.
+    """
+
+    backbone_id: str = "Qwen/Qwen2.5-7B"
+    backbone_dtype: str = "bfloat16"
+
+    # Layer to read activations from (1-indexed into output_hidden_states,
+    # which is length num_hidden_layers + 1 with index 0 = embeddings).
+    extraction_layer: int = 20
+
+    # Vector dimensionality. None = use backbone hidden_size.
+    d_vector: int | None = None
+
+    # AV generation
+    max_new_tokens: int = 128
+    temperature: float = 1.0
+    top_p: float = 1.0
+
+    # Pooling for AR
+    pool: str = "last"  # "last" | "mean" | "first"
+
+    # Reward weights (see docs/SRT_NLA_PLAN.md §5)
+    lambda_mag: float = 0.1  # magnitude-match penalty
+    beta_kl: float = 0.05  # legibility KL vs base
+    gamma_entropy: float = 0.05  # entropy hinge coefficient (both sides)
+    delta_community: float = 0.2  # community-consistency bonus
+    h_min: float = 1.5  # minimum acceptable token-level entropy (nats)
+    h_max: float = 3.0  # maximum acceptable token-level entropy (nats); above this the upper hinge pulls H back down to prevent uniform-noise runaway
+
+    # Injection layout — number of learned prefix tokens after the injected
+    # vector. The injected vector occupies slot 0; this many extra learned
+    # embeddings (BOS-like) follow before the model autoregresses.
+    num_prefix_tokens: int = 1
+
+    # Prefix mode:
+    #   "static" — `prefix_embeds` is a (P, d_embed) learned tensor shared
+    #              across all inputs (original design).
+    #   "mlp"    — a 2-layer MLP maps v → (P, d_embed) per-input, so the
+    #              prefix is conditioned on the target vector. Strictly
+    #              more expressive; same trainable shape elsewhere.
+    prefix_mode: str = "static"
+    prefix_mlp_hidden: int = 256
+
+    # Multi-position injection: number of slots at which proj(v) is injected
+    # (with independent linear projections per slot). 1 = original single-slot
+    # design. >1 gives the backbone several independent views of v before the
+    # learned prefix; intended to break the single-slot information bottleneck.
+    num_inject_slots: int = 1
+
+    # ---- (de)serialization helpers --------------------------------------
+
+    def to_json(self, path: str | Path) -> None:
+        """Write this config to a JSON file (used by HF model cards)."""
+        Path(path).write_text(json.dumps(asdict(self), indent=2, sort_keys=True))
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> "NLAConfig":
+        """Load a config from JSON. Unknown keys are ignored for
+        forward-compatibility with newer checkpoints."""
+        data = json.loads(Path(path).read_text())
+        known = {f.name for f in fields(cls)}
+        return cls(**{k: v for k, v in data.items() if k in known})

srt/nla/decoding.py

@@ -0,0 +1,249 @@
+"""Deployable decoding methods for the SRT-NLA verbalizer.
+
+The headline result of `scripts/rerank_eval.py` is that **oracle-scored
+best-of-K** closes the greedy → ceiling gap on every checkpoint we have:
+greedy ρ_norm ≈ 0.26, K=64 oracle ρ_norm ≈ 0.92 against a paraphrase
+ceiling of 1.00. Because the target activation ``v`` is, by construction,
+available at deploy time in the SRT-NLA setup (the system was given ``v``
+in order to verbalize it), oracle scoring is a *free* deploy-time signal,
+not an unfair eval-time cheat.
+
+This module packages that decoding method into a clean reusable function
+so it can be called from train scripts, eval scripts, demos, or product
+code. The cost is K-way sampling + one batched scoring forward pass.
+
+Two scoring options:
+- **centered** (recommended): ``0.5 * (1 + cos(h_last - mu, v - mu))``
+  with ``mu`` the per-coordinate mean of a held-out pool of real
+  backbone-L activations. Anisotropy-corrected.
+- **raw**: ``0.5 * (1 + cos(h_last, v))``. Inflated by Qwen L20's
+  cos≈0.24 anisotropy floor; reported only for back-compatibility with
+  the pre-`centered_eval.py` literature.
+
+ρ_norm conversion (the paper's units) is provided as a helper:
+``rho_norm(cen) = (cen - random_floor) / (paraphrase_ceiling - random_floor)``
+with the calibrated constants ``random_floor = 0.510`` and
+``paraphrase_ceiling = 0.799``.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+from srt.nla.verbalizer import ActivationVerbalizer
+
+# Calibrated anchors from `scripts/centered_eval.py` on Qwen2.5-7B L20
+# with pool_size=2000 of real Qwen activations. See SESSION_HANDOFF.md.
+RANDOM_FLOOR_CEN = 0.510
+PARAPHRASE_CEILING_CEN = 0.799
+
+
+def rho_norm(cen: float | torch.Tensor) -> float | torch.Tensor:
+    """Convert centered_fve to the paper's normalized ρ_norm units.
+
+    ρ_norm = 0.0 corresponds to random retrieval, 1.0 to the paraphrase
+    ceiling. Values can fall slightly below 0 or above 1 due to sampling
+    noise on small M or unusually-clean checkpoints.
+    """
+    return (cen - RANDOM_FLOOR_CEN) / (PARAPHRASE_CEILING_CEN - RANDOM_FLOOR_CEN)
+
+
+def _fve_from_cos(c: torch.Tensor) -> torch.Tensor:
+    return 0.5 * (1.0 + c)
+
+
+def _cos_pairs(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    an = F.normalize(a.float(), dim=-1)
+    bn = F.normalize(b.float(), dim=-1)
+    return (an * bn).sum(-1)
+
+
+@torch.no_grad()
+def _rollout(
+    av: ActivationVerbalizer,
+    v: torch.Tensor,
+    *,
+    max_new_tokens: int,
+    do_sample: bool,
+    temperature: float,
+    eos_id: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Generate (B, T) token ids and (B, T) attention mask truncated at first EOS."""
+    gen_ids = av.generate(
+        v,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        temperature=max(temperature, 1e-6),
+        top_p=1.0,
+    )
+    B, T = gen_ids.shape
+    device = gen_ids.device
+    is_eos = gen_ids == eos_id
+    pos = torch.arange(T, device=device).unsqueeze(0).expand(B, T)
+    first_eos = torch.where(
+        is_eos.any(dim=-1, keepdim=True),
+        is_eos.float().argmax(dim=-1, keepdim=True).long(),
+        torch.full((B, 1), T - 1, device=device, dtype=torch.long),
+    )
+    gen_attn = (pos <= first_eos).long()
+    return gen_ids, gen_attn
+
+
+@torch.no_grad()
+def _h_last_prefix_free(
+    backbone,
+    gen_ids: torch.Tensor,
+    gen_attn: torch.Tensor,
+    layer: int,
+) -> torch.Tensor:
+    """Read h_last from a *prefix-free* re-forward of the generated tokens.
+
+    Critical: scoring must NOT include the proj(v) prefix in the forward
+    used to extract h_last, because the prefix would leak v directly into
+    the activation and inflate the metric. This matches centered_eval.py
+    and the deployment-realistic measurement (text → embed → read activation).
+    """
+    out = backbone(
+        input_ids=gen_ids,
+        attention_mask=gen_attn,
+        output_hidden_states=True,
+        use_cache=False,
+    )
+    full_h = out.hidden_states[layer]
+    B = gen_ids.size(0)
+    last_idx = (gen_attn.sum(-1) - 1).clamp(min=0).long()
+    rows = torch.arange(B, device=full_h.device)
+    return full_h[rows, last_idx].float()
+
+
+@dataclass
+class OracleRerankResult:
+    """Result of oracle-rerank decoding for a batch of B target vectors."""
+
+    best_ids: torch.Tensor          # (B, T) chosen rollout token ids
+    best_attn: torch.Tensor         # (B, T) attention mask
+    best_h_last: torch.Tensor       # (B, d) chosen rollout h_last
+    best_cen: torch.Tensor          # (B,) centered_fve scores of the chosen rollout
+    best_idx: torch.Tensor          # (B,) which of the K candidates was chosen
+    all_cen: torch.Tensor | None    # (B, K) scores for every candidate (if return_all)
+    all_h_last: torch.Tensor | None # (B, K, d) (if return_all)
+    all_ids: torch.Tensor | None    # (B, K, T) (if return_all)
+
+
+@torch.no_grad()
+def oracle_rerank_decode(
+    av: ActivationVerbalizer,
+    backbone,
+    v: torch.Tensor,
+    *,
+    layer: int,
+    K: int = 64,
+    max_new_tokens: int = 64,
+    temperature: float = 1.0,
+    mu: torch.Tensor | None = None,
+    score_batch_size: int = 64,
+    eos_id: int | None = None,
+    return_all: bool = False,
+) -> OracleRerankResult:
+    """Sample K rollouts per target, score with centered_fve, return the argmax.
+
+    This is the recommended deploy-time decoding method for SRT-NLA when
+    ``v`` is available (which it is, by construction).
+
+    Args:
+        av: an ``ActivationVerbalizer`` already loaded with weights and on
+            the same device as ``v``.
+        backbone: the frozen causal-LM backbone (same one wrapped by ``av``).
+        v: target activation vectors, shape ``(B, d)``.
+        layer: which hidden-state index to read for h_last. The AR/AV
+            extraction layer (typically 20 for Qwen2.5-7B).
+        K: number of samples per target.
+        max_new_tokens: rollout length cap.
+        temperature: sampling temperature (top-p is fixed to 1.0 here; the
+            K-curve is dominated by K, not nucleus).
+        mu: anisotropy mean ``(1, d)``. If ``None``, falls back to **raw**
+            scoring (cos to v with no centering); strongly recommended to
+            pass an actual pool-mean for anisotropy-correct comparison
+            against published numbers.
+        score_batch_size: cap concurrent (candidate, prefix) pairs in the
+            scoring forward to bound memory.
+        eos_id: token id used to truncate rollouts. Defaults to the AV
+            tokenizer's eos.
+        return_all: if True, also return per-candidate scores + h_last + ids.
+
+    Returns:
+        ``OracleRerankResult``. ``best_text`` is not auto-decoded — call
+        ``av.tokenizer.batch_decode(result.best_ids, skip_special_tokens=True)``.
+    """
+    if v.dim() != 2:
+        raise ValueError(f"v must be (B, d); got shape {tuple(v.shape)}")
+    B, d = v.shape
+    device = v.device
+    if eos_id is None:
+        eos_id = int(av.tokenizer.eos_token_id)
+
+    # 1) K rollouts per target — flatten to (B*K, ...) for one big sampling call.
+    v_rep = v.unsqueeze(1).expand(B, K, d).reshape(B * K, d)
+    s_ids, s_attn = _rollout(
+        av, v_rep,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        temperature=temperature,
+        eos_id=eos_id,
+    )
+
+    # 2) Score in chunks to bound memory.
+    h_chunks: list[torch.Tensor] = []
+    for j in range(0, s_ids.size(0), score_batch_size):
+        h = _h_last_prefix_free(
+            backbone,
+            s_ids[j : j + score_batch_size],
+            s_attn[j : j + score_batch_size],
+            layer,
+        )
+        h_chunks.append(h)
+    h_all = torch.cat(h_chunks, dim=0)                              # (B*K, d)
+
+    # 3) Score: centered if mu given, raw otherwise.
+    if mu is not None:
+        if mu.dim() == 1:
+            mu = mu.unsqueeze(0)
+        mu_dev = mu.to(device=device, dtype=torch.float32)
+        cen = _fve_from_cos(_cos_pairs(h_all - mu_dev, v_rep - mu_dev))
+    else:
+        cen = _fve_from_cos(_cos_pairs(h_all, v_rep))
+
+    cen = cen.view(B, K)
+    h_BK = h_all.view(B, K, d)
+    ids_BK = s_ids.view(B, K, -1)
+    attn_BK = s_attn.view(B, K, -1)
+
+    # 4) Argmax over K per target.
+    best_idx = cen.argmax(dim=1)                                    # (B,)
+    rows = torch.arange(B, device=device)
+    best_cen = cen[rows, best_idx]
+    best_h = h_BK[rows, best_idx]
+    best_ids = ids_BK[rows, best_idx]
+    best_attn = attn_BK[rows, best_idx]
+
+    return OracleRerankResult(
+        best_ids=best_ids,
+        best_attn=best_attn,
+        best_h_last=best_h,
+        best_cen=best_cen,
+        best_idx=best_idx,
+        all_cen=cen if return_all else None,
+        all_h_last=h_BK if return_all else None,
+        all_ids=ids_BK if return_all else None,
+    )
+
+
+__all__ = [
+    "RANDOM_FLOOR_CEN",
+    "PARAPHRASE_CEILING_CEN",
+    "rho_norm",
+    "OracleRerankResult",
+    "oracle_rerank_decode",
+]

srt/nla/loss.py

@@ -0,0 +1,143 @@
+"""Losses and metrics for SRT-NLA.
+
+The canonical NLA quality metric is direction MSE on L2-normalised
+vectors (``mse_nrm``); ``fve_nrm`` is the fraction of variance explained
+relative to a uniform random unit baseline (which has ``E[mse_nrm] = 2``).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+
+def mse_nrm(
+    v_hat: torch.Tensor, v_target: torch.Tensor, *, reduce: bool = True
+) -> torch.Tensor:
+    """Squared L2 distance between unit-normalised vectors. Lower is better.
+
+    Both tensors are ``(..., d)``. With ``reduce=True`` (default) returns a
+    scalar (mean over leading dims); with ``reduce=False`` returns ``(...,)``
+    per-sample values — needed by REINFORCE so the advantage baseline
+    ``r - r.mean()`` is not identically zero.
+    """
+    a = F.normalize(v_hat, dim=-1)
+    b = F.normalize(v_target, dim=-1)
+    per_sample = ((a - b) ** 2).sum(dim=-1)
+    return per_sample.mean() if reduce else per_sample
+
+
+def cosine_similarity(v_hat: torch.Tensor, v_target: torch.Tensor) -> torch.Tensor:
+    """Mean cosine similarity. ``= 1 - mse_nrm/2``."""
+    return F.cosine_similarity(v_hat, v_target, dim=-1).mean()
+
+
+def fraction_variance_explained(
+    v_hat: torch.Tensor, v_target: torch.Tensor
+) -> torch.Tensor:
+    """FVE on L2-normalised vectors.
+
+    Random uniform unit vectors have ``E[||u - v||^2] = 2``, so
+    ``fve_nrm = 1 - mse_nrm / 2``. The Qwen2.5-7B NLA reference reports
+    ``fve_nrm = 0.752`` at L=20.
+    """
+    return 1.0 - mse_nrm(v_hat, v_target) / 2.0
+
+
+def magnitude_penalty(
+    v_hat: torch.Tensor, v_target: torch.Tensor, *, reduce: bool = True
+) -> torch.Tensor:
+    """Scale-invariant magnitude mismatch.
+
+    Returns the squared *relative* error of the L2 norms:
+
+        ((||v_hat|| - ||v_target||) / ||v_target||) ** 2
+
+    Hidden states at deep layers of large LMs can have L2 norms in the
+    1e3-1e4 range; the previous formulation ``(||v_hat|| - ||v_target||)^2``
+    produced reward magnitudes of order 1e8, swamping ``mse_nrm`` (which is
+    bounded in [0, 4]) by 8 orders of magnitude regardless of
+    ``lambda_mag``. The relative form is bounded in roughly [0, 1] for
+    sensible predictions, matching ``mse_nrm``'s scale.
+    """
+    target_norm = v_target.norm(dim=-1).clamp(min=1e-6)
+    per_sample = ((v_hat.norm(dim=-1) - v_target.norm(dim=-1)) / target_norm).pow(2)
+    return per_sample.mean() if reduce else per_sample
+
+
+def entropy_floor_hinge(
+    token_entropy: torch.Tensor, h_min: float, *, reduce: bool = True
+) -> torch.Tensor:
+    """Hinge that fires when per-token entropy drops below ``h_min`` nats."""
+    per_sample = torch.clamp(h_min - token_entropy, min=0.0)
+    return per_sample.mean() if reduce else per_sample
+
+
+@dataclass
+class RewardComponents:
+    mse: torch.Tensor
+    mag: torch.Tensor
+    kl: torch.Tensor
+    entropy_hinge: torch.Tensor
+    community: torch.Tensor
+    total: torch.Tensor
+
+
+def nla_reward(
+    v_hat: torch.Tensor,
+    v_target: torch.Tensor,
+    *,
+    kl_av_vs_base: torch.Tensor | None = None,
+    token_entropy: torch.Tensor | None = None,
+    community_consistency: torch.Tensor | None = None,
+    lambda_mag: float = 0.1,
+    beta_kl: float = 0.05,
+    gamma_entropy: float = 0.05,
+    delta_community: float = 0.2,
+    h_min: float = 1.5,
+    reduce: bool = True,
+) -> RewardComponents:
+    """5-term NLA reward (higher is better).
+
+    ``r = -mse_nrm - lambda_mag * mag - beta_kl * KL - gamma * H_hinge + delta * comm``
+
+    Any optional term left as ``None`` is treated as zero (so N0/N1 can
+    compose the reward incrementally as the surrounding modules come
+    online).
+
+    With ``reduce=False`` every component is per-sample ``(B,)``; with
+    ``reduce=True`` (default) every component is a scalar. REINFORCE must
+    use ``reduce=False`` so ``advantage = r - r.mean()`` is non-degenerate.
+    """
+    device = v_hat.device
+    dtype = v_hat.dtype
+    B = v_hat.shape[0] if v_hat.dim() >= 1 else 1
+    zero_shape = () if reduce else (B,)
+    zero = torch.zeros(zero_shape, device=device, dtype=dtype)
+
+    def _maybe_per_sample(x: torch.Tensor) -> torch.Tensor:
+        if reduce:
+            return x.mean() if x.dim() > 0 else x
+        # caller passed a scalar but we want per-sample → broadcast.
+        return x.expand(B) if x.dim() == 0 else x
+
+    mse = mse_nrm(v_hat, v_target, reduce=reduce)
+    mag = magnitude_penalty(v_hat, v_target, reduce=reduce)
+    kl = _maybe_per_sample(kl_av_vs_base) if kl_av_vs_base is not None else zero
+    ent = (
+        entropy_floor_hinge(token_entropy, h_min, reduce=reduce)
+        if token_entropy is not None
+        else zero
+    )
+    comm = (
+        _maybe_per_sample(community_consistency)
+        if community_consistency is not None
+        else zero
+    )
+
+    total = -mse - lambda_mag * mag - beta_kl * kl - gamma_entropy * ent + delta_community * comm
+    return RewardComponents(
+        mse=mse, mag=mag, kl=kl, entropy_hinge=ent, community=comm, total=total
+    )

srt/nla/metrics.py

@@ -0,0 +1,89 @@
+"""Canonical NLA evaluation metrics.
+
+Single source of truth for fve_nrm, the centered (anisotropy-corrected)
+variant, and the rho_norm normalization used in paper_nla.md.
+
+Previously these helpers were duplicated in 6+ scripts. New code should
+import from here. The duplicates remain in legacy scripts for now to
+preserve bit-exact reproducibility of historical artifacts/nla/*.json.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+
+# Anchors from artifacts/nla/oracle_ceiling_30k_v2.json (Qwen2.5-7B L20,
+# 200 held-out targets, pool=2000). Used to normalize centered fve_nrm
+# into rho_norm ∈ [0, 1].
+RANDOM_FLOOR_CEN: float = 0.510
+PARAPHRASE_CEILING_CEN: float = 0.799
+RHO_DENOM: float = PARAPHRASE_CEILING_CEN - RANDOM_FLOOR_CEN  # 0.289
+
+
+def fve_nrm(h: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """fve_nrm(h, v) = 0.5 * (1 + cos(h, v))  in [0, 1].
+
+    Linear remap of cosine similarity. Operates on the last dim.
+    Shapes: h, v broadcast-compatible; returns batch shape.
+    """
+    return 0.5 * (1.0 + F.cosine_similarity(h.float(), v.float(), dim=-1))
+
+
+def fve_nrm_centered(
+    h: torch.Tensor, v: torch.Tensor, mu: torch.Tensor
+) -> torch.Tensor:
+    """Anisotropy-corrected fve_nrm: subtract pool mean mu before cosine.
+
+    mu is the (d,) mean activation of the held-out pool at the same
+    layer/backbone. Without this correction, Qwen2.5-7B L20 activations
+    sit at cos≈0.24 baseline (||mu||≈55) which inflates raw fve_nrm by
+    ~0.11 and compresses dynamic range.
+    """
+    return fve_nrm(h - mu, v - mu)
+
+
+def rho_norm(cen: torch.Tensor | float) -> torch.Tensor | float:
+    """Normalize centered fve_nrm to [0, 1] using paper anchors.
+
+    rho = (cen - random_floor) / (paraphrase_ceiling - random_floor)
+
+    rho = 0 ↔ unrelated; rho = 1 ↔ saturates Qwen paraphrase ceiling.
+    """
+    if isinstance(cen, torch.Tensor):
+        return (cen - RANDOM_FLOOR_CEN) / RHO_DENOM
+    return (cen - RANDOM_FLOOR_CEN) / RHO_DENOM
+
+
+def anisotropy_mu(pool: torch.Tensor) -> torch.Tensor:
+    """Compute the anisotropy mean mu = pool.mean(0).
+
+    pool: (N, d) of last-token L20 activations from a held-out target file.
+    Returns (d,) float32.
+    """
+    return pool.float().mean(0)
+
+
+@dataclass(frozen=True)
+class Anchors:
+    """Reference points reported in paper_nla.md §3."""
+
+    random_floor_cen: float = RANDOM_FLOOR_CEN
+    paraphrase_ceiling_cen: float = PARAPHRASE_CEILING_CEN
+    nn_in_pool_cen: float = 0.663  # pool=200
+    nn_retrieval_cen: float = 0.714  # pool=2000
+    replay_cen: float = 0.968
+
+
+__all__ = [
+    "fve_nrm",
+    "fve_nrm_centered",
+    "rho_norm",
+    "anisotropy_mu",
+    "Anchors",
+    "RANDOM_FLOOR_CEN",
+    "PARAPHRASE_CEILING_CEN",
+    "RHO_DENOM",
+]

srt/nla/reconstructor.py

@@ -0,0 +1,122 @@
+"""ActivationReconstructor (AR): frozen-backbone hidden-state reader.
+
+Zero learned parameters in the canonical setup — the "reconstruction" is
+just the forward pass through layers 0..L of the frozen backbone.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Literal
+
+import torch
+import torch.nn as nn
+from transformers import AutoModelForCausalLM
+
+from srt.nla.config import NLAConfig
+
+logger = logging.getLogger(__name__)
+
+_DTYPE_MAP = {
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+}
+
+
+class ActivationReconstructor(nn.Module):
+    """Reads pooled hidden states at ``cfg.extraction_layer`` from a text input.
+
+    Pass an existing ``backbone`` to share weights with an
+    ``ActivationVerbalizer`` (recommended — saves 14GB of RAM).
+    """
+
+    def __init__(
+        self,
+        cfg: NLAConfig,
+        backbone: AutoModelForCausalLM | None = None,
+    ) -> None:
+        super().__init__()
+        self.cfg = cfg
+        if backbone is None:
+            dtype = _DTYPE_MAP.get(cfg.backbone_dtype, torch.bfloat16)
+            logger.info("AR: loading backbone %s (%s)", cfg.backbone_id, cfg.backbone_dtype)
+            backbone = AutoModelForCausalLM.from_pretrained(
+                cfg.backbone_id, torch_dtype=dtype
+            )
+            for p in backbone.parameters():
+                p.requires_grad = False
+            backbone.eval()
+        self.backbone = backbone
+        self._d = self.backbone.config.hidden_size
+        self._num_layers = self.backbone.config.num_hidden_layers
+        if not (0 < cfg.extraction_layer <= self._num_layers):
+            raise ValueError(
+                f"extraction_layer={cfg.extraction_layer} out of range "
+                f"(backbone has {self._num_layers} hidden layers)"
+            )
+
+    @property
+    def d_vector(self) -> int:
+        return self.cfg.d_vector or self._d
+
+    def _pool(
+        self,
+        h_layer: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        mode: Literal["last", "mean", "first"],
+    ) -> torch.Tensor:
+        if mode == "first":
+            return h_layer[:, 0]
+        if mode == "mean":
+            if attention_mask is None:
+                return h_layer.mean(dim=1)
+            mask = attention_mask.unsqueeze(-1).to(h_layer.dtype)
+            return (h_layer * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+        # default: "last" non-pad token
+        if attention_mask is None:
+            return h_layer[:, -1]
+        lengths = (attention_mask.sum(dim=1).long() - 1).clamp(min=0)
+        idx = torch.arange(h_layer.size(0), device=h_layer.device)
+        return h_layer[idx, lengths]
+
+    @torch.no_grad()
+    def reconstruct(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        pool: str | None = None,
+    ) -> torch.Tensor:
+        """Return pooled hidden state ``(B, d)`` at layer ``cfg.extraction_layer``."""
+        out = self.backbone(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            use_cache=False,
+        )
+        # hidden_states is a tuple of length num_hidden_layers + 1
+        h_layer = out.hidden_states[self.cfg.extraction_layer]
+        return self._pool(h_layer, attention_mask, pool or self.cfg.pool)  # type: ignore[arg-type]
+
+    def reconstruct_from_embeds(
+        self,
+        inputs_embeds: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        pool: str | None = None,
+    ) -> torch.Tensor:
+        """Differentiable AR forward from soft input embeddings.
+
+        Used by the Phase-2 soft-embedding bridge. Caller passes
+        ``softmax(logits / tau) @ E_token`` as ``inputs_embeds`` so that
+        the entire pipeline (AV -> soft embeds -> backbone -> pool) has
+        a per-token-per-dim gradient w.r.t. AV parameters. Backbone
+        weights stay frozen (requires_grad=False set in __init__).
+        """
+        out = self.backbone(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            use_cache=False,
+        )
+        h_layer = out.hidden_states[self.cfg.extraction_layer]
+        return self._pool(h_layer, attention_mask, pool or self.cfg.pool)  # type: ignore[arg-type]

srt/nla/sidecar.py

@@ -0,0 +1,56 @@
+"""Sidecar metadata for an NLA checkpoint (``nla_meta.yaml``)."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class NLAMeta:
+    """What an NLA checkpoint represents.
+
+    Persisted alongside the weights so a consumer can answer:
+    "which backbone, which layer, what vector convention".
+    """
+
+    backbone_id: str
+    backbone_revision: str | None
+    extraction_layer: int
+    d_vector: int
+    pool: str
+    norm_convention: str = "raw"  # "raw" | "unit" | "z"
+    pair: str = "av"  # "av" | "ar" | "av_ar"
+    phase: str = "N0"
+    notes: str = ""
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def save_meta(meta: NLAMeta, path: str | Path) -> None:
+    """Write ``meta`` as YAML if PyYAML is available, else JSON."""
+    p = Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    payload = meta.to_dict()
+    try:
+        import yaml  # type: ignore[import-not-found]
+
+        p.write_text(yaml.safe_dump(payload, sort_keys=False))
+    except ImportError:
+        p.write_text(json.dumps(payload, indent=2))
+
+
+def load_meta(path: str | Path) -> NLAMeta:
+    p = Path(path)
+    text = p.read_text()
+    try:
+        import yaml  # type: ignore[import-not-found]
+
+        data = yaml.safe_load(text)
+    except ImportError:
+        data = json.loads(text)
+    return NLAMeta(**data)

srt/nla/targets_check.py

@@ -0,0 +1,75 @@
+"""Shared precondition checks + fingerprint for NLA target pools.
+
+A whole iteration was wasted in May 2026 training on a degenerate target
+pool where every sample collapsed to the same (~all-zeros) activation
+because Qwen2.5's ``bos_token_id == eos_token_id`` caused the BOS slot
+to be picked as the "last token" by ``sample_targets.py``. Loss
+nonetheless decreased and ``fve_nrm`` plateaued at a fake ~0.62
+because the AV learned to emit one constant string. Nothing in any
+train script noticed.
+
+These helpers fail loudly *before* the optimizer runs if the pool is
+degenerate or constant-norm, and emit a one-line fingerprint
+(sha + N + per-elem std + norm mean/std) into the train log header so
+every run is traceable to the exact data version it consumed.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+from pathlib import Path
+
+import torch
+
+logger = logging.getLogger("nla.targets_check")
+
+
+def fingerprint_targets(path: Path, pool: torch.Tensor) -> dict:
+    """Compute sha256 (first 16 hex) + summary stats of the pool."""
+    p = Path(path)
+    h = hashlib.sha256()
+    with p.open("rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    sha16 = h.hexdigest()[:16]
+    pool_f = pool.float()
+    per_elem_std = float(pool_f.std(dim=0).mean().item())
+    norms = pool_f.norm(dim=-1)
+    return {
+        "path": str(p),
+        "sha16": sha16,
+        "n": int(pool.size(0)),
+        "d": int(pool.size(1)),
+        "per_elem_std": per_elem_std,
+        "norm_mean": float(norms.mean().item()),
+        "norm_std": float(norms.std().item()),
+    }
+
+
+def assert_targets_healthy(
+    path: Path,
+    pool: torch.Tensor,
+    *,
+    min_per_elem_std: float = 0.1,
+    min_norm_std: float = 1.0,
+) -> dict:
+    """Fingerprint + log + assert pool is non-degenerate. Returns fingerprint."""
+    fp = fingerprint_targets(path, pool)
+    logger.info(
+        "TARGETS sha=%s N=%d d=%d per_elem_std=%.4f norm=%.3f±%.3f",
+        fp["sha16"], fp["n"], fp["d"],
+        fp["per_elem_std"], fp["norm_mean"], fp["norm_std"],
+    )
+    if fp["per_elem_std"] < min_per_elem_std:
+        raise RuntimeError(
+            f"DEGENERATE target pool {path}: per-elem std={fp['per_elem_std']:.4f} "
+            f"< {min_per_elem_std}. Almost certainly a sampling bug (e.g. "
+            f"BOS==EOS collapse). Refusing to train."
+        )
+    if fp["norm_std"] < min_norm_std:
+        raise RuntimeError(
+            f"CONSTANT-NORM target pool {path}: norm std={fp['norm_std']:.4f} "
+            f"< {min_norm_std}. Targets are not diverse. Refusing to train."
+        )
+    return fp

srt/nla/verbalizer.py

@@ -0,0 +1,213 @@
+"""ActivationVerbalizer (AV): inject a target vector, decode text.
+
+The frozen backbone runs over a prefix consisting of:
+
+    [proj(v) , learned_prefix_embeds..., generated_tokens...]
+
+``proj`` is the only trainable module in N0–N2 (a linear map from the
+vector space to the backbone embedding space). After N2 the SRT adapter
+replaces ``proj`` with the full 12.7M-parameter adapter pipeline so the
+divergence and community channels condition the generation.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+import torch.nn as nn
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from srt.nla.config import NLAConfig
+
+logger = logging.getLogger(__name__)
+
+_DTYPE_MAP = {
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+}
+
+
+class ActivationVerbalizer(nn.Module):
+    """Injection-prefix language generator.
+
+    ``forward`` returns logits over a teacher-forced sequence (for training).
+    ``generate`` returns sampled token ids (for evaluation / RL rollout).
+    """
+
+    def __init__(
+        self,
+        cfg: NLAConfig,
+        backbone: AutoModelForCausalLM | None = None,
+        tokenizer: AutoTokenizer | None = None,
+    ) -> None:
+        super().__init__()
+        self.cfg = cfg
+        dtype = _DTYPE_MAP.get(cfg.backbone_dtype, torch.bfloat16)
+        if backbone is None:
+            logger.info("AV: loading backbone %s (%s)", cfg.backbone_id, cfg.backbone_dtype)
+            backbone = AutoModelForCausalLM.from_pretrained(
+                cfg.backbone_id, torch_dtype=dtype
+            )
+            for p in backbone.parameters():
+                p.requires_grad = False
+            backbone.eval()
+        self.backbone = backbone
+        if tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained(cfg.backbone_id)
+        self.tokenizer = tokenizer
+
+        self._d_embed = self.backbone.config.hidden_size
+        self._backbone_dtype = dtype
+        d_vec = cfg.d_vector or self._d_embed
+
+        # Trainable projection from vector space → backbone embedding space.
+        # Initialised to identity when shapes match; small random otherwise.
+        # NOTE: adapter params are kept in float32 even when the backbone is
+        # bf16/fp16 — AdamW on low-precision params is numerically unstable
+        # and the eye-init drifts to noise within a few hundred steps. We
+        # cast to backbone dtype only at the boundary inside _inject_prefix.
+        self.proj = nn.Linear(d_vec, self._d_embed, bias=False)
+        with torch.no_grad():
+            if d_vec == self._d_embed:
+                nn.init.eye_(self.proj.weight)
+            else:
+                nn.init.xavier_uniform_(self.proj.weight)
+
+        # Multi-position injection: M independent projections of v placed at
+        # the first M slots of the prefix. proj_extra holds slots 1..M-1
+        # (proj covers slot 0). Initialised small so the model starts near
+        # the single-slot baseline and the extra slots add information as
+        # training progresses.
+        n_inject = max(1, getattr(cfg, "num_inject_slots", 1))
+        self._n_inject = n_inject
+        if n_inject > 1:
+            self.proj_extra = nn.ModuleList(
+                [nn.Linear(d_vec, self._d_embed, bias=False) for _ in range(n_inject - 1)]
+            )
+            with torch.no_grad():
+                for lin in self.proj_extra:
+                    nn.init.xavier_uniform_(lin.weight, gain=0.1)
+        else:
+            self.proj_extra = None
+
+        # Trainable extra prefix embeddings that follow the injected vector.
+        # These give the model a small "header" before it starts generating
+        # describing text.
+        #
+        # Two modes:
+        #   "static" — learned (P, d_embed) tensor, shared across inputs.
+        #              Initialised from the backbone's BOS embedding.
+        #   "mlp"    — 2-layer MLP maps v → (P, d_embed) per-input. Init
+        #              second layer to zeros so prefix starts equal to the
+        #              BOS embedding (the bias of layer 2), then drifts.
+        # Adapter params are float32 (see note above).
+        n_pref = max(0, cfg.num_prefix_tokens)
+        self._n_pref = n_pref
+        self._prefix_mode = getattr(cfg, "prefix_mode", "static")
+        if n_pref > 0:
+            bos_id = (
+                self.backbone.config.bos_token_id
+                or self.tokenizer.bos_token_id
+                or 0
+            )
+            with torch.no_grad():
+                embed = self.backbone.get_input_embeddings()
+                bos_embed = (
+                    embed(torch.tensor([bos_id], device=embed.weight.device))
+                    .detach()
+                    .clone()
+                    .float()
+                )  # (1, d) float32
+            if self._prefix_mode == "mlp":
+                hidden = getattr(cfg, "prefix_mlp_hidden", 256)
+                out_dim = n_pref * self._d_embed
+                self.prefix_mlp = nn.Sequential(
+                    nn.Linear(d_vec, hidden),
+                    nn.GELU(),
+                    nn.Linear(hidden, out_dim),
+                )
+                with torch.no_grad():
+                    # Zero output weights so prefix = bias (set to tiled BOS)
+                    # — preserves warm-start behaviour at init.
+                    nn.init.xavier_uniform_(self.prefix_mlp[0].weight)
+                    nn.init.zeros_(self.prefix_mlp[0].bias)
+                    nn.init.zeros_(self.prefix_mlp[2].weight)
+                    self.prefix_mlp[2].bias.copy_(
+                        bos_embed.expand(n_pref, -1).reshape(-1)
+                    )
+                self.register_parameter("prefix_embeds", None)
+            else:
+                init = bos_embed.expand(n_pref, -1).clone()
+                self.prefix_embeds = nn.Parameter(init)
+                self.prefix_mlp = None
+        else:
+            self.register_parameter("prefix_embeds", None)
+            self.prefix_mlp = None
+
+    # ─────────────────────────── helpers ────────────────────────────
+
+    def _inject_prefix(self, v: torch.Tensor) -> torch.Tensor:
+        """Build ``(B, P, d_embed)`` injection prefix from vectors ``(B, d_vec)``.
+
+        Adapter params are float32; the output is cast to the backbone's
+        dtype at this boundary so ``inputs_embeds`` matches the frozen
+        backbone weights without forcing the optimizer onto low precision.
+        """
+        if v.dim() == 1:
+            v = v.unsqueeze(0)
+        v32 = v.float()
+        # Slot 0 always uses self.proj; slots 1..M-1 use proj_extra.
+        inject_list = [self.proj(v32).unsqueeze(1)]
+        if self.proj_extra is not None:
+            for lin in self.proj_extra:
+                inject_list.append(lin(v32).unsqueeze(1))
+        inject = torch.cat(inject_list, dim=1)  # (B, M, d) float32
+        if self._n_pref == 0:
+            return inject.to(self._backbone_dtype)
+        if self.prefix_mlp is not None:
+            extra = self.prefix_mlp(v32).view(v.size(0), self._n_pref, self._d_embed)
+        else:
+            extra = self.prefix_embeds.unsqueeze(0).expand(v.size(0), -1, -1)
+        return torch.cat([inject, extra], dim=1).to(self._backbone_dtype)
+
+    @property
+    def prefix_length(self) -> int:
+        return self._n_inject + self._n_pref
+
+    # ───────────────────────── generation ───────────────────────────
+
+    @torch.no_grad()
+    def generate(
+        self,
+        v: torch.Tensor,
+        max_new_tokens: int | None = None,
+        do_sample: bool = True,
+        temperature: float | None = None,
+        top_p: float | None = None,
+    ) -> torch.Tensor:
+        """Return generated token ids ``(B, T_new)`` (without the prefix)."""
+        inputs_embeds = self._inject_prefix(v)
+        attn = torch.ones(inputs_embeds.shape[:2], dtype=torch.long, device=inputs_embeds.device)
+        out = self.backbone.generate(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attn,
+            max_new_tokens=max_new_tokens or self.cfg.max_new_tokens,
+            do_sample=do_sample,
+            temperature=temperature if temperature is not None else self.cfg.temperature,
+            top_p=top_p if top_p is not None else self.cfg.top_p,
+            pad_token_id=self.tokenizer.pad_token_id or self.tokenizer.eos_token_id,
+        )
+        # When inputs_embeds is used, HF returns only the new tokens.
+        return out
+
+    @torch.no_grad()
+    def verbalize(
+        self,
+        v: torch.Tensor,
+        **gen_kwargs,
+    ) -> list[str]:
+        """Convenience: vectors → text strings."""
+        ids = self.generate(v, **gen_kwargs)
+        return self.tokenizer.batch_decode(ids, skip_special_tokens=True)

srt/training/__init__.py

@@ -0,0 +1 @@
+"""SRT training utilities."""

srt/training/losses.py

@@ -0,0 +1,560 @@
+"""Semiotic loss functions for SRT Adapter.
+
+All losses operate on the adapter's intermediate outputs — divergence vectors,
+meta-state, r̂, regime, community assignments.  The CE loss comes directly from
+the backbone's native LM head and is not computed here.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+from srt.adapter import SRTAdapterOutput
+from srt.config import LossConfig
+
+
+def chain_loss(
+    divergences: list[torch.Tensor],
+    chain_predictor: torch.nn.Module,
+    attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Self-supervised chain-of-interpretants loss.
+
+    Each MAH layer's divergence should be predictable from the previous layer's
+    divergence.  This is Peirce's chain of interpretants: each interpretation
+    leads to the next in a predictable way when meaning is stable.
+
+    Args:
+        divergences: list of (B, T, d_div) tensors from successive MAH layers.
+        chain_predictor: nn.Linear that predicts next divergence from current.
+        attention_mask: (B, T) padding mask (1 = real, 0 = pad). Optional.
+
+    Returns:
+        Scalar loss.
+    """
+    if len(divergences) < 2:
+        return torch.tensor(0.0, device=divergences[0].device)
+
+    loss = torch.tensor(0.0, device=divergences[0].device)
+    for i in range(len(divergences) - 1):
+        pred = chain_predictor(divergences[i])
+        target = divergences[i + 1].detach()
+        per_pos = (pred - target).pow(2).mean(dim=-1)  # (B, T)
+        if attention_mask is not None:
+            mask = attention_mask.to(per_pos.dtype)
+            loss = loss + (per_pos * mask).sum() / mask.sum().clamp(min=1)
+        else:
+            loss = loss + per_pos.mean()
+    return loss / (len(divergences) - 1)
+
+
+def bifurcation_loss(
+    r_hat: torch.Tensor,
+    r_true: torch.Tensor,
+    r_mask: torch.Tensor,
+) -> torch.Tensor:
+    """Smooth L1 loss between predicted and log-compressed true reflexivity.
+
+    r_true spans [0, ~13] but BEN outputs Tanh ∈ [-1, 1].  Log-compressing
+    the target via  sign(r) * log(1 + |r|)  maps the bulk of r_true into
+    roughly [-1.5, 1.5], much better aligned with the output range.
+
+    Uses mild focal weighting (1.0×) for supercritical tokens.
+
+    Args:
+        r_hat: (B, T) predicted reflexivity from BEN.
+        r_true: (B, T) ground-truth reflexivity.
+        r_mask: (B, T) bool mask — True where r_true is valid.
+
+    Returns:
+        Scalar loss.
+    """
+    if r_mask.sum() == 0:
+        return torch.tensor(0.0, device=r_hat.device)
+
+    r_hat_masked = r_hat[r_mask]
+    r_true_masked = r_true[r_mask]
+
+    # Log-compress targets to match BEN's Tanh output range
+    r_true_compressed = r_true_masked.sign() * (1.0 + r_true_masked.abs()).log()
+
+    # Mild focal weighting on compressed scale
+    focal_weight = 1.0 + 1.0 * r_true_compressed.abs()
+    diff = F.smooth_l1_loss(r_hat_masked, r_true_compressed, reduction="none")
+    return (diff * focal_weight).mean()
+
+
+def regime_loss(
+    regime_logits: torch.Tensor,
+    r_true: torch.Tensor,
+    r_mask: torch.Tensor,
+) -> torch.Tensor:
+    """Cross-entropy loss for regime classification.
+
+    Regime is derived from r_true:
+      - subcritical (class 0): r_true <= 0
+      - supercritical (class 1): r_true > 0
+
+    Args:
+        regime_logits: (B, T, 2) from BEN.
+        r_true: (B, T) ground-truth reflexivity.
+        r_mask: (B, T) bool mask.
+
+    Returns:
+        Scalar loss.
+    """
+    if r_mask.sum() == 0:
+        return torch.tensor(0.0, device=regime_logits.device)
+
+    regime_targets = (r_true > 0).long()  # (B, T)
+    logits_masked = regime_logits[r_mask]  # (N, 2)
+    targets_masked = regime_targets[r_mask]  # (N,)
+    return F.cross_entropy(logits_masked, targets_masked)
+
+
+def divergence_alive_loss(
+    divergences: list[torch.Tensor],
+    attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Prevent divergence vectors from collapsing to zero.
+
+    Encourages divergence norms to stay near a target value (1.0).
+
+    Args:
+        divergences: list of (B, T, d_div) tensors.
+        attention_mask: (B, T) padding mask (1 = real, 0 = pad). Optional.
+
+    Returns:
+        Scalar loss.
+    """
+    if not divergences:
+        return torch.tensor(0.0, device=divergences[0].device if divergences else "cpu")
+
+    total = torch.tensor(0.0, device=divergences[0].device)
+    for d in divergences:
+        norms = d.norm(dim=-1)  # (B, T)
+        if attention_mask is not None:
+            mask = attention_mask.to(norms.dtype)
+            mean_norm = (norms * mask).sum() / mask.sum().clamp(min=1)
+        else:
+            mean_norm = norms.mean()
+        total = total + (1.0 - mean_norm).abs()
+    return total / len(divergences)
+
+
+def injection_regularization(
+    injections: list[torch.Tensor],
+    attention_mask: torch.Tensor | None = None,
+    target_norm: float = 1.0,
+) -> torch.Tensor:
+    """Penalize injection norms that deviate from a target.
+
+    Uses (||inj|| - target)^2 per position, which pulls norms toward the
+    target rather than toward zero.  This lets the RRM contribute useful
+    signal at norm ~1 while strongly penalizing the 6-8 norms seen in v2.
+
+    Args:
+        injections: list of (B, T, d_backbone) injection vectors.
+        attention_mask: (B, T) padding mask (1 = real, 0 = pad). Optional.
+        target_norm: desired L2 norm for injection vectors.
+
+    Returns:
+        Scalar loss.
+    """
+    if not injections:
+        return torch.tensor(0.0, device=injections[0].device if injections else "cpu")
+
+    total = torch.tensor(0.0, device=injections[0].device)
+    for inj in injections:
+        norms = inj.norm(dim=-1)  # (B, T) — L2 norm per position
+        per_pos = (norms - target_norm).pow(2)  # (B, T)
+        if attention_mask is not None:
+            mask = attention_mask.to(per_pos.dtype)
+            total = total + (per_pos * mask).sum() / mask.sum().clamp(min=1)
+        else:
+            total = total + per_pos.mean()
+    return total / len(injections)
+
+
+def community_entropy_loss(community_weights: torch.Tensor) -> torch.Tensor:
+    """Encourage diverse community usage across the batch.
+
+    Maximizes entropy of the average community assignment distribution.
+    Without this, the model might collapse all inputs to one community.
+
+    Args:
+        community_weights: (B, K) soft community assignments.
+
+    Returns:
+        Scalar loss (lower = more diverse).
+    """
+    avg_dist = community_weights.mean(dim=0)  # (K,)
+    entropy = -(avg_dist * (avg_dist + 1e-8).log()).sum()
+    max_entropy = math.log(community_weights.shape[-1])
+    return max_entropy - entropy
+
+
+def community_supcon_loss(
+    community_vectors: torch.Tensor,
+    community_ids: torch.Tensor,
+    temperature: float = 0.1,
+) -> tuple[torch.Tensor, dict[str, float]]:
+    """Supervised contrastive loss on per-sample community vectors.
+
+    The v3 entropy regularizer kept the prototype distribution from collapsing
+    to a single mode but produced congruent collapse instead: pairwise cosine
+    similarity between learned prototypes converged to ~0.99 and the assignment
+    head learned to be near-uniform.  SupCon (Khosla et al. 2020) gives the
+    encoder direct gradient pressure to put samples from the same source into
+    a tight neighborhood and push different sources apart, which forces
+    prototypes to diversify because that is the only way to satisfy the loss.
+
+    v5 note: pass `encoded` (pre-mixing) rather than `vector` here. When the
+    assignment head collapses, `vector ≈ prototype[k*]` is constant across
+    the batch and this loss is identically `log(B-1)` with zero gradient.
+
+    Args:
+        community_vectors: (B, d) per-sample vectors to contrast.
+        community_ids: (B,) integer ids — same id = positive pair.
+        temperature: softmax temperature.
+
+    Returns:
+        (loss, diagnostics) where diagnostics carries the number of positive
+        pairs and unique classes seen this batch (for sanity checks).
+    """
+    B = community_vectors.shape[0]
+    device = community_vectors.device
+    n_unique = int(community_ids.unique().numel())
+    if B < 2:
+        return torch.tensor(0.0, device=device), {
+            "pos_pairs": 0.0, "unique_classes": float(n_unique),
+        }
+
+    # Cast to float32 for numerical stability of the contrastive softmax;
+    # the input vector may be bf16 because the adapter modules run in bf16.
+    z = F.normalize(community_vectors.float(), dim=-1)
+    sim = (z @ z.T) / temperature  # (B, B)
+
+    # Mask: True where i != j and ids match
+    eye = torch.eye(B, dtype=torch.bool, device=z.device)
+    pos_mask = (community_ids.view(-1, 1) == community_ids.view(1, -1)) & ~eye
+
+    # If no positive pairs in this batch, loss is zero (skip rather than NaN)
+    pos_count = pos_mask.sum(dim=1)
+    n_pos_pairs = float(pos_mask.sum().item())
+    if pos_count.sum() == 0:
+        return torch.tensor(0.0, device=device), {
+            "pos_pairs": 0.0, "unique_classes": float(n_unique),
+        }
+
+    # Mask self-similarity from denominator
+    sim = sim.masked_fill(eye, float("-inf"))
+    log_prob = sim - torch.logsumexp(sim, dim=1, keepdim=True)  # (B, B)
+    # Diagonal entries are -inf; zero them so they don't contaminate the
+    # masked sum below (0 * -inf = NaN otherwise).
+    log_prob = log_prob.masked_fill(eye, 0.0)
+
+    # Per-anchor average log-prob over its positives; rows with no positives
+    # contribute 0 (avoid divide-by-zero with clamp).
+    per_anchor = -(log_prob * pos_mask.float()).sum(dim=1) / pos_count.clamp(min=1).float()
+    valid = pos_count > 0
+    if not valid.any():
+        return torch.tensor(0.0, device=device), {
+            "pos_pairs": 0.0, "unique_classes": float(n_unique),
+        }
+    return per_anchor[valid].mean(), {
+        "pos_pairs": n_pos_pairs, "unique_classes": float(n_unique),
+    }
+
+
+def archetype_supcon_loss(
+    community_vectors: torch.Tensor,
+    archetype_ids: torch.Tensor,
+    temperature: float = 0.1,
+) -> tuple[torch.Tensor, dict[str, float]]:
+    """Supervised contrastive loss keyed by archetype_id.
+
+    v9 addition. Uses the same supcon kernel as community_supcon_loss but
+    masks out anchors with archetype_id == -1 (Reddit-corpus rows that
+    carry no archetype label). Designed to be applied to the same
+    `community_output.encoded` representation that community_supcon
+    operates on, so both signals shape the same encoder geometry.
+
+    Args:
+        community_vectors: (B, d) per-sample encoder output.
+        archetype_ids: (B,) ints in [1, 33] for archetype rows, -1 for
+            Reddit rows. Anchors with id == -1 are dropped.
+        temperature: softmax temperature.
+
+    Returns:
+        (loss, diagnostics).
+    """
+    device = community_vectors.device
+    valid = archetype_ids >= 0
+    n_valid = int(valid.sum().item())
+    if n_valid < 2:
+        return torch.tensor(0.0, device=device), {
+            "pos_pairs": 0.0, "unique_classes": 0.0, "n_valid": float(n_valid),
+        }
+    sub_vec = community_vectors[valid]
+    sub_ids = archetype_ids[valid]
+    loss, diag = community_supcon_loss(
+        sub_vec, sub_ids, temperature=temperature,
+    )
+    diag["n_valid"] = float(n_valid)
+    return loss, diag
+
+
+def divergence_supcon_loss(
+    divergences: list[torch.Tensor],
+    community_ids: torch.Tensor,
+    attention_mask: torch.Tensor | None = None,
+    temperature: float = 0.1,
+) -> tuple[torch.Tensor, dict[str, float]]:
+    """Supervised contrastive loss on per-sample mean divergence vectors.
+
+    v6 extension of the v5 community-SupCon idea applied to MAH divergence:
+    the last MAH layer's divergence is mean-pooled (masked) to a per-sample
+    vector, then contrasted by community id. Same lesson as v5 — operate on
+    a representation that is bijective with the encoder input rather than on
+    a quantity that can collapse to a constant across the batch.
+
+    Args:
+        divergences: list of (B, T, d_div) tensors from successive MAH layers.
+            The last entry is used as the contrastive representation.
+        community_ids: (B,) integer ids — same id = positive pair.
+        attention_mask: (B, T) padding mask (1 = real, 0 = pad). Optional.
+        temperature: softmax temperature.
+
+    Returns:
+        (loss, diagnostics).
+    """
+    if not divergences:
+        device = community_ids.device if community_ids is not None else "cpu"
+        return torch.tensor(0.0, device=device), {
+            "pos_pairs": 0.0, "unique_classes": 0.0,
+        }
+    div = divergences[-1]  # (B, T, d_div)
+    if attention_mask is not None:
+        m = attention_mask.to(div.dtype).unsqueeze(-1)  # (B, T, 1)
+        pooled = (div * m).sum(dim=1) / m.sum(dim=1).clamp(min=1)
+    else:
+        pooled = div.mean(dim=1)
+    # Reuse the same SupCon kernel as community_supcon_loss.
+    return community_supcon_loss(pooled, community_ids, temperature=temperature)
+
+
+def listnet_loss(
+    r_hat: torch.Tensor,
+    r_true: torch.Tensor,
+    r_mask: torch.Tensor,
+    temperature: float = 1.0,
+) -> torch.Tensor:
+    """ListNet ranking loss for r̂ over each sequence.
+
+    Cross-entropy between softmax(r_true) and softmax(r_hat) treated as
+    rankings over the valid positions of each sequence. Complements the
+    pointwise smooth-L1 bifurcation loss by giving direct gradient on the
+    *ordering* of r̂ within a passage, which is what downstream uses
+    (top-k attention probes, heatmap visualization, percentile thresholds).
+
+    Args:
+        r_hat: (B, T) predicted reflexivity.
+        r_true: (B, T) ground-truth reflexivity (same scale).
+        r_mask: (B, T) bool mask.
+        temperature: softmax temperature.
+
+    Returns:
+        Scalar loss averaged over sequences with >=2 valid positions.
+    """
+    B, T = r_hat.shape
+    device = r_hat.device
+    losses: list[torch.Tensor] = []
+    # Compress true reflexivity to match r_hat scale (same transform as bif loss).
+    r_true_c = r_true.sign() * (1.0 + r_true.abs()).log()
+    for b in range(B):
+        m = r_mask[b]
+        n = int(m.sum())
+        if n < 2:
+            continue
+        rh = r_hat[b][m].float() / temperature
+        rt = r_true_c[b][m].float() / temperature
+        # Softmax over the valid positions; mask out -inf done implicitly
+        # because we only index the valid slice.
+        log_p_hat = rh - torch.logsumexp(rh, dim=0)
+        p_true = torch.softmax(rt, dim=0)
+        losses.append(-(p_true * log_p_hat).sum())
+    if not losses:
+        return torch.tensor(0.0, device=device)
+    return torch.stack(losses).mean()
+
+
+def chain_residual_aux_loss(
+    chain_residual: torch.Tensor,
+    attention_mask: torch.Tensor | None = None,
+    target: float = 0.5,
+) -> torch.Tensor:
+    """Auxiliary penalty pulling mean per-token chain residual toward a target.
+
+    The chain prediction loss already minimizes the residual directly, but
+    that signal is averaged across all positions and dimensions. This term
+    keeps the per-token residual at a non-trivial level so it remains a
+    useful inference-time signal (used by the hallucination probe) rather
+    than collapsing to ~0 everywhere as training proceeds. With a small
+    weight this acts as a soft floor, not a primary objective.
+
+    Args:
+        chain_residual: (B, T) per-token mean chain residual.
+        attention_mask: (B, T) padding mask.
+        target: desired mean residual on real tokens.
+
+    Returns:
+        Scalar loss.
+    """
+    if attention_mask is not None:
+        m = attention_mask.to(chain_residual.dtype)
+        mean = (chain_residual * m).sum() / m.sum().clamp(min=1)
+    else:
+        mean = chain_residual.mean()
+    return (mean - target).pow(2)
+
+
+def compute_total_loss(
+    output: SRTAdapterOutput,
+    chain_predictor: torch.nn.Module,
+    r_true: torch.Tensor | None,
+    r_mask: torch.Tensor | None,
+    config: LossConfig,
+    attention_mask: torch.Tensor | None = None,
+    community_ids: torch.Tensor | None = None,
+    archetype_ids: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, dict[str, float]]:
+    """Compute combined loss from adapter output.
+
+    Args:
+        output: SRTAdapterOutput from adapter forward pass.
+        chain_predictor: the chain predictor module from the adapter.
+        r_true: (B, T) ground-truth reflexivity, or None.
+        r_mask: (B, T) bool mask for valid r_true positions, or None.
+        config: LossConfig with weights.
+        attention_mask: (B, T) padding mask (1 = real, 0 = pad). Optional.
+
+    Returns:
+        (total_loss, metrics_dict) where metrics_dict has per-component values.
+    """
+    device = output.logits.device
+    metrics: dict[str, float] = {}
+    total = torch.tensor(0.0, device=device)
+
+    # CE loss (from backbone)
+    if output.ce_loss is not None:
+        total = total + config.ce_weight * output.ce_loss
+        metrics["ce"] = output.ce_loss.item()
+
+    # Chain-of-interpretants loss
+    if output.divergences:
+        l_chain = chain_loss(output.divergences, chain_predictor, attention_mask)
+        total = total + config.chain_weight * l_chain
+        metrics["chain"] = l_chain.item()
+
+    # Bifurcation + regime losses
+    if output.ben_output is not None and r_true is not None and r_mask is not None:
+        l_bif = bifurcation_loss(output.ben_output.r_hat, r_true, r_mask)
+        total = total + config.bif_weight * l_bif
+        metrics["bif"] = l_bif.item()
+
+        l_regime = regime_loss(output.ben_output.regime_logits, r_true, r_mask)
+        total = total + config.regime_weight * l_regime
+        metrics["regime"] = l_regime.item()
+
+        # v6: ListNet ranking loss on r̂ within each sequence.
+        if config.listnet_weight > 0:
+            l_listnet = listnet_loss(
+                output.ben_output.r_hat, r_true, r_mask,
+                temperature=config.listnet_temperature,
+            )
+            total = total + config.listnet_weight * l_listnet
+            metrics["listnet"] = l_listnet.item()
+
+    # Divergence alive
+    if output.divergences:
+        l_alive = divergence_alive_loss(output.divergences, attention_mask)
+        total = total + config.div_alive_weight * l_alive
+        metrics["div_alive"] = l_alive.item()
+
+        # v6: SupCon on mean-pooled last-MAH divergence.
+        if (community_ids is not None
+                and config.divergence_supcon_weight > 0):
+            l_div_sup, div_sup_diag = divergence_supcon_loss(
+                output.divergences, community_ids,
+                attention_mask=attention_mask,
+                temperature=config.divergence_supcon_temperature,
+            )
+            total = total + config.divergence_supcon_weight * l_div_sup
+            metrics["div_supcon"] = l_div_sup.item()
+            metrics["div_supcon_pos_pairs"] = div_sup_diag["pos_pairs"]
+
+    # v6: chain-residual auxiliary floor (keeps inference signal alive).
+    if (output.chain_residual_per_token is not None
+            and config.chain_residual_aux_weight > 0):
+        l_chain_aux = chain_residual_aux_loss(
+            output.chain_residual_per_token,
+            attention_mask=attention_mask,
+            target=config.chain_residual_aux_target,
+        )
+        total = total + config.chain_residual_aux_weight * l_chain_aux
+        metrics["chain_aux"] = l_chain_aux.item()
+
+    # Injection regularization (target-norm penalty)
+    if output.injections:
+        l_inject = injection_regularization(
+            output.injections, attention_mask, target_norm=config.inject_target_norm,
+        )
+        total = total + config.inject_reg_weight * l_inject
+        metrics["inject_reg"] = l_inject.item()
+
+    # Community entropy (only meaningful when there is a soft assignment)
+    if output.community_output is not None:
+        if output.community_output.weights is not None:
+            l_comm = community_entropy_loss(output.community_output.weights)
+            total = total + config.community_entropy_weight * l_comm
+            metrics["comm_entropy"] = l_comm.item()
+
+        # Community SupCon (v5 — applied to pre-mixing encoder output, not
+        # the prototype-weighted vector. The vector is a convex combination
+        # of prototypes; if the assignment head collapses to one prototype
+        # then `vector` becomes constant across the batch and SupCon's
+        # gradient is zero by symmetry — this is exactly what killed v4.
+        # `encoded` is the encoder's bijective image of the pooled hidden
+        # state, so it always varies per-sample, giving SupCon non-zero
+        # gradient even from a degenerate warm-start.)
+        if community_ids is not None and config.community_supcon_weight > 0:
+            l_supcon, supcon_diag = community_supcon_loss(
+                output.community_output.encoded,
+                community_ids,
+                temperature=config.community_supcon_temperature,
+            )
+            total = total + config.community_supcon_weight * l_supcon
+            metrics["comm_supcon"] = l_supcon.item()
+            metrics["comm_supcon_pos_pairs"] = supcon_diag["pos_pairs"]
+            metrics["comm_supcon_unique_classes"] = supcon_diag["unique_classes"]
+
+        # v9: archetype-keyed SupCon on the same encoder representation.
+        # Skipped silently if no archetype rows are in the batch (n_valid<2).
+        if (archetype_ids is not None
+                and config.archetype_supcon_weight > 0):
+            l_arch, arch_diag = archetype_supcon_loss(
+                output.community_output.encoded,
+                archetype_ids,
+                temperature=config.archetype_supcon_temperature,
+            )
+            total = total + config.archetype_supcon_weight * l_arch
+            metrics["arch_supcon"] = l_arch.item()
+            metrics["arch_supcon_pos_pairs"] = arch_diag["pos_pairs"]
+            metrics["arch_supcon_n_valid"] = arch_diag["n_valid"]
+
+    metrics["total"] = total.item()
+    return total, metrics