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；
-- 原始 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 不向后兼容，直接拒绝加载。