"""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()