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

Refactor DeepHealth model and related components

Browse Source 
- 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:
Jiarui Li
2026-06-17 11:05:10 +08:00
parent 27aefb2f90
commit 1757bcd25b
7 changed files with 435 additions and 493 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
README.md
View File

@@ -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`

352
backbones.py
View File

@@ -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):

1
evaluate_auc.py
View File

@@ -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")),

1
evaluate_auc_v2.py
View File

@@ -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")),

304
models.py
View File

@@ -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)

3
train_all_future.py
View File

@@ -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,

160
train_next_step.py
View File

@@ -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,
)