Remove legacy event and mixed distribution paths

README.md

@@ -6,7 +6,9 @@
 疾病序列 stream + 统一的额外信息 token stream
 ```
 
-疾病、死亡、checkup 事件保存在预处理事件文件中；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。dataset 仅在实验选择了至少一种 extra-info type 时保留 checkup；显式传入空列表时，模型输入是没有 checkup 的纯疾病历史。
+疾病和死亡事件保存在预处理事件文件中；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。每个额外信息 token 自带测量时间并直接与疾病 token 拼接，不再生成或使用 assessment/checkup 事件 token。dataset 会无条件清除旧预处理文件中遗留的 `label=1` 事件。
+
+所有连续额外信息强制使用训练子集拟合的 RobustScale（median/IQR）。center 和 scale 保存为模型 buffer，并用于验证集、测试集和推理；代码不提供未标准化模式，缺少 scaler buffer 的连续变量 checkpoint 不受支持。该规则同时适用于 `next_token` 和 `all_future`。
 
 ## 数据准备
 
@@ -28,7 +30,7 @@ python prepare_data.py
 - `ukb_event_data.npy`
   - 形状为 `(N, 3)`
   - 每行是 `(eid, days, label)`
-  - 包含疾病、死亡、checkup 事件
+  - 包含疾病和死亡事件
 
 - `ukb_basic_info.csv`
   - index 为 `eid`
@@ -64,7 +66,7 @@ python prepare_data.py
 
 - `AllFutureHealthDataset`
   - 用于 query-conditioned all-future 监督
-  - 对应 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`
+  - 对应 `ExponentialLoss`、`WeibullLoss`
 
 为了兼容旧训练入口：
 
@@ -215,7 +217,6 @@ all-future / query-conditioned 监督：
 
 - `ExponentialLoss`
 - `WeibullLoss`
-- `MixedLoss`
 
 all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-info tokens 作为主序列上下文输入，但不会被单独读出，也不会被纳入 loss 监督。
 
@@ -232,7 +233,7 @@ all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-in
 - `train_all_future.py`
   - 使用 `AllFutureHealthDataset`
   - 不使用 readout，直接对 query hidden 计算风险
-  - `--dist_mode exponential/weibull/mixed` 分别搭配 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`
+  - `--dist_mode exponential/weibull` 分别搭配 `ExponentialLoss`、`WeibullLoss`
   - 展开的 extra-info tokens 只作为 query 上下文，不单独监督
 
 当前 `train_next_step.py` / `train_all_future.py` 支持所有已有训练目标定义的组合：
@@ -242,7 +243,6 @@ all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-in
 | `next_token` | `absolute` | `target_mode=delphi2m`, `dist_mode=exponential` | `Delphi2MLoss` + `token` |
 | `all_future` | `relative`, `absolute` | `dist_mode=exponential` | `ExponentialLoss`，无 readout |
 | `all_future` | `relative`, `absolute` | `dist_mode=weibull` | `WeibullLoss`，无 readout |
-| `all_future` | `relative`, `absolute` | `dist_mode=mixed` | `MixedLoss`，无 readout |
 
 示例：
 
@@ -366,7 +366,7 @@ python evaluate_auc.py \
 - 为每个患者和 landmark age 构造 landmark query 样本。
 - 根据模型模式插入 `<NO_EVENT>` token 或直接传 `t_query`，取 landmark/query hidden。
 - 对疾病 token 分块投影到 `risk_head`；`score_mode="risk"` 时会根据 `dist_mode` 把线性输出转换为固定 horizon 风险概率。
-- 分布转换规则与 all-future 训练损失一致：`exponential` 使用 `1 - exp(-rate * horizon)`；`weibull` 使用 `1 - exp(-rate * horizon ** rho)`；`mixed` 中普通疾病使用 exponential，死亡 endpoint 使用 Weibull death rho。
+- 分布转换规则与 all-future 训练损失一致：`exponential` 使用 `1 - exp(-rate * horizon)`；`weibull` 使用 `1 - exp(-rate * horizon ** rho)`。
 - `score_mode="eta"` 是诊断用排序分数，不使用 `rho`，因此不区分不同分布的风险曲线。
 - 按疾病、性别、landmark age、horizon 计算 incident disease AUC。
 - 可选择排除 horizon 内先于目标疾病发生的死亡竞争风险。
@@ -402,7 +402,7 @@ python evaluate_auc_v2.py \
 | 评估口径 | next-step/token-level 预测点 | landmark fixed-horizon incident risk |
 | 查询位置 | next-token 用满足 offset 条件的最新 readout token；all-future 直接用该预测点年龄作为 `t_query` | next-token 用人工插入的 `<NO_EVENT>` landmark token；all-future 直接用 `t_query` |
 | 时间参数 | `offsets`：预测点至少早于目标事件多少年 | `landmark_*` 和 `horizons`：固定年龄点与未来窗口 |
-| score 与分布 | 使用 eta/logit 排序分数；不按 `dist_mode` 转换风险概率 | `score_mode="risk"` 按 `dist_mode` 区分 exponential / Weibull / mixed；`score_mode="eta"` 不区分分布 |
+| score 与分布 | 使用 eta/logit 排序分数；不按 `dist_mode` 转换风险概率 | `score_mode="risk"` 按 `dist_mode` 区分 exponential / Weibull；`score_mode="eta"` 不区分分布 |
 | 病例定义 | target table 中出现目标疾病的患者/事件 | landmark 后 horizon 内首次发生目标疾病 |
 | 对照定义 | 从未出现该疾病的患者的 eligible target occurrence | landmark 时未患病，且 horizon 内未发病并有足够随访 |
 | 分层 | sex + age bracket + offset | sex + landmark age + horizon |
@@ -450,4 +450,4 @@ python evaluate_auc_v2.py \
   - landmark fixed-horizon incident disease AUC 评估
   - next-token 模型通过插入 `<NO_EVENT>` landmark token 查询固定年龄点风险
   - all-future 模型直接通过 `t_query` 查询固定年龄点风险
-  - `score_mode="risk"` 按 exponential / Weibull / mixed 分布计算固定 horizon 风险
+  - `score_mode="risk"` 按 exponential / Weibull 分布计算固定 horizon 风险

dataset.py

@@ -10,10 +10,10 @@ from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import Dataset
 
 from targets import (
-    CHECKUP_IDX,
     DAYS_PER_YEAR,
     NO_EVENT_IDX,
     PAD_IDX,
+    RESERVED_IDX,
     build_next_token_targets,
 )
 
@@ -188,12 +188,12 @@ def load_label_vocab(
 ) -> Tuple[Dict[str, int], Dict[int, str]]:
     label_id_to_code: Dict[int, str] = {
         PAD_IDX: "<PAD>",
-        CHECKUP_IDX: "<CHECKUP>",
+        RESERVED_IDX: "<RESERVED>",
     }
     if include_no_event:
         label_id_to_code[NO_EVENT_IDX] = "<NO_EVENT>"
 
-    offset = NO_EVENT_IDX + 1 if include_no_event else CHECKUP_IDX + 1
+    offset = NO_EVENT_IDX + 1 if include_no_event else RESERVED_IDX + 1
     label_code_to_id: Dict[str, int] = {}
     with open(labels_file, encoding="utf-8") as f:
         for i, line in enumerate(f):
@@ -416,12 +416,10 @@ class _ExpoBaseDataset(Dataset):
             times_days_raw = rows[:, 1].astype(np.float32)
             labels_raw = rows[:, 2].astype(np.int64)
 
-            # CHECKUP is the assessment landmark for selected extra-info tokens.
+            # Label 1 was emitted as a CHECKUP event by older prepared files.
-            # An explicitly empty selection represents a disease-only history,
+            # It is now an unused reserved slot and must never enter either the
-            # so retaining CHECKUP in that case would introduce an empty
+            # next-token or all-future disease sequence.
-            # landmark token that is not part of the disease sequence.
+            keep = labels_raw != RESERVED_IDX
-            if not self.extra_info_types:
-                keep = labels_raw != CHECKUP_IDX
             times_days_raw = times_days_raw[keep]
             labels_raw = labels_raw[keep]
 
@@ -615,7 +613,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
         labels = patient["labels"]
         real_event_mask = ~np.isin(
             labels,
-            np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
+            np.array([PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64),
         )
         n_hist = int((times <= t_query).sum())
         n_future = int(((times > t_query) & real_event_mask).sum())
@@ -634,7 +632,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
         labels = np.asarray(patient["labels"], dtype=np.int64)
         real_event_mask = ~np.isin(
             labels,
-            np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
+            np.array([PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64),
         )
         real_times = np.sort(times[real_event_mask].astype(np.float32, copy=False))
         n_real_events = int(real_times.size)

eval_data.py

@@ -151,6 +151,33 @@ def build_model_from_dataset(
             f"{model_target_mode!r}"
         )
     model_architecture = resolve_model_architecture(cfg, state_dict)
+    continuous_value_center = None
+    continuous_value_scale = None
+    if dataset.n_cont_types > 0:
+        scaling = str(cfg.get("continuous_value_scaling", "")).lower()
+        if scaling != "robust":
+            raise RuntimeError(
+                "Continuous-variable checkpoints must declare "
+                "continuous_value_scaling='robust'; unscaled checkpoints are "
+                "not supported"
+            )
+        if state_dict is None:
+            raise RuntimeError(
+                "A checkpoint state_dict is required to restore RobustScale buffers"
+            )
+        center_key = "tokenizer.continuous_value_center"
+        scale_key = "tokenizer.continuous_value_scale"
+        missing = [
+            key for key in (center_key, scale_key)
+            if key not in state_dict
+        ]
+        if missing:
+            raise RuntimeError(
+                "Checkpoint is missing required RobustScale buffers: "
+                + ", ".join(missing)
+            )
+        continuous_value_center = state_dict[center_key]
+        continuous_value_scale = state_dict[scale_key]
     return DeepHealth(
         vocab_size=dataset.vocab_size,
         n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
@@ -161,9 +188,8 @@ def build_model_from_dataset(
         n_categories=dataset.n_categories,
         cont_type_ids=dataset.cont_type_ids,
         n_bins=int(cfg_get(args, cfg, "n_bins", 16)),
-        continuous_value_scaling=str(
+        continuous_value_center=continuous_value_center,
-            cfg_get(args, cfg, "continuous_value_scaling", "none")
+        continuous_value_scale=continuous_value_scale,
-        ),
         extra_pool_reduce=str(
             cfg_get(args, cfg, "extra_pool_reduce", "mean")
         ),
@@ -251,9 +277,9 @@ class AllFutureSequenceEvalDataset:
     Eval-only sequence view for all-future checkpoints.
 
     All-future training uses the observed history without reusing the
-    next-step view that contains imputed <NO_EVENT> gap tokens. CHECKUP is
+    next-step view that contains imputed <NO_EVENT> gap tokens. Legacy label-1
-    retained only when the experiment selects at least one extra-info type;
+    assessment events are removed by the shared base dataset for every
-    an explicitly empty selection is a disease-only history.
+    extra-info selection.
     """
 
     def __init__(

evaluate_auc.py

@@ -63,7 +63,7 @@ from eval_data import (
 )
 from model_architectures import resolve_model_architecture
 from models import DeepHealth
-from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
+from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
 
 
 # ---------------------------------------------------------------------------
@@ -168,7 +168,7 @@ def get_auc_delong_var(control_scores: np.ndarray, case_scores: np.ndarray) -> T
 # Disease selection
 # ---------------------------------------------------------------------------
 
-SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
+SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
 
 
 def _get_death_token_ids(dataset: HealthDataset) -> List[int]:
@@ -249,27 +249,20 @@ def load_checkpoint_state_dict(checkpoint_path: str, map_location: str | torch.d
 
 def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str:
     mode = str(cfg_dist_mode).lower()
+    if mode not in {"exponential", "weibull"}:
+        raise ValueError(
+            f"Unsupported dist_mode={mode!r}; expected exponential or weibull."
+        )
     has_rho_head = any(str(k).startswith("rho_head.")
                        for k in state_dict.keys())
-    has_rho_death_head = any(str(k).startswith("rho_death_head.")
+    if mode == "weibull" and not has_rho_head:
-                             for k in state_dict.keys())
+        raise RuntimeError(
-
+            "Weibull checkpoint is missing rho_head parameters."
-    if has_rho_head and mode != "weibull":
+        )
-        print(
+    if mode == "exponential" and has_rho_head:
-            "[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
+        raise RuntimeError(
-        return "weibull"
+            "Exponential checkpoint unexpectedly contains rho_head parameters."
-    if has_rho_death_head and mode != "mixed":
+        )
-        print(
-            "[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
-        return "mixed"
-    if (not has_rho_head) and mode == "weibull":
-        print(
-            "[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
-        return "exponential"
-    if (not has_rho_death_head) and mode == "mixed":
-        print(
-            "[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.")
-        return "exponential"
     return mode
 
 
@@ -977,7 +970,7 @@ def evaluate_auc_pipeline(
     sex_items = [("female", 0), ("male", 1)]
     all_rows: List[Dict[str, Any]] = []
 
-    valid_target_min_id = CHECKUP_IDX if NO_EVENT_IDX >= dataset.vocab_size else CHECKUP_IDX
+    valid_target_min_id = RESERVED_IDX
     # If NO_EVENT exists and should not be a disease/control target, require target > NO_EVENT_IDX.
     if NO_EVENT_IDX in dataset.label_id_to_code and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>":
         valid_target_min_id = NO_EVENT_IDX

evaluate_auc_v2.py

@@ -1,7 +1,7 @@
 """Evaluate landmark fixed-horizon incident disease AUC for DeepHealth.
 
-This script supports DeepHealth fixed-horizon risk scores for exponential,
+This script supports DeepHealth fixed-horizon risk scores for exponential and
-Weibull, and mixed all-future distributions.
+Weibull all-future distributions.
 
 The default horizons are 0.1, 1, 5, and 10 years.  As in Delphi2M, 0.1 years
 is reported as the no-gap evaluation.
@@ -53,9 +53,9 @@ from eval_data import (
 )
 from model_architectures import resolve_model_architecture
 from models import DeepHealth
-from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
+from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
 
-SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
+SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
 
 
 def parse_int_list(value: Any) -> Optional[List[int]]:
@@ -137,29 +137,19 @@ def load_checkpoint_state_dict(checkpoint_path: Path, map_location: str | torch.
 
 def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str:
     mode = str(cfg_dist_mode).lower()
+    if mode not in {"exponential", "weibull"}:
+        raise ValueError(
+            f"Unsupported dist_mode={mode!r}; expected exponential or weibull."
+        )
     has_rho_head = any(str(k).startswith("rho_head.")
                        for k in state_dict.keys())
-    has_rho_death_head = any(str(k).startswith("rho_death_head.")
+    if mode == "weibull" and not has_rho_head:
-                             for k in state_dict.keys())
+        raise RuntimeError("Weibull checkpoint is missing rho_head parameters.")
-    if has_rho_head:
+    if mode == "exponential" and has_rho_head:
-        if mode != "weibull":
+        raise RuntimeError(
-            print(
+            "Exponential checkpoint unexpectedly contains rho_head parameters."
-                "[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
+        )
-        return "weibull"
+    return mode
-    if has_rho_death_head:
-        if mode != "mixed":
-            print(
-                "[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
-        return "mixed"
-    if mode == "weibull":
-        print(
-            "[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
-        return "exponential"
-    if mode == "mixed":
-        print(
-            "[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.")
-        return "exponential"
-    return mode if mode in {"exponential", "weibull", "mixed"} else "exponential"
 
 
 def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
@@ -290,7 +280,7 @@ def _metadata_count_map(dataset: HealthDataset, labels_meta: Optional[pd.DataFra
     return out
 
 
-def _get_death_token_ids(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> List[int]:
+def _get_death_token_ids(dataset: HealthDataset) -> List[int]:
     return [int(dataset.vocab_size) - 1]
 
 
@@ -418,7 +408,7 @@ class LandmarkDataset(Dataset):
                 prefix_times = full_time[prefix_mask]
 
                 valid_history_mask = ~np.isin(prefix_events, np.array(
-                    [PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64))
+                    [PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64))
                 if valid_history_mask.sum() < self.min_history_events:
                     continue
 
@@ -671,24 +661,11 @@ def project_distribution_chunk(
             device=device, dtype=compute_dtype)
     rho_weight = None
     rho_bias = None
-    death_rho_weight = None
-    death_rho_bias = None
-    mixed_death_cols: List[int] = []
-    death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1))
-
     if dist_mode == "weibull":
         rho_weight = model.rho_head.weight[disease_ids].detach().to(
             device=device, dtype=compute_dtype)
         rho_bias = model.rho_head.bias[disease_ids].detach().to(
             device=device, dtype=compute_dtype)
-    elif dist_mode == "mixed":
-        mixed_death_cols = [j for j, token in enumerate(disease_ids)
-                            if int(token) == death_idx]
-        if mixed_death_cols:
-            death_rho_weight = model.rho_death_head.weight.detach().to(
-                device=device, dtype=compute_dtype)
-            death_rho_bias = model.rho_death_head.bias.detach().to(
-                device=device, dtype=compute_dtype)
 
     out_parts: List[np.ndarray] = []
     rho_parts: List[np.ndarray] = []
@@ -703,14 +680,6 @@ def project_distribution_chunk(
         if dist_mode == "weibull":
             assert rho_weight is not None and rho_bias is not None
             rho = F.softplus(torch.matmul(h, rho_weight.t()) + rho_bias) + 1e-6
-        elif dist_mode == "mixed" and mixed_death_cols:
-            assert death_rho_weight is not None and death_rho_bias is not None
-            rho = torch.ones_like(logits)
-            death_rho = F.softplus(
-                torch.matmul(h, death_rho_weight.t()).squeeze(-1) + death_rho_bias.squeeze(0)
-            ) + 1e-6
-            for col in mixed_death_cols:
-                rho[:, int(col)] = death_rho
 
         out_parts.append(logits.float().cpu(
         ).numpy().astype(np.float32, copy=False))
@@ -747,7 +716,6 @@ def _init_worker(
     exclude_death_competing: bool,
     death_token_ids: np.ndarray,
     dist_mode: str,
-    model_death_idx: int,
 ) -> None:
     os.environ.setdefault("OMP_NUM_THREADS", "1")
     os.environ.setdefault("MKL_NUM_THREADS", "1")
@@ -772,7 +740,6 @@ def _init_worker(
             "exclude_death_competing": bool(exclude_death_competing),
             "death_token_ids": set(int(x) for x in np.asarray(death_token_ids, dtype=np.int64).tolist()),
             "dist_mode": str(dist_mode).lower(),
-            "model_death_idx": int(model_death_idx),
             "first_time_cache": {},
         }
     )
@@ -798,8 +765,6 @@ def _score_to_probability(
     score_mode: str,
     horizon: float,
     dist_mode: str,
-    token: int,
-    death_idx: int,
 ) -> np.ndarray:
     if score_mode == "eta":
         return logits.astype(np.float64, copy=False)
@@ -811,11 +776,6 @@ def _score_to_probability(
             raise RuntimeError("Weibull risk scoring requires rho parameters.")
         exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False))
         return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False)
-    if dist_mode == "mixed" and int(token) == int(death_idx):
-        if rho is None:
-            raise RuntimeError("Mixed death risk scoring requires death rho parameters.")
-        exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False))
-        return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False)
     return (-np.expm1(-rate * np.float32(horizon))).astype(np.float64, copy=False)
 
 
@@ -832,7 +792,6 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
     rho_chunk = _WORKER["rho_chunk"]
     rho_token = None if rho_chunk is None else rho_chunk[:, int(j)]
     dist_mode = _WORKER["dist_mode"]
-    model_death_idx = int(_WORKER["model_death_idx"])
 
     first_time_patient = _first_time_by_patient(token)
     is_death_target = token in _WORKER["death_token_ids"]
@@ -891,8 +850,6 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
                     score_mode=score_mode,
                     horizon=horizon,
                     dist_mode=dist_mode,
-                    token=token,
-                    death_idx=model_death_idx,
                 )
                 control_scores = _score_to_probability(
                     logits_token[idx[control_idx]],
@@ -900,8 +857,6 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
                     score_mode=score_mode,
                     horizon=horizon,
                     dist_mode=dist_mode,
-                    token=token,
-                    death_idx=model_death_idx,
                 )
 
                 auc, auc_var = get_auc_delong_var(case_scores, control_scores)
@@ -1019,8 +974,6 @@ def evaluate_landmark_auc(
                 death_token_ids=np.asarray(
                     landmark_dataset.death_token_ids, dtype=np.int64),
                 dist_mode=dist_mode,
-                model_death_idx=int(getattr(
-                    model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
             )
             nested = [_eval_token(t) for t in tqdm(
                 tasks, desc=f"AUC chunk {chunk_idx}", leave=False, dynamic_ncols=True)]
@@ -1048,8 +1001,6 @@ def evaluate_landmark_auc(
                     np.asarray(landmark_dataset.death_token_ids,
                                dtype=np.int64),
                     dist_mode,
-                    int(getattr(
-                        model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
                 ),
             ) as ex:
                 nested = list(
@@ -1253,10 +1204,6 @@ def main() -> None:
 
     state_dict = load_checkpoint_state_dict(model_ckpt_path, map_location="cpu")
     dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict)
-    if dist_mode not in {"exponential", "weibull", "mixed"}:
-        raise ValueError(
-            f"Unsupported dist_mode={dist_mode!r}; expected exponential, weibull, or mixed."
-        )
 
     if score_mode == "eta":
         print(
@@ -1302,7 +1249,7 @@ def main() -> None:
             "Please use a checkpoint trained with the same no-event vocabulary configuration."
         )
 
-    death_token_ids = _get_death_token_ids(dataset, labels_meta)
+    death_token_ids = _get_death_token_ids(dataset)
     min_history_events = int(cfg_get(args, cfg, "min_history_events", 1))
     landmark_dataset = LandmarkDataset(
         dataset=dataset,

evaluate_calibration.py

@@ -23,12 +23,11 @@ from __future__ import annotations
 import argparse
 import contextlib
 import json
-import math
 import os
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 import numpy as np
 import pandas as pd
@@ -67,7 +66,7 @@ from evaluate_auc_v2 import (
 )
 from losses import build_loss
 from model_architectures import resolve_model_architecture
-from targets import CHECKUP_IDX, PAD_IDX
+from targets import PAD_IDX, RESERVED_IDX
 from train_util import load_eid_file
 
 
@@ -971,8 +970,6 @@ def _risk_probability_matrix(
     rho: Optional[np.ndarray],
     horizons: np.ndarray,
     dist_mode: str,
-    token: int,
-    death_idx: int,
 ) -> np.ndarray:
     """Convert one token's logits to all horizon risks at once."""
     logits = np.asarray(logits, dtype=np.float32)
@@ -982,13 +979,7 @@ def _risk_probability_matrix(
         + np.maximum(logits, np.float32(0.0))
         + np.float32(1e-8)
     )
-    use_weibull = (
+    use_weibull = str(dist_mode).lower() == "weibull"
-        str(dist_mode).lower() == "weibull"
-        or (
-            str(dist_mode).lower() == "mixed"
-            and int(token) == int(death_idx)
-        )
-    )
     if use_weibull:
         if rho is None:
             raise RuntimeError(
@@ -1047,7 +1038,6 @@ def _evaluate_calibration_token(
     label_id_to_code: Dict[int, str],
     dist_mode: str,
     horizons: np.ndarray,
-    death_index: int,
     min_cases: int,
     min_controls: int,
     max_ipcw_weight: float,
@@ -1114,8 +1104,6 @@ def _evaluate_calibration_token(
             ),
             horizons=horizons,
             dist_mode=dist_mode,
-            token=token,
-            death_idx=death_index,
         )
         results = compute_ipcw_horizons(
             probabilities=probabilities,
@@ -1200,9 +1188,6 @@ def evaluate_landmark_calibration(
 
     patient_count = len(landmark_dataset.subset_indices)
     death_tokens = set(int(value) for value in landmark_dataset.death_token_ids)
-    death_index = int(
-        getattr(model, "death_idx", getattr(model, "vocab_size", 1) - 1)
-    )
     strata = _build_calibration_strata(
         row_arrays["sex"],
         row_arrays["landmark_age"],
@@ -1265,7 +1250,6 @@ def evaluate_landmark_calibration(
                     ),
                     "dist_mode": dist_mode,
                     "horizons": horizons,
-                    "death_index": death_index,
                     "min_cases": min_cases,
                     "min_controls": min_controls,
                     "max_ipcw_weight": max_ipcw_weight,
@@ -1498,19 +1482,12 @@ def build_calibration_summary(metrics: pd.DataFrame) -> pd.DataFrame:
 
 def _build_point_process_criterion(
     dist_mode: str,
-    death_index: int,
 ) -> Any:
-    ignored = {PAD_IDX, CHECKUP_IDX}
+    ignored = {PAD_IDX, RESERVED_IDX}
     if dist_mode == "exponential":
         return build_loss("exponential", ignored_idx=ignored)
     if dist_mode == "weibull":
         return build_loss("weibull", ignored_idx=ignored)
-    if dist_mode == "mixed":
-        return build_loss(
-            "mixed",
-            death_idx=death_index,
-            ignored_idx=ignored,
-        )
     raise ValueError(f"Unsupported dist_mode: {dist_mode!r}")
 
 
@@ -1523,16 +1500,7 @@ def evaluate_point_process_nll(
     device: torch.device,
     use_amp: bool,
 ) -> Dict[str, Any]:
-    criterion = _build_point_process_criterion(
+    criterion = _build_point_process_criterion(dist_mode)
-        dist_mode,
-        death_index=int(
-            getattr(
-                model,
-                "death_idx",
-                int(getattr(model, "vocab_size", 1)) - 1,
-            )
-        ),
-    )
     model.eval().to(device)
     total_nll = 0.0
     query_count = 0
@@ -1587,13 +1555,7 @@ def evaluate_point_process_nll(
                     exposure=batch_device["exposure"],
                 )
             else:
-                loss = criterion(
+                raise ValueError(f"Unsupported dist_mode: {dist_mode!r}")
-                    logits=logits,
-                    death_rho=model.calc_death_rho(hidden),
-                    targets=batch_device["future_targets"],
-                    dt=batch_device["future_dt"],
-                    exposure=batch_device["exposure"],
-                )
 
         if not torch.isfinite(loss):
             raise RuntimeError("Non-finite point-process NLL encountered.")
@@ -1601,7 +1563,7 @@ def evaluate_point_process_nll(
         total_nll += float(loss.detach().cpu()) * batch_size
         query_count += batch_size
         valid_targets = batch["future_targets"] > PAD_IDX
-        valid_targets &= batch["future_targets"] != CHECKUP_IDX
+        valid_targets &= batch["future_targets"] != RESERVED_IDX
         future_event_count += int(valid_targets.sum().item())
         exposure_sum += float(batch["exposure"].sum().item())
 

extra_info_assessment_all_analysis_plan.md

@@ -4,10 +4,13 @@
 
 补充现有四级 extra-information 证据链：
 
-1. `disease_only`：疾病事件、相对患病时间和 sex；无 CHECKUP、无 extra-info token。
+1. `disease_only`：疾病事件、相对患病时间和 sex；无 extra-info token。
-2. `smoking_alcohol_bmi`：疾病史、sex、CHECKUP、smoking/alcohol/BMI。
+2. `smoking_alcohol_bmi`：疾病史、sex、smoking/alcohol/BMI token。
-3. `assessment_only`：疾病史、sex、CHECKUP、65项常规体格、肺功能、血液、尿液和生化指标。
+3. `assessment_only`：疾病史、sex、65项常规体格、肺功能、血液、尿液和生化指标 token。
-4. `all`：疾病史、sex、CHECKUP、全部265项体检和暴露信息。
+4. `all`：疾病史、sex、全部265项体检和暴露信息 token。
+
+所有配置均不使用 CHECKUP。额外信息以独立 token 注入，并使用各自的 assessment 时间。
+所有连续变量均强制使用训练子集拟合的 RobustScale；训练与评估不设置未标准化对照或兼容模式。
 
 目标是区分：
 
@@ -63,7 +66,7 @@ assessment_only − disease_only
 all − assessment_only
 ```
 
-回答生活方式、社会经济、心理和环境暴露是否在常规体检之后仍有增量价值。这是新增实验中最干净的主要比较，因为两组都保留 CHECKUP，Landmark 和随访边界应一致。
+回答生活方式、社会经济、心理和环境暴露是否在常规体检之后仍有增量价值。两组使用完全相同的疾病事件、Landmark 和随访边界，仅 extra-info token 集合不同。
 
 ### 4.3 全体检相对紧凑变量集
 
@@ -130,15 +133,9 @@ smoking_alcohol_bmi − disease_only
 - Brier、NLL及绝对校准偏差越低越好；
 - 不以单个 seed 或单个 horizon 决定模型。
 
-## 7. disease_only 比较的评估限制
+## 7. disease_only 比较的评估边界
 
-`disease_only` 按设计删除 CHECKUP，其他三组保留 CHECKUP。当前评估实现会使两类模型的随访终点和 `n_at_risk` 略有差异。
+所有配置均从同一疾病/死亡事件流构造历史、查询点、随访终点和 censoring，并且都不使用 CHECKUP。不同配置只改变 extra-info token，因此可以在共同支持集上把差异解释为额外信息的增量价值。
-
-因此：
-
-- `assessment_only`、`smoking_alcohol_bmi`、`all` 三者之间可以直接比较；
-- 它们与 `disease_only` 的比较应使用固定的原始随访终点、Landmark 和 censoring；
-- 在共享风险集评估完成前，不能把与 `disease_only` 的全部差异严格归因于 extra-info 数值。
 
 ## 8. 决策规则
 

losses.py

@@ -8,7 +8,7 @@ import torch.nn.functional as F
 
 
 PAD_IDX = 0
-CHECKUP_IDX = 1
+RESERVED_IDX = 1
 NO_EVENT_IDX = 2
 
 
@@ -52,7 +52,7 @@ class Delphi2MLoss(nn.Module):
         super().__init__()
         self.t_min = float(t_min)
         self.ignored_tokens = (
-            [PAD_IDX, CHECKUP_IDX]
+            [PAD_IDX, RESERVED_IDX]
             if ignored_tokens is None
             else [int(x) for x in ignored_tokens]
         )
@@ -150,7 +150,7 @@ class ExponentialLoss(nn.Module):
 
     def __init__(
         self,
-        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
+        ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX),
         eps: float = 1e-8,
     ):
         super().__init__()
@@ -183,7 +183,7 @@ class WeibullLoss(nn.Module):
 
     def __init__(
         self,
-        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
+        ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX),
         eps: float = 1e-8,
     ):
         super().__init__()
@@ -232,78 +232,14 @@ class WeibullLoss(nn.Module):
         return (-observed + penalty).mean()
 
 
-class MixedLoss(nn.Module):
-    """Exponential diseases plus one Weibull death endpoint."""
-
-    def __init__(
-        self,
-        death_idx: int,
-        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
-        eps: float = 1e-8,
-    ):
-        super().__init__()
-        self.death_idx = int(death_idx)
-        self.ignored_idx = tuple(int(i) for i in ignored_idx)
-        self.eps = eps
-
-    def forward(
-        self,
-        logits: torch.Tensor,
-        death_rho: torch.Tensor,
-        targets: torch.Tensor,
-        dt: torch.Tensor,
-        exposure: torch.Tensor,
-    ) -> torch.Tensor:
-        _, vocab_size = logits.shape
-        dtype = logits.dtype
-        rate = F.softplus(logits) + self.eps
-
-        if death_rho.dim() == 2:
-            death_rho = death_rho.squeeze(-1)
-        death_rho = death_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps)
-
-        valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
-        valid_disease_vocab = valid_vocab.clone()
-        valid_disease_vocab[self.death_idx] = False
-
-        t_exp = exposure.to(dtype).clamp_min(self.eps)
-        disease_penalty = t_exp * rate[:, valid_disease_vocab].sum(dim=-1)
-        death_rate = rate[:, self.death_idx]
-        death_penalty = death_rate * torch.pow(t_exp, death_rho)
-        penalty = disease_penalty + death_penalty
-
-        target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
-        for idx in self.ignored_idx:
-            target_valid &= targets != idx
-
-        disease_event_mask = target_valid & (targets != self.death_idx)
-        safe_targets = targets.clamp(min=0, max=vocab_size - 1)
-        disease_log_rate = rate.log().gather(1, safe_targets)
-        observed_disease = (disease_log_rate * disease_event_mask.to(dtype)).sum(dim=-1)
-
-        death_event_mask = target_valid & (targets == self.death_idx)
-        death_observed = death_event_mask.any(dim=1)
-        death_dt = (dt.to(dtype).clamp_min(self.eps) * death_event_mask.to(dtype)).sum(dim=1)
-        death_log_intensity = (
-            death_rate.log()
-            + death_rho.log()
-            + (death_rho - 1.0) * death_dt.clamp_min(self.eps).log()
-        )
-        observed_death = death_log_intensity * death_observed.to(dtype)
-
-        return (-observed_disease - observed_death + penalty).mean()
-
-
 def build_loss(name: str, **kwargs) -> nn.Module:
     name = name.lower()
     if name == "delphi2m":
         return Delphi2MLoss(**kwargs)
-    if name in {"exponential", "query_exponential"}:
+    if name == "exponential":
         return ExponentialLoss(**kwargs)
-    if name in {"weibull", "query_weibull"}:
+    if name == "weibull":
         return WeibullLoss(**kwargs)
-    if name in {"mixed", "query_mixed"}:
-        return MixedLoss(**kwargs)
     raise ValueError(
-        f"Unknown loss {name!r}. Available: delphi2m, exponential, weibull, mixed."
+        f"Unknown loss {name!r}. Available: delphi2m, exponential, weibull."
     )

models.py

@@ -37,7 +37,6 @@ class OtherInfoTokenizer(nn.Module):
         cont_type_ids: list[int],
         n_value_kinds: int = 3,
         n_bins: int = 16,
-        continuous_value_scaling: str = "none",
         continuous_value_center: torch.Tensor | list[float] | None = None,
         continuous_value_scale: torch.Tensor | list[float] | None = None,
     ):
@@ -57,13 +56,6 @@ class OtherInfoTokenizer(nn.Module):
             raise ValueError(
                 f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}"
             )
-        continuous_value_scaling = str(continuous_value_scaling).lower()
-        if continuous_value_scaling not in {"none", "robust"}:
-            raise ValueError(
-                "continuous_value_scaling must be either 'none' or 'robust', "
-                f"got {continuous_value_scaling!r}"
-            )
-
         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 = (
@@ -80,18 +72,20 @@ class OtherInfoTokenizer(nn.Module):
             n_embd,
             padding_idx=0,
         )
-        self.continuous_value_scaling = continuous_value_scaling
+        if n_cont_types > 0:
-        if continuous_value_scaling == "robust" and n_cont_types > 0:
+            if continuous_value_center is None or continuous_value_scale is None:
+                raise ValueError(
+                    "Continuous values require train-split RobustScale center "
+                    "and scale statistics"
+                )
             center = self._coerce_scaler_buffer(
                 continuous_value_center,
                 n_cont_types=n_cont_types,
-                default=0.0,
                 name="continuous_value_center",
             )
             scale = self._coerce_scaler_buffer(
                 continuous_value_scale,
                 n_cont_types=n_cont_types,
-                default=1.0,
                 name="continuous_value_scale",
             )
             if not torch.isfinite(center).all():
@@ -130,11 +124,10 @@ class OtherInfoTokenizer(nn.Module):
         value: torch.Tensor | list[float] | None,
         *,
         n_cont_types: int,
-        default: float,
         name: str,
     ) -> torch.Tensor:
         if value is None:
-            return torch.full((n_cont_types,), float(default), dtype=torch.float32)
+            raise ValueError(f"{name} is required")
         tensor = torch.as_tensor(value, dtype=torch.float32).detach().clone()
         if tensor.shape != (n_cont_types,):
             raise ValueError(
@@ -184,12 +177,11 @@ class OtherInfoTokenizer(nn.Module):
                     "cont_type_ids"
                 )
             cont_value = other_value[cont_pos].to(type_emb.dtype)
-            if self.continuous_value_scaling == "robust":
             if (
                 self.continuous_value_center is None
                 or self.continuous_value_scale is None
             ):
-                    raise RuntimeError("Robust continuous-value scaler buffers are missing")
+                raise RuntimeError("RobustScale buffers are missing")
             center = self.continuous_value_center[cont_idx].to(type_emb.dtype)
             scale = self.continuous_value_scale[cont_idx].to(type_emb.dtype)
             cont_value = (cont_value - center) / scale
@@ -221,12 +213,11 @@ class DeepHealth(nn.Module):
         cont_type_ids: list[int],
         n_value_kinds: int = 3,
         n_bins: int = 16,
-        continuous_value_scaling: str = "none",
         continuous_value_center: torch.Tensor | list[float] | None = None,
         continuous_value_scale: torch.Tensor | list[float] | None = None,
         target_mode: str = "next_token",  # "next_token" or "all_future"
         time_mode: str = "absolute",  # next_token requires absolute
-        dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
+        dist_mode: str = "exponential",  # "exponential" or "weibull"
         extra_pool_reduce: str = "mean",
         dropout: float = 0.0,
         model_architecture: str | None = None,
@@ -243,9 +234,9 @@ class DeepHealth(nn.Module):
                 "next_token is reserved for Delphi2M reproduction and "
                 "requires time_mode='absolute'"
             )
-        if dist_mode not in ["exponential", "weibull", "mixed"]:
+        if dist_mode not in ["exponential", "weibull"]:
             raise ValueError(
-                "dist_mode must be either 'exponential', 'weibull' or 'mixed'")
+                "dist_mode must be either 'exponential' or 'weibull'")
         if extra_pool_reduce not in {"mean", "sum"}:
             raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
         if n_layer < 1:
@@ -262,7 +253,6 @@ class DeepHealth(nn.Module):
             cont_type_ids=cont_type_ids,
             n_value_kinds=n_value_kinds,
             n_bins=n_bins,
-            continuous_value_scaling=continuous_value_scaling,
             continuous_value_center=continuous_value_center,
             continuous_value_scale=continuous_value_scale,
         )
@@ -270,7 +260,6 @@ class DeepHealth(nn.Module):
         self.time_mode = time_mode
         self.dist_mode = dist_mode
         self.extra_pool_reduce = extra_pool_reduce
-        self.continuous_value_scaling = str(continuous_value_scaling).lower()
         self.model_architecture = model_architecture
         self.n_layer = n_layer
         self.n_embd = n_embd
@@ -283,12 +272,6 @@ class DeepHealth(nn.Module):
             nn.init.zeros_(self.rho_head.weight)
             nn.init.constant_(self.rho_head.bias, 0.5413)
 
-        if dist_mode == "mixed":
-            self.death_idx = vocab_size - 1
-            self.rho_death_head = nn.Linear(n_embd, 1)
-            nn.init.zeros_(self.rho_death_head.weight)
-            nn.init.constant_(self.rho_death_head.bias, 0.5413)
-
         if time_mode == "absolute":
             self.age_encoding = AgeSinusoidalEncoding(n_embd)
             self.blocks = nn.ModuleList([
@@ -551,10 +534,3 @@ class DeepHealth(nn.Module):
                 f"calc_weibull_rho called with dist_mode={self.dist_mode!r}"
             )
         return F.softplus(self.rho_head(x)) + 1e-6
-
-    def calc_death_rho(self, x: torch.Tensor) -> torch.Tensor:
-        if self.dist_mode != "mixed":
-            raise RuntimeError(
-                f"calc_death_rho called with dist_mode={self.dist_mode!r}"
-            )
-        return F.softplus(self.rho_death_head(x)).squeeze(-1) + 1e-6

prepare_data.py

@@ -4,7 +4,7 @@ This script converts raw UK Biobank CSV exports into the artefacts consumed by
 DeepHealth:
 
 * ``ukb_event_data.npy``: ``(N, 3)`` uint32 array of ``(eid, days, label)``
-  disease/death/checkup events sorted by patient then time.
+  disease/death events sorted by patient then time.
 * ``ukb_basic_info.csv``: basic patient table indexed by ``eid`` with ``sex``.
 * ``ukb_other_info.npy``: ``(M, 5)`` float64 array of
   ``(eid, type, value, value_kind, time)`` rows. ``type=0`` is reserved for
@@ -219,7 +219,8 @@ with open(labels_file, encoding="utf-8") as f:  # Open labels file
     for idx, line in enumerate(f):  # Enumerate to assign incremental label IDs
         parts = line.strip().split(" ")  # Split by space
         if parts and parts[0]:  # Guard against empty lines
-            # Start labels from 1 to reserve 0 for padding, 1 for checkup
+            # Keep raw disease ids at 2+ so existing prepared data and model
+            # vocabulary indices remain stable; raw id 1 is unused.
             label_dict[parts[0]] = idx + 2
 
 # Pre-build lookup: ICD/Death column name -> integer label for fast per-column extraction
@@ -327,19 +328,6 @@ for ukb_chunk in tqdm.tqdm(ukb_iterator, desc="Processing UK Biobank data"):
     if cancer_frames:
         event_list.append(np.vstack(cancer_frames))
 
-    # Add checkup events with label=1 using date_of_assessment (already in days from dob)
-    if "date_of_assessment" in ukb_chunk.columns:
-        doa_series = ukb_chunk["date_of_assessment"].dropna()
-        if not doa_series.empty:
-            checkup_data = np.column_stack(
-                (
-                    doa_series.index.values,
-                    doa_series.values.astype(int),
-                    np.ones(len(doa_series), dtype=int),
-                )
-            )
-            event_list.append(checkup_data)
-
 # Combine tabular chunks
 final_tabular = pd.concat(tabular_list, axis=0, ignore_index=False)
 final_tabular.index.name = "eid"  # Ensure index named consistently

prepare_event_dates.py

@@ -9,12 +9,12 @@ with exactly three fields:
     token       int32
 
 ``token`` follows the existing ``labels.csv`` convention used by
-``prepare_data.py``: padding=0, checkup=1 (not emitted here), and the first
+``prepare_data.py``: padding=0, token 1 is reserved and unused, and the first
 label in ``labels.csv`` receives token 2. Each ``(eid, token)`` is deduplicated
 to the first known event date.
 
 The output is intended for calendar-indexed temperature and air-pollution
-queries.  It contains no date of birth, sex, covariates, or checkup events.
+queries. It contains no date of birth, sex, covariates, or assessment events.
 
 Usage
 -----

targets.py

@@ -6,7 +6,7 @@ import numpy as np
 
 
 PAD_IDX = 0
-CHECKUP_IDX = 1
+RESERVED_IDX = 1
 NO_EVENT_IDX = 2
 DAYS_PER_YEAR = 365.25
 
@@ -105,7 +105,7 @@ def build_next_token_targets(
         target_events:      [x1, x2, ..., xN-1]
         target_times_years: [t1, t2, ..., tN-1] / 365.25
 
-    This function does not ignore PAD/CHECKUP/NO_EVENT. Ignoring belongs to
+    This function does not ignore PAD/RESERVED/NO_EVENT. Ignoring belongs to
     the loss function because different objectives may use different ignore ids.
     """
     labels = _as_numpy_1d(labels, "labels", np.int64)

tests/test_calibration_metrics.py

@@ -207,19 +207,15 @@ class IPCWCalibrationMetricTests(unittest.TestCase):
         rho = np.asarray([0.8, 1.0, 1.2, 1.5], dtype=np.float32)
         horizons = np.asarray([0.1, 1.0, 5.0], dtype=np.float32)
 
-        for dist_mode, token, death_idx, selected_rho in (
+        for dist_mode, selected_rho in (
-            ("exponential", 4, 9, None),
+            ("exponential", None),
-            ("weibull", 4, 9, rho),
+            ("weibull", rho),
-            ("mixed", 9, 9, rho),
-            ("mixed", 4, 9, None),
         ):
             actual = _risk_probability_matrix(
                 logits=logits,
                 rho=selected_rho,
                 horizons=horizons,
                 dist_mode=dist_mode,
-                token=token,
-                death_idx=death_idx,
             )
             expected = np.vstack(
                 [
@@ -229,8 +225,6 @@ class IPCWCalibrationMetricTests(unittest.TestCase):
                         score_mode="risk",
                         horizon=float(horizon),
                         dist_mode=dist_mode,
-                        token=token,
-                        death_idx=death_idx,
                     )
                     for horizon in horizons
                 ]
@@ -271,7 +265,6 @@ class IPCWCalibrationMetricTests(unittest.TestCase):
             "label_id_to_code": {4: "D4", 5: "D5"},
             "dist_mode": "exponential",
             "horizons": np.asarray([1.0, 5.0], dtype=np.float32),
-            "death_index": 9,
             "min_cases": 1,
             "min_controls": 1,
             "max_ipcw_weight": 0.0,

tests/test_continuous_value_scaling.py

@@ -5,7 +5,8 @@ import torch
 import torch.nn as nn
 from torch.utils.data import Subset
 
-from models import OtherInfoTokenizer
+from eval_data import build_model_from_dataset
+from models import DeepHealth, OtherInfoTokenizer
 from train_util import fit_continuous_robust_scaler
 
 
@@ -59,6 +60,17 @@ class ContinuousValueScalingTests(unittest.TestCase):
         np.testing.assert_allclose(stats.center, np.asarray([2.0, 10.0]))
         np.testing.assert_allclose(stats.scale, np.asarray([2.0, 1.0]))
 
+    def test_fit_supports_next_step_sample_storage(self):
+        dataset = _ToyAllFutureDataset()
+        dataset.samples = dataset.patients
+        del dataset.patients
+        train_subset = Subset(dataset, np.asarray([0, 1, 2, 3, 4]))
+
+        stats = fit_continuous_robust_scaler(dataset, train_subset)
+
+        np.testing.assert_allclose(stats.center, np.asarray([2.0, 10.0]))
+        np.testing.assert_allclose(stats.scale, np.asarray([2.0, 1.0]))
+
     def test_tokenizer_standardizes_only_continuous_values(self):
         tokenizer = OtherInfoTokenizer(
             n_embd=4,
@@ -66,7 +78,6 @@ class ContinuousValueScalingTests(unittest.TestCase):
             n_cont_types=2,
             n_categories=3,
             cont_type_ids=[1, 3],
-            continuous_value_scaling="robust",
             continuous_value_center=[10.0, 100.0],
             continuous_value_scale=[2.0, 20.0],
         )
@@ -89,7 +100,6 @@ class ContinuousValueScalingTests(unittest.TestCase):
             n_cont_types=2,
             n_categories=2,
             cont_type_ids=[1, 3],
-            continuous_value_scaling="robust",
             continuous_value_center=[2.0, 10.0],
             continuous_value_scale=[1.5, 4.0],
         )
@@ -103,7 +113,8 @@ class ContinuousValueScalingTests(unittest.TestCase):
             n_cont_types=2,
             n_categories=2,
             cont_type_ids=[1, 3],
-            continuous_value_scaling="robust",
+            continuous_value_center=[0.0, 0.0],
+            continuous_value_scale=[1.0, 1.0],
         )
         restored.load_state_dict(state, strict=True)
 
@@ -116,27 +127,99 @@ class ContinuousValueScalingTests(unittest.TestCase):
             torch.tensor([1.5, 4.0]),
         )
 
-    def test_legacy_none_mode_keeps_old_state_dict_schema(self):
+    def test_continuous_tokenizer_rejects_missing_scaler_statistics(self):
-        tokenizer = OtherInfoTokenizer(
+        with self.assertRaisesRegex(ValueError, "require train-split RobustScale"):
+            OtherInfoTokenizer(
                 n_embd=4,
                 n_types=4,
                 n_cont_types=2,
                 n_categories=2,
                 cont_type_ids=[1, 3],
             )
-        state = tokenizer.state_dict()
 
-        self.assertNotIn("continuous_value_center", state)
+    def test_evaluation_rejects_unscaled_continuous_checkpoint(self):
-        self.assertNotIn("continuous_value_scale", state)
+        dataset = type(
-        restored = OtherInfoTokenizer(
+            "DatasetMetadata",
+            (),
+            {
+                "vocab_size": 8,
+                "n_types": 4,
+                "n_cont_types": 2,
+                "n_categories": 2,
+                "cont_type_ids": [1, 3],
+            },
+        )()
+        cfg = {
+            "model_target_mode": "all_future",
+            "target_mode": "all_future",
+            "model_architecture": "transformer_ffn_v1",
+            "n_layer": 1,
+            "time_mode": "absolute",
+            "dist_mode": "exponential",
+        }
+
+        with self.assertRaisesRegex(RuntimeError, "unscaled checkpoints are not supported"):
+            build_model_from_dataset(
+                None,
+                cfg,
+                dataset,
+                state_dict={"blocks.0.mlp.w1.weight": torch.zeros(1)},
+            )
+
+    def test_evaluation_restores_required_scaler_buffers(self):
+        dataset = type(
+            "DatasetMetadata",
+            (),
+            {
+                "vocab_size": 8,
+                "n_types": 4,
+                "n_cont_types": 2,
+                "n_categories": 2,
+                "cont_type_ids": [1, 3],
+            },
+        )()
+        source = DeepHealth(
+            vocab_size=8,
             n_embd=4,
+            n_head=1,
+            n_layer=1,
             n_types=4,
             n_cont_types=2,
             n_categories=2,
             cont_type_ids=[1, 3],
+            continuous_value_center=[2.0, 10.0],
+            continuous_value_scale=[1.5, 4.0],
+            target_mode="all_future",
+            time_mode="absolute",
+            dist_mode="exponential",
+            model_architecture="transformer_ffn_v1",
         )
+        state = source.state_dict()
+        cfg = {
+            "model_target_mode": "all_future",
+            "target_mode": "all_future",
+            "model_architecture": "transformer_ffn_v1",
+            "n_embd": 4,
+            "n_head": 1,
+            "n_layer": 1,
+            "n_bins": 16,
+            "time_mode": "absolute",
+            "dist_mode": "exponential",
+            "continuous_value_scaling": "robust",
+        }
+
+        restored = build_model_from_dataset(None, cfg, dataset, state_dict=state)
         restored.load_state_dict(state, strict=True)
 
+        torch.testing.assert_close(
+            restored.tokenizer.continuous_value_center,
+            torch.tensor([2.0, 10.0]),
+        )
+        torch.testing.assert_close(
+            restored.tokenizer.continuous_value_scale,
+            torch.tensor([1.5, 4.0]),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_dataset_checkup.py

@@ -1,57 +0,0 @@
-import unittest
-from pathlib import Path
-
-import numpy as np
-
-from dataset import _ExpoBaseDataset
-from targets import CHECKUP_IDX
-from train_util import load_extra_info_types_file
-
-
-class CheckupSelectionTests(unittest.TestCase):
-    @staticmethod
-    def _base(extra_info_types):
-        dataset = _ExpoBaseDataset.__new__(_ExpoBaseDataset)
-        dataset.extra_info_types = list(extra_info_types)
-        dataset.event_data = np.asarray(
-            [
-                [101, 10, CHECKUP_IDX],
-                [101, 20, 2],
-                [101, 30, 3],
-            ],
-            dtype=np.float64,
-        )
-        return dataset
-
-    def test_explicit_empty_extra_info_removes_checkup(self):
-        project_root = Path(__file__).resolve().parents[1]
-        selected_types = load_extra_info_types_file(
-            str(project_root / "extra_info_types_none.txt")
-        )
-        self.assertEqual(selected_types, [])
-        dataset = self._base(selected_types)
-
-        rows = list(dataset._iter_patient_events(impute_no_event_gaps=False))
-
-        self.assertEqual(len(rows), 1)
-        eid, times, labels = rows[0]
-        self.assertEqual(eid, 101)
-        np.testing.assert_array_equal(times, np.asarray([20, 30], dtype=np.float32))
-        self.assertNotIn(CHECKUP_IDX, labels.tolist())
-
-    def test_selected_extra_info_keeps_checkup(self):
-        dataset = self._base([11])
-
-        rows = list(dataset._iter_patient_events(impute_no_event_gaps=False))
-
-        self.assertEqual(len(rows), 1)
-        _, times, labels = rows[0]
-        np.testing.assert_array_equal(
-            times,
-            np.asarray([10, 20, 30], dtype=np.float32),
-        )
-        self.assertEqual(labels[0], CHECKUP_IDX)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/test_dataset_reserved_event.py

@@ -0,0 +1,44 @@
+import unittest
+
+import numpy as np
+
+from dataset import _ExpoBaseDataset
+from targets import RESERVED_IDX
+
+
+class ReservedEventFilteringTests(unittest.TestCase):
+    @staticmethod
+    def _base(extra_info_types):
+        dataset = _ExpoBaseDataset.__new__(_ExpoBaseDataset)
+        dataset.extra_info_types = list(extra_info_types)
+        dataset.event_data = np.asarray(
+            [
+                [101, 10, RESERVED_IDX],
+                [101, 20, 2],
+                [101, 30, 3],
+            ],
+            dtype=np.float64,
+        )
+        return dataset
+
+    def _assert_reserved_event_removed(self, extra_info_types):
+        dataset = self._base(extra_info_types)
+
+        rows = list(dataset._iter_patient_events(impute_no_event_gaps=False))
+
+        self.assertEqual(len(rows), 1)
+        eid, times, labels = rows[0]
+        self.assertEqual(eid, 101)
+        np.testing.assert_array_equal(times, np.asarray([20, 30], dtype=np.float32))
+        np.testing.assert_array_equal(labels, np.asarray([3, 4], dtype=np.int64))
+        self.assertNotIn(RESERVED_IDX, labels.tolist())
+
+    def test_empty_extra_info_removes_legacy_reserved_event(self):
+        self._assert_reserved_event_removed([])
+
+    def test_selected_extra_info_removes_legacy_reserved_event(self):
+        self._assert_reserved_event_removed([11])
+
+
+if __name__ == "__main__":
+    unittest.main()

train_all_future.py

@@ -37,7 +37,7 @@ from model_architectures import (
     SUPPORTED_MODEL_ARCHITECTURES,
 )
 from models import DeepHealth
-from targets import CHECKUP_IDX, PAD_IDX
+from targets import PAD_IDX, RESERVED_IDX
 from train_util import (
     ContinuousRobustScalerStats,
     configure_torch_for_training,
@@ -106,22 +106,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--n_head", type=int, default=10)
     parser.add_argument("--n_layer", type=int, default=12)
     parser.add_argument("--n_bins", type=int, default=16)
-    parser.add_argument(
-        "--continuous_value_scaling",
-        type=str,
-        default="robust",
-        choices=["none", "robust"],
-        help=(
-            "Continuous extra-info scaling. 'robust' fits the median and IQR "
-            "on the complete training subset and stores them in the checkpoint."
-        ),
-    )
     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",
-                        choices=["exponential", "weibull", "mixed"])
+                        choices=["exponential", "weibull"])
     parser.add_argument("--dropout", type=float, default=0.0)
     parser.add_argument(
         "--model_architecture",
@@ -188,19 +178,14 @@ def parse_args() -> argparse.Namespace:
 def build_model(
     args: argparse.Namespace,
     dataset: AllFutureHealthDataset,
-    scaler_stats: ContinuousRobustScalerStats | None = None,
+    scaler_stats: ContinuousRobustScalerStats,
 ) -> DeepHealth:
-    if (
+    if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids:
-        args.continuous_value_scaling == "robust"
-        and dataset.n_cont_types > 0
-        and scaler_stats is None
-    ):
         raise ValueError(
-            "Robust continuous-value scaling requires statistics fitted on the "
+            "RobustScale statistics are not aligned with dataset.cont_type_ids"
-            "training subset"
         )
-    center = None if scaler_stats is None else scaler_stats.center
+    center = scaler_stats.center if dataset.n_cont_types > 0 else None
-    scale = None if scaler_stats is None else scaler_stats.scale
+    scale = scaler_stats.scale if dataset.n_cont_types > 0 else None
     return DeepHealth(
         vocab_size=dataset.vocab_size,
         n_embd=args.n_embd,
@@ -211,7 +196,6 @@ def build_model(
         n_categories=dataset.n_categories,
         cont_type_ids=dataset.cont_type_ids,
         n_bins=args.n_bins,
-        continuous_value_scaling=args.continuous_value_scaling,
         continuous_value_center=center,
         continuous_value_scale=scale,
         extra_pool_reduce=args.extra_pool_reduce,
@@ -223,18 +207,12 @@ def build_model(
     )
 
 
-def build_criterion(args: argparse.Namespace, dataset: AllFutureHealthDataset):
+def build_criterion(args: argparse.Namespace):
-    ignored_idx = {PAD_IDX, CHECKUP_IDX}
+    ignored_idx = {PAD_IDX, RESERVED_IDX}
     if args.dist_mode == "exponential":
         return build_loss("exponential", ignored_idx=ignored_idx)
     if args.dist_mode == "weibull":
         return build_loss("weibull", ignored_idx=ignored_idx)
-    if args.dist_mode == "mixed":
-        return build_loss(
-            "mixed",
-            death_idx=dataset.vocab_size - 1,
-            ignored_idx=ignored_idx,
-        )
     raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
 
 
@@ -247,7 +225,7 @@ def compute_all_future_loss(
 ) -> torch.Tensor:
     required_keys = set(MODEL_INPUT_KEYS)
     required_keys.update(("future_targets", "exposure"))
-    if args.dist_mode in {"weibull", "mixed"}:
+    if args.dist_mode == "weibull":
         required_keys.add("future_dt")
     batch = move_batch_to_device(
         {key: batch[key] for key in required_keys},
@@ -282,13 +260,7 @@ def compute_all_future_loss(
             exposure=batch["exposure"],
         )
     else:
-        loss = criterion(
+        raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
-            logits=logits,
-            death_rho=model.calc_death_rho(hidden),
-            targets=batch["future_targets"],
-            dt=batch["future_dt"],
-            exposure=batch["exposure"],
-        )
 
     if not torch.isfinite(loss):
         raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
@@ -352,15 +324,8 @@ def build_metadata(
     train_subset,
     val_subset,
     test_subset,
-    scaler_stats: ContinuousRobustScalerStats | None,
+    scaler_stats: ContinuousRobustScalerStats,
 ) -> Dict[str, Any]:
-    scaler_metadata: Dict[str, Any]
-    if scaler_stats is None:
-        scaler_metadata = {
-            "method": "none",
-            "fitted_on": None,
-        }
-    else:
     scaler_metadata = scaler_stats.as_metadata()
     return {
         "run_name": run_name,
@@ -370,6 +335,8 @@ def build_metadata(
         "model_architecture": args.model_architecture,
         "model_target_mode": "all_future",
         "target_mode": "all_future",
+        "event_stream_version": "disease_death_only_v1",
+        "uses_assessment_event_token": False,
         "dist_mode": args.dist_mode,
         "disease_history_mode": args.disease_history_mode,
         "all_future_min_history_events": int(args.min_history_events),
@@ -381,6 +348,7 @@ def build_metadata(
             else None
         ),
         "extra_info_types": [int(x) for x in dataset.extra_info_types],
+        "continuous_value_scaling": "robust",
         "continuous_value_scaler": scaler_metadata,
         "dataset_metadata": {
             "vocab_size": int(dataset.vocab_size),
@@ -389,6 +357,8 @@ def build_metadata(
             "n_categories": int(dataset.n_categories),
             "cont_type_ids": [int(x) for x in dataset.cont_type_ids],
             "extra_info_types": [int(x) for x in dataset.extra_info_types],
+            "event_stream_version": "disease_death_only_v1",
+            "uses_assessment_event_token": False,
         },
         "split_sizes": {
             "train": int(len(train_subset)),
@@ -425,7 +395,7 @@ def main() -> None:
     logger.info(f"Model architecture: {args.model_architecture}")
     logger.info(f"Disease history mode: {args.disease_history_mode}")
     logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
-    logger.info(f"Continuous value scaling: {args.continuous_value_scaling}")
+    logger.info("Continuous value scaling: RobustScale (required)")
 
     logger.info("Loading all-future datasets...")
     train_dataset = AllFutureHealthDataset(
@@ -489,8 +459,7 @@ def main() -> None:
         f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
     )
 
-    scaler_stats = None
+    if train_dataset.n_cont_types > 0:
-    if args.continuous_value_scaling == "robust" and train_dataset.n_cont_types > 0:
         logger.info(
             "Fitting continuous RobustScaler on the complete training subset: "
             f"patients={len(train_subset):,}, features={train_dataset.n_cont_types}"
@@ -499,6 +468,7 @@ def main() -> None:
         train_dataset,
         train_subset,
     )
+    if train_dataset.n_cont_types > 0:
         logger.info(
             "Continuous RobustScaler fitted: "
             f"observations={int(scaler_stats.observation_count.sum()):,}, "
@@ -550,7 +520,7 @@ def main() -> None:
         betas=tuple(args.betas),
         weight_decay=args.weight_decay,
     )
-    criterion = build_criterion(args, train_dataset)
+    criterion = build_criterion(args)
     adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
 
     train_metadata = build_metadata(

train_batch_linux.sh

@@ -4,7 +4,7 @@
 #
 # The matrix contains:
 #   1. FFN + Delphi2M next-token reproduction
-#   2. FFN/TrajMixer x absolute/relative x exponential/Weibull/mixed
+#   2. FFN/TrajMixer x absolute/relative x exponential/Weibull
 #
 # Every task uses one GPU. Each selected GPU runs its assigned tasks
 # sequentially, while different GPUs run in parallel.
@@ -47,6 +47,7 @@ Options:
 Fixed experiment settings:
   batch_size             256
   extra_info_types       extra_info_types_smoking_alcohol_bmi.txt
+  continuous scaling     required train-split RobustScale
   model size             Defaults from the individual training entrypoints
 
 Outputs:
@@ -180,7 +181,7 @@ add_job \
 
 for architecture in transformer_ffn_v1 traj_mixer_v5; do
     for time_mode in absolute relative; do
-        for dist_mode in exponential weibull mixed; do
+        for dist_mode in exponential weibull; do
             add_job \
                 "${architecture}_all_future_${time_mode}_${dist_mode}" \
                 "train_all_future.py" \

train_extra_info_assessment_all_multiseed_linux.sh

@@ -236,7 +236,6 @@ run_job() {
         --time_mode relative
         --dist_mode weibull
         --disease_history_mode timed
-        --continuous_value_scaling robust
         --extra_info_types_file "$EXTRA_INFO_TYPES_FILE"
     )
 

train_key_models_multiseed_linux.sh

@@ -59,6 +59,7 @@ Options:
 Fixed experiment settings:
   batch_size             256
   extra_info_types       extra_info_types_smoking_alcohol_bmi.txt
+  continuous scaling     required train-split RobustScale
   model size             Defaults from the training entrypoints
   tasks per seed         6
 

train_next_step.py

@@ -23,10 +23,12 @@ from model_architectures import (
     SUPPORTED_MODEL_ARCHITECTURES,
 )
 from models import DeepHealth, DeepHealthOutput
-from targets import CHECKUP_IDX, PAD_IDX
+from targets import PAD_IDX, RESERVED_IDX
 from train_util import (
+    ContinuousRobustScalerStats,
     configure_torch_for_training,
     create_unique_run_dir,
+    fit_continuous_robust_scaler,
     format_extra_info_types,
     get_lr,
     get_model_parameter_counts,
@@ -120,7 +122,17 @@ def parse_args() -> argparse.Namespace:
     return args
 
 
-def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
+def build_model(
+    args: argparse.Namespace,
+    dataset: HealthDataset,
+    scaler_stats: ContinuousRobustScalerStats,
+) -> DeepHealth:
+    if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids:
+        raise ValueError(
+            "RobustScale statistics are not aligned with dataset.cont_type_ids"
+        )
+    center = scaler_stats.center if dataset.n_cont_types > 0 else None
+    scale = scaler_stats.scale if dataset.n_cont_types > 0 else None
     return DeepHealth(
         vocab_size=dataset.vocab_size,
         n_embd=args.n_embd,
@@ -131,6 +143,8 @@ 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,
+        continuous_value_center=center,
+        continuous_value_scale=scale,
         extra_pool_reduce=args.extra_pool_reduce,
         target_mode="next_token",
         time_mode="absolute",
@@ -143,7 +157,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
 def build_next_step_loss(args: argparse.Namespace):
     return build_loss(
         "delphi2m",
-        ignored_tokens={PAD_IDX, CHECKUP_IDX},
+        ignored_tokens={PAD_IDX, RESERVED_IDX},
         t_min=args.t_min,
         max_exp_input=args.max_exp_input,
         ce_weight=args.ce_weight,
@@ -244,7 +258,6 @@ def build_augmented_next_step_targets(
 
 
 def compute_next_step_loss(
-    args: argparse.Namespace,
     model: DeepHealth,
     criterion,
     batch: Dict[str, torch.Tensor],
@@ -309,7 +322,7 @@ def run_epoch(
     for batch_idx, batch in enumerate(progress):
         try:
             loss, parts = compute_next_step_loss(
-                args, model, criterion, batch, device
+                model, criterion, batch, device
             )
             if is_train:
                 if optimizer is None:
@@ -352,6 +365,7 @@ def build_metadata(
     train_subset,
     val_subset,
     test_subset,
+    scaler_stats: ContinuousRobustScalerStats,
 ) -> Dict[str, Any]:
     return {
         "run_name": run_name,
@@ -361,6 +375,8 @@ def build_metadata(
         "model_architecture": args.model_architecture,
         "model_target_mode": "next_token",
         "target_mode": "delphi2m",
+        "event_stream_version": "disease_death_only_v1",
+        "uses_assessment_event_token": False,
         "time_mode": "absolute",
         "dist_mode": "exponential",
         "extra_info_types_file": (
@@ -369,6 +385,8 @@ def build_metadata(
             else None
         ),
         "extra_info_types": [int(x) for x in dataset.extra_info_types],
+        "continuous_value_scaling": "robust",
+        "continuous_value_scaler": scaler_stats.as_metadata(),
         "dataset_metadata": {
             "vocab_size": int(dataset.vocab_size),
             "n_types": int(dataset.n_types),
@@ -376,6 +394,8 @@ def build_metadata(
             "n_categories": int(dataset.n_categories),
             "cont_type_ids": [int(x) for x in dataset.cont_type_ids],
             "extra_info_types": [int(x) for x in dataset.extra_info_types],
+            "event_stream_version": "disease_death_only_v1",
+            "uses_assessment_event_token": False,
         },
         "split_sizes": {
             "train": int(len(train_subset)),
@@ -406,6 +426,7 @@ def main() -> None:
     logger.info(f"Device: {device}")
     logger.info(f"Model architecture: {args.model_architecture}")
     logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
+    logger.info("Continuous value scaling: RobustScale (required)")
     logger.info("time_mode=absolute, readout=token, target_mode=delphi2m")
 
     dataset = HealthDataset(
@@ -441,6 +462,20 @@ def main() -> None:
         f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
     )
 
+    if dataset.n_cont_types > 0:
+        logger.info(
+            "Fitting continuous RobustScaler on the complete training subset: "
+            f"patients={len(train_subset):,}, features={dataset.n_cont_types}"
+        )
+    scaler_stats = fit_continuous_robust_scaler(dataset, train_subset)
+    if dataset.n_cont_types > 0:
+        logger.info(
+            "Continuous RobustScaler fitted: "
+            f"observations={int(scaler_stats.observation_count.sum()):,}, "
+            f"min_per_feature={int(scaler_stats.observation_count.min()):,}, "
+            f"max_per_feature={int(scaler_stats.observation_count.max()):,}"
+        )
+
     train_loader = DataLoader(
         train_subset,
         batch_size=args.batch_size,
@@ -472,7 +507,7 @@ def main() -> None:
         prefetch_factor=2 if args.num_workers > 0 else None,
     )
 
-    model = build_model(args, dataset).to(device)
+    model = build_model(args, dataset, scaler_stats).to(device)
     parameter_counts = get_model_parameter_counts(model)
     logger.info(
         "Model parameters: "
@@ -489,7 +524,13 @@ def main() -> None:
     adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
 
     train_metadata = build_metadata(
-        args, dataset, run_name, train_subset, val_subset, test_subset
+        args,
+        dataset,
+        run_name,
+        train_subset,
+        val_subset,
+        test_subset,
+        scaler_stats,
     )
     train_metadata.update(parameter_counts)
     save_config(

train_util.py

@@ -43,7 +43,7 @@ class ContinuousRobustScalerStats:
 
 
 def fit_continuous_robust_scaler(
-    dataset: AllFutureHealthDataset,
+    dataset: HealthDataset | AllFutureHealthDataset,
     subset: Subset,
     *,
     quantile_range: tuple[float, float] = (25.0, 75.0),
@@ -90,13 +90,22 @@ def fit_continuous_robust_scaler(
             )
         type_to_column[type_id] = column
 
+    if hasattr(dataset, "patients"):
+        records = dataset.patients
+    elif hasattr(dataset, "samples"):
+        records = dataset.samples
+    else:
+        raise TypeError(
+            "dataset must expose patient records through .patients or .samples"
+        )
+
     values = np.full(
         (int(subset_indices.size), n_cont_types),
         np.nan,
         dtype=np.float32,
     )
     for row, patient_index in enumerate(subset_indices.tolist()):
-        patient = dataset.patients[int(patient_index)]
+        patient = records[int(patient_index)]
         other_type = np.asarray(patient["other_type"], dtype=np.int64)
         other_value = np.asarray(patient["other_value"], dtype=np.float32)
         other_kind = np.asarray(patient["other_value_kind"], dtype=np.int64)