DeepHealth/Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md

# 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{多轮状态依赖推理}
}
\]