Add train-split robust scaling for continuous values

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