Implement shared event-trajectory reasoning backbone

Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md

@@ -0,0 +1,386 @@
+# Event–Trajectory Shared Reasoning Backbone
+
+> 状态：**Frozen implementation baseline**
+> 架构标识：`event_trajectory_shared_v1`
+> 固化日期：**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 | 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 默认使用8个：
+
+\[
+S\in\mathbb{R}^{B\times Q\times8\times32}.
+\]
+
+其中：
+
+- all-future：\(Q=1\)；
+- next-token：\(Q=L\)，所有查询位置并行计算。
+
+定义可学习原型：
+
+\[
+P\in\mathbb{R}^{8\times32}.
+\]
+
+查询上下文经过投影并 reshape：
+
+\[
+C_Q
+=\operatorname{QueryProjection}(\text{query features})
+\in\mathbb{R}^{B\times Q\times8\times32},
+\]
+
+\[
+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_v1
+model_size: nano
+d_model: 256
+n_trajectory: 8
+trajectory_dim: 32
+traj_hidden: 32
+n_reasoning_rounds: 12
+model_parameter_count: <runtime count>
+trainable_parameter_count: <runtime count>
+```
+
+评估和导出入口必须同时验证：
+
+1. `model_architecture` 完全匹配；
+2. `model_size` 属于 `nano / 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{多轮状态依赖推理}
+}
+\]