Create analyze_soup.py

analyze_soup.py

@@ -0,0 +1,498 @@
+#!/usr/bin/env python3
+"""
+BASE TIER SOUP ANALYSIS
+========================
+Load the trained 800K param soup and examine:
+  - Anchor geometry on the 128-d hypersphere
+  - Projector alignment (do the 3 experts converge?)
+  - Triangulation patterns (which anchors are used?)
+  - Patchwork compartment activation profiles
+  - Per-expert projected distributions
+  - CV and volume geometry of the learned space
+  - Per-class anchor affinity (which anchors serve which COCO classes?)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import math
+import os
+import gc
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+D_EXPERT = 768
+D_ANCHOR = 128
+N_ANCHORS = 256
+N_CLASSES = 80
+N_COMP = 8
+D_COMP = 64
+EXPERTS = ["clip_l14_openai", "dinov2_b14", "siglip_b16_384"]
+
+print("=" * 65)
+print("BASE TIER SOUP ANALYSIS")
+print(f"  Device: {DEVICE}")
+print("=" * 65)
+
+
+# ══════════════════════════════════════════════════════════════════
+# LOAD MODEL + DATA
+# ══════════════════════════════════════════════════════════════════
+
+# Rebuild model class (minimal, for loading)
+class ExpertProjector(nn.Module):
+    def __init__(self, d_in=D_EXPERT, d_out=D_ANCHOR):
+        super().__init__()
+        self.proj = nn.Sequential(nn.Linear(d_in, d_out), nn.LayerNorm(d_out))
+    def forward(self, x):
+        return F.normalize(self.proj(x), dim=-1)
+
+class Constellation(nn.Module):
+    def __init__(self, n_anchors=N_ANCHORS, d=D_ANCHOR):
+        super().__init__()
+        self.n_anchors = n_anchors
+        self.anchors = nn.Parameter(F.normalize(torch.randn(n_anchors, d), dim=-1))
+    def triangulate(self, emb):
+        a = F.normalize(self.anchors, dim=-1)
+        cos = emb @ a.T
+        return 1.0 - cos, cos.argmax(dim=-1)
+
+class Patchwork(nn.Module):
+    def __init__(self, n_anchors=N_ANCHORS, n_comp=N_COMP, d_comp=D_COMP):
+        super().__init__()
+        self.n_comp = n_comp
+        asgn = torch.arange(n_anchors) % n_comp
+        self.register_buffer("asgn", asgn)
+        self.comps = nn.ModuleList([nn.Sequential(
+            nn.Linear((asgn == k).sum().item(), d_comp * 2), nn.GELU(),
+            nn.Linear(d_comp * 2, d_comp), nn.LayerNorm(d_comp))
+            for k in range(n_comp)])
+    def forward(self, tri):
+        return torch.cat([self.comps[k](tri[:, self.asgn == k])
+                         for k in range(self.n_comp)], -1)
+
+class BaseTierSoup(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.n_experts = 3
+        self.projectors = nn.ModuleList([ExpertProjector() for _ in range(3)])
+        self.constellation = Constellation()
+        self.patchwork = Patchwork()
+        pw_dim = N_COMP * D_COMP
+        self.classifier = nn.Sequential(
+            nn.Linear(pw_dim + D_ANCHOR, pw_dim), nn.GELU(),
+            nn.LayerNorm(pw_dim), nn.Dropout(0.1),
+            nn.Linear(pw_dim, N_CLASSES))
+    def forward(self, expert_embeddings, apply_autograd=False):
+        projected = [self.projectors[i](expert_embeddings[i]) for i in range(3)]
+        fused = F.normalize(sum(projected) / 3, dim=-1)
+        tri, nearest = self.constellation.triangulate(fused)
+        pw = self.patchwork(tri)
+        logits = self.classifier(torch.cat([pw, fused], -1))
+        return logits, fused, tri, nearest, projected
+
+print(f"\n  Loading checkpoint...")
+ckpt = torch.load("checkpoints/base_tier_best.pt", map_location="cpu", weights_only=False)
+model = BaseTierSoup()
+model.load_state_dict(ckpt["state_dict"])
+model = model.eval().to(DEVICE)
+print(f"  Loaded: mAP={ckpt['mAP']:.3f} cv={ckpt['cv']:.4f} epoch={ckpt['epoch']}")
+
+# Load val data
+from datasets import load_dataset
+ref = load_dataset("AbstractPhil/bulk-coco-features", EXPERTS[0], split="val")
+val_ids = ref["image_id"]; N_val = len(val_ids)
+id_map = {iid: i for i, iid in enumerate(val_ids)}
+val_labels = torch.zeros(N_val, N_CLASSES)
+for i, labs in enumerate(ref["labels"]):
+    for l in labs:
+        if l < N_CLASSES: val_labels[i, l] = 1.0
+
+val_feats = []
+for name in EXPERTS:
+    ds = load_dataset("AbstractPhil/bulk-coco-features", name, split="val")
+    feats = torch.zeros(N_val, D_EXPERT)
+    for row in ds:
+        if row["image_id"] in id_map:
+            feats[id_map[row["image_id"]]] = torch.tensor(row["features"], dtype=torch.float32)
+    val_feats.append(feats.to(DEVICE))
+    print(f"  {name} loaded")
+    del ds; gc.collect()
+
+# Run full val through model
+print(f"\n  Running inference on {N_val} val images...")
+all_logits, all_fused, all_tri, all_nearest, all_proj = [], [], [], [], [[], [], []]
+BATCH = 256
+with torch.no_grad():
+    for j in range(0, N_val, BATCH):
+        end = min(j + BATCH, N_val)
+        batch = [val_feats[e][j:end] for e in range(3)]
+        lo, fu, tr, ne, pr = model(batch)
+        all_logits.append(lo.cpu())
+        all_fused.append(fu.cpu())
+        all_tri.append(tr.cpu())
+        all_nearest.append(ne.cpu())
+        for e in range(3):
+            all_proj[e].append(pr[e].cpu())
+
+logits = torch.cat(all_logits)
+fused = torch.cat(all_fused)
+tri = torch.cat(all_tri)
+nearest = torch.cat(all_nearest)
+proj = [torch.cat(all_proj[e]) for e in range(3)]
+print(f"  Done: fused={fused.shape} tri={tri.shape}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# SCAN 1: ANCHOR GEOMETRY
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("SCAN 1: ANCHOR GEOMETRY")
+print(f"{'='*65}")
+
+anchors = F.normalize(model.constellation.anchors.detach().cpu(), dim=-1)
+
+# Pairwise cosine
+anchor_sim = anchors @ anchors.T
+anchor_sim.fill_diagonal_(0)
+
+print(f"  Anchor pairwise cosine:")
+print(f"    mean={anchor_sim.mean():.4f} std={anchor_sim.std():.4f}")
+print(f"    max={anchor_sim.max():.4f} min={anchor_sim.min():.4f}")
+
+# Distribution of max-neighbor cosine
+max_neighbor = anchor_sim.max(dim=1).values
+print(f"  Max neighbor cosine per anchor:")
+print(f"    mean={max_neighbor.mean():.4f} std={max_neighbor.std():.4f}")
+print(f"    max={max_neighbor.max():.4f} min={max_neighbor.min():.4f}")
+
+# Anchor norms (should be ~1.0 after normalize)
+anchor_norms = anchors.norm(dim=-1)
+print(f"  Anchor norms: mean={anchor_norms.mean():.6f} std={anchor_norms.std():.6f}")
+
+# SVD of anchor matrix
+sv = torch.linalg.svdvals(anchors)
+eff_rank = ((sv.sum()**2) / (sv.pow(2).sum() + 1e-12)).item()
+print(f"  Anchor spectral: eff_rank={eff_rank:.1f}/{min(anchors.shape)}")
+print(f"    sv_max={sv[0]:.4f} sv_10={sv[9]:.4f} sv_50={sv[49]:.4f} sv_min={sv[-1]:.6f}")
+
+# Volume CV of anchors
+def cayley_menger_vol2(pts):
+    pts = pts.float()
+    diff = pts.unsqueeze(-2) - pts.unsqueeze(-3)
+    d2 = (diff * diff).sum(-1)
+    B, V, _ = d2.shape
+    cm = torch.zeros(B, V+1, V+1, device=d2.device, dtype=torch.float32)
+    cm[:, 0, 1:] = 1; cm[:, 1:, 0] = 1; cm[:, 1:, 1:] = d2
+    s = (-1.0)**V; f = math.factorial(V-1)
+    return s / ((2.0**(V-1)) * f*f) * torch.linalg.det(cm)
+
+vols = []
+for _ in range(500):
+    idx = torch.randperm(N_ANCHORS)[:5]
+    v2 = cayley_menger_vol2(anchors[idx].unsqueeze(0))
+    v = torch.sqrt(F.relu(v2[0]) + 1e-12).item()
+    if v > 0: vols.append(v)
+anchor_cv = np.std(vols) / (np.mean(vols) + 1e-8)
+print(f"  Anchor pentachoron CV: {anchor_cv:.4f}")
+print(f"    mean_vol={np.mean(vols):.6f} std_vol={np.std(vols):.6f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# SCAN 2: ANCHOR UTILIZATION
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("SCAN 2: ANCHOR UTILIZATION")
+print(f"{'='*65}")
+
+# How many images use each anchor as nearest
+anchor_counts = torch.bincount(nearest, minlength=N_ANCHORS).float()
+active = (anchor_counts > 0).sum().item()
+print(f"  Active anchors: {active}/{N_ANCHORS} ({active/N_ANCHORS*100:.1f}%)")
+print(f"  Visit counts: mean={anchor_counts.mean():.1f} std={anchor_counts.std():.1f}")
+print(f"    max={anchor_counts.max():.0f} min={anchor_counts.min():.0f}")
+print(f"    top 10: {anchor_counts.topk(10).values.long().tolist()}")
+print(f"    bottom 10: {anchor_counts.sort().values[:10].long().tolist()}")
+
+# Entropy of anchor distribution
+probs = anchor_counts / anchor_counts.sum()
+entropy = -(probs[probs > 0] * probs[probs > 0].log()).sum().item()
+max_entropy = math.log(N_ANCHORS)
+print(f"  Anchor entropy: {entropy:.4f} / {max_entropy:.4f} ({entropy/max_entropy*100:.1f}%)")
+
+# Per-anchor mean cosine to fused embeddings
+print(f"\n  Per-anchor embedding density:")
+anchor_mean_cos = []
+for a_idx in range(N_ANCHORS):
+    mask = nearest == a_idx
+    if mask.sum() < 2:
+        anchor_mean_cos.append(0.0)
+        continue
+    cluster_embs = fused[mask]
+    mean_cos = F.cosine_similarity(
+        cluster_embs.unsqueeze(0), cluster_embs.unsqueeze(1), dim=-1)
+    mean_cos.fill_diagonal_(0)
+    n = cluster_embs.shape[0]
+    avg = mean_cos.sum().item() / max(n * (n-1), 1)
+    anchor_mean_cos.append(avg)
+amc = np.array(anchor_mean_cos)
+print(f"    Intra-cluster cosine: mean={amc[amc>0].mean():.4f} std={amc[amc>0].std():.4f}")
+
+
+# ═══════════════════════════════════════════════════════════���══════
+# SCAN 3: PROJECTOR ANALYSIS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("SCAN 3: PROJECTOR ANALYSIS")
+print(f"{'='*65}")
+
+expert_names = ["clip_l14", "dinov2_b14", "siglip_b16"]
+
+# Per-expert projection stats
+for e, name in enumerate(expert_names):
+    p = proj[e]
+    print(f"\n  {name}:")
+    print(f"    norm: mean={p.norm(dim=-1).mean():.6f} (should be 1.0)")
+    print(f"    self-sim off-diag: {(F.normalize(p,dim=-1) @ F.normalize(p,dim=-1).T).fill_diagonal_(0).mean():.4f}")
+
+    # SVD of projected embeddings
+    pc = p.float() - p.float().mean(0, keepdim=True)
+    sv = torch.linalg.svdvals(pc)
+    eff_dim = ((sv.sum()**2) / (sv.pow(2).sum() + 1e-12)).item()
+    print(f"    eff_dim: {eff_dim:.1f}/{D_ANCHOR}")
+
+# Pairwise agreement
+print(f"\n  Expert agreement (cosine in 128-d):")
+for i in range(3):
+    for j in range(i+1, 3):
+        cos = F.cosine_similarity(proj[i], proj[j], dim=-1)
+        print(f"    {expert_names[i]:<15} × {expert_names[j]:<15}: "
+              f"mean={cos.mean():.4f} std={cos.std():.4f} min={cos.min():.4f}")
+
+# How different are the nearest anchors per expert?
+print(f"\n  Per-expert nearest anchor agreement:")
+expert_nearest = []
+for e in range(3):
+    a = F.normalize(anchors, dim=-1)
+    cos = proj[e] @ a.T
+    en = cos.argmax(dim=-1)
+    expert_nearest.append(en)
+for i in range(3):
+    for j in range(i+1, 3):
+        agree = (expert_nearest[i] == expert_nearest[j]).float().mean().item()
+        print(f"    {expert_names[i]:<15} × {expert_names[j]:<15}: "
+              f"same_anchor={agree:.4f} ({agree*100:.1f}%)")
+
+# Projector weight analysis
+print(f"\n  Projector weight comparison:")
+proj_weights = []
+for e in range(3):
+    w = model.projectors[e].proj[0].weight.detach().float()  # (128, 768)
+    proj_weights.append(w)
+    sv = torch.linalg.svdvals(w)
+    eff_r = ((sv.sum()**2) / (sv.pow(2).sum() + 1e-12)).item()
+    print(f"    {expert_names[e]:<15}: norm={w.norm():.4f} eff_rank={eff_r:.1f}/{min(w.shape)}")
+
+# Cross-projector cosine
+for i in range(3):
+    for j in range(i+1, 3):
+        cos = F.cosine_similarity(
+            proj_weights[i].reshape(-1).unsqueeze(0),
+            proj_weights[j].reshape(-1).unsqueeze(0)).item()
+        print(f"    {expert_names[i]:<15} × {expert_names[j]:<15} weight_cos={cos:.4f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# SCAN 4: PATCHWORK COMPARTMENT ANALYSIS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("SCAN 4: PATCHWORK COMPARTMENTS")
+print(f"{'='*65}")
+
+# Which anchors are in which compartment
+asgn = model.patchwork.asgn.cpu()
+for k in range(N_COMP):
+    anchor_ids = (asgn == k).nonzero(as_tuple=True)[0]
+    print(f"  Comp {k}: {len(anchor_ids)} anchors")
+
+# Patchwork output analysis
+with torch.no_grad():
+    pw_all = []
+    for j in range(0, N_val, BATCH):
+        end = min(j + BATCH, N_val)
+        pw = model.patchwork(tri[j:end].to(DEVICE))
+        pw_all.append(pw.cpu())
+    pw_cat = torch.cat(pw_all)
+
+print(f"\n  Patchwork output: {pw_cat.shape}")
+print(f"    norm: mean={pw_cat.norm(dim=-1).mean():.4f} std={pw_cat.norm(dim=-1).std():.4f}")
+
+# Per-compartment output magnitude
+for k in range(N_COMP):
+    comp_out = pw_cat[:, k*D_COMP:(k+1)*D_COMP]
+    print(f"    comp {k}: norm={comp_out.norm(dim=-1).mean():.4f} "
+          f"std_across_dims={comp_out.std(dim=0).mean():.4f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# SCAN 5: TRIANGULATION PATTERN ANALYSIS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("SCAN 5: TRIANGULATION PATTERNS")
+print(f"{'='*65}")
+
+# Triangulation distance stats
+print(f"  Triangulation distances (1-cosine):")
+print(f"    mean={tri.mean():.4f} std={tri.std():.4f}")
+print(f"    min={tri.min():.4f} max={tri.max():.4f}")
+
+# Nearest anchor distance
+nearest_dist = tri.gather(1, nearest.unsqueeze(1)).squeeze(1)
+print(f"  Nearest anchor distance:")
+print(f"    mean={nearest_dist.mean():.4f} std={nearest_dist.std():.4f}")
+print(f"    max={nearest_dist.max():.4f} min={nearest_dist.min():.4f}")
+
+# How many anchors are "close" (cosine > 0.5, i.e. dist < 0.5)
+close_count = (tri < 0.5).float().sum(dim=1)
+print(f"  Anchors within cos>0.5 per image:")
+print(f"    mean={close_count.mean():.1f} std={close_count.std():.1f}")
+
+# Top-k nearest anchors — how spread are they?
+topk_dists = tri.topk(10, dim=1, largest=False)
+print(f"  Top-10 nearest anchor distances:")
+for k_idx in range(10):
+    d = topk_dists.values[:, k_idx]
+    print(f"    k={k_idx}: mean={d.mean():.4f} std={d.std():.4f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# SCAN 6: PER-CLASS ANCHOR AFFINITY
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("SCAN 6: PER-CLASS ANCHOR AFFINITY")
+print(f"{'='*65}")
+
+# COCO class names (subset)
+coco_names = ["person", "bicycle", "car", "motorcycle", "airplane",
+              "bus", "train", "truck", "boat", "traffic light",
+              "fire hydrant", "stop sign", "parking meter", "bench", "bird",
+              "cat", "dog", "horse", "sheep", "cow"]
+
+# For each class, which anchors are most associated?
+print(f"\n  Top-3 anchors per class (first 20 classes):")
+for c in range(min(20, N_CLASSES)):
+    mask = val_labels[:, c] > 0
+    if mask.sum() < 5: continue
+    class_nearest = nearest[mask]
+    counts = torch.bincount(class_nearest, minlength=N_ANCHORS)
+    top3 = counts.topk(3)
+    name = coco_names[c] if c < len(coco_names) else f"class_{c}"
+    total = mask.sum().item()
+    pcts = [f"{top3.indices[k]}({top3.values[k].item()}/{total})" for k in range(3)]
+    print(f"    {name:<15} (n={total:4d}): {' '.join(pcts)}")
+
+# Anchor specialization: how many classes does each anchor serve?
+anchor_class_count = torch.zeros(N_ANCHORS)
+for a in range(N_ANCHORS):
+    mask = nearest == a
+    if mask.sum() < 1: continue
+    class_present = val_labels[mask].sum(0) > 0
+    anchor_class_count[a] = class_present.sum().item()
+print(f"\n  Anchor specialization:")
+print(f"    classes per anchor: mean={anchor_class_count[anchor_class_count>0].mean():.1f} "
+      f"std={anchor_class_count[anchor_class_count>0].std():.1f}")
+print(f"    max={anchor_class_count.max():.0f} min={anchor_class_count[anchor_class_count>0].min():.0f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# SCAN 7: FUSED EMBEDDING GEOMETRY
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("SCAN 7: FUSED EMBEDDING GEOMETRY")
+print(f"{'='*65}")
+
+# Norms (should be 1.0)
+fused_norms = fused.norm(dim=-1)
+print(f"  Norms: mean={fused_norms.mean():.6f} std={fused_norms.std():.6f}")
+
+# Self-similarity
+fused_n = F.normalize(fused, dim=-1)
+self_sim = fused_n @ fused_n.T
+self_sim_off = (self_sim.sum() - self_sim.diag().sum()) / (N_val**2 - N_val)
+print(f"  Self-sim (off-diag): {self_sim_off:.4f}")
+
+# SVD
+fc = fused.float() - fused.float().mean(0, keepdim=True)
+sv = torch.linalg.svdvals(fc)
+eff_dim = ((sv.sum()**2) / (sv.pow(2).sum() + 1e-12)).item()
+print(f"  Effective dim: {eff_dim:.1f}/{D_ANCHOR}")
+cumvar = sv.pow(2).cumsum(0) / sv.pow(2).sum()
+for k in [5, 10, 20, 50, 100]:
+    if k-1 < len(cumvar):
+        print(f"    top-{k} SVs explain {cumvar[k-1]*100:.1f}%")
+
+# CV
+vols = []
+for _ in range(500):
+    idx = torch.randperm(N_val)[:5]
+    v2 = cayley_menger_vol2(fused_n[idx].unsqueeze(0))
+    v = torch.sqrt(F.relu(v2[0]) + 1e-12).item()
+    if v > 0: vols.append(v)
+fused_cv = np.std(vols) / (np.mean(vols) + 1e-8)
+print(f"  Pentachoron CV: {fused_cv:.4f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# SCAN 8: EXPERT CONTRIBUTION ANALYSIS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'='*65}")
+print("SCAN 8: EXPERT CONTRIBUTION")
+print(f"{'='*65}")
+
+# How much does each expert contribute to the fused embedding?
+# cos(expert_proj, fused) tells us alignment
+for e, name in enumerate(expert_names):
+    cos = F.cosine_similarity(proj[e], fused, dim=-1)
+    print(f"  {name:<15}: cos_to_fused mean={cos.mean():.4f} std={cos.std():.4f}")
+
+# Residual after removing each expert
+for e, name in enumerate(expert_names):
+    others = [proj[i] for i in range(3) if i != e]
+    fused_without = F.normalize(sum(others) / 2, dim=-1)
+    delta = F.cosine_similarity(fused, fused_without, dim=-1)
+    print(f"  Without {name:<15}: cos_to_full={delta.mean():.4f} (uniqueness={1-delta.mean():.4f})")
+
+# Per-image expert disagreement
+print(f"\n  Per-image expert disagreement:")
+all_cos = []
+for i in range(3):
+    for j in range(i+1, 3):
+        cos = F.cosine_similarity(proj[i], proj[j], dim=-1)
+        all_cos.append(cos)
+stacked = torch.stack(all_cos, dim=1)  # (N, 3)
+per_image_agree = stacked.mean(dim=1)
+per_image_disagree = stacked.std(dim=1)
+print(f"  Agreement: mean={per_image_agree.mean():.4f} std={per_image_agree.std():.4f}")
+print(f"  Disagreement: mean={per_image_disagree.mean():.4f} std={per_image_disagree.std():.4f}")
+
+# Most agreed and disagreed images
+most_agree_idx = per_image_agree.argmax().item()
+most_disagree_idx = per_image_agree.argmin().item()
+print(f"\n  Most agreed image ({most_agree_idx}): agreement={per_image_agree[most_agree_idx]:.4f}")
+print(f"    labels: {val_labels[most_agree_idx].nonzero(as_tuple=True)[0].tolist()}")
+print(f"  Most disagreed image ({most_disagree_idx}): agreement={per_image_agree[most_disagree_idx]:.4f}")
+print(f"    labels: {val_labels[most_disagree_idx].nonzero(as_tuple=True)[0].tolist()}")
+
+
+print(f"\n{'='*65}")
+print("ANALYSIS COMPLETE")
+print(f"{'='*65}")