Add train-split robust scaling for continuous values

eval_data.py

@@ -161,6 +161,9 @@ 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(
+            cfg_get(args, cfg, "continuous_value_scaling", "none")
+        ),
         extra_pool_reduce=str(
             cfg_get(args, cfg, "extra_pool_reduce", "mean")
         ),

models.py

@@ -37,6 +37,9 @@ 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,
     ):
         super().__init__()
         if len(cont_type_ids) != n_cont_types:
@@ -54,6 +57,12 @@ 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)
@@ -71,6 +80,36 @@ class OtherInfoTokenizer(nn.Module):
             n_embd,
             padding_idx=0,
         )
+        self.continuous_value_scaling = continuous_value_scaling
+        if continuous_value_scaling == "robust" and n_cont_types > 0:
+            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():
+                raise ValueError(
+                    "continuous_value_center must contain only finite values"
+                )
+            if not torch.isfinite(scale).all() or torch.any(scale <= 0):
+                raise ValueError(
+                    "continuous_value_scale must be finite and strictly positive"
+                )
+            self.register_buffer("continuous_value_center", center)
+            self.register_buffer("continuous_value_scale", scale)
+        else:
+            # ``None`` buffers are omitted from state_dict. This preserves the
+            # exact checkpoint schema used by models trained before continuous
+            # value scaling was introduced.
+            self.register_buffer("continuous_value_center", None)
+            self.register_buffer("continuous_value_scale", None)
 
         cont_type_index = torch.full((n_types,), -1, dtype=torch.long)
         for idx, type_id in enumerate(cont_type_ids):
@@ -86,6 +125,23 @@ class OtherInfoTokenizer(nn.Module):
         )
         self.reset_parameters()
 
+    @staticmethod
+    def _coerce_scaler_buffer(
+        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)
+        tensor = torch.as_tensor(value, dtype=torch.float32).detach().clone()
+        if tensor.shape != (n_cont_types,):
+            raise ValueError(
+                f"{name} must have shape ({n_cont_types},), got {tuple(tensor.shape)}"
+            )
+        return tensor
+
     def reset_parameters(self) -> None:
         nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02)
         nn.init.zeros_(self.type_emb.weight[0])
@@ -127,9 +183,19 @@ class OtherInfoTokenizer(nn.Module):
                     f"type_id={bad_type} is marked continuous but is not in "
                     "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")
+                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
             value_emb[cont_pos] = self.cont_value_encoder(
                 cont_type_idx=cont_idx,
-                value=other_value[cont_pos].to(type_emb.dtype),
+                value=cont_value,
             )
 
         cate_pos = other_valid & (other_value_kind == self.CATE_KIND)
@@ -155,6 +221,9 @@ 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"
@@ -193,11 +262,15 @@ 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,
         )
         self.target_mode = target_mode
         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

tests/test_continuous_value_scaling.py

@@ -0,0 +1,142 @@
+import unittest
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import Subset
+
+from models import OtherInfoTokenizer
+from train_util import fit_continuous_robust_scaler
+
+
+class _ToyAllFutureDataset:
+    def __init__(self):
+        self.cont_type_ids = [1, 3]
+        self.n_types = 4
+        self.patients = [
+            self._patient([1, 3], [0.0, 10.0]),
+            self._patient([1, 3], [1.0, 10.0]),
+            self._patient([1, 3], [2.0, 10.0]),
+            self._patient([1, 3], [3.0, 10.0]),
+            self._patient([1, 3], [4.0, 10.0]),
+            self._patient([1, 3], [1000.0, 999.0]),
+        ]
+
+    @staticmethod
+    def _patient(types, values):
+        return {
+            "other_type": np.asarray(types, dtype=np.int64),
+            "other_value": np.asarray(values, dtype=np.float32),
+            "other_value_kind": np.ones(len(types), dtype=np.int64),
+        }
+
+    def __len__(self):
+        return len(self.patients)
+
+
+class _CaptureContinuousEncoder(nn.Module):
+    def __init__(self, n_embd):
+        super().__init__()
+        self.n_embd = n_embd
+        self.last_type = None
+        self.last_value = None
+
+    def forward(self, cont_type_idx, value):
+        self.last_type = cont_type_idx.detach().clone()
+        self.last_value = value.detach().clone()
+        return value[:, None].expand(-1, self.n_embd)
+
+
+class ContinuousValueScalingTests(unittest.TestCase):
+    def test_fit_uses_only_training_subset_and_handles_constant_features(self):
+        dataset = _ToyAllFutureDataset()
+        train_subset = Subset(dataset, np.asarray([0, 1, 2, 3, 4]))
+
+        stats = fit_continuous_robust_scaler(dataset, train_subset)
+
+        self.assertEqual(stats.cont_type_ids, (1, 3))
+        np.testing.assert_array_equal(stats.observation_count, np.asarray([5, 5]))
+        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,
+            n_types=4,
+            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],
+        )
+        capture = _CaptureContinuousEncoder(n_embd=4)
+        tokenizer.cont_value_encoder = capture
+
+        tokenizer(
+            other_type=torch.tensor([[1, 2, 3]], dtype=torch.long),
+            other_value=torch.tensor([[14.0, 1.0, 80.0]]),
+            other_value_kind=torch.tensor([[1, 2, 1]], dtype=torch.long),
+        )
+
+        torch.testing.assert_close(capture.last_type, torch.tensor([0, 1]))
+        torch.testing.assert_close(capture.last_value, torch.tensor([2.0, -1.0]))
+
+    def test_scaler_buffers_round_trip_in_new_checkpoint(self):
+        tokenizer = OtherInfoTokenizer(
+            n_embd=4,
+            n_types=4,
+            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],
+        )
+        state = tokenizer.state_dict()
+        self.assertIn("continuous_value_center", state)
+        self.assertIn("continuous_value_scale", state)
+
+        restored = OtherInfoTokenizer(
+            n_embd=4,
+            n_types=4,
+            n_cont_types=2,
+            n_categories=2,
+            cont_type_ids=[1, 3],
+            continuous_value_scaling="robust",
+        )
+        restored.load_state_dict(state, strict=True)
+
+        torch.testing.assert_close(
+            restored.continuous_value_center,
+            torch.tensor([2.0, 10.0]),
+        )
+        torch.testing.assert_close(
+            restored.continuous_value_scale,
+            torch.tensor([1.5, 4.0]),
+        )
+
+    def test_legacy_none_mode_keeps_old_state_dict_schema(self):
+        tokenizer = 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)
+        self.assertNotIn("continuous_value_scale", state)
+        restored = OtherInfoTokenizer(
+            n_embd=4,
+            n_types=4,
+            n_cont_types=2,
+            n_categories=2,
+            cont_type_ids=[1, 3],
+        )
+        restored.load_state_dict(state, strict=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

train_all_future.py

@@ -39,9 +39,11 @@ from model_architectures import (
 from models import DeepHealth
 from targets import CHECKUP_IDX, PAD_IDX
 from train_util import (
+    ContinuousRobustScalerStats,
     configure_torch_for_training,
     create_unique_run_dir,
     format_extra_info_types,
+    fit_continuous_robust_scaler,
     get_lr,
     get_model_parameter_counts,
     load_extra_info_types_file,
@@ -104,6 +106,16 @@ 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",
@@ -173,7 +185,22 @@ def parse_args() -> argparse.Namespace:
     return args
 
 
-def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth:
+def build_model(
+    args: argparse.Namespace,
+    dataset: AllFutureHealthDataset,
+    scaler_stats: ContinuousRobustScalerStats | None = None,
+) -> DeepHealth:
+    if (
+        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 "
+            "training subset"
+        )
+    center = None if scaler_stats is None else scaler_stats.center
+    scale = None if scaler_stats is None else scaler_stats.scale
     return DeepHealth(
         vocab_size=dataset.vocab_size,
         n_embd=args.n_embd,
@@ -184,6 +211,9 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De
         n_categories=dataset.n_categories,
         cont_type_ids=dataset.cont_type_ids,
         n_bins=args.n_bins,
+        continuous_value_scaling=args.continuous_value_scaling,
+        continuous_value_center=center,
+        continuous_value_scale=scale,
         extra_pool_reduce=args.extra_pool_reduce,
         target_mode="all_future",
         time_mode=args.time_mode,
@@ -322,7 +352,16 @@ def build_metadata(
     train_subset,
     val_subset,
     test_subset,
+    scaler_stats: ContinuousRobustScalerStats | None,
 ) -> 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,
         "dataset_class": "AllFutureHealthDataset",
@@ -342,6 +381,7 @@ def build_metadata(
             else None
         ),
         "extra_info_types": [int(x) for x in dataset.extra_info_types],
+        "continuous_value_scaler": scaler_metadata,
         "dataset_metadata": {
             "vocab_size": int(dataset.vocab_size),
             "n_types": int(dataset.n_types),
@@ -385,6 +425,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("Loading all-future datasets...")
     train_dataset = AllFutureHealthDataset(
@@ -448,6 +489,23 @@ def main() -> None:
         f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
     )
 
+    scaler_stats = None
+    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}"
+        )
+        scaler_stats = fit_continuous_robust_scaler(
+            train_dataset,
+            train_subset,
+        )
+        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,
@@ -479,7 +537,7 @@ def main() -> None:
         prefetch_factor=2 if args.num_workers > 0 else None,
     )
 
-    model = build_model(args, train_dataset).to(device)
+    model = build_model(args, train_dataset, scaler_stats=scaler_stats).to(device)
     parameter_counts = get_model_parameter_counts(model)
     logger.info(
         "Model parameters: "
@@ -496,7 +554,13 @@ def main() -> None:
     adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
 
     train_metadata = build_metadata(
-        args, train_dataset, run_name, train_subset, val_subset, test_subset
+        args,
+        train_dataset,
+        run_name,
+        train_subset,
+        val_subset,
+        test_subset,
+        scaler_stats,
     )
     train_metadata.update(parameter_counts)
     save_config(

train_extra_info_assessment_all_multiseed_linux.sh

@@ -10,6 +10,7 @@
 #   - smoking
 #   - alcohol
 #   - BMI is already included in the assessment variables
+#   - continuous values use train-split RobustScaler statistics
 #
 # A6000 48 GB default:
 #   batch_size=256
@@ -34,7 +35,7 @@ SEED_CSV="42,43,44"
 NUM_WORKERS=4
 BATCH_SIZE=256
 PYTHON_BIN="${PYTHON_BIN:-python}"
-CAMPAIGN_NAME="extra_info_assessment_smoking_alcohol_multiseed"
+CAMPAIGN_NAME="extra_info_assessment_smoking_alcohol_robust_multiseed"
 DRY_RUN=0
 
 ENTRYPOINT="$SCRIPT_DIR/train_all_future.py"
@@ -66,6 +67,7 @@ Fixed experiment settings:
   disease history        timed
   sex                    enabled by the model
   extra information      extra_info_types_assessment_smoking_alcohol.txt
+  continuous scaling     robust (training-subset median/IQR)
 
 A6000 48 GB default:
   batch_size             256
@@ -234,6 +236,7 @@ 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_util.py

@@ -5,6 +5,7 @@ import logging
 import sys
 import time
 import csv
+from dataclasses import dataclass
 from datetime import datetime
 import math
 from pathlib import Path
@@ -19,6 +20,134 @@ from dataset import AllFutureHealthDataset, HealthDataset
 from models import DeepHealth
 
 
+@dataclass(frozen=True)
+class ContinuousRobustScalerStats:
+    """Train-split robust scaling statistics aligned to ``cont_type_ids``."""
+
+    cont_type_ids: tuple[int, ...]
+    center: np.ndarray
+    scale: np.ndarray
+    observation_count: np.ndarray
+    quantile_range: tuple[float, float] = (25.0, 75.0)
+
+    def as_metadata(self) -> Dict[str, Any]:
+        return {
+            "method": "robust",
+            "fitted_on": "train_subset",
+            "quantile_range": [float(x) for x in self.quantile_range],
+            "cont_type_ids": [int(x) for x in self.cont_type_ids],
+            "observation_count": [int(x) for x in self.observation_count.tolist()],
+            "center_buffer": "tokenizer.continuous_value_center",
+            "scale_buffer": "tokenizer.continuous_value_scale",
+        }
+
+
+def fit_continuous_robust_scaler(
+    dataset: AllFutureHealthDataset,
+    subset: Subset,
+    *,
+    quantile_range: tuple[float, float] = (25.0, 75.0),
+    scale_epsilon: float = 1e-6,
+) -> ContinuousRobustScalerStats:
+    """Fit median/IQR statistics using only patients in the training subset.
+
+    The prepared arrays remain unchanged. Scaling is performed later inside the
+    model tokenizer so the fitted center and scale can live in the checkpoint.
+    """
+
+    if subset.dataset is not dataset:
+        raise ValueError("subset must reference the dataset used to fit the scaler")
+    low, high = (float(quantile_range[0]), float(quantile_range[1]))
+    if not (0.0 <= low < high <= 100.0):
+        raise ValueError(
+            "quantile_range must satisfy 0 <= low < high <= 100, got "
+            f"{quantile_range!r}"
+        )
+    if scale_epsilon <= 0:
+        raise ValueError("scale_epsilon must be > 0")
+
+    cont_type_ids = tuple(int(x) for x in dataset.cont_type_ids)
+    n_cont_types = len(cont_type_ids)
+    if n_cont_types == 0:
+        empty = np.zeros(0, dtype=np.float32)
+        return ContinuousRobustScalerStats(
+            cont_type_ids=cont_type_ids,
+            center=empty.copy(),
+            scale=empty.copy(),
+            observation_count=np.zeros(0, dtype=np.int64),
+            quantile_range=(low, high),
+        )
+
+    subset_indices = np.asarray(subset.indices, dtype=np.int64)
+    if subset_indices.ndim != 1 or subset_indices.size == 0:
+        raise ValueError("training subset must contain at least one patient")
+
+    type_to_column = np.full(int(dataset.n_types), -1, dtype=np.int64)
+    for column, type_id in enumerate(cont_type_ids):
+        if type_id <= 0 or type_id >= len(type_to_column):
+            raise ValueError(
+                f"continuous type id {type_id} is outside [1, {len(type_to_column)})"
+            )
+        type_to_column[type_id] = column
+
+    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)]
+        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)
+        continuous = other_kind == 1
+        if not np.any(continuous):
+            continue
+        selected_type = other_type[continuous]
+        selected_value = other_value[continuous]
+        valid_type = (selected_type > 0) & (selected_type < len(type_to_column))
+        columns = np.full(selected_type.shape, -1, dtype=np.int64)
+        columns[valid_type] = type_to_column[selected_type[valid_type]]
+        valid = (columns >= 0) & np.isfinite(selected_value)
+        values[row, columns[valid]] = selected_value[valid]
+
+    observation_count = np.isfinite(values).sum(axis=0).astype(np.int64)
+    missing_types = [
+        type_id
+        for type_id, count in zip(cont_type_ids, observation_count.tolist())
+        if count == 0
+    ]
+    if missing_types:
+        raise ValueError(
+            "Training subset has no finite observations for continuous type ids: "
+            f"{missing_types}"
+        )
+
+    low_value, center, high_value = np.nanpercentile(
+        values,
+        [low, 50.0, high],
+        axis=0,
+    )
+    scale = high_value - low_value
+    near_constant = (~np.isfinite(scale)) | (np.abs(scale) <= float(scale_epsilon))
+    scale[near_constant] = 1.0
+
+    center = np.asarray(center, dtype=np.float32)
+    scale = np.asarray(scale, dtype=np.float32)
+    if not np.isfinite(center).all():
+        raise RuntimeError("Robust scaler produced non-finite center values")
+    if not np.isfinite(scale).all() or np.any(scale <= 0):
+        raise RuntimeError("Robust scaler produced invalid scale values")
+
+    return ContinuousRobustScalerStats(
+        cont_type_ids=cont_type_ids,
+        center=center,
+        scale=scale,
+        observation_count=observation_count,
+        quantile_range=(low, high),
+    )
+
+
 def create_unique_run_dir(name_fn, runs_root: Path = Path("runs")) -> tuple[Path, str]:
     while True:
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")