Create model_v1.py

model_v1.py

@@ -0,0 +1,218 @@
+"""
+SpectralViT — Pure SpectralCell Transformer
+=============================================
+No conv backbone. No external attention. Stacked SpectralCells are the model.
+
+Architecture:
+    Image → PatchEmbed(4×4) → 64 tokens × embed_dim
+    → Cayley hypersphere positional encoding (multi-plane rotations on S^{d-1})
+    → SpectralCell × depth
+    → LayerNorm → mean pool → classify
+
+Positional encoding on the hypersphere:
+    Each position has K learnable rotation angles in K fixed 2D planes.
+    Rotation in plane (2k, 2k+1) by angle θ:
+        x[2k]   = cos(θ) · x[2k] - sin(θ) · x[2k+1]
+        x[2k+1] = sin(θ) · x[2k] + cos(θ) · x[2k+1]
+    Composing K plane rotations = rich orthogonal rotation.
+    Preserves norms. Operates naturally on S^{d-1}.
+    Learnable angles, not fixed sinusoidal.
+
+SpectralCell and cv_of are in namespace from prior cell execution.
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ── Cayley Hypersphere Positional Encoding ───────────────────────
+
+class CayleyPositionalEncoding(nn.Module):
+    """Multi-plane rotation positional encoding on the hypersphere.
+
+    Each position gets K learnable rotation angles applied in K paired
+    dimension planes. Composing K Givens rotations produces a rich
+    orthogonal transformation that preserves embedding norm.
+
+    For embed_dim=256: K=128 planes, each position has 128 angles.
+    64 positions × 128 angles = 8,192 learnable parameters.
+
+    This is geometrically natural — the SpectralCell projects onto S^{D-1},
+    and Cayley rotations are the native transformations of the hypersphere.
+    """
+    def __init__(self, n_positions, embed_dim):
+        super().__init__()
+        assert embed_dim % 2 == 0, "embed_dim must be even for paired rotations"
+        self.n_positions = n_positions
+        self.embed_dim = embed_dim
+        self.n_planes = embed_dim // 2
+
+        # Learnable rotation angles: (n_positions, n_planes)
+        # Initialize small — near-identity rotation at start
+        self.angles = nn.Parameter(torch.randn(n_positions, self.n_planes) * 0.02)
+
+    def forward(self, x):
+        """x: (B, N, D) → (B, N, D) with position-dependent rotation."""
+        B, N, D = x.shape
+        angles = self.angles[:N]  # (N, K)
+
+        cos_a = angles.cos()  # (N, K)
+        sin_a = angles.sin()  # (N, K)
+
+        # Split into even/odd dimension pairs
+        x_even = x[:, :, 0::2]   # (B, N, K)
+        x_odd = x[:, :, 1::2]    # (B, N, K)
+
+        # Givens rotation per plane per position
+        x_rot_even = cos_a.unsqueeze(0) * x_even - sin_a.unsqueeze(0) * x_odd
+        x_rot_odd = sin_a.unsqueeze(0) * x_even + cos_a.unsqueeze(0) * x_odd
+
+        # Interleave back
+        out = torch.stack([x_rot_even, x_rot_odd], dim=-1)  # (B, N, K, 2)
+        return out.reshape(B, N, D)
+
+
+# ── Patch Embedding ──────────────────────────────────────────────
+
+class PatchEmbed(nn.Module):
+    """Image → patches → linear projection.
+    32×32 with patch_size=4 → 8×8 = 64 tokens.
+    """
+    def __init__(self, img_size=32, patch_size=4, in_channels=3, embed_dim=256):
+        super().__init__()
+        self.n_patches = (img_size // patch_size) ** 2
+        self.proj = nn.Conv2d(in_channels, embed_dim,
+                              kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        # x: (B, 3, H, W) → (B, embed_dim, H/ps, W/ps) → (B, N, embed_dim)
+        return self.proj(x).flatten(2).transpose(1, 2)
+
+
+# ── SpectralViT ─────────────────────────────────────────────────
+
+class SpectralViT(nn.Module):
+    """Pure SpectralCell vision transformer.
+
+    No conv backbone. No external attention.
+    Stacked SpectralCells with Cayley hypersphere positional encoding.
+
+    Args:
+        img_size:    input image size (32 for CIFAR)
+        patch_size:  patch size (4 → 64 tokens)
+        in_channels: input channels (3)
+        embed_dim:   token embedding dimension
+        depth:       number of SpectralCell blocks
+        cell_V:      V parameter for SpectralCell
+        cell_D:      D parameter for SpectralCell
+        cell_hidden: hidden dimension inside each cell
+        cell_depth:  residual MLP depth inside each cell
+        n_cross:     cross-attention layers per cell
+        n_heads:     attention heads in cell cross-attention
+        n_classes:   classification output
+        dropout:     classifier dropout
+    """
+    def __init__(
+        self,
+        img_size=32,
+        patch_size=4,
+        in_channels=3,
+        embed_dim=256,
+        depth=6,
+        cell_V=16,
+        cell_D=16,
+        cell_hidden=256,
+        cell_depth=2,
+        n_cross=2,
+        n_heads=4,
+        n_classes=100,
+        dropout=0.1,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.depth = depth
+        n_patches = (img_size // patch_size) ** 2
+
+        # Patch embedding
+        self.patch_embed = PatchEmbed(img_size, patch_size, in_channels, embed_dim)
+
+        # Cayley hypersphere positional encoding
+        self.pos_enc = CayleyPositionalEncoding(n_patches, embed_dim)
+
+        # Stacked SpectralCells — the entire backbone
+        self.cells = nn.ModuleList([
+            SpectralCell(
+                token_dim=embed_dim, V=cell_V, D=cell_D,
+                hidden=cell_hidden, depth=cell_depth,
+                n_cross=n_cross, n_heads=n_heads,
+                max_alpha=0.2,
+            ) for _ in range(depth)
+        ])
+
+        # Pre-norm before each cell
+        self.norms = nn.ModuleList([
+            nn.LayerNorm(embed_dim) for _ in range(depth)
+        ])
+
+        # Final norm + classifier
+        self.final_norm = nn.LayerNorm(embed_dim)
+        self.classifier = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(embed_dim, n_classes),
+        )
+
+    def forward(self, x):
+        """x: (B, 3, H, W) → dict with logits and cell outputs."""
+        # Patch embed → positional encoding
+        tokens = self.patch_embed(x)          # (B, N, embed_dim)
+        tokens = self.pos_enc(tokens)          # rotated on hypersphere
+
+        # Stacked SpectralCells with residual connections
+        cell_outputs = []
+        for i, (cell, norm) in enumerate(zip(self.cells, self.norms)):
+            normed = norm(tokens)
+            cell_out = cell.format(normed)
+            tokens = tokens + cell_out['output']  # residual
+            cell_outputs.append(cell_out)
+
+        # Pool + classify
+        tokens = self.final_norm(tokens)
+        pooled = tokens.mean(dim=1)               # (B, embed_dim)
+        logits = self.classifier(pooled)
+
+        return {
+            'logits': logits,
+            'cell_outputs': cell_outputs,
+            'tokens': tokens,
+        }
+
+    def get_cross_attn_params(self):
+        """Cross-attention params for separate grad clipping."""
+        params = []
+        for name, p in self.named_parameters():
+            if 'cross_attn' in name:
+                params.append(p)
+        return params
+
+    def summary(self):
+        n_params = sum(p.numel() for p in self.parameters())
+        n_embed = sum(p.numel() for p in self.patch_embed.parameters())
+        n_pos = sum(p.numel() for p in self.pos_enc.parameters())
+        n_cells = sum(p.numel() for p in self.cells.parameters())
+        n_norms = sum(p.numel() for p in self.norms.parameters()) + sum(p.numel() for p in self.final_norm.parameters())
+        n_head = sum(p.numel() for p in self.classifier.parameters())
+        n_cross = sum(p.numel() for p in self.get_cross_attn_params())
+
+        print(f"SpectralViT:")
+        print(f"  Patch embed:  {n_embed:,}")
+        print(f"  Cayley PE:    {n_pos:,} ({self.pos_enc.n_planes} rotation planes × {self.pos_enc.n_positions} positions)")
+        print(f"  Cells ({self.depth}×):  {n_cells:,}  ({n_cells // self.depth:,} per cell)")
+        print(f"  LayerNorms:   {n_norms:,}")
+        print(f"  Classifier:   {n_head:,}")
+        print(f"  Cross-attn:   {n_cross:,} (clipped at 0.5)")
+        print(f"  Total:        {n_params:,}")
+        print(f"  Architecture: PatchEmbed(4×4) → CayleyPE → {self.depth}× SpectralCell → pool → classify")