From c3e49db859e44dba06b78f95bdb34800f38d9561 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Mon, 15 Jun 2026 14:10:09 +0800 Subject: [PATCH] Enhance DeepHealth model to incorporate CHECKUP state tokens in next-step training and evaluation, update dataset cache versioning, and improve handling of observed event histories. --- backbones.py | 24 ++++++++++++----- dataset.py | 16 ++++++----- eval_data.py | 6 ++--- evaluate_auc.py | 5 +--- evaluate_auc_v2.py | 9 +++---- evaluate_doa_auc.py | 65 +++++++++++++++++++++----------------------- models.py | 66 +++++++++++++++++++++++++++++---------------- train_next_step.py | 6 ++--- 8 files changed, 111 insertions(+), 86 deletions(-) diff --git a/backbones.py b/backbones.py index edf0dcb..2db4c1c 100644 --- a/backbones.py +++ b/backbones.py @@ -333,6 +333,7 @@ class BaselineEncoder(nn.Module): ) 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 = ( @@ -376,6 +377,7 @@ class BaselineEncoder(nn.Module): 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) @@ -439,17 +441,27 @@ class BaselineEncoder(nn.Module): f = type_emb + kind_emb + value_emb f = f * other_valid.unsqueeze(-1).to(f.dtype) - if f.size(1) == 0: - return f, other_valid + 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) - attn_mask = self._make_attn_mask(other_valid, f.dtype) + attn_mask = self._make_attn_mask(full_valid, f.dtype) for block in self.blocks: f = block(f, attn_mask=attn_mask) - f = f * other_valid.unsqueeze(-1).to(f.dtype) + f = f * full_valid.unsqueeze(-1).to(f.dtype) h = self.ln(f) - h = h * other_valid.unsqueeze(-1).to(h.dtype) - return h, other_valid + 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): diff --git a/dataset.py b/dataset.py index f50df3e..7c45028 100644 --- a/dataset.py +++ b/dataset.py @@ -110,7 +110,7 @@ def _cache_file_path( selected.append(type_id) selected_types = ",".join(str(t) for t in selected) signature_parts = [ - "deephealthnew_dataset_cache_v2", + "deephealthnew_dataset_cache_v3_checkup_state", dataset_kind, split or "", event_path, @@ -320,9 +320,6 @@ class _ExpoBaseDataset(Dataset): times_days_raw = rows[:, 1].astype(np.float32) labels_raw = rows[:, 2].astype(np.int64) - disease_mask = labels_raw != CHECKUP_IDX - times_days_raw = times_days_raw[disease_mask] - labels_raw = labels_raw[disease_mask] if len(labels_raw) == 0: yield eid, times_days_raw, labels_raw continue @@ -392,7 +389,7 @@ class NextStepHealthDataset(_ExpoBaseDataset): - UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot """ - CACHE_VERSION = 2 + CACHE_VERSION = 3 def __init__( self, @@ -488,7 +485,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset): time range, with at least one future event after every query. """ - CACHE_VERSION = 4 + CACHE_VERSION = 5 def __init__( self, @@ -582,8 +579,13 @@ class AllFutureHealthDataset(_ExpoBaseDataset): def _is_valid_query(self, patient: Dict, t_query: float) -> bool: times = patient["times"] + labels = patient["labels"] + real_event_mask = ~np.isin( + labels, + np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64), + ) n_hist = int((times <= t_query).sum()) - n_future = int((times > t_query).sum()) + n_future = int(((times > t_query) & real_event_mask).sum()) return ( n_hist >= self.min_history_events and n_future >= self.min_future_events diff --git a/eval_data.py b/eval_data.py index 5dc1a51..f139b3c 100644 --- a/eval_data.py +++ b/eval_data.py @@ -14,9 +14,9 @@ class AllFutureSequenceEvalDataset: """ Eval-only sequence view for all-future checkpoints. - All-future training uses pure disease histories, so token-level and landmark - evaluation should not reuse the next-step dataset view that contains - imputed gap tokens. + All-future training uses the observed history, including CHECKUP state + tokens, without reusing the next-step view that contains imputed + gap tokens. """ def __init__( diff --git a/evaluate_auc.py b/evaluate_auc.py index 47af00f..2dd4bb6 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -386,10 +386,7 @@ def load_model_state( state = state_dict if state_dict is not None else load_checkpoint_state_dict( checkpoint_path, map_location=device) - missing, unexpected = model.load_state_dict(state, strict=False) - if missing or unexpected: - print( - f"[WARN] load_state_dict strict=False: missing={missing[:10]}, unexpected={unexpected[:10]}") + model.load_state_dict(state, strict=True) def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any]) -> Tuple[Subset, np.ndarray]: diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index bbd922b..1bc5f68 100644 --- a/evaluate_auc_v2.py +++ b/evaluate_auc_v2.py @@ -192,10 +192,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None: - missing, unexpected = model.load_state_dict(state_dict, strict=False) - if missing or unexpected: - print( - f"[WARN] load_state_dict strict=False: missing={missing[:10]}, unexpected={unexpected[:10]}") + model.load_state_dict(state_dict, strict=True) def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None: @@ -1495,8 +1492,8 @@ def main() -> None: load_model_state(model, state_dict) except RuntimeError as exc: raise RuntimeError( - "Checkpoint vocabulary shape is incompatible with the no-event dataset/model setup. " - "Please ensure this run was trained with the current no-event vocabulary and matching labels file." + "Checkpoint vocabulary shape is incompatible with the dataset/model setup. " + "Please ensure this run was trained with the same special-token vocabulary and labels file." ) from exc if model.token_embedding.num_embeddings != dataset.vocab_size or model.risk_head.out_features != dataset.vocab_size: diff --git a/evaluate_doa_auc.py b/evaluate_doa_auc.py index 1a09118..36d4441 100644 --- a/evaluate_doa_auc.py +++ b/evaluate_doa_auc.py @@ -6,10 +6,9 @@ full observed record. Patients prevalent at/before DOA or incident after the horizon are not used for that disease-horizon AUC. The script adapts automatically to checkpoint target mode: - - next_token: use the DOA token position, inserting at DOA when no - real disease token exists at DOA; - - all_future: query the model directly with t_query=DOA, allowing empty - disease history because other-info tokens still describe the DOA state. + - next_token: use the CHECKUP token position at DOA; + - all_future: query the model directly with t_query=DOA. The history includes + the CHECKUP token at DOA. """ from __future__ import annotations @@ -163,17 +162,20 @@ class DOAStatusDataset(_ExpoBaseDataset): doa_days = float(np.min(checkup_rows[:, 1].astype(np.float32))) doa_years = np.float32(doa_days / DAYS_PER_YEAR) - disease_rows = rows[rows[:, 2].astype(np.int64) != CHECKUP_IDX] - disease_times = disease_rows[:, 1].astype(np.float32) / DAYS_PER_YEAR - disease_labels_raw = disease_rows[:, 2].astype(np.int64) - disease_labels = np.where( - disease_labels_raw >= NO_EVENT_IDX, - disease_labels_raw + 1, - disease_labels_raw, + raw_times = rows[:, 1].astype(np.float32) / DAYS_PER_YEAR + raw_labels = rows[:, 2].astype(np.int64) + shifted_labels = np.where( + raw_labels >= NO_EVENT_IDX, + raw_labels + 1, + raw_labels, ).astype(np.int64) - order = np.lexsort((disease_labels, disease_times)) - disease_times = disease_times[order].astype(np.float32) - disease_labels = disease_labels[order].astype(np.int64) + order = np.lexsort((shifted_labels, raw_times)) + event_times = raw_times[order].astype(np.float32) + event_labels = shifted_labels[order].astype(np.int64) + + disease_mask = event_labels != CHECKUP_IDX + disease_times = event_times[disease_mask] + disease_labels = event_labels[disease_mask] patient_id = len(self.records) for token in np.unique(disease_labels).tolist(): @@ -186,25 +188,20 @@ class DOAStatusDataset(_ExpoBaseDataset): (patient_id, float(disease_times[int(hit[0])])) ) - hist = disease_times <= doa_years - hist_events = disease_labels[hist] - hist_times = disease_times[hist] + hist = event_times <= doa_years + hist_events = event_labels[hist] + hist_times = event_times[hist] if self.model_target_mode == "next_token": - at_doa = np.isclose(hist_times, doa_years, rtol=0.0, atol=1e-6) - if hist_events.size == 0 or not np.any(at_doa): - event_seq = np.concatenate([ - hist_events, - np.array([NO_EVENT_IDX], dtype=np.int64), - ]) - time_seq = np.concatenate([ - hist_times, - np.array([doa_years], dtype=np.float32), - ]) - else: - event_seq = hist_events - time_seq = hist_times - readout_pos = int(len(event_seq) - 1) + checkup_at_doa = ( + (hist_events == CHECKUP_IDX) + & np.isclose(hist_times, doa_years, rtol=0.0, atol=1e-6) + ) + if not np.any(checkup_at_doa): + raise RuntimeError(f"Missing CHECKUP token at DOA for eid={eid}") + event_seq = hist_events + time_seq = hist_times + readout_pos = int(np.where(checkup_at_doa)[0][-1]) else: event_seq = hist_events time_seq = hist_times @@ -682,10 +679,10 @@ def main() -> None: model.eval() if model_target_mode == "next_token" and ( - model.token_embedding.num_embeddings <= NO_EVENT_IDX - or model.risk_head.out_features <= NO_EVENT_IDX + model.token_embedding.num_embeddings <= CHECKUP_IDX + or model.risk_head.out_features <= CHECKUP_IDX ): - raise RuntimeError("Next-token DOA evaluation requires in the model vocabulary.") + raise RuntimeError("Next-token DOA evaluation requires in the model vocabulary.") eval_dataset = Subset(dataset, eval_indices) loader = DataLoader( diff --git a/models.py b/models.py index e74db4b..2abcd96 100644 --- a/models.py +++ b/models.py @@ -10,6 +10,7 @@ from backbones import ( GaussianRBFTimeBasis, TimeRoPE, ) +from targets import CHECKUP_IDX, PAD_IDX class DeepHealth(nn.Module): @@ -128,12 +129,24 @@ class DeepHealth(nn.Module): dtype=dtype, ).masked_fill(~valid, -1e4)[:, None, :, :] + def _insert_baseline_summary( + 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 _encode_other_tokens( self, other_type: torch.LongTensor, other_value: torch.Tensor, other_value_kind: torch.LongTensor, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: return self.token_encoder( other_type=other_type, other_value=other_value, @@ -173,12 +186,40 @@ class DeepHealth(nn.Module): ) if padding_mask is None: - padding_mask = event_seq > 0 + padding_mask = event_seq > PAD_IDX else: padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool) + h_disease = self.token_embedding(event_seq) t_disease = time_seq + h_token, token_mask, baseline_summary = self._encode_other_tokens( + other_type=other_type, + other_value=other_value, + other_value_kind=other_value_kind, + ) + 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)}" + ) + 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 = 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) @@ -222,27 +263,6 @@ class DeepHealth(nn.Module): h_disease = self.final_ln(h_disease) h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype) - h_token, token_mask = self._encode_other_tokens( - other_type=other_type, - other_value=other_value, - other_value_kind=other_value_kind, - ) - 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)}" - ) - token_time = other_time.to(device=h_token.device, dtype=t_disease.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) - if mode == "all_future": return h_disease[:, -1, :] return h_disease diff --git a/train_next_step.py b/train_next_step.py index 02e044c..71957c6 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -1,9 +1,9 @@ """ Train DeepHealth with next-token / next-time-point supervision. -The dataset remains the current next-step construction: pure disease events plus -optional gap imputation are shifted into autoregressive inputs and -targets. UTS training reads out only same-time group ends. +The next-step dataset uses observed event histories, including CHECKUP state +tokens, plus optional gap imputation. UTS training reads out only +same-time group ends. """ from __future__ import annotations