DeepHealth/Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md

Event–Trajectory Shared Reasoning Backbone

状态：Frozen implementation baseline 架构标识：event_trajectory_shared_v2 固化日期：2026-07-23

1. 核心定义

使用一个共享的 Attention–TrajMixer 推理核心，对固定 Event Memory 进行多轮读取，并持续更新 Trajectory State。

模型只实例化：

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}}}. ]

事件特征由以下信息相加：

disease / covariate embedding
+ age/time encoding
+ sex context

然后只编码一次：

[ E=\operatorname{EventNorm} \left(\operatorname{EventProjection}(X_E)\right). ]

进入 reasoning loop 后，(E) 的数值保持不变，但不执行 detach，梯度仍可回传到事件编码器。

Key 和 Value 同样每次 forward 只投影一次：

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. ]

形状为：

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)}. ]

参数形状：

W_g: [trajectory_dim, n_trajectory, traj_hidden]
W_v: [trajectory_dim, n_trajectory, traj_hidden]
W_o: [trajectory_dim, traj_hidden, n_trajectory]

默认：

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. 多轮参数共享

同一个核心重复运行：

for _ in range(n_reasoning_rounds):
    S = self.reasoning_core(...)

所有轮次共享：

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}}}}. ]

两个标量可学习，并由全部轮次共享。

第一版不加入：

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 约束

训练配置必须写入：

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. 信息流

E ─────────────┬──────────────┬──────────────┬──────────────┐
               │              │              │              │
               ▼              ▼              ▼              ▼
S0 -> Shared Core -> S1 -> Shared Core -> S2 -> ... -> Shared Core -> S12
       同一套参数          同一套参数                   同一套参数

整体定义：

[ \boxed{ \text{一个共享 Event–Trajectory 推理核心} \times \text{多轮状态依赖推理} } ]