exp015 shipped: address-as-key falsified with corrected instruments (zero routing excess, zero match advantage) - the dispatch-key form closes; TE+AMOE joint training is the open form

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)** |
+| exp015_router_v2 | corrected router test: **address-as-key FALSIFIED** (routing excess 2.5e-06 over constant-key null; match advantage −0.0); open form = TE+AMOE joint training | **shipped (candidate, s0)** |
 | 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)** |
 

exp015_router_v2/README.md

@@ -0,0 +1,32 @@
+# exp015_router_v2 — the corrected router test: address-as-key FALSIFIED
+# (CANDIDATE, s0 — the window's final training run)
+
+**Fixes applied from exp014's instrument findings:** the address enters as
+32 slots → shared 128→4 projection → per-slot dense signed aleph
+coefficients → MEAN composition (low-D before addressing — the text line's
+cure for the high-D flattening disease); null = REPEATED KEY (constant
+address for all prompts — isolates text-specific routing from state-driven
+variance); plus a matched-vs-mismatched key val delta (the correctness
+gauge v1 lacked).
+
+**Results** (`results.json`, flow trunk, fused cache):
+
+| arm | val | usage-var (real) | usage-var (repeated null) | match advantage |
+|---|---|---|---|---|
+| state_sigma (control) | **0.57310** | 1.36e-02 | — | — |
+| mhat_addr (v2) | 0.57491 | 1.13e-02 | 1.13e-02 | **−0.0** |
+
+**Verdict.** Routing excess over the constant-key null: **2.5e-06 — zero.**
+Matched and mismatched address keys produce identical val. With correct
+geometry AND correct instruments, the frozen text address contributes no
+prompt-specific routing on this substrate at this scale; all per-prompt
+usage structure is state-driven (consistent with exp014's finding that the
+hidden-state key already routes by prompt). Two-experiment closure of the
+dispatch-key form of the router-solidifier idea; the open form is the full
+TE+AMOE joint-training collective (docket), where the text encoder learns
+THROUGH the diffusion loss rather than entering as a frozen key.
+
+**Caveats.** s0; one substrate (lune flow), one scale (A=4); the fused
+captions are stylistically homogeneous (fashion photography) — a
+low-text-diversity regime; a multi-domain caption mixture is the one
+remaining escape hatch for the frozen-key form.

exp015_router_v2/dexp015_router_v2.py

@@ -0,0 +1,287 @@
+"""dexp015_router_v2.py — exp015: ROUTER v2 — the M-hat-bottlenecked address
+key + the corrected null (exp014's two instrument findings applied).
+
+v1 failures fixed: (1) the raw-flattened [32*128] address killed routing
+(high-D disease geometry) -> v2 reads the address AS 32 SLOTS, projects each
+slot 128->4 (orthogonal, shared) BEFORE addressing, takes per-slot dense
+signed aleph coefficients against the bank codebook, and MEAN-COMPOSES over
+slots (the aleph's own order-insensitive composition rule) — low-D before
+addressing, per the text line's cure. (2) the shuffled-key null measured
+diversity, not correctness -> v2 null = REPEATED KEY (every prompt gets
+row-0's address; any residual per-prompt usage variance is state-driven).
+
+Arms: state_sigma (control = v1's winner) | mhat_addr (v2). Judged: common
+flow-MSE; HIGH-band blob gauge; per-prompt usage variance with the REPEATED
+null; matched-vs-mismatched key val delta (the correctness gauge v1 lacked).
+Prereg: P1 mhat_addr routing variance > 3x its REPEATED null EXCESS over
+state-driven variance; P2 common val >= control - 0.5%; P3 toggle.
+
+Pod: bash pod2/run_exp015.sh
+"""
+from __future__ import annotations
+
+import json
+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
+from aleph_diffusion_core import derangement
+from dexp013_blob_flow import load_lune, blob_lp_err
+from dexp009_bandroles import lp
+from dexp014_te_dispatch import (TeDispatchBank, Wrap, attach_banks,
+                                 build_text_feats, fourier_sigma,
+                                 N_TRAIN, N_VAL, SHIFT)
+
+BATCH = int(os.environ.get("DEXP15_BATCH", "16"))
+STEPS = int(os.environ.get("DEXP15_STEPS", "3000"))
+SEED = int(os.environ.get("DEXP15_SEED", "0"))
+LR, CFG_DROPOUT = 1e-3, 0.1
+D11 = ("/workspace/data/dexp011" if os.path.isdir("/workspace")
+       else "./data/dexp011")
+D14 = ("/workspace/data/dexp014" if os.path.isdir("/workspace")
+       else "./data/dexp014")
+DATA_DIR = ("/workspace/data/dexp015" if os.path.isdir("/workspace")
+            else "./data/dexp015")
+CKPT_DIR = ("/workspace/ckpts2/dexp015" if os.path.isdir("/workspace")
+            else DATA_DIR)
+
+
+class MhatAddrBank(TeDispatchBank):
+    """v2: address read as 32 slots -> shared 128->4 projection -> per-slot
+    signed aleph vs the bank codebook -> MEAN composition over slots."""
+
+    def __init__(self, d: int, **kw):
+        super().__init__(d, txt_dim=0, **kw)
+        self.slot_proj = nn.Linear(128, 4, bias=False)
+        nn.init.orthogonal_(self.slot_proj.weight)
+        self.addr_mix = nn.Parameter(torch.tensor(0.0))   # sigmoid=0.5 mix
+
+    def addr_weights(self, addr_flat):
+        slots = addr_flat.view(addr_flat.shape[0], 32, 128)
+        keys = self.slot_proj(slots)                       # (B, 32, 4)
+        A = F.normalize(self.codebook, dim=-1)
+        u = (F.normalize(keys, dim=-1) @ A.transpose(-1, -2)) / self.tau
+        m = u.abs().amax(dim=-1, keepdim=True)
+        ep, en = torch.exp(u - m), torch.exp(-u - m)
+        w = (ep - en) / (ep + en).sum(dim=-1, keepdim=True)  # (B, 32, A)
+        return w.mean(dim=1)                               # mean composition
+
+    def forward(self, x, sig_feat, txt_feat=None):
+        if not self.enabled:
+            return x
+        w_state = self.dispatch(x, sig_feat, None)          # (..., A)
+        if txt_feat is not None:
+            wa = self.addr_weights(txt_feat)                # (B, A)
+            wa = wa.view(wa.shape[0], *([1] * (x.ndim - 2)), self.A)
+            mix = torch.sigmoid(self.addr_mix)
+            w = (1 - mix) * w_state + mix * wa
+        else:
+            w = w_state
+        with torch.no_grad():
+            self.last_usage = w.abs().mean(dim=tuple(
+                range(1, w.ndim - 1))).detach().cpu()
+        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
+
+
+def run(device="cuda"):
+    torch.cuda.set_per_process_memory_fraction(MEM_FRACTION, 0)
+    os.makedirs(CKPT_DIR, exist_ok=True)
+    os.makedirs(DATA_DIR, exist_ok=True)
+    cache = torch.load(os.path.join(D11, "cache.pt"), map_location="cpu",
+                       weights_only=True)
+    tf = torch.load(os.path.join(D14, "textfeats.pt"), map_location="cpu",
+                    weights_only=True)
+    u = torch.linspace(0.02, 0.98, N_VAL)
+    val_sigma = (SHIFT * u) / (1 + (SHIFT - 1) * u)
+    ARMS = {"state_sigma": (TeDispatchBank, None),
+            "mhat_addr": (MhatAddrBank, "addr")}
+
+    def attach_v2(unet, cls):
+        from d1_substrate import enumerate_sd15_sites
+        sites = enumerate_sd15_sites(unet)
+        mods, wraps = nn.ModuleList(), []
+        for name, block, d in sites:
+            p0 = next(block.parameters())
+            m = (cls(d, txt_dim=0) if cls is TeDispatchBank
+                 else cls(d)).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 txt_of(kind, sel, repeat=False, mismatch=None):
+        if kind is None:
+            return None
+        t = tf[kind]
+        if repeat:
+            return t[0:1].expand(len(sel), -1).float().to(device)
+        if mismatch is not None:
+            return t[mismatch].float().to(device)
+        return t[sel].float().to(device)
+
+    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, mode="real"):
+        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])
+            txt = txt_of(kind, sel, repeat=(mode == "repeat"),
+                         mismatch=(N_TRAIN + perm[i:i + lat.shape[0]])
+                         if mode == "mismatch" else None)
+            sf = fourier_sigma(s)
+            for wr in wraps:
+                wr.sig_feat = sf
+                wr.txt_feat = txt
+            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()
+            bg = blob_lp_err(x_t - s4 * pred, 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)
+        uvar = float(torch.cat(usages).var(dim=0).mean()) if usages else 0
+        return (sum(tot) / len(tot),
+                round(sum(blob_high) / max(len(blob_high), 1), 6),
+                round(uvar, 8))
+
+    results = {"config": {"steps": STEPS, "batch": BATCH, "seed": SEED}}
+    with ledger_run(f"dexp015 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, (cls, kind) in ARMS.items():
+        with ledger_run(f"dexp015 {arm} s{SEED}", budget_h=2.2) as h:
+            unet = load_lune(device)
+            mods, wraps = attach_v2(unet, cls)
+            for m in mods:
+                m.assert_zero_init()
+            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, uvar = val(unet, wraps, kind)
+            entry = {"val": v, "blob_high": bg, "uvar_real": uvar}
+            if kind:
+                _, _, uvar_rep = val(unet, wraps, kind, mode="repeat")
+                v_mis, _, _ = val(unet, wraps, kind, mode="mismatch")
+                entry.update({"uvar_repeated_null": uvar_rep,
+                              "val_mismatched_key": v_mis,
+                              "match_advantage": round(v_mis - v, 6)})
+            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] = entry
+            del unet, mods
+            torch.cuda.empty_cache()
+            h["verdict"] = f"val {v:.5f}"
+
+    ma, ss = results["mhat_addr"], results["state_sigma"]
+    excess = ma["uvar_real"] - ma["uvar_repeated_null"]
+    results["verdict"] = {
+        "P1_routing_excess_over_repeated_null": round(excess, 8),
+        "P1_hit": excess > 2 * ma["uvar_repeated_null"],
+        "P2_val_vs_control": {"mhat": ma["val"], "state_sigma": ss["val"],
+                              "ok": ma["val"] <= ss["val"] * 1.005},
+        "match_advantage_val": ma.get("match_advantage"),
+        "note": "router v2 (M-hat slots, mean composition, repeated null); "
+                "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"dexp015: {json.dumps(results['verdict'])}")
+    print(json.dumps(results["verdict"], indent=2))
+    burn_down()
+    return results
+
+
+def smoke():
+    b = MhatAddrBank(320)
+    b.assert_zero_init()
+    x = torch.randn(2, 9, 320)
+    sf = fourier_sigma(torch.tensor([0.3, 0.8]))
+    addr = torch.randn(2, 32 * 128)
+    assert torch.equal(b(x, sf, addr), x), "zero-init must be exact"
+    w = b.addr_weights(addr)
+    w2 = b.addr_weights(torch.randn(2, 32 * 128))
+    assert w.shape == (2, 4) and not torch.equal(w, w2)
+    assert torch.equal(b(x, sf, None), x)
+    print("dexp015 smoke PASSED (M-hat slot key, mean composition, "
+          "zero-init/toggle)")
+
+
+if __name__ == "__main__":
+    if "--run" in sys.argv:
+        run()
+    else:
+        smoke()

exp015_router_v2/results.json

@@ -0,0 +1,35 @@
+{
+  "config": {
+    "steps": 3000,
+    "batch": 16,
+    "seed": 0
+  },
+  "frozen": {
+    "val": 0.6298933834768832,
+    "blob_high": 0.184471
+  },
+  "state_sigma": {
+    "val": 0.5730983872199431,
+    "blob_high": 0.126756,
+    "uvar_real": 0.01355188
+  },
+  "mhat_addr": {
+    "val": 0.5749124474241398,
+    "blob_high": 0.128959,
+    "uvar_real": 0.01133719,
+    "uvar_repeated_null": 0.01133466,
+    "val_mismatched_key": 0.5749121565022506,
+    "match_advantage": -0.0
+  },
+  "verdict": {
+    "P1_routing_excess_over_repeated_null": 2.53e-06,
+    "P1_hit": false,
+    "P2_val_vs_control": {
+      "mhat": 0.5749124474241398,
+      "state_sigma": 0.5730983872199431,
+      "ok": true
+    },
+    "match_advantage_val": -0.0,
+    "note": "router v2 (M-hat slots, mean composition, repeated null); 1-seed CANDIDATE"
+  }
+}