Refactor DeepHealth model and related components
- Removed BaselineEncoder and CrossAttention classes from models.py. - Introduced OtherInfoTokenizer for handling additional token types. - Updated DeepHealth class to integrate OtherInfoTokenizer and manage extra pooling logic. - Added support for extra_pool_reduce parameter to control pooling behavior. - Modified forward methods to return structured output using DeepHealthOutput dataclass. - Updated training scripts to accommodate changes in model architecture and output handling. - Enhanced error handling and validation for input shapes and types.
This commit is contained in:
107
README.md
107
README.md
@@ -112,18 +112,17 @@ dataset.vocab_size
|
||||
|
||||
模型主体定义在 `models.py`，通用网络模块定义在 `backbones.py`。
|
||||
|
||||
### BaselineEncoder
|
||||
### OtherInfoTokenizer
|
||||
|
||||
`BaselineEncoder` 编码统一的 other-info token：
|
||||
`DeepHealth` 内部的 `OtherInfoTokenizer` 编码统一的 other-info token：
|
||||
|
||||
```python
|
||||
other_type # (B, K)
|
||||
other_value # (B, K)
|
||||
other_value_kind # (B, K)
|
||||
other_time # (B, K)
|
||||
```
|
||||
|
||||
它暂时不直接使用 `other_time`。时间信息保留给后续 `CrossAttention`，用于建模疾病/query 与 other-info token 的相对时间关系。
|
||||
|
||||
连续值使用 `TokenAutoDiscretization`：
|
||||
|
||||
```text
|
||||
@@ -136,57 +135,50 @@ type_id -> continuous type index -> soft bin embedding
|
||||
selected type offsets + local category id
|
||||
```
|
||||
|
||||
### CrossAttention
|
||||
同一患者、同一 `other_time` 的 extra-info token 会先被池化为一个 token，再并入主序列。默认池化方式是平均池化：
|
||||
|
||||
`CrossAttention` 让 disease-side hidden state 注意到 other-info token：
|
||||
|
||||
```python
|
||||
h_disease # (B, L, D)
|
||||
t_disease # (B, L)
|
||||
h_token # (B, K, D)
|
||||
t_token # (B, K)
|
||||
```bash
|
||||
--extra_pool_reduce mean
|
||||
```
|
||||
|
||||
时间信息通过两种方式进入注意力：
|
||||
|
||||
- `TimeRoPE`
|
||||
- 使用 query time 和 key time 旋转 q/k
|
||||
- 让 q-k 相似度带有时间位置信息
|
||||
|
||||
- `GaussianRBFTimeBasis`
|
||||
- 对 `t_disease - t_token` 做 RBF 编码
|
||||
- 投影成每个 attention head 的时间 bias
|
||||
|
||||
注意力是时间因果的：
|
||||
|
||||
```text
|
||||
other_info_time <= disease_or_query_time
|
||||
```
|
||||
|
||||
如果某个 disease/query 位置没有任何可见 other-info token，则该位置保持原 hidden 不变。
|
||||
也可以设为 `sum`。池化后，每个 extra-info 时间点最多产生一个主序列 token。
|
||||
|
||||
### DeepHealth
|
||||
|
||||
`DeepHealth` 的统一路径是：
|
||||
|
||||
```text
|
||||
disease-side sequence
|
||||
-> disease temporal backbone
|
||||
-> CrossAttention 到 other-info tokens
|
||||
disease tokens + pooled extra-info tokens
|
||||
-> temporal backbone
|
||||
-> risk head
|
||||
```
|
||||
|
||||
两种目标模式共用同一套语义：
|
||||
extra-info 不再通过独立的 `BaselineEncoder` 或 `CrossAttention` 注入；它们作为主序列 token 直接参与同一个 temporal transformer。相对时间模式下，主序列内所有 token 共用 `TimeRoPE` 和 `GaussianRBFTimeBasis`。
|
||||
|
||||
两种目标模式的读出语义不同：
|
||||
|
||||
- `next_token`
|
||||
- `h_disease` 长度为 `L`
|
||||
- 输出 `(B, L, D)`
|
||||
- 模型内部序列包含 disease tokens 和 pooled extra-info tokens
|
||||
- 默认 Tensor 返回值仍只返回 disease token hidden，形状为 `(B, L, D)`，兼容评估和旧调用
|
||||
- 训练时使用 `return_output=True` 取完整主序列输出；pooled extra-info tokens 也会产生 logits，并在有未来 disease target 时参与 prediction/loss 监督
|
||||
- next-token 模式下 `risk_head.weight` 与 `token_embedding.weight` 使用 weight tying
|
||||
|
||||
- `all_future`
|
||||
- 在 disease-side 序列末尾拼接一个 query token
|
||||
- 在合并后的主序列末尾拼接一个 query token
|
||||
- query token 的时间是 `t_query`
|
||||
- `h_disease` 长度为 `L + 1`
|
||||
- 输出最后一个 query hidden，形状为 `(B, D)`
|
||||
- 只读出最后一个 query hidden，形状为 `(B, D)`
|
||||
- pooled 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.time_seq # 与 hidden 对齐的时间
|
||||
out.padding_mask # 与 hidden 对齐的有效位置
|
||||
out.event_len # 原 disease token 长度
|
||||
```
|
||||
|
||||
模型初始化示例：
|
||||
|
||||
@@ -196,11 +188,12 @@ model = DeepHealth(
|
||||
n_embd=120,
|
||||
n_head=10,
|
||||
n_hist_layer=12,
|
||||
n_tab_layer=4,
|
||||
n_tab_layer=4, # 兼容旧配置；当前不再创建独立 tabular transformer
|
||||
n_types=dataset.n_types,
|
||||
n_cont_types=dataset.n_cont_types,
|
||||
n_categories=dataset.n_categories,
|
||||
cont_type_ids=dataset.cont_type_ids,
|
||||
extra_pool_reduce="mean",
|
||||
)
|
||||
```
|
||||
|
||||
@@ -213,27 +206,38 @@ next-token 监督：
|
||||
- `Delphi2MLoss`
|
||||
- `UniqueTimeSetExponentialLoss`
|
||||
|
||||
next-token 训练中，模型会请求 `return_output=True`，因此 loss 的预测位置包括：
|
||||
|
||||
- 原 disease token readout 位置
|
||||
- 同一时间点 extra-info 池化后的 pooled extra-info 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。
|
||||
|
||||
all-future / query-conditioned 监督：
|
||||
|
||||
- `ExponentialLoss`
|
||||
- `WeibullLoss`
|
||||
- `MixedLoss`
|
||||
|
||||
all-future 训练只读出 `t_query` 对应的 query hidden。pooled extra-info tokens 作为主序列上下文输入，但不会被单独读出，也不会被纳入 loss 监督。
|
||||
|
||||
`UniqueTimeSetExponentialLoss` 的 observed term 固定使用 sum reduction，不再暴露旧的 `observed_reduction` 参数。
|
||||
|
||||
## 训练
|
||||
|
||||
当前 `train_next_step.py` / `train_all_future.py` 支持 next-token 和 all-future 两类训练入口：
|
||||
当前提供两类训练入口：
|
||||
|
||||
- `--model_target_mode next_token`
|
||||
- `train_next_step.py`
|
||||
- 使用 `NextStepHealthDataset`
|
||||
- `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout
|
||||
- `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout
|
||||
- 当前 next-token 训练只支持 `--dist_mode exponential`
|
||||
- `--model_target_mode all_future`
|
||||
- 当前 next-token 训练只支持 exponential time loss
|
||||
- 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 上下文，不单独监督
|
||||
|
||||
当前 `train_next_step.py` / `train_all_future.py` 支持所有已有训练目标定义的组合：
|
||||
|
||||
@@ -251,23 +255,23 @@ all-future / query-conditioned 监督：
|
||||
python train_next_step.py \
|
||||
--data_prefix ukb \
|
||||
--labels_file labels.csv \
|
||||
--model_target_mode next_token \
|
||||
--target_mode uts \
|
||||
--n_embd 120 \
|
||||
--n_head 10 \
|
||||
--n_hist_layer 12 \
|
||||
--n_tab_layer 4
|
||||
--n_tab_layer 4 \
|
||||
--extra_pool_reduce mean
|
||||
```
|
||||
|
||||
all-future 示例：
|
||||
|
||||
```bash
|
||||
python train_next_step.py \
|
||||
python train_all_future.py \
|
||||
--data_prefix ukb \
|
||||
--labels_file labels.csv \
|
||||
--model_target_mode all_future \
|
||||
--dist_mode weibull \
|
||||
--time_mode relative
|
||||
--time_mode relative \
|
||||
--extra_pool_reduce mean
|
||||
```
|
||||
|
||||
选择额外信息变量：
|
||||
@@ -282,12 +286,15 @@ python train_next_step.py --extra_info_types_file extra_info_types_smoking_alcoh
|
||||
|
||||
- `extra_info_types_file`：训练时使用的列表文件名
|
||||
- `extra_info_types`：解析后的实际 type id 列表，用于评估脚本复现变量选择
|
||||
- `extra_pool_reduce`：同一 `other_time` 的 extra-info tokens 池化方式，默认为 `mean`
|
||||
- `model_target_mode`、`time_mode`、`dist_mode`、`dataset_class`、`collate_fn`、`resolved_loss_name`：用于评估脚本重建模型和输入方式
|
||||
|
||||
## 评估 AUC
|
||||
|
||||
当前提供两个 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 评估位置单独输出。
|
||||
|
||||
### `evaluate_auc.py`
|
||||
|
||||
`evaluate_auc.py` 评估的是 **next-step / token-level 预测位置上的疾病 AUC**。
|
||||
@@ -403,6 +410,8 @@ python evaluate_auc_v2.py \
|
||||
|
||||
- `models.py`
|
||||
- `DeepHealth`
|
||||
- `DeepHealthOutput`
|
||||
- `OtherInfoTokenizer`
|
||||
|
||||
- `backbones.py`
|
||||
- `TimeRoPE`
|
||||
@@ -410,8 +419,6 @@ python evaluate_auc_v2.py \
|
||||
- `TemporalAttention`
|
||||
- `GPTBlock`
|
||||
- `TokenAutoDiscretization`
|
||||
- `BaselineEncoder`
|
||||
- `CrossAttention`
|
||||
- `AgeSinusoidalEncoding`
|
||||
|
||||
- `losses.py`
|
||||
|
||||
Reference in New Issue
Block a user