Add two-stage TrajMixer mixing

backbones.py

@@ -177,7 +177,7 @@ class TemporalAttention(nn.Module):
 
 
 class TrajMixer(nn.Module):
-    """Lightweight gated interaction across latent residual-space groups.
+    """Two-stage gated mixing within and across latent trajectory groups.
 
     The groups are contiguous partitions of the post-``W_O`` residual
     representation. They are deliberately not treated as attention heads.
@@ -205,8 +205,24 @@ class TrajMixer(nn.Module):
         # are still residual-space partitions rather than attention heads.
         self.n_group = n_head
         self.d_group = n_embd // n_head
+        self.intra_hidden = 4 * self.d_group
         self.hidden_group = 4 * n_head
 
+        # Stage 1: each group independently mixes its internal features.
+        self.intra_norm = nn.LayerNorm(self.d_group)
+        self.intra_gate_proj = nn.Parameter(
+            torch.empty(self.n_group, self.d_group, self.intra_hidden)
+        )
+        self.intra_value_proj = nn.Parameter(
+            torch.empty(self.n_group, self.d_group, self.intra_hidden)
+        )
+        self.intra_output_proj = nn.Parameter(
+            torch.empty(self.n_group, self.intra_hidden, self.d_group)
+        )
+
+        # Stage 2: each internal coordinate independently mixes groups.
+        self.cross_norm = nn.LayerNorm(self.n_group)
+
         # Per-group feature alignment: [group, input feature, output feature].
         self.group_align = nn.Parameter(
             torch.empty(self.n_group, self.d_group, self.d_group)
@@ -227,6 +243,11 @@ class TrajMixer(nn.Module):
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
+        for group_idx in range(self.n_group):
+            nn.init.xavier_uniform_(self.intra_gate_proj[group_idx])
+            nn.init.xavier_uniform_(self.intra_value_proj[group_idx])
+        nn.init.normal_(self.intra_output_proj, mean=0.0, std=1e-3)
+
         with torch.no_grad():
             identity = torch.eye(
                 self.d_group,
@@ -243,7 +264,7 @@ class TrajMixer(nn.Module):
         nn.init.normal_(self.output_proj, mean=0.0, std=1e-3)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Map ``(B, L, n_embd)`` to an equally shaped residual update."""
+        """Apply two PreNorm residual stages without mixing sequence positions."""
         if x.ndim != 3:
             raise ValueError(f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}")
         if x.size(-1) != self.n_embd:
@@ -255,8 +276,27 @@ class TrajMixer(nn.Module):
         grouped = x.reshape(
             batch_size, seq_len, self.n_group, self.d_group
         )
+
+        # Stage 1: d_group -> 4*d_group -> d_group, independently per group.
+        intra_input = self.intra_norm(grouped)
+        intra_gate = torch.einsum(
+            "blgd,gdh->blgh", intra_input, self.intra_gate_proj
+        )
+        intra_value = torch.einsum(
+            "blgd,gdh->blgh", intra_input, self.intra_value_proj
+        )
+        intra_hidden = F.silu(intra_gate) * intra_value
+        intra_update = torch.einsum(
+            "blgh,ghd->blgd", intra_hidden, self.intra_output_proj
+        )
+        grouped = grouped + self.drop(intra_update)
+
+        # Stage 2: n_group -> 4*n_group -> n_group for each coordinate.
+        cross_input = self.cross_norm(
+            grouped.transpose(-1, -2)
+        ).transpose(-1, -2)
         aligned = torch.einsum(
-            "blgd,gde->blge", grouped, self.group_align
+            "blgd,gde->blge", cross_input, self.group_align
         )
 
         gate = torch.einsum(
@@ -269,7 +309,8 @@ class TrajMixer(nn.Module):
         mixed = torch.einsum(
             "blhr,rhg->blgr", hidden, self.output_proj
         )
-        return self.drop(mixed.reshape(batch_size, seq_len, self.n_embd))
+        grouped = grouped + self.drop(mixed)
+        return grouped.reshape(batch_size, seq_len, self.n_embd)
 
 
 class GPTBlock(nn.Module):
@@ -299,7 +340,6 @@ class GPTBlock(nn.Module):
             dropout=mlp_dropout,
         )
         self.ln1 = nn.LayerNorm(n_embd)
-        self.ln2 = nn.LayerNorm(n_embd)
 
     def forward(
         self,
@@ -309,8 +349,7 @@ class GPTBlock(nn.Module):
         attn_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
-        x = x + self.mlp(self.ln2(x))
-        return x
+        return self.mlp(x)
 
 
 class TokenAutoDiscretization(nn.Module):