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