Refactor TrajMixer to single residual

TrajMixer_设计方案.md

@@ -2,176 +2,178 @@
 
 > 状态：**Frozen implementation baseline**
 >
-> 版本：**v2.0 / traj_mixer_v3**
+> 版本：**v3.0 / traj_mixer_v5**
 >
 > 固化日期：**2026-07-24**
 
-本文档是 TrajMixer 后续实现与实验的结构基线。本版本将原先仅含跨轨迹交互的 TrajMixer 扩展为“组内混合 + 跨组混合”的两阶段结构。
+本文档是当前 TrajMixer 的实现与实验基线。本版本采用单 PreNorm、单外层 residual、静态门控组内融合和跨 group SwiGLU。
 
 ## 1. 目标
 
-在保持原始 Delphi Transformer Attention 结构不变的前提下，用轻量、可并行的两阶段 TrajMixer 替换 FFN。
+在不改变 Delphi Transformer Attention 的前提下，用轻量、完全并行的 TrajMixer 替换 FFN。
 
-保持不变的组件包括：
+保持不变：
 
-- 原始 causal mask；
+- causal mask；
-- 原始 TimeRoPE / Relative Time Attention Bias；
+- TimeRoPE；
-- 原始 Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)；
+- Relative Time Attention Bias；
-- 原始序列建模和训练目标。
+- Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)；
+- 序列建模和训练目标。
 
 TrajMixer 不沿序列维度混合，也不引入时间递归。
 
-## 2. Block 总体结构
+## 2. Block 结构
 
 ```text
 PreNorm Causal Multi-Head Attention
 → Attention Residual
+→ Full-width TrajMixer PreNorm
 → reshape [B, L, n_group, d_group]
-→ Intra-Group PreNorm
 → Per-Group SwiGLU: d_group → 4d_group → d_group
-→ Intra-Group Residual
+→ Static Gated Fusion
-→ Cross-Group PreNorm
-→ Group-wise Feature Alignment
 → Cross-Group SwiGLU: n_group → 4n_group → n_group
-→ Cross-Group Residual
+→ reshape [B, L, n_embd]
-→ reshape [B, L, d]
+→ Dropout
+→ One TrajMixer Residual
 ```
 
-Attention 阶段保持原样：
+Attention 阶段：
 
 \[
-U=X^{(l)}+\operatorname{CausalMHA}
+X
+=X^{(l)}
++\operatorname{CausalMHA}
 \left(\operatorname{LN}_{\mathrm{attn}}(X^{(l)})\right).
 \]
 
-随后：
+TrajMixer 阶段：
 
 \[
-H^{(0)}
+N=\operatorname{LN}_{d}(X),
-=\operatorname{reshape}(U)
-\in\mathbb{R}^{B\times L\times G\times D},
-\]
-
-其中 \(G=n_{\mathrm{group}}\)，\(D=d_{\mathrm{group}}\)。
-
-两阶段 TrajMixer 为：
-
-\[
-H^{(1)}
-=H^{(0)}
-+\operatorname{Dropout}
-\left(\operatorname{IntraMixer}
-\left(\operatorname{LN}_{D}(H^{(0)})\right)\right),
 \]
 
 \[
-H^{(2)}
+G=\operatorname{reshape}(N)
-=H^{(1)}
+\in\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}},
-+\operatorname{Dropout}
-\left(\operatorname{CrossMixer}
-\left(\operatorname{LN}_{G}(H^{(1)})\right)\right),
 \]
 
 \[
-X^{(l+1)}=\operatorname{reshape}(H^{(2)})
+P=\operatorname{IntraMixer}(G),
-\in\mathbb{R}^{B\times L\times d}.
 \]
 
-`TrajMixer.forward()` 返回的是已经完成两次 residual update 的完整状态，而不是单个 residual delta。因此 `GPTBlock` 在 Attention residual 后直接返回 `TrajMixer(U)`，不得再写成 `U + TrajMixer(U)`。
+\[
+U=G+\sigma(\Theta)\odot P,
+\]
 
-## 3. Latent Trajectory Group 定义
+\[
+\Delta=\operatorname{reshape}
+\left(\operatorname{CrossMixer}(U)\right),
+\]
+
+\[
+X^{(l+1)}=X+\operatorname{Dropout}(\Delta).
+\]
+
+整个 TrajMixer 只有最后一次 `X + update` 是 residual。`U=G+\sigma(\Theta)\odot P` 是 update 分支内部的静态门控特征融合，不是相对于主 residual stream 的独立 residual stage。
+
+## 3. Group 定义
 
 Attention 输出经过 \(W_O\) 后仍是标准 residual representation：
 
 \[
-U\in\mathbb{R}^{B\times L\times d}.
+X\in\mathbb{R}^{B\times L\times d}.
 \]
 
-固定：
+定义：
 
 \[
-G:=n_{\mathrm{group}}=n_{\mathrm{head}},
+n_{\mathrm{group}}:=n_{\mathrm{head}},
 \qquad
-D:=d_{\mathrm{group}}=\frac{d}{G},
+d_{\mathrm{group}}=\frac{d}{n_{\mathrm{group}}},
 \qquad
-d=GD.
+d=n_{\mathrm{group}}d_{\mathrm{group}}.
 \]
 
-默认配置：
+默认：
 
 \[
-d=120,\qquad G=10,\qquad D=12.
+d=120,\qquad
+n_{\mathrm{group}}=10,\qquad
+d_{\mathrm{group}}=12.
 \]
 
-reshape 后：
+这些 group 是 residual space 的连续分区，不等同于 Attention heads；二者只共享数量。
+
+## 4. 唯一的 Full-Width PreNorm
+
+TrajMixer 只使用一个：
+
+```text
+norm: LayerNorm(n_embd)
+```
+
+LayerNorm 作用于完整 \(d\) 维 residual representation，然后才 reshape：
 
 \[
-H^{(0)}\in\mathbb{R}^{B\times L\times G\times D}.
+G=\operatorname{reshape}
+\left(\operatorname{LN}_{d}(X)\right).
 \]
 
-`n_group` 由 `n_head` 决定，但 residual groups 只是 residual space 的连续分区，不等同于 Attention heads。
+本版本明确删除：
 
-## 4. 第一阶段：组内 SwiGLU Mixer
+```text
+intra_norm
+cross_norm
+group_align
+```
 
-第一阶段对每个 group 独立进行特征变换。不同 group 使用各自的投影参数，不发生 group 间信息交换。
+不得在组内或跨组阶段再增加额外 LayerNorm。
 
-先对每个 \((b,t,g)\) 的 \(D\) 维向量独立执行 LayerNorm：
+## 5. 组内 SwiGLU
+
+每个 group 使用独立参数，对其 \(d_{\mathrm{group}}\) 维内部特征执行：
 
 \[
-\widetilde H^{(0)}
+d_{\mathrm{group}}
-=\operatorname{LN}_{D}(H^{(0)}).
+\rightarrow
+4d_{\mathrm{group}}
+\rightarrow
+d_{\mathrm{group}}.
 \]
 
-归一化统计量在每个 group 内独立计算；为保持轻量，LayerNorm 的 affine 参数在各 group 间共享。
+对 group \(g\)：
-
-对于 \(g=1,\ldots,G\)，定义：
 
 \[
 W_{g,\mathrm{intra}}^{(g)},
 W_{v,\mathrm{intra}}^{(g)}
-\in\mathbb{R}^{D\times 4D},
+\in
+\mathbb{R}^{d_{\mathrm{group}}\times4d_{\mathrm{group}}},
 \]
 
 \[
 W_{o,\mathrm{intra}}^{(g)}
-\in\mathbb{R}^{4D\times D}.
+\in
+\mathbb{R}^{4d_{\mathrm{group}}\times d_{\mathrm{group}}}.
 \]
 
 计算：
 
 \[
-P_g
+H_g
-=\operatorname{SiLU}
+=
-\left(\widetilde H^{(0)}_g
+\operatorname{SiLU}
-W_{g,\mathrm{intra}}^{(g)}\right)
+\left(G_gW_{g,\mathrm{intra}}^{(g)}\right)
 \odot
-\left(\widetilde H^{(0)}_g
+\left(G_gW_{v,\mathrm{intra}}^{(g)}\right),
-W_{v,\mathrm{intra}}^{(g)}\right),
 \]
 
 \[
-\Delta_{\mathrm{intra},g}
+P_g=H_gW_{o,\mathrm{intra}}^{(g)}.
-=P_gW_{o,\mathrm{intra}}^{(g)},
 \]
 
-\[
+实现形状：
-H^{(1)}
-=H^{(0)}
-+\operatorname{Dropout}(\Delta_{\mathrm{intra}}).
-\]
-
-该阶段完成：
-
-\[
-D\rightarrow4D\rightarrow D,
-\]
-
-用于增强每条潜在轨迹内部的非线性特征组合能力。
-
-实现张量形状：
 
 ```text
-intra_norm:        LayerNorm(d_group)
 intra_gate_proj:   [n_group, d_group, 4 * d_group]
 intra_value_proj:  [n_group, d_group, 4 * d_group]
 intra_output_proj: [n_group, 4 * d_group, d_group]
@@ -179,222 +181,205 @@ intra_output_proj: [n_group, 4 * d_group, d_group]
 
 三个 projection 均不带 bias。
 
-## 5. 第二阶段：跨组 TrajMixer
+## 6. 静态门控融合
 
-第二阶段沿 group 维度进行交互。对于每个内部坐标 \(r\)，独立执行：
+定义可学习 gate logits：
 
 \[
-G\rightarrow4G\rightarrow G.
+\Theta\in
+\mathbb{R}^{n_{\mathrm{group}}\times d_{\mathrm{group}}}.
 \]
 
-首先将 \(H^{(1)}\) 的最后两个维度交换，并在 group 维度执行 LayerNorm：
+实际门值为：
 
 \[
-\widetilde H^{(1)}_{b,t,:,r}
+\Gamma=\sigma(\Theta).
-=\operatorname{LN}_{G}
-\left(H^{(1)}_{b,t,:,r}\right).
 \]
 
-归一化统计量对每个内部坐标 \(r\) 独立计算；LayerNorm 的 affine 参数在各内部坐标间共享。
+初始化：
-
-### 5.1 Group-wise Feature Alignment
-
-沿用现有的可学习 group 特征对齐矩阵：
 
 \[
-B_g\in\mathbb{R}^{D\times D},
+\Theta_{g,r}
-\qquad g=1,\ldots,G,
+=\operatorname{logit}(0.1)
+=\log\frac{0.1}{0.9}
+\approx-2.1972,
 \]
 
+因此：
+
 \[
-Z_{b,t,g,:}
+\Gamma_{g,r}\approx0.1.
-=\widetilde H^{(1)}_{b,t,g,:}B_g.
 \]
 
-\(B_g\) 不带 bias，并使用单位矩阵初始化。
+融合：
 
-### 5.2 Cross-Group SwiGLU
+\[
+U=G+\Gamma\odot P.
+\]
 
-对每个内部坐标 \(r=1,\ldots,D\)，定义：
+\(\Gamma\) 对 batch 和序列位置共享，但每个 group、每个内部坐标拥有独立可学习值。
+
+## 7. 跨 Group SwiGLU
+
+对于每个内部坐标 \(r\)，独立沿 group 维度执行：
+
+\[
+n_{\mathrm{group}}
+\rightarrow
+4n_{\mathrm{group}}
+\rightarrow
+n_{\mathrm{group}}.
+\]
+
+定义：
 
 \[
 A_g^{(r)},A_v^{(r)}
-\in\mathbb{R}^{G\times4G},
+\in
-\qquad
+\mathbb{R}^{n_{\mathrm{group}}\times4n_{\mathrm{group}}},
+\]
+
+\[
 A_o^{(r)}
-\in\mathbb{R}^{4G\times G}.
+\in
+\mathbb{R}^{4n_{\mathrm{group}}\times n_{\mathrm{group}}}.
 \]
 
 计算：
 
 \[
-Q_{b,t,:,r}
+Q_{:,r}
-=\operatorname{SiLU}
+=
-\left(Z_{b,t,:,r}A_g^{(r)}\right)
+\operatorname{SiLU}\left(U_{:,r}A_g^{(r)}\right)
 \odot
-\left(Z_{b,t,:,r}A_v^{(r)}\right),
+\left(U_{:,r}A_v^{(r)}\right),
 \]
 
 \[
-\Delta_{\mathrm{cross},b,t,:,r}
+\Delta_{:,r}=Q_{:,r}A_o^{(r)}.
-=Q_{b,t,:,r}A_o^{(r)},
 \]
 
-\[
+实现形状：
-H^{(2)}
-=H^{(1)}
-+\operatorname{Dropout}(\Delta_{\mathrm{cross}}).
-\]
-
-实现张量形状：
 
 ```text
-cross_norm:  LayerNorm(n_group)
-group_align: [n_group, d_group, d_group]
 gate_proj:   [d_group, n_group, 4 * n_group]
 value_proj:  [d_group, n_group, 4 * n_group]
 output_proj: [d_group, 4 * n_group, n_group]
 ```
 
-三个 projection 均不带 bias。
+三个 projection 均不带 bias。不同内部坐标拥有独立的跨 group 参数，且不沿序列维度交互。
 
-## 6. PreNorm 与残差约束
+## 8. 唯一的外层 Residual
 
-本版本固定使用两个独立的 PreNorm residual stage：
+跨 group 输出 reshape 回：
-
-1. `intra_norm` 只服务于组内 Mixer；
-2. `cross_norm` 只服务于跨组 Mixer；
-3. 第一阶段 residual 的输出是第二阶段的输入；
-4. 两个 residual 都在 `TrajMixer` 内部完成；
-5. 不再保留 block 外部的全维度 `ln2` 或额外 Mixer residual。
-
-因此信息流必须是：
-
-```text
-U
-→ U + IntraMixer(IntraNorm(U))
-→ H1 + CrossMixer(CrossNorm(H1))
-→ output
-```
-
-## 7. 参数量
-
-默认 \(d=120,G=10,D=12\)。
-
-### 7.1 组内阶段
-
-投影权重：
 
 \[
-3G D(4D)
+\Delta\in\mathbb{R}^{B\times L\times d}.
-=12GD^2
-=17{,}280.
 \]
 
-`LayerNorm(D)`：
+最终：
 
 \[
-2D=24.
+\operatorname{TrajMixer}(X)
+=X+\operatorname{Dropout}(\Delta).
 \]
 
-### 7.2 跨组阶段
+固定约束：
 
-跨组投影权重：
+- 组内阶段后不执行独立 residual；
+- 跨组阶段后不执行独立 residual；
+- `GPTBlock` 不再额外执行 `X + TrajMixer(X)`；
+- 整个 TrajMixer 只有一次主 residual。
 
-\[
+## 9. 初始化
-3D G(4G)
-=12DG^2
-=14{,}400.
-\]
 
-Group Feature Alignment：
+固定初始化：
 
-\[
+- `intra_gate_proj/intra_value_proj`：每个 group 独立 Xavier uniform；
-GD^2
+- `intra_output_proj`：每个 group 独立 Xavier uniform；
-=1{,}440.
+- `intra_gate_logits`：初始化为 \(\operatorname{logit}(0.1)\)；
-\]
-
-`LayerNorm(G)`：
-
-\[
-2G=20.
-\]
-
-### 7.3 每个 TrajMixer 合计
-
-\[
-17{,}280+24+14{,}400+1{,}440+20
-=\boxed{33{,}164}.
-\]
-
-相对于 `traj_mixer_v2` 的跨组单阶段结构 \(15{,}840\)，每层增加 \(17{,}324\) 个参数。作为历史实现对照，代码库原全维度 SwiGLU FFN 每层为 \(108{,}720\) 个参数。
-
-## 8. 初始化
-
-固定初始化约定：
-
-- 组内 `intra_gate_proj/intra_value_proj`：每个 group 独立 Xavier uniform；
-- 组内 `intra_output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布；
-- Group Alignment：单位矩阵；
 - 跨组 `gate_proj/value_proj`：每个内部坐标独立 Xavier uniform；
-- 跨组 `output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布；
+- 最终跨组 `output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布；
-- 两个 LayerNorm：PyTorch 默认 affine 初始化；
+- Full-width LayerNorm：PyTorch 默认 affine 初始化；
-- 两个 residual stage 的 Dropout 均沿用 `mlp_dropout`。
+- Dropout：沿用 `mlp_dropout`。
 
-两个 output projection 的小方差初始化使两阶段在训练初期都接近恒等 residual update。
+组内输出使用正常 Xavier 初始化以保证其具有完整表达能力；静态门控将其初始贡献限制在约 0.1。最终跨 group 输出投影保持小值初始化，使整个 TrajMixer residual update 在训练初期接近零。
 
-Relative Time Attention Bias 的初始化固定为：
+Relative Time Attention Bias 初始化固定为：
 
 - `rbf_proj.weight`：零初始化；
 - `time_bias_scale`：初始化为 \(1.0\)；
-- 初始 RBF attention bias 仍严格为零；
+- 初始 RBF attention bias 严格为零；
 - `rbf_proj.weight` 从第一个优化步骤即可获得梯度。
 
-不得同时将 `rbf_proj.weight` 和 `time_bias_scale` 初始化为零，否则两个相乘分支的梯度都会为零，RBF 时间偏置将无法开始学习。
+## 10. 参数量
 
-## 9. 信息流与语义
+默认 \(d=120\)、\(n_{\mathrm{group}}=10\)、\(d_{\mathrm{group}}=12\)。
 
-**Attention**：从历史疾病事件中选择和整合相关信息。
+Full-width LayerNorm：
 
-**Intra-Group Mixer**：学习每条潜在轨迹内部的非线性特征组合。
+\[
+2d=240.
+\]
 
-**Group Feature Alignment**：对齐不同潜在轨迹的内部坐标。
+组内 projections：
 
-**Cross-Group Mixer**：学习不同潜在轨迹在相同内部坐标上的门控交互。
+\[
+3n_{\mathrm{group}}d_{\mathrm{group}}
+\left(4d_{\mathrm{group}}\right)
+=17{,}280.
+\]
 
-整个模块保持：
+静态门控：
 
-- 无时间递归；
+\[
-- 不混合序列位置；
+n_{\mathrm{group}}d_{\mathrm{group}}
-- 序列维度完全并行；
+=120.
-- 保留原始因果 Attention；
+\]
-- residual groups 不等同于 Attention heads。
 
-## 10. 固定配置与 checkpoint 约束
+跨 group projections：
+
+\[
+3d_{\mathrm{group}}n_{\mathrm{group}}
+\left(4n_{\mathrm{group}}\right)
+=14{,}400.
+\]
+
+每层 TrajMixer 合计：
+
+\[
+240+17{,}280+120+14{,}400
+=\boxed{32{,}040}.
+\]
+
+默认 relative-time、12 层、`vocab_size=1256`、无额外信息类型时，完整模型参数量为：
+
+\[
+\boxed{1{,}232{,}428}.
+\]
+
+## 11. 固定配置与 checkpoint 约束
 
 ```yaml
-model_architecture: traj_mixer_v3
+model_architecture: traj_mixer_v5
 d_model: 120
 n_head: 10
 n_group_rule: n_head
 d_group_rule: d_model / n_group
+traj_mixer_norm: layer_norm_over_n_embd
 intra_hidden_rule: 4 * d_group
+intra_gate_shape: [n_group, d_group]
+intra_gate_initial_sigmoid: 0.1
 cross_hidden_rule: 4 * n_group
-attention: unchanged
+group_alignment: false
-intra_norm: layer_norm_over_d_group
+intra_residual: false
-cross_norm: layer_norm_over_n_group
+cross_residual: false
-group_alignment: per_group_d_group_x_d_group
+traj_mixer_outer_residual: true
 projection_bias: false
-gate_value_init: xavier_uniform
+intra_output_init: xavier_uniform
-output_init_std: 0.001
+cross_output_init_std: 0.001
 ```
 
-必须满足：
+训练时必须将 `model_architecture: traj_mixer_v5`、`model_parameter_count` 和 `trainable_parameter_count` 写入 `train_config.json`，并在日志中打印参数量。
 
-\[
+评估和导出入口只接受 `traj_mixer_v5` checkpoint，并检查 Full-width LayerNorm、组内 projections、静态门控和跨 group projections 是否齐全。`traj_mixer_v4` 及更早 checkpoint 不向后兼容，直接拒绝加载。
-d=n_{\mathrm{group}}d_{\mathrm{group}}.
-\]
-
-训练时必须将 `model_architecture: traj_mixer_v3`、`model_parameter_count` 和 `trainable_parameter_count` 写入 `train_config.json`，并在训练日志中显式打印参数量。
-
-本分支的评估和导出入口只接受 `traj_mixer_v3` checkpoint，并检查两阶段 Norm、组内 projection、Group Alignment 和跨组 projection 参数是否齐全。`traj_mixer_v2` 及更早 checkpoint 不向后兼容，直接拒绝加载。

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)
         )
-
+        self.intra_gate_logits = nn.Parameter(
-        # Stage 2: each internal coordinate independently mixes groups.
+            torch.empty(self.n_group, self.d_group)
-        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)
         )
 
         # 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)
+            nn.init.xavier_uniform_(self.intra_output_proj[group_idx])
-
+        nn.init.constant_(
-        with torch.no_grad():
+            self.intra_gate_logits,
-            identity = torch.eye(
+            math.log(0.1 / 0.9),
-                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.
@@ -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.
+        # The static per-channel gate starts at sigmoid(logit) ~= 0.1.
-        intra_input = self.intra_norm(grouped)
+        intra_output = self._intra_mix(grouped)
-        intra_gate = torch.einsum(
+        intra_gate = torch.sigmoid(self.intra_gate_logits).view(
-            "blgd,gdh->blgh", intra_input, self.intra_gate_proj
+            1, 1, self.n_group, self.d_group
         )
-        intra_value = torch.einsum(
+        mixed_input = grouped + intra_gate * intra_output
-            "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(
+        update = self._cross_mix(mixed_input).reshape(
-            grouped.transpose(-1, -2)
+            batch_size, seq_len, self.n_embd
-        ).transpose(-1, -2)
-        aligned = torch.einsum(
-            "blgd,gde->blge", cross_input, self.group_align
         )
-
+        return x + self.drop(update)
-        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)
 
 
 class GPTBlock(nn.Module):

models.py

@@ -15,7 +15,7 @@ from backbones import (
 from targets import PAD_IDX
 
 
-TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v3"
+TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v5"
 
 
 def validate_traj_mixer_config(config: Mapping[str, object]) -> None:
@@ -29,14 +29,12 @@ def validate_traj_mixer_config(config: Mapping[str, object]) -> None:
 
 def validate_traj_mixer_state_dict(state_dict: Mapping[str, object]) -> None:
     required_keys = {
-        "blocks.0.mlp.intra_norm.weight",
+        "blocks.0.mlp.norm.weight",
-        "blocks.0.mlp.intra_norm.bias",
+        "blocks.0.mlp.norm.bias",
         "blocks.0.mlp.intra_gate_proj",
         "blocks.0.mlp.intra_value_proj",
         "blocks.0.mlp.intra_output_proj",
-        "blocks.0.mlp.cross_norm.weight",
+        "blocks.0.mlp.intra_gate_logits",
-        "blocks.0.mlp.cross_norm.bias",
-        "blocks.0.mlp.group_align",
         "blocks.0.mlp.gate_proj",
         "blocks.0.mlp.value_proj",
         "blocks.0.mlp.output_proj",

test_traj_mixer.py

@@ -57,10 +57,16 @@ class TrajMixerTest(unittest.TestCase):
 
         x = torch.randn(2, 7, 120)
         self.assertEqual(mixer(x).shape, x.shape)
-        self.assertEqual(sum(p.numel() for p in mixer.parameters()), 33_164)
+        self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040)
-
+        self.assertFalse(hasattr(mixer, "group_align"))
-        expected = torch.eye(12).expand(10, 12, 12)
+        self.assertFalse(hasattr(mixer, "intra_norm"))
-        torch.testing.assert_close(mixer.group_align.detach(), expected)
+        self.assertFalse(hasattr(mixer, "cross_norm"))
+        self.assertEqual(tuple(mixer.norm.normalized_shape), (120,))
+        self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12))
+        torch.testing.assert_close(
+            torch.sigmoid(mixer.intra_gate_logits.detach()),
+            torch.full((10, 12), 0.1),
+        )
         self.assertEqual(mixer.intra_hidden, 48)
         self.assertEqual(
             tuple(mixer.intra_gate_proj.shape),
@@ -78,35 +84,42 @@ class TrajMixerTest(unittest.TestCase):
         self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40))
         self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 40))
         self.assertEqual(tuple(mixer.output_proj.shape), (12, 40, 10))
-        self.assertEqual(tuple(mixer.intra_norm.normalized_shape), (12,))
-        self.assertEqual(tuple(mixer.cross_norm.normalized_shape), (10,))
 
-    def test_zero_output_projections_make_both_stages_identity(self) -> None:
+    def test_zero_final_output_projection_makes_mixer_identity(self) -> None:
         torch.manual_seed(0)
         mixer = TrajMixer(120, n_head=10, dropout=0.0)
         with torch.no_grad():
-            mixer.intra_output_proj.zero_()
             mixer.output_proj.zero_()
         x = torch.randn(2, 5, 120)
         torch.testing.assert_close(mixer(x), x)
 
+    def test_forward_matches_single_outer_residual_formula(self) -> None:
+        torch.manual_seed(0)
+        mixer = TrajMixer(120, n_head=10, dropout=0.0)
+        mixer.eval()
+        x = torch.randn(2, 5, 120)
+
+        grouped = mixer.norm(x).reshape(2, 5, 10, 12)
+        intra_output = mixer._intra_mix(grouped)
+        static_gate = torch.sigmoid(mixer.intra_gate_logits).view(
+            1, 1, 10, 12
+        )
+        mixed_input = grouped + static_gate * intra_output
+        update = mixer._cross_mix(mixed_input).reshape(2, 5, 120)
+
+        torch.testing.assert_close(mixer(x), x + update)
+
     def test_intra_stage_is_independent_across_groups(self) -> None:
         torch.manual_seed(0)
         mixer = TrajMixer(120, n_head=10, dropout=0.0)
         mixer.eval()
-        with torch.no_grad():
-            mixer.output_proj.zero_()
 
         grouped = torch.randn(2, 4, 10, 12)
         changed = grouped.clone()
         changed[:, :, 3, :] += torch.randn_like(changed[:, :, 3, :])
 
-        original_out = mixer(grouped.reshape(2, 4, 120)).reshape(
+        original_out = mixer._intra_mix(grouped)
-            2, 4, 10, 12
+        changed_out = mixer._intra_mix(changed)
-        )
-        changed_out = mixer(changed.reshape(2, 4, 120)).reshape(
-            2, 4, 10, 12
-        )
         unchanged_groups = torch.tensor([0, 1, 2, 4, 5, 6, 7, 8, 9])
         torch.testing.assert_close(
             original_out.index_select(2, unchanged_groups),
@@ -117,7 +130,6 @@ class TrajMixerTest(unittest.TestCase):
         mixer = TrajMixer(6, n_head=3, dropout=0.0)
         mixer.eval()
         with torch.no_grad():
-            mixer.intra_output_proj.zero_()
             mixer.gate_proj.zero_()
             mixer.value_proj.zero_()
             mixer.output_proj.zero_()
@@ -138,8 +150,8 @@ class TrajMixerTest(unittest.TestCase):
         changed = grouped.clone()
         changed[0, 0, 0, 0] = 2.0
 
-        original_out = mixer(grouped.reshape(1, 1, 6)).reshape(1, 1, 3, 2)
+        original_out = mixer._cross_mix(grouped)
-        changed_out = mixer(changed.reshape(1, 1, 6)).reshape(1, 1, 3, 2)
+        changed_out = mixer._cross_mix(changed)
 
         self.assertNotEqual(
             original_out[0, 0, 1, 0].item(),
@@ -178,12 +190,13 @@ class TrajMixerTest(unittest.TestCase):
             self.assertIsNotNone(parameter.grad, name)
             self.assertTrue(torch.isfinite(parameter.grad).all(), name)
 
-    def test_gpt_block_delegates_both_mixer_residuals_to_traj_mixer(self) -> None:
+    def test_gpt_block_delegates_single_mixer_residual_to_traj_mixer(self) -> None:
         block = GPTBlock(n_embd=120, n_head=10)
         self.assertIsInstance(block.mlp, TrajMixer)
         self.assertFalse(hasattr(block, "ln2"))
-        self.assertIsInstance(block.mlp.intra_norm, torch.nn.LayerNorm)
+        self.assertIsInstance(block.mlp.norm, torch.nn.LayerNorm)
-        self.assertIsInstance(block.mlp.cross_norm, torch.nn.LayerNorm)
+        self.assertFalse(hasattr(block.mlp, "intra_norm"))
+        self.assertFalse(hasattr(block.mlp, "cross_norm"))
 
         x = torch.randn(2, 6, 120)
         self.assertEqual(block(x).shape, x.shape)
@@ -200,6 +213,14 @@ class TrajMixerTest(unittest.TestCase):
             validate_traj_mixer_config(
                 {"model_architecture": "traj_mixer_v2"}
             )
+        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
+            validate_traj_mixer_config(
+                {"model_architecture": "traj_mixer_v3"}
+            )
+        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
+            validate_traj_mixer_config(
+                {"model_architecture": "traj_mixer_v4"}
+            )
 
     def test_checkpoint_must_contain_traj_mixer_parameters(self) -> None:
         block = GPTBlock(n_embd=120, n_head=10)
@@ -222,8 +243,8 @@ class TrajMixerTest(unittest.TestCase):
         self.assertEqual(
             get_model_parameter_counts(mixer),
             {
-                "model_parameter_count": 33_164,
+                "model_parameter_count": 32_040,
-                "trainable_parameter_count": 33_164,
+                "trainable_parameter_count": 32_040,
             },
         )