Implement TrajMixer block

backbones.py

@@ -176,38 +176,106 @@ class TemporalAttention(nn.Module):
         return self.resid_drop(self.out_proj(out))
 
 
-class SwiGLU(nn.Module):
+class TrajMixer(nn.Module):
+    """Lightweight gated interaction across latent residual-space groups.
+
+    The groups are contiguous partitions of the post-``W_O`` residual
+    representation. They are deliberately not treated as attention heads.
+    All operations are position-wise, so the sequence dimension remains fully
+    parallel and no temporal information can leak between positions here.
+    """
+
     def __init__(
         self,
         n_embd: int,
-        hidden_dim: int | None = None,
+        n_head: int = 10,
+        hidden_group: int = 20,
         dropout: float = 0.0,
-        bias: bool = True,
     ):
         super().__init__()
-        hidden_dim = hidden_dim if hidden_dim is not None else int(
-            n_embd * 2.5)
+        if n_embd <= 0:
+            raise ValueError(f"n_embd must be > 0, got {n_embd}")
+        if n_head <= 0:
+            raise ValueError(f"n_head must be > 0, got {n_head}")
+        if n_embd % n_head != 0:
+            raise ValueError(
+                f"n_embd must be divisible by n_head, got {n_embd} and {n_head}"
+            )
+        if hidden_group <= 0:
+            raise ValueError(
+                f"hidden_group must be > 0, got {hidden_group}"
+            )
 
-        self.w1 = nn.Linear(n_embd, hidden_dim, bias=bias)  # gate path
-        self.w2 = nn.Linear(n_embd, hidden_dim, bias=bias)  # value path
-        # output projection
-        self.w3 = nn.Linear(hidden_dim, n_embd, bias=bias)
+        self.n_embd = n_embd
+        # The residual-group count is tied to n_head, but the resulting groups
+        # are still residual-space partitions rather than attention heads.
+        self.n_group = n_head
+        self.d_group = n_embd // n_head
+        self.hidden_group = hidden_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)
+        )
+
+        # Per-feature cross-group projections. The feature index is kept
+        # independent, exactly as specified by the TrajMixer baseline.
+        self.gate_proj = nn.Parameter(
+            torch.empty(self.d_group, self.n_group, hidden_group)
+        )
+        self.value_proj = nn.Parameter(
+            torch.empty(self.d_group, self.n_group, hidden_group)
+        )
+        self.output_proj = nn.Parameter(
+            torch.empty(self.d_group, hidden_group, self.n_group)
+        )
         self.drop = nn.Dropout(dropout)
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
-        """GPT-style parameter initialization for MLP paths."""
-        nn.init.normal_(self.w1.weight, mean=0.0, std=0.02)
-        nn.init.normal_(self.w2.weight, mean=0.0, std=0.02)
-        nn.init.normal_(self.w3.weight, mean=0.0, std=0.02)
-        if self.w1.bias is not None:
-            nn.init.zeros_(self.w1.bias)
-            nn.init.zeros_(self.w2.bias)
-            nn.init.zeros_(self.w3.bias)
+        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))
+
+        # Initialise each feature-specific matrix independently so Xavier's
+        # fan-in/fan-out calculation sees a two-dimensional matrix.
+        for feature_idx in range(self.d_group):
+            nn.init.xavier_uniform_(self.gate_proj[feature_idx])
+            nn.init.xavier_uniform_(self.value_proj[feature_idx])
+        nn.init.normal_(self.output_proj, mean=0.0, std=1e-3)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """``(B, L, n_embd) -> (B, L, n_embd)``."""
-        return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))
+        """Map ``(B, L, n_embd)`` to an equally shaped 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:
+            raise ValueError(
+                f"Expected hidden size {self.n_embd}, got {x.size(-1)}"
+            )
+
+        batch_size, seq_len, _ = x.shape
+        grouped = x.reshape(
+            batch_size, seq_len, self.n_group, self.d_group
+        )
+        aligned = torch.einsum(
+            "blgd,gde->blge", grouped, self.group_align
+        )
+
+        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
+        )
+        return self.drop(mixed.reshape(batch_size, seq_len, self.n_embd))
 
 
 class GPTBlock(nn.Module):
@@ -218,6 +286,7 @@ class GPTBlock(nn.Module):
 
         attn_dropout: float = 0.0,
         mlp_dropout: float = 0.0,
+        hidden_group: int = 20,
         use_time_rope: bool = False,
         use_rbf_bias: bool = False,
         n_rbf_bases: int = 16,
@@ -231,7 +300,12 @@ class GPTBlock(nn.Module):
             use_time_rope=use_time_rope,
             use_rbf_bias=use_rbf_bias,
         )
-        self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout)
+        self.mlp = TrajMixer(
+            n_embd=n_embd,
+            n_head=n_head,
+            hidden_group=hidden_group,
+            dropout=mlp_dropout,
+        )
         self.ln1 = nn.LayerNorm(n_embd)
         self.ln2 = nn.LayerNorm(n_embd)