Refactor TrajMixer to single residual

2026-07-24 15:25:02 +08:00
parent 7b48cb8425
commit 8d0d71292e
4 changed files with 304 additions and 303 deletions
--- a/TrajMixer_设计方案.md
+++ b/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；
- 原始 TimeRoPE / Relative Time Attention Bias；
- 原始 Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)；
- 原始序列建模和训练目标。
+- causal mask；
+- TimeRoPE；
+- 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
-→ Cross-Group PreNorm
-→ Group-wise Feature Alignment
+→ Static Gated Fusion
 → Cross-Group SwiGLU: n_group → 4n_group → n_group
-→ Cross-Group Residual
-→ reshape [B, L, d]
+→ reshape [B, L, n_embd]
+→ 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)}
-=\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),
+N=\operatorname{LN}_{d}(X),
 \]

 \[
-H^{(2)}
-=H^{(1)}
-+\operatorname{Dropout}
-\left(\operatorname{CrossMixer}
-\left(\operatorname{LN}_{G}(H^{(1)})\right)\right),
+G=\operatorname{reshape}(N)
+\in\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}},
 \]

 \[
-X^{(l+1)}=\operatorname{reshape}(H^{(2)})
-\in\mathbb{R}^{B\times L\times d}.
+P=\operatorname{IntraMixer}(G),
 \]

-`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)}
-=\operatorname{LN}_{D}(H^{(0)}).
+d_{\mathrm{group}}
+\rightarrow
+4d_{\mathrm{group}}
+\rightarrow
+d_{\mathrm{group}}.
 \]

-归一化统计量在每个 group 内独立计算；为保持轻量，LayerNorm 的 affine 参数在各 group 间共享。
-
-对于 \(g=1,\ldots,G\)，定义：
+对 group \(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
-=\operatorname{SiLU}
-\left(\widetilde H^{(0)}_g
-W_{g,\mathrm{intra}}^{(g)}\right)
+H_g
+=
+\operatorname{SiLU}
+\left(G_gW_{g,\mathrm{intra}}^{(g)}\right)
 \odot
-\left(\widetilde H^{(0)}_g
-W_{v,\mathrm{intra}}^{(g)}\right),
+\left(G_gW_{v,\mathrm{intra}}^{(g)}\right),
 \]

 \[
-\Delta_{\mathrm{intra},g}
-=P_gW_{o,\mathrm{intra}}^{(g)},
+P_g=H_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}
-=\operatorname{LN}_{G}
-\left(H^{(1)}_{b,t,:,r}\right).
+\Gamma=\sigma(\Theta).
 \]

-归一化统计量对每个内部坐标 \(r\) 独立计算；LayerNorm 的 affine 参数在各内部坐标间共享。
-
-### 5.1 Group-wise Feature Alignment
-
-沿用现有的可学习 group 特征对齐矩阵：
+初始化：

 \[
-B_g\in\mathbb{R}^{D\times D},
-\qquad g=1,\ldots,G,
+\Theta_{g,r}
+=\operatorname{logit}(0.1)
+=\log\frac{0.1}{0.9}
+\approx-2.1972,
 \]

+因此：
+
 \[
-Z_{b,t,g,:}
-=\widetilde H^{(1)}_{b,t,g,:}B_g.
+\Gamma_{g,r}\approx0.1.
 \]

-\(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},
-\qquad
+\in
+\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}
-=\operatorname{SiLU}
-\left(Z_{b,t,:,r}A_g^{(r)}\right)
+Q_{:,r}
+=
+\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}
-=Q_{b,t,:,r}A_o^{(r)},
+\Delta_{:,r}=Q_{:,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：
-
-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 组内阶段
-
-投影权重：
+跨 group 输出 reshape 回：

 \[
-3G D(4D)
-=12GD^2
-=17{,}280.
+\Delta\in\mathbb{R}^{B\times L\times d}.
 \]

-`LayerNorm(D)`：
+最终：

 \[
-2D=24.
+\operatorname{TrajMixer}(X)
+=X+\operatorname{Dropout}(\Delta).
 \]

-### 7.2 跨组阶段
+固定约束：

-跨组投影权重：
+- 组内阶段后不执行独立 residual；
+- 跨组阶段后不执行独立 residual；
+- `GPTBlock` 不再额外执行 `X + TrajMixer(X)`；
+- 整个 TrajMixer 只有一次主 residual。

-\[
-3D G(4G)
-=12DG^2
-=14{,}400.
-\]
+## 9. 初始化

-Group Feature Alignment：
+固定初始化：

-\[
-GD^2
-=1{,}440.
-\]
-
-`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：单位矩阵；
+- `intra_gate_proj/intra_value_proj`：每个 group 独立 Xavier uniform；
+- `intra_output_proj`：每个 group 独立 Xavier uniform；
+- `intra_gate_logits`：初始化为 \(\operatorname{logit}(0.1)\)；
 - 跨组 `gate_proj/value_proj`：每个内部坐标独立 Xavier uniform；
- 跨组 `output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布；
- 两个 LayerNorm：PyTorch 默认 affine 初始化；
- 两个 residual stage 的 Dropout 均沿用 `mlp_dropout`。
+- 最终跨组 `output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布；
+- Full-width LayerNorm：PyTorch 默认 affine 初始化；
+- 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.
+\]

-整个模块保持：
+静态门控：

- 无时间递归；
- 不混合序列位置；
- 序列维度完全并行；
- 保留原始因果 Attention；
- residual groups 不等同于 Attention heads。
+\[
+n_{\mathrm{group}}d_{\mathrm{group}}
+=120.
+\]

-## 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
-intra_norm: layer_norm_over_d_group
-cross_norm: layer_norm_over_n_group
-group_alignment: per_group_d_group_x_d_group
+group_alignment: false
+intra_residual: false
+cross_residual: false
+traj_mixer_outer_residual: true
 projection_bias: false
-gate_value_init: xavier_uniform
-output_init_std: 0.001
+intra_output_init: xavier_uniform
+cross_output_init_std: 0.001
 ```

-必须满足：
+训练时必须将 `model_architecture: traj_mixer_v5`、`model_parameter_count` 和 `trainable_parameter_count` 写入 `train_config.json`，并在日志中打印参数量。

-\[
-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 不向后兼容，直接拒绝加载。
+评估和导出入口只接受 `traj_mixer_v5` checkpoint，并检查 Full-width LayerNorm、组内 projections、静态门控和跨 group projections 是否齐全。`traj_mixer_v4` 及更早 checkpoint 不向后兼容，直接拒绝加载。
--- a/backbones.py
+++ b/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):
--- a/models.py
+++ b/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.intra_norm.bias",
+        "blocks.0.mlp.norm.weight",
+        "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.cross_norm.bias",
-        "blocks.0.mlp.group_align",
+        "blocks.0.mlp.intra_gate_logits",
        "blocks.0.mlp.gate_proj",
        "blocks.0.mlp.value_proj",
        "blocks.0.mlp.output_proj",
--- a/test_traj_mixer.py
+++ b/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)
-
-        expected = torch.eye(12).expand(10, 12, 12)
-        torch.testing.assert_close(mixer.group_align.detach(), expected)
+        self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040)
+        self.assertFalse(hasattr(mixer, "group_align"))
+        self.assertFalse(hasattr(mixer, "intra_norm"))
+        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(
-            2, 4, 10, 12
-        )
-        changed_out = mixer(changed.reshape(2, 4, 120)).reshape(
-            2, 4, 10, 12
-        )
+        original_out = mixer._intra_mix(grouped)
+        changed_out = mixer._intra_mix(changed)
        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)
-        changed_out = mixer(changed.reshape(1, 1, 6)).reshape(1, 1, 3, 2)
+        original_out = mixer._cross_mix(grouped)
+        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.cross_norm, torch.nn.LayerNorm)
+        self.assertIsInstance(block.mlp.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,
-                "trainable_parameter_count": 33_164,
+                "model_parameter_count": 32_040,
+                "trainable_parameter_count": 32_040,
            },
        )