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

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

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

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 <NO_EVENT> gap tokens.
+    All-future training uses the observed history, including CHECKUP state
+    tokens, without reusing the next-step view that contains imputed
+    <NO_EVENT> gap tokens.
     """
 
     def __init__(

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

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:

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 <NO_EVENT> 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 <NO_EVENT> in the model vocabulary.")
+        raise RuntimeError("Next-token DOA evaluation requires <CHECKUP> in the model vocabulary.")
 
     eval_dataset = Subset(dataset, eval_indices)
     loader = DataLoader(

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

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 <NO_EVENT> 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 <NO_EVENT> imputation. UTS training reads out only
+same-time group ends.
 """
 from __future__ import annotations