From 1ea72e9133e80c9039dbefa5e5be2c0c258157bf Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 17 Jun 2026 14:27:07 +0800 Subject: [PATCH] Refactor DeepHealth model to expand extra-info token handling and update output structure --- README.md | 26 +++++++++++++------------- models.py | 21 ++++++++++++--------- 2 files changed, 25 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index c0f47bf..17081f6 100644 --- a/README.md +++ b/README.md @@ -135,20 +135,20 @@ type_id -> continuous type index -> soft bin embedding 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 --extra_pool_reduce mean ``` -也可以设为 `sum`。池化后，每个 extra-info 时间点最多产生一个主序列 token。 +也可以设为 `sum`。池化后，每个 extra-info 时间点最多产生一个用于预测和 loss 的 readout token。 ### DeepHealth `DeepHealth` 的统一路径是： ```text -disease tokens + pooled extra-info tokens +disease tokens + expanded extra-info tokens -> temporal backbone -> risk head ``` @@ -158,23 +158,23 @@ extra-info 不再通过独立的 `BaselineEncoder` 或 `CrossAttention` 注入 两种目标模式的读出语义不同： - `next_token` - - 模型内部序列包含 disease tokens 和 pooled extra-info tokens + - 模型内部序列包含 disease tokens 和展开的单个 extra-info tokens - 默认 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 - `all_future` - 在合并后的主序列末尾拼接一个 query token - query token 的时间是 `t_query` - 只读出最后一个 query hidden，形状为 `(B, D)` - - pooled extra-info tokens 只作为 query 的上下文输入，不单独读出、不参与 loss 监督 + - 展开的 extra-info tokens 只作为 query 的上下文输入，不单独读出、不参与 loss 监督 - all-future 模式不使用 weight tying 如果需要拿到完整 next-token 输出，可使用结构化返回： ```python 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.padding_mask # 与 hidden 对齐的有效位置 out.event_len # 原 disease token 长度 @@ -209,9 +209,9 @@ next-token 监督： next-token 训练中，模型会请求 `return_output=True`，因此 loss 的预测位置包括： - 原 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 监督： @@ -219,7 +219,7 @@ all-future / query-conditioned 监督： - `WeibullLoss` - `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` 参数。 @@ -232,12 +232,12 @@ all-future 训练只读出 `t_query` 对应的 query hidden。pooled extra-info - `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout - `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout - 当前 next-token 训练只支持 exponential time loss - - pooled extra-info tokens 会加入 prediction/loss 监督 + - 展开的 extra-info tokens 进入主序列；读出端 pooled extra-info tokens 会加入 prediction/loss 监督 - `train_all_future.py` - 使用 `AllFutureHealthDataset` - 不使用 readout，直接对 query hidden 计算风险 - `--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` 支持所有已有训练目标定义的组合： @@ -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 计算、病例/对照筛选和分层聚合逻辑保持原评估脚本口径。 -评估脚本默认使用常规 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` diff --git a/models.py b/models.py index 8521849..6d409da 100644 --- a/models.py +++ b/models.py @@ -322,7 +322,7 @@ class DeepHealth(nn.Module): other_time: torch.FloatTensor | None = None, return_output: bool = False, **unused_kwargs, - ) -> torch.Tensor: + ) -> torch.Tensor | DeepHealthOutput: if unused_kwargs: unknown = ", ".join(sorted(unused_kwargs)) 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) 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) t_disease = torch.cat([t_disease, other_time], dim=1) @@ -433,10 +428,18 @@ class DeepHealth(nn.Module): ) return hidden 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( - hidden=h_disease, - time_seq=t_disease, - padding_mask=padding_mask, + hidden=torch.cat([h_event, h_extra], dim=1), + time_seq=torch.cat([t_event, t_extra], dim=1), + padding_mask=torch.cat([event_mask, extra_mask], dim=1), event_len=event_len, ) return h_disease[:, :event_len, :]