Upload 6 files

spectral/notebooks/experiment_2_manifold_structures.ipynb

@@ -1766,47 +1766,41 @@
    "source": [
     "# @title Experiment 2.2 \u2014 Grassmannian Subspace Features\n",
     "class GrassmannianFrontEnd(nn.Module):\n",
-    "    \"\"\"Grassmannian subspace features via SVD computed through 3\u00d73 eigh path.\n",
-    "    X = U S V^T. Since X^T X = V S\u00b2 V^T, we get S,V from eigh(X^T X),\n",
-    "    then U = X V S\u207b\u00b9. Full SVD result, but via 3\u00d73 eigendecomposition\n",
-    "    instead of iterative cusolver on (64,3). Instant on GPU.\"\"\"\n",
+    "    \"\"\"Grassmannian subspace features via direct SVD.\n",
+    "    X = U S Vh. Features: singular values (spectral profile),\n",
+    "    log singular value ratios (relative spectrum), and right singular\n",
+    "    vectors V (subspace orientation in channel space \u2014 the actual\n",
+    "    Grassmannian coordinate). V varies per patch and encodes which\n",
+    "    linear combinations of RGB correspond to principal directions.\"\"\"\n",
     "    def __init__(self, patch_size=8, k=3, input_size=32):\n",
     "        super().__init__()\n",
     "        self.patch_size = patch_size\n",
     "        self.k = k\n",
+    "        self.C = 3  # RGB\n",
     "        n_patches_h = input_size // patch_size\n",
     "        self.n_patches = n_patches_h * n_patches_h\n",
-    "        # k singular values + k log-ratios + k*(k+1)/2 gram upper-tri of U\n",
-    "        self.features_per_patch = k + k + k * (k + 1) // 2\n",
+    "        # k singular values + k log-ratios + k*C right singular vector entries\n",
+    "        self.features_per_patch = k + k + k * self.C\n",
     "        self.output_dim = self.n_patches * self.features_per_patch\n",
-    "        print(f\"[GRASS] {self.n_patches} patches, k={k}, dim={self.output_dim} (SVD via 3\u00d73 eigh)\")\n",
+    "        print(f\"[GRASS] {self.n_patches} patches, k={k}, dim={self.output_dim} (direct SVD)\")\n",
     "\n",
+    "    @torch.amp.custom_fwd(device_type='cuda', cast_inputs=torch.float32)\n",
     "    def forward(self, x):\n",
     "        B, C, H, W = x.shape\n",
     "        ps = self.patch_size\n",
     "        patches = x.unfold(2, ps, ps).unfold(3, ps, ps)\n",
     "        n_p = patches.shape[2] * patches.shape[3]\n",
     "        # X: (B*n_p, ps*ps, C) \u2014 each patch as a tall-skinny matrix\n",
-    "        X = patches.permute(0, 2, 3, 1, 4, 5).reshape(B * n_p, C, ps * ps).permute(0, 2, 1).float()\n",
-    "        # X^T X: (B*n_p, C, C) = (B*n_p, 3, 3)\n",
-    "        XtX = torch.bmm(X.transpose(1, 2), X)\n",
-    "        # eigh on 3\u00d73 symmetric PSD \u2192 eigenvalues \u03bb (ascending), eigenvectors V\n",
-    "        eigvals, V = torch.linalg.eigh(XtX)  # (B*n_p, C), (B*n_p, C, C)\n",
-    "        # Flip to descending order\n",
-    "        eigvals = eigvals.flip(-1)\n",
-    "        V = V.flip(-1)\n",
-    "        # S = sqrt(\u03bb), these ARE the singular values\n",
-    "        S = torch.sqrt(eigvals.clamp(min=1e-10))[:, :self.k]\n",
-    "        # U = X V S\u207b\u00b9 (recover left singular vectors)\n",
-    "        Vk = V[:, :, :self.k]  # (B*n_p, C, k)\n",
-    "        XV = torch.bmm(X, Vk)  # (B*n_p, ps*ps, k)\n",
-    "        U = XV / (S.unsqueeze(1) + 1e-10)  # (B*n_p, ps*ps, k)\n",
-    "        # Features: singular values + log-ratios + gram of U\n",
+    "        X = patches.permute(0, 2, 3, 1, 4, 5).reshape(B * n_p, C, ps * ps).permute(0, 2, 1)\n",
+    "        # Direct thin SVD: X = U S Vh, U:(N,64,3) S:(N,3) Vh:(N,3,3)\n",
+    "        U, S, Vh = torch.linalg.svd(X, full_matrices=False)\n",
+    "        S = S[:, :self.k]\n",
+    "        # Log singular value ratios (scale-invariant spectrum)\n",
     "        sv_ratios = torch.log(S / (S[:, -1:] + 1e-8) + 1e-8)\n",
-    "        gram = torch.bmm(U.transpose(1, 2), U)  # (B*n_p, k, k)\n",
-    "        triu_idx = torch.triu_indices(self.k, self.k)\n",
-    "        gram_feat = gram[:, triu_idx[0], triu_idx[1]]\n",
-    "        feats = torch.cat([S, sv_ratios, gram_feat], dim=-1)\n",
+    "        # Right singular vectors Vh[:k]: subspace orientation in channel space\n",
+    "        # This IS the Grassmannian coordinate \u2014 varies meaningfully per patch\n",
+    "        V_feat = Vh[:, :self.k, :].reshape(-1, self.k * C)\n",
+    "        feats = torch.cat([S, sv_ratios, V_feat], dim=-1)\n",
     "        return feats.reshape(B, -1)\n",
     "\n",
     "front = GrassmannianFrontEnd(patch_size=8, k=3).to(device)\n",
@@ -1840,9 +1834,10 @@
    "source": [
     "# @title Experiment 2.3 \u2014 Flag Manifold\n",
     "class FlagManifoldFrontEnd(nn.Module):\n",
-    "    \"\"\"Cascading SVD at multiple truncation levels via 3\u00d73 eigh path.\n",
+    "    \"\"\"Cascading SVD at multiple truncation levels via direct SVD.\n",
     "    Nested subspace features: singular values + projection norms at each flag level.\n",
-    "    SVD computed through eigh(X^T X) \u2014 no iterative cusolver.\"\"\"\n",
+    "    The flag structure captures how information distributes across\n",
+    "    nested subspace hierarchies \u2014 a genuine flag manifold signature.\"\"\"\n",
     "    def __init__(self, patch_size=8, levels=(1, 2, 3), input_size=32):\n",
     "        super().__init__()\n",
     "        self.patch_size = patch_size\n",
@@ -1852,23 +1847,18 @@
     "        max_sv = min(3, patch_size * patch_size)\n",
     "        self.features_per_patch = sum(min(k, max_sv) * 2 for k in levels)\n",
     "        self.output_dim = self.n_patches * self.features_per_patch\n",
-    "        print(f\"[FLAG] {self.n_patches} patches, levels={levels}, dim={self.output_dim} (SVD via 3\u00d73 eigh)\")\n",
+    "        print(f\"[FLAG] {self.n_patches} patches, levels={levels}, dim={self.output_dim} (direct SVD)\")\n",
     "\n",
+    "    @torch.amp.custom_fwd(device_type='cuda', cast_inputs=torch.float32)\n",
     "    def forward(self, x):\n",
     "        B, C, H, W = x.shape\n",
     "        ps = self.patch_size\n",
     "        patches = x.unfold(2, ps, ps).unfold(3, ps, ps)\n",
     "        n_p = patches.shape[2] * patches.shape[3]\n",
     "        # X: (B*n_p, ps*ps, C)\n",
-    "        X = patches.permute(0, 2, 3, 1, 4, 5).reshape(B * n_p, C, ps * ps).permute(0, 2, 1).float()\n",
-    "        # X^T X: (B*n_p, 3, 3)\n",
-    "        XtX = torch.bmm(X.transpose(1, 2), X)\n",
-    "        eigvals, V = torch.linalg.eigh(XtX)\n",
-    "        eigvals = eigvals.flip(-1)\n",
-    "        V = V.flip(-1)\n",
-    "        S = torch.sqrt(eigvals.clamp(min=1e-10))  # (B*n_p, C)\n",
-    "        # U = X V S\u207b\u00b9\n",
-    "        U = torch.bmm(X, V) / (S.unsqueeze(1) + 1e-10)  # (B*n_p, ps*ps, C)\n",
+    "        X = patches.permute(0, 2, 3, 1, 4, 5).reshape(B * n_p, C, ps * ps).permute(0, 2, 1)\n",
+    "        # Direct thin SVD\n",
+    "        U, S, Vh = torch.linalg.svd(X, full_matrices=False)\n",
     "        # Features at each flag level\n",
     "        feats = []\n",
     "        for k in self.levels:\n",

spectral/notebooks/experiment_4_invertible_transforms.ipynb

@@ -1977,12 +1977,13 @@
     "        self.output_dim = n_templates * (2 + min(3, patch_dim))\n",
     "        print(f\"[PROCRUSTES] {n_templates} templates, dim={self.output_dim}\")\n",
     "\n",
+    "    @torch.amp.custom_fwd(device_type='cuda', cast_inputs=torch.float32)\n",
     "    def forward(self, x):\n",
     "        B, C, H, W = x.shape\n",
     "        ps = self.patch_size\n",
     "        patches = x.unfold(2, ps, ps).unfold(3, ps, ps)\n",
     "        patches = patches.contiguous().reshape(B, self.n_patches, -1)\n",
-    "        patches_n = F.normalize(patches.float(), dim=-1)\n",
+    "        patches_n = F.normalize(patches, dim=-1)\n",
     "\n",
     "        results = []\n",
     "        for t in range(self.n_templates):\n",
@@ -1990,16 +1991,10 @@
     "            # Cross-covariance M: (B, D, D) where D = patch_dim\n",
     "            M = torch.bmm(patches_n.transpose(1, 2),\n",
     "                          template.unsqueeze(0).expand(B, -1, -1))\n",
-    "            # Full SVD via eigh path: M = U S V^T\n",
-    "            # M^T M = V S\u00b2 V^T \u2192 eigh gives S, V\n",
-    "            MtM = torch.bmm(M.transpose(1, 2), M)  # (B, D, D) symmetric PSD\n",
-    "            eigvals, V = torch.linalg.eigh(MtM)\n",
-    "            eigvals = eigvals.flip(-1); V = V.flip(-1)\n",
-    "            S = torch.sqrt(eigvals.clamp(min=1e-10))\n",
-    "            # Recover U = M V S\u207b\u00b9\n",
-    "            U = torch.bmm(M, V) / (S.unsqueeze(1) + 1e-10)\n",
-    "            # Optimal Procrustes rotation R = U V^T\n",
-    "            R_opt = torch.bmm(U, V.transpose(1, 2))  # (B, D, D)\n",
+    "            # Direct SVD of cross-covariance: M = U S Vh\n",
+    "            U, S, Vh = torch.linalg.svd(M, full_matrices=False)\n",
+    "            # Optimal Procrustes rotation R = U Vh (= U V^T)\n",
+    "            R_opt = torch.bmm(U, Vh)  # (B, D, D)\n",
     "            # Features: alignment quality + top singular values + rotation trace\n",
     "            align_quality = S.sum(dim=-1, keepdim=True)\n",
     "            top_s = S[:, :min(3, S.shape[1])]\n",

spectral/notebooks/experiment_5_matrix_decompositions.ipynb

@@ -1716,13 +1716,14 @@
     "        self.output_dim = n_patches * n_upper\n",
     "        print(f\"[QR] {n_patches} patches, k={k}, dim={self.output_dim} (via 3\u00d73 Cholesky)\")\n",
     "\n",
+    "    @torch.amp.custom_fwd(device_type='cuda', cast_inputs=torch.float32)\n",
     "    def forward(self, x):\n",
     "        B, C, H, W = x.shape\n",
     "        ps = self.patch_size\n",
     "        patches = x.unfold(2, ps, ps).unfold(3, ps, ps)\n",
     "        n_p = patches.shape[2] * patches.shape[3]\n",
     "        # X: (B*n_p, ps*ps, C)\n",
-    "        X = patches.permute(0, 2, 3, 1, 4, 5).reshape(B * n_p, C, ps * ps).permute(0, 2, 1).float()\n",
+    "        X = patches.permute(0, 2, 3, 1, 4, 5).reshape(B * n_p, C, ps * ps).permute(0, 2, 1)\n",
     "        # X^T X: (B*n_p, 3, 3)\n",
     "        XtX = torch.bmm(X.transpose(1, 2), X)\n",
     "        # Add small diagonal for numerical stability\n",