Scale TrajMixer hidden width with head count

backbones.py

@@ -189,7 +189,6 @@ class TrajMixer(nn.Module):
         self,
         n_embd: int,
         n_head: int = 10,
-        hidden_group: int = 20,
         dropout: float = 0.0,
     ):
         super().__init__()
@@ -201,17 +200,12 @@ class TrajMixer(nn.Module):
             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.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
+        self.hidden_group = 4 * n_head
 
         # Per-group feature alignment: [group, input feature, output feature].
         self.group_align = nn.Parameter(
@@ -221,13 +215,13 @@ class TrajMixer(nn.Module):
         # 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)
+            torch.empty(self.d_group, self.n_group, self.hidden_group)
         )
         self.value_proj = nn.Parameter(
-            torch.empty(self.d_group, self.n_group, hidden_group)
+            torch.empty(self.d_group, self.n_group, self.hidden_group)
         )
         self.output_proj = nn.Parameter(
-            torch.empty(self.d_group, hidden_group, self.n_group)
+            torch.empty(self.d_group, self.hidden_group, self.n_group)
         )
         self.drop = nn.Dropout(dropout)
         self.reset_parameters()
@@ -286,7 +280,6 @@ 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,
@@ -303,7 +296,6 @@ class GPTBlock(nn.Module):
         self.mlp = TrajMixer(
             n_embd=n_embd,
             n_head=n_head,
-            hidden_group=hidden_group,
             dropout=mlp_dropout,
         )
         self.ln1 = nn.LayerNorm(n_embd)