Upload 7 files

spectral/notebooks/experiment_1_signal_decomposition.ipynb

@@ -83,8 +83,16 @@
     "from collections import defaultdict\n",
     "\n",
     "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "\n",
+    "# \u2500\u2500 Performance: TF32 + cudnn benchmark \u2500\u2500\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "torch.backends.cudnn.allow_tf32 = True\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "\n",
     "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "if device.type == \"cuda\":\n",
+    "    print(f\"[PERF] TF32={torch.backends.cuda.matmul.allow_tf32}, cudnn.benchmark={torch.backends.cudnn.benchmark}\")\n",
     "\n",
     "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
     "# GEOLIP CORE \u2014 Geometric Building Blocks\n",
@@ -454,8 +462,10 @@
     "    return torch.stack(vols)\n",
     "\n",
     "\n",
-    "def cv_loss(emb, target=0.22, n_samples=64, n_points=5, batched=True):\n",
-    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\"\"\"\n",
+    "def cv_loss(emb, target=0.22, n_samples=32, n_points=5, batched=True):\n",
+    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\n",
+    "    Default n_samples=32 for training speed (141x faster than sequential).\n",
+    "    Use n_samples=200 for monitoring/metrics only.\"\"\"\n",
     "    if emb.shape[0] < n_points:\n",
     "        return torch.tensor(0.0, device=emb.device, requires_grad=True)\n",
     "    vols = _batch_pentachoron_volumes(emb, n_samples, n_points) if batched else _sequential_pentachoron_volumes(emb, n_samples, n_points)\n",
@@ -668,8 +678,9 @@
     "    l_spread = spread_loss(constellation.anchors)\n",
     "    ld['spread'] = l_spread\n",
     "\n",
-    "    # \u2500\u2500 kNN \u2500\u2500\n",
-    "    ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "    # \u2500\u2500 kNN (skip during training for speed \u2014 only compute when explicitly needed) \u2500\u2500\n",
+    "    if targets is not None and emb1.shape[0] <= 512:\n",
+    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "\n",
     "    # \u2500\u2500 TOTAL \u2500\u2500\n",
     "    loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
@@ -962,19 +973,22 @@
     "                     w_assign=0.5, w_assign_nce=0.25,\n",
     "                     w_nce_tri=0.5, w_attract=0.25,\n",
     "                     w_cv=0.01, w_spread=0.01,\n",
-    "                     cv_batched=True):\n",
-    "        \"\"\"Three-domain cooperative loss.\n",
+    "                     cv_batched=True, compute_knn=False):\n",
+    "        \"\"\"Three-domain cooperative loss \u2014 fully batched, zero Python loops.\n",
+    "        Args:\n",
+    "            compute_knn: if False (default), skip kNN during training for speed.\n",
+    "                         Set True during validation or every N steps.\n",
     "        Returns:\n",
     "            total_loss, loss_dict\n",
     "        \"\"\"\n",
     "        ld = {}\n",
     "        emb1, emb2 = output['embedding'], output['embedding_aug']\n",
-    "        # \u2500\u2500 EXTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 EXTERNAL (batched matmul) \u2500\u2500\n",
     "        l_ce, acc = ce_loss_paired(output['logits'], output['logits_aug'], targets)\n",
     "        ld['ce'], ld['acc'] = l_ce, acc\n",
     "        l_nce_emb, nce_emb_acc = nce_loss(emb1, emb2, self.infonce_temp, normalize=False)\n",
     "        ld['nce_emb'], ld['nce_emb_acc'] = l_nce_emb, nce_emb_acc\n",
-    "        # \u2500\u2500 GEOMETRIC \u2500\u2500\n",
+    "        # \u2500\u2500 GEOMETRIC (batched matmul) \u2500\u2500\n",
     "        l_nce_pw, nce_pw_acc = nce_loss(\n",
     "            output['patchwork1'], output['patchwork1_aug'], self.assign_temp, normalize=True)\n",
     "        ld['nce_pw'], ld['nce_pw_acc'] = l_nce_pw, nce_pw_acc\n",
@@ -982,7 +996,7 @@
     "            output['bridge1'], output['bridge2'],\n",
     "            output['assign1'], output['assign2'])\n",
     "        ld['bridge'], ld['bridge_acc'] = l_bridge, bridge_acc\n",
-    "        # \u2500\u2500 INTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 INTERNAL (batched \u2014 no Python loops) \u2500\u2500\n",
     "        l_assign, assign_ent = assign_bce_loss(output['assign1'], output['cos1'])\n",
     "        ld['assign'], ld['assign_entropy'] = l_assign, assign_ent\n",
     "        l_assign_nce, assign_nce_acc = assign_nce_loss(\n",
@@ -993,12 +1007,14 @@
     "        ld['nce_tri'], ld['nce_tri_acc'] = l_nce_tri, nce_tri_acc\n",
     "        l_attract, nearest_cos = attraction_loss(output['cos1'])\n",
     "        ld['attract'], ld['nearest_cos'] = l_attract, nearest_cos\n",
-    "        l_cv = cv_loss(emb1, target=self.cv_target, batched=cv_batched)\n",
+    "        # CV: batched Cayley-Menger, n_samples=32 for training speed\n",
+    "        l_cv = cv_loss(emb1, target=self.cv_target, n_samples=32, batched=cv_batched)\n",
     "        ld['cv'] = l_cv\n",
     "        l_spread = spread_loss(self.constellation.anchors)\n",
     "        ld['spread'] = l_spread\n",
-    "        # \u2500\u2500 kNN \u2500\u2500\n",
-    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "        # \u2500\u2500 kNN (SKIP during training \u2014 B\u00d7B matmul is expensive every batch) \u2500\u2500\n",
+    "        if compute_knn:\n",
+    "            ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "        # \u2500\u2500 TOTAL \u2500\u2500\n",
     "        loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
     "        loss_geometric = w_nce_pw * l_nce_pw + w_bridge * l_bridge\n",
@@ -1335,9 +1351,11 @@
     "\n",
     "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,\n",
     "                              num_workers=num_workers, pin_memory=True,\n",
-    "                              drop_last=True, collate_fn=paired_collate)\n",
+    "                              drop_last=True, collate_fn=paired_collate,\n",
+    "                              persistent_workers=(num_workers > 0))\n",
     "    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,\n",
-    "                            num_workers=num_workers, pin_memory=True)\n",
+    "                            num_workers=num_workers, pin_memory=True,\n",
+    "                            persistent_workers=(num_workers > 0))\n",
     "    print(f\"[DATA] CIFAR-10 paired: {len(train_ds)} train, {len(val_ds)} val, bs={batch_size}\")\n",
     "    return train_loader, val_loader\n",
     "\n",
@@ -1468,7 +1486,7 @@
     "            optimizer.zero_grad()\n",
     "\n",
     "            output = model.forward_paired(v1, v2)\n",
-    "            loss, ld = model.compute_loss(output, labels, **lw)\n",
+    "            loss, ld = model.compute_loss(output, labels, compute_knn=False, **lw)\n",
     "\n",
     "            loss.backward()\n",
     "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",

spectral/notebooks/experiment_2_manifold_structures.ipynb

@@ -75,8 +75,16 @@
     "from collections import defaultdict\n",
     "\n",
     "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "\n",
+    "# \u2500\u2500 Performance: TF32 + cudnn benchmark \u2500\u2500\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "torch.backends.cudnn.allow_tf32 = True\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "\n",
     "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "if device.type == \"cuda\":\n",
+    "    print(f\"[PERF] TF32={torch.backends.cuda.matmul.allow_tf32}, cudnn.benchmark={torch.backends.cudnn.benchmark}\")\n",
     "\n",
     "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
     "# GEOLIP CORE \u2014 Geometric Building Blocks\n",
@@ -446,8 +454,10 @@
     "    return torch.stack(vols)\n",
     "\n",
     "\n",
-    "def cv_loss(emb, target=0.22, n_samples=64, n_points=5, batched=True):\n",
-    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\"\"\"\n",
+    "def cv_loss(emb, target=0.22, n_samples=32, n_points=5, batched=True):\n",
+    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\n",
+    "    Default n_samples=32 for training speed (141x faster than sequential).\n",
+    "    Use n_samples=200 for monitoring/metrics only.\"\"\"\n",
     "    if emb.shape[0] < n_points:\n",
     "        return torch.tensor(0.0, device=emb.device, requires_grad=True)\n",
     "    vols = _batch_pentachoron_volumes(emb, n_samples, n_points) if batched else _sequential_pentachoron_volumes(emb, n_samples, n_points)\n",
@@ -660,8 +670,9 @@
     "    l_spread = spread_loss(constellation.anchors)\n",
     "    ld['spread'] = l_spread\n",
     "\n",
-    "    # \u2500\u2500 kNN \u2500\u2500\n",
-    "    ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "    # \u2500\u2500 kNN (skip during training for speed \u2014 only compute when explicitly needed) \u2500\u2500\n",
+    "    if targets is not None and emb1.shape[0] <= 512:\n",
+    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "\n",
     "    # \u2500\u2500 TOTAL \u2500\u2500\n",
     "    loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
@@ -954,19 +965,22 @@
     "                     w_assign=0.5, w_assign_nce=0.25,\n",
     "                     w_nce_tri=0.5, w_attract=0.25,\n",
     "                     w_cv=0.01, w_spread=0.01,\n",
-    "                     cv_batched=True):\n",
-    "        \"\"\"Three-domain cooperative loss.\n",
+    "                     cv_batched=True, compute_knn=False):\n",
+    "        \"\"\"Three-domain cooperative loss \u2014 fully batched, zero Python loops.\n",
+    "        Args:\n",
+    "            compute_knn: if False (default), skip kNN during training for speed.\n",
+    "                         Set True during validation or every N steps.\n",
     "        Returns:\n",
     "            total_loss, loss_dict\n",
     "        \"\"\"\n",
     "        ld = {}\n",
     "        emb1, emb2 = output['embedding'], output['embedding_aug']\n",
-    "        # \u2500\u2500 EXTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 EXTERNAL (batched matmul) \u2500\u2500\n",
     "        l_ce, acc = ce_loss_paired(output['logits'], output['logits_aug'], targets)\n",
     "        ld['ce'], ld['acc'] = l_ce, acc\n",
     "        l_nce_emb, nce_emb_acc = nce_loss(emb1, emb2, self.infonce_temp, normalize=False)\n",
     "        ld['nce_emb'], ld['nce_emb_acc'] = l_nce_emb, nce_emb_acc\n",
-    "        # \u2500\u2500 GEOMETRIC \u2500\u2500\n",
+    "        # \u2500\u2500 GEOMETRIC (batched matmul) \u2500\u2500\n",
     "        l_nce_pw, nce_pw_acc = nce_loss(\n",
     "            output['patchwork1'], output['patchwork1_aug'], self.assign_temp, normalize=True)\n",
     "        ld['nce_pw'], ld['nce_pw_acc'] = l_nce_pw, nce_pw_acc\n",
@@ -974,7 +988,7 @@
     "            output['bridge1'], output['bridge2'],\n",
     "            output['assign1'], output['assign2'])\n",
     "        ld['bridge'], ld['bridge_acc'] = l_bridge, bridge_acc\n",
-    "        # \u2500\u2500 INTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 INTERNAL (batched \u2014 no Python loops) \u2500\u2500\n",
     "        l_assign, assign_ent = assign_bce_loss(output['assign1'], output['cos1'])\n",
     "        ld['assign'], ld['assign_entropy'] = l_assign, assign_ent\n",
     "        l_assign_nce, assign_nce_acc = assign_nce_loss(\n",
@@ -985,12 +999,14 @@
     "        ld['nce_tri'], ld['nce_tri_acc'] = l_nce_tri, nce_tri_acc\n",
     "        l_attract, nearest_cos = attraction_loss(output['cos1'])\n",
     "        ld['attract'], ld['nearest_cos'] = l_attract, nearest_cos\n",
-    "        l_cv = cv_loss(emb1, target=self.cv_target, batched=cv_batched)\n",
+    "        # CV: batched Cayley-Menger, n_samples=32 for training speed\n",
+    "        l_cv = cv_loss(emb1, target=self.cv_target, n_samples=32, batched=cv_batched)\n",
     "        ld['cv'] = l_cv\n",
     "        l_spread = spread_loss(self.constellation.anchors)\n",
     "        ld['spread'] = l_spread\n",
-    "        # \u2500\u2500 kNN \u2500\u2500\n",
-    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "        # \u2500\u2500 kNN (SKIP during training \u2014 B\u00d7B matmul is expensive every batch) \u2500\u2500\n",
+    "        if compute_knn:\n",
+    "            ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "        # \u2500\u2500 TOTAL \u2500\u2500\n",
     "        loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
     "        loss_geometric = w_nce_pw * l_nce_pw + w_bridge * l_bridge\n",
@@ -1327,9 +1343,11 @@
     "\n",
     "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,\n",
     "                              num_workers=num_workers, pin_memory=True,\n",
-    "                              drop_last=True, collate_fn=paired_collate)\n",
+    "                              drop_last=True, collate_fn=paired_collate,\n",
+    "                              persistent_workers=(num_workers > 0))\n",
     "    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,\n",
-    "                            num_workers=num_workers, pin_memory=True)\n",
+    "                            num_workers=num_workers, pin_memory=True,\n",
+    "                            persistent_workers=(num_workers > 0))\n",
     "    print(f\"[DATA] CIFAR-10 paired: {len(train_ds)} train, {len(val_ds)} val, bs={batch_size}\")\n",
     "    return train_loader, val_loader\n",
     "\n",
@@ -1460,7 +1478,7 @@
     "            optimizer.zero_grad()\n",
     "\n",
     "            output = model.forward_paired(v1, v2)\n",
-    "            loss, ld = model.compute_loss(output, labels, **lw)\n",
+    "            loss, ld = model.compute_loss(output, labels, compute_knn=False, **lw)\n",
     "\n",
     "            loss.backward()\n",
     "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",

spectral/notebooks/experiment_3_compact_representations.ipynb

@@ -74,8 +74,16 @@
     "from collections import defaultdict\n",
     "\n",
     "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "\n",
+    "# \u2500\u2500 Performance: TF32 + cudnn benchmark \u2500\u2500\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "torch.backends.cudnn.allow_tf32 = True\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "\n",
     "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "if device.type == \"cuda\":\n",
+    "    print(f\"[PERF] TF32={torch.backends.cuda.matmul.allow_tf32}, cudnn.benchmark={torch.backends.cudnn.benchmark}\")\n",
     "\n",
     "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
     "# GEOLIP CORE \u2014 Geometric Building Blocks\n",
@@ -445,8 +453,10 @@
     "    return torch.stack(vols)\n",
     "\n",
     "\n",
-    "def cv_loss(emb, target=0.22, n_samples=64, n_points=5, batched=True):\n",
-    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\"\"\"\n",
+    "def cv_loss(emb, target=0.22, n_samples=32, n_points=5, batched=True):\n",
+    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\n",
+    "    Default n_samples=32 for training speed (141x faster than sequential).\n",
+    "    Use n_samples=200 for monitoring/metrics only.\"\"\"\n",
     "    if emb.shape[0] < n_points:\n",
     "        return torch.tensor(0.0, device=emb.device, requires_grad=True)\n",
     "    vols = _batch_pentachoron_volumes(emb, n_samples, n_points) if batched else _sequential_pentachoron_volumes(emb, n_samples, n_points)\n",
@@ -659,8 +669,9 @@
     "    l_spread = spread_loss(constellation.anchors)\n",
     "    ld['spread'] = l_spread\n",
     "\n",
-    "    # \u2500\u2500 kNN \u2500\u2500\n",
-    "    ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "    # \u2500\u2500 kNN (skip during training for speed \u2014 only compute when explicitly needed) \u2500\u2500\n",
+    "    if targets is not None and emb1.shape[0] <= 512:\n",
+    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "\n",
     "    # \u2500\u2500 TOTAL \u2500\u2500\n",
     "    loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
@@ -953,19 +964,22 @@
     "                     w_assign=0.5, w_assign_nce=0.25,\n",
     "                     w_nce_tri=0.5, w_attract=0.25,\n",
     "                     w_cv=0.01, w_spread=0.01,\n",
-    "                     cv_batched=True):\n",
-    "        \"\"\"Three-domain cooperative loss.\n",
+    "                     cv_batched=True, compute_knn=False):\n",
+    "        \"\"\"Three-domain cooperative loss \u2014 fully batched, zero Python loops.\n",
+    "        Args:\n",
+    "            compute_knn: if False (default), skip kNN during training for speed.\n",
+    "                         Set True during validation or every N steps.\n",
     "        Returns:\n",
     "            total_loss, loss_dict\n",
     "        \"\"\"\n",
     "        ld = {}\n",
     "        emb1, emb2 = output['embedding'], output['embedding_aug']\n",
-    "        # \u2500\u2500 EXTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 EXTERNAL (batched matmul) \u2500\u2500\n",
     "        l_ce, acc = ce_loss_paired(output['logits'], output['logits_aug'], targets)\n",
     "        ld['ce'], ld['acc'] = l_ce, acc\n",
     "        l_nce_emb, nce_emb_acc = nce_loss(emb1, emb2, self.infonce_temp, normalize=False)\n",
     "        ld['nce_emb'], ld['nce_emb_acc'] = l_nce_emb, nce_emb_acc\n",
-    "        # \u2500\u2500 GEOMETRIC \u2500\u2500\n",
+    "        # \u2500\u2500 GEOMETRIC (batched matmul) \u2500\u2500\n",
     "        l_nce_pw, nce_pw_acc = nce_loss(\n",
     "            output['patchwork1'], output['patchwork1_aug'], self.assign_temp, normalize=True)\n",
     "        ld['nce_pw'], ld['nce_pw_acc'] = l_nce_pw, nce_pw_acc\n",
@@ -973,7 +987,7 @@
     "            output['bridge1'], output['bridge2'],\n",
     "            output['assign1'], output['assign2'])\n",
     "        ld['bridge'], ld['bridge_acc'] = l_bridge, bridge_acc\n",
-    "        # \u2500\u2500 INTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 INTERNAL (batched \u2014 no Python loops) \u2500\u2500\n",
     "        l_assign, assign_ent = assign_bce_loss(output['assign1'], output['cos1'])\n",
     "        ld['assign'], ld['assign_entropy'] = l_assign, assign_ent\n",
     "        l_assign_nce, assign_nce_acc = assign_nce_loss(\n",
@@ -984,12 +998,14 @@
     "        ld['nce_tri'], ld['nce_tri_acc'] = l_nce_tri, nce_tri_acc\n",
     "        l_attract, nearest_cos = attraction_loss(output['cos1'])\n",
     "        ld['attract'], ld['nearest_cos'] = l_attract, nearest_cos\n",
-    "        l_cv = cv_loss(emb1, target=self.cv_target, batched=cv_batched)\n",
+    "        # CV: batched Cayley-Menger, n_samples=32 for training speed\n",
+    "        l_cv = cv_loss(emb1, target=self.cv_target, n_samples=32, batched=cv_batched)\n",
     "        ld['cv'] = l_cv\n",
     "        l_spread = spread_loss(self.constellation.anchors)\n",
     "        ld['spread'] = l_spread\n",
-    "        # \u2500\u2500 kNN \u2500\u2500\n",
-    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "        # \u2500\u2500 kNN (SKIP during training \u2014 B\u00d7B matmul is expensive every batch) \u2500\u2500\n",
+    "        if compute_knn:\n",
+    "            ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "        # \u2500\u2500 TOTAL \u2500\u2500\n",
     "        loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
     "        loss_geometric = w_nce_pw * l_nce_pw + w_bridge * l_bridge\n",
@@ -1326,9 +1342,11 @@
     "\n",
     "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,\n",
     "                              num_workers=num_workers, pin_memory=True,\n",
-    "                              drop_last=True, collate_fn=paired_collate)\n",
+    "                              drop_last=True, collate_fn=paired_collate,\n",
+    "                              persistent_workers=(num_workers > 0))\n",
     "    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,\n",
-    "                            num_workers=num_workers, pin_memory=True)\n",
+    "                            num_workers=num_workers, pin_memory=True,\n",
+    "                            persistent_workers=(num_workers > 0))\n",
     "    print(f\"[DATA] CIFAR-10 paired: {len(train_ds)} train, {len(val_ds)} val, bs={batch_size}\")\n",
     "    return train_loader, val_loader\n",
     "\n",
@@ -1459,7 +1477,7 @@
     "            optimizer.zero_grad()\n",
     "\n",
     "            output = model.forward_paired(v1, v2)\n",
-    "            loss, ld = model.compute_loss(output, labels, **lw)\n",
+    "            loss, ld = model.compute_loss(output, labels, compute_knn=False, **lw)\n",
     "\n",
     "            loss.backward()\n",
     "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",

spectral/notebooks/experiment_4_invertible_transforms.ipynb

@@ -75,8 +75,16 @@
     "from collections import defaultdict\n",
     "\n",
     "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "\n",
+    "# \u2500\u2500 Performance: TF32 + cudnn benchmark \u2500\u2500\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "torch.backends.cudnn.allow_tf32 = True\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "\n",
     "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "if device.type == \"cuda\":\n",
+    "    print(f\"[PERF] TF32={torch.backends.cuda.matmul.allow_tf32}, cudnn.benchmark={torch.backends.cudnn.benchmark}\")\n",
     "\n",
     "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
     "# GEOLIP CORE \u2014 Geometric Building Blocks\n",
@@ -446,8 +454,10 @@
     "    return torch.stack(vols)\n",
     "\n",
     "\n",
-    "def cv_loss(emb, target=0.22, n_samples=64, n_points=5, batched=True):\n",
-    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\"\"\"\n",
+    "def cv_loss(emb, target=0.22, n_samples=32, n_points=5, batched=True):\n",
+    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\n",
+    "    Default n_samples=32 for training speed (141x faster than sequential).\n",
+    "    Use n_samples=200 for monitoring/metrics only.\"\"\"\n",
     "    if emb.shape[0] < n_points:\n",
     "        return torch.tensor(0.0, device=emb.device, requires_grad=True)\n",
     "    vols = _batch_pentachoron_volumes(emb, n_samples, n_points) if batched else _sequential_pentachoron_volumes(emb, n_samples, n_points)\n",
@@ -660,8 +670,9 @@
     "    l_spread = spread_loss(constellation.anchors)\n",
     "    ld['spread'] = l_spread\n",
     "\n",
-    "    # \u2500\u2500 kNN \u2500\u2500\n",
-    "    ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "    # \u2500\u2500 kNN (skip during training for speed \u2014 only compute when explicitly needed) \u2500\u2500\n",
+    "    if targets is not None and emb1.shape[0] <= 512:\n",
+    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "\n",
     "    # \u2500\u2500 TOTAL \u2500\u2500\n",
     "    loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
@@ -954,19 +965,22 @@
     "                     w_assign=0.5, w_assign_nce=0.25,\n",
     "                     w_nce_tri=0.5, w_attract=0.25,\n",
     "                     w_cv=0.01, w_spread=0.01,\n",
-    "                     cv_batched=True):\n",
-    "        \"\"\"Three-domain cooperative loss.\n",
+    "                     cv_batched=True, compute_knn=False):\n",
+    "        \"\"\"Three-domain cooperative loss \u2014 fully batched, zero Python loops.\n",
+    "        Args:\n",
+    "            compute_knn: if False (default), skip kNN during training for speed.\n",
+    "                         Set True during validation or every N steps.\n",
     "        Returns:\n",
     "            total_loss, loss_dict\n",
     "        \"\"\"\n",
     "        ld = {}\n",
     "        emb1, emb2 = output['embedding'], output['embedding_aug']\n",
-    "        # \u2500\u2500 EXTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 EXTERNAL (batched matmul) \u2500\u2500\n",
     "        l_ce, acc = ce_loss_paired(output['logits'], output['logits_aug'], targets)\n",
     "        ld['ce'], ld['acc'] = l_ce, acc\n",
     "        l_nce_emb, nce_emb_acc = nce_loss(emb1, emb2, self.infonce_temp, normalize=False)\n",
     "        ld['nce_emb'], ld['nce_emb_acc'] = l_nce_emb, nce_emb_acc\n",
-    "        # \u2500\u2500 GEOMETRIC \u2500\u2500\n",
+    "        # \u2500\u2500 GEOMETRIC (batched matmul) \u2500\u2500\n",
     "        l_nce_pw, nce_pw_acc = nce_loss(\n",
     "            output['patchwork1'], output['patchwork1_aug'], self.assign_temp, normalize=True)\n",
     "        ld['nce_pw'], ld['nce_pw_acc'] = l_nce_pw, nce_pw_acc\n",
@@ -974,7 +988,7 @@
     "            output['bridge1'], output['bridge2'],\n",
     "            output['assign1'], output['assign2'])\n",
     "        ld['bridge'], ld['bridge_acc'] = l_bridge, bridge_acc\n",
-    "        # \u2500\u2500 INTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 INTERNAL (batched \u2014 no Python loops) \u2500\u2500\n",
     "        l_assign, assign_ent = assign_bce_loss(output['assign1'], output['cos1'])\n",
     "        ld['assign'], ld['assign_entropy'] = l_assign, assign_ent\n",
     "        l_assign_nce, assign_nce_acc = assign_nce_loss(\n",
@@ -985,12 +999,14 @@
     "        ld['nce_tri'], ld['nce_tri_acc'] = l_nce_tri, nce_tri_acc\n",
     "        l_attract, nearest_cos = attraction_loss(output['cos1'])\n",
     "        ld['attract'], ld['nearest_cos'] = l_attract, nearest_cos\n",
-    "        l_cv = cv_loss(emb1, target=self.cv_target, batched=cv_batched)\n",
+    "        # CV: batched Cayley-Menger, n_samples=32 for training speed\n",
+    "        l_cv = cv_loss(emb1, target=self.cv_target, n_samples=32, batched=cv_batched)\n",
     "        ld['cv'] = l_cv\n",
     "        l_spread = spread_loss(self.constellation.anchors)\n",
     "        ld['spread'] = l_spread\n",
-    "        # \u2500\u2500 kNN \u2500\u2500\n",
-    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "        # \u2500\u2500 kNN (SKIP during training \u2014 B\u00d7B matmul is expensive every batch) \u2500\u2500\n",
+    "        if compute_knn:\n",
+    "            ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "        # \u2500\u2500 TOTAL \u2500\u2500\n",
     "        loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
     "        loss_geometric = w_nce_pw * l_nce_pw + w_bridge * l_bridge\n",
@@ -1327,9 +1343,11 @@
     "\n",
     "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,\n",
     "                              num_workers=num_workers, pin_memory=True,\n",
-    "                              drop_last=True, collate_fn=paired_collate)\n",
+    "                              drop_last=True, collate_fn=paired_collate,\n",
+    "                              persistent_workers=(num_workers > 0))\n",
     "    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,\n",
-    "                            num_workers=num_workers, pin_memory=True)\n",
+    "                            num_workers=num_workers, pin_memory=True,\n",
+    "                            persistent_workers=(num_workers > 0))\n",
     "    print(f\"[DATA] CIFAR-10 paired: {len(train_ds)} train, {len(val_ds)} val, bs={batch_size}\")\n",
     "    return train_loader, val_loader\n",
     "\n",
@@ -1460,7 +1478,7 @@
     "            optimizer.zero_grad()\n",
     "\n",
     "            output = model.forward_paired(v1, v2)\n",
-    "            loss, ld = model.compute_loss(output, labels, **lw)\n",
+    "            loss, ld = model.compute_loss(output, labels, compute_knn=False, **lw)\n",
     "\n",
     "            loss.backward()\n",
     "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",

spectral/notebooks/experiment_5_matrix_decompositions.ipynb

@@ -75,8 +75,16 @@
     "from collections import defaultdict\n",
     "\n",
     "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "\n",
+    "# \u2500\u2500 Performance: TF32 + cudnn benchmark \u2500\u2500\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "torch.backends.cudnn.allow_tf32 = True\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "\n",
     "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "if device.type == \"cuda\":\n",
+    "    print(f\"[PERF] TF32={torch.backends.cuda.matmul.allow_tf32}, cudnn.benchmark={torch.backends.cudnn.benchmark}\")\n",
     "\n",
     "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
     "# GEOLIP CORE \u2014 Geometric Building Blocks\n",
@@ -446,8 +454,10 @@
     "    return torch.stack(vols)\n",
     "\n",
     "\n",
-    "def cv_loss(emb, target=0.22, n_samples=64, n_points=5, batched=True):\n",
-    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\"\"\"\n",
+    "def cv_loss(emb, target=0.22, n_samples=32, n_points=5, batched=True):\n",
+    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\n",
+    "    Default n_samples=32 for training speed (141x faster than sequential).\n",
+    "    Use n_samples=200 for monitoring/metrics only.\"\"\"\n",
     "    if emb.shape[0] < n_points:\n",
     "        return torch.tensor(0.0, device=emb.device, requires_grad=True)\n",
     "    vols = _batch_pentachoron_volumes(emb, n_samples, n_points) if batched else _sequential_pentachoron_volumes(emb, n_samples, n_points)\n",
@@ -660,8 +670,9 @@
     "    l_spread = spread_loss(constellation.anchors)\n",
     "    ld['spread'] = l_spread\n",
     "\n",
-    "    # \u2500\u2500 kNN \u2500\u2500\n",
-    "    ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "    # \u2500\u2500 kNN (skip during training for speed \u2014 only compute when explicitly needed) \u2500\u2500\n",
+    "    if targets is not None and emb1.shape[0] <= 512:\n",
+    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "\n",
     "    # \u2500\u2500 TOTAL \u2500\u2500\n",
     "    loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
@@ -954,19 +965,22 @@
     "                     w_assign=0.5, w_assign_nce=0.25,\n",
     "                     w_nce_tri=0.5, w_attract=0.25,\n",
     "                     w_cv=0.01, w_spread=0.01,\n",
-    "                     cv_batched=True):\n",
-    "        \"\"\"Three-domain cooperative loss.\n",
+    "                     cv_batched=True, compute_knn=False):\n",
+    "        \"\"\"Three-domain cooperative loss \u2014 fully batched, zero Python loops.\n",
+    "        Args:\n",
+    "            compute_knn: if False (default), skip kNN during training for speed.\n",
+    "                         Set True during validation or every N steps.\n",
     "        Returns:\n",
     "            total_loss, loss_dict\n",
     "        \"\"\"\n",
     "        ld = {}\n",
     "        emb1, emb2 = output['embedding'], output['embedding_aug']\n",
-    "        # \u2500\u2500 EXTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 EXTERNAL (batched matmul) \u2500\u2500\n",
     "        l_ce, acc = ce_loss_paired(output['logits'], output['logits_aug'], targets)\n",
     "        ld['ce'], ld['acc'] = l_ce, acc\n",
     "        l_nce_emb, nce_emb_acc = nce_loss(emb1, emb2, self.infonce_temp, normalize=False)\n",
     "        ld['nce_emb'], ld['nce_emb_acc'] = l_nce_emb, nce_emb_acc\n",
-    "        # \u2500\u2500 GEOMETRIC \u2500\u2500\n",
+    "        # \u2500\u2500 GEOMETRIC (batched matmul) \u2500\u2500\n",
     "        l_nce_pw, nce_pw_acc = nce_loss(\n",
     "            output['patchwork1'], output['patchwork1_aug'], self.assign_temp, normalize=True)\n",
     "        ld['nce_pw'], ld['nce_pw_acc'] = l_nce_pw, nce_pw_acc\n",
@@ -974,7 +988,7 @@
     "            output['bridge1'], output['bridge2'],\n",
     "            output['assign1'], output['assign2'])\n",
     "        ld['bridge'], ld['bridge_acc'] = l_bridge, bridge_acc\n",
-    "        # \u2500\u2500 INTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 INTERNAL (batched \u2014 no Python loops) \u2500\u2500\n",
     "        l_assign, assign_ent = assign_bce_loss(output['assign1'], output['cos1'])\n",
     "        ld['assign'], ld['assign_entropy'] = l_assign, assign_ent\n",
     "        l_assign_nce, assign_nce_acc = assign_nce_loss(\n",
@@ -985,12 +999,14 @@
     "        ld['nce_tri'], ld['nce_tri_acc'] = l_nce_tri, nce_tri_acc\n",
     "        l_attract, nearest_cos = attraction_loss(output['cos1'])\n",
     "        ld['attract'], ld['nearest_cos'] = l_attract, nearest_cos\n",
-    "        l_cv = cv_loss(emb1, target=self.cv_target, batched=cv_batched)\n",
+    "        # CV: batched Cayley-Menger, n_samples=32 for training speed\n",
+    "        l_cv = cv_loss(emb1, target=self.cv_target, n_samples=32, batched=cv_batched)\n",
     "        ld['cv'] = l_cv\n",
     "        l_spread = spread_loss(self.constellation.anchors)\n",
     "        ld['spread'] = l_spread\n",
-    "        # \u2500\u2500 kNN \u2500\u2500\n",
-    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "        # \u2500\u2500 kNN (SKIP during training \u2014 B\u00d7B matmul is expensive every batch) \u2500\u2500\n",
+    "        if compute_knn:\n",
+    "            ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "        # \u2500\u2500 TOTAL \u2500\u2500\n",
     "        loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
     "        loss_geometric = w_nce_pw * l_nce_pw + w_bridge * l_bridge\n",
@@ -1327,9 +1343,11 @@
     "\n",
     "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,\n",
     "                              num_workers=num_workers, pin_memory=True,\n",
-    "                              drop_last=True, collate_fn=paired_collate)\n",
+    "                              drop_last=True, collate_fn=paired_collate,\n",
+    "                              persistent_workers=(num_workers > 0))\n",
     "    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,\n",
-    "                            num_workers=num_workers, pin_memory=True)\n",
+    "                            num_workers=num_workers, pin_memory=True,\n",
+    "                            persistent_workers=(num_workers > 0))\n",
     "    print(f\"[DATA] CIFAR-10 paired: {len(train_ds)} train, {len(val_ds)} val, bs={batch_size}\")\n",
     "    return train_loader, val_loader\n",
     "\n",
@@ -1460,7 +1478,7 @@
     "            optimizer.zero_grad()\n",
     "\n",
     "            output = model.forward_paired(v1, v2)\n",
-    "            loss, ld = model.compute_loss(output, labels, **lw)\n",
+    "            loss, ld = model.compute_loss(output, labels, compute_knn=False, **lw)\n",
     "\n",
     "            loss.backward()\n",
     "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",

spectral/notebooks/experiment_6_losses_and_anchors.ipynb

@@ -77,8 +77,16 @@
     "from collections import defaultdict\n",
     "\n",
     "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "\n",
+    "# \u2500\u2500 Performance: TF32 + cudnn benchmark \u2500\u2500\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "torch.backends.cudnn.allow_tf32 = True\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "\n",
     "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "if device.type == \"cuda\":\n",
+    "    print(f\"[PERF] TF32={torch.backends.cuda.matmul.allow_tf32}, cudnn.benchmark={torch.backends.cudnn.benchmark}\")\n",
     "\n",
     "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
     "# GEOLIP CORE \u2014 Geometric Building Blocks\n",
@@ -448,8 +456,10 @@
     "    return torch.stack(vols)\n",
     "\n",
     "\n",
-    "def cv_loss(emb, target=0.22, n_samples=64, n_points=5, batched=True):\n",
-    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\"\"\"\n",
+    "def cv_loss(emb, target=0.22, n_samples=32, n_points=5, batched=True):\n",
+    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\n",
+    "    Default n_samples=32 for training speed (141x faster than sequential).\n",
+    "    Use n_samples=200 for monitoring/metrics only.\"\"\"\n",
     "    if emb.shape[0] < n_points:\n",
     "        return torch.tensor(0.0, device=emb.device, requires_grad=True)\n",
     "    vols = _batch_pentachoron_volumes(emb, n_samples, n_points) if batched else _sequential_pentachoron_volumes(emb, n_samples, n_points)\n",
@@ -662,8 +672,9 @@
     "    l_spread = spread_loss(constellation.anchors)\n",
     "    ld['spread'] = l_spread\n",
     "\n",
-    "    # \u2500\u2500 kNN \u2500\u2500\n",
-    "    ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "    # \u2500\u2500 kNN (skip during training for speed \u2014 only compute when explicitly needed) \u2500\u2500\n",
+    "    if targets is not None and emb1.shape[0] <= 512:\n",
+    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "\n",
     "    # \u2500\u2500 TOTAL \u2500\u2500\n",
     "    loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
@@ -956,19 +967,22 @@
     "                     w_assign=0.5, w_assign_nce=0.25,\n",
     "                     w_nce_tri=0.5, w_attract=0.25,\n",
     "                     w_cv=0.01, w_spread=0.01,\n",
-    "                     cv_batched=True):\n",
-    "        \"\"\"Three-domain cooperative loss.\n",
+    "                     cv_batched=True, compute_knn=False):\n",
+    "        \"\"\"Three-domain cooperative loss \u2014 fully batched, zero Python loops.\n",
+    "        Args:\n",
+    "            compute_knn: if False (default), skip kNN during training for speed.\n",
+    "                         Set True during validation or every N steps.\n",
     "        Returns:\n",
     "            total_loss, loss_dict\n",
     "        \"\"\"\n",
     "        ld = {}\n",
     "        emb1, emb2 = output['embedding'], output['embedding_aug']\n",
-    "        # \u2500\u2500 EXTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 EXTERNAL (batched matmul) \u2500\u2500\n",
     "        l_ce, acc = ce_loss_paired(output['logits'], output['logits_aug'], targets)\n",
     "        ld['ce'], ld['acc'] = l_ce, acc\n",
     "        l_nce_emb, nce_emb_acc = nce_loss(emb1, emb2, self.infonce_temp, normalize=False)\n",
     "        ld['nce_emb'], ld['nce_emb_acc'] = l_nce_emb, nce_emb_acc\n",
-    "        # \u2500\u2500 GEOMETRIC \u2500\u2500\n",
+    "        # \u2500\u2500 GEOMETRIC (batched matmul) \u2500\u2500\n",
     "        l_nce_pw, nce_pw_acc = nce_loss(\n",
     "            output['patchwork1'], output['patchwork1_aug'], self.assign_temp, normalize=True)\n",
     "        ld['nce_pw'], ld['nce_pw_acc'] = l_nce_pw, nce_pw_acc\n",
@@ -976,7 +990,7 @@
     "            output['bridge1'], output['bridge2'],\n",
     "            output['assign1'], output['assign2'])\n",
     "        ld['bridge'], ld['bridge_acc'] = l_bridge, bridge_acc\n",
-    "        # \u2500\u2500 INTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 INTERNAL (batched \u2014 no Python loops) \u2500\u2500\n",
     "        l_assign, assign_ent = assign_bce_loss(output['assign1'], output['cos1'])\n",
     "        ld['assign'], ld['assign_entropy'] = l_assign, assign_ent\n",
     "        l_assign_nce, assign_nce_acc = assign_nce_loss(\n",
@@ -987,12 +1001,14 @@
     "        ld['nce_tri'], ld['nce_tri_acc'] = l_nce_tri, nce_tri_acc\n",
     "        l_attract, nearest_cos = attraction_loss(output['cos1'])\n",
     "        ld['attract'], ld['nearest_cos'] = l_attract, nearest_cos\n",
-    "        l_cv = cv_loss(emb1, target=self.cv_target, batched=cv_batched)\n",
+    "        # CV: batched Cayley-Menger, n_samples=32 for training speed\n",
+    "        l_cv = cv_loss(emb1, target=self.cv_target, n_samples=32, batched=cv_batched)\n",
     "        ld['cv'] = l_cv\n",
     "        l_spread = spread_loss(self.constellation.anchors)\n",
     "        ld['spread'] = l_spread\n",
-    "        # \u2500\u2500 kNN \u2500\u2500\n",
-    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "        # \u2500\u2500 kNN (SKIP during training \u2014 B\u00d7B matmul is expensive every batch) \u2500\u2500\n",
+    "        if compute_knn:\n",
+    "            ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "        # \u2500\u2500 TOTAL \u2500\u2500\n",
     "        loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
     "        loss_geometric = w_nce_pw * l_nce_pw + w_bridge * l_bridge\n",
@@ -1329,9 +1345,11 @@
     "\n",
     "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,\n",
     "                              num_workers=num_workers, pin_memory=True,\n",
-    "                              drop_last=True, collate_fn=paired_collate)\n",
+    "                              drop_last=True, collate_fn=paired_collate,\n",
+    "                              persistent_workers=(num_workers > 0))\n",
     "    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,\n",
-    "                            num_workers=num_workers, pin_memory=True)\n",
+    "                            num_workers=num_workers, pin_memory=True,\n",
+    "                            persistent_workers=(num_workers > 0))\n",
     "    print(f\"[DATA] CIFAR-10 paired: {len(train_ds)} train, {len(val_ds)} val, bs={batch_size}\")\n",
     "    return train_loader, val_loader\n",
     "\n",
@@ -1462,7 +1480,7 @@
     "            optimizer.zero_grad()\n",
     "\n",
     "            output = model.forward_paired(v1, v2)\n",
-    "            loss, ld = model.compute_loss(output, labels, **lw)\n",
+    "            loss, ld = model.compute_loss(output, labels, compute_knn=False, **lw)\n",
     "\n",
     "            loss.backward()\n",
     "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",

spectral/notebooks/experiment_7_composite_pipelines.ipynb

@@ -75,8 +75,16 @@
     "from collections import defaultdict\n",
     "\n",
     "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "\n",
+    "# \u2500\u2500 Performance: TF32 + cudnn benchmark \u2500\u2500\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "torch.backends.cudnn.allow_tf32 = True\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "\n",
     "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "if device.type == \"cuda\":\n",
+    "    print(f\"[PERF] TF32={torch.backends.cuda.matmul.allow_tf32}, cudnn.benchmark={torch.backends.cudnn.benchmark}\")\n",
     "\n",
     "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
     "# GEOLIP CORE \u2014 Geometric Building Blocks\n",
@@ -446,8 +454,10 @@
     "    return torch.stack(vols)\n",
     "\n",
     "\n",
-    "def cv_loss(emb, target=0.22, n_samples=64, n_points=5, batched=True):\n",
-    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\"\"\"\n",
+    "def cv_loss(emb, target=0.22, n_samples=32, n_points=5, batched=True):\n",
+    "    \"\"\"Differentiable CV loss. Returns (CV - target)\u00b2.\n",
+    "    Default n_samples=32 for training speed (141x faster than sequential).\n",
+    "    Use n_samples=200 for monitoring/metrics only.\"\"\"\n",
     "    if emb.shape[0] < n_points:\n",
     "        return torch.tensor(0.0, device=emb.device, requires_grad=True)\n",
     "    vols = _batch_pentachoron_volumes(emb, n_samples, n_points) if batched else _sequential_pentachoron_volumes(emb, n_samples, n_points)\n",
@@ -660,8 +670,9 @@
     "    l_spread = spread_loss(constellation.anchors)\n",
     "    ld['spread'] = l_spread\n",
     "\n",
-    "    # \u2500\u2500 kNN \u2500\u2500\n",
-    "    ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "    # \u2500\u2500 kNN (skip during training for speed \u2014 only compute when explicitly needed) \u2500\u2500\n",
+    "    if targets is not None and emb1.shape[0] <= 512:\n",
+    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "\n",
     "    # \u2500\u2500 TOTAL \u2500\u2500\n",
     "    loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
@@ -954,19 +965,22 @@
     "                     w_assign=0.5, w_assign_nce=0.25,\n",
     "                     w_nce_tri=0.5, w_attract=0.25,\n",
     "                     w_cv=0.01, w_spread=0.01,\n",
-    "                     cv_batched=True):\n",
-    "        \"\"\"Three-domain cooperative loss.\n",
+    "                     cv_batched=True, compute_knn=False):\n",
+    "        \"\"\"Three-domain cooperative loss \u2014 fully batched, zero Python loops.\n",
+    "        Args:\n",
+    "            compute_knn: if False (default), skip kNN during training for speed.\n",
+    "                         Set True during validation or every N steps.\n",
     "        Returns:\n",
     "            total_loss, loss_dict\n",
     "        \"\"\"\n",
     "        ld = {}\n",
     "        emb1, emb2 = output['embedding'], output['embedding_aug']\n",
-    "        # \u2500\u2500 EXTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 EXTERNAL (batched matmul) \u2500\u2500\n",
     "        l_ce, acc = ce_loss_paired(output['logits'], output['logits_aug'], targets)\n",
     "        ld['ce'], ld['acc'] = l_ce, acc\n",
     "        l_nce_emb, nce_emb_acc = nce_loss(emb1, emb2, self.infonce_temp, normalize=False)\n",
     "        ld['nce_emb'], ld['nce_emb_acc'] = l_nce_emb, nce_emb_acc\n",
-    "        # \u2500\u2500 GEOMETRIC \u2500\u2500\n",
+    "        # \u2500\u2500 GEOMETRIC (batched matmul) \u2500\u2500\n",
     "        l_nce_pw, nce_pw_acc = nce_loss(\n",
     "            output['patchwork1'], output['patchwork1_aug'], self.assign_temp, normalize=True)\n",
     "        ld['nce_pw'], ld['nce_pw_acc'] = l_nce_pw, nce_pw_acc\n",
@@ -974,7 +988,7 @@
     "            output['bridge1'], output['bridge2'],\n",
     "            output['assign1'], output['assign2'])\n",
     "        ld['bridge'], ld['bridge_acc'] = l_bridge, bridge_acc\n",
-    "        # \u2500\u2500 INTERNAL \u2500\u2500\n",
+    "        # \u2500\u2500 INTERNAL (batched \u2014 no Python loops) \u2500\u2500\n",
     "        l_assign, assign_ent = assign_bce_loss(output['assign1'], output['cos1'])\n",
     "        ld['assign'], ld['assign_entropy'] = l_assign, assign_ent\n",
     "        l_assign_nce, assign_nce_acc = assign_nce_loss(\n",
@@ -985,12 +999,14 @@
     "        ld['nce_tri'], ld['nce_tri_acc'] = l_nce_tri, nce_tri_acc\n",
     "        l_attract, nearest_cos = attraction_loss(output['cos1'])\n",
     "        ld['attract'], ld['nearest_cos'] = l_attract, nearest_cos\n",
-    "        l_cv = cv_loss(emb1, target=self.cv_target, batched=cv_batched)\n",
+    "        # CV: batched Cayley-Menger, n_samples=32 for training speed\n",
+    "        l_cv = cv_loss(emb1, target=self.cv_target, n_samples=32, batched=cv_batched)\n",
     "        ld['cv'] = l_cv\n",
     "        l_spread = spread_loss(self.constellation.anchors)\n",
     "        ld['spread'] = l_spread\n",
-    "        # \u2500\u2500 kNN \u2500\u2500\n",
-    "        ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
+    "        # \u2500\u2500 kNN (SKIP during training \u2014 B\u00d7B matmul is expensive every batch) \u2500\u2500\n",
+    "        if compute_knn:\n",
+    "            ld['knn_acc'] = knn_accuracy(emb1, targets)\n",
     "        # \u2500\u2500 TOTAL \u2500\u2500\n",
     "        loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb\n",
     "        loss_geometric = w_nce_pw * l_nce_pw + w_bridge * l_bridge\n",
@@ -1327,9 +1343,11 @@
     "\n",
     "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,\n",
     "                              num_workers=num_workers, pin_memory=True,\n",
-    "                              drop_last=True, collate_fn=paired_collate)\n",
+    "                              drop_last=True, collate_fn=paired_collate,\n",
+    "                              persistent_workers=(num_workers > 0))\n",
     "    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,\n",
-    "                            num_workers=num_workers, pin_memory=True)\n",
+    "                            num_workers=num_workers, pin_memory=True,\n",
+    "                            persistent_workers=(num_workers > 0))\n",
     "    print(f\"[DATA] CIFAR-10 paired: {len(train_ds)} train, {len(val_ds)} val, bs={batch_size}\")\n",
     "    return train_loader, val_loader\n",
     "\n",
@@ -1460,7 +1478,7 @@
     "            optimizer.zero_grad()\n",
     "\n",
     "            output = model.forward_paired(v1, v2)\n",
-    "            loss, ld = model.compute_loss(output, labels, **lw)\n",
+    "            loss, ld = model.compute_loss(output, labels, compute_knn=False, **lw)\n",
     "\n",
     "            loss.backward()\n",
     "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",