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