Add two-stage TrajMixer mixing

TrajMixer_设计方案.md

@@ -2,332 +2,390 @@
 
 > 状态：**Frozen implementation baseline**
 >
-> 版本：**v1.0**
+> 版本：**v2.0 / traj_mixer_v3**
 >
-> 固化日期：**2026-07-22**
+> 固化日期：**2026-07-24**
 
-本文档是 TrajMixer 后续实现与实验的唯一结构基线。除显式标记为消融项的配置外，所有实现均应遵循本文档；若结构发生变化，应先更新版本和实验记录。
+本文档是 TrajMixer 后续实现与实验的结构基线。本版本将原先仅含跨轨迹交互的 TrajMixer 扩展为“组内混合 + 跨组混合”的两阶段结构。
 
 ## 1. 目标
 
-在保持原始 Delphi Transformer Attention 结构不变的前提下，用轻量、可并行的轨迹交互模块替换 FFN。
+在保持原始 Delphi Transformer Attention 结构不变的前提下，用轻量、可并行的两阶段 TrajMixer 替换 FFN。
 
 保持不变的组件包括：
 
 - 原始 causal mask；
 - 原始 TimeRoPE / Relative Time Attention Bias；
 - 原始 Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)；
-- 原始序列建模与训练目标。
+- 原始序列建模和训练目标。
 
-TrajMixer 不修改 Attention，只替换每个 Transformer block 中的 FFN residual branch。
+TrajMixer 不沿序列维度混合，也不引入时间递归。
 
 ## 2. Block 总体结构
 
-概念结构：
-
 ```text
 PreNorm Causal Multi-Head Attention
-→ Residual
+→ Attention Residual
-→ Standard Mixer 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
-→ SwiGLU Cross-Group Mixer
+→ Cross-Group SwiGLU: n_group → 4n_group → n_group
-→ Residual
+→ Cross-Group Residual
+→ reshape [B, L, d]
 ```
 
-完整计算为：
+Attention 阶段保持原样：
 
 \[
-U = X^{(l)} + \operatorname{Dropout}\!\left(
+U=X^{(l)}+\operatorname{CausalMHA}
-\operatorname{CausalMHA}\left(
+\left(\operatorname{LN}_{\mathrm{attn}}(X^{(l)})\right).
-\operatorname{LN}_{\mathrm{attn}}(X^{(l)}),
+\]
-\text{time information}
+
-\right)\right),
+随后：
+
+\[
+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}_{\mathrm{mixer}}(U),
+H^{(2)}
+=H^{(1)}
++\operatorname{Dropout}
+\left(\operatorname{CrossMixer}
+\left(\operatorname{LN}_{G}(H^{(1)})\right)\right),
 \]
 
 \[
-\Delta = \operatorname{TrajMixer}(N),
+X^{(l+1)}=\operatorname{reshape}(H^{(2)})
+\in\mathbb{R}^{B\times L\times d}.
 \]
 
-\[
+`TrajMixer.forward()` 返回的是已经完成两次 residual update 的完整状态，而不是单个 residual delta。因此 `GPTBlock` 在 Attention residual 后直接返回 `TrajMixer(U)`，不得再写成 `U + TrajMixer(U)`。
-X^{(l+1)} = U + \operatorname{Dropout}(\Delta).
-\]
-
-首版中的 \(\operatorname{LN}_{\mathrm{mixer}}\) 是作用于完整 \(d=120\) 维 residual representation 的标准 LayerNorm。
 
 ## 3. Latent Trajectory Group 定义
 
 Attention 输出经过 \(W_O\) 后仍是标准 residual representation：
 
 \[
-N\in\mathbb{R}^{B\times L\times d},\qquad d=120.
+U\in\mathbb{R}^{B\times L\times d}.
 \]
 
-将 hidden dimension 划分为与 Attention head 数量相同的 group 数量：
+固定：
 
 \[
-n_{\mathrm{group}}:=n_{\mathrm{head}}=10,
+G:=n_{\mathrm{group}}=n_{\mathrm{head}},
-\qquad d_{\mathrm{group}}=\frac{d}{n_{\mathrm{head}}}=12,
+\qquad
+D:=d_{\mathrm{group}}=\frac{d}{G},
+\qquad
+d=GD.
 \]
 
-`n_group` 不再是独立超参数，代码统一使用 `n_head` 确定 residual group 数量。二者只共享数量；这些 residual groups 在语义和张量来源上仍不等同于原始 Attention heads。
+默认配置：
-
-并 reshape 为：
 
 \[
-N_{\mathrm{group}}in
+d=120,\qquad G=10,\qquad D=12.
-\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}.
 \]
 
-这些 group 是 residual space 中的 **latent trajectory groups**，不等同于原始 Attention heads。本文中的 group、trajectory group 均指这一 residual-channel partition。
+reshape 后：
-
-## 4. Group-wise Feature Alignment
-
-为缓解不同 group 内部坐标不对齐的问题，每个 group 使用独立的小矩阵：
 
 \[
-B_i\in\mathbb{R}^{d_{\mathrm{group}}\times d_{\mathrm{group}}},
+H^{(0)}\in\mathbb{R}^{B\times L\times G\times D}.
-\qquad i=1,\ldots,n_{\mathrm{group}}.
 \]
 
-对每个 group 内的特征进行可学习对齐：
+`n_group` 由 `n_head` 决定，但 residual groups 只是 residual space 的连续分区，不等同于 Attention heads。
+
+## 4. 第一阶段：组内 SwiGLU Mixer
+
+第一阶段对每个 group 独立进行特征变换。不同 group 使用各自的投影参数，不发生 group 间信息交换。
+
+先对每个 \((b,t,g)\) 的 \(D\) 维向量独立执行 LayerNorm：
 
 \[
-Z_{b,t,i,:}=N_{\mathrm{group},b,t,i,:}B_i.
+\widetilde H^{(0)}
+=\operatorname{LN}_{D}(H^{(0)}).
 \]
 
-因此：
+归一化统计量在每个 group 内独立计算；为保持轻量，LayerNorm 的 affine 参数在各 group 间共享。
+
+对于 \(g=1,\ldots,G\)，定义：
 
 \[
-Z\in
+W_{g,\mathrm{intra}}^{(g)},
-\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}.
+W_{v,\mathrm{intra}}^{(g)}
+\in\mathbb{R}^{D\times 4D},
 \]
 
-首版实现约定：
-
-- \(B_i\) 不带 bias；
-- \(B_i\) 使用单位矩阵初始化；
-- Alignment 只作用于 Mixer residual branch，不改变 Attention residual stream；
-- 首版不增加逆变换或额外的 group 内输出投影。
-
-Alignment 每层权重参数量为：
-
 \[
-n_{\mathrm{group}}d_{\mathrm{group}}^2
+W_{o,\mathrm{intra}}^{(g)}
-=10\times12^2
+\in\mathbb{R}^{4D\times D}.
-=1{,}440.
 \]
 
-## 5. SwiGLU Cross-Group Mixer
+计算：
-
-Mixer 只沿 group 维度交互，不沿序列维度交互，因此不会引入时间递归或未来信息泄漏。
-
-对于每个 group 内特征维度：
 
 \[
-r=1,\ldots,d_{\mathrm{group}},
+P_g
+=\operatorname{SiLU}
+\left(\widetilde H^{(0)}_g
+W_{g,\mathrm{intra}}^{(g)}\right)
+\odot
+\left(\widetilde H^{(0)}_g
+W_{v,\mathrm{intra}}^{(g)}\right),
 \]
 
-定义：
+\[
+\Delta_{\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]
+```
+
+三个 projection 均不带 bias。
+
+## 5. 第二阶段：跨组 TrajMixer
+
+第二阶段沿 group 维度进行交互。对于每个内部坐标 \(r\)，独立执行：
+
+\[
+G\rightarrow4G\rightarrow G.
+\]
+
+首先将 \(H^{(1)}\) 的最后两个维度交换，并在 group 维度执行 LayerNorm：
+
+\[
+\widetilde H^{(1)}_{b,t,:,r}
+=\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},
+\qquad g=1,\ldots,G,
+\]
+
+\[
+Z_{b,t,g,:}
+=\widetilde H^{(1)}_{b,t,g,:}B_g.
+\]
+
+\(B_g\) 不带 bias，并使用单位矩阵初始化。
+
+### 5.2 Cross-Group SwiGLU
+
+对每个内部坐标 \(r=1,\ldots,D\)，定义：
 
 \[
 A_g^{(r)},A_v^{(r)}
-\in\mathbb{R}^{n_{\mathrm{group}}\times h_{\mathrm{group}}},
+\in\mathbb{R}^{G\times4G},
-\]
+\qquad
-
-\[
 A_o^{(r)}
-\in\mathbb{R}^{h_{\mathrm{group}}\times n_{\mathrm{group}}}.
+\in\mathbb{R}^{4G\times G}.
 \]
 
-隐藏宽度不再独立配置，固定为：
+计算：
 
 \[
-h_{\mathrm{group}}=4n_{\mathrm{head}}
+Q_{b,t,:,r}
-=4n_{\mathrm{group}}.
+=\operatorname{SiLU}
-\]
+\left(Z_{b,t,:,r}A_g^{(r)}\right)
-
+\odot
-当前 \(n_{\mathrm{head}}=10\)，因此 \(h_{\mathrm{group}}=40\)。
+\left(Z_{b,t,:,r}A_v^{(r)}\right),
-
-对固定的 batch、时间位置和内部特征维度 \(r\)，将：
-
-\[
-Z_{b,t,:,r}\in\mathbb{R}^{n_{\mathrm{group}}}
-\]
-
-视为 row vector，计算：
-
-\[
-G_{b,t,:,r}=Z_{b,t,:,r}A_g^{(r)},
 \]
 
 \[
-V_{b,t,:,r}=Z_{b,t,:,r}A_v^{(r)},
+\Delta_{\mathrm{cross},b,t,:,r}
+=Q_{b,t,:,r}A_o^{(r)},
 \]
 
 \[
-M_{b,t,:,r}=\operatorname{SiLU}(G_{b,t,:,r})\odot V_{b,t,:,r},
+H^{(2)}
+=H^{(1)}
++\operatorname{Dropout}(\Delta_{\mathrm{cross}}).
 \]
 
-\[
+实现张量形状：
-Y_{b,t,:,r}=M_{b,t,:,r}A_o^{(r)}.
-\]
-
-其中：
-
-- gate 分支控制信息写入；
-- value 分支提供交互内容；
-- output matrix 将隐藏 group 表示投影回原始 group 数量；
-- hidden group 表示固定扩展为 group 数量的 4 倍。
-
-所有 \(r\) 的输出组合为：
-
-\[
-Y\in
-\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}},
-\]
-
-再 reshape 为：
-
-\[
-\Delta\in\mathbb{R}^{B\times L\times d}.
-\]
-
-## 6. 参数张量与无歧义索引
-
-建议的实现存储形状为：
 
 ```text
+cross_norm:  LayerNorm(n_group)
 group_align: [n_group, d_group, d_group]
-gate_proj:   [d_group, n_group, hidden_group]
+gate_proj:   [d_group, n_group, 4 * n_group]
-value_proj:  [d_group, n_group, hidden_group]
+value_proj:  [d_group, n_group, 4 * n_group]
-output_proj: [d_group, hidden_group, n_group]
+output_proj: [d_group, 4 * n_group, n_group]
 ```
 
-对应的索引公式为：
+三个 projection 均不带 bias。
+
+## 6. PreNorm 与残差约束
+
+本版本固定使用两个独立的 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 组内阶段
+
+投影权重：
 
 \[
-G_{b,t,q,r}
+3G D(4D)
-=\sum_i Z_{b,t,i,r}\,A_{g,r,i,q},
+=12GD^2
+=17{,}280.
 \]
 
+`LayerNorm(D)`：
+
 \[
-V_{b,t,q,r}
+2D=24.
-=\sum_i Z_{b,t,i,r}\,A_{v,r,i,q},
 \]
 
-\[
+### 7.2 跨组阶段
-Y_{b,t,i,r}
-=\sum_q
-\left[\operatorname{SiLU}(G_{b,t,q,r})V_{b,t,q,r}\right]
-A_{o,r,q,i}.
-\]
 
-首版的三个 Mixer projection 均不带 bias。
+跨组投影权重：
-
-## 7. Mixer Hidden Width 与参数量
-
-Mixer 的 group 维变换为：
 
 \[
-n_{\mathrm{group}}
+3D G(4G)
-\rightarrow
+=12DG^2
-h_{\mathrm{group}}
-\rightarrow
-n_{\mathrm{group}}.
-\]
-
-固定 \(h_{\mathrm{group}}=4n_{\mathrm{group}}=40\) 时，Mixer 每层权重参数量为：
-
-\[
-3d_{\mathrm{group}}n_{\mathrm{group}}h_{\mathrm{group}}
-=3\times12\times10\times40
 =14{,}400.
 \]
 
-加上 Group Feature Alignment 后，TrajMixer residual branch 每层共有：
+Group Feature Alignment：
 
 \[
-14{,}400+1{,}440=15{,}840
+GD^2
+=1{,}440.
 \]
 
-个主要权重参数。作为对照，原始 \(120\rightarrow480\rightarrow120\) FFN 每层约有 115,800 个参数。
+`LayerNorm(G)`：
 
-参数对照口径说明：上面的 115,800 对应结构方案中的标准两层 FFN。当前代码库在 TrajMixer 替换前实际使用的是隐藏宽度 300 的全维度 SwiGLU（gate/value/output 三个线性层），每层共有 108,720 个参数（含 bias）。代码实验和 checkpoint 参数量比较必须以 108,720 作为历史实现基线，不能与概念方案中的标准 FFN 参数量混用。
+\[
+2G=20.
+\]
 
-## 8. LayerNorm 基线与消融
+### 7.3 每个 TrajMixer 合计
 
-为保持与原始 Transformer 的可比性，首版固定使用：
+\[
+17{,}280+24+14{,}400+1{,}440+20
+=\boxed{33{,}164}.
+\]
 
-```text
+相对于 `traj_mixer_v2` 的跨组单阶段结构 \(15{,}840\)，每层增加 \(17{,}324\) 个参数。作为历史实现对照，代码库原全维度 SwiGLU FFN 每层为 \(108{,}720\) 个参数。
-原始 FFN baseline：FFN + 标准 LayerNorm
-TrajMixer baseline：Mixer + 标准 LayerNorm
-```
 
-以下配置不属于首版主实验，只作为独立消融：
+## 8. 初始化
 
-```text
+固定初始化约定：
-Mixer + Group-wise LayerNorm
-```
 
-不得将 Group-wise LayerNorm 的结果直接作为“仅替换 FFN”的对照结果。
+- 组内 `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}\) 的正态分布；
+- 两个 LayerNorm：PyTorch 默认 affine 初始化；
+- 两个 residual stage 的 Dropout 均沿用 `mlp_dropout`。
 
-## 9. 初始化
+两个 output projection 的小方差初始化使两阶段在训练初期都接近恒等 residual update。
 
-首版初始化约定：
+## 9. 信息流与语义
 
-- Group Alignment \(B_i\)：单位矩阵初始化；
+**Attention**：从历史疾病事件中选择和整合相关信息。
-- \(A_g/A_v\)：Xavier uniform 初始化；
-- \(A_o\)：均值为 0、标准差为 \(10^{-3}\) 的正态初始化；
-- Dropout 概率沿用原始 FFN residual branch 的配置。
 
-小初始化的 \(A_o\) 使新增分支在训练初期接近恒等残差更新，同时允许模型逐步学习轨迹交互。
+**Intra-Group Mixer**：学习每条潜在轨迹内部的非线性特征组合。
 
-## 10. 核心设计思想
+**Group Feature Alignment**：对齐不同潜在轨迹的内部坐标。
 
-**Attention**：负责从历史疾病序列中选择并整合相关信息。
+**Cross-Group Mixer**：学习不同潜在轨迹在相同内部坐标上的门控交互。
-
-**Group Feature Alignment**：负责学习不同 latent trajectory groups 的内部特征对齐。
-
-**Cross-Group Mixer**：负责不同潜在疾病轨迹之间的非线性门控交互。
 
 整个模块保持：
 
 - 无时间递归；
+- 不混合序列位置；
 - 序列维度完全并行；
-- 参数量远低于原始 FFN；
+- 保留原始因果 Attention；
-- 保留 Transformer 的因果历史建模能力；
+- residual groups 不等同于 Attention heads。
-- 不把 residual groups 误解释为原始 Attention heads。
 
-## 11. 首版固定配置
+## 10. 固定配置与 checkpoint 约束
 
 ```yaml
-model_architecture: traj_mixer_v2
+model_architecture: traj_mixer_v3
 d_model: 120
-n_head: 10  # 同时决定 residual group 数量
+n_head: 10
-d_group: 12
+n_group_rule: n_head
-hidden_group_rule: 4 * n_head  # 不单独配置
+d_group_rule: d_model / n_group
+intra_hidden_rule: 4 * d_group
+cross_hidden_rule: 4 * n_group
 attention: unchanged
-attention_output_projection: unchanged
+intra_norm: layer_norm_over_d_group
-mixer_norm: standard_layer_norm
+cross_norm: layer_norm_over_n_group
-group_alignment: per_group_12x12
+group_alignment: per_group_d_group_x_d_group
-group_alignment_bias: false
+projection_bias: false
-group_alignment_init: identity
-mixer_bias: false
 gate_value_init: xavier_uniform
 output_init_std: 0.001
-group_wise_layer_norm: false
 ```
 
-训练时必须将 `model_architecture: traj_mixer_v2`、`model_parameter_count` 和 `trainable_parameter_count` 写入 `train_config.json`，并在训练日志中显式打印总参数量与可训练参数量。本分支的评估和导出入口只接受带有该架构标识、且 checkpoint 中包含 TrajMixer 参数张量的模型；其他版本或分支生成的模型应直接拒绝加载。
-
 必须满足：
 
 \[
 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

@@ -177,7 +177,7 @@ class TemporalAttention(nn.Module):
 
 
 class TrajMixer(nn.Module):
-    """Lightweight gated interaction across latent residual-space groups.
+    """Two-stage 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.
@@ -205,8 +205,24 @@ class TrajMixer(nn.Module):
         # are still residual-space partitions rather than attention heads.
         self.n_group = n_head
         self.d_group = n_embd // n_head
+        self.intra_hidden = 4 * self.d_group
         self.hidden_group = 4 * n_head
 
+        # 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)
+        )
+        self.intra_value_proj = nn.Parameter(
+            torch.empty(self.n_group, self.d_group, self.intra_hidden)
+        )
+        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)
@@ -227,6 +243,11 @@ class TrajMixer(nn.Module):
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
+        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,
@@ -243,7 +264,7 @@ class TrajMixer(nn.Module):
         nn.init.normal_(self.output_proj, mean=0.0, std=1e-3)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Map ``(B, L, n_embd)`` to an equally shaped residual update."""
+        """Apply two PreNorm residual stages without mixing sequence positions."""
         if x.ndim != 3:
             raise ValueError(f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}")
         if x.size(-1) != self.n_embd:
@@ -255,8 +276,27 @@ class TrajMixer(nn.Module):
         grouped = 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
+        )
+        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)
+
+        # 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", grouped, self.group_align
+            "blgd,gde->blge", cross_input, self.group_align
         )
 
         gate = torch.einsum(
@@ -269,7 +309,8 @@ class TrajMixer(nn.Module):
         mixed = torch.einsum(
             "blhr,rhg->blgr", hidden, self.output_proj
         )
-        return self.drop(mixed.reshape(batch_size, seq_len, self.n_embd))
+        grouped = grouped + self.drop(mixed)
+        return grouped.reshape(batch_size, seq_len, self.n_embd)
 
 
 class GPTBlock(nn.Module):
@@ -299,7 +340,6 @@ class GPTBlock(nn.Module):
             dropout=mlp_dropout,
         )
         self.ln1 = nn.LayerNorm(n_embd)
-        self.ln2 = nn.LayerNorm(n_embd)
 
     def forward(
         self,
@@ -309,8 +349,7 @@ class GPTBlock(nn.Module):
         attn_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
-        x = x + self.mlp(self.ln2(x))
+        return self.mlp(x)
-        return x
 
 
 class TokenAutoDiscretization(nn.Module):

models.py

@@ -15,7 +15,7 @@ from backbones import (
 from targets import PAD_IDX
 
 
-TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v2"
+TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v3"
 
 
 def validate_traj_mixer_config(config: Mapping[str, object]) -> None:
@@ -29,6 +29,13 @@ 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.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.gate_proj",
         "blocks.0.mlp.value_proj",

test_traj_mixer.py

@@ -21,14 +21,98 @@ 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()), 15_840)
+        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(mixer.intra_hidden, 48)
+        self.assertEqual(
+            tuple(mixer.intra_gate_proj.shape),
+            (10, 12, 48),
+        )
+        self.assertEqual(
+            tuple(mixer.intra_value_proj.shape),
+            (10, 12, 48),
+        )
+        self.assertEqual(
+            tuple(mixer.intra_output_proj.shape),
+            (10, 48, 12),
+        )
         self.assertEqual(mixer.hidden_group, 40)
         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:
+        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_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
+        )
+        unchanged_groups = torch.tensor([0, 1, 2, 4, 5, 6, 7, 8, 9])
+        torch.testing.assert_close(
+            original_out.index_select(2, unchanged_groups),
+            changed_out.index_select(2, unchanged_groups),
+        )
+
+    def test_cross_stage_mixes_groups_without_mixing_coordinates(self) -> None:
+        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_()
+
+            # For coordinate 0 only, read group 0 through hidden unit 0 and
+            # write the resulting gated value into group 1.
+            mixer.gate_proj[0, 0, 0] = 1.0
+            mixer.value_proj[0, 0, 0] = 1.0
+            mixer.output_proj[0, 0, 1] = 1.0
+
+        grouped = torch.tensor(
+            [[[
+                [-1.0, 4.0],
+                [0.0, 5.0],
+                [1.0, 6.0],
+            ]]]
+        )
+        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)
+
+        self.assertNotEqual(
+            original_out[0, 0, 1, 0].item(),
+            changed_out[0, 0, 1, 0].item(),
+        )
+        torch.testing.assert_close(
+            original_out[..., 1],
+            changed_out[..., 1],
+        )
 
     def test_mixer_does_not_mix_sequence_positions(self) -> None:
         torch.manual_seed(0)
@@ -58,11 +142,12 @@ class TrajMixerTest(unittest.TestCase):
             self.assertIsNotNone(parameter.grad, name)
             self.assertTrue(torch.isfinite(parameter.grad).all(), name)
 
-    def test_gpt_block_defaults_to_traj_mixer_and_standard_layer_norm(self) -> None:
+    def test_gpt_block_delegates_both_mixer_residuals_to_traj_mixer(self) -> None:
         block = GPTBlock(n_embd=120, n_head=10)
         self.assertIsInstance(block.mlp, TrajMixer)
-        self.assertIsInstance(block.ln2, torch.nn.LayerNorm)
+        self.assertFalse(hasattr(block, "ln2"))
-        self.assertEqual(tuple(block.ln2.normalized_shape), (120,))
+        self.assertIsInstance(block.mlp.intra_norm, torch.nn.LayerNorm)
+        self.assertIsInstance(block.mlp.cross_norm, torch.nn.LayerNorm)
 
         x = torch.randn(2, 6, 120)
         self.assertEqual(block(x).shape, x.shape)
@@ -75,6 +160,10 @@ class TrajMixerTest(unittest.TestCase):
             validate_traj_mixer_config({})
         with self.assertRaisesRegex(ValueError, "only accepts models trained"):
             validate_traj_mixer_config({"model_architecture": "delphi_swiglu"})
+        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
+            validate_traj_mixer_config(
+                {"model_architecture": "traj_mixer_v2"}
+            )
 
     def test_checkpoint_must_contain_traj_mixer_parameters(self) -> None:
         block = GPTBlock(n_embd=120, n_head=10)
@@ -84,7 +173,7 @@ class TrajMixerTest(unittest.TestCase):
         }
         validate_traj_mixer_state_dict(state_dict)
 
-        state_dict.pop("blocks.0.mlp.group_align")
+        state_dict.pop("blocks.0.mlp.intra_gate_proj")
         with self.assertRaisesRegex(ValueError, "not a TrajMixer checkpoint"):
             validate_traj_mixer_state_dict(state_dict)
 
@@ -97,8 +186,8 @@ class TrajMixerTest(unittest.TestCase):
         self.assertEqual(
             get_model_parameter_counts(mixer),
             {
-                "model_parameter_count": 15_840,
+                "model_parameter_count": 33_164,
-                "trainable_parameter_count": 15_840,
+                "trainable_parameter_count": 33_164,
             },
         )