exp008 shipped: multiband mechanism certifies (3/3 surgical band lesions; LOW expert wins its band; uniform objective doesn't pay - roles enter exp009)
8f831c4 verified | """dexp008_multiband.py — exp008: MULTIBAND STEPWISE AMoE, the band-assignment | |
| mechanism test (stage plan: history/plans/2026-07-17_multiband_stage_plan.md). | |
| Phil's band-role map (roles enter as OBJECTIVES in exp009; this bed certifies | |
| the MECHANISM): HIGH (0.75,1] diversity/blobbing | MID (0.35,0.75] | |
| continuity/semantics | LOW [0,0.35] fidelity/detail. | |
| Per site (16): three rank-16 zero-init deltas (W+b, own gates -3.0), combined | |
| by SMOOTH COSINE CROSSFADE WINDOWS on s01 (width 0.06, sum to 1 everywhere — | |
| STRUCTURAL/positional gating, dense + differentiable, no selection event). | |
| Per-band gradients arise from the windows themselves: a sample at s01=0.9 | |
| trains (almost) only the HIGH expert. Falsifier: matched monolith rank-48. | |
| Prereg (stage plan): | |
| P1 multiband >= monolith on overall paired val; | |
| P2 SPECIALIZATION SIGNATURE — each band expert beats the monolith IN ITS | |
| OWN band (per-band paired val); | |
| P3 toggle bit-exact post-train; | |
| P4 BAND LESION — disabling band-b's experts damages band-b val >> other | |
| bands (surgical decoupling on the sigma axis). | |
| Substrate: core epred SD1.5 (2-seed certified). Reuses the dexp006 cache. | |
| Pod: bash pod2/run_exp008.sh [DEXP8_SEED=0 DEXP8_STEPS=3000] | |
| """ | |
| 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 enumerate_sd15_sites, MEM_FRACTION | |
| from dexp006_sd15core_relay import make_schedule, add_noise | |
| SD_BASE = "stable-diffusion-v1-5/stable-diffusion-v1-5" | |
| N_TRAIN, N_VAL = 4096, 256 | |
| RANK, N_BANDS = 16, 3 | |
| BAND_EDGES = (0.35, 0.75) # LOW [0,.35] | MID (.35,.75] | HIGH (.75,1] | |
| XFADE = 0.06 # cosine crossfade half-width | |
| BATCH = int(os.environ.get("DEXP8_BATCH", "16")) | |
| STEPS = int(os.environ.get("DEXP8_STEPS", "3000")) | |
| SEED = int(os.environ.get("DEXP8_SEED", "0")) | |
| LR, CFG_DROPOUT = 1e-3, 0.1 | |
| DATA_DIR = ("/workspace/data/dexp008" if os.path.isdir("/workspace") | |
| else os.path.join(os.environ.get("GEOLIP_DATA", "./data"), | |
| "dexp008")) | |
| DEXP6_DIR = ("/workspace/data/dexp006" if os.path.isdir("/workspace") | |
| else os.path.join(os.environ.get("GEOLIP_DATA", "./data"), | |
| "dexp006")) | |
| CKPT_DIR = ("/workspace/ckpts2/dexp008" if os.path.isdir("/workspace") | |
| else DATA_DIR) | |
| def band_weights(s01: torch.Tensor) -> torch.Tensor: | |
| """Smooth structural band windows: (B, 3), sum to 1 everywhere. | |
| Cosine crossfade of half-width XFADE around each edge — positional | |
| gating on the sigma axis, never comparative.""" | |
| def ramp(x): # 0 below -XFADE, 1 above +XFADE, smooth | |
| t = ((x / XFADE).clamp(-1, 1) + 1) / 2 | |
| return 0.5 - 0.5 * torch.cos(t * math.pi) | |
| e1, e2 = BAND_EDGES | |
| up1, up2 = ramp(s01 - e1), ramp(s01 - e2) | |
| low = 1 - up1 | |
| mid = up1 * (1 - up2) | |
| high = up1 * up2 | |
| return torch.stack([low, mid, high], dim=-1) | |
| def band_of(s01: float) -> int: | |
| e1, e2 = BAND_EDGES | |
| return 0 if s01 <= e1 else (1 if s01 <= e2 else 2) | |
| class MultibandDelta(nn.Module): | |
| """Three band experts per site, window-combined; per-expert enable flags | |
| (band lesions) + global enable (toggle law).""" | |
| def __init__(self, d: int, r: int = RANK): | |
| super().__init__() | |
| self.d, self.r = d, r | |
| self.down = nn.ModuleList(nn.Linear(d, r, bias=False) | |
| for _ in range(N_BANDS)) | |
| self.up = nn.ModuleList(nn.Linear(r, d) for _ in range(N_BANDS)) | |
| 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.gates = nn.Parameter(torch.full((N_BANDS,), -3.0)) | |
| self.enabled = True | |
| self.band_enabled = [True] * N_BANDS | |
| 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 forward(self, x, w_bands): # w_bands: (B, 3) | |
| if not self.enabled: | |
| return x | |
| g = torch.sigmoid(self.gates) | |
| delta = 0 | |
| w = w_bands.view(w_bands.shape[0], | |
| *([1] * (x.ndim - 2)), N_BANDS) | |
| for b in range(N_BANDS): | |
| if not self.band_enabled[b]: | |
| continue | |
| delta = delta + g[b] * w[..., b:b + 1] * self.up[b]( | |
| self.down[b](x)) | |
| return x + delta if not isinstance(delta, int) else x | |
| class MonoDelta(nn.Module): | |
| def __init__(self, d: int, r: int = RANK * N_BANDS): | |
| super().__init__() | |
| self.down = nn.Linear(d, r, bias=False) | |
| nn.init.orthogonal_(self.down.weight) | |
| self.up = nn.Linear(r, d) | |
| nn.init.zeros_(self.up.weight) | |
| nn.init.zeros_(self.up.bias) | |
| self.gate = nn.Parameter(torch.tensor(-3.0)) | |
| self.enabled = True | |
| def forward(self, x, w_bands=None): | |
| if not self.enabled: | |
| return x | |
| return x + torch.sigmoid(self.gate) * self.up(self.down(x)) | |
| class BlockWrap(nn.Module): | |
| def __init__(self, block, mod): | |
| super().__init__() | |
| self.block = block | |
| self.mod = mod | |
| self.w_bands = None | |
| def forward(self, *args, **kwargs): | |
| out = self.block(*args, **kwargs) | |
| h = out[0] if isinstance(out, tuple) else out | |
| h = self.mod(h, self.w_bands) | |
| return (h,) + out[1:] if isinstance(out, tuple) else h | |
| def load_unet(device): | |
| from diffusers import UNet2DConditionModel | |
| unet = UNet2DConditionModel.from_pretrained( | |
| SD_BASE, subfolder="unet", torch_dtype=torch.float32).to(device) | |
| unet.requires_grad_(False) | |
| unet.eval() | |
| unet.enable_gradient_checkpointing() | |
| return unet | |
| def attach(unet, mk): | |
| sites = enumerate_sd15_sites(unet) | |
| assert len(sites) == 16 | |
| mods, wraps = nn.ModuleList(), [] | |
| for name, block, d in sites: | |
| p0 = next(block.parameters()) | |
| m = mk(d).to(device=p0.device, dtype=p0.dtype) | |
| w = BlockWrap(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 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) | |
| acp = make_schedule(device) | |
| cache_f = os.path.join(DEXP6_DIR, "cache.pt") | |
| assert os.path.exists(cache_f), "needs the dexp006 cache" | |
| cache = torch.load(cache_f, map_location="cpu", weights_only=True) | |
| def set_w(wraps, s01): | |
| w = band_weights(s01) | |
| for wr in wraps: | |
| wr.w_bands = w | |
| 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 | |
| t = torch.randint(0, 1000, (bsz,), generator=gen, device=device) | |
| set_w(wraps, t.float() / 1000.0) | |
| noise = torch.randn(lat.shape, generator=gen, device=device) | |
| pred = unet(add_noise(lat, noise, t, acp), t, ehs, | |
| return_dict=False)[0] | |
| return F.mse_loss(pred, noise) | |
| def val(unet, wraps): | |
| tot, per_band = [], {0: [], 1: [], 2: []} | |
| 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) | |
| t = cache["val_t"][i:i + 32].to(device) | |
| set_w(wraps, t.float() / 1000.0) | |
| pred = unet(add_noise(lat, noise, t, acp), t, ehs, | |
| return_dict=False)[0] | |
| mse = ((pred - noise) ** 2).mean(dim=(1, 2, 3)) | |
| tot += mse.tolist() | |
| for j, tv in enumerate((t.float() / 1000.0).tolist()): | |
| per_band[band_of(tv)].append(mse[j].item()) | |
| return (sum(tot) / len(tot), | |
| {f"band{b}": round(sum(v) / max(len(v), 1), 6) | |
| for b, v in per_band.items()}, | |
| {f"band{b}_n": len(v) for b, v in per_band.items()}) | |
| def train(unet, wraps, params, label): | |
| opt = torch.optim.Adam(params, lr=LR, weight_decay=0.0) | |
| gen = torch.Generator(device=device).manual_seed(SEED + 42) | |
| idx = torch.Generator().manual_seed(SEED + 7) | |
| torch.cuda.reset_peak_memory_stats() | |
| t0 = time.time() | |
| for step in range(1, STEPS + 1): | |
| sel = torch.randint(0, N_TRAIN, (BATCH,), generator=idx) | |
| 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"[{label}] step {step}: loss {loss.item():.4f} | " | |
| f"{(time.time() - t0) / step:.2f}s/step | peak " | |
| f"{torch.cuda.max_memory_allocated() / 2**30:.1f}GB", | |
| flush=True) | |
| return round((time.time() - t0) / STEPS, 3) | |
| results = {"config": {"bands": 3, "edges": BAND_EDGES, "xfade": XFADE, | |
| "rank": RANK, "steps": STEPS, "batch": BATCH, | |
| "lr": LR, "seed": SEED, "substrate": "core epred"}} | |
| with ledger_run(f"dexp008 frozen s{SEED}", budget_h=0.2) as h: | |
| unet = load_unet(device) | |
| v, pb, pn = val(unet, []) | |
| results["frozen"] = {"val": v, "per_band": pb, "band_n": pn} | |
| del unet | |
| torch.cuda.empty_cache() | |
| h["verdict"] = f"val {v:.5f}" | |
| with ledger_run(f"dexp008 multiband3 s{SEED}", budget_h=2.5) as h: | |
| unet = load_unet(device) | |
| mods, wraps = attach(unet, lambda d: MultibandDelta(d)) | |
| for m in mods: | |
| m.assert_zero_init() | |
| n_params = sum(p.numel() for p in mods.parameters()) | |
| print(f"[mb3] 16 sites x 3 band experts, {n_params:,} trainable", | |
| flush=True) | |
| spd = train(unet, wraps, mods.parameters(), "mb3") | |
| v_on, pb_on, _ = val(unet, wraps) | |
| for m in mods: # P3 toggle | |
| m.enabled = False | |
| v_off, _, _ = val(unet, wraps) | |
| d = abs(v_off - results["frozen"]["val"]) | |
| assert d < 1e-9, f"toggle parity broken: {d}" | |
| for m in mods: | |
| m.enabled = True | |
| # P4 band lesions: disable band b everywhere, re-val per band | |
| lesions = {} | |
| for b in range(N_BANDS): | |
| for m in mods: | |
| m.band_enabled[b] = False | |
| _, pb_les, _ = val(unet, wraps) | |
| lesions[f"lesion_band{b}"] = pb_les | |
| for m in mods: | |
| m.band_enabled[b] = True | |
| torch.save({"mods": [m.state_dict() for m in mods]}, | |
| os.path.join(CKPT_DIR, f"mb3_s{SEED}.pt")) | |
| results["multiband3"] = { | |
| "n_params": n_params, "val": v_on, "per_band": pb_on, | |
| "val_toggled_off": v_off, "s_per_step": spd, | |
| "lesions": lesions, | |
| "gates": [[round(torch.sigmoid(g).item(), 4) | |
| for g in m.gates] for m in mods[8:9]][0]} | |
| del unet, mods | |
| torch.cuda.empty_cache() | |
| h["verdict"] = f"val {v_on:.5f}" | |
| with ledger_run(f"dexp008 monolith r48 s{SEED}", budget_h=2.5) as h: | |
| unet = load_unet(device) | |
| mods, wraps = attach(unet, lambda d: MonoDelta(d)) | |
| n_params = sum(p.numel() for p in mods.parameters()) | |
| print(f"[mono48] {n_params:,} trainable", flush=True) | |
| spd = train(unet, wraps, mods.parameters(), "mono48") | |
| v, pb, _ = val(unet, wraps) | |
| torch.save({"mods": [m.state_dict() for m in mods]}, | |
| os.path.join(CKPT_DIR, f"mono48_s{SEED}.pt")) | |
| results["monolith"] = {"n_params": n_params, "val": v, | |
| "per_band": pb, "s_per_step": spd} | |
| del unet, mods | |
| torch.cuda.empty_cache() | |
| h["verdict"] = f"val {v:.5f}" | |
| mb, mo = results["multiband3"], results["monolith"] | |
| per_band_wins = {b: mb["per_band"][b] < mo["per_band"][b] | |
| for b in ("band0", "band1", "band2")} | |
| lesion_asym = {} | |
| for b in range(N_BANDS): | |
| les = mb["lesions"][f"lesion_band{b}"] | |
| own = les[f"band{b}"] - mb["per_band"][f"band{b}"] | |
| other = sum(les[f"band{o}"] - mb["per_band"][f"band{o}"] | |
| for o in range(N_BANDS) if o != b) / (N_BANDS - 1) | |
| lesion_asym[f"band{b}"] = {"own_damage": round(own, 6), | |
| "other_damage": round(other, 6), | |
| "surgical": own > 3 * abs(other) and own > 0} | |
| results["verdict"] = { | |
| "P1_multiband_vs_monolith": "multiband" if mb["val"] <= mo["val"] | |
| else "monolith", | |
| "P2_per_band_wins": per_band_wins, | |
| "P4_lesion_asymmetry": lesion_asym, | |
| "param_ratio": round(mb["n_params"] / mo["n_params"], 3), | |
| "note": f"seed {SEED}; mechanism test — roles enter in exp009", | |
| } | |
| 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"dexp008 verdict: {json.dumps(results['verdict'])}") | |
| print(json.dumps(results["verdict"], indent=2)) | |
| burn_down() | |
| return results | |
| def smoke(): | |
| w = band_weights(torch.tensor([0.1, 0.35, 0.55, 0.75, 0.9, 0.72, 0.78])) | |
| assert w.shape == (7, 3) | |
| assert torch.allclose(w.sum(-1), torch.ones(7), atol=1e-6), "windows must sum to 1" | |
| assert w[0].argmax() == 0 and w[2].argmax() == 1 and w[4].argmax() == 2 | |
| m = MultibandDelta(320) | |
| m.assert_zero_init() | |
| x = torch.randn(2, 9, 320) | |
| wb = band_weights(torch.tensor([0.2, 0.9])) | |
| assert torch.equal(m(x, wb), x), "zero-init must be exact" | |
| m.enabled = False | |
| assert m(x, wb) is x | |
| m.enabled = True | |
| m.band_enabled = [False, False, False] | |
| assert m(x, wb) is x, "all-band lesion must be code-path exact" | |
| m.band_enabled = [True] * 3 | |
| nn.init.normal_(m.up[0].weight, std=0.02) | |
| assert not torch.equal(m(x, wb), x) | |
| mb = sum(p.numel() for p in MultibandDelta(320).parameters()) | |
| mo = sum(p.numel() for p in MonoDelta(320).parameters()) | |
| print(f"dexp008 smoke PASSED (windows sum-1 + band argmax, zero-init, " | |
| f"toggle, lesion path; params mb {mb:,} vs mono {mo:,} @d=320)") | |
| if __name__ == "__main__": | |
| if "--run" in sys.argv: | |
| run() | |
| else: | |
| smoke() | |