Create cell2_model_v10.py

cell2_model_v10.py

@@ -0,0 +1,350 @@
+"""
+Superposition Patch Classifier - Two-Tier Gated Transformer
+=============================================================
+Colab Cell 2 of 3 - depends on Cell 1 (generator.py) namespace.
+
+Architecture:
+  voxels → patch_embed → e₀
+
+  Stage 0 (local gates): From raw embeddings, no attention
+    e₀ → local_dim_head    → dim_soft    ─┐
+    e₀ → local_curv_head   → curv_soft   ─┤ LOCAL_GATE_DIM = 11
+    e₀ → local_bound_head  → bound_soft  ─┤
+    e₀ → local_axis_head   → axis_soft   ─┘→ local_gates (detached)
+
+  Stage 1 (bootstrap): Attention sees local gates
+    proj([e₀, local_gates]) → bootstrap_block × N → h
+
+  Stage 1.5 (structural gates): From h, after cross-patch context
+    h → struct_topo_head     → topo_soft    ─┐
+    h → struct_neighbor_head → neighbor_soft ─┤ STRUCTURAL_GATE_DIM = 6
+    h → struct_role_head     → role_soft     ─┘→ structural_gates (detached)
+
+  Stage 2 (geometric routing): Both gate tiers
+    (h, local_gates, structural_gates) → geometric_block × N → h'
+
+  Stage 3 (classification): Gated shape heads
+    [h', local_gates, structural_gates] → shape_heads
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Cell 1 provides: all constants including LOCAL_GATE_DIM, STRUCTURAL_GATE_DIM, TOTAL_GATE_DIM
+
+
+# === Patch Embedding ==========================================================
+
+class PatchEmbedding3D(nn.Module):
+    def __init__(self, patch_dim=64):
+        super().__init__()
+        self.proj = nn.Linear(PATCH_VOL, patch_dim)
+        pz = torch.arange(MACRO_Z).float() / MACRO_Z
+        py = torch.arange(MACRO_Y).float() / MACRO_Y
+        px = torch.arange(MACRO_X).float() / MACRO_X
+        pos = torch.stack(torch.meshgrid(pz, py, px, indexing='ij'), dim=-1).reshape(MACRO_N, 3)
+        self.register_buffer('pos_embed', pos)
+        self.pos_proj = nn.Linear(3, patch_dim)
+
+    def forward(self, x):
+        B = x.shape[0]
+        patches = x.view(B, MACRO_Z, PATCH_Z, MACRO_Y, PATCH_Y, MACRO_X, PATCH_X)
+        patches = patches.permute(0, 1, 3, 5, 2, 4, 6).contiguous().view(B, MACRO_N, PATCH_VOL)
+        return self.proj(patches) + self.pos_proj(self.pos_embed)
+
+
+# === Standard Transformer Block ===============================================
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, n_heads, dropout=0.1):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True)
+        self.ff = nn.Sequential(
+            nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim), nn.Dropout(dropout)
+        )
+        self.ln1, self.ln2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x), self.ln1(x), self.ln1(x))[0]
+        return x + self.ff(self.ln2(x))
+
+
+# === Geometric Gated Attention ================================================
+
+class GatedGeometricAttention(nn.Module):
+    """
+    Multi-head attention with two-tier gate modulation.
+    Q, K see both local and structural gates.
+    V modulated by combined gate vector.
+    Per-head compatibility bias from gate interactions.
+    """
+
+    def __init__(self, embed_dim, gate_dim, n_heads, dropout=0.1):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.n_heads = n_heads
+        self.head_dim = embed_dim // n_heads
+
+        # Q, K from [h, all_gates]
+        self.q_proj = nn.Linear(embed_dim + gate_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim + gate_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+
+        # Per-head gate compatibility
+        self.gate_q = nn.Linear(gate_dim, n_heads)
+        self.gate_k = nn.Linear(gate_dim, n_heads)
+
+        # Value modulation by gates
+        self.v_gate = nn.Sequential(nn.Linear(gate_dim, embed_dim), nn.Sigmoid())
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.attn_drop = nn.Dropout(dropout)
+        self.scale = math.sqrt(self.head_dim)
+
+    def forward(self, h, gate_features):
+        B, N, _ = h.shape
+        hg = torch.cat([h, gate_features], dim=-1)
+        Q = self.q_proj(hg).view(B, N, self.n_heads, self.head_dim).transpose(1, 2)
+        K = self.k_proj(hg).view(B, N, self.n_heads, self.head_dim).transpose(1, 2)
+
+        V = self.v_proj(h)
+        V = (V * self.v_gate(gate_features)).view(B, N, self.n_heads, self.head_dim).transpose(1, 2)
+
+        content_scores = (Q @ K.transpose(-2, -1)) / self.scale
+        gq = self.gate_q(gate_features)
+        gk = self.gate_k(gate_features)
+        compat = torch.einsum('bih,bjh->bhij', gq, gk)
+
+        attn = F.softmax(content_scores + compat, dim=-1)
+        attn = self.attn_drop(attn)
+
+        out = (attn @ V).transpose(1, 2).reshape(B, N, self.embed_dim)
+        return self.out_proj(out)
+
+
+class GeometricTransformerBlock(nn.Module):
+    def __init__(self, embed_dim, gate_dim, n_heads, dropout=0.1, ff_mult=4):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(embed_dim)
+        self.attn = GatedGeometricAttention(embed_dim, gate_dim, n_heads, dropout)
+        self.ln2 = nn.LayerNorm(embed_dim)
+        self.ff = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim * ff_mult), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(embed_dim * ff_mult, embed_dim), nn.Dropout(dropout)
+        )
+
+    def forward(self, h, gate_features):
+        h = h + self.attn(self.ln1(h), gate_features)
+        h = h + self.ff(self.ln2(h))
+        return h
+
+
+# === Main Classifier ==========================================================
+
+class SuperpositionPatchClassifier(nn.Module):
+    """
+    Two-tier gated transformer for multi-shape superposition.
+
+    Tier 1 (local): Gates from raw patch embeddings — what IS in this patch
+    Tier 2 (structural): Gates from post-attention h — what ROLE this patch plays
+
+    Both tiers feed into geometric attention and classification.
+    """
+
+    def __init__(self, embed_dim=128, patch_dim=64, n_bootstrap=2, n_geometric=2,
+                 n_heads=4, dropout=0.1):
+        super().__init__()
+        self.embed_dim = embed_dim
+
+        # Patch embedding
+        self.patch_embed = PatchEmbedding3D(patch_dim)
+
+        # === Stage 0: Local encoder + gate heads (pre-attention) ===
+        # Shared MLP gives local heads enough capacity to extract
+        # dims/curvature/boundary from 32 voxels without cross-patch info
+        local_hidden = patch_dim * 2  # 128
+        self.local_encoder = nn.Sequential(
+            nn.Linear(patch_dim, local_hidden), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(local_hidden, local_hidden), nn.GELU(), nn.Dropout(dropout),
+        )
+        self.local_dim_head = nn.Linear(local_hidden, NUM_LOCAL_DIMS)
+        self.local_curv_head = nn.Linear(local_hidden, NUM_LOCAL_CURVS)
+        self.local_bound_head = nn.Linear(local_hidden, NUM_LOCAL_BOUNDARY)
+        self.local_axis_head = nn.Linear(local_hidden, NUM_LOCAL_AXES)
+
+        # Project [embedding, local_gates] → embed_dim for bootstrap
+        self.proj = nn.Linear(patch_dim + LOCAL_GATE_DIM, embed_dim)
+
+        # === Stage 1: Bootstrap blocks (attention with local gate context) ===
+        self.bootstrap_blocks = nn.ModuleList([
+            TransformerBlock(embed_dim, n_heads, dropout)
+            for _ in range(n_bootstrap)
+        ])
+
+        # === Stage 1.5: Structural gate heads (from h, post-attention) ===
+        self.struct_topo_head = nn.Linear(embed_dim, NUM_STRUCT_TOPO)
+        self.struct_neighbor_head = nn.Linear(embed_dim, NUM_STRUCT_NEIGHBOR)
+        self.struct_role_head = nn.Linear(embed_dim, NUM_STRUCT_ROLE)
+
+        # === Stage 2: Geometric gated blocks (see both gate tiers) ===
+        self.geometric_blocks = nn.ModuleList([
+            GeometricTransformerBlock(embed_dim, TOTAL_GATE_DIM, n_heads, dropout)
+            for _ in range(n_geometric)
+        ])
+
+        # === Stage 3: Gated classification ===
+        gated_dim = embed_dim + TOTAL_GATE_DIM
+
+        self.patch_shape_head = nn.Sequential(
+            nn.Linear(gated_dim, embed_dim), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(embed_dim, NUM_CLASSES)
+        )
+
+        self.global_pool = nn.Sequential(
+            nn.Linear(gated_dim, embed_dim), nn.GELU(),
+            nn.Linear(embed_dim, embed_dim)
+        )
+        self.global_gate_head = nn.Linear(embed_dim, NUM_GATES)
+        self.global_shape_head = nn.Linear(embed_dim, NUM_CLASSES)
+
+    def forward(self, x):
+        # === Raw patch embedding ===
+        e = self.patch_embed(x)  # (B, 64, patch_dim)
+
+        # === Stage 0: Local gates from raw embedding via local encoder ===
+        e_local = self.local_encoder(e)  # (B, 64, local_hidden)
+        local_dim_logits = self.local_dim_head(e_local)
+        local_curv_logits = self.local_curv_head(e_local)
+        local_bound_logits = self.local_bound_head(e_local)
+        local_axis_logits = self.local_axis_head(e_local)
+
+        local_gates = torch.cat([
+            F.softmax(local_dim_logits, dim=-1),
+            F.softmax(local_curv_logits, dim=-1),
+            torch.sigmoid(local_bound_logits),
+            torch.sigmoid(local_axis_logits),
+        ], dim=-1)  # (B, 64, 11)
+
+        # === Stage 1: Bootstrap with local gate context ===
+        h = self.proj(torch.cat([e, local_gates], dim=-1))
+        for blk in self.bootstrap_blocks:
+            h = blk(h)
+
+        # === Stage 1.5: Structural gates from h (after cross-patch context) ===
+        struct_topo_logits = self.struct_topo_head(h)
+        struct_neighbor_logits = self.struct_neighbor_head(h)
+        struct_role_logits = self.struct_role_head(h)
+
+        structural_gates = torch.cat([
+            F.softmax(struct_topo_logits, dim=-1),
+            torch.sigmoid(struct_neighbor_logits),
+            F.softmax(struct_role_logits, dim=-1),
+        ], dim=-1)  # (B, 64, 6)
+
+        # === Combined gate vector ===
+        all_gates = torch.cat([local_gates, structural_gates], dim=-1)  # (B, 64, 17)
+
+        # === Stage 2: Geometric gated transformer ===
+        for blk in self.geometric_blocks:
+            h = blk(h, all_gates)
+
+        # === Stage 3: Classification from gated representations ===
+        h_gated = torch.cat([h, all_gates], dim=-1)
+        shape_logits = self.patch_shape_head(h_gated)
+        g = self.global_pool(h_gated.mean(dim=1))
+
+        return {
+            # Local gate predictions (Stage 0)
+            "local_dim_logits": local_dim_logits,
+            "local_curv_logits": local_curv_logits,
+            "local_bound_logits": local_bound_logits,
+            "local_axis_logits": local_axis_logits,
+
+            # Structural gate predictions (Stage 1.5)
+            "struct_topo_logits": struct_topo_logits,
+            "struct_neighbor_logits": struct_neighbor_logits,
+            "struct_role_logits": struct_role_logits,
+
+            # Shape predictions (Stage 3)
+            "patch_shape_logits": shape_logits,
+            "patch_features": h,
+            "global_features": g,
+            "global_gates": self.global_gate_head(g),
+            "global_shapes": self.global_shape_head(g),
+        }
+
+
+# === Loss =====================================================================
+
+class SuperpositionLoss(nn.Module):
+    def __init__(self, local_weight=1.0, struct_weight=1.0, shape_weight=1.0, global_weight=0.5):
+        super().__init__()
+        self.lw, self.sw, self.shw, self.gw = local_weight, struct_weight, shape_weight, global_weight
+
+    def forward(self, outputs, targets):
+        occ_mask = targets["patch_occupancy"] > 0.01
+        n_occ = occ_mask.sum().clamp(min=1)
+
+        # --- Local gate losses ---
+        dim_loss = F.cross_entropy(
+            outputs["local_dim_logits"].view(-1, NUM_LOCAL_DIMS),
+            targets["patch_dims"].clamp(0, NUM_LOCAL_DIMS - 1).view(-1),
+            reduction='none').view_as(occ_mask)
+        curv_loss = F.cross_entropy(
+            outputs["local_curv_logits"].view(-1, NUM_LOCAL_CURVS),
+            targets["patch_curvature"].clamp(0, NUM_LOCAL_CURVS - 1).view(-1),
+            reduction='none').view_as(occ_mask)
+        bound_loss = F.binary_cross_entropy_with_logits(
+            outputs["local_bound_logits"].squeeze(-1),
+            targets["patch_boundary"],
+            reduction='none')
+        axis_loss = F.binary_cross_entropy_with_logits(
+            outputs["local_axis_logits"],
+            targets["patch_axis_active"],
+            reduction='none').mean(dim=-1)
+
+        local_loss = ((dim_loss + curv_loss + bound_loss + axis_loss) * occ_mask.float()).sum() / n_occ
+
+        # --- Structural gate losses ---
+        topo_loss = F.cross_entropy(
+            outputs["struct_topo_logits"].view(-1, NUM_STRUCT_TOPO),
+            targets["patch_topology"].clamp(0, NUM_STRUCT_TOPO - 1).view(-1),
+            reduction='none').view_as(occ_mask)
+        neighbor_loss = F.mse_loss(
+            torch.sigmoid(outputs["struct_neighbor_logits"].squeeze(-1)),
+            targets["patch_neighbor_count"],
+            reduction='none')
+        role_loss = F.cross_entropy(
+            outputs["struct_role_logits"].view(-1, NUM_STRUCT_ROLE),
+            targets["patch_surface_role"].clamp(0, NUM_STRUCT_ROLE - 1).view(-1),
+            reduction='none').view_as(occ_mask)
+
+        struct_loss = ((topo_loss + neighbor_loss + role_loss) * occ_mask.float()).sum() / n_occ
+
+        # --- Shape losses ---
+        shape_loss = F.binary_cross_entropy_with_logits(
+            outputs["patch_shape_logits"],
+            targets["patch_shape_membership"],
+            reduction='none').mean(dim=-1)
+        shape_loss = (shape_loss * occ_mask.float()).sum() / n_occ
+
+        # --- Global losses ---
+        global_gate_loss = F.binary_cross_entropy_with_logits(outputs["global_gates"], targets["global_gates"])
+        global_shape_loss = F.binary_cross_entropy_with_logits(outputs["global_shapes"], targets["global_shapes"])
+        global_loss = global_gate_loss + global_shape_loss
+
+        total = self.lw * local_loss + self.sw * struct_loss + self.shw * shape_loss + self.gw * global_loss
+
+        return {
+            "total": total,
+            "local": local_loss,
+            "struct": struct_loss,
+            "shape": shape_loss,
+            "global": global_loss,
+        }
+
+
+print("✓ Model ready (Two-Tier Gated Transformer)")