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

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`
 

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, :]