Report AUCs in Delphi2M format

This reverts commit 06f29c0f0a.

Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md

@@ -1,389 +0,0 @@
-# Event–Trajectory Shared Reasoning Backbone
-
-> 状态：**Frozen implementation baseline**
-> 架构标识：`event_trajectory_shared_v2`
-> 固化日期：**2026-07-23**
-
-## 1. 核心定义
-
-使用一个共享的 Attention–TrajMixer 推理核心，对固定 Event Memory 进行多轮读取，并持续更新 Trajectory State。
-
-模型只实例化：
-
-```python
-self.reasoning_core = SharedEventTrajectoryCore(...)
-```
-
-禁止为不同推理轮创建独立 Transformer blocks。参数只保存一套，计算上顺序运行多轮。
-
-模型规模固定为五档：
-
-| model_size | d_model | n_trajectory | trajectory_dim | traj_hidden |
-|---|---:|---:|---:|---:|
-| nano | 120 | 6 | 20 | 24 |
-| tiny | 256 | 8 | 32 | 32 |
-| small | 512 | 8 | 64 | 32 |
-| medium | 768 | 12 | 64 | 48 |
-| huge | 1024 | 16 | 64 | 64 |
-
-默认使用 `model_size=nano`。`n_reasoning_rounds` 是独立参数，默认值为12，
-不属于模型规模预设；任意模型规模均可单独指定推理轮数。
-
-必须满足：
-
-\[
-d_{\mathrm{model}}
-=n_{\mathrm{trajectory}}d_{\mathrm{trajectory}}.
-\]
-
-## 2. 固定 Event Memory
-
-疾病事件与可选协变量首先组成事件序列：
-
-\[
-X_E\in\mathbb{R}^{B\times L\times d_{\mathrm{model}}}.
-\]
-
-事件特征由以下信息相加：
-
-```text
-disease / covariate embedding
-+ age/time encoding
-+ sex context
-```
-
-然后只编码一次：
-
-\[
-E=\operatorname{EventNorm}
-\left(\operatorname{EventProjection}(X_E)\right).
-\]
-
-进入 reasoning loop 后，\(E\) 的数值保持不变，但不执行 `detach`，梯度仍可回传到事件编码器。
-
-Key 和 Value 同样每次 forward 只投影一次：
-
-```python
-event_key_value = reasoning_core.project_event_memory(E)
-```
-
-12 轮共享并复用该结果。
-
-## 3. Trajectory State
-
-每个查询维护 `n_trajectory` 个显式 trajectory slots；nano 默认使用6个：
-
-\[
-S\in\mathbb{R}^{B\times Q\times n_{\mathrm{trajectory}}\times
-d_{\mathrm{trajectory}}}.
-\]
-
-其中：
-
-- all-future：\(Q=1\)；
-- next-token：\(Q=L\)，所有查询位置并行计算。
-
-定义可学习原型：
-
-\[
-P\in\mathbb{R}^{n_{\mathrm{trajectory}}\times d_{\mathrm{trajectory}}}.
-\]
-
-查询上下文经过投影并 reshape：
-
-\[
-C_Q
-=\operatorname{QueryProjection}(\text{query features})
-\in\mathbb{R}^{B\times Q\times n_{\mathrm{trajectory}}\times
-d_{\mathrm{trajectory}}},
-\]
-
-\[
-S^{(0)}=P+C_Q.
-\]
-
-all-future 的 query features 包含可学习 query token、查询年龄和性别；next-token 的 query features 使用当前位置的事件、时间和性别表示，以保留 token-level 预测语义。
-
-## 4. 共享 Trajectory-to-Event Attention
-
-Trajectory State 作为 Query，固定 Event Memory 作为 Key 和 Value：
-
-\[
-Q^{(r)}
-=W_Q\operatorname{LN}_{A}(S^{(r)}),
-\]
-
-\[
-K=W_KE,\qquad V=W_VE.
-\]
-
-形状为：
-
-```text
-Q: [B, query, trajectory, trajectory_dim]
-K: [B, trajectory, event, trajectory_dim]
-V: [B, trajectory, event, trajectory_dim]
-```
-
-Attention：
-
-\[
-\operatorname{score}_{b,q,h,l}
-=
-\frac{
-\left\langle Q_{b,q,h,:},K_{b,h,l,:}\right\rangle
-}{
-\sqrt{d_{\mathrm{trajectory}}}
-}.
-\]
-
-每个 trajectory slot 独立读取整段 Event Memory。Attention 不包含跨 trajectory 的完整输出投影；trajectory 之间的交互只由后续 TrajMixer 完成。
-
-外部 `padding_mask` 的语义固定为 `True = valid`。
-
-all-future 的内部 mask 必须满足：
-
-\[
-\operatorname{valid}_{b,q,l}
-=
-\operatorname{eventValid}_{b,l}
-\land
-(t_l\le t_q).
-\]
-
-next-token 还必须对相同时间戳加入位置因果约束：
-
-\[
-\operatorname{valid}_{b,q,l}
-=
-\operatorname{eventValid}_{b,l}
-\land
-\left[
-(t_l<t_q)
-\lor
-\left((t_l=t_q)\land(l\le q)\right)
-\right].
-\]
-
-这可以阻止前一个 token 直接读到同时间的后续目标 token；它只改变并行
-Attention 的可见性矩阵，不沿疾病时间轴递归。
-
-两种 mask 均不能读取未来事件或协变量。
-
-全 masked query 的 Attention readout 必须显式返回零，不能产生 NaN。
-
-## 5. 时间信息
-
-所有模式都在 Event Memory 和 query context 中加入 age/time encoding。
-
-当 `time_mode=relative` 时，shared cross-attention 额外使用：
-
-- query-time 对 event-time 的 Cross-TimeRoPE；
-- query–event 时间差的 Gaussian RBF bias。
-
-对应缓存形状为：
-
-\[
-\text{RBF cache}\in\mathbb{R}^{B\times Q\times L\times n_{\mathrm{rbf}}}.
-\]
-
-## 6. 共享 TrajMixer
-
-TrajMixer 输入：
-
-\[
-U\in\mathbb{R}^{B\times Q\times H\times D_h}.
-\]
-
-它只沿 trajectory 轴交互，不使用普通全维度 FFN，也不包含旧版 Group Alignment。
-
-对每个内部坐标 \(r\)：
-
-\[
-G_r=U_rW_g^{(r)},\qquad
-V_r=U_rW_v^{(r)},
-\]
-
-\[
-M_r
-=
-\operatorname{SiLU}(G_r)\odot V_r,
-\]
-
-\[
-Y_r=M_rW_o^{(r)}.
-\]
-
-参数形状：
-
-```text
-W_g: [trajectory_dim, n_trajectory, traj_hidden]
-W_v: [trajectory_dim, n_trajectory, traj_hidden]
-W_o: [trajectory_dim, traj_hidden, n_trajectory]
-```
-
-默认：
-
-```text
-n_trajectory -> 4 * n_trajectory -> n_trajectory
-```
-
-其中：
-
-\[
-d_{\mathrm{trajHidden}}=4n_{\mathrm{trajectory}}.
-\]
-
-## 7. 单轮共享核心
-
-单轮计算：
-
-\[
-R^{(r)}
-=
-A_\theta\left(
-\operatorname{LN}_A(S^{(r)}),E
-\right),
-\]
-
-\[
-U^{(r)}
-=
-S^{(r)}
-+\alpha_A\operatorname{Dropout}(R^{(r)}),
-\]
-
-\[
-S^{(r+1)}
-=
-U^{(r)}
-+\alpha_M\operatorname{Dropout}
-\left(
-M_\phi(\operatorname{LN}_M(U^{(r)}))
-\right).
-\]
-
-残差统一使用加法。
-
-## 8. 多轮参数共享
-
-同一个核心重复运行：
-
-```python
-for _ in range(n_reasoning_rounds):
-    S = self.reasoning_core(...)
-```
-
-所有轮次共享：
-
-```text
-q_proj / k_proj / v_proj
-relative-time projection
-TrajMixer parameters
-LayerNorm parameters
-attn_scale / mixer_scale
-```
-
-因此推理轮数不改变模型参数量：
-
-\[
-A^{(1)}=\cdots=A^{(12)}=A_\theta,
-\]
-
-\[
-M^{(1)}=\cdots=M^{(12)}=M_\phi.
-\]
-
-但每轮 state 不同，因此 Query 与 Attention weights 也不同。
-
-## 9. 稳定性设计
-
-共享残差缩放初始化为：
-
-\[
-\alpha_A=\alpha_M
-=
-\frac{1}{\sqrt{n_{\mathrm{reasoningRounds}}}}.
-\]
-
-两个标量可学习，并由全部轮次共享。
-
-第一版不加入：
-
-```text
-round-specific parameters
-round embedding
-每轮独立 LayerNorm
-每轮独立 residual scale
-GRU 或其他时间递归
-```
-
-## 10. 输出接口
-
-推理结束后按固定顺序 flatten trajectory slots：
-
-\[
-H
-=
-\operatorname{FinalNorm}
-\left(
-\operatorname{Flatten}(S^{(R)})
-\right).
-\]
-
-- all-future 输出：`[B, d_model]`；
-- next-token 输出：`[B, L, d_model]`；
-- next-token 的 risk-head weight tying 保持不变；
-- Weibull 与 mixed heads 继续使用同一最终 hidden。
-
-next-token 的 query 位置全部并行，只有 reasoning rounds 顺序执行，因此不存在沿疾病时间轴的状态递归。
-
-## 11. 配置与 checkpoint 约束
-
-训练配置必须写入：
-
-```yaml
-model_architecture: event_trajectory_shared_v2
-model_size: nano
-d_model: 120
-n_trajectory: 6
-trajectory_dim: 20
-traj_hidden: 24
-n_reasoning_rounds: 12
-model_parameter_count: <runtime count>
-trainable_parameter_count: <runtime count>
-```
-
-评估和导出入口必须同时验证：
-
-1. `model_architecture` 完全匹配；
-2. `model_size` 属于 `nano / tiny / small / medium / huge`；
-3. `d_model`、`n_trajectory`、`trajectory_dim` 和 `traj_hidden`
-   与对应规模预设完全匹配；
-4. checkpoint 包含一套且仅一套 `reasoning_core` 关键参数；
-5. checkpoint 内持久化的 `d_model`、`n_trajectory` 和
-   `n_reasoning_rounds` 架构指纹与训练配置完全一致；
-6. 不接受旧 `traj_mixer_v2` checkpoint。
-
-其中 `n_reasoning_rounds` 必须进入 checkpoint 架构指纹，因为改变轮数
-不会改变参数 shape，不能仅依赖 `load_state_dict(strict=True)` 检出错配。
-
-## 12. 信息流
-
-```text
-E ─────────────┬──────────────┬──────────────┬──────────────┐
-               │              │              │              │
-               ▼              ▼              ▼              ▼
-S0 -> Shared Core -> S1 -> Shared Core -> S2 -> ... -> Shared Core -> S12
-       同一套参数          同一套参数                   同一套参数
-```
-
-整体定义：
-
-\[
-\boxed{
-\text{一个共享 Event–Trajectory 推理核心}
-\times
-\text{多轮状态依赖推理}
-}
-\]

TrajMixer_设计方案.md

@@ -0,0 +1,385 @@
+# TrajMixer Block 最终设计方案
+
+> 状态：**Frozen implementation baseline**
+>
+> 版本：**v3.0 / traj_mixer_v5**
+>
+> 固化日期：**2026-07-24**
+
+本文档是当前 TrajMixer 的实现与实验基线。本版本采用单 PreNorm、单外层 residual、静态门控组内融合和跨 group SwiGLU。
+
+## 1. 目标
+
+在不改变 Delphi Transformer Attention 的前提下，用轻量、完全并行的 TrajMixer 替换 FFN。
+
+保持不变：
+
+- causal mask；
+- TimeRoPE；
+- Relative Time Attention Bias；
+- Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)；
+- 序列建模和训练目标。
+
+TrajMixer 不沿序列维度混合，也不引入时间递归。
+
+## 2. Block 结构
+
+```text
+PreNorm Causal Multi-Head Attention
+→ Attention Residual
+→ Full-width TrajMixer PreNorm
+→ reshape [B, L, n_group, d_group]
+→ Per-Group SwiGLU: d_group → 4d_group → d_group
+→ Static Gated Fusion
+→ Cross-Group SwiGLU: n_group → 4n_group → n_group
+→ reshape [B, L, n_embd]
+→ Dropout
+→ One TrajMixer Residual
+```
+
+Attention 阶段：
+
+\[
+X
+=X^{(l)}
++\operatorname{CausalMHA}
+\left(\operatorname{LN}_{\mathrm{attn}}(X^{(l)})\right).
+\]
+
+TrajMixer 阶段：
+
+\[
+N=\operatorname{LN}_{d}(X),
+\]
+
+\[
+G=\operatorname{reshape}(N)
+\in\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}},
+\]
+
+\[
+P=\operatorname{IntraMixer}(G),
+\]
+
+\[
+U=G+\sigma(\Theta)\odot P,
+\]
+
+\[
+\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：
+
+\[
+X\in\mathbb{R}^{B\times L\times d}.
+\]
+
+定义：
+
+\[
+n_{\mathrm{group}}:=n_{\mathrm{head}},
+\qquad
+d_{\mathrm{group}}=\frac{d}{n_{\mathrm{group}}},
+\qquad
+d=n_{\mathrm{group}}d_{\mathrm{group}}.
+\]
+
+默认：
+
+\[
+d=120,\qquad
+n_{\mathrm{group}}=10,\qquad
+d_{\mathrm{group}}=12.
+\]
+
+这些 group 是 residual space 的连续分区，不等同于 Attention heads；二者只共享数量。
+
+## 4. 唯一的 Full-Width PreNorm
+
+TrajMixer 只使用一个：
+
+```text
+norm: LayerNorm(n_embd)
+```
+
+LayerNorm 作用于完整 \(d\) 维 residual representation，然后才 reshape：
+
+\[
+G=\operatorname{reshape}
+\left(\operatorname{LN}_{d}(X)\right).
+\]
+
+本版本明确删除：
+
+```text
+intra_norm
+cross_norm
+group_align
+```
+
+不得在组内或跨组阶段再增加额外 LayerNorm。
+
+## 5. 组内 SwiGLU
+
+每个 group 使用独立参数，对其 \(d_{\mathrm{group}}\) 维内部特征执行：
+
+\[
+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,
+\]
+
+因此：
+
+\[
+\Gamma_{g,r}\approx0.1.
+\]
+
+融合：
+
+\[
+U=G+\Gamma\odot P.
+\]
+
+\(\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}^{n_{\mathrm{group}}\times4n_{\mathrm{group}}},
+\]
+
+\[
+A_o^{(r)}
+\in
+\mathbb{R}^{4n_{\mathrm{group}}\times n_{\mathrm{group}}}.
+\]
+
+计算：
+
+\[
+Q_{:,r}
+=
+\operatorname{SiLU}\left(U_{:,r}A_g^{(r)}\right)
+\odot
+\left(U_{:,r}A_v^{(r)}\right),
+\]
+
+\[
+\Delta_{:,r}=Q_{:,r}A_o^{(r)}.
+\]
+
+实现形状：
+
+```text
+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 参数，且不沿序列维度交互。
+
+## 8. 唯一的外层 Residual
+
+跨 group 输出 reshape 回：
+
+\[
+\Delta\in\mathbb{R}^{B\times L\times d}.
+\]
+
+最终：
+
+\[
+\operatorname{TrajMixer}(X)
+=X+\operatorname{Dropout}(\Delta).
+\]
+
+固定约束：
+
+- 组内阶段后不执行独立 residual；
+- 跨组阶段后不执行独立 residual；
+- `GPTBlock` 不再额外执行 `X + TrajMixer(X)`；
+- 整个 TrajMixer 只有一次主 residual。
+
+## 9. 初始化
+
+固定初始化：
+
+- `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}\) 的正态分布；
+- Full-width LayerNorm：PyTorch 默认 affine 初始化；
+- Dropout：沿用 `mlp_dropout`。
+
+组内输出使用正常 Xavier 初始化以保证其具有完整表达能力；静态门控将其初始贡献限制在约 0.1。最终跨 group 输出投影保持小值初始化，使整个 TrajMixer residual update 在训练初期接近零。
+
+Relative Time Attention Bias 初始化固定为：
+
+- `rbf_proj.weight`：零初始化；
+- `time_bias_scale`：初始化为 \(1.0\)；
+- 初始 RBF attention bias 严格为零；
+- `rbf_proj.weight` 从第一个优化步骤即可获得梯度。
+
+## 10. 参数量
+
+默认 \(d=120\)、\(n_{\mathrm{group}}=10\)、\(d_{\mathrm{group}}=12\)。
+
+Full-width LayerNorm：
+
+\[
+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

@@ -29,14 +29,6 @@ class TimeRoPE(nn.Module):
         x2 = x[..., 1::2]
         return torch.stack((-x2, x1), dim=-1).flatten(-2)
 
-    @staticmethod
-    def apply_single_from_cache(
-        x: torch.Tensor,
-        rope_cache: tuple[torch.Tensor, torch.Tensor],
-    ) -> torch.Tensor:
-        cos, sin = rope_cache
-        return x * cos + TimeRoPE._rotate_half(x) * sin
-
     @staticmethod
     def apply_from_cache(
         q: torch.Tensor,
@@ -93,280 +85,271 @@ class GaussianRBFTimeBasis(nn.Module):
         )
         return rbf_acts
 
-    def precompute_cross_cache(
-        self,
-        query_tau: torch.Tensor,
-        key_tau: torch.Tensor,
-    ) -> torch.Tensor:
-        """Return RBF activations for query-time minus event-time."""
-        diff = query_tau.float().unsqueeze(2) - key_tau.float().unsqueeze(1)
-        widths = self.log_widths.exp()
-        return torch.exp(
-            -0.5
-            * (
-                (diff.unsqueeze(-1) - self.centers)
-                / widths
-            ).square()
-        )
-
-
-class TrajectoryCrossAttention(nn.Module):
-    """Shared trajectory-slot queries reading a fixed event memory."""
 
+class TemporalAttention(nn.Module):
     def __init__(
         self,
-        d_model: int,
+        n_embd: int,
-        n_trajectory: int,
+        n_head: int,
         n_rbf_bases: int = 16,
-        use_time_rope: bool = False,
+        dropout: float = 0.0,
-        use_rbf_bias: bool = False,
+        use_time_rope: bool = True,
+        use_rbf_bias: bool = True,
     ):
         super().__init__()
-        if d_model <= 0 or n_trajectory <= 0:
+        assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
-            raise ValueError("d_model and n_trajectory must be positive")
+        self.n_head = n_head
-        if d_model % n_trajectory != 0:
+        self.d_head = n_embd // n_head
-            raise ValueError(
+        self.scale = 1.0 / math.sqrt(self.d_head)
-                "d_model must be divisible by n_trajectory, got "
-                f"{d_model} and {n_trajectory}"
-            )
-        self.d_model = d_model
-        self.n_trajectory = n_trajectory
-        self.trajectory_dim = d_model // n_trajectory
-        self.scale = self.trajectory_dim ** -0.5
         self.use_time_rope = use_time_rope
         self.use_rbf_bias = use_rbf_bias
 
-        # q_proj acts on each slot independently and is shared across slots.
+        # QKV projection (fused for efficiency)
-        self.q_proj = nn.Linear(
+        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
-            self.trajectory_dim,
+        # Output projection
-            self.trajectory_dim,
+        self.out_proj = nn.Linear(n_embd, n_embd, bias=False)
-            bias=False,
+
-        )
+        # Layer-specific projection from shared RBF basis activations to per-head attention bias.
-        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
-        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+        # Keep the initial RBF attention bias exactly zero through the
-        if use_rbf_bias:
+        # zero-initialized projection, while leaving that projection with a
-            self.rbf_proj = nn.Linear(
+        # live gradient from the first optimization step.
-                n_rbf_bases,
+        self.time_bias_scale = nn.Parameter(torch.tensor(1.0))
-                n_trajectory,
+
-                bias=False,
+        self.resid_drop = nn.Dropout(dropout)
-            )
-            self.time_bias_scale = nn.Parameter(torch.tensor(0.0))
-        else:
-            self.rbf_proj = None
-            self.register_parameter("time_bias_scale", None)
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
-        nn.init.normal_(self.q_proj.weight, mean=0.0, std=0.02)
+        """Match the previous version's GPT-style weight initialization."""
-        nn.init.normal_(self.k_proj.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.qkv.weight, mean=0.0, std=0.02)
-        nn.init.normal_(self.v_proj.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
-        if self.rbf_proj is not None:
+        nn.init.zeros_(self.rbf_proj.weight)
-            # The scalar gate starts at zero, so the relative-time bias still
-            # starts disabled. A nonzero projection is necessary for the gate
-            # itself to receive a gradient on the first optimization step.
-            nn.init.xavier_uniform_(self.rbf_proj.weight)
-
-    def project_event_memory(
-        self,
-        event_memory: torch.Tensor,
-        event_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Project K/V once for reuse by every reasoning round."""
-        batch_size, memory_len, _ = event_memory.shape
-        key = self.k_proj(event_memory).reshape(
-            batch_size,
-            memory_len,
-            self.n_trajectory,
-            self.trajectory_dim,
-        ).transpose(1, 2)
-        value = self.v_proj(event_memory).reshape(
-            batch_size,
-            memory_len,
-            self.n_trajectory,
-            self.trajectory_dim,
-        ).transpose(1, 2)
-        if self.use_time_rope:
-            if event_rope_cache is None:
-                raise ValueError(
-                    "event_rope_cache is required when TimeRoPE is enabled"
-                )
-            key = TimeRoPE.apply_single_from_cache(key, event_rope_cache)
-        return key, value
 
     def forward(
         self,
-        trajectory_state: torch.Tensor,
+        x: torch.Tensor,
-        event_key_value: tuple[torch.Tensor, torch.Tensor],
+        rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
-        event_invalid_mask: torch.Tensor,
-        query_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
         rbf_cache: torch.Tensor | None = None,
+        attn_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """Read memory for states shaped ``(B, Q, H, Dh)``."""
-        if trajectory_state.ndim != 4:
-            raise ValueError(
-                "trajectory_state must have shape (B, Q, H, Dh), got "
-                f"{tuple(trajectory_state.shape)}"
-            )
-        batch_size, n_query, n_trajectory, trajectory_dim = (
-            trajectory_state.shape
-        )
-        if (n_trajectory, trajectory_dim) != (
-            self.n_trajectory,
-            self.trajectory_dim,
-        ):
-            raise ValueError(
-                "Unexpected trajectory shape: "
-                f"{(n_trajectory, trajectory_dim)}"
-            )
-        key, value = event_key_value
-        memory_len = key.size(2)
-        if event_invalid_mask.shape != (batch_size, n_query, memory_len):
-            raise ValueError(
-                "event_invalid_mask must have shape "
-                f"{(batch_size, n_query, memory_len)}, got "
-                f"{tuple(event_invalid_mask.shape)}"
-            )
-
-        query = self.q_proj(trajectory_state).transpose(1, 2)
         if self.use_time_rope:
-            if query_rope_cache is None:
+            assert rope_cache is not None, "rope_cache must be provided when use_time_rope is True"
-                raise ValueError(
-                    "query_rope_cache is required when TimeRoPE is enabled"
-                )
-            query = TimeRoPE.apply_single_from_cache(query, query_rope_cache)
-        query = query.transpose(1, 2)
-
-        scores = torch.einsum("bqhd,bhld->bqhl", query, key) * self.scale
         if self.use_rbf_bias:
-            if rbf_cache is None or self.rbf_proj is None:
+            assert rbf_cache is not None, "rbf_cache must be provided when use_rbf_bias is True"
-                raise ValueError(
-                    "rbf_cache is required when relative time bias is enabled"
-                )
-            time_bias = self.rbf_proj(rbf_cache).permute(0, 1, 3, 2)
-            scores = scores + self.time_bias_scale.tanh() * time_bias
 
-        mask = event_invalid_mask.unsqueeze(2)
+        B, L, _ = x.shape
-        min_value = torch.finfo(scores.dtype).min
+        H, D = self.n_head, self.d_head
-        masked_scores = scores.masked_fill(mask, min_value)
+
-        weights = torch.softmax(masked_scores.float(), dim=-1).to(scores.dtype)
+        # --- QKV ----------------------------------------------------------
-        weights = weights.masked_fill(mask, 0.0)
+        qkv = self.qkv(x).reshape(B, L, 3, H, D).permute(2, 0, 3, 1, 4)
-        denominator = weights.sum(dim=-1, keepdim=True)
+        q, k, v = qkv.unbind(0)  # each (B, H, L, D)
-        weights = weights / denominator.clamp_min(
+
-            torch.finfo(weights.dtype).eps
+        # --- Apply RoPE (from shared cache) --------------------------------
+        if self.use_time_rope:
+            q, k = TimeRoPE.apply_from_cache(q, k, rope_cache)
+
+        # Build additive attention bias mask: time bias + causal/padding mask.
+        time_bias = None
+        if self.use_rbf_bias:
+            time_bias = self.rbf_proj(rbf_cache).permute(
+                0, 3, 1, 2)  # (B, H, L, L)
+            time_bias = self.time_bias_scale.tanh() * time_bias
+
+        if time_bias is not None and attn_mask is not None:
+            attn_bias = time_bias + attn_mask.to(time_bias.dtype)
+        elif time_bias is not None:
+            attn_bias = time_bias
+        elif attn_mask is not None:
+            attn_bias = attn_mask
+        else:
+            attn_bias = None
+
+        out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_bias,
+            dropout_p=0.0,
+            is_causal=False,
+            scale=self.scale,
         )
-        return torch.einsum("bqhl,bhld->bqhd", weights, value)
+
+        # --- Aggregate & project out --------------------------------------
+        out = out.transpose(1, 2).reshape(B, L, H * D)
+        return self.resid_drop(self.out_proj(out))
 
 
-class SharedTrajectoryMixer(nn.Module):
+class TrajMixer(nn.Module):
-    """SwiGLU interaction along the trajectory axis only."""
+    """PreNorm gated mixing within and across latent trajectory groups.
 
-    def __init__(self, n_trajectory: int, trajectory_dim: int):
+    The groups are contiguous partitions of the post-``W_O`` residual
+    representation. They are deliberately not treated as attention heads.
+    All operations are position-wise, so the sequence dimension remains fully
+    parallel and no temporal information can leak between positions here.
+    """
+
+    def __init__(
+        self,
+        n_embd: int,
+        n_head: int = 10,
+        dropout: float = 0.0,
+    ):
         super().__init__()
-        if n_trajectory <= 0 or trajectory_dim <= 0:
+        if n_embd <= 0:
-            raise ValueError("trajectory dimensions must be positive")
+            raise ValueError(f"n_embd must be > 0, got {n_embd}")
-        self.n_trajectory = n_trajectory
+        if n_head <= 0:
-        self.trajectory_dim = trajectory_dim
+            raise ValueError(f"n_head must be > 0, got {n_head}")
-        self.traj_hidden = 4 * n_trajectory
+        if n_embd % n_head != 0:
+            raise ValueError(
+                f"n_embd must be divisible by n_head, got {n_embd} and {n_head}"
+            )
+        self.n_embd = n_embd
+        # The residual-group count is tied to n_head, but the resulting groups
+        # 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
+
+        # 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_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
+        # independent, exactly as specified by the TrajMixer baseline.
         self.gate_proj = nn.Parameter(
-            torch.empty(trajectory_dim, n_trajectory, self.traj_hidden)
+            torch.empty(self.d_group, self.n_group, self.hidden_group)
         )
         self.value_proj = nn.Parameter(
-            torch.empty(trajectory_dim, n_trajectory, self.traj_hidden)
+            torch.empty(self.d_group, self.n_group, self.hidden_group)
         )
         self.output_proj = nn.Parameter(
-            torch.empty(trajectory_dim, self.traj_hidden, n_trajectory)
+            torch.empty(self.d_group, self.hidden_group, self.n_group)
         )
+        self.drop = nn.Dropout(dropout)
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
-        for feature_idx in range(self.trajectory_dim):
+        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.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.
+        for feature_idx in range(self.d_group):
             nn.init.xavier_uniform_(self.gate_proj[feature_idx])
             nn.init.xavier_uniform_(self.value_proj[feature_idx])
         nn.init.normal_(self.output_proj, mean=0.0, std=1e-3)
 
-    def forward(self, state: torch.Tensor) -> torch.Tensor:
+    def _intra_mix(self, grouped: torch.Tensor) -> torch.Tensor:
-        if state.shape[-2:] != (self.n_trajectory, self.trajectory_dim):
+        """Mix features independently inside each residual-space group."""
-            raise ValueError(
+        intra_gate = torch.einsum(
-                "Expected trailing trajectory shape "
+            "blgd,gdh->blgh", grouped, self.intra_gate_proj
-                f"{(self.n_trajectory, self.trajectory_dim)}, got "
+        )
-                f"{tuple(state.shape[-2:])}"
+        intra_value = torch.einsum(
-            )
+            "blgd,gdh->blgh", grouped, self.intra_value_proj
-        gate = torch.einsum("...hr,rhk->...kr", state, self.gate_proj)
+        )
-        value = torch.einsum("...hr,rhk->...kr", state, self.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("...kr,rkh->...hr", hidden, self.output_proj)
+        return torch.einsum(
+            "blhr,rhg->blgr", hidden, self.output_proj
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """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:
+            raise ValueError(
+                f"Expected hidden size {self.n_embd}, got {x.size(-1)}"
+            )
+
+        batch_size, seq_len, _ = x.shape
+        grouped = self.norm(x).reshape(
+            batch_size, seq_len, self.n_group, self.d_group
+        )
+
+        # 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
+        )
+        mixed_input = grouped + intra_gate * intra_output
+
+        # 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
+        )
+        return x + self.drop(update)
 
 
-class SharedEventTrajectoryCore(nn.Module):
+class GPTBlock(nn.Module):
-    """One parameter-shared reasoning core reused across all rounds."""
-
     def __init__(
         self,
-        d_model: int,
+        n_embd: int,
-        n_trajectory: int,
+        n_head: int,
-        n_reasoning_rounds: int,
+
-        dropout: float = 0.0,
+        attn_dropout: float = 0.0,
-        n_rbf_bases: int = 16,
+        mlp_dropout: float = 0.0,
         use_time_rope: bool = False,
         use_rbf_bias: bool = False,
+        n_rbf_bases: int = 16,
     ):
         super().__init__()
-        if n_reasoning_rounds <= 0:
+        self.attn = TemporalAttention(
-            raise ValueError("n_reasoning_rounds must be positive")
+            n_embd=n_embd,
-        if d_model <= 0 or n_trajectory <= 0:
+            n_head=n_head,
-            raise ValueError("d_model and n_trajectory must be positive")
-        if d_model % n_trajectory != 0:
-            raise ValueError("d_model must equal n_trajectory * trajectory_dim")
-        trajectory_dim = d_model // n_trajectory
-        self.norm_attn = nn.LayerNorm(trajectory_dim)
-        self.cross_attention = TrajectoryCrossAttention(
-            d_model=d_model,
-            n_trajectory=n_trajectory,
             n_rbf_bases=n_rbf_bases,
+            dropout=attn_dropout,
             use_time_rope=use_time_rope,
             use_rbf_bias=use_rbf_bias,
         )
-        self.norm_mixer = nn.LayerNorm(trajectory_dim)
+        self.mlp = TrajMixer(
-        self.traj_mixer = SharedTrajectoryMixer(
+            n_embd=n_embd,
-            n_trajectory=n_trajectory,
+            n_head=n_head,
-            trajectory_dim=trajectory_dim,
+            dropout=mlp_dropout,
-        )
-        initial_scale = 1.0 / math.sqrt(n_reasoning_rounds)
-        self.attn_scale = nn.Parameter(torch.tensor(initial_scale))
-        self.mixer_scale = nn.Parameter(torch.tensor(initial_scale))
-        self.dropout = nn.Dropout(dropout)
-
-    def project_event_memory(
-        self,
-        event_memory: torch.Tensor,
-        event_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        return self.cross_attention.project_event_memory(
-            event_memory,
-            event_rope_cache=event_rope_cache,
         )
+        self.ln1 = nn.LayerNorm(n_embd)
 
     def forward(
         self,
-        trajectory_state: torch.Tensor,
+        x: torch.Tensor,
-        event_key_value: tuple[torch.Tensor, torch.Tensor],
+        rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
-        event_invalid_mask: torch.Tensor,
-        query_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
         rbf_cache: torch.Tensor | None = None,
+        attn_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        readout = self.cross_attention(
+        x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
-            trajectory_state=self.norm_attn(trajectory_state),
+        return self.mlp(x)
-            event_key_value=event_key_value,
-            event_invalid_mask=event_invalid_mask,
-            query_rope_cache=query_rope_cache,
-            rbf_cache=rbf_cache,
-        )
-        updated = (
-            trajectory_state
-            + self.attn_scale * self.dropout(readout)
-        )
-        mixed = self.traj_mixer(self.norm_mixer(updated))
-        return updated + self.mixer_scale * self.dropout(mixed)
 
 
 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,11 +40,15 @@ 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,
-    validate_event_trajectory_config,
+    validate_traj_mixer_config,
-    validate_event_trajectory_state_dict,
+    validate_traj_mixer_state_dict,
 )
 from readouts import build_readout
 from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
@@ -313,7 +318,7 @@ def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: floa
 
 
 def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
-    validate_event_trajectory_config(cfg)
+    validate_traj_mixer_config(cfg)
     model_target_mode = str(cfg_get(
         args, cfg, "model_target_mode", "next_token")).lower()
     if model_target_mode not in {"next_token", "all_future"}:
@@ -322,10 +327,10 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
         )
     return DeepHealth(
         vocab_size=dataset.vocab_size,
-        model_size=str(cfg_get(args, cfg, "model_size", "nano")),
+        n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
-        n_reasoning_rounds=int(
+        n_head=int(cfg_get(args, cfg, "n_head", 10)),
-            cfg_get(args, cfg, "n_reasoning_rounds", 12)
+        n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
-        ),
+        n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -391,12 +396,7 @@ def load_model_state(
     state = state_dict if state_dict is not None else load_checkpoint_state_dict(
         checkpoint_path, map_location=device)
 
-    validate_event_trajectory_state_dict(
+    validate_traj_mixer_state_dict(state)
-        state,
-        expected_d_model=model.d_model,
-        expected_n_trajectory=model.n_trajectory,
-        expected_n_reasoning_rounds=model.n_reasoning_rounds,
-    )
     model.load_state_dict(state, strict=True)
 
 
@@ -533,7 +533,7 @@ def infer_readout_hidden(
                 hidden = torch.zeros(
                     batch_size,
                     seq_len,
-                    model.d_model,
+                    model.n_embd,
                     device=event_seq.device,
                     dtype=torch.float32,
                 )
@@ -1169,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
 
 
 # ---------------------------------------------------------------------------
@@ -1242,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()
@@ -1291,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,11 +31,15 @@ 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,
-    validate_event_trajectory_config,
+    validate_traj_mixer_config,
-    validate_event_trajectory_state_dict,
+    validate_traj_mixer_state_dict,
 )
 from readouts import build_readout
 from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
@@ -182,7 +189,7 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
 
 
 def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
-    validate_event_trajectory_config(cfg)
+    validate_traj_mixer_config(cfg)
     model_target_mode = str(cfg_get(
         args, cfg, "model_target_mode", "next_token")).lower()
     if model_target_mode not in {"next_token", "all_future"}:
@@ -191,10 +198,10 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
         )
     return DeepHealth(
         vocab_size=dataset.vocab_size,
-        model_size=str(cfg_get(args, cfg, "model_size", "nano")),
+        n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
-        n_reasoning_rounds=int(
+        n_head=int(cfg_get(args, cfg, "n_head", 10)),
-            cfg_get(args, cfg, "n_reasoning_rounds", 12)
+        n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
-        ),
+        n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -209,12 +216,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
 
 
 def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None:
-    validate_event_trajectory_state_dict(
+    validate_traj_mixer_state_dict(state_dict)
-        state_dict,
-        expected_d_model=model.d_model,
-        expected_n_trajectory=model.n_trajectory,
-        expected_n_reasoning_rounds=model.n_reasoning_rounds,
-    )
     model.load_state_dict(state_dict, strict=True)
 
 
@@ -335,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 {}
@@ -1112,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,
@@ -1129,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)
 
@@ -1246,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:
@@ -1319,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)
@@ -1439,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:
@@ -1531,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))
@@ -1564,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,
@@ -1601,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,
     )
 
 

export_tquery_logits_hidden.py

@@ -205,7 +205,7 @@ def main() -> None:
 
     n_rows = len(landmark_dataset)
     vocab_size = int(dataset.vocab_size)
-    hidden_dim = int(model.d_model)
+    hidden_dim = int(getattr(model, "n_embd", cfg_get(args, cfg_model, "n_embd", 120)))
     logits_dtype = numpy_float_dtype(args.logits_dtype)
     hidden_dtype = numpy_float_dtype(args.hidden_dtype)
 

models.py

@@ -7,198 +7,44 @@ import torch.nn.functional as F
 
 from backbones import (
     AgeSinusoidalEncoding,
+    GPTBlock,
     GaussianRBFTimeBasis,
-    SharedEventTrajectoryCore,
     TimeRoPE,
     TokenAutoDiscretization,
 )
 from targets import PAD_IDX
 
 
-EVENT_TRAJECTORY_ARCHITECTURE = "event_trajectory_shared_v2"
+TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v5"
 
 
-@dataclass(frozen=True)
+def validate_traj_mixer_config(config: Mapping[str, object]) -> None:
-class EventTrajectoryModelSize:
-    d_model: int
-    n_trajectory: int
-
-    @property
-    def trajectory_dim(self) -> int:
-        return self.d_model // self.n_trajectory
-
-    @property
-    def traj_hidden(self) -> int:
-        return 4 * self.n_trajectory
-
-
-MODEL_SIZE_PRESETS = {
-    "nano": EventTrajectoryModelSize(d_model=120, n_trajectory=6),
-    "tiny": EventTrajectoryModelSize(d_model=256, n_trajectory=8),
-    "small": EventTrajectoryModelSize(d_model=512, n_trajectory=8),
-    "medium": EventTrajectoryModelSize(d_model=768, n_trajectory=12),
-    "huge": EventTrajectoryModelSize(d_model=1024, n_trajectory=16),
-}
-MODEL_SIZE_NAMES = tuple(MODEL_SIZE_PRESETS)
-
-
-def resolve_model_size(model_size: str) -> EventTrajectoryModelSize:
-    if not isinstance(model_size, str):
-        raise ValueError(
-            f"model_size must be a string, got {type(model_size).__name__}"
-        )
-    normalized = model_size.strip().lower()
-    try:
-        return MODEL_SIZE_PRESETS[normalized]
-    except KeyError as exc:
-        choices = ", ".join(MODEL_SIZE_NAMES)
-        raise ValueError(
-            f"Unknown model_size {model_size!r}; expected one of: {choices}"
-        ) from exc
-
-
-def _required_config_int(
-    config: Mapping[str, object],
-    key: str,
-) -> int:
-    raw_value = config.get(key)
-    if isinstance(raw_value, bool):
-        raise ValueError(f"Config field {key!r} must be an integer")
-    try:
-        value = int(raw_value)
-    except (TypeError, ValueError) as exc:
-        raise ValueError(
-            f"Config field {key!r} must be present and integer-valued; "
-            f"got {raw_value!r}"
-        ) from exc
-    if isinstance(raw_value, float) and not raw_value.is_integer():
-        raise ValueError(f"Config field {key!r} must be an integer")
-    return value
-
-
-def validate_event_trajectory_config(config: Mapping[str, object]) -> None:
     actual = config.get("model_architecture")
-    if actual != EVENT_TRAJECTORY_ARCHITECTURE:
+    if actual != TRAJ_MIXER_ARCHITECTURE:
         raise ValueError(
-            "This branch only accepts models trained with the shared "
+            "This branch only accepts models trained with the TrajMixer "
-            "event-trajectory architecture marker "
+            f"architecture marker {TRAJ_MIXER_ARCHITECTURE!r}; got {actual!r}."
-            f"{EVENT_TRAJECTORY_ARCHITECTURE!r}; got {actual!r}."
-        )
-    raw_model_size = config.get("model_size")
-    if not isinstance(raw_model_size, str):
-        raise ValueError(
-            "Config field 'model_size' must be one of: "
-            + ", ".join(MODEL_SIZE_NAMES)
-        )
-    model_size = raw_model_size.strip().lower()
-    preset = resolve_model_size(model_size)
-    d_model = _required_config_int(config, "d_model")
-    n_trajectory = _required_config_int(config, "n_trajectory")
-    n_reasoning_rounds = _required_config_int(
-        config,
-        "n_reasoning_rounds",
-    )
-    trajectory_dim = _required_config_int(config, "trajectory_dim")
-    traj_hidden = _required_config_int(config, "traj_hidden")
-    if n_reasoning_rounds <= 0:
-        raise ValueError(
-            "n_reasoning_rounds must be positive"
-        )
-    expected_values = {
-        "d_model": preset.d_model,
-        "n_trajectory": preset.n_trajectory,
-        "trajectory_dim": preset.trajectory_dim,
-        "traj_hidden": preset.traj_hidden,
-    }
-    actual_values = {
-        "d_model": d_model,
-        "n_trajectory": n_trajectory,
-        "trajectory_dim": trajectory_dim,
-        "traj_hidden": traj_hidden,
-    }
-    mismatches = [
-        f"{key}: expected {expected}, got {actual_values[key]}"
-        for key, expected in expected_values.items()
-        if actual_values[key] != expected
-    ]
-    if mismatches:
-        raise ValueError(
-            f"Config does not match model_size={model_size!r}: "
-            + "; ".join(mismatches)
         )
 
 
-def _checkpoint_scalar_int(
+def validate_traj_mixer_state_dict(state_dict: Mapping[str, object]) -> None:
-    state_dict: Mapping[str, object],
-    key: str,
-) -> int:
-    value = state_dict[key]
-    if not isinstance(value, torch.Tensor) or value.numel() != 1:
-        raise ValueError(
-            f"Checkpoint architecture field {key!r} must be a scalar tensor"
-        )
-    return int(value.detach().cpu().item())
-
-
-def validate_event_trajectory_state_dict(
-    state_dict: Mapping[str, object],
-    *,
-    expected_d_model: int | None = None,
-    expected_n_trajectory: int | None = None,
-    expected_n_reasoning_rounds: int | None = None,
-) -> None:
     required_keys = {
-        "architecture_d_model",
+        "blocks.0.mlp.norm.weight",
-        "architecture_n_trajectory",
+        "blocks.0.mlp.norm.bias",
-        "architecture_n_reasoning_rounds",
+        "blocks.0.mlp.intra_gate_proj",
-        "event_projection.weight",
+        "blocks.0.mlp.intra_value_proj",
-        "trajectory_prototypes",
+        "blocks.0.mlp.intra_output_proj",
-        "query_projection.weight",
+        "blocks.0.mlp.intra_gate_logits",
-        "reasoning_core.cross_attention.q_proj.weight",
+        "blocks.0.mlp.gate_proj",
-        "reasoning_core.cross_attention.k_proj.weight",
+        "blocks.0.mlp.value_proj",
-        "reasoning_core.cross_attention.v_proj.weight",
+        "blocks.0.mlp.output_proj",
-        "reasoning_core.traj_mixer.gate_proj",
-        "reasoning_core.traj_mixer.value_proj",
-        "reasoning_core.traj_mixer.output_proj",
-        "reasoning_core.attn_scale",
-        "reasoning_core.mixer_scale",
     }
     missing = sorted(required_keys.difference(state_dict))
     if missing:
         raise ValueError(
-            "Checkpoint is not a shared event-trajectory checkpoint; "
+            "Checkpoint is not a TrajMixer checkpoint; missing required "
-            "missing required "
             f"parameters: {', '.join(missing)}"
         )
-    checkpoint_values = {
-        "d_model": _checkpoint_scalar_int(
-            state_dict,
-            "architecture_d_model",
-        ),
-        "n_trajectory": _checkpoint_scalar_int(
-            state_dict,
-            "architecture_n_trajectory",
-        ),
-        "n_reasoning_rounds": _checkpoint_scalar_int(
-            state_dict,
-            "architecture_n_reasoning_rounds",
-        ),
-    }
-    expected_values = {
-        "d_model": expected_d_model,
-        "n_trajectory": expected_n_trajectory,
-        "n_reasoning_rounds": expected_n_reasoning_rounds,
-    }
-    mismatches = [
-        f"{name}: checkpoint={checkpoint_values[name]}, expected={expected}"
-        for name, expected in expected_values.items()
-        if expected is not None and checkpoint_values[name] != expected
-    ]
-    if mismatches:
-        raise ValueError(
-            "Checkpoint architecture does not match the constructed model: "
-            + "; ".join(mismatches)
-        )
 
 
 @dataclass
@@ -332,8 +178,10 @@ class DeepHealth(nn.Module):
     def __init__(
         self,
         vocab_size: int,
-        model_size: str,
+        n_embd: int,
-        n_reasoning_rounds: int,
+        n_head: int,
+        n_hist_layer: int,
+        n_tab_layer: int,
         n_types: int,
         n_cont_types: int,
         n_categories: int,
@@ -358,21 +206,11 @@ class DeepHealth(nn.Module):
                 "dist_mode must be either 'exponential', 'weibull' or 'mixed'")
         if extra_pool_reduce not in {"mean", "sum"}:
             raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
-        if n_reasoning_rounds <= 0:
+        self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
-            raise ValueError(
-                "n_reasoning_rounds must be positive, got "
-                f"{n_reasoning_rounds}"
-            )
-        size_config = resolve_model_size(model_size)
-        normalized_model_size = model_size.strip().lower()
-        d_model = size_config.d_model
-        n_trajectory = size_config.n_trajectory
-
-        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
         self.gender_embedding = nn.Embedding(
-            2, d_model)  # Assuming binary gender
+            2, n_embd)  # Assuming binary gender
         self.tokenizer = OtherInfoTokenizer(
-            n_embd=d_model,
+            n_embd=n_embd,
             n_types=n_types,
             n_cont_types=n_cont_types,
             n_categories=n_categories,
@@ -384,101 +222,70 @@ class DeepHealth(nn.Module):
         self.time_mode = time_mode
         self.dist_mode = dist_mode
         self.extra_pool_reduce = extra_pool_reduce
-        self.model_size = normalized_model_size
+        self.n_embd = n_embd
-        self.d_model = d_model
-        self.n_trajectory = n_trajectory
-        self.trajectory_dim = d_model // n_trajectory
-        self.traj_hidden = 4 * n_trajectory
-        self.n_reasoning_rounds = n_reasoning_rounds
         self.vocab_size = vocab_size
-        self.register_buffer(
-            "architecture_d_model",
-            torch.tensor(d_model, dtype=torch.int64),
-        )
-        self.register_buffer(
-            "architecture_n_trajectory",
-            torch.tensor(n_trajectory, dtype=torch.int64),
-        )
-        self.register_buffer(
-            "architecture_n_reasoning_rounds",
-            torch.tensor(n_reasoning_rounds, dtype=torch.int64),
-        )
         nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
         nn.init.zeros_(self.token_embedding.weight[0])
         nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02)
         if dist_mode == "weibull":
-            self.rho_head = nn.Linear(d_model, vocab_size)
+            self.rho_head = nn.Linear(n_embd, vocab_size)
             nn.init.zeros_(self.rho_head.weight)
             nn.init.constant_(self.rho_head.bias, 0.5413)
 
         if dist_mode == "mixed":
             self.death_idx = vocab_size - 1
-            self.rho_death_head = nn.Linear(d_model, 1)
+            self.rho_death_head = nn.Linear(n_embd, 1)
             nn.init.zeros_(self.rho_death_head.weight)
             nn.init.constant_(self.rho_death_head.bias, 0.5413)
 
-        # Event and query time are encoded once before shared reasoning. In
+        if time_mode == "absolute":
-        # relative mode, cross-attention additionally uses TimeRoPE and RBF.
+            self.age_encoding = AgeSinusoidalEncoding(n_embd)
-        self.age_encoding = AgeSinusoidalEncoding(d_model)
+            self.blocks = nn.ModuleList([
-        self.event_projection = nn.Linear(d_model, d_model, bias=False)
+                GPTBlock(
-        self.event_norm = nn.LayerNorm(d_model)
+                    n_embd=n_embd,
-        self.query_projection = nn.Linear(d_model, d_model, bias=False)
+                    n_head=n_head,
-        self.trajectory_prototypes = nn.Parameter(
+                    use_time_rope=False,
-            torch.empty(n_trajectory, self.trajectory_dim)
+                    use_rbf_bias=False,
-        )
+                    mlp_dropout=dropout,
-        self.query_token = nn.Parameter(torch.empty(d_model))
+                ) for _ in range(n_hist_layer)
-        nn.init.normal_(self.event_projection.weight, mean=0.0, std=0.02)
+            ])
-        nn.init.normal_(self.query_projection.weight, mean=0.0, std=0.02)
-        nn.init.normal_(self.trajectory_prototypes, mean=0.0, std=0.02)
-        nn.init.normal_(self.query_token, mean=0.0, std=0.02)
-
-        use_relative_time = time_mode == "relative"
-        self.reasoning_core = SharedEventTrajectoryCore(
-            d_model=d_model,
-            n_trajectory=n_trajectory,
-            n_reasoning_rounds=n_reasoning_rounds,
-            dropout=dropout,
-            n_rbf_bases=16,
-            use_time_rope=use_relative_time,
-            use_rbf_bias=use_relative_time,
-        )
-        if use_relative_time:
-            self.rope = TimeRoPE(self.trajectory_dim)
-            self.rbf = GaussianRBFTimeBasis(
-                n_bases=16,
-                max_time_diff=40.0,
-            )
-        else:
             self.rope = None
             self.rbf = None
+        elif time_mode == "relative":
+            self.age_encoding = None
+            self.blocks = nn.ModuleList([
+                GPTBlock(
+                    n_embd=n_embd,
+                    n_head=n_head,
+                    use_time_rope=True,
+                    use_rbf_bias=True,
+                    mlp_dropout=dropout,
+                ) for _ in range(n_hist_layer)
+            ])
+            self.rope = TimeRoPE(n_embd // n_head)
+            self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)
 
-        self.final_ln = nn.LayerNorm(d_model)
+        self.final_ln = nn.LayerNorm(n_embd)
-        self.risk_head = nn.Linear(d_model, vocab_size, bias=False)
+        self.risk_head = nn.Linear(n_embd, vocab_size, bias=False)
         if target_mode == "next_token":
             self.risk_head.weight = self.token_embedding.weight
+        self.query_token = nn.Parameter(torch.zeros(n_embd))
+        nn.init.normal_(self.query_token, mean=0.0, std=0.02)
 
-    def _make_event_invalid_mask(
+    def _make_history_attn_mask(
         self,
-        event_valid_mask: torch.Tensor,
+        padding_mask: torch.Tensor,
-        event_time: torch.Tensor,
+        time_seq: torch.Tensor,
-        query_time: torch.Tensor,
+        dtype: torch.dtype,
-        query_position: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        valid_key = event_valid_mask[:, None, :]
+        valid_key = padding_mask[:, None, :]  # (B, 1, L)
-        key_time = event_time[:, None, :]
+        visible_by_time = time_seq[:, None, :] <= time_seq[:, :, None]
-        query_time = query_time[:, :, None]
+        valid = valid_key & visible_by_time
-        if query_position is None:
+        return torch.zeros(
-            visible_by_time = key_time <= query_time
+            valid.shape,
-        else:
+            device=valid.device,
-            key_position = torch.arange(
+            dtype=dtype,
-                event_time.size(1),
+        ).masked_fill(~valid, -1e4)[:, None, :, :]
-                device=event_time.device,
-            ).view(1, 1, -1)
-            visible_by_time = (key_time < query_time) | (
-                (key_time == query_time)
-                & (key_position <= query_position[:, :, None])
-            )
-        return ~(valid_key & visible_by_time)
 
     def _pool_other_by_time(
         self,
@@ -582,8 +389,8 @@ class DeepHealth(nn.Module):
             padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
 
         event_len = event_seq.size(1)
-        event_features = self.token_embedding(event_seq)
+        h_disease = self.token_embedding(event_seq)
-        event_time = time_seq
+        t_disease = time_seq
 
         if other_time.shape != other_type.shape:
             raise ValueError(
@@ -591,120 +398,64 @@ class DeepHealth(nn.Module):
                 f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
             )
         other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype)
-        other_features, other_mask = self.tokenizer(
+        h_other, other_mask = self.tokenizer(
             other_type=other_type,
             other_value=other_value,
             other_value_kind=other_value_kind,
         )
-        other_features = other_features.to(device=event_seq.device)
+        h_other = h_other.to(device=event_seq.device)
         other_mask = other_mask.to(device=event_seq.device, dtype=torch.bool)
 
-        event_features = torch.cat([event_features, other_features], dim=1)
+        h_disease = torch.cat([h_disease, h_other], dim=1)
-        event_time = torch.cat([event_time, other_time], dim=1)
+        t_disease = torch.cat([t_disease, other_time], dim=1)
-        event_valid_mask = torch.cat([padding_mask, other_mask], dim=1)
+        padding_mask = torch.cat([padding_mask, other_mask], dim=1)
-
+        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
-        batch_size = event_seq.size(0)
-        sex_context = self.gender_embedding(sex)[:, None, :]
-        event_features = (
-            event_features
-            + sex_context
-            + self.age_encoding(event_time)
-        )
-        event_features = event_features * event_valid_mask.unsqueeze(-1).to(
-            event_features.dtype
-        )
-        event_memory = self.event_norm(
-            self.event_projection(event_features)
-        )
-        event_memory = event_memory * event_valid_mask.unsqueeze(-1).to(
-            event_memory.dtype
-        )
 
         if mode == "all_future":
-            query_time = t_query[:, None]
+            batch_size = event_seq.size(0)
-            query_position = None
+            query = self.query_token.view(1, 1, -1).expand(batch_size, 1, -1)
-            query_features = (
+            h_disease = torch.cat([h_disease, query], dim=1)
-                self.query_token.view(1, 1, -1)
+            t_disease = torch.cat([t_disease, t_query[:, None]], dim=1)
-                + sex_context
+            query_mask = torch.ones(
-                + self.age_encoding(query_time)
-            )
-            query_valid_mask = torch.ones(
                 batch_size,
                 1,
                 dtype=torch.bool,
                 device=event_seq.device,
             )
-        else:
+            padding_mask = torch.cat([padding_mask, query_mask], dim=1)
-            # Each event position is an independent parallel query. Including
-            # its event feature preserves token-level next-step semantics.
-            # Equal-time memory is additionally position-causal so a token
-            # cannot read a later token that may be its Delphi2M target.
-            query_time = event_time
-            query_position = torch.arange(
-                event_time.size(1),
-                device=event_time.device,
-            ).view(1, -1).expand(batch_size, -1)
-            query_features = event_features
-            query_valid_mask = event_valid_mask
 
-        n_query = query_time.size(1)
+        sex_emb = self.gender_embedding(sex)[:, None, :]
-        query_context = self.query_projection(query_features).reshape(
+        h_disease = h_disease + sex_emb
-            batch_size,
+        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
-            n_query,
-            self.n_trajectory,
-            self.trajectory_dim,
-        )
-        trajectory_state = (
-            self.trajectory_prototypes.view(
-                1,
-                1,
-                self.n_trajectory,
-                self.trajectory_dim,
-            )
-            + query_context
-        )
-        event_invalid_mask = self._make_event_invalid_mask(
-            event_valid_mask=event_valid_mask,
-            event_time=event_time,
-            query_time=query_time,
-            query_position=query_position,
-        )
 
-        event_rope_cache = None
+        rope_cache = None
-        query_rope_cache = None
         rbf_cache = None
-        if self.time_mode == "relative":
+        if self.time_mode == "absolute":
-            if self.rope is None or self.rbf is None:
+            h_disease = h_disease + self.age_encoding(t_disease)
-                raise RuntimeError("Relative-time modules are not initialized")
+            h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
-            event_rope_cache = self.rope.precompute_cache(event_time)
+        elif self.time_mode == "relative":
-            query_rope_cache = self.rope.precompute_cache(query_time)
+            rope_cache = self.rope.precompute_cache(t_disease)
-            rbf_cache = self.rbf.precompute_cross_cache(
+            rbf_cache = self.rbf.precompute_cache(t_disease)
-                query_time,
-                event_time,
-            )
 
-        event_key_value = self.reasoning_core.project_event_memory(
+        attn_mask = self._make_history_attn_mask(
-            event_memory,
+            padding_mask=padding_mask,
-            event_rope_cache=event_rope_cache,
+            time_seq=t_disease,
+            dtype=h_disease.dtype,
         )
-        for _ in range(self.n_reasoning_rounds):
+        for block in self.blocks:
-            trajectory_state = self.reasoning_core(
+            h_disease = block(
-                trajectory_state=trajectory_state,
+                h_disease,
-                event_key_value=event_key_value,
+                rope_cache=rope_cache,
-                event_invalid_mask=event_invalid_mask,
-                query_rope_cache=query_rope_cache,
                 rbf_cache=rbf_cache,
+                attn_mask=attn_mask,
             )
+            h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
 
-        hidden_sequence = self.final_ln(
+        h_disease = self.final_ln(h_disease)
-            trajectory_state.reshape(batch_size, n_query, self.d_model)
+        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
-        )
-        hidden_sequence = hidden_sequence * query_valid_mask.unsqueeze(-1).to(
-            hidden_sequence.dtype
-        )
 
         if mode == "all_future":
-            hidden = hidden_sequence[:, 0, :]
+            hidden = h_disease[:, -1, :]
             if return_output:
                 return DeepHealthOutput(
                     hidden=hidden,
@@ -719,13 +470,13 @@ class DeepHealth(nn.Module):
                 )
             return hidden
         if return_output:
-            h_event = hidden_sequence[:, :event_len, :]
+            h_event = h_disease[:, :event_len, :]
-            t_event = event_time[:, :event_len]
+            t_event = t_disease[:, :event_len]
-            event_mask = event_valid_mask[:, :event_len]
+            event_mask = padding_mask[:, :event_len]
             h_extra, t_extra, extra_mask = self._pool_other_by_time(
-                h_other=hidden_sequence[:, event_len:, :],
+                h_other=h_disease[:, event_len:, :],
-                other_time=event_time[:, event_len:],
+                other_time=t_disease[:, event_len:],
-                other_mask=event_valid_mask[:, event_len:],
+                other_mask=padding_mask[:, event_len:],
             )
             return DeepHealthOutput(
                 hidden=torch.cat([h_event, h_extra], dim=1),
@@ -733,7 +484,7 @@ class DeepHealth(nn.Module):
                 padding_mask=torch.cat([event_mask, extra_mask], dim=1),
                 event_len=event_len,
             )
-        return hidden_sequence[:, :event_len, :]
+        return h_disease[:, :event_len, :]
 
     def forward_next_token(self, **kwargs) -> torch.Tensor:
         return self._forward_shared(mode="next_token", **kwargs)

test_event_trajectory_backbone.py

@@ -1,330 +0,0 @@
-import math
-import unittest
-
-import torch
-
-from backbones import (
-    SharedEventTrajectoryCore,
-    SharedTrajectoryMixer,
-    TrajectoryCrossAttention,
-)
-from models import (
-    EVENT_TRAJECTORY_ARCHITECTURE,
-    MODEL_SIZE_PRESETS,
-    DeepHealth,
-    resolve_model_size,
-    validate_event_trajectory_config,
-    validate_event_trajectory_state_dict,
-)
-from train_util import get_model_parameter_counts
-
-
-def build_test_model(
-    *,
-    target_mode: str = "next_token",
-    time_mode: str = "absolute",
-    n_reasoning_rounds: int = 3,
-) -> DeepHealth:
-    return DeepHealth(
-        vocab_size=32,
-        model_size="nano",
-        n_reasoning_rounds=n_reasoning_rounds,
-        n_types=2,
-        n_cont_types=0,
-        n_categories=2,
-        cont_type_ids=[],
-        target_mode=target_mode,
-        time_mode=time_mode,
-    )
-
-
-def model_inputs() -> dict[str, torch.Tensor]:
-    return {
-        "event_seq": torch.tensor([[1, 2, 3, 4], [5, 6, 0, 0]]),
-        "time_seq": torch.tensor(
-            [[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 0.0, 0.0]]
-        ),
-        "sex": torch.tensor([0, 1]),
-        "padding_mask": torch.tensor(
-            [[True, True, True, True], [True, True, False, False]]
-        ),
-        "other_type": torch.zeros(2, 1, dtype=torch.long),
-        "other_value": torch.zeros(2, 1),
-        "other_value_kind": torch.zeros(2, 1, dtype=torch.long),
-        "other_time": torch.zeros(2, 1),
-    }
-
-
-class EventTrajectoryBackboneTest(unittest.TestCase):
-    def test_model_size_presets(self) -> None:
-        expected = {
-            "nano": (120, 6, 20, 24),
-            "tiny": (256, 8, 32, 32),
-            "small": (512, 8, 64, 32),
-            "medium": (768, 12, 64, 48),
-            "huge": (1024, 16, 64, 64),
-        }
-        self.assertEqual(set(MODEL_SIZE_PRESETS), set(expected))
-        for name, values in expected.items():
-            preset = resolve_model_size(name)
-            self.assertEqual(
-                (
-                    preset.d_model,
-                    preset.n_trajectory,
-                    preset.trajectory_dim,
-                    preset.traj_hidden,
-                ),
-                values,
-            )
-
-    def test_default_mixer_shapes_and_parameter_count(self) -> None:
-        mixer = SharedTrajectoryMixer(
-            n_trajectory=8,
-            trajectory_dim=32,
-        )
-        state = torch.randn(2, 5, 8, 32)
-        self.assertEqual(mixer(state).shape, state.shape)
-        self.assertEqual(mixer.traj_hidden, 32)
-        self.assertEqual(tuple(mixer.gate_proj.shape), (32, 8, 32))
-        self.assertEqual(tuple(mixer.value_proj.shape), (32, 8, 32))
-        self.assertEqual(tuple(mixer.output_proj.shape), (32, 32, 8))
-        self.assertEqual(
-            get_model_parameter_counts(mixer),
-            {
-                "model_parameter_count": 24_576,
-                "trainable_parameter_count": 24_576,
-            },
-        )
-
-    def test_reasoning_rounds_share_one_core_parameter_set(self) -> None:
-        core_one = SharedEventTrajectoryCore(
-            d_model=256,
-            n_trajectory=8,
-            n_reasoning_rounds=1,
-        )
-        core_twelve = SharedEventTrajectoryCore(
-            d_model=256,
-            n_trajectory=8,
-            n_reasoning_rounds=12,
-        )
-        self.assertEqual(
-            sum(p.numel() for p in core_one.parameters()),
-            sum(p.numel() for p in core_twelve.parameters()),
-        )
-        self.assertAlmostEqual(core_one.attn_scale.item(), 1.0)
-        self.assertAlmostEqual(
-            core_twelve.attn_scale.item(),
-            1.0 / math.sqrt(12),
-            places=6,
-        )
-        self.assertAlmostEqual(
-            core_twelve.mixer_scale.item(),
-            1.0 / math.sqrt(12),
-            places=6,
-        )
-
-    def test_all_masked_attention_is_finite_and_zero(self) -> None:
-        attention = TrajectoryCrossAttention(
-            d_model=32,
-            n_trajectory=4,
-        )
-        memory = torch.randn(2, 3, 32)
-        key_value = attention.project_event_memory(memory)
-        state = torch.randn(2, 2, 4, 8)
-        invalid_mask = torch.ones(2, 2, 3, dtype=torch.bool)
-        output = attention(
-            trajectory_state=state,
-            event_key_value=key_value,
-            event_invalid_mask=invalid_mask,
-        )
-        self.assertTrue(torch.isfinite(output).all())
-        torch.testing.assert_close(output, torch.zeros_like(output))
-
-    def test_next_token_future_events_do_not_change_earlier_query(self) -> None:
-        torch.manual_seed(0)
-        model = build_test_model(n_reasoning_rounds=2)
-        model.eval()
-        inputs = model_inputs()
-        original = model(**inputs)
-        changed_inputs = dict(inputs)
-        changed_inputs["event_seq"] = inputs["event_seq"].clone()
-        changed_inputs["event_seq"][0, 3] = 9
-        changed = model(**changed_inputs)
-        torch.testing.assert_close(original[0, 1], changed[0, 1])
-
-    def test_next_token_later_equal_time_event_is_not_visible(self) -> None:
-        torch.manual_seed(0)
-        model = build_test_model(n_reasoning_rounds=2)
-        model.eval()
-        inputs = model_inputs()
-        inputs["time_seq"] = inputs["time_seq"].clone()
-        inputs["time_seq"][0] = torch.tensor([1.0, 1.0, 2.0, 3.0])
-        original = model(**inputs)
-        changed_inputs = dict(inputs)
-        changed_inputs["event_seq"] = inputs["event_seq"].clone()
-        changed_inputs["event_seq"][0, 1] = 9
-        changed = model(**changed_inputs)
-        torch.testing.assert_close(original[0, 0], changed[0, 0])
-
-    def test_padding_content_does_not_change_valid_queries(self) -> None:
-        torch.manual_seed(0)
-        model = build_test_model(
-            time_mode="relative",
-            n_reasoning_rounds=2,
-        )
-        model.eval()
-        inputs = model_inputs()
-        original = model(**inputs)
-        changed_inputs = dict(inputs)
-        changed_inputs["event_seq"] = inputs["event_seq"].clone()
-        changed_inputs["time_seq"] = inputs["time_seq"].clone()
-        changed_inputs["event_seq"][1, 2:] = torch.tensor([9, 10])
-        changed_inputs["time_seq"][1, 2:] = torch.tensor([30.0, 40.0])
-        changed = model(**changed_inputs)
-        torch.testing.assert_close(original[1, :2], changed[1, :2])
-
-    def test_next_token_and_all_future_output_contracts(self) -> None:
-        inputs = model_inputs()
-        next_model = build_test_model(target_mode="next_token")
-        next_hidden = next_model(**inputs)
-        self.assertEqual(tuple(next_hidden.shape), (2, 4, 120))
-        next_output = next_model(**inputs, return_output=True)
-        self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 120))
-        self.assertEqual(tuple(next_output.padding_mask.shape), (2, 4))
-
-        future_model = build_test_model(target_mode="all_future")
-        future_hidden = future_model(
-            **inputs,
-            t_query=torch.tensor([5.0, 3.0]),
-        )
-        self.assertEqual(tuple(future_hidden.shape), (2, 120))
-
-    def test_model_contains_one_shared_core_and_no_block_stack(self) -> None:
-        model = build_test_model(n_reasoning_rounds=12)
-        self.assertFalse(hasattr(model, "blocks"))
-        reasoning_keys = [
-            key
-            for key in model.state_dict()
-            if key.startswith("reasoning_core.")
-        ]
-        self.assertTrue(reasoning_keys)
-        self.assertFalse(any("blocks." in key for key in model.state_dict()))
-        self.assertFalse(any("out_proj" in key for key in reasoning_keys))
-        self.assertFalse(any("group_align" in key for key in reasoning_keys))
-
-    def test_event_key_and_value_are_projected_once_per_forward(self) -> None:
-        model = build_test_model(n_reasoning_rounds=12)
-        call_counts = {"key": 0, "value": 0}
-
-        def count_key(*_args) -> None:
-            call_counts["key"] += 1
-
-        def count_value(*_args) -> None:
-            call_counts["value"] += 1
-
-        key_handle = (
-            model.reasoning_core.cross_attention.k_proj
-            .register_forward_hook(count_key)
-        )
-        value_handle = (
-            model.reasoning_core.cross_attention.v_proj
-            .register_forward_hook(count_value)
-        )
-        try:
-            model(**model_inputs())
-        finally:
-            key_handle.remove()
-            value_handle.remove()
-        self.assertEqual(call_counts, {"key": 1, "value": 1})
-
-    def test_relative_time_forward_and_backward_are_finite(self) -> None:
-        torch.manual_seed(0)
-        model = build_test_model(
-            target_mode="all_future",
-            time_mode="relative",
-            n_reasoning_rounds=2,
-        )
-        hidden = model(
-            **model_inputs(),
-            t_query=torch.tensor([5.0, 3.0]),
-        )
-        (hidden * torch.randn_like(hidden)).sum().backward()
-        self.assertTrue(torch.isfinite(hidden).all())
-        self.assertIsNotNone(model.event_projection.weight.grad)
-        self.assertTrue(torch.isfinite(model.event_projection.weight.grad).all())
-        time_scale = model.reasoning_core.cross_attention.time_bias_scale
-        self.assertIsNotNone(time_scale)
-        self.assertIsNotNone(time_scale.grad)
-        self.assertGreater(abs(float(time_scale.grad)), 0.0)
-
-    def test_architecture_marker_and_checkpoint_are_required(self) -> None:
-        validate_event_trajectory_config(
-            {
-                "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
-                "model_size": "nano",
-                "d_model": 120,
-                "n_trajectory": 6,
-                "trajectory_dim": 20,
-                "traj_hidden": 24,
-                "n_reasoning_rounds": 3,
-            }
-        )
-        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
-            validate_event_trajectory_config(
-                {"model_architecture": "traj_mixer_v2"}
-            )
-        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
-            validate_event_trajectory_config(
-                {"model_architecture": "event_trajectory_shared_v1"}
-            )
-        with self.assertRaisesRegex(ValueError, "trajectory_dim"):
-            validate_event_trajectory_config(
-                {
-                    "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
-                    "model_size": "nano",
-                    "d_model": 120,
-                    "n_trajectory": 6,
-                    "trajectory_dim": 10,
-                    "traj_hidden": 24,
-                    "n_reasoning_rounds": 3,
-                }
-            )
-
-        model = build_test_model()
-        state_dict = model.state_dict()
-        validate_event_trajectory_state_dict(
-            state_dict,
-            expected_d_model=120,
-            expected_n_trajectory=6,
-            expected_n_reasoning_rounds=3,
-        )
-        with self.assertRaisesRegex(
-            ValueError,
-            "Checkpoint architecture does not match",
-        ):
-            validate_event_trajectory_state_dict(
-                state_dict,
-                expected_n_reasoning_rounds=12,
-            )
-        state_dict.pop("reasoning_core.attn_scale")
-        with self.assertRaisesRegex(
-            ValueError,
-            "not a shared event-trajectory checkpoint",
-        ):
-            validate_event_trajectory_state_dict(state_dict)
-
-    def test_unknown_model_size_is_rejected(self) -> None:
-        with self.assertRaisesRegex(ValueError, "Unknown model_size"):
-            DeepHealth(
-                vocab_size=32,
-                model_size="giant",
-                n_reasoning_rounds=2,
-                n_types=2,
-                n_cont_types=0,
-                n_categories=2,
-                cont_type_ids=[],
-            )
-
-
-if __name__ == "__main__":
-    unittest.main()

test_traj_mixer.py

@@ -0,0 +1,253 @@
+import unittest
+
+import torch
+
+from backbones import GPTBlock, TemporalAttention, TrajMixer
+from models import (
+    TRAJ_MIXER_ARCHITECTURE,
+    validate_traj_mixer_config,
+    validate_traj_mixer_state_dict,
+)
+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,
+            n_head=10,
+            dropout=0.0,
+        )
+
+        x = torch.randn(2, 7, 120)
+        self.assertEqual(mixer(x).shape, x.shape)
+        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),
+            (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)
+        mixer.eval()
+        x = torch.randn(2, 5, 120)
+        changed = x.clone()
+        changed[:, 3, :] += torch.randn_like(changed[:, 3, :])
+
+        original_out = mixer(x)
+        changed_out = mixer(changed)
+        unchanged_positions = torch.tensor([0, 1, 2, 4])
+        torch.testing.assert_close(
+            original_out.index_select(1, unchanged_positions),
+            changed_out.index_select(1, unchanged_positions),
+        )
+
+    def test_gradients_reach_all_projection_families(self) -> None:
+        torch.manual_seed(1)
+        mixer = TrajMixer(120, n_head=10, dropout=0.0)
+        x = torch.randn(2, 4, 120, requires_grad=True)
+
+        mixer(x).square().mean().backward()
+
+        self.assertIsNotNone(x.grad)
+        for name, parameter in mixer.named_parameters():
+            self.assertIsNotNone(parameter.grad, name)
+            self.assertTrue(torch.isfinite(parameter.grad).all(), name)
+
+    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.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)
+
+    def test_architecture_marker_is_required(self) -> None:
+        validate_traj_mixer_config(
+            {"model_architecture": TRAJ_MIXER_ARCHITECTURE}
+        )
+        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
+            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)
+        state_dict = {
+            f"blocks.0.{key}": value
+            for key, value in block.state_dict().items()
+        }
+        validate_traj_mixer_state_dict(state_dict)
+
+        state_dict.pop("blocks.0.mlp.intra_gate_proj")
+        with self.assertRaisesRegex(ValueError, "not a TrajMixer checkpoint"):
+            validate_traj_mixer_state_dict(state_dict)
+
+    def test_invalid_group_partition_is_rejected(self) -> None:
+        with self.assertRaisesRegex(ValueError, "divisible"):
+            TrajMixer(n_embd=121, n_head=10)
+
+    def test_parameter_counts_match_traj_mixer_parameters(self) -> None:
+        mixer = TrajMixer(n_embd=120, n_head=10)
+        self.assertEqual(
+            get_model_parameter_counts(mixer),
+            {
+                "model_parameter_count": 32_040,
+                "trainable_parameter_count": 32_040,
+            },
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

train_all_future.py

@@ -27,12 +27,7 @@ from tqdm.auto import tqdm
 
 from dataset import AllFutureHealthDataset, all_future_collate_fn
 from losses import build_loss
-from models import (
+from models import TRAJ_MIXER_ARCHITECTURE, DeepHealth
-    EVENT_TRAJECTORY_ARCHITECTURE,
-    MODEL_SIZE_NAMES,
-    DeepHealth,
-    resolve_model_size,
-)
 from targets import CHECKUP_IDX, PAD_IDX
 from train_util import (
     configure_torch_for_training,
@@ -83,13 +78,10 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--min_future_events", type=int, default=1)
     parser.add_argument("--validation_query_seed", type=int, default=None)
 
-    parser.add_argument(
+    parser.add_argument("--n_embd", type=int, default=120)
-        "--model_size",
+    parser.add_argument("--n_head", type=int, default=10)
-        type=str,
+    parser.add_argument("--n_hist_layer", type=int, default=12)
-        default="nano",
+    parser.add_argument("--n_tab_layer", type=int, default=4)
-        choices=MODEL_SIZE_NAMES,
-    )
-    parser.add_argument("--n_reasoning_rounds", type=int, default=12)
     parser.add_argument("--n_bins", type=int, default=16)
     parser.add_argument("--extra_pool_reduce", type=str, default="mean",
                         choices=["mean", "sum"])
@@ -154,8 +146,10 @@ def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -
 def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth:
     return DeepHealth(
         vocab_size=dataset.vocab_size,
-        model_size=args.model_size,
+        n_embd=args.n_embd,
-        n_reasoning_rounds=args.n_reasoning_rounds,
+        n_head=args.n_head,
+        n_hist_layer=args.n_hist_layer,
+        n_tab_layer=args.n_tab_layer,
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -300,17 +294,12 @@ def build_metadata(
     val_subset,
     test_subset,
 ) -> Dict[str, Any]:
-    size_config = resolve_model_size(args.model_size)
     return {
         "run_name": run_name,
         "dataset_class": "AllFutureHealthDataset",
         "collate_fn": "all_future_collate_fn",
         "model_class": "DeepHealth",
-        "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
+        "model_architecture": TRAJ_MIXER_ARCHITECTURE,
-        "d_model": size_config.d_model,
-        "n_trajectory": size_config.n_trajectory,
-        "trajectory_dim": size_config.trajectory_dim,
-        "traj_hidden": size_config.traj_hidden,
         "model_target_mode": "all_future",
         "target_mode": "all_future",
         "dist_mode": args.dist_mode,
@@ -348,26 +337,12 @@ def main() -> None:
     configure_torch_for_training(device)
 
     run_dir, run_name = create_unique_run_dir(
-        lambda timestamp: (
+        lambda timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}"
-            f"{args.model_size}_r{args.n_reasoning_rounds}_"
-            f"{args.time_mode}_{args.dist_mode}_"
-            f"all_future_pure_disease_{timestamp}"
-        )
     )
     logger = setup_logging(run_dir)
 
     logger.info(f"Starting all-future training run: {run_name}")
     logger.info(f"Device: {device}")
-    size_config = resolve_model_size(args.model_size)
-    logger.info(
-        "Model size: "
-        f"{args.model_size} "
-        f"(d_model={size_config.d_model}, "
-        f"n_trajectory={size_config.n_trajectory}, "
-        f"trajectory_dim={size_config.trajectory_dim}, "
-        f"traj_hidden={size_config.traj_hidden}); "
-        f"reasoning_rounds={args.n_reasoning_rounds}"
-    )
     logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
 
     logger.info("Loading all-future datasets...")

train_next_step.py

@@ -24,13 +24,7 @@ from tqdm.auto import tqdm
 
 from dataset import HealthDataset, collate_fn
 from losses import build_loss
-from models import (
+from models import TRAJ_MIXER_ARCHITECTURE, DeepHealth, DeepHealthOutput
-    EVENT_TRAJECTORY_ARCHITECTURE,
-    MODEL_SIZE_NAMES,
-    DeepHealth,
-    DeepHealthOutput,
-    resolve_model_size,
-)
 from readouts import build_readout
 from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
 from train_util import (
@@ -80,13 +74,10 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
     parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
 
-    parser.add_argument(
+    parser.add_argument("--n_embd", type=int, default=120)
-        "--model_size",
+    parser.add_argument("--n_head", type=int, default=10)
-        type=str,
+    parser.add_argument("--n_hist_layer", type=int, default=12)
-        default="nano",
+    parser.add_argument("--n_tab_layer", type=int, default=4)
-        choices=MODEL_SIZE_NAMES,
-    )
-    parser.add_argument("--n_reasoning_rounds", type=int, default=12)
     parser.add_argument("--n_bins", type=int, default=16)
     parser.add_argument("--extra_pool_reduce", type=str, default="mean",
                         choices=["mean", "sum"])
@@ -160,8 +151,10 @@ def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -
 def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
     return DeepHealth(
         vocab_size=dataset.vocab_size,
-        model_size=args.model_size,
+        n_embd=args.n_embd,
-        n_reasoning_rounds=args.n_reasoning_rounds,
+        n_head=args.n_head,
+        n_hist_layer=args.n_hist_layer,
+        n_tab_layer=args.n_tab_layer,
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -487,17 +480,12 @@ def build_metadata(
     val_subset,
     test_subset,
 ) -> Dict[str, Any]:
-    size_config = resolve_model_size(args.model_size)
     return {
         "run_name": run_name,
         "dataset_class": "NextStepHealthDataset",
         "collate_fn": "next_step_collate_fn",
         "model_class": "DeepHealth",
-        "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
+        "model_architecture": TRAJ_MIXER_ARCHITECTURE,
-        "d_model": size_config.d_model,
-        "n_trajectory": size_config.n_trajectory,
-        "trajectory_dim": size_config.trajectory_dim,
-        "traj_hidden": size_config.traj_hidden,
         "model_target_mode": "next_token",
         "target_mode": args.target_mode,
         "dist_mode": "exponential",
@@ -533,9 +521,7 @@ def main() -> None:
 
     run_dir, run_name = create_unique_run_dir(
         lambda timestamp: (
-            f"{args.model_size}_r{args.n_reasoning_rounds}_"
+            f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
-            f"{args.time_mode}_exponential_"
-            f"next_token_{args.target_mode}_"
             f"gap_{args.no_event_interval_years:g}y_{timestamp}"
         )
     )
@@ -543,16 +529,6 @@ def main() -> None:
 
     logger.info(f"Starting next-step training run: {run_name}")
     logger.info(f"Device: {device}")
-    size_config = resolve_model_size(args.model_size)
-    logger.info(
-        "Model size: "
-        f"{args.model_size} "
-        f"(d_model={size_config.d_model}, "
-        f"n_trajectory={size_config.n_trajectory}, "
-        f"trajectory_dim={size_config.trajectory_dim}, "
-        f"traj_hidden={size_config.traj_hidden}); "
-        f"reasoning_rounds={args.n_reasoning_rounds}"
-    )
     logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
     logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")