Create analysis_v1.py

analysis_v1.py

@@ -0,0 +1,505 @@
+"""
+Constellation Bottleneck — Full Analysis
+==========================================
+Paste directly after the training cell.
+Uses `model` already in memory.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import math
+import os
+from torchvision import datasets, transforms
+from torchvision.utils import save_image, make_grid
+
+DEVICE = "cuda"
+os.makedirs("analysis_bn", exist_ok=True)
+
+def compute_cv(points, n_samples=1500, n_points=5):
+    N = points.shape[0]
+    if N < n_points: return float('nan')
+    points = F.normalize(points.to(DEVICE).float(), dim=-1)
+    vols = []
+    for _ in range(n_samples):
+        idx = torch.randperm(min(N, 5000), device=DEVICE)[:n_points]
+        pts = points[idx].unsqueeze(0)
+        gram = torch.bmm(pts, pts.transpose(1, 2))
+        norms = torch.diagonal(gram, dim1=1, dim2=2)
+        d2 = norms.unsqueeze(2) + norms.unsqueeze(1) - 2 * gram
+        d2 = F.relu(d2)
+        cm = torch.zeros(1, 6, 6, device=DEVICE, dtype=torch.float32)
+        cm[:, 0, 1:] = 1; cm[:, 1:, 0] = 1; cm[:, 1:, 1:] = d2
+        v2 = -torch.linalg.det(cm) / 9216
+        if v2[0].item() > 1e-20:
+            vols.append(v2[0].sqrt().cpu())
+    if len(vols) < 50: return float('nan')
+    vt = torch.stack(vols)
+    return (vt.std() / (vt.mean() + 1e-8)).item()
+
+def eff_dim(x):
+    x_c = x - x.mean(0, keepdim=True)
+    n = min(512, x.shape[0])
+    _, S, _ = torch.linalg.svd(x_c[:n].float(), full_matrices=False)
+    p = S / S.sum()
+    return p.pow(2).sum().reciprocal().item()
+
+CLASS_NAMES = ['plane','auto','bird','cat','deer','dog','frog','horse','ship','truck']
+
+model.eval()
+bn = model.bottleneck
+
+print("=" * 80)
+print("CONSTELLATION BOTTLENECK — FULL ANALYSIS")
+print(f"  Params: {sum(p.numel() for p in model.parameters()):,}")
+print(f"  Bottleneck: {sum(p.numel() for p in bn.parameters()):,}")
+print("=" * 80)
+
+# Load test data
+transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.5,)*3, (0.5,)*3),
+])
+test_ds = datasets.CIFAR10('./data', train=False, download=True, transform=transform)
+test_loader = torch.utils.data.DataLoader(test_ds, batch_size=256, shuffle=False)
+images_test, labels_test = next(iter(test_loader))
+images_test = images_test.to(DEVICE)
+labels_test = labels_test.to(DEVICE)
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 1: BOTTLENECK DIAGNOSTICS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 1: Bottleneck Diagnostics")
+print(f"{'━'*80}")
+
+drift = bn.drift().detach()
+home = F.normalize(bn.home, dim=-1).detach()
+curr = F.normalize(bn.anchors, dim=-1).detach()
+P, A, d = home.shape
+
+print(f"  Patches: {P}, Anchors/patch: {A}, Patch dim: {d}")
+print(f"  Drift: mean={drift.mean():.6f} rad ({math.degrees(drift.mean()):.2f}°)")
+print(f"         std={drift.std():.6f}  min={drift.min():.6f}  max={drift.max():.6f}")
+print(f"         max degrees: {math.degrees(drift.max()):.2f}°")
+print(f"  Skip gate: {bn.skip_gate.sigmoid().item():.4f}")
+print(f"  Near 0.29154: {(drift - 0.29154).abs().lt(0.05).float().mean().item():.1%}")
+
+# Per-patch drift
+print(f"\n  Per-patch drift:")
+for p in range(P):
+    d_p = drift[p].mean().item()
+    d_max = drift[p].max().item()
+    marker = " ◄ 0.29" if abs(d_p - 0.29154) < 0.05 else ""
+    marker2 = " ◄ MAX near 0.29" if abs(d_max - 0.29154) < 0.05 else ""
+    print(f"    P{p:2d}: mean={d_p:.4f} ({math.degrees(d_p):.1f}°) "
+          f"max={d_max:.4f} ({math.degrees(d_max):.1f}°){marker}{marker2}")
+
+# Anchor pairwise spread
+print(f"\n  Anchor spread per patch:")
+for p in range(min(8, P)):
+    sim = (curr[p] @ curr[p].T)
+    sim.fill_diagonal_(0)
+    print(f"    P{p}: mean_cos={sim.mean():.4f} max={sim.max():.4f} min={sim.min():.4f}")
+
+# Anchor effective dimensionality
+print(f"\n  Anchor effective dimensionality:")
+for p in range(min(8, P)):
+    _, S, _ = torch.linalg.svd(curr[p].float(), full_matrices=False)
+    pr = S / S.sum()
+    ed = pr.pow(2).sum().reciprocal().item()
+    print(f"    P{p}: eff_dim={ed:.1f} / {A}")
+
+# Drift histogram — where do anchors cluster?
+all_drifts = drift.flatten().cpu().numpy()
+print(f"\n  Drift distribution:")
+bins = [0.0, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40]
+hist, _ = np.histogram(all_drifts, bins=bins)
+for i in range(len(bins)-1):
+    bar = "█" * hist[i]
+    print(f"    {bins[i]:.2f}-{bins[i+1]:.2f}: {hist[i]:3d} {bar}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 2: SPHERE REPRESENTATION — CV OF BOTTLENECK EMBEDDINGS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 2: Sphere Representation — CV of bottleneck embeddings")
+print(f"  These live on S^15. Does CV approach 0.20?")
+print(f"{'━'*80}")
+
+# Hook to capture sphere embeddings
+sphere_embeddings = {}
+tri_profiles = {}
+
+def hook_sphere(module, input, output):
+    # The forward method: proj_in → norm → reshape → normalize
+    # We need to grab AFTER L2 norm. Hook the full bottleneck
+    # and manually compute the sphere embedding.
+    pass
+
+# Manually extract sphere embeddings at different timesteps
+print(f"\n  {'t':>6} {'CV_sphere':>10} {'CV_tri':>10} {'eff_d_sph':>10} "
+      f"{'eff_d_tri':>10} {'sph_norm':>10}")
+
+for t_val in [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]:
+    B = images_test.shape[0]
+    t = torch.full((B,), t_val, device=DEVICE)
+    eps = torch.randn_like(images_test)
+    t_b = t.view(B, 1, 1, 1)
+    x_t = (1 - t_b) * images_test + t_b * eps
+
+    with torch.no_grad():
+        # Run encoder manually
+        cond = model.time_emb(t) + model.class_emb(labels_test)
+        h = model.in_conv(x_t)
+        skips = [h]
+        for i in range(len(model.channel_mults)):
+            for block in model.enc[i]:
+                if isinstance(block, nn.Sequential):
+                    h = block[0](h); h = block[1](h, cond)
+                else:
+                    h = block(h, cond)
+            skips.append(h)
+            if i < len(model.enc_down):
+                h = model.enc_down[i](h)
+
+        # Get sphere embedding
+        h_flat = h.reshape(B, -1)
+        emb = bn.proj_in(h_flat)
+        emb = bn.proj_in_norm(emb)
+        patches = emb.reshape(B, bn.n_patches, bn.patch_dim)
+        patches_n = F.normalize(patches, dim=-1)
+
+        # CV of sphere embeddings (flatten patches back to one vector)
+        sphere_flat = patches_n.reshape(B, -1)  # (B, 256) on product of spheres
+        cv_sphere = compute_cv(sphere_flat, n_samples=1000)
+        ed_sphere = eff_dim(sphere_flat)
+        norm_sph = sphere_flat.norm(dim=-1).mean().item()
+
+        # Triangulation profile
+        tri = bn.triangulate(patches_n)  # (B, 768)
+        cv_tri = compute_cv(tri, n_samples=1000)
+        ed_tri = eff_dim(tri)
+
+        # Per-patch CV
+        if t_val == 0.0:
+            print(f"\n  Per-patch CV at t=0 (should be ≈0.20 if d=16):")
+            for p in range(min(8, bn.n_patches)):
+                patch_p = patches_n[:, p, :]  # (B, 16) on S^15
+                cv_p = compute_cv(patch_p, n_samples=1000)
+                print(f"    Patch {p}: CV={cv_p:.4f}")
+            print()
+
+    print(f"  {t_val:>6.2f} {cv_sphere:>10.4f} {cv_tri:>10.4f} {ed_sphere:>10.1f} "
+          f"{ed_tri:>10.1f} {norm_sph:>10.4f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 3: PER-CLASS ANCHOR ROUTING
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 3: Per-Class Anchor Routing")
+print(f"{'━'*80}")
+
+# Collect per-class nearest anchors across all patches
+class_nearest = {c: [] for c in range(10)}
+anchors_n = F.normalize(bn.anchors.detach(), dim=-1)
+
+for images_b, labels_b in test_loader:
+    images_b = images_b.to(DEVICE)
+    labels_b = labels_b.to(DEVICE)
+    B = images_b.shape[0]
+    t = torch.zeros(B, device=DEVICE)  # clean images
+
+    with torch.no_grad():
+        cond = model.time_emb(t) + model.class_emb(labels_b)
+        h = model.in_conv(images_b)
+        for i in range(len(model.channel_mults)):
+            for block in model.enc[i]:
+                if isinstance(block, nn.Sequential):
+                    h = block[0](h); h = block[1](h, cond)
+                else:
+                    h = block(h, cond)
+            if i < len(model.enc_down):
+                h = model.enc_down[i](h)
+
+        h_flat = h.reshape(B, -1)
+        emb = bn.proj_in_norm(bn.proj_in(h_flat))
+        patches = F.normalize(emb.reshape(B, bn.n_patches, bn.patch_dim), dim=-1)
+
+        # Nearest anchor per patch
+        cos = torch.einsum('bpd,pad->bpa', patches, anchors_n)  # (B, P, A)
+        nearest = cos.argmax(dim=-1)  # (B, P)
+
+        for i in range(B):
+            c = labels_b[i].item()
+            class_nearest[c].append(nearest[i].cpu())
+
+    if sum(len(v) for v in class_nearest.values()) > 5000:
+        break
+
+# Show routing for first 4 patches
+for p_idx in range(min(4, bn.n_patches)):
+    print(f"\n  Patch {p_idx} — nearest anchor per class:")
+    print(f"  {'class':>10}", end="")
+    for a in range(A):
+        print(f" {a:>4}", end="")
+    print()
+
+    for c in range(10):
+        if not class_nearest[c]:
+            continue
+        nearest_all = torch.stack(class_nearest[c])  # (N, P)
+        nearest_p = nearest_all[:, p_idx]
+        counts = torch.bincount(nearest_p, minlength=A).float()
+        counts = counts / counts.sum()
+        row = f"  {CLASS_NAMES[c]:>10}"
+        for a in range(A):
+            pct = counts[a].item()
+            if pct > 0.15:
+                row += f" {pct:>3.0%}█"
+            elif pct > 0.05:
+                row += f" {pct:>3.0%}░"
+            else:
+                row += f"  {pct:>3.0%}"
+            #row += f" {pct:>3.0%}"
+        print(row)
+
+# Are anchor patterns class-specific?
+print(f"\n  Anchor routing entropy per class (lower = more concentrated):")
+for c in range(10):
+    if not class_nearest[c]:
+        continue
+    nearest_all = torch.stack(class_nearest[c])
+    # Average across patches
+    total_entropy = 0
+    for p_idx in range(bn.n_patches):
+        counts = torch.bincount(nearest_all[:, p_idx], minlength=A).float()
+        counts = counts / counts.sum()
+        entropy = -(counts * (counts + 1e-8).log()).sum().item()
+        total_entropy += entropy
+    avg_entropy = total_entropy / bn.n_patches
+    max_entropy = math.log(A)
+    print(f"    {CLASS_NAMES[c]:>10}: H={avg_entropy:.3f} / {max_entropy:.3f} "
+          f"({avg_entropy/max_entropy:.1%} of max)")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 4: SKIP GATE ANALYSIS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 4: Skip Gate — how much goes through constellation vs skip?")
+print(f"{'━'*80}")
+
+gate = bn.skip_gate.sigmoid().item()
+print(f"  Skip gate value: {gate:.4f}")
+print(f"  Skip path:          {gate:.1%}")
+print(f"  Constellation path: {1-gate:.1%}")
+print(f"  Skip proj params:   {sum(p.numel() for p in [bn.skip_proj.weight, bn.skip_proj.bias]):,}")
+print(f"  Patchwork params:   {sum(p.numel() for p in bn.patchwork.parameters()):,}")
+print(f"\n  ⚠ skip_proj is Linear(16384, 16384) = "
+      f"{bn.skip_proj.weight.numel():,} params")
+print(f"  ⚠ This single layer is {bn.skip_proj.weight.numel()/1e6:.0f}M params — "
+      f"larger than the rest of the model combined")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 5: GENERATION — PER CLASS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 5: Generation Quality")
+print(f"{'━'*80}")
+
+print(f"  {'class':>10} {'intra_cos':>10} {'std':>8} {'CV':>8} {'norm':>8}")
+
+all_gen = []
+for c in range(10):
+    imgs, _ = sample(model, 64, 50, class_label=c)
+    imgs = (imgs + 1) / 2  # to [0,1]
+    all_gen.append(imgs)
+
+    flat = imgs.reshape(64, -1)
+    flat_n = F.normalize(flat, dim=-1)
+    sim = flat_n @ flat_n.T
+    mask = ~torch.eye(64, device=DEVICE, dtype=torch.bool)
+    intra = sim[mask].mean().item()
+    std = sim[mask].std().item()
+    cv = compute_cv(flat, 500)
+    norm = flat.norm(dim=-1).mean().item()
+    print(f"  {CLASS_NAMES[c]:>10} {intra:>10.4f} {std:>8.4f} {cv:>8.4f} {norm:>8.2f}")
+
+    save_image(make_grid(imgs[:16], nrow=4), f"analysis_bn/class_{CLASS_NAMES[c]}.png")
+
+# All classes grid
+all_grid = torch.cat([g[:4] for g in all_gen])
+save_image(make_grid(all_grid, nrow=10), "analysis_bn/all_classes.png")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 6: ABLATION — SKIP ONLY vs CONSTELLATION ONLY
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 6: Ablation — Skip-only vs Constellation-only")
+print(f"{'━'*80}")
+
+original_gate = bn.skip_gate.data.clone()
+
+# A) Full model (as trained)
+torch.manual_seed(999)
+with torch.no_grad():
+    imgs_full, _ = sample(model, 32, 50, class_label=3)
+
+# B) Skip only (gate → +100, sigmoid ≈ 1.0)
+bn.skip_gate.data.fill_(100.0)
+torch.manual_seed(999)
+with torch.no_grad():
+    imgs_skip, _ = sample(model, 32, 50, class_label=3)
+
+# C) Constellation only (gate → -100, sigmoid ≈ 0.0)
+bn.skip_gate.data.fill_(-100.0)
+torch.manual_seed(999)
+with torch.no_grad():
+    imgs_const, _ = sample(model, 32, 50, class_label=3)
+
+# Restore
+bn.skip_gate.data.copy_(original_gate)
+
+imgs_full_01 = (imgs_full + 1) / 2
+imgs_skip_01 = (imgs_skip + 1) / 2
+imgs_const_01 = (imgs_const + 1) / 2
+
+# Compare
+for name, imgs in [('skip_only', imgs_skip), ('const_only', imgs_const)]:
+    delta = (imgs_full - imgs).abs()
+    pixel_diff = delta.mean().item()
+    cos = F.cosine_similarity(
+        imgs_full.reshape(32, -1), imgs.reshape(32, -1)).mean().item()
+    print(f"  {name:>15}: pixel_Δ={pixel_diff:.6f}  cos_sim={cos:.6f}  "
+          f"max_Δ={delta.max():.4f}")
+
+# Save comparison: top=full, mid=skip_only, bot=constellation_only
+comparison = torch.cat([imgs_full_01[:8], imgs_skip_01[:8], imgs_const_01[:8]])
+save_image(make_grid(comparison, nrow=8), "analysis_bn/ablation_skip_vs_const.png")
+print(f"  ✓ Saved (top=full, mid=skip_only, bot=constellation_only)")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 7: VELOCITY FIELD
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 7: Velocity Field Quality")
+print(f"{'━'*80}")
+
+print(f"  {'t':>6} {'v_norm':>10} {'v·target':>10} {'mse':>10}")
+
+for t_val in [0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95]:
+    B = 128
+    imgs_v = images_test[:B]
+    labs_v = labels_test[:B]
+    t = torch.full((B,), t_val, device=DEVICE)
+    eps = torch.randn_like(imgs_v)
+    t_b = t.view(B, 1, 1, 1)
+    x_t = (1 - t_b) * imgs_v + t_b * eps
+    v_target = eps - imgs_v
+
+    with torch.no_grad():
+        v_pred = model(x_t, t, labs_v)
+
+    v_norm = v_pred.reshape(B, -1).norm(dim=-1).mean().item()
+    v_cos = F.cosine_similarity(
+        v_pred.reshape(B, -1), v_target.reshape(B, -1)).mean().item()
+    mse = F.mse_loss(v_pred, v_target).item()
+    print(f"  {t_val:>6.2f} {v_norm:>10.2f} {v_cos:>10.4f} {mse:>10.4f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 8: ODE TRAJECTORY — CV THROUGH GENERATION
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 8: ODE Trajectory — geometry through generation")
+print(f"{'━'*80}")
+
+n_steps = 50
+B_traj = 256
+x = torch.randn(B_traj, 3, 32, 32, device=DEVICE)
+labels_traj = torch.randint(0, 10, (B_traj,), device=DEVICE)
+dt = 1.0 / n_steps
+
+print(f"  {'step':>6} {'t':>6} {'x_norm':>10} {'x_std':>10} {'CV':>8}")
+
+for step in range(n_steps):
+    t_val = 1.0 - step * dt
+    t = torch.full((B_traj,), t_val, device=DEVICE)
+    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        v = model(x, t, labels_traj)
+    x = x - v.float() * dt
+
+    if step in [0, 1, 5, 10, 20, 30, 40, 49]:
+        xf = x.reshape(B_traj, -1)
+        print(f"  {step:>6} {t_val:>6.2f} {xf.norm(dim=-1).mean().item():>10.2f} "
+              f"{x.std().item():>10.4f} {compute_cv(xf, 500):>8.4f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 9: INTER vs INTRA CLASS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 9: Inter vs Intra Class Separation")
+print(f"{'━'*80}")
+
+intra_sims = []
+inter_sims = []
+for c in range(10):
+    flat = F.normalize(all_gen[c].reshape(64, -1), dim=-1)
+    sim = flat @ flat.T
+    mask = ~torch.eye(64, device=DEVICE, dtype=torch.bool)
+    intra_sims.append(sim[mask].mean().item())
+
+for i in range(10):
+    for j in range(i+1, 10):
+        fi = F.normalize(all_gen[i].reshape(64, -1), dim=-1)
+        fj = F.normalize(all_gen[j].reshape(64, -1), dim=-1)
+        inter_sims.append((fi @ fj.T).mean().item())
+
+print(f"  Intra-class cos: {np.mean(intra_sims):.4f} ± {np.std(intra_sims):.4f}")
+print(f"  Inter-class cos: {np.mean(inter_sims):.4f} ± {np.std(inter_sims):.4f}")
+ratio = np.mean(intra_sims) / (np.mean(inter_sims) + 1e-8)
+print(f"  Separation ratio: {ratio:.3f}×")
+
+
+# ══════════════════════════════════════════════════════════════════
+# SUMMARY
+# ═════════════════════════════════════════════════════════���════════
+
+print(f"\n{'='*80}")
+print("ANALYSIS COMPLETE")
+print(f"{'='*80}")
+print(f"""
+  Files in analysis_bn/:
+    class_*.png              per-class samples
+    all_classes.png          4 per class grid
+    ablation_skip_vs_const.png  top=full, mid=skip, bot=constellation
+
+  Key questions answered:
+    1. Does per-patch CV ≈ 0.20? (Test 2)
+       → If yes, the bottleneck lives at the natural S^15 dimension
+    2. Is anchor routing class-specific? (Test 3)
+       → If entropy varies by class, constellation routes differently
+    3. Does the skip path dominate? (Tests 4 & 6)
+       → If skip_only ≈ full, the 268M skip_proj IS the model
+    4. Does constellation-only work at all? (Test 6)
+       → The real test of whether geometric encoding carries signal
+""")
+print("=" * 80)