JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor DeepHealth model to expand extra-info token handling and update output structure

Browse Source
This commit is contained in:
Jiarui Li
2026-06-17 14:27:07 +08:00
parent 1757bcd25b
commit 1ea72e9133
2 changed files with 25 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
README.md
View File

@@ -135,20 +135,20 @@ type_id -> continuous type index -> soft bin embedding
selected type offsets + local category id selected type offsets + local category id
``` ```
同一患者、同一 `other_time` 的 extra-info token 会先被池化为一个 token，再并入主序列。默认池化方式是平均池化： 主序列内部保留展开的单个 extra-info token；疾病 token 因此可以直接关注到每个额外信息变量。只有在 next-token 训练需要把 extra-info 位置接入 readout/loss 时，同一患者、同一 `other_time` 的 extra hidden 才会被池化成一个监督用 token。默认池化方式是平均池化：
```bash ```bash
--extra_pool_reduce mean --extra_pool_reduce mean
``` ```
也可以设为 `sum`。池化后，每个 extra-info 时间点最多产生一个主序列 token。 也可以设为 `sum`。池化后，每个 extra-info 时间点最多产生一个用于预测和 loss 的 readout token。
### DeepHealth ### DeepHealth
`DeepHealth` 的统一路径是： `DeepHealth` 的统一路径是：
```text ```text
disease tokens + pooled extra-info tokens disease tokens + expanded extra-info tokens
-> temporal backbone -> temporal backbone
-> risk head -> risk head
``` ```
@@ -158,23 +158,23 @@ extra-info 不再通过独立的 `BaselineEncoder` 或 `CrossAttention` 注入
两种目标模式的读出语义不同： 两种目标模式的读出语义不同：
- `next_token` - `next_token`
- 模型内部序列包含 disease tokens 和 pooled extra-info tokens - 模型内部序列包含 disease tokens 和展开的单个 extra-info tokens
- 默认 Tensor 返回值仍只返回 disease token hidden，形状为 `(B, L, D)`，兼容评估和旧调用 - 默认 Tensor 返回值仍只返回 disease token hidden，形状为 `(B, L, D)`，兼容评估和旧调用
- 训练时使用 `return_output=True` 取完整主序列输出；pooled extra-info tokens 也会产生 logits，并在有未来 disease target 时参与 prediction/loss 监督 - 训练时使用 `return_output=True`；此时同一 `other_time` 的 extra hidden 会在读出端池化，pooled extra readout tokens 也会产生 logits，并在有未来 disease target 时参与 prediction/loss 监督
- next-token 模式下 `risk_head.weight` 与 `token_embedding.weight` 使用 weight tying - next-token 模式下 `risk_head.weight` 与 `token_embedding.weight` 使用 weight tying
- `all_future` - `all_future`
- 在合并后的主序列末尾拼接一个 query token - 在合并后的主序列末尾拼接一个 query token
- query token 的时间是 `t_query` - query token 的时间是 `t_query`
- 只读出最后一个 query hidden，形状为 `(B, D)` - 只读出最后一个 query hidden，形状为 `(B, D)`
- pooled extra-info tokens 只作为 query 的上下文输入，不单独读出、不参与 loss 监督 - 展开的 extra-info tokens 只作为 query 的上下文输入，不单独读出、不参与 loss 监督
- all-future 模式不使用 weight tying - all-future 模式不使用 weight tying
如果需要拿到完整 next-token 输出，可使用结构化返回： 如果需要拿到完整 next-token 输出，可使用结构化返回：
```python ```python
out = model(..., target_mode="next_token", return_output=True) out = model(..., target_mode="next_token", return_output=True)
out.hidden # disease tokens + pooled extra-info tokens out.hidden # disease tokens + pooled extra-info readout tokens
out.time_seq # 与 hidden 对齐的时间 out.time_seq # 与 hidden 对齐的时间
out.padding_mask # 与 hidden 对齐的有效位置 out.padding_mask # 与 hidden 对齐的有效位置
out.event_len # 原 disease token 长度 out.event_len # 原 disease token 长度
@@ -209,9 +209,9 @@ next-token 监督：
next-token 训练中，模型会请求 `return_output=True`，因此 loss 的预测位置包括： next-token 训练中，模型会请求 `return_output=True`，因此 loss 的预测位置包括：
- 原 disease token readout 位置 - 原 disease token readout 位置
- 同一时间点 extra-info 池化后的 pooled extra-info token - 同一时间点 extra hidden 池化后的 pooled extra-info readout token
pooled extra-info token 的监督目标在训练时动态构造：对 pooled extra-info token 的时间 `t`，寻找该患者 `t` 之后的下一个 disease 事件时间；`delphi2m` 使用第一个未来事件作为 next-token target，`uts` 使用下一唯一时间点上的事件集合做 multi-hot target。若该 extra-info 时间点之后没有未来 disease target，则该位置不参与 loss。 pooled extra-info readout token 的监督目标在训练时动态构造：对 pooled extra-info readout token 的时间 `t`，寻找该患者 `t` 之后的下一个 disease 事件时间；`delphi2m` 使用第一个未来事件作为 next-token target，`uts` 使用下一唯一时间点上的事件集合做 multi-hot target。若该 extra-info 时间点之后没有未来 disease target，则该位置不参与 loss。
all-future / query-conditioned 监督： all-future / query-conditioned 监督：
@@ -219,7 +219,7 @@ all-future / query-conditioned 监督：
- `WeibullLoss` - `WeibullLoss`
- `MixedLoss` - `MixedLoss`
all-future 训练只读出 `t_query` 对应的 query hidden。pooled extra-info tokens 作为主序列上下文输入，但不会被单独读出，也不会被纳入 loss 监督。 all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-info tokens 作为主序列上下文输入，但不会被单独读出，也不会被纳入 loss 监督。
`UniqueTimeSetExponentialLoss` 的 observed term 固定使用 sum reduction，不再暴露旧的 `observed_reduction` 参数。 `UniqueTimeSetExponentialLoss` 的 observed term 固定使用 sum reduction，不再暴露旧的 `observed_reduction` 参数。
@@ -232,12 +232,12 @@ all-future 训练只读出 `t_query` 对应的 query hidden。pooled extra-info
- `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout - `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout
- `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout - `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout
- 当前 next-token 训练只支持 exponential time loss - 当前 next-token 训练只支持 exponential time loss
- pooled extra-info tokens 会加入 prediction/loss 监督 - 展开的 extra-info tokens 进入主序列；读出端 pooled extra-info tokens 会加入 prediction/loss 监督
- `train_all_future.py` - `train_all_future.py`
- 使用 `AllFutureHealthDataset` - 使用 `AllFutureHealthDataset`
- 不使用 readout，直接对 query hidden 计算风险 - 不使用 readout，直接对 query hidden 计算风险
- `--dist_mode exponential/weibull/mixed` 分别搭配 `ExponentialLoss`、`WeibullLoss`、`MixedLoss` - `--dist_mode exponential/weibull/mixed` 分别搭配 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`
- pooled extra-info tokens 只作为 query 上下文，不单独监督 - 展开的 extra-info tokens 只作为 query 上下文，不单独监督
当前 `train_next_step.py` / `train_all_future.py` 支持所有已有训练目标定义的组合： 当前 `train_next_step.py` / `train_all_future.py` 支持所有已有训练目标定义的组合：
@@ -293,7 +293,7 @@ python train_next_step.py --extra_info_types_file extra_info_types_smoking_alcoh
当前提供两个 AUC 评估入口，二者都已适配新的 `DeepHealth` 模型和统一的 other-info token 输入；AUC 的 DeLong 计算、病例/对照筛选和分层聚合逻辑保持原评估脚本口径。 当前提供两个 AUC 评估入口，二者都已适配新的 `DeepHealth` 模型和统一的 other-info token 输入；AUC 的 DeLong 计算、病例/对照筛选和分层聚合逻辑保持原评估脚本口径。
评估脚本默认使用常规 Tensor 返回值：next-token checkpoint 只缓存 disease token/readout 位置的 hidden；all-future checkpoint 只缓存 `t_query` query hidden。训练中额外加入监督的 pooled extra-info tokens 不作为 AUC 评估位置单独输出。 评估脚本默认使用常规 Tensor 返回值：next-token checkpoint 只缓存 disease token/readout 位置的 hidden；all-future checkpoint 只缓存 `t_query` query hidden。训练中额外加入监督的 pooled extra-info readout tokens 不作为 AUC 评估位置单独输出。
### `evaluate_auc.py` ### `evaluate_auc.py`

21
models.py
View File

@@ -322,7 +322,7 @@ class DeepHealth(nn.Module):
other_time: torch.FloatTensor | None = None, other_time: torch.FloatTensor | None = None,
return_output: bool = False, return_output: bool = False,
**unused_kwargs, **unused_kwargs,
) -> torch.Tensor: ) -> torch.Tensor | DeepHealthOutput:
if unused_kwargs: if unused_kwargs:
unknown = ", ".join(sorted(unused_kwargs)) unknown = ", ".join(sorted(unused_kwargs))
raise TypeError(f"Unexpected DeepHealth forward arguments: {unknown}") raise TypeError(f"Unexpected DeepHealth forward arguments: {unknown}")
@@ -363,11 +363,6 @@ class DeepHealth(nn.Module):
) )
h_other = h_other.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) other_mask = other_mask.to(device=event_seq.device, dtype=torch.bool)
h_other, other_time, other_mask = self._pool_other_by_time(
h_other=h_other,
other_time=other_time,
other_mask=other_mask,
)
h_disease = torch.cat([h_disease, h_other], dim=1) h_disease = torch.cat([h_disease, h_other], dim=1)
t_disease = torch.cat([t_disease, other_time], dim=1) t_disease = torch.cat([t_disease, other_time], dim=1)
@@ -433,10 +428,18 @@ class DeepHealth(nn.Module):
) )
return hidden return hidden
if return_output: if return_output:
h_event = h_disease[:, :event_len, :]
t_event = t_disease[:, :event_len]
event_mask = padding_mask[:, :event_len]
h_extra, t_extra, extra_mask = self._pool_other_by_time(
h_other=h_disease[:, event_len:, :],
other_time=t_disease[:, event_len:],
other_mask=padding_mask[:, event_len:],
)
return DeepHealthOutput( return DeepHealthOutput(
hidden=h_disease, hidden=torch.cat([h_event, h_extra], dim=1),
time_seq=t_disease, time_seq=torch.cat([t_event, t_extra], dim=1),
padding_mask=padding_mask, padding_mask=torch.cat([event_mask, extra_mask], dim=1),
event_len=event_len, event_len=event_len,
) )
return h_disease[:, :event_len, :] return h_disease[:, :event_len, :]