exp014 shipped: router v1 honest negative (state keys already route by prompt; raw-flat address = disease geometry; null lesson -> v2)

README.md

@@ -82,6 +82,7 @@ resolution) so gaps live in a narrow band — the paired design is load-bearing.
 | exp009_bandroles | role objectives: **directional hit 4/4 but noise-adjacent — frequency reweighting too collinear; needs qualitatively different supervision + generation-side gauges (exp010)** | **shipped (2 seeds + rejudge)** |
 | exp010_controller | **StepGatedSampler ships**: controller lifts grounding +0.089 over frozen; monotonic lesion ladder; HIGH lesion 14x LP-dominant (coarse-to-fine confirmed in image space); eps-trained HIGH band concentrates (diversity = open training goal) | **shipped (candidate)** |
 | exp011a_fused_multiband | multiband on REAL fused data: **adapters pay 2-3x more; structural story replicates (3/3 surgical, monolith edge persists); blob targets built 100% after a schema lesson** | **shipped (candidate, s0; s1 running)** |
+| exp014_te_dispatch | router-solidifier v1: **honest negative** — state keys already route by prompt; raw-flattened address kills routing (the known high-D disease geometry); null-design lesson → v2 | **shipped (candidate, s0; s1 running)** |
 | exp013_blob_flow | **CONDITIONING HYPOTHESIS CONFIRMED**: same blob coupling, ~200x the eps effect on the flow substrate (−5.9% vs +0.03%); blob supervision belongs on flow/v-pred trunks | **shipped (candidate, s0; s1 running)** |
 
 Numbering: a number binds to a preregistered design at assignment and is never reassigned — gaps mean reserved-not-yet-run (exp004 = the Anima relay train, integration-certified by R0b; exp005 = the SDXL Tree-4a battery, designed). Sub-letters (000b) extend a shipped experiment without disturbing its rows.

exp014_te_dispatch/README.md

@@ -0,0 +1,36 @@
+# exp014_te_dispatch — the router-solidifier v1 (CANDIDATE s0; s1 running)
+
+**Question.** Does text-side addressing in the DISPATCH KEY (a trained
+text-encoder adapter, or the frozen canonical byte-trigram address) route
+the expert bank better than state+sigma keys alone?
+
+**Arms** (flow trunk, fused cache, A=4 r=8 banks, dense signed dispatch):
+sigma_key (control) | addr_key (frozen [32,128] address flattened → key) |
+te_key (trained CLIP-pooler adapter → key).
+
+**Results** (`results.json`):
+
+| arm | val flow-MSE | per-prompt usage var | shuffled-key null |
+|---|---|---|---|
+| sigma_key | **0.57371** | 2.36e-02 | 0 (no text) |
+| addr_key | 0.57755 | **1.40e-05 (dead)** | 1.40e-05 |
+| te_key | 0.57422 | 2.60e-02 | 2.71e-02 |
+
+**Verdict (honest negative + two instrument findings).**
+1. P1 MISS: no val win for text keys at this scale (all within 0.07%).
+2. **The hidden-state key already routes by prompt** (2.36e-02 usage
+   variance with zero text input) — prompt-level structure reaches the
+   usage surface through x itself; exp007's flatness was across SIGMA
+   BANDS, a different axis. The TE key adds ~nothing on top.
+3. **The raw-flattened address kills routing** (1.4e-05) and dents val —
+   reproducing the text line's documented disease geometry (addressing raw
+   high-D flattens the oriented softmax; healthy forms bottleneck to low-D
+   FIRST). The canonical address needs a structured reduction (slots/M̂
+   read) before it can serve as a router key — v2's design, not a defect
+   of the address itself.
+4. **Confessed instrument gap:** the shuffled-key null measures key
+   DIVERSITY, not routing CORRECTNESS (te null ≈ te real). v2's null:
+   repeated-key, plus matched-vs-mismatched key val deltas.
+
+**Caveats.** s0 (s1 running); one scale (A=4, 3k steps); pooler-only TE
+feature; flow trunk only. Cost ≈ 4.6 GPU-h incl. text-feature build.

exp014_te_dispatch/dexp014_te_dispatch.py

@@ -0,0 +1,383 @@
+"""dexp014_te_dispatch.py — exp014: TEXT-ADDRESSED DISPATCH (Phil's "router
+solidifier": the MoE accepts a trained text-encoder adapter as an addressing
+conditioning router).
+
+The junction of three findings: diffusion-side codes compress caption
+structure to ~0 (exp003) | the text-address channel is live where not
+redundant (exp002) | state+sigma dispatch alone finds nothing to route on
+(exp007). Here the TEXT side enters the DISPATCH KEY, not the content path.
+
+Per site: AmoeBank-style expert bank (A=4, rank-8 zero-init deltas, dense
+signed aleph dispatch, no selectors). Dispatch-key arms:
+  sigma_key      — key = key_proj(x) + sig_proj(fourier(sigma))  [exp007 ctrl]
+  addr_key       — + txt_proj(frozen byte-trigram aleph address of the
+                   caption, flattened [32*128]) — the CANONICAL address as
+                   router (frozen address, trained projection).
+  te_key         — + te_adapter(CLIP text pooler_output [768] -> D) — THE
+                   TRAINED TE ADAPTER AS ROUTER SOLIDIFIER (trained jointly
+                   through the diffusion loss).
+Substrate: base_lune flow + the fused cache (exp013's regime — largest
+adapter headroom; blob gauge carried as the role-aligned instrument).
+
+Judged (all arms, in-bed): common flow-MSE per band; HIGH-band blob gauge;
+ROUTER SIGNATURE = per-prompt usage variance vs a SHUFFLED-KEY null (val
+pass rerun with deranged text keys; text-specific routing must beat the
+null); toggle bit-exact.
+Prereg: P1 te_key >= sigma_key on common val; P2 router signature: usage-
+by-prompt variance(te_key) > 3x shuffled null; P3 toggle; P4 addr_key vs
+te_key ordering = which text representation routes better (no prediction —
+first measurement). 1-seed CANDIDATE.
+
+Pod: bash pod2/run_exp014.sh   [DEXP14_SEED=0]
+"""
+from __future__ import annotations
+
+import json
+import math
+import os
+import sys
+import time
+
+sys.path[:0] = ["pod2", "."]
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from pod_ledger import ledger_run, note, burn_down
+from d1_substrate import MEM_FRACTION, enumerate_sd15_sites
+from aleph_diffusion_core import derangement
+from dexp009_bandroles import lp
+from dexp013_blob_flow import load_lune, blob_lp_err
+
+SHIFT = 2.5
+N_TRAIN, N_VAL = 4096, 256
+A_EXPERTS, RANK = 4, 8
+BATCH = int(os.environ.get("DEXP14_BATCH", "16"))
+STEPS = int(os.environ.get("DEXP14_STEPS", "3000"))
+SEED = int(os.environ.get("DEXP14_SEED", "0"))
+LR, CFG_DROPOUT = 1e-3, 0.1
+D11 = ("/workspace/data/dexp011" if os.path.isdir("/workspace")
+       else "./data/dexp011")
+DATA_DIR = ("/workspace/data/dexp014" if os.path.isdir("/workspace")
+            else os.path.join(os.environ.get("GEOLIP_DATA", "./data"),
+                              "dexp014"))
+CKPT_DIR = ("/workspace/ckpts2/dexp014" if os.path.isdir("/workspace")
+            else DATA_DIR)
+
+
+def fourier_sigma(s, dim=8):
+    freqs = torch.pow(2.0, torch.arange(dim // 2, device=s.device))
+    ang = 2 * math.pi * s[:, None] * freqs[None]
+    return torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1)
+
+
+class TeDispatchBank(nn.Module):
+    """Expert bank with a text-conditioned dense signed dispatch key."""
+
+    def __init__(self, d: int, txt_dim: int = 0, A: int = A_EXPERTS,
+                 r: int = RANK, D: int = 4, tau: float = 0.1):
+        super().__init__()
+        self.d, self.A, self.txt_dim = d, A, txt_dim
+        self.tau = tau
+        self.down = nn.ModuleList(nn.Linear(d, r, bias=False)
+                                  for _ in range(A))
+        self.up = nn.ModuleList(nn.Linear(r, d) for _ in range(A))
+        for dn, up in zip(self.down, self.up):
+            nn.init.orthogonal_(dn.weight)
+            nn.init.zeros_(up.weight)
+            nn.init.zeros_(up.bias)
+        self.key_proj = nn.Linear(d, D, bias=False)
+        self.sig_proj = nn.Linear(8, D, bias=False)
+        nn.init.orthogonal_(self.key_proj.weight)
+        nn.init.orthogonal_(self.sig_proj.weight)
+        if txt_dim > 0:
+            self.txt_proj = nn.Linear(txt_dim, D, bias=False)
+            nn.init.orthogonal_(self.txt_proj.weight)
+        self.codebook = nn.Parameter(F.normalize(torch.randn(A, D), dim=-1))
+        self.gates = nn.Parameter(torch.full((A,), -3.0))
+        self.enabled = True
+        self.last_usage = None
+
+    def assert_zero_init(self):
+        for up in self.up:
+            assert up.weight.abs().max().item() == 0.0
+            assert up.bias.abs().max().item() == 0.0
+
+    def dispatch(self, x, sig_feat, txt_feat):
+        key = self.key_proj(x)
+        sig = self.sig_proj(sig_feat)
+        sig = sig.view(sig.shape[0], *([1] * (x.ndim - 2)), sig.shape[-1])
+        key = key + sig
+        if self.txt_dim > 0 and txt_feat is not None:
+            tx = self.txt_proj(txt_feat)
+            tx = tx.view(tx.shape[0], *([1] * (x.ndim - 2)), tx.shape[-1])
+            key = key + tx
+        Ax = F.normalize(self.codebook, dim=-1)
+        u = (F.normalize(key, dim=-1) @ Ax.transpose(-1, -2)) / self.tau
+        m = u.abs().amax(dim=-1, keepdim=True)
+        ep, en = torch.exp(u - m), torch.exp(-u - m)
+        return (ep - en) / (ep + en).sum(dim=-1, keepdim=True)
+
+    def forward(self, x, sig_feat, txt_feat=None):
+        if not self.enabled:
+            return x
+        w = self.dispatch(x, sig_feat, txt_feat)
+        with torch.no_grad():
+            self.last_usage = w.abs().mean(dim=tuple(
+                range(1, w.ndim - 1))).detach().cpu()   # (B, A) per prompt
+        g = torch.sigmoid(self.gates)
+        delta = 0
+        for k in range(self.A):
+            delta = delta + g[k] * w[..., k:k + 1] * self.up[k](
+                self.down[k](x))
+        return x + delta
+
+
+class Wrap(nn.Module):
+    def __init__(self, block, bank):
+        super().__init__()
+        self.block = block
+        self.bank = bank
+        self.sig_feat = None
+        self.txt_feat = None
+
+    def forward(self, *args, **kwargs):
+        out = self.block(*args, **kwargs)
+        h = out[0] if isinstance(out, tuple) else out
+        h = self.bank(h, self.sig_feat, self.txt_feat)
+        return (h,) + out[1:] if isinstance(out, tuple) else h
+
+
+def attach_banks(unet, txt_dim):
+    sites = enumerate_sd15_sites(unet)
+    assert len(sites) == 16
+    mods, wraps = nn.ModuleList(), []
+    for name, block, d in sites:
+        p0 = next(block.parameters())
+        m = TeDispatchBank(d, txt_dim).to(device=p0.device, dtype=p0.dtype)
+        w = Wrap(block, m)
+        parent = unet
+        parts = name.split(".")
+        for p in parts[:-1]:
+            parent = getattr(parent, p) if not p.isdigit() else parent[int(p)]
+        if parts[-1].isdigit():
+            parent[int(parts[-1])] = w
+        else:
+            setattr(parent, parts[-1], w)
+        mods.append(m)
+        wraps.append(w)
+    return mods, wraps
+
+
+def build_text_feats(device):
+    """addr: frozen byte-trigram aleph address of each fused caption
+    (canonical recipe, flattened 4096); pooler: CLIP text pooler_output."""
+    f = os.path.join(DATA_DIR, "textfeats.pt")
+    if os.path.exists(f):
+        return torch.load(f, map_location="cpu", weights_only=True)
+    os.makedirs(DATA_DIR, exist_ok=True)
+    from dexp011_fused_multiband import _iter_fused
+    from transformers import CLIPTextModel, CLIPTokenizer
+    caps = []
+    for r in _iter_fused(N_TRAIN + N_VAL):
+        caps.append(str(r.get("caption_joycaption")
+                        or r.get("prompt_fused") or ""))
+    from dexp002_sd15_addrcond import Addr
+    enc = Addr(device)
+    addr = enc.encode(caps).reshape(len(caps), -1)      # (N, 4096)
+    del enc
+    torch.cuda.empty_cache()
+    tok = CLIPTokenizer.from_pretrained(
+        "stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="tokenizer")
+    te = CLIPTextModel.from_pretrained(
+        "stable-diffusion-v1-5/stable-diffusion-v1-5",
+        subfolder="text_encoder", torch_dtype=torch.float32).to(device).eval()
+    pools = []
+    with torch.no_grad():
+        for i in range(0, len(caps), 64):
+            ids = tok(caps[i:i + 64], padding="max_length", max_length=77,
+                      truncation=True, return_tensors="pt").input_ids
+            pools.append(te(ids.to(device)).pooler_output.cpu())
+    del te
+    torch.cuda.empty_cache()
+    out = {"addr": addr, "pooler": torch.cat(pools)}
+    torch.save(out, f)
+    print(f"[textfeats] built: addr {out['addr'].shape}, "
+          f"pooler {out['pooler'].shape}", flush=True)
+    return out
+
+
+def run(device="cuda"):
+    torch.cuda.set_per_process_memory_fraction(MEM_FRACTION, 0)
+    os.makedirs(CKPT_DIR, exist_ok=True)
+    cache = torch.load(os.path.join(D11, "cache.pt"), map_location="cpu",
+                       weights_only=True)
+    with ledger_run("dexp014 text feats", budget_h=0.5):
+        tf = build_text_feats(device)
+    gv = torch.Generator().manual_seed(SEED + 11)
+    u = torch.linspace(0.02, 0.98, N_VAL)
+    val_sigma = (SHIFT * u) / (1 + (SHIFT - 1) * u)
+    ARMS = {"sigma_key": (0, None),
+            "addr_key": (4096, "addr"),
+            "te_key": (768, "pooler")}
+
+    def txt_of(kind, sel):
+        if kind is None:
+            return None
+        return tf[kind][sel].float().to("cuda")
+
+    def set_feats(wraps, s01, txt):
+        sf = fourier_sigma(s01)
+        for wr in wraps:
+            wr.sig_feat = sf
+            wr.txt_feat = txt
+
+    def loss_of(unet, wraps, lat, ehs, gen):
+        bsz = lat.shape[0]
+        drop = torch.rand(bsz, generator=gen, device=device) < CFG_DROPOUT
+        ehs = ehs.clone()
+        ehs[drop] = 0
+        s = torch.rand(bsz, generator=gen, device=device)
+        s = (SHIFT * s) / (1 + (SHIFT - 1) * s)
+        noise = torch.randn(lat.shape, generator=gen, device=device)
+        s4 = s[:, None, None, None]
+        x_t, v = noise * s4 + lat * (1 - s4), noise - lat
+        for wr in wraps:
+            wr.sig_feat = fourier_sigma(s)
+        pred = unet(x_t, s * 1000, ehs, return_dict=False)[0]
+        return F.mse_loss(pred, v)
+
+    @torch.no_grad()
+    def val(unet, wraps, kind, shuffle_txt=False):
+        tot, blob_high, usages = [], [], []
+        perm = derangement(N_VAL, seed=SEED + 3)
+        for i in range(0, N_VAL, 32):
+            lat = cache["val_lat"][i:i + 32].to(device)
+            ehs = cache["val_ehs"][i:i + 32].to(device)
+            noise = cache["val_noise"][i:i + 32].to(device)
+            blob = cache["val_blob"][i:i + 32].float().to(device)
+            s = val_sigma[i:i + 32].to(device)
+            sel = torch.arange(N_TRAIN + i, N_TRAIN + i + lat.shape[0])
+            if shuffle_txt:
+                sel = N_TRAIN + perm[i:i + lat.shape[0]]
+            set_feats(wraps, s, txt_of(kind, sel))
+            s4 = s[:, None, None, None]
+            x_t, v = noise * s4 + lat * (1 - s4), noise - lat
+            pred = unet(x_t, s * 1000, ehs, return_dict=False)[0]
+            tot += ((pred - v) ** 2).mean(dim=(1, 2, 3)).tolist()
+            x0_hat = x_t - s4 * pred
+            bg = blob_lp_err(x0_hat, lat, blob)
+            for j, sv in enumerate(s.tolist()):
+                if sv > 0.75:
+                    blob_high.append(bg[j].item())
+            for wr in wraps[8:9]:
+                if wr.bank.last_usage is not None:
+                    usages.append(wr.bank.last_usage)
+        usage = torch.cat(usages) if usages else None
+        usage_var = float(usage.var(dim=0).mean()) if usage is not None else 0
+        return (sum(tot) / len(tot),
+                round(sum(blob_high) / max(len(blob_high), 1), 6),
+                round(usage_var, 8))
+
+    results = {"config": {"steps": STEPS, "batch": BATCH, "seed": SEED,
+                          "trunk": "base_lune flow", "A": A_EXPERTS,
+                          "r": RANK}}
+    with ledger_run(f"dexp014 frozen s{SEED}", budget_h=0.2) as h:
+        unet = load_lune(device)
+        v, bg, _ = val(unet, [], None)
+        results["frozen"] = {"val": v, "blob_high": bg}
+        del unet
+        torch.cuda.empty_cache()
+        h["verdict"] = f"val {v:.5f}"
+
+    for arm, (txt_dim, kind) in ARMS.items():
+        with ledger_run(f"dexp014 {arm} s{SEED}", budget_h=2.2) as h:
+            unet = load_lune(device)
+            mods, wraps = attach_banks(unet, txt_dim)
+            for m in mods:
+                m.assert_zero_init()
+            # bind txt feats into the TRAIN loss via a wrapper closure
+            opt = torch.optim.Adam(mods.parameters(), lr=LR, weight_decay=0.0)
+            gen = torch.Generator(device=device).manual_seed(SEED + 42)
+            idx = torch.Generator().manual_seed(SEED + 7)
+            t0 = time.time()
+            for step in range(1, STEPS + 1):
+                sel = torch.randint(0, N_TRAIN, (BATCH,), generator=idx)
+                for wr in wraps:
+                    wr.txt_feat = txt_of(kind, sel)
+                loss = loss_of(unet, wraps, cache["lat"][sel].to(device),
+                               cache["ehs"][sel].to(device), gen)
+                loss.backward()
+                opt.step()
+                opt.zero_grad(set_to_none=True)
+                if step == 50 or step % 500 == 0:
+                    print(f"[{arm}] step {step}: loss {loss.item():.4f} | "
+                          f"{(time.time() - t0) / step:.2f}s/step",
+                          flush=True)
+            v, bg, uv = val(unet, wraps, kind)
+            _, _, uv_null = val(unet, wraps, kind, shuffle_txt=True) \
+                if kind else (0, 0, 0)
+            for m in mods:
+                m.enabled = False
+            v_off, _, _ = val(unet, wraps, kind)
+            d = abs(v_off - results["frozen"]["val"])
+            assert d < 1e-9, f"toggle parity broken: {d}"
+            for m in mods:
+                m.enabled = True
+            torch.save({"mods": [m.state_dict() for m in mods]},
+                       os.path.join(CKPT_DIR, f"{arm}_s{SEED}.pt"))
+            results[arm] = {"val": v, "blob_high": bg,
+                            "usage_var_by_prompt": uv,
+                            "usage_var_shuffled_null": uv_null,
+                            "val_toggled_off": v_off}
+            del unet, mods
+            torch.cuda.empty_cache()
+            h["verdict"] = f"val {v:.5f} uvar {uv:.2e}"
+
+    sk, ak, tk = (results[a] for a in ("sigma_key", "addr_key", "te_key"))
+    results["verdict"] = {
+        "P1_te_vs_sigma": "te" if tk["val"] <= sk["val"] else "sigma",
+        "P2_router_signature": {
+            "te_var": tk["usage_var_by_prompt"],
+            "te_null": tk["usage_var_shuffled_null"],
+            "hit_3x": tk["usage_var_by_prompt"]
+            > 3 * max(tk["usage_var_shuffled_null"], 1e-12)},
+        "P4_text_rep_ordering": {
+            "addr_key_val": ak["val"], "te_key_val": tk["val"],
+            "addr_key_var": ak["usage_var_by_prompt"],
+            "te_key_var": tk["usage_var_by_prompt"]},
+        "common": {a: results[a]["val"] for a in ARMS},
+        "note": "router-solidifier v1; 1-seed CANDIDATE",
+    }
+    with open(os.path.join(DATA_DIR, "results.json" if SEED == 0
+                           else f"results_s{SEED}.json"), "w") as f:
+        json.dump(results, f, indent=2)
+    note(f"dexp014: {json.dumps(results['verdict']['P2_router_signature'])}")
+    print(json.dumps(results["verdict"], indent=2))
+    burn_down()
+    return results
+
+
+def smoke():
+    b = TeDispatchBank(320, txt_dim=768)
+    b.assert_zero_init()
+    x = torch.randn(2, 9, 320)
+    sf = fourier_sigma(torch.tensor([0.3, 0.8]))
+    tx = torch.randn(2, 768)
+    assert torch.equal(b(x, sf, tx), x), "zero-init must be exact"
+    w = b.dispatch(x, sf, tx)
+    w2 = b.dispatch(x, sf, torch.randn(2, 768))
+    assert w.shape == (2, 9, 4) and not torch.equal(w, w2), \
+        "text key must move the dispatch"
+    b0 = TeDispatchBank(320, txt_dim=0)
+    assert torch.equal(b0(x, sf, None), x)
+    print("dexp014 smoke PASSED (zero-init, text-sensitive dispatch, "
+          "no-text arm)")
+
+
+if __name__ == "__main__":
+    if "--run" in sys.argv:
+        run()
+    else:
+        smoke()

exp014_te_dispatch/results.json

@@ -0,0 +1,55 @@
+{
+  "config": {
+    "steps": 3000,
+    "batch": 16,
+    "seed": 0,
+    "trunk": "base_lune flow",
+    "A": 4,
+    "r": 8
+  },
+  "frozen": {
+    "val": 0.6298933834768832,
+    "blob_high": 0.184471
+  },
+  "sigma_key": {
+    "val": 0.5737131804926321,
+    "blob_high": 0.127068,
+    "usage_var_by_prompt": 0.02356817,
+    "usage_var_shuffled_null": 0,
+    "val_toggled_off": 0.6298933834768832
+  },
+  "addr_key": {
+    "val": 0.5775512405671179,
+    "blob_high": 0.127947,
+    "usage_var_by_prompt": 1.403e-05,
+    "usage_var_shuffled_null": 1.4e-05,
+    "val_toggled_off": 0.6298933834768832
+  },
+  "te_key": {
+    "val": 0.5742165236733854,
+    "blob_high": 0.12704,
+    "usage_var_by_prompt": 0.02597812,
+    "usage_var_shuffled_null": 0.02707456,
+    "val_toggled_off": 0.6298933834768832
+  },
+  "verdict": {
+    "P1_te_vs_sigma": "sigma",
+    "P2_router_signature": {
+      "te_var": 0.02597812,
+      "te_null": 0.02707456,
+      "hit_3x": false
+    },
+    "P4_text_rep_ordering": {
+      "addr_key_val": 0.5775512405671179,
+      "te_key_val": 0.5742165236733854,
+      "addr_key_var": 1.403e-05,
+      "te_key_var": 0.02597812
+    },
+    "common": {
+      "sigma_key": 0.5737131804926321,
+      "addr_key": 0.5775512405671179,
+      "te_key": 0.5742165236733854
+    },
+    "note": "router-solidifier v1; 1-seed CANDIDATE"
+  }
+}