Create colab_cv_sweep_batched.py

colab_cv_sweep_batched.py

@@ -0,0 +1,456 @@
+"""
+CV Loss Sweep — Pure Noise Prediction
+========================================
+Random inputs → MLP encoder → S^(d-1) → constellation → predict 10 random labels.
+
+No dataset. No structure. No signal. The model memorizes random noise→label
+mappings. Any geometric regularity (CV convergence) is purely from:
+  - The unit hypersphere S^(d-1)
+  - The smooth optimizer (AdamW)
+  - The CV loss pressure (or lack thereof)
+
+If CV ≈ 0.20 with zero CV loss on pure noise, the constant is the sphere's
+property, not a training artifact and not a data property.
+
+Each run: 200 steps, ~2 seconds. Full sweep: ~1 minute.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import time
+import json
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+
+# ═══════════════════════════════════════════════════════════════
+# Noise Dataset — pure random, zero structure
+# ═══════════════════════════════════════════════════════════════
+
+class NoiseDataset(torch.utils.data.Dataset):
+    """Random Gaussian inputs with random labels. No signal."""
+    def __init__(self, n_samples=5000, input_dim=128, num_classes=10, seed=0):
+        torch.manual_seed(seed)
+        self.data = torch.randn(n_samples, input_dim)
+        self.labels = torch.randint(0, num_classes, (n_samples,))
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx], self.labels[idx]
+
+
+# ═══════════════════════════════════════════════════════════════
+# Minimal MLP Encoder → S^(d-1)
+# ═══════════════════════════════════════════════════════════════
+
+class NoiseEncoder(nn.Module):
+    """MLP → sphere. No convolutions, no structure."""
+    def __init__(self, input_dim=128, hidden_dim=256, output_dim=128):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, output_dim),
+            nn.LayerNorm(output_dim),
+        )
+
+    def forward(self, x):
+        return F.normalize(self.net(x), dim=-1)
+
+
+# ═══════════════════════════════════════════════════════════════
+# Minimal Constellation + Classifier
+# ═══════════════════════════════════════════════════════════════
+
+class NoiseConstellation(nn.Module):
+    """Minimal: anchors + patchwork + classifier. No bridge, no push, no magnitude."""
+    def __init__(self, dim=128, n_anchors=64, n_comp=8, num_classes=10):
+        super().__init__()
+        self.n_anchors = n_anchors
+        self.n_comp = n_comp
+
+        anchors = F.normalize(torch.randn(n_anchors, dim), dim=-1)
+        self.anchors = nn.Parameter(anchors)
+
+        apc = n_anchors // n_comp
+        self.patchwork = nn.ModuleList([
+            nn.Sequential(nn.Linear(apc, 64), nn.GELU(), nn.Linear(64, 64))
+            for _ in range(n_comp)
+        ])
+        self.classifier = nn.Linear(n_comp * 64 + dim, num_classes)
+
+    def forward(self, emb):
+        anchors_n = F.normalize(self.anchors, dim=-1)
+        tri = emb @ anchors_n.T
+        apc = self.n_anchors // self.n_comp
+        pw_parts = []
+        for k in range(self.n_comp):
+            pw_parts.append(self.patchwork[k](tri[:, k*apc:(k+1)*apc]))
+        pw = torch.cat(pw_parts, dim=-1)
+        logits = self.classifier(torch.cat([pw, emb], dim=-1))
+        return logits, emb
+
+
+# ═══════════════════════════════════════════════════════════════
+# CV Computation
+# ═══════════════════════════════════════════════════════════════
+
+def _batch_volumes(emb, n_samples=200, n_points=5):
+    """Batched pentachoron volumes — zero Python loops."""
+    N, D = emb.shape
+    device, dtype = emb.device, emb.dtype
+    pool = min(N, 512)
+
+    # Batched randperm via argsort on random values
+    rand_keys = torch.rand(n_samples, pool, device=device)
+    indices = rand_keys.argsort(dim=1)[:, :n_points]  # (n_samples, n_points)
+
+    # Gather: (n_samples, n_points, D)
+    pts = emb[:pool][indices]
+
+    # Gram → squared distances: all batched
+    gram = torch.bmm(pts, pts.transpose(1, 2))
+    norms = torch.diagonal(gram, dim1=1, dim2=2)
+    d2 = F.relu(norms.unsqueeze(2) + norms.unsqueeze(1) - 2 * gram)
+
+    # Cayley-Menger: (n_samples, n_points+1, n_points+1)
+    M = n_points + 1
+    cm = torch.zeros(n_samples, M, M, device=device, dtype=dtype)
+    cm[:, 0, 1:] = 1.0
+    cm[:, 1:, 0] = 1.0
+    cm[:, 1:, 1:] = d2
+
+    k = n_points - 1
+    pf = ((-1.0) ** (k + 1)) / ((2.0 ** k) * (math.factorial(k) ** 2))
+
+    # Single batched det call
+    dets = pf * torch.linalg.det(cm.float())
+    valid = dets > 1e-20
+    return dets[valid].to(dtype).sqrt()
+
+
+def cv_loss(emb, target=0.22, n_samples=32, n_points=5):
+    """Differentiable CV loss — batched."""
+    if emb.shape[0] < n_points:
+        return torch.tensor(0.0, device=emb.device, requires_grad=True)
+    vols = _batch_volumes(emb, n_samples=n_samples, n_points=n_points)
+    if vols.shape[0] < 5:
+        return torch.tensor(0.0, device=emb.device, requires_grad=True)
+    cv = vols.std() / (vols.mean() + 1e-8)
+    return (cv - target).pow(2)
+
+
+def cv_metric(emb, n_samples=200, n_points=5):
+    """Non-differentiable CV — batched."""
+    with torch.no_grad():
+        vols = _batch_volumes(emb, n_samples=n_samples, n_points=n_points)
+        if vols.shape[0] < 10:
+            return 0.0
+        return (vols.std() / (vols.mean() + 1e-8)).item()
+
+
+# ═══════════════════════════════════════════════════════════════
+# Single Run
+# ═══════════════════════════════════════════════════════════════
+
+def run_experiment(cv_weight, cv_target, n_steps=200, dim=128, n_anchors=64,
+                   batch_size=256, n_samples=5000, seed=42, pure_cv=False):
+    """One configuration. Returns results dict. ~2 seconds."""
+    torch.manual_seed(seed)
+
+    ds = NoiseDataset(n_samples=n_samples, input_dim=dim, num_classes=10, seed=seed + 1000)
+    loader = torch.utils.data.DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=True)
+
+    encoder = NoiseEncoder(input_dim=dim, hidden_dim=256, output_dim=dim).to(DEVICE)
+    constellation = NoiseConstellation(dim=dim, n_anchors=n_anchors).to(DEVICE)
+
+    params = list(encoder.parameters()) + list(constellation.parameters())
+    optimizer = torch.optim.AdamW(params, lr=0.001, weight_decay=0.05)
+
+    step = 0
+    cv_history = []
+    ce_history = []
+    acc_history = []
+
+    while step < n_steps:
+        for data, labels in loader:
+            if step >= n_steps:
+                break
+            data, labels = data.to(DEVICE), labels.to(DEVICE)
+            emb = encoder(data)
+            logits, _ = constellation(emb)
+
+            l_ce = F.cross_entropy(logits, labels)
+
+            if cv_weight > 0:
+                l_cv = cv_loss(emb, target=cv_target, n_samples=32)
+            else:
+                l_cv = torch.tensor(0.0, device=DEVICE)
+
+            if pure_cv:
+                loss = cv_weight * l_cv  # NO CE, pure geometric pressure
+            else:
+                loss = l_ce + cv_weight * l_cv
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(params, 1.0)
+            optimizer.step()
+
+            acc = (logits.argmax(-1) == labels).float().mean().item()
+
+            # Measure CV every 50 steps
+            if step % 50 == 0 or step == n_steps - 1:
+                with torch.no_grad():
+                    # Collect embeddings for CV measurement
+                    all_emb = []
+                    for d_batch, _ in loader:
+                        all_emb.append(encoder(d_batch.to(DEVICE)))
+                        if len(all_emb) * batch_size >= 2000:
+                            break
+                    all_emb = torch.cat(all_emb)[:2000]
+                    v_cv = cv_metric(all_emb, n_samples=200)
+
+                    # Effective dim
+                    centered = all_emb[:1000] - all_emb[:1000].mean(0)
+                    s = torch.linalg.svdvals(centered.float())
+                    s_n = s / s.sum()
+                    eff_dim = (1.0 / (s_n ** 2).sum()).item()
+
+                cv_history.append({'step': step, 'cv': round(v_cv, 4), 'eff_dim': round(eff_dim, 1)})
+
+            ce_history.append(l_ce.item())
+            acc_history.append(acc)
+            step += 1
+
+    # Final measurement
+    with torch.no_grad():
+        all_emb = []
+        for d_batch, _ in loader:
+            all_emb.append(encoder(d_batch.to(DEVICE)))
+            if len(all_emb) * batch_size >= 2000:
+                break
+        all_emb = torch.cat(all_emb)[:2000]
+        final_cv = cv_metric(all_emb, n_samples=300)
+        centered = all_emb[:1000] - all_emb[:1000].mean(0)
+        s = torch.linalg.svdvals(centered.float())
+        s_n = s / s.sum()
+        final_dim = (1.0 / (s_n ** 2).sum()).item()
+
+    return {
+        'cv_weight': cv_weight,
+        'cv_target': cv_target,
+        'pure_cv': pure_cv,
+        'seed': seed,
+        'n_steps': n_steps,
+        'dim': dim,
+        'final_cv': round(final_cv, 4),
+        'final_dim': round(final_dim, 1),
+        'final_ce': round(sum(ce_history[-20:]) / 20, 4),
+        'final_acc': round(sum(acc_history[-20:]) / 20 * 100, 1),
+        'cv_trajectory': cv_history,
+    }
+
+
+# ═══════════════════════════════════════════════════════════════
+# SWEEP
+# ═══════════════════════════════════════════════════════════════
+
+print("=" * 80)
+print("CV LOSS SWEEP — PURE NOISE PREDICTION")
+print("  Random inputs → MLP → S^(d-1) → constellation → 10 random labels")
+print("  No data structure. No signal. Pure sphere geometry + optimizer.")
+print("  200 steps per run, ~2s each.")
+print("=" * 80)
+
+# (cv_weight, cv_target, label, seed)
+configs = [
+    # ── NO CV LOSS — baseline ──
+    (0.0,    0.0,  "no_cv",           42),
+    (0.0,    0.0,  "no_cv_s2",        123),
+    (0.0,    0.0,  "no_cv_s3",        456),
+    (0.0,    0.0,  "no_cv_s4",        789),
+    (0.0,    0.0,  "no_cv_s5",        1337),
+
+    # ── CORRECT TARGET, VARYING WEIGHT ──
+    (0.001,  0.22, "w0.001_t0.22",    42),
+    (0.01,   0.22, "w0.01_t0.22",     42),
+    (0.1,    0.22, "w0.1_t0.22",      42),
+    (0.5,    0.22, "w0.5_t0.22",      42),
+    (1.0,    0.22, "w1.0_t0.22",      42),
+    (5.0,    0.22, "w5.0_t0.22",      42),
+    (10.0,   0.22, "w10_t0.22",       42),
+    (50.0,   0.22, "w50_t0.22",       42),
+    (100.0,  0.22, "w100_t0.22",      42),
+
+    # ── WRONG TARGETS, LOW WEIGHT (gentle push) ──
+    (0.01,   0.00, "w0.01_t0.00",     42),
+    (0.01,   0.05, "w0.01_t0.05",     42),
+    (0.01,   0.10, "w0.01_t0.10",     42),
+    (0.01,   0.30, "w0.01_t0.30",     42),
+    (0.01,   0.50, "w0.01_t0.50",     42),
+    (0.01,   0.80, "w0.01_t0.80",     42),
+    (0.01,   1.00, "w0.01_t1.00",     42),
+    (0.01,   2.00, "w0.01_t2.00",     42),
+
+    # ── WRONG TARGETS, MEDIUM WEIGHT (strong push) ──
+    (1.0,    0.00, "w1_t0.00",        42),
+    (1.0,    0.05, "w1_t0.05",        42),
+    (1.0,    0.50, "w1_t0.50",        42),
+    (1.0,    0.80, "w1_t0.80",        42),
+    (1.0,    1.00, "w1_t1.00",        42),
+
+    # ── WRONG TARGETS, EXTREME WEIGHT (maximum force) ──
+    (100.0,  0.00, "w100_t0.00",      42),
+    (100.0,  0.05, "w100_t0.05",      42),
+    (100.0,  0.10, "w100_t0.10",      42),
+    (100.0,  0.50, "w100_t0.50",      42),
+    (100.0,  0.80, "w100_t0.80",      42),
+    (100.0,  1.00, "w100_t1.00",      42),
+
+    # ── CV LOSS ONLY, NO CE (pure geometric pressure) ──
+    (1.0,    0.22, "pure_cv_t0.22",   42),  # mark for CE override
+    (1.0,    0.05, "pure_cv_t0.05",   42),
+    (1.0,    0.50, "pure_cv_t0.50",   42),
+    (1.0,    0.80, "pure_cv_t0.80",   42),
+    (1.0,    1.00, "pure_cv_t1.00",   42),
+
+    # ── DIMENSION SWEEP (does dim change the constant?) ──
+    (0.0,    0.0,  "dim16",           42),  # mark for dim override
+    (0.0,    0.0,  "dim32",           42),
+    (0.0,    0.0,  "dim64",           42),
+    (0.0,    0.0,  "dim256",          42),
+    (0.0,    0.0,  "dim512",          42),
+]
+
+# Special handling
+pure_cv_labels = {l for _, _, l, _ in configs if l.startswith("pure_cv")}
+dim_overrides = {"dim16": 16, "dim32": 32, "dim64": 64, "dim256": 256, "dim512": 512}
+
+all_results = []
+total = len(configs)
+
+print(f"\n  Running {total} configurations, 200 steps each")
+print(f"  Estimated time: ~{total * 3}s\n")
+
+for i, (cv_w, cv_t, label, seed) in enumerate(configs):
+    t0 = time.time()
+    dim = dim_overrides.get(label, 128)
+    is_pure_cv = label in pure_cv_labels
+
+    print(f"[{i+1:2d}/{total}] {label:20s}  w={cv_w:<8.3f} t={cv_t:<5.2f} d={dim:<4d}", end=" ", flush=True)
+
+    result = run_experiment(
+        cv_weight=cv_w,
+        cv_target=cv_t,
+        n_steps=200,
+        dim=dim,
+        seed=seed,
+        pure_cv=is_pure_cv,
+    )
+    result['label'] = label
+
+    elapsed = time.time() - t0
+    print(f"→ CV={result['final_cv']:.4f} dim={result['final_dim']:.0f} "
+          f"acc={result['final_acc']:.0f}% ({elapsed:.1f}s)")
+
+    all_results.append(result)
+
+
+# ═══════════════════════════════════════════════════════════════
+# SUMMARY TABLE
+# ═════════════════════════════════════════��═════════════════════
+
+print(f"\n\n{'='*90}")
+print(f"{'LABEL':20s} {'CV_W':>8s} {'CV_T':>6s} {'DIM':>5s} {'FINAL_CV':>9s} {'EFF_DIM':>8s} {'ACC%':>6s} {'CE':>8s}")
+print(f"{'─'*90}")
+
+for r in all_results:
+    cv_mark = "✓" if 0.17 <= r['final_cv'] <= 0.24 else "~" if 0.15 <= r['final_cv'] <= 0.27 else "✗"
+    print(f"{r['label']:20s} {r['cv_weight']:>8.3f} {r['cv_target']:>6.2f} {r['dim']:>5d} "
+          f"{r['final_cv']:>8.4f}{cv_mark} {r['final_dim']:>7.0f} {r['final_acc']:>5.0f}% {r['final_ce']:>8.4f}")
+
+
+# ═══════════════════════════════════════════════════════════════
+# ANALYSIS
+# ═══════════════════════════════════════════════════════════════
+
+print(f"\n\n{'='*90}")
+print("ANALYSIS")
+print(f"{'='*90}")
+
+# 1. Baseline: no CV loss
+no_cv = [r for r in all_results if r['cv_weight'] == 0 and r['dim'] == 128]
+if no_cv:
+    cvs = [r['final_cv'] for r in no_cv]
+    dims = [r['final_dim'] for r in no_cv]
+    print(f"\n  [1] NO CV LOSS, PURE NOISE (d=128, {len(no_cv)} seeds):")
+    print(f"      CV:  mean={sum(cvs)/len(cvs):.4f}  min={min(cvs):.4f}  max={max(cvs):.4f}  spread={max(cvs)-min(cvs):.4f}")
+    print(f"      Dim: mean={sum(dims)/len(dims):.1f}")
+    within_band = sum(1 for c in cvs if 0.17 <= c <= 0.24)
+    print(f"      Within [0.17, 0.24]: {within_band}/{len(cvs)}")
+
+# 2. Weight sweep at correct target
+weight_sweep = [r for r in all_results if r['cv_target'] == 0.22 and r['dim'] == 128 and not r['pure_cv']]
+if weight_sweep:
+    print(f"\n  [2] WEIGHT SWEEP (target=0.22, d=128):")
+    for r in sorted(weight_sweep, key=lambda x: x['cv_weight']):
+        print(f"      w={r['cv_weight']:>8.3f} → CV={r['final_cv']:.4f}  acc={r['final_acc']:.0f}%")
+
+# 3. Target sweep at fixed weight
+for w in [0.01, 1.0, 100.0]:
+    target_runs = [r for r in all_results if r['cv_weight'] == w and r['dim'] == 128 and not r['pure_cv']]
+    if len(target_runs) > 2:
+        print(f"\n  [3] TARGET SWEEP (w={w}, d=128):")
+        for r in sorted(target_runs, key=lambda x: x['cv_target']):
+            cv_mark = "✓" if 0.17 <= r['final_cv'] <= 0.24 else "✗"
+            print(f"      target={r['cv_target']:.2f} → CV={r['final_cv']:.4f}{cv_mark}  acc={r['final_acc']:.0f}%")
+
+# 4. Dimension sweep
+dim_runs = [r for r in all_results if r['label'].startswith('dim')]
+if dim_runs:
+    print(f"\n  [4] DIMENSION SWEEP (no CV loss):")
+    for r in sorted(dim_runs, key=lambda x: x['dim']):
+        print(f"      d={r['dim']:>4d} → CV={r['final_cv']:.4f}  eff_dim={r['final_dim']:.0f}")
+
+# 5. Key question: can extreme weight move CV?
+extreme = [r for r in all_results if r['cv_weight'] >= 100 and r['dim'] == 128]
+if extreme:
+    print(f"\n  [5] EXTREME FORCE (w≥100, d=128):")
+    for r in sorted(extreme, key=lambda x: x['cv_target']):
+        delta = abs(r['final_cv'] - 0.20)
+        print(f"      target={r['cv_target']:.2f} → CV={r['final_cv']:.4f}  (Δ from 0.20: {delta:.4f})  acc={r['final_acc']:.0f}%")
+
+# 5b. Pure CV — no CE, only geometric pressure
+pure_runs = [r for r in all_results if r['pure_cv']]
+if pure_runs:
+    print(f"\n  [5b] PURE CV (no CE loss, only geometric pressure):")
+    for r in sorted(pure_runs, key=lambda x: x['cv_target']):
+        delta = abs(r['final_cv'] - 0.20)
+        print(f"      target={r['cv_target']:.2f} → CV={r['final_cv']:.4f}  (Δ from 0.20: {delta:.4f})  dim={r['final_dim']:.0f}")
+
+# 6. CV trajectory analysis — does it start elsewhere and converge?
+print(f"\n  [6] CV TRAJECTORIES (step 0 → step 200):")
+for r in all_results[:5]:  # first 5 runs
+    traj = r.get('cv_trajectory', [])
+    if len(traj) >= 2:
+        first = traj[0]['cv']
+        last = traj[-1]['cv']
+        print(f"      {r['label']:20s}: {first:.4f} → {last:.4f} (Δ={last-first:+.4f})")
+
+# Save
+with open('cv_sweep_results.json', 'w') as f:
+    json.dump(all_results, f, indent=2, default=str)
+print(f"\n  Raw results saved to cv_sweep_results.json")
+
+print(f"\n{'='*80}")
+print("CV SWEEP COMPLETE")
+print(f"{'='*80}")