From 1757bcd25bb9eea8e4ae91468dd9260fa8de0819 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 17 Jun 2026 11:05:10 +0800 Subject: [PATCH] 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. --- README.md | 107 +++++++------- backbones.py | 352 +------------------------------------------- evaluate_auc.py | 1 + evaluate_auc_v2.py | 1 + models.py | 304 +++++++++++++++++++++++++++----------- train_all_future.py | 3 + train_next_step.py | 160 ++++++++++++++++++-- 7 files changed, 435 insertions(+), 493 deletions(-) diff --git a/README.md b/README.md index 143942d..c0f47bf 100644 --- a/README.md +++ b/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` diff --git a/backbones.py b/backbones.py index 899e0bf..ea61077 100644 --- a/backbones.py +++ b/backbones.py @@ -10,7 +10,7 @@ class TimeRoPE(nn.Module): super().__init__() assert dim % 2 == 0, "RoPE dim must be even" self.dim = dim - # inv_freq: (dim // 2,) — not trainable, but should move with device + # inv_freq is not trainable, but should move with device. inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) @@ -73,7 +73,7 @@ class GaussianRBFTimeBasis(nn.Module): def precompute_cache(self, tau: torch.Tensor) -> torch.Tensor: time_coord = tau.float() # (B, L) - # Pairwise signed difference: query_i − key_j + # Pairwise signed difference: query_i - key_j. diff = time_coord.unsqueeze( 2) - time_coord.unsqueeze(1) # (B, L_q, L_k) # Gaussian RBF: exp(-0.5 * ((diff - c) / w)^2) @@ -297,354 +297,6 @@ class TokenAutoDiscretization(nn.Module): return torch.einsum("nb,nbd->nd", probs, e) -class BaselineEncoder(nn.Module): - PAD_KIND = 0 - CONT_KIND = 1 - CATE_KIND = 2 - - def __init__( - self, - n_embd: int, - n_head: int, - n_types: int, - n_cont_types: int, - n_categories: int, - cont_type_ids: list[int], - n_value_kinds: int = 3, - n_bins: int = 16, - n_tab_layer: int = 2, - dropout: float = 0.0, - ): - super().__init__() - if len(cont_type_ids) != n_cont_types: - raise ValueError( - "cont_type_ids length must match n_cont_types, got " - f"{len(cont_type_ids)} vs {n_cont_types}" - ) - if n_types <= 0: - raise ValueError(f"n_types must include PAD and be > 0, got {n_types}") - if n_categories <= 0: - raise ValueError( - f"n_categories must include PAD and be > 0, got {n_categories}" - ) - if n_value_kinds <= self.CATE_KIND: - raise ValueError( - f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}" - ) - - self.n_embd = n_embd - self.cls_token = nn.Parameter(torch.zeros(1, 1, n_embd)) - self.type_emb = nn.Embedding(n_types, n_embd, padding_idx=0) - self.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0) - self.cont_value_encoder = ( - TokenAutoDiscretization( - n_cont_types=n_cont_types, - n_bins=n_bins, - n_embd=n_embd, - ) - if n_cont_types > 0 - else None - ) - self.cate_value_emb = nn.Embedding( - n_categories, - n_embd, - padding_idx=0, - ) - - cont_type_index = torch.full((n_types,), -1, dtype=torch.long) - for idx, type_id in enumerate(cont_type_ids): - if type_id <= 0 or type_id >= n_types: - raise ValueError( - f"continuous type id {type_id} must be in [1, {n_types})" - ) - cont_type_index[type_id] = idx - self.register_buffer( - "cont_type_index", - cont_type_index, - persistent=False, - ) - - self.blocks = nn.ModuleList([ - GPTBlock( - n_embd=n_embd, - n_head=n_head, - use_time_rope=False, - use_rbf_bias=False, - mlp_dropout=dropout, - ) for _ in range(n_tab_layer) - ]) - self.ln = nn.LayerNorm(n_embd) - self.reset_parameters() - - def reset_parameters(self) -> None: - nn.init.normal_(self.cls_token, mean=0.0, std=0.02) - nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02) - nn.init.zeros_(self.type_emb.weight[0]) - nn.init.normal_(self.kind_emb.weight, mean=0.0, std=0.02) - nn.init.zeros_(self.kind_emb.weight[0]) - nn.init.normal_(self.cate_value_emb.weight, mean=0.0, std=0.02) - nn.init.zeros_(self.cate_value_emb.weight[0]) - - def _make_attn_mask(self, mask: torch.Tensor, dtype: torch.dtype): - return torch.zeros( - mask.size(0), - 1, - 1, - mask.size(1), - device=mask.device, - dtype=dtype, - ).masked_fill(~mask[:, None, None, :], -1e4) - - def _make_time_attn_mask( - self, - mask: torch.Tensor, - time: torch.Tensor, - dtype: torch.dtype, - ): - valid_key = mask[:, None, :] - visible_by_time = time[:, None, :] <= time[:, :, None] - valid = valid_key & visible_by_time - return torch.zeros( - valid.shape, - device=valid.device, - dtype=dtype, - ).masked_fill(~valid, -1e4)[:, None, :, :] - - def forward( - self, - other_type: torch.LongTensor, # (B, K), 0 = padding - other_value: torch.Tensor, # (B, K), cate stores global id - other_value_kind: torch.LongTensor, # (B, K), 0=PAD, 1=CONT, 2=CATE - other_time: torch.Tensor | None = None, # (B, K) - cls_time: torch.Tensor | None = None, # (B,) - ): - if other_type.shape != other_value.shape: - raise ValueError( - "other_type and other_value must have the same shape, got " - f"{tuple(other_type.shape)} vs {tuple(other_value.shape)}" - ) - if other_type.shape != other_value_kind.shape: - raise ValueError( - "other_type and other_value_kind must have the same shape, got " - f"{tuple(other_type.shape)} vs {tuple(other_value_kind.shape)}" - ) - - other_valid = other_type > 0 - type_emb = self.type_emb(other_type) - kind_emb = self.kind_emb(other_value_kind) - value_emb = torch.zeros_like(type_emb) - - cont_pos = other_valid & (other_value_kind == self.CONT_KIND) - if cont_pos.any(): - if self.cont_value_encoder is None: - raise ValueError("continuous tokens found but n_cont_types is 0") - cont_idx = self.cont_type_index[other_type[cont_pos]] - if (cont_idx < 0).any(): - bad_type = other_type[cont_pos][cont_idx < 0][0].item() - raise ValueError( - f"type_id={bad_type} is marked continuous but is not in " - "cont_type_ids" - ) - value_emb[cont_pos] = self.cont_value_encoder( - cont_type_idx=cont_idx, - value=other_value[cont_pos].to(type_emb.dtype), - ) - - cate_pos = other_valid & (other_value_kind == self.CATE_KIND) - if cate_pos.any(): - cate_id = other_value[cate_pos].long() - value_emb[cate_pos] = self.cate_value_emb(cate_id) - - f = type_emb + kind_emb + value_emb - f = f * other_valid.unsqueeze(-1).to(f.dtype) - - cls = self.cls_token.expand(f.size(0), -1, -1) - f = torch.cat([cls, f], dim=1) - cls_valid = torch.ones( - other_valid.size(0), - 1, - device=other_valid.device, - dtype=torch.bool, - ) - full_valid = torch.cat([cls_valid, other_valid], dim=1) - - if other_time is None: - attn_mask = self._make_attn_mask(full_valid, f.dtype) - else: - if other_time.shape != other_type.shape: - raise ValueError( - "other_time must have the same shape as other_type, got " - f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}" - ) - if cls_time is None: - raise ValueError("cls_time is required when other_time is provided") - full_time = torch.cat( - [ - cls_time.to(device=other_time.device, dtype=other_time.dtype)[:, None], - other_time, - ], - dim=1, - ) - attn_mask = self._make_time_attn_mask(full_valid, full_time, f.dtype) - for block in self.blocks: - f = block(f, attn_mask=attn_mask) - f = f * full_valid.unsqueeze(-1).to(f.dtype) - - h = self.ln(f) - h = h * full_valid.unsqueeze(-1).to(h.dtype) - cls_summary = h[:, 0, :] - token_h = h[:, 1:, :] - token_h = token_h * other_valid.unsqueeze(-1).to(token_h.dtype) - return token_h, other_valid, cls_summary - - -class CrossAttention(nn.Module): - def __init__( - self, - n_embd: int, - n_head: int, - dropout: float = 0.0, - n_rbf_bases: int = 16, - max_time_diff: float = 40.0, - ): - super().__init__() - assert n_embd % n_head == 0, "n_embd must be divisible by n_head" - self.n_head = n_head - self.d_head = n_embd // n_head - self.scale = 1.0 / math.sqrt(self.d_head) - self.dropout = dropout - self.mask_value = -1e4 - - self.q_proj = nn.Linear(n_embd, n_embd, bias=False) - self.k_proj = nn.Linear(n_embd, n_embd, bias=False) - self.v_proj = nn.Linear(n_embd, n_embd, bias=False) - self.out_proj = nn.Linear(n_embd, n_embd, bias=False) - - self.time_rope = TimeRoPE(self.d_head) - self.rbf_time_basis = GaussianRBFTimeBasis( - n_bases=n_rbf_bases, - max_time_diff=max_time_diff, - ) - self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False) - self.time_bias_scale = nn.Parameter(torch.tensor(0.0)) - self.resid_drop = nn.Dropout(dropout) - self.ln = nn.LayerNorm(n_embd) - self.out_ln = nn.LayerNorm(n_embd) - self.reset_parameters() - - def reset_parameters(self) -> None: - nn.init.normal_(self.q_proj.weight, mean=0.0, std=0.02) - nn.init.normal_(self.k_proj.weight, mean=0.0, std=0.02) - nn.init.normal_(self.v_proj.weight, mean=0.0, std=0.02) - nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02) - nn.init.normal_(self.rbf_proj.weight, mean=0.0, std=0.02) - - def _make_attn_mask( - self, - token_mask: torch.Tensor, # (B, K), True = valid - t_disease: torch.Tensor, # (B, L) - t_token: torch.Tensor, # (B, K) - dtype: torch.dtype, - ): - valid_token = token_mask[:, None, :] # (B, 1, K) - visible_by_time = t_token[:, None, :] <= t_disease[:, :, None] - valid = visible_by_time & valid_token # (B, L, K) - - has_any_valid = valid.any(dim=-1, keepdim=True) - safe_valid = valid.clone() - safe_valid[..., 0:1] = torch.where( - has_any_valid, - safe_valid[..., 0:1], - torch.ones_like(safe_valid[..., 0:1]), - ) - - attn_mask = torch.zeros( - safe_valid.shape, - device=safe_valid.device, - dtype=dtype, - ).masked_fill(~safe_valid, self.mask_value) - return attn_mask[:, None, :, :], valid - - def _cross_rbf_cache( - self, - t_disease: torch.Tensor, - t_token: torch.Tensor, - ) -> torch.Tensor: - tau = torch.cat([t_disease, t_token], dim=1) - rbf_cache = self.rbf_time_basis.precompute_cache(tau) - n_disease = t_disease.size(1) - return rbf_cache[:, :n_disease, n_disease:, :] - - def forward( - self, - h_disease: torch.Tensor, # (B, L, D) - t_disease: torch.Tensor, # (B, L) - h_token: torch.Tensor, # (B, K, D) - t_token: torch.Tensor, # (B, K) - token_mask: torch.Tensor, # (B, K), True = valid - need_weights: bool = False, - ): - B, L, _ = h_disease.shape - K = h_token.size(1) - H, Dh = self.n_head, self.d_head - if K == 0: - empty_weights = h_disease.new_zeros(B, L, 0) - if need_weights: - return h_disease, empty_weights - return h_disease - - attn_mask, valid = self._make_attn_mask( - token_mask=token_mask.to(device=h_disease.device, dtype=torch.bool), - t_disease=t_disease, - t_token=t_token, - dtype=h_disease.dtype, - ) - - q = self.q_proj(h_disease).reshape(B, L, H, Dh).transpose(1, 2) - k = self.k_proj(h_token).reshape(B, K, H, Dh).transpose(1, 2) - v = self.v_proj(h_token).reshape(B, K, H, Dh).transpose(1, 2) - - q_cos, q_sin = self.time_rope.precompute_cache(t_disease) - k_cos, k_sin = self.time_rope.precompute_cache(t_token) - q = q * q_cos + TimeRoPE._rotate_half(q) * q_sin - k = k * k_cos + TimeRoPE._rotate_half(k) * k_sin - - rbf_cache = self._cross_rbf_cache(t_disease, t_token) # (B, L, K, R) - time_bias = self.rbf_proj(rbf_cache).permute(0, 3, 1, 2) - time_bias = self.time_bias_scale.tanh() * time_bias - - attn_bias = attn_mask.to(time_bias.dtype) + time_bias - attn_out = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=attn_bias, - dropout_p=self.dropout if self.training else 0.0, - is_causal=False, - scale=self.scale, - ) - attn_out = attn_out.transpose(1, 2).reshape(B, L, H * Dh) - attn_out = self.resid_drop(self.out_proj(attn_out)) - - has_any_valid = valid.any(dim=-1) # (B, L) - attn_out = attn_out * has_any_valid.unsqueeze(-1).to(attn_out.dtype) - - attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale - attn_scores = attn_scores + attn_bias - attn_weights = torch.softmax(attn_scores, dim=-1).mean(dim=1) - attn_weights = attn_weights * valid.to(attn_weights.dtype) - weight_sum = attn_weights.sum(dim=-1, keepdim=True).clamp_min(1e-12) - norm_weights = attn_weights / weight_sum - norm_weights = norm_weights * has_any_valid.unsqueeze(-1).to( - norm_weights.dtype - ) - out = self.out_ln(h_disease + attn_out) - out = torch.where(has_any_valid.unsqueeze(-1), out, h_disease) - - if need_weights: - return out, norm_weights - return out - class AgeSinusoidalEncoding(nn.Module): diff --git a/evaluate_auc.py b/evaluate_auc.py index 2dd4bb6..6a3fbe1 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -327,6 +327,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data n_categories=dataset.n_categories, cont_type_ids=dataset.cont_type_ids, n_bins=int(cfg_get(args, cfg, "n_bins", 16)), + extra_pool_reduce=str(cfg_get(args, cfg, "extra_pool_reduce", "mean")), target_mode=model_target_mode, time_mode=str(cfg_get(args, cfg, "time_mode", "relative")), dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index 1bc5f68..507b98f 100644 --- a/evaluate_auc_v2.py +++ b/evaluate_auc_v2.py @@ -184,6 +184,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data n_categories=dataset.n_categories, cont_type_ids=dataset.cont_type_ids, n_bins=int(cfg_get(args, cfg, "n_bins", 16)), + extra_pool_reduce=str(cfg_get(args, cfg, "extra_pool_reduce", "mean")), target_mode=model_target_mode, time_mode=str(cfg_get(args, cfg, "time_mode", "relative")), dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), diff --git a/models.py b/models.py index bbac0c3..8521849 100644 --- a/models.py +++ b/models.py @@ -1,16 +1,144 @@ +from dataclasses import dataclass + import torch import torch.nn as nn import torch.nn.functional as F from backbones import ( AgeSinusoidalEncoding, - BaselineEncoder, - CrossAttention, GPTBlock, GaussianRBFTimeBasis, TimeRoPE, + TokenAutoDiscretization, ) -from targets import CHECKUP_IDX, PAD_IDX +from targets import PAD_IDX + + +@dataclass +class DeepHealthOutput: + hidden: torch.Tensor + time_seq: torch.Tensor + padding_mask: torch.Tensor + event_len: int + + +class OtherInfoTokenizer(nn.Module): + PAD_KIND = 0 + CONT_KIND = 1 + CATE_KIND = 2 + + def __init__( + self, + n_embd: int, + n_types: int, + n_cont_types: int, + n_categories: int, + cont_type_ids: list[int], + n_value_kinds: int = 3, + n_bins: int = 16, + ): + super().__init__() + if len(cont_type_ids) != n_cont_types: + raise ValueError( + "cont_type_ids length must match n_cont_types, got " + f"{len(cont_type_ids)} vs {n_cont_types}" + ) + if n_types <= 0: + raise ValueError(f"n_types must include PAD and be > 0, got {n_types}") + if n_categories <= 0: + raise ValueError( + f"n_categories must include PAD and be > 0, got {n_categories}" + ) + if n_value_kinds <= self.CATE_KIND: + raise ValueError( + f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}" + ) + + self.type_emb = nn.Embedding(n_types, n_embd, padding_idx=0) + self.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0) + self.cont_value_encoder = ( + TokenAutoDiscretization( + n_cont_types=n_cont_types, + n_bins=n_bins, + n_embd=n_embd, + ) + if n_cont_types > 0 + else None + ) + self.cate_value_emb = nn.Embedding( + n_categories, + n_embd, + padding_idx=0, + ) + + cont_type_index = torch.full((n_types,), -1, dtype=torch.long) + for idx, type_id in enumerate(cont_type_ids): + if type_id <= 0 or type_id >= n_types: + raise ValueError( + f"continuous type id {type_id} must be in [1, {n_types})" + ) + cont_type_index[type_id] = idx + self.register_buffer( + "cont_type_index", + cont_type_index, + persistent=False, + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.type_emb.weight[0]) + nn.init.normal_(self.kind_emb.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.kind_emb.weight[0]) + nn.init.normal_(self.cate_value_emb.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.cate_value_emb.weight[0]) + + def forward( + self, + other_type: torch.LongTensor, + other_value: torch.Tensor, + other_value_kind: torch.LongTensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if other_type.shape != other_value.shape: + raise ValueError( + "other_type and other_value must have the same shape, got " + f"{tuple(other_type.shape)} vs {tuple(other_value.shape)}" + ) + if other_type.shape != other_value_kind.shape: + raise ValueError( + "other_type and other_value_kind must have the same shape, got " + f"{tuple(other_type.shape)} vs {tuple(other_value_kind.shape)}" + ) + + other_valid = other_type > 0 + type_emb = self.type_emb(other_type) + kind_emb = self.kind_emb(other_value_kind) + value_emb = torch.zeros_like(type_emb) + + cont_pos = other_valid & (other_value_kind == self.CONT_KIND) + if cont_pos.any(): + if self.cont_value_encoder is None: + raise ValueError("continuous tokens found but n_cont_types is 0") + cont_idx = self.cont_type_index[other_type[cont_pos]] + if (cont_idx < 0).any(): + bad_type = other_type[cont_pos][cont_idx < 0][0].item() + raise ValueError( + f"type_id={bad_type} is marked continuous but is not in " + "cont_type_ids" + ) + value_emb[cont_pos] = self.cont_value_encoder( + cont_type_idx=cont_idx, + value=other_value[cont_pos].to(type_emb.dtype), + ) + + cate_pos = other_valid & (other_value_kind == self.CATE_KIND) + if cate_pos.any(): + cate_id = other_value[cate_pos].long() + value_emb[cate_pos] = self.cate_value_emb(cate_id) + + out = type_emb + kind_emb + value_emb + out = out * other_valid.unsqueeze(-1).to(out.dtype) + return out, other_valid class DeepHealth(nn.Module): @@ -30,6 +158,7 @@ class DeepHealth(nn.Module): target_mode: str = "next_token", # "next_token" or "all_future" time_mode: str = "relative", # "relative" or "absolute" dist_mode: str = "exponential", # "exponential", "weibull" or "mixed" + extra_pool_reduce: str = "mean", dropout: float = 0.0, ): super().__init__() @@ -42,30 +171,24 @@ class DeepHealth(nn.Module): if dist_mode not in ["exponential", "weibull", "mixed"]: raise ValueError( "dist_mode must be either 'exponential', 'weibull' or 'mixed'") + if extra_pool_reduce not in {"mean", "sum"}: + raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'") self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0) self.gender_embedding = nn.Embedding( 2, n_embd) # Assuming binary gender - self.token_encoder = BaselineEncoder( + self.tokenizer = OtherInfoTokenizer( n_embd=n_embd, - n_head=n_head, n_types=n_types, n_cont_types=n_cont_types, n_categories=n_categories, cont_type_ids=cont_type_ids, n_value_kinds=n_value_kinds, n_bins=n_bins, - n_tab_layer=n_tab_layer, - dropout=dropout, - ) - self.cross_attention = CrossAttention( - n_embd=n_embd, - n_head=n_head, - dropout=dropout, - n_rbf_bases=16, ) self.target_mode = target_mode self.time_mode = time_mode self.dist_mode = dist_mode + self.extra_pool_reduce = extra_pool_reduce self.n_embd = n_embd self.vocab_size = vocab_size nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02) @@ -111,6 +234,8 @@ class DeepHealth(nn.Module): self.final_ln = nn.LayerNorm(n_embd) self.risk_head = nn.Linear(n_embd, vocab_size, bias=False) + if target_mode == "next_token": + self.risk_head.weight = self.token_embedding.weight self.query_token = nn.Parameter(torch.zeros(n_embd)) nn.init.normal_(self.query_token, mean=0.0, std=0.02) @@ -129,55 +254,59 @@ class DeepHealth(nn.Module): dtype=dtype, ).masked_fill(~valid, -1e4)[:, None, :, :] - def _insert_baseline_summary( + def _pool_other_by_time( self, - h_disease: torch.Tensor, - event_seq: torch.Tensor, - baseline_summary: torch.Tensor, - ) -> torch.Tensor: - checkup_mask = event_seq == CHECKUP_IDX - if not checkup_mask.any(): - return h_disease - summary = baseline_summary.to(device=h_disease.device, dtype=h_disease.dtype) - return torch.where(checkup_mask.unsqueeze(-1), summary[:, None, :], h_disease) - - def _baseline_cls_time( - self, - event_seq: torch.Tensor, - time_seq: torch.Tensor, - padding_mask: torch.Tensor, - ) -> torch.Tensor: - checkup_mask = event_seq == CHECKUP_IDX - inf = torch.full_like(time_seq, float("inf")) - first_checkup = torch.where(checkup_mask, time_seq, inf).min(dim=1).values - has_checkup = torch.isfinite(first_checkup) - fallback_time = torch.where( - padding_mask, - time_seq, - torch.full_like(time_seq, float("-inf")), - ).max(dim=1).values - fallback_time = torch.where( - torch.isfinite(fallback_time), - fallback_time, - torch.zeros_like(fallback_time), - ) - return torch.where(has_checkup, first_checkup, fallback_time) - - def _encode_other_tokens( - self, - other_type: torch.LongTensor, - other_value: torch.Tensor, - other_value_kind: torch.LongTensor, + h_other: torch.Tensor, other_time: torch.Tensor, - cls_time: torch.Tensor, + other_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return self.token_encoder( - other_type=other_type, - other_value=other_value, - other_value_kind=other_value_kind, - other_time=other_time, - cls_time=cls_time, + batch_size, _n_other, n_embd = h_other.shape + group_counts = [ + int(torch.unique(other_time[b, other_mask[b]], sorted=True).numel()) + for b in range(batch_size) + ] + max_groups = max(group_counts, default=0) + pooled_h = h_other.new_zeros(batch_size, max_groups, n_embd) + pooled_time = other_time.new_zeros(batch_size, max_groups) + pooled_mask = torch.zeros( + batch_size, + max_groups, + dtype=torch.bool, + device=h_other.device, ) + if max_groups == 0: + return pooled_h, pooled_time, pooled_mask + + for b in range(batch_size): + valid_time = other_time[b, other_mask[b]] + if valid_time.numel() == 0: + continue + valid_h = h_other[b, other_mask[b]] + unique_time, inverse = torch.unique( + valid_time, + sorted=True, + return_inverse=True, + ) + n_groups = unique_time.numel() + group_h = valid_h.new_zeros(n_groups, n_embd) + group_h.scatter_add_( + 0, + inverse[:, None].expand(-1, n_embd), + valid_h, + ) + if self.extra_pool_reduce == "mean": + counts = valid_h.new_zeros(n_groups, 1) + counts.scatter_add_( + 0, + inverse[:, None], + torch.ones_like(valid_h[:, :1]), + ) + group_h = group_h / counts.clamp_min(1.0) + + pooled_h[b, :n_groups] = group_h + pooled_time[b, :n_groups] = unique_time + pooled_mask[b, :n_groups] = True + return pooled_h, pooled_time, pooled_mask def _forward_shared( self, @@ -191,6 +320,7 @@ class DeepHealth(nn.Module): other_value: torch.Tensor | None = None, other_value_kind: torch.LongTensor | None = None, other_time: torch.FloatTensor | None = None, + return_output: bool = False, **unused_kwargs, ) -> torch.Tensor: if unused_kwargs: @@ -216,6 +346,7 @@ class DeepHealth(nn.Module): else: padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool) + event_len = event_seq.size(1) h_disease = self.token_embedding(event_seq) t_disease = time_seq @@ -225,40 +356,29 @@ class DeepHealth(nn.Module): f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}" ) other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype) - cls_time = self._baseline_cls_time( - event_seq=event_seq, - time_seq=time_seq, - padding_mask=padding_mask, - ) - h_token, token_mask, baseline_summary = self._encode_other_tokens( + h_other, other_mask = self.tokenizer( other_type=other_type, other_value=other_value, other_value_kind=other_value_kind, + ) + 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, - cls_time=cls_time, + other_mask=other_mask, ) - token_time = other_time.to(device=h_token.device, dtype=time_seq.dtype) - h_disease = self.cross_attention( - h_disease=h_disease, - t_disease=t_disease, - h_token=h_token, - t_token=token_time, - token_mask=token_mask, - ) - h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype) - h_disease = self._insert_baseline_summary( - h_disease=h_disease, - event_seq=event_seq, - baseline_summary=baseline_summary, - ) + h_disease = torch.cat([h_disease, h_other], dim=1) + t_disease = torch.cat([t_disease, other_time], dim=1) + padding_mask = torch.cat([padding_mask, other_mask], dim=1) h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype) if mode == "all_future": batch_size = event_seq.size(0) query = self.query_token.view(1, 1, -1).expand(batch_size, 1, -1) h_disease = torch.cat([h_disease, query], dim=1) - t_disease = torch.cat([time_seq, t_query[:, None]], dim=1) + t_disease = torch.cat([t_disease, t_query[:, None]], dim=1) query_mask = torch.ones( batch_size, 1, @@ -298,8 +418,28 @@ class DeepHealth(nn.Module): h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype) if mode == "all_future": - return h_disease[:, -1, :] - return h_disease + hidden = h_disease[:, -1, :] + if return_output: + return DeepHealthOutput( + hidden=hidden, + time_seq=t_query[:, None], + padding_mask=torch.ones( + hidden.size(0), + 1, + dtype=torch.bool, + device=hidden.device, + ), + event_len=event_len, + ) + return hidden + if return_output: + return DeepHealthOutput( + hidden=h_disease, + time_seq=t_disease, + padding_mask=padding_mask, + event_len=event_len, + ) + return h_disease[:, :event_len, :] def forward_next_token(self, **kwargs) -> torch.Tensor: return self._forward_shared(mode="next_token", **kwargs) diff --git a/train_all_future.py b/train_all_future.py index 9d35ca5..6407906 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -64,6 +64,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--n_hist_layer", type=int, default=12) parser.add_argument("--n_tab_layer", type=int, default=4) parser.add_argument("--n_bins", type=int, default=16) + parser.add_argument("--extra_pool_reduce", type=str, default="mean", + choices=["mean", "sum"]) parser.add_argument("--time_mode", type=str, default="relative", choices=["relative", "absolute"]) parser.add_argument("--dist_mode", type=str, default="exponential", @@ -129,6 +131,7 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De n_categories=dataset.n_categories, cont_type_ids=dataset.cont_type_ids, n_bins=args.n_bins, + extra_pool_reduce=args.extra_pool_reduce, target_mode="all_future", time_mode=args.time_mode, dist_mode=args.dist_mode, diff --git a/train_next_step.py b/train_next_step.py index 718bc9e..92c2b35 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -24,7 +24,7 @@ from tqdm.auto import tqdm from dataset import HealthDataset, collate_fn from losses import build_loss -from models import DeepHealth +from models import DeepHealth, DeepHealthOutput from readouts import build_readout from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX from train_util import ( @@ -61,6 +61,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--n_hist_layer", type=int, default=12) parser.add_argument("--n_tab_layer", type=int, default=4) parser.add_argument("--n_bins", type=int, default=16) + parser.add_argument("--extra_pool_reduce", type=str, default="mean", + choices=["mean", "sum"]) parser.add_argument("--time_mode", type=str, default="relative", choices=["relative", "absolute"]) parser.add_argument("--dropout", type=float, default=0.0) @@ -135,6 +137,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: n_categories=dataset.n_categories, cont_type_ids=dataset.cont_type_ids, n_bins=args.n_bins, + extra_pool_reduce=args.extra_pool_reduce, target_mode="next_token", time_mode=args.time_mode, dist_mode="exponential", @@ -169,6 +172,137 @@ def build_next_step_loss(args: argparse.Namespace): ) +def build_augmented_next_step_targets( + batch: Dict[str, torch.Tensor], + model_out: DeepHealthOutput, +) -> Dict[str, torch.Tensor]: + hidden_len = model_out.hidden.size(1) + event_len = int(model_out.event_len) + extra_len = hidden_len - event_len + if extra_len <= 0: + return { + "target_event_seq": batch["target_event_seq"], + "target_time_seq": batch["target_time_seq"], + "readout_mask": batch["readout_mask"], + "target_dt_unique": batch["target_dt_unique"], + "target_multi_hot": batch["target_multi_hot"], + } + + device = model_out.hidden.device + bsz, _seq_len, vocab_size = batch["target_multi_hot"].shape + extra_mask = model_out.padding_mask[:, event_len:] + extra_time = model_out.time_seq[:, event_len:] + + target_event_seq = torch.cat( + [ + batch["target_event_seq"], + torch.full( + (bsz, extra_len), + PAD_IDX, + dtype=batch["target_event_seq"].dtype, + device=device, + ), + ], + dim=1, + ) + target_time_seq = torch.cat( + [ + batch["target_time_seq"], + torch.zeros( + bsz, + extra_len, + dtype=batch["target_time_seq"].dtype, + device=device, + ), + ], + dim=1, + ) + readout_mask = torch.cat([batch["readout_mask"], extra_mask], dim=1) + target_dt_unique = torch.cat( + [ + batch["target_dt_unique"], + torch.zeros( + bsz, + extra_len, + dtype=batch["target_dt_unique"].dtype, + device=device, + ), + ], + dim=1, + ) + target_multi_hot = torch.cat( + [ + batch["target_multi_hot"], + torch.zeros( + bsz, + extra_len, + vocab_size, + dtype=batch["target_multi_hot"].dtype, + device=device, + ), + ], + dim=1, + ) + + for b in range(bsz): + valid_event = batch["padding_mask"][b].bool() + if not valid_event.any(): + continue + n_event = int(valid_event.sum().item()) + events = torch.cat( + [ + batch["event_seq"][b, :n_event], + batch["target_event_seq"][b, n_event - 1:n_event], + ] + ) + times = torch.cat( + [ + batch["time_seq"][b, :n_event], + batch["target_time_seq"][b, n_event - 1:n_event], + ] + ) + valid_full = events > PAD_IDX + events = events[valid_full] + times = times[valid_full] + if events.numel() == 0: + continue + + for j in range(extra_len): + if not bool(extra_mask[b, j]): + continue + pos = event_len + j + t = extra_time[b, j] + future = times > t + if not future.any(): + readout_mask[b, pos] = False + continue + + first_idx = int(torch.nonzero(future, as_tuple=False)[0].item()) + next_time = times[first_idx] + next_event = events[first_idx] + target_event_seq[b, pos] = next_event + target_time_seq[b, pos] = next_time + + same_next_time = times == next_time + next_events = events[same_next_time] + valid_next_events = next_events[ + (next_events > PAD_IDX) & (next_events < vocab_size) + ].long() + if valid_next_events.numel() == 0: + readout_mask[b, pos] = False + continue + target_multi_hot[b, pos, valid_next_events] = True + target_dt_unique[b, pos] = next_time - t + + return { + "target_event_seq": target_event_seq, + "target_time_seq": target_time_seq, + "readout_mask": readout_mask, + "target_dt_unique": target_dt_unique, + "target_multi_hot": target_multi_hot, + } + + def compute_next_step_loss( args: argparse.Namespace, model: DeepHealth, @@ -178,7 +312,7 @@ def compute_next_step_loss( device: torch.device, ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]: batch = move_batch_to_device(batch, device) - hidden = model( + model_out = model( event_seq=batch["event_seq"], time_seq=batch["time_seq"], sex=batch["sex"], @@ -188,12 +322,16 @@ def compute_next_step_loss( other_value_kind=batch["other_value_kind"], other_time=batch["other_time"], target_mode="next_token", + return_output=True, ) + if not isinstance(model_out, DeepHealthOutput): + raise TypeError("DeepHealth return_output=True must return DeepHealthOutput") + targets = build_augmented_next_step_targets(batch=batch, model_out=model_out) readout_out = readout( - hidden=hidden, - time_seq=batch["time_seq"], - padding_mask=batch["padding_mask"], - readout_mask=batch["readout_mask"] + hidden=model_out.hidden, + time_seq=model_out.time_seq, + padding_mask=model_out.padding_mask, + readout_mask=targets["readout_mask"] if args.readout_name == "same_time_group_end" else None, ) @@ -202,17 +340,17 @@ def compute_next_step_loss( if args.target_mode == "delphi2m": loss, parts = criterion( logits=logits, - target_events=batch["target_event_seq"], - target_times=batch["target_time_seq"], - current_times=batch["time_seq"], + target_events=targets["target_event_seq"], + target_times=targets["target_time_seq"], + current_times=model_out.time_seq, padding_mask=readout_out.readout_mask, return_components=True, ) else: loss, parts = criterion( logits=logits, - target_multi_hot=batch["target_multi_hot"], - target_dt_unique=batch["target_dt_unique"], + target_multi_hot=targets["target_multi_hot"], + target_dt_unique=targets["target_dt_unique"], readout_mask=readout_out.readout_mask, return_components=True, )