Refactor TrajMixer to single residual

backbones.py

@@ -180,7 +180,7 @@ class TemporalAttention(nn.Module):
 
 
 class TrajMixer(nn.Module):
-    """Two-stage gated mixing within and across latent trajectory groups.
+    """PreNorm 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.
@@ -211,8 +211,10 @@ class TrajMixer(nn.Module):
         self.intra_hidden = 4 * self.d_group
         self.hidden_group = 4 * n_head
 
+        # A single full-width PreNorm serves the entire TrajMixer branch.
+        self.norm = nn.LayerNorm(self.n_embd)
+
         # 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)
         )
@@ -222,13 +224,8 @@ class TrajMixer(nn.Module):
         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)
+        self.intra_gate_logits = nn.Parameter(
+            torch.empty(self.n_group, self.d_group)
         )
 
         # Per-feature cross-group projections. The feature index is kept
@@ -249,15 +246,11 @@ class TrajMixer(nn.Module):
         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,
-                dtype=self.group_align.dtype,
-                device=self.group_align.device,
-            )
-            self.group_align.copy_(identity.unsqueeze(0).expand_as(self.group_align))
+            nn.init.xavier_uniform_(self.intra_output_proj[group_idx])
+        nn.init.constant_(
+            self.intra_gate_logits,
+            math.log(0.1 / 0.9),
+        )
 
         # Initialise each feature-specific matrix independently so Xavier's
         # fan-in/fan-out calculation sees a two-dimensional matrix.
@@ -266,8 +259,34 @@ class TrajMixer(nn.Module):
             nn.init.xavier_uniform_(self.value_proj[feature_idx])
         nn.init.normal_(self.output_proj, mean=0.0, std=1e-3)
 
+    def _intra_mix(self, grouped: torch.Tensor) -> torch.Tensor:
+        """Mix features independently inside each residual-space group."""
+        intra_gate = torch.einsum(
+            "blgd,gdh->blgh", grouped, self.intra_gate_proj
+        )
+        intra_value = torch.einsum(
+            "blgd,gdh->blgh", grouped, self.intra_value_proj
+        )
+        intra_hidden = F.silu(intra_gate) * intra_value
+        return torch.einsum(
+            "blgh,ghd->blgd", intra_hidden, self.intra_output_proj
+        )
+
+    def _cross_mix(self, grouped: torch.Tensor) -> torch.Tensor:
+        """Mix groups independently for each within-group coordinate."""
+        gate = torch.einsum(
+            "blgr,rgh->blhr", grouped, self.gate_proj
+        )
+        value = torch.einsum(
+            "blgr,rgh->blhr", grouped, self.value_proj
+        )
+        hidden = F.silu(gate) * value
+        return torch.einsum(
+            "blhr,rhg->blgr", hidden, self.output_proj
+        )
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Apply two PreNorm residual stages without mixing sequence positions."""
+        """Apply one full-width PreNorm and one outer residual update."""
         if x.ndim != 3:
             raise ValueError(f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}")
         if x.size(-1) != self.n_embd:
@@ -276,44 +295,22 @@ class TrajMixer(nn.Module):
             )
 
         batch_size, seq_len, _ = x.shape
-        grouped = x.reshape(
+        grouped = self.norm(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
+        # The static per-channel gate starts at sigmoid(logit) ~= 0.1.
+        intra_output = self._intra_mix(grouped)
+        intra_gate = torch.sigmoid(self.intra_gate_logits).view(
+            1, 1, self.n_group, self.d_group
         )
-        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)
+        mixed_input = grouped + intra_gate * intra_output
 
         # 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", cross_input, self.group_align
+        update = self._cross_mix(mixed_input).reshape(
+            batch_size, seq_len, self.n_embd
         )
-
-        gate = torch.einsum(
-            "blgr,rgh->blhr", aligned, self.gate_proj
-        )
-        value = torch.einsum(
-            "blgr,rgh->blhr", aligned, self.value_proj
-        )
-        hidden = F.silu(gate) * value
-        mixed = torch.einsum(
-            "blhr,rhg->blgr", hidden, self.output_proj
-        )
-        grouped = grouped + self.drop(mixed)
-        return grouped.reshape(batch_size, seq_len, self.n_embd)
+        return x + self.drop(update)
 
 
 class GPTBlock(nn.Module):