Update train_cell_cifar_10_resnet.py

train_cell_cifar_10_resnet.py

@@ -29,6 +29,31 @@ import torchvision.transforms as T
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+# ── Soft Hand CV ─────────────────────────────────────────────────
+
+@torch.no_grad()
+def compute_target_cv(V, D, n_trials=20, n_samples=200, device='cuda'):
+    """Geometric attractor CV for (V, D) on S^{D-1}.
+    Random unit vectors, pentachoron sampling, averaged.
+    """
+    cvs = []
+    for _ in range(n_trials):
+        pts = F.normalize(torch.randn(V, D, device=device), dim=-1)
+        cv = cv_of(pts, n_samples=n_samples)
+        if cv > 0:
+            cvs.append(cv)
+    return sum(cvs) / len(cvs) if cvs else 0.0
+
+import math as _math
+
+def cv_proximity(measured, target, sigma=0.15):
+    """Gaussian proximity: 1.0 at target, decays with distance."""
+    return _math.exp(-(measured - target) ** 2 / (2 * sigma ** 2))
+
+def soft_hand_weights(proximity, boost=0.5, penalty_weight=0.3):
+    """Near target: boost primary, penalty→0. Far: primary baseline, penalty active."""
+    return 1.0 + boost * proximity, penalty_weight * (1.0 - proximity)
+
 # ── Conv building blocks ─────────────────────────────────────────
 
 class ConvBNAct(nn.Module):
@@ -228,6 +253,18 @@ print(f"  Label smooth:  {LABEL_SMOOTH}")
 criterion = nn.CrossEntropyLoss(label_smoothing=LABEL_SMOOTH)
 opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WD)
 
+# Soft hand config
+SIGMA = 0.15
+BOOST = 0.5
+PENALTY_WEIGHT = 0.3
+CV_MEASURE_EVERY = 50  # batches
+
+print(f"\nComputing target CV for V=16, D=16 on S^15...")
+target_cv = compute_target_cv(16, 16, n_trials=20, device=device)
+print(f"  target_cv = {target_cv:.4f}")
+print(f"  soft hand: σ={SIGMA}, boost={BOOST}, penalty={PENALTY_WEIGHT}")
+print(f"  CV measured every {CV_MEASURE_EVERY} batches")
+
 def cosine_lr(epoch):
     t = epoch / float(max(1, EPOCHS - 1))
     min_ratio = 1e-5 / LR
@@ -242,15 +279,35 @@ print(f"{'=' * 70}")
 best_acc = 0
 t0 = time.time()
 
+# Soft hand state
+last_cv = target_cv
+last_prox = 1.0
+
 for epoch in range(1, EPOCHS + 1):
     model.train()
     correct, total = 0, 0
     loss_sum = 0
+    ep_prox_sum, ep_prox_n = 0, 0
 
-    for images, labels in tqdm(train_loader, desc=f"Ep {epoch:3d}", leave=False):
+    for batch_idx, (images, labels) in enumerate(tqdm(train_loader, desc=f"Ep {epoch:3d}", leave=False)):
         images, labels = images.to(device), labels.to(device)
         out = model(images)
-        loss = criterion(out['logits'], labels)
+        ce_loss = criterion(out['logits'], labels)
+
+        # Measure CV periodically from cell M rows
+        with torch.no_grad():
+            if batch_idx % CV_MEASURE_EVERY == 0:
+                # Sample from cell3 M rows (deeper features)
+                M_sample = out['cell3']['M'][0, 0]  # (V, D)
+                current_cv = cv_of(M_sample, n_samples=100)
+                if current_cv > 0:
+                    last_cv = current_cv
+                last_prox = cv_proximity(last_cv, target_cv, SIGMA)
+
+        # Soft hand adaptive weights
+        primary_w, cv_w = soft_hand_weights(last_prox, BOOST, PENALTY_WEIGHT)
+        cv_l = (last_cv - target_cv) ** 2
+        loss = primary_w * ce_loss + cv_w * cv_l
 
         opt.zero_grad(set_to_none=True)
         loss.backward()
@@ -259,17 +316,21 @@ for epoch in range(1, EPOCHS + 1):
 
         correct += (out['logits'].argmax(-1) == labels).sum().item()
         total += images.shape[0]
-        loss_sum += loss.item()
+        loss_sum += ce_loss.item()
+        ep_prox_sum += last_prox
+        ep_prox_n += 1
 
     sched.step()
     train_acc = correct / total
     avg_loss = loss_sum / len(train_loader)
+    avg_prox = ep_prox_sum / max(1, ep_prox_n)
 
     # Val
     model.eval()
     val_correct, val_total = 0, 0
     pcc = torch.zeros(10)
     pct = torch.zeros(10)
+    val_cv_samples = []
 
     with torch.no_grad():
         for images, labels in test_loader:
@@ -283,7 +344,14 @@ for epoch in range(1, EPOCHS + 1):
                 pcc[c] += (preds[m] == labels[m]).sum().item()
                 pct[c] += m.sum().item()
 
+        # Measure val CV
+        for i in range(min(4, out['cell3']['M'].shape[0])):
+            cv_val = cv_of(out['cell3']['M'][i, 0], n_samples=100)
+            if cv_val > 0:
+                val_cv_samples.append(cv_val)
+
     val_acc = val_correct / val_total
+    val_cv = sum(val_cv_samples) / len(val_cv_samples) if val_cv_samples else 0.0
     star = ''
     if val_acc > best_acc:
         best_acc = val_acc
@@ -292,7 +360,6 @@ for epoch in range(1, EPOCHS + 1):
     lr = opt.param_groups[0]['lr']
 
     if epoch <= 5 or epoch % 5 == 0 or epoch == EPOCHS:
-        # Cell spectral diagnostics
         with torch.no_grad():
             S2 = out['cell2']['S_orig'].mean(dim=(0, 1))
             S3 = out['cell3']['S_orig'].mean(dim=(0, 1))
@@ -300,7 +367,7 @@ for epoch in range(1, EPOCHS + 1):
             s3_str = ', '.join(f'{v:.2f}' for v in S3.tolist()[:3]) + f'...{S3[-1]:.2f}'
 
         print(f"  ep{epoch:3d} loss={avg_loss:.4f} acc={val_acc:.1%}{star} "
-              f"train={train_acc:.1%} lr={lr:.6f}")
+              f"train={train_acc:.1%} cv={val_cv:.4f} prox={avg_prox:.3f} lr={lr:.6f}")
         print(f"    S2=[{s2_str}]  S3=[{s3_str}]")
 
     if epoch <= 3 or epoch % 20 == 0 or epoch == EPOCHS: