Create mlp_ablation.py

mlp_ablation.py

@@ -0,0 +1,433 @@
+# Run cell 1 and cell 2, shape factory and model then run this to continue.
+# ablation showed random chance without the full geometric architecture.
+
+# =============================================================================
+# CELL 5: Architecture Ablation — MLP Baseline with Same Loss
+# Requires: Cell 1 + Cell 2 already executed (constants, generator, deform_grid)
+# Question: Is the loss creating the behavior, or the architecture?
+#
+# Same composite loss, same data, same hyperparams.
+# Plain MLP replaces: tracer attention, capacity cascade,
+# differentiation gate, curvature head, rectified flow arbiter.
+# =============================================================================
+
+import math, time, numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from pathlib import Path
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+if device.type == "cuda":
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+
+use_amp = device.type == "cuda"
+amp_dtype = (torch.bfloat16 if (device.type == "cuda" and
+             torch.cuda.is_bf16_supported()) else torch.float16)
+
+
+# =============================================================================
+# MLP Baseline — same output contract as GeometricShapeClassifier
+# =============================================================================
+
+class MLPBaseline(nn.Module):
+    """Plain MLP producing the same output dict as GeometricShapeClassifier.
+    No geometric inductive bias. Same loss surface."""
+
+    def __init__(self, grid_size=GS, n_classes=NUM_CLASSES,
+                 n_curvatures=NUM_CURVATURES, trunk_dim=256):
+        super().__init__()
+        inp = grid_size ** 3  # 125
+
+        self.trunk = nn.Sequential(
+            nn.Linear(inp, 512), nn.GELU(),
+            nn.Linear(512, 512), nn.GELU(),
+            nn.Linear(512, trunk_dim), nn.GELU(),
+            nn.Linear(trunk_dim, trunk_dim), nn.GELU(),
+        )
+
+        # Primary classifier
+        self.classifier = nn.Sequential(
+            nn.Linear(trunk_dim, 128), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(128, n_classes))
+
+        # Capacity analog: fill ratios (4 dims, sigmoid)
+        self.fill_head = nn.Sequential(
+            nn.Linear(trunk_dim, 64), nn.GELU(),
+            nn.Linear(64, 4), nn.Sigmoid())
+
+        # Learned capacities for diversity loss
+        self.cap_head = nn.Sequential(
+            nn.Linear(trunk_dim, 32), nn.GELU(),
+            nn.Linear(32, 4), nn.Softplus())
+
+        # Peak dimension (4-class)
+        self.peak_head = nn.Sequential(
+            nn.Linear(trunk_dim, 32), nn.GELU(), nn.Linear(32, 4))
+
+        # Overflow (4 dims, sigmoid)
+        self.overflow_head = nn.Sequential(
+            nn.Linear(trunk_dim, 32), nn.GELU(),
+            nn.Linear(32, 4), nn.Sigmoid())
+
+        # Volume regression
+        self.volume_head = nn.Sequential(
+            nn.Linear(trunk_dim, 64), nn.GELU(), nn.Linear(64, 1))
+
+        # Cayley-Menger determinant sign
+        self.cm_head = nn.Sequential(
+            nn.Linear(trunk_dim, 64), nn.GELU(),
+            nn.Linear(64, 1), nn.Tanh())
+
+        # Curvature binary
+        self.curved_head = nn.Sequential(
+            nn.Linear(trunk_dim, 32), nn.GELU(),
+            nn.Linear(32, 1), nn.Sigmoid())
+
+        # Curvature type (8-class)
+        self.curv_type_head = nn.Sequential(
+            nn.Linear(trunk_dim, 64), nn.GELU(),
+            nn.Linear(64, n_curvatures))
+
+        # Second classifier (arbiter analog)
+        self.refiner = nn.Sequential(
+            nn.Linear(trunk_dim, 128), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(128, n_classes))
+
+        # Confidence and blend
+        self.confidence_head = nn.Sequential(
+            nn.Linear(trunk_dim, 32), nn.GELU(),
+            nn.Linear(32, 1), nn.Sigmoid())
+        self.blend_head = nn.Sequential(
+            nn.Linear(trunk_dim, 32), nn.GELU(),
+            nn.Linear(32, 1), nn.Sigmoid())
+
+    def forward(self, grid, labels=None):
+        B = grid.shape[0]
+        x = grid.reshape(B, -1).float()
+        feat = self.trunk(x)
+
+        initial_logits = self.classifier(feat)
+        refined_logits = self.refiner(feat)
+
+        blend = self.blend_head(feat).squeeze(-1)
+        class_logits = (blend.unsqueeze(-1) * initial_logits +
+                        (1 - blend.unsqueeze(-1)) * refined_logits)
+
+        conf = self.confidence_head(feat).squeeze(-1)
+
+        return {
+            "class_logits": class_logits,
+            "initial_logits": initial_logits,
+            "refined_logits": refined_logits,
+            "fill_ratios": self.fill_head(feat),
+            "peak_logits": self.peak_head(feat),
+            "overflows": self.overflow_head(feat),
+            "capacities": self.cap_head(feat),
+            "volume_pred": self.volume_head(feat).squeeze(-1),
+            "cm_pred": self.cm_head(feat).squeeze(-1),
+            "is_curved_pred": self.curved_head(feat),
+            "curv_type_logits": self.curv_type_head(feat),
+            "trajectory_logits": [refined_logits],
+            "flow_loss": torch.tensor(0.0, device=grid.device),
+            "refined_confidence": self.confidence_head(feat),
+            "blend_weight": blend,
+            "confidence": conf,
+            "alternation": torch.zeros(B, device=grid.device),
+            "features": feat,
+        }
+
+
+# =============================================================================
+# Loss Functions (identical to Cell 3)
+# =============================================================================
+
+def _safe_bce(inp, tgt):
+    with torch.amp.autocast('cuda', enabled=False):
+        return F.binary_cross_entropy(inp.float(), tgt.float())
+
+def capacity_fill_loss(fr, dt): return _safe_bce(fr, dt)
+
+def overflow_reg(on, dt):
+    pk = dt.sum(dim=-1).long() - 1
+    loss = sum(on[b, pk[b].item():].sum() for b in range(on.shape[0]))
+    return loss / (on.shape[0] + 1e-8)
+
+def cap_diversity(c): return -c.var()
+def peak_loss(l, t): return F.cross_entropy(l, t)
+def cm_loss(p, t): return F.mse_loss(p, torch.sign(t))
+def curved_bce(p, t): return _safe_bce(p.squeeze(-1), t)
+def ctype_loss(l, t): return F.cross_entropy(l, t)
+
+
+# =============================================================================
+# Data — load cached or generate + cache
+# =============================================================================
+
+DATASET_PATH = Path("./cached_dataset.pt")
+N_SAMPLES = 500000
+SEED = 42
+
+if DATASET_PATH.exists():
+    print(f"Loading cached dataset from {DATASET_PATH}...")
+    t0 = time.time()
+    _cached = torch.load(DATASET_PATH, weights_only=True)
+    if _cached["n_samples"] == N_SAMPLES and _cached["seed"] == SEED:
+        train_ds = ShapeDataset.__new__(ShapeDataset)
+        val_ds = ShapeDataset.__new__(ShapeDataset)
+        for k in ["grids", "labels", "dim_conf", "peak_dim", "volume",
+                  "cm_det", "is_curved", "curvature"]:
+            setattr(train_ds, k, _cached["train"][k])
+            setattr(val_ds, k, _cached["val"][k])
+        print(f"Loaded {len(train_ds)} train + {len(val_ds)} val in {time.time()-t0:.1f}s")
+    else:
+        print(f"Cache mismatch — regenerating")
+        DATASET_PATH.unlink()
+
+if not DATASET_PATH.exists():
+    print("Generating dataset...")
+    all_samples = generate_parallel(N_SAMPLES, seed=SEED, n_workers=8)
+    n_train = int(len(all_samples) * 0.8)
+    train_ds = ShapeDataset(all_samples[:n_train])
+    val_ds = ShapeDataset(all_samples[n_train:])
+
+    print(f"Caching to {DATASET_PATH}...")
+    cache_data = {
+        "n_samples": N_SAMPLES, "seed": SEED,
+        "train": {k: getattr(train_ds, k) for k in ["grids", "labels", "dim_conf",
+                  "peak_dim", "volume", "cm_det", "is_curved", "curvature"]},
+        "val": {k: getattr(val_ds, k) for k in ["grids", "labels", "dim_conf",
+                "peak_dim", "volume", "cm_det", "is_curved", "curvature"]},
+    }
+    torch.save(cache_data, DATASET_PATH)
+    size_mb = DATASET_PATH.stat().st_size / 1e6
+    print(f"Cached: {size_mb:.0f}MB | {len(train_ds)} train + {len(val_ds)} val")
+
+train_loader = torch.utils.data.DataLoader(
+    train_ds, batch_size=4096, shuffle=True,
+    num_workers=4, pin_memory=True, persistent_workers=True)
+val_loader = torch.utils.data.DataLoader(
+    val_ds, batch_size=4096, shuffle=False,
+    num_workers=4, pin_memory=True, persistent_workers=True)
+
+
+# =============================================================================
+# Train
+# =============================================================================
+
+model = MLPBaseline().to(device)
+n_params = sum(p.numel() for p in model.parameters())
+print(f"MLPBaseline: {n_params:,} params")
+print(f"(GeometricShapeClassifier was 1,852,870 params)")
+
+if device.type == "cuda" and hasattr(torch, 'compile'):
+    try:
+        model = torch.compile(model, mode="default")
+        print("torch.compile: enabled")
+    except Exception as e:
+        print(f"torch.compile: skipped ({e})")
+
+epochs = 80
+lr = 3e-3
+optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
+warmup_epochs = 5
+def lr_lambda(ep):
+    if ep < warmup_epochs: return (ep + 1) / warmup_epochs
+    return 0.5 * (1 + math.cos(math.pi * (ep - warmup_epochs) / (epochs - warmup_epochs)))
+scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+w = {"cls": 1.0, "fill": 0.3, "peak": 0.3, "ovf": 0.05,
+     "div": 0.02, "vol": 0.1, "cm": 0.1, "curved": 0.2, "ctype": 0.2,
+     "arb_cls": 0.8, "arb_traj": 0.2, "arb_conf": 0.1, "flow": 0.5}
+
+use_scaler = use_amp and amp_dtype == torch.float16
+scaler = torch.amp.GradScaler('cuda', enabled=use_scaler)
+
+print(f"\nAblation: MLPBaseline vs GeometricShapeClassifier")
+print(f"Same loss ({len(w)} terms), same data, same schedule")
+print(f"{'='*70}")
+
+best_val_acc = 0
+t_start = time.time()
+
+for epoch in range(epochs):
+    t0 = time.time()
+    model.train()
+    correct, total = 0, 0
+    correct_init, correct_ref = 0, 0
+
+    for grid, label, dc, pd, vol, cm, ic, ct in train_loader:
+        grid = grid.to(device, non_blocking=True)
+        label = label.to(device, non_blocking=True)
+        dc = dc.to(device, non_blocking=True)
+        pd = pd.to(device, non_blocking=True)
+        vol = vol.to(device, non_blocking=True)
+        cm = cm.to(device, non_blocking=True)
+        ic = ic.to(device, non_blocking=True)
+        ct = ct.to(device, non_blocking=True)
+
+        grid = deform_grid(grid, p_dropout=0.05, p_add=0.05, p_shift=0.08)
+        optimizer.zero_grad(set_to_none=True)
+
+        with torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
+            out = model(grid, labels=label)
+
+            loss_first = (w["cls"] * F.cross_entropy(out["initial_logits"], label) +
+                          w["fill"] * capacity_fill_loss(out["fill_ratios"], dc) +
+                          w["peak"] * peak_loss(out["peak_logits"], pd) +
+                          w["ovf"] * overflow_reg(out["overflows"], dc) +
+                          w["div"] * cap_diversity(out["capacities"]) +
+                          w["vol"] * F.mse_loss(out["volume_pred"], torch.log1p(vol)) +
+                          w["cm"] * cm_loss(out["cm_pred"], cm) +
+                          w["curved"] * curved_bce(out["is_curved_pred"], ic) +
+                          w["ctype"] * ctype_loss(out["curv_type_logits"], ct))
+
+            loss_arb = w["arb_cls"] * F.cross_entropy(out["refined_logits"], label)
+            traj_loss = 0
+            for step_i, step_logits in enumerate(out["trajectory_logits"]):
+                step_weight = (step_i + 1) / len(out["trajectory_logits"])
+                traj_loss += step_weight * F.cross_entropy(step_logits, label)
+            traj_loss /= len(out["trajectory_logits"])
+            loss_arb += w["arb_traj"] * traj_loss
+            loss_arb += w["flow"] * out["flow_loss"]
+
+            with torch.no_grad():
+                is_correct = (out["refined_logits"].argmax(1) == label).float()
+            loss_arb += w["arb_conf"] * _safe_bce(
+                out["refined_confidence"].squeeze(-1), is_correct)
+
+            with torch.no_grad():
+                init_correct = (out["initial_logits"].argmax(1) == label).float()
+                ref_correct = (out["refined_logits"].argmax(1) == label).float()
+                blend_target = torch.where(init_correct >= ref_correct,
+                    torch.ones_like(init_correct) * 0.8,
+                    torch.ones_like(init_correct) * 0.2)
+            loss_arb += 0.1 * _safe_bce(out["blend_weight"], blend_target)
+
+            loss_blend = w["cls"] * F.cross_entropy(out["class_logits"], label)
+            loss = loss_first + loss_arb + loss_blend
+
+        scaler.scale(loss).backward()
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        scaler.step(optimizer)
+        scaler.update()
+
+        correct += (out["class_logits"].argmax(1) == label).sum().item()
+        correct_init += (out["initial_logits"].argmax(1) == label).sum().item()
+        correct_ref += (out["refined_logits"].argmax(1) == label).sum().item()
+        total += grid.size(0)
+
+    scheduler.step()
+
+    if epoch == 0 and device.type == "cuda":
+        peak = torch.cuda.max_memory_allocated() / 1e9
+        print(f"VRAM peak: {peak:.2f}GB | throughput: {total/(time.time()-t0):.0f} samples/s")
+
+    # Validate
+    model.eval()
+    vc, vt, vcc, vct = 0, 0, 0, 0
+    vc_init, vc_ref = 0, 0
+
+    with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
+        for grid, label, dc, pd, vol, cm, ic, ct in val_loader:
+            grid = grid.to(device, non_blocking=True)
+            label = label.to(device, non_blocking=True)
+            ic = ic.to(device, non_blocking=True)
+            out = model(grid)
+            vc += (out["class_logits"].argmax(1) == label).sum().item()
+            vc_init += (out["initial_logits"].argmax(1) == label).sum().item()
+            vc_ref += (out["refined_logits"].argmax(1) == label).sum().item()
+            vt += grid.size(0)
+            vcc += ((out["is_curved_pred"].squeeze(-1) > 0.5).float() == ic).sum().item()
+            vct += grid.size(0)
+
+    val_acc = vc / vt
+    val_init = vc_init / vt
+    val_ref = vc_ref / vt
+    curved_acc = vcc / vct
+    marker = " *" if val_acc > best_val_acc else ""
+    if val_acc > best_val_acc:
+        best_val_acc = val_acc
+
+    dt = time.time() - t0
+    if (epoch + 1) % 10 == 0 or epoch == 0 or marker:
+        print(f"Ep {epoch+1:3d}/{epochs} [{dt:.1f}s] | "
+              f"blend {val_acc:.3f} init {val_init:.3f} arb {val_ref:.3f} | "
+              f"curved {curved_acc:.3f}{marker}")
+
+total_time = time.time() - t_start
+print(f"\nDone in {total_time:.0f}s ({total_time/60:.1f}min)")
+
+
+# =============================================================================
+# Per-Class Breakdown
+# =============================================================================
+
+print(f"\n{'='*70}")
+print(f"Per-Class Results — MLPBaseline")
+print(f"{'='*70}")
+
+model.eval()
+cc_b = {n: 0 for n in CLASS_NAMES}
+cc_i = {n: 0 for n in CLASS_NAMES}
+cc_r = {n: 0 for n in CLASS_NAMES}
+ct_c = {n: 0 for n in CLASS_NAMES}
+
+with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
+    for grid, label, *_ in val_loader:
+        grid = grid.to(device, non_blocking=True)
+        label = label.to(device, non_blocking=True)
+        out = model(grid)
+        pb = out["class_logits"].argmax(1)
+        pi = out["initial_logits"].argmax(1)
+        pr = out["refined_logits"].argmax(1)
+        for k in range(len(label)):
+            name = CLASS_NAMES[label[k].item()]
+            cc_b[name] += (pb[k] == label[k]).item()
+            cc_i[name] += (pi[k] == label[k]).item()
+            cc_r[name] += (pr[k] == label[k]).item()
+            ct_c[name] += 1
+
+print(f"\n{'Class':22s} | {'Blend':>5s} {'Init':>5s} {'Arb':>5s} | "
+      f"{'Corr':>4s}/{'Tot':>4s} | {'Type':8s} Curvature")
+print("-" * 85)
+for name in CLASS_NAMES:
+    if ct_c[name] == 0: continue
+    ab = cc_b[name]/ct_c[name]
+    ai = cc_i[name]/ct_c[name]
+    ar = cc_r[name]/ct_c[name]
+    info = SHAPE_CATALOG[name]
+    print(f"  {name:20s} | {ab:.3f} {ai:.3f} {ar:.3f} | "
+          f"{cc_b[name]:4d}/{ct_c[name]:4d} | "
+          f"{'CURVED' if info['curved'] else 'rigid':8s} {info['curvature']}")
+
+
+# =============================================================================
+# Summary Comparison
+# =============================================================================
+
+print(f"\n{'='*70}")
+print(f"ABLATION SUMMARY")
+print(f"{'='*70}")
+print(f"  MLPBaseline:              {n_params:>10,} params | best val acc: {best_val_acc:.4f}")
+print(f"  GeometricShapeClassifier: 1,852,870 params | best val acc: 0.9022")
+print(f"  Delta:                    {n_params - 1852870:>+10,} params | "
+      f"delta acc: {best_val_acc - 0.9022:+.4f}")
+print()
+if best_val_acc >= 0.89:
+    print("  -> Loss is doing most of the work.")
+    print("     The composite multi-task signal is sufficient to discover")
+    print("     geometric structure without architectural inductive bias.")
+elif best_val_acc >= 0.80:
+    print("  -> Architecture contributes meaningfully.")
+    print("     The loss provides signal but the geometric inductive bias")
+    print("     (capacity cascade, tracers, flow arbiter) adds real value.")
+else:
+    print("  -> Architecture is critical.")
+    print("     The MLP cannot recover the same behavior from loss alone.")
+    print("     Geometric inductive bias is doing the heavy lifting.")
+print(f"\n  Curved detection: {curved_acc:.3f}")
+print(f"  Training time: {total_time:.0f}s ({total_time/60:.1f}min)")