Update trainer.py

trainer.py

@@ -1,10 +1,12 @@
 #!/usr/bin/env python3
 """
-CIFAR-10 — Tri-Stream GeoLIP ViT — Experiment 7
-==================================================
-Stream A (CE), Stream B (BCE), GAL (geometric arbitration).
-GAL anchors updated via Procrustes every M batches.
-Whitened Procrustes toggleable.
+CIFAR-10 — Tri-Stream GeoLIP ViT v8
+=====================================
+v7→v8 changes:
+  1. GAL_UPDATE_INTERVAL: 50 → 25 (2× more frequent)
+  2. GAL_LR: 0.01 → 0.015 (+50% response)
+  3. Tracks nce_b and geo_nce_acc separately
+  4. stream_b_nce_weight=0.5, geo_nce_weight=0.5
 """
 
 import torch
@@ -29,8 +31,8 @@ STREAM_DIM = 192
 N_BLOCKS = 9
 N_HEADS = 8
 OUTPUT_DIM = 256
-N_ANCHORS = 128          # constellation (sphere space)
-N_GAL_ANCHORS = 64       # GAL anchors (stream space)
+N_ANCHORS = 128
+N_GAL_ANCHORS = 64
 N_COMP = 16
 D_COMP = 128
 ANCHOR_DROP = 0.10
@@ -47,11 +49,15 @@ BCE_WEIGHT = 1.0
 CM_WEIGHT = 0.1
 INFONCE_TEMP = 0.07
 
-# ── GAL ──
-GAL_UPDATE_INTERVAL = 50   # batches between Procrustes updates
-GAL_LR = 0.01             # step size for anchor rotation
+# ── v8: Stream B + Geo NCE weights ──
+STREAM_B_NCE_WEIGHT = 0.5
+GEO_NCE_WEIGHT = 0.5
+
+# ── v8: GAL — faster updates, stronger response ──
+GAL_UPDATE_INTERVAL = 25    # was 50
+GAL_LR = 0.015              # was 0.01 (+50%)
 GAL_BUFFER_SIZE = 50000
-USE_WHITENED_PROCRUSTES = False  # toggle for benchmarking
+USE_WHITENED_PROCRUSTES = False
 
 # ── Mastery queue ──
 MASTERY_PATIENCE = 50
@@ -71,18 +77,17 @@ EPOCHS = 100
 LR = 3e-4
 WARMUP = 5
 GRAD_CLIP = 1.0
-V1_CKPT = ""
+V1_CKPT = ""  # set to checkpoint path for warm start
 
 print("=" * 60)
-print("CIFAR-10 — Tri-Stream GeoLIP ViT — EXP 7 (GAL)")
+print("CIFAR-10 — Tri-Stream GeoLIP ViT v8")
 print(f"  Architecture: {N_BLOCKS}× TriStreamBlock")
 print(f"  Sphere: {OUTPUT_DIM}-d, {N_ANCHORS} anchors, {N_COMP}×{D_COMP} pw")
 print(f"  GAL: {N_GAL_ANCHORS} anchors, Procrustes every {GAL_UPDATE_INTERVAL} "
       f"batches (lr={GAL_LR}, whiten={USE_WHITENED_PROCRUSTES})")
-print(f"  GAL buffer: {GAL_BUFFER_SIZE}")
-print(f"  InfoNCE={INFONCE_WEIGHT} on emb+geo")
-print(f"  CE(stream A) + BCE(stream B) + BCE(geo)")
-print(f"  LS={LABEL_SMOOTHING}, CV={CV_WEIGHT}")
+print(f"  v8 fixes: uniform hypersphere init, gate_init=1/(2×{N_BLOCKS})")
+print(f"  v8 fixes: InfoNCE on emb_b (w={STREAM_B_NCE_WEIGHT}) "
+      f"+ geo_emb (w={GEO_NCE_WEIGHT})")
 print(f"  Device: {DEVICE}")
 print("=" * 60)
 
@@ -145,12 +150,17 @@ model = create_tri_stream_vit(
     autograd_tang=AUTOGRAD_TANG, autograd_sep=AUTOGRAD_SEP,
     enable_autograd=ENABLE_AUTOGRAD,
     label_smoothing=LABEL_SMOOTHING,
+    stream_b_nce_weight=STREAM_B_NCE_WEIGHT,
+    geo_nce_weight=GEO_NCE_WEIGHT,
 ).to(DEVICE)
 
 if V1_CKPT and os.path.exists(V1_CKPT):
     ckpt = torch.load(V1_CKPT, map_location="cpu", weights_only=False)
-    model.load_state_dict(ckpt["state_dict"], strict=False)
-    print(f"  ✓ Loaded weights: epoch {ckpt['epoch']}")
+    missing, unexpected = model.load_state_dict(
+        ckpt["state_dict"], strict=False)
+    print(f"  ✓ Loaded weights: epoch {ckpt.get('epoch', '?')}")
+    if missing:
+        print(f"    New params (expected): {len(missing)}")
 else:
     print(f"  Training from scratch")
 
@@ -165,9 +175,6 @@ print(f"\n{'='*60}")
 print(f"TRAINING — {EPOCHS} epochs, lr={LR}, batch={BATCH}")
 print(f"  GAL Procrustes: every {GAL_UPDATE_INTERVAL} batches, "
       f"lr={GAL_LR}, whiten={USE_WHITENED_PROCRUSTES}")
-print(f"  Mastery: patience={MASTERY_PATIENCE}, queue adaptive "
-      f"[{MASTERY_MIN_SIZE}–{MASTERY_MAX_SIZE}]")
-print(f"  Optimizer: Adam")
 print(f"{'='*60}")
 
 optimizer = torch.optim.Adam(model.parameters(), lr=LR)
@@ -184,7 +191,7 @@ scheduler = torch.optim.lr_scheduler.SequentialLR(
 
 scaler = torch.amp.GradScaler("cuda")
 os.makedirs("checkpoints", exist_ok=True)
-writer = SummaryWriter("runs/cifar10_tri_stream_v7_gal")
+writer = SummaryWriter("runs/cifar10_tri_stream_v8")
 best_acc = 0.0
 gs = 0
 
@@ -201,7 +208,7 @@ mastery = MasteryQueue(
 simplex_buf = SimplexBuffer(
     dim=STREAM_DIM, max_size=GAL_BUFFER_SIZE, device=DEVICE)
 
-gal_update_count = 0  # track Procrustes updates
+gal_update_count = 0
 
 # ══════════════════════════════════════════════════════════════════
 # TRAINING LOOP
@@ -214,7 +221,9 @@ for epoch in range(EPOCHS):
     acc_dict = {
         "loss": 0, "ce": 0, "bce": 0, "geo_bce": 0,
         "acc_a": 0, "acc_b": 0, "geo_acc": 0,
-        "nce": 0, "nce_acc": 0, "geo_nce": 0, "geo_nce_acc": 0,
+        "nce": 0, "nce_acc": 0,
+        "nce_b": 0, "nce_b_acc": 0,
+        "geo_nce": 0, "geo_nce_acc": 0,
         "cm": 0, "cm_valid": 0, "cv": 0, "cv_main": 0, "cv_geo": 0,
         "spread": 0, "mastery": 0, "hard_neg": 0, "hard_pos": 0,
         "correct": 0, "total": 0, "n": 0}
@@ -240,31 +249,28 @@ for epoch in range(EPOCHS):
         scaler.step(optimizer); scaler.update()
         scheduler.step()
 
-        # ── Mastery activation check ──
         mastery.check_activation(ld.get('nce_acc', 0))
 
-        # ── Accumulate geo features into simplex buffer ──
         pool_geo = out1.get('pool_geo')
         if pool_geo is not None:
             simplex_buf.push(pool_geo.float(), targets)
 
-        # ── Periodic GAL Procrustes update ──
         gs += 1
-        if gs % GAL_UPDATE_INTERVAL == 0 and simplex_buf.size > 1000:
+        if gs % GAL_UPDATE_INTERVAL == 0 and simplex_buf.size > 500:
             score = model.update_gal_anchors(
                 simplex_buf, lr=GAL_LR, whiten=USE_WHITENED_PROCRUSTES)
             if score is not None:
                 gal_update_count += 1
                 writer.add_scalar("step/procrustes_score", score, gs)
 
-        # ── Track metrics ──
+        # Track
         preds = out1['logits_a'].argmax(-1)
         correct = (preds == targets).sum().item()
         acc_dict["correct"] += correct
         acc_dict["total"] += targets.shape[0]
         acc_dict["loss"] += loss.item()
 
-        for k in ["ce", "bce", "geo_bce", "nce", "geo_nce",
+        for k in ["ce", "bce", "geo_bce", "nce", "nce_b", "geo_nce",
                    "cm", "cv", "spread", "mastery"]:
             v = ld.get(k, 0)
             acc_dict[k] += v.item() if torch.is_tensor(v) else v
@@ -273,6 +279,7 @@ for epoch in range(EPOCHS):
         acc_dict["acc_b"] += ld.get("acc_b", 0)
         acc_dict["geo_acc"] += ld.get("geo_acc", 0)
         acc_dict["nce_acc"] += ld.get("nce_acc", 0)
+        acc_dict["nce_b_acc"] += ld.get("nce_b_acc", 0)
         acc_dict["geo_nce_acc"] += ld.get("geo_nce_acc", 0)
         acc_dict["cm_valid"] += ld.get("cm_valid", 0)
         acc_dict["cv_main"] += ld.get("cv_main", 0)
@@ -285,22 +292,22 @@ for epoch in range(EPOCHS):
             d = acc_dict["n"]
             ta = 100 * acc_dict["correct"] / acc_dict["total"]
             ga = 100 * acc_dict["geo_acc"] / d
-            mst = acc_dict["mastery"] / d
+            nb = acc_dict["nce_b_acc"] / d
             stg = "M" if mastery.active else "S1"
             pbar.set_postfix(
                 loss=f"{acc_dict['loss']/d:.4f}",
                 a=f"{ta:.0f}%",
                 ga=f"{ga:.0f}%",
-                mst=f"{mst:.3f}",
+                nb=f"{nb:.2f}",
                 stg=stg,
                 gal=gal_update_count,
                 ordered=True)
 
-        # Step-level TB
         if gs % 20 == 0:
             writer.add_scalar("step/loss", loss.item(), gs)
             writer.add_scalar("step/geo_acc", ld.get("geo_acc", 0), gs)
-            # Log gate values
+            writer.add_scalar("step/nce_b_acc", ld.get("nce_b_acc", 0), gs)
+            writer.add_scalar("step/geo_nce_acc", ld.get("geo_nce_acc", 0), gs)
             gates_a = out1.get('gates_a', [])
             if gates_a:
                 writer.add_scalar("step/gate_a_mean",
@@ -319,19 +326,18 @@ for epoch in range(EPOCHS):
     writer.add_scalar("epoch/acc_b", 100 * acc_dict["acc_b"] / d, epoch + 1)
     writer.add_scalar("epoch/geo_acc", 100 * acc_dict["geo_acc"] / d, epoch + 1)
     writer.add_scalar("epoch/nce_acc", acc_dict["nce_acc"] / d, epoch + 1)
+    writer.add_scalar("epoch/nce_b_acc", acc_dict["nce_b_acc"] / d, epoch + 1)
     writer.add_scalar("epoch/geo_nce_acc", acc_dict["geo_nce_acc"] / d, epoch + 1)
     writer.add_scalar("epoch/cv_main", acc_dict["cv_main"] / d, epoch + 1)
     writer.add_scalar("epoch/cv_geo", acc_dict["cv_geo"] / d, epoch + 1)
     writer.add_scalar("epoch/cm_valid", acc_dict["cm_valid"] / d, epoch + 1)
-    writer.add_scalar("epoch/margin", mastery.current_margin, epoch + 1)
-    writer.add_scalar("epoch/queue_max", mastery._current_max, epoch + 1)
-    writer.add_scalar("epoch/simplex_buf", simplex_buf.size, epoch + 1)
     writer.add_scalar("epoch/gal_updates", gal_update_count, epoch + 1)
 
     # ── Validation ──
     model.eval()
     val_correct, val_total, val_loss_sum, val_n = 0, 0, 0, 0
     val_geo_correct = 0
+    val_b_correct = 0
     all_embs = []
 
     with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
@@ -341,8 +347,8 @@ for epoch in range(EPOCHS):
             out = model(images, apply_autograd=False)
             preds = out['logits_a'].argmax(dim=-1)
             val_correct += (preds == labels_v).sum().item()
-            geo_preds = out['geo_logits'].argmax(dim=-1)
-            val_geo_correct += (geo_preds == labels_v).sum().item()
+            val_b_correct += (out['logits_b'].argmax(-1) == labels_v).sum().item()
+            val_geo_correct += (out['geo_logits'].argmax(-1) == labels_v).sum().item()
             val_total += labels_v.shape[0]
             loss_v = F.cross_entropy(out['logits_a'], labels_v)
             val_loss_sum += loss_v.item()
@@ -350,10 +356,11 @@ for epoch in range(EPOCHS):
             all_embs.append(out['embedding'].float().cpu())
 
     val_acc = 100 * val_correct / val_total
+    val_b_acc = 100 * val_b_correct / val_total
     val_geo_acc = 100 * val_geo_correct / val_total
     val_loss = val_loss_sum / max(val_n, 1)
 
-    # CV on val embeddings
+    # ── Val embedding diagnostics ──
     embs = torch.cat(all_embs)
     with torch.no_grad():
         sample = embs[:2000].to(DEVICE)
@@ -372,18 +379,17 @@ for epoch in range(EPOCHS):
                 vols.append(v2[0].sqrt())
         v_cv = (torch.stack(vols).std() / (torch.stack(vols).mean() + 1e-8)).item() if len(vols) > 10 else 0.0
 
-    # Anchor utilization
     with torch.no_grad():
         _, v_np = model.constellation.triangulate(
             embs[:2000].to(DEVICE), training=False)
         n_active = v_np.cpu().unique().numel()
 
     writer.add_scalar("epoch/val_acc", val_acc, epoch + 1)
+    writer.add_scalar("epoch/val_b_acc", val_b_acc, epoch + 1)
     writer.add_scalar("epoch/val_geo_acc", val_geo_acc, epoch + 1)
     writer.add_scalar("epoch/val_cv", v_cv, epoch + 1)
     writer.add_scalar("epoch/val_anchors", n_active, epoch + 1)
 
-    # ── Adaptive queue resize ──
     mastery.update_size(train_acc, val_acc, epoch + 1)
 
     # ── Checkpoint ──
@@ -395,10 +401,11 @@ for epoch in range(EPOCHS):
             "config": model.config,
             "epoch": epoch + 1,
             "val_acc": val_acc,
+            "val_b_acc": val_b_acc,
             "val_geo_acc": val_geo_acc,
             "mastery": mastery.state_dict(),
             "gal_updates": gal_update_count,
-        }, "checkpoints/tri_stream_v7_best.pt")
+        }, "checkpoints/tri_stream_v8_best.pt")
         mk = " ★"
 
     if (epoch + 1) % 10 == 0:
@@ -408,39 +415,37 @@ for epoch in range(EPOCHS):
             "epoch": epoch + 1,
             "val_acc": val_acc,
             "optimizer": optimizer.state_dict(),
-        }, f"checkpoints/tri_stream_v7_e{epoch+1:03d}.pt")
+        }, f"checkpoints/tri_stream_v8_e{epoch+1:03d}.pt")
 
-    # ── Epoch print ──
+    # ── Epoch print — v8: shows B acc + nce_b + geo_nce ──
     ga = 100 * acc_dict["geo_acc"] / d
     ab = 100 * acc_dict["acc_b"] / d
-    mst_m = acc_dict["mastery"] / d
-    hn = acc_dict["hard_neg"] / d if mastery.active else 0
-    hp = acc_dict["hard_pos"] / d if mastery.active else 0
+    nb_acc = acc_dict["nce_b_acc"] / d
+    gn_acc = acc_dict["geo_nce_acc"] / d
     cvf = acc_dict["cv_main"] / d
     cvg = acc_dict["cv_geo"] / d
     cmv = acc_dict["cm_valid"] / d
     stage = "MASTERY" if mastery.active else "stage1"
 
-    # Mean gate values
-    last_out_gates = []
+    # Gate check
+    last_gates = []
     try:
         model.eval()
         with torch.no_grad():
             sample_imgs = next(iter(val_loader))[0][:4].to(DEVICE)
             sample_out = model(sample_imgs, apply_autograd=False)
-            last_out_gates = sample_out.get('gates_a', [])
+            last_gates = sample_out.get('gates_a', [])
     except:
         pass
-    gate_str = f"g={np.mean(last_out_gates):.3f}" if last_out_gates else "g=?"
+    gate_str = f"g={np.mean(last_gates):.4f}" if last_gates else "g=?"
 
     print(f"  E{epoch+1:3d}: A={train_acc:.1f}% B={ab:.0f}% "
-          f"val={val_acc:.1f}% geo={ga:.0f}/{val_geo_acc:.0f}% "
+          f"val={val_acc:.1f}%/{val_b_acc:.1f}%/{val_geo_acc:.1f}% "
           f"loss={acc_dict['loss']/d:.4f}/{val_loss:.4f} "
+          f"nb={nb_acc:.2f} gn={gn_acc:.2f} "
           f"cv={v_cv:.4f}(m={cvf:.5f} g={cvg:.5f}) "
           f"cm={cmv:.0%} anch={n_active}/{N_ANCHORS} "
-          f"[{stage}] mst={mst_m:.3f} {gate_str} "
-          f"hn={hn:.3f} hp={hp:.3f} "
-          f"q={mastery.size}/{mastery._current_max} "
+          f"[{stage}] {gate_str} "
           f"gal={gal_update_count} ({elapsed:.0f}s){mk}")
 
 writer.close()