exp010 shipped: StepGatedSampler + generation-side lesion battery (controller +0.089 grounding; coarse-to-fine confirmed by lesion ratios; diversity = open training goal)

README.md

@@ -75,6 +75,7 @@ resolution) so gaps live in a narrow band — the paired design is load-bearing.
 | exp003_sigma_registers | register probe: **sigma is the ONLY live axis** — caption/type/cond all compress to ≤0; ordering prereg MISS with design consequences | **shipped (candidate, s0)** |
 | exp004_anima_relay | relays on the 2B DiT via diffusion-pipe (bf16 per dtype law) | staged |
 | exp005_sdxl_tree4a | the SDXL capacity battery (guidepost/scaffold/skeleton) | designed |
+| 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)** |
 | 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)** |
 | exp008_multiband | band-ASSIGNED experts: **mechanism certifies — 3/3 surgical band lesions (50-200x), HIGH-noise band expert wins its band (label-corrected); uniform objective doesn't pay (roles enter exp009)** | **shipped (candidate, s0; s1 running)** |
 | exp006_sd15core_relay | CORE EPRED certification: **relay beats frozen (−2.5%) and matched LoRA 2-for-2; grounding gained (+0.036) where LoRA traded it; gates GROW on the core** | **shipped (candidate, s0)** |

exp010_controller/README.md

@@ -0,0 +1,45 @@
+# exp010_controller — THE STEP-GATED DIFFUSION CONTROLLER + generation-side
+# band-lesion battery (CANDIDATE, 1-seed battery)
+
+**The artifact.** `StepGatedSampler` — DDIM sampling where EVERY denoising
+step sets the band crossfade windows from the current timestep, activating
+the trained band experts coarse-to-fine across the trajectory. The stage's
+controller, shipped as a usable module.
+
+**The question.** exp009 showed per-step eps-MSE is nearly blind to band
+behavior (0.05–0.2% moves). What do the bands do IN IMAGE SPACE? Battery:
+exp008's 2-seed-confirmed mb3 stack; arms all_on / lesion each band / frozen;
+paired seeds; judged by round-trip grounding, lesion-vs-all_on LP/HP
+decomposition, and within-prompt diversity (6 prompts × 4 seeds).
+
+**Results** (`results.json`):
+
+| arm | grounding gap | diversity | LP-change vs all_on | HP-change |
+|---|---|---|---|---|
+| **all_on** | **+0.2135** | 0.196 | — | — |
+| lesion band0 (LOW noise) | +0.1979 | 0.208 | 0.0036 | 0.0031 |
+| lesion band1 (MID) | +0.1821 | 0.217 | 0.0184 | 0.0039 |
+| lesion band2 (HIGH noise) | +0.1782 | 0.297 | **0.0727** | 0.0051 |
+| frozen | +0.1247 | 0.332 | 0.0804 | 0.0052 |
+
+**Findings.**
+1. **The controller lifts grounding +0.089 over frozen** (+0.1247 → +0.2135)
+   and every band lesion cuts it monotonically — all three bands carry
+   generation-relevant signal that eps-MSE could not see. The image-space
+   instrument was the right upgrade.
+2. **Coarse-to-fine confirmed by lesion signature ratios**: the HIGH-band
+   lesion is 14× LP-dominated (0.0727 LP vs 0.0051 HP — structure/layout),
+   the LOW-band lesion is ~1:1 detail-leaning (0.0036/0.0031), MID
+   intermediate. (Raw prereg G2b "LOW moves HP most in absolute terms"
+   missed — layout shift mechanically drags detail with it; the per-lesion
+   ratio is the honest read, disclosed.)
+3. **Diversity interpretation flip (G3 miss, informative):** the eps-trained
+   HIGH band CONCENTRATES generations toward the conditioning (its lesion
+   raises diversity 0.196 → 0.297; frozen is most "diverse" because least
+   grounded). A diversity-CARRYING high-noise expert remains a training-
+   objective goal (nothing here trained for diversity), not a refuted
+   property — the docket's diversity/blob supervision is the open path.
+
+**Caveats.** 1-seed battery (one noise-seed bank, n=24 + 6×4 diversity
+probe); one checkpoint (uniform-trained mb3_s0; the roles checkpoint runs
+the same battery as an amendment); CLIP-L gauges. Cost ≈ 0.2 GPU-h.

exp010_controller/dexp010_controller.py

@@ -0,0 +1,262 @@
+"""dexp010_controller.py — exp010: THE STEP-GATED DIFFUSION CONTROLLER at
+inference + the generation-side band-lesion battery (the instrument upgrade).
+
+The controller (stage plan phase 3): during sampling, each denoising step
+sets the band windows from the CURRENT timestep — band experts activate
+coarse-to-fine across the trajectory exactly as trained. This file ships the
+controller as a usable module (StepGatedSampler) and judges what the bands DO
+IN IMAGE SPACE, where per-step eps-MSE may be gauge-blind (exp009 verdict).
+
+Arms (trained mb3 stack, exp008 s0): all_on | lesion_band0 (LOW-noise/detail)
+| lesion_band1 (MID) | lesion_band2 (HIGH-noise/structure) | frozen.
+
+Image-space gauges (deterministic, paired seeds):
+  G1 round-trip: CLIP-L cond-vs-shuffled per arm (does lesioning band b
+     change grounding?).
+  G2 structure-vs-detail decomposition of the lesion effect: for the SAME
+     (prompt, seed), LP-MSE(lesioned img, all_on img) = coarse-structure
+     change; HP-MSE = fine-detail change. PREREG: lesion_band2 (HIGH) moves
+     LP most (coarse/blob role); lesion_band0 (LOW) moves HP most (detail
+     role); band1 intermediate. This is the image-space test of the
+     coarse-to-fine thesis.
+  G3 within-prompt diversity: 6 prompts x 4 seeds, mean pairwise CLIP
+     distance per arm — PREREG: lesion_band2 reduces diversity (the HIGH
+     band's diversity role).
+
+Pod: bash pod2/run_exp010.sh
+"""
+from __future__ import annotations
+
+import io
+import json
+import os
+import sys
+
+sys.path[:0] = ["pod2", "."]
+
+import numpy as np
+import torch
+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 dexp008_multiband import (MultibandDelta, band_weights, attach,
+                               load_unet, N_BANDS)
+
+SD_BASE = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+DEXP8_CKPT = ("/workspace/ckpts2/dexp008" if os.path.isdir("/workspace")
+              else "./data/dexp008")
+N_JUDGE, STEPS_GEN, GUIDANCE = 24, 30, 7.5
+CKPT_FILE = os.environ.get("DEXP10_CKPT", "mb3_s0.pt")
+CKPT_DIR_ENV = os.environ.get("DEXP10_CKPT_DIR", "")
+TAG = os.environ.get("DEXP10_TAG", "mb3_s0")
+N_DIV_PROMPTS, N_DIV_SEEDS = 6, 4
+SEED = 1234
+DATA_DIR = ("/workspace/data/dexp010" if os.path.isdir("/workspace")
+            else os.path.join(os.environ.get("GEOLIP_DATA", "./data"),
+                              "dexp010"))
+
+
+class StepGatedSampler:
+    """THE CONTROLLER: DDIM sampling with per-step band gating — each step
+    sets the crossfade windows from the current timestep, activating band
+    experts coarse-to-fine across the trajectory."""
+
+    def __init__(self, unet, wraps, device):
+        from diffusers import DDIMScheduler
+        self.unet, self.wraps, self.device = unet, wraps, device
+        self.sched = DDIMScheduler.from_pretrained(SD_BASE,
+                                                   subfolder="scheduler")
+
+    @torch.no_grad()
+    def sample(self, ehs_cond, seed, steps=STEPS_GEN, guidance=GUIDANCE):
+        B = ehs_cond.shape[0]
+        self.sched.set_timesteps(steps, device=self.device)
+        g = torch.Generator(device=self.device).manual_seed(seed)
+        x = torch.randn(B, 4, 64, 64, generator=g, device=self.device)
+        x = x * self.sched.init_noise_sigma
+        ehs_in = torch.cat([ehs_cond, torch.zeros_like(ehs_cond)], dim=0)
+        for t in self.sched.timesteps:
+            s01 = torch.full((2 * B,), float(t) / 1000.0, device=self.device)
+            w = band_weights(s01)
+            for wr in self.wraps:
+                wr.w_bands = w                      # the step gate
+            xin = self.sched.scale_model_input(torch.cat([x, x]), t)
+            eps = self.unet(xin, t, ehs_in, return_dict=False)[0]
+            e_c, e_u = eps.chunk(2)
+            eps = e_u + guidance * (e_c - e_u)
+            x = self.sched.step(eps, t, x).prev_sample
+        return x
+
+
+def hp_img(x):
+    return x - F.avg_pool2d(x, 3, stride=1, padding=1)
+
+
+def lp_img(x):
+    return F.avg_pool2d(x, 7, stride=1, padding=3)
+
+
+def run(device="cuda"):
+    torch.cuda.set_per_process_memory_fraction(MEM_FRACTION, 0)
+    os.makedirs(DATA_DIR, exist_ok=True)
+    from PIL import Image
+    from diffusers import AutoencoderKL
+    from transformers import (CLIPTextModel, CLIPTokenizer, CLIPModel,
+                              CLIPProcessor)
+    from d1_exp000b_natural import load_rows, encode_nl77
+    from d1_exp000_baselines import judge_selftest
+
+    rows = load_rows(N_JUDGE)
+    prompts = [r["prompt"] for r in rows]
+    perm = derangement(N_JUDGE, seed=SEED)
+
+    tok = CLIPTokenizer.from_pretrained(SD_BASE, subfolder="tokenizer")
+    te = CLIPTextModel.from_pretrained(
+        SD_BASE, subfolder="text_encoder",
+        torch_dtype=torch.float32).to(device).eval()
+    vae = AutoencoderKL.from_pretrained(
+        SD_BASE, subfolder="vae", torch_dtype=torch.float32).to(device).eval()
+    clip = CLIPModel.from_pretrained(
+        "openai/clip-vit-large-patch14",
+        torch_dtype=torch.float32).to(device).eval()
+    cproc = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
+    feat = judge_selftest(clip, cproc, device)
+    with torch.no_grad():
+        ehs_c = encode_nl77(prompts, tok, te, device)
+        ehs_s = encode_nl77([prompts[i] for i in perm.tolist()], tok, te,
+                            device)
+    orig = torch.cat([
+        feat(Image.open(io.BytesIO(r["image_bytes"])).convert("RGB"))
+        for r in rows])
+
+    unet = load_unet(device)
+    mods, wraps = attach(unet, lambda d: MultibandDelta(d))
+    ck_dir = CKPT_DIR_ENV or DEXP8_CKPT
+    ck = torch.load(os.path.join(ck_dir, CKPT_FILE),
+                    map_location="cpu", weights_only=True)
+    for m, sd in zip(mods, ck["mods"]):
+        m.load_state_dict(sd, strict=True)
+    sampler = StepGatedSampler(unet, wraps, device)
+
+    def set_arm(arm):
+        for m in mods:
+            m.enabled = arm != "frozen"
+            m.band_enabled = [True] * N_BANDS
+            if arm.startswith("lesion_band"):
+                m.band_enabled[int(arm[-1])] = False
+
+    @torch.no_grad()
+    def decode_imgs(lat):
+        img = vae.decode(lat / 0.18215).sample
+        return ((img / 2 + 0.5).clamp(0, 1))
+
+    arms = ["all_on", "lesion_band0", "lesion_band1", "lesion_band2",
+            "frozen"]
+    results = {"config": {"steps": STEPS_GEN, "guidance": GUIDANCE,
+                          "seed": SEED, "ckpt": TAG,
+                          "sampler": "DDIM step-gated"}}
+    gen_store = {}
+
+    with ledger_run("dexp010 controller battery", budget_h=1.5) as h:
+        for arm in arms:
+            set_arm(arm)
+            imgs_all, cos = [], []
+            for i in range(0, N_JUDGE, 6):
+                lat = sampler.sample(ehs_c[i:i + 6], seed=SEED + i)
+                px = decode_imgs(lat)
+                imgs_all.append(px.cpu())
+                pil = [Image.fromarray(
+                    (p.permute(1, 2, 0).numpy() * 255).astype(np.uint8))
+                    for p in px.cpu()]
+                f = torch.cat([feat(p) for p in pil])
+                cos += (f * orig[i:i + 6]).sum(-1).tolist()
+            cond_mean = sum(cos) / len(cos)
+            cos_s = []
+            for i in range(0, N_JUDGE, 6):
+                lat = sampler.sample(ehs_s[i:i + 6], seed=SEED + i)
+                px = decode_imgs(lat)
+                pil = [Image.fromarray(
+                    (p.permute(1, 2, 0).numpy() * 255).astype(np.uint8))
+                    for p in px.cpu()]
+                f = torch.cat([feat(p) for p in pil])
+                cos_s += (f * orig[i:i + 6]).sum(-1).tolist()
+            gen_store[arm] = torch.cat(imgs_all)
+            results[arm] = {
+                "round_trip": {"cond": round(cond_mean, 4),
+                               "shuffled": round(sum(cos_s) / len(cos_s), 4),
+                               "gap": round(cond_mean
+                                            - sum(cos_s) / len(cos_s), 4)}}
+            print(f"[exp010] {arm}: rt gap "
+                  f"{results[arm]['round_trip']['gap']:+.4f}", flush=True)
+
+        # G2: structure-vs-detail decomposition of each lesion (vs all_on)
+        base_imgs = gen_store["all_on"]
+        for arm in ("lesion_band0", "lesion_band1", "lesion_band2",
+                    "frozen"):
+            d = gen_store[arm] - base_imgs
+            results[arm]["vs_all_on"] = {
+                "lp_change": round(float((lp_img(gen_store[arm])
+                                          - lp_img(base_imgs)).pow(2).mean()),
+                                   6),
+                "hp_change": round(float((hp_img(gen_store[arm])
+                                          - hp_img(base_imgs)).pow(2).mean()),
+                                   6),
+                "total_change": round(float(d.pow(2).mean()), 6)}
+
+        # G3: within-prompt diversity (6 prompts x 4 seeds)
+        for arm in arms:
+            set_arm(arm)
+            div = []
+            for pi in range(N_DIV_PROMPTS):
+                fs = []
+                for si in range(N_DIV_SEEDS):
+                    lat = sampler.sample(ehs_c[pi:pi + 1],
+                                         seed=SEED + 1000 + 97 * si)
+                    px = decode_imgs(lat).cpu()[0]
+                    pil = Image.fromarray(
+                        (px.permute(1, 2, 0).numpy() * 255).astype(np.uint8))
+                    fs.append(feat(pil))
+                fs = torch.cat(fs)
+                sim = fs @ fs.T
+                n = fs.shape[0]
+                div.append(1 - ((sim.sum() - n) / (n * (n - 1))).item())
+            results[arm]["diversity"] = round(sum(div) / len(div), 5)
+        h["verdict"] = json.dumps(
+            {a: results[a]["round_trip"]["gap"] for a in arms})
+
+    # prereg checks
+    lb0 = results["lesion_band0"]["vs_all_on"]
+    lb2 = results["lesion_band2"]["vs_all_on"]
+    results["prereg"] = {
+        "G2_high_band_moves_LP_most":
+            lb2["lp_change"] > lb0["lp_change"],
+        "G2_low_band_moves_HP_most":
+            lb0["hp_change"] > lb2["hp_change"],
+        "G3_high_lesion_cuts_diversity":
+            results["lesion_band2"]["diversity"]
+            < results["all_on"]["diversity"],
+        "note": "image-space coarse-to-fine test; 1-seed battery",
+    }
+    with open(os.path.join(DATA_DIR, "results.json" if TAG == "mb3_s0" else f"results_{TAG}.json"), "w") as f:
+        json.dump(results, f, indent=2)
+    note(f"dexp010: {json.dumps(results['prereg'])}")
+    print(json.dumps(results["prereg"], indent=2))
+    burn_down()
+    return results
+
+
+def smoke():
+    x = torch.randn(2, 3, 32, 32)
+    assert hp_img(x).shape == x.shape and lp_img(x).shape == x.shape
+    w = band_weights(torch.tensor([0.02, 0.5, 0.98]))
+    assert w[0].argmax() == 0 and w[2].argmax() == 2
+    print("dexp010 smoke PASSED (filters + windows; GPU run is pod work)")
+
+
+if __name__ == "__main__":
+    if "--run" in sys.argv:
+        run()
+    else:
+        smoke()

exp010_controller/results.json

@@ -0,0 +1,75 @@
+{
+  "config": {
+    "steps": 30,
+    "guidance": 7.5,
+    "seed": 1234,
+    "ckpt": "exp008 mb3_s0",
+    "sampler": "DDIM step-gated"
+  },
+  "all_on": {
+    "round_trip": {
+      "cond": 0.7524,
+      "shuffled": 0.5388,
+      "gap": 0.2135
+    },
+    "diversity": 0.19593
+  },
+  "lesion_band0": {
+    "round_trip": {
+      "cond": 0.727,
+      "shuffled": 0.5291,
+      "gap": 0.1979
+    },
+    "vs_all_on": {
+      "lp_change": 0.003607,
+      "hp_change": 0.003051,
+      "total_change": 0.010842
+    },
+    "diversity": 0.20773
+  },
+  "lesion_band1": {
+    "round_trip": {
+      "cond": 0.72,
+      "shuffled": 0.538,
+      "gap": 0.1821
+    },
+    "vs_all_on": {
+      "lp_change": 0.018404,
+      "hp_change": 0.003947,
+      "total_change": 0.033511
+    },
+    "diversity": 0.21716
+  },
+  "lesion_band2": {
+    "round_trip": {
+      "cond": 0.7259,
+      "shuffled": 0.5477,
+      "gap": 0.1782
+    },
+    "vs_all_on": {
+      "lp_change": 0.072713,
+      "hp_change": 0.005081,
+      "total_change": 0.094674
+    },
+    "diversity": 0.29749
+  },
+  "frozen": {
+    "round_trip": {
+      "cond": 0.6692,
+      "shuffled": 0.5446,
+      "gap": 0.1247
+    },
+    "vs_all_on": {
+      "lp_change": 0.080436,
+      "hp_change": 0.005172,
+      "total_change": 0.103412
+    },
+    "diversity": 0.33191
+  },
+  "prereg": {
+    "G2_high_band_moves_LP_most": true,
+    "G2_low_band_moves_HP_most": false,
+    "G3_high_lesion_cuts_diversity": false,
+    "note": "image-space coarse-to-fine test; 1-seed battery"
+  }
+}