Report AUCs in Delphi2M format

TrajMixer_设计方案.md

@@ -2,332 +2,384 @@
 
 > 状态：**Frozen implementation baseline**
 >
-> 版本：**v1.0**
+> 版本：**v3.0 / traj_mixer_v5**
 >
-> 固化日期：**2026-07-22**
+> 固化日期：**2026-07-24**
 
-本文档是 TrajMixer 后续实现与实验的唯一结构基线。除显式标记为消融项的配置外，所有实现均应遵循本文档；若结构发生变化，应先更新版本和实验记录。
+本文档是当前 TrajMixer 的实现与实验基线。本版本采用单 PreNorm、单外层 residual、静态门控组内融合和跨 group SwiGLU。
 
 ## 1. 目标
 
-在保持原始 Delphi Transformer Attention 结构不变的前提下，用轻量、可并行的轨迹交互模块替换 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 不修改 Attention，只替换每个 Transformer block 中的 FFN residual branch。
+TrajMixer 不沿序列维度混合，也不引入时间递归。
 
-## 2. Block 总体结构
+## 2. Block 结构
-
-概念结构：
 
 ```text
 PreNorm Causal Multi-Head Attention
-→ Residual
+→ Attention Residual
-→ Standard Mixer PreNorm
+→ Full-width TrajMixer PreNorm
-→ Group-wise Feature Alignment
+→ reshape [B, L, n_group, d_group]
-→ SwiGLU Cross-Group Mixer
+→ Per-Group SwiGLU: d_group → 4d_group → d_group
-→ Residual
+→ Static Gated Fusion
+→ Cross-Group SwiGLU: n_group → 4n_group → n_group
+→ reshape [B, L, n_embd]
+→ Dropout
+→ One TrajMixer Residual
 ```
 
-完整计算为：
+Attention 阶段：
 
 \[
-U = X^{(l)} + \operatorname{Dropout}\!\left(
+X
-\operatorname{CausalMHA}\left(
+=X^{(l)}
-\operatorname{LN}_{\mathrm{attn}}(X^{(l)}),
++\operatorname{CausalMHA}
-\text{time information}
+\left(\operatorname{LN}_{\mathrm{attn}}(X^{(l)})\right).
-\right)\right),
+\]
+
+TrajMixer 阶段：
+
+\[
+N=\operatorname{LN}_{d}(X),
 \]
 
 \[
-N = \operatorname{LN}_{\mathrm{mixer}}(U),
+G=\operatorname{reshape}(N)
+\in\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}},
 \]
 
 \[
-\Delta = \operatorname{TrajMixer}(N),
+P=\operatorname{IntraMixer}(G),
 \]
 
 \[
-X^{(l+1)} = U + \operatorname{Dropout}(\Delta).
+U=G+\sigma(\Theta)\odot P,
 \]
 
-首版中的 \(\operatorname{LN}_{\mathrm{mixer}}\) 是作用于完整 \(d=120\) 维 residual representation 的标准 LayerNorm。
+\[
+\Delta=\operatorname{reshape}
+\left(\operatorname{CrossMixer}(U)\right),
+\]
 
-## 3. Latent Trajectory Group 定义
+\[
+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：
 
 \[
-N\in\mathbb{R}^{B\times L\times d},\qquad d=120.
+X\in\mathbb{R}^{B\times L\times d}.
 \]
 
-将 hidden dimension 划分为与 Attention head 数量相同的 group 数量：
+定义：
 
 \[
-n_{\mathrm{group}}:=n_{\mathrm{head}}=10,
+n_{\mathrm{group}}:=n_{\mathrm{head}},
-\qquad d_{\mathrm{group}}=\frac{d}{n_{\mathrm{head}}}=12,
+\qquad
+d_{\mathrm{group}}=\frac{d}{n_{\mathrm{group}}},
+\qquad
+d=n_{\mathrm{group}}d_{\mathrm{group}}.
 \]
 
-`n_group` 不再是独立超参数，代码统一使用 `n_head` 确定 residual group 数量。二者只共享数量；这些 residual groups 在语义和张量来源上仍不等同于原始 Attention heads。
+默认：
-
-并 reshape 为：
 
 \[
-N_{\mathrm{group}}in
+d=120,\qquad
-\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}.
+n_{\mathrm{group}}=10,\qquad
+d_{\mathrm{group}}=12.
 \]
 
-这些 group 是 residual space 中的 **latent trajectory groups**，不等同于原始 Attention heads。本文中的 group、trajectory group 均指这一 residual-channel partition。
+这些 group 是 residual space 的连续分区，不等同于 Attention heads；二者只共享数量。
 
-## 4. Group-wise Feature Alignment
+## 4. 唯一的 Full-Width PreNorm
 
-为缓解不同 group 内部坐标不对齐的问题，每个 group 使用独立的小矩阵：
+TrajMixer 只使用一个：
+
+```text
+norm: LayerNorm(n_embd)
+```
+
+LayerNorm 作用于完整 \(d\) 维 residual representation，然后才 reshape：
 
 \[
-B_i\in\mathbb{R}^{d_{\mathrm{group}}\times d_{\mathrm{group}}},
+G=\operatorname{reshape}
-\qquad i=1,\ldots,n_{\mathrm{group}}.
+\left(\operatorname{LN}_{d}(X)\right).
 \]
 
-对每个 group 内的特征进行可学习对齐：
+本版本明确删除：
+
+```text
+intra_norm
+cross_norm
+group_align
+```
+
+不得在组内或跨组阶段再增加额外 LayerNorm。
+
+## 5. 组内 SwiGLU
+
+每个 group 使用独立参数，对其 \(d_{\mathrm{group}}\) 维内部特征执行：
 
 \[
-Z_{b,t,i,:}=N_{\mathrm{group},b,t,i,:}B_i.
+d_{\mathrm{group}}
+\rightarrow
+4d_{\mathrm{group}}
+\rightarrow
+d_{\mathrm{group}}.
+\]
+
+对 group \(g\)：
+
+\[
+W_{g,\mathrm{intra}}^{(g)},
+W_{v,\mathrm{intra}}^{(g)}
+\in
+\mathbb{R}^{d_{\mathrm{group}}\times4d_{\mathrm{group}}},
+\]
+
+\[
+W_{o,\mathrm{intra}}^{(g)}
+\in
+\mathbb{R}^{4d_{\mathrm{group}}\times d_{\mathrm{group}}}.
+\]
+
+计算：
+
+\[
+H_g
+=
+\operatorname{SiLU}
+\left(G_gW_{g,\mathrm{intra}}^{(g)}\right)
+\odot
+\left(G_gW_{v,\mathrm{intra}}^{(g)}\right),
+\]
+
+\[
+P_g=H_gW_{o,\mathrm{intra}}^{(g)}.
+\]
+
+实现形状：
+
+```text
+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。
+
+## 6. 静态门控融合
+
+定义可学习 gate logits：
+
+\[
+\Theta\in
+\mathbb{R}^{n_{\mathrm{group}}\times d_{\mathrm{group}}}.
+\]
+
+实际门值为：
+
+\[
+\Gamma=\sigma(\Theta).
+\]
+
+初始化：
+
+\[
+\Theta_{g,r}
+=\operatorname{logit}(0.1)
+=\log\frac{0.1}{0.9}
+\approx-2.1972,
 \]
 
 因此：
 
 \[
-Z\in
+\Gamma_{g,r}\approx0.1.
-\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}.
 \]
 
-首版实现约定：
+融合：
-
-- \(B_i\) 不带 bias；
-- \(B_i\) 使用单位矩阵初始化；
-- Alignment 只作用于 Mixer residual branch，不改变 Attention residual stream；
-- 首版不增加逆变换或额外的 group 内输出投影。
-
-Alignment 每层权重参数量为：
 
 \[
-n_{\mathrm{group}}d_{\mathrm{group}}^2
+U=G+\Gamma\odot P.
-=10\times12^2
-=1{,}440.
 \]
 
-## 5. SwiGLU Cross-Group Mixer
+\(\Gamma\) 对 batch 和序列位置共享，但每个 group、每个内部坐标拥有独立可学习值。
 
-Mixer 只沿 group 维度交互，不沿序列维度交互，因此不会引入时间递归或未来信息泄漏。
+## 7. 跨 Group SwiGLU
 
-对于每个 group 内特征维度：
+对于每个内部坐标 \(r\)，独立沿 group 维度执行：
 
 \[
-r=1,\ldots,d_{\mathrm{group}},
+n_{\mathrm{group}}
+\rightarrow
+4n_{\mathrm{group}}
+\rightarrow
+n_{\mathrm{group}}.
 \]
 
 定义：
 
 \[
 A_g^{(r)},A_v^{(r)}
-\in\mathbb{R}^{n_{\mathrm{group}}\times h_{\mathrm{group}}},
+\in
+\mathbb{R}^{n_{\mathrm{group}}\times4n_{\mathrm{group}}},
 \]
 
 \[
 A_o^{(r)}
-\in\mathbb{R}^{h_{\mathrm{group}}\times n_{\mathrm{group}}}.
+\in
+\mathbb{R}^{4n_{\mathrm{group}}\times n_{\mathrm{group}}}.
 \]
 
-隐藏宽度不再独立配置，固定为：
+计算：
 
 \[
-h_{\mathrm{group}}=4n_{\mathrm{head}}
+Q_{:,r}
-=4n_{\mathrm{group}}.
+=
-\]
+\operatorname{SiLU}\left(U_{:,r}A_g^{(r)}\right)
-
+\odot
-当前 \(n_{\mathrm{head}}=10\)，因此 \(h_{\mathrm{group}}=40\)。
+\left(U_{:,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_{:,r}=Q_{:,r}A_o^{(r)}.
 \]
 
-\[
+实现形状：
-M_{b,t,:,r}=\operatorname{SiLU}(G_{b,t,:,r})\odot V_{b,t,:,r},
-\]
 
-\[
+```text
-Y_{b,t,:,r}=M_{b,t,:,r}A_o^{(r)}.
+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。不同内部坐标拥有独立的跨 group 参数，且不沿序列维度交互。
 
-- gate 分支控制信息写入；
+## 8. 唯一的外层 Residual
-- value 分支提供交互内容；
-- output matrix 将隐藏 group 表示投影回原始 group 数量；
-- hidden group 表示固定扩展为 group 数量的 4 倍。
 
-所有 \(r\) 的输出组合为：
+跨 group 输出 reshape 回：
-
-\[
-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
-group_align: [n_group, d_group, d_group]
-gate_proj:   [d_group, n_group, hidden_group]
-value_proj:  [d_group, n_group, hidden_group]
-output_proj: [d_group, hidden_group, n_group]
-```
-
-对应的索引公式为：
 
 \[
-G_{b,t,q,r}
+\operatorname{TrajMixer}(X)
-=\sum_i Z_{b,t,i,r}\,A_{g,r,i,q},
+=X+\operatorname{Dropout}(\Delta).
 \]
 
-\[
+固定约束：
-V_{b,t,q,r}
-=\sum_i Z_{b,t,i,r}\,A_{v,r,i,q},
-\]
 
-\[
+- 组内阶段后不执行独立 residual；
-Y_{b,t,i,r}
+- 跨组阶段后不执行独立 residual；
-=\sum_q
+- `GPTBlock` 不再额外执行 `X + TrajMixer(X)`；
-\left[\operatorname{SiLU}(G_{b,t,q,r})V_{b,t,q,r}\right]
+- 整个 TrajMixer 只有一次主 residual。
-A_{o,r,q,i}.
-\]
-
-首版的三个 Mixer projection 均不带 bias。
-
-## 7. Mixer Hidden Width 与参数量
-
-Mixer 的 group 维变换为：
-
-\[
-n_{\mathrm{group}}
-\rightarrow
-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 每层共有：
-
-\[
-14{,}400+1{,}440=15{,}840
-\]
-
-个主要权重参数。作为对照，原始 \(120\rightarrow480\rightarrow120\) FFN 每层约有 115,800 个参数。
-
-参数对照口径说明：上面的 115,800 对应结构方案中的标准两层 FFN。当前代码库在 TrajMixer 替换前实际使用的是隐藏宽度 300 的全维度 SwiGLU（gate/value/output 三个线性层），每层共有 108,720 个参数（含 bias）。代码实验和 checkpoint 参数量比较必须以 108,720 作为历史实现基线，不能与概念方案中的标准 FFN 参数量混用。
-
-## 8. LayerNorm 基线与消融
-
-为保持与原始 Transformer 的可比性，首版固定使用：
-
-```text
-原始 FFN baseline：FFN + 标准 LayerNorm
-TrajMixer baseline：Mixer + 标准 LayerNorm
-```
-
-以下配置不属于首版主实验，只作为独立消融：
-
-```text
-Mixer + Group-wise LayerNorm
-```
-
-不得将 Group-wise LayerNorm 的结果直接作为“仅替换 FFN”的对照结果。
 
 ## 9. 初始化
 
-首版初始化约定：
+固定初始化：
 
-- Group Alignment \(B_i\)：单位矩阵初始化；
+- `intra_gate_proj/intra_value_proj`：每个 group 独立 Xavier uniform；
-- \(A_g/A_v\)：Xavier uniform 初始化；
+- `intra_output_proj`：每个 group 独立 Xavier uniform；
-- \(A_o\)：均值为 0、标准差为 \(10^{-3}\) 的正态初始化；
+- `intra_gate_logits`：初始化为 \(\operatorname{logit}(0.1)\)；
-- Dropout 概率沿用原始 FFN residual branch 的配置。
+- 跨组 `gate_proj/value_proj`：每个内部坐标独立 Xavier uniform；
+- 最终跨组 `output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布；
+- Full-width LayerNorm：PyTorch 默认 affine 初始化；
+- Dropout：沿用 `mlp_dropout`。
 
-小初始化的 \(A_o\) 使新增分支在训练初期接近恒等残差更新，同时允许模型逐步学习轨迹交互。
+组内输出使用正常 Xavier 初始化以保证其具有完整表达能力；静态门控将其初始贡献限制在约 0.1。最终跨 group 输出投影保持小值初始化，使整个 TrajMixer residual update 在训练初期接近零。
 
-## 10. 核心设计思想
+Relative Time Attention Bias 初始化固定为：
 
-**Attention**：负责从历史疾病序列中选择并整合相关信息。
+- `rbf_proj.weight`：零初始化；
+- `time_bias_scale`：初始化为 \(1.0\)；
+- 初始 RBF attention bias 严格为零；
+- `rbf_proj.weight` 从第一个优化步骤即可获得梯度。
 
-**Group Feature Alignment**：负责学习不同 latent trajectory groups 的内部特征对齐。
+## 10. 参数量
 
-**Cross-Group Mixer**：负责不同潜在疾病轨迹之间的非线性门控交互。
+默认 \(d=120\)、\(n_{\mathrm{group}}=10\)、\(d_{\mathrm{group}}=12\)。
 
-整个模块保持：
+Full-width LayerNorm：
-
-- 无时间递归；
-- 序列维度完全并行；
-- 参数量远低于原始 FFN；
-- 保留 Transformer 的因果历史建模能力；
-- 不把 residual groups 误解释为原始 Attention heads。
-
-## 11. 首版固定配置
-
-```yaml
-model_architecture: traj_mixer_v2
-d_model: 120
-n_head: 10  # 同时决定 residual group 数量
-d_group: 12
-hidden_group_rule: 4 * n_head  # 不单独配置
-attention: unchanged
-attention_output_projection: unchanged
-mixer_norm: standard_layer_norm
-group_alignment: per_group_12x12
-group_alignment_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}}.
+2d=240.
 \]
 
-后续实现、单元测试、参数量核验和主实验均以以上配置为默认基线。
+组内 projections：
+
+\[
+3n_{\mathrm{group}}d_{\mathrm{group}}
+\left(4d_{\mathrm{group}}\right)
+=17{,}280.
+\]
+
+静态门控：
+
+\[
+n_{\mathrm{group}}d_{\mathrm{group}}
+=120.
+\]
+
+跨 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_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
+group_alignment: false
+intra_residual: false
+cross_residual: false
+traj_mixer_outer_residual: true
+projection_bias: false
+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`，并在日志中打印参数量。
+
+评估和导出入口只接受 `traj_mixer_v5` checkpoint，并检查 Full-width LayerNorm、组内 projections、静态门控和跨 group projections 是否齐全。`traj_mixer_v4` 及更早 checkpoint 不向后兼容，直接拒绝加载。

backbones.py

@@ -111,7 +111,10 @@ class TemporalAttention(nn.Module):
 
         # Layer-specific projection from shared RBF basis activations to per-head attention bias.
         self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
-        self.time_bias_scale = nn.Parameter(torch.tensor(0.0))
+        # Keep the initial RBF attention bias exactly zero through the
+        # zero-initialized projection, while leaving that projection with a
+        # live gradient from the first optimization step.
+        self.time_bias_scale = nn.Parameter(torch.tensor(1.0))
 
         self.resid_drop = nn.Dropout(dropout)
         self.reset_parameters()
@@ -177,7 +180,7 @@ class TemporalAttention(nn.Module):
 
 
 class TrajMixer(nn.Module):
-    """Lightweight gated interaction across latent residual-space 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.
@@ -205,11 +208,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
 
-        # Per-group feature alignment: [group, input feature, output feature].
+        # A single full-width PreNorm serves the entire TrajMixer branch.
-        self.group_align = nn.Parameter(
+        self.norm = nn.LayerNorm(self.n_embd)
-            torch.empty(self.n_group, self.d_group, self.d_group)
+
+        # Stage 1: each group independently mixes its internal features.
+        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)
+        )
+        self.intra_gate_logits = nn.Parameter(
+            torch.empty(self.n_group, self.d_group)
         )
 
         # Per-feature cross-group projections. The feature index is kept
@@ -227,13 +243,14 @@ class TrajMixer(nn.Module):
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
-        with torch.no_grad():
+        for group_idx in range(self.n_group):
-            identity = torch.eye(
+            nn.init.xavier_uniform_(self.intra_gate_proj[group_idx])
-                self.d_group,
+            nn.init.xavier_uniform_(self.intra_value_proj[group_idx])
-                dtype=self.group_align.dtype,
+            nn.init.xavier_uniform_(self.intra_output_proj[group_idx])
-                device=self.group_align.device,
+        nn.init.constant_(
+            self.intra_gate_logits,
+            math.log(0.1 / 0.9),
         )
-            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.
@@ -242,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:
-        """Map ``(B, L, n_embd)`` to an equally shaped residual update."""
+        """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:
@@ -252,24 +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
         )
-        aligned = torch.einsum(
-            "blgd,gde->blge", grouped, self.group_align
-        )
 
-        gate = torch.einsum(
+        # The static per-channel gate starts at sigmoid(logit) ~= 0.1.
-            "blgr,rgh->blhr", aligned, self.gate_proj
+        intra_output = self._intra_mix(grouped)
+        intra_gate = torch.sigmoid(self.intra_gate_logits).view(
+            1, 1, self.n_group, self.d_group
         )
-        value = torch.einsum(
+        mixed_input = grouped + intra_gate * intra_output
-            "blgr,rgh->blhr", aligned, self.value_proj
+
+        # Stage 2: n_group -> 4*n_group -> n_group for each coordinate.
+        update = self._cross_mix(mixed_input).reshape(
+            batch_size, seq_len, self.n_embd
         )
-        hidden = F.silu(gate) * value
+        return x + self.drop(update)
-        mixed = torch.einsum(
-            "blhr,rhg->blgr", hidden, self.output_proj
-        )
-        return self.drop(mixed.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):

delphi2m_auc_report.py

@@ -0,0 +1,183 @@
+"""Build Delphi2M-style sex-specific AUC reports.
+
+The Delphi2M evaluation code uses 0.1 years for the no-gap evaluation.  The
+published report displays that point as 0 months, while retaining the actual
+0.1-year evaluation period in this project's report output.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict, Optional
+
+import numpy as np
+import pandas as pd
+
+
+DEFAULT_DELPHI2M_PERIODS_YEARS = (0.1, 1.0, 5.0, 10.0)
+
+_CHAPTER_SHORT_NAMES = {
+    "I": "I. Infectious Diseases",
+    "II": "II. Neoplasms",
+    "III": "III. Blood & Immune Disorders",
+    "IV": "IV. Metabolic Diseases",
+    "V": "V. Mental Disorders",
+    "VI": "VI. Nervous System Diseases",
+    "VII": "VII. Eye Diseases",
+    "VIII": "VIII. Ear Diseases",
+    "IX": "IX. Circulatory Diseases",
+    "X": "X. Respiratory Diseases",
+    "XI": "XI. Digestive Diseases",
+    "XII": "XII. Skin Diseases",
+    "XIII": "XIII. Musculoskeletal Diseases",
+    "XIV": "XIV. Genitourinary Diseases",
+    "XV": "XV. Pregnancy & Childbirth",
+    "XVI": "XVI. Perinatal Conditions",
+    "XVII": "XVII. Congenital Abnormalities",
+    "XVIII": "XVIII. Symptoms & Signs",
+    "XIX": "XIX. Injury & Poisoning",
+    "XX": "XX. External Causes",
+    "XXI": "XXI. Health Services",
+    "XXII": "XXII. Special Purposes",
+    "Death": "Death",
+    "Unmapped": "Unmapped",
+}
+
+
+def _is_no_gap(period_years: float) -> bool:
+    return bool(np.isclose(float(period_years), 0.1, rtol=0.0, atol=1e-8))
+
+
+def _canonical_period_years(period_years: float) -> float:
+    value = float(period_years)
+    for canonical in DEFAULT_DELPHI2M_PERIODS_YEARS:
+        if np.isclose(value, canonical, rtol=0.0, atol=1e-6):
+            return float(canonical)
+    return value
+
+
+def _gap_months(period_years: float) -> int:
+    if _is_no_gap(period_years):
+        return 0
+    return int(round(float(period_years) * 12.0))
+
+
+def _gap_label(period_years: float) -> str:
+    if _is_no_gap(period_years):
+        return "No gap"
+    value = float(period_years)
+    value_text = f"{value:g}"
+    unit = "year" if np.isclose(value, 1.0) else "years"
+    return f"{value_text} {unit}"
+
+
+def _load_chapter_by_code(
+    chapter_mapping_path: Optional[str | Path] = None,
+) -> Dict[str, str]:
+    if chapter_mapping_path is None:
+        chapter_mapping_path = Path(__file__).with_name(
+            "icd10_chapter_organ_mapping.csv"
+        )
+    path = Path(chapter_mapping_path)
+    if not path.exists():
+        return {}
+
+    mapping = pd.read_csv(
+        path,
+        usecols=["code", "icd10_chapter"],
+        dtype={"code": str, "icd10_chapter": str},
+    )
+    mapping["code"] = mapping["code"].str.strip()
+    mapping["chapter"] = (
+        mapping["icd10_chapter"]
+        .str.strip()
+        .map(_CHAPTER_SHORT_NAMES)
+        .fillna("Unmapped")
+    )
+    return dict(zip(mapping["code"], mapping["chapter"]))
+
+
+def build_delphi2m_auc_report(
+    df_unpooled: pd.DataFrame,
+    *,
+    period_col: str,
+    chapter_mapping_path: Optional[str | Path] = None,
+) -> pd.DataFrame:
+    """Aggregate age strata by sex and return a Delphi2M-style AUC report.
+
+    Required input columns are ``token``, ``label_code``, ``sex``,
+    ``auc_delong``, and the supplied ``period_col`` (``offset`` or
+    ``horizon``).  The output begins with the five columns used by Delphi2M
+    Fig. 2e and then records the actual evaluation period and ICD-10 code.
+    """
+    required = {"token", "label_code", "sex", "auc_delong", period_col}
+    missing = sorted(required - set(df_unpooled.columns))
+    if missing:
+        raise ValueError(
+            "Cannot build Delphi2M AUC report; missing columns: "
+            + ", ".join(missing)
+        )
+
+    source = df_unpooled.loc[
+        :,
+        ["token", "label_code", "sex", "auc_delong", period_col],
+    ].copy()
+    source["sex"] = source["sex"].astype(str).str.strip().str.lower()
+    source = source[source["sex"].isin(["female", "male"])]
+    source["auc_delong"] = pd.to_numeric(
+        source["auc_delong"], errors="coerce"
+    )
+    source[period_col] = pd.to_numeric(source[period_col], errors="coerce")
+    source = source.dropna(subset=[period_col, "auc_delong"])
+    source[period_col] = source[period_col].map(_canonical_period_years)
+
+    if source.empty:
+        raise ValueError("Cannot build Delphi2M AUC report from empty AUC data.")
+
+    grouped = (
+        source.groupby(
+            ["token", "label_code", period_col, "sex"],
+            dropna=False,
+            as_index=False,
+        )
+        .agg(auc=("auc_delong", "mean"))
+    )
+    report = (
+        grouped.pivot(
+            index=["token", "label_code", period_col],
+            columns="sex",
+            values="auc",
+        )
+        .reset_index()
+        .rename_axis(columns=None)
+        .rename(columns={"female": "Female", "male": "Male"})
+    )
+    for col in ["Female", "Male"]:
+        if col not in report.columns:
+            report[col] = np.nan
+
+    chapter_by_code = _load_chapter_by_code(chapter_mapping_path)
+    report["chapter"] = (
+        report["label_code"].astype(str).map(chapter_by_code).fillna("Unmapped")
+    )
+    report["Gap, months"] = report[period_col].map(_gap_months).astype("Int64")
+    report["Gap label"] = report[period_col].map(_gap_label)
+    report["icd10"] = pd.to_numeric(report["token"], errors="coerce").astype(
+        "Int64"
+    )
+
+    report = report.sort_values(
+        ["icd10", period_col], kind="stable", ignore_index=True
+    )
+    return report.loc[
+        :,
+        [
+            "Gap, months",
+            "chapter",
+            "icd10",
+            "Female",
+            "Male",
+            period_col,
+            "Gap label",
+            "label_code",
+        ],
+    ]

evaluate_auc.py

@@ -7,7 +7,8 @@ This script follows the logic of the Delphi evaluation script supplied by the us
      at least `offset` years before the target time;
   3. run model inference by disease chunks to avoid materializing all logits;
   4. compute AUC separately by sex and age bracket;
-  5. aggregate age brackets with DeLong variance.
+  5. average age-bracket AUCs within each sex and write a Delphi2M-style
+     Female/Male report.
 
 Efficiency notes:
   - transformer/readout inference is executed once and cached;
@@ -39,6 +40,10 @@ from torch.utils.data import DataLoader, Subset
 from tqdm.auto import tqdm
 
 from dataset import HealthDataset
+from delphi2m_auc_report import (
+    DEFAULT_DELPHI2M_PERIODS_YEARS,
+    build_delphi2m_auc_report,
+)
 from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
 from models import (
     DeepHealth,
@@ -1164,30 +1169,23 @@ def evaluate_auc_pipeline(
     df_auc_unpooled["label_code"] = df_auc_unpooled["token"].map(
         dataset.label_id_to_code)
 
-    print("Using DeLong method to calculate AUC confidence intervals.")
+    print(
-    grouped = df_auc_unpooled.groupby(
+        "Building Delphi2M-style report: mean AUC across age strata, "
-        ["token", "label_code", "offset"], dropna=False, as_index=False)
+        "reported separately for Female and Male."
-    df_auc = grouped.agg(
-        auc=("auc_delong", "mean"),
-        n_strata=("auc_delong", "size"),
-        n_diseased=("n_diseased", "sum"),
-        n_healthy=("n_healthy", "sum"),
-        auc_variance_sum=("auc_variance_delong", "sum"),
     )
-    df_auc["auc_variance_delong"] = (
+    df_report = build_delphi2m_auc_report(
-        df_auc["auc_variance_sum"]
+        df_auc_unpooled,
-        / (df_auc["n_strata"].clip(lower=1).astype(np.float64) ** 2)
+        period_col="offset",
     )
-    df_auc = df_auc.drop(columns=["auc_variance_sum"])
 
     if output_path is not None:
         out_dir = Path(output_path)
         out_dir.mkdir(parents=True, exist_ok=True)
-        df_auc.to_csv(out_dir / "df_both.csv", index=False)
+        report_path = out_dir / "df_auc_delphi2m_report.csv"
-        df_auc_unpooled.to_csv(
+        df_report.to_csv(report_path, index=False)
-            out_dir / "df_auc_unpooled.csv", index=False)
+        print(f"Saved Delphi2M-style AUC report: {report_path}")
 
-    return df_auc_unpooled, df_auc
+    return df_auc_unpooled, df_report
 
 
 # ---------------------------------------------------------------------------
@@ -1237,8 +1235,18 @@ def make_auc_offsets(args: argparse.Namespace, cfg: Dict[str, Any]) -> List[floa
     if explicit_offsets is not None:
         base_offsets = explicit_offsets
     else:
-        next_token_offset = float(cfg_get(args, cfg, "offset", 0.1))
+        next_token_offset = float(
-        base_offsets = [next_token_offset, 1.0, 5.0, 10.0]
+            cfg_get(
+                args,
+                cfg,
+                "offset",
+                DEFAULT_DELPHI2M_PERIODS_YEARS[0],
+            )
+        )
+        base_offsets = [
+            next_token_offset,
+            *DEFAULT_DELPHI2M_PERIODS_YEARS[1:],
+        ]
 
     offsets: List[float] = []
     seen = set()
@@ -1286,9 +1294,9 @@ def main() -> None:
     parser.add_argument("--filter_min_total", type=int, default=None,
                         help="Minimum metadata count for disease selection; default 0.")
     parser.add_argument("--offset", type=float, default=None,
-                        help="Next-token prediction offset in years; preserved and evaluated alongside 1, 5, and 10 years by default.")
+                        help="Next-token prediction offset in years; 0.1 is Delphi2M no gap and is evaluated alongside 1, 5, and 10 years by default.")
     parser.add_argument("--offsets", type=str, default=None,
-                        help="Comma-separated prediction offsets in years. Overrides the default set of offset,1,5,10.")
+                        help="Comma-separated prediction offsets in years. Overrides the default set of 0.1,1,5,10.")
     parser.add_argument("--age_start", type=float, default=None)
     parser.add_argument("--age_stop", type=float, default=None)
     parser.add_argument("--age_step", type=float, default=None)

evaluate_auc_v2.py

@@ -3,6 +3,9 @@
 This script supports DeepHealth fixed-horizon risk scores for exponential,
 Weibull, and mixed all-future distributions.
 
+The default horizons are 0.1, 1, 5, and 10 years.  As in Delphi2M, 0.1 years
+is reported as the no-gap evaluation.
+
 Landmark querying depends on the model target mode saved in train_config.json:
     - next_token: insert a <NO_EVENT> token at landmark age and read it out;
     - all_future: pass landmark age directly as t_query.
@@ -28,6 +31,10 @@ from torch.utils.data import DataLoader, Dataset
 from tqdm.auto import tqdm
 
 from dataset import HealthDataset
+from delphi2m_auc_report import (
+    DEFAULT_DELPHI2M_PERIODS_YEARS,
+    build_delphi2m_auc_report,
+)
 from eval_data import load_sequence_eval_dataset
 from models import (
     DeepHealth,
@@ -330,44 +337,6 @@ def _first_existing_column(df: pd.DataFrame, candidates: Sequence[str]) -> Optio
     return None
 
 
-def build_metadata_for_merge(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> pd.DataFrame:
-    base_rows = []
-    for token, code in dataset.label_id_to_code.items():
-        token = int(token)
-        code_text = str(code)
-        if token in SPECIAL_TOKENS or code_text.startswith("<"):
-            continue
-        base_rows.append({"token": token, "label_code": code_text})
-    base = pd.DataFrame(base_rows)
-    if labels_meta is None or labels_meta.empty:
-        return base
-
-    meta = labels_meta.copy()
-    code_col = _first_existing_column(
-        meta, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"])
-    if code_col is not None:
-        meta["_label_code"] = meta[code_col].astype(
-            str).map(lambda s: s.split()[0].strip())
-        merged = base.merge(meta, left_on="label_code",
-                            right_on="_label_code", how="left")
-        return merged.drop(columns=["_label_code"], errors="ignore")
-
-    if "index" in meta.columns:
-        idx = pd.to_numeric(meta["index"], errors="coerce")
-        has_no_event = (
-            NO_EVENT_IDX in dataset.label_id_to_code
-            and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>"
-        )
-        if has_no_event:
-            idx = idx.where(idx < NO_EVENT_IDX, idx + 1)
-        meta["_index_int"] = idx.astype("Int64")
-        merged = base.merge(meta, left_on="token",
-                            right_on="_index_int", how="left")
-        return merged.drop(columns=["_index_int"], errors="ignore")
-
-    return base
-
-
 def _metadata_count_map(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> Dict[int, float]:
     if labels_meta is None or labels_meta.empty or "count" not in labels_meta.columns:
         return {}
@@ -1107,7 +1076,6 @@ def evaluate_landmark_auc(
     loader: DataLoader,
     landmark_dataset: LandmarkDataset,
     output_path: Path,
-    labels_meta: Optional[pd.DataFrame],
     disease_ids: Sequence[int],
     disease_chunk_size: int,
     score_mode: str,
@@ -1124,7 +1092,6 @@ def evaluate_landmark_auc(
     use_amp: bool,
     hidden_cache_dtype: str,
     logit_batch_size: int,
-    meta_info: Dict[str, Any],
 ) -> Tuple[pd.DataFrame, pd.DataFrame]:
     model.eval().to(device)
 
@@ -1241,54 +1208,21 @@ def evaluate_landmark_auc(
     df_unpooled["label_code"] = df_unpooled["token"].map(
         landmark_dataset.dataset.label_id_to_code)
 
-    for k, v in meta_info.items():
+    print(
-        df_unpooled[k] = v
+        "Building Delphi2M-style report: mean AUC across landmark-age "
-
+        "strata, reported separately for Female and Male."
-    meta_table = build_metadata_for_merge(landmark_dataset.dataset, labels_meta)
-    df_unpooled = df_unpooled.merge(
-        meta_table, on=["token", "label_code"], how="left")
-
-    grouped = df_unpooled.groupby(
-        ["token", "label_code", "horizon"], dropna=False, as_index=False)
-    df_merged = grouped.agg(
-        auc=("auc_delong", "mean"),
-        n_strata=("auc_delong", "size"),
-        n_diseased=("n_diseased", "sum"),
-        n_healthy=("n_healthy", "sum"),
-        auc_variance_sum=("auc_variance_delong", "sum"),
     )
-    df_merged["auc_variance_delong"] = (
+    df_report = build_delphi2m_auc_report(
-        df_merged["auc_variance_sum"]
+        df_unpooled,
-        / (df_merged["n_strata"].clip(lower=1).astype(np.float64) ** 2)
+        period_col="horizon",
     )
-    df_merged = df_merged.drop(columns=["auc_variance_sum"])
-
-    keep_meta = [
-        c for c in [
-            "model_ckpt_path",
-            "config_path",
-            "target_mode",
-            "model_target_mode",
-            "dist_mode",
-            "time_mode",
-            "attn_mask_mode",
-            "readout_name",
-            "landmark_query_mode",
-            "landmark_token_mode",
-            "score_mode",
-            "eval_split",
-        ]
-        if c in df_unpooled.columns
-    ]
-    for col in keep_meta:
-        df_merged[col] = meta_info[col]
 
     output_path.mkdir(parents=True, exist_ok=True)
-    df_unpooled.to_csv(
+    report_path = output_path / "df_auc_landmark_delphi2m_report.csv"
-        output_path / "df_auc_landmark_unpooled.csv", index=False)
+    df_report.to_csv(report_path, index=False)
-    df_merged.to_csv(output_path / "df_auc_landmark.csv", index=False)
+    print(f"Saved Delphi2M-style landmark AUC report: {report_path}")
 
-    return df_unpooled, df_merged
+    return df_unpooled, df_report
 
 
 def main() -> None:
@@ -1314,7 +1248,12 @@ def main() -> None:
     parser.add_argument("--landmark_start", type=float, default=None)
     parser.add_argument("--landmark_stop", type=float, default=None)
     parser.add_argument("--landmark_step", type=float, default=None)
-    parser.add_argument("--horizons", type=str, default=None)
+    parser.add_argument(
+        "--horizons",
+        type=str,
+        default=None,
+        help="Comma-separated horizons in years; defaults to 0.1,1,5,10, where 0.1 is Delphi2M no gap.",
+    )
 
     parser.add_argument("--min_cases", type=int, default=None)
     parser.add_argument("--min_history_events", type=int, default=None)
@@ -1434,8 +1373,9 @@ def main() -> None:
             "Landmark ages are empty. Check landmark_start/landmark_stop/landmark_step.")
 
     horizons = np.asarray(
-        parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [
+        parse_float_list(
-            1.0, 5.0, 10.0],
+            cfg_get(args, cfg, "horizons", "0.1,1,5,10")
+        ) or list(DEFAULT_DELPHI2M_PERIODS_YEARS),
         dtype=np.float32,
     )
     if horizons.size == 0:
@@ -1526,8 +1466,6 @@ def main() -> None:
         if model_target_mode == "next_token"
         else "direct_t_query"
     )
-    score_mode_out = f"{landmark_query_mode}_{score_mode}"
-
     num_workers_auc = int(
         cfg_get(args, cfg, "num_workers_auc", max(1, (os.cpu_count() or 2) - 1)))
     auc_task_chunk_size = int(cfg_get(args, cfg, "auc_task_chunk_size", 0))
@@ -1559,27 +1497,11 @@ def main() -> None:
     print(f"AUC workers: {num_workers_auc}")
     print(f"Output path: {output_path}")
 
-    meta_info = {
-        "score_mode": score_mode_out,
-        "eval_split": eval_split,
-        "model_ckpt_path": str(model_ckpt_path),
-        "config_path": str(config_path),
-        "target_mode": str(target_mode),
-        "model_target_mode": str(model_target_mode),
-        "dist_mode": str(dist_mode),
-        "time_mode": str(time_mode),
-        "attn_mask_mode": str(attn_mask_mode),
-        "readout_name": str(readout_name),
-        "landmark_query_mode": landmark_query_mode,
-        "landmark_token_mode": "no_event" if model_target_mode == "next_token" else "none",
-    }
-
     evaluate_landmark_auc(
         model=model,
         loader=loader,
         landmark_dataset=landmark_dataset,
         output_path=output_path,
-        labels_meta=labels_meta,
         disease_ids=disease_ids,
         disease_chunk_size=disease_chunk_size,
         score_mode=score_mode,
@@ -1596,7 +1518,6 @@ def main() -> None:
         use_amp=use_amp,
         hidden_cache_dtype=hidden_cache_dtype,
         logit_batch_size=logit_batch_size,
-        meta_info=meta_info,
     )
 
 

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_v5"
 
 
 def validate_traj_mixer_config(config: Mapping[str, object]) -> None:
@@ -29,7 +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.group_align",
+        "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.intra_gate_logits",
         "blocks.0.mlp.gate_proj",
         "blocks.0.mlp.value_proj",
         "blocks.0.mlp.output_proj",

test_traj_mixer.py

@@ -2,7 +2,7 @@ import unittest
 
 import torch
 
-from backbones import GPTBlock, TrajMixer
+from backbones import GPTBlock, TemporalAttention, TrajMixer
 from models import (
     TRAJ_MIXER_ARCHITECTURE,
     validate_traj_mixer_config,
@@ -12,6 +12,42 @@ from train_util import get_model_parameter_counts
 
 
 class TrajMixerTest(unittest.TestCase):
+    def test_zero_rbf_bias_has_live_projection_gradient(self) -> None:
+        attention = TemporalAttention(
+            n_embd=12,
+            n_head=3,
+            use_time_rope=False,
+            use_rbf_bias=True,
+        )
+        features = torch.randn(2, 4, 4, 16)
+        target = torch.randn(2, 4, 4, 3)
+
+        initial_bias = (
+            attention.time_bias_scale.tanh()
+            * attention.rbf_proj(features)
+        )
+        torch.testing.assert_close(initial_bias, torch.zeros_like(initial_bias))
+
+        loss = (initial_bias * target).sum()
+        loss.backward()
+
+        projection_grad = attention.rbf_proj.weight.grad
+        self.assertIsNotNone(projection_grad)
+        self.assertGreater(projection_grad.abs().sum().item(), 0.0)
+
+        with torch.no_grad():
+            attention.rbf_proj.weight.add_(projection_grad, alpha=-1e-3)
+        attention.zero_grad(set_to_none=True)
+        updated_bias = (
+            attention.time_bias_scale.tanh()
+            * attention.rbf_proj(features)
+        )
+        (updated_bias * target).sum().backward()
+
+        scale_grad = attention.time_bias_scale.grad
+        self.assertIsNotNone(scale_grad)
+        self.assertGreater(scale_grad.abs().item(), 0.0)
+
     def test_default_shape_parameters_and_initialization(self) -> None:
         mixer = TrajMixer(
             n_embd=120,
@@ -21,15 +57,111 @@ 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()), 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),
+            (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))
 
+    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.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()
+
+        grouped = torch.randn(2, 4, 10, 12)
+        changed = grouped.clone()
+        changed[:, :, 3, :] += torch.randn_like(changed[:, :, 3, :])
+
+        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),
+            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.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._cross_mix(grouped)
+        changed_out = mixer._cross_mix(changed)
+
+        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)
         mixer = TrajMixer(120, n_head=10, dropout=0.0)
@@ -58,11 +190,13 @@ 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_single_mixer_residual_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.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)
@@ -75,6 +209,18 @@ 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"}
+            )
+        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)
@@ -84,7 +230,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 +243,8 @@ class TrajMixerTest(unittest.TestCase):
         self.assertEqual(
             get_model_parameter_counts(mixer),
             {
-                "model_parameter_count": 15_840,
+                "model_parameter_count": 32_040,
-                "trainable_parameter_count": 15_840,
+                "trainable_parameter_count": 32_040,
             },
         )