"""dist_bed.py — the CLIP-distillation training bed that produced every released checkpoint. Student: 8.66M-parameter ViT (d=240, depth 12, heads 4, patch 16, 160px, CLS readout, linear head to the target dimension). Optimizer: Adam lr 3e-4, weight_decay 0, fp32, TF32 off, batch 256, fixed CRC-derived seeds. The one exception is the cbert_full replication arm, which reproduces its source recipe exactly (AdamW wd 0.01 + 1000-step linear warmup + cosine to 1e-6 + grad-clip 1.0) — flagged in its startup banner. ARMS feature_mse per-element MSE vs the deployment teacher (mimicry) infonce symmetric InfoNCE vs the deployment teacher, temp 0.07 siglip_pairwise decoupled pairwise sigmoid vs the SigLIP teacher (768-d) affinity_kl similarity-row KL (affinity mimicking) blueprint infonce + 0.3*procrustes_sq + 1e-3 spread force consensus_gpa per-sample MSE vs the 5-teacher GPA consensus mean consensus_nce_mse infonce + per-sample MSE vs the GPA mean (champion) x3_full/x3_bce/x3_autograd/nce_mse_cv/cbert_full battery variants zs_floor / teacher_ceiling controls EVALS (identical for every arm; deterministic — re-runs reproduce to 4dp) CIFAR-10/100 zero-shot via the matched text tower; COCO val retrieval R@1/R@5; cosine agreement to the arm target; embedding-geometry readouts (spread CV on a fixed 16-d projection; effective rank). JSONL ledgers. DATA LAYOUT (set DIST_DATA_ROOT; defaults to ./data) $DIST_DATA_ROOT/dist/bank/{config}/{split}-*.parquet COCO teacher bank $DIST_DATA_ROOT/dist/coco/{val2017,train2017,annotations} $DIST_DATA_ROOT/dist/ckpts/ outputs CC12M split: see cc12m_data.py (tar-offset index + feature shards). RUN python dist_bed.py --smoke [--arm feature_mse] python dist_bed.py --gates python dist_bed.py --arm infonce --seed 0 --steps 88000 --split cc12m python dist_bed.py --eval-only [--extra-evals] """ import argparse import json import math import os import sys import time import zlib import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from torch.utils.data import DataLoader, Dataset from torchvision import datasets as tvds from torchvision import transforms as T from torchvision.transforms import InterpolationMode def _root(): d = os.path.abspath(os.getcwd()) while True: if os.path.exists(os.path.join(d, "TRAINING.md")): return d p = os.path.dirname(d) if p == d: return os.getcwd() d = p ROOT = _root() if os.path.join(ROOT, "tools") not in sys.path: sys.path.insert(0, os.path.join(ROOT, "tools")) from dist_align_gate import (CFG_MODEL, CLIP_MEAN, CLIP_STD, PRIMARY, # noqa: E402 _feat, first_captions, laion_local_dir, load_bank, load_siglip_text_tokenizer, load_vision) from vitals import pentachoron_cv # noqa: E402 from losses import (EmbeddingAutograd, a7_grid_infonce, # noqa: E402 affinity_kl_loss, collinearity_novelty, prim_sq, procrustes_sq, siglip_pairwise_loss, x3_centered_cos_align, x3_cv_loss) torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False DEV = "cuda" if torch.cuda.is_available() else "cpu" DATA_ROOT = os.environ.get("DIST_DATA_ROOT", "./data") DATA = os.path.join(DATA_ROOT, "dist") COCO = os.path.join(DATA, "coco") RUNS = os.path.join(ROOT, "tools", "dist_runs") CKPTS = os.path.join(DATA, "ckpts") BATCH = 256 SIGLIP_CFG = "siglip_b16_384" CONSENSUS = ("clip_b16_laion2b", "clip_b32_openai", "clip_b32_laion2b", "clip_b32_datacomp", "clip_b16_openai") TRAIN_ARMS = ("feature_mse", "infonce", "siglip_pairwise", "affinity_kl", "blueprint", "consensus_gpa", "consensus_nce_mse", "x3_full", "x3_bce", "x3_autograd", "nce_mse_cv", "cbert_full") X3_ARMS = ("x3_full", "x3_bce", "x3_autograd", "nce_mse_cv") ARM_TARGET_GPA_EXTRA = ("cbert_full",) ARMS = TRAIN_ARMS + ("zs_floor", "teacher_ceiling", "hybrid_best") ARM_TARGET = {a: PRIMARY for a in ARMS} ARM_TARGET["siglip_pairwise"] = SIGLIP_CFG ARM_TARGET["consensus_gpa"] = "gpa" ARM_TARGET["consensus_nce_mse"] = "gpa" for _a in X3_ARMS + ARM_TARGET_GPA_EXTRA: ARM_TARGET[_a] = "gpa" CIFAR10_CLASSES = ("airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck") BLUEPRINT_CV_WEIGHT = 1e-3 EXTRA_EVALS = False def seed_for(name): return zlib.crc32(name.encode()) & 0x7FFFFFFF # ================================================================== STUDENT class Block(nn.Module): def __init__(self, d, heads): super().__init__() self.n1 = nn.LayerNorm(d) self.qkv = nn.Linear(d, 3 * d) self.proj = nn.Linear(d, d) self.n2 = nn.LayerNorm(d) self.fc1 = nn.Linear(d, 4 * d) self.fc2 = nn.Linear(4 * d, d) self.heads = heads def forward(self, x): B, N, C = x.shape q, k, v = (self.qkv(self.n1(x)) .reshape(B, N, 3, self.heads, C // self.heads) .permute(2, 0, 3, 1, 4)) a = F.scaled_dot_product_attention(q, k, v) x = x + self.proj(a.transpose(1, 2).reshape(B, N, C)) return x + self.fc2(F.gelu(self.fc1(self.n2(x)))) class Student(nn.Module): """ViT-Ti-class: d=240 depth 12 heads 4 patch 16 img 160. CLS-token readout, Linear head to the arm's target dim.""" def __init__(self, out_dim=512, d=240, depth=12, heads=4, patch=16, img=160): super().__init__() self.patch = nn.Conv2d(3, d, patch, patch) self.cls = nn.Parameter(torch.zeros(1, 1, d)) self.pos = nn.Parameter(torch.zeros(1, (img // patch) ** 2 + 1, d)) self.blocks = nn.ModuleList(Block(d, heads) for _ in range(depth)) self.norm = nn.LayerNorm(d) self.head = nn.Linear(d, out_dim) self.apply(self._init) nn.init.trunc_normal_(self.cls, std=0.02) nn.init.trunc_normal_(self.pos, std=0.02) @staticmethod def _init(m): if isinstance(m, (nn.Linear, nn.Conv2d)): nn.init.trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) def forward_features(self, x): x = self.patch(x).flatten(2).transpose(1, 2) x = torch.cat([self.cls.expand(x.shape[0], -1, -1), x], 1) + self.pos for b in self.blocks: x = b(x) return self.norm(x)[:, 0] def forward(self, x): return self.head(self.forward_features(x)) def build_student(arm, seed, out_dim): torch.manual_seed(seed_for(f"dist:{arm}:{seed}")) m = Student(out_dim=out_dim).to(DEV) tot = sum(p.numel() for p in m.parameters()) hd = sum(p.numel() for p in m.head.parameters()) print(f"[{arm} s{seed}] params total {tot:,} trunk {tot - hd:,} " f"head({out_dim}) {hd:,}", flush=True) return m # ===================================================================== DATA STUDENT_TF = T.Compose([ T.Resize(182, interpolation=InterpolationMode.BICUBIC), T.CenterCrop(160), T.ToTensor(), T.Normalize(CLIP_MEAN, CLIP_STD)]) class CocoImages(Dataset): """image_id-joined COCO jpgs -> (tensor, row_idx). Targets are indexed by row_idx in the main process (keeps DataLoader workers light).""" def __init__(self, split, ids): root = os.path.join(COCO, f"{split}2017") if not os.path.isdir(root): raise FileNotFoundError( f"COCO {split}2017 images absent at {root} — still " f"downloading? The bed cannot run this split yet.") self.paths = [os.path.join(root, f"{int(i):012d}.jpg") for i in ids] missing = [p for p in self.paths if not os.path.isfile(p)] if missing: raise FileNotFoundError( f"{len(missing)}/{len(ids)} {split}2017 jpgs missing " f"(e.g. {missing[0]}) — download still landing?") self.tf = STUDENT_TF def __len__(self): return len(self.paths) def __getitem__(self, i): return self.tf(Image.open(self.paths[i]).convert("RGB")), i def build_gpa(split): """GPA mean shape over the consensus five (all 512-d), cached once. L2-normalize towers -> 5 iters of { orthogonal-Procrustes each tower to the mean over a crc32-fixed 10k subsample (SVD fp64), re-mean, renorm }. Rows in the PRIMARY parquet order (towers reindexed by image_id).""" path = os.path.join(DATA, f"gpa_targets_{split}.pt") if os.path.isfile(path): return path towers, ids0 = [], None for cfg in CONSENSUS: ids, feats = load_bank(cfg, split) z = F.normalize(feats, dim=-1) if ids0 is None: ids0 = ids else: pos = {int(v): i for i, v in enumerate(ids)} z = z[torch.tensor([pos[int(v)] for v in ids0])] towers.append(z) X = torch.stack(towers) # (5, N, 512) N = X.shape[1] g = torch.Generator().manual_seed(seed_for(f"dist-gpa-sub:{split}")) sub = torch.randperm(N, generator=g)[:min(10000, N)] mean = F.normalize(X.mean(0), dim=-1) aligned = X.clone() print(f"[gpa {split}] towers {X.shape[0]} rows {N} subsample {len(sub)}") for it in range(5): for t in range(X.shape[0]): U, _, Vt = torch.linalg.svd( X[t][sub].t().double() @ mean[sub].double()) aligned[t] = X[t] @ (U @ Vt).float() mean = F.normalize(aligned.mean(0), dim=-1) cos = (F.normalize(aligned, dim=-1) * mean.unsqueeze(0)).sum(-1).mean(1) print(f"[gpa {split}] iter {it + 1} align cos " + " ".join(f"{float(c):.4f}" for c in cos) + f" | mean {float(cos.mean()):.4f}", flush=True) torch.save({"image_ids": torch.from_numpy(np.array(ids0, dtype=np.int64)), "targets": mean, "towers": list(CONSENSUS), "iters": 5}, path) print(f"[gpa {split}] cached -> {path}") return path def load_targets(arm, split): """-> (ids int64 np, targets fp32 L2-normalized, row-aligned to ids).""" t = ARM_TARGET[arm] if t == "gpa": d = torch.load(build_gpa(split), map_location="cpu", weights_only=True) return d["image_ids"].numpy(), d["targets"] ids, feats = load_bank(t, split) return ids, F.normalize(feats, dim=-1) # ============================================================ TEACHER HEADS def _text_tower(tower): from transformers import AutoTokenizer from dist_align_gate import _from_pretrained_fp32 if tower == "laion": d = laion_local_dir() return (_from_pretrained_fp32(d).to(DEV).eval(), AutoTokenizer.from_pretrained(d), False) m = _from_pretrained_fp32(CFG_MODEL[SIGLIP_CFG]).to(DEV).eval() return m, load_siglip_text_tokenizer(), True @torch.no_grad() def _embed_text(model, tok, texts, siglip, bs=128): outs = [] for i in range(0, len(texts), bs): kw = (dict(padding="max_length", truncation=True, max_length=64) if siglip else dict(padding=True, truncation=True)) tk = tok(texts[i:i + bs], return_tensors="pt", **kw) args = {"input_ids": tk["input_ids"].to(DEV)} if not siglip: # canonical SigLIP: full attention over pads args["attention_mask"] = tk["attention_mask"].to(DEV) outs.append(_feat(model.get_text_features(**args)).float().cpu()) return F.normalize(torch.cat(outs), dim=-1) def text_val(tower): """First-caption text embeddings for the 5k val images, cached.""" path = os.path.join(DATA, f"text_val_{tower}.pt") if os.path.isfile(path): return torch.load(path, map_location="cpu", weights_only=True) ids, _ = load_bank(PRIMARY if tower == "laion" else SIGLIP_CFG, "val") caps = first_captions() model, tok, sg = _text_tower(tower) emb = _embed_text(model, tok, [caps[int(i)] for i in ids], sg) del model torch.cuda.empty_cache() obj = {"image_ids": torch.from_numpy(np.array(ids, dtype=np.int64)), "emb": emb} torch.save(obj, path) print(f"[text_val {tower}] cached {tuple(emb.shape)} -> {path}") return obj def text_classes(tower, name, classes): path = os.path.join(DATA, f"text_{name}_{tower}.pt") if os.path.isfile(path): return torch.load(path, map_location="cpu", weights_only=True) model, tok, sg = _text_tower(tower) emb = _embed_text(model, tok, [f"a photo of a {c}" for c in classes], sg) del model torch.cuda.empty_cache() obj = {"classes": list(classes), "emb": emb} torch.save(obj, path) return obj # ===================================================================== LOSS def _diff_cv(z16, cv_target=0.20, weight=BLUEPRINT_CV_WEIGHT, n_sets=64): """Differentiable pentachoron-CV band term on PROJECTED student embeddings (16-d rows on the unit sphere) — same statistic as exp017:154-186 but fp32 and on embeddings, per the campaign spec. The weight is capped at 1e-3; subset draw crc32-fixed so the subset pattern is deterministic across steps.""" assert weight <= 1e-3, "CV force above 1e-3 is prohibited" x = F.normalize(z16, dim=-1) n = x.shape[0] g = torch.Generator().manual_seed(seed_for("dist-cv-subsets")) idx = torch.stack([torch.randperm(n, generator=g)[:5] for _ in range(n_sets)]).to(x.device) pts = x[idx] # (n_sets, 5, 16) d2 = torch.cdist(pts, pts).pow(2) cm = torch.ones(n_sets, 6, 6, dtype=pts.dtype, device=x.device) cm[:, 0, 0] = 0.0 cm[:, 1:, 1:] = d2 v = (-torch.linalg.det(cm) / 9216.0).clamp_min(1e-24).sqrt() cv = v.std() / v.mean().clamp_min(1e-12) return weight * (cv - cv_target).abs() def _fit_teacher_head(targets, labels, epochs=5, bs=2048, lr=1e-3): """Frozen teacher-space classifier for the x3 BCE term — the analog of x3's frozen calibrated soup head: a linear 80-class head fit ON THE GPA TARGETS (teacher side), then frozen. Deterministic (crc32 seeds).""" torch.manual_seed(seed_for("x3-teacher-head")) head = nn.Linear(targets.shape[-1], labels.shape[-1]).to(DEV) opt = torch.optim.Adam(head.parameters(), lr=lr, weight_decay=0.0) Tt, Y = targets.float().to(DEV), labels.float().to(DEV) n = Tt.shape[0] for ep in range(epochs): g = torch.Generator().manual_seed(seed_for(f"x3-head-ep:{ep}")) perm = torch.randperm(n, generator=g) for i in range(0, n - bs + 1, bs): j = perm[i:i + bs] opt.zero_grad(set_to_none=True) F.binary_cross_entropy_with_logits(head(Tt[j]), Y[j]).backward() opt.step() head.eval() for p in head.parameters(): p.requires_grad_(False) return head def loss_of(arm, zs, zt, ctx): """zs, zt L2-normalized. consensus_gpa uses the feature_mse form against the GPA mean (the target changes, not the primitive).""" if arm in ("feature_mse", "consensus_gpa"): return prim_sq(zs, zt).mean() if arm == "consensus_nce_mse": # CaptionBert Stage-A composite on the GPA target. Per-SAMPLE MSE # (= ||zs-zt||^2 = 2(1-cos)) so the 1.0/1.0 weighting is real — # per-element mean is ~1000x smaller than the NCE term at init. return (a7_grid_infonce(zs, zt, temp=0.07) + prim_sq(zs, zt).sum(-1).mean()) if arm == "cbert_full": # Source recipe VERBATIM (replication arm): per-ELEMENT # MSE at 1.0 and CV weight 0.1 — 100x the usual ceiling, # sanctioned ONLY as a flagged replication exception. return (a7_grid_infonce(zs, zt, temp=0.07) + prim_sq(zs, zt).mean() + 0.1 * x3_cv_loss(zs, ctx["cv_target"])) if arm in X3_ARMS: if arm in ("x3_full", "x3_autograd"): zs = EmbeddingAutograd.apply(zs, zs, ctx["anchors"], 0.01, 1.0) if arm == "x3_full": # faithful x3 stack incl. its per-ELEMENT MSE at 0.5 (measured # ~inert vs the NCE term — disclosed, kept for fidelity) and # its CV at 0.001 = the house ceiling. return (a7_grid_infonce(zs, zt, temp=0.07) + 0.5 * prim_sq(zs, zt).mean() + 0.3 * F.binary_cross_entropy_with_logits( ctx["head"](zs), ctx["yb"]) + 0.5 * x3_centered_cos_align(zs, zt) + 1e-3 * x3_cv_loss(zs, ctx["cv_target"])) base = (a7_grid_infonce(zs, zt, temp=0.07) + prim_sq(zs, zt).sum(-1).mean()) if arm == "x3_bce": return base + 0.3 * F.binary_cross_entropy_with_logits( ctx["head"](zs), ctx["yb"]) if arm == "nce_mse_cv": return base + 1e-3 * x3_cv_loss(zs, ctx["cv_target"]) return base # x3_autograd: EA-wrapped base if arm == "infonce": return a7_grid_infonce(zs, zt, temp=0.07) if arm == "siglip_pairwise": return siglip_pairwise_loss(zs, zt) if arm == "affinity_kl": return affinity_kl_loss(zs, zt, temp=0.07) if arm == "blueprint": return (a7_grid_infonce(zs, zt, temp=0.07) + 0.3 * procrustes_sq(zs, zt) + _diff_cv(zs @ ctx["P16"])) raise ValueError(arm) # ==================================================================== EVALS def student_embed01(model): mean = torch.tensor(CLIP_MEAN).view(1, 3, 1, 1).to(DEV) std = torch.tensor(CLIP_STD).view(1, 3, 1, 1).to(DEV) @torch.no_grad() def f(x01): x = F.interpolate(x01, size=(160, 160), mode="bicubic", align_corners=False).clamp(0, 1) return F.normalize(model((x - mean) / std), dim=-1) return f def teacher_embed01(model): """LAION B/16 vision tower on [0,1] tensors at its native 224 (tensor bicubic upscale — CIFAR is not in the bank, so the ceiling embeds it with the real tower).""" mean = torch.tensor(CLIP_MEAN).view(1, 3, 1, 1).to(DEV) std = torch.tensor(CLIP_STD).view(1, 3, 1, 1).to(DEV) @torch.no_grad() def f(x01): x = F.interpolate(x01, size=(224, 224), mode="bicubic", align_corners=False).clamp(0, 1) return F.normalize(_feat(model.get_image_features( pixel_values=(x - mean) / std)), dim=-1) return f @torch.no_grad() def zeroshot_classes(embed01, txt, dataset, subset_seed, subset=None, bs=500): x, y = dataset if subset: g = torch.Generator().manual_seed(seed_for(subset_seed)) keep = torch.randperm(len(y), generator=g)[:subset] x, y = x[keep], y[keep] txt = txt.to(DEV) ok = 0 for i in range(0, len(y), bs): z = embed01(x[i:i + bs].to(DEV)) ok += int((z @ txt.t()).argmax(-1).cpu().eq(y[i:i + bs]).sum()) return ok / len(y) def _cifar10(): ds = tvds.CIFAR10(root=DATA_ROOT, train=False, download=True) x = torch.from_numpy(ds.data).permute(0, 3, 1, 2).contiguous() return x.float().div_(255.0), torch.tensor(ds.targets) @torch.no_grad() def embed_val_student(model, ids, subset=None, workers=6): if subset is not None: g = torch.Generator().manual_seed(seed_for("dist-eval-sub")) keep = torch.randperm(len(ids), generator=g)[:subset].numpy() else: keep = np.arange(len(ids)) ds = CocoImages("val", ids[keep]) dl = DataLoader(ds, batch_size=BATCH, shuffle=False, num_workers=workers, pin_memory=True) emb = torch.empty(len(ds), model.head.out_features) for x, idx in dl: emb[idx] = F.normalize(model(x.to(DEV, non_blocking=True)), dim=-1).float().cpu() return emb, ids[keep] @torch.no_grad() def coco_retrieval(emb, row_ids, tower): """Image->text R@1/R@5 against the cached first-caption embeddings.""" tc = text_val(tower) pos = {int(v): i for i, v in enumerate(tc["image_ids"].tolist())} txt = tc["emb"][torch.tensor([pos[int(v)] for v in row_ids])].to(DEV) sim = emb.to(DEV) @ txt.t() ranks = (sim > sim.diag().unsqueeze(1)).sum(1) return (float((ranks == 0).float().mean()), float((ranks < 5).float().mean())) @torch.no_grad() def agreement(emb, row_ids, t_ids, targets): pos = {int(v): i for i, v in enumerate(t_ids)} idx = torch.tensor([pos[int(v)] for v in row_ids]) return float((emb * targets[idx]).sum(-1).mean()) @torch.no_grad() def geometry(emb): """Readout channel (CV is a readout here — the only CV FORCE in this bed is the blueprint arm's gated 1e-3 term): pentachoron CV on a fixed random D->16 projection + erank of the val matrix (fp64 gauges).""" D = emb.shape[1] gp = torch.Generator().manual_seed(seed_for(f"dist-proj16-{D}")) P = torch.randn(D, 16, generator=gp) cv = pentachoron_cv(emb.cpu() @ P, generator=torch.Generator().manual_seed( seed_for("dist-cv-eval"))) s = torch.linalg.svdvals(emb.double().cpu()) p = (s / s.sum()).clamp_min(1e-24) return float(cv), float(torch.exp(-(p * p.log()).sum())) def evaluate_student(arm, model, tower, smoke, workers=6): model.eval() ids_v, targets_v = load_targets(arm, "val") emb, row_ids = embed_val_student(model, ids_v, subset=500 if smoke else None, workers=workers) txt10 = text_classes(tower, "cifar10", CIFAR10_CLASSES)["emb"] zs = zeroshot_classes(student_embed01(model), txt10, _cifar10(), "dist-cifar-sub", subset=1000 if smoke else None) r1, r5 = coco_retrieval(emb, row_ids, tower) ag = agreement(emb, row_ids, ids_v, targets_v) cv, er = geometry(emb) rec = {"zs_cifar10": round(zs, 4), "coco_r1": round(r1, 4), "coco_r5": round(r5, 4), "agree_cos": round(ag, 4), "cv16": round(cv, 4), "erank": round(er, 2), "n_eval": int(len(row_ids)), "target": ARM_TARGET[arm]} if EXTRA_EVALS: try: ds = tvds.CIFAR100(root=DATA_ROOT, train=False, download=False) x = torch.from_numpy(ds.data).permute(0, 3, 1, 2).contiguous() data = (x.float().div_(255.0), torch.tensor(ds.targets)) cls = [c.replace("_", " ") for c in ds.classes] txt100 = text_classes(tower, "cifar100", cls)["emb"] rec["zs_cifar100"] = round(zeroshot_classes( student_embed01(model), txt100, data, "dist-c100-sub", subset=1000 if smoke else None), 4) except RuntimeError as e: print(f"[extra-evals] CIFAR-100 SKIP: {e}") model.train() return rec # ================================================================= TRAINING def _ledger(rec, smoke): os.makedirs(RUNS, exist_ok=True) name = "smoke.jsonl" if smoke else f"{rec['arm']}_s{rec['seed']}.jsonl" out = os.path.join(RUNS, name) with open(out, "a", encoding="utf-8") as f: f.write(json.dumps(rec) + "\n") return out def train_arm(arm, seed, steps, split="val", smoke=False, workers=6, bench=False): assert arm in ARMS, f"unknown arm {arm} (see --list)" if arm == "hybrid_best": raise SystemExit( "hybrid_best is RESERVED: minted only after the Night-0 verdicts " "(best vector arm + best grid arm at 0.5/0.5, and only if the " "gates show novelty >= 0.3 between them).") if DEV != "cuda": raise SystemExit("no CUDA device — this bed never CPU-trains " "(smoke mode checks shapes/parse only).") torch.cuda.set_per_process_memory_fraction(0.73) torch.cuda.reset_peak_memory_stats() t0 = time.time() if arm == "teacher_ceiling": return teacher_ceiling_arm(seed, smoke, t0) tower = "siglip" if arm == "siglip_pairwise" else "laion" out_dim = 768 if arm == "siglip_pairwise" else 512 cc_keys = None if split == "cc12m": from cc12m_data import (load_cc12m_gpa, load_cc12m_targets, load_cc12m_tower) assert arm not in ("x3_full", "x3_bce"), \ (f"{arm} needs the COCO 80-class multilabel BCE channel — " f"CC12M has no gold labels (COCO-scoped arm)") if ARM_TARGET[arm] == "gpa": cc_keys, t16 = load_cc12m_gpa() elif arm == "siglip_pairwise": cc_keys, t16 = load_cc12m_tower(SIGLIP_CFG) else: assert arm in ("feature_mse", "infonce", "affinity_kl", "blueprint"), \ f"{arm} has no CC12M target path yet" cc_keys, t16 = load_cc12m_targets() # fp16-RESIDENT, normalized per-batch on GPU (row-wise normalize is # batch-independent — numerically identical, HALF the cgroup # residency; two siglip OOM kills taught this) targets = t16 else: ids, targets = load_targets(arm, split) model = build_student(arm, seed, out_dim) ctx = {} if arm == "blueprint": gp = torch.Generator().manual_seed(seed_for("dist-proj16-512")) ctx["P16"] = torch.randn(512, 16, generator=gp).to(DEV) if arm == "cbert_full": from bank_utils import _sample_cv g = torch.Generator().manual_seed(seed_for("x3-cv-target")) ctx["cv_target"] = _sample_cv(F.normalize(targets[:20000].float(), dim=-1), g) print(f"[{arm} s{seed}] REPLICATION ARM (prototype-replication " f"law): CaptionBert Stage-A verbatim — per-element MSE 1.0, " f"CV 0.1 (ceiling EXCEPTION, flagged), AdamW+warmup+cosine+" f"clip. cv_target {ctx['cv_target']:.4f}", flush=True) if arm in X3_ARMS: from bank_utils import _sample_cv, coco_multilabel ai = torch.linspace(0, len(targets) - 1, 512).long() ctx["anchors"] = F.normalize(targets[ai].float(), dim=-1).to(DEV) g = torch.Generator().manual_seed(seed_for("x3-cv-target")) ctx["cv_target"] = _sample_cv(F.normalize(targets[:20000].float(), dim=-1), g) if arm in ("x3_full", "x3_bce"): # only these consume BCE ctx["labels"] = coco_multilabel(COCO, split, ids) ctx["head"] = _fit_teacher_head(targets, ctx["labels"]) print(f"[{arm} s{seed}] x3-ctx: cv_target " f"{ctx['cv_target']:.4f}" + (" head frozen (teacher-space)" if "head" in ctx else ""), flush=True) train_steps = 0 if arm == "zs_floor" else steps stats, s_per_step = {}, 0.0 if train_steps: if split == "cc12m": from cc12m_data import Cc12mImages ds = Cc12mImages(cc_keys, STUDENT_TF) else: ds = CocoImages(split, ids) print(f"[{arm} s{seed}] split {split} rows {len(ds)} steps " f"{train_steps} (~{train_steps * BATCH / len(ds):.1f} epochs)", flush=True) g = torch.Generator().manual_seed(seed_for(f"dist-data:{arm}:{seed}")) dl = DataLoader(ds, batch_size=BATCH, shuffle=True, generator=g, num_workers=workers, persistent_workers=workers > 0, pin_memory=True, drop_last=True) sched = None if arm == "cbert_full": # REPLICATION EXCEPTION (prototype-replication law): the # source recipe's AdamW(wd .01) + 1000-step linear warmup + # cosine to 1e-6 + grad-clip 1.0. Variants use pure Adam. opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01) warm = torch.optim.lr_scheduler.LinearLR( opt, start_factor=0.01, end_factor=1.0, total_iters=1000) cos = torch.optim.lr_scheduler.CosineAnnealingLR( opt, T_max=max(1, train_steps - 1000), eta_min=1e-6) sched = torch.optim.lr_scheduler.SequentialLR( opt, [warm, cos], [1000]) else: opt = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=0.0) it = iter(dl) t_train = time.time() for step in range(train_steps): try: x, idx = next(it) except StopIteration: it = iter(dl) x, idx = next(it) zt = targets[idx].to(DEV, non_blocking=True) if zt.dtype == torch.half: # cc12m fp16-resident path zt = F.normalize(zt.float(), dim=-1) if "labels" in ctx: ctx["yb"] = ctx["labels"][idx].to(DEV, non_blocking=True) opt.zero_grad(set_to_none=True) zs = F.normalize(model(x.to(DEV, non_blocking=True)), dim=-1) L = loss_of(arm, zs, zt, ctx) L.backward() if arm == "cbert_full": torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() if sched is not None: sched.step() lf = float(L.detach()) if step == 0: stats["loss_step1"] = round(lf, 4) if step == 10: print(f"[{arm} s{seed}] step10 loss {lf:.4f} peak " f"{torch.cuda.max_memory_allocated() / 2**30:.2f}GB " f"{(time.time() - t_train) / 11:.2f}s/step", flush=True) ema = lf if step == 0 else 0.98 * ema + 0.02 * lf if step and step % 2000 == 0: print(f"[{arm} s{seed}] step{step}/{train_steps} " f"loss {lf:.4f} ema {ema:.4f} " f"{(time.time() - t_train) / (step + 1):.2f}s/step", flush=True) if not smoke and step and step % 20000 == 0: os.makedirs(CKPTS, exist_ok=True) mid = os.path.join(CKPTS, f"{arm}_s{seed}_mid.pt") torch.save({"arm": arm, "seed": seed, "steps": step, "split": split, "out_dim": out_dim, "state_dict": model.state_dict()}, mid + ".tmp") os.replace(mid + ".tmp", mid) stats["loss_final"] = round(lf, 4) s_per_step = (time.time() - t_train) / train_steps del dl, it if bench: print(f"[BENCH {arm} s{seed}] {train_steps} steps " f"{s_per_step:.3f}s/step peak " f"{torch.cuda.max_memory_allocated() / 2**30:.2f}GB") return None ckpt = None if not smoke and train_steps: os.makedirs(CKPTS, exist_ok=True) ckpt = os.path.join(CKPTS, f"{arm}_s{seed}_t{train_steps}.pt") torch.save({"arm": arm, "seed": seed, "steps": train_steps, "split": split, "out_dim": out_dim, "state_dict": model.state_dict()}, ckpt) rec = {"arm": arm, "seed": seed, "steps": train_steps, "split": split, **stats, **evaluate_student(arm, model, tower, smoke, workers), "params": sum(p.numel() for p in model.parameters()), "s_per_step": round(s_per_step, 3), "wall_s": round(time.time() - t0, 1), "peak_gb": round(torch.cuda.max_memory_allocated() / 2**30, 2), "ckpt": ckpt, "smoke": smoke} out = _ledger(rec, smoke) print(f"[DONE {arm} s{seed}] zs {rec['zs_cifar10']:.4f} " f"r@1 {rec['coco_r1']:.4f} r@5 {rec['coco_r5']:.4f} " f"agree {rec['agree_cos']:.4f} cv16 {rec['cv16']:.3f} " f"erank {rec['erank']:.1f} ({rec['wall_s']}s) -> {out}", flush=True) return rec def teacher_ceiling_arm(seed, smoke, t0): """Night-0 reference row: stored laion b16 bank val features stand in as the student for retrieval/agreement/geometry; the real LAION tower embeds CIFAR (the bank holds no CIFAR rows).""" ids, targets = load_targets("teacher_ceiling", "val") if smoke: g = torch.Generator().manual_seed(seed_for("dist-eval-sub")) keep = torch.randperm(len(ids), generator=g)[:500].numpy() else: keep = np.arange(len(ids)) emb, row_ids = targets[keep], ids[keep] model, _ = load_vision(PRIMARY) txt10 = text_classes("laion", "cifar10", CIFAR10_CLASSES)["emb"] zs = zeroshot_classes(teacher_embed01(model), txt10, _cifar10(), "dist-cifar-sub", subset=1000 if smoke else None) del model torch.cuda.empty_cache() r1, r5 = coco_retrieval(emb, row_ids, "laion") cv, er = geometry(emb) rec = {"arm": "teacher_ceiling", "seed": seed, "steps": 0, "split": "val", "zs_cifar10": round(zs, 4), "coco_r1": round(r1, 4), "coco_r5": round(r5, 4), "agree_cos": 1.0, "cv16": round(cv, 4), "erank": round(er, 2), "n_eval": int(len(row_ids)), "target": PRIMARY, "params": 0, "s_per_step": 0.0, "wall_s": round(time.time() - t0, 1), "peak_gb": round(torch.cuda.max_memory_allocated() / 2**30, 2), "ckpt": None, "smoke": smoke} out = _ledger(rec, smoke) print(f"[DONE teacher_ceiling] zs {zs:.4f} r@1 {r1:.4f} r@5 {r5:.4f} " f"cv16 {cv:.3f} erank {er:.1f} ({rec['wall_s']}s) -> {out}", flush=True) return rec # ==================================================================== GATES def run_gates(workers=0): """Collinearity matrix across all six trainable losses on ONE fixed val batch at init state (reuses loss_forms.collinearity_novelty; base = feature_mse). Shared trunk + both heads; grads zero-filled on the shared parameter support.""" if DEV != "cuda": raise SystemExit("gates need the GPU (one forward on a 256 batch)") torch.cuda.set_per_process_memory_fraction(0.73) ids_l, t_l = load_targets("feature_mse", "val") ids_s, t_s = load_targets("siglip_pairwise", "val") ids_g, t_g = load_targets("consensus_gpa", "val") g = torch.Generator().manual_seed(seed_for("dist-gates-batch")) keep = torch.randperm(len(ids_l), generator=g)[:BATCH].numpy() ds = CocoImages("val", ids_l[keep]) x = torch.stack([ds[i][0] for i in range(len(ds))]).to(DEV) def rows(t_ids, T_): pos = {int(v): i for i, v in enumerate(t_ids)} return T_[torch.tensor([pos[int(v)] for v in ids_l[keep]])].to(DEV) zt_l, zt_s, zt_g = rows(ids_l, t_l), rows(ids_s, t_s), rows(ids_g, t_g) model = build_student("gates", 0, 512) torch.manual_seed(seed_for("dist:gates-head768")) head768 = nn.Linear(240, 768).to(DEV) gp = torch.Generator().manual_seed(seed_for("dist-proj16-512")) P16 = torch.randn(512, 16, generator=gp).to(DEV) f = model.forward_features(x) zs512 = F.normalize(model.head(f), dim=-1) zs768 = F.normalize(head768(f), dim=-1) L = {"feature_mse": prim_sq(zs512, zt_l).mean(), "infonce": a7_grid_infonce(zs512, zt_l, temp=0.07), "siglip_pairwise": siglip_pairwise_loss(zs768, zt_s), "affinity_kl": affinity_kl_loss(zs512, zt_l, temp=0.07), "blueprint": (a7_grid_infonce(zs512, zt_l, temp=0.07) + 0.3 * procrustes_sq(zs512, zt_l) + _diff_cv(zs512 @ P16)), "consensus_gpa": prim_sq(zs512, zt_g).mean()} params = list(model.parameters()) + list(head768.parameters()) names = list(L) print(f"=== DIST GATES: collinearity novelty 1-|cos(grad_i,grad_j)| " f"(init state, fixed val batch n={BATCH}) ===") print("losses: " + " ".join(f"{k}={float(L[k]):.4f}" for k in names)) nov = {} for i, a in enumerate(names): for b in names[i + 1:]: nov[(a, b)] = collinearity_novelty(L[a], L[b], params) base = "feature_mse" print("novelty vs feature_mse (base): " + " | ".join(f"{b} {nov[(base, b)]:.4f}" for b in names[1:])) col = {"feature_mse": "featmse", "infonce": "infonce", "siglip_pairwise": "siglip", "affinity_kl": "affkl", "blueprint": "blprnt", "consensus_gpa": "gpa"} print(f"{'':<16}" + "".join(f"{col[b]:>9}" for b in names)) for i, a in enumerate(names): cells = [] for j, b in enumerate(names): if i == j: cells.append(f"{'--':>9}") else: key = (a, b) if (a, b) in nov else (b, a) cells.append(f"{nov[key]:>9.4f}") print(f"{col[a]:<16}" + "".join(cells)) return nov # ==================================================================== SMOKE def smoke_summary(rows): print("\n=== DIST SMOKE TABLE (25 steps on the val2017 join; tiny eval: " "CIFAR 1000, retrieval 500) ===") hdr = (f"{'arm':<16}{'loss1':>9}{'loss25':>9}{'s/step':>8}{'peakGB':>8}" f"{'zs@1k':>8}{'r@1':>8}{'r@5':>8}{'agree':>8}{'cv16':>7}" f"{'erank':>8} verdict") print(hdr) for r in rows: if r is None: continue l1, l2 = r.get("loss_step1", float("nan")), r.get("loss_final", float("nan")) if "loss_final" not in r: # eval-only control arms verdict = "PASS" if r["peak_gb"] < 18.0 else "FAIL" else: ok = (math.isfinite(l2) and r["peak_gb"] < 18.0) fell = math.isfinite(l1) and math.isfinite(l2) and l2 < l1 verdict = "PASS" if ok and fell else ("WARN" if ok else "FAIL") print(f"{r['arm']:<16}{l1:>9.4f}{l2:>9.4f}{r['s_per_step']:>8.3f}" f"{r['peak_gb']:>8.2f}{r['zs_cifar10']:>8.4f}" f"{r['coco_r1']:>8.4f}{r['coco_r5']:>8.4f}" f"{r['agree_cos']:>8.4f}{r['cv16']:>7.3f}{r['erank']:>8.1f}" f" {verdict}") def print_launch_matrix(): print("\n=== OVERNIGHT LAUNCH MATRIX (after train2017 lands; sequential; " "~12000 steps ~ 26 epochs of train2017 at batch 256) ===") for arm in TRAIN_ARMS: for seed in (0, 1): print(f"python tools/dist_bed.py --arm {arm} --seed {seed} " f"--steps 12000 --split train") print("python tools/dist_bed.py --arm zs_floor --seed 0") print("python tools/dist_bed.py --arm teacher_ceiling --seed 0") # ===================================================================== MAIN if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("--arm", default=None) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--steps", type=int, default=4000) ap.add_argument("--split", default="val", choices=("val", "train", "cc12m")) ap.add_argument("--workers", type=int, default=6) ap.add_argument("--smoke", action="store_true") ap.add_argument("--smoke-all", action="store_true", dest="smoke_all") ap.add_argument("--gates", action="store_true") ap.add_argument("--bench", action="store_true") ap.add_argument("--list", action="store_true") ap.add_argument("--eval-only", default=None, dest="eval_only") ap.add_argument("--extra-evals", action="store_true", dest="extra") a, _ = ap.parse_known_args() EXTRA_EVALS = a.extra if a.list: print("arms:", " ".join(ARMS)) sys.exit(0) if a.gates: run_gates() sys.exit(0) if a.smoke_all: rows = [train_arm(arm, a.seed, 25, "val", smoke=True, workers=a.workers) for arm in TRAIN_ARMS] smoke_summary(rows) print_launch_matrix() sys.exit(0) if a.smoke: rec = train_arm(a.arm or "feature_mse", a.seed, 25, "val", smoke=True, workers=a.workers) smoke_summary([rec]) print_launch_matrix() sys.exit(0) if a.eval_only: ck = torch.load(a.eval_only, map_location="cpu", weights_only=True) torch.cuda.set_per_process_memory_fraction(0.73) torch.cuda.reset_peak_memory_stats() t0 = time.time() model = build_student(ck["arm"], ck["seed"], ck["out_dim"]) model.load_state_dict(ck["state_dict"]) tower = "siglip" if ck["arm"] == "siglip_pairwise" else "laion" rec = {"arm": ck["arm"], "seed": ck["seed"], "steps": ck["steps"], "split": ck.get("split", "?"), "eval_only": True, **evaluate_student(ck["arm"], model, tower, False, a.workers), "wall_s": round(time.time() - t0, 1)} out = _ledger(rec, False) print(f"[DONE eval-only {ck['arm']} s{ck['seed']}] " f"zs {rec['zs_cifar10']:.4f} r@1 {rec['coco_r1']:.4f} -> {out}") sys.exit(0) if a.bench: train_arm(a.arm or "feature_mse", a.seed, 10, a.split, smoke=True, workers=a.workers, bench=True) sys.exit(0) if not a.arm: print("arms:", " ".join(ARMS)) sys.exit(0) train_arm(a.arm, a.seed, a.steps, a.split, workers=a.workers)